| # 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 deque |
| from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar |
| |
| import structlog |
| |
| if TYPE_CHECKING: |
| from collections.abc import Collection, Iterable, Iterator |
| from typing import Self |
| |
| from ..logging.types import Logger |
| |
| |
| class DagProtocol(Protocol): |
| """Protocol defining the minimum interface required for Dag generic type.""" |
| |
| dag_id: str |
| task_dict: dict[str, Any] |
| |
| def get_task(self, tid: str) -> Any: |
| """Retrieve a task by its task ID.""" |
| ... |
| |
| |
| class TaskProtocol(Protocol): |
| """Protocol defining the minimum interface required for Task generic type.""" |
| |
| task_id: str |
| is_setup: bool |
| is_teardown: bool |
| downstream_list: Iterable[Self] |
| downstream_task_ids: set[str] |
| |
| |
| class TaskGroupProtocol(Protocol): |
| """Protocol defining the minimum interface required for TaskGroup generic type.""" |
| |
| node_id: str |
| prefix_group_id: bool |
| |
| |
| Dag = TypeVar("Dag", bound=DagProtocol) |
| Task = TypeVar("Task", bound=TaskProtocol) |
| TaskGroup = TypeVar("TaskGroup", bound=TaskGroupProtocol) |
| Node = TypeVar("Node") |
| |
| |
| class GenericDAGNode(Generic[Dag, Task, TaskGroup]): |
| """ |
| Generic class for a node in the graph of a workflow. |
| |
| A node may be an operator or task group, either mapped or unmapped. |
| """ |
| |
| dag: Dag | None |
| task_group: TaskGroup | None |
| downstream_group_ids: set[str | None] |
| upstream_task_ids: set[str] |
| downstream_task_ids: set[str] |
| |
| _log_config_logger_name: str | None = None |
| _logger_name: str | None = None |
| _cached_logger: Logger | None = None |
| |
| def __init__(self): |
| super().__init__() |
| self.upstream_task_ids = set() |
| self.downstream_task_ids = set() |
| |
| @property |
| def log(self) -> Logger: |
| if self._cached_logger is not None: |
| return self._cached_logger |
| |
| typ = type(self) |
| |
| logger_name: str = ( |
| self._logger_name if self._logger_name is not None else f"{typ.__module__}.{typ.__qualname__}" |
| ) |
| |
| if self._log_config_logger_name: |
| logger_name = ( |
| f"{self._log_config_logger_name}.{logger_name}" |
| if logger_name |
| else self._log_config_logger_name |
| ) |
| |
| self._cached_logger = structlog.get_logger(logger_name) |
| return self._cached_logger |
| |
| @property |
| def dag_id(self) -> str: |
| if self.dag: |
| return self.dag.dag_id |
| return "_in_memory_dag_" |
| |
| @property |
| def node_id(self) -> str: |
| raise NotImplementedError() |
| |
| @property |
| def label(self) -> str | None: |
| tg = self.task_group |
| if tg and tg.node_id and tg.prefix_group_id: |
| # "task_group_id.task_id" -> "task_id" |
| return self.node_id[len(tg.node_id) + 1 :] |
| return self.node_id |
| |
| @property |
| def upstream_list(self) -> Iterable[Task]: |
| if not self.dag: |
| raise RuntimeError(f"Operator {self} has not been assigned to a Dag yet") |
| return [self.dag.get_task(tid) for tid in self.upstream_task_ids] |
| |
| @property |
| def downstream_list(self) -> Iterable[Task]: |
| if not self.dag: |
| raise RuntimeError(f"Operator {self} has not been assigned to a Dag yet") |
| return [self.dag.get_task(tid) for tid in self.downstream_task_ids] |
| |
| @property |
| def _topological_upstream_ids(self) -> Collection[str]: |
| """ |
| Node ids this node must be ordered after within its parent group. |
| |
| A plain task depends only on its direct upstream tasks. Task groups override |
| this to also cover group-to-group and cross-group edges. |
| """ |
| return self.upstream_task_ids |
| |
| def has_dag(self) -> bool: |
| return self.dag is not None |
| |
| def get_dag(self) -> Dag | None: |
| return self.dag |
| |
| def get_direct_relative_ids(self, upstream: bool = False) -> set[str]: |
| """Get set of the direct relative ids to the current task, upstream or downstream.""" |
| if upstream: |
| return self.upstream_task_ids |
| return self.downstream_task_ids |
| |
| def get_direct_relatives(self, upstream: bool = False) -> Iterable[Task]: |
| """Get list of the direct relatives to the current task, upstream or downstream.""" |
| if upstream: |
| return self.upstream_list |
| return self.downstream_list |
| |
| def get_flat_relative_ids(self, *, upstream: bool = False, depth: int | None = None) -> set[str]: |
| """ |
| Get a flat set of relative IDs, upstream or downstream. |
| |
| Will recurse each relative found in the direction specified. |
| |
| :param upstream: Whether to look for upstream or downstream relatives. |
| :param depth: Maximum number of levels to traverse. If None, traverses all levels. |
| Must be non-negative. |
| """ |
| if depth is not None and depth < 0: |
| raise ValueError(f"depth must be non-negative, got {depth}") |
| |
| dag = self.get_dag() |
| if not dag: |
| return set() |
| |
| relatives: set[str] = set() |
| |
| # This is intentionally implemented as a loop, instead of calling |
| # get_direct_relative_ids() recursively, since Python has significant |
| # limitation on stack level, and a recursive implementation can blow up |
| # if a DAG contains very long routes. |
| task_ids_to_trace = self.get_direct_relative_ids(upstream) |
| levels_remaining = depth |
| while task_ids_to_trace: |
| # if depth is set we have bounded traversal and should break when |
| # there are no more levels remaining |
| if levels_remaining is not None and levels_remaining <= 0: |
| break |
| task_ids_to_trace_next: set[str] = set() |
| for task_id in task_ids_to_trace: |
| if task_id in relatives: |
| continue |
| task_ids_to_trace_next.update(dag.task_dict[task_id].get_direct_relative_ids(upstream)) |
| relatives.add(task_id) |
| task_ids_to_trace = task_ids_to_trace_next |
| if levels_remaining is not None: |
| levels_remaining -= 1 |
| |
| return relatives |
| |
| def get_flat_relatives(self, upstream: bool = False, depth: int | None = None) -> Collection[Task]: |
| """ |
| Get a flat list of relatives, either upstream or downstream. |
| |
| :param upstream: Whether to look for upstream or downstream relatives. |
| :param depth: Maximum number of levels to traverse. If None, traverses all levels. |
| Must be non-negative. |
| """ |
| dag = self.get_dag() |
| if not dag: |
| return set() |
| return [ |
| dag.task_dict[task_id] for task_id in self.get_flat_relative_ids(upstream=upstream, depth=depth) |
| ] |
| |
| def get_upstreams_follow_setups(self, depth: int | None = None) -> Iterable[Task]: |
| """ |
| All upstreams and, for each upstream setup, its respective teardowns. |
| |
| :param depth: Maximum number of levels to traverse. If None, traverses all levels. |
| Must be non-negative. |
| """ |
| for task in self.get_flat_relatives(upstream=True, depth=depth): |
| yield task |
| if task.is_setup: |
| for t in task.downstream_list: |
| if t.is_teardown and t != self: |
| yield t |
| |
| def get_upstreams_only_setups_and_teardowns(self) -> Iterable[Task]: |
| """ |
| Only *relevant* upstream setups and their teardowns. |
| |
| This method is meant to be used when we are clearing the task (non-upstream) and we need |
| to add in the *relevant* setups and their teardowns. |
| |
| Relevant in this case means, the setup has a teardown that is downstream of ``self``, |
| or the setup has no teardowns. |
| """ |
| downstream_teardown_ids = { |
| x.task_id for x in self.get_flat_relatives(upstream=False) if x.is_teardown |
| } |
| for task in self.get_flat_relatives(upstream=True): |
| if not task.is_setup: |
| continue |
| has_no_teardowns = not any(x.is_teardown for x in task.downstream_list) |
| # if task has no teardowns or has teardowns downstream of self |
| if has_no_teardowns or task.downstream_task_ids.intersection(downstream_teardown_ids): |
| yield task |
| for t in task.downstream_list: |
| if t.is_teardown and t != self: |
| yield t |
| |
| def get_upstreams_only_setups(self) -> Iterable[Task]: |
| """ |
| Return relevant upstream setups. |
| |
| This method is meant to be used when we are checking task dependencies where we need |
| to wait for all the upstream setups to complete before we can run the task. |
| """ |
| for task in self.get_upstreams_only_setups_and_teardowns(): |
| if task.is_setup: |
| yield task |
| |
| |
| class TaskGroupMixin: |
| """Mixin to host common logic between authored and serialized task group classes.""" |
| |
| upstream_task_ids: set[str] |
| upstream_group_ids: set[str | None] |
| |
| def get_roots(self) -> Iterable[GenericDAGNode]: |
| raise NotImplementedError() |
| |
| @property |
| def _topological_upstream_ids(self) -> Collection[str]: |
| """ |
| Node ids this node must be ordered after within its parent group. |
| |
| A group's upstream_task_ids only reflects direct task-to-group edges |
| (e.g. ``task >> this_group``). This explicitly pulls in two more cases: |
| |
| * Group-to-group edges (e.g. ``another_group >> this_group``). |
| * Task-level edges crossing into the group. |
| """ |
| return self.upstream_task_ids.union( |
| (gid for gid in self.upstream_group_ids if gid is not None), |
| (t for root in self.get_roots() for t in root.upstream_task_ids), |
| ) |
| |
| @staticmethod |
| def _compute_pass_order(projected: list[tuple[int, ...]]) -> list[int]: |
| """ |
| Order child indices by (pass number, insertion index). |
| |
| pass(X) is the earliest pass at which a greedy sweep in insertion order would emit X: |
| ``max over deps d of (pass(d) if idx(d) < idx(X) else pass(d) + 1)``. A dep declared |
| before X can be emitted in the same pass; a dep declared after X forces X into the next |
| pass. Computed via Kahn's traversal in O(V + E). |
| |
| :param projected: for each child, the indices of the siblings it depends on |
| :return: the ordered indices, leaving out every child on or downstream of a cycle |
| """ |
| n = len(projected) |
| in_degree = [len(deps) for deps in projected] |
| successors: list[list[int]] = [[] for _ in range(n)] |
| for i, deps in enumerate(projected): |
| for d in deps: |
| successors[d].append(i) |
| |
| pass_of = [0] * n |
| queue: deque[int] = deque(i for i in range(n) if in_degree[i] == 0) |
| processed: list[int] = [] |
| while queue: |
| i = queue.popleft() |
| my_pass = 1 |
| for d in projected[i]: |
| d_pass = pass_of[d] |
| if d < i: |
| if d_pass > my_pass: |
| my_pass = d_pass |
| elif d_pass + 1 > my_pass: |
| my_pass = d_pass + 1 |
| pass_of[i] = my_pass |
| processed.append(i) |
| for s in successors[i]: |
| in_degree[s] -= 1 |
| if in_degree[s] == 0: |
| queue.append(s) |
| |
| return sorted(processed, key=lambda i: (pass_of[i], i)) |
| |
| def _sort_cyclic_projection(self, nodes: list[Node], projected: list[tuple[int, ...]]) -> list[Node]: |
| """ |
| Order children whose projection has a cycle, ordering the siblings on each cycle as one unit. |
| |
| ``DAG.check_cycle`` rejects task-level cycles, but a task with no upstream inside its own |
| group counts as a root of that group, so edges routed through tasks outside the group can |
| still make siblings depend on each other. ``partial_subset`` can also create such a cycle |
| by dropping a task's in-group upstream. Grid and Graph must still render these Dags. |
| |
| Each unit is placed by its first child and follows the same pass ordering as the acyclic |
| path; a unit's children keep insertion order. |
| """ |
| component_of = self._find_projection_components(projected) |
| count = max(component_of) + 1 |
| members: list[list[int]] = [[] for _ in range(count)] |
| component_deps: list[set[int]] = [set() for _ in range(count)] |
| for i, deps in enumerate(projected): |
| c = component_of[i] |
| members[c].append(i) |
| component_deps[c].update(component_of[d] for d in deps if component_of[d] != c) |
| component_order = self._compute_pass_order([tuple(deps) for deps in component_deps]) |
| return [nodes[i] for c in component_order for i in members[c]] |
| |
| @staticmethod |
| def _find_projection_components(projected: list[tuple[int, ...]]) -> list[int]: |
| """Return each child's strongly connected component, numbered in order of its first child.""" |
| n = len(projected) |
| successors: list[list[int]] = [[] for _ in range(n)] |
| for i, deps in enumerate(projected): |
| for d in deps: |
| successors[d].append(i) |
| |
| # Kosaraju: finish order along successors, then collect components along dependencies. |
| visited = bytearray(n) |
| finish_order: list[int] = [] |
| for start in range(n): |
| if visited[start]: |
| continue |
| visited[start] = 1 |
| stack: list[tuple[int, Iterator[int]]] = [(start, iter(successors[start]))] |
| while stack: |
| node, remaining = stack[-1] |
| for s in remaining: |
| if not visited[s]: |
| visited[s] = 1 |
| stack.append((s, iter(successors[s]))) |
| break |
| else: |
| stack.pop() |
| finish_order.append(node) |
| |
| root_of = [-1] * n |
| for start in reversed(finish_order): |
| if root_of[start] != -1: |
| continue |
| root_of[start] = start |
| to_visit = [start] |
| while to_visit: |
| node = to_visit.pop() |
| for d in projected[node]: |
| if root_of[d] == -1: |
| root_of[d] = start |
| to_visit.append(d) |
| |
| numbering: dict[int, int] = {} |
| return [numbering.setdefault(root, len(numbering)) for root in root_of] |
| |
| # TODO: Move more duplicated logic between Core and SDK task group types. |
