AWS环境下Airflow执行元数据批量依赖任务时状态更新报错

PostgreSQL元数据中配置了数百个带依赖关系的任务，需要从元数据拉取任务并按依赖执行，同时跟踪任务状态。涉及技术包括Airflow、AWS、PostgreSQL，任务类型涵盖EMR、Lambda、Python脚本等。

已尝试方案及问题 #

• 方案1：在DAG文件中循环任务生成TaskGroup，每个Group包含获取EMR应用ID、执行任务、更新任务状态、更新对比状态的Operator。
  问题：job_status_complete执行成功，但comp_status_ready失败，原因是DAG多次解析导致任务从列表中被移除。
• 方案2：在PythonOperator的generate_test_tasks函数内循环任务并创建Operator链，但Operator嵌套调用不生效。

相关代码 #

方案1代码片段 #

jobs = get_pending_test_cases(batch_id)
task_map = {}

task_logger.info(jobs)
for idx, job in enumerate(jobs): 

    task_logger.info(f"Job-{idx}: {job}") 
    task_logger.info(f"Job type: {type(job)}")

    test_id, job_id, job_args, service_type, job_status, comparison_status = job
    with TaskGroup(f"job_tasks_{test_id}") as job_tasks:
        emr_app_id = PythonOperator(
           task_id="get_emr_app_id",
           python_callable=lookup_emr_app_id,
           op_args=['qa100'],
        )

        job_execution_task = PythonOperator(
           task_id=f"run_{test_id}",
           python_callable=execute_test_job,
           op_kwargs={'test_id': test_id},
           dag=dag
        )

        job_status_complete = MetaOperator(
           task_id=f"update_job_status_{test_id}",
           trigger_rule="all_done",
           cmd="update_job_status",
           op_kwargs={
               "test_id": test_id,
               "batch_id": batch_id,
               "job_status": "COMPLETE",
           },
        )

        comp_status_ready = MetaOperator(
           task_id=f"update_comp_status_{test_id}",
           trigger_rule="all_done",
           cmd="update_comp_status",
           op_kwargs={
               "test_id": test_id,
               "batch_id": batch_id,
               "comp_status": "READY",
           },
        )

        (
               emr_app_id
               >> job_execution_task
               >> job_status_complete
               >> comp_status_ready
        )

   job_tasks

方案2代码片段 #

def generate_test_tasks(batch_id: str, dag, **context):

    test_jobs = context["ti"].xcom_pull(task_ids="get_pending_test_cases")
    task_logger.info(f"test_jobs: {test_jobs}")

    task_map = {}
    for idx, job in enumerate(test_jobs):
        task_logger.info(f"Job-{idx}: {job}")
        task_logger.info(f"Job type: {type(job)}")

        test_id, job_id, job_args, service_type, job_status, comparison_status = job
        task_logger.info(f"{test_id}, {job_id}, {job_args},{service_type}, {job_status}, {comparison_status}")

        job_execution_task = PythonOperator(
            task_id=f"run_{test_id}",
            python_callable=execute_test_job,
            op_kwargs={'test_id': test_id},
            dag=dag
        )

        job_status_complete = MetaOperator(
            task_id=f"update_job_status_{test_id}",
            trigger_rule="all_done",
            cmd="update_job_status",
            op_kwargs={
                "test_id": test_id,
                "batch_id": batch_id,
                "job_status": "COMPLETE",
            },
        )

        job_execution_task.execute(dict())

    return True

元数据获取相关代码 #

# pending_test_cases: Meta operator 返回如下格式的列表
test_cases.append([test_id, job_id, job_args, service_type, job_status, comparison_status])

test_cases = MetaOperator(
        task_id="get_pending_test_cases",
        cmd="pending_test_cases",
        op_kwargs={
            "batch_id": batch_id,
            "is_prev_batch_incomplete": is_prev_batch_incomplete,
        }
    )


generate_tasks_pipeline = PythonOperator(
        task_id="gen_test_tasks",
        python_callable=generate_test_tasks,
        provide_context=True,
        op_kwargs={
            "batch_id": batch_id,
        },
        dag=dag
    )

test_cases >> generate_tasks_pipeline

针对方案1的问题修复 #

方案1中任务被移除的核心原因是DAG解析时任务列表动态变化：Airflow周期性解析DAG，若每次解析时get_pending_test_cases返回的任务列表不一致，会导致之前生成的TaskGroup/Operator被丢弃，进而出现后续任务找不到的错误。

修复步骤：

1. 固定任务标识：确保每个TaskGroup和Operator的ID基于元数据中稳定的test_id生成，避免因任务顺序变化导致ID变更。
2. 静态结构+动态触发：用DynamicTaskMapping（Airflow 2.2+支持）替代手动循环，将任务拉取和实例生成分开：
   • 第一步：用Operator拉取待执行任务并存入XCom。
   • 第二步：基于XCom结果动态生成TaskGroup实例。

示例代码调整：

from airflow.decorators import task_group

# 定义单个任务的TaskGroup模板
@task_group
def job_task_group(test_id, batch_id):
    emr_app_id = PythonOperator(
       task_id="get_emr_app_id",
       python_callable=lookup_emr_app_id,
       op_args=['qa100'],
    )

    job_execution_task = PythonOperator(
       task_id=f"run_{test_id}",
       python_callable=execute_test_job,
       op_kwargs={'test_id': test_id},
    )

    job_status_complete = MetaOperator(
       task_id=f"update_job_status_{test_id}",
       trigger_rule="all_done",
       cmd="update_job_status",
       op_kwargs={
           "test_id": test_id,
           "batch_id": batch_id,
           "job_status": "COMPLETE",
       },
    )

    comp_status_ready = MetaOperator(
       task_id=f"update_comp_status_{test_id}",
       trigger_rule="all_done",
       cmd="update_comp_status",
       op_kwargs={
           "test_id": test_id,
           "batch_id": batch_id,
           "comp_status": "READY",
       },
    )

    emr_app_id >> job_execution_task >> job_status_complete >> comp_status_ready

# 在DAG中使用动态映射
test_cases = MetaOperator(
    task_id="get_pending_test_cases",
    cmd="pending_test_cases",
    op_kwargs={
        "batch_id": batch_id,
        "is_prev_batch_incomplete": is_prev_batch_incomplete,
    },
    do_xcom_push=True
)

# 动态生成TaskGroup实例
job_tasks_mapped = job_task_group.expand(
    test_id=test_cases.output.map(lambda x: x[0]),
    batch_id=batch_id
)

test_cases >> job_tasks_mapped

针对方案2的问题修复 #

方案2中直接调用job_execution_task.execute(dict())是错误的：Airflow的任务定义必须在DAG解析阶段完成，运行阶段仅执行预定义任务，不能动态创建并执行Operator。

正确做法：

• 若使用Airflow 2.2+，直接采用上述DynamicTaskMapping方案。
• 若版本低于2.2，可将任务生成逻辑移到DAG解析阶段，添加缓存机制保证每次解析返回的任务列表一致；或用BranchPythonOperator结合TriggerDagRunOperator，为每个任务触发独立子DAG。

通用注意事项 #

• 幂等性设计：确保update_job_status和update_comp_status操作是幂等的，避免任务重试导致状态更新错误。
• DAG解析稳定性：避免在DAG解析阶段执行实时数据库查询，若必须查询，需添加缓存保证结果一致。
• 依赖处理：拉取任务时同步获取依赖信息，通过>>或set_upstream动态构建任务依赖链。

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