diff --git a/shared/dagnode/src/airflow_shared/dagnode/cycle.py b/shared/dagnode/src/airflow_shared/dagnode/cycle.py new file mode 100644 index 0000000000000..b913bb2f5320e --- /dev/null +++ b/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 diff --git a/shared/dagnode/tests/dagnode/test_cycle.py b/shared/dagnode/tests/dagnode/test_cycle.py new file mode 100644 index 0000000000000..fd890b3fce51c --- /dev/null +++ b/shared/dagnode/tests/dagnode/test_cycle.py @@ -0,0 +1,79 @@ +# 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"), + pytest.param({"a": [""], "": ["a"]}, "", id="empty-task-id"), + ], +) +def test_returns_the_node_whose_edge_closes_the_cycle(graph, expected): + assert detect_cycle(graph, edges_of(graph)) == expected + + +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) diff --git a/task-sdk/src/airflow/sdk/definitions/dag.py b/task-sdk/src/airflow/sdk/definitions/dag.py index c3d485eb91b21..4a5764e800e5e 100644 --- a/task-sdk/src/airflow/sdk/definitions/dag.py +++ b/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."""