Airflow动态任务映射引发DagBag导入超时问题求解

问题根源 #

你的代码在DAG顶层使用了while True循环，Airflow加载DAG时会执行所有顶层代码，直接进入无限循环，触发DagBag导入超时。Airflow要求顶层代码仅用于定义DAG结构，不能包含运行时的循环逻辑，循环必须放到任务执行阶段处理。

解决方案：用TaskFlow API实现循环+动态任务映射 #

通过分支任务控制循环流程，结合动态任务映射处理每个批次条目，核心思路是：

1. 定义获取数据批次的任务，返回当前批次条目和下一轮的cursor；
2. 用动态映射批量处理当前批次的所有条目；
3. 用分支任务判断是否需要继续获取下一批数据，若需要则重复执行上述流程，否则结束。

完整代码示例

from airflow.decorators import dag, task, task_group
from datetime import datetime

@dag(
    start_date=datetime(2024, 1, 1),
    schedule_interval=None,
    catchup=False,
    tags=["batch_processing"]
)
def batch_item_processor():
    BATCH_LIMIT = 10

    @task(multiple_outputs=True)
    def fetch_batch(cursor=None):
        # 替换为实际的API调用逻辑
        if cursor is None:
            new_cursor = BATCH_LIMIT + 1
            items = list(range(0, BATCH_LIMIT))
        else:
            # 模拟API返回：当cursor超过20时，返回None表示无更多数据
            if cursor > 20:
                new_cursor = None
                items = []
            else:
                new_cursor = cursor + BATCH_LIMIT + 1
                items = list(range(cursor, cursor + BATCH_LIMIT))
        
        return {"next_cursor": new_cursor, "batch_items": items}

    @task
    def process_single_item(item):
        print(f"Processing item ID: {item}")
        # 这里写实际的业务处理逻辑

    @task.branch
    def check_continue_processing(next_cursor):
        # 判断是否还有下一批数据
        if next_cursor is not None:
            return "batch_processing_group.fetch_batch"
        else:
            return "finalize_processing"

    @task
    def finalize_processing():
        print("All batches have been processed successfully.")

    # 定义任务组，封装单次批次的获取和处理流程
    @task_group
    def batch_processing_group():
        batch_data = fetch_batch()
        # 动态映射处理当前批次的所有条目
        process_single_item.expand(item=batch_data["batch_items"])
        return batch_data["next_cursor"]

    # 初始化流程
    initial_cursor = batch_processing_group()
    branch = check_continue_processing(initial_cursor)

    # 设置循环依赖：分支指向任务组或结束任务
    branch >> [batch_processing_group(), finalize_processing()]

# 实例化DAG
batch_item_processor()

代码说明

1. 任务组（TaskGroup）：将fetch_batch和process_single_item.expand封装成可重复执行的单元，简化循环逻辑的依赖管理；
2. 分支任务：根据fetch_batch返回的next_cursor判断是否继续循环，要么再次执行任务组，要么跳转到结束任务；
3. 动态任务映射（expand）：基于fetch_batch返回的batch_items自动生成多个process_single_item任务，实现批量处理；
4. 避免顶层循环：所有循环逻辑都在任务执行阶段通过分支判断实现，DAG解析时不会触发无限循环，彻底解决导入超时问题。

注意事项 #

• 实际使用时，需将fetch_batch中的模拟逻辑替换为真实API调用，确保正确返回批次数据和下一轮cursor（或None表示结束）；
• 需使用Airflow 2.2+版本，该版本开始支持动态任务映射；
• 可根据业务需求调整BATCH_LIMIT值及分支判断条件。

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