Airflow中如何将KubernetesPodOperator推送的XCom用于TaskGroup循环？

我有一个使用KubernetesPodOperator的Airflow DAG，其中get_train_test_model_task_count任务推送了一个XCom变量，希望在后续任务中使用该变量。

以下BashOperator任务可以正常工作，成功打印出ti_key=24：

run_this = BashOperator(
    task_id="also_run_this",
    bash_command='echo "ti_key={{ ti.xcom_pull(task_ids=\"get_train_test_model_task_count\", key=\"return_value\")[\"models_count\"] }}"',
)

尝试将该变量用于TaskGroup中的循环生成任务，代码如下：

with TaskGroup("train_test_model_config") as train_test_model_config:
    models_count = "{{ ti.xcom_pull(task_ids=\"get_train_test_model_task_count\", key=\"return_value\")[\"models_count\"] }}"
    print(models_count)
    for task_num in range(0, int(models_count)):
        generate_train_test_model_config_task(task_num)

但int(models_count)执行失败，抛出错误：

ValueError: invalid literal for int() with base 10: '{{ ti.xcom_pull(task_ids="get_train_test_model_task_count", key="return_value")["models_count"] }}'

generate_train_test_model_config_task定义如下：

def generate_train_test_model_config_task(task_num):
    task = KubernetesPodOperator(
        name=f"train_test_model_config_{task_num}",
        image=build_model_image,
        labels=labels,
        cmds=[
            "python3",
            "-m",
            "src.models.train_test_model_config",
            "--tenant=neu",
            f"--model_tag_id={task_num}",
            "--line_plan={{ ti.xcom_pull(key=\"file_name\", task_ids=\"extract_file_name\") }}",
            "--staging_bucket=cs-us-ds"
        ],
        task_id=f"train_test_model_config_{task_num}",
        do_xcom_push=False,
        namespace="airflow",
        service_account_name="airflow-worker",
        get_logs=True,
        startup_timeout_seconds=300,
        container_resources={"request_memory": "29G", "request_cpu": "7000m"},
        node_selector={"cloud.google.com/gke-nodepool": NODE_POOL},
        tolerations=[
            {
                "key": NODE_POOL,
                "operator": "Equal",
                "value": "true",
                "effect": "NoSchedule",
            }
        ],
    )

    return task

如何解决这个问题，让XCom变量可以正常用于循环生成任务？

错误原因 #

报错核心是Airflow的DAG解析阶段与任务运行阶段完全分离：

• 循环for task_num in range(0, int(models_count))在DAG解析时执行，此时Jinja模板字符串还未被渲染，仍是原始文本，无法转成整数。
• BashOperator中的Jinja模板是在任务运行时才会被替换为实际XCom值，因此能正常工作。

正确实现：使用动态任务映射（Dynamic Task Mapping） #

Airflow 2.2+支持动态任务映射，可在运行时根据XCom值动态生成任务实例，适配这类场景。修改步骤如下：

1. 调整任务生成函数，改用模板参数区分实例：

def generate_train_test_model_config_task():
    task = KubernetesPodOperator(
        name="train_test_model_config_{{ task_instance_key_str }}",
        image=build_model_image,
        labels=labels,
        cmds=[
            "python3",
            "-m",
            "src.models.train_test_model_config",
            "--tenant=neu",
            "--model_tag_id={{ task_instance_key_str }}",
            "--line_plan={{ ti.xcom_pull(key=\"file_name\", task_ids=\"extract_file_name\") }}",
            "--staging_bucket=cs-us-ds"
        ],
        task_id="train_test_model_config",
        do_xcom_push=False,
        namespace="airflow",
        service_account_name="airflow-worker",
        get_logs=True,
        startup_timeout_seconds=300,
        container_resources={"request_memory": "29G", "request_cpu": "7000m"},
        node_selector={"cloud.google.com/gke-nodepool": NODE_POOL},
        tolerations=[
            {
                "key": NODE_POOL,
                "operator": "Equal",
                "value": "true",
                "effect": "NoSchedule",
            }
        ],
    )

    return task

2. 在TaskGroup中用expand实现动态映射：

with TaskGroup("train_test_model_config") as train_test_model_config:
    # 从XCom获取models_count并生成数字列表
    def get_model_numbers(**context):
        models_count = context["ti"].xcom_pull(task_ids="get_train_test_model_task_count", key="return_value")["models_count"]
        return list(range(models_count))

    # 动态生成任务实例
    generate_train_test_model_config_task().expand(task_instance_key_str=get_model_numbers())

简化写法（Airflow 2.3+） #

如果使用Airflow 2.3及以上版本，可直接用Jinja模板在expand中引用XCom，无需额外函数：

with TaskGroup("train_test_model_config") as train_test_model_config:
    generate_train_test_model_config_task().expand(
        task_instance_key_str="{{ ti.xcom_pull(task_ids='get_train_test_model_task_count', key='return_value')['models_count'] | range | list }}"
    )

内容的提问来源于stack exchange，提问作者Tom J Muthirenthi