如何从ECS传递参数至Airflow XCom？

不需要依赖CloudWatch、S3等中间服务，核心思路是让ECS任务主动调用Airflow的REST API，将结果直接写入XCom。具体实现步骤如下：

1. 配置Airflow API访问权限 #

• 在Airflow中创建具备Edit权限的用户，并生成对应的API Token（可在Airflow UI的Profile > API Keys中创建）。
• 确保ECS任务所在的VPC能够访问Airflow Webserver（如果Airflow部署在私有网络，需配置VPC对等或安全组规则允许ECS的IP访问）。

2. 在ECS任务中添加推送XCom的逻辑 #

在ECS任务的业务代码末尾，添加调用Airflow API的逻辑，将需要传递的参数推送到XCom。以下是Python示例：

import requests
import os

# 从环境变量读取Airflow相关配置（这些变量将由Airflow在启动ECS任务时传入）
airflow_api_url = os.environ["AIRFLOW_API_URL"]
airflow_token = os.environ["AIRFLOW_TOKEN"]
dag_id = os.environ["AIRFLOW_DAG_ID"]
dag_run_id = os.environ["AIRFLOW_DAG_RUN_ID"]
task_id = os.environ["AIRFLOW_TASK_ID"]

# 模拟ECS任务生成的结果参数
ecs_task_result = {
    "processing_status": "success",
    "processed_count": 100,
    "output_file_prefix": "20240520_batch"
}

# 调用Airflow API推送XCom
try:
    response = requests.post(
        f"{airflow_api_url}/api/v1/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/xcomEntries",
        headers={
            "Authorization": f"Bearer {airflow_token}",
            "Content-Type": "application/json"
        },
        json={
            "key": "ecs_task_output",
            "value": ecs_task_result
        }
    )
    response.raise_for_status()
    print("XCom推送成功")
except Exception as e:
    print(f"XCom推送失败: {str(e)}")
    raise  # 可选：如果推送失败，让ECS任务标记为失败

3. 在Airflow DAG中配置EcsRunTaskOperator #

启动ECS任务时，通过overrides参数将Airflow的相关信息作为环境变量传递给ECS容器：

from airflow import DAG
from airflow.providers.amazon.aws.operators.ecs import EcsRunTaskOperator
from airflow.operators.python import PythonOperator
from datetime import datetime

with DAG(
    dag_id="ecs_xcom_demo",
    schedule_interval=None,
    start_date=datetime(2024, 5, 20),
    catchup=False
) as dag:
    run_ecs_task = EcsRunTaskOperator(
        task_id="trigger_ecs_task",
        cluster="your-ecs-cluster-name",
        task_definition="your-task-definition-name",
        overrides={
            "containerOverrides": [
                {
                    "name": "your-container-name",
                    "environment": [
                        {"name": "AIRFLOW_API_URL", "value": "http://airflow-webserver:8080"},  # 根据实际部署调整
                        {"name": "AIRFLOW_TOKEN", "value": "{{ conn.airflow_api.password }}"},  # 建议将Token存在Airflow连接中
                        {"name": "AIRFLOW_DAG_ID", "value": "{{ dag.dag_id }}"},
                        {"name": "AIRFLOW_DAG_RUN_ID", "value": "{{ run_id }}"},
                        {"name": "AIRFLOW_TASK_ID", "value": "{{ task.task_id }}"}
                    ]
                }
            ]
        },
        aws_conn_id="aws_default",
        wait_for_completion=True,  # 等待ECS任务完成后再执行后续任务
        dag=dag
    )

    # 后续任务通过XCom获取ECS返回的参数
    def process_ecs_result(**context):
        ecs_output = context["ti"].xcom_pull(task_ids="trigger_ecs_task", key="ecs_task_output")
        print(f"Received ECS result: {ecs_output}")

    process_result = PythonOperator(
        task_id="process_ecs_result",
        python_callable=process_ecs_result,
        provide_context=True,
        dag=dag
    )

    run_ecs_task >> process_result

关键注意事项 #

• XCom大小限制：Airflow默认限制XCom的value大小为48KB左右，若ECS返回的参数过大，此方案不适用（此时仍需借助中间存储）。
• API安全：避免硬编码Airflow Token，建议将Token存储在Airflow的Connections中，通过{{ conn.airflow_api.password }}引用。
• ECS权限：确保ECS任务的IAM角色允许出站访问Airflow Webserver的端口（默认8080）。

内容的提问来源于stack exchange，提问作者leonidas efrem