Share Dag cycle detection between the Task SDK and core

Cycle detection only runs for Dags that go through the Python SDK's
DAG object. Dags authored with non-Python language SDKs arrive at core
already serialized and never pass through that object, so nothing
rejects a cycle in them today.

Core cannot reuse the SDK implementation as-is: it is a method on the
authoring class and reaches for task objects, while core holds raw
serialized data it should not have to hydrate just to run a check.
Moving the traversal into the shared distribution both distributions
already depend on, and describing the graph through callbacks instead
of node objects, lets core add that check against the serialized form
in a follow-up without duplicating the algorithm or inverting the
dependency by importing the SDK's exception.

Author: jason810496

shared/dagnode/src/airflow_shared/dagnode/cycle.py

@@ -0,0 +1,73 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from collections import defaultdict, deque
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterable
+
+# default of int is 0 which corresponds to _NEW
+_NEW = 0
+_IN_PROGRESS = 1
+_DONE = 2
+
+__all__ = ["detect_cycle"]
+
+
+def detect_cycle(node_ids: Iterable[str], downstream_of: Callable[[str], Iterable[str]]) -> str | None:
+    """
+    Search the graph for a cycle, following downstream edges only.
+
+    The graph is described by callbacks rather than by node objects so that callers
+    holding a serialized Dag can answer from raw data without hydrating it.
+
+    :param node_ids: Every node in the graph, so that disconnected components are covered too.
+    :param downstream_of: Returns the ids directly downstream of the given node id.
+    :return: The id of the node whose downstream edge closes a cycle, or None when acyclic.
+    """
+    visited: dict[str, int] = defaultdict(int)
+    path_stack: deque[str] = deque()
+
+    for node_id in node_ids:
+        if visited[node_id] == _DONE:
+            continue
+        path_stack.append(node_id)
+        # Iterative rather than recursive: Python's stack depth limit is easily
+        # reached by a graph with long chains.
+        while path_stack:
+            current_id = path_stack[-1]
+            if visited[current_id] == _NEW:
+                visited[current_id] = _IN_PROGRESS
+
+            child_to_check = None
+            for downstream_id in downstream_of(current_id):
+                if visited[downstream_id] == _IN_PROGRESS:
+                    return current_id
+                if visited[downstream_id] == _NEW:
+                    child_to_check = downstream_id
+                    break
+
+            if child_to_check is None:
+                visited[current_id] = _DONE
+                path_stack.pop()
+            else:
+                path_stack.append(child_to_check)
+
+    return None

shared/dagnode/tests/dagnode/test_cycle.py

@@ -0,0 +1,87 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import sys
+
+import pytest
+
+from airflow_shared.dagnode.cycle import detect_cycle
+
+
+def edges_of(graph: dict[str, list[str]]):
+    return lambda node_id: graph[node_id]
+
+
+@pytest.mark.parametrize(
+    "graph",
+    [
+        pytest.param({}, id="empty"),
+        pytest.param({"a": []}, id="single-node"),
+        pytest.param({"a": ["b"], "b": ["c"], "c": []}, id="chain"),
+        pytest.param({"a": ["b", "c"], "b": ["d"], "c": ["d"], "d": []}, id="diamond"),
+        pytest.param({"a": ["b"], "b": [], "c": ["d"], "d": []}, id="disconnected-components"),
+    ],
+)
+def test_returns_none_when_acyclic(graph):
+    assert detect_cycle(graph, edges_of(graph)) is None
+
+
+@pytest.mark.parametrize(
+    ("graph", "expected"),
+    [
+        pytest.param({"a": ["a"]}, "a", id="self-loop"),
+        pytest.param({"a": ["b"], "b": ["a"]}, "b", id="two-node-cycle"),
+        pytest.param({"a": ["b"], "b": ["c"], "c": ["a"]}, "c", id="three-node-cycle"),
+        pytest.param({"a": ["b"], "b": ["c"], "c": ["b"]}, "c", id="cycle-below-an-acyclic-root"),
+        pytest.param({"a": [], "b": ["c"], "c": ["b"]}, "c", id="cycle-outside-the-first-component"),
+    ],
+)
+def test_returns_the_node_whose_edge_closes_the_cycle(graph, expected):
+    assert detect_cycle(graph, edges_of(graph)) == expected
+
+
+def test_follows_downstream_edges_only():
+    """An edge is only traversed in the direction the callback reports."""
+    downstream = {"a": ["b"], "b": []}
+    upstream = {"a": [], "b": ["a"]}
+
+    assert detect_cycle(downstream, lambda node_id: downstream[node_id]) is None
+    assert detect_cycle(upstream, lambda node_id: upstream[node_id]) is None
+
+
+def test_reads_the_graph_only_through_the_callback():
+    """Callers may answer from any shape, such as raw serialized Dag data."""
+    serialized = {
+        "tasks": [
+            {"task_id": "a", "downstream_task_ids": ["b"]},
+            {"task_id": "b", "downstream_task_ids": ["a"]},
+        ]
+    }
+    by_id = {task["task_id"]: task for task in serialized["tasks"]}
+
+    assert detect_cycle(by_id, lambda node_id: by_id[node_id]["downstream_task_ids"]) == "b"
+
+
+@pytest.mark.parametrize("cyclic", [False, True])
+def test_handles_a_graph_deeper_than_the_recursion_limit(cyclic):
+    depth = sys.getrecursionlimit() * 2
+    graph = {str(i): [str(i + 1)] for i in range(depth)}
+    graph[str(depth)] = ["0"] if cyclic else []
+
+    assert detect_cycle(graph, edges_of(graph)) == (str(depth) if cyclic else None)

task-sdk/src/airflow/sdk/definitions/dag.py

@@ -27,7 +27,7 @@
 import sys
 import warnings
 import weakref
-from collections import abc, defaultdict, deque
+from collections import abc
 from collections.abc import Callable, Collection, Iterable, MutableSet
 from datetime import datetime, timedelta
 from inspect import signature
@@ -41,6 +41,7 @@
 
 from airflow import settings
 from airflow.sdk import TaskInstanceState, TriggerRule, XComArg
+from airflow.sdk._shared.dagnode.cycle import detect_cycle
 from airflow.sdk.bases.operator import BaseOperator
 from airflow.sdk.bases.timetable import BaseTimetable
 from airflow.sdk.definitions._internal.node import DAGNode, validate_key
@@ -1142,40 +1143,12 @@ def check_cycle(self) -> None:
 
         :raises AirflowDagCycleException: If cycle is found in the Dag.
         """
-        # default of int is 0 which corresponds to CYCLE_NEW
-        CYCLE_NEW = 0
-        CYCLE_IN_PROGRESS = 1
-        CYCLE_DONE = 2
-
-        visited: dict[str, int] = defaultdict(int)
-        path_stack: deque[str] = deque()
         task_dict = self.task_dict
-
-        def _check_adjacent_tasks(task_id, current_task):
-            """Return first untraversed child task, else None if all tasks traversed."""
-            for adjacent_task in current_task.get_direct_relative_ids():
-                if visited[adjacent_task] == CYCLE_IN_PROGRESS:
-                    msg = f"Cycle detected in Dag: {self.dag_id}. Faulty task: {task_id}"
-                    raise AirflowDagCycleException(msg)
-                if visited[adjacent_task] == CYCLE_NEW:
-                    return adjacent_task
-            return None
-
-        for dag_task_id in self.task_dict.keys():
-            if visited[dag_task_id] == CYCLE_DONE:
-                continue
-            path_stack.append(dag_task_id)
-            while path_stack:
-                current_task_id = path_stack[-1]
-                if visited[current_task_id] == CYCLE_NEW:
-                    visited[current_task_id] = CYCLE_IN_PROGRESS
-                task = task_dict[current_task_id]
-                child_to_check = _check_adjacent_tasks(current_task_id, task)
-                if not child_to_check:
-                    visited[current_task_id] = CYCLE_DONE
-                    path_stack.pop()
-                else:
-                    path_stack.append(child_to_check)
+        faulty_task_id = detect_cycle(task_dict, lambda tid: task_dict[tid].get_direct_relative_ids())
+        if faulty_task_id is not None:
+            raise AirflowDagCycleException(
+                f"Cycle detected in Dag: {self.dag_id}. Faulty task: {faulty_task_id}"
+            )
 
     def cli(self):
         """Exposes a CLI specific to this Dag."""