pydantic-graph 1.2.1__py3-none-any.whl → 1.22.0__py3-none-any.whl

pydantic_graph/beta/node_types.py

@@ -0,0 +1,90 @@
+"""Type definitions for graph node categories.
+
+This module defines type aliases and utilities for categorizing nodes in the
+graph execution system. It provides clear distinctions between source nodes,
+destination nodes, and middle nodes, along with type guards for validation.
+"""
+
+from __future__ import annotations
+
+from typing import Any, TypeGuard
+
+from typing_extensions import TypeAliasType, TypeVar
+
+from pydantic_graph.beta.decision import Decision
+from pydantic_graph.beta.join import Join
+from pydantic_graph.beta.node import EndNode, Fork, StartNode
+from pydantic_graph.beta.step import Step
+
+StateT = TypeVar('StateT', infer_variance=True)
+DepsT = TypeVar('DepsT', infer_variance=True)
+InputT = TypeVar('InputT', infer_variance=True)
+OutputT = TypeVar('OutputT', infer_variance=True)
+
+MiddleNode = TypeAliasType(
+    'MiddleNode',
+    Step[StateT, DepsT, InputT, OutputT] | Join[StateT, DepsT, InputT, OutputT] | Fork[InputT, OutputT],
+    type_params=(StateT, DepsT, InputT, OutputT),
+)
+"""Type alias for nodes that can appear in the middle of a graph execution path.
+
+Middle nodes can both receive input and produce output, making them suitable
+for intermediate processing steps in the graph.
+"""
+SourceNode = TypeAliasType(
+    'SourceNode', MiddleNode[StateT, DepsT, Any, OutputT] | StartNode[OutputT], type_params=(StateT, DepsT, OutputT)
+)
+"""Type alias for nodes that can serve as sources in a graph execution path.
+
+Source nodes produce output data and can be the starting point for data flow
+in the graph. This includes start nodes and middle nodes configured as sources.
+"""
+DestinationNode = TypeAliasType(
+    'DestinationNode',
+    MiddleNode[StateT, DepsT, InputT, Any] | Decision[StateT, DepsT, InputT] | EndNode[InputT],
+    type_params=(StateT, DepsT, InputT),
+)
+"""Type alias for nodes that can serve as destinations in a graph execution path.
+
+Destination nodes consume input data and can be the ending point for data flow
+in the graph. This includes end nodes, decision nodes, and middle nodes configured as destinations.
+"""
+
+AnySourceNode = TypeAliasType('AnySourceNode', SourceNode[Any, Any, Any])
+"""Type alias for source nodes with any type parameters."""
+
+AnyDestinationNode = TypeAliasType('AnyDestinationNode', DestinationNode[Any, Any, Any])
+"""Type alias for destination nodes with any type parameters."""
+
+AnyNode = TypeAliasType('AnyNode', AnySourceNode | AnyDestinationNode)
+"""Type alias for any node in the graph, regardless of its role or type parameters."""
+
+
+def is_source(node: AnyNode) -> TypeGuard[AnySourceNode]:
+    """Check if a node can serve as a source in the graph.
+
+    Source nodes are capable of producing output data and can be the starting
+    point for data flow in graph execution paths.
+
+    Args:
+        node: The node to check
+
+    Returns:
+        True if the node can serve as a source, False otherwise
+    """
+    return isinstance(node, StartNode | Step | Join)
+
+
+def is_destination(node: AnyNode) -> TypeGuard[AnyDestinationNode]:
+    """Check if a node can serve as a destination in the graph.
+
+    Destination nodes are capable of consuming input data and can be the ending
+    point for data flow in graph execution paths.
+
+    Args:
+        node: The node to check
+
+    Returns:
+        True if the node can serve as a destination, False otherwise
+    """
+    return isinstance(node, EndNode | Step | Join | Decision)

pydantic_graph/beta/parent_forks.py

@@ -0,0 +1,232 @@
+"""Parent fork identification and deadlock avoidance in parallel graph execution.
+
+This module provides functionality to identify "parent forks" in a graph, which are dominating
+fork nodes that control access to join nodes. A parent fork is a fork node that:
+
+1. Dominates a join node (all paths to the join must pass through the fork)
+2. Does not participate in cycles that bypass it to reach the join
+
+Identifying parent forks is crucial for deadlock avoidance in parallel execution. When a join
+node waits for all its incoming branches, knowing the parent fork helps determine when it's
+safe to proceed without risking deadlock.
+
+In most typical graphs, such dominating forks exist naturally. However, when there are multiple
+subsequent forks, the choice of parent fork can be ambiguous and may need to be specified by
+the graph designer.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Hashable
+from dataclasses import dataclass
+from functools import cached_property
+from typing import Generic
+
+from typing_extensions import TypeVar
+
+from pydantic_graph.exceptions import GraphBuildingError
+
+T = TypeVar('T', bound=Hashable, infer_variance=True, default=str)
+
+
+@dataclass
+class ParentFork(Generic[T]):
+    """Represents a parent fork node and its relationship to a join node.
+
+    A parent fork is a dominating fork that controls the execution flow to a join node.
+    It tracks which nodes lie between the fork and the join, which is essential for
+    determining when it's safe to proceed past the join point.
+    """
+
+    fork_id: T
+    """The identifier of the fork node that serves as the parent."""
+
+    intermediate_nodes: set[T]
+    """The set of node IDs of nodes upstream of the join and downstream of the parent fork.
+
+    If there are no graph walkers in these nodes that were a part of a previous fork, it is safe to proceed downstream
+    of the join.
+    """
+
+
+@dataclass
+class ParentForkFinder(Generic[T]):
+    """Analyzes graph structure to identify parent forks for join nodes.
+
+    This class implements algorithms to find dominating forks in a directed graph,
+    which is essential for coordinating parallel execution and avoiding deadlocks.
+    """
+
+    nodes: set[T]
+    """All node identifiers in the graph."""
+
+    start_ids: set[T]
+    """Node identifiers that serve as entry points to the graph."""
+
+    fork_ids: set[T]
+    """Node identifiers that represent fork nodes (nodes that create parallel branches)."""
+
+    edges: dict[T, list[T]]  # source_id to list of destination_ids
+    """Graph edges represented as adjacency list mapping source nodes to destinations."""
+
+    def find_parent_fork(
+        self, join_id: T, *, parent_fork_id: T | None = None, prefer_closest: bool = False
+    ) -> ParentFork[T] | None:
+        """Find the parent fork for a given join node.
+
+        Searches for the _most_ ancestral dominating fork that can serve as a parent fork
+        for the specified join node. A valid parent fork must dominate the join without
+        allowing cycles that bypass it.
+
+        Args:
+            join_id: The identifier of the join node to analyze.
+            parent_fork_id: Optional manually selected node ID to attempt to use as the parent fork node.
+            prefer_closest: If no explicit fork is specified, this argument is used to determine
+                whether to find the closest or farthest (i.e., most ancestral) dominating fork.
+
+        Returns:
+            A ParentFork object containing the fork ID and intermediate nodes if a valid
+            parent fork exists, or None if no valid parent fork can be found (which would
+            indicate potential deadlock risk).
+
+        Note:
+            If every dominating fork of the join lets it participate in a cycle that avoids
+            the fork, None is returned since no valid "parent fork" exists.
+        """
+        if parent_fork_id is not None:
+            # A fork was manually specified; we still verify it's a valid dominating fork
+            upstream_nodes = self._get_upstream_nodes_if_parent(join_id, parent_fork_id)
+            if upstream_nodes is None:
+                raise GraphBuildingError(
+                    f'There is a cycle in the graph passing through {join_id!r} that does not include {parent_fork_id!r}.'
+                    f' Parent forks of a join must be a part of any cycles involving that join.'
+                )
+            return ParentFork[T](parent_fork_id, upstream_nodes)
+
+        visited: set[str] = set()
+        cur = join_id  # start at J and walk up the immediate dominator chain
+
+        parent_fork: ParentFork[T] | None = None
+        while True:
+            cur = self._immediate_dominator(cur)
+            if cur is None:  # reached the root
+                break
+
+            # The visited-tracking shouldn't be necessary, but I included it to prevent infinite loops if there are bugs
+            assert cur not in visited, f'Cycle detected in dominator tree: {join_id} → {cur} → {visited}'
+            visited.add(cur)
+
+            if cur not in self.fork_ids:
+                continue  # not a fork, so keep climbing
+
+            upstream_nodes = self._get_upstream_nodes_if_parent(join_id, cur)
+            if upstream_nodes is not None:  # found upstream nodes without a cycle
+                parent_fork = ParentFork[T](cur, upstream_nodes)
+                if prefer_closest:
+                    return parent_fork
+            elif parent_fork is not None:
+                # We reached a fork that is an ancestor of a parent fork but is not itself a parent fork.
+                # This means there is a cycle to J that is downstream of `cur`, and so any node further upstream
+                # will fail to be a parent fork for the same reason. So we can stop here and just return `parent_fork`.
+                return parent_fork
+
+        # No dominating fork passed the cycle test to be a "parent" fork
+        return parent_fork
+
+    @cached_property
+    def _predecessors(self) -> dict[T, list[T]]:
+        """Compute and cache the predecessor mapping for all nodes.
+
+        Returns:
+            A dictionary mapping each node to a list of its immediate predecessors.
+        """
+        predecessors: dict[T, list[T]] = {n: [] for n in self.nodes}
+        for source_id in self.nodes:
+            for destination_id in self.edges.get(source_id, []):
+                predecessors[destination_id].append(source_id)
+        return predecessors
+
+    @cached_property
+    def _dominators(self) -> dict[T, set[T]]:
+        """Compute the dominator sets for all nodes using iterative dataflow analysis.
+
+        A node D dominates node N if every path from a start node to N must pass through D.
+        This is computed using a fixed-point iteration algorithm.
+
+        Returns:
+            A dictionary mapping each node to its set of dominators.
+        """
+        node_ids = set(self.nodes)
+        start_ids = self.start_ids
+
+        dom: dict[T, set[T]] = {n: set(node_ids) for n in node_ids}
+        for s in start_ids:
+            dom[s] = {s}
+
+        changed = True
+        while changed:
+            changed = False
+            for n in node_ids - start_ids:
+                preds = self._predecessors[n]
+                if not preds:  # unreachable from any start
+                    continue
+                intersection = set[T].intersection(*(dom[p] for p in preds)) if preds else set[T]()
+                new_dom = {n} | intersection
+                if new_dom != dom[n]:
+                    dom[n] = new_dom
+                    changed = True
+        return dom
+
+    def _immediate_dominator(self, node_id: T) -> T | None:
+        """Find the immediate dominator of a node.
+
+        The immediate dominator is the closest dominator to a node (other than itself)
+        in the dominator tree.
+
+        Args:
+            node_id: The node to find the immediate dominator for.
+
+        Returns:
+            The immediate dominator's ID if one exists, None otherwise.
+        """
+        dom = self._dominators
+        candidates = dom[node_id] - {node_id}
+        for c in candidates:
+            if all((c == d) or (c not in dom[d]) for d in candidates):
+                return c
+        return None
+
+    def _get_upstream_nodes_if_parent(self, join_id: T, fork_id: T) -> set[T] | None:
+        """Check if a fork is a valid parent and return upstream nodes.
+
+        Tests whether the given fork can serve as a parent fork for the join by checking
+        for cycles that bypass the fork. If valid, returns all nodes that can reach the
+        join without going through the fork.
+
+        Args:
+            join_id: The join node being analyzed.
+            fork_id: The potential parent fork to test.
+
+        Returns:
+            The set of node IDs upstream of the join (excluding the fork) if the fork is
+            a valid parent, or None if a cycle exists that bypasses the fork (making it
+            invalid as a parent fork).
+
+        Note:
+            If, in the graph with fork_id removed, a path exists that starts and ends at
+            the join (i.e., join is on a cycle avoiding the fork), we return None because
+            the fork would not be a valid "parent fork".
+        """
+        upstream: set[T] = set()
+        stack = [join_id]
+        while stack:
+            v = stack.pop()
+            for p in self._predecessors[v]:
+                if p == fork_id:
+                    continue
+                if p == join_id:
+                    return None  # J sits on a cycle w/out the specified node
+                if p not in upstream:
+                    upstream.add(p)
+                    stack.append(p)
+        return upstream