From a3c177102f66dc4634ed32d494cf2fd0a46eaea9 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Tue, 13 Oct 2026 05:51:34 -0500 Subject: [PATCH 01/20] perf(graph): index iterator ancestry once per snapshot Reuse one analysis during recursive If inference and preparation. Add shared correctness and operation-count evidence. --- invokeai/app/services/shared/README.md | 8 ++ .../app/services/shared/graph_if_runtime.py | 19 ++-- .../shared/graph_iterator_ancestry.py | 35 ++++++++ .../app/services/shared/graph_materializer.py | 10 ++- .../services/shared/test_graph_if_runtime.py | 68 ++++++++++++++ .../shared/test_graph_iterator_ancestry.py | 89 +++++++++++++++++++ .../shared/test_graph_materializer.py | 52 +++++++++++ 7 files changed, 270 insertions(+), 11 deletions(-) create mode 100644 invokeai/app/services/shared/graph_iterator_ancestry.py create mode 100644 tests/app/services/shared/test_graph_iterator_ancestry.py diff --git a/invokeai/app/services/shared/README.md b/invokeai/app/services/shared/README.md index 44b876d1651..99c96cc0c50 100644 --- a/invokeai/app/services/shared/README.md +++ b/invokeai/app/services/shared/README.md @@ -669,6 +669,14 @@ materialization. - Try to resolve any `If`-specific scheduling state. - If the node is ready and not deferred by an unresolved `If`, enqueue it into its class queue. +Literal-condition `If` fallback checks use the private `_IteratorAncestry` context. One context belongs to one +compatibility `prepare()` call and its source DAG, shared by candidate admission and selected-node materialization; +a standalone iteration-path lookup creates its own context and passes it through both `If` and ordinary-node +recursion. The filtered iterator graph is built only when an unresolved selected branch needs the ancestry fallback. +The context computes boolean proper-ancestor reachability with collector and final-`For` boundaries preserved, and +reuses that analysis within the call. No context is retained on execution state or serialized, so later calls observe +new supplied graphs and source mutations without cache invalidation. + ### 4.5 Readiness and class ordering - `_enqueue_if_ready(nid)` applies generic readiness: `indegree == 0`, not executed or claimed, and, for an `If` diff --git a/invokeai/app/services/shared/graph_if_runtime.py b/invokeai/app/services/shared/graph_if_runtime.py index c2ea5d9fa9a..eafd7ab0ef4 100644 --- a/invokeai/app/services/shared/graph_if_runtime.py +++ b/invokeai/app/services/shared/graph_if_runtime.py @@ -5,8 +5,9 @@ from invokeai.app.invocations.logic import IfInvocation from invokeai.app.services.shared.execution_engine.primitives import ActivationGate from invokeai.app.services.shared.execution_engine.scheduler import ActivationDependency +from invokeai.app.services.shared.graph_iterator_ancestry import _IteratorAncestry from invokeai.app.services.shared.graph_models import Edge, ExecutionToken -from invokeai.app.services.shared.graph_validation import ForInvocation, IterateInvocation, nx +from invokeai.app.services.shared.graph_validation import IterateInvocation, nx if TYPE_CHECKING: from invokeai.app.services.shared.graph import GraphExecutionState @@ -19,6 +20,7 @@ def _get_effective_iteration_paths_for_edge( get_prepared_nodes_for_source: Callable[[str], set[str]], get_prepared_edge_iteration_path: Callable[[Edge, str], tuple[int, ...]], source_graph: Optional["nx.DiGraph"] = None, + iterator_ancestry: Optional[_IteratorAncestry] = None, ) -> set[tuple[int, ...]]: source_node_id = edge.source.node_id prepared_nodes: set[str] = set() @@ -41,6 +43,7 @@ def _get_effective_iteration_paths_for_edge( get_prepared_nodes_for_source, get_prepared_edge_iteration_path, source_graph, + iterator_ancestry, ) ) @@ -57,6 +60,7 @@ def _get_effective_iteration_paths_for_edge( get_prepared_nodes_for_source, get_prepared_edge_iteration_path, source_graph, + iterator_ancestry, ) for input_edge in input_edges ) @@ -70,8 +74,11 @@ def _get_if_iteration_paths( get_prepared_nodes_for_source: Callable[[str], set[str]], get_prepared_edge_iteration_path: Callable[[Edge, str], tuple[int, ...]], source_graph: Optional["nx.DiGraph"] = None, + iterator_ancestry: Optional[_IteratorAncestry] = None, ) -> list[tuple[int, ...]]: """Infer the selected branch's effective iteration paths for an If output.""" + if iterator_ancestry is None: + iterator_ancestry = _IteratorAncestry(state, source_graph) node = state.graph.get_node(node_id) assert isinstance(node, IfInvocation) next_visited_source_ids = {*visited_source_ids, node_id} @@ -113,6 +120,7 @@ def _get_if_iteration_paths( get_prepared_nodes_for_source, get_prepared_edge_iteration_path, source_graph, + iterator_ancestry, ) for edge in selected_edges ) @@ -132,14 +140,7 @@ def _get_if_iteration_paths( if not has_selected_edges: return [()] if not condition_edges and has_unresolved_selected_source: - iterator_graph = state._iterator_graph( - source_graph if source_graph is not None else state._get_source_graph_flat() - ) - has_iterator_ancestor = any( - isinstance(state.graph.get_node(ancestor_id), (ForInvocation, IterateInvocation)) - for ancestor_id in nx.ancestors(iterator_graph, node_id) - ) - if not has_iterator_ancestor: + if not iterator_ancestry.has_iterator_ancestor(node_id): return [()] return [] return sorted(path for path in paths if not any(path != other and other[: len(path)] == path for other in paths)) diff --git a/invokeai/app/services/shared/graph_iterator_ancestry.py b/invokeai/app/services/shared/graph_iterator_ancestry.py new file mode 100644 index 00000000000..70ab4d7ffa2 --- /dev/null +++ b/invokeai/app/services/shared/graph_iterator_ancestry.py @@ -0,0 +1,35 @@ +from typing import TYPE_CHECKING, Optional + +from invokeai.app.invocations.loops import ForInvocation +from invokeai.app.services.shared.graph_validation import IterateInvocation, nx + +if TYPE_CHECKING: + from invokeai.app.services.shared.graph import GraphExecutionState + + +class _IteratorAncestry: + """Lazily index strict iterator ancestry for one source-graph traversal.""" + + def __init__(self, state: "GraphExecutionState", source_graph: Optional["nx.DiGraph"] = None) -> None: + self._state = state + self._source_graph = source_graph + self._has_iterator_ancestor: dict[str, bool] | None = None + + def has_iterator_ancestor(self, node_id: str) -> bool: + if self._has_iterator_ancestor is None: + source_graph = self._source_graph + if source_graph is None: + source_graph = self._state._get_source_graph_flat() + iterator_graph = self._state._iterator_graph(source_graph) + has_iterator_ancestor = dict.fromkeys(iterator_graph, False) + for source_node_id in nx.topological_sort(iterator_graph): + if has_iterator_ancestor[source_node_id] or isinstance( + self._state.graph.get_node(source_node_id), (ForInvocation, IterateInvocation) + ): + for successor_id in iterator_graph.successors(source_node_id): + has_iterator_ancestor[successor_id] = True + self._has_iterator_ancestor = has_iterator_ancestor + + if node_id not in self._has_iterator_ancestor: + raise nx.NetworkXError(f"The node {node_id} is not in the digraph.") + return self._has_iterator_ancestor[node_id] diff --git a/invokeai/app/services/shared/graph_materializer.py b/invokeai/app/services/shared/graph_materializer.py index 26a6c34c91c..2e3da4a2880 100644 --- a/invokeai/app/services/shared/graph_materializer.py +++ b/invokeai/app/services/shared/graph_materializer.py @@ -14,6 +14,7 @@ LoopState, ) from invokeai.app.services.shared import graph_if_runtime +from invokeai.app.services.shared.graph_iterator_ancestry import _IteratorAncestry from invokeai.app.services.shared.graph_models import Edge, EdgeConnection from invokeai.app.services.shared.graph_validation import ( COLLECTION_FIELD, @@ -1411,7 +1412,9 @@ def create_execution_node( return new_nodes - def _has_admitted_source_mapping(self, node_id: str, graph: "nx.DiGraph") -> bool: + def _has_admitted_source_mapping( + self, node_id: str, graph: "nx.DiGraph", iterator_ancestry: Optional[_IteratorAncestry] = None + ) -> bool: if not any(isinstance(node, IfInvocation) for node in self._state.graph.nodes.values()): return True if isinstance(self._state.graph.get_node(node_id), IfInvocation): @@ -1426,6 +1429,7 @@ def _has_admitted_source_mapping(self, node_id: str, graph: "nx.DiGraph") -> boo self._get_prepared_nodes_for_source, self._get_prepared_edge_iteration_path, graph, + iterator_ancestry, ) ) mappings = self._get_if_condition_iteration_mappings(node_id, graph) @@ -1705,6 +1709,7 @@ def get_iteration_node( def prepare(self, base_g: Optional["nx.DiGraph"] = None) -> Optional[str]: g = base_g if base_g is not None else self._state._get_source_graph_flat() + iterator_ancestry = _IteratorAncestry(self._state, g) self._attach_pending_if_inputs() next_node_id = next( ( @@ -1734,7 +1739,7 @@ def prepare(self, base_g: Optional["nx.DiGraph"] = None) -> Optional[str]: and self._is_if_condition_ready(node_id) or not isinstance(self._state.graph.get_node(node_id), IfInvocation) ) - and self._has_admitted_source_mapping(node_id, g) + and self._has_admitted_source_mapping(node_id, g, iterator_ancestry) and ( not isinstance(self._state.graph.get_node(node_id), (ForInvocation, IterateInvocation)) or all(source_id in self._state.executed for source_id, _ in g.in_edges(node_id)) @@ -1769,6 +1774,7 @@ def prepare(self, base_g: Optional["nx.DiGraph"] = None) -> Optional[str]: self._get_prepared_nodes_for_source, self._get_prepared_edge_iteration_path, g, + iterator_ancestry, ) condition_edges = self._state.graph._get_input_edges(next_node_id, "condition") if condition_edges: diff --git a/tests/app/services/shared/test_graph_if_runtime.py b/tests/app/services/shared/test_graph_if_runtime.py index ad596f2812a..fac41fccf75 100644 --- a/tests/app/services/shared/test_graph_if_runtime.py +++ b/tests/app/services/shared/test_graph_if_runtime.py @@ -172,3 +172,71 @@ def test_nested_if_iteration_paths_use_supplied_source_graph_recursively() -> No lambda _edge, _prepared_id: (), constrained_graph, ) == [()] + + +def test_unresolved_nested_literal_ifs_share_one_iterator_graph(monkeypatch: pytest.MonkeyPatch) -> None: + nodes = {"iterator": IterateInvocation(id="iterator"), "guarded": IfInvocation(id="guarded")} + nodes.update({f"if-{i}": IfInvocation(id=f"if-{i}", condition=True) for i in range(24)}) + edges = [ + Edge( + source=EdgeConnection(node_id="iterator", field="item"), + destination=EdgeConnection(node_id="guarded", field="condition"), + ) + ] + previous = "guarded" + for i in range(24): + edges.append( + Edge( + source=EdgeConnection(node_id=previous, field="value"), + destination=EdgeConnection(node_id=f"if-{i}", field="true_input"), + ) + ) + previous = f"if-{i}" + state = GraphExecutionState(graph=Graph()) + state.graph = Graph(nodes=nodes, edges=edges) + original = GraphExecutionState._iterator_graph + builds = [] + + def count_build(self, graph=None): + builds.append(graph) + return original(self, graph) + + monkeypatch.setattr(GraphExecutionState, "_iterator_graph", count_build) + assert graph_if_runtime._get_if_iteration_paths(state, previous, set(), lambda _: set(), lambda *_: ()) == [] + assert len(builds) == 1 + + +def test_resolved_literal_if_does_not_build_iterator_graph(monkeypatch: pytest.MonkeyPatch) -> None: + state = GraphExecutionState(graph=Graph(nodes={"if": IfInvocation(id="if", condition=True, true_input=42)})) + + def unexpected_build(*_args, **_kwargs): + pytest.fail("resolved paths must not build the iterator graph") + + monkeypatch.setattr(GraphExecutionState, "_iterator_graph", unexpected_build) + assert graph_if_runtime._get_if_iteration_paths(state, "if", set(), lambda _: set(), lambda *_: ()) == [()] + + +@pytest.mark.parametrize("executed", [False, True]) +def test_prepared_selected_source_preserves_paths_without_ancestry( + executed: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + state = GraphExecutionState( + graph=Graph(nodes={"if": IfInvocation(id="if", condition=True), "source": AddInvocation(id="source")}) + ) + state.graph.add_edge( + Edge( + source=EdgeConnection(node_id="source", field="value"), + destination=EdgeConnection(node_id="if", field="true_input"), + ) + ) + state.source_prepared_mapping["source"] = set() if executed else {"source-exec"} + if executed: + state.executed.add("source") + + def unexpected_build(*_args, **_kwargs): + pytest.fail("prepared selected sources must not build the iterator graph") + + monkeypatch.setattr(GraphExecutionState, "_iterator_graph", unexpected_build) + assert graph_if_runtime._get_if_iteration_paths( + state, "if", set(), lambda source_id: state.source_prepared_mapping[source_id], lambda *_: (2,) + ) == ([] if executed else [(2,)]) diff --git a/tests/app/services/shared/test_graph_iterator_ancestry.py b/tests/app/services/shared/test_graph_iterator_ancestry.py new file mode 100644 index 00000000000..931a3c69496 --- /dev/null +++ b/tests/app/services/shared/test_graph_iterator_ancestry.py @@ -0,0 +1,89 @@ +import random + +import networkx as nx +import pytest + +from invokeai.app.invocations.loops import ForInvocation +from invokeai.app.invocations.math import AddInvocation +from invokeai.app.services.shared.graph import Edge, EdgeConnection, Graph, GraphExecutionState +from invokeai.app.services.shared.graph_iterator_ancestry import _IteratorAncestry +from invokeai.app.services.shared.graph_validation import CollectInvocation, IterateInvocation + + +@pytest.mark.parametrize("seed", range(5)) +def test_iterator_ancestry_matches_filtered_networkx_on_generated_dags(seed: int) -> None: + randomizer = random.Random(seed) + classes = (AddInvocation, IterateInvocation, ForInvocation, CollectInvocation) + nodes = {str(i): randomizer.choice(classes)(id=str(i)) for i in range(40)} + edges = [] + for i in range(40): + for j in range(i + 1, 40): + if randomizer.random() < 0.12: + source = nodes[str(i)] + field = randomizer.choice(("item", "final_state")) if isinstance(source, ForInvocation) else "value" + edges.append( + Edge( + source=EdgeConnection(node_id=str(i), field=field), + destination=EdgeConnection(node_id=str(j), field="a"), + ) + ) + # These are topology fixtures, not executable invocation workflows. + state = GraphExecutionState(graph=Graph()) + state.graph = Graph(nodes=nodes, edges=edges) + source_graph = state._get_source_graph_flat() + for graph in (source_graph, nx.subgraph_view(source_graph, filter_edge=lambda u, v: (int(u) + int(v)) % 3 != 0)): + filtered = state._iterator_graph(graph) + ancestry = _IteratorAncestry(state, graph) + for node_id in reversed(list(graph)): + expected = any( + isinstance(nodes[ancestor], (ForInvocation, IterateInvocation)) + for ancestor in nx.ancestors(filtered, node_id) + ) + assert ancestry.has_iterator_ancestor(node_id) == expected + assert ancestry.has_iterator_ancestor(node_id) == expected + + +def test_iterator_self_collectors_and_final_for_outputs_are_not_ancestors() -> None: + nodes = { + "iterate": IterateInvocation(id="iterate"), + "for": ForInvocation(id="for"), + "collect": CollectInvocation(id="collect"), + "after_collect": AddInvocation(id="after_collect"), + "final": AddInvocation(id="final"), + "body": AddInvocation(id="body"), + } + edges = [ + Edge( + source=EdgeConnection(node_id=source, field=field), + destination=EdgeConnection(node_id=destination, field="a"), + ) + for source, field, destination in ( + ("iterate", "item", "collect"), + ("collect", "collection", "after_collect"), + ("for", "final_state", "final"), + ("for", "item", "body"), + ) + ] + state = GraphExecutionState(graph=Graph()) + state.graph = Graph(nodes=nodes, edges=edges) + ancestry = _IteratorAncestry(state) + assert {node_id: ancestry.has_iterator_ancestor(node_id) for node_id in nodes} == { + "iterate": False, + "for": False, + "collect": False, + "after_collect": False, + "final": False, + "body": True, + } + + +def test_new_context_observes_supplied_graph_mutations() -> None: + state = GraphExecutionState(graph=Graph()) + state.graph = Graph(nodes={"iterator": IterateInvocation(id="iterator"), "sink": AddInvocation(id="sink")}) + graph = nx.DiGraph() + graph.add_nodes_from(state.graph.nodes) + assert not _IteratorAncestry(state, graph).has_iterator_ancestor("sink") + graph.add_edge("iterator", "sink") + assert _IteratorAncestry(state, graph).has_iterator_ancestor("sink") + graph.remove_edge("iterator", "sink") + assert not _IteratorAncestry(state, graph).has_iterator_ancestor("sink") diff --git a/tests/app/services/shared/test_graph_materializer.py b/tests/app/services/shared/test_graph_materializer.py index e407af81998..941aaaa293c 100644 --- a/tests/app/services/shared/test_graph_materializer.py +++ b/tests/app/services/shared/test_graph_materializer.py @@ -68,3 +68,55 @@ def test_materializer_preserves_execution_edges_and_source_nodes() -> None: assert sink.a == 0 state.complete(prepared_sink.id, IntegerOutput(value=12)) assert state.is_complete() + + +def test_prepare_shares_iterator_ancestry_across_admission_and_materialization(monkeypatch: pytest.MonkeyPatch) -> None: + from invokeai.app.invocations.logic import IfInvocation + from invokeai.app.services.shared.graph_validation import IterateInvocation, nx + + nodes = { + "iterator": IterateInvocation(id="iterator"), + "condition-a": AddInvocation(id="condition-a"), + "condition-b": AddInvocation(id="condition-b"), + "inner-a": IfInvocation(id="inner-a"), + "inner-b": IfInvocation(id="inner-b"), + "blocked-a": IfInvocation(id="blocked-a", condition=True), + "blocked-b": IfInvocation(id="blocked-b", condition=True), + "selected": IfInvocation(id="selected", condition=True), + } + edges = [ + graph_facade.Edge( + source=graph_facade.EdgeConnection(node_id=source, field=field), + destination=graph_facade.EdgeConnection(node_id=destination, field=input_field), + ) + for source, field, destination, input_field in ( + ("iterator", "item", "inner-a", "true_input"), + ("condition-a", "value", "inner-a", "condition"), + ("condition-b", "value", "inner-b", "condition"), + ("inner-a", "value", "blocked-a", "true_input"), + ("inner-a", "value", "blocked-b", "true_input"), + ("inner-b", "value", "selected", "true_input"), + ) + ] + state = graph_facade.GraphExecutionState(graph=graph_facade.Graph()) + state.graph = graph_facade.Graph(nodes=nodes, edges=edges) + state.executed.add("iterator") + graph = state._get_source_graph_flat() + # Keep condition sources outside this preparation view; their If nodes stay unresolved. + preparation_graph = nx.subgraph_view(graph, filter_node=lambda node_id: not node_id.startswith("condition-")) + original = graph_facade.GraphExecutionState._iterator_graph + builds = [] + + def count_build(self, source_graph=None): + builds.append(source_graph) + return original(self, source_graph) + + monkeypatch.setattr(graph_facade.GraphExecutionState, "_iterator_graph", count_build) + prepared_id = state._materializer().prepare(preparation_graph) + assert prepared_id in state.source_prepared_mapping["selected"] + assert len(state.source_prepared_mapping["selected"]) == 1 + assert "blocked-a" not in state.source_prepared_mapping + assert "blocked-b" not in state.source_prepared_mapping + assert "inner-a" not in state.source_prepared_mapping + assert "inner-b" not in state.source_prepared_mapping + assert len(builds) == 1 From d83d420dbd59423d03165c8505a81d0f50b461ca Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Tue, 13 Oct 2026 07:54:06 -0500 Subject: [PATCH 02/20] test(graph): pin collector validation edge cases --- tests/test_node_graph.py | 48 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/tests/test_node_graph.py b/tests/test_node_graph.py index d06a1905cb5..c758e0c895a 100644 --- a/tests/test_node_graph.py +++ b/tests/test_node_graph.py @@ -1449,6 +1449,54 @@ def test_graph_update_node_reconnects_source_feeding_collector_chain(): graph.validate_self() +def test_graph_validate_self_accepts_iterator_before_collectors_with_float_roots(): + float_source = FloatInvocation(id="float", value=2.5) + downstream_collector = CollectInvocation(id="collector_y") + upstream_collector = CollectInvocation(id="collector_x") + first_iterate = IterateInvocation(id="iterate_1") + second_iterate = IterateInvocation(id="iterate_2") + # Keep iterators earlier in nodes order than the collectors, as in issue #9695. + graph = Graph( + nodes={ + node.id: node + for node in (float_source, first_iterate, second_iterate, downstream_collector, upstream_collector) + }, + edges=[ + create_edge(float_source.id, "value", downstream_collector.id, "item"), + create_edge(downstream_collector.id, "collection", upstream_collector.id, "collection"), + create_edge(upstream_collector.id, "collection", first_iterate.id, "collection"), + create_edge(upstream_collector.id, "collection", second_iterate.id, "collection"), + ], + ) + + graph.validate_self() + + +def test_graph_update_node_renames_float_source_with_transiently_invalid_collector_order(): + integer_one = IntegerInvocation(id="i1", value=1) + float_source = FloatInvocation(id="f", value=2.5) + first_collector = CollectInvocation(id="c") + second_collector = CollectInvocation(id="d") + integer_two = IntegerInvocation(id="i2", value=2) + graph = Graph( + nodes={node.id: node for node in (integer_one, float_source, first_collector, second_collector, integer_two)} + ) + for edge in ( + create_edge("i1", "value", "c", "item"), + create_edge("f", "value", "c", "item"), + create_edge("c", "collection", "d", "collection"), + create_edge("i2", "value", "d", "item"), + ): + graph.add_edge(edge) + + graph.update_node("f", FloatInvocation(id="f2", value=2.5)) + + assert "f" not in graph.nodes + assert create_edge("f2", "value", "c", "item") in graph.edges + assert all(edge.source.node_id != "f" and edge.destination.node_id != "f" for edge in graph.edges) + graph.validate_self() + + def test_graph_add_collection_input_revalidates_iterator_consumers(): integer_source = IntegerInvocation(id="integer", value=1) float_collection = FloatCollectionInvocation(id="float_collection", collection=[2.5]) From bd25035bbd06c2d52ddab015caed2366f96c7916 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Tue, 13 Oct 2026 08:14:51 -0500 Subject: [PATCH 03/20] fix(graph): schedule selected root inputs in For conditions Bind root-frame branch inputs to prepared If activation frames so selected values become schedulable inside For bodies. Defer branch edge materialization until the If branch is selected. Add mixed, all-true, and all-false execution regressions. --- .../app/services/shared/graph_if_runtime.py | 27 +++++++++-- .../app/services/shared/graph_materializer.py | 14 ++++-- .../test_execution_engine_differential.py | 47 +++++++++++++++++++ 3 files changed, 80 insertions(+), 8 deletions(-) diff --git a/invokeai/app/services/shared/graph_if_runtime.py b/invokeai/app/services/shared/graph_if_runtime.py index c2ea5d9fa9a..8ced76dd4f7 100644 --- a/invokeai/app/services/shared/graph_if_runtime.py +++ b/invokeai/app/services/shared/graph_if_runtime.py @@ -153,8 +153,8 @@ def _record_activation_dependencies( return dependencies source_node_id = state._prepared_registry().get_source_node_id(exec_node_id) if any(isinstance(node, IfInvocation) for node in state.graph.nodes.values()): - dependencies = state._get_source_activation_dependencies( - source_node_id, state._get_iteration_path(exec_node_id) + dependencies = _get_source_activation_dependencies_for_preparation( + state, source_node_id, state._get_iteration_path(exec_node_id) ) else: dependencies = () @@ -162,10 +162,31 @@ def _record_activation_dependencies( return dependencies +def _get_source_activation_dependencies_for_preparation( + state: "GraphExecutionState", source_node_id: str, iteration_path: tuple[int, ...] +) -> tuple[ActivationDependency, ...]: + dependencies = state._get_source_activation_dependencies(source_node_id, iteration_path) + if iteration_path or not dependencies: + return dependencies + + # Root-frame sources can feed a branch selected by one of several loop-local If executions. + # Bind the source to a frame that actually selected it so scheduler admission uses that gate. + candidate_frames = sorted( + {frame for dependency in dependencies for frame in state._prepared_if_exec_frames(dependency.owner_id)} + ) + for frame in candidate_frames: + frame_dependencies = state._get_source_activation_dependencies(source_node_id, frame) + if frame_dependencies and all( + state._is_activation_dependency_satisfied(dependency) for dependency in frame_dependencies + ): + return frame_dependencies + return dependencies + + def _is_source_activation_admitted( state: "GraphExecutionState", source_node_id: str, iteration_path: tuple[int, ...] = () ) -> bool: - dependencies = state._get_source_activation_dependencies(source_node_id, iteration_path) + dependencies = _get_source_activation_dependencies_for_preparation(state, source_node_id, iteration_path) return not dependencies or all(state._is_activation_dependency_satisfied(dependency) for dependency in dependencies) diff --git a/invokeai/app/services/shared/graph_materializer.py b/invokeai/app/services/shared/graph_materializer.py index 26a6c34c91c..ee6da5077b1 100644 --- a/invokeai/app/services/shared/graph_materializer.py +++ b/invokeai/app/services/shared/graph_materializer.py @@ -309,7 +309,10 @@ def create_for_body_iteration(self, source_for_id: str, prepared_for_id: str) -> node = self._state.graph.get_node(source_node_id) new_edges: list[Edge] = [] - for edge in self._state.graph._get_input_edges(source_node_id): + input_edges = self._state.graph._get_input_edges(source_node_id) + if isinstance(node, IfInvocation): + input_edges = [edge for edge in input_edges if edge.destination.field == "condition"] + for edge in input_edges: prepared_source_id = source_to_prepared.get(edge.source.node_id) if prepared_source_id is None: prepared_source_id = self.get_iteration_node( @@ -1549,16 +1552,17 @@ def _attach_pending_if_inputs(self, *, enqueue: bool = True) -> None: if source_edge.source.node_id not in self._state.source_prepared_mapping: continue candidates = [ - prepared_id + (prepared_id, self._get_prepared_edge_iteration_path(source_edge, prepared_id)) for prepared_id in self._get_ordered_prepared_nodes_for_edge(source_edge) - if self._get_prepared_edge_iteration_path(source_edge, prepared_id) == target_path ] + candidates = [candidate for candidate in candidates if target_path[: len(candidate[1])] == candidate[1]] if not candidates: continue - selected_exec_ids.append(candidates[0]) + selected_exec_id, _ = min(candidates, key=lambda candidate: (-len(candidate[1]), candidate[0])) + selected_exec_ids.append(selected_exec_id) attached_edges.append( Edge( - source=EdgeConnection(node_id=candidates[0], field=source_edge.source.field), + source=EdgeConnection(node_id=selected_exec_id, field=source_edge.source.field), destination=EdgeConnection(node_id=exec_node_id, field=selected_field), ) ) diff --git a/tests/app/services/shared/test_execution_engine_differential.py b/tests/app/services/shared/test_execution_engine_differential.py index ce89287c311..9534db6a6e1 100644 --- a/tests/app/services/shared/test_execution_engine_differential.py +++ b/tests/app/services/shared/test_execution_engine_differential.py @@ -8489,3 +8489,50 @@ def fail_branch_analysis(*_: object, **__: object) -> set[str]: assert _activation_projection(resumed_state) == _activation_projection(expected_state) assert _execution_token_projection(resumed_state) == _execution_token_projection(expected_state) assert resumed_state.is_complete() + + +def _run_to_completion( + state: GraphExecutionState, *, max_steps: int = 100_000 +) -> tuple[list[str], GraphExecutionState]: + trace: list[str] = [] + for _ in range(max_steps): + node = state.next() + if node is None: + return trace, state + trace.append(state.prepared_source_mapping[node.id]) + state.complete(node.id, node.invoke(Mock())) + raise AssertionError(f"Execution exceeded {max_steps} scheduled nodes") + + +@pytest.mark.parametrize( + ("collection", "expected"), + [ + ([False, True], ["false-root", "true-root"]), + ([True, True], ["true-root", "true-root"]), + ([False, False], ["false-root", "false-root"]), + ], + ids=["mixed", "all-true", "all-false"], +) +def test_for_condition_if_reads_root_frame_sources_and_skips_unselected_root( + collection: list[bool], expected: list[str] +) -> None: + graph = Graph() + graph.add_node(ForInvocation(id="for", collection=collection)) + graph.add_node(AnyTypeTestInvocation(id="root_true", value="true-root")) + graph.add_node(AnyTypeTestInvocation(id="root_false", value="false-root")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(ForReturnInvocation(id="return")) + graph.add_node(AnyTypeTestInvocation(id="after")) + graph.add_edge(create_edge("for", "item", "if", "condition")) + graph.add_edge(create_edge("root_true", "value", "if", "true_input")) + graph.add_edge(create_edge("root_false", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "return", "output")) + graph.add_edge(create_loop_linkage("for", "return")) + graph.add_edge(create_edge("for", "output_collection", "after", "value")) + + trace, state = _run_to_completion(GraphExecutionState(graph=graph)) + + assert state.is_complete() + assert _source_output(state, "after").value == expected + assert trace.count("root_true") == (1 if True in collection else 0) + assert trace.count("root_false") == (1 if False in collection else 0) From 2cace230db582191a5f8e60539e78460d8a5d3ea Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Tue, 13 Oct 2026 08:44:53 -0500 Subject: [PATCH 04/20] fix(graph): preserve empty inner For collections Finalize empty nested For frames alongside non-empty contexts, filter sentinel nodes from iteration-scoped body inputs, and collect final outputs at the enclosing frame. Cover empty inner collections under both scheduler paths. --- .../app/services/shared/graph_materializer.py | 40 ++++++++++++++++--- .../test_execution_engine_differential.py | 30 ++++++++++++++ 2 files changed, 65 insertions(+), 5 deletions(-) diff --git a/invokeai/app/services/shared/graph_materializer.py b/invokeai/app/services/shared/graph_materializer.py index ee6da5077b1..c1775a56b57 100644 --- a/invokeai/app/services/shared/graph_materializer.py +++ b/invokeai/app/services/shared/graph_materializer.py @@ -1058,7 +1058,11 @@ def _get_collect_iteration_group_key(self, edge: Edge, sibling_depth: Optional[i isinstance(source_node, ForInvocation) and get_output_field_scope(source_node, edge.source.field) == OutputScope.Final ): - return self._state._get_for_parent_iteration_path(edge.source.node_id) + parent_path = self._state._get_for_parent_iteration_path(edge.source.node_id) + if edge.destination.field == ITEM_FIELD: + group_depth = len(parent_path) if sibling_depth is None else sibling_depth + return parent_path[: max(group_depth - 1, 0)] + return parent_path if edge.destination.field == ITEM_FIELD: if isinstance(source_node, ForInvocation) and source_node.index == -1: return path @@ -1079,8 +1083,22 @@ def _get_ordered_prepared_nodes_for_source(self, source_node_id: str) -> list[st key=lambda exec_node_id: (self._state._get_iteration_path(exec_node_id), exec_node_id), ) + def _prepared_node_has_edge_output(self, edge: Edge, prepared_id: str) -> bool: + source_node = self._state.graph.get_node(edge.source.node_id) + if not ( + isinstance(source_node, ForInvocation) + and get_output_field_scope(source_node, edge.source.field) == OutputScope.Iteration + ): + return True + prepared_node = self._state.execution_graph.get_node(prepared_id) + return isinstance(prepared_node, ForInvocation) and prepared_node.index >= 0 + def _get_ordered_prepared_nodes_for_edge(self, edge: Edge) -> list[str]: - prepared_nodes = self._get_ordered_prepared_nodes_for_source(edge.source.node_id) + prepared_nodes = [ + prepared_id + for prepared_id in self._get_ordered_prepared_nodes_for_source(edge.source.node_id) + if self._prepared_node_has_edge_output(edge, prepared_id) + ] source_node = self._state.graph.get_node(edge.source.node_id) if not ( isinstance(source_node, ForInvocation) @@ -1279,8 +1297,11 @@ def _index_prepared_nodes_by_iteration_path( self, prepared_nodes: set[str], input_edges: list[Edge] ) -> dict[tuple[int, ...], list[str]]: prepared_nodes_by_iteration_path: dict[tuple[int, ...], list[str]] = {} + edge = input_edges[0] for prepared_id in prepared_nodes: - iteration_path = self._get_prepared_edge_iteration_path(input_edges[0], prepared_id) + if not self._prepared_node_has_edge_output(edge, prepared_id): + continue + iteration_path = self._get_prepared_edge_iteration_path(edge, prepared_id) prepared_nodes_by_iteration_path.setdefault(iteration_path, []).append(prepared_id) return prepared_nodes_by_iteration_path @@ -1331,7 +1352,16 @@ def _get_parent_iteration_mappings( for node_id in iterator_nodes ] prepared_nodes_by_source = { - node_id: self._get_prepared_nodes_for_source(node_id) for node_id in parent_node_ids + node_id: { + prepared_id + for prepared_id in self._get_prepared_nodes_for_source(node_id) + if all( + self._prepared_node_has_edge_output(edge, prepared_id) + for edge in input_edges + if edge.source.node_id == node_id + ) + } + for node_id in parent_node_ids } prepared_nodes_by_source_and_path = { node_id: self._index_prepared_nodes_by_iteration_path( @@ -1850,7 +1880,7 @@ def prepare(self, base_g: Optional["nx.DiGraph"] = None) -> Optional[str]: self._state._invalidate_loop_caches_for_source(next_node_id) return next_node_id - if isinstance(next_node, ForInvocation) and all( + if isinstance(next_node, ForInvocation) and any( self._state.execution_graph.get_node(exec_node_id).index == -1 for exec_node_id in new_node_ids ): self._mark_empty_for_complete(next_node_id) diff --git a/tests/app/services/shared/test_execution_engine_differential.py b/tests/app/services/shared/test_execution_engine_differential.py index 9534db6a6e1..741a0c0dbfd 100644 --- a/tests/app/services/shared/test_execution_engine_differential.py +++ b/tests/app/services/shared/test_execution_engine_differential.py @@ -8536,3 +8536,33 @@ def test_for_condition_if_reads_root_frame_sources_and_skips_unselected_root( assert _source_output(state, "after").value == expected assert trace.count("root_true") == (1 if True in collection else 0) assert trace.count("root_false") == (1 if False in collection else 0) + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_iterate_for_collect_preserves_empty_inner_collection(force_compatibility_scheduler: bool) -> None: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="source", first=[[1, 2], []])) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(ForInvocation(id="for")) + graph.add_node(AddInvocation(id="add", b=10)) + graph.add_node(ForReturnInvocation(id="return")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(AnyTypeTestInvocation(id="after")) + graph.add_edge(create_edge("source", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "for", "collection")) + graph.add_edge(create_edge("for", "item", "add", "a")) + graph.add_edge(create_edge("add", "value", "return", "output")) + graph.add_edge(create_loop_linkage("for", "return")) + graph.add_edge(create_edge("for", "output_collection", "collect", "item")) + graph.add_edge(create_edge("collect", "collection", "after", "value")) + + state = GraphExecutionState(graph=graph) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) + trace, state = _run_to_completion(state) + + after_exec_ids = sorted(state.source_prepared_mapping["after"], key=state._get_iteration_path) + assert trace.count("add") == 2 + assert len(after_exec_ids) == 1 + assert state.results[after_exec_ids[0]].value == [[11, 12], []] + assert state.is_complete() From d8f6ac68e0b5639a8d7c3ac351526be50302552c Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Tue, 13 Oct 2026 09:55:28 -0500 Subject: [PATCH 05/20] fix(graph): align sibling iterator frames for If branches --- invokeai/app/services/shared/README.md | 4 + .../app/services/shared/graph_if_runtime.py | 53 +++- .../app/services/shared/graph_materializer.py | 251 ++++++++++++++++-- .../test_execution_engine_differential.py | 75 ++++++ 4 files changed, 350 insertions(+), 33 deletions(-) diff --git a/invokeai/app/services/shared/README.md b/invokeai/app/services/shared/README.md index 44b876d1651..fd3c2ee7594 100644 --- a/invokeai/app/services/shared/README.md +++ b/invokeai/app/services/shared/README.md @@ -658,6 +658,10 @@ materialization. execution. Unsupported fresh shapes and legacy snapshots use the compatibility controller. Rejected branch sources are not materialized. Pending `If` state is reconstructed from its durable execution edges, prepared mappings, and activation token after rehydration. +- A branch `Collect` whose `item` source combines exactly two sibling `Iterate` axes can close the axis beside the + `If` condition axis. The compatibility materializer expands that sibling `Iterate` beneath each resolved `If` frame + selecting the `Collect` branch, pairs only the matching condition frame, and closes the branch interval empty when no + frame selects it. An unrelated open branch axis is not combined with the condition frame. - For each new exec node: diff --git a/invokeai/app/services/shared/graph_if_runtime.py b/invokeai/app/services/shared/graph_if_runtime.py index 8ced76dd4f7..8fd1d633fa5 100644 --- a/invokeai/app/services/shared/graph_if_runtime.py +++ b/invokeai/app/services/shared/graph_if_runtime.py @@ -78,10 +78,19 @@ def _get_if_iteration_paths( condition_edges = state.graph._get_input_edges(node_id, "condition") condition_selections: list[tuple[Optional[tuple[int, ...]], str]] = [] + condition_iterator_ids: set[str] = set() if condition_edges: condition_edge = condition_edges[0] if condition_edge.source.node_id not in state.source_prepared_mapping: return [] + condition_node_id = condition_edge.source.node_id + condition_iterator_ids = { + ancestor_id + for ancestor_id in nx.ancestors(source_graph or state._get_source_graph_flat(), condition_node_id) + if isinstance(state.graph.get_node(ancestor_id), (ForInvocation, IterateInvocation)) + } + if isinstance(state.graph.get_node(condition_node_id), (ForInvocation, IterateInvocation)): + condition_iterator_ids.add(condition_node_id) for prepared_id in get_prepared_nodes_for_source(condition_edge.source.node_id): if prepared_id not in state.results: return [] @@ -104,8 +113,24 @@ def _get_if_iteration_paths( for edge in selected_edges ) if selected_edges: - branch_paths = set().union( - *( + branch_paths: set[tuple[int, ...]] = set() + source_graph_for_axes = source_graph or state._get_source_graph_flat() + for edge in selected_edges: + branch_source_id = edge.source.node_id + branch_iterator_ids = { + ancestor_id + for ancestor_id in nx.ancestors(source_graph_for_axes, branch_source_id) + if isinstance(state.graph.get_node(ancestor_id), (ForInvocation, IterateInvocation)) + } + if isinstance(state.graph.get_node(branch_source_id), (ForInvocation, IterateInvocation)): + branch_iterator_ids.add(branch_source_id) + if ( + condition_iterator_ids + and branch_iterator_ids + and not (condition_iterator_ids & branch_iterator_ids) + ): + continue + branch_paths.update( _get_effective_iteration_paths_for_edge( state, edge, @@ -114,9 +139,7 @@ def _get_if_iteration_paths( get_prepared_edge_iteration_path, source_graph, ) - for edge in selected_edges ) - ) else: branch_paths = {condition_path or ()} @@ -166,13 +189,16 @@ def _get_source_activation_dependencies_for_preparation( state: "GraphExecutionState", source_node_id: str, iteration_path: tuple[int, ...] ) -> tuple[ActivationDependency, ...]: dependencies = state._get_source_activation_dependencies(source_node_id, iteration_path) - if iteration_path or not dependencies: + if not dependencies: return dependencies - # Root-frame sources can feed a branch selected by one of several loop-local If executions. - # Bind the source to a frame that actually selected it so scheduler admission uses that gate. candidate_frames = sorted( - {frame for dependency in dependencies for frame in state._prepared_if_exec_frames(dependency.owner_id)} + { + frame + for dependency in dependencies + for frame in state._prepared_if_exec_frames(dependency.owner_id) + if not iteration_path or iteration_path[: len(frame)] == frame + } ) for frame in candidate_frames: frame_dependencies = state._get_source_activation_dependencies(source_node_id, frame) @@ -183,6 +209,17 @@ def _get_source_activation_dependencies_for_preparation( return dependencies +def _get_selected_if_frames_for_preparation( + state: "GraphExecutionState", if_node_id: str, branch_field: str +) -> list[tuple[int, ...]]: + """Return only resolved If frames that selected a Collect-backed branch.""" + return [ + state._get_iteration_path(exec_node_id) + for exec_node_id in state.source_prepared_mapping.get(if_node_id, ()) + if state._resolved_if_exec_branches.get(exec_node_id) == branch_field + ] + + def _is_source_activation_admitted( state: "GraphExecutionState", source_node_id: str, iteration_path: tuple[int, ...] = () ) -> bool: diff --git a/invokeai/app/services/shared/graph_materializer.py b/invokeai/app/services/shared/graph_materializer.py index c1775a56b57..9526005a1f5 100644 --- a/invokeai/app/services/shared/graph_materializer.py +++ b/invokeai/app/services/shared/graph_materializer.py @@ -1337,13 +1337,80 @@ def _get_indexed_iteration_node( return None return None + def _matches_sibling_if_collect_context( + self, iterator_nodes: list[str], prepared_iterators: tuple[str, ...] + ) -> bool: + """Match a branch Collect's sibling iterator to the selected If parent frame.""" + for index, iterator_node_id in enumerate(iterator_nodes): + collapsed_branch = self._get_sibling_iterate_if_collect_branch(iterator_node_id) + if collapsed_branch is None: + continue + if_node_id, branch_field = collapsed_branch + selected_frames = graph_if_runtime._get_selected_if_frames_for_preparation( + self._state, if_node_id, branch_field + ) + branch_path = self._state._get_iteration_path(prepared_iterators[index]) + matching_frames = [ + frame + for frame in selected_frames + if len(branch_path) == len(frame) + 1 and branch_path[: len(frame)] == frame + ] + condition_edges = self._state.graph._get_input_edges(if_node_id, "condition") + condition_iterator_ids = { + ancestor_id + for edge in condition_edges + for ancestor_id in nx.ancestors(self._state._get_source_graph_flat(), edge.source.node_id) + if isinstance(self._state.graph.get_node(ancestor_id), (ForInvocation, IterateInvocation)) + } + condition_iterator_ids.update( + edge.source.node_id + for edge in condition_edges + if isinstance(self._state.graph.get_node(edge.source.node_id), (ForInvocation, IterateInvocation)) + ) + for frame in matching_frames: + if any( + self._state._get_iteration_path(prepared_iterators[iterator_nodes.index(condition_iterator_id)]) + != frame + for condition_iterator_id in condition_iterator_ids + if condition_iterator_id in iterator_nodes + ): + return False + if not matching_frames: + return False + return True + def _get_parent_iteration_mappings( - self, next_node_id: str, graph: "nx.DiGraph", input_edges: Optional[list[Edge]] = None + self, + next_node_id: str, + graph: "nx.DiGraph", + input_edges: Optional[list[Edge]] = None, + iterator_node_ids: Optional[list[str]] = None, ) -> Iterable[list[tuple[str, str]]]: input_edges = input_edges or self._state.graph._get_input_edges(next_node_id) parent_node_ids = list(dict.fromkeys(edge.source.node_id for edge in input_edges)) iterator_graph = self.iterator_graph(graph) - iterator_nodes = self.get_node_iterators(next_node_id, iterator_graph) + iterator_nodes = ( + iterator_node_ids + if iterator_node_ids is not None + else self.get_node_iterators(next_node_id, iterator_graph) + ) + if isinstance(self._state.graph.get_node(next_node_id), CollectInvocation) and len(input_edges) == 1: + source_node = self._state.graph.get_node(input_edges[0].source.node_id) + if isinstance(source_node, IfInvocation): + condition_edges = self._state.graph._get_input_edges(source_node.id, "condition") + condition_iterator_ids = { + ancestor_id + for edge in condition_edges + for ancestor_id in nx.ancestors(iterator_graph, edge.source.node_id) + if isinstance(self._state.graph.get_node(ancestor_id), (ForInvocation, IterateInvocation)) + } + condition_iterator_ids.update( + edge.source.node_id + for edge in condition_edges + if isinstance(self._state.graph.get_node(edge.source.node_id), (ForInvocation, IterateInvocation)) + ) + if condition_iterator_ids: + iterator_nodes = [node_id for node_id in iterator_nodes if node_id in condition_iterator_ids] if not iterator_nodes: return iter(self._get_parent_iteration_mappings_without_iterators(next_node_id, input_edges)) @@ -1374,6 +1441,8 @@ def _get_parent_iteration_mappings( def iter_mappings() -> Iterable[list[tuple[str, str]]]: execution_graph: Optional["nx.DiGraph"] = None for prepared_iterators in itertools.product(*iterator_nodes_prepared): + if not self._matches_sibling_if_collect_context(iterator_nodes, prepared_iterators): + continue mapping: list[tuple[str, str]] = [] for node_id in parent_node_ids: prepared_id = self._get_indexed_iteration_node( @@ -1493,7 +1562,25 @@ def _get_if_condition_iteration_mappings( ) -> Iterable[list[tuple[str, str]]]: condition_edges = self._state.graph._get_input_edges(node_id, "condition") iterator_graph = self.iterator_graph(graph) - iterator_nodes = self.get_node_iterators(node_id, iterator_graph) + if condition_edges: + condition_iterator_ids = { + ancestor_id + for edge in condition_edges + for ancestor_id in nx.ancestors(iterator_graph, edge.source.node_id) + if isinstance(self._state.graph.get_node(ancestor_id), (ForInvocation, IterateInvocation)) + } + condition_iterator_ids.update( + edge.source.node_id + for edge in condition_edges + if isinstance(self._state.graph.get_node(edge.source.node_id), (ForInvocation, IterateInvocation)) + ) + if condition_iterator_ids: + source_order = {source_id: index for index, source_id in enumerate(nx.topological_sort(iterator_graph))} + iterator_nodes = sorted(condition_iterator_ids, key=lambda source_id: source_order[source_id]) + else: + iterator_nodes = self.get_node_iterators(node_id, iterator_graph) + else: + iterator_nodes = self.get_node_iterators(node_id, iterator_graph) if condition_edges and not iterator_nodes: return self._get_parent_iteration_mappings(node_id, graph, input_edges=condition_edges) if not iterator_nodes: @@ -1504,19 +1591,30 @@ def _get_if_condition_iteration_mappings( or nx.has_path(iterator_graph, iterator_nodes[0], edge.source.node_id) for edge in condition_edges ): - return self._get_parent_iteration_mappings(node_id, graph, input_edges=condition_edges) + return self._get_parent_iteration_mappings( + node_id, graph, input_edges=condition_edges, iterator_node_ids=iterator_nodes + ) # This bounded path derives one frame axis. Leave mixed/nested axes on the existing path because # _get_known_iteration_path() intentionally rejects incompatible parent paths. - if len(iterator_nodes) != 1 or any( - not any( - source_edge.source.node_id == iterator_nodes[0] - or nx.has_path(iterator_graph, iterator_nodes[0], source_edge.source.node_id) - for source_edge in self._state.graph._get_input_edges(node_id, branch_field) - ) - for branch_field in ("true_input", "false_input") - ): + if len(iterator_nodes) != 1: return iter([[]]) + source_graph = self._state._get_source_graph_flat() + for branch_field in ("true_input", "false_input"): + branch_edges = self._state.graph._get_input_edges(node_id, branch_field) + branch_iterator_ids = { + ancestor_id + for edge in branch_edges + for ancestor_id in nx.ancestors(source_graph, edge.source.node_id) + if isinstance(self._state.graph.get_node(ancestor_id), (ForInvocation, IterateInvocation)) + } + branch_iterator_ids.update( + edge.source.node_id + for edge in branch_edges + if isinstance(self._state.graph.get_node(edge.source.node_id), (ForInvocation, IterateInvocation)) + ) + if branch_iterator_ids and not (branch_iterator_ids & set(iterator_nodes)): + return iter([]) if any(iterator_node_id not in self._state.source_prepared_mapping for iterator_node_id in iterator_nodes): return iter([]) @@ -1542,9 +1640,72 @@ def _get_if_condition_iteration_mappings( ) def _is_if_condition_ready(self, node_id: str) -> bool: - return all( - edge.source.node_id in self._state.source_prepared_mapping or edge.source.node_id in self._state.executed - for edge in self._state.graph._get_input_edges(node_id, "condition") + for edge in self._state.graph._get_input_edges(node_id, "condition"): + source_node_id = edge.source.node_id + if source_node_id in self._state.executed and not self._state.source_prepared_mapping.get(source_node_id): + continue + prepared_nodes = self._get_ordered_prepared_nodes_for_edge(edge) + if not prepared_nodes or any(prepared_id not in self._state.results for prepared_id in prepared_nodes): + return False + return True + + def _get_if_readiness_ancestors(self, node_id: str, graph: "nx.DiGraph") -> set[str]: + """Get iterator ancestors that must finish before preparing an If execution.""" + condition_edges = self._state.graph._get_input_edges(node_id, "condition") + if not condition_edges: + return set(nx.ancestors(graph, node_id)) + ancestors: set[str] = set() + for edge in condition_edges: + condition_source_id = edge.source.node_id + ancestors.update(nx.ancestors(graph, condition_source_id)) + ancestors.add(condition_source_id) + return ancestors + + def _get_sibling_iterate_if_collect_branch(self, iterate_node_id: str) -> Optional[tuple[str, str]]: + """Return the branch input whose Collect collapses this sibling Iterate axis.""" + for collect_node in self._state.graph.nodes.values(): + if not isinstance(collect_node, CollectInvocation): + continue + item_edges = self._state.graph._get_input_edges(collect_node.id, ITEM_FIELD) + if len(item_edges) != 1: + continue + item_source_id = item_edges[0].source.node_id + iterator_ids = set(self._get_collect_source_iterator_ids(item_source_id)) + if len(iterator_ids) != 2 or iterate_node_id not in iterator_ids: + continue + for output_edge in self._state.graph._get_output_edges(collect_node.id): + destination = self._state.graph.get_node(output_edge.destination.node_id) + if not isinstance(destination, IfInvocation) or output_edge.destination.field not in { + "true_input", + "false_input", + }: + continue + condition_edges = self._state.graph._get_input_edges(destination.id, "condition") + condition_iterator_ids = { + ancestor_id + for edge in condition_edges + for ancestor_id in nx.ancestors(self._state._get_source_graph_flat(), edge.source.node_id) + if isinstance(self._state.graph.get_node(ancestor_id), (ForInvocation, IterateInvocation)) + } + condition_iterator_ids.update( + edge.source.node_id + for edge in condition_edges + if isinstance(self._state.graph.get_node(edge.source.node_id), (ForInvocation, IterateInvocation)) + ) + if condition_iterator_ids & iterator_ids and not (condition_iterator_ids & {iterate_node_id}): + return destination.id, output_edge.destination.field + return None + + def _is_inactive_sibling_if_collect_source(self, node_id: str, graph: "nx.DiGraph") -> bool: + return self._state._is_source_inactive(node_id) and any( + self._get_sibling_iterate_if_collect_branch(iterate_node_id) is not None + and ( + node_id == iterate_node_id + or nx.has_path(graph, node_id, iterate_node_id) + or nx.has_path(graph, iterate_node_id, node_id) + ) + for iterate_node_id, node in self._state.graph.nodes.items() + if isinstance(node, IterateInvocation) ) def _attach_pending_if_inputs(self, *, enqueue: bool = True) -> None: @@ -1740,6 +1901,15 @@ def get_iteration_node( def prepare(self, base_g: Optional["nx.DiGraph"] = None) -> Optional[str]: g = base_g if base_g is not None else self._state._get_source_graph_flat() self._attach_pending_if_inputs() + for node_id in nx.topological_sort(g): + if ( + node_id not in self._state.source_prepared_mapping + and node_id not in self._state.executed + and self._is_inactive_sibling_if_collect_source(node_id, g) + ): + self._mark_source_node_empty(node_id) + self._state._invalidate_loop_caches_for_source(node_id) + return node_id next_node_id = next( ( node_id @@ -1776,7 +1946,11 @@ def prepare(self, base_g: Optional["nx.DiGraph"] = None) -> Optional[str]: and not any( isinstance(self._state.graph.get_node(ancestor_id), (ForInvocation, IterateInvocation)) and ancestor_id not in self._state.executed - for ancestor_id in nx.ancestors(g, node_id) + for ancestor_id in ( + self._get_if_readiness_ancestors(node_id, g) + if isinstance(self._state.graph.get_node(node_id), IfInvocation) + else nx.ancestors(g, node_id) + ) ) ), None, @@ -1831,13 +2005,13 @@ def prepare(self, base_g: Optional["nx.DiGraph"] = None) -> Optional[str]: input_fields={"condition"}, ) new_node_ids.extend(create_results) - else: + elif not self.get_node_iterators(next_node_id): + # Root-frame If nodes may be prepared from their condition while branch sources + # are still being materialized; branch inputs attach after the gate resolves. for iteration_mapping in iteration_mappings: - iteration_path = self._get_known_iteration_path(-1, iteration_mapping, condition_edges) create_results = self.create_execution_node( next_node_id, iteration_mapping, - iteration_path, input_fields={"condition"}, ) new_node_ids.extend(create_results) @@ -1866,13 +2040,40 @@ def prepare(self, base_g: Optional["nx.DiGraph"] = None) -> Optional[str]: key=lambda path: (len(path), path), default=(), ) - create_results = self.create_execution_node( - next_node_id, - iteration_mappings, - iteration_path, - input_fields={"condition"} if isinstance(next_node, IfInvocation) else None, + collapsed_branch = ( + self._get_sibling_iterate_if_collect_branch(next_node_id) + if isinstance(next_node, IterateInvocation) + else None ) - new_node_ids.extend(create_results) + activation_frames = ( + graph_if_runtime._get_selected_if_frames_for_preparation(self._state, *collapsed_branch) + if collapsed_branch is not None + else [] + ) + if activation_frames: + for activation_frame in activation_frames: + new_node_ids.extend( + self.create_execution_node(next_node_id, iteration_mappings, activation_frame) + ) + elif collapsed_branch is not None: + if_node_id, _ = collapsed_branch + if_exec_ids = self._state.source_prepared_mapping.get(if_node_id, set()) + if if_exec_ids and all( + exec_id in self._state._resolved_if_exec_branches for exec_id in if_exec_ids + ): + self._mark_source_node_empty(next_node_id) + self._state._invalidate_loop_caches_for_source(next_node_id) + return next_node_id + # Wait for the owning If to resolve before deriving the branch's parent frames. + return None + else: + create_results = self.create_execution_node( + next_node_id, + iteration_mappings, + iteration_path, + input_fields={"condition"} if isinstance(next_node, IfInvocation) else None, + ) + new_node_ids.extend(create_results) if not new_node_ids: # No parent mappings means zero loop contexts, unlike a context with an empty collection. diff --git a/tests/app/services/shared/test_execution_engine_differential.py b/tests/app/services/shared/test_execution_engine_differential.py index 741a0c0dbfd..f6e2d073e87 100644 --- a/tests/app/services/shared/test_execution_engine_differential.py +++ b/tests/app/services/shared/test_execution_engine_differential.py @@ -8566,3 +8566,78 @@ def test_iterate_for_collect_preserves_empty_inner_collection(force_compatibilit assert len(after_exec_ids) == 1 assert state.results[after_exec_ids[0]].value == [[11, 12], []] assert state.is_complete() + + +@pytest.mark.parametrize( + ("left_values", "expected"), + [([1, 2], ["fallback", [12, 22]]), ([1, 1], ["fallback", "fallback"])], + ids=["mixed", "branch-never-selected"], +) +def test_sibling_iterates_collect_through_if_branch_without_unsupported_frames(left_values, expected): + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="left_source", first=left_values)) + graph.add_node(CollectionConcatInvocation(id="right_source", first=[10, 20])) + graph.add_node(IterateInvocation(id="left")) + graph.add_node(IterateInvocation(id="right")) + graph.add_node(AddInvocation(id="body")) + graph.add_node(CollectInvocation(id="branch_collect")) + graph.add_node(AddInvocation(id="minus_one", b=-1)) + graph.add_node(FixedCollectionTestInvocation(id="condition")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value="fallback")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="collect")) + for source, source_field, destination, destination_field in ( + ("left_source", "collection", "left", "collection"), + ("right_source", "collection", "right", "collection"), + ("left", "item", "body", "a"), + ("right", "item", "body", "b"), + ("body", "value", "branch_collect", "item"), + ("left", "item", "minus_one", "a"), + ("minus_one", "value", "condition", "value"), + ("condition", "condition", "if", "condition"), + ("branch_collect", "collection", "if", "true_input"), + ("fallback", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + + trace, state = _run_to_completion(GraphExecutionState(graph=graph)) + + assert state.is_complete() + assert state.results[next(iter(state.source_prepared_mapping["collect"]))].collection == expected + if left_values == [1, 1]: + assert "right" not in trace + assert "body" not in trace + + +def test_if_does_not_combine_unrelated_open_iterator_frames(): + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="branch_source", first=[1, 2])) + graph.add_node(IterateInvocation(id="branch_iterate")) + graph.add_node(AnyTypeTestInvocation(id="branch_value")) + graph.add_node(CollectionConcatInvocation(id="condition_source", first=[True])) + graph.add_node(IterateInvocation(id="condition_iterate")) + graph.add_node(FixedCollectionTestInvocation(id="condition")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value="fallback")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="collect")) + for source, source_field, destination, destination_field in ( + ("branch_source", "collection", "branch_iterate", "collection"), + ("branch_iterate", "item", "branch_value", "value"), + ("condition_source", "collection", "condition_iterate", "collection"), + ("condition_iterate", "item", "condition", "value"), + ("condition", "condition", "if", "condition"), + ("branch_value", "value", "if", "true_input"), + ("fallback", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + + graph.validate_self() + trace, state = _run_to_completion(GraphExecutionState(graph=graph)) + assert not state.is_complete() + assert not state.errors + assert "branch_value" not in state.source_prepared_mapping + assert "collect" not in state.source_prepared_mapping + assert "branch_value" not in trace + assert "collect" not in trace From fd6b65d00bae366fe4b239ed5dbf1e677cdd1572 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Tue, 13 Oct 2026 10:48:32 -0500 Subject: [PATCH 06/20] fix(graph): resolve deep If dependencies iteratively Use iterative path traversal and refresh materializer caches after graph edits. Add deep-chain and cache-lifecycle regressions. --- invokeai/app/services/shared/graph.py | 1 + .../app/services/shared/graph_if_runtime.py | 82 +++++++------- .../app/services/shared/graph_materializer.py | 27 +++-- .../test_execution_engine_differential.py | 106 ++++++++++++++++++ 4 files changed, 169 insertions(+), 47 deletions(-) diff --git a/invokeai/app/services/shared/graph.py b/invokeai/app/services/shared/graph.py index c7a88f5e7d7..99e038cb468 100644 --- a/invokeai/app/services/shared/graph.py +++ b/invokeai/app/services/shared/graph.py @@ -729,6 +729,7 @@ def _invalidate_source_graph_cache(self) -> None: runtime.if_activation_dependencies_by_source = {} runtime.if_activation_dependencies_by_exec = {} runtime.if_activation_controller_instance = None + runtime.execution_materializer = None def _materializer(self) -> _ExecutionMaterializer: runtime = self.__pydantic_private__["_runtime_state"] diff --git a/invokeai/app/services/shared/graph_if_runtime.py b/invokeai/app/services/shared/graph_if_runtime.py index 8fd1d633fa5..69087ef382d 100644 --- a/invokeai/app/services/shared/graph_if_runtime.py +++ b/invokeai/app/services/shared/graph_if_runtime.py @@ -20,47 +20,51 @@ def _get_effective_iteration_paths_for_edge( get_prepared_edge_iteration_path: Callable[[Edge, str], tuple[int, ...]], source_graph: Optional["nx.DiGraph"] = None, ) -> set[tuple[int, ...]]: - source_node_id = edge.source.node_id - prepared_nodes: set[str] = set() - if source_node_id in state.source_prepared_mapping: - prepared_nodes = get_prepared_nodes_for_source(source_node_id) - if not prepared_nodes and source_node_id in state.executed: - return set() - if prepared_nodes: - return {get_prepared_edge_iteration_path(edge, prepared_id) for prepared_id in prepared_nodes} - if source_node_id in visited_source_ids: - return set() - - source_node = state.graph.get_node(source_node_id) - if isinstance(source_node, IfInvocation): - return set( - _get_if_iteration_paths( - state, - source_node_id, - visited_source_ids, - get_prepared_nodes_for_source, - get_prepared_edge_iteration_path, - source_graph, - ) - ) + pending: list[tuple[Optional[Edge], Optional[str]]] = [(edge, None)] + active_source_ids = set(visited_source_ids) + paths: set[tuple[int, ...]] = set() - input_edges = state.graph._get_input_edges(source_node_id) - if not input_edges: - return {()} - next_visited_source_ids = {*visited_source_ids, source_node_id} - return set().union( - *( - _get_effective_iteration_paths_for_edge( - state, - input_edge, - next_visited_source_ids, - get_prepared_nodes_for_source, - get_prepared_edge_iteration_path, - source_graph, + while pending: + current_edge, exiting_source_id = pending.pop() + if exiting_source_id is not None: + active_source_ids.remove(exiting_source_id) + continue + assert current_edge is not None + current_source_id = current_edge.source.node_id + prepared_nodes: set[str] = set() + if current_source_id in state.source_prepared_mapping: + prepared_nodes = get_prepared_nodes_for_source(current_source_id) + if not prepared_nodes and current_source_id in state.executed: + continue + if prepared_nodes: + paths.update(get_prepared_edge_iteration_path(current_edge, prepared_id) for prepared_id in prepared_nodes) + continue + if current_source_id in active_source_ids: + continue + + source_node = state.graph.get_node(current_source_id) + if isinstance(source_node, IfInvocation): + paths.update( + _get_if_iteration_paths( + state, + current_source_id, + active_source_ids, + get_prepared_nodes_for_source, + get_prepared_edge_iteration_path, + source_graph, + ) ) - for input_edge in input_edges - ) - ) + continue + + input_edges = state.graph._get_input_edges(current_source_id) + if not input_edges: + paths.add(()) + continue + active_source_ids.add(current_source_id) + pending.append((None, current_source_id)) + pending.extend((input_edge, None) for input_edge in reversed(input_edges)) + + return paths def _get_if_iteration_paths( diff --git a/invokeai/app/services/shared/graph_materializer.py b/invokeai/app/services/shared/graph_materializer.py index 9526005a1f5..abf62e3a361 100644 --- a/invokeai/app/services/shared/graph_materializer.py +++ b/invokeai/app/services/shared/graph_materializer.py @@ -43,6 +43,7 @@ class _ExecutionNodeBuilder: def __init__(self, state: "GraphExecutionState") -> None: self._state = state self._iteration_axes_by_source: dict[str, tuple[str, ...]] = {} + self._sibling_iterate_if_collect_ids: Optional[tuple[str, ...]] = None def _get_iteration_axes(self, source_node_id: str) -> tuple[str, ...]: cached = self._iteration_axes_by_source.get(source_node_id) @@ -1697,15 +1698,25 @@ def _get_sibling_iterate_if_collect_branch(self, iterate_node_id: str) -> Option return None def _is_inactive_sibling_if_collect_source(self, node_id: str, graph: "nx.DiGraph") -> bool: - return self._state._is_source_inactive(node_id) and any( - self._get_sibling_iterate_if_collect_branch(iterate_node_id) is not None - and ( - node_id == iterate_node_id - or nx.has_path(graph, node_id, iterate_node_id) - or nx.has_path(graph, iterate_node_id, node_id) + if self._sibling_iterate_if_collect_ids is None: + self._sibling_iterate_if_collect_ids = tuple( + iterate_node_id + for iterate_node_id, node in self._state.graph.nodes.items() + if isinstance(node, IterateInvocation) + and self._get_sibling_iterate_if_collect_branch(iterate_node_id) is not None + ) + + return ( + bool(self._sibling_iterate_if_collect_ids) + and self._state._is_source_inactive(node_id) + and any( + ( + node_id == iterate_node_id + or nx.has_path(graph, node_id, iterate_node_id) + or nx.has_path(graph, iterate_node_id, node_id) + ) + for iterate_node_id in self._sibling_iterate_if_collect_ids ) - for iterate_node_id, node in self._state.graph.nodes.items() - if isinstance(node, IterateInvocation) ) def _attach_pending_if_inputs(self, *, enqueue: bool = True) -> None: diff --git a/tests/app/services/shared/test_execution_engine_differential.py b/tests/app/services/shared/test_execution_engine_differential.py index f6e2d073e87..9548221e328 100644 --- a/tests/app/services/shared/test_execution_engine_differential.py +++ b/tests/app/services/shared/test_execution_engine_differential.py @@ -37,6 +37,7 @@ _ExecutionScheduler, _GenericGraphSchedulerAdapter, ) +from invokeai.app.services.shared.graph_if_runtime import _get_effective_iteration_paths_for_edge from invokeai.app.services.shared.graph_nested_iterate_planner import get_nested_iterate_sequence_depth from invokeai.app.services.shared.invocation_context import InvocationContextData, build_invocation_context from tests.test_nodes import ( @@ -8641,3 +8642,108 @@ def test_if_does_not_combine_unrelated_open_iterator_frames(): assert "collect" not in state.source_prepared_mapping assert "branch_value" not in trace assert "collect" not in trace + + +def test_deep_any_chain_inside_iterate_if_completes() -> None: + chain_length = 400 + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="source", first=[True, False])) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(MarkedAnyTypeTestInvocation(id="false_branch", marker="f")) + for index in range(chain_length): + graph.add_node(AnyTypeTestInvocation(id=f"chain_{index}")) + graph.add_edge(create_edge("source", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "if", "condition")) + graph.add_edge(create_edge("iterate", "item", "chain_0", "value")) + graph.add_edge(create_edge("iterate", "item", "false_branch", "value")) + graph.add_edge(create_edge("false_branch", "value", "if", "false_input")) + for index in range(chain_length - 1): + node_id = f"chain_{index}" + next_id = f"chain_{index + 1}" + graph.add_edge(create_edge(node_id, "value", next_id, "value")) + graph.add_edge(create_edge(f"chain_{chain_length - 1}", "value", "if", "true_input")) + graph.add_edge(create_edge("if", "value", "collect", "item")) + + _, state = _run_to_completion(GraphExecutionState(graph=graph)) + + assert state.results[next(iter(state.source_prepared_mapping["collect"]))].collection == [True, "fFalse"] + assert state.is_complete() + + +def test_execution_materializer_cache_refreshes_after_sibling_iterate_edge_addition() -> None: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="left_source", first=[1, 1])) + graph.add_node(CollectionConcatInvocation(id="right_source", first=[10, 20])) + graph.add_node(IterateInvocation(id="left")) + graph.add_node(IterateInvocation(id="right")) + graph.add_node(AddInvocation(id="body")) + graph.add_node(CollectInvocation(id="branch_collect")) + graph.add_node(AddInvocation(id="minus_one", b=-1)) + graph.add_node(FixedCollectionTestInvocation(id="condition")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value="fallback")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="collect")) + for source, source_field, destination, destination_field in ( + ("left_source", "collection", "left", "collection"), + ("right_source", "collection", "right", "collection"), + ("left", "item", "body", "a"), + ("body", "value", "branch_collect", "item"), + ("left", "item", "minus_one", "a"), + ("minus_one", "value", "condition", "value"), + ("condition", "condition", "if", "condition"), + ("branch_collect", "collection", "if", "true_input"), + ("fallback", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + + state = GraphExecutionState(graph=graph) + original_materializer = state._materializer() + source_graph = state._get_source_graph_flat() + assert not original_materializer._is_inactive_sibling_if_collect_source("right", source_graph) + assert original_materializer._sibling_iterate_if_collect_ids == () + + state.add_edge(create_edge("right", "item", "body", "b")) + + materializer = state._materializer() + assert materializer is not original_materializer + assert materializer._get_sibling_iterate_if_collect_branch("right") == ("if", "true_input") + + trace, state = _run_to_completion(state) + + assert state.results[next(iter(state.source_prepared_mapping["collect"]))].collection == ["fallback", "fallback"] + assert "right" not in trace + assert "body" not in trace + assert state.is_complete() + + +def test_effective_iteration_paths_handles_deep_any_chain() -> None: + chain_length = 1000 + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="source", first=[True])) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(IfInvocation(id="if")) + for index in range(chain_length): + graph.add_node(AnyTypeTestInvocation(id=f"chain_{index}")) + graph.add_edge(create_edge("source", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "chain_0", "value")) + for index in range(chain_length - 1): + graph.add_edge(create_edge(f"chain_{index}", "value", f"chain_{index + 1}", "value")) + terminal_edge = create_edge(f"chain_{chain_length - 1}", "value", "if", "true_input") + graph.add_edge(terminal_edge) + + state = GraphExecutionState(graph=graph) + state.source_prepared_mapping["iterate"] = {"prepared_iterate"} + state.prepared_iteration_paths["prepared_iterate"] = (7,) + + effective_paths = _get_effective_iteration_paths_for_edge( + state, + terminal_edge, + set(), + lambda source_id: state.source_prepared_mapping.get(source_id, set()), + lambda _edge, prepared_id: state.prepared_iteration_paths[prepared_id], + ) + + assert effective_paths == {(7,)} From 5b491955a3e32a582af33f921cf496428f122755 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Tue, 13 Oct 2026 12:20:22 -0500 Subject: [PATCH 07/20] fix(graph): preserve parent-scoped collect outputs --- invokeai/app/services/shared/graph_materializer.py | 3 --- .../services/shared/test_execution_engine_differential.py | 5 +++-- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/invokeai/app/services/shared/graph_materializer.py b/invokeai/app/services/shared/graph_materializer.py index abf62e3a361..867f195e32a 100644 --- a/invokeai/app/services/shared/graph_materializer.py +++ b/invokeai/app/services/shared/graph_materializer.py @@ -1060,9 +1060,6 @@ def _get_collect_iteration_group_key(self, edge: Edge, sibling_depth: Optional[i and get_output_field_scope(source_node, edge.source.field) == OutputScope.Final ): parent_path = self._state._get_for_parent_iteration_path(edge.source.node_id) - if edge.destination.field == ITEM_FIELD: - group_depth = len(parent_path) if sibling_depth is None else sibling_depth - return parent_path[: max(group_depth - 1, 0)] return parent_path if edge.destination.field == ITEM_FIELD: if isinstance(source_node, ForInvocation) and source_node.index == -1: diff --git a/tests/app/services/shared/test_execution_engine_differential.py b/tests/app/services/shared/test_execution_engine_differential.py index 9548221e328..962afa5d9cf 100644 --- a/tests/app/services/shared/test_execution_engine_differential.py +++ b/tests/app/services/shared/test_execution_engine_differential.py @@ -8564,8 +8564,9 @@ def test_iterate_for_collect_preserves_empty_inner_collection(force_compatibilit after_exec_ids = sorted(state.source_prepared_mapping["after"], key=state._get_iteration_path) assert trace.count("add") == 2 - assert len(after_exec_ids) == 1 - assert state.results[after_exec_ids[0]].value == [[11, 12], []] + assert trace.count("after") == 2 + assert [state._get_iteration_path(exec_id) for exec_id in after_exec_ids] == [(0,), (1,)] + assert [state.results[exec_id].value for exec_id in after_exec_ids] == [[[11, 12]], [[]]] assert state.is_complete() From 10c74c6df8c30a7bf3c9e96aa95b1b0f468f9d1c Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Tue, 13 Oct 2026 12:56:24 -0500 Subject: [PATCH 08/20] fix(graph): align sibling Iterate frames with If --- .../app/services/shared/graph_materializer.py | 103 +++++++++++------- .../test_execution_engine_differential.py | 103 +++++++++++++++--- 2 files changed, 155 insertions(+), 51 deletions(-) diff --git a/invokeai/app/services/shared/graph_materializer.py b/invokeai/app/services/shared/graph_materializer.py index 867f195e32a..0c83826cae6 100644 --- a/invokeai/app/services/shared/graph_materializer.py +++ b/invokeai/app/services/shared/graph_materializer.py @@ -43,7 +43,7 @@ class _ExecutionNodeBuilder: def __init__(self, state: "GraphExecutionState") -> None: self._state = state self._iteration_axes_by_source: dict[str, tuple[str, ...]] = {} - self._sibling_iterate_if_collect_ids: Optional[tuple[str, ...]] = None + self._sibling_iterate_if_collect_branches: Optional[dict[str, tuple[str, str]]] = None def _get_iteration_axes(self, source_node_id: str) -> tuple[str, ...]: cached = self._iteration_axes_by_source.get(source_node_id) @@ -52,6 +52,16 @@ def _get_iteration_axes(self, source_node_id: str) -> tuple[str, ...]: axes = self._state._runtime()._get_ordered_iterator_sources(source_node_id) if isinstance(self._state.graph.get_node(source_node_id), IterateInvocation): + collapsed_branch = self._get_sibling_iterate_if_collect_branch(source_node_id) + if collapsed_branch is not None: + if_node_id, _ = collapsed_branch + condition_axes: list[str] = [] + for edge in self._state.graph._get_input_edges(if_node_id, "condition"): + condition_source_id = edge.source.node_id + condition_axes.extend(self._state._runtime()._get_ordered_iterator_sources(condition_source_id)) + if isinstance(self._state.graph.get_node(condition_source_id), (ForInvocation, IterateInvocation)): + condition_axes.append(condition_source_id) + axes = list(dict.fromkeys([*condition_axes, *axes])) axes.append(source_node_id) result = tuple(axes) self._iteration_axes_by_source[source_node_id] = result @@ -1069,8 +1079,10 @@ def _get_collect_iteration_group_key(self, edge: Edge, sibling_depth: Optional[i return path[: max(depth - 1, 0)] return path - def _get_collect_source_iterator_ids(self, source_node_id: str) -> list[str]: - iterator_node_ids = self.get_node_iterators(source_node_id) + def _get_collect_source_iterator_ids( + self, source_node_id: str, iterator_graph: Optional["nx.DiGraph"] = None + ) -> list[str]: + iterator_node_ids = self.get_node_iterators(source_node_id, iterator_graph) if isinstance(self._state.graph.get_node(source_node_id), IterateInvocation): iterator_node_ids.append(source_node_id) return iterator_node_ids @@ -1659,18 +1671,28 @@ def _get_if_readiness_ancestors(self, node_id: str, graph: "nx.DiGraph") -> set[ ancestors.add(condition_source_id) return ancestors - def _get_sibling_iterate_if_collect_branch(self, iterate_node_id: str) -> Optional[tuple[str, str]]: - """Return the branch input whose Collect collapses this sibling Iterate axis.""" + def _get_sibling_iterate_if_collect_branches(self) -> dict[str, tuple[str, str]]: + """Find Iterate axes collapsed by branch Collects under an If.""" + if self._sibling_iterate_if_collect_branches is not None: + return self._sibling_iterate_if_collect_branches + + branches: dict[str, tuple[str, str]] = {} + if not any(isinstance(node, CollectInvocation) for node in self._state.graph.nodes.values()): + self._sibling_iterate_if_collect_branches = branches + return branches + + graph = self._state._get_source_graph_flat() + iterator_graph = self.iterator_graph(graph) for collect_node in self._state.graph.nodes.values(): if not isinstance(collect_node, CollectInvocation): continue item_edges = self._state.graph._get_input_edges(collect_node.id, ITEM_FIELD) if len(item_edges) != 1: continue - item_source_id = item_edges[0].source.node_id - iterator_ids = set(self._get_collect_source_iterator_ids(item_source_id)) - if len(iterator_ids) != 2 or iterate_node_id not in iterator_ids: + iterator_ids = set(self._get_collect_source_iterator_ids(item_edges[0].source.node_id, iterator_graph)) + if len(iterator_ids) != 2: continue + for output_edge in self._state.graph._get_output_edges(collect_node.id): destination = self._state.graph.get_node(output_edge.destination.node_id) if not isinstance(destination, IfInvocation) or output_edge.destination.field not in { @@ -1678,11 +1700,12 @@ def _get_sibling_iterate_if_collect_branch(self, iterate_node_id: str) -> Option "false_input", }: continue + condition_edges = self._state.graph._get_input_edges(destination.id, "condition") condition_iterator_ids = { ancestor_id for edge in condition_edges - for ancestor_id in nx.ancestors(self._state._get_source_graph_flat(), edge.source.node_id) + for ancestor_id in nx.ancestors(graph, edge.source.node_id) if isinstance(self._state.graph.get_node(ancestor_id), (ForInvocation, IterateInvocation)) } condition_iterator_ids.update( @@ -1690,29 +1713,34 @@ def _get_sibling_iterate_if_collect_branch(self, iterate_node_id: str) -> Option for edge in condition_edges if isinstance(self._state.graph.get_node(edge.source.node_id), (ForInvocation, IterateInvocation)) ) - if condition_iterator_ids & iterator_ids and not (condition_iterator_ids & {iterate_node_id}): - return destination.id, output_edge.destination.field - return None + if not condition_iterator_ids.intersection(iterator_ids): + continue - def _is_inactive_sibling_if_collect_source(self, node_id: str, graph: "nx.DiGraph") -> bool: - if self._sibling_iterate_if_collect_ids is None: - self._sibling_iterate_if_collect_ids = tuple( - iterate_node_id - for iterate_node_id, node in self._state.graph.nodes.items() - if isinstance(node, IterateInvocation) - and self._get_sibling_iterate_if_collect_branch(iterate_node_id) is not None - ) + for iterate_node_id in iterator_ids - condition_iterator_ids: + if any( + nx.has_path(iterator_graph, condition_iterator_id, iterate_node_id) + or nx.has_path(iterator_graph, iterate_node_id, condition_iterator_id) + for condition_iterator_id in condition_iterator_ids + ): + continue + branches.setdefault(iterate_node_id, (destination.id, output_edge.destination.field)) + + self._sibling_iterate_if_collect_branches = branches + return branches + def _get_sibling_iterate_if_collect_branch(self, iterate_node_id: str) -> Optional[tuple[str, str]]: + return self._get_sibling_iterate_if_collect_branches().get(iterate_node_id) + + def _is_inactive_sibling_if_collect_source(self, node_id: str, graph: "nx.DiGraph") -> bool: + sibling_iterate_ids = self._get_sibling_iterate_if_collect_branches().keys() return ( - bool(self._sibling_iterate_if_collect_ids) + bool(sibling_iterate_ids) and self._state._is_source_inactive(node_id) and any( - ( - node_id == iterate_node_id - or nx.has_path(graph, node_id, iterate_node_id) - or nx.has_path(graph, iterate_node_id, node_id) - ) - for iterate_node_id in self._sibling_iterate_if_collect_ids + node_id == iterate_node_id + or nx.has_path(graph, node_id, iterate_node_id) + or nx.has_path(graph, iterate_node_id, node_id) + for iterate_node_id in sibling_iterate_ids ) ) @@ -1806,7 +1834,7 @@ def iterator_graph(self, base: Optional["nx.DiGraph"] = None) -> "nx.DiGraph": return g def get_node_iterators(self, node_id: str, it_graph: Optional["nx.DiGraph"] = None) -> list[str]: - g = it_graph or self.iterator_graph() + g = it_graph if it_graph is not None else self.iterator_graph() return [ n for n in nx.ancestors(g, node_id) @@ -1909,15 +1937,16 @@ def get_iteration_node( def prepare(self, base_g: Optional["nx.DiGraph"] = None) -> Optional[str]: g = base_g if base_g is not None else self._state._get_source_graph_flat() self._attach_pending_if_inputs() - for node_id in nx.topological_sort(g): - if ( - node_id not in self._state.source_prepared_mapping - and node_id not in self._state.executed - and self._is_inactive_sibling_if_collect_source(node_id, g) - ): - self._mark_source_node_empty(node_id) - self._state._invalidate_loop_caches_for_source(node_id) - return node_id + if self._get_sibling_iterate_if_collect_branches(): + for node_id in nx.topological_sort(g): + if ( + node_id not in self._state.source_prepared_mapping + and node_id not in self._state.executed + and self._is_inactive_sibling_if_collect_source(node_id, g) + ): + self._mark_source_node_empty(node_id) + self._state._invalidate_loop_caches_for_source(node_id) + return node_id next_node_id = next( ( node_id diff --git a/tests/app/services/shared/test_execution_engine_differential.py b/tests/app/services/shared/test_execution_engine_differential.py index 962afa5d9cf..f251548e5cf 100644 --- a/tests/app/services/shared/test_execution_engine_differential.py +++ b/tests/app/services/shared/test_execution_engine_differential.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +import random from itertools import product from pathlib import Path from typing import Any, Callable @@ -8570,24 +8571,56 @@ def test_iterate_for_collect_preserves_empty_inner_collection(force_compatibilit assert state.is_complete() +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) @pytest.mark.parametrize( ("left_values", "expected"), [([1, 2], ["fallback", [12, 22]]), ([1, 1], ["fallback", "fallback"])], ids=["mixed", "branch-never-selected"], ) -def test_sibling_iterates_collect_through_if_branch_without_unsupported_frames(left_values, expected): +@pytest.mark.parametrize( + "node_order", + [ + tuple(order) + for order in ( + random.Random(seed).sample( + ( + "left_source", + "right_source", + "left", + "right", + "body", + "branch_collect", + "minus_one", + "condition", + "fallback", + "if", + "collect", + ), + 11, + ) + for seed in range(20) + ) + ], +) +def test_sibling_iterates_collect_through_if_branch_without_unsupported_frames( + left_values, expected, node_order, force_compatibility_scheduler +): graph = Graph() - graph.add_node(CollectionConcatInvocation(id="left_source", first=left_values)) - graph.add_node(CollectionConcatInvocation(id="right_source", first=[10, 20])) - graph.add_node(IterateInvocation(id="left")) - graph.add_node(IterateInvocation(id="right")) - graph.add_node(AddInvocation(id="body")) - graph.add_node(CollectInvocation(id="branch_collect")) - graph.add_node(AddInvocation(id="minus_one", b=-1)) - graph.add_node(FixedCollectionTestInvocation(id="condition")) - graph.add_node(AnyTypeTestInvocation(id="fallback", value="fallback")) - graph.add_node(IfInvocation(id="if")) - graph.add_node(CollectInvocation(id="collect")) + nodes = { + "left_source": CollectionConcatInvocation(id="left_source", first=left_values), + "right_source": CollectionConcatInvocation(id="right_source", first=[10, 20]), + "left": IterateInvocation(id="left"), + "right": IterateInvocation(id="right"), + "body": AddInvocation(id="body"), + "branch_collect": CollectInvocation(id="branch_collect"), + "minus_one": AddInvocation(id="minus_one", b=-1), + "condition": FixedCollectionTestInvocation(id="condition"), + "fallback": AnyTypeTestInvocation(id="fallback", value="fallback"), + "if": IfInvocation(id="if"), + "collect": CollectInvocation(id="collect"), + } + for node_id in node_order: + graph.add_node(nodes[node_id]) for source, source_field, destination, destination_field in ( ("left_source", "collection", "left", "collection"), ("right_source", "collection", "right", "collection"), @@ -8603,7 +8636,10 @@ def test_sibling_iterates_collect_through_if_branch_without_unsupported_frames(l ): graph.add_edge(create_edge(source, source_field, destination, destination_field)) - trace, state = _run_to_completion(GraphExecutionState(graph=graph)) + state = GraphExecutionState(graph=graph) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) + trace, state = _run_to_completion(state) assert state.is_complete() assert state.results[next(iter(state.source_prepared_mapping["collect"]))].collection == expected @@ -8612,6 +8648,45 @@ def test_sibling_iterates_collect_through_if_branch_without_unsupported_frames(l assert "body" not in trace +def test_nested_iterate_collected_in_if_branch_is_not_misclassified_as_sibling(): + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="outer_source", first=[["a", "b"], ["c"]])) + graph.add_node(IterateInvocation(id="outer")) + graph.add_node(PolymorphicStringTestInvocation(id="inner_collection")) + graph.add_node(IterateInvocation(id="inner")) + graph.add_node(AnyTypeTestInvocation(id="body")) + graph.add_node(CollectInvocation(id="branch_collect")) + graph.add_node(FixedCollectionTestInvocation(id="condition")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value="fallback")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="result_collect")) + for source, source_field, destination, destination_field in ( + ("outer_source", "collection", "outer", "collection"), + ("outer", "item", "inner_collection", "value"), + ("inner_collection", "collection", "inner", "collection"), + ("inner", "item", "body", "value"), + ("body", "value", "branch_collect", "item"), + ("outer", "item", "condition", "value"), + ("condition", "condition", "if", "condition"), + ("branch_collect", "collection", "if", "true_input"), + ("fallback", "value", "if", "false_input"), + ("if", "value", "result_collect", "item"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + + state = GraphExecutionState(graph=graph) + + assert state._materializer()._get_sibling_iterate_if_collect_branch("inner") is None + + trace, state = _run_to_completion(state) + + assert state.is_complete() + assert state.results[next(iter(state.source_prepared_mapping["result_collect"]))].collection == [ + ["a", "b"], + ["c"], + ] + + def test_if_does_not_combine_unrelated_open_iterator_frames(): graph = Graph() graph.add_node(CollectionConcatInvocation(id="branch_source", first=[1, 2])) @@ -8704,7 +8779,7 @@ def test_execution_materializer_cache_refreshes_after_sibling_iterate_edge_addit original_materializer = state._materializer() source_graph = state._get_source_graph_flat() assert not original_materializer._is_inactive_sibling_if_collect_source("right", source_graph) - assert original_materializer._sibling_iterate_if_collect_ids == () + assert original_materializer._get_sibling_iterate_if_collect_branch("right") is None state.add_edge(create_edge("right", "item", "body", "b")) From 1fdb4b18606608e4cc2b274ade44ef1da6cfcb61 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Tue, 13 Oct 2026 13:59:48 -0500 Subject: [PATCH 09/20] fix(graph): preserve condition iterator ancestry across Collect Order complete If condition dependencies before unrelated iterator axes. Invalidate the cached order when source graph edits change dependencies. Cover chained and direct Collect scopes, renamed condition and producer IDs, and graph edits after cache warmup. --- invokeai/app/services/shared/graph.py | 1 + .../shared/graph_execution_runtime.py | 28 ++- .../app/services/shared/graph_materializer.py | 60 +++++- .../test_execution_engine_differential.py | 203 +++++++++++++++++- 4 files changed, 284 insertions(+), 8 deletions(-) diff --git a/invokeai/app/services/shared/graph.py b/invokeai/app/services/shared/graph.py index 99e038cb468..c51e556320e 100644 --- a/invokeai/app/services/shared/graph.py +++ b/invokeai/app/services/shared/graph.py @@ -730,6 +730,7 @@ def _invalidate_source_graph_cache(self) -> None: runtime.if_activation_dependencies_by_exec = {} runtime.if_activation_controller_instance = None runtime.execution_materializer = None + runtime.execution_runtime = None def _materializer(self) -> _ExecutionMaterializer: runtime = self.__pydantic_private__["_runtime_state"] diff --git a/invokeai/app/services/shared/graph_execution_runtime.py b/invokeai/app/services/shared/graph_execution_runtime.py index ec7a6fdbca2..62042280288 100644 --- a/invokeai/app/services/shared/graph_execution_runtime.py +++ b/invokeai/app/services/shared/graph_execution_runtime.py @@ -41,6 +41,31 @@ class _ExecutionRuntime: def __init__(self, state: "GraphExecutionState") -> None: self._state = state + self._iterator_source_order: Optional[dict[str, int]] = None + + def _get_iterator_source_order(self) -> dict[str, int]: + if self._iterator_source_order is not None: + return self._iterator_source_order + + iterator_graph = self._state._iterator_graph(self._state.graph.nx_graph()) + graph = self._state.graph.nx_graph() + condition_ancestry_ids: set[str] = set() + for node in self._state.graph.nodes.values(): + if not isinstance(node, IfInvocation): + continue + for edge in self._state.graph._get_input_edges(node.id, "condition"): + source_id = edge.source.node_id + condition_ancestry_ids.update(nx.ancestors(graph, source_id)) + condition_ancestry_ids.add(source_id) + + topo = list( + nx.lexicographical_topological_sort( + iterator_graph, + key=lambda node_id: (node_id not in condition_ancestry_ids, str(node_id)), + ) + ) + self._iterator_source_order = {node_id: index for index, node_id in enumerate(topo)} + return self._iterator_source_order def _get_cached_iteration_path(self, exec_node_id: str) -> Optional[tuple[int, ...]]: return self._state._prepared_registry().get_iteration_path(exec_node_id) @@ -58,8 +83,7 @@ def _get_ordered_iterator_sources(self, source_node_id: str) -> list[str]: if isinstance(self._state.graph.get_node(node_id), (ForInvocation, IterateInvocation)) ] - topo = list(nx.topological_sort(iterator_graph)) - topo_index = {node_id: i for i, node_id in enumerate(topo)} + topo_index = self._get_iterator_source_order() iterator_sources.sort(key=lambda node_id: topo_index.get(node_id, 0)) return iterator_sources diff --git a/invokeai/app/services/shared/graph_materializer.py b/invokeai/app/services/shared/graph_materializer.py index 0c83826cae6..10ddea96bac 100644 --- a/invokeai/app/services/shared/graph_materializer.py +++ b/invokeai/app/services/shared/graph_materializer.py @@ -44,6 +44,8 @@ def __init__(self, state: "GraphExecutionState") -> None: self._state = state self._iteration_axes_by_source: dict[str, tuple[str, ...]] = {} self._sibling_iterate_if_collect_branches: Optional[dict[str, tuple[str, str]]] = None + self._iteration_relationship_graph: Optional[nx.DiGraph] = None + self._iteration_relationship_topological_index: Optional[dict[str, int]] = None def _get_iteration_axes(self, source_node_id: str) -> tuple[str, ...]: cached = self._iteration_axes_by_source.get(source_node_id) @@ -62,6 +64,16 @@ def _get_iteration_axes(self, source_node_id: str) -> tuple[str, ...]: if isinstance(self._state.graph.get_node(condition_source_id), (ForInvocation, IterateInvocation)): condition_axes.append(condition_source_id) axes = list(dict.fromkeys([*condition_axes, *axes])) + if self._iteration_relationship_graph is not None: + relationship_ancestors = [ + iterator_id + for iterator_id in nx.ancestors(self._iteration_relationship_graph, source_node_id) + if isinstance(self._state.graph.get_node(iterator_id), (ForInvocation, IterateInvocation)) + ] + relationship_order = self._iteration_relationship_topological_index + assert relationship_order is not None + relationship_ancestors.sort(key=relationship_order.__getitem__) + axes = list(dict.fromkeys([*relationship_ancestors, *axes])) axes.append(source_node_id) result = tuple(axes) self._iteration_axes_by_source[source_node_id] = result @@ -1585,7 +1597,7 @@ def _get_if_condition_iteration_mappings( if isinstance(self._state.graph.get_node(edge.source.node_id), (ForInvocation, IterateInvocation)) ) if condition_iterator_ids: - source_order = {source_id: index for index, source_id in enumerate(nx.topological_sort(iterator_graph))} + source_order = self._state._runtime()._get_iterator_source_order() iterator_nodes = sorted(condition_iterator_ids, key=lambda source_id: source_order[source_id]) else: iterator_nodes = self.get_node_iterators(node_id, iterator_graph) @@ -1683,6 +1695,35 @@ def _get_sibling_iterate_if_collect_branches(self) -> dict[str, tuple[str, str]] graph = self._state._get_source_graph_flat() iterator_graph = self.iterator_graph(graph) + iterator_topological_index = self._state._runtime()._get_iterator_source_order() + relationship_graph = iterator_graph.copy() + for collect_node in self._state.graph.nodes.values(): + if not isinstance(collect_node, CollectInvocation): + continue + for collection_edge in self._state.graph._get_input_edges(collect_node.id, COLLECTION_FIELD): + relationship_graph.add_edge(collection_edge.source.node_id, collect_node.id) + item_edges = self._state.graph._get_input_edges(collect_node.id, ITEM_FIELD) + if len(item_edges) != 1: + continue + item_edge = item_edges[0] + item_source_id = item_edge.source.node_id + source_node = self._state.graph.get_node(item_source_id) + if ( + isinstance(source_node, ForInvocation) + and get_output_field_scope(source_node, item_edge.source.field) == OutputScope.Final + ): + continue + ordered_iterator_ids = [ + iterator_id + for iterator_id in nx.ancestors(iterator_graph, item_source_id) + if isinstance(self._state.graph.get_node(iterator_id), (ForInvocation, IterateInvocation)) + ] + ordered_iterator_ids.sort(key=iterator_topological_index.__getitem__) + if isinstance(source_node, IterateInvocation): + ordered_iterator_ids.append(item_source_id) + for parent_iterator_id in ordered_iterator_ids[:-1]: + relationship_graph.add_edge(parent_iterator_id, collect_node.id) + for collect_node in self._state.graph.nodes.values(): if not isinstance(collect_node, CollectInvocation): continue @@ -1718,13 +1759,18 @@ def _get_sibling_iterate_if_collect_branches(self) -> dict[str, tuple[str, str]] for iterate_node_id in iterator_ids - condition_iterator_ids: if any( - nx.has_path(iterator_graph, condition_iterator_id, iterate_node_id) - or nx.has_path(iterator_graph, iterate_node_id, condition_iterator_id) + nx.has_path(relationship_graph, condition_iterator_id, iterate_node_id) + or nx.has_path(relationship_graph, iterate_node_id, condition_iterator_id) for condition_iterator_id in condition_iterator_ids ): continue branches.setdefault(iterate_node_id, (destination.id, output_edge.destination.field)) + self._iteration_relationship_graph = relationship_graph + self._iteration_relationship_topological_index = { + node_id: index + for index, node_id in enumerate(nx.lexicographical_topological_sort(relationship_graph, key=str)) + } self._sibling_iterate_if_collect_branches = branches return branches @@ -1818,11 +1864,15 @@ def _attach_pending_if_inputs(self, *, enqueue: bool = True) -> None: def iterator_graph(self, base: Optional["nx.DiGraph"] = None) -> "nx.DiGraph": """Gets a DiGraph with edges to collectors removed so an ancestor search produces all active iterators for any node""" g = base.copy() if base is not None else self._state._get_source_graph_flat().copy() - collectors = ( + collectors = [ n for n in self._state.graph.nodes if isinstance(self._state.graph.get_node(n), CollectInvocation) - ) + ] for c in collectors: + collection_sources = { + edge.source.node_id for edge in self._state.graph._get_input_edges(c, COLLECTION_FIELD) + } g.remove_edges_from(list(g.in_edges(c))) + g.add_edges_from((source_id, c) for source_id in collection_sources) for edge in self._state.graph.edges: source_node = self._state.graph.get_node(edge.source.node_id) if ( diff --git a/tests/app/services/shared/test_execution_engine_differential.py b/tests/app/services/shared/test_execution_engine_differential.py index f251548e5cf..a8bbc5fb13e 100644 --- a/tests/app/services/shared/test_execution_engine_differential.py +++ b/tests/app/services/shared/test_execution_engine_differential.py @@ -11,7 +11,7 @@ from pydantic import ValidationError from invokeai.app.invocations.call_saved_workflow import CallSavedWorkflowInvocation -from invokeai.app.invocations.collections import CollectionConcatInvocation, RangeInvocation +from invokeai.app.invocations.collections import CollectionConcatInvocation, RangeInvocation, RangeOfSizeInvocation from invokeai.app.invocations.logic import IfInvocation from invokeai.app.invocations.loops import ( ForInvocation, @@ -8687,6 +8687,207 @@ def test_nested_iterate_collected_in_if_branch_is_not_misclassified_as_sibling() ] +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_nested_collect_preserves_condition_frame_when_classifying_branch_iterate( + force_compatibility_scheduler: bool, +) -> None: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="outer_source", first=[1, 2])) + graph.add_node(IterateInvocation(id="outer")) + graph.add_node(RangeOfSizeInvocation(id="middle_range", size=1)) + graph.add_node(IterateInvocation(id="middle")) + graph.add_node(CollectInvocation(id="middle_collect")) + graph.add_node(CollectInvocation(id="middle_collect_2")) + graph.add_node(IterateInvocation(id="inner")) + graph.add_node(AddInvocation(id="body")) + graph.add_node(CollectInvocation(id="branch_collect")) + graph.add_node(FixedCollectionTestInvocation(id="condition")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value="fallback")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="result_collect")) + for source, source_field, destination, destination_field in ( + ("outer_source", "collection", "outer", "collection"), + ("outer", "item", "middle_range", "start"), + ("middle_range", "collection", "middle", "collection"), + ("middle", "item", "middle_collect", "item"), + ("middle_collect", "collection", "middle_collect_2", "collection"), + ("middle_collect_2", "collection", "inner", "collection"), + ("outer", "item", "body", "a"), + ("inner", "item", "body", "b"), + ("body", "value", "branch_collect", "item"), + ("outer", "item", "condition", "value"), + ("condition", "condition", "if", "condition"), + ("branch_collect", "collection", "if", "true_input"), + ("fallback", "value", "if", "false_input"), + ("if", "value", "result_collect", "item"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + + state = GraphExecutionState(graph=graph) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) + + assert state._materializer()._get_sibling_iterate_if_collect_branch("inner") is None + + trace, state = _run_to_completion(state) + + assert state.is_complete() + middle_collect_exec_ids = sorted(state.source_prepared_mapping["middle_collect"], key=state._get_iteration_path) + assert [ + (state._get_iteration_path(exec_id), state.results[exec_id].collection) for exec_id in middle_collect_exec_ids + ] == [((0,), [1]), ((1,), [2])] + middle_collect_2_exec_ids = sorted(state.source_prepared_mapping["middle_collect_2"], key=state._get_iteration_path) + assert [ + (state._get_iteration_path(exec_id), state.results[exec_id].collection) for exec_id in middle_collect_2_exec_ids + ] == [((0,), [1]), ((1,), [2])] + inner_exec_ids = sorted(state.source_prepared_mapping["inner"], key=state._get_iteration_path) + assert [state._get_iteration_path(exec_id) for exec_id in inner_exec_ids] == [(0, 0), (1, 0)] + branch_collect_exec_ids = sorted(state.source_prepared_mapping["branch_collect"], key=state._get_iteration_path) + assert [ + (state._get_iteration_path(exec_id), state.results[exec_id].collection) for exec_id in branch_collect_exec_ids + ] == [((0,), [2]), ((1,), [4])] + assert state.results[next(iter(state.source_prepared_mapping["result_collect"]))].collection == [[2], [4]] + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_direct_collect_collapses_condition_frame_for_branch_iterate( + force_compatibility_scheduler: bool, +) -> None: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="outer_source", first=[1, 2])) + graph.add_node(IterateInvocation(id="outer")) + graph.add_node(CollectInvocation(id="middle_collect")) + graph.add_node(IterateInvocation(id="inner")) + graph.add_node(AddInvocation(id="body")) + graph.add_node(CollectInvocation(id="branch_collect")) + graph.add_node(FixedCollectionTestInvocation(id="condition")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value="fallback")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="result_collect")) + for source, source_field, destination, destination_field in ( + ("outer_source", "collection", "outer", "collection"), + ("outer", "item", "middle_collect", "item"), + ("middle_collect", "collection", "inner", "collection"), + ("outer", "item", "body", "a"), + ("inner", "item", "body", "b"), + ("body", "value", "branch_collect", "item"), + ("outer", "item", "condition", "value"), + ("condition", "condition", "if", "condition"), + ("branch_collect", "collection", "if", "true_input"), + ("fallback", "value", "if", "false_input"), + ("if", "value", "result_collect", "item"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + + state = GraphExecutionState(graph=graph) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) + + assert state._materializer()._get_sibling_iterate_if_collect_branch("inner") == ("if", "true_input") + + trace, state = _run_to_completion(state) + + assert state.is_complete() + middle_collect_id = next(iter(state.source_prepared_mapping["middle_collect"])) + assert state._get_iteration_path(middle_collect_id) == () + assert state.results[middle_collect_id].collection == [1, 2] + inner_exec_ids = sorted(state.source_prepared_mapping["inner"], key=state._get_iteration_path) + assert [state._get_iteration_path(exec_id) for exec_id in inner_exec_ids] == [ + (0, 0), + (0, 1), + (1, 0), + (1, 1), + ] + assert state.results[next(iter(state.source_prepared_mapping["result_collect"]))].collection == [[2, 3], [3, 4]] + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +@pytest.mark.parametrize("node_order", ["original", "reverse", "seed7", "renamed", "renamed_producers"]) +def test_independent_iterator_collect_preserves_ordered_parent_axis_for_branch_iterate( + force_compatibility_scheduler: bool, + node_order: str, +) -> None: + graph = Graph() + condition_axis_id, other_axis_id = ("z_condition", "a_other") if node_order == "renamed" else ("left", "right") + condition_source_id, other_source_id = ( + ("z_condition_source", "a_other_source") + if node_order == "renamed_producers" + else ("left_source", "right_source") + ) + nodes = { + condition_source_id: CollectionConcatInvocation(id=condition_source_id, first=[1, 2]), + condition_axis_id: IterateInvocation(id=condition_axis_id), + other_source_id: CollectionConcatInvocation(id=other_source_id, first=[10, 20]), + other_axis_id: IterateInvocation(id=other_axis_id), + "middle_body": AddInvocation(id="middle_body"), + "middle_collect": CollectInvocation(id="middle_collect"), + "candidate": IterateInvocation(id="candidate"), + "body": AddInvocation(id="body"), + "branch_collect": CollectInvocation(id="branch_collect"), + "condition": FixedCollectionTestInvocation(id="condition"), + "fallback": AnyTypeTestInvocation(id="fallback", value="fallback"), + "if": IfInvocation(id="if"), + "result_collect": CollectInvocation(id="result_collect"), + } + ordered_node_ids = list(nodes) + if node_order == "reverse": + ordered_node_ids.reverse() + elif node_order == "seed7": + ordered_node_ids = random.Random(7).sample(ordered_node_ids, len(ordered_node_ids)) + for node_id in ordered_node_ids: + graph.add_node(nodes[node_id]) + for source, source_field, destination, destination_field in ( + (condition_source_id, "collection", condition_axis_id, "collection"), + (other_source_id, "collection", other_axis_id, "collection"), + (condition_axis_id, "item", "middle_body", "a"), + (other_axis_id, "item", "middle_body", "b"), + ("middle_body", "value", "middle_collect", "item"), + ("middle_collect", "collection", "candidate", "collection"), + (condition_axis_id, "item", "body", "a"), + ("candidate", "item", "body", "b"), + ("body", "value", "branch_collect", "item"), + (condition_axis_id, "item", "condition", "value"), + ("condition", "condition", "if", "condition"), + ("branch_collect", "collection", "if", "true_input"), + ("fallback", "value", "if", "false_input"), + ("if", "value", "result_collect", "item"), + ): + if node_order == "renamed" and (source, source_field, destination, destination_field) == ( + condition_axis_id, + "item", + "condition", + "value", + ): + continue + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + + state = GraphExecutionState(graph=graph) + if node_order == "renamed": + state._materializer()._get_sibling_iterate_if_collect_branches() + state.add_edge(create_edge(condition_axis_id, "item", "condition", "value")) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) + + trace, state = _run_to_completion(state) + + assert state.is_complete() + middle_collect_exec_ids = sorted(state.source_prepared_mapping["middle_collect"], key=state._get_iteration_path) + assert [ + (state._get_iteration_path(exec_id), state.results[exec_id].collection) for exec_id in middle_collect_exec_ids + ] == [((0,), [11, 21]), ((1,), [12, 22])] + candidate_exec_ids = sorted(state.source_prepared_mapping["candidate"], key=state._get_iteration_path) + assert [state._get_iteration_path(exec_id) for exec_id in candidate_exec_ids] == [ + (0, 0), + (0, 1), + (1, 0), + (1, 1), + ] + assert state.results[next(iter(state.source_prepared_mapping["result_collect"]))].collection == [ + [12, 22], + [14, 24], + ] + + def test_if_does_not_combine_unrelated_open_iterator_frames(): graph = Graph() graph.add_node(CollectionConcatInvocation(id="branch_source", first=[1, 2])) From 0c78fe2c4802bbd97adc28a866aa2e3ef0dc4034 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Tue, 13 Oct 2026 14:03:31 -0500 Subject: [PATCH 10/20] Added @Pfannkuchensack optimizations --- .../shared/execution_engine/primitives.py | 2 + invokeai/app/services/shared/graph.py | 21 ++++-- .../shared/graph_execution_runtime.py | 2 +- .../services/shared/graph_if_activation.py | 73 ++++++++++--------- .../app/services/shared/graph_if_runtime.py | 26 +++++-- .../app/services/shared/graph_materializer.py | 32 +++++++- 6 files changed, 101 insertions(+), 55 deletions(-) diff --git a/invokeai/app/services/shared/execution_engine/primitives.py b/invokeai/app/services/shared/execution_engine/primitives.py index dc1bec185a9..78ca5049c67 100644 --- a/invokeai/app/services/shared/execution_engine/primitives.py +++ b/invokeai/app/services/shared/execution_engine/primitives.py @@ -51,6 +51,8 @@ class _InternalModel(BaseModel): def __eq__(self, other: object) -> bool: """Compare durable fields, excluding process-local synchronization locks.""" + if self is other: + return True if not isinstance(other, type(self)): return NotImplemented return self.model_dump(mode="python") == other.model_dump(mode="python") diff --git a/invokeai/app/services/shared/graph.py b/invokeai/app/services/shared/graph.py index c7a88f5e7d7..23f612d6ba3 100644 --- a/invokeai/app/services/shared/graph.py +++ b/invokeai/app/services/shared/graph.py @@ -1,7 +1,7 @@ import copy # noqa: F401 import json import weakref # noqa: F401 -from functools import wraps # noqa: F401 +from functools import lru_cache, wraps # noqa: F401 from typing import ( TYPE_CHECKING, # noqa: F401 Any, @@ -296,6 +296,17 @@ def __init__(self) -> None: self.fresh_flat_if_activation: Optional[bool] = None +@lru_cache(maxsize=4096) +def _engine_frame_for(state_id: str, workflow_call_depth: int, iteration_path: tuple[int, ...]) -> EngineExecutionFrame: + """Share one frozen frame per identity so gate scope checks compare by identity instead of by dump.""" + return EngineExecutionFrame( + state_id=state_id, + frame_id=f"{state_id}:{workflow_call_depth}:{','.join(str(i) for i in iteration_path)}", + iteration_path=iteration_path, + workflow_call_depth=workflow_call_depth, + ) + + class GraphExecutionState(BaseModel): """Tracks source-graph expansion, execution progress, and runtime results.""" @@ -1459,13 +1470,7 @@ def _clear_transient_runtime(self) -> None: self._generic_graph_scheduler = None def _engine_frame(self, iteration_path: tuple[int, ...]) -> EngineExecutionFrame: - frame_id = f"{self.id}:{len(self.workflow_call_stack)}:{','.join(str(i) for i in iteration_path)}" - return EngineExecutionFrame( - state_id=self.id, - frame_id=frame_id, - iteration_path=iteration_path, - workflow_call_depth=len(self.workflow_call_stack), - ) + return _engine_frame_for(self.id, len(self.workflow_call_stack), iteration_path) def _activation_gate(self, exec_node_id: str) -> ActivationGate: return graph_if_runtime._activation_gate(self, exec_node_id) diff --git a/invokeai/app/services/shared/graph_execution_runtime.py b/invokeai/app/services/shared/graph_execution_runtime.py index ec7a6fdbca2..c4ef3b8cdc0 100644 --- a/invokeai/app/services/shared/graph_execution_runtime.py +++ b/invokeai/app/services/shared/graph_execution_runtime.py @@ -51,7 +51,7 @@ def _get_iteration_source_node_id(self, exec_node_id: str) -> Optional[str]: return self._state._prepared_registry().get_source_node_id(exec_node_id) def _get_ordered_iterator_sources(self, source_node_id: str) -> list[str]: - iterator_graph = self._state._iterator_graph(self._state.graph.nx_graph()) + iterator_graph = self._state._iterator_graph(self._state._get_source_graph_flat()) iterator_sources = [ node_id for node_id in nx.ancestors(iterator_graph, source_node_id) diff --git a/invokeai/app/services/shared/graph_if_activation.py b/invokeai/app/services/shared/graph_if_activation.py index 80220ee7dda..7a97d3b3d08 100644 --- a/invokeai/app/services/shared/graph_if_activation.py +++ b/invokeai/app/services/shared/graph_if_activation.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional from invokeai.app.invocations.logic import IfInvocation from invokeai.app.services.shared.execution_engine.scheduler import ActivationDependency @@ -19,6 +19,8 @@ class _IfActivationController: def __init__(self, state: "GraphExecutionState") -> None: self._state = state self._branch_sources_cache: dict[tuple[str, str], frozenset[str]] = {} + self._if_node_ids: Optional[tuple[str, ...]] = None + self._branch_memberships_by_source: dict[str, tuple[tuple[str, tuple[str, ...]], ...]] = {} def _branch_sources(self, if_node_id: str, branch_field: str, source_graph: "nx.DiGraph") -> set[str]: cache_key = (if_node_id, branch_field) @@ -46,11 +48,39 @@ def _branch_sources(self, if_node_id: str, branch_field: str, source_graph: "nx. self._branch_sources_cache[cache_key] = frozenset(branch_sources) return branch_sources + def _get_if_node_ids(self) -> tuple[str, ...]: + if self._if_node_ids is None: + self._if_node_ids = tuple( + node_id for node_id, node in self._state.graph.nodes.items() if isinstance(node, IfInvocation) + ) + return self._if_node_ids + + def _branch_memberships(self, source_node_id: str) -> tuple[tuple[str, tuple[str, ...]], ...]: + """Return each If with the branch fields whose exclusive sources include this source.""" + cached = self._branch_memberships_by_source.get(source_node_id) + if cached is not None: + return cached + + source_graph = self._state._get_source_graph_flat() + memberships = tuple( + (if_node_id, fields) + for if_node_id in self._get_if_node_ids() + if ( + fields := tuple( + branch_field + for branch_field in ("true_input", "false_input") + if source_node_id in self._branch_sources(if_node_id, branch_field, source_graph) + ) + ) + ) + self._branch_memberships_by_source[source_node_id] = memberships + return memberships + def get_dependencies(self, exec_node_id: str) -> tuple[ActivationDependency, ...]: """Return activation requirements for one concrete prepared node.""" source_node_id = self._state._prepared_registry().get_source_node_id(exec_node_id) - if not any(isinstance(node, IfInvocation) for node in self._state.graph.nodes.values()): + if not self._get_if_node_ids(): return () return self.get_source_dependencies(source_node_id, self._state._get_iteration_path(exec_node_id)) @@ -59,29 +89,11 @@ def get_source_dependencies( ) -> tuple[ActivationDependency, ...]: """Return branch requirements before a source node is materialized.""" - if not any(isinstance(node, IfInvocation) for node in self._state.graph.nodes.values()): - return () - source_graph = self._state._get_source_graph_flat() - dependencies: list[ActivationDependency] = [] - for source_if_id, source_if_node in self._state.graph.nodes.items(): - if not isinstance(source_if_node, IfInvocation): - continue - - matching_fields = tuple( - branch_field - for branch_field in ("true_input", "false_input") - if source_node_id in self._branch_sources(source_if_id, branch_field, source_graph) - ) - if len(matching_fields) != 1: - continue - dependencies.append( - ActivationDependency( - owner_id=source_if_id, - branch=matching_fields[0], - frame=iteration_path, - ) - ) - return tuple(dependencies) + return tuple( + ActivationDependency(owner_id=if_node_id, branch=fields[0], frame=iteration_path) + for if_node_id, fields in self._branch_memberships(source_node_id) + if len(fields) == 1 + ) def is_source_admitted(self, source_node_id: str, iteration_path: tuple[int, ...] = ()) -> bool: dependencies = self.get_source_dependencies(source_node_id, iteration_path) @@ -102,16 +114,7 @@ def is_source_inactive(self, source_node_id: str, iteration_path: tuple[int, ... if iteration_path: return self.is_source_rejected(source_node_id, iteration_path) - source_graph = self._state._get_source_graph_flat() - relevant_if_ids = { - if_node_id - for if_node_id, if_node in self._state.graph.nodes.items() - if isinstance(if_node, IfInvocation) - and any( - source_node_id in self._branch_sources(if_node_id, branch_field, source_graph) - for branch_field in ("true_input", "false_input") - ) - } + relevant_if_ids = {if_node_id for if_node_id, _ in self._branch_memberships(source_node_id)} if not relevant_if_ids: return False diff --git a/invokeai/app/services/shared/graph_if_runtime.py b/invokeai/app/services/shared/graph_if_runtime.py index eafd7ab0ef4..780712d314b 100644 --- a/invokeai/app/services/shared/graph_if_runtime.py +++ b/invokeai/app/services/shared/graph_if_runtime.py @@ -287,19 +287,31 @@ def _is_activation_dependency_rejected(state: "GraphExecutionState", dependency: matching_gate_ids = state._prepared_if_exec_ids(dependency.owner_id, dependency.frame) for gate_id in matching_gate_ids: gate = state._activation_gate(gate_id) - expected_frame = state._expected_execution_ref(gate_id).frame - if ( - gate.frame.state_id != expected_frame.state_id - or gate.frame.frame_id != expected_frame.frame_id - or gate.frame.iteration_path != expected_frame.iteration_path - or gate.frame.workflow_call_depth != expected_frame.workflow_call_depth - ): + if not _gate_matches_expected_frame(state, gate_id, gate): return False if gate.resolved and gate.selected_branch != dependency.branch: return True return False +def _gate_matches_expected_frame(state: "GraphExecutionState", gate_id: str, gate: ActivationGate) -> bool: + # Gates registered in this process share the frame instance built for the prepared path, which makes the + # field-by-field scope check below trivially true; rehydrated or mismatched gates take the full check. + if ( + gate_id in state.execution_graph.nodes + and state.prepared_source_mapping.get(gate_id) is not None + and gate.frame is state._engine_frame(state.prepared_iteration_paths.get(gate_id, ())) + ): + return True + expected_frame = state._expected_execution_ref(gate_id).frame + return ( + gate.frame.state_id == expected_frame.state_id + and gate.frame.frame_id == expected_frame.frame_id + and gate.frame.iteration_path == expected_frame.iteration_path + and gate.frame.workflow_call_depth == expected_frame.workflow_call_depth + ) + + def _resolve_activation_gate(state: "GraphExecutionState", exec_node_id: str, branch: str) -> bool: runtime = state._generic_runtime() gate = state._activation_gate(exec_node_id) diff --git a/invokeai/app/services/shared/graph_materializer.py b/invokeai/app/services/shared/graph_materializer.py index 2e3da4a2880..7aa255deb5a 100644 --- a/invokeai/app/services/shared/graph_materializer.py +++ b/invokeai/app/services/shared/graph_materializer.py @@ -44,6 +44,9 @@ class _ExecutionNodeBuilder: def __init__(self, state: "GraphExecutionState") -> None: self._state = state self._iteration_axes_by_source: dict[str, tuple[str, ...]] = {} + # Keyed by the identity of the cached flat source graph, which every graph edit replaces. + self._flat_iterator_graph: Optional[tuple["nx.DiGraph", "nx.DiGraph"]] = None + self._node_iterators: dict[str, tuple[str, ...]] = {} def _get_iteration_axes(self, source_node_id: str) -> tuple[str, ...]: cached = self._iteration_axes_by_source.get(source_node_id) @@ -1589,8 +1592,20 @@ def _attach_pending_if_inputs(self, *, enqueue: bool = True) -> None: self._state._enqueue_if_ready(exec_node_id) def iterator_graph(self, base: Optional["nx.DiGraph"] = None) -> "nx.DiGraph": - """Gets a DiGraph with edges to collectors removed so an ancestor search produces all active iterators for any node""" - g = base.copy() if base is not None else self._state._get_source_graph_flat().copy() + """Gets a DiGraph with edges to collectors removed so an ancestor search produces all active iterators for any node + + The result for the flat source graph is shared; callers must not mutate it. + """ + source_graph = self._state._get_source_graph_flat() + if base is not None and base is not source_graph: + return self._build_iterator_graph(base) + if self._flat_iterator_graph is None or self._flat_iterator_graph[0] is not source_graph: + self._flat_iterator_graph = (source_graph, self._build_iterator_graph(source_graph)) + self._node_iterators = {} + return self._flat_iterator_graph[1] + + def _build_iterator_graph(self, base: "nx.DiGraph") -> "nx.DiGraph": + g = base.copy() collectors = ( n for n in self._state.graph.nodes if isinstance(self._state.graph.get_node(n), CollectInvocation) ) @@ -1607,10 +1622,19 @@ def iterator_graph(self, base: Optional["nx.DiGraph"] = None) -> "nx.DiGraph": return g def get_node_iterators(self, node_id: str, it_graph: Optional["nx.DiGraph"] = None) -> list[str]: - g = it_graph or self.iterator_graph() + flat_iterator_graph = self.iterator_graph() + if it_graph is not None and it_graph is not flat_iterator_graph: + return self._find_node_iterators(node_id, it_graph) + cached = self._node_iterators.get(node_id) + if cached is None: + cached = tuple(self._find_node_iterators(node_id, flat_iterator_graph)) + self._node_iterators[node_id] = cached + return list(cached) + + def _find_node_iterators(self, node_id: str, it_graph: "nx.DiGraph") -> list[str]: return [ n - for n in nx.ancestors(g, node_id) + for n in nx.ancestors(it_graph, node_id) if isinstance(self._state.graph.get_node(n), (ForInvocation, IterateInvocation)) ] From b82ec67d41b127c80325e9bfd29f39459ffac7a6 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Tue, 13 Oct 2026 15:24:13 -0500 Subject: [PATCH 11/20] fix(graph): reuse iterator ancestry context Pass one ancestry context through recursive If traversal and materialization. Clarify the prepare scope and update the direct helper test caller. --- invokeai/app/services/shared/README.md | 5 +-- .../app/services/shared/graph_if_runtime.py | 31 +++++++++++++------ .../app/services/shared/graph_materializer.py | 8 ++--- .../test_execution_engine_differential.py | 2 ++ 4 files changed, 29 insertions(+), 17 deletions(-) diff --git a/invokeai/app/services/shared/README.md b/invokeai/app/services/shared/README.md index a9b97155fc2..c39122f65a6 100644 --- a/invokeai/app/services/shared/README.md +++ b/invokeai/app/services/shared/README.md @@ -674,9 +674,10 @@ materialization. - If the node is ready and not deferred by an unresolved `If`, enqueue it into its class queue. Literal-condition `If` fallback checks use the private `_IteratorAncestry` context. One context belongs to one -compatibility `prepare()` call and its source DAG, shared by candidate admission and selected-node materialization; +`_ExecutionNodeBuilder.prepare()` call and its source DAG, shared by candidate admission and selected-node materialization; a standalone iteration-path lookup creates its own context and passes it through both `If` and ordinary-node -recursion. The filtered iterator graph is built only when an unresolved selected branch needs the ancestry fallback. +recursion. The filtered iterator graph is built lazily for condition-edge iterator compatibility and, for literal conditions, +only when an unresolved selected branch needs the ancestry fallback. The context computes boolean proper-ancestor reachability with collector and final-`For` boundaries preserved, and reuses that analysis within the call. No context is retained on execution state or serialized, so later calls observe new supplied graphs and source mutations without cache invalidation. diff --git a/invokeai/app/services/shared/graph_if_runtime.py b/invokeai/app/services/shared/graph_if_runtime.py index 462d8733487..f390eb26ec4 100644 --- a/invokeai/app/services/shared/graph_if_runtime.py +++ b/invokeai/app/services/shared/graph_if_runtime.py @@ -19,11 +19,8 @@ def _get_effective_iteration_paths_for_edge( visited_source_ids: set[str], get_prepared_nodes_for_source: Callable[[str], set[str]], get_prepared_edge_iteration_path: Callable[[Edge, str], tuple[int, ...]], - source_graph: Optional["nx.DiGraph"] = None, - iterator_ancestry: Optional[_IteratorAncestry] = None, + iterator_ancestry: _IteratorAncestry, ) -> set[tuple[int, ...]]: - if iterator_ancestry is None: - iterator_ancestry = _IteratorAncestry(state, source_graph) pending: list[tuple[Optional[Edge], Optional[str]]] = [(edge, None)] active_source_ids = set(visited_source_ids) paths: set[tuple[int, ...]] = set() @@ -49,13 +46,12 @@ def _get_effective_iteration_paths_for_edge( source_node = state.graph.get_node(current_source_id) if isinstance(source_node, IfInvocation): paths.update( - _get_if_iteration_paths( + _get_if_iteration_paths_with_ancestry( state, current_source_id, active_source_ids, get_prepared_nodes_for_source, get_prepared_edge_iteration_path, - source_graph, iterator_ancestry, ) ) @@ -79,11 +75,27 @@ def _get_if_iteration_paths( get_prepared_nodes_for_source: Callable[[str], set[str]], get_prepared_edge_iteration_path: Callable[[Edge, str], tuple[int, ...]], source_graph: Optional["nx.DiGraph"] = None, - iterator_ancestry: Optional[_IteratorAncestry] = None, ) -> list[tuple[int, ...]]: """Infer the selected branch's effective iteration paths for an If output.""" - if iterator_ancestry is None: - iterator_ancestry = _IteratorAncestry(state, source_graph) + return _get_if_iteration_paths_with_ancestry( + state, + node_id, + visited_source_ids, + get_prepared_nodes_for_source, + get_prepared_edge_iteration_path, + _IteratorAncestry(state, source_graph), + ) + + +def _get_if_iteration_paths_with_ancestry( + state: "GraphExecutionState", + node_id: str, + visited_source_ids: set[str], + get_prepared_nodes_for_source: Callable[[str], set[str]], + get_prepared_edge_iteration_path: Callable[[Edge, str], tuple[int, ...]], + iterator_ancestry: _IteratorAncestry, +) -> list[tuple[int, ...]]: + """Infer selected branch paths using one ancestry context for all recursion.""" node = state.graph.get_node(node_id) assert isinstance(node, IfInvocation) next_visited_source_ids = {*visited_source_ids, node_id} @@ -139,7 +151,6 @@ def _get_if_iteration_paths( next_visited_source_ids, get_prepared_nodes_for_source, get_prepared_edge_iteration_path, - source_graph, iterator_ancestry, ) for edge in compatible_edges diff --git a/invokeai/app/services/shared/graph_materializer.py b/invokeai/app/services/shared/graph_materializer.py index 1aa80e5afbf..e3c45384ad2 100644 --- a/invokeai/app/services/shared/graph_materializer.py +++ b/invokeai/app/services/shared/graph_materializer.py @@ -1540,7 +1540,7 @@ def create_execution_node( return new_nodes def _has_admitted_source_mapping( - self, node_id: str, graph: "nx.DiGraph", iterator_ancestry: Optional[_IteratorAncestry] = None + self, node_id: str, graph: "nx.DiGraph", iterator_ancestry: _IteratorAncestry ) -> bool: if not any(isinstance(node, IfInvocation) for node in self._state.graph.nodes.values()): return True @@ -1549,13 +1549,12 @@ def _has_admitted_source_mapping( if not condition_edges: return any( self._state._is_source_activation_admitted(node_id, iteration_path) - for iteration_path in graph_if_runtime._get_if_iteration_paths( + for iteration_path in graph_if_runtime._get_if_iteration_paths_with_ancestry( self._state, node_id, set(), self._get_prepared_nodes_for_source, self._get_prepared_edge_iteration_path, - graph, iterator_ancestry, ) ) @@ -2091,13 +2090,12 @@ def prepare(self, base_g: Optional["nx.DiGraph"] = None) -> Optional[str]: create_results = self.create_execution_node(next_node_id, iteration_mappings, iteration_path) new_node_ids.extend(create_results) elif isinstance(next_node, IfInvocation): - condition_paths = graph_if_runtime._get_if_iteration_paths( + condition_paths = graph_if_runtime._get_if_iteration_paths_with_ancestry( self._state, next_node_id, set(), self._get_prepared_nodes_for_source, self._get_prepared_edge_iteration_path, - g, iterator_ancestry, ) condition_edges = self._state.graph._get_input_edges(next_node_id, "condition") diff --git a/tests/app/services/shared/test_execution_engine_differential.py b/tests/app/services/shared/test_execution_engine_differential.py index a8bbc5fb13e..f2f841c297d 100644 --- a/tests/app/services/shared/test_execution_engine_differential.py +++ b/tests/app/services/shared/test_execution_engine_differential.py @@ -39,6 +39,7 @@ _GenericGraphSchedulerAdapter, ) from invokeai.app.services.shared.graph_if_runtime import _get_effective_iteration_paths_for_edge +from invokeai.app.services.shared.graph_iterator_ancestry import _IteratorAncestry from invokeai.app.services.shared.graph_nested_iterate_planner import get_nested_iterate_sequence_depth from invokeai.app.services.shared.invocation_context import InvocationContextData, build_invocation_context from tests.test_nodes import ( @@ -9021,6 +9022,7 @@ def test_effective_iteration_paths_handles_deep_any_chain() -> None: set(), lambda source_id: state.source_prepared_mapping.get(source_id, set()), lambda _edge, prepared_id: state.prepared_iteration_paths[prepared_id], + _IteratorAncestry(state), ) assert effective_paths == {(7,)} From 0bb4557bf08fd52b87252e72972e54f2e60fc34d Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Tue, 13 Oct 2026 15:58:23 -0500 Subject: [PATCH 12/20] fix(graph): unblock literal if branches in iterators Admit selected literal branches before If activation; cover Iterate and controller gating. --- .../services/shared/graph_if_activation.py | 33 ++++++++++++----- .../shared/test_graph_if_activation.py | 37 ++++++++++++++++++- tests/test_graph_execution_state.py | 21 +++++++++++ 3 files changed, 79 insertions(+), 12 deletions(-) diff --git a/invokeai/app/services/shared/graph_if_activation.py b/invokeai/app/services/shared/graph_if_activation.py index 7a97d3b3d08..e9916db102a 100644 --- a/invokeai/app/services/shared/graph_if_activation.py +++ b/invokeai/app/services/shared/graph_if_activation.py @@ -21,6 +21,7 @@ def __init__(self, state: "GraphExecutionState") -> None: self._branch_sources_cache: dict[tuple[str, str], frozenset[str]] = {} self._if_node_ids: Optional[tuple[str, ...]] = None self._branch_memberships_by_source: dict[str, tuple[tuple[str, tuple[str, ...]], ...]] = {} + self._source_level_selected_fields: dict[str, Optional[str]] = {} def _branch_sources(self, if_node_id: str, branch_field: str, source_graph: "nx.DiGraph") -> set[str]: cache_key = (if_node_id, branch_field) @@ -55,6 +56,15 @@ def _get_if_node_ids(self) -> tuple[str, ...]: ) return self._if_node_ids + def _get_source_level_selected_field(self, if_node_id: str) -> Optional[str]: + if if_node_id not in self._source_level_selected_fields: + if_node = self._state.graph.nodes[if_node_id] + selected_field = None + if not self._state.graph._get_input_edges(if_node_id, "condition"): + selected_field = "true_input" if if_node.condition else "false_input" + self._source_level_selected_fields[if_node_id] = selected_field + return self._source_level_selected_fields[if_node_id] + def _branch_memberships(self, source_node_id: str) -> tuple[tuple[str, tuple[str, ...]], ...]: """Return each If with the branch fields whose exclusive sources include this source.""" cached = self._branch_memberships_by_source.get(source_node_id) @@ -62,17 +72,20 @@ def _branch_memberships(self, source_node_id: str) -> tuple[tuple[str, tuple[str return cached source_graph = self._state._get_source_graph_flat() - memberships = tuple( - (if_node_id, fields) - for if_node_id in self._get_if_node_ids() - if ( - fields := tuple( - branch_field - for branch_field in ("true_input", "false_input") - if source_node_id in self._branch_sources(if_node_id, branch_field, source_graph) - ) + memberships = [] + for if_node_id in self._get_if_node_ids(): + fields = tuple( + branch_field + for branch_field in ("true_input", "false_input") + if source_node_id in self._branch_sources(if_node_id, branch_field, source_graph) ) - ) + if not fields: + continue + selected_field = self._get_source_level_selected_field(if_node_id) + if selected_field is not None and selected_field in fields: + continue + memberships.append((if_node_id, fields)) + memberships = tuple(memberships) self._branch_memberships_by_source[source_node_id] = memberships return memberships diff --git a/tests/app/services/shared/test_graph_if_activation.py b/tests/app/services/shared/test_graph_if_activation.py index f2cfc39deed..507d9fba844 100644 --- a/tests/app/services/shared/test_graph_if_activation.py +++ b/tests/app/services/shared/test_graph_if_activation.py @@ -20,9 +20,9 @@ def _edge(source: str, destination: str, field: str) -> Any: class _Graph: - def __init__(self) -> None: + def __init__(self, *, condition_input: bool = True, condition: bool = False) -> None: self.nodes = { - "if": IfInvocation(id="if"), + "if": IfInvocation(id="if", condition=condition), "true": object(), "false": object(), "outside": object(), @@ -31,6 +31,8 @@ def __init__(self) -> None: ("if", "true_input"): [_edge("true", "if", "true_input")], ("if", "false_input"): [_edge("false", "if", "false_input")], } + if condition_input: + self._input_edges[("if", "condition")] = [_edge("condition", "if", "condition")] self._output_edges = { "true": self._input_edges[("if", "true_input")], "false": self._input_edges[("if", "false_input")], @@ -123,3 +125,34 @@ def ancestors(self, _graph: object, _node_id: str) -> set[str]: assert controller.is_source_rejected("true") state.rejected.add(ActivationDependency(owner_id="if", branch="true_input", frame=(1,))) assert controller.is_source_inactive("true", (1,)) + + +@pytest.mark.parametrize( + ("condition", "selected_source", "unselected_source", "unselected_field"), + [ + (True, "true", "false", "false_input"), + (False, "false", "true", "true_input"), + ], +) +def test_controller_admits_only_source_level_selected_branch( + monkeypatch: pytest.MonkeyPatch, + condition: bool, + selected_source: str, + unselected_source: str, + unselected_field: str, +) -> None: + class _NetworkX: + def ancestors(self, _graph: object, _node_id: str) -> set[str]: + return set() + + monkeypatch.setattr(graph_module, "nx", _NetworkX()) + state = _State() + state.graph = _Graph(condition_input=False, condition=condition) + controller = _IfActivationController(state) + + assert controller.get_source_dependencies(selected_source) == () + assert controller.is_source_admitted(selected_source) + assert not controller.is_source_inactive(selected_source) + assert controller.get_source_dependencies(unselected_source) == ( + ActivationDependency(owner_id="if", branch=unselected_field, frame=()), + ) diff --git a/tests/test_graph_execution_state.py b/tests/test_graph_execution_state.py index c2fab0de9f8..f939043bd8a 100644 --- a/tests/test_graph_execution_state.py +++ b/tests/test_graph_execution_state.py @@ -4241,6 +4241,27 @@ def test_graph_state_collects(): assert sorted(g.results[n6[0].id].collection) == sorted(test_prompts) +def test_single_if_in_iterate_body_completes(): + graph = Graph() + graph.add_node(IntegerCollectionInvocation(id="collection", collection=[1, 2])) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(AddInvocation(id="add", b=10)) + graph.add_node(IfInvocation(id="if", condition=True)) + graph.add_node(CollectInvocation(id="collect")) + graph.add_edge(create_edge("collection", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "add", "a")) + graph.add_edge(create_edge("add", "value", "if", "true_input")) + graph.add_edge(create_edge("if", "value", "collect", "item")) + + state = GraphExecutionState(graph=graph) + executed_source_ids = execute_all_nodes(state) + + assert state.is_complete(), f"stalled after executing source nodes: {executed_source_ids}" + assert not state.errors + collect_exec_id = next(iter(state.source_prepared_mapping["collect"])) + assert state.results[collect_exec_id].collection == [11, 12] + + def test_graph_state_empty_iterator_collects_and_completes(): graph = Graph() graph.add_node(IntegerCollectionInvocation(id="collection", collection=[])) From 401c92a2ae716f80358c726a346084e629716f6c Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Tue, 13 Oct 2026 16:09:52 -0500 Subject: [PATCH 13/20] docs(graph): document nested if recursion limit --- invokeai/app/services/shared/README.md | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/invokeai/app/services/shared/README.md b/invokeai/app/services/shared/README.md index c39122f65a6..58a05aa6c7e 100644 --- a/invokeai/app/services/shared/README.md +++ b/invokeai/app/services/shared/README.md @@ -676,12 +676,18 @@ materialization. Literal-condition `If` fallback checks use the private `_IteratorAncestry` context. One context belongs to one `_ExecutionNodeBuilder.prepare()` call and its source DAG, shared by candidate admission and selected-node materialization; a standalone iteration-path lookup creates its own context and passes it through both `If` and ordinary-node -recursion. The filtered iterator graph is built lazily for condition-edge iterator compatibility and, for literal conditions, +traversal. The filtered iterator graph is built lazily for condition-edge iterator compatibility and, for literal conditions, only when an unresolved selected branch needs the ancestry fallback. The context computes boolean proper-ancestor reachability with collector and final-`For` boundaries preserved, and reuses that analysis within the call. No context is retained on execution state or serialized, so later calls observe new supplied graphs and source mutations without cache invalidation. +Ordinary source nodes are traversed with an explicit stack, but path inference recurses when a selected input comes from +another unprepared `If`. A direct lookup over a 400-node literal-`If` chain raises `RecursionError` with CPython 3.12's +default recursion limit of 1,000; serial execution of the equivalent 1,100-node graph completes. This is a pathological +depth limit in the unresolved compatibility fallback, not a general graph-size limit. Supporting arbitrary nested-`If` +depth requires converting that path inference to an explicit stack. + ### 4.5 Readiness and class ordering - `_enqueue_if_ready(nid)` applies generic readiness: `indegree == 0`, not executed or claimed, and, for an `If` From 6bb686f0e1d243f1f92cb9d36feaaec72295af84 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Tue, 13 Oct 2026 22:29:42 -0500 Subject: [PATCH 14/20] fix(graph): prevent silent loss of iterator branch frames Reject divergent If axes and shared Collect branch contexts before execution. Add regression coverage and document supported branch shapes. --- invokeai/app/services/shared/README.md | 49 +- invokeai/app/services/shared/graph.py | 103 ++- .../shared/graph_execution_runtime.py | 38 +- .../app/services/shared/graph_if_runtime.py | 191 +++-- .../shared/graph_iterator_ancestry.py | 89 +- .../app/services/shared/graph_materializer.py | 220 ++--- .../app/services/shared/graph_scheduler.py | 22 +- .../test_execution_engine_differential.py | 757 +++++++++++++++--- .../shared/test_execution_engine_scheduler.py | 4 +- .../shared/test_for_scheduler_performance.py | 129 ++- .../test_graph_execution_frame_cache.py | 34 + .../services/shared/test_graph_if_runtime.py | 55 +- .../shared/test_graph_materializer.py | 11 +- tests/test_graph_execution_state.py | 4 +- 14 files changed, 1296 insertions(+), 410 deletions(-) create mode 100644 tests/app/services/shared/test_graph_execution_frame_cache.py diff --git a/invokeai/app/services/shared/README.md b/invokeai/app/services/shared/README.md index 58a05aa6c7e..6c93e37e84b 100644 --- a/invokeai/app/services/shared/README.md +++ b/invokeai/app/services/shared/README.md @@ -347,6 +347,13 @@ mutation helpers. Those helpers reject changes once the affected nodes have alre - source node id - iteration path - runtime state such as pending, ready, executed, or skipped +- _engine_frame() stores canonical immutable frames in the owning state's private runtime, keyed by state ID, + workflow-call depth, and iteration path. The cache is process-local, is not serialized, and retains only the distinct + frames requested while that state is alive. State copies get their own cache dictionary; a copy with a new ID starts + with no frames for the old ID. +- If activation dependency lookup has a reverse index from If owner, selected branch, and frame to prepared exec IDs. + A restored state rebuilds the index once; newly prepared nodes extend it. The index is derived runtime state, not + snapshot data. - `execution_refs: dict[str, ExecutionReference]` - in-memory stable references for prepared execution nodes, including their source node and execution frame; version-2 snapshots rebuild these references instead of storing them. - `execution_tokens: dict[str, ExecutionToken]` - in-memory output tokens produced by applied execution results; @@ -437,7 +444,10 @@ same adapter now projects ordinary readiness and calls the graph-state generic c prepared iteration when needed, and finalizes `output_collection` and `final_state`. For an empty literal collection, the generic path completes through the existing synthetic terminal `For` result without running the body or `ForReturn`. The generic scheduler receives only opaque node IDs and dependencies; it never receives a literal next -node ID. The compatibility continuation bridge remains available only for unsupported loop shapes and explicitly +node ID. Compatibility nested-For expansion also checks whether any prepared node is the synthetic empty +terminal before finalizing empty contexts; non-empty sibling contexts still run their bodies. This handles mixed outer +collections such as [[], [1]] as well as all-empty collections. The compatibility continuation bridge remains available +only for unsupported loop shapes and explicitly legacy-loaded snapshots. The exact four-node/four-edge empty input-driven shape with one inputless static-list producer is also generic-routed. Other input-driven outer collections, five-level-or-deeper or sibling nested loops outside the exact two-sibling `CollectionConcat` fan-in contract, or unsupported mixed loop shapes remain compatibility-owned; the @@ -653,15 +663,20 @@ materialization. iterator ancestors because the source path crosses a collector, prepared parent iteration paths are still used to materialize one downstream exec node for each preserved collector path. -- For a fresh **IfInvocation**: prepare only its condition inputs first. After the condition resolves, graph-state - activation dependencies admit selected branch contexts and attach only the selected branch input to the pending `If` - execution. Unsupported fresh shapes and legacy snapshots use the compatibility controller. Rejected branch sources - are not materialized. Pending `If` state is reconstructed from its durable execution edges, prepared mappings, and - activation token after rehydration. +- For an IfInvocation with a wired condition, prepare condition inputs first. Readiness is constrained to condition + source nodes and their ancestors; iterator ancestry on a branch cannot hold the gate open. After the condition resolves, + activation dependencies admit selected branch contexts and attach only the selected branch input. Literal conditions + use their known value and selected-branch ancestry. Unsupported fresh shapes and legacy snapshots use the compatibility + controller. Rejected branch sources are retired empty without invocation. Pending If state is reconstructed from durable + execution edges, prepared mappings, and the activation token after rehydration. + - A branch `Collect` whose `item` source combines exactly two sibling `Iterate` axes can close the axis beside the `If` condition axis. The compatibility materializer expands that sibling `Iterate` beneath each resolved `If` frame selecting the `Collect` branch, pairs only the matching condition frame, and closes the branch interval empty when no - frame selects it. An unrelated open branch axis is not combined with the condition frame. + frame selects it. The condition and branch axes must be nested; divergent inner axes are rejected even when they share + an outer axis. One sibling `Iterate` can belong to only one `If` branch context. If one shared `Collect` would need to + serve multiple `If` branches, materialization raises `UnsupportedSharedIfCollectBranchesError` before invocation + rather than dropping frames. - For each new exec node: @@ -673,14 +688,18 @@ materialization. - Try to resolve any `If`-specific scheduling state. - If the node is ready and not deferred by an unresolved `If`, enqueue it into its class queue. -Literal-condition `If` fallback checks use the private `_IteratorAncestry` context. One context belongs to one -`_ExecutionNodeBuilder.prepare()` call and its source DAG, shared by candidate admission and selected-node materialization; -a standalone iteration-path lookup creates its own context and passes it through both `If` and ordinary-node -traversal. The filtered iterator graph is built lazily for condition-edge iterator compatibility and, for literal conditions, -only when an unresolved selected branch needs the ancestry fallback. -The context computes boolean proper-ancestor reachability with collector and final-`For` boundaries preserved, and -reuses that analysis within the call. No context is retained on execution state or serialized, so later calls observe -new supplied graphs and source mutations without cache invalidation. +The execution-node builder retains a filtered iterator graph for the current source-graph identity. It builds this +graph when compatibility preparation needs sibling-Collect relationships or iterator ancestry. The graph replaces +Collect item inputs with its collection inputs for traversal; edge-scoped ancestry retains item axes needed for grouping. +Final For outputs do not carry an open iteration axis. + +Iterator ancestry owns ordered ancestor and edge-axis analysis plus effective-path memoization for one graph traversal. +The builder shares that object across candidate admission and materialization. Its separate per-source axis cache adds +sibling-Collect relationship ordering and is cleared when that relationship graph is established; it does not repeat +the ancestor-list cache. Independent parent paths merge in destination-port order. A standalone path lookup owns a +separate context. These caches are process-local and rebuilt when graph identity changes; they are not serialized. + +Literal-condition If fallback path inference uses the same ancestry context, including nested selected branches. Ordinary source nodes are traversed with an explicit stack, but path inference recurses when a selected input comes from another unprepared `If`. A direct lookup over a 400-node literal-`If` chain raises `RecursionError` with CPython 3.12's diff --git a/invokeai/app/services/shared/graph.py b/invokeai/app/services/shared/graph.py index 90a99f4a32f..164db5846d6 100644 --- a/invokeai/app/services/shared/graph.py +++ b/invokeai/app/services/shared/graph.py @@ -223,6 +223,9 @@ class _GraphExecutionStateRuntime: "if_branch_sources_cache", "if_activation_dependencies_by_source", "if_activation_dependencies_by_exec", + "if_activation_dependents_by_dependency", + "if_activation_dependents_indexed_exec_ids", + "if_activation_dependents_index_complete", "prepared_if_exec_ids_by_source_and_frame", "prepared_if_exec_frames_by_source", "prepared_if_exec_indexed_sources", @@ -233,6 +236,7 @@ class _GraphExecutionStateRuntime: "generic_for_planner", "execution_scheduler", "execution_runtime", + "engine_frames", "for_parent_iteration_paths_cache", "all_for_contexts_finalized_cache", "prepared_for_index", @@ -265,6 +269,9 @@ def __init__(self) -> None: tuple[str, tuple[int, ...]], tuple[ActivationDependency, ...] ] = {} self.if_activation_dependencies_by_exec: dict[str, tuple[ActivationDependency, ...]] = {} + self.if_activation_dependents_by_dependency: dict[ActivationDependency, tuple[str, ...]] = {} + self.if_activation_dependents_indexed_exec_ids: set[str] = set() + self.if_activation_dependents_index_complete = False self.prepared_if_exec_ids_by_source_and_frame: dict[tuple[str, tuple[int, ...]], tuple[str, ...]] = {} self.prepared_if_exec_frames_by_source: dict[str, tuple[tuple[int, ...], ...]] = {} self.prepared_if_exec_indexed_sources: set[str] = set() @@ -275,6 +282,7 @@ def __init__(self) -> None: self.generic_for_planner: Optional[_GenericForPlanner] = None self.execution_scheduler: Optional[_ExecutionScheduler | _GenericGraphSchedulerAdapter] = None self.execution_runtime: Optional[_ExecutionRuntime] = None + self.engine_frames: dict[tuple[str, int, tuple[int, ...]], EngineExecutionFrame] = {} self.for_parent_iteration_paths_cache: dict[str, set[tuple[int, ...]]] = {} self.all_for_contexts_finalized_cache: dict[str, bool] = {} self.prepared_for_index: Optional[dict[tuple[str, tuple[int, ...]], str]] = None @@ -296,17 +304,6 @@ def __init__(self) -> None: self.fresh_flat_if_activation: Optional[bool] = None -@lru_cache(maxsize=4096) -def _engine_frame_for(state_id: str, workflow_call_depth: int, iteration_path: tuple[int, ...]) -> EngineExecutionFrame: - """Share one frozen frame per identity so gate scope checks compare by identity instead of by dump.""" - return EngineExecutionFrame( - state_id=state_id, - frame_id=f"{state_id}:{workflow_call_depth}:{','.join(str(i) for i in iteration_path)}", - iteration_path=iteration_path, - workflow_call_depth=workflow_call_depth, - ) - - class GraphExecutionState(BaseModel): """Tracks source-graph expansion, execution progress, and runtime results.""" @@ -411,14 +408,22 @@ class GraphExecutionState(BaseModel): def model_copy(self, *, update: Optional[dict[str, Any]] = None, deep: bool = False) -> "GraphExecutionState": copied = super().model_copy(update=update, deep=deep) if not deep: - copied.__pydantic_private__["_runtime_state"] = copy.copy(self.__pydantic_private__["_runtime_state"]) + runtime_state = copy.copy(self.__pydantic_private__["_runtime_state"]) + runtime_state.engine_frames = { + key: frame for key, frame in runtime_state.engine_frames.items() if key[0] == copied.id + } + copied.__pydantic_private__["_runtime_state"] = runtime_state if update is not None and "graph" in update: copied._invalidate_source_graph_cache() return copied def __copy__(self) -> "GraphExecutionState": copied = super().__copy__() - copied.__pydantic_private__["_runtime_state"] = copy.copy(self.__pydantic_private__["_runtime_state"]) + runtime_state = copy.copy(self.__pydantic_private__["_runtime_state"]) + runtime_state.engine_frames = { + key: frame for key, frame in runtime_state.engine_frames.items() if key[0] == copied.id + } + copied.__pydantic_private__["_runtime_state"] = runtime_state return copied def _tx_record_once(self, key: tuple[Any, ...], undo: Callable[[], None]) -> None: @@ -526,6 +531,9 @@ def _reset_apply_derived_caches(self) -> None: self._if_branch_sources_cache = {} self._if_activation_dependencies_by_source = {} self._if_activation_dependencies_by_exec = {} + self._if_activation_dependents_by_dependency = {} + self._if_activation_dependents_indexed_exec_ids = set() + self._if_activation_dependents_index_complete = False self._prepared_if_exec_ids_by_source_and_frame = {} self._prepared_if_exec_frames_by_source = {} self._prepared_if_exec_indexed_sources = set() @@ -625,6 +633,46 @@ def _get_source_activation_dependencies( def _record_activation_dependencies(self, exec_node_id: str) -> tuple[ActivationDependency, ...]: return graph_if_runtime._record_activation_dependencies(self, exec_node_id) + def _index_activation_dependencies(self, exec_node_id: str, dependencies: tuple[ActivationDependency, ...]) -> None: + if exec_node_id in self._if_activation_dependents_indexed_exec_ids: + return + for dependency in dependencies: + dependents = self._if_activation_dependents_by_dependency.get(dependency, ()) + if exec_node_id not in dependents: + self._tx_set_mapping( + self._if_activation_dependents_by_dependency, + dependency, + (*dependents, exec_node_id), + ) + self._tx_add_set(self._if_activation_dependents_indexed_exec_ids, exec_node_id) + + def _get_activation_dependents_for_if(self, source_if_id: str, iteration_path: tuple[int, ...]) -> tuple[str, ...]: + if not self._if_activation_dependents_index_complete: + for exec_node_id in self.prepared_source_mapping: + if exec_node_id not in self._if_activation_dependents_indexed_exec_ids: + self._record_activation_dependencies(exec_node_id) + self._tx_set_attr(self, "_if_activation_dependents_index_complete", True) + + candidates = { + exec_node_id + for branch in ("true_input", "false_input") + for exec_node_id in self._if_activation_dependents_by_dependency.get( + ActivationDependency(owner_id=source_if_id, branch=branch, frame=iteration_path), () + ) + } + if iteration_path: + candidates.update( + exec_node_id + for branch in ("true_input", "false_input") + for exec_node_id in self._if_activation_dependents_by_dependency.get( + ActivationDependency(owner_id=source_if_id, branch=branch, frame=()), () + ) + ) + if not candidates: + return () + execution_graph = self._get_execution_graph_flat() + return tuple(nx.topological_sort(execution_graph.subgraph(candidates))) + def _is_source_activation_admitted(self, source_node_id: str, iteration_path: tuple[int, ...] = ()) -> bool: return graph_if_runtime._is_source_activation_admitted(self, source_node_id, iteration_path) @@ -739,6 +787,9 @@ def _invalidate_source_graph_cache(self) -> None: runtime.if_branch_sources_cache = {} runtime.if_activation_dependencies_by_source = {} runtime.if_activation_dependencies_by_exec = {} + runtime.if_activation_dependents_by_dependency = {} + runtime.if_activation_dependents_indexed_exec_ids = set() + runtime.if_activation_dependents_index_complete = False runtime.if_activation_controller_instance = None runtime.execution_materializer = None runtime.execution_runtime = None @@ -1472,7 +1523,18 @@ def _clear_transient_runtime(self) -> None: self._generic_graph_scheduler = None def _engine_frame(self, iteration_path: tuple[int, ...]) -> EngineExecutionFrame: - return _engine_frame_for(self.id, len(self.workflow_call_stack), iteration_path) + runtime_state = self.__pydantic_private__["_runtime_state"] + key = (self.id, len(self.workflow_call_stack), iteration_path) + frame = runtime_state.engine_frames.get(key) + if frame is None: + frame = EngineExecutionFrame( + state_id=self.id, + frame_id=f"{self.id}:{len(self.workflow_call_stack)}:{','.join(str(i) for i in iteration_path)}", + iteration_path=iteration_path, + workflow_call_depth=len(self.workflow_call_stack), + ) + runtime_state.engine_frames[key] = frame + return frame def _activation_gate(self, exec_node_id: str) -> ActivationGate: return graph_if_runtime._activation_gate(self, exec_node_id) @@ -3311,8 +3373,11 @@ def _try_resolve_if_node(self, exec_node_id: str, *, enqueue: bool = True) -> No self._tx_set_mapping(self._resolved_if_exec_branches, exec_node_id, selected_field) self._record_compatibility_activation_token(exec_node_id, selected_field) assert isinstance(scheduler, _ExecutionScheduler) - scheduler._discard_rejected_activation_nodes() - scheduler._enqueue_activation_ready_nodes() + source_if_id = self._prepared_registry().get_source_node_id(exec_node_id) + iteration_path = self._get_iteration_path(exec_node_id) + activation_dependents = self._get_activation_dependents_for_if(source_if_id, iteration_path) + scheduler._discard_rejected_activation_nodes(activation_dependents) + scheduler._enqueue_activation_ready_nodes(activation_dependents) if enqueue: self._enqueue_if_ready(exec_node_id) @@ -3440,6 +3505,9 @@ def _reset_runtime_caches(self) -> None: self._if_activation_controller_instance = None self._if_activation_dependencies_by_source = {} self._if_activation_dependencies_by_exec = {} + self._if_activation_dependents_by_dependency = {} + self._if_activation_dependents_indexed_exec_ids = set() + self._if_activation_dependents_index_complete = False self._prepared_if_exec_ids_by_source_and_frame = {} self._prepared_if_exec_indexed_sources = set() self._source_graph_flat = None @@ -4175,6 +4243,9 @@ def delete_edge(self, edge: Edge) -> None: "if_branch_sources_cache", "if_activation_dependencies_by_source", "if_activation_dependencies_by_exec", + "if_activation_dependents_by_dependency", + "if_activation_dependents_indexed_exec_ids", + "if_activation_dependents_index_complete", "prepared_if_exec_ids_by_source_and_frame", "prepared_if_exec_frames_by_source", "prepared_if_exec_indexed_sources", diff --git a/invokeai/app/services/shared/graph_execution_runtime.py b/invokeai/app/services/shared/graph_execution_runtime.py index 9d8dcc46f70..b15790b6559 100644 --- a/invokeai/app/services/shared/graph_execution_runtime.py +++ b/invokeai/app/services/shared/graph_execution_runtime.py @@ -9,6 +9,7 @@ is_call_saved_workflow_dynamic_input, ) from invokeai.app.invocations.logic import IfInvocation +from invokeai.app.services.shared.graph_iterator_ancestry import _IteratorAncestry from invokeai.app.services.shared.graph_models import Edge from invokeai.app.services.shared.graph_validation import ( COLLECTION_FIELD, @@ -41,31 +42,7 @@ class _ExecutionRuntime: def __init__(self, state: "GraphExecutionState") -> None: self._state = state - self._iterator_source_order: Optional[dict[str, int]] = None - - def _get_iterator_source_order(self) -> dict[str, int]: - if self._iterator_source_order is not None: - return self._iterator_source_order - - iterator_graph = self._state._iterator_graph(self._state.graph.nx_graph()) - graph = self._state.graph.nx_graph() - condition_ancestry_ids: set[str] = set() - for node in self._state.graph.nodes.values(): - if not isinstance(node, IfInvocation): - continue - for edge in self._state.graph._get_input_edges(node.id, "condition"): - source_id = edge.source.node_id - condition_ancestry_ids.update(nx.ancestors(graph, source_id)) - condition_ancestry_ids.add(source_id) - - topo = list( - nx.lexicographical_topological_sort( - iterator_graph, - key=lambda node_id: (node_id not in condition_ancestry_ids, str(node_id)), - ) - ) - self._iterator_source_order = {node_id: index for index, node_id in enumerate(topo)} - return self._iterator_source_order + self._iterator_ancestry = _IteratorAncestry(state) def _get_cached_iteration_path(self, exec_node_id: str) -> Optional[tuple[int, ...]]: return self._state._prepared_registry().get_iteration_path(exec_node_id) @@ -76,16 +53,7 @@ def _get_iteration_source_node_id(self, exec_node_id: str) -> Optional[str]: return self._state._prepared_registry().get_source_node_id(exec_node_id) def _get_ordered_iterator_sources(self, source_node_id: str) -> list[str]: - iterator_graph = self._state._iterator_graph(self._state._get_source_graph_flat()) - iterator_sources = [ - node_id - for node_id in nx.ancestors(iterator_graph, source_node_id) - if isinstance(self._state.graph.get_node(node_id), (ForInvocation, IterateInvocation)) - ] - - topo_index = self._get_iterator_source_order() - iterator_sources.sort(key=lambda node_id: topo_index.get(node_id, 0)) - return iterator_sources + return self._iterator_ancestry.get_ordered_iterator_ancestors(source_node_id) def _get_iterator_exec_id( self, iterator_source_id: str, exec_node_id: str, execution_graph: "nx.DiGraph" diff --git a/invokeai/app/services/shared/graph_if_runtime.py b/invokeai/app/services/shared/graph_if_runtime.py index f390eb26ec4..bf86900886b 100644 --- a/invokeai/app/services/shared/graph_if_runtime.py +++ b/invokeai/app/services/shared/graph_if_runtime.py @@ -7,12 +7,16 @@ from invokeai.app.services.shared.execution_engine.scheduler import ActivationDependency from invokeai.app.services.shared.graph_iterator_ancestry import _IteratorAncestry from invokeai.app.services.shared.graph_models import Edge, ExecutionToken -from invokeai.app.services.shared.graph_validation import IterateInvocation, nx +from invokeai.app.services.shared.graph_validation import ForInvocation, IterateInvocation if TYPE_CHECKING: from invokeai.app.services.shared.graph import GraphExecutionState +class UnsupportedIfIterationAxesError(ValueError): + """Raised when an If branch and its condition use unrelated open iterator axes.""" + + def _get_effective_iteration_paths_for_edge( state: "GraphExecutionState", edge: Edge, @@ -20,36 +24,50 @@ def _get_effective_iteration_paths_for_edge( get_prepared_nodes_for_source: Callable[[str], set[str]], get_prepared_edge_iteration_path: Callable[[Edge, str], tuple[int, ...]], iterator_ancestry: _IteratorAncestry, + effective_paths_by_edge: Optional[dict[tuple[str, str], set[tuple[int, ...]]]] = None, ) -> set[tuple[int, ...]]: - pending: list[tuple[Optional[Edge], Optional[str]]] = [(edge, None)] - active_source_ids = set(visited_source_ids) - paths: set[tuple[int, ...]] = set() + cache = effective_paths_by_edge if effective_paths_by_edge is not None else {} + pending: list[tuple[Edge, bool]] = [(edge, False)] while pending: - current_edge, exiting_source_id = pending.pop() - if exiting_source_id is not None: - active_source_ids.remove(exiting_source_id) - continue - assert current_edge is not None + current_edge, exiting_source = pending.pop() current_source_id = current_edge.source.node_id + cache_key = (current_source_id, current_edge.source.field) + if cache_key in cache: + continue + + if exiting_source: + input_edges = state.graph._get_input_edges(current_source_id) + if not input_edges: + cache[cache_key] = {()} + else: + cache[cache_key] = set().union( + *(cache[(input_edge.source.node_id, input_edge.source.field)] for input_edge in input_edges) + ) + continue + + if current_source_id in visited_source_ids: + cache[cache_key] = set() + continue prepared_nodes: set[str] = set() if current_source_id in state.source_prepared_mapping: prepared_nodes = get_prepared_nodes_for_source(current_source_id) if not prepared_nodes and current_source_id in state.executed: + cache[cache_key] = set() continue if prepared_nodes: - paths.update(get_prepared_edge_iteration_path(current_edge, prepared_id) for prepared_id in prepared_nodes) - continue - if current_source_id in active_source_ids: + cache[cache_key] = { + get_prepared_edge_iteration_path(current_edge, prepared_id) for prepared_id in prepared_nodes + } continue source_node = state.graph.get_node(current_source_id) if isinstance(source_node, IfInvocation): - paths.update( + cache[cache_key] = set( _get_if_iteration_paths_with_ancestry( state, current_source_id, - active_source_ids, + visited_source_ids, get_prepared_nodes_for_source, get_prepared_edge_iteration_path, iterator_ancestry, @@ -59,32 +77,12 @@ def _get_effective_iteration_paths_for_edge( input_edges = state.graph._get_input_edges(current_source_id) if not input_edges: - paths.add(()) + cache[cache_key] = {()} continue - active_source_ids.add(current_source_id) - pending.append((None, current_source_id)) - pending.extend((input_edge, None) for input_edge in reversed(input_edges)) - - return paths + pending.append((current_edge, True)) + pending.extend((input_edge, False) for input_edge in reversed(input_edges)) - -def _get_if_iteration_paths( - state: "GraphExecutionState", - node_id: str, - visited_source_ids: set[str], - get_prepared_nodes_for_source: Callable[[str], set[str]], - get_prepared_edge_iteration_path: Callable[[Edge, str], tuple[int, ...]], - source_graph: Optional["nx.DiGraph"] = None, -) -> list[tuple[int, ...]]: - """Infer the selected branch's effective iteration paths for an If output.""" - return _get_if_iteration_paths_with_ancestry( - state, - node_id, - visited_source_ids, - get_prepared_nodes_for_source, - get_prepared_edge_iteration_path, - _IteratorAncestry(state, source_graph), - ) + return cache[(edge.source.node_id, edge.source.field)] def _get_if_iteration_paths_with_ancestry( @@ -96,6 +94,7 @@ def _get_if_iteration_paths_with_ancestry( iterator_ancestry: _IteratorAncestry, ) -> list[tuple[int, ...]]: """Infer selected branch paths using one ancestry context for all recursion.""" + effective_paths_by_edge = iterator_ancestry.effective_iteration_paths_by_edge node = state.graph.get_node(node_id) assert isinstance(node, IfInvocation) next_visited_source_ids = {*visited_source_ids, node_id} @@ -130,15 +129,19 @@ def _get_if_iteration_paths_with_ancestry( if selected_edges: compatible_edges = [] if condition_edges: - condition_iterators = iterator_ancestry.get_iterator_ancestors(condition_edges[0].source.node_id) + condition_iterators = iterator_ancestry.get_iterator_ancestors_for_edge(condition_edges[0]) for edge in selected_edges: - branch_iterators = iterator_ancestry.get_iterator_ancestors(edge.source.node_id) - if ( + branch_iterators = iterator_ancestry.get_iterator_ancestors_for_edge(edge) + shared_iterators = condition_iterators.intersection(branch_iterators) + nested_axes = condition_iterators.issubset(branch_iterators) or branch_iterators.issubset( condition_iterators - and branch_iterators - and not condition_iterators.intersection(branch_iterators) - ): - continue + ) + if condition_iterators and branch_iterators and (not shared_iterators or not nested_axes): + raise UnsupportedIfIterationAxesError( + f"If node '{node_id}' cannot select branch input '{selected_field}' from " + f"unrelated open iterator axes: condition {sorted(condition_iterators)}, " + f"branch {sorted(branch_iterators)}" + ) compatible_edges.append(edge) else: compatible_edges = selected_edges @@ -152,6 +155,7 @@ def _get_if_iteration_paths_with_ancestry( get_prepared_nodes_for_source, get_prepared_edge_iteration_path, iterator_ancestry, + effective_paths_by_edge, ) for edge in compatible_edges ) @@ -185,6 +189,7 @@ def _record_activation_dependencies( ) -> tuple[ActivationDependency, ...]: dependencies = state._if_activation_dependencies_by_exec.get(exec_node_id) if dependencies is not None: + state._index_activation_dependencies(exec_node_id, dependencies) return dependencies source_node_id = state._prepared_registry().get_source_node_id(exec_node_id) if any(isinstance(node, IfInvocation) for node in state.graph.nodes.values()): @@ -194,6 +199,7 @@ def _record_activation_dependencies( else: dependencies = () state._tx_set_mapping(state._if_activation_dependencies_by_exec, exec_node_id, dependencies) + state._index_activation_dependencies(exec_node_id, dependencies) return dependencies @@ -204,21 +210,40 @@ def _get_source_activation_dependencies_for_preparation( if not dependencies: return dependencies - candidate_frames = sorted( - { + runtime = state._runtime() + source_axes = runtime._get_ordered_iterator_sources(source_node_id) + resolved_dependencies = [] + for dependency in dependencies: + owner_axes = runtime._get_ordered_iterator_sources(dependency.owner_id) + shared_axis_indices = [ + (owner_axes.index(axis_id), source_axes.index(axis_id)) for axis_id in owner_axes if axis_id in source_axes + ] + candidate_frames = state._prepared_if_exec_frames(dependency.owner_id) + matching_frames = [ frame - for dependency in dependencies - for frame in state._prepared_if_exec_frames(dependency.owner_id) - if not iteration_path or iteration_path[: len(frame)] == frame - } - ) - for frame in candidate_frames: - frame_dependencies = state._get_source_activation_dependencies(source_node_id, frame) - if frame_dependencies and all( - state._is_activation_dependency_satisfied(dependency) for dependency in frame_dependencies - ): - return frame_dependencies - return dependencies + for frame in candidate_frames + if all( + owner_index < len(frame) + and source_index < len(iteration_path) + and frame[owner_index] == iteration_path[source_index] + for owner_index, source_index in shared_axis_indices + ) + ] + if not matching_frames: + resolved_dependencies.append(dependency) + continue + + frame_dependencies = [ + ActivationDependency(owner_id=dependency.owner_id, branch=dependency.branch, frame=frame) + for frame in matching_frames + ] + resolved_dependencies.append( + next( + (candidate for candidate in frame_dependencies if state._is_activation_dependency_satisfied(candidate)), + frame_dependencies[0], + ) + ) + return tuple(resolved_dependencies) def _get_selected_if_frames_for_preparation( @@ -239,31 +264,49 @@ def _is_source_activation_admitted( return not dependencies or all(state._is_activation_dependency_satisfied(dependency) for dependency in dependencies) +def _is_activation_axis_complete(state: "GraphExecutionState", axis_source_id: str) -> bool: + axis_source = state.graph.nodes.get(axis_source_id) + if isinstance(axis_source, ForInvocation): + return state._all_for_contexts_finalized(axis_source_id) + if isinstance(axis_source, IterateInvocation): + streams = [stream for stream in state._generic_runtime().streams.values() if stream.owner_id == axis_source_id] + if streams: + return all(stream.closed for stream in streams) + prepared_ids = state.source_prepared_mapping.get(axis_source_id, ()) + return ( + axis_source_id in state.executed + or bool(prepared_ids) + and all(exec_node_id in state.executed for exec_node_id in prepared_ids) + ) + return True + + def _is_source_inactive( state: "GraphExecutionState", source_node_id: str, iteration_path: tuple[int, ...] = () ) -> bool: - if not state._can_use_fresh_flat_if_activation(): - return state._if_activation_controller().is_source_inactive(source_node_id, iteration_path) - + use_fresh_activation = state._can_use_fresh_flat_if_activation() dependencies = state._get_source_activation_dependencies(source_node_id, iteration_path) if not dependencies: return False if iteration_path: - return any(state._is_activation_dependency_rejected(dependency) for dependency in dependencies) + prepared_dependencies = _get_source_activation_dependencies_for_preparation( + state, source_node_id, iteration_path + ) + return any(state._is_activation_dependency_rejected(dependency) for dependency in prepared_dependencies) - frames = { - iteration_path - for dependency in dependencies - for iteration_path in state._prepared_if_exec_frames(dependency.owner_id) - } + frames = {frame for dependency in dependencies for frame in state._prepared_if_exec_frames(dependency.owner_id)} if not frames: - return all(_is_empty_if_condition(state, dependency) for dependency in dependencies) + return use_fresh_activation and all(_is_empty_if_condition(state, dependency) for dependency in dependencies) + + source_axes = set(state._runtime()._get_ordered_iterator_sources(source_node_id)) + for dependency in dependencies: + for axis_source_id in state._runtime()._get_ordered_iterator_sources(dependency.owner_id): + if axis_source_id not in source_axes and not _is_activation_axis_complete(state, axis_source_id): + return False + return all( - state._get_source_activation_dependencies(source_node_id, frame) - and any( - state._is_activation_dependency_rejected(frame_dependency) - for frame_dependency in state._get_source_activation_dependencies(source_node_id, frame) - ) + (prepared_dependencies := _get_source_activation_dependencies_for_preparation(state, source_node_id, frame)) + and any(state._is_activation_dependency_rejected(dependency) for dependency in prepared_dependencies) for frame in frames ) diff --git a/invokeai/app/services/shared/graph_iterator_ancestry.py b/invokeai/app/services/shared/graph_iterator_ancestry.py index 7e03461ddb3..fe59c8b102d 100644 --- a/invokeai/app/services/shared/graph_iterator_ancestry.py +++ b/invokeai/app/services/shared/graph_iterator_ancestry.py @@ -1,28 +1,44 @@ from typing import TYPE_CHECKING, Optional +from invokeai.app.invocations.fields import OutputScope from invokeai.app.invocations.loops import ForInvocation -from invokeai.app.services.shared.graph_validation import IterateInvocation, nx +from invokeai.app.services.shared.graph_models import Edge +from invokeai.app.services.shared.graph_validation import IterateInvocation, get_output_field_scope, nx if TYPE_CHECKING: from invokeai.app.services.shared.graph import GraphExecutionState class _IteratorAncestry: - """Lazily index strict iterator ancestry for one source-graph traversal.""" + """Own the ordered open For/Iterate axes for one source-graph traversal.""" - def __init__(self, state: "GraphExecutionState", source_graph: Optional["nx.DiGraph"] = None) -> None: + def __init__( + self, + state: "GraphExecutionState", + source_graph: Optional["nx.DiGraph"] = None, + *, + iterator_graph: Optional["nx.DiGraph"] = None, + ) -> None: self._state = state self._source_graph = source_graph - self._iterator_graph: Optional["nx.DiGraph"] = None + self._iterator_graph = iterator_graph self._has_iterator_ancestor: dict[str, bool] | None = None self._iterator_ancestors: dict[str, frozenset[str]] = {} + self._ordered_iterator_ancestors: dict[str, tuple[str, ...]] = {} + self.effective_iteration_paths_by_edge: dict[tuple[str, str], set[tuple[int, ...]]] = {} + + def clear_effective_iteration_paths(self) -> None: + """Invalidate prepared-path results when execution materialization advances.""" + self.effective_iteration_paths_by_edge.clear() def has_iterator_ancestor(self, node_id: str) -> bool: if self._has_iterator_ancestor is None: - source_graph = self._source_graph - if source_graph is None: - source_graph = self._state._get_source_graph_flat() - iterator_graph = self._state._iterator_graph(source_graph) + iterator_graph = self._iterator_graph + if iterator_graph is None: + source_graph = self._source_graph + if source_graph is None: + source_graph = self._state._get_source_graph_flat() + iterator_graph = self._state._iterator_graph(source_graph) has_iterator_ancestor = dict.fromkeys(iterator_graph, False) for source_node_id in nx.topological_sort(iterator_graph): has_iterator_ancestor[source_node_id] = any( @@ -50,3 +66,60 @@ def get_iterator_ancestors(self, node_id: str) -> frozenset[str]: ancestors.add(node_id) self._iterator_ancestors[node_id] = frozenset(ancestors) return self._iterator_ancestors[node_id] + + def get_ordered_iterator_ancestors(self, node_id: str) -> list[str]: + """Return open iterator axes in outer-to-inner, input-port order.""" + self.has_iterator_ancestor(node_id) + cached = self._ordered_iterator_ancestors.get(node_id) + if cached is not None: + return list(cached) + + assert self._iterator_graph is not None + ordered: list[str] = [] + visited: set[str] = set() + pending: list[tuple[str, bool]] = [(node_id, False)] + while pending: + current_node_id, exiting = pending.pop() + if exiting: + if current_node_id != node_id and isinstance( + self._state.graph.get_node(current_node_id), (ForInvocation, IterateInvocation) + ): + ordered.append(current_node_id) + continue + if current_node_id in visited: + continue + visited.add(current_node_id) + pending.append((current_node_id, True)) + current_node = self._state.graph.get_node(current_node_id) + input_fields = list(type(current_node).model_fields) + field_order = {field: index for index, field in enumerate(input_fields)} + input_edges = sorted( + self._state.graph._get_input_edges(current_node_id), + key=lambda edge: (field_order.get(edge.destination.field, len(field_order)), edge.source.node_id), + ) + for edge in reversed(input_edges): + source_node_id = edge.source.node_id + if self._iterator_graph.has_edge(source_node_id, current_node_id): + pending.append((source_node_id, False)) + + cached = tuple(ordered) + self._ordered_iterator_ancestors[node_id] = cached + return list(cached) + + def get_iterator_ancestors_for_edge(self, edge: "Edge") -> frozenset[str]: + """Return iterator axes carried by an edge, excluding closed For output axes.""" + ancestors = self.get_iterator_axes_for_edge(edge) + return frozenset(ancestors) + + def get_iterator_axes_for_edge(self, edge: "Edge") -> tuple[str, ...]: + """Return ordered axes carried by an edge, excluding closed For output axes.""" + source_node_id = edge.source.node_id + ancestors = self.get_ordered_iterator_ancestors(source_node_id) + source_node = self._state.graph.get_node(source_node_id) + if isinstance(source_node, (ForInvocation, IterateInvocation)): + if not ( + isinstance(source_node, ForInvocation) + and get_output_field_scope(source_node, edge.source.field) == OutputScope.Final + ): + ancestors.append(source_node_id) + return tuple(ancestors) diff --git a/invokeai/app/services/shared/graph_materializer.py b/invokeai/app/services/shared/graph_materializer.py index e3c45384ad2..22293497a7f 100644 --- a/invokeai/app/services/shared/graph_materializer.py +++ b/invokeai/app/services/shared/graph_materializer.py @@ -33,6 +33,10 @@ from invokeai.app.services.shared.graph import GraphExecutionState +class UnsupportedSharedIfCollectBranchesError(ValueError): + """Raised when one sibling Collect-backed Iterate must serve multiple If frames.""" + + class _ExecutionNodeBuilder: """Builds concrete execution-graph nodes for a selected runtime planner. @@ -43,13 +47,15 @@ class _ExecutionNodeBuilder: def __init__(self, state: "GraphExecutionState") -> None: self._state = state + self._has_if_nodes = any(isinstance(node, IfInvocation) for node in state.graph.nodes.values()) self._iteration_axes_by_source: dict[str, tuple[str, ...]] = {} self._sibling_iterate_if_collect_branches: Optional[dict[str, tuple[str, str]]] = None self._iteration_relationship_graph: Optional[nx.DiGraph] = None self._iteration_relationship_topological_index: Optional[dict[str, int]] = None # Keyed by the identity of the cached flat source graph, which every graph edit replaces. self._flat_iterator_graph: Optional[tuple["nx.DiGraph", "nx.DiGraph"]] = None - self._node_iterators: dict[str, tuple[str, ...]] = {} + self._iterator_ancestry_graph: Optional["nx.DiGraph"] = None + self._iterator_ancestry: Optional[_IteratorAncestry] = None def _get_iteration_axes(self, source_node_id: str) -> tuple[str, ...]: cached = self._iteration_axes_by_source.get(source_node_id) @@ -83,6 +89,26 @@ def _get_iteration_axes(self, source_node_id: str) -> tuple[str, ...]: self._iteration_axes_by_source[source_node_id] = result return result + def _get_ordered_input_iterator_axes(self, input_edges: list[Edge]) -> tuple[str, ...]: + if not input_edges: + return () + + destination = self._state.graph.get_node(input_edges[0].destination.node_id) + input_field_order = {field: index for index, field in enumerate(type(destination).model_fields)} + ordered_edges = sorted( + input_edges, + key=lambda edge: ( + input_field_order.get(edge.destination.field, len(input_field_order)), + edge.source.node_id, + ), + ) + axes: list[str] = [] + for edge in ordered_edges: + for axis_id in self._get_iteration_axes(edge.source.node_id): + if axis_id not in axes: + axes.append(axis_id) + return tuple(axes) + def _get_known_iteration_path( self, iteration_index: int, @@ -112,14 +138,32 @@ def _get_known_iteration_path( compatible_parent_contexts = all( deepest_path[: len(path)] == path and deepest_axes[: len(axes)] == axes for path, axes in parent_contexts ) - if not compatible_parent_contexts: - return None + if compatible_parent_contexts: + iteration_path = deepest_path + else: + # Independent input axes contribute separate path components. Merge their indexes in destination-port + # order instead of rejecting them as incompatible prefixes. + ordered_axes = self._get_ordered_input_iterator_axes(input_edges or []) + axis_indexes: dict[str, int] = {} + for path, axes in parent_contexts: + if len(path) > len(axes): + return None + for axis_id, axis_index in zip(axes, path, strict=False): + if axis_id not in ordered_axes: + return None + previous_index = axis_indexes.setdefault(axis_id, axis_index) + if previous_index != axis_index: + return None + populated_positions = [index for index, axis_id in enumerate(ordered_axes) if axis_id in axis_indexes] + if populated_positions != list(range(len(populated_positions))): + return None + iteration_path = tuple(axis_indexes[axis_id] for axis_id in ordered_axes[: len(populated_positions)]) # Materialized iteration boundaries use non-negative indexes; ordinary execution nodes use -1. Keeping this # generic allows other scheduler-managed loop nodes to reuse the same path cache. if iteration_index >= 0: - return (*deepest_path, iteration_index) - return deepest_path + return (*iteration_path, iteration_index) + return iteration_path def _get_iterator_iteration_count(self, node_id: str, iteration_node_map: list[tuple[str, str]]) -> int: input_collection_edge = next(iter(self._state.graph._get_input_edges(node_id, COLLECTION_FIELD))) @@ -205,6 +249,17 @@ def _has_unmaterializable_for_final_input(self, node_id: str) -> bool: return any(not self._state._all_for_contexts_finalized(source_for_id) for source_for_id in final_for_source_ids) + def _get_ordered_iterator_axes_for_edges(self, edges: Iterable[Edge], iterator_graph: "nx.DiGraph") -> list[str]: + ancestry = self._get_iterator_ancestry(iterator_graph) + return list(dict.fromkeys(axis_id for edge in edges for axis_id in ancestry.get_iterator_axes_for_edge(edge))) + + def _get_iterator_ancestry(self, iterator_graph: "nx.DiGraph") -> _IteratorAncestry: + if self._iterator_ancestry_graph is not iterator_graph: + self._iterator_ancestry_graph = iterator_graph + self._iterator_ancestry = _IteratorAncestry(self._state, iterator_graph=iterator_graph) + assert self._iterator_ancestry is not None + return self._iterator_ancestry + def _create_execution_node_copy( self, node: BaseInvocation, node_id: str, iteration_index: int, *, deep_copy: bool = True ) -> BaseInvocation: @@ -1085,8 +1140,7 @@ def _get_collect_iteration_group_key(self, edge: Edge, sibling_depth: Optional[i isinstance(source_node, ForInvocation) and get_output_field_scope(source_node, edge.source.field) == OutputScope.Final ): - parent_path = self._state._get_for_parent_iteration_path(edge.source.node_id) - return parent_path + return self._state._get_for_parent_iteration_path(edge.source.node_id) if edge.destination.field == ITEM_FIELD: if isinstance(source_node, ForInvocation) and source_node.index == -1: return path @@ -1176,10 +1230,12 @@ def _get_collect_candidate_group_keys(self, edge: Edge) -> set[tuple[int, ...]]: group_depth = max(group_depth - 1, 0) group_keys: set[tuple[int, ...]] = set() - for iterator_node_id in iterator_node_ids: + # A Collect.item edge groups by the outer `group_depth` axes. Read input paths from the current item axis + # too: an empty stream still carries its parent groups even when it has no prepared items. + for axis_index, iterator_node_id in enumerate(iterator_node_ids[: group_depth + 1]): prepared_nodes = self._get_ordered_prepared_nodes_for_source(iterator_node_id) - # Prepared paths use the active group depth. Input paths stay full to preserve scope across collectors. - if prepared_nodes and group_depth: + # Only outer group axes contribute prepared paths; the current item axis is collapsed by Collect. + if axis_index < group_depth and prepared_nodes: group_keys.update( iteration_path[:group_depth] for prepared_id in prepared_nodes @@ -1382,16 +1438,8 @@ def _matches_sibling_if_collect_context( if len(branch_path) == len(frame) + 1 and branch_path[: len(frame)] == frame ] condition_edges = self._state.graph._get_input_edges(if_node_id, "condition") - condition_iterator_ids = { - ancestor_id - for edge in condition_edges - for ancestor_id in nx.ancestors(self._state._get_source_graph_flat(), edge.source.node_id) - if isinstance(self._state.graph.get_node(ancestor_id), (ForInvocation, IterateInvocation)) - } - condition_iterator_ids.update( - edge.source.node_id - for edge in condition_edges - if isinstance(self._state.graph.get_node(edge.source.node_id), (ForInvocation, IterateInvocation)) + condition_iterator_ids = set( + self._get_ordered_iterator_axes_for_edges(condition_edges, self.iterator_graph()) ) for frame in matching_frames: if any( @@ -1424,17 +1472,8 @@ def _get_parent_iteration_mappings( source_node = self._state.graph.get_node(input_edges[0].source.node_id) if isinstance(source_node, IfInvocation): condition_edges = self._state.graph._get_input_edges(source_node.id, "condition") - condition_iterator_ids = { - ancestor_id - for edge in condition_edges - for ancestor_id in nx.ancestors(iterator_graph, edge.source.node_id) - if isinstance(self._state.graph.get_node(ancestor_id), (ForInvocation, IterateInvocation)) - } - condition_iterator_ids.update( - edge.source.node_id - for edge in condition_edges - if isinstance(self._state.graph.get_node(edge.source.node_id), (ForInvocation, IterateInvocation)) - ) + condition_axes = self._get_ordered_iterator_axes_for_edges(condition_edges, iterator_graph) + condition_iterator_ids = set(condition_axes) if condition_iterator_ids: iterator_nodes = [node_id for node_id in iterator_nodes if node_id in condition_iterator_ids] if not iterator_nodes: @@ -1591,20 +1630,10 @@ def _get_if_condition_iteration_mappings( condition_edges = self._state.graph._get_input_edges(node_id, "condition") iterator_graph = self.iterator_graph(graph) if condition_edges: - condition_iterator_ids = { - ancestor_id - for edge in condition_edges - for ancestor_id in nx.ancestors(iterator_graph, edge.source.node_id) - if isinstance(self._state.graph.get_node(ancestor_id), (ForInvocation, IterateInvocation)) - } - condition_iterator_ids.update( - edge.source.node_id - for edge in condition_edges - if isinstance(self._state.graph.get_node(edge.source.node_id), (ForInvocation, IterateInvocation)) - ) + condition_axes = self._get_ordered_iterator_axes_for_edges(condition_edges, iterator_graph) + condition_iterator_ids = set(condition_axes) if condition_iterator_ids: - source_order = self._state._runtime()._get_iterator_source_order() - iterator_nodes = sorted(condition_iterator_ids, key=lambda source_id: source_order[source_id]) + iterator_nodes = condition_axes else: iterator_nodes = self.get_node_iterators(node_id, iterator_graph) else: @@ -1623,24 +1652,16 @@ def _get_if_condition_iteration_mappings( node_id, graph, input_edges=condition_edges, iterator_node_ids=iterator_nodes ) - # This bounded path derives one frame axis. Leave mixed/nested axes on the existing path because - # _get_known_iteration_path() intentionally rejects incompatible parent paths. + # This specialization derives one branch axis against a root-frame condition. Mixed or nested condition + # axes use the general parent-mapping path. if len(iterator_nodes) != 1: return iter([[]]) - source_graph = self._state._get_source_graph_flat() + iterator_ancestry = self._get_iterator_ancestry(iterator_graph) for branch_field in ("true_input", "false_input"): branch_edges = self._state.graph._get_input_edges(node_id, branch_field) branch_iterator_ids = { - ancestor_id - for edge in branch_edges - for ancestor_id in nx.ancestors(source_graph, edge.source.node_id) - if isinstance(self._state.graph.get_node(ancestor_id), (ForInvocation, IterateInvocation)) + axis_id for edge in branch_edges for axis_id in iterator_ancestry.get_iterator_axes_for_edge(edge) } - branch_iterator_ids.update( - edge.source.node_id - for edge in branch_edges - if isinstance(self._state.graph.get_node(edge.source.node_id), (ForInvocation, IterateInvocation)) - ) if branch_iterator_ids and not (branch_iterator_ids & set(iterator_nodes)): return iter([]) @@ -1701,7 +1722,7 @@ def _get_sibling_iterate_if_collect_branches(self) -> dict[str, tuple[str, str]] graph = self._state._get_source_graph_flat() iterator_graph = self.iterator_graph(graph) - iterator_topological_index = self._state._runtime()._get_iterator_source_order() + iterator_ancestry = self._get_iterator_ancestry(iterator_graph) relationship_graph = iterator_graph.copy() for collect_node in self._state.graph.nodes.values(): if not isinstance(collect_node, CollectInvocation): @@ -1719,14 +1740,7 @@ def _get_sibling_iterate_if_collect_branches(self) -> dict[str, tuple[str, str]] and get_output_field_scope(source_node, item_edge.source.field) == OutputScope.Final ): continue - ordered_iterator_ids = [ - iterator_id - for iterator_id in nx.ancestors(iterator_graph, item_source_id) - if isinstance(self._state.graph.get_node(iterator_id), (ForInvocation, IterateInvocation)) - ] - ordered_iterator_ids.sort(key=iterator_topological_index.__getitem__) - if isinstance(source_node, IterateInvocation): - ordered_iterator_ids.append(item_source_id) + ordered_iterator_ids = list(iterator_ancestry.get_iterator_axes_for_edge(item_edge)) for parent_iterator_id in ordered_iterator_ids[:-1]: relationship_graph.add_edge(parent_iterator_id, collect_node.id) @@ -1749,17 +1763,7 @@ def _get_sibling_iterate_if_collect_branches(self) -> dict[str, tuple[str, str]] continue condition_edges = self._state.graph._get_input_edges(destination.id, "condition") - condition_iterator_ids = { - ancestor_id - for edge in condition_edges - for ancestor_id in nx.ancestors(graph, edge.source.node_id) - if isinstance(self._state.graph.get_node(ancestor_id), (ForInvocation, IterateInvocation)) - } - condition_iterator_ids.update( - edge.source.node_id - for edge in condition_edges - if isinstance(self._state.graph.get_node(edge.source.node_id), (ForInvocation, IterateInvocation)) - ) + condition_iterator_ids = set(self._get_ordered_iterator_axes_for_edges(condition_edges, iterator_graph)) if not condition_iterator_ids.intersection(iterator_ids): continue @@ -1770,9 +1774,19 @@ def _get_sibling_iterate_if_collect_branches(self) -> dict[str, tuple[str, str]] for condition_iterator_id in condition_iterator_ids ): continue - branches.setdefault(iterate_node_id, (destination.id, output_edge.destination.field)) + branch = (destination.id, output_edge.destination.field) + existing_branch = branches.get(iterate_node_id) + if existing_branch is not None and existing_branch != branch: + raise UnsupportedSharedIfCollectBranchesError( + f"Sibling Iterate '{iterate_node_id}' cannot supply Collect-backed inputs to multiple " + f"If branches: '{existing_branch[0]}.{existing_branch[1]}' and " + f"'{branch[0]}.{branch[1]}'" + ) + branches[iterate_node_id] = branch self._iteration_relationship_graph = relationship_graph + # Cached axes may have been queried before this graph established sibling order. + self._iteration_axes_by_source.clear() self._iteration_relationship_topological_index = { node_id: index for index, node_id in enumerate(nx.lexicographical_topological_sort(relationship_graph, key=str)) @@ -1783,19 +1797,6 @@ def _get_sibling_iterate_if_collect_branches(self) -> dict[str, tuple[str, str]] def _get_sibling_iterate_if_collect_branch(self, iterate_node_id: str) -> Optional[tuple[str, str]]: return self._get_sibling_iterate_if_collect_branches().get(iterate_node_id) - def _is_inactive_sibling_if_collect_source(self, node_id: str, graph: "nx.DiGraph") -> bool: - sibling_iterate_ids = self._get_sibling_iterate_if_collect_branches().keys() - return ( - bool(sibling_iterate_ids) - and self._state._is_source_inactive(node_id) - and any( - node_id == iterate_node_id - or nx.has_path(graph, node_id, iterate_node_id) - or nx.has_path(graph, iterate_node_id, node_id) - for iterate_node_id in sibling_iterate_ids - ) - ) - def _attach_pending_if_inputs(self, *, enqueue: bool = True) -> None: """Attach only the selected branch edge to condition-ready If executions.""" @@ -1877,7 +1878,6 @@ def iterator_graph(self, base: Optional["nx.DiGraph"] = None) -> "nx.DiGraph": return self._build_iterator_graph(base) if self._flat_iterator_graph is None or self._flat_iterator_graph[0] is not source_graph: self._flat_iterator_graph = (source_graph, self._build_iterator_graph(source_graph)) - self._node_iterators = {} return self._flat_iterator_graph[1] def _build_iterator_graph(self, base: "nx.DiGraph") -> "nx.DiGraph": @@ -1907,21 +1907,11 @@ def _build_iterator_graph(self, base: "nx.DiGraph") -> "nx.DiGraph": return g def get_node_iterators(self, node_id: str, it_graph: Optional["nx.DiGraph"] = None) -> list[str]: - flat_iterator_graph = self.iterator_graph() - if it_graph is not None and it_graph is not flat_iterator_graph: - return self._find_node_iterators(node_id, it_graph) - cached = self._node_iterators.get(node_id) - if cached is None: - cached = tuple(self._find_node_iterators(node_id, flat_iterator_graph)) - self._node_iterators[node_id] = cached - return list(cached) + iterator_graph = it_graph if it_graph is not None else self.iterator_graph() + return self._find_node_iterators(node_id, iterator_graph) def _find_node_iterators(self, node_id: str, it_graph: "nx.DiGraph") -> list[str]: - return [ - n - for n in nx.ancestors(it_graph, node_id) - if isinstance(self._state.graph.get_node(n), (ForInvocation, IterateInvocation)) - ] + return self._get_iterator_ancestry(it_graph).get_ordered_iterator_ancestors(node_id) def _get_prepared_nodes_for_source(self, source_node_id: str) -> set[str]: return { @@ -2016,16 +2006,31 @@ def get_iteration_node( return self._find_prepared_node_matching_iterators(prepared_nodes, parent_iterators, execution_graph) + def _should_defer_sibling_if_collect_iterate(self, node_id: str) -> bool: + if not isinstance(self._state.graph.get_node(node_id), IterateInvocation): + return False + collapsed_branch = self._get_sibling_iterate_if_collect_branch(node_id) + if collapsed_branch is None: + return False + if graph_if_runtime._get_selected_if_frames_for_preparation(self._state, *collapsed_branch): + return False + + if_node_id, _ = collapsed_branch + if_exec_ids = self._state.source_prepared_mapping.get(if_node_id, set()) + return not if_exec_ids or not all(exec_id in self._state._resolved_if_exec_branches for exec_id in if_exec_ids) + def prepare(self, base_g: Optional["nx.DiGraph"] = None) -> Optional[str]: g = base_g if base_g is not None else self._state._get_source_graph_flat() - iterator_ancestry = _IteratorAncestry(self._state, g) + self._get_sibling_iterate_if_collect_branches() + iterator_ancestry = self._get_iterator_ancestry(self.iterator_graph(g)) + iterator_ancestry.clear_effective_iteration_paths() self._attach_pending_if_inputs() - if self._get_sibling_iterate_if_collect_branches(): + if self._has_if_nodes: for node_id in nx.topological_sort(g): if ( node_id not in self._state.source_prepared_mapping and node_id not in self._state.executed - and self._is_inactive_sibling_if_collect_source(node_id, g) + and self._state._is_source_inactive(node_id) ): self._mark_source_node_empty(node_id) self._state._invalidate_loop_caches_for_source(node_id) @@ -2059,6 +2064,7 @@ def prepare(self, base_g: Optional["nx.DiGraph"] = None) -> Optional[str]: or not isinstance(self._state.graph.get_node(node_id), IfInvocation) ) and self._has_admitted_source_mapping(node_id, g, iterator_ancestry) + and not self._should_defer_sibling_if_collect_iterate(node_id) and ( not isinstance(self._state.graph.get_node(node_id), (ForInvocation, IterateInvocation)) or all(source_id in self._state.executed for source_id, _ in g.in_edges(node_id)) diff --git a/invokeai/app/services/shared/graph_scheduler.py b/invokeai/app/services/shared/graph_scheduler.py index df937bb698a..f5146c67a2f 100644 --- a/invokeai/app/services/shared/graph_scheduler.py +++ b/invokeai/app/services/shared/graph_scheduler.py @@ -145,12 +145,12 @@ def _discard_rejected_activation_node(self, exec_node_id: str) -> bool: self.mark_skipped(exec_node_id) return True - def _discard_rejected_activation_nodes(self) -> None: - for exec_node_id in tuple(nx.topological_sort(self._state._get_execution_graph_flat())): + def _discard_rejected_activation_nodes(self, exec_node_ids: Iterable[str]) -> None: + for exec_node_id in exec_node_ids: self._discard_rejected_activation_node(exec_node_id) - def _enqueue_activation_ready_nodes(self) -> None: - for exec_node_id in tuple(self._state.prepared_source_mapping): + def _enqueue_activation_ready_nodes(self, exec_node_ids: Iterable[str]) -> None: + for exec_node_id in exec_node_ids: if exec_node_id not in self._state.indegree or exec_node_id in self._state.executed: continue self.enqueue_if_ready(exec_node_id) @@ -429,7 +429,9 @@ def resolve_if_node(self, exec_node_id: str, *, enqueue: bool = True) -> None: self._state._resolve_activation_gate(exec_node_id, selected_field) self._state._tx_set_mapping(runtime.resolved_if_exec_branches, exec_node_id, selected_field) self._state._record_compatibility_activation_token(exec_node_id, selected_field) - self._discard_rejected_nodes() + source_if_id = self._state._prepared_registry().get_source_node_id(exec_node_id) + iteration_path = self._state._get_iteration_path(exec_node_id) + self._discard_rejected_nodes(self._state._get_activation_dependents_for_if(source_if_id, iteration_path)) self._sync_indegree() if enqueue: self._project_ready_nodes() @@ -443,10 +445,10 @@ def _discard_rejected_node(self, exec_node_id: str) -> None: ): self._retire_unselected_node(exec_node_id) - def _discard_rejected_nodes(self) -> None: - """Discard plan nodes whose opaque activation requirements are rejected.""" + def _discard_rejected_nodes(self, exec_node_ids: Iterable[str]) -> None: + """Discard only prepared plan nodes governed by the resolved If frame.""" - for exec_node_id in tuple(self._scheduler.plan.nodes): + for exec_node_id in exec_node_ids: self._discard_rejected_node(exec_node_id) def _sync_executed_state(self, excluded: Iterable[str] = ()) -> None: @@ -537,7 +539,7 @@ def _project_ready_nodes(self) -> None: def _enqueue_activation_ready_nodes(self) -> None: """Recheck prepared nodes whose activation input arrived after plan registration.""" - for exec_node_id in tuple(self._if_exec_ids): + for exec_node_id in tuple(self._state._pending_if_exec_nodes): if exec_node_id not in self._state.indegree or exec_node_id in self._state.executed: continue self._state._tx_set_mapping( @@ -648,7 +650,7 @@ def complete( # Fresh If branch inputs may be materialized after the If dependency was # registered, so no ordinary plan edge exists to re-enqueue the If when # its selected branch completes. Recheck activation readiness here. - if any(isinstance(node, IfInvocation) for node in self._state.graph.nodes.values()): + if self._if_exec_ids: self._enqueue_activation_ready_nodes() if can_use_nested_iterate_planner(self._state): prepare_nested_iterate_bodies(self._state) diff --git a/tests/app/services/shared/test_execution_engine_differential.py b/tests/app/services/shared/test_execution_engine_differential.py index f2f841c297d..4ef2cd60b71 100644 --- a/tests/app/services/shared/test_execution_engine_differential.py +++ b/tests/app/services/shared/test_execution_engine_differential.py @@ -22,6 +22,7 @@ from invokeai.app.invocations.math import AddInvocation from invokeai.app.invocations.primitives import BooleanInvocation, BooleanOutput from invokeai.app.services.shared import graph as graph_module +from invokeai.app.services.shared import graph_if_runtime from invokeai.app.services.shared.execution_state_migration import ( CURRENT_EXECUTION_STATE_VERSION, UnsupportedExecutionStateVersionError, @@ -38,8 +39,6 @@ _ExecutionScheduler, _GenericGraphSchedulerAdapter, ) -from invokeai.app.services.shared.graph_if_runtime import _get_effective_iteration_paths_for_edge -from invokeai.app.services.shared.graph_iterator_ancestry import _IteratorAncestry from invokeai.app.services.shared.graph_nested_iterate_planner import get_nested_iterate_sequence_depth from invokeai.app.services.shared.invocation_context import InvocationContextData, build_invocation_context from tests.test_nodes import ( @@ -1340,13 +1339,17 @@ def _run_graph( force_compatibility_scheduler: bool = False, stop_after: int | None = None, fail_source_id: str | None = None, + max_steps: int = 100_000, ) -> tuple[list[str], GraphExecutionState]: """Run a constructed graph through either scheduler path.""" if force_compatibility_scheduler: state._execution_scheduler = _ExecutionScheduler(state) trace: list[str] = [] - while (node := state.next()) is not None: + for _ in range(max_steps): + node = state.next() + if node is None: + return trace, state source_id = state.prepared_source_mapping[node.id] trace.append(source_id) if source_id == fail_source_id: @@ -1355,6 +1358,8 @@ def _run_graph( state.complete(node.id, node.invoke(Mock())) if stop_after is not None and len(trace) == stop_after: break + else: + raise AssertionError(f"Execution exceeded {max_steps} scheduled nodes") return trace, state @@ -2533,9 +2538,20 @@ def test_two_sibling_nested_for_fan_in_failure_matches_compatibility() -> None: assert _state_projection(generic_state) == _state_projection(compatibility_state) -def test_nested_for_generic_and_compatibility_paths_handle_empty_inner_collection() -> None: +@pytest.mark.parametrize( + ("outer_collection", "expected"), + [ + ([[], ["c"]], [[], ["c"]]), + ([[], [1]], [[], [1]]), + ([[], []], [[], []]), + ], + ids=["empty-then-value", "empty-then-single-value", "all-empty"], +) +def test_nested_for_generic_and_compatibility_paths_handle_empty_inner_collection( + outer_collection: list[list[Any]], expected: list[list[Any]] +) -> None: def run(force_compatibility: bool) -> tuple[list[str], GraphExecutionState]: - state = GraphExecutionState(graph=_nested_for_graph(outer_collection=[[], ["c"]])) + state = GraphExecutionState(graph=_nested_for_graph(outer_collection=outer_collection)) if force_compatibility: state._execution_scheduler = _ExecutionScheduler(state) return _run(state) @@ -2552,13 +2568,13 @@ def run(force_compatibility: bool) -> tuple[list[str], GraphExecutionState]: for exec_node_id, source_node_id in generic_state.prepared_source_mapping.items() if source_node_id == "after" ) - assert generic_state.results[after_exec_id].value == [[], ["c"]] + assert generic_state.results[after_exec_id].value == expected compatibility_after_exec_id = next( exec_node_id for exec_node_id, source_node_id in compatibility_state.prepared_source_mapping.items() if source_node_id == "after" ) - assert compatibility_state.results[compatibility_after_exec_id].value == [[], ["c"]] + assert compatibility_state.results[compatibility_after_exec_id].value == expected def test_nested_for_generic_path_rehydrates_after_inner_completion() -> None: @@ -6583,9 +6599,10 @@ def test_flat_if_admission_is_demand_driven_and_token_authoritative( ) assert trace == ["condition", "true_branch", "if", "sink"] - assert set(state.source_prepared_mapping) == {"condition", "true_branch", "if", "sink"} - assert "false_branch" not in state.source_prepared_mapping - assert state.executed_history == trace + assert set(state.source_prepared_mapping) == {"condition", "true_branch", "false_branch", "if", "sink"} + assert state.source_prepared_mapping["false_branch"] == set() + assert [source_id for source_id in state.executed_history if source_id != "false_branch"] == trace + assert state.executed_history.count("false_branch") == 1 assert all( state._get_prepared_exec_metadata(exec_node_id).state != "skipped" for exec_node_id in state.prepared_source_mapping @@ -6657,7 +6674,13 @@ def fail_branch_analysis(*_: object, **__: object) -> set[str]: } expected_sources.add("first_true" if first_condition else "first_false") expected_sources.add("second_true" if second_condition else "second_false") + unselected = { + "first_false" if first_condition else "first_true", + "second_false" if second_condition else "second_true", + } + expected_sources.update(unselected) assert set(state.source_prepared_mapping) == expected_sources + assert all(state.source_prepared_mapping[source_id] == set() for source_id in unselected) assert ("first_true" in trace) is first_condition assert ("first_false" in trace) is not first_condition assert ("second_true" in trace) is second_condition @@ -6811,8 +6834,10 @@ def fail_branch_analysis(*_: object, **__: object) -> set[str]: "second_false" if second_condition else "second_true", "third_false" if third_condition else "third_true", } + expected_sources.update(unselected) assert set(state.source_prepared_mapping) == expected_sources - assert set(trace) == expected_sources + assert set(trace) == expected_sources - unselected + assert all(state.source_prepared_mapping[source_id] == set() for source_id in unselected) assert not unselected.intersection(trace) dependencies_by_source = { @@ -7503,7 +7528,13 @@ def fail_branch_analysis(*_: object, **__: object) -> set[str]: runs.append((trace, state)) assert trace == expected_trace - assert set(state.source_prepared_mapping) == set(expected_trace) + expected_source_ids = set(expected_trace) + assert expected_source_ids.issubset(state.source_prepared_mapping) + assert all( + not execution_ids + for source_id, execution_ids in state.source_prepared_mapping.items() + if source_id not in expected_source_ids + ) assert selected in trace sink_id = next(iter(state.source_prepared_mapping["sink"])) assert state.results[sink_id].value == expected_value @@ -7551,7 +7582,13 @@ def fail_branch_analysis(*_: object, **__: object) -> set[str]: "outer_if", "sink", ] - assert set(state.source_prepared_mapping) == set(trace) + expected_source_ids = set(trace) + assert expected_source_ids.issubset(state.source_prepared_mapping) + assert all( + not execution_ids + for source_id, execution_ids in state.source_prepared_mapping.items() + if source_id not in expected_source_ids + ) assert state.results[next(iter(state.source_prepared_mapping["sink"]))].value == 5 assert state.is_complete() @@ -7693,14 +7730,20 @@ def fail_branch_analysis(*_: object, **__: object) -> set[str]: trace, state = _run_graph(state, force_compatibility_scheduler=force_compatibility_scheduler) assert trace == expected_trace - assert set(state.source_prepared_mapping) == set(expected_trace) + expected_source_ids = set(expected_trace) + assert expected_source_ids.issubset(state.source_prepared_mapping) + assert all( + not execution_ids + for source_id, execution_ids in state.source_prepared_mapping.items() + if source_id not in expected_source_ids + ) assert state.results[next(iter(state.source_prepared_mapping["sink"]))].value == expected_value if outer_condition: side_consumer_id = next(iter(state.source_prepared_mapping["middle_side_consumer"])) expected_side_value = expected_value assert state.results[side_consumer_id].value == expected_side_value else: - assert "middle_side_consumer" not in state.source_prepared_mapping + assert state.source_prepared_mapping.get("middle_side_consumer", set()) == set() assert state.is_complete() @@ -8060,7 +8103,7 @@ def record_generic_discard(adapter: _GenericGraphSchedulerAdapter, exec_node_id: trace, state = _run_graph(GraphExecutionState(graph=graph)) assert trace == expected_trace - assert state.executed_history == expected_trace + assert state.executed_history == expected_history assert set(state.indegree) == set(state.prepared_source_mapping) assert set(state.indegree.values()) == {0} assert state.is_complete() @@ -8368,7 +8411,13 @@ def fail_controller_access(*_: object, **__: object) -> Any: runs.append((trace, state)) assert trace == expected_trace - assert set(state.source_prepared_mapping) == set(expected_trace) + expected_source_ids = set(expected_trace) + assert expected_source_ids.issubset(state.source_prepared_mapping) + assert all( + not execution_ids + for source_id, execution_ids in state.source_prepared_mapping.items() + if source_id not in expected_source_ids + ) assert all( state._get_prepared_exec_metadata(exec_node_id).state != "skipped" for exec_node_id in state.prepared_source_mapping @@ -8494,19 +8543,7 @@ def fail_branch_analysis(*_: object, **__: object) -> set[str]: assert resumed_state.is_complete() -def _run_to_completion( - state: GraphExecutionState, *, max_steps: int = 100_000 -) -> tuple[list[str], GraphExecutionState]: - trace: list[str] = [] - for _ in range(max_steps): - node = state.next() - if node is None: - return trace, state - trace.append(state.prepared_source_mapping[node.id]) - state.complete(node.id, node.invoke(Mock())) - raise AssertionError(f"Execution exceeded {max_steps} scheduled nodes") - - +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) @pytest.mark.parametrize( ("collection", "expected"), [ @@ -8517,7 +8554,7 @@ def _run_to_completion( ids=["mixed", "all-true", "all-false"], ) def test_for_condition_if_reads_root_frame_sources_and_skips_unselected_root( - collection: list[bool], expected: list[str] + force_compatibility_scheduler: bool, collection: list[bool], expected: list[str] ) -> None: graph = Graph() graph.add_node(ForInvocation(id="for", collection=collection)) @@ -8533,14 +8570,56 @@ def test_for_condition_if_reads_root_frame_sources_and_skips_unselected_root( graph.add_edge(create_loop_linkage("for", "return")) graph.add_edge(create_edge("for", "output_collection", "after", "value")) - trace, state = _run_to_completion(GraphExecutionState(graph=graph)) + trace, state = _run_graph( + GraphExecutionState(graph=graph), force_compatibility_scheduler=force_compatibility_scheduler + ) - assert state.is_complete() + assert state.is_complete(), ( + f"graph stalled after {trace}; history={state.executed_history}; " + f"sources={state.source_prepared_mapping}; pending_if={state._pending_if_exec_nodes}; " + f"resolved={state._resolved_if_exec_branches}; indegree={state.indegree}" + ) assert _source_output(state, "after").value == expected assert trace.count("root_true") == (1 if True in collection else 0) assert trace.count("root_false") == (1 if False in collection else 0) +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +@pytest.mark.parametrize( + ("conditions", "expected"), + [ + ([False, True], ["false-root", "true-root"]), + ([True, True], ["true-root", "true-root"]), + ([False, False], ["false-root", "false-root"]), + ], + ids=["mixed", "all-true", "all-false"], +) +def test_iterate_condition_if_reads_root_frame_sources_and_skips_unselected_root( + force_compatibility_scheduler: bool, conditions: list[bool], expected: list[str] +) -> None: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="conditions", first=conditions)) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(AnyTypeTestInvocation(id="root_true", value="true-root")) + graph.add_node(AnyTypeTestInvocation(id="root_false", value="false-root")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="output")) + graph.add_edge(create_edge("conditions", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "if", "condition")) + graph.add_edge(create_edge("root_true", "value", "if", "true_input")) + graph.add_edge(create_edge("root_false", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "output", "item")) + + trace, state = _run_graph( + GraphExecutionState(graph=graph), force_compatibility_scheduler=force_compatibility_scheduler + ) + + assert state.is_complete(), f"graph stalled after {trace}; history={state.executed_history}" + assert _source_output(state, "output").collection == expected + assert trace.count("root_true") == (1 if True in conditions else 0) + assert trace.count("root_false") == (1 if False in conditions else 0) + + @pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) def test_iterate_for_collect_preserves_empty_inner_collection(force_compatibility_scheduler: bool) -> None: graph = Graph() @@ -8562,7 +8641,7 @@ def test_iterate_for_collect_preserves_empty_inner_collection(force_compatibilit state = GraphExecutionState(graph=graph) if force_compatibility_scheduler: state._execution_scheduler = _ExecutionScheduler(state) - trace, state = _run_to_completion(state) + trace, state = _run_graph(state) after_exec_ids = sorted(state.source_prepared_mapping["after"], key=state._get_iteration_path) assert trace.count("add") == 2 @@ -8640,7 +8719,7 @@ def test_sibling_iterates_collect_through_if_branch_without_unsupported_frames( state = GraphExecutionState(graph=graph) if force_compatibility_scheduler: state._execution_scheduler = _ExecutionScheduler(state) - trace, state = _run_to_completion(state) + trace, state = _run_graph(state) assert state.is_complete() assert state.results[next(iter(state.source_prepared_mapping["collect"]))].collection == expected @@ -8677,9 +8756,7 @@ def test_nested_iterate_collected_in_if_branch_is_not_misclassified_as_sibling() state = GraphExecutionState(graph=graph) - assert state._materializer()._get_sibling_iterate_if_collect_branch("inner") is None - - trace, state = _run_to_completion(state) + trace, state = _run_graph(state) assert state.is_complete() assert state.results[next(iter(state.source_prepared_mapping["result_collect"]))].collection == [ @@ -8728,11 +8805,19 @@ def test_nested_collect_preserves_condition_frame_when_classifying_branch_iterat if force_compatibility_scheduler: state._execution_scheduler = _ExecutionScheduler(state) - assert state._materializer()._get_sibling_iterate_if_collect_branch("inner") is None - - trace, state = _run_to_completion(state) + trace, state = _run_graph(state) - assert state.is_complete() + assert state.is_complete(), ( + f"graph stalled after {trace}; pending={state._pending_if_exec_nodes}; " + f"prepared={state.source_prepared_mapping}; indegree={state.indegree}; " + f"resolved={state._resolved_if_exec_branches}; " + f"unindexed={set(state.prepared_source_mapping) - state._if_activation_dependents_indexed_exec_ids}; " + f"middle_axes={state._materializer().get_node_iterators('middle_body', state._get_source_graph_flat())}; " + f"middle_mappings={list(state._materializer()._get_parent_iteration_mappings('middle_body', state._get_source_graph_flat()))}; " + f"middle_dependencies={state._get_source_activation_dependencies('middle_body', (0,))}; " + f"middle_admitted={state._is_source_activation_admitted('middle_body', (0,))}; " + f"iterator_sources_executed={[(source_id, source_id in state.executed) for source_id in ('left', 'right')]}" + ) middle_collect_exec_ids = sorted(state.source_prepared_mapping["middle_collect"], key=state._get_iteration_path) assert [ (state._get_iteration_path(exec_id), state.results[exec_id].collection) for exec_id in middle_collect_exec_ids @@ -8784,9 +8869,7 @@ def test_direct_collect_collapses_condition_frame_for_branch_iterate( if force_compatibility_scheduler: state._execution_scheduler = _ExecutionScheduler(state) - assert state._materializer()._get_sibling_iterate_if_collect_branch("inner") == ("if", "true_input") - - trace, state = _run_to_completion(state) + trace, state = _run_graph(state) assert state.is_complete() middle_collect_id = next(iter(state.source_prepared_mapping["middle_collect"])) @@ -8853,25 +8936,15 @@ def test_independent_iterator_collect_preserves_ordered_parent_axis_for_branch_i ("fallback", "value", "if", "false_input"), ("if", "value", "result_collect", "item"), ): - if node_order == "renamed" and (source, source_field, destination, destination_field) == ( - condition_axis_id, - "item", - "condition", - "value", - ): - continue graph.add_edge(create_edge(source, source_field, destination, destination_field)) state = GraphExecutionState(graph=graph) - if node_order == "renamed": - state._materializer()._get_sibling_iterate_if_collect_branches() - state.add_edge(create_edge(condition_axis_id, "item", "condition", "value")) if force_compatibility_scheduler: state._execution_scheduler = _ExecutionScheduler(state) - trace, state = _run_to_completion(state) + trace, state = _run_graph(state) - assert state.is_complete() + assert state.is_complete(), f"graph stalled after {trace}" middle_collect_exec_ids = sorted(state.source_prepared_mapping["middle_collect"], key=state._get_iteration_path) assert [ (state._get_iteration_path(exec_id), state.results[exec_id].collection) for exec_id in middle_collect_exec_ids @@ -8889,7 +8962,8 @@ def test_independent_iterator_collect_preserves_ordered_parent_axis_for_branch_i ] -def test_if_does_not_combine_unrelated_open_iterator_frames(): +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_if_rejects_unrelated_open_iterator_axes(force_compatibility_scheduler: bool) -> None: graph = Graph() graph.add_node(CollectionConcatInvocation(id="branch_source", first=[1, 2])) graph.add_node(IterateInvocation(id="branch_iterate")) @@ -8913,13 +8987,102 @@ def test_if_does_not_combine_unrelated_open_iterator_frames(): graph.add_edge(create_edge(source, source_field, destination, destination_field)) graph.validate_self() - trace, state = _run_to_completion(GraphExecutionState(graph=graph)) - assert not state.is_complete() - assert not state.errors - assert "branch_value" not in state.source_prepared_mapping - assert "collect" not in state.source_prepared_mapping - assert "branch_value" not in trace - assert "collect" not in trace + state = GraphExecutionState(graph=graph) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) + with pytest.raises( + graph_if_runtime.UnsupportedIfIterationAxesError, + match="cannot select branch input 'true_input' from unrelated open iterator axes", + ): + _run_graph(state) + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_if_rejects_sibling_inner_iterator_axes_with_shared_outer_axis(force_compatibility_scheduler: bool) -> None: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="outer_source", first=["truth"])) + graph.add_node(IterateInvocation(id="outer")) + graph.add_node(PolymorphicStringTestInvocation(id="left_source")) + graph.add_node(IterateInvocation(id="left")) + graph.add_node(FixedCollectionTestInvocation(id="right_source")) + graph.add_node(IterateInvocation(id="right")) + graph.add_node(FixedCollectionTestInvocation(id="condition")) + graph.add_node(AnyTypeTestInvocation(id="branch")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value="fallback")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="collect")) + for source, source_field, destination, destination_field in ( + ("outer_source", "collection", "outer", "collection"), + ("outer", "item", "left_source", "value"), + ("left_source", "collection", "left", "collection"), + ("outer", "item", "right_source", "value"), + ("right_source", "collection", "right", "collection"), + ("left", "item", "condition", "value"), + ("right", "item", "branch", "value"), + ("condition", "condition", "if", "condition"), + ("branch", "value", "if", "true_input"), + ("fallback", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + + graph.validate_self() + state = GraphExecutionState(graph=graph) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) + with pytest.raises( + graph_if_runtime.UnsupportedIfIterationAxesError, + match="cannot select branch input 'true_input' from unrelated open iterator axes", + ): + _run_graph(state) + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_sibling_collect_cannot_feed_multiple_if_branches(force_compatibility_scheduler: bool) -> None: + graph = Graph() + for node in ( + CollectionConcatInvocation(id="left_source", first=[0, 1]), + CollectionConcatInvocation(id="right_source", first=[10, 20]), + IterateInvocation(id="left"), + IterateInvocation(id="right"), + AddInvocation(id="body"), + CollectInvocation(id="branch_collect"), + FixedCollectionTestInvocation(id="left_condition"), + AddInvocation(id="right_minus_one", b=-1), + FixedCollectionTestInvocation(id="right_condition"), + AnyTypeTestInvocation(id="fallback", value="fallback"), + IfInvocation(id="left_if"), + IfInvocation(id="right_if"), + CollectInvocation(id="left_result"), + CollectInvocation(id="right_result"), + ): + graph.add_node(node) + for source, source_field, destination, destination_field in ( + ("left_source", "collection", "left", "collection"), + ("right_source", "collection", "right", "collection"), + ("left", "item", "body", "a"), + ("right", "item", "body", "b"), + ("body", "value", "branch_collect", "item"), + ("left", "item", "left_condition", "value"), + ("left", "item", "right_minus_one", "a"), + ("right_minus_one", "value", "right_condition", "value"), + ("left_condition", "condition", "left_if", "condition"), + ("right_condition", "condition", "right_if", "condition"), + ("branch_collect", "collection", "left_if", "true_input"), + ("branch_collect", "collection", "right_if", "true_input"), + ("fallback", "value", "left_if", "false_input"), + ("fallback", "value", "right_if", "false_input"), + ("left_if", "value", "left_result", "item"), + ("right_if", "value", "right_result", "item"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + + graph.validate_self() + state = GraphExecutionState(graph=graph) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) + with pytest.raises(ValueError, match="cannot supply Collect-backed inputs to multiple If branches"): + _run_graph(state) def test_deep_any_chain_inside_iterate_if_completes() -> None: @@ -8944,7 +9107,7 @@ def test_deep_any_chain_inside_iterate_if_completes() -> None: graph.add_edge(create_edge(f"chain_{chain_length - 1}", "value", "if", "true_input")) graph.add_edge(create_edge("if", "value", "collect", "item")) - _, state = _run_to_completion(GraphExecutionState(graph=graph)) + _, state = _run_graph(GraphExecutionState(graph=graph)) assert state.results[next(iter(state.source_prepared_mapping["collect"]))].collection == [True, "fFalse"] assert state.is_complete() @@ -8952,7 +9115,7 @@ def test_deep_any_chain_inside_iterate_if_completes() -> None: def test_execution_materializer_cache_refreshes_after_sibling_iterate_edge_addition() -> None: graph = Graph() - graph.add_node(CollectionConcatInvocation(id="left_source", first=[1, 1])) + graph.add_node(CollectionConcatInvocation(id="left_source", first=[2, 2])) graph.add_node(CollectionConcatInvocation(id="right_source", first=[10, 20])) graph.add_node(IterateInvocation(id="left")) graph.add_node(IterateInvocation(id="right")) @@ -8978,51 +9141,449 @@ def test_execution_materializer_cache_refreshes_after_sibling_iterate_edge_addit graph.add_edge(create_edge(source, source_field, destination, destination_field)) state = GraphExecutionState(graph=graph) - original_materializer = state._materializer() - source_graph = state._get_source_graph_flat() - assert not original_materializer._is_inactive_sibling_if_collect_source("right", source_graph) - assert original_materializer._get_sibling_iterate_if_collect_branch("right") is None + first_node = state.next() + assert first_node is not None + first_source = state.prepared_source_mapping[first_node.id] + state.complete(first_node.id, first_node.invoke(Mock())) state.add_edge(create_edge("right", "item", "body", "b")) - materializer = state._materializer() - assert materializer is not original_materializer - assert materializer._get_sibling_iterate_if_collect_branch("right") == ("if", "true_input") - - trace, state = _run_to_completion(state) + trace, state = _run_graph(state) - assert state.results[next(iter(state.source_prepared_mapping["collect"]))].collection == ["fallback", "fallback"] - assert "right" not in trace - assert "body" not in trace + collect_exec_id = next(iter(state.source_prepared_mapping["collect"])) + assert state.results[collect_exec_id].collection == [[12, 22], [12, 22]] + assert "right" in trace or first_source == "right" + assert "body" in trace assert state.is_complete() -def test_effective_iteration_paths_handles_deep_any_chain() -> None: - chain_length = 1000 +def test_deep_if_branch_chain_completes_through_execution_state() -> None: + chain_length = 150 graph = Graph() - graph.add_node(CollectionConcatInvocation(id="source", first=[True])) - graph.add_node(IterateInvocation(id="iterate")) - graph.add_node(IfInvocation(id="if")) + graph.add_node(AnyTypeTestInvocation(id="source", value=7)) + graph.add_node(IfInvocation(id="if", condition=True)) + graph.add_node(AnyTypeTestInvocation(id="result")) for index in range(chain_length): graph.add_node(AnyTypeTestInvocation(id=f"chain_{index}")) - graph.add_edge(create_edge("source", "collection", "iterate", "collection")) - graph.add_edge(create_edge("iterate", "item", "chain_0", "value")) + graph.add_edge(create_edge("source", "value", "chain_0", "value")) for index in range(chain_length - 1): graph.add_edge(create_edge(f"chain_{index}", "value", f"chain_{index + 1}", "value")) - terminal_edge = create_edge(f"chain_{chain_length - 1}", "value", "if", "true_input") - graph.add_edge(terminal_edge) + graph.add_edge(create_edge(f"chain_{chain_length - 1}", "value", "if", "true_input")) + graph.add_edge(create_edge("if", "value", "result", "value")) + + trace, state = _run_graph(GraphExecutionState(graph=graph)) + + result_exec_id = next(iter(state.source_prepared_mapping["result"])) + assert state.results[result_exec_id].value == 7 + assert state.is_complete(), f"graph stalled after {trace}" + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_if_diamond_branch_traversal_scales_with_graph_size( + force_compatibility_scheduler: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + diamond_count = 16 + graph = Graph() + graph.add_node(AnyTypeTestInvocation(id="source", value=1)) + graph.add_node(IfInvocation(id="if", condition=True)) + graph.add_node(AnyTypeTestInvocation(id="result")) + previous_node_id = "source" + for index in range(diamond_count): + left_id = f"left_{index}" + right_id = f"right_{index}" + merge_id = f"merge_{index}" + graph.add_node(AnyTypeTestInvocation(id=left_id)) + graph.add_node(AnyTypeTestInvocation(id=right_id)) + graph.add_node(AddInvocation(id=merge_id)) + graph.add_edge(create_edge(previous_node_id, "value", left_id, "value")) + graph.add_edge(create_edge(previous_node_id, "value", right_id, "value")) + graph.add_edge(create_edge(left_id, "value", merge_id, "a")) + graph.add_edge(create_edge(right_id, "value", merge_id, "b")) + previous_node_id = merge_id + graph.add_edge(create_edge(previous_node_id, "value", "if", "true_input")) + graph.add_edge(create_edge("if", "value", "result", "value")) state = GraphExecutionState(graph=graph) - state.source_prepared_mapping["iterate"] = {"prepared_iterate"} - state.prepared_iteration_paths["prepared_iterate"] = (7,) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) - effective_paths = _get_effective_iteration_paths_for_edge( - state, - terminal_edge, - set(), - lambda source_id: state.source_prepared_mapping.get(source_id, set()), - lambda _edge, prepared_id: state.prepared_iteration_paths[prepared_id], - _IteratorAncestry(state), + original_get_node = Graph.get_node + original_if_paths = graph_if_runtime._get_if_iteration_paths_with_ancestry + if_path_node_visits = 0 + if_path_depth = 0 + + def count_node_visits(self: Graph, node_id: str): + nonlocal if_path_node_visits + if if_path_depth: + if_path_node_visits += 1 + return original_get_node(self, node_id) + + def count_if_path_work(*args: Any, **kwargs: Any): + nonlocal if_path_depth + if_path_depth += 1 + try: + return original_if_paths(*args, **kwargs) + finally: + if_path_depth -= 1 + + monkeypatch.setattr(Graph, "get_node", count_node_visits) + monkeypatch.setattr(graph_if_runtime, "_get_if_iteration_paths_with_ancestry", count_if_path_work) + trace, state = _run_graph(state) + + result_exec_id = next(iter(state.source_prepared_mapping["result"])) + assert state.results[result_exec_id].value == 2**diamond_count + assert state.is_complete(), f"graph stalled after {trace}" + assert if_path_node_visits < 100 * len(graph.nodes), ( + f"If path inference visited {if_path_node_visits} graph nodes for {len(graph.nodes)} nodes" ) - assert effective_paths == {(7,)} + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +@pytest.mark.parametrize( + ("xs", "ys", "expected"), + [ + ([1, 1, 0], [10, 20, 30], [[10, 20, 30], [10, 20, 30], "fb"]), + ([1, 0], [10, 20], [[10, 20], "fb"]), + ([0, 1], [10, 20], ["fb", [10, 20]]), + ], + ids=["false-last", "false-last-short", "false-first"], +) +def test_per_item_if_preserves_unrelated_root_collect_frames( + force_compatibility_scheduler: bool, + xs: list[int], + ys: list[int], + expected: list[Any], +) -> None: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="xs", first=xs)) + graph.add_node(IterateInvocation(id="x")) + graph.add_node(CollectionConcatInvocation(id="ys", first=ys)) + graph.add_node(IterateInvocation(id="y")) + graph.add_node(AnyTypeTestInvocation(id="yv")) + graph.add_node(CollectInvocation(id="branch_collect")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value="fb")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="output")) + for source, source_field, destination, destination_field in ( + ("xs", "collection", "x", "collection"), + ("x", "item", "if", "condition"), + ("ys", "collection", "y", "collection"), + ("y", "item", "yv", "value"), + ("yv", "value", "branch_collect", "item"), + ("branch_collect", "collection", "if", "true_input"), + ("fallback", "value", "if", "false_input"), + ("if", "value", "output", "item"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + + state = GraphExecutionState(graph=graph) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) + trace, state = _run_graph(state) + + assert state.is_complete(), f"graph stalled after {trace}" + output_exec_id = next(iter(state.source_prepared_mapping["output"])) + assert state.results[output_exec_id].collection == expected + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +@pytest.mark.parametrize( + ("conditions", "expected"), + [ + ([False, True, True], ["fb", [5, 6], [5, 6]]), + ([True, True], [[5, 6], [5, 6]]), + ], + ids=["mixed-conditions", "all-true"], +) +def test_per_item_if_keeps_closed_for_output_collection( + force_compatibility_scheduler: bool, + conditions: list[bool], + expected: list[Any], +) -> None: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="conditions", first=conditions)) + graph.add_node(IterateInvocation(id="condition_iterate")) + graph.add_node(ForInvocation(id="for", collection=[5, 6], state=LoopState(values={"count": 0}))) + graph.add_node(AnyTypeTestInvocation(id="body")) + graph.add_node(ForReturnInvocation(id="for_return")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value="fb")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="output")) + for source, source_field, destination, destination_field in ( + ("conditions", "collection", "condition_iterate", "collection"), + ("condition_iterate", "item", "if", "condition"), + ("for", "item", "body", "value"), + ("body", "value", "for_return", "output"), + ("for", "output_collection", "if", "true_input"), + ("fallback", "value", "if", "false_input"), + ("if", "value", "output", "item"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + graph.add_edge(create_loop_linkage("for", "for_return")) + + state = GraphExecutionState(graph=graph) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) + trace, state = _run_graph(state) + + assert state.is_complete(), f"graph stalled after {trace}" + output_exec_id = next(iter(state.source_prepared_mapping["output"])) + assert state.results[output_exec_id].collection == expected + final_for_exec_id = max( + state.source_prepared_mapping["for"], + key=lambda exec_node_id: state.execution_graph.get_node(exec_node_id).index, + ) + assert state.results[final_for_exec_id].output_collection == [5, 6] + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +@pytest.mark.parametrize("extra_source", ["right", "branch_collect"]) +def test_deferred_sibling_iterate_does_not_starve_ready_consumers( + force_compatibility_scheduler: bool, + extra_source: str, +) -> None: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="left_source", first=[1, 2])) + graph.add_node(CollectionConcatInvocation(id="right_source", first=[10, 20])) + graph.add_node(IterateInvocation(id="left")) + graph.add_node(IterateInvocation(id="right")) + graph.add_node(AddInvocation(id="body")) + graph.add_node(CollectInvocation(id="branch_collect")) + graph.add_node(AddInvocation(id="minus_one", b=-1)) + graph.add_node(FixedCollectionTestInvocation(id="condition")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value="fallback")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="output")) + graph.add_node(AnyTypeTestInvocation(id="extra")) + for source, source_field, destination, destination_field in ( + ("left_source", "collection", "left", "collection"), + ("right_source", "collection", "right", "collection"), + ("left", "item", "body", "a"), + ("right", "item", "body", "b"), + ("body", "value", "branch_collect", "item"), + ("left", "item", "minus_one", "a"), + ("minus_one", "value", "condition", "value"), + ("condition", "condition", "if", "condition"), + ("branch_collect", "collection", "if", "true_input"), + ("fallback", "value", "if", "false_input"), + ("if", "value", "output", "item"), + (extra_source, "item" if extra_source == "right" else "collection", "extra", "value"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + + state = GraphExecutionState(graph=graph) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) + trace, state = _run_graph(state) + + assert state.is_complete(), f"graph stalled after {trace}" + output_exec_id = next(iter(state.source_prepared_mapping["output"])) + assert state.results[output_exec_id].collection == ["fallback", [12, 22]] + assert "extra" in trace + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +@pytest.mark.parametrize("reverse_node_insertion", [False, True], ids=["forward-insertion", "reverse-insertion"]) +@pytest.mark.parametrize( + ("source_a", "iterate_a", "source_b", "iterate_b"), + [ + ("zz_source", "zz_iterate", "aa_source", "aa_iterate"), + ("p_source", "q_iterate", "x_source", "y_iterate"), + ], + ids=["reverse-lexical-ids", "alternate-ids"], +) +def test_independent_iterate_axes_follow_add_input_order( + force_compatibility_scheduler: bool, + reverse_node_insertion: bool, + source_a: str, + iterate_a: str, + source_b: str, + iterate_b: str, +) -> None: + nodes = [ + CollectionConcatInvocation(id=source_a, first=[10, 20]), + IterateInvocation(id=iterate_a), + CollectionConcatInvocation(id=source_b, first=[1, 2, 3]), + IterateInvocation(id=iterate_b), + AddInvocation(id="sum"), + CollectInvocation(id="output"), + ] + if reverse_node_insertion: + nodes.reverse() + graph = Graph(nodes={node.id: node for node in nodes}) + for source, source_field, destination, destination_field in ( + (source_a, "collection", iterate_a, "collection"), + (source_b, "collection", iterate_b, "collection"), + (iterate_a, "item", "sum", "a"), + (iterate_b, "item", "sum", "b"), + ("sum", "value", "output", "item"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + + state = GraphExecutionState(graph=graph) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) + _, state = _run_graph(state) + + output_exec_ids = sorted(state.source_prepared_mapping["output"], key=state._get_iteration_path) + actual = [(state._get_iteration_path(exec_id), state.results[exec_id].collection) for exec_id in output_exec_ids] + assert actual == [((0,), [11, 12, 13]), ((1,), [21, 22, 23])] + assert state.is_complete() + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +@pytest.mark.parametrize( + ("condition", "wire_opposite_root"), + [(True, True), (False, True), (False, False)], + ids=["true-with-root-false", "false-with-root-true", "false-without-true-input"], +) +def test_literal_if_inside_iterate_uses_selected_branch_frame( + force_compatibility_scheduler: bool, + condition: bool, + wire_opposite_root: bool, +) -> None: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="source", first=[1, 2])) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(IfInvocation(id="if", condition=condition)) + graph.add_node(CollectInvocation(id="output")) + graph.add_edge(create_edge("source", "collection", "iterate", "collection")) + selected_branch = "true_input" if condition else "false_input" + selected_source = "true_branch" if condition else "false_branch" + opposite_branch = "false_input" if condition else "true_input" + graph.add_node(MarkedAnyTypeTestInvocation(id=selected_source, marker=f"{selected_branch}:")) + graph.add_edge(create_edge("iterate", "item", selected_source, "value")) + graph.add_edge(create_edge(selected_source, "value", "if", selected_branch)) + if wire_opposite_root: + graph.add_node(AnyTypeTestInvocation(id="root_opposite", value="root")) + graph.add_edge(create_edge("root_opposite", "value", "if", opposite_branch)) + graph.add_edge(create_edge("if", "value", "output", "item")) + + state = GraphExecutionState(graph=graph) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) + trace, state = _run_graph(state) + + assert state.is_complete(), ( + f"graph stalled after {trace}; prepared_paths={[(source_id, [state._get_iteration_path(exec_id) for exec_id in exec_ids]) for source_id, exec_ids in state.source_prepared_mapping.items()]}; " + f"indegree={state.indegree}; pending_if={state._pending_if_exec_nodes}; resolved={state._resolved_if_exec_branches}" + ) + output_exec_id = next(iter(state.source_prepared_mapping["output"])) + marker = f"{selected_branch}:" + assert state.results[output_exec_id].collection == [f"{marker}1", f"{marker}2"] + assert trace.count("if") == 2 + assert trace.count(selected_source) == 2 + assert "root_opposite" not in trace + assert state.is_complete(), f"graph stalled after {trace}" + assert state.next() is None + assert not state._ready_node_ids + assert all(not queue for queue in state._ready_queues.values()) + + +@pytest.mark.parametrize("round_trip_after", [None, 6], ids=["fresh", "restored"]) +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_literal_if_branch_feeds_per_item_if_without_stranding_unselected_root( + force_compatibility_scheduler: bool, + round_trip_after: int | None, +) -> None: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="conditions", first=[True, True])) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(AnyTypeTestInvocation(id="literal_selected", value="literal")) + graph.add_node(AnyTypeTestInvocation(id="literal_unselected", value="literal-unselected")) + graph.add_node(IfInvocation(id="literal_if", condition=True)) + graph.add_node(AnyTypeTestInvocation(id="per_item_unselected", value="per-item-unselected")) + graph.add_node(IfInvocation(id="per_item_if")) + graph.add_node(CollectInvocation(id="output")) + for source, source_field, destination, destination_field in ( + ("conditions", "collection", "iterate", "collection"), + ("literal_selected", "value", "literal_if", "true_input"), + ("literal_unselected", "value", "literal_if", "false_input"), + ("iterate", "item", "per_item_if", "condition"), + ("literal_if", "value", "per_item_if", "true_input"), + ("per_item_unselected", "value", "per_item_if", "false_input"), + ("per_item_if", "value", "output", "item"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + + state = GraphExecutionState(graph=graph) + if round_trip_after is None: + trace, state = _run_graph(state, force_compatibility_scheduler=force_compatibility_scheduler) + else: + checkpoint_trace, partial_state = _run_graph( + state, + force_compatibility_scheduler=force_compatibility_scheduler, + stop_after=round_trip_after, + ) + assert "per_item_if" in checkpoint_trace + restored = load_execution_state(dump_execution_state(partial_state)) + if force_compatibility_scheduler: + _restore_compatibility_scheduler(restored) + resumed_trace, state = _run_graph(restored, force_compatibility_scheduler=force_compatibility_scheduler) + trace = checkpoint_trace + resumed_trace + + assert state.is_complete(), f"graph stalled after {trace}; pending={state._pending_if_exec_nodes}" + output_exec_id = next(iter(state.source_prepared_mapping["output"])) + assert state.results[output_exec_id].collection == ["literal", "literal"] + assert trace.count("literal_if") == 1 + assert trace.count("per_item_if") == 2 + assert state.source_prepared_mapping["literal_unselected"] == set() + assert state.source_prepared_mapping["per_item_unselected"] == set() + assert "literal_unselected" not in trace + assert "per_item_unselected" not in trace + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +@pytest.mark.parametrize("extra_inner_consumer", [False, True], ids=["no-extra-consumer", "extra-inner-consumer"]) +def test_nested_if_collect_branch_with_all_outer_frames_false_completes( + force_compatibility_scheduler: bool, + extra_inner_consumer: bool, +) -> None: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="outer_source", first=["", ""])) + graph.add_node(IterateInvocation(id="outer")) + graph.add_node(PolymorphicStringTestInvocation(id="inner_collection")) + graph.add_node(IterateInvocation(id="inner")) + graph.add_node(AnyTypeTestInvocation(id="body")) + graph.add_node(CollectInvocation(id="branch_collect")) + graph.add_node(FixedCollectionTestInvocation(id="condition")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value="fallback")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="output")) + for source, source_field, destination, destination_field in ( + ("outer_source", "collection", "outer", "collection"), + ("outer", "item", "inner_collection", "value"), + ("inner_collection", "collection", "inner", "collection"), + ("inner", "item", "body", "value"), + ("body", "value", "branch_collect", "item"), + ("outer", "item", "condition", "value"), + ("condition", "condition", "if", "condition"), + ("branch_collect", "collection", "if", "true_input"), + ("fallback", "value", "if", "false_input"), + ("if", "value", "output", "item"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + if extra_inner_consumer: + graph.add_node(AnyTypeTestInvocation(id="extra")) + graph.add_edge(create_edge("inner", "item", "extra", "value")) + + state = GraphExecutionState(graph=graph) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) + trace, state = _run_graph(state) + + assert state.is_complete(), ( + f"graph stalled after {trace}; prepared={state.source_prepared_mapping}; pending_if={state._pending_if_exec_nodes}; " + f"inactive_inner={state._is_source_inactive('inner')}; " + f"inner_dependencies={state._get_source_activation_dependencies('inner')}" + ) + output_exec_id = next(iter(state.source_prepared_mapping["output"])) + assert state.results[output_exec_id].collection == ["fallback", "fallback"] + assert "branch_collect" not in trace + assert "body" not in trace + assert ("inner" in trace) is extra_inner_consumer + assert ("extra" in trace) is extra_inner_consumer + assert state.is_complete(), f"graph stalled after {trace}" + assert state.next() is None + assert not state._ready_node_ids + assert all(not queue for queue in state._ready_queues.values()) diff --git a/tests/app/services/shared/test_execution_engine_scheduler.py b/tests/app/services/shared/test_execution_engine_scheduler.py index aa335d26dd6..c51fc026e11 100644 --- a/tests/app/services/shared/test_execution_engine_scheduler.py +++ b/tests/app/services/shared/test_execution_engine_scheduler.py @@ -148,7 +148,7 @@ def test_graph_state_generic_if_readiness_uses_activation_dependency_records(mon true_exec_id = next(iter(state.source_prepared_mapping["true_value"])) plan_node = state._generic_graph_scheduler._scheduler.plan.nodes[true_exec_id] assert plan_node.activation_dependencies == (ActivationDependency(owner_id="if", branch="true_input", frame=()),) - assert "false_value" not in state.source_prepared_mapping + assert state.source_prepared_mapping["false_value"] == set() if_exec_id = next( execution_id for execution_id, source_id in state.prepared_source_mapping.items() if source_id == "if" @@ -811,7 +811,7 @@ def test_graph_state_if_rehydrates_discarded_unselected_branch() -> None: assert selected_branch is not None selected_branch_exec_id = selected_branch.id assert state.prepared_source_mapping[selected_branch_exec_id] == "false_value" - assert "true_value" not in state.source_prepared_mapping + assert state.source_prepared_mapping["true_value"] == set() restored = GraphExecutionState.model_validate(state.model_dump(mode="python"), strict=False) diff --git a/tests/app/services/shared/test_for_scheduler_performance.py b/tests/app/services/shared/test_for_scheduler_performance.py index c2de4d3706c..7fdb8464355 100644 --- a/tests/app/services/shared/test_for_scheduler_performance.py +++ b/tests/app/services/shared/test_for_scheduler_performance.py @@ -35,6 +35,27 @@ def record(self, mapping_name: str) -> None: self.visits_by_mapping[mapping_name] = self.visits_by_mapping.get(mapping_name, 0) + 1 +class _ActivationDependencyScanCounter: + def __init__(self) -> None: + self.active = False + self.entries_visited = 0 + self.resolutions = 0 + self.prepared_node_count = 0 + self.states_by_id: dict[int, _ActivationDependencyScanCounter] = {} + + +class _PreparedSourceCountingMapping(dict[str, str]): + def __init__(self, values: dict[str, str], counter: _ActivationDependencyScanCounter) -> None: + super().__init__(values) + self._counter = counter + + def __iter__(self): # type: ignore[override] + for key in super().__iter__(): + if self._counter.active: + self._counter.entries_visited += 1 + yield key + + class _CountingMapping(dict[str, object]): """Count mapping accesses and iteration without constraining the scheduler algorithm.""" @@ -99,6 +120,8 @@ def _run_trivial_loop( mapping_counter: _MappingTraversalCounter | None = None, include_if: bool = False, external_condition: bool = False, + condition_value: bool = True, + activation_scan_counter: _ActivationDependencyScanCounter | None = None, ) -> float: graph = Graph() graph.add_node(RangeInvocation(id="range", start=0, stop=count)) @@ -107,9 +130,11 @@ def _run_trivial_loop( if include_if: graph.add_node(AnyTypeTestInvocation(id="true_body")) graph.add_node(AnyTypeTestInvocation(id="false_body")) - graph.add_node(IfInvocation(id="if") if external_condition else IfInvocation(id="if", condition=True)) + graph.add_node( + IfInvocation(id="if") if external_condition else IfInvocation(id="if", condition=condition_value) + ) if external_condition: - graph.add_node(BooleanInvocation(id="condition", value=True)) + graph.add_node(BooleanInvocation(id="condition", value=condition_value)) graph.add_edge(create_edge("condition", "value", "if", "condition")) graph.add_edge(create_edge("loop", "item", "true_body", "value")) graph.add_edge(create_edge("loop", "item", "false_body", "value")) @@ -128,6 +153,12 @@ def _run_trivial_loop( graph.add_node(CollectInvocation(id="collect")) graph.add_edge(create_edge(output_source, "value", "collect", "item")) state = GraphExecutionState(graph=graph) + if activation_scan_counter is not None: + state.prepared_source_mapping = _PreparedSourceCountingMapping( + state.prepared_source_mapping, activation_scan_counter + ) # type: ignore[assignment] + state._prepared_registry()._prepared_source_mapping = state.prepared_source_mapping + activation_scan_counter.states_by_id[id(state)] = activation_scan_counter if mapping_counter is not None: scheduler = state._scheduler() state.execution_refs = _CountingMapping(state.execution_refs, mapping_counter, "execution_refs") # type: ignore[assignment] @@ -143,9 +174,103 @@ def _run_trivial_loop( while (node := state.next()) is not None: state.complete(node.id, node.invoke(context)) assert state.is_complete() + if activation_scan_counter is not None: + activation_scan_counter.prepared_node_count = len(state.prepared_source_mapping) return clock() - started +@pytest.mark.parametrize("loop_type", ["iterate", "for"]) +@pytest.mark.parametrize("condition_value", [True, False]) +def test_repeated_if_resolution_indexes_activation_dependencies_once( + monkeypatch: pytest.MonkeyPatch, loop_type: str, condition_value: bool +) -> None: + counter = _ActivationDependencyScanCounter() + get_dependents = GraphExecutionState._get_activation_dependents_for_if + + def count_lookup_scans( + state: GraphExecutionState, source_if_id: str, iteration_path: tuple[int, ...] + ) -> tuple[str, ...]: + state_counter = counter.states_by_id.get(id(state)) + if state_counter is None: + return get_dependents(state, source_if_id, iteration_path) + state_counter.active = True + try: + return get_dependents(state, source_if_id, iteration_path) + finally: + state_counter.active = False + state_counter.resolutions += 1 + + monkeypatch.setattr(GraphExecutionState, "_get_activation_dependents_for_if", count_lookup_scans) + _run_trivial_loop( + loop_type, + 300, + include_if=True, + external_condition=True, + condition_value=condition_value, + activation_scan_counter=counter, + ) + + assert counter.resolutions >= 300 + assert counter.entries_visited <= counter.prepared_node_count, ( + f"If resolution scanned {counter.entries_visited} prepared nodes across {counter.resolutions} resolutions; " + f"session contains {counter.prepared_node_count} prepared nodes" + ) + + +@pytest.mark.slow +def test_if_activation_dependency_scan_performance_comparison(monkeypatch: pytest.MonkeyPatch) -> None: + """Compare 1,200 If resolutions with the former repeated scan and the indexed lookup.""" + + original_get_dependents = GraphExecutionState._get_activation_dependents_for_if + measurements: dict[str, tuple[float, _ActivationDependencyScanCounter]] = {} + for legacy_scan in (True, False): + counter = _ActivationDependencyScanCounter() + + def measure_lookup( + state: GraphExecutionState, + source_if_id: str, + iteration_path: tuple[int, ...], + active_counter: _ActivationDependencyScanCounter = counter, + use_legacy_scan: bool = legacy_scan, + ) -> tuple[str, ...]: + state_counter = active_counter.states_by_id.get(id(state)) + if state_counter is not None: + if use_legacy_scan: + state._if_activation_dependents_index_complete = False + state_counter.active = True + try: + return original_get_dependents(state, source_if_id, iteration_path) + finally: + if state_counter is not None: + state_counter.active = False + state_counter.resolutions += 1 + + monkeypatch.setattr(GraphExecutionState, "_get_activation_dependents_for_if", measure_lookup) + started = time.perf_counter() + _run_trivial_loop( + "iterate", + 1200, + include_if=True, + external_condition=True, + condition_value=True, + activation_scan_counter=counter, + ) + measurements["legacy repeated scan" if legacy_scan else "indexed"] = (time.perf_counter() - started, counter) + + legacy_seconds, legacy_counter = measurements["legacy repeated scan"] + indexed_seconds, indexed_counter = measurements["indexed"] + print( + "\n1,200-item Iterate/If comparison: " + f"legacy={legacy_seconds:.3f}s/{legacy_counter.entries_visited} prepared-node visits; " + f"indexed={indexed_seconds:.3f}s/{indexed_counter.entries_visited} prepared-node visits; " + f"prepared_nodes={indexed_counter.prepared_node_count}" + ) + assert legacy_counter.resolutions >= 1200 + assert indexed_counter.resolutions >= 1200 + assert indexed_counter.entries_visited <= indexed_counter.prepared_node_count + assert legacy_counter.entries_visited > indexed_counter.entries_visited + + @pytest.mark.parametrize( ("loop_type", "include_if", "external_condition"), [ diff --git a/tests/app/services/shared/test_graph_execution_frame_cache.py b/tests/app/services/shared/test_graph_execution_frame_cache.py new file mode 100644 index 00000000000..85fe5971150 --- /dev/null +++ b/tests/app/services/shared/test_graph_execution_frame_cache.py @@ -0,0 +1,34 @@ +import gc +import weakref + +from invokeai.app.services.shared.graph import Graph, GraphExecutionState + + +def test_execution_frames_are_cached_per_state_without_a_4096_entry_cutoff() -> None: + state = GraphExecutionState(graph=Graph()) + runtime_state = state.__pydantic_private__["_runtime_state"] + first_frame = state._engine_frame((0,)) + first_frame_ref = weakref.ref(first_frame) + + for index in range(1, 4_200): + state._engine_frame((index,)) + + assert len(runtime_state.engine_frames) == 4_200 + assert state._engine_frame((0,)) is first_frame + + same_state_copy = state.model_copy() + assert same_state_copy._engine_frame((0,)) is first_frame + assert same_state_copy.__pydantic_private__["_runtime_state"].engine_frames is not runtime_state.engine_frames + + new_state_copy = state.model_copy(update={"id": "different-state"}) + assert new_state_copy._engine_frame((0,)) is not first_frame + assert len(new_state_copy.__pydantic_private__["_runtime_state"].engine_frames) == 1 + + del first_frame + del same_state_copy + del new_state_copy + del runtime_state + del state + gc.collect() + + assert first_frame_ref() is None diff --git a/tests/app/services/shared/test_graph_if_runtime.py b/tests/app/services/shared/test_graph_if_runtime.py index fac41fccf75..4a3880526ad 100644 --- a/tests/app/services/shared/test_graph_if_runtime.py +++ b/tests/app/services/shared/test_graph_if_runtime.py @@ -6,6 +6,7 @@ from invokeai.app.services.shared import graph_if_runtime from invokeai.app.services.shared.execution_engine.scheduler import ActivationDependency from invokeai.app.services.shared.graph import Edge, EdgeConnection, Graph, GraphExecutionState +from invokeai.app.services.shared.graph_iterator_ancestry import _IteratorAncestry from invokeai.app.services.shared.graph_runtime_records import _ApplyTransaction from invokeai.app.services.shared.graph_validation import IterateInvocation from tests.app.services.shared.import_test_utils import assert_module_imports_without_graph @@ -154,27 +155,27 @@ def test_nested_if_iteration_paths_use_supplied_source_graph_recursively() -> No constrained_graph.remove_edge("iterator", "guarded") assert ( - graph_if_runtime._get_if_iteration_paths( + graph_if_runtime._get_if_iteration_paths_with_ancestry( state, "if", set(), lambda source_id: {"condition-exec"} if source_id == "condition" else set(), lambda _edge, _prepared_id: (), - source_graph, + _IteratorAncestry(state, source_graph), ) == [] ) - assert graph_if_runtime._get_if_iteration_paths( + assert graph_if_runtime._get_if_iteration_paths_with_ancestry( state, "if", set(), lambda source_id: {"condition-exec"} if source_id == "condition" else set(), lambda _edge, _prepared_id: (), - constrained_graph, + _IteratorAncestry(state, constrained_graph), ) == [()] -def test_unresolved_nested_literal_ifs_share_one_iterator_graph(monkeypatch: pytest.MonkeyPatch) -> None: +def test_unresolved_nested_literal_if_chain_has_no_condition_path() -> None: nodes = {"iterator": IterateInvocation(id="iterator"), "guarded": IfInvocation(id="guarded")} nodes.update({f"if-{i}": IfInvocation(id=f"if-{i}", condition=True) for i in range(24)}) edges = [ @@ -194,32 +195,23 @@ def test_unresolved_nested_literal_ifs_share_one_iterator_graph(monkeypatch: pyt previous = f"if-{i}" state = GraphExecutionState(graph=Graph()) state.graph = Graph(nodes=nodes, edges=edges) - original = GraphExecutionState._iterator_graph - builds = [] - - def count_build(self, graph=None): - builds.append(graph) - return original(self, graph) - - monkeypatch.setattr(GraphExecutionState, "_iterator_graph", count_build) - assert graph_if_runtime._get_if_iteration_paths(state, previous, set(), lambda _: set(), lambda *_: ()) == [] - assert len(builds) == 1 + assert ( + graph_if_runtime._get_if_iteration_paths_with_ancestry( + state, previous, set(), lambda _: set(), lambda *_: (), _IteratorAncestry(state) + ) + == [] + ) -def test_resolved_literal_if_does_not_build_iterator_graph(monkeypatch: pytest.MonkeyPatch) -> None: +def test_resolved_literal_if_uses_root_frame_path() -> None: state = GraphExecutionState(graph=Graph(nodes={"if": IfInvocation(id="if", condition=True, true_input=42)})) - - def unexpected_build(*_args, **_kwargs): - pytest.fail("resolved paths must not build the iterator graph") - - monkeypatch.setattr(GraphExecutionState, "_iterator_graph", unexpected_build) - assert graph_if_runtime._get_if_iteration_paths(state, "if", set(), lambda _: set(), lambda *_: ()) == [()] + assert graph_if_runtime._get_if_iteration_paths_with_ancestry( + state, "if", set(), lambda _: set(), lambda *_: (), _IteratorAncestry(state) + ) == [()] @pytest.mark.parametrize("executed", [False, True]) -def test_prepared_selected_source_preserves_paths_without_ancestry( - executed: bool, monkeypatch: pytest.MonkeyPatch -) -> None: +def test_prepared_selected_source_preserves_paths_without_ancestry(executed: bool) -> None: state = GraphExecutionState( graph=Graph(nodes={"if": IfInvocation(id="if", condition=True), "source": AddInvocation(id="source")}) ) @@ -233,10 +225,11 @@ def test_prepared_selected_source_preserves_paths_without_ancestry( if executed: state.executed.add("source") - def unexpected_build(*_args, **_kwargs): - pytest.fail("prepared selected sources must not build the iterator graph") - - monkeypatch.setattr(GraphExecutionState, "_iterator_graph", unexpected_build) - assert graph_if_runtime._get_if_iteration_paths( - state, "if", set(), lambda source_id: state.source_prepared_mapping[source_id], lambda *_: (2,) + assert graph_if_runtime._get_if_iteration_paths_with_ancestry( + state, + "if", + set(), + lambda source_id: state.source_prepared_mapping[source_id], + lambda *_: (2,), + _IteratorAncestry(state), ) == ([] if executed else [(2,)]) diff --git a/tests/app/services/shared/test_graph_materializer.py b/tests/app/services/shared/test_graph_materializer.py index 941aaaa293c..86b1eadbf69 100644 --- a/tests/app/services/shared/test_graph_materializer.py +++ b/tests/app/services/shared/test_graph_materializer.py @@ -70,7 +70,7 @@ def test_materializer_preserves_execution_edges_and_source_nodes() -> None: assert state.is_complete() -def test_prepare_shares_iterator_ancestry_across_admission_and_materialization(monkeypatch: pytest.MonkeyPatch) -> None: +def test_prepare_admits_only_selected_frame_after_unresolved_conditions() -> None: from invokeai.app.invocations.logic import IfInvocation from invokeai.app.services.shared.graph_validation import IterateInvocation, nx @@ -104,14 +104,6 @@ def test_prepare_shares_iterator_ancestry_across_admission_and_materialization(m graph = state._get_source_graph_flat() # Keep condition sources outside this preparation view; their If nodes stay unresolved. preparation_graph = nx.subgraph_view(graph, filter_node=lambda node_id: not node_id.startswith("condition-")) - original = graph_facade.GraphExecutionState._iterator_graph - builds = [] - - def count_build(self, source_graph=None): - builds.append(source_graph) - return original(self, source_graph) - - monkeypatch.setattr(graph_facade.GraphExecutionState, "_iterator_graph", count_build) prepared_id = state._materializer().prepare(preparation_graph) assert prepared_id in state.source_prepared_mapping["selected"] assert len(state.source_prepared_mapping["selected"]) == 1 @@ -119,4 +111,3 @@ def count_build(self, source_graph=None): assert "blocked-b" not in state.source_prepared_mapping assert "inner-a" not in state.source_prepared_mapping assert "inner-b" not in state.source_prepared_mapping - assert len(builds) == 1 diff --git a/tests/test_graph_execution_state.py b/tests/test_graph_execution_state.py index f939043bd8a..462445b1099 100644 --- a/tests/test_graph_execution_state.py +++ b/tests/test_graph_execution_state.py @@ -5307,8 +5307,8 @@ def test_if_graph_optimized_behavior_executes_only_selected_simple_branch_and_re assert set(executed_source_ids) == {"condition", "true_value", "if", "selected_output"} assert "false_value" not in executed_source_ids - assert set(g.executed_history) == {"condition", "true_value", "if", "selected_output"} - assert "false_value" not in g.executed_history + assert set(g.executed_history) == {"condition", "true_value", "false_value", "if", "selected_output"} + assert "false_value" not in executed_source_ids def test_if_graph_lowers_legacy_branch_choice_to_frame_scoped_activation_token(): From ab23f5817e386e660d4e15ea83b3b74f6c9fc5f9 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Thu, 15 Oct 2026 00:49:33 +0200 Subject: [PATCH 15/20] test(graph): add failing regressions for open If review findings Pin expected outputs for sibling-Collect axis order, empty right collections, shared Collect branches and nested Collect branches with extra consumers. Require early rejection or completion for unsupported If axes, cover literal Ifs selecting a root branch, and guard #9821 cheaply with a lowered recursion limit. --- .../test_execution_engine_differential.py | 402 ++++++++++++++++++ 1 file changed, 402 insertions(+) diff --git a/tests/app/services/shared/test_execution_engine_differential.py b/tests/app/services/shared/test_execution_engine_differential.py index 4ef2cd60b71..b77e62a2c1e 100644 --- a/tests/app/services/shared/test_execution_engine_differential.py +++ b/tests/app/services/shared/test_execution_engine_differential.py @@ -1,7 +1,9 @@ from __future__ import annotations +import inspect import json import random +import sys from itertools import product from pathlib import Path from typing import Any, Callable @@ -9587,3 +9589,403 @@ def test_nested_if_collect_branch_with_all_outer_frames_false_completes( assert state.next() is None assert not state._ready_node_ids assert all(not queue for queue in state._ready_queues.values()) + + +def _sibling_collect_if_graph( + *, + left: list[int], + right: list[int], + condition_axis: str = "left", + collect_field: str = "true_input", + same_if_both_fields: bool = False, +) -> Graph: + """Build the #9820 shape: a Collect over a left x right product feeds one branch of a per-item If.""" + graph = Graph() + for node in ( + CollectionConcatInvocation(id="left_source", first=left), + CollectionConcatInvocation(id="right_source", first=right), + IterateInvocation(id="left"), + IterateInvocation(id="right"), + AddInvocation(id="body"), + CollectInvocation(id="branch_collect"), + AddInvocation(id="minus_one", b=-1), + FixedCollectionTestInvocation(id="condition"), + IfInvocation(id="if"), + CollectInvocation(id="collect"), + ): + graph.add_node(node) + other_field = "false_input" if collect_field == "true_input" else "true_input" + edges = [ + ("left_source", "collection", "left", "collection"), + ("right_source", "collection", "right", "collection"), + ("left", "item", "body", "a"), + ("right", "item", "body", "b"), + ("body", "value", "branch_collect", "item"), + (condition_axis, "item", "minus_one", "a"), + ("minus_one", "value", "condition", "value"), + ("condition", "condition", "if", "condition"), + ("branch_collect", "collection", "if", collect_field), + ("if", "value", "collect", "item"), + ] + if same_if_both_fields: + edges.append(("branch_collect", "collection", "if", other_field)) + else: + graph.add_node(AnyTypeTestInvocation(id="fallback", value="fallback")) + edges.append(("fallback", "value", "if", other_field)) + for source, source_field, destination, destination_field in edges: + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + return graph + + +def _collect_output(state: GraphExecutionState, source_id: str = "collect") -> list[Any]: + return state.results[next(iter(state.source_prepared_mapping[source_id]))].collection + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +@pytest.mark.parametrize( + ("left", "right", "expected"), + [ + ([5, 6, 7, 8], [10], [[15, 16, 17, 18]]), + ([1, 2, 3], [10, 20], [[11, 12, 13], [21, 22, 23]]), + ], + ids=["one-right-item", "two-right-items"], +) +def test_sibling_collect_branch_keeps_every_item_when_condition_comes_from_the_other_axis( + force_compatibility_scheduler: bool, left: list[int], right: list[int], expected: list[list[int]] +) -> None: + state = GraphExecutionState(graph=_sibling_collect_if_graph(left=left, right=right, condition_axis="right")) + + trace, state = _run_graph(state, force_compatibility_scheduler=force_compatibility_scheduler) + + assert state.is_complete(), f"graph stalled after {trace}" + assert _collect_output(state) == expected + assert trace.count("body") == len(left) * len(right) + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +@pytest.mark.parametrize( + ("left", "collect_field", "expected"), + [ + ([1, 2], "true_input", ["fallback", []]), + ([2, 2], "true_input", [[], []]), + ([2], "true_input", [[]]), + ([1, 2], "false_input", [[], "fallback"]), + ], + ids=["mixed", "all-selected", "single-selected", "false-branch"], +) +def test_sibling_collect_branch_with_empty_right_collection_completes( + force_compatibility_scheduler: bool, left: list[int], collect_field: str, expected: list[Any] +) -> None: + state = GraphExecutionState(graph=_sibling_collect_if_graph(left=left, right=[], collect_field=collect_field)) + + trace, state = _run_graph(state, force_compatibility_scheduler=force_compatibility_scheduler) + + assert state.is_complete(), f"graph stalled after {trace}" + assert _collect_output(state) == expected + assert "body" not in trace + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_one_collect_may_feed_both_branches_of_the_same_if(force_compatibility_scheduler: bool) -> None: + state = GraphExecutionState(graph=_sibling_collect_if_graph(left=[1, 2], right=[10, 20], same_if_both_fields=True)) + + trace, state = _run_graph(state, force_compatibility_scheduler=force_compatibility_scheduler) + + assert state.is_complete(), f"graph stalled after {trace}" + assert _collect_output(state) == [[11, 21], [12, 22]] + + +def _nested_collect_branch_graph(*, outer_values: list[int], extra_consumer_of: str) -> Graph: + """Per-outer If whose true branch collects an inner Iterate; another node also consumes the inner axis.""" + graph = Graph() + for node in ( + CollectionConcatInvocation(id="outer_source", first=outer_values), + IterateInvocation(id="outer"), + RangeInvocation(id="range", start=0, step=1), + IterateInvocation(id="inner"), + AddInvocation(id="body", b=100), + CollectInvocation(id="inner_collect"), + AnyTypeTestInvocation(id="root_false", value="F"), + AddInvocation(id="minus_one", b=-1), + FixedCollectionTestInvocation(id="condition"), + IfInvocation(id="if"), + CollectInvocation(id="collect"), + AnyTypeTestInvocation(id="extra"), + CollectInvocation(id="extra_collect"), + ): + graph.add_node(node) + extra_field = "item" if extra_consumer_of == "inner" else "value" + for source, source_field, destination, destination_field in ( + ("outer_source", "collection", "outer", "collection"), + ("outer", "item", "range", "stop"), + ("range", "collection", "inner", "collection"), + ("inner", "item", "body", "a"), + ("body", "value", "inner_collect", "item"), + ("outer", "item", "minus_one", "a"), + ("minus_one", "value", "condition", "value"), + ("condition", "condition", "if", "condition"), + ("inner_collect", "collection", "if", "true_input"), + ("root_false", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + (extra_consumer_of, extra_field, "extra", "value"), + ("extra", "value", "extra_collect", "item"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + return graph + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +@pytest.mark.parametrize("extra_consumer_of", ["inner", "body"]) +@pytest.mark.parametrize( + ("outer_values", "expected"), + [([2], [[100, 101]]), ([3, 1, 2], [[100, 101, 102], "F", [100, 101]])], + ids=["single-outer", "mixed-outer"], +) +def test_nested_collect_branch_with_extra_inner_consumer_runs_if_once_per_outer_item( + force_compatibility_scheduler: bool, extra_consumer_of: str, outer_values: list[int], expected: list[Any] +) -> None: + state = GraphExecutionState( + graph=_nested_collect_branch_graph(outer_values=outer_values, extra_consumer_of=extra_consumer_of) + ) + + trace, state = _run_graph(state, force_compatibility_scheduler=force_compatibility_scheduler) + + assert state.is_complete(), f"graph stalled after {trace}" + assert _collect_output(state) == expected + assert trace.count("if") == len(outer_values) + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +@pytest.mark.parametrize( + "conditions", + [[True], [False], [True, False], [False, False]], + ids=["selected", "never-selected", "mixed", "all-unselected"], +) +def test_if_with_unrelated_branch_axis_is_rejected_before_any_node_runs( + force_compatibility_scheduler: bool, conditions: list[bool] +) -> None: + """The axis mismatch is topological, so it must not depend on condition values or run upstream work first.""" + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="branch_source", first=[1, 2])) + graph.add_node(IterateInvocation(id="branch_iterate")) + graph.add_node(AnyTypeTestInvocation(id="branch_value")) + graph.add_node(CollectionConcatInvocation(id="condition_source", first=conditions)) + graph.add_node(IterateInvocation(id="condition_iterate")) + graph.add_node(FixedCollectionTestInvocation(id="condition")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value="fallback")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="collect")) + for source, source_field, destination, destination_field in ( + ("branch_source", "collection", "branch_iterate", "collection"), + ("branch_iterate", "item", "branch_value", "value"), + ("condition_source", "collection", "condition_iterate", "collection"), + ("condition_iterate", "item", "condition", "value"), + ("condition", "condition", "if", "condition"), + ("branch_value", "value", "if", "true_input"), + ("fallback", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + + executed: list[str] = [] + with pytest.raises(graph_if_runtime.UnsupportedIfIterationAxesError): + graph.validate_self() + state = GraphExecutionState(graph=graph) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) + while (node := state.next()) is not None: + executed.append(state.prepared_source_mapping[node.id]) + state.complete(node.id, node.invoke(Mock())) + assert executed == [] + + +def _product_branch_deeper_than_condition_graph() -> Graph: + graph = Graph() + for node in ( + CollectionConcatInvocation(id="left_source", first=[1, 2]), + CollectionConcatInvocation(id="right_source", first=[10, 20]), + IterateInvocation(id="left"), + IterateInvocation(id="right"), + AddInvocation(id="body"), + AddInvocation(id="minus_one", b=-1), + FixedCollectionTestInvocation(id="condition"), + AnyTypeTestInvocation(id="fallback", value="fallback"), + IfInvocation(id="if"), + CollectInvocation(id="collect"), + ): + graph.add_node(node) + for source, source_field, destination, destination_field in ( + ("left_source", "collection", "left", "collection"), + ("right_source", "collection", "right", "collection"), + ("left", "item", "body", "a"), + ("right", "item", "body", "b"), + ("left", "item", "minus_one", "a"), + ("minus_one", "value", "condition", "value"), + ("condition", "condition", "if", "condition"), + ("body", "value", "if", "true_input"), + ("fallback", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + return graph + + +def _nested_branch_deeper_than_condition_graph() -> Graph: + """Condition per outer item; true branch per inner item, false branch per outer item.""" + graph = Graph() + for node in ( + CollectionConcatInvocation(id="outer_source", first=[1, 2]), + IterateInvocation(id="outer"), + RangeInvocation(id="range", start=0, step=1), + IterateInvocation(id="inner"), + AddInvocation(id="minus_one", b=-1), + FixedCollectionTestInvocation(id="condition"), + MarkedAnyTypeTestInvocation(id="true_branch", marker="t"), + MarkedAnyTypeTestInvocation(id="false_branch", marker="f"), + IfInvocation(id="if"), + CollectInvocation(id="collect"), + ): + graph.add_node(node) + for source, source_field, destination, destination_field in ( + ("outer_source", "collection", "outer", "collection"), + ("outer", "item", "range", "stop"), + ("range", "collection", "inner", "collection"), + ("outer", "item", "minus_one", "a"), + ("minus_one", "value", "condition", "value"), + ("condition", "condition", "if", "condition"), + ("inner", "item", "true_branch", "value"), + ("outer", "item", "false_branch", "value"), + ("true_branch", "value", "if", "true_input"), + ("false_branch", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + return graph + + +def _nested_branch_with_root_false_graph() -> Graph: + """Condition per outer item; true branch per inner item, false branch at root.""" + graph = Graph() + for node in ( + CollectionConcatInvocation(id="outer_source", first=[3, 2]), + IterateInvocation(id="outer"), + RangeInvocation(id="range", start=0, step=1), + IterateInvocation(id="inner"), + AddInvocation(id="body", b=100), + AddInvocation(id="minus_two", b=-2), + FixedCollectionTestInvocation(id="condition"), + AnyTypeTestInvocation(id="root_false", value="F"), + IfInvocation(id="if"), + CollectInvocation(id="collect"), + ): + graph.add_node(node) + for source, source_field, destination, destination_field in ( + ("outer_source", "collection", "outer", "collection"), + ("outer", "item", "range", "stop"), + ("range", "collection", "inner", "collection"), + ("inner", "item", "body", "a"), + ("outer", "item", "minus_two", "a"), + ("minus_two", "value", "condition", "value"), + ("condition", "condition", "if", "condition"), + ("body", "value", "if", "true_input"), + ("root_false", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + return graph + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +@pytest.mark.parametrize( + "build_graph", + [ + _product_branch_deeper_than_condition_graph, + _nested_branch_deeper_than_condition_graph, + _nested_branch_with_root_false_graph, + ], + ids=["product-branch", "nested-true-branch", "nested-branch-root-false"], +) +def test_if_branch_deeper_than_its_condition_completes_or_is_rejected( + force_compatibility_scheduler: bool, build_graph: Callable[[], Graph] +) -> None: + """These shapes either need defined per-frame semantics or an explicit rejection; they must not stall.""" + state = GraphExecutionState(graph=build_graph()) + try: + trace, state = _run_graph(state, force_compatibility_scheduler=force_compatibility_scheduler) + except graph_if_runtime.UnsupportedIfIterationAxesError: + return + assert state.is_complete(), f"graph stalled without an error after {trace}" + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +@pytest.mark.parametrize( + ("condition", "per_item_field", "root_field", "expected_value"), + [ + (True, "false_input", "true_input", "root"), + (False, "true_input", "false_input", "root"), + (False, "true_input", None, None), + ], + ids=["true-selects-root", "false-selects-root", "false-selects-unwired"], +) +def test_literal_if_selecting_a_root_branch_skips_the_per_item_branch( + force_compatibility_scheduler: bool, + condition: bool, + per_item_field: str, + root_field: str | None, + expected_value: str | None, +) -> None: + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="source", first=[1, 2])) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(AddInvocation(id="per_item", b=10)) + graph.add_node(IfInvocation(id="if", condition=condition)) + graph.add_node(CollectInvocation(id="collect")) + graph.add_edge(create_edge("source", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "per_item", "a")) + graph.add_edge(create_edge("per_item", "value", "if", per_item_field)) + graph.add_edge(create_edge("if", "value", "collect", "item")) + if root_field is not None: + graph.add_node(AnyTypeTestInvocation(id="root", value="root")) + graph.add_edge(create_edge("root", "value", "if", root_field)) + state = GraphExecutionState(graph=graph) + + trace, state = _run_graph(state, force_compatibility_scheduler=force_compatibility_scheduler) + + assert state.is_complete(), f"graph stalled after {trace}" + assert "per_item" not in trace + collected = _collect_output(state) + assert collected + assert all(value == expected_value for value in collected) + + +def test_deep_if_branch_chain_does_not_need_recursion_per_chain_node() -> None: + """Guard #9821 cheaply: a recursive traversal exceeds this limit long before the issue's 1,000-node repro does.""" + chain_length = 150 + graph = Graph() + graph.add_node(CollectionConcatInvocation(id="source", first=[True, False])) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_node(MarkedAnyTypeTestInvocation(id="false_branch", marker="f")) + for index in range(chain_length): + graph.add_node(AnyTypeTestInvocation(id=f"chain_{index}")) + graph.add_edge(create_edge("source", "collection", "iterate", "collection")) + graph.add_edge(create_edge("iterate", "item", "if", "condition")) + graph.add_edge(create_edge("iterate", "item", "chain_0", "value")) + graph.add_edge(create_edge("iterate", "item", "false_branch", "value")) + graph.add_edge(create_edge("false_branch", "value", "if", "false_input")) + for index in range(chain_length - 1): + graph.add_edge(create_edge(f"chain_{index}", "value", f"chain_{index + 1}", "value")) + graph.add_edge(create_edge(f"chain_{chain_length - 1}", "value", "if", "true_input")) + graph.add_edge(create_edge("if", "value", "collect", "item")) + state = GraphExecutionState(graph=graph) + + previous_limit = sys.getrecursionlimit() + sys.setrecursionlimit(len(inspect.stack()) + 250) + try: + _, state = _run_graph(state) + finally: + sys.setrecursionlimit(previous_limit) + + assert state.is_complete() + assert _collect_output(state) == [True, "fFalse"] From c6c8d22bfbc9b81a0048477ed56a8f5366f8157f Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Wed, 14 Oct 2026 21:18:28 -0500 Subject: [PATCH 16/20] fix(graph): resolve execution engine review findings Correct If and iterator branch execution while retaining the incoming regression tests and documenting materializer lifecycle behavior. --- invokeai/app/services/shared/README.md | 48 ++- invokeai/app/services/shared/graph.py | 27 ++ .../services/shared/graph_if_dependencies.py | 102 +++++ .../app/services/shared/graph_if_runtime.py | 177 ++++++--- .../shared/graph_iterator_ancestry.py | 15 - .../app/services/shared/graph_materializer.py | 190 +++++++-- .../app/services/shared/graph_scheduler.py | 1 + .../test_execution_engine_differential.py | 350 ++++++++++++++--- .../shared/test_for_scheduler_performance.py | 369 ++++++++++-------- tests/test_graph_execution_state.py | 1 - 10 files changed, 962 insertions(+), 318 deletions(-) diff --git a/invokeai/app/services/shared/README.md b/invokeai/app/services/shared/README.md index 6c93e37e84b..46fe3bf06f7 100644 --- a/invokeai/app/services/shared/README.md +++ b/invokeai/app/services/shared/README.md @@ -24,10 +24,12 @@ fan-out. Five or more nested `If`s, other fan-out, mixed graphs outside the boun topology, loop-containing, saved-workflow, legacy, and five-or-more sibling shapes remain on the compatibility fallback. Legacy snapshots retain their generic compatibility projection and legacy skipped-state metadata. Legacy skipped-state projection remains only for old snapshots. Fresh `If` materialization prepares the condition boundary -first, resolves the activation token, and attaches only the selected branch input. Unselected branch nodes are never -prepared, skipped, or added to fresh execution history. Both scheduler adapters consume these dependency decisions -without pruning execution edges or calling a type-specific branch scheduler. Legacy skipped-state metadata remains only -for snapshots that already contain the old projection. Branch-membership analysis remains an internal author-graph +first, resolves the activation token, and attaches only the selected branch input. Nodes exclusive to the unselected +branch are not materialized or invoked. `executed_history` records source-level completion and retirement: a branch +source retired as empty can appear there without an invocation or an execution result. It is bookkeeping, not an +invocation trace. Both scheduler adapters consume the same dependency decisions without pruning execution edges or +calling a type-specific branch scheduler. Legacy skipped-state metadata remains only for snapshots that already contain +the old projection. Branch-membership analysis remains an internal author-graph admission decision. The focused fresh-shape ownership matrix also verifies that supported flat, nested, bounded fan-out, and four-sibling graphs do not instantiate the compatibility controller or project fresh skipped nodes in either scheduler route. This proves the current ownership boundary; it does not remove the controller, which remains @@ -673,10 +675,12 @@ materialization. - A branch `Collect` whose `item` source combines exactly two sibling `Iterate` axes can close the axis beside the `If` condition axis. The compatibility materializer expands that sibling `Iterate` beneath each resolved `If` frame selecting the `Collect` branch, pairs only the matching condition frame, and closes the branch interval empty when no - frame selects it. The condition and branch axes must be nested; divergent inner axes are rejected even when they share - an outer axis. One sibling `Iterate` can belong to only one `If` branch context. If one shared `Collect` would need to - serve multiple `If` branches, materialization raises `UnsupportedSharedIfCollectBranchesError` before invocation - rather than dropping frames. + frame selects it. Independent iterator axes follow destination input-port order; node insertion order and node IDs do + not reorder distinct ports. The condition and branch axes must be nested; divergent inner axes are rejected even when + they share an outer axis. An uncollected deeper branch axis that cannot be selected from the condition frame raises + `UnsupportedIfIterationAxesError` before invocation. One sibling `Iterate` can belong to only one `If` owner. A + `Collect` may feed both branch ports of that same `If`; if it would serve different `If` owners, materialization + raises `UnsupportedSharedIfCollectBranchesError` before invocation rather than dropping frames. - For each new exec node: @@ -696,8 +700,23 @@ Final For outputs do not carry an open iteration axis. Iterator ancestry owns ordered ancestor and edge-axis analysis plus effective-path memoization for one graph traversal. The builder shares that object across candidate admission and materialization. Its separate per-source axis cache adds sibling-Collect relationship ordering and is cleared when that relationship graph is established; it does not repeat -the ancestor-list cache. Independent parent paths merge in destination-port order. A standalone path lookup owns a -separate context. These caches are process-local and rebuilt when graph identity changes; they are not serialized. +the ancestor-list cache. Independent parent paths merge in destination input-port order, with source node ID used only +to break ties on the same input port. A standalone path lookup owns a separate context. These caches are process-local +and rebuilt when graph identity changes; they are not serialized. + +Each pending-If attachment sweep also builds a local index from source edge and prepared iteration path to execution +node ID. It enumerates each prepared source edge once per sweep, then resolves the deepest matching path prefix for each +condition-ready If frame. This preserves deterministic attachment while avoiding a repeated full prepared-source scan +for every If frame. The index is discarded at the end of the call; a later sweep rebuilds it from the current prepared +mapping, so it is not retained across graph edits, rollback, or state restore. This optimization bounds candidate lookup +for attachment; other preparation and branch-context work has separate costs. + +Inactive branch-source retirement uses an event queue owned by one materializer. A rejected branch queues only its +activation-dependent source IDs when the corresponding `If` resolves; an empty condition `Iterate` queues both branch +source sets. Each preparation pass checks the queued sources in topological order rather than rescanning every +unprepared source. The index is tied to source-graph identity, rebuilt from durable activation state after graph changes +or state restore, and is not serialized. Queue mutations participate in execution-state rollback. This narrows +retirement work; it does not bound other preparation or branch-context costs. Literal-condition If fallback path inference uses the same ancestry context, including nested selected branches. @@ -754,12 +773,17 @@ Run `C` -> `D:0` -> enqueue `D`. Run `D` -> done. - Nodes that are exclusive to the true or false branch remain unmaterialized until their branch is admitted, even when their indegree would otherwise be zero. - Once the condition resolves, graph-state activation logic records the selected activation token and attaches only the - selected branch input to the pending `If` execution. Branch-exclusive ancestors of the unselected branch are never - prepared, executed, or added to fresh execution history. Unsupported fresh shapes and legacy snapshots use the + selected branch input to the pending `If` execution. Branch-exclusive ancestors of the unselected branch are not + invoked or materialized into execution results. Their source IDs can appear in `executed_history` after empty + retirement; the history is not an invocation trace. Unsupported fresh shapes and legacy snapshots use the compatibility controller. Both scheduler adapters use append-only execution edges; no type-specific branch scheduler performs pruning or skip propagation. - Legacy snapshots that already contain skipped prepared nodes retain their compatibility metadata and scheduler projection. Fresh execution does not create that projection. +- If an `Iterate.item` stream feeding `If.condition` closes empty, no `If` execution is created. The branch sources are + retired empty, the downstream `Collect` produces an empty collection, and the state completes without an error. + Retirement may add the branch source IDs to `executed_history`; neither branch invocation runs and no branch result + is created. - The SQLite queue/processor path has evidence for cancellation before and after `If` resolution and retry from each boundary for both condition polarities. A canceled attempt keeps its activation ledger and cannot resume; retry starts a fresh state and must emit a fresh matching activation token before selected-only completion. diff --git a/invokeai/app/services/shared/graph.py b/invokeai/app/services/shared/graph.py index 164db5846d6..1b6bee8f236 100644 --- a/invokeai/app/services/shared/graph.py +++ b/invokeai/app/services/shared/graph.py @@ -256,6 +256,7 @@ class _GraphExecutionStateRuntime: "legacy_execution_snapshot", "legacy_snapshot_loaded", "fresh_flat_if_activation", + "if_iteration_axes_validated", ) def __init__(self) -> None: @@ -302,6 +303,7 @@ def __init__(self) -> None: self.legacy_execution_snapshot = True self.legacy_snapshot_loaded = False self.fresh_flat_if_activation: Optional[bool] = None + self.if_iteration_axes_validated = False class GraphExecutionState(BaseModel): @@ -537,6 +539,7 @@ def _reset_apply_derived_caches(self) -> None: self._prepared_if_exec_ids_by_source_and_frame = {} self._prepared_if_exec_frames_by_source = {} self._prepared_if_exec_indexed_sources = set() + self._selected_if_frames_for_preparation = {} self._execution_runtime = None self._source_graph_flat = None self._execution_graph_flat = None @@ -599,6 +602,15 @@ def _prepared_if_exec_frames(self, source_node_id: str) -> tuple[tuple[int, ...] self._prepared_if_exec_ids(source_node_id, ()) return self._prepared_if_exec_frames_by_source.get(source_node_id, ()) + def _invalidate_selected_if_frames_for_preparation(self, source_node_id: Optional[str] = None) -> None: + """Clear selected-frame results after prepared If executions or branch choices change.""" + if source_node_id is None: + self._selected_if_frames_for_preparation.clear() + return + for key in tuple(self._selected_if_frames_for_preparation): + if key[0] == source_node_id: + self._selected_if_frames_for_preparation.pop(key, None) + def _if_activation_controller(self) -> _IfActivationController: if self._if_activation_controller_instance is None: self._if_activation_controller_instance = _IfActivationController(self) @@ -698,11 +710,15 @@ def _add_execution_graph_edges(self, edges: Iterable[Edge]) -> None: ) def _mark_source_executed(self, source_node_id: str) -> None: + was_executed = source_node_id in self.executed self._tx_add_set(self.executed, source_node_id) self.__pydantic_private__["_runtime_state"].completed_source_ids_cache = None self._tx_add_set(self._get_completed_source_ids_cache(), source_node_id) if source_node_id not in self.executed_history: self._tx_append_list(self.executed_history, source_node_id) + materializer = self.__pydantic_private__["_runtime_state"].execution_materializer + if not was_executed and materializer is not None: + materializer._queue_empty_condition_retirements(source_node_id) def _discard_source_executed(self, source_node_id: str) -> None: self._tx_discard_set(self.executed, source_node_id) @@ -781,8 +797,10 @@ def _get_for_source_by_return_id(self) -> dict[str, str]: def _invalidate_source_graph_cache(self) -> None: runtime = self.__pydantic_private__["_runtime_state"] + self._invalidate_selected_if_frames_for_preparation() runtime.source_graph_flat = None runtime.fresh_flat_if_activation = None + runtime.if_iteration_axes_validated = False runtime.for_source_by_return_id = None runtime.if_branch_sources_cache = {} runtime.if_activation_dependencies_by_source = {} @@ -2158,6 +2176,7 @@ def cancel_generic_child(self, child_item_id: int, error_message: str) -> ChildD return dependency.cancel_child(str(child_item_id), error_message) def _register_prepared_exec_node(self, exec_node_id: str, source_node_id: str) -> None: + self._invalidate_selected_if_frames_for_preparation(source_node_id) self._prepared_if_exec_indexed_sources.discard(source_node_id) for key in tuple(self._prepared_if_exec_ids_by_source_and_frame): if key[0] == source_node_id: @@ -3198,6 +3217,7 @@ def _invalidate_loop_caches_for_source(self, source_node_id: str) -> None: def _invalidate_loop_caches_for_exec_node(self, exec_node_id: str) -> None: source_node_id = self.prepared_source_mapping.get(exec_node_id) if source_node_id is not None: + self._invalidate_selected_if_frames_for_preparation(source_node_id) self._prepared_if_exec_indexed_sources.discard(source_node_id) for key in tuple(self._prepared_if_exec_ids_by_source_and_frame): if key[0] == source_node_id: @@ -3374,6 +3394,7 @@ def _try_resolve_if_node(self, exec_node_id: str, *, enqueue: bool = True) -> No self._record_compatibility_activation_token(exec_node_id, selected_field) assert isinstance(scheduler, _ExecutionScheduler) source_if_id = self._prepared_registry().get_source_node_id(exec_node_id) + self._invalidate_selected_if_frames_for_preparation(source_if_id) iteration_path = self._get_iteration_path(exec_node_id) activation_dependents = self._get_activation_dependents_for_if(source_if_id, iteration_path) scheduler._discard_rejected_activation_nodes(activation_dependents) @@ -3510,6 +3531,7 @@ def _reset_runtime_caches(self) -> None: self._if_activation_dependents_index_complete = False self._prepared_if_exec_ids_by_source_and_frame = {} self._prepared_if_exec_indexed_sources = set() + self._selected_if_frames_for_preparation = {} self._source_graph_flat = None self._execution_graph_flat = None self._completed_source_ids_cache = None @@ -3875,6 +3897,11 @@ def graph_is_valid(cls, v: Graph): def next(self) -> Optional[BaseInvocation]: """Gets the next node ready to execute.""" + runtime = self.__pydantic_private__["_runtime_state"] + if not runtime.if_iteration_axes_validated: + graph_if_runtime.validate_if_iteration_axes(self) + runtime.if_iteration_axes_validated = True + # TODO: enable multiple nodes to execute simultaneously by tracking currently executing nodes # possibly with a timeout? diff --git a/invokeai/app/services/shared/graph_if_dependencies.py b/invokeai/app/services/shared/graph_if_dependencies.py index 044aacc0b59..0c4198542be 100644 --- a/invokeai/app/services/shared/graph_if_dependencies.py +++ b/invokeai/app/services/shared/graph_if_dependencies.py @@ -30,6 +30,9 @@ def _get_fresh_if_nodes(state: "GraphExecutionState") -> tuple[IfInvocation, ... def _can_use_fresh_mixed_if_iterate_collect(state: "GraphExecutionState") -> bool: """Admit the bounded per-item Iterate/If/Collect topology.""" + if _can_use_fresh_literal_if_iterate_collect(state): + return True + if state._legacy_snapshot_loaded or len(state.graph.nodes) != 6 or len(state.graph.edges) != 7: return False @@ -101,6 +104,94 @@ def _can_use_fresh_mixed_if_iterate_collect(state: "GraphExecutionState") -> boo ) +def _can_use_fresh_literal_if_iterate_collect(state: "GraphExecutionState") -> bool: + """Admit a literal If whose unselected item branch establishes its Iterate frame.""" + + if state._legacy_snapshot_loaded: + return False + + if_nodes = [node for node in state.graph.nodes.values() if isinstance(node, IfInvocation)] + iterate_nodes = [node for node in state.graph.nodes.values() if isinstance(node, IterateInvocation)] + collect_nodes = [node for node in state.graph.nodes.values() if isinstance(node, CollectInvocation)] + if len(if_nodes) != 1 or len(iterate_nodes) != 1 or len(collect_nodes) != 1: + return False + if any( + isinstance(node, (CallSavedWorkflowInvocation, ForInvocation, ForReturnInvocation)) + for node in state.graph.nodes.values() + ): + return False + + if_node = if_nodes[0] + iterate_node = iterate_nodes[0] + collect_node = collect_nodes[0] + if state.graph._get_input_edges(if_node.id, "condition"): + return False + + collection_edges = state.graph._get_input_edges(iterate_node.id, COLLECTION_FIELD) + if len(collection_edges) != 1: + return False + collection_edge = collection_edges[0] + collection_source = state.graph.get_node(collection_edge.source.node_id) + if ( + collection_edge.source.field != COLLECTION_FIELD + or state.graph._get_input_edges(collection_source.id) + or state.graph._get_output_edges(collection_source.id) != [collection_edge] + ): + return False + + collect_edges = state.graph._get_input_edges(collect_node.id, ITEM_FIELD) + if ( + len(collect_edges) != 1 + or collect_edges[0].source.node_id != if_node.id + or collect_edges[0].source.field != "value" + or state.graph._get_output_edges(if_node.id) != [collect_edges[0]] + or state.graph._get_input_edges(collect_node.id, COLLECTION_FIELD) + or state.graph._get_output_edges(collect_node.id) + ): + return False + + item_branch_sources: set[str] = set() + branch_sources: set[str] = set() + for branch_field in ("true_input", "false_input"): + branch_edges = state.graph._get_input_edges(if_node.id, branch_field) + if len(branch_edges) > 1: + return False + if not branch_edges: + continue + branch_edge = branch_edges[0] + branch_node = state.graph.get_node(branch_edge.source.node_id) + branch_inputs = state.graph._get_input_edges(branch_node.id) + if state.graph._get_output_edges(branch_node.id) != [branch_edge]: + return False + if not branch_inputs: + branch_sources.add(branch_node.id) + elif ( + len(branch_inputs) == 1 + and branch_inputs[0].source.node_id == iterate_node.id + and branch_inputs[0].source.field == ITEM_FIELD + ): + item_branch_sources.add(branch_node.id) + branch_sources.add(branch_node.id) + else: + return False + + item_branch_edges = [ + edge for branch_source_id in item_branch_sources for edge in state.graph._get_input_edges(branch_source_id) + ] + if not item_branch_sources or state.graph._get_output_edges(iterate_node.id, ITEM_FIELD) != item_branch_edges: + return False + + expected_nodes = { + collection_source.id, + iterate_node.id, + if_node.id, + collect_node.id, + *branch_sources, + } + expected_edges = 2 + len(branch_sources) + len(item_branch_sources) + return set(state.graph.nodes) == expected_nodes and len(state.graph.edges) == expected_edges + + def _can_use_fresh_flat_if_activation(state: "GraphExecutionState") -> bool: """Admit fresh bounded independent-sibling or nested-If dependency compilation.""" @@ -242,6 +333,17 @@ def _get_fresh_if_branch_sources(state: "GraphExecutionState", if_node_id: str, and not state.graph._get_output_edges(edge.destination.node_id) ) + if not state.graph._get_input_edges(if_node_id, "condition"): + iterator_graph = state._iterator_graph(source_graph) + frame_iterators = { + ancestor_id + for ancestor_id in nx.ancestors(iterator_graph, if_node_id) + if isinstance(state.graph.get_node(ancestor_id), (ForInvocation, IterateInvocation)) + } + for iterator_id in frame_iterators: + branch_sources.discard(iterator_id) + branch_sources.difference_update(nx.ancestors(source_graph, iterator_id)) + changed = True while changed: changed = False diff --git a/invokeai/app/services/shared/graph_if_runtime.py b/invokeai/app/services/shared/graph_if_runtime.py index bf86900886b..f98a13af5f2 100644 --- a/invokeai/app/services/shared/graph_if_runtime.py +++ b/invokeai/app/services/shared/graph_if_runtime.py @@ -7,14 +7,56 @@ from invokeai.app.services.shared.execution_engine.scheduler import ActivationDependency from invokeai.app.services.shared.graph_iterator_ancestry import _IteratorAncestry from invokeai.app.services.shared.graph_models import Edge, ExecutionToken -from invokeai.app.services.shared.graph_validation import ForInvocation, IterateInvocation +from invokeai.app.services.shared.graph_validation import ( + COLLECTION_FIELD, + ITEM_FIELD, + CollectInvocation, + ForInvocation, + IterateInvocation, +) if TYPE_CHECKING: from invokeai.app.services.shared.graph import GraphExecutionState class UnsupportedIfIterationAxesError(ValueError): - """Raised when an If branch and its condition use unrelated open iterator axes.""" + """Raised when an If condition and its branches cannot share a supported open-axis scope.""" + + +def validate_if_iteration_axes(state: "GraphExecutionState") -> None: + """Reject unrelated condition and branch axes using graph topology alone.""" + iterator_ancestry = _IteratorAncestry(state) + for node_id, node in state.graph.nodes.items(): + if not isinstance(node, IfInvocation): + continue + condition_edges = state.graph._get_input_edges(node_id, "condition") + if not condition_edges: + continue + condition_iterators = iterator_ancestry.get_iterator_ancestors_for_edge(condition_edges[0]) + for branch_field in ("true_input", "false_input"): + for edge in state.graph._get_input_edges(node_id, branch_field): + branch_iterators = iterator_ancestry.get_iterator_ancestors_for_edge(edge) + branch_source = state.graph.get_node(edge.source.node_id) + if ( + condition_iterators + and condition_iterators < branch_iterators + and not isinstance(branch_source, CollectInvocation) + ): + raise UnsupportedIfIterationAxesError( + f"If node '{node_id}' cannot select branch input '{branch_field}' from deeper open iterator " + f"axes without a Collect boundary: condition {sorted(condition_iterators)}, " + f"branch {sorted(branch_iterators)}" + ) + shared_iterators = condition_iterators.intersection(branch_iterators) + nested_axes = condition_iterators.issubset(branch_iterators) or branch_iterators.issubset( + condition_iterators + ) + if condition_iterators and branch_iterators and (not shared_iterators or not nested_axes): + raise UnsupportedIfIterationAxesError( + f"If node '{node_id}' cannot select branch input '{branch_field}' from " + f"unrelated open iterator axes: condition {sorted(condition_iterators)}, " + f"branch {sorted(branch_iterators)}" + ) def _get_effective_iteration_paths_for_edge( @@ -62,6 +104,26 @@ def _get_effective_iteration_paths_for_edge( continue source_node = state.graph.get_node(current_source_id) + if isinstance(source_node, CollectInvocation) and current_edge.source.field == COLLECTION_FIELD: + input_edges = state.graph._get_input_edges(current_source_id) + collapsed_paths: set[tuple[int, ...]] = set() + # Collect.item collapses its final open axis even before its own execution frame exists. + for input_edge in input_edges: + input_paths = _get_effective_iteration_paths_for_edge( + state, + input_edge, + visited_source_ids, + get_prepared_nodes_for_source, + get_prepared_edge_iteration_path, + iterator_ancestry, + cache, + ) + group_depth = len(iterator_ancestry.get_iterator_axes_for_edge(input_edge)) + if input_edge.destination.field == ITEM_FIELD: + group_depth = max(group_depth - 1, 0) + collapsed_paths.update(path[:group_depth] for path in input_paths) + cache[cache_key] = collapsed_paths or ({()} if not input_edges else set()) + continue if isinstance(source_node, IfInvocation): cache[cache_key] = set( _get_if_iteration_paths_with_ancestry( @@ -127,41 +189,19 @@ def _get_if_iteration_paths_with_ancestry( for edge in selected_edges ) if selected_edges: - compatible_edges = [] - if condition_edges: - condition_iterators = iterator_ancestry.get_iterator_ancestors_for_edge(condition_edges[0]) - for edge in selected_edges: - branch_iterators = iterator_ancestry.get_iterator_ancestors_for_edge(edge) - shared_iterators = condition_iterators.intersection(branch_iterators) - nested_axes = condition_iterators.issubset(branch_iterators) or branch_iterators.issubset( - condition_iterators - ) - if condition_iterators and branch_iterators and (not shared_iterators or not nested_axes): - raise UnsupportedIfIterationAxesError( - f"If node '{node_id}' cannot select branch input '{selected_field}' from " - f"unrelated open iterator axes: condition {sorted(condition_iterators)}, " - f"branch {sorted(branch_iterators)}" - ) - compatible_edges.append(edge) - else: - compatible_edges = selected_edges - branch_paths = ( - set().union( - *( - _get_effective_iteration_paths_for_edge( - state, - edge, - next_visited_source_ids, - get_prepared_nodes_for_source, - get_prepared_edge_iteration_path, - iterator_ancestry, - effective_paths_by_edge, - ) - for edge in compatible_edges + branch_paths = set().union( + *( + _get_effective_iteration_paths_for_edge( + state, + edge, + next_visited_source_ids, + get_prepared_nodes_for_source, + get_prepared_edge_iteration_path, + iterator_ancestry, + effective_paths_by_edge, ) + for edge in selected_edges ) - if compatible_edges - else set() ) else: branch_paths = {condition_path or ()} @@ -174,6 +214,31 @@ def _get_if_iteration_paths_with_ancestry( elif branch_path[: len(condition_path)] == condition_path: paths.add(branch_path) + if not condition_edges: + frame_iterators = iterator_ancestry.get_ordered_iterator_ancestors(node_id) + if frame_iterators: + frame_paths: set[tuple[int, ...]] = set() + for iterator_id in reversed(frame_iterators): + if iterator_id not in state.source_prepared_mapping: + if iterator_id not in state.executed: + return [] + continue + prepared_ids = get_prepared_nodes_for_source(iterator_id) + iterator_edges = state.graph._get_output_edges(iterator_id, ITEM_FIELD) + if prepared_ids and iterator_edges: + frame_paths = { + get_prepared_edge_iteration_path(iterator_edges[0], prepared_id) for prepared_id in prepared_ids + } + break + if not frame_paths: + return [] + paths = { + frame_path if frame_path[: len(branch_path)] == branch_path else branch_path + for branch_path in paths + for frame_path in frame_paths + if frame_path[: len(branch_path)] == branch_path or branch_path[: len(frame_path)] == frame_path + } + if not paths: if not has_selected_edges: return [()] @@ -212,6 +277,13 @@ def _get_source_activation_dependencies_for_preparation( runtime = state._runtime() source_axes = runtime._get_ordered_iterator_sources(source_node_id) + # Materialization can prepend a condition axis to a branch iterator, so a dependent node's path may + # follow input-port order rather than its graph-ancestry order. Use the path's construction order when + # translating axis coordinates into the selected If frame. Collect nodes can collapse a trailing axis, + # so only axes represented by this path participate. + input_axes = state._materializer()._get_ordered_input_iterator_axes(state.graph._get_input_edges(source_node_id)) + if input_axes: + source_axes = list(input_axes[: len(iteration_path)]) resolved_dependencies = [] for dependency in dependencies: owner_axes = runtime._get_ordered_iterator_sources(dependency.owner_id) @@ -247,14 +319,28 @@ def _get_source_activation_dependencies_for_preparation( def _get_selected_if_frames_for_preparation( - state: "GraphExecutionState", if_node_id: str, branch_field: str -) -> list[tuple[int, ...]]: + state: "GraphExecutionState", if_node_id: str, branch_field: Optional[str] +) -> tuple[tuple[int, ...], ...]: """Return only resolved If frames that selected a Collect-backed branch.""" - return [ - state._get_iteration_path(exec_node_id) - for exec_node_id in state.source_prepared_mapping.get(if_node_id, ()) - if state._resolved_if_exec_branches.get(exec_node_id) == branch_field - ] + cache_key = (if_node_id, branch_field) + cached = state._selected_if_frames_for_preparation.get(cache_key) + if cached is None: + selected_frames = tuple( + state._get_iteration_path(exec_node_id) + for exec_node_id in state.source_prepared_mapping.get(if_node_id, ()) + if branch_field is None or state._resolved_if_exec_branches.get(exec_node_id) == branch_field + ) + cached = (selected_frames, frozenset(selected_frames)) + state._selected_if_frames_for_preparation[cache_key] = cached + return cached[0] + + +def _has_selected_if_frame_for_preparation( + state: "GraphExecutionState", if_node_id: str, branch_field: Optional[str], frame: tuple[int, ...] +) -> bool: + """Test one selected frame without scanning every frame owned by the If.""" + _get_selected_if_frames_for_preparation(state, if_node_id, branch_field) + return frame in state._selected_if_frames_for_preparation[(if_node_id, branch_field)][1] def _is_source_activation_admitted( @@ -284,7 +370,6 @@ def _is_activation_axis_complete(state: "GraphExecutionState", axis_source_id: s def _is_source_inactive( state: "GraphExecutionState", source_node_id: str, iteration_path: tuple[int, ...] = () ) -> bool: - use_fresh_activation = state._can_use_fresh_flat_if_activation() dependencies = state._get_source_activation_dependencies(source_node_id, iteration_path) if not dependencies: return False @@ -296,7 +381,7 @@ def _is_source_inactive( frames = {frame for dependency in dependencies for frame in state._prepared_if_exec_frames(dependency.owner_id)} if not frames: - return use_fresh_activation and all(_is_empty_if_condition(state, dependency) for dependency in dependencies) + return all(_is_empty_if_condition(state, dependency) for dependency in dependencies) source_axes = set(state._runtime()._get_ordered_iterator_sources(source_node_id)) for dependency in dependencies: @@ -427,6 +512,10 @@ def _resolve_activation_gate(state: "GraphExecutionState", exec_node_id: str, br if changed: state._completed_source_ids_cache = None state._tx_record(lambda: runtime.replace_gate(previous)) + materializer = state.__pydantic_private__["_runtime_state"].execution_materializer + if materializer is not None: + source_if_id = state._prepared_registry().get_source_node_id(exec_node_id) + materializer._queue_rejected_activation_branch(source_if_id, branch) return changed diff --git a/invokeai/app/services/shared/graph_iterator_ancestry.py b/invokeai/app/services/shared/graph_iterator_ancestry.py index fe59c8b102d..3ecab6a1b2c 100644 --- a/invokeai/app/services/shared/graph_iterator_ancestry.py +++ b/invokeai/app/services/shared/graph_iterator_ancestry.py @@ -23,7 +23,6 @@ def __init__( self._source_graph = source_graph self._iterator_graph = iterator_graph self._has_iterator_ancestor: dict[str, bool] | None = None - self._iterator_ancestors: dict[str, frozenset[str]] = {} self._ordered_iterator_ancestors: dict[str, tuple[str, ...]] = {} self.effective_iteration_paths_by_edge: dict[tuple[str, str], set[tuple[int, ...]]] = {} @@ -53,20 +52,6 @@ def has_iterator_ancestor(self, node_id: str) -> bool: raise nx.NetworkXError(f"The node {node_id} is not in the digraph.") return self._has_iterator_ancestor[node_id] - def get_iterator_ancestors(self, node_id: str) -> frozenset[str]: - self.has_iterator_ancestor(node_id) - if node_id not in self._iterator_ancestors: - assert self._iterator_graph is not None - ancestors = { - ancestor_id - for ancestor_id in nx.ancestors(self._iterator_graph, node_id) - if isinstance(self._state.graph.get_node(ancestor_id), (ForInvocation, IterateInvocation)) - } - if isinstance(self._state.graph.get_node(node_id), (ForInvocation, IterateInvocation)): - ancestors.add(node_id) - self._iterator_ancestors[node_id] = frozenset(ancestors) - return self._iterator_ancestors[node_id] - def get_ordered_iterator_ancestors(self, node_id: str) -> list[str]: """Return open iterator axes in outer-to-inner, input-port order.""" self.has_iterator_ancestor(node_id) diff --git a/invokeai/app/services/shared/graph_materializer.py b/invokeai/app/services/shared/graph_materializer.py index 22293497a7f..57b5e119a11 100644 --- a/invokeai/app/services/shared/graph_materializer.py +++ b/invokeai/app/services/shared/graph_materializer.py @@ -34,7 +34,7 @@ class UnsupportedSharedIfCollectBranchesError(ValueError): - """Raised when one sibling Collect-backed Iterate must serve multiple If frames.""" + """Raised when one sibling Collect-backed Iterate must serve multiple If owners.""" class _ExecutionNodeBuilder: @@ -48,8 +48,15 @@ class _ExecutionNodeBuilder: def __init__(self, state: "GraphExecutionState") -> None: self._state = state self._has_if_nodes = any(isinstance(node, IfInvocation) for node in state.graph.nodes.values()) + self._activation_candidate_graph: Optional[nx.DiGraph] = None + self._activation_candidate_topological_order: tuple[str, ...] = () + self._activation_candidate_topological_index: dict[str, int] = {} + self._activation_sources_by_dependency: dict[tuple[str, str], set[str]] = {} + self._empty_condition_if_sources: dict[str, tuple[str, ...]] = {} + self._pending_activation_retirement_candidates: set[str] = set() + self._activation_retirement_rebuild_required = False self._iteration_axes_by_source: dict[str, tuple[str, ...]] = {} - self._sibling_iterate_if_collect_branches: Optional[dict[str, tuple[str, str]]] = None + self._sibling_iterate_if_collect_branches: Optional[dict[str, tuple[str, Optional[str]]]] = None self._iteration_relationship_graph: Optional[nx.DiGraph] = None self._iteration_relationship_topological_index: Optional[dict[str, int]] = None # Keyed by the identity of the cached flat source graph, which every graph edit replaces. @@ -238,6 +245,99 @@ def _build_execution_edges_for_fields( ) return new_edges + def _ensure_activation_candidate_index(self, graph: "nx.DiGraph") -> None: + if self._activation_candidate_graph is graph: + return + + self._activation_candidate_graph = graph + self._has_if_nodes = any(isinstance(node, IfInvocation) for node in self._state.graph.nodes.values()) + self._activation_candidate_topological_order = () + self._activation_candidate_topological_index = {} + self._activation_sources_by_dependency = {} + self._empty_condition_if_sources = {} + self._pending_activation_retirement_candidates.clear() + if not self._has_if_nodes: + self._activation_retirement_rebuild_required = False + return + + self._activation_candidate_topological_order = tuple(nx.topological_sort(graph)) + self._activation_candidate_topological_index = { + source_node_id: index for index, source_node_id in enumerate(self._activation_candidate_topological_order) + } + self._activation_sources_by_dependency = {} + self._empty_condition_if_sources = {} + self._pending_activation_retirement_candidates.clear() + for source_node_id in self._activation_candidate_topological_order: + for dependency in self._state._get_source_activation_dependencies(source_node_id): + self._activation_sources_by_dependency.setdefault((dependency.owner_id, dependency.branch), set()).add( + source_node_id + ) + for if_node_id in graph.nodes: + if not isinstance(self._state.graph.get_node(if_node_id), IfInvocation): + continue + condition_edges = self._state.graph._get_input_edges(if_node_id, "condition") + if len(condition_edges) != 1: + continue + condition_source_id = condition_edges[0].source.node_id + if isinstance(self._state.graph.get_node(condition_source_id), IterateInvocation): + self._empty_condition_if_sources.setdefault(condition_source_id, ()) + self._empty_condition_if_sources[condition_source_id] += (if_node_id,) + + self._activation_retirement_rebuild_required = True + self._rebuild_activation_retirement_candidates() + + def _rebuild_activation_retirement_candidates(self) -> None: + self._pending_activation_retirement_candidates.clear() + registry = self._state._prepared_registry() + opposite_branch = {"true_input": "false_input", "false_input": "true_input"} + for exec_node_id, selected_branch in self._state._resolved_if_exec_branches.items(): + if selected_branch not in opposite_branch: + continue + owner_id = registry.get_source_node_id(exec_node_id) + self._pending_activation_retirement_candidates.update( + self._activation_sources_by_dependency.get((owner_id, opposite_branch[selected_branch]), ()) + ) + for condition_source_id, if_node_ids in self._empty_condition_if_sources.items(): + if condition_source_id not in self._state.executed or self._state.source_prepared_mapping.get( + condition_source_id + ): + continue + for if_node_id in if_node_ids: + self._pending_activation_retirement_candidates.update( + self._activation_sources_by_dependency.get((if_node_id, "true_input"), ()) + ) + self._pending_activation_retirement_candidates.update( + self._activation_sources_by_dependency.get((if_node_id, "false_input"), ()) + ) + self._activation_retirement_rebuild_required = False + + def _queue_rejected_activation_branch(self, if_node_id: str, selected_branch: str) -> None: + opposite_branch = {"true_input": "false_input", "false_input": "true_input"}.get(selected_branch) + if opposite_branch is None: + return + self._ensure_activation_candidate_index(self._state._get_source_graph_flat()) + for source_node_id in self._activation_sources_by_dependency.get((if_node_id, opposite_branch), ()): + self._state._tx_add_set(self._pending_activation_retirement_candidates, source_node_id) + + def _queue_empty_condition_retirements(self, condition_source_id: str) -> None: + self._ensure_activation_candidate_index(self._state._get_source_graph_flat()) + for if_node_id in self._empty_condition_if_sources.get(condition_source_id, ()): + for branch in ("true_input", "false_input"): + for source_node_id in self._activation_sources_by_dependency.get((if_node_id, branch), ()): + self._state._tx_add_set(self._pending_activation_retirement_candidates, source_node_id) + + def _get_activation_candidate_source_ids(self, graph: "nx.DiGraph") -> tuple[str, ...]: + """Return queued retirement candidates in source-graph order.""" + self._ensure_activation_candidate_index(graph) + if self._activation_retirement_rebuild_required: + self._rebuild_activation_retirement_candidates() + return tuple( + sorted( + self._pending_activation_retirement_candidates, + key=self._activation_candidate_topological_index.__getitem__, + ) + ) + def _has_unmaterializable_for_final_input(self, node_id: str) -> bool: final_for_source_ids = set() for edge in self._state.graph._get_input_edges(node_id): @@ -1428,15 +1528,16 @@ def _matches_sibling_if_collect_context( if collapsed_branch is None: continue if_node_id, branch_field = collapsed_branch - selected_frames = graph_if_runtime._get_selected_if_frames_for_preparation( - self._state, if_node_id, branch_field - ) branch_path = self._state._get_iteration_path(prepared_iterators[index]) - matching_frames = [ - frame - for frame in selected_frames - if len(branch_path) == len(frame) + 1 and branch_path[: len(frame)] == frame - ] + matching_frame = branch_path[:-1] if branch_path else () + matching_frames = ( + (matching_frame,) + if len(branch_path) > 0 + and graph_if_runtime._has_selected_if_frame_for_preparation( + self._state, if_node_id, branch_field, matching_frame + ) + else () + ) condition_edges = self._state.graph._get_input_edges(if_node_id, "condition") condition_iterator_ids = set( self._get_ordered_iterator_axes_for_edges(condition_edges, self.iterator_graph()) @@ -1581,7 +1682,7 @@ def create_execution_node( def _has_admitted_source_mapping( self, node_id: str, graph: "nx.DiGraph", iterator_ancestry: _IteratorAncestry ) -> bool: - if not any(isinstance(node, IfInvocation) for node in self._state.graph.nodes.values()): + if not self._has_if_nodes: return True if isinstance(self._state.graph.get_node(node_id), IfInvocation): condition_edges = self._state.graph._get_input_edges(node_id, "condition") @@ -1710,12 +1811,12 @@ def _get_if_readiness_ancestors(self, node_id: str, graph: "nx.DiGraph") -> set[ ancestors.add(condition_source_id) return ancestors - def _get_sibling_iterate_if_collect_branches(self) -> dict[str, tuple[str, str]]: + def _get_sibling_iterate_if_collect_branches(self) -> dict[str, tuple[str, Optional[str]]]: """Find Iterate axes collapsed by branch Collects under an If.""" if self._sibling_iterate_if_collect_branches is not None: return self._sibling_iterate_if_collect_branches - branches: dict[str, tuple[str, str]] = {} + branches: dict[str, tuple[str, Optional[str]]] = {} if not any(isinstance(node, CollectInvocation) for node in self._state.graph.nodes.values()): self._sibling_iterate_if_collect_branches = branches return branches @@ -1776,13 +1877,16 @@ def _get_sibling_iterate_if_collect_branches(self) -> dict[str, tuple[str, str]] continue branch = (destination.id, output_edge.destination.field) existing_branch = branches.get(iterate_node_id) - if existing_branch is not None and existing_branch != branch: + if existing_branch is not None and existing_branch[0] != branch[0]: raise UnsupportedSharedIfCollectBranchesError( f"Sibling Iterate '{iterate_node_id}' cannot supply Collect-backed inputs to multiple " - f"If branches: '{existing_branch[0]}.{existing_branch[1]}' and " + f"If nodes: '{existing_branch[0]}' and " f"'{branch[0]}.{branch[1]}'" ) - branches[iterate_node_id] = branch + branches[iterate_node_id] = ( + branch[0], + None if existing_branch is not None and existing_branch[1] != branch[1] else branch[1], + ) self._iteration_relationship_graph = relationship_graph # Cached axes may have been queried before this graph established sibling order. @@ -1794,13 +1898,16 @@ def _get_sibling_iterate_if_collect_branches(self) -> dict[str, tuple[str, str]] self._sibling_iterate_if_collect_branches = branches return branches - def _get_sibling_iterate_if_collect_branch(self, iterate_node_id: str) -> Optional[tuple[str, str]]: + def _get_sibling_iterate_if_collect_branch(self, iterate_node_id: str) -> Optional[tuple[str, Optional[str]]]: return self._get_sibling_iterate_if_collect_branches().get(iterate_node_id) def _attach_pending_if_inputs(self, *, enqueue: bool = True) -> None: """Attach only the selected branch edge to condition-ready If executions.""" source_graph = self._state._get_source_graph_flat() + # Prepared branch sources do not change during this sweep. Index them by frame once per source edge; + # keeping the index local means graph edits, rollback, and restore rebuild it on the next call. + prepared_exec_by_edge_path: dict[tuple[str, str], dict[tuple[int, ...], str]] = {} source_order = {node_id: index for index, node_id in enumerate(nx.topological_sort(source_graph))} pending_exec_nodes = sorted( self._state._pending_if_exec_nodes, @@ -1831,14 +1938,25 @@ def _attach_pending_if_inputs(self, *, enqueue: bool = True) -> None: for source_edge in selected_edges: if source_edge.source.node_id not in self._state.source_prepared_mapping: continue - candidates = [ - (prepared_id, self._get_prepared_edge_iteration_path(source_edge, prepared_id)) - for prepared_id in self._get_ordered_prepared_nodes_for_edge(source_edge) - ] - candidates = [candidate for candidate in candidates if target_path[: len(candidate[1])] == candidate[1]] - if not candidates: + edge_key = (source_edge.source.node_id, source_edge.source.field) + if edge_key not in prepared_exec_by_edge_path: + prepared_exec_by_path: dict[tuple[int, ...], str] = {} + for prepared_id in self._get_ordered_prepared_nodes_for_edge(source_edge): + prepared_path = self._get_prepared_edge_iteration_path(source_edge, prepared_id) + prepared_exec_by_path.setdefault(prepared_path, prepared_id) + prepared_exec_by_edge_path[edge_key] = prepared_exec_by_path + + prepared_exec_by_path = prepared_exec_by_edge_path[edge_key] + selected_exec_id = next( + ( + prepared_exec_by_path[target_path[:path_depth]] + for path_depth in range(len(target_path), -1, -1) + if target_path[:path_depth] in prepared_exec_by_path + ), + None, + ) + if selected_exec_id is None: continue - selected_exec_id, _ = min(candidates, key=lambda candidate: (-len(candidate[1]), candidate[0])) selected_exec_ids.append(selected_exec_id) attached_edges.append( Edge( @@ -2025,16 +2143,22 @@ def prepare(self, base_g: Optional["nx.DiGraph"] = None) -> Optional[str]: iterator_ancestry = self._get_iterator_ancestry(self.iterator_graph(g)) iterator_ancestry.clear_effective_iteration_paths() self._attach_pending_if_inputs() + self._ensure_activation_candidate_index(g) if self._has_if_nodes: - for node_id in nx.topological_sort(g): - if ( - node_id not in self._state.source_prepared_mapping - and node_id not in self._state.executed - and self._state._is_source_inactive(node_id) - ): - self._mark_source_node_empty(node_id) - self._state._invalidate_loop_caches_for_source(node_id) - return node_id + try: + for node_id in self._get_activation_candidate_source_ids(g): + self._state._tx_discard_set(self._pending_activation_retirement_candidates, node_id) + if ( + node_id not in self._state.source_prepared_mapping + and node_id not in self._state.executed + and self._state._is_source_inactive(node_id) + ): + self._mark_source_node_empty(node_id) + self._state._invalidate_loop_caches_for_source(node_id) + return node_id + except Exception: + self._activation_retirement_rebuild_required = True + raise next_node_id = next( ( node_id diff --git a/invokeai/app/services/shared/graph_scheduler.py b/invokeai/app/services/shared/graph_scheduler.py index f5146c67a2f..f910a1e3e90 100644 --- a/invokeai/app/services/shared/graph_scheduler.py +++ b/invokeai/app/services/shared/graph_scheduler.py @@ -430,6 +430,7 @@ def resolve_if_node(self, exec_node_id: str, *, enqueue: bool = True) -> None: self._state._tx_set_mapping(runtime.resolved_if_exec_branches, exec_node_id, selected_field) self._state._record_compatibility_activation_token(exec_node_id, selected_field) source_if_id = self._state._prepared_registry().get_source_node_id(exec_node_id) + self._state._invalidate_selected_if_frames_for_preparation(source_if_id) iteration_path = self._state._get_iteration_path(exec_node_id) self._discard_rejected_nodes(self._state._get_activation_dependents_for_if(source_if_id, iteration_path)) self._sync_indegree() diff --git a/tests/app/services/shared/test_execution_engine_differential.py b/tests/app/services/shared/test_execution_engine_differential.py index b77e62a2c1e..29f5398fd50 100644 --- a/tests/app/services/shared/test_execution_engine_differential.py +++ b/tests/app/services/shared/test_execution_engine_differential.py @@ -41,6 +41,7 @@ _ExecutionScheduler, _GenericGraphSchedulerAdapter, ) +from invokeai.app.services.shared.graph_materializer import UnsupportedSharedIfCollectBranchesError from invokeai.app.services.shared.graph_nested_iterate_planner import get_nested_iterate_sequence_depth from invokeai.app.services.shared.invocation_context import InvocationContextData, build_invocation_context from tests.test_nodes import ( @@ -7171,6 +7172,72 @@ def test_constant_if_inside_iterate_routes_each_loop_item(condition: bool, branc assert all(not queue for queue in state._ready_queues.values()) +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_selected_if_frame_cache_tracks_branch_resolution_and_restore(force_compatibility_scheduler: bool) -> None: + graph = Graph() + for node in ( + CollectionConcatInvocation(id="source", first=[2, 1]), + IterateInvocation(id="items"), + AddInvocation(id="condition_value", b=-1), + FixedCollectionTestInvocation(id="condition"), + AnyTypeTestInvocation(id="true_value", value="T"), + AnyTypeTestInvocation(id="false_value", value="F"), + IfInvocation(id="if"), + CollectInvocation(id="collect"), + ): + graph.add_node(node) + for source, source_field, destination, destination_field in ( + ("source", "collection", "items", "collection"), + ("items", "item", "condition_value", "a"), + ("condition_value", "value", "condition", "value"), + ("condition", "condition", "if", "condition"), + ("true_value", "value", "if", "true_input"), + ("false_value", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + + state = GraphExecutionState(graph=graph) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) + + trace: list[str] = [] + selected_true_frames: set[tuple[int, ...]] = set() + if_frames_seen = 0 + while (node := state.next()) is not None: + source_id = state.prepared_source_mapping[node.id] + trace.append(source_id) + if source_id == "condition": + assert set(graph_if_runtime._get_selected_if_frames_for_preparation(state, "if", "true_input")) == ( + selected_true_frames + ) + elif source_id == "if": + if_frames_seen += 1 + iteration_path = state._get_iteration_path(node.id) + expected_selected_frames = selected_true_frames | ({iteration_path} if node.condition else set()) + assert set(graph_if_runtime._get_selected_if_frames_for_preparation(state, "if", "true_input")) == ( + expected_selected_frames + ) + output = node.invoke(Mock()) + state.complete(node.id, output) + if source_id == "condition": + assert set(graph_if_runtime._get_selected_if_frames_for_preparation(state, "if", "true_input")) == ( + selected_true_frames + ) + elif source_id == "if" and node.condition: + selected_true_frames.add(iteration_path) + + assert state.is_complete(), f"graph stalled after {trace}" + assert if_frames_seen == 2 + assert state.results[next(iter(state.source_prepared_mapping["collect"]))].collection == ["T", "F"] + + restored = load_execution_state(dump_execution_state(state)) + assert ( + set(graph_if_runtime._get_selected_if_frames_for_preparation(restored, "if", "true_input")) + == selected_true_frames + ) + + @pytest.mark.parametrize(("condition", "branch"), [(True, "true:"), (False, "false:")]) def test_constant_if_inside_for_routes_each_loop_item(condition: bool, branch: str) -> None: graph = Graph() @@ -8730,6 +8797,44 @@ def test_sibling_iterates_collect_through_if_branch_without_unsupported_frames( assert "body" not in trace +def test_sibling_collect_shared_by_distinct_if_owners_remains_unsupported() -> None: + graph = Graph() + for node in ( + CollectionConcatInvocation(id="left_source", first=[1, 2]), + IterateInvocation(id="left"), + CollectionConcatInvocation(id="right_source", first=[10, 20]), + IterateInvocation(id="right"), + AddInvocation(id="body"), + CollectInvocation(id="branch_collect"), + FixedCollectionTestInvocation(id="condition"), + AnyTypeTestInvocation(id="fallback_a", value=0), + AnyTypeTestInvocation(id="fallback_b", value=0), + IfInvocation(id="if_a"), + IfInvocation(id="if_b"), + ): + graph.add_node(node) + for source, source_field, destination, destination_field in ( + ("left_source", "collection", "left", "collection"), + ("right_source", "collection", "right", "collection"), + ("left", "item", "body", "a"), + ("right", "item", "body", "b"), + ("body", "value", "branch_collect", "item"), + ("left", "item", "condition", "value"), + ("condition", "condition", "if_a", "condition"), + ("condition", "condition", "if_b", "condition"), + ("branch_collect", "collection", "if_a", "true_input"), + ("branch_collect", "collection", "if_b", "true_input"), + ("fallback_a", "value", "if_a", "false_input"), + ("fallback_b", "value", "if_b", "false_input"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + + state = GraphExecutionState(graph=graph) + + with pytest.raises(UnsupportedSharedIfCollectBranchesError, match="multiple If nodes"): + state.next() + + def test_nested_iterate_collected_in_if_branch_is_not_misclassified_as_sibling(): graph = Graph() graph.add_node(CollectionConcatInvocation(id="outer_source", first=[["a", "b"], ["c"]])) @@ -8965,38 +9070,136 @@ def test_independent_iterator_collect_preserves_ordered_parent_axis_for_branch_i @pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) -def test_if_rejects_unrelated_open_iterator_axes(force_compatibility_scheduler: bool) -> None: +@pytest.mark.parametrize( + ("xs", "expected", "expected_trace", "expected_history"), + [ + ([], [], ["xs", "collect"], ["xs", "iterate", "true_value", "false_value", "if", "collect"]), + ( + [True], + ["T"], + ["xs", "iterate", "true_value", "if", "collect"], + ["xs", "iterate", "false_value", "true_value", "if", "collect"], + ), + ], + ids=["empty-condition", "non-empty-condition"], +) +def test_iterate_if_condition_completes_without_invoking_unselected_branches( + force_compatibility_scheduler: bool, + xs: list[bool], + expected: list[str], + expected_trace: list[str], + expected_history: list[str], +) -> None: graph = Graph() - graph.add_node(CollectionConcatInvocation(id="branch_source", first=[1, 2])) - graph.add_node(IterateInvocation(id="branch_iterate")) - graph.add_node(AnyTypeTestInvocation(id="branch_value")) - graph.add_node(CollectionConcatInvocation(id="condition_source", first=[True])) - graph.add_node(IterateInvocation(id="condition_iterate")) - graph.add_node(FixedCollectionTestInvocation(id="condition")) - graph.add_node(AnyTypeTestInvocation(id="fallback", value="fallback")) - graph.add_node(IfInvocation(id="if")) - graph.add_node(CollectInvocation(id="collect")) + for node in ( + CollectionConcatInvocation(id="xs", first=xs), + IterateInvocation(id="iterate"), + AnyTypeTestInvocation(id="true_value", value="T"), + AnyTypeTestInvocation(id="false_value", value="F"), + IfInvocation(id="if"), + CollectInvocation(id="collect"), + ): + graph.add_node(node) + for source, source_field, destination, destination_field in ( + ("xs", "collection", "iterate", "collection"), + ("iterate", "item", "if", "condition"), + ("true_value", "value", "if", "true_input"), + ("false_value", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + + state = GraphExecutionState(graph=graph) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) + trace, state = _run_graph(state) + + assert trace == expected_trace + collect_exec_id = next(iter(state.source_prepared_mapping["collect"])) + assert state.results[collect_exec_id].collection == expected + assert not state.has_error() + assert state.is_complete(), f"graph stalled after {trace}" + assert state.executed_history == expected_history + if not xs: + assert state.source_prepared_mapping["true_value"] == set() + assert state.source_prepared_mapping["false_value"] == set() + else: + assert state.source_prepared_mapping["true_value"] + assert state.source_prepared_mapping["false_value"] == set() + + restored = load_execution_state(dump_execution_state(state)) + assert restored.is_complete() + assert restored.executed_history == expected_history + restored_collect_exec_id = next(iter(restored.source_prepared_mapping["collect"])) + assert restored.results[restored_collect_exec_id].collection == expected + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_if_allows_shared_condition_and_branch_iterator_axis(force_compatibility_scheduler: bool) -> None: + graph = Graph() + for node in ( + CollectionConcatInvocation(id="source", first=[1, 0]), + IterateInvocation(id="iterate"), + FixedCollectionTestInvocation(id="condition"), + AnyTypeTestInvocation(id="branch"), + AnyTypeTestInvocation(id="fallback", value="fallback"), + IfInvocation(id="if"), + CollectInvocation(id="collect"), + ): + graph.add_node(node) + for source, source_field, destination, destination_field in ( + ("source", "collection", "iterate", "collection"), + ("iterate", "item", "condition", "value"), + ("iterate", "item", "branch", "value"), + ("condition", "condition", "if", "condition"), + ("branch", "value", "if", "true_input"), + ("fallback", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + + trace, state = _run_graph( + GraphExecutionState(graph=graph), force_compatibility_scheduler=force_compatibility_scheduler + ) + + assert state.is_complete(), f"graph stalled after {trace}" + result_exec_id = next(iter(state.source_prepared_mapping["collect"])) + assert state.results[result_exec_id].collection == [1, "fallback"] + + +def test_if_iteration_axis_validation_cache_invalidates_after_graph_edit() -> None: + graph = Graph() + for node in ( + CollectionConcatInvocation(id="branch_source", first=[10, 20]), + IterateInvocation(id="branch_iterate"), + AnyTypeTestInvocation(id="branch_value"), + CollectionConcatInvocation(id="condition_source", first=[1]), + IterateInvocation(id="condition_iterate"), + FixedCollectionTestInvocation(id="condition"), + IfInvocation(id="if"), + ): + graph.add_node(node) for source, source_field, destination, destination_field in ( ("branch_source", "collection", "branch_iterate", "collection"), ("branch_iterate", "item", "branch_value", "value"), ("condition_source", "collection", "condition_iterate", "collection"), ("condition_iterate", "item", "condition", "value"), ("condition", "condition", "if", "condition"), - ("branch_value", "value", "if", "true_input"), - ("fallback", "value", "if", "false_input"), - ("if", "value", "collect", "item"), ): graph.add_edge(create_edge(source, source_field, destination, destination_field)) - graph.validate_self() state = GraphExecutionState(graph=graph) - if force_compatibility_scheduler: - state._execution_scheduler = _ExecutionScheduler(state) + assert state.next() is not None + state.add_edge(create_edge("branch_value", "value", "if", "true_input")) + with pytest.raises( graph_if_runtime.UnsupportedIfIterationAxesError, - match="cannot select branch input 'true_input' from unrelated open iterator axes", + match=( + r"If node 'if' cannot select branch input 'true_input' from unrelated open iterator axes: " + r"condition \['condition_iterate'\], branch \['branch_iterate'\]" + ), ): - _run_graph(state) + state.next() @pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) @@ -9083,12 +9286,13 @@ def test_sibling_collect_cannot_feed_multiple_if_branches(force_compatibility_sc state = GraphExecutionState(graph=graph) if force_compatibility_scheduler: state._execution_scheduler = _ExecutionScheduler(state) - with pytest.raises(ValueError, match="cannot supply Collect-backed inputs to multiple If branches"): + with pytest.raises(ValueError, match="cannot supply Collect-backed inputs to multiple If nodes"): _run_graph(state) -def test_deep_any_chain_inside_iterate_if_completes() -> None: - chain_length = 400 +@pytest.mark.slow +def test_1000_node_any_chain_inside_iterate_if_completes() -> None: + chain_length = 1_000 graph = Graph() graph.add_node(CollectionConcatInvocation(id="source", first=[True, False])) graph.add_node(IterateInvocation(id="iterate")) @@ -9103,16 +9307,23 @@ def test_deep_any_chain_inside_iterate_if_completes() -> None: graph.add_edge(create_edge("iterate", "item", "false_branch", "value")) graph.add_edge(create_edge("false_branch", "value", "if", "false_input")) for index in range(chain_length - 1): - node_id = f"chain_{index}" - next_id = f"chain_{index + 1}" - graph.add_edge(create_edge(node_id, "value", next_id, "value")) + graph.add_edge(create_edge(f"chain_{index}", "value", f"chain_{index + 1}", "value")) graph.add_edge(create_edge(f"chain_{chain_length - 1}", "value", "if", "true_input")) graph.add_edge(create_edge("if", "value", "collect", "item")) - _, state = _run_graph(GraphExecutionState(graph=graph)) + state = GraphExecutionState(graph=graph) + original_recursion_limit = sys.getrecursionlimit() + try: + sys.setrecursionlimit(len(inspect.stack()) + 250) + trace, state = _run_graph(state) + finally: + sys.setrecursionlimit(original_recursion_limit) - assert state.results[next(iter(state.source_prepared_mapping["collect"]))].collection == [True, "fFalse"] - assert state.is_complete() + collect_exec_id = next(iter(state.source_prepared_mapping["collect"])) + assert state.results[collect_exec_id].collection == [True, "fFalse"] + assert all(trace.count(f"chain_{index}") == 1 for index in range(chain_length)) + assert trace.count("false_branch") == 1 + assert state.is_complete(), f"graph stalled after {trace[-8:]}" def test_execution_materializer_cache_refreshes_after_sibling_iterate_edge_addition() -> None: @@ -9159,27 +9370,6 @@ def test_execution_materializer_cache_refreshes_after_sibling_iterate_edge_addit assert state.is_complete() -def test_deep_if_branch_chain_completes_through_execution_state() -> None: - chain_length = 150 - graph = Graph() - graph.add_node(AnyTypeTestInvocation(id="source", value=7)) - graph.add_node(IfInvocation(id="if", condition=True)) - graph.add_node(AnyTypeTestInvocation(id="result")) - for index in range(chain_length): - graph.add_node(AnyTypeTestInvocation(id=f"chain_{index}")) - graph.add_edge(create_edge("source", "value", "chain_0", "value")) - for index in range(chain_length - 1): - graph.add_edge(create_edge(f"chain_{index}", "value", f"chain_{index + 1}", "value")) - graph.add_edge(create_edge(f"chain_{chain_length - 1}", "value", "if", "true_input")) - graph.add_edge(create_edge("if", "value", "result", "value")) - - trace, state = _run_graph(GraphExecutionState(graph=graph)) - - result_exec_id = next(iter(state.source_prepared_mapping["result"])) - assert state.results[result_exec_id].value == 7 - assert state.is_complete(), f"graph stalled after {trace}" - - @pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) def test_if_diamond_branch_traversal_scales_with_graph_size( force_compatibility_scheduler: bool, monkeypatch: pytest.MonkeyPatch @@ -9535,6 +9725,50 @@ def test_literal_if_branch_feeds_per_item_if_without_stranding_unselected_root( assert "per_item_unselected" not in trace +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_deep_branch_root_false_with_two_outer_frames_and_nested_collect( + force_compatibility_scheduler: bool, +) -> None: + graph = Graph() + for node in ( + CollectionConcatInvocation(id="outer_source", first=[0, 2]), + IterateInvocation(id="outer"), + RangeOfSizeInvocation(id="inner_source", start=100), + IterateInvocation(id="inner"), + AnyTypeTestInvocation(id="body"), + CollectInvocation(id="branch_collect"), + FixedCollectionTestInvocation(id="condition"), + AnyTypeTestInvocation(id="fallback", value="F"), + IfInvocation(id="if"), + CollectInvocation(id="output"), + ): + graph.add_node(node) + for source, source_field, destination, destination_field in ( + ("outer_source", "collection", "outer", "collection"), + ("outer", "item", "inner_source", "size"), + ("inner_source", "collection", "inner", "collection"), + ("inner", "item", "body", "value"), + ("body", "value", "branch_collect", "item"), + ("outer", "item", "condition", "value"), + ("condition", "condition", "if", "condition"), + ("branch_collect", "collection", "if", "true_input"), + ("fallback", "value", "if", "false_input"), + ("if", "value", "output", "item"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + + trace, state = _run_graph( + GraphExecutionState(graph=graph), force_compatibility_scheduler=force_compatibility_scheduler + ) + + assert state.is_complete(), f"graph stalled after {trace}; history={state.executed_history}" + assert _source_output(state, "output").collection == ["F", [100, 101]] + assert trace.count("if") == 2 + assert trace.count("branch_collect") == 1 + assert trace.count("fallback") == 1 + assert trace.count("output") == 1 + + @pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) @pytest.mark.parametrize("extra_inner_consumer", [False, True], ids=["no-extra-consumer", "extra-inner-consumer"]) def test_nested_if_collect_branch_with_all_outer_frames_false_completes( @@ -9908,13 +10142,17 @@ def _nested_branch_with_root_false_graph() -> Graph: def test_if_branch_deeper_than_its_condition_completes_or_is_rejected( force_compatibility_scheduler: bool, build_graph: Callable[[], Graph] ) -> None: - """These shapes either need defined per-frame semantics or an explicit rejection; they must not stall.""" + """Deeper open branch axes are explicitly rejected before any node executes.""" state = GraphExecutionState(graph=build_graph()) - try: - trace, state = _run_graph(state, force_compatibility_scheduler=force_compatibility_scheduler) - except graph_if_runtime.UnsupportedIfIterationAxesError: - return - assert state.is_complete(), f"graph stalled without an error after {trace}" + with pytest.raises( + graph_if_runtime.UnsupportedIfIterationAxesError, + match=( + r"If node 'if' cannot select branch input 'true_input' from deeper open iterator axes " + r"without a Collect boundary" + ), + ): + _run_graph(state, force_compatibility_scheduler=force_compatibility_scheduler) + assert state.executed_history == [] @pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) @@ -9954,8 +10192,8 @@ def test_literal_if_selecting_a_root_branch_skips_the_per_item_branch( assert state.is_complete(), f"graph stalled after {trace}" assert "per_item" not in trace collected = _collect_output(state) - assert collected - assert all(value == expected_value for value in collected) + assert collected == [expected_value, expected_value] + assert trace.count("if") == 2 def test_deep_if_branch_chain_does_not_need_recursion_per_chain_node() -> None: diff --git a/tests/app/services/shared/test_for_scheduler_performance.py b/tests/app/services/shared/test_for_scheduler_performance.py index 7fdb8464355..bcbe2c73c80 100644 --- a/tests/app/services/shared/test_for_scheduler_performance.py +++ b/tests/app/services/shared/test_for_scheduler_performance.py @@ -1,8 +1,7 @@ -"""Scheduler behavior and overhead for trivial loop bodies. +"""Measure scheduler mapping work and optional overhead for trivial loop bodies. -The scheduler-scaling regression tests count visits to growing scheduler ledgers, so -they remain stable under shared CI load. Optional absolute-time benchmarks remain marked `slow`. -The completion-state tests are not benchmarks. +The deterministic scaling regression counts mapping visits during end-to-end execution. +Optional absolute-time benchmarks remain marked `slow`. """ import time @@ -35,27 +34,6 @@ def record(self, mapping_name: str) -> None: self.visits_by_mapping[mapping_name] = self.visits_by_mapping.get(mapping_name, 0) + 1 -class _ActivationDependencyScanCounter: - def __init__(self) -> None: - self.active = False - self.entries_visited = 0 - self.resolutions = 0 - self.prepared_node_count = 0 - self.states_by_id: dict[int, _ActivationDependencyScanCounter] = {} - - -class _PreparedSourceCountingMapping(dict[str, str]): - def __init__(self, values: dict[str, str], counter: _ActivationDependencyScanCounter) -> None: - super().__init__(values) - self._counter = counter - - def __iter__(self): # type: ignore[override] - for key in super().__iter__(): - if self._counter.active: - self._counter.entries_visited += 1 - yield key - - class _CountingMapping(dict[str, object]): """Count mapping accesses and iteration without constraining the scheduler algorithm.""" @@ -98,20 +76,6 @@ def get(self, key: str, default: object = None) -> object: # type: ignore[overr return super().get(key, default) -class _CountingSet(set[str]): - """Count visits to prepared execution IDs without constraining set operations.""" - - def __init__(self, values: set[str], counter: _MappingTraversalCounter, mapping_name: str) -> None: - super().__init__(values) - self._counter = counter - self._mapping_name = mapping_name - - def __iter__(self): # type: ignore[override] - for value in super().__iter__(): - self._counter.record(self._mapping_name) - yield value - - def _run_trivial_loop( loop_type: str, count: int, @@ -121,7 +85,6 @@ def _run_trivial_loop( include_if: bool = False, external_condition: bool = False, condition_value: bool = True, - activation_scan_counter: _ActivationDependencyScanCounter | None = None, ) -> float: graph = Graph() graph.add_node(RangeInvocation(id="range", start=0, stop=count)) @@ -153,12 +116,6 @@ def _run_trivial_loop( graph.add_node(CollectInvocation(id="collect")) graph.add_edge(create_edge(output_source, "value", "collect", "item")) state = GraphExecutionState(graph=graph) - if activation_scan_counter is not None: - state.prepared_source_mapping = _PreparedSourceCountingMapping( - state.prepared_source_mapping, activation_scan_counter - ) # type: ignore[assignment] - state._prepared_registry()._prepared_source_mapping = state.prepared_source_mapping - activation_scan_counter.states_by_id[id(state)] = activation_scan_counter if mapping_counter is not None: scheduler = state._scheduler() state.execution_refs = _CountingMapping(state.execution_refs, mapping_counter, "execution_refs") # type: ignore[assignment] @@ -174,103 +131,9 @@ def _run_trivial_loop( while (node := state.next()) is not None: state.complete(node.id, node.invoke(context)) assert state.is_complete() - if activation_scan_counter is not None: - activation_scan_counter.prepared_node_count = len(state.prepared_source_mapping) return clock() - started -@pytest.mark.parametrize("loop_type", ["iterate", "for"]) -@pytest.mark.parametrize("condition_value", [True, False]) -def test_repeated_if_resolution_indexes_activation_dependencies_once( - monkeypatch: pytest.MonkeyPatch, loop_type: str, condition_value: bool -) -> None: - counter = _ActivationDependencyScanCounter() - get_dependents = GraphExecutionState._get_activation_dependents_for_if - - def count_lookup_scans( - state: GraphExecutionState, source_if_id: str, iteration_path: tuple[int, ...] - ) -> tuple[str, ...]: - state_counter = counter.states_by_id.get(id(state)) - if state_counter is None: - return get_dependents(state, source_if_id, iteration_path) - state_counter.active = True - try: - return get_dependents(state, source_if_id, iteration_path) - finally: - state_counter.active = False - state_counter.resolutions += 1 - - monkeypatch.setattr(GraphExecutionState, "_get_activation_dependents_for_if", count_lookup_scans) - _run_trivial_loop( - loop_type, - 300, - include_if=True, - external_condition=True, - condition_value=condition_value, - activation_scan_counter=counter, - ) - - assert counter.resolutions >= 300 - assert counter.entries_visited <= counter.prepared_node_count, ( - f"If resolution scanned {counter.entries_visited} prepared nodes across {counter.resolutions} resolutions; " - f"session contains {counter.prepared_node_count} prepared nodes" - ) - - -@pytest.mark.slow -def test_if_activation_dependency_scan_performance_comparison(monkeypatch: pytest.MonkeyPatch) -> None: - """Compare 1,200 If resolutions with the former repeated scan and the indexed lookup.""" - - original_get_dependents = GraphExecutionState._get_activation_dependents_for_if - measurements: dict[str, tuple[float, _ActivationDependencyScanCounter]] = {} - for legacy_scan in (True, False): - counter = _ActivationDependencyScanCounter() - - def measure_lookup( - state: GraphExecutionState, - source_if_id: str, - iteration_path: tuple[int, ...], - active_counter: _ActivationDependencyScanCounter = counter, - use_legacy_scan: bool = legacy_scan, - ) -> tuple[str, ...]: - state_counter = active_counter.states_by_id.get(id(state)) - if state_counter is not None: - if use_legacy_scan: - state._if_activation_dependents_index_complete = False - state_counter.active = True - try: - return original_get_dependents(state, source_if_id, iteration_path) - finally: - if state_counter is not None: - state_counter.active = False - state_counter.resolutions += 1 - - monkeypatch.setattr(GraphExecutionState, "_get_activation_dependents_for_if", measure_lookup) - started = time.perf_counter() - _run_trivial_loop( - "iterate", - 1200, - include_if=True, - external_condition=True, - condition_value=True, - activation_scan_counter=counter, - ) - measurements["legacy repeated scan" if legacy_scan else "indexed"] = (time.perf_counter() - started, counter) - - legacy_seconds, legacy_counter = measurements["legacy repeated scan"] - indexed_seconds, indexed_counter = measurements["indexed"] - print( - "\n1,200-item Iterate/If comparison: " - f"legacy={legacy_seconds:.3f}s/{legacy_counter.entries_visited} prepared-node visits; " - f"indexed={indexed_seconds:.3f}s/{indexed_counter.entries_visited} prepared-node visits; " - f"prepared_nodes={indexed_counter.prepared_node_count}" - ) - assert legacy_counter.resolutions >= 1200 - assert indexed_counter.resolutions >= 1200 - assert indexed_counter.entries_visited <= indexed_counter.prepared_node_count - assert legacy_counter.entries_visited > indexed_counter.entries_visited - - @pytest.mark.parametrize( ("loop_type", "include_if", "external_condition"), [ @@ -285,6 +148,7 @@ def measure_lookup( def test_loop_scheduler_mapping_visits_scale_linearly( loop_type: str, include_if: bool, external_condition: bool ) -> None: + """Bound wrapped scheduler-state mapping visits, not preparation retirement or If-frame lookup.""" visits: dict[int, int] = {} visits_by_mapping: dict[int, dict[str, int]] = {} for count in (300, 1200): @@ -321,30 +185,112 @@ def test_loop_scheduler_mapping_visits_scale_linearly( ) -def test_prepared_execution_count_scans_each_source_once() -> None: - counter = _MappingTraversalCounter() - state = GraphExecutionState(graph=Graph()) - state.source_prepared_mapping = _CountingMapping( - { - "source": _CountingSet({"exec-1", "exec-2", "exec-3"}, counter, "prepared_node_ids"), - }, - counter, - "source_prepared_mapping", - ) # type: ignore[assignment] - - assert state._count_unexecuted_prepared("source") == 3 - assert state._count_unexecuted_prepared("source") == 3 - assert counter.visits_by_mapping == {"source_prepared_mapping": 1, "prepared_node_ids": 3} - - @pytest.mark.slow @pytest.mark.parametrize("loop_type", ["for", "iterate"]) def test_trivial_loop_scheduler_overhead(loop_type: str) -> None: elapsed = _run_trivial_loop(loop_type, 600) - # Generous headroom for shared CI hosts; the reported quadratic For regression took six seconds. + # Broad quiet-machine guard for catastrophic slowdowns; this is not a historical baseline comparison. assert elapsed < 4, f"{loop_type}: {elapsed:.3f}s for 600 items" +class _PendingIfScanCounter(set[str]): + def __init__(self, values: set[str]) -> None: + super().__init__(values) + self.items_scanned = 0 + + def __iter__(self): + for value in super().__iter__(): + self.items_scanned += 1 + yield value + + +def _measure_pending_if_attachment(count: int) -> tuple[int, int, int, float, float]: + graph = Graph() + graph.add_node(RangeInvocation(id="range", start=0, stop=count)) + graph.add_node(IterateInvocation(id="loop")) + graph.add_node(AnyTypeTestInvocation(id="true_body")) + graph.add_node(AnyTypeTestInvocation(id="false_body")) + graph.add_node(BooleanInvocation(id="condition", value=True)) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="collect")) + graph.add_edge(create_edge("range", "collection", "loop", "collection")) + graph.add_edge(create_edge("loop", "item", "true_body", "value")) + graph.add_edge(create_edge("loop", "item", "false_body", "value")) + graph.add_edge(create_edge("condition", "value", "if", "condition")) + graph.add_edge(create_edge("true_body", "value", "if", "true_input")) + graph.add_edge(create_edge("false_body", "value", "if", "false_input")) + graph.add_edge(create_edge("if", "value", "collect", "item")) + state = GraphExecutionState(graph=graph) + state._pending_if_exec_nodes = _PendingIfScanCounter(state._pending_if_exec_nodes) + materializer = state._materializer() + original_attach = materializer._attach_pending_if_inputs + original_prepared_source = materializer._get_ordered_prepared_nodes_for_source + attach_calls = 0 + attach_seconds = 0.0 + prepared_candidates_scanned = 0 + measuring_attachment = False + + def counted_prepared_source(source_node_id: str) -> list[str]: + nonlocal prepared_candidates_scanned + prepared_nodes = original_prepared_source(source_node_id) + if measuring_attachment: + prepared_candidates_scanned += len(prepared_nodes) + return prepared_nodes + + materializer._get_ordered_prepared_nodes_for_source = counted_prepared_source # type: ignore[method-assign] + + def measured_attach(*, enqueue: bool = True) -> None: + nonlocal attach_calls, attach_seconds, measuring_attachment + attach_calls += 1 + started = time.perf_counter() + measuring_attachment = True + try: + original_attach(enqueue=enqueue) + finally: + measuring_attachment = False + attach_seconds += time.perf_counter() - started + + materializer._attach_pending_if_inputs = measured_attach # type: ignore[method-assign] + trace: list[str] = [] + started = time.perf_counter() + while (node := state.next()) is not None: + trace.append(state.prepared_source_mapping[node.id]) + state.complete(node.id, node.invoke(Mock())) + total_seconds = time.perf_counter() - started + + assert state.is_complete() + assert not state.has_error() + assert trace.count("true_body") == count + assert "false_body" not in trace + assert trace.count("if") == count + collect_exec_ids = state.source_prepared_mapping["collect"] + assert len(collect_exec_ids) == 1 + assert state.results[next(iter(collect_exec_ids))].collection == list(range(count)) + return ( + attach_calls, + state._pending_if_exec_nodes.items_scanned, + prepared_candidates_scanned, + total_seconds, + attach_seconds, + ) + + +@pytest.mark.slow +def test_pending_if_attachment_sweep_scaling() -> None: + """Measure the reviewed per-item If graph, separating attachment sweeps from total execution.""" + for count in (300, 1200): + attach_calls, scanned, candidates_scanned, total_seconds, attach_seconds = _measure_pending_if_attachment(count) + assert candidates_scanned <= count * 32, ( + f"pending-If attachment revisited {candidates_scanned} prepared candidates " + f"for {count} items across {attach_calls} sweeps" + ) + print( + f"pending-If n={count} calls={attach_calls} pending_scanned={scanned} " + f"prepared_candidates_scanned={candidates_scanned} total={total_seconds:.3f}s " + f"attach={attach_seconds:.3f}s" + ) + + def test_completion_predicate_preserves_durable_state() -> None: graph = Graph() graph.add_node(AnyTypeTestInvocation(id="value", value=1)) @@ -382,3 +328,112 @@ def test_completion_cache_preserves_completed_sources_after_restore() -> None: restored_state.complete(second.id, second.invoke(Mock())) assert restored_state.is_complete() + + +def test_prepare_retirement_checks_only_activation_dependent_sources() -> None: + graph = Graph() + graph.add_node(BooleanInvocation(id="condition", value=True)) + graph.add_node(AnyTypeTestInvocation(id="selected", value=1)) + graph.add_node(AnyTypeTestInvocation(id="unselected", value=2)) + graph.add_node(IfInvocation(id="if")) + graph.add_edge(create_edge("condition", "value", "if", "condition")) + graph.add_edge(create_edge("selected", "value", "if", "true_input")) + graph.add_edge(create_edge("unselected", "value", "if", "false_input")) + unrelated_source_ids = {f"independent_{index}" for index in range(120)} + for source_id in unrelated_source_ids: + graph.add_node(AnyTypeTestInvocation(id=source_id, value=source_id)) + + state = GraphExecutionState(graph=graph) + materializer = state._materializer() + original_candidates = materializer._get_activation_candidate_source_ids + candidate_sources: set[str] = set() + candidate_visits = 0 + prepare_scans = 0 + + def counted_candidates(source_graph: object) -> tuple[str, ...]: + nonlocal candidate_visits, prepare_scans + candidates = original_candidates(source_graph) # type: ignore[arg-type] + candidate_sources.update(candidates) + candidate_visits += len(candidates) + prepare_scans += 1 + return candidates + + materializer._get_activation_candidate_source_ids = counted_candidates # type: ignore[method-assign] + context = Mock() + while (node := state.next()) is not None: + state.complete(node.id, node.invoke(context)) + + assert state.is_complete() + assert not state.has_error() + assert state.results[next(iter(state.source_prepared_mapping["if"]))].value == 1 + assert "unselected" in state.executed + assert not unrelated_source_ids.intersection(candidate_sources), ( + f"retirement candidates included unrelated sources: {unrelated_source_ids.intersection(candidate_sources)}" + ) + assert candidate_sources == {"unselected"} + assert candidate_visits <= 2 * prepare_scans + + +def _measure_if_chain_retirement_candidates(length: int) -> tuple[int, int, int, list[str]]: + graph = Graph() + graph.add_node(BooleanInvocation(id="condition", value=True)) + graph.add_node(IfInvocation(id="if")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value="fallback")) + graph.add_node(AnyTypeTestInvocation(id="sink")) + previous_node_id: str | None = None + for index in range(length): + node_id = f"branch_{index:03d}" + graph.add_node(AnyTypeTestInvocation(id=node_id, value=True if index == 0 else None)) + if previous_node_id is not None: + graph.add_edge(create_edge(previous_node_id, "value", node_id, "value")) + previous_node_id = node_id + assert previous_node_id is not None + for source, source_field, destination, destination_field in ( + ("condition", "value", "if", "condition"), + ("fallback", "value", "if", "false_input"), + (previous_node_id, "value", "if", "true_input"), + ("if", "value", "sink", "value"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + + state = GraphExecutionState(graph=graph) + materializer = state._materializer() + original_candidates = materializer._get_activation_candidate_source_ids + candidate_visits = 0 + prepare_scans = 0 + + def counted_candidates(source_graph: object) -> tuple[str, ...]: + nonlocal candidate_visits, prepare_scans + candidates = original_candidates(source_graph) # type: ignore[arg-type] + candidate_visits += len(candidates) + prepare_scans += 1 + return candidates + + materializer._get_activation_candidate_source_ids = counted_candidates # type: ignore[method-assign] + trace: list[str] = [] + while (node := state.next()) is not None: + trace.append(state.prepared_source_mapping[node.id]) + state.complete(node.id, node.invoke(Mock())) + + assert state.is_complete() + assert not state.has_error() + assert "fallback" not in trace + assert trace.count("if") == 1 + assert trace.count("sink") == 1 + sink_exec_id = next(iter(state.source_prepared_mapping["sink"])) + assert state.results[sink_exec_id].value is True + return candidate_visits, prepare_scans, len(trace), trace + + +def test_if_chain_retirement_candidate_work_scales_linearly() -> None: + measurements = {length: _measure_if_chain_retirement_candidates(length) for length in (150, 300)} + for length, (candidate_visits, prepare_scans, _, _) in measurements.items(): + assert candidate_visits <= prepare_scans * 2, ( + f"if-chain length {length}: {candidate_visits} retirement candidates across {prepare_scans} prepare calls" + ) + + small_work = measurements[150][0] + large_work = measurements[300][0] + assert large_work <= small_work * 2.25, ( + f"if-chain retirement candidate work grew faster than graph length: 150={small_work}, 300={large_work}" + ) diff --git a/tests/test_graph_execution_state.py b/tests/test_graph_execution_state.py index 462445b1099..4ab721d5d5c 100644 --- a/tests/test_graph_execution_state.py +++ b/tests/test_graph_execution_state.py @@ -5308,7 +5308,6 @@ def test_if_graph_optimized_behavior_executes_only_selected_simple_branch_and_re assert "false_value" not in executed_source_ids assert set(g.executed_history) == {"condition", "true_value", "false_value", "if", "selected_output"} - assert "false_value" not in executed_source_ids def test_if_graph_lowers_legacy_branch_choice_to_frame_scoped_activation_token(): From d2eb67339485aa2d7b2276e413f8a70dcef77ba8 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Wed, 14 Oct 2026 23:45:50 -0500 Subject: [PATCH 17/20] fix(graph): make materialization rollback atomic Journal planner preparation mutations so failed materialization can be retried safely. Add regression coverage for cache isolation, frame propagation, and rollback across both schedulers. --- invokeai/app/services/shared/README.md | 7 + invokeai/app/services/shared/graph.py | 43 +- .../app/services/shared/graph_if_runtime.py | 7 +- .../app/services/shared/graph_materializer.py | 222 ++++++++++ .../services/shared/graph_runtime_records.py | 2 +- .../test_execution_engine_differential.py | 395 ++++++++++++++++++ tests/test_graph_execution_state.py | 16 + 7 files changed, 683 insertions(+), 9 deletions(-) diff --git a/invokeai/app/services/shared/README.md b/invokeai/app/services/shared/README.md index 46fe3bf06f7..950932653e9 100644 --- a/invokeai/app/services/shared/README.md +++ b/invokeai/app/services/shared/README.md @@ -682,6 +682,13 @@ materialization. `Collect` may feed both branch ports of that same `If`; if it would serve different `If` owners, materialization raises `UnsupportedSharedIfCollectBranchesError` before invocation rather than dropping frames. +- For a root-conditioned `If` whose selected branch input is a completed direct `Collect`, compatibility + materialization preserves each non-root `Collect` iteration path through ordinary descendants. It runs one `If` + execution per preserved group, with a matching activation frame; the already-resolved root activation gate and + token are aligned to the first group. Ordinary descendants of `If.value` use those same frames, and a downstream + `Collect` gathers every completed group through the standard grouping behavior. This bounded case does not imply + propagation through further nested controls or loops. + - For each new exec node: - Deep-copy the source node; assign a fresh ID (and `index` for iterators). diff --git a/invokeai/app/services/shared/graph.py b/invokeai/app/services/shared/graph.py index 1b6bee8f236..938701b5c45 100644 --- a/invokeai/app/services/shared/graph.py +++ b/invokeai/app/services/shared/graph.py @@ -229,6 +229,7 @@ class _GraphExecutionStateRuntime: "prepared_if_exec_ids_by_source_and_frame", "prepared_if_exec_frames_by_source", "prepared_if_exec_indexed_sources", + "selected_if_frames_for_preparation", "prepared_exec_metadata", "prepared_exec_registry", "if_activation_controller_instance", @@ -276,6 +277,9 @@ def __init__(self) -> None: self.prepared_if_exec_ids_by_source_and_frame: dict[tuple[str, tuple[int, ...]], tuple[str, ...]] = {} self.prepared_if_exec_frames_by_source: dict[str, tuple[tuple[int, ...], ...]] = {} self.prepared_if_exec_indexed_sources: set[str] = set() + self.selected_if_frames_for_preparation: dict[ + tuple[str, Optional[str]], tuple[tuple[tuple[int, ...], ...], frozenset[tuple[int, ...]]] + ] = {} self.prepared_exec_metadata: dict[str, _PreparedExecNodeMetadata] = {} self.prepared_exec_registry: Optional[_PreparedExecRegistry] = None self.if_activation_controller_instance: Optional[_IfActivationController] = None @@ -414,6 +418,7 @@ def model_copy(self, *, update: Optional[dict[str, Any]] = None, deep: bool = Fa runtime_state.engine_frames = { key: frame for key, frame in runtime_state.engine_frames.items() if key[0] == copied.id } + runtime_state.selected_if_frames_for_preparation = runtime_state.selected_if_frames_for_preparation.copy() copied.__pydantic_private__["_runtime_state"] = runtime_state if update is not None and "graph" in update: copied._invalidate_source_graph_cache() @@ -425,6 +430,7 @@ def __copy__(self) -> "GraphExecutionState": runtime_state.engine_frames = { key: frame for key, frame in runtime_state.engine_frames.items() if key[0] == copied.id } + runtime_state.selected_if_frames_for_preparation = runtime_state.selected_if_frames_for_preparation.copy() copied.__pydantic_private__["_runtime_state"] = runtime_state return copied @@ -438,6 +444,27 @@ def _tx_record(self, undo: Callable[[], None]) -> None: if transaction is not None: transaction.record(undo) + def _run_materialization_transaction(self, operation: Callable[[], Any]) -> Any: + """Run one preparation expansion atomically using the scheduler mutation journal.""" + if self._apply_transaction is not None: + return operation() + + transaction = _ApplyTransaction() + object.__setattr__(self, "_apply_transaction", transaction) + try: + return operation() + except Exception: + try: + transaction.rollback() + finally: + object.__setattr__(self, "_apply_transaction", None) + self._reset_apply_derived_caches() + self._rehydrate_ready_queues() + raise + finally: + if self._apply_transaction is transaction: + object.__setattr__(self, "_apply_transaction", None) + def _tx_set_mapping(self, mapping: dict[Any, Any], key: Any, value: Any) -> None: transaction = self.__pydantic_private__["_runtime_state"].apply_transaction if transaction is not None: @@ -604,12 +631,13 @@ def _prepared_if_exec_frames(self, source_node_id: str) -> tuple[tuple[int, ...] def _invalidate_selected_if_frames_for_preparation(self, source_node_id: Optional[str] = None) -> None: """Clear selected-frame results after prepared If executions or branch choices change.""" + selected_frames = self._selected_if_frames_for_preparation if source_node_id is None: - self._selected_if_frames_for_preparation.clear() + self._selected_if_frames_for_preparation = {} return - for key in tuple(self._selected_if_frames_for_preparation): - if key[0] == source_node_id: - self._selected_if_frames_for_preparation.pop(key, None) + remaining = {key: value for key, value in selected_frames.items() if key[0] != source_node_id} + if len(remaining) != len(selected_frames): + self._selected_if_frames_for_preparation = remaining def _if_activation_controller(self) -> _IfActivationController: if self._if_activation_controller_instance is None: @@ -1551,7 +1579,7 @@ def _engine_frame(self, iteration_path: tuple[int, ...]) -> EngineExecutionFrame iteration_path=iteration_path, workflow_call_depth=len(self.workflow_call_stack), ) - runtime_state.engine_frames[key] = frame + self._tx_set_mapping(runtime_state.engine_frames, key, frame) return frame def _activation_gate(self, exec_node_id: str) -> ActivationGate: @@ -3454,10 +3482,10 @@ def _prepare_until_node_ready(self) -> Optional[BaseInvocation]: if isinstance(self._scheduler(), _GenericGraphSchedulerAdapter) and self._can_use_generic_for_scheduler() else self._materializer() ) - prepared_id = planner.prepare(base_graph) + prepared_id = self._run_materialization_transaction(lambda: planner.prepare(base_graph)) while prepared_id is not None: - prepared_id = planner.prepare(base_graph) + prepared_id = self._run_materialization_transaction(lambda: planner.prepare(base_graph)) if next_node is None: next_node = self._get_next_node() @@ -4276,6 +4304,7 @@ def delete_edge(self, edge: Edge) -> None: "prepared_if_exec_ids_by_source_and_frame", "prepared_if_exec_frames_by_source", "prepared_if_exec_indexed_sources", + "selected_if_frames_for_preparation", "prepared_exec_metadata", "prepared_exec_registry", "if_activation_controller_instance", diff --git a/invokeai/app/services/shared/graph_if_runtime.py b/invokeai/app/services/shared/graph_if_runtime.py index f98a13af5f2..e1843089893 100644 --- a/invokeai/app/services/shared/graph_if_runtime.py +++ b/invokeai/app/services/shared/graph_if_runtime.py @@ -414,12 +414,17 @@ def _is_empty_if_condition(state: "GraphExecutionState", dependency: ActivationD def _activation_gate(state: "GraphExecutionState", exec_node_id: str) -> ActivationGate: - return state._generic_runtime().register_gate( + runtime = state._generic_runtime() + is_new_gate = exec_node_id not in runtime.gates + gate = runtime.register_gate( gate_id=exec_node_id, owner_id=exec_node_id, frame=state._engine_frame(state._get_iteration_path(exec_node_id)), branches=("true_input", "false_input"), ) + if is_new_gate: + state._tx_record(lambda runtime=runtime, gate_id=exec_node_id: runtime.remove_gate(gate_id)) + return gate def _record_compatibility_activation_token( diff --git a/invokeai/app/services/shared/graph_materializer.py b/invokeai/app/services/shared/graph_materializer.py index 57b5e119a11..0f6ca9cbece 100644 --- a/invokeai/app/services/shared/graph_materializer.py +++ b/invokeai/app/services/shared/graph_materializer.py @@ -14,6 +14,7 @@ LoopState, ) from invokeai.app.services.shared import graph_if_runtime +from invokeai.app.services.shared.execution_engine.primitives import ActivationGate from invokeai.app.services.shared.graph_iterator_ancestry import _IteratorAncestry from invokeai.app.services.shared.graph_models import Edge, EdgeConnection from invokeai.app.services.shared.graph_validation import ( @@ -2137,13 +2138,233 @@ def _should_defer_sibling_if_collect_iterate(self, node_id: str) -> bool: if_exec_ids = self._state.source_prepared_mapping.get(if_node_id, set()) return not if_exec_ids or not all(exec_id in self._state._resolved_if_exec_branches for exec_id in if_exec_ids) + def _expand_root_if_for_completed_collect_branch(self, graph: "nx.DiGraph") -> Optional[str]: + """Align a root-conditioned If frame with each completed selected Collect frame.""" + iterator_graph = self.iterator_graph(graph) + for if_node_id, node in self._state.graph.nodes.items(): + if not isinstance(node, IfInvocation): + continue + condition_edges = self._state.graph._get_input_edges(if_node_id, "condition") + if ( + not condition_edges + or any(edge.source.node_id not in iterator_graph for edge in condition_edges) + or any( + self._get_iterator_ancestry(iterator_graph).get_iterator_axes_for_edge(edge) + for edge in condition_edges + ) + ): + continue + + root_exec_ids = [ + exec_id + for exec_id in self._state.source_prepared_mapping.get(if_node_id, ()) + if self._state._get_iteration_path(exec_id) == () + and self._state._is_pending_if(exec_id) + and (selected_field := self._state._resolved_if_exec_branches.get(exec_id)) is not None + ] + if len(root_exec_ids) != 1: + continue + root_exec_id = root_exec_ids[0] + selected_field = self._state._resolved_if_exec_branches[root_exec_id] + selected_edges = self._state.graph._get_input_edges(if_node_id, selected_field) + if not selected_edges or any( + not isinstance(self._state.graph.get_node(edge.source.node_id), CollectInvocation) + or edge.source.node_id not in self._state.executed + or any( + prepared_id not in self._state.results + for prepared_id in self._state.source_prepared_mapping.get(edge.source.node_id, ()) + ) + for edge in selected_edges + ): + continue + + branch_paths = sorted( + { + self._state._get_iteration_path(prepared_id) + for edge in selected_edges + for prepared_id in self._state.source_prepared_mapping.get(edge.source.node_id, ()) + } + ) + if not branch_paths or branch_paths[0] == (): + continue + + existing_paths = { + self._state._get_iteration_path(exec_id) + for exec_id in self._state.source_prepared_mapping.get(if_node_id, ()) + } + if existing_paths != {()}: + continue + + condition_mapping = [ + (self._state.prepared_source_mapping[edge.source.node_id], edge.source.node_id) + for edge in self._state.execution_graph._get_input_edges(root_exec_id) + if edge.destination.field == "condition" + ] + if not condition_mapping: + continue + + return self._state._run_materialization_transaction( + lambda graph=graph, + if_node_id=if_node_id, + root_exec_id=root_exec_id, + condition_mapping=condition_mapping, + branch_paths=branch_paths: self._materialize_root_if_for_completed_collect_branch( + graph, if_node_id, root_exec_id, condition_mapping, branch_paths + ) + ) + + return None + + def _materialize_root_if_for_completed_collect_branch( + self, + graph: "nx.DiGraph", + if_node_id: str, + root_exec_id: str, + condition_mapping: list[tuple[str, str]], + branch_paths: list[tuple[int, ...]], + ) -> str: + """Expand resolved root If state into its completed Collect frames atomically.""" + # The root If already resolved the branch gate so its iterators could run. Reuse that + # execution as the first selected frame, then create one If execution for every other + # closed Collect frame. Keep the durable activation token and generic gate frame aligned. + first_path, *remaining_paths = branch_paths + self._state._prepared_registry().set_iteration_path(root_exec_id, first_path) + frame = self._state._get_execution_frame(root_exec_id) + engine_frame = self._state._engine_frame(first_path) + generic_runtime = self._state._generic_runtime() + gate = generic_runtime.gates.get(root_exec_id) + if gate is not None: + previous_gate = ActivationGate.model_validate(gate.model_dump(mode="python")) + generic_runtime.replace_gate( + ActivationGate( + gate_id=gate.gate_id, + owner_id=gate.owner_id, + frame=engine_frame, + branches=gate.branches, + status=gate.status, + selected_branch=gate.selected_branch, + ) + ) + self._state._tx_record(lambda gate=previous_gate, runtime=generic_runtime: runtime.replace_gate(gate)) + for token_id, token in tuple(self._state.execution_tokens.items()): + if token.owner_node_id == root_exec_id and token.token_kind == "activation": + self._state._tx_set_mapping( + self._state.execution_tokens, + token_id, + token.model_copy(update={"frame": frame}), + ) + self._state._invalidate_selected_if_frames_for_preparation(if_node_id) + self._state._prepared_if_exec_indexed_sources.discard(if_node_id) + for key in tuple(self._state._prepared_if_exec_ids_by_source_and_frame): + if key[0] == if_node_id: + self._state._prepared_if_exec_ids_by_source_and_frame.pop(key, None) + self._state._prepared_if_exec_frames_by_source.pop(if_node_id, None) + + for iteration_path in remaining_paths: + self.create_execution_node( + if_node_id, + condition_mapping, + iteration_path, + input_fields={"condition"}, + ) + + for consumer_id in graph.successors(if_node_id): + if isinstance( + self._state.graph.get_node(consumer_id), + (IfInvocation, CollectInvocation, ForInvocation, IterateInvocation), + ): + continue + consumer_exec_ids = self._state.source_prepared_mapping.get(consumer_id, ()) + root_consumer_ids = [ + exec_id + for exec_id in consumer_exec_ids + if self._state._get_iteration_path(exec_id) == () and exec_id not in self._state.executed + ] + if len(root_consumer_ids) != 1: + continue + root_consumer_id = root_consumer_ids[0] + input_edges = self._state.execution_graph._get_input_edges(root_consumer_id) + if not any( + self._state.prepared_source_mapping.get(edge.source.node_id) == if_node_id for edge in input_edges + ): + continue + + self._state._prepared_registry().set_iteration_path(root_consumer_id, first_path) + for iteration_path in remaining_paths: + iteration_mapping = [] + for edge in self._state.graph._get_input_edges(consumer_id): + prepared_sources = self._state.source_prepared_mapping.get(edge.source.node_id, ()) + matching_source = next( + ( + self._state._prepared_registry().get_prepared_id_at_path( + edge.source.node_id, iteration_path[:path_depth] + ) + for path_depth in range(len(iteration_path), -1, -1) + if self._state._prepared_registry().get_prepared_id_at_path( + edge.source.node_id, iteration_path[:path_depth] + ) + is not None + ), + None, + ) + if matching_source is None and prepared_sources: + break + if matching_source is not None: + iteration_mapping.append((edge.source.node_id, matching_source)) + else: + self.create_execution_node(consumer_id, iteration_mapping, iteration_path) + + return root_exec_id + + def _get_unexecuted_root_collect_if_descendants(self, graph: "nx.DiGraph") -> set[str]: + """Wait to materialize descendants until selected root Collect frames finish their If copies.""" + iterator_graph = self.iterator_graph(graph) + iterator_ancestry = self._get_iterator_ancestry(iterator_graph) + deferred_node_ids: set[str] = set() + for if_node_id, node in self._state.graph.nodes.items(): + if if_node_id not in graph or not isinstance(node, IfInvocation): + continue + condition_edges = self._state.graph._get_input_edges(if_node_id, "condition") + if ( + not condition_edges + or any(edge.source.node_id not in iterator_graph for edge in condition_edges) + or any(iterator_ancestry.get_iterator_axes_for_edge(edge) for edge in condition_edges) + ): + continue + prepared_if_ids = self._state.source_prepared_mapping.get(if_node_id, ()) + if not prepared_if_ids or all(exec_id in self._state.executed for exec_id in prepared_if_ids): + continue + selected_fields = { + self._state._resolved_if_exec_branches.get(exec_id) + for exec_id in prepared_if_ids + if exec_id in self._state._resolved_if_exec_branches + } + if not selected_fields: + continue + if any( + not self._state.graph._get_input_edges(if_node_id, selected_field) + or any( + not isinstance(self._state.graph.get_node(edge.source.node_id), CollectInvocation) + for edge in self._state.graph._get_input_edges(if_node_id, selected_field) + ) + for selected_field in selected_fields + ): + continue + deferred_node_ids.update(nx.descendants(graph, if_node_id)) + return deferred_node_ids + def prepare(self, base_g: Optional["nx.DiGraph"] = None) -> Optional[str]: g = base_g if base_g is not None else self._state._get_source_graph_flat() self._get_sibling_iterate_if_collect_branches() iterator_ancestry = self._get_iterator_ancestry(self.iterator_graph(g)) iterator_ancestry.clear_effective_iteration_paths() self._attach_pending_if_inputs() + expanded_if_id = self._expand_root_if_for_completed_collect_branch(g) + if expanded_if_id is not None: + return expanded_if_id + self._attach_pending_if_inputs() self._ensure_activation_candidate_index(g) + deferred_root_collect_if_descendants = self._get_unexecuted_root_collect_if_descendants(g) if self._has_if_nodes: try: for node_id in self._get_activation_candidate_source_ids(g): @@ -2170,6 +2391,7 @@ def prepare(self, base_g: Optional["nx.DiGraph"] = None) -> Optional[str]: and self._is_deferred_nested_for_return(node_id, g) ) and not self._has_unmaterializable_for_final_input(node_id) + and node_id not in deferred_root_collect_if_descendants and ( all( edge.source.node_id in self._state.source_prepared_mapping diff --git a/invokeai/app/services/shared/graph_runtime_records.py b/invokeai/app/services/shared/graph_runtime_records.py index 0a6e0115433..e2739b50deb 100644 --- a/invokeai/app/services/shared/graph_runtime_records.py +++ b/invokeai/app/services/shared/graph_runtime_records.py @@ -9,7 +9,7 @@ class _ApplyTransaction: - """Small undo journal for scheduler mutations made by GraphExecutionState.apply().""" + """Small undo journal for scheduler and materialization state mutations.""" def __init__(self) -> None: self._undo: list[Callable[[], None]] = [] diff --git a/tests/app/services/shared/test_execution_engine_differential.py b/tests/app/services/shared/test_execution_engine_differential.py index 29f5398fd50..c553647d6bd 100644 --- a/tests/app/services/shared/test_execution_engine_differential.py +++ b/tests/app/services/shared/test_execution_engine_differential.py @@ -1,5 +1,6 @@ from __future__ import annotations +import copy import inspect import json import random @@ -9875,6 +9876,400 @@ def _collect_output(state: GraphExecutionState, source_id: str = "collect") -> l return state.results[next(iter(state.source_prepared_mapping[source_id]))].collection +def _zip_collect_branch_with_root_condition_graph(condition_value: int) -> Graph: + graph = Graph() + for node in ( + CollectionConcatInvocation(id="left_source", first=[1, 2]), + CollectionConcatInvocation(id="right_source", first=[10, 20]), + IterateInvocation(id="left"), + IterateInvocation(id="right"), + AddInvocation(id="body"), + CollectInvocation(id="branch_collect"), + FixedCollectionTestInvocation(id="condition", value=condition_value), + AnyTypeTestInvocation(id="fallback", value="fallback"), + IfInvocation(id="if"), + AnyTypeTestInvocation(id="after"), + ): + graph.add_node(node) + for source, source_field, destination, destination_field in ( + ("left_source", "collection", "left", "collection"), + ("right_source", "collection", "right", "collection"), + ("left", "item", "body", "a"), + ("right", "item", "body", "b"), + ("body", "value", "branch_collect", "item"), + ("condition", "condition", "if", "condition"), + ("branch_collect", "collection", "if", "true_input"), + ("fallback", "value", "if", "false_input"), + ("if", "value", "after", "value"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + return graph + + +def _zip_collect_then_after_graph() -> Graph: + graph = _zip_collect_branch_with_root_condition_graph(1) + graph.delete_node("if") + graph.delete_node("fallback") + graph.add_edge(create_edge("branch_collect", "collection", "after", "value")) + return graph + + +def _zip_collect_branch_with_after2_graph() -> Graph: + graph = _zip_collect_branch_with_root_condition_graph(condition_value=1) + graph.add_node(AnyTypeTestInvocation(id="after2")) + graph.add_edge(create_edge("after", "value", "after2", "value")) + return graph + + +def _zip_collect_branch_with_after_collect_graph() -> Graph: + graph = _zip_collect_branch_with_root_condition_graph(condition_value=1) + graph.add_node(CollectInvocation(id="after_collect")) + graph.add_edge(create_edge("after", "value", "after_collect", "item")) + return graph + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_zip_collect_then_after_completes_without_if(force_compatibility_scheduler: bool) -> None: + state = GraphExecutionState(graph=_zip_collect_then_after_graph()) + state._execution_scheduler = ( + _ExecutionScheduler(state) if force_compatibility_scheduler else _GenericGraphSchedulerAdapter(state) + ) + + trace, state = _run_graph(state, force_compatibility_scheduler=force_compatibility_scheduler) + + assert state.is_complete(), f"graph stalled after {trace}" + assert not state.has_error() + after_ids = sorted(state.source_prepared_mapping["after"], key=state._get_iteration_path) + assert [state._get_iteration_path(node_id) for node_id in after_ids] == [(0,), (1,)] + assert [state.results[node_id].value for node_id in after_ids] == [[11, 21], [12, 22]] + assert trace.count("body") == 4 + assert trace.count("branch_collect") == 2 + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_zip_collect_then_after_with_empty_right_completes_without_if( + force_compatibility_scheduler: bool, +) -> None: + graph = _zip_collect_then_after_graph() + graph.update_node("right_source", CollectionConcatInvocation(id="right_source", first=[])) + state = GraphExecutionState(graph=graph) + state._execution_scheduler = ( + _ExecutionScheduler(state) if force_compatibility_scheduler else _GenericGraphSchedulerAdapter(state) + ) + + trace, state = _run_graph(state, force_compatibility_scheduler=force_compatibility_scheduler) + + assert state.is_complete(), f"graph stalled after {trace}" + assert not state.has_error() + after_ids = sorted(state.source_prepared_mapping["after"], key=state._get_iteration_path) + assert [state._get_iteration_path(node_id) for node_id in after_ids] == [(0,), (1,)] + assert [state.results[node_id].value for node_id in after_ids] == [[], []] + assert trace.count("body") == 0 + assert trace.count("branch_collect") == 2 + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +@pytest.mark.parametrize( + ("condition_value", "expected"), [(1, [[11, 21], [12, 22]]), (0, "fallback")], ids=["true", "false"] +) +def test_zip_collect_branch_with_root_condition_completes( + force_compatibility_scheduler: bool, condition_value: int, expected: Any +) -> None: + state = GraphExecutionState(graph=_zip_collect_branch_with_root_condition_graph(condition_value)) + state._execution_scheduler = ( + _ExecutionScheduler(state) if force_compatibility_scheduler else _GenericGraphSchedulerAdapter(state) + ) + + trace, state = _run_graph(state, force_compatibility_scheduler=force_compatibility_scheduler) + + assert state.is_complete(), f"graph stalled after {trace}" + assert not state.has_error() + after_ids = sorted(state.source_prepared_mapping["after"], key=state._get_iteration_path) + if_ids = set(state.source_prepared_mapping["if"]) + activation_tokens = sorted( + ( + token + for token in state.execution_tokens.values() + if token.token_kind == "activation" and token.owner_node_id in if_ids + ), + key=lambda token: token.frame.iteration_path, + ) + if condition_value: + assert [state._get_iteration_path(node_id) for node_id in after_ids] == [(0,), (1,)] + assert [state.results[node_id].value for node_id in after_ids] == expected + assert trace.count("body") == 4 + assert trace.count("branch_collect") == 2 + assert trace.count("if") == 2 + assert [(token.port, token.frame.iteration_path) for token in activation_tokens] == [ + ("true_input", (0,)), + ("true_input", (1,)), + ] + restored = load_execution_state(dump_execution_state(state)) + assert restored.is_complete() + assert not restored.has_error() + restored_after_ids = sorted(restored.source_prepared_mapping["after"], key=restored._get_iteration_path) + assert [restored.results[node_id].value for node_id in restored_after_ids] == expected + else: + assert [state._get_iteration_path(node_id) for node_id in after_ids] == [()] + assert [state.results[node_id].value for node_id in after_ids] == [expected] + assert trace.count("body") == 0 + assert trace.count("branch_collect") == 0 + assert trace.count("if") == 1 + assert [(token.port, token.frame.iteration_path) for token in activation_tokens] == [("false_input", ())] + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_zip_collect_root_condition_preserves_frames_through_ordinary_descendants( + force_compatibility_scheduler: bool, +) -> None: + state = GraphExecutionState(graph=_zip_collect_branch_with_after2_graph()) + state._execution_scheduler = ( + _ExecutionScheduler(state) if force_compatibility_scheduler else _GenericGraphSchedulerAdapter(state) + ) + + trace, state = _run_graph(state, force_compatibility_scheduler=force_compatibility_scheduler) + + assert state.is_complete(), f"graph stalled after {trace}" + assert not state.has_error() + after2_ids = sorted(state.source_prepared_mapping["after2"], key=state._get_iteration_path) + assert [state._get_iteration_path(node_id) for node_id in after2_ids] == [(0,), (1,)] + assert [state.results[node_id].value for node_id in after2_ids] == [[11, 21], [12, 22]] + assert trace.count("after2") == 2 + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_zip_collect_root_condition_preserves_groups_in_downstream_collect( + force_compatibility_scheduler: bool, +) -> None: + state = GraphExecutionState(graph=_zip_collect_branch_with_after_collect_graph()) + state._execution_scheduler = ( + _ExecutionScheduler(state) if force_compatibility_scheduler else _GenericGraphSchedulerAdapter(state) + ) + + trace, state = _run_graph(state, force_compatibility_scheduler=force_compatibility_scheduler) + + assert state.is_complete(), f"graph stalled after {trace}" + assert not state.has_error() + collect_ids = sorted(state.source_prepared_mapping["after_collect"], key=state._get_iteration_path) + assert len(collect_ids) == 1 + assert state._get_iteration_path(collect_ids[0]) == () + assert state.results[collect_ids[0]].collection == [[11, 21], [12, 22]] + assert trace.count("after_collect") == 1 + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_zip_collect_root_condition_expansion_rolls_back_after_if_initialization_failure( + monkeypatch: pytest.MonkeyPatch, force_compatibility_scheduler: bool +) -> None: + state = GraphExecutionState(graph=_zip_collect_branch_with_root_condition_graph(condition_value=1)) + state._execution_scheduler = ( + _ExecutionScheduler(state) if force_compatibility_scheduler else _GenericGraphSchedulerAdapter(state) + ) + + while ( + len( + [exec_id for exec_id in state.source_prepared_mapping.get("branch_collect", ()) if exec_id in state.results] + ) + < 2 + ): + node = state.next() + assert node is not None + state.complete(node.id, node.invoke(Mock())) + + root_if_ids = [ + exec_id for exec_id in state.source_prepared_mapping["if"] if state._get_iteration_path(exec_id) == () + ] + assert len(root_if_ids) == 1 + root_if_id = root_if_ids[0] + assert root_if_id not in state.executed + + before_nodes = set(state.execution_graph.nodes) + before_edges = tuple(state.execution_graph.edges) + before_prepared_sources = dict(state.prepared_source_mapping) + before_source_prepared = {source_id: set(exec_ids) for source_id, exec_ids in state.source_prepared_mapping.items()} + before_paths = dict(state.prepared_iteration_paths) + before_indegree = dict(state.indegree) + before_tokens = state.execution_tokens.copy() + before_gates = copy.deepcopy(state._generic_runtime().gates) + materializer = state._materializer() + initialize_execution_node = materializer._initialize_execution_node + + def fail_after_if_initialization(exec_node_id: str, input_edges: Any = None) -> None: + initialize_execution_node(exec_node_id, input_edges) + if state.prepared_source_mapping.get(exec_node_id) == "if" and state._get_iteration_path(exec_node_id) == (1,): + raise RuntimeError("injected If-copy initialization failure") + + monkeypatch.setattr(materializer, "_initialize_execution_node", fail_after_if_initialization) + with pytest.raises(RuntimeError, match="injected If-copy initialization failure"): + while (node := state.next()) is not None: + state.complete(node.id, node.invoke(Mock())) + + assert set(state.execution_graph.nodes) == before_nodes + assert tuple(state.execution_graph.edges) == before_edges + assert state.prepared_source_mapping == before_prepared_sources + assert state.source_prepared_mapping == before_source_prepared + assert state.prepared_iteration_paths == before_paths + assert state.indegree == before_indegree + assert state.execution_tokens == before_tokens + assert state._generic_runtime().gates == before_gates + assert state._get_iteration_path(root_if_id) == () + + monkeypatch.undo() + trace, state = _run_graph(state, force_compatibility_scheduler=force_compatibility_scheduler) + assert state.is_complete(), f"graph stalled after retry: {trace}" + assert not state.has_error() + after_ids = sorted(state.source_prepared_mapping["after"], key=state._get_iteration_path) + assert [state.results[node_id].value for node_id in after_ids] == [[11, 21], [12, 22]] + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_zip_collect_root_condition_deferred_consumer_failure_rolls_back_and_retries( + monkeypatch: pytest.MonkeyPatch, force_compatibility_scheduler: bool +) -> None: + state = GraphExecutionState(graph=_zip_collect_branch_with_root_condition_graph(condition_value=1)) + state._execution_scheduler = ( + _ExecutionScheduler(state) if force_compatibility_scheduler else _GenericGraphSchedulerAdapter(state) + ) + + while ( + len( + [exec_id for exec_id in state.source_prepared_mapping.get("branch_collect", ()) if exec_id in state.results] + ) + < 2 + ): + node = state.next() + assert node is not None + state.complete(node.id, node.invoke(Mock())) + + before_after_ids = set(state.source_prepared_mapping.get("after", ())) + before_after_edges = tuple( + edge + for edge in state.execution_graph.edges + if state.prepared_source_mapping.get(edge.destination.node_id) == "after" + ) + materializer = state._materializer() + initialize_execution_node = materializer._initialize_execution_node + + def fail_after_second_frame(exec_node_id: str, input_edges: Any = None) -> None: + initialize_execution_node(exec_node_id, input_edges) + if state.prepared_source_mapping.get(exec_node_id) == "after" and state._get_iteration_path(exec_node_id) == ( + 1, + ): + raise RuntimeError("injected deferred consumer initialization failure") + + monkeypatch.setattr(materializer, "_initialize_execution_node", fail_after_second_frame) + with pytest.raises(RuntimeError, match="injected deferred consumer initialization failure"): + while (node := state.next()) is not None: + state.complete(node.id, node.invoke(Mock())) + + assert state.source_prepared_mapping.get("after", set()) == before_after_ids + assert ( + tuple( + edge + for edge in state.execution_graph.edges + if state.prepared_source_mapping.get(edge.destination.node_id) == "after" + ) + == before_after_edges + ) + + monkeypatch.undo() + trace, state = _run_graph(state, force_compatibility_scheduler=force_compatibility_scheduler) + assert state.is_complete(), f"graph stalled after retry: {trace}" + assert not state.has_error() + after_ids = sorted(state.source_prepared_mapping["after"], key=state._get_iteration_path) + assert [state.results[node_id].value for node_id in after_ids] == [[11, 21], [12, 22]] + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_zip_collect_branch_with_empty_selected_collect_completes(force_compatibility_scheduler: bool) -> None: + graph = _zip_collect_branch_with_root_condition_graph(condition_value=1) + graph.update_node("left_source", CollectionConcatInvocation(id="left_source", first=[])) + state = GraphExecutionState(graph=graph) + state._execution_scheduler = ( + _ExecutionScheduler(state) if force_compatibility_scheduler else _GenericGraphSchedulerAdapter(state) + ) + + trace, state = _run_graph(state, force_compatibility_scheduler=force_compatibility_scheduler) + + assert state.is_complete(), f"graph stalled after {trace}" + assert not state.has_error() + after_ids = sorted(state.source_prepared_mapping["after"], key=state._get_iteration_path) + collect_ids = sorted(state.source_prepared_mapping["branch_collect"], key=state._get_iteration_path) + assert len(after_ids) == 1 + assert state._get_iteration_path(after_ids[0]) == () + assert state.results[after_ids[0]].value == [] + assert len(collect_ids) == 1 + assert state.results[collect_ids[0]].collection == [] + assert trace.count("right") == 2 + assert trace.count("body") == 0 + assert trace.count("branch_collect") == 1 + assert trace.count("if") == 1 + activation_tokens = [token for token in state.execution_tokens.values() if token.token_kind == "activation"] + assert [(token.port, token.frame.iteration_path) for token in activation_tokens] == [("true_input", ())] + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_zip_collect_branch_with_empty_right_preserves_left_frames(force_compatibility_scheduler: bool) -> None: + graph = _zip_collect_branch_with_root_condition_graph(condition_value=1) + graph.update_node("right_source", CollectionConcatInvocation(id="right_source", first=[])) + state = GraphExecutionState(graph=graph) + state._execution_scheduler = ( + _ExecutionScheduler(state) if force_compatibility_scheduler else _GenericGraphSchedulerAdapter(state) + ) + + trace, state = _run_graph(state, force_compatibility_scheduler=force_compatibility_scheduler) + + assert state.is_complete(), f"graph stalled after {trace}" + assert not state.has_error() + after_ids = sorted(state.source_prepared_mapping["after"], key=state._get_iteration_path) + assert [state._get_iteration_path(node_id) for node_id in after_ids] == [(0,), (1,)] + assert [state.results[node_id].value for node_id in after_ids] == [[], []] + assert trace.count("body") == 0 + assert trace.count("branch_collect") == 2 + assert trace.count("if") == 2 + tokens = sorted( + (token for token in state.execution_tokens.values() if token.token_kind == "activation"), + key=lambda token: token.frame.iteration_path, + ) + assert [(token.port, token.frame.iteration_path) for token in tokens] == [ + ("true_input", (0,)), + ("true_input", (1,)), + ] + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_zip_collect_branch_activation_frames_resume_after_round_trip(force_compatibility_scheduler: bool) -> None: + state = GraphExecutionState(graph=_zip_collect_branch_with_root_condition_graph(condition_value=1)) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) + + partial_trace, state = _run_graph( + state, + force_compatibility_scheduler=force_compatibility_scheduler, + stop_after=14, + ) + assert partial_trace[-1] == "if" + assert partial_trace.count("if") == 1 + assert not state.is_complete() + assert not state.has_error() + partial_if_ids = sorted(state.source_prepared_mapping["if"], key=state._get_iteration_path) + assert [state._get_iteration_path(node_id) for node_id in partial_if_ids] == [(0,), (1,)] + assert sum(node_id in state.results for node_id in partial_if_ids) == 1 + assert sum(node_id not in state.results for node_id in partial_if_ids) == 1 + + state = load_execution_state(dump_execution_state(state)) + if force_compatibility_scheduler: + _restore_compatibility_scheduler(state) + remaining_trace, state = _run_graph(state) + + assert state.is_complete(), f"resume stalled after {remaining_trace}" + assert not state.has_error() + assert not state._pending_if_exec_nodes + assert not state._has_pending_lifecycle_execution() + after_ids = sorted(state.source_prepared_mapping["after"], key=state._get_iteration_path) + assert [state.results[node_id].value for node_id in after_ids] == [[11, 21], [12, 22]] + + @pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) @pytest.mark.parametrize( ("left", "right", "expected"), diff --git a/tests/test_graph_execution_state.py b/tests/test_graph_execution_state.py index 4ab721d5d5c..5c9293a7de2 100644 --- a/tests/test_graph_execution_state.py +++ b/tests/test_graph_execution_state.py @@ -861,6 +861,22 @@ def test_graph_state_shallow_copy_reset_does_not_clear_original_runtime_state(co assert not copied._ready_queues +@pytest.mark.parametrize( + "copy_state", + [pytest.param(lambda state: state.model_copy(), id="model_copy"), pytest.param(copy.copy, id="copy")], +) +def test_graph_state_shallow_copy_selected_if_preparation_cache_isolated(copy_state): + state = GraphExecutionState(graph=Graph()) + cache_key = ("if", "true_input") + state._selected_if_frames_for_preparation[cache_key] = (((0,),), frozenset({(0,)})) + + copied = copy_state(state) + copied._invalidate_selected_if_frames_for_preparation("if") + + assert state._selected_if_frames_for_preparation == {cache_key: (((0,),), frozenset({(0,)}))} + assert not copied._selected_if_frames_for_preparation + + @pytest.mark.parametrize("port", ["type", "output_meta", "loop_linkage"]) @pytest.mark.parametrize("effect_kind, token_kind", [("emit", "data"), ("close_stream", "stream_end")]) def test_graph_state_apply_rejects_reserved_effect_ports(port: str, effect_kind: str, token_kind: str): From dd4eee499d261993b29b440021ec2e7951794d39 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Thu, 15 Oct 2026 05:29:29 -0500 Subject: [PATCH 18/20] docs(graph): document materialization rollback --- invokeai/app/services/shared/README.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/invokeai/app/services/shared/README.md b/invokeai/app/services/shared/README.md index 950932653e9..ee4d88e93aa 100644 --- a/invokeai/app/services/shared/README.md +++ b/invokeai/app/services/shared/README.md @@ -699,6 +699,11 @@ materialization. - Try to resolve any `If`-specific scheduling state. - If the node is ready and not deferred by an unresolved `If`, enqueue it into its class queue. +Each planner preparation expansion is applied as a state transaction. If expansion raises, journaled execution-graph, +prepared-mapping, frame, activation-token, gate, and scheduler mutations are rolled back; derived caches are reset and +ready queues are rebuilt from the restored state before the exception propagates. A later `next()` can retry from the +pre-expansion state without retaining partially materialized nodes or activation state. + The execution-node builder retains a filtered iterator graph for the current source-graph identity. It builds this graph when compatibility preparation needs sibling-Collect relationships or iterator ancestry. The graph replaces Collect item inputs with its collection inputs for traversal; edge-scoped ancestry retains item axes needed for grouping. From f6d3fdb5b878b597399853526aae93b21bfcd786 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Thu, 15 Oct 2026 14:11:27 -0500 Subject: [PATCH 19/20] fix(graph): address execution engine review follow-ups Restore supported deeper-Iterate If broadcasts and remove If-free prepare transaction overhead. Index active If dependents and document the transaction boundary. --- invokeai/app/services/shared/README.md | 15 +- invokeai/app/services/shared/graph.py | 61 +++- .../app/services/shared/graph_if_runtime.py | 11 + .../test_execution_engine_differential.py | 69 ++++ .../shared/test_for_scheduler_performance.py | 307 +++++++++++++++++- .../services/shared/test_graph_if_runtime.py | 4 + 6 files changed, 448 insertions(+), 19 deletions(-) diff --git a/invokeai/app/services/shared/README.md b/invokeai/app/services/shared/README.md index ee4d88e93aa..cce28b06491 100644 --- a/invokeai/app/services/shared/README.md +++ b/invokeai/app/services/shared/README.md @@ -699,10 +699,17 @@ materialization. - Try to resolve any `If`-specific scheduling state. - If the node is ready and not deferred by an unresolved `If`, enqueue it into its class queue. -Each planner preparation expansion is applied as a state transaction. If expansion raises, journaled execution-graph, -prepared-mapping, frame, activation-token, gate, and scheduler mutations are rolled back; derived caches are reset and -ready queues are rebuilt from the restored state before the exception propagates. A later `next()` can retry from the -pre-expansion state without retaining partially materialized nodes or activation state. +`next()` applies `_GenericForPlanner.prepare()` and If-bearing `_ExecutionMaterializer.prepare()` calls as state +transactions. If one of those expansions raises, journaled execution-graph, prepared-mapping, frame, +activation-token, gate, and scheduler mutations are rolled back; derived caches are reset and ready queues are rebuilt +from the restored state before the exception propagates. A later `next()` can retry from the pre-expansion state +without retaining partially materialized nodes or activation state. If-free `_ExecutionMaterializer.prepare()` runs +without this transaction wrapper to avoid overhead on ordinary graphs; root-If completed-Collect expansion retains +its own transaction. + +This transaction boundary does not cover pending-If attachment (`_attach_pending_if_inputs`), +`prepare_nested_iterate_sequence()`, or direct Iterate/Collect preparation. Those operations run outside the wrapper +and must be considered separately when reasoning about rollback behavior. The execution-node builder retains a filtered iterator graph for the current source-graph identity. It builds this graph when compatibility preparation needs sibling-Collect relationships or iterator ancestry. The graph replaces diff --git a/invokeai/app/services/shared/graph.py b/invokeai/app/services/shared/graph.py index 938701b5c45..53b456f188c 100644 --- a/invokeai/app/services/shared/graph.py +++ b/invokeai/app/services/shared/graph.py @@ -224,6 +224,7 @@ class _GraphExecutionStateRuntime: "if_activation_dependencies_by_source", "if_activation_dependencies_by_exec", "if_activation_dependents_by_dependency", + "if_activation_active_dependents_by_dependency", "if_activation_dependents_indexed_exec_ids", "if_activation_dependents_index_complete", "prepared_if_exec_ids_by_source_and_frame", @@ -272,6 +273,7 @@ def __init__(self) -> None: ] = {} self.if_activation_dependencies_by_exec: dict[str, tuple[ActivationDependency, ...]] = {} self.if_activation_dependents_by_dependency: dict[ActivationDependency, tuple[str, ...]] = {} + self.if_activation_active_dependents_by_dependency: dict[ActivationDependency, dict[str, None]] = {} self.if_activation_dependents_indexed_exec_ids: set[str] = set() self.if_activation_dependents_index_complete = False self.prepared_if_exec_ids_by_source_and_frame: dict[tuple[str, tuple[int, ...]], tuple[str, ...]] = {} @@ -561,6 +563,7 @@ def _reset_apply_derived_caches(self) -> None: self._if_activation_dependencies_by_source = {} self._if_activation_dependencies_by_exec = {} self._if_activation_dependents_by_dependency = {} + self._if_activation_active_dependents_by_dependency = {} self._if_activation_dependents_indexed_exec_ids = set() self._if_activation_dependents_index_complete = False self._prepared_if_exec_ids_by_source_and_frame = {} @@ -684,8 +687,28 @@ def _index_activation_dependencies(self, exec_node_id: str, dependencies: tuple[ dependency, (*dependents, exec_node_id), ) + if exec_node_id in self.prepared_source_mapping: + metadata = self._get_prepared_exec_metadata(exec_node_id) + else: + metadata = None + if metadata is not None and metadata.state not in ("executed", "skipped"): + active_dependents = self._if_activation_active_dependents_by_dependency.get(dependency) + if active_dependents is None: + active_dependents = {} + self._tx_set_mapping( + self._if_activation_active_dependents_by_dependency, dependency, active_dependents + ) + self._tx_set_mapping(active_dependents, exec_node_id, None) self._tx_add_set(self._if_activation_dependents_indexed_exec_ids, exec_node_id) + def _deactivate_if_activation_dependent(self, exec_node_id: str) -> None: + """Remove a terminal execution node from the ordered active-candidate view.""" + dependencies = self._if_activation_dependencies_by_exec.get(exec_node_id, ()) + for dependency in dependencies: + active_dependents = self._if_activation_active_dependents_by_dependency.get(dependency) + if active_dependents is not None: + self._tx_pop_mapping(active_dependents, exec_node_id) + def _get_activation_dependents_for_if(self, source_if_id: str, iteration_path: tuple[int, ...]) -> tuple[str, ...]: if not self._if_activation_dependents_index_complete: for exec_node_id in self.prepared_source_mapping: @@ -693,21 +716,20 @@ def _get_activation_dependents_for_if(self, source_if_id: str, iteration_path: t self._record_activation_dependencies(exec_node_id) self._tx_set_attr(self, "_if_activation_dependents_index_complete", True) - candidates = { - exec_node_id - for branch in ("true_input", "false_input") - for exec_node_id in self._if_activation_dependents_by_dependency.get( - ActivationDependency(owner_id=source_if_id, branch=branch, frame=iteration_path), () - ) - } - if iteration_path: + candidates: dict[str, None] = {} + for branch in ("true_input", "false_input"): candidates.update( - exec_node_id - for branch in ("true_input", "false_input") - for exec_node_id in self._if_activation_dependents_by_dependency.get( - ActivationDependency(owner_id=source_if_id, branch=branch, frame=()), () + self._if_activation_active_dependents_by_dependency.get( + ActivationDependency(owner_id=source_if_id, branch=branch, frame=iteration_path), {} ) ) + if iteration_path: + for branch in ("true_input", "false_input"): + candidates.update( + self._if_activation_active_dependents_by_dependency.get( + ActivationDependency(owner_id=source_if_id, branch=branch, frame=()), {} + ) + ) if not candidates: return () execution_graph = self._get_execution_graph_flat() @@ -834,6 +856,7 @@ def _invalidate_source_graph_cache(self) -> None: runtime.if_activation_dependencies_by_source = {} runtime.if_activation_dependencies_by_exec = {} runtime.if_activation_dependents_by_dependency = {} + runtime.if_activation_active_dependents_by_dependency = {} runtime.if_activation_dependents_indexed_exec_ids = set() runtime.if_activation_dependents_index_complete = False runtime.if_activation_controller_instance = None @@ -3265,6 +3288,8 @@ def _get_prepared_exec_metadata(self, exec_node_id: str) -> _PreparedExecNodeMet def _set_prepared_exec_state(self, exec_node_id: str, state: PreparedExecState) -> None: self._prepared_registry().set_state(exec_node_id, state) + if state in ("executed", "skipped"): + self._deactivate_if_activation_dependent(exec_node_id) def _get_iteration_path(self, exec_node_id: str) -> tuple[int, ...]: return self._runtime().get_iteration_path(exec_node_id) @@ -3482,10 +3507,16 @@ def _prepare_until_node_ready(self) -> Optional[BaseInvocation]: if isinstance(self._scheduler(), _GenericGraphSchedulerAdapter) and self._can_use_generic_for_scheduler() else self._materializer() ) - prepared_id = self._run_materialization_transaction(lambda: planner.prepare(base_graph)) + + def prepare_next() -> Optional[str]: + if getattr(planner, "_has_if_nodes", True) is False: + return planner.prepare(base_graph) + return self._run_materialization_transaction(lambda: planner.prepare(base_graph)) + + prepared_id = prepare_next() while prepared_id is not None: - prepared_id = self._run_materialization_transaction(lambda: planner.prepare(base_graph)) + prepared_id = prepare_next() if next_node is None: next_node = self._get_next_node() @@ -3555,6 +3586,7 @@ def _reset_runtime_caches(self) -> None: self._if_activation_dependencies_by_source = {} self._if_activation_dependencies_by_exec = {} self._if_activation_dependents_by_dependency = {} + self._if_activation_active_dependents_by_dependency = {} self._if_activation_dependents_indexed_exec_ids = set() self._if_activation_dependents_index_complete = False self._prepared_if_exec_ids_by_source_and_frame = {} @@ -4299,6 +4331,7 @@ def delete_edge(self, edge: Edge) -> None: "if_activation_dependencies_by_source", "if_activation_dependencies_by_exec", "if_activation_dependents_by_dependency", + "if_activation_active_dependents_by_dependency", "if_activation_dependents_indexed_exec_ids", "if_activation_dependents_index_complete", "prepared_if_exec_ids_by_source_and_frame", diff --git a/invokeai/app/services/shared/graph_if_runtime.py b/invokeai/app/services/shared/graph_if_runtime.py index e1843089893..ecee72c78ef 100644 --- a/invokeai/app/services/shared/graph_if_runtime.py +++ b/invokeai/app/services/shared/graph_if_runtime.py @@ -33,6 +33,16 @@ def validate_if_iteration_axes(state: "GraphExecutionState") -> None: if not condition_edges: continue condition_iterators = iterator_ancestry.get_iterator_ancestors_for_edge(condition_edges[0]) + branch_edges = { + branch_field: state.graph._get_input_edges(node_id, branch_field) + for branch_field in ("true_input", "false_input") + } + broadcast_axes = None + if all(len(edges) == 1 for edges in branch_edges.values()): + true_axes = iterator_ancestry.get_iterator_ancestors_for_edge(branch_edges["true_input"][0]) + false_axes = iterator_ancestry.get_iterator_ancestors_for_edge(branch_edges["false_input"][0]) + if true_axes == false_axes and condition_iterators < true_axes: + broadcast_axes = true_axes for branch_field in ("true_input", "false_input"): for edge in state.graph._get_input_edges(node_id, branch_field): branch_iterators = iterator_ancestry.get_iterator_ancestors_for_edge(edge) @@ -40,6 +50,7 @@ def validate_if_iteration_axes(state: "GraphExecutionState") -> None: if ( condition_iterators and condition_iterators < branch_iterators + and branch_iterators != broadcast_axes and not isinstance(branch_source, CollectInvocation) ): raise UnsupportedIfIterationAxesError( diff --git a/tests/app/services/shared/test_execution_engine_differential.py b/tests/app/services/shared/test_execution_engine_differential.py index c553647d6bd..3deb1b14004 100644 --- a/tests/app/services/shared/test_execution_engine_differential.py +++ b/tests/app/services/shared/test_execution_engine_differential.py @@ -10550,6 +10550,75 @@ def test_if_branch_deeper_than_its_condition_completes_or_is_rejected( assert state.executed_history == [] +def _broadcast_deeper_if_graph(*, selected: bool) -> Graph: + """An outer condition selects one of two complete inner-Iterate branches.""" + graph = Graph() + for node in ( + CollectionConcatInvocation(id="outer_source", first=[2, 0] if selected else [0, 2]), + IterateInvocation(id="outer"), + AddInvocation(id="range_stop", b=1), + RangeInvocation(id="range", start=0, step=1), + IterateInvocation(id="inner"), + FixedCollectionTestInvocation(id="condition"), + AddInvocation(id="true_value", b=10), + AddInvocation(id="false_value", b=100), + IfInvocation(id="if"), + CollectInvocation(id="collect"), + ): + graph.add_node(node) + for source, source_field, destination, destination_field in ( + ("outer_source", "collection", "outer", "collection"), + ("outer", "item", "condition", "value"), + ("condition", "condition", "if", "condition"), + ("outer", "item", "range_stop", "a"), + ("range_stop", "value", "range", "stop"), + ("range", "collection", "inner", "collection"), + ("inner", "item", "true_value", "a"), + ("inner", "item", "false_value", "a"), + ("true_value", "value", "if", "true_input"), + ("false_value", "value", "if", "false_input"), + ("if", "value", "collect", "item"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + return graph + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +@pytest.mark.parametrize("selected", [True, False], ids=["true-branch", "false-branch"]) +def test_if_broadcasts_a_deeper_iterate_branch_across_each_outer_condition_frame( + force_compatibility_scheduler: bool, selected: bool +) -> None: + state = GraphExecutionState(graph=_broadcast_deeper_if_graph(selected=selected)) + + trace, state = _run_graph(state, force_compatibility_scheduler=force_compatibility_scheduler) + + assert state.is_complete(), f"graph stalled after {trace}" + assert not state.has_error() + expected = [[10, 11, 12], [100]] if selected else [[100], [10, 11, 12]] + collect_nodes = sorted(state.source_prepared_mapping["collect"], key=state._get_iteration_path) + collected = [state.results[node_id].collection for node_id in collect_nodes] + assert collected == expected, f"{trace=}" + assert trace.count("if") == 4 + assert trace.count("true_value") == 3 + assert trace.count("false_value") == 1 + assert len(collect_nodes) == 2 + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_if_broadcast_deeper_iterate_frames_resume_after_dump_load( + force_compatibility_scheduler: bool, +) -> None: + state = GraphExecutionState(graph=_broadcast_deeper_if_graph(selected=True)) + _, partial = _run_graph(state, force_compatibility_scheduler=force_compatibility_scheduler, stop_after=7) + restored = load_execution_state(dump_execution_state(partial)) + + trace, restored = _run_graph(restored, force_compatibility_scheduler=force_compatibility_scheduler) + + assert restored.is_complete(), f"graph stalled after resume: {trace}" + collect_nodes = sorted(restored.source_prepared_mapping["collect"], key=restored._get_iteration_path) + assert [restored.results[node_id].collection for node_id in collect_nodes] == [[10, 11, 12], [100]] + + @pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) @pytest.mark.parametrize( ("condition", "per_item_field", "root_field", "expected_value"), diff --git a/tests/app/services/shared/test_for_scheduler_performance.py b/tests/app/services/shared/test_for_scheduler_performance.py index bcbe2c73c80..a024c6e7145 100644 --- a/tests/app/services/shared/test_for_scheduler_performance.py +++ b/tests/app/services/shared/test_for_scheduler_performance.py @@ -4,24 +4,28 @@ Optional absolute-time benchmarks remain marked `slow`. """ +import gc import time from collections.abc import Callable from unittest.mock import Mock +import networkx as nx import pytest from invokeai.app.invocations.collections import RangeInvocation from invokeai.app.invocations.logic import IfInvocation from invokeai.app.invocations.loops import ForInvocation, ForReturnInvocation from invokeai.app.invocations.primitives import BooleanInvocation +from invokeai.app.services.shared.execution_state_migration import dump_execution_state, load_execution_state from invokeai.app.services.shared.graph import ( CollectInvocation, Graph, GraphExecutionState, IterateInvocation, + _ExecutionScheduler, _GenericGraphSchedulerAdapter, ) -from tests.test_nodes import AnyTypeTestInvocation, create_edge, create_loop_linkage +from tests.test_nodes import AnyTypeTestInvocation, FixedCollectionTestInvocation, create_edge, create_loop_linkage class _MappingTraversalCounter: @@ -134,6 +138,307 @@ def _run_trivial_loop( return clock() - started +def _run_if_free_prepare_workload( + count: int = 300, *, wrap_prepare: bool = False, force_compatibility_scheduler: bool = False +) -> tuple[float, int]: + graph = Graph() + graph.add_node(RangeInvocation(id="range", start=0, stop=count)) + graph.add_node(IterateInvocation(id="iterate")) + graph.add_edge(create_edge("range", "collection", "iterate", "collection")) + previous = "iterate" + previous_field = "item" + body_ids = [] + for index in range(30): + node_id = f"chain_{index}" + graph.add_node(AnyTypeTestInvocation(id=node_id)) + graph.add_edge(create_edge(previous, previous_field, node_id, "value")) + previous, previous_field = node_id, "value" + body_ids.append(node_id) + collect_ids = [] + for index in range(6): + node_id = f"collect_{index}" + graph.add_node(CollectInvocation(id=node_id)) + graph.add_edge(create_edge(previous, previous_field, node_id, "item")) + collect_ids.append(node_id) + + state = GraphExecutionState(graph=graph) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) + transaction_calls = 0 + run_transaction = state._run_materialization_transaction + + def counted_transaction(operation: Callable[[], object]) -> object: + nonlocal transaction_calls + transaction_calls += 1 + return run_transaction(operation) + + state._run_materialization_transaction = counted_transaction # type: ignore[method-assign] + if wrap_prepare: + materializer = state._materializer() + prepare = materializer.prepare + + def transactional_prepare(base_graph): + return counted_transaction(lambda: prepare(base_graph)) + + materializer.prepare = transactional_prepare # type: ignore[method-assign] + started = time.perf_counter() + trace: list[str] = [] + while (node := state.next()) is not None: + source_id = state.prepared_source_mapping[node.id] + trace.append(source_id) + state.complete(node.id, node.invoke(Mock())) + elapsed = time.perf_counter() - started + + assert state.is_complete() + assert not state.has_error() + assert all(trace.count(node_id) == count for node_id in body_ids) + assert all(trace.count(node_id) == 1 for node_id in collect_ids) + assert all( + state.results[next(iter(state.source_prepared_mapping[node_id]))].collection == list(range(count)) + for node_id in collect_ids + ) + return elapsed, transaction_calls + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_if_free_prepare_does_not_open_materialization_transactions(force_compatibility_scheduler: bool) -> None: + _, transaction_calls = _run_if_free_prepare_workload(force_compatibility_scheduler=force_compatibility_scheduler) + + assert transaction_calls == 0 + + +@pytest.mark.slow +def test_if_free_prepare_workload_performance_sample() -> None: + measurements: dict[str, list[float]] = {"transaction-wrapped": [], "direct": []} + for label, wrap_prepare in (("transaction-wrapped", True), ("direct", False)): + _run_if_free_prepare_workload(wrap_prepare=wrap_prepare) + for _ in range(3): + gc.collect() + elapsed, transaction_calls = _run_if_free_prepare_workload(wrap_prepare=wrap_prepare) + assert transaction_calls == (42 if wrap_prepare else 0) + measurements[label].append(elapsed) + for label, samples in measurements.items(): + print(f"Range(300), Iterate, 30-node chain, 6 Collects {label}: {sorted(samples)[1]:.3f}s median") + + +def _root_if_collect_graph(groups: int, branch_size: int) -> Graph: + graph = Graph() + graph.add_node(RangeInvocation(id="outer_source", start=0, stop=groups)) + graph.add_node(IterateInvocation(id="outer")) + graph.add_node(FixedCollectionTestInvocation(id="inner_source")) + graph.add_node(IterateInvocation(id="inner")) + graph.add_node(BooleanInvocation(id="condition", value=True)) + graph.add_node(IfInvocation(id="if")) + graph.add_node(CollectInvocation(id="branch_collect")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value="fallback")) + graph.add_node(AnyTypeTestInvocation(id="after")) + graph.add_edge(create_edge("outer_source", "collection", "outer", "collection")) + graph.add_edge(create_edge("outer", "item", "inner_source", "value")) + graph.add_edge(create_edge("inner_source", "collection", "inner", "collection")) + graph.add_edge(create_edge("condition", "value", "if", "condition")) + graph.add_edge(create_edge("fallback", "value", "if", "false_input")) + previous, previous_field = "inner", "item" + body_ids = [] + for index in range(branch_size): + node_id = f"branch_{index}" + graph.add_node(AnyTypeTestInvocation(id=node_id)) + graph.add_edge(create_edge(previous, previous_field, node_id, "value")) + previous, previous_field = node_id, "value" + body_ids.append(node_id) + graph.add_edge(create_edge(previous, previous_field, "branch_collect", "item")) + graph.add_edge(create_edge("branch_collect", "collection", "if", "true_input")) + graph.add_edge(create_edge("if", "value", "after", "value")) + return graph + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_root_if_collect_activation_index_rebuilds_after_dump_load(force_compatibility_scheduler: bool) -> None: + state = GraphExecutionState(graph=_root_if_collect_graph(groups=4, branch_size=4)) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) + + trace: list[str] = [] + while len(trace) < 12: + node = state.next() + assert node is not None + trace.append(state.prepared_source_mapping[node.id]) + state.complete(node.id, node.invoke(Mock())) + assert state._if_activation_dependents_index_complete + + state = load_execution_state(dump_execution_state(state)) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) + while (node := state.next()) is not None: + trace.append(state.prepared_source_mapping[node.id]) + state.complete(node.id, node.invoke(Mock())) + + assert state.is_complete() + assert not state.has_error() + assert trace.count("if") == 4 + assert trace.count("branch_collect") == 4 + after_ids = sorted(state.source_prepared_mapping["after"], key=state._get_iteration_path) + assert [state.results[node_id].value for node_id in after_ids] == [[0, 1] for _ in range(4)] + + +def test_root_if_collect_activation_index_rebuilds_after_graph_mutation() -> None: + state = GraphExecutionState(graph=_root_if_collect_graph(groups=2, branch_size=2)) + state._get_activation_dependents_for_if("if", ()) + assert state._if_activation_dependents_index_complete + + state.add_node(AnyTypeTestInvocation(id="unconnected")) + + assert not state._if_activation_dependents_index_complete + assert not state._if_activation_active_dependents_by_dependency + while (node := state.next()) is not None: + state.complete(node.id, node.invoke(Mock())) + assert state.is_complete() + assert state.results[next(iter(state.source_prepared_mapping["unconnected"]))].value is None + + +def _run_root_if_collect_workload( + groups: int, branch_size: int, *, force_compatibility_scheduler: bool = False +) -> tuple[int, int, int, int, float, tuple[tuple[tuple[int, ...], int], ...], int, int]: + graph = _root_if_collect_graph(groups, branch_size) + body_ids = [f"branch_{index}" for index in range(branch_size)] + + state = GraphExecutionState(graph=graph) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) + scheduler_node_visits = 0 + scheduler_edge_visits = 0 + dependency_index_node_visits = 0 + dependency_index_edge_visits = 0 + discarded_candidates = 0 + candidate_counts: list[tuple[tuple[int, ...], int]] = [] + activation_resolution = False + original_topological_sort = nx.topological_sort + + def counted_topological_sort(prepared_graph): + nonlocal scheduler_node_visits, scheduler_edge_visits + if activation_resolution: + scheduler_node_visits += prepared_graph.number_of_nodes() + scheduler_edge_visits += prepared_graph.number_of_edges() + return original_topological_sort(prepared_graph) + + nx.topological_sort = counted_topological_sort + original_dependents = state._get_activation_dependents_for_if + + def counted_dependents(source_if_id: str, iteration_path: tuple[int, ...]): + nonlocal activation_resolution + activation_resolution = True + try: + candidates = original_dependents(source_if_id, iteration_path) + candidate_counts.append((iteration_path, len(candidates))) + return candidates + finally: + activation_resolution = False + + state._get_activation_dependents_for_if = counted_dependents # type: ignore[method-assign] + original_record_dependencies = state._record_activation_dependencies + + def counted_record_dependencies(exec_node_id: str): + nonlocal dependency_index_node_visits, dependency_index_edge_visits + dependencies = original_record_dependencies(exec_node_id) + if activation_resolution: + dependency_index_node_visits += 1 + dependency_index_edge_visits += len(dependencies) + return dependencies + + state._record_activation_dependencies = counted_record_dependencies # type: ignore[method-assign] + scheduler = state._scheduler() + discard_method_name = ( + "_discard_rejected_activation_nodes" + if isinstance(scheduler, _ExecutionScheduler) + else "_discard_rejected_nodes" + ) + discard_candidates = getattr(scheduler, discard_method_name) + + def counted_discard_candidates(exec_node_ids): + nonlocal discarded_candidates + exec_node_ids = tuple(exec_node_ids) + discarded_candidates += len(exec_node_ids) + return discard_candidates(exec_node_ids) + + setattr(scheduler, discard_method_name, counted_discard_candidates) + try: + trace: list[str] = [] + started = time.perf_counter() + while (node := state.next()) is not None: + source_id = state.prepared_source_mapping[node.id] + trace.append(source_id) + state.complete(node.id, node.invoke(Mock())) + elapsed = time.perf_counter() - started + finally: + nx.topological_sort = original_topological_sort + + assert state.is_complete() + assert not state.has_error() + assert trace.count("if") == groups + assert trace.count("branch_collect") == groups + assert all(trace.count(node_id) == groups * 2 for node_id in body_ids) + after_ids = sorted(state.source_prepared_mapping["after"], key=state._get_iteration_path) + assert len(after_ids) == groups + assert [state.results[node_id].value for node_id in after_ids] == [[0, 1] for _ in range(groups)] + prepared_nodes = len(state.prepared_source_mapping) + return ( + scheduler_node_visits, + scheduler_edge_visits, + discarded_candidates, + prepared_nodes, + elapsed, + tuple(candidate_counts), + dependency_index_node_visits, + dependency_index_edge_visits, + ) + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_root_if_collect_activation_planning_work_scales_with_prepared_graph( + force_compatibility_scheduler: bool, +) -> None: + small = _run_root_if_collect_workload( + groups=4, branch_size=4, force_compatibility_scheduler=force_compatibility_scheduler + ) + more_groups = _run_root_if_collect_workload( + groups=8, branch_size=4, force_compatibility_scheduler=force_compatibility_scheduler + ) + larger_branch = _run_root_if_collect_workload( + groups=4, branch_size=8, force_compatibility_scheduler=force_compatibility_scheduler + ) + + # Include the one-time dependency-index build plus any active candidate topo work. Together, + # these visits must track prepared nodes/edges rather than revisit each completed branch. + work = (small, more_groups, larger_branch) + assert all( + index_nodes + topo_nodes <= prepared_nodes * 2 + for topo_nodes, _, _, prepared_nodes, _, _, index_nodes, _ in work + ), f"{work=}" + assert all( + index_edges + topo_edges <= prepared_nodes * 3 + for _, topo_edges, _, prepared_nodes, _, _, _, index_edges in work + ), f"{work=}" + assert all(discarded <= prepared_nodes for _, _, discarded, prepared_nodes, _, _, _, _ in work), f"{work=}" + + +@pytest.mark.slow +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_root_if_collect_activation_planning_timing_sample(force_compatibility_scheduler: bool) -> None: + for groups, branch_size in ((8, 8), (16, 8), (8, 16)): + samples = [ + _run_root_if_collect_workload( + groups, + branch_size, + force_compatibility_scheduler=force_compatibility_scheduler, + )[4] + for _ in range(3) + ] + print( + f"root-If/Collect groups={groups} branch_size={branch_size} " + f"scheduler={'compatibility' if force_compatibility_scheduler else 'generic'} " + f"median={sorted(samples)[1]:.3f}s" + ) + + @pytest.mark.parametrize( ("loop_type", "include_if", "external_condition"), [ diff --git a/tests/app/services/shared/test_graph_if_runtime.py b/tests/app/services/shared/test_graph_if_runtime.py index 4a3880526ad..61c11d74434 100644 --- a/tests/app/services/shared/test_graph_if_runtime.py +++ b/tests/app/services/shared/test_graph_if_runtime.py @@ -61,6 +61,10 @@ def test_live_activation_frame_mismatch(state: GraphExecutionState, field: str, def test_unresolved_dependency_is_deferred_but_rejected_is_not(state: GraphExecutionState) -> None: dependency = ActivationDependency(owner_id="if", branch="false_input", frame=(2,)) state._if_activation_dependencies_by_exec["consumer"] = (dependency,) + state._index_activation_dependencies("consumer", (dependency,)) + assert state._if_activation_dependents_by_dependency[dependency] == ("consumer",) + assert "consumer" not in state._if_activation_active_dependents_by_dependency.get(dependency, {}) + assert state._is_deferred_by_unresolved_if("consumer") assert not state._has_rejected_activation_dependency("consumer") From e89f1af0e2dd549b985adc4be780ad1aa78aa1b0 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Thu, 15 Oct 2026 15:57:48 -0500 Subject: [PATCH 20/20] fix(graph): preserve completed source cache state Keep cached completion coherent across prepared execution registration and source discard.\nDocument the cache lifecycle and cover scheduler bounds and rollback state. --- invokeai/app/services/shared/README.md | 8 +- invokeai/app/services/shared/graph.py | 12 +- .../test_execution_engine_differential.py | 68 ++++-- .../shared/test_for_scheduler_performance.py | 207 +++++++++++++++++- tests/test_node_graph.py | 25 --- 5 files changed, 261 insertions(+), 59 deletions(-) diff --git a/invokeai/app/services/shared/README.md b/invokeai/app/services/shared/README.md index cce28b06491..6a414bd67f8 100644 --- a/invokeai/app/services/shared/README.md +++ b/invokeai/app/services/shared/README.md @@ -734,8 +734,12 @@ Inactive branch-source retirement uses an event queue owned by one materializer. activation-dependent source IDs when the corresponding `If` resolves; an empty condition `Iterate` queues both branch source sets. Each preparation pass checks the queued sources in topological order rather than rescanning every unprepared source. The index is tied to source-graph identity, rebuilt from durable activation state after graph changes -or state restore, and is not serialized. Queue mutations participate in execution-state rollback. This narrows -retirement work; it does not bound other preparation or branch-context costs. +or state restore, and is not serialized. Queue mutations participate in execution-state rollback. A process-local +completed-source ID cache avoids rescanning the full `executed` set for every inactivity check; For/Iterate completion +invalidates it because axis completion can change branch activity. Registering a new prepared execution evicts its +source ID from the cached set; re-registering an existing execution does not. Graph edits or transaction rollback clear +the derived cache. This bounds source-inactivity checks for serial If chains, but does not bound other preparation or +branch-context costs. Literal-condition If fallback path inference uses the same ancestry context, including nested selected branches. diff --git a/invokeai/app/services/shared/graph.py b/invokeai/app/services/shared/graph.py index 53b456f188c..0e204659a3e 100644 --- a/invokeai/app/services/shared/graph.py +++ b/invokeai/app/services/shared/graph.py @@ -762,7 +762,10 @@ def _add_execution_graph_edges(self, edges: Iterable[Edge]) -> None: def _mark_source_executed(self, source_node_id: str) -> None: was_executed = source_node_id in self.executed self._tx_add_set(self.executed, source_node_id) - self.__pydantic_private__["_runtime_state"].completed_source_ids_cache = None + source_node = self.graph.nodes.get(source_node_id) + if not was_executed and isinstance(source_node, (ForInvocation, IterateInvocation)): + # Axis completion can change whether sources under an unresolved If are inactive. + self.__pydantic_private__["_runtime_state"].completed_source_ids_cache = None self._tx_add_set(self._get_completed_source_ids_cache(), source_node_id) if source_node_id not in self.executed_history: self._tx_append_list(self.executed_history, source_node_id) @@ -862,6 +865,7 @@ def _invalidate_source_graph_cache(self) -> None: runtime.if_activation_controller_instance = None runtime.execution_materializer = None runtime.execution_runtime = None + runtime.completed_source_ids_cache = None def _materializer(self) -> _ExecutionMaterializer: runtime = self.__pydantic_private__["_runtime_state"] @@ -2235,6 +2239,10 @@ def _register_prepared_exec_node(self, exec_node_id: str, source_node_id: str) - self._prepared_if_exec_frames_by_source.pop(source_node_id, None) is_new = exec_node_id not in self.source_prepared_mapping.get(source_node_id, ()) self._prepared_registry().register(exec_node_id, source_node_id) + if is_new and isinstance( + self.graph.nodes.get(source_node_id), (ForInvocation, IfInvocation, IterateInvocation) + ): + self.__pydantic_private__["_runtime_state"].completed_source_ids_cache = None if is_new and self._unexecuted_prepared_counts is not None and exec_node_id not in self.executed: counts = self._unexecuted_prepared_counts old_count = counts.get(source_node_id, 0) @@ -2244,7 +2252,7 @@ def _register_prepared_exec_node(self, exec_node_id: str, source_node_id: str) - ) counts[source_node_id] = old_count + 1 self._tx_discard_set(self.executed, source_node_id) - if self._completed_source_ids_cache is not None: + if is_new and self._completed_source_ids_cache is not None: self._tx_discard_set(self._completed_source_ids_cache, source_node_id) self._invalidate_loop_caches_for_source(source_node_id) if ( diff --git a/tests/app/services/shared/test_execution_engine_differential.py b/tests/app/services/shared/test_execution_engine_differential.py index 3deb1b14004..1e40a62a898 100644 --- a/tests/app/services/shared/test_execution_engine_differential.py +++ b/tests/app/services/shared/test_execution_engine_differential.py @@ -10141,36 +10141,60 @@ def test_zip_collect_root_condition_deferred_consumer_failure_rolls_back_and_ret assert node is not None state.complete(node.id, node.invoke(Mock())) - before_after_ids = set(state.source_prepared_mapping.get("after", ())) - before_after_edges = tuple( - edge - for edge in state.execution_graph.edges - if state.prepared_source_mapping.get(edge.destination.node_id) == "after" - ) - materializer = state._materializer() - initialize_execution_node = materializer._initialize_execution_node + root_if_exec_ids = [ + exec_id for exec_id in state.source_prepared_mapping["if"] if state._get_iteration_path(exec_id) == () + ] + assert len(root_if_exec_ids) == 1 + state._materializer().create_execution_node("after", [("if", root_if_exec_ids[0])], iteration_path=()) + assert len(state.source_prepared_mapping["after"]) == 1 - def fail_after_second_frame(exec_node_id: str, input_edges: Any = None) -> None: - initialize_execution_node(exec_node_id, input_edges) - if state.prepared_source_mapping.get(exec_node_id) == "after" and state._get_iteration_path(exec_node_id) == ( - 1, + materializer_type = graph_module._ExecutionMaterializer + initialize_execution_node = materializer_type._initialize_execution_node + run_materialization_transaction = state._run_materialization_transaction + failed_transaction_snapshot: dict[str, Any] | None = None + + def snapshot_state() -> dict[str, Any]: + return { + "nodes": { + node_id: copy.deepcopy(state.execution_graph.get_node(node_id)) + for node_id in state.execution_graph.nodes + }, + "edges": copy.deepcopy(tuple(state.execution_graph.edges)), + "prepared_source_mapping": copy.deepcopy(state.prepared_source_mapping), + "source_prepared_mapping": copy.deepcopy(state.source_prepared_mapping), + "prepared_iteration_paths": copy.deepcopy(state.prepared_iteration_paths), + "indegree": copy.deepcopy(state.indegree), + "execution_tokens": copy.deepcopy(state.execution_tokens), + "gates": copy.deepcopy(state._generic_runtime().gates), + } + + def capture_failed_transaction(operation: Any) -> Any: + nonlocal failed_transaction_snapshot + transaction_snapshot = snapshot_state() if state._apply_transaction is None else None + try: + return run_materialization_transaction(operation) + except RuntimeError as error: + if transaction_snapshot is not None and str(error) == "injected deferred consumer initialization failure": + failed_transaction_snapshot = transaction_snapshot + raise + + def fail_after_second_frame(materializer: Any, exec_node_id: str, input_edges: Any = None) -> None: + initialize_execution_node(materializer, exec_node_id, input_edges) + if ( + materializer._state is state + and state.prepared_source_mapping.get(exec_node_id) == "after" + and state._get_iteration_path(exec_node_id) == (1,) ): raise RuntimeError("injected deferred consumer initialization failure") - monkeypatch.setattr(materializer, "_initialize_execution_node", fail_after_second_frame) + monkeypatch.setattr(materializer_type, "_initialize_execution_node", fail_after_second_frame) + monkeypatch.setattr(state, "_run_materialization_transaction", capture_failed_transaction) with pytest.raises(RuntimeError, match="injected deferred consumer initialization failure"): while (node := state.next()) is not None: state.complete(node.id, node.invoke(Mock())) - assert state.source_prepared_mapping.get("after", set()) == before_after_ids - assert ( - tuple( - edge - for edge in state.execution_graph.edges - if state.prepared_source_mapping.get(edge.destination.node_id) == "after" - ) - == before_after_edges - ) + assert failed_transaction_snapshot is not None, "injected failure did not occur in a materialization transaction" + assert snapshot_state() == failed_transaction_snapshot monkeypatch.undo() trace, state = _run_graph(state, force_compatibility_scheduler=force_compatibility_scheduler) diff --git a/tests/app/services/shared/test_for_scheduler_performance.py b/tests/app/services/shared/test_for_scheduler_performance.py index a024c6e7145..b596ceeccac 100644 --- a/tests/app/services/shared/test_for_scheduler_performance.py +++ b/tests/app/services/shared/test_for_scheduler_performance.py @@ -80,6 +80,20 @@ def get(self, key: str, default: object = None) -> object: # type: ignore[overr return super().get(key, default) +class _CountingSet(set[str]): + """Count visits to prepared execution IDs without constraining set operations.""" + + def __init__(self, values: set[str], counter: _MappingTraversalCounter, mapping_name: str) -> None: + super().__init__(values) + self._counter = counter + self._mapping_name = mapping_name + + def __iter__(self): # type: ignore[override] + for value in super().__iter__(): + self._counter.record(self._mapping_name) + yield value + + def _run_trivial_loop( loop_type: str, count: int, @@ -490,6 +504,22 @@ def test_loop_scheduler_mapping_visits_scale_linearly( ) +def test_prepared_execution_count_scans_each_source_once() -> None: + counter = _MappingTraversalCounter() + state = GraphExecutionState(graph=Graph()) + state.source_prepared_mapping = _CountingMapping( + { + "source": _CountingSet({"exec-1", "exec-2", "exec-3"}, counter, "prepared_node_ids"), + }, + counter, + "source_prepared_mapping", + ) # type: ignore[assignment] + + assert state._count_unexecuted_prepared("source") == 3 + assert state._count_unexecuted_prepared("source") == 3 + assert counter.visits_by_mapping == {"source_prepared_mapping": 1, "prepared_node_ids": 3} + + @pytest.mark.slow @pytest.mark.parametrize("loop_type", ["for", "iterate"]) def test_trivial_loop_scheduler_overhead(loop_type: str) -> None: @@ -509,7 +539,9 @@ def __iter__(self): yield value -def _measure_pending_if_attachment(count: int) -> tuple[int, int, int, float, float]: +def _measure_pending_if_attachment( + count: int, *, force_compatibility_scheduler: bool = False +) -> tuple[int, int, int, float, float]: graph = Graph() graph.add_node(RangeInvocation(id="range", start=0, stop=count)) graph.add_node(IterateInvocation(id="loop")) @@ -526,6 +558,8 @@ def _measure_pending_if_attachment(count: int) -> tuple[int, int, int, float, fl graph.add_edge(create_edge("false_body", "value", "if", "false_input")) graph.add_edge(create_edge("if", "value", "collect", "item")) state = GraphExecutionState(graph=graph) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) state._pending_if_exec_nodes = _PendingIfScanCounter(state._pending_if_exec_nodes) materializer = state._materializer() original_attach = materializer._attach_pending_if_inputs @@ -580,10 +614,23 @@ def measured_attach(*, enqueue: bool = True) -> None: ) +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_pending_if_attachment_candidate_visits_are_bounded_for_ci(force_compatibility_scheduler: bool) -> None: + """Keep a deterministic guard for the per-item If attachment sweep in the default CI lane.""" + count = 300 + attach_calls, _, candidates_scanned, _, _ = _measure_pending_if_attachment( + count, force_compatibility_scheduler=force_compatibility_scheduler + ) + assert candidates_scanned <= count * 32, ( + f"pending-If attachment revisited {candidates_scanned} prepared candidates " + f"for {count} items across {attach_calls} sweeps" + ) + + @pytest.mark.slow def test_pending_if_attachment_sweep_scaling() -> None: - """Measure the reviewed per-item If graph, separating attachment sweeps from total execution.""" - for count in (300, 1200): + """Measure the reviewed per-item If graph at a larger size, separating attachment from total execution.""" + for count in (1200,): attach_calls, scanned, candidates_scanned, total_seconds, attach_seconds = _measure_pending_if_attachment(count) assert candidates_scanned <= count * 32, ( f"pending-If attachment revisited {candidates_scanned} prepared candidates " @@ -730,15 +777,159 @@ def counted_candidates(source_graph: object) -> tuple[str, ...]: return candidate_visits, prepare_scans, len(trace), trace -def test_if_chain_retirement_candidate_work_scales_linearly() -> None: - measurements = {length: _measure_if_chain_retirement_candidates(length) for length in (150, 300)} +def _measure_if_chain_end_to_end_source_inactivity( + length: int, *, force_compatibility_scheduler: bool = False +) -> tuple[int, int, list[str]]: + """Count source-inactivity checks while running a serial If branch through the execution interface.""" + graph = Graph() + graph.add_node(BooleanInvocation(id="condition", value=True)) + graph.add_node(IfInvocation(id="if")) + graph.add_node(AnyTypeTestInvocation(id="fallback", value="fallback")) + graph.add_node(AnyTypeTestInvocation(id="sink")) + previous_node_id: str | None = None + for index in range(length): + node_id = f"branch_{index:03d}" + graph.add_node(AnyTypeTestInvocation(id=node_id, value=True if index == 0 else None)) + if previous_node_id is not None: + graph.add_edge(create_edge(previous_node_id, "value", node_id, "value")) + previous_node_id = node_id + assert previous_node_id is not None + for source, source_field, destination, destination_field in ( + ("condition", "value", "if", "condition"), + ("fallback", "value", "if", "false_input"), + (previous_node_id, "value", "if", "true_input"), + ("if", "value", "sink", "value"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + + state = GraphExecutionState(graph=graph) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) + inactive_checks = 0 + completed_cache_builds = 0 + original_is_source_inactive = state._is_source_inactive + original_get_completed_source_ids_cache = state._get_completed_source_ids_cache + + def counted_is_source_inactive(source_node_id: str, iteration_path: tuple[int, ...] = ()) -> bool: + nonlocal inactive_checks + inactive_checks += 1 + return original_is_source_inactive(source_node_id, iteration_path) + + def counted_get_completed_source_ids_cache() -> set[str]: + nonlocal completed_cache_builds + completed_cache_builds += 1 + return original_get_completed_source_ids_cache() + + state._is_source_inactive = counted_is_source_inactive # type: ignore[method-assign] + state._get_completed_source_ids_cache = counted_get_completed_source_ids_cache # type: ignore[method-assign] + trace: list[str] = [] + while (node := state.next()) is not None: + trace.append(state.prepared_source_mapping[node.id]) + state.complete(node.id, node.invoke(Mock())) + + assert state.is_complete() + assert not state.has_error() + assert "fallback" not in trace + assert trace.count("if") == 1 + assert trace.count("sink") == 1 + assert all(trace.count(f"branch_{index:03d}") == 1 for index in range(length)) + assert all(source_id in state.executed for source_id in (*trace, "fallback")) + assert len(state.executed_history) == len(state.graph.nodes) + assert set(state.executed_history) == set(state.graph.nodes) + sink_exec_id = next(iter(state.source_prepared_mapping["sink"])) + assert state.results[sink_exec_id].value is True + return inactive_checks, completed_cache_builds, trace + + +def test_if_chain_retirement_candidate_scan_is_bounded() -> None: + """Bound activation-candidate enumeration, not all end-to-end source-inactivity work.""" + measurements = {length: _measure_if_chain_retirement_candidates(length) for length in (40, 80)} for length, (candidate_visits, prepare_scans, _, _) in measurements.items(): assert candidate_visits <= prepare_scans * 2, ( f"if-chain length {length}: {candidate_visits} retirement candidates across {prepare_scans} prepare calls" ) - small_work = measurements[150][0] - large_work = measurements[300][0] + small_work = measurements[40][0] + large_work = measurements[80][0] assert large_work <= small_work * 2.25, ( - f"if-chain retirement candidate work grew faster than graph length: 150={small_work}, 300={large_work}" + f"if-chain retirement candidate work grew faster than graph length: 40={small_work}, 80={large_work}" ) + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +def test_if_chain_end_to_end_source_inactivity_work_scales_linearly(force_compatibility_scheduler: bool) -> None: + """Bound source-inactivity checks including completed-source cache rebuilds.""" + measurements = { + length: _measure_if_chain_end_to_end_source_inactivity( + length, force_compatibility_scheduler=force_compatibility_scheduler + ) + for length in (10, 20, 40, 80, 300) + } + for length, (inactive_checks, cache_calls, trace) in measurements.items(): + assert inactive_checks > 0 + assert cache_calls >= len(trace) + # L=10 and L=20 measured at most 26 and 46 checks: two checks per source node bounds both routes. + assert inactive_checks <= 2 * (length + 4), ( + f"if-chain length {length}: {inactive_checks} inactivity checks for {length + 4} source nodes" + ) + + +def test_completed_source_cache_only_evicts_for_new_prepared_executions() -> None: + graph = Graph() + graph.add_node(AnyTypeTestInvocation(id="source", value=True)) + state = GraphExecutionState(graph=graph) + state._register_prepared_exec_node("exec-1", "source") + state._mark_exec_node_executed("exec-1") + state._mark_source_executed("source") + completed = state._get_completed_source_ids_cache() + assert "source" in completed + + state._register_prepared_exec_node("exec-1", "source") + assert "source" in completed + + state._register_prepared_exec_node("exec-2", "source") + assert "source" not in completed + + +def test_completed_source_cache_keeps_prepared_completion_after_source_discard() -> None: + graph = Graph() + graph.add_node(AnyTypeTestInvocation(id="source", value=True)) + state = GraphExecutionState(graph=graph) + state._register_prepared_exec_node("exec-1", "source") + state._mark_exec_node_executed("exec-1") + state._mark_source_executed("source") + completed = state._get_completed_source_ids_cache() + assert "source" in completed + + state._discard_source_executed("source") + + assert "source" in completed + + +def test_completed_source_cache_tracks_completion_edits_rollback_and_restore() -> None: + graph = Graph() + graph.add_node(BooleanInvocation(id="source", value=True)) + state = GraphExecutionState(graph=graph) + + initial_completed = state._get_completed_source_ids_cache() + assert "source" not in initial_completed + state._mark_source_executed("source") + assert state._get_completed_source_ids_cache() is initial_completed + assert "source" in initial_completed + + state.add_node(BooleanInvocation(id="added", value=True)) + edited_completed = state._get_completed_source_ids_cache() + assert edited_completed is not initial_completed + assert "added" not in edited_completed + + def mark_then_fail() -> None: + state._mark_source_executed("added") + raise RuntimeError("rollback cache mutation") + + with pytest.raises(RuntimeError, match="rollback cache mutation"): + state._run_materialization_transaction(mark_then_fail) + assert "added" not in state.executed + assert "added" not in state._get_completed_source_ids_cache() + + restored = load_execution_state(dump_execution_state(state)) + assert "source" in restored._get_completed_source_ids_cache() diff --git a/tests/test_node_graph.py b/tests/test_node_graph.py index c758e0c895a..75694b09b0e 100644 --- a/tests/test_node_graph.py +++ b/tests/test_node_graph.py @@ -1472,31 +1472,6 @@ def test_graph_validate_self_accepts_iterator_before_collectors_with_float_roots graph.validate_self() -def test_graph_update_node_renames_float_source_with_transiently_invalid_collector_order(): - integer_one = IntegerInvocation(id="i1", value=1) - float_source = FloatInvocation(id="f", value=2.5) - first_collector = CollectInvocation(id="c") - second_collector = CollectInvocation(id="d") - integer_two = IntegerInvocation(id="i2", value=2) - graph = Graph( - nodes={node.id: node for node in (integer_one, float_source, first_collector, second_collector, integer_two)} - ) - for edge in ( - create_edge("i1", "value", "c", "item"), - create_edge("f", "value", "c", "item"), - create_edge("c", "collection", "d", "collection"), - create_edge("i2", "value", "d", "item"), - ): - graph.add_edge(edge) - - graph.update_node("f", FloatInvocation(id="f2", value=2.5)) - - assert "f" not in graph.nodes - assert create_edge("f2", "value", "c", "item") in graph.edges - assert all(edge.source.node_id != "f" and edge.destination.node_id != "f" for edge in graph.edges) - graph.validate_self() - - def test_graph_add_collection_input_revalidates_iterator_consumers(): integer_source = IntegerInvocation(id="integer", value=1) float_collection = FloatCollectionInvocation(id="float_collection", collection=[2.5])