Performance enhancement: index iterator ancestry once per snapshot

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`

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")

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)

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)

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
 

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))
@@ -286,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)

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]