Airflow中如何基于上游任务输出执行并行任务

该需求完全可以在Airflow 2.0+版本中实现，XCom也是实现该需求的核心依赖组件，你提到的典型分治场景刚好匹配Airflow 2.3版本后推出的原生**动态任务映射（Dynamic Task Mapping）**能力，不需要额外的第三方工具支持，具体实现逻辑如下：

核心实现步骤 #

• 第一步：定义拉取条目任务：调用检索API获取符合条件的卫星影像条目列表，将列表结果通过XCom推送存储。如果返回的条目数量极大，也可以将列表存入对象存储/共享文件系统，XCom仅存储文件读取路径，避免给元数据库带来过大压力。
• 第二步：并行处理任务绑定：定义单个影像的下载、处理任务，通过.expand()方法绑定上游输出的条目列表，Airflow会自动根据列表长度生成对应数量的并行任务实例，每个实例自动分配单个条目作为入参。如果搭配KubernetesExecutor使用，还可以为每个处理任务单独配置CPU、GPU资源，自动调度到符合算力要求的工作节点运行，完全满足卫星数据处理的高算力需求。
• 第三步：结果汇总任务：Airflow默认会等待所有并行的处理任务全部执行成功后，才触发后续的汇总任务，你可以直接读取所有处理任务输出的XCom结果，做统一的归档存储。

简化代码示例 #

from airflow.decorators import dag, task
from datetime import datetime

@dag(start_date=datetime(2024, 1, 1), schedule=None, catchup=False)
def sentinel_data_process_dag():
    @task
    def query_satellite_images():
        # 调用卫星影像检索API，返回符合查询条件的影像ID列表
        return [
            "S1A_IW_GRDH_1SDV_20240501_000000_000A",
            "S1A_IW_GRDH_1SDV_20240502_000000_000B",
            "S1A_IW_GRDH_1SDV_20240503_000000_000C"
        ]

    @task(executor_config={
        "KubernetesExecutor": {
            "resources": {
                "requests": {"cpu": "8", "memory": "32Gi", "nvidia.com/gpu": "1"}
            }
        }
    })
    def process_single_image(image_id):
        # 单景影像的下载、辐射校正、滤波等处理逻辑
        return f"item {image_id} successfully processed"

    @task
    def archive_process_results(results):
        # 汇总所有处理结果，统一写入元数据库或对象存储
        for result in results:
            print(result)
        # 后续持久化逻辑
    
    # 定义任务依赖
    image_list = query_satellite_images()
    process_results = process_single_image.expand(image_id=image_list)
    archive_process_results(process_results)

sentinel_data_process_dag = sentinel_data_process_dag()

补充优化建议 #

• 如果你使用的Airflow版本低于2.3，也可以通过循环遍历上游返回的列表手动生成并行任务，只是代码复杂度会稍高。
• 大体积的中间处理结果不建议直接存入XCom，可以将结果写入对象存储后，仅将存储路径通过XCom传递，降低元数据库存储压力。
• 涉及GPU调度的场景建议搭配KubernetesExecutor使用，资源隔离性和调度灵活度更高，使用体验和Argo工作流基本一致。

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