Fix graph execution bug with multiple joins downstream of same fork (pydantic#3337)

Author: dmontagu

pydantic_graph/pydantic_graph/beta/graph.py

@@ -148,6 +148,12 @@ class Graph(Generic[StateT, DepsT, InputT, OutputT]):
     parent_forks: dict[JoinID, ParentFork[NodeID]]
     """Parent fork information for each join node."""
 
+    intermediate_join_nodes: dict[JoinID, set[JoinID]]
+    """For each join, the set of other joins that appear between it and its parent fork.
+
+    Used to determine which joins are "final" (have no other joins as intermediates) and
+    which joins should preserve fork stacks when proceeding downstream."""
+
     def get_parent_fork(self, join_id: JoinID) -> ParentFork[NodeID]:
         """Get the parent fork information for a join node.
 
@@ -165,6 +171,24 @@ def get_parent_fork(self, join_id: JoinID) -> ParentFork[NodeID]:
             raise RuntimeError(f'Node {join_id} is not a join node or did not have a dominating fork (this is a bug)')
         return result
 
+    def is_final_join(self, join_id: JoinID) -> bool:
+        """Check if a join is 'final' (has no downstream joins with the same parent fork).
+
+        A join is non-final if it appears as an intermediate node for another join
+        with the same parent fork.
+
+        Args:
+            join_id: The ID of the join node
+
+        Returns:
+            True if the join is final, False if it's non-final
+        """
+        # Check if this join appears in any other join's intermediate_join_nodes
+        for intermediate_joins in self.intermediate_join_nodes.values():
+            if join_id in intermediate_joins:
+                return False
+        return True
+
     async def run(
         self,
         *,
@@ -517,7 +541,14 @@ async def iter_graph(  # noqa C901
                                     parent_fork_id = self.graph.get_parent_fork(result.join_id).fork_id
                                     for i, x in enumerate(result.fork_stack[::-1]):
                                         if x.fork_id == parent_fork_id:
-                                            downstream_fork_stack = result.fork_stack[: len(result.fork_stack) - i]
+                                            # For non-final joins (those that are intermediate nodes of other joins),
+                                            # preserve the fork stack so downstream joins can still associate with the same fork run
+                                            if self.graph.is_final_join(result.join_id):
+                                                # Final join: remove the parent fork from the stack
+                                                downstream_fork_stack = result.fork_stack[: len(result.fork_stack) - i]
+                                            else:
+                                                # Non-final join: preserve the fork stack
+                                                downstream_fork_stack = result.fork_stack
                                             fork_run_id = x.node_run_id
                                             break
                                     else:  # pragma: no cover
@@ -535,13 +566,9 @@ async def iter_graph(  # noqa C901
                                     join_state.current = join_node.reduce(context, join_state.current, result.inputs)
                                     if join_state.cancelled_sibling_tasks:
                                         await self._cancel_sibling_tasks(parent_fork_id, fork_run_id)
-                                    if task_result.source_is_finished:  # pragma: no branch
-                                        await self._finish_task(task_result.source.task_id)
                                 else:
                                     for new_task in maybe_overridden_result:
                                         self.active_tasks[new_task.task_id] = new_task
-                                    if task_result.source_is_finished:
-                                        await self._finish_task(task_result.source.task_id)
 
                                 tasks_by_id_values = list(self.active_tasks.values())
                                 join_tasks: list[GraphTask] = []
@@ -566,28 +593,61 @@ async def iter_graph(  # noqa C901
                                         new_task_ids = {t.task_id for t in maybe_overridden_result}
                                         for t in task_result.result:
                                             if t.task_id not in new_task_ids:
-                                                await self._finish_task(t.task_id)
+                                                await self._finish_task(t.task_id, t.node_id)
                                     self._handle_execution_request(maybe_overridden_result)
 
+                                if task_result.source_is_finished:
+                                    await self._finish_task(task_result.source.task_id, task_result.source.node_id)
+
                                 if not self.active_tasks:
                                     # if there are no active tasks, we'll be waiting forever for the next result..
                                     break
 
                             if self.active_reducers:  # pragma: no branch
-                                # In this case, there are no pending tasks. We can therefore finalize all active reducers whose
-                                # downstream fork stacks are not a strict "prefix" of another active reducer. (If it was, finalizing the
-                                # deeper reducer could produce new tasks in the "prefix" reducer.)
-                                active_fork_stacks = [
-                                    join_state.downstream_fork_stack for join_state in self.active_reducers.values()
-                                ]
+                                # In this case, there are no pending tasks. We can therefore finalize all active reducers
+                                # that don't have intermediate joins which are also active reducers. If a join J2 has an
+                                # intermediate join J1 that shares the same parent fork run, we must finalize J1 first
+                                # because it might produce items that feed into J2.
                                 for (join_id, fork_run_id), join_state in list(self.active_reducers.items()):
-                                    fork_stack = join_state.downstream_fork_stack
-                                    if any(
-                                        len(afs) > len(fork_stack) and fork_stack == afs[: len(fork_stack)]
-                                        for afs in active_fork_stacks
-                                    ):
-                                        # this join_state is a strict prefix for one of the other active join_states
-                                        continue  # pragma: no cover  # It's difficult to cover this
+                                    # Check if this join has any intermediate joins that are also active reducers
+                                    should_skip = False
+                                    intermediate_joins = self.graph.intermediate_join_nodes.get(join_id, set())
+
+                                    # Get the parent fork for this join to use for comparison
+                                    join_parent_fork = self.graph.get_parent_fork(join_id)
+
+                                    for intermediate_join_id in intermediate_joins:
+                                        # Check if the intermediate join is also an active reducer with matching fork run
+                                        for (other_join_id, _), other_join_state in self.active_reducers.items():
+                                            if other_join_id == intermediate_join_id:
+                                                # Check if they share the same fork run for this join's parent fork
+                                                # by finding the parent fork's node_run_id in both fork stacks
+                                                join_parent_fork_run_id = None
+                                                other_parent_fork_run_id = None
+
+                                                for fsi in join_state.downstream_fork_stack:  # pragma: no branch
+                                                    if fsi.fork_id == join_parent_fork.fork_id:
+                                                        join_parent_fork_run_id = fsi.node_run_id
+                                                        break
+
+                                                for fsi in other_join_state.downstream_fork_stack:  # pragma: no branch
+                                                    if fsi.fork_id == join_parent_fork.fork_id:
+                                                        other_parent_fork_run_id = fsi.node_run_id
+                                                        break
+
+                                                if (
+                                                    join_parent_fork_run_id
+                                                    and other_parent_fork_run_id
+                                                    and join_parent_fork_run_id == other_parent_fork_run_id
+                                                ):  # pragma: no branch
+                                                    should_skip = True
+                                                    break
+                                        if should_skip:
+                                            break
+
+                                    if should_skip:
+                                        continue
+
                                     self.active_reducers.pop(
                                         (join_id, fork_run_id)
                                     )  # we're handling it now, so we can pop it
@@ -610,7 +670,7 @@ async def iter_graph(  # noqa C901
                                         # Same note as above about how this is theoretically reachable but we should
                                         # just get coverage by unifying the code paths
                                         if t.task_id not in new_task_ids:  # pragma: no cover
-                                            await self._finish_task(t.task_id)
+                                            await self._finish_task(t.task_id, t.node_id)
                                     self._handle_execution_request(maybe_overridden_result)
                 except GeneratorExit:
                     self._task_group.cancel_scope.cancel()
@@ -620,7 +680,8 @@ async def iter_graph(  # noqa C901
             'Graph run completed, but no result was produced. This is either a bug in the graph or a bug in the graph runner.'
         )
 
-    async def _finish_task(self, task_id: TaskID) -> None:
+    async def _finish_task(self, task_id: TaskID, node_id: str) -> None:
+        # node_id is just included for debugging right now
         scope = self.cancel_scopes.pop(task_id, None)
         if scope is not None:
             scope.cancel()
@@ -837,7 +898,7 @@ async def _cancel_sibling_tasks(self, parent_fork_id: ForkID, node_run_id: NodeR
                 else:
                     pass
         for task_id in task_ids_to_cancel:
-            await self._finish_task(task_id)
+            await self._finish_task(task_id, 'sibling')
 
 
 def _is_any_iterable(x: Any) -> TypeGuard[Iterable[Any]]:

pydantic_graph/pydantic_graph/beta/graph_builder.py

@@ -658,6 +658,7 @@ def build(self, validate_graph_structure: bool = True) -> Graph[StateT, DepsT, G
         if validate_graph_structure:
             _validate_graph_structure(nodes, edges_by_source)
         parent_forks = _collect_dominating_forks(nodes, edges_by_source)
+        intermediate_join_nodes = _compute_intermediate_join_nodes(nodes, parent_forks)
 
         return Graph[StateT, DepsT, GraphInputT, GraphOutputT](
             name=self.name,
@@ -668,6 +669,7 @@ def build(self, validate_graph_structure: bool = True) -> Graph[StateT, DepsT, G
             nodes=nodes,
             edges_by_source=edges_by_source,
             parent_forks=parent_forks,
+            intermediate_join_nodes=intermediate_join_nodes,
             auto_instrument=self.auto_instrument,
         )
 
@@ -948,6 +950,40 @@ def _handle_path(path: Path, last_source_id: NodeID):
     return dominating_forks
 
 
+def _compute_intermediate_join_nodes(
+    nodes: dict[NodeID, AnyNode], parent_forks: dict[JoinID, ParentFork[NodeID]]
+) -> dict[JoinID, set[JoinID]]:
+    """Compute which joins have other joins as intermediate nodes.
+
+    A join J1 is an intermediate node of join J2 if J1 appears in J2's intermediate_nodes
+    (as computed relative to J2's parent fork).
+
+    This information is used to determine:
+    1. Which joins are "final" (have no other joins in their intermediate_nodes)
+    2. When selecting which reducer to proceed with when there are no active tasks
+
+    Args:
+        nodes: All nodes in the graph
+        parent_forks: Parent fork information for each join
+
+    Returns:
+        A mapping from each join to the set of joins that are intermediate to it
+    """
+    intermediate_join_nodes: dict[JoinID, set[JoinID]] = {}
+
+    for join_id, parent_fork in parent_forks.items():
+        intermediate_joins = set[JoinID]()
+        for intermediate_node_id in parent_fork.intermediate_nodes:
+            # Check if this intermediate node is also a join
+            intermediate_node = nodes.get(intermediate_node_id)
+            if isinstance(intermediate_node, Join):
+                # Add it regardless of whether it has the same parent fork
+                intermediate_joins.add(JoinID(intermediate_node_id))
+        intermediate_join_nodes[join_id] = intermediate_joins
+
+    return intermediate_join_nodes
+
+
 def _replace_placeholder_node_ids(nodes: dict[NodeID, AnyNode], edges_by_source: dict[NodeID, list[Path]]):
     node_id_remapping = _build_placeholder_node_id_remapping(nodes)
     replaced_nodes = {

tests/graph/beta/test_graph_execution.py

@@ -7,7 +7,7 @@
 import pytest
 
 from pydantic_graph.beta import GraphBuilder, StepContext
-from pydantic_graph.beta.join import ReduceFirstValue, reduce_list_append
+from pydantic_graph.beta.join import ReduceFirstValue, reduce_list_append, reduce_list_extend
 
 pytestmark = pytest.mark.anyio
 
@@ -343,3 +343,40 @@ async def step3(ctx: StepContext[ExecutionState, None, list[int]]) -> str:
     assert 'parallel-100' in state.log
     assert 'step3' in state.log
     assert result == 'Result: 152'  # (50+1) + (100+1) = 152
+
+
+async def test_multiple_sequential_joins():
+    g = GraphBuilder(output_type=list[int])
+
+    @g.step
+    async def source(ctx: StepContext[None, None, None]) -> int:
+        return 10
+
+    @g.step
+    async def add_one(ctx: StepContext[None, None, int]) -> list[int]:
+        return [ctx.inputs + 1]
+
+    @g.step
+    async def add_two(ctx: StepContext[None, None, int]) -> list[int]:
+        return [ctx.inputs + 2]
+
+    @g.step
+    async def add_three(ctx: StepContext[None, None, int]) -> list[int]:
+        return [ctx.inputs + 3]
+
+    collect = g.join(reduce_list_extend, initial_factory=list[int], parent_fork_id='source_fork', node_id='collect')
+    mediator = g.join(reduce_list_extend, initial_factory=list[int], node_id='mediator')
+
+    # Broadcasting: send the value from source to all three steps
+    g.add(
+        g.edge_from(g.start_node).to(source),
+        g.edge_from(source).to(add_one, add_two, add_three, fork_id='source_fork'),
+        g.edge_from(add_one, add_two).to(mediator),
+        g.edge_from(mediator).to(collect),
+        g.edge_from(add_three).to(collect),
+        g.edge_from(collect).to(g.end_node),
+    )
+
+    graph = g.build()
+    result = await graph.run()
+    assert sorted(result) == [11, 12, 13]