UNDF: UNDF-2026-000000347

airflow-0001 — O(N²) Topological Sort in TaskGroup

Severity: HIGH Complexity: O(N²) → O(N) CWE: CWE-407 (Algorithmic Complexity)

Affected Files

File	Lines	Notes
`task-sdk/src/airflow/sdk/definitions/taskgroup.py`	536–568	Runtime task group sort
`airflow-core/src/airflow/serialization/definitions/taskgroup.py`	230–248	Serialized task group sort

Defective Code

`task-sdk/src/airflow/sdk/definitions/taskgroup.py` lines 536–568

while graph_unsorted:
    acyclic = False
    for node in list(graph_unsorted.values()):    # O(N) scan every round
        for edge in node.upstream_list:
            if edge.node_id in graph_unsorted:
                break
            tg = edge.task_group
            while tg:
                if tg.node_id in graph_unsorted:
                    break
                tg = tg.parent_group
            if tg:
                break
        else:
            acyclic = True
            del graph_unsorted[node.node_id]
            graph_sorted.append(node)

    if not acyclic:
        raise AirflowDagCycleException(...)

Why it is O(N²): On each pass of the outer while graph_unsorted: loop the inner for node in list(graph_unsorted.values()) rescans every remaining unsorted node from scratch. In the worst case (a linear chain of N tasks) only one node is resolved per round, so the total work is N + (N−1) + … + 1 = O(N²/2).

Same pattern is copy-pasted verbatim into the serialization layer (airflow-core/src/airflow/serialization/definitions/taskgroup.py:230–248).

Fixed Code

Standard Kahn's algorithm: pre-compute in-degrees, maintain a ready queue, process each node exactly once — O(N + E).

def topological_sort(self):
    """
    Sorts children in topographical order, such that a task comes after any of
    its upstream dependencies.

    Uses Kahn's algorithm with an explicit ready-queue: O(N + E).
    """
    from collections import deque

    graph_unsorted = dict(self.children)
    graph_sorted: list[DAGNode] = []

    if not graph_unsorted:
        return graph_sorted

    # Build in-degree counts and adjacency (upstream → dependents).
    in_degree: dict[str, int] = {node_id: 0 for node_id in graph_unsorted}
    dependents: dict[str, list[str]] = {node_id: [] for node_id in graph_unsorted}

    for node_id, node in graph_unsorted.items():
        for edge in node.upstream_list:
            upstream_id = edge.node_id
            if upstream_id in graph_unsorted:
                in_degree[node_id] += 1
                dependents[upstream_id].append(node_id)
            else:
                # Check if task's group is a child of this TG
                tg = edge.task_group
                while tg:
                    if tg.node_id in graph_unsorted:
                        in_degree[node_id] += 1
                        dependents[tg.node_id].append(node_id)
                        break
                    tg = tg.parent_group

    ready: deque[str] = deque(
        node_id for node_id, deg in in_degree.items() if deg == 0
    )

    while ready:
        node_id = ready.popleft()
        graph_sorted.append(graph_unsorted[node_id])
        for dep_id in dependents[node_id]:
            in_degree[dep_id] -= 1
            if in_degree[dep_id] == 0:
                ready.append(dep_id)

    if len(graph_sorted) != len(graph_unsorted):
        raise AirflowDagCycleException(
            f"A cyclic dependency occurred in dag: {self.dag_id}"
        )

    return graph_sorted

Apply the same fix to the copy in airflow-core/src/airflow/serialization/definitions/taskgroup.py.

Complexity Analysis

Version	Time	Space
Defective	O(N²)	O(N)
Fixed	O(N + E)	O(N + E)

For a DAG with N=1000 tasks and E=2000 edges, the defective version performs ~500,000 operations per sort; the fixed version performs ~3,000.

Trigger Path

topological_sort() is called:

airflow-core/.../routes/ui/structure.py:93 on every UI structure render
airflow-core/.../services/ui/task_group.py:36 via task_group_to_dict
Indirectly by DAG.topological_sort() which delegates to this method

Production Airflow deployments with large DAGs (hundreds of tasks) will see O(N²) CPU consumption on every API call that renders DAG structure.

4.3 KiB Raw Permalink Blame History Unescape Escape