Airflow任务组动态映射中如何跳过指定索引的后续任务

使用Airflow 2.9.1及Taskflow API，通过任务装饰器创建带动态任务映射的任务组时遇到以下问题：

• 当某一映射任务失败时，希望对应索引的后续任务被跳过，但其他索引的任务若上游成功则继续执行，且任务间无数据传递。
• 已将下游任务触发规则设为all_done，使其能在部分索引任务失败时启动，但由于下游任务未接收上游输入，即便对应上游任务失败仍会运行。
• 预期对应索引的下游任务随上游失败而失败或跳过，而非全部任务被跳过。

已考虑的方案：

• 使用XCom传递任务状态给下游任务：可行但增加复杂度；
• 自定义触发规则：Airflow内置触发规则无法直接满足需求。

疑问：是否有无需自定义XCom处理的替代方案？或能否更优雅地利用Airflow现有功能实现该场景？

附示例DAG代码：

from airflow.decorators import task, dag, task_group
from airflow.utils.dates import days_ago
from airflow.operators.python import get_current_context
from airflow.exceptions import AirflowSkipException

@dag(schedule_interval=None, start_date=days_ago(1), catchup=False)
def skip_subsequent_tasks():
    @task
    def start_task():
        return ['a', 'b', 'c']  # 动态任务映射的参数列表

    @task_group
    def my_task_group(value):
        @task
        def process_task(value):
            if value == 'b':
                raise ValueError("Intentional Failure")
            if value == 'c':
                raise AirflowSkipException("Intentional Skip")
            return f"Processed {value}"

        @task(trigger_rule='all_done')
        def end_task():
            context = get_current_context()
            ti = context['ti']
            current_task_index = ti.map_index
            # TODO: 索引1（value='b'）的任务应失败
            # TODO: 索引2（value='c'）的任务应跳过
            upstream_task_val = ti.xcom_pull(task_ids=f'my_task_group.process_task', map_indexes=current_task_index)
            return f"Ended '{current_task_index}' with output '{upstream_task_val}'"

        process_task(value) >> end_task()

    start = start_task()
    my_task_group.expand(value=start)

dag = skip_subsequent_tasks()

参考状态：

• DAG图显示3组映射任务，每组包含process_task和end_task
• process_task状态：索引0成功，索引1失败，索引2跳过
• end_task状态：当前全部成功（不符合预期）

可以通过利用TaskFlow API的自动依赖关联实现需求，无需手动处理XCom，具体实现如下：

核心思路 #

1. 让下游任务接收上游任务的输出（无需实际使用）：借助TaskFlow的自动绑定逻辑，将同一map_index的上游和下游任务关联，避免手动指定索引匹配。
2. 基于上游任务状态控制下游行为：通过任务上下文获取对应索引的上游任务实例，根据其状态决定当前任务是继续执行、跳过还是失败。

修改后的DAG代码：

from airflow.decorators import task, dag, task_group
from airflow.utils.dates import days_ago
from airflow.operators.python import get_current_context
from airflow.exceptions import AirflowSkipException
from airflow.utils.state import TaskInstanceState

@dag(schedule_interval=None, start_date=days_ago(1), catchup=False)
def skip_subsequent_tasks():
    @task
    def start_task():
        return ['a', 'b', 'c']  # 动态任务映射的参数列表

    @task_group
    def my_task_group(value):
        @task
        def process_task(value):
            if value == 'b':
                raise ValueError("Intentional Failure")
            if value == 'c':
                raise AirflowSkipException("Intentional Skip")
            return f"Processed {value}"

        @task(trigger_rule='all_done')
        def end_task(upstream_output):
            context = get_current_context()
            ti = context['ti']
            # 获取对应索引的上游任务实例
            upstream_ti = ti.get_dagrun().get_task_instance(
                task_id=f"{ti.task_group_id}.process_task",
                map_index=ti.map_index
            )
            
            # 根据上游状态处理当前任务
            if upstream_ti.state == TaskInstanceState.FAILED:
                raise ValueError(f"Upstream task for index {ti.map_index} failed")
            elif upstream_ti.state == TaskInstanceState.SKIPPED:
                raise AirflowSkipException(f"Upstream task for index {ti.map_index} was skipped")
            
            return f"Ended '{ti.map_index}' with output '{upstream_output}'"

        # 建立依赖并传递上游输出，自动关联同索引任务
        process = process_task(value)
        process >> end_task(process)

    start = start_task()
    my_task_group.expand(value=start)

dag = skip_subsequent_tasks()

方案优势 #

• 无需手动管理XCom拉取和索引匹配，利用TaskFlow原生依赖逻辑简化代码；
• 通过Airflow内置的TaskInstanceState枚举判断状态，代码更规范易维护；
• 精准控制单索引任务行为：上游失败则下游失败，上游跳过则下游跳过，上游成功则下游正常执行，不影响其他索引任务流程。

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