Airflow：如何处理动态Task Group映射内的任务依赖关系

我在Airflow中实现了包含Python Operator和自定义可延迟Operator的动态映射Task Group。当前遇到的问题是：当可延迟Operator进入延迟状态时，系统会触发下一个映射实例的Task Group任务，但我需要等待当前实例的可延迟Operator完全执行完成后，再启动下一个映射实例。

相关代码 #

import logging
from airflow.decorators import dag, task_group, task
from airflow.utils.weight_rule import WeightRule
from pendulum import datetime
from airflow.operators.python import get_current_context, PythonOperator
from airflow.utils.context import context_merge
from airflow.triggers.temporal import TimeDeltaTrigger
from datetime import timedelta
from typing import Any, Mapping
from airflow.operators.empty import EmptyOperator


class DeferrableOperator(PythonOperator):
    """
    Deferrable Sensor class
    """

    def __init__(self, op_kwargs: Mapping[str, Any] | None = None, poke_interval=30, **kwargs):
        super().__init__(**kwargs)
        self.op_kwargs = op_kwargs or {}
        self.poke_interval = poke_interval

    def execute(self, context):
        """
        Method will be invoked automatically
        :param context:
        """
        # Unpacking Value like super class
        context_merge(context, self.op_kwargs, templates_dict=self.templates_dict)
        self.op_kwargs = self.determine_kwargs(context)
        logging.info(f"Arguments {self.op_kwargs}")
        self.defer(
            trigger=TimeDeltaTrigger(delta=timedelta(seconds=self.poke_interval)),
            method_name="execute_complete",
        )

    def execute_complete(self, context, event=None):
        """
        Method will be triggered by execute method
        :param context:
        :param event:
        """
        context.update(self.op_kwargs)
        num = self.python_callable(**self.op_kwargs)
        if not num:
            self.execute(context)


def add(num):
    return int(num) + 10


@dag(
    start_date=datetime(2022, 12, 1),
    schedule=None,
    catchup=False,
    max_active_tasks=1
)
def task_group_mapping_example3():

    @task
    def push_xcom():
        context = get_current_context()
        ti = context["ti"]
        ti.xcom_push(key="batches", value=[19, 23, 42])
        return [19, 23, 42]

    @task_group(group_id="group1")
    def tg1(my_num):

        @task(weight_rule=WeightRule.ABSOLUTE)
        def print_num(num):
            return num

        print_num = print_num(my_num)

        add_num = DeferrableOperator(
            task_id="add",
            poke_interval=30,
            op_kwargs={"num": print_num},
            python_callable=add,
            weight_rule=WeightRule.ABSOLUTE
        )

        end = EmptyOperator(task_id="batch_execution_completed", weight_rule=WeightRule.ABSOLUTE, priority_weight=1)

        print_num >> add_num >> end

    # a downstream task to print out resulting XComs
    @task
    def pull_xcom(**context):
        pulled_xcom = context["ti"].xcom_pull(
            # reference a task in a task group with task_group_id.task_id
            task_ids=["group1.add_42"],
            # only pull Xcom from specific mapped task group instances (2.5 feature)
            map_indexes=[2, 3],
            key="return_value",
        )

        # will print out a list of results from map index 2 and 3 of the add_42 task
        print(pulled_xcom)

    tg1.expand(my_num=push_xcom()) >> pull_xcom()


task_group_mapping_example3()

解决方案 #

问题根源在于：当可延迟Operator进入延迟状态时，Airflow会释放占用的Worker资源，此时max_active_tasks=1无法阻止下一个映射实例启动。要实现串行执行，需让每个映射实例等待前一个实例完全完成后再启动，可通过以下方式实现：

方法：添加前置等待任务

在Task Group内部新增一个任务，检查当前实例的前一个实例是否已完成，未完成则重试，直到前一个实例成功结束。

修改后的tg1代码如下：

@task_group(group_id="group1")
def tg1(my_num):
    from airflow.models import TaskInstance

    @task(weight_rule=WeightRule.ABSOLUTE, retries=100, retry_delay=timedelta(seconds=30))
    def wait_for_previous_batch(**context):
        map_index = context["ti"].map_index
        if map_index == 0:
            return
        # 获取前一个实例的end任务实例
        prev_end_ti = TaskInstance(
            task_id="group1.batch_execution_completed",
            dag_id=context["dag"].dag_id,
            execution_date=context["execution_date"],
            map_index=map_index - 1
        )
        prev_end_ti.refresh_from_db()
        if prev_end_ti.state != "success":
            raise Exception(f"等待前一批次（索引：{map_index-1}）完成...")

    @task(weight_rule=WeightRule.ABSOLUTE)
    def print_num(num):
        return num

    wait_task = wait_for_previous_batch()
    print_num = print_num(my_num)

    add_num = DeferrableOperator(
        task_id="add",
        poke_interval=30,
        op_kwargs={"num": print_num},
        python_callable=add,
        weight_rule=WeightRule.ABSOLUTE
    )

    end = EmptyOperator(task_id="batch_execution_completed", weight_rule=WeightRule.ABSOLUTE, priority_weight=1)

    # 任务依赖：等待前一批次完成 → 打印数字 → 延迟计算 → 标记完成
    wait_task >> print_num >> add_num >> end

关键说明

• wait_for_previous_batch任务通过map_index识别当前实例位置，仅当不是第一个实例时才触发等待逻辑。
• 利用TaskInstance查询前一个实例的end任务状态，若未成功则抛出异常触发重试，直到前一个实例完成。
• 设置足够的retries和retry_delay，确保能等待前一个实例完成。

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