Apache Airflow 任务组拓扑排序优化：pass-number 遍历如何消除反向声明 DAG 的 O(N²) 最坏情况

Apache Airflow 在最新改进中进一步优化了 TaskGroup.topological_sort：针对"反向声明"（reverse-declared）DAG，通过 pass-number 遍历算法，将原先在对抗性结构（如反向插入链）下高达 O(N²) 的最坏情况时间复杂度大幅改善。本文以该改进为核心，结合 Airflow 核心仓库与 task-sdk 的源码实现和单元测试，深入讲解排序算法分支的调度逻辑、两种算法各自的适用场景，以及这一优化对大型 DAG 调度与 UI 渲染的实际价值。

从 Newsfragment 看本次改进的定位

本次改进的原始记录位于 airflow-core/newsfragments/67688.improvement.rst：

Further optimize TaskGroup.topological_sort for reverse-declared DAGs via pass-number traversal; dramatically improves the O(N²) worst-case for adversarial shapes (e.g., reverse-insertion chains).

它是 Airflow 变更日志体系（newsfragments）中的一条 improvement 条目。在 Airflow 的贡献流程中，每次改动都会在 newsfragments 目录下添加独立片段文件，最终由发布工具合并进 RELEASE_NOTES.rst。这条片段揭示了本版本（主分支 3.x 演进阶段）对任务组排序性能的一次关键改进：排序算法不再是一条固定路径，而是根据依赖结构动态选择"扫描"与"pass-number 遍历"两种策略。

TaskGroup 拓扑排序解决什么问题

TaskGroup 是 Airflow 中"视觉上相关的一组任务"的组织单元，用于在 DAG 中创建分组层级。当 DAG 被调度器序列化（serialize）后，SerializedTaskGroup 会在受保护进程（scheduler 等）中重建整个任务图。

topological_sort 的核心职责是：对某个 TaskGroup 的直接子节点（子任务或子 TaskGroup）进行排序，保证任何任务都排在其上游依赖之后。这一排序结果直接决定：

• DAG 图中各节点在 UI 上的展示顺序；
• 序列化后的任务组在调度器侧的重建顺序；
• 后续遍历（如 Graph/Grid 视图构建）的一致性基础。

排序对象是 children 字典的值（dict[str, DAGNode]），因此排序还隐含地需要处理"兄弟节点之间"与"跨 TaskGroup 祖先层级"的依赖关系——这正是 _project_child_deps 要做的投影工作（见下文）。

性能瓶颈：反向声明 DAG 的 O(N²) 最坏情况

什么是"反向声明"结构

所谓反向声明（reverse-declared），指的是先声明的节点反而依赖后声明的节点。测试用例 task-sdk/tests/task_sdk/definitions/test_taskgroup.py 中的 _make_reverse_chain 精确刻画了这种形状：

def _make_reverse_chain(n: int) -> DAG:
    with DAG(f"reverse_chain_{n}", schedule=None, start_date=DEFAULT_DATE) as dag:
        tasks = [EmptyOperator(task_id=f"t{n - 1 - i}") for i in range(n)]
        by_id = {t.task_id: t for t in tasks}
        for i in range(n - 1):
            by_id[f"t{i}"] >> by_id[f"t{i + 1}"]
    return dag

这里先创建 t{n-1-i}（即先创建 t99、t98……再创建 t0），但边 t0 -> t1 -> ... -> t99 意味着声明顺序恰好与依赖方向相反：位于 children 字典前面的节点反而依赖位于后面的节点。

为什么朴素扫描会退化为 O(N²)

在改进之前，排序核心是 task-sdk/src/airflow/sdk/definitions/taskgroup.py 中的 _sweep_projection（多轮贪心扫描）：

• 第一遍按索引顺序检查每个节点，若其所有依赖都已发射（emitted）则立即输出；否则进入 pending 列表等待下一轮；
• 每一轮只重新检查 pending 中的节点。

对于反向链 t0 -> t1 -> ... -> t99，第一轮只有 t99 能发射（因为 t99 无依赖……实际上整条链的入口是 t0 的被依赖方 t99）。更一般地，每一轮扫描只能向前推进一个"层级"，而每轮都要遍历剩余所有 pending 节点，于是总代价接近 O(N²)（N 为子节点数）。当链特别长、且混入大量独立节点"填充"（padded）时，扫描轮次与每轮开销叠加，最坏情况呈平方级增长。

核心实现中的注释也印证了这一点：

# The ratio catches dense back-heavy groups; a 32-node absolute cutoff keeps
# padded reverse-declared runs on the fast path once sweep rescans overtake pass-numbering.

改进方案：双分支调度 + pass-number 遍历

统一的依赖投影（projection）

两条算法分支共享同一个预处理：_project_child_deps 将每个子节点的上游依赖投影为兄弟节点层面的整数索引集合：

def _project_child_deps(self, child_idx, child, id_to_idx, group_dict) -> tuple[int, ...]:
    upstream_ids = child._topological_upstream_ids
    if not upstream_ids:
        return ()
    sib_deps: set[int] = set()
    for edge_id in upstream_ids:
        j = id_to_idx.get(edge_id)
        if j is not None:
            if j != child_idx:
                sib_deps.add(j)
            continue
        tg = group_dict.get(edge_id)
        if tg is None:
            edge = self.dag.get_task(edge_id)
            tg = edge.task_group
        while tg is not None:
            anc_idx = id_to_idx.get(tg.node_id)
            if anc_idx is not None:
                if anc_idx != child_idx:
                    sib_deps.add(anc_idx)
                break
            tg = tg.parent_group
    return tuple(sib_deps)

这段逻辑（airflow-core/src/airflow/serialization/definitions/taskgroup.py）做三件事：

1. 直接兄弟依赖：上游任务 id 恰好是当前组的兄弟节点，直接记入索引；
2. 跨组依赖：上游任务属于其他 TaskGroup，则沿其 parent_group 链向上查找，直到命中当前组的一个祖先兄弟节点（anc_idx）；
3. 忽略自环：j != child_idx / anc_idx != child_idx 排除指向自身的边。

投影完成后，问题被规约为"对 N 个整数索引节点做拓扑排序"，两个算法分支都在这个紧凑的数据结构上运行。

分支选择：基于"反向边"的启发式

airflow-core/src/airflow/serialization/definitions/taskgroup.py 中，topological_sort 的主流程统计了"反向边"数量 nodes_with_back_edge（即 deps 中存在索引 d > i 的节点数，表示有依赖指向声明顺序更靠后的节点）：

projected: list[tuple[int, ...]] = [()] * n
nodes_with_back_edge = 0
for i, child in enumerate(nodes):
    deps = self._project_child_deps(i, child, id_to_idx, group_dict)
    if deps:
        projected[i] = deps
        if any(d > i for d in deps):
            nodes_with_back_edge += 1

# 分支调度：满足任一条件即走 pass-numbering
if nodes_with_back_edge >= 32 or nodes_with_back_edge * 2 > n:
    return self._sort_via_pass_numbering(nodes, projected)
return self._sweep_projection(nodes, projected)

调度规则清晰且可复述：

条件	含义	走哪个分支
nodes_with_back_edge >= 32	反向节点数达到绝对阈值	pass-number 遍历
nodes_with_back_edge * 2 > n	反向节点占比超过一半（密集反向结构）	pass-number 遍历
其余情况（正向声明或少量反向边）	常见形态	多轮扫描（sweep）

阈值 32 是"绝对截止"，用于保证即使反向节点比例未过半（如被大量独立节点填充的 padded 反向链），一旦规模足够大也切到 pass-numbering——因为此时扫描的重复轮次成本已经超过 pass-numbering 的常数开销。

pass-number 遍历的算法细节

_sort_via_pass_numbering（task-sdk/src/airflow/sdk/definitions/taskgroup.py）采用 Kahn 式遍历为每个节点计算"最早可发射轮次"（pass number）：

• 维护 in_degree（投影后依赖数）与 successors（后继表）；
• 从零入度节点出发 BFS，pass_of[i] 的递推式为：

pass(i) = max over deps d of (
    pass(d)            if idx(d) < idx(i)   # 依赖声明在前，可同轮发射
    pass(d) + 1        if idx(d) > idx(i)   # 依赖声明在后，必须推迟到下一轮
)

• 最终按 (pass_of[i], i) 排序输出，即"轮次级升序，同轮按插入顺序"。

这个递推式的妙处在于：它精确模拟了贪心扫描的发射时序，却不需要真的多轮扫描。反向声明链在扫描分支下每轮只前进一层，而 pass-numbering 用一次 O(V + E) 的 Kahn 遍历就把所有轮次算出来，再以排序收尾，总复杂度约为 O((V + E) log V)，彻底绕开 O(N²)。

为什么两个分支的输出顺序一致

算法注释明确说明两条分支产生完全一致的发射顺序："level-by-legacy-pass，平局按 children 插入顺序打破"（见 task-sdk/src/airflow/sdk/definitions/taskgroup.py）。由于 pass-number 的递推式是对扫描行为的精确建模，两者输出的节点序列相同，因此切换分支不会改变排序结果，只改变计算成本——这是该优化能够安全合入的关键前提。

环检测与防御性处理

拓扑排序天然会遇到环。两条分支对环的处理保持一致且显式：

• _sweep_projection 中，若某一轮 next_pending 的长度与 pending 相同（没有任何节点得到推进），立即抛出异常；
• _sort_via_pass_numbering 中，若 processed != n（BFS 无法覆盖全部节点），同样抛错。

在 airflow-core 序列化侧抛出的异常信息为 ValueError(f"A cyclic dependency occurred in dag: {self.dag_id}")，而在 task-sdk 侧抛出的是 AirflowDagCycleException。从源码注释看，设计意图是：DAG 在序列化前已由 DAG.check_cycle 拒绝环形图，因此排序阶段遇到环说明序列化数据本身损坏，此时宁可显式报错也不静默死循环（见 airflow-core/src/airflow/serialization/definitions/taskgroup.py）。

测试验证：对抗性结构与分支触发

优化是否真正生效，取决于单元测试能否构造出"逼退"扫描分支的反向结构。核心测试用例在 airflow-core/tests/unit/utils/test_task_group.py：

def test_topological_sort_serialized_padded_reverse_chain_uses_pass_numbering(monkeypatch):
    dag = _make_padded_reverse_chain(chain_length=80, independent_count=80)
    serialized = create_scheduler_dag(dag)
    serialized.task_group.children = {
        **{f"r{i}": serialized.task_group.children[f"r{i}"] for i in range(79, -1, -1)},
        **{f"i{i}": serialized.task_group.children[f"i{i}"] for i in range(80)},
    }

    called = {"value": False}
    serialized_task_group_cls = type(serialized.task_group)
    original = serialized_task_group_cls._sort_via_pass_numbering

    def spy(self, nodes, projected):
        called["value"] = True
        return original(self, nodes, projected)

    monkeypatch.setattr(serialized_task_group_cls, "_sort_via_pass_numbering", spy)

    order = [node.node_id for node in serialized.task_group.topological_sort()]
    ...

测试要点：

1. 构造 padded 反向链：chain_length=80, independent_count=80，反向链 r0 -> r1 -> ... -> r79（声明顺序被打乱为从 r79 到 r0）外加 80 个独立节点 i0..i79，再通过重排 children 字典把反向链置于独立节点之前——这是"扫描重复轮次成本反超 pass-numbering"的典型对抗结构；
2. 用 spy 断言分支被触发：monkeypatch 包装 _sort_via_pass_numbering 并断言其确实被调用，证明启发式调度把该结构判给了 pass-number 分支；
3. 校验排序正确性：80 个反向链节点 + 80 个独立节点全部出现，且 position[r_i] < position[r_{i+1}] 对所有 i 成立，即依赖序保持。

同样的构造还出现在 task-sdk 测试 task-sdk/tests/task_sdk/definitions/test_taskgroup.py（_make_padded_reverse_chain）与 test_topological_sort_padded_reverse_chain_uses_pass_numbering（第 1106 行起）中，说明该算法在 task-sdk 与 airflow-core 序列化侧是成对实现、成对验证的。

此外，airflow-core/tests/unit/utils/test_task_group.py 还覆盖了序列化往返后的分组间任务级依赖（b_end >> a_start 跨组依赖仍保证 pos["stage_b"] < pos["stage_a"]），确保优化没有牺牲依赖语义的完整性。

双份实现：airflow-core 与 task-sdk 的对应关系

值得读者注意，本次优化同时在两个包中存在镜像实现：

• task-sdk（任务执行 SDK，供分布式/远程任务执行使用）：task-sdk/src/airflow/sdk/definitions/taskgroup.py，排序基于 TaskGroup 对象；
• airflow-core（调度器等受保护进程）：airflow-core/src/airflow/serialization/definitions/taskgroup.py，排序基于 SerializedTaskGroup 反序列化对象。

两者共享相同的双分支启发式、相同的 _project_child_deps 投影逻辑、相同的 pass-number 递推式，仅异常类型（AirflowDagCycleException vs ValueError）与注释措辞略有差异。这样设计是为了保证无论排序发生在 DAG 构建阶段还是序列化重建阶段，输出顺序完全一致，避免调度器与执行端对任务顺序的认知产生偏差。

实际收益与适用场景

综合源码与测试证据，本次改进的收益可以总结为：

• 消除最坏情况复杂度：对反向插入链等对抗性结构，排序从接近 O(N²) 降至约 O((V + E) log V)；
• 常见形态零回退：正向声明或少量反向边的 DAG 仍走轻量的多轮扫描，不引入额外常数开销；分支阈值的比例判据（反向节点过半）与绝对判据（≥32）共同保证决策稳健；
• 行为完全兼容：两条分支输出相同顺序，下游依赖该顺序的 UI 渲染、序列化重建与任务遍历逻辑无需任何改动；
• 成对验证：task-sdk 与 airflow-core 两侧均有用例证明分支触发与结果正确性，方便后续维护者在此基础上扩展更多对抗性用例。

如果你的 DAG 中存在大量"先定义任务再声明其依赖"的写法（例如用循环按业务倒序拼接任务链、或通过脚本程序化生成反向依赖的任务图），本次优化将直接受益——排序阶段不会再在这些结构上产生平方级耗时。对于绝大多数按正向顺序声明依赖的常规 DAG，算法行为与此前完全一致，无需任何迁移或配置调整。

相关代码与测试索引

• 改进记录：airflow-core/newsfragments/67688.improvement.rst
• airflow-core 实现（双分支调度 + 投影 + 扫描 + pass-numbering）：airflow-core/src/airflow/serialization/definitions/taskgroup.py
• task-sdk 实现：task-sdk/src/airflow/sdk/definitions/taskgroup.py
• airflow-core 单元测试（含 padded 反向链触发分支断言）：airflow-core/tests/unit/utils/test_task_group.py
• task-sdk 单元测试（反向链/padded 反向链构造与验证）：task-sdk/tests/task_sdk/definitions/test_taskgroup.py