Airflow需求：任务组内任务重试时，重试整个任务组TG

背景与需求 #

我有一个自定义内部任务组TG，被用于多个DAG中，该任务组包含A、B、C、D等多个任务。目前任务D因意外故障触发重试。
需求是：不单独重试任务D，而是重试整个任务组TG，即清理并重新运行整个TG（包括下游任务，递归执行）。

原尝试方案 #

我尝试在任务D中传入回调函数：

on_retry_callback=lambda context: reset_and_retry_tg(context, group_id)

实现代码如下：

@provide_session
def reset_and_retry_tg(context: Any, task_group_id: str, session: Any = None) -> None:
    dag_run = context["dag_run"]
    dag = dag_run.get_dag()
    task_instances = dag_run.get_task_instances(session=session)
    tasks_to_clear = set()
    for task in dag.tasks:
        if task.task_id.startswith(f"{task_group_id}."):
            tasks_to_clear.add(task.task_id)
            tasks_to_clear.update(dag.get_task(task.task_id).get_flat_relatives(upstream=False))

    LOG.info("Clearing tasks: %s", tasks_to_clear)

    tis_to_clear = [ti for ti in task_instances if ti.task_id in tasks_to_clear]
    if not tis_to_clear:
        LOG.warning("No tasks found in task group %s to clear.", task_group_id)
        return

    clear_task_instances(
        tis=tis_to_clear,
        session=session,
        dag=dag,
    )

    LOG.info("Cleared all tasks in task group %s to retry.", task_group_id)
    raise Exception(  # pylint: disable = broad-exception-raised
        f"Retrying entire task group '{task_group_id}' because task retries."
    )

但该方案未达到预期效果，仍仅任务D在重试。

原方案问题分析 #

1. 重试流程冲突：on_retry_callback是在任务D的重试流程中触发的，此时Airflow已经调度了任务D的重试，即使清除其他任务实例，也会继续执行任务D的重试，不会触发整个任务组的重新运行。
2. 缺乏主动触发逻辑：清除任务实例仅重置状态，但没有主动触发这些任务重新运行，依赖DAG调度机制无法及时响应。
3. 异常抛出的副作用：回调中抛出的Exception会标记任务D为失败，但不会触发整个任务组的重新执行，反而中断了当前重试流程。

可行解决方案与思路 #

方案一：改用失败回调+主动触发任务重跑 #

关闭任务D的自动重试，让失败直接触发回调，在回调中清除任务实例并主动触发重跑：

1. 修改任务D的定义，关闭自动重试：

task_d = PythonOperator(
    task_id="D",
    python_callable=your_task_function,
    retries=0,  # 关闭自动重试，失败直接触发回调
    on_failure_callback=lambda context: reset_and_retry_tg(context, "TG"),
    dag=dag
)

2. 更新回调函数，增加触发任务重跑的逻辑：

from airflow.api.common.experimental.trigger_task import trigger_task

@provide_session
def reset_and_retry_tg(context: Any, task_group_id: str, session: Any = None) -> None:
    dag_run = context["dag_run"]
    dag = dag_run.get_dag()
    task_instances = dag_run.get_task_instances(session=session)
    tasks_to_clear = set()
    
    # 收集任务组内所有任务及下游任务
    for task in dag.tasks:
        if task.task_id.startswith(f"{task_group_id}."):
            tasks_to_clear.add(task.task_id)
            tasks_to_clear.update(dag.get_task(task.task_id).get_flat_relatives(upstream=False))

    LOG.info("Clearing and retrying tasks: %s", tasks_to_clear)

    tis_to_clear = [ti for ti in task_instances if ti.task_id in tasks_to_clear]
    if not tis_to_clear:
        LOG.warning("No tasks found in task group %s to clear.", task_group_id)
        return

    # 清除任务实例状态
    clear_task_instances(
        tis=tis_to_clear,
        session=session,
        dag=dag,
        reset_dag_runs=False,
    )

    # 逐个触发任务重新运行
    for task_id in tasks_to_clear:
        trigger_task(
            dag_id=dag.dag_id,
            task_id=task_id,
            run_id=dag_run.run_id,
            session=session
        )

    LOG.info("Triggered retry for entire task group %s and downstream tasks.", task_group_id)

方案二：利用Airflow原生TaskGroup重试（2.2+版本） #

Airflow 2.2及以上版本支持为TaskGroup直接设置重试策略，当组内任意任务失败时，整个TaskGroup会自动重置并重跑：

from airflow.utils.task_group import TaskGroup
from datetime import timedelta

with TaskGroup("TG", retries=3, retry_delay=timedelta(minutes=5)) as tg:
    task_a = PythonOperator(task_id="A", python_callable=...)
    task_b = PythonOperator(task_id="B", python_callable=...)
    task_c = PythonOperator(task_id="C", python_callable=...)
    task_d = PythonOperator(task_id="D", python_callable=...)
    
    task_a >> task_b >> task_c >> task_d

该方案无需自定义回调，依赖Airflow原生特性，简洁高效。

方案三：DAG级别统一监听处理 #

在DAG层面设置任务失败回调，统一判断失败任务是否属于目标TaskGroup，再执行重试逻辑：

def dag_level_failure_callback(context):
    failed_task_id = context["task_instance"].task_id
    target_tg_id = "TG"
    
    if failed_task_id.startswith(f"{target_tg_id}."):
        reset_and_retry_tg(context, target_tg_id)

dag = DAG(
    dag_id="your_dag_id",
    default_args=default_args,
    on_task_failure=dag_level_failure_callback,
    schedule_interval="@daily"
)

这种方式无需单独给每个任务配置回调，适合多任务组的场景。

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