diff --git a/invokeai/app/services/shared/README.md b/invokeai/app/services/shared/README.md index 44b876d1651..6a414bd67f8 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 @@ -347,6 +349,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 +446,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,11 +665,29 @@ 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. 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 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: @@ -669,6 +699,56 @@ 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. +`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 +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 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. 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. + +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` @@ -716,12 +796,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/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..0e204659a3e 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, @@ -223,9 +223,14 @@ class _GraphExecutionStateRuntime: "if_branch_sources_cache", "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", "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", @@ -233,6 +238,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", @@ -252,6 +258,7 @@ class _GraphExecutionStateRuntime: "legacy_execution_snapshot", "legacy_snapshot_loaded", "fresh_flat_if_activation", + "if_iteration_axes_validated", ) def __init__(self) -> None: @@ -265,9 +272,16 @@ 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_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, ...]] = {} 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 @@ -275,6 +289,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 @@ -294,6 +309,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): @@ -400,14 +416,24 @@ 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 + } + 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() 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 + } + runtime_state.selected_if_frames_for_preparation = runtime_state.selected_if_frames_for_preparation.copy() + copied.__pydantic_private__["_runtime_state"] = runtime_state return copied def _tx_record_once(self, key: tuple[Any, ...], undo: Callable[[], None]) -> None: @@ -420,6 +446,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: @@ -515,9 +562,14 @@ 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_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 = {} 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 @@ -580,6 +632,16 @@ 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.""" + selected_frames = self._selected_if_frames_for_preparation + if source_node_id is None: + self._selected_if_frames_for_preparation = {} + return + 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: self._if_activation_controller_instance = _IfActivationController(self) @@ -614,6 +676,65 @@ 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), + ) + 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: + 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: dict[str, None] = {} + 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=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() + 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) @@ -639,11 +760,18 @@ 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) + 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) @@ -722,13 +850,22 @@ 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 = {} 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 + runtime.execution_materializer = None + runtime.execution_runtime = None + runtime.completed_source_ids_cache = None def _materializer(self) -> _ExecutionMaterializer: runtime = self.__pydantic_private__["_runtime_state"] @@ -1459,13 +1596,18 @@ 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), - ) + 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), + ) + self._tx_set_mapping(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) @@ -2089,6 +2231,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: @@ -2096,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) @@ -2105,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 ( @@ -3129,6 +3276,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: @@ -3148,6 +3296,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) @@ -3304,8 +3454,12 @@ 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) + 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) + scheduler._enqueue_activation_ready_nodes(activation_dependents) if enqueue: self._enqueue_if_ready(exec_node_id) @@ -3361,10 +3515,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 = 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 = planner.prepare(base_graph) + prepared_id = prepare_next() if next_node is None: next_node = self._get_next_node() @@ -3433,8 +3593,13 @@ 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_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 = {} 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 @@ -3800,6 +3965,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? @@ -4168,9 +4338,14 @@ 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_active_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", + "selected_if_frames_for_preparation", "prepared_exec_metadata", "prepared_exec_registry", "if_activation_controller_instance", diff --git a/invokeai/app/services/shared/graph_execution_runtime.py b/invokeai/app/services/shared/graph_execution_runtime.py index ec7a6fdbca2..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,6 +42,7 @@ class _ExecutionRuntime: def __init__(self, state: "GraphExecutionState") -> None: self._state = state + 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) @@ -51,17 +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.graph.nx_graph()) - 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 = list(nx.topological_sort(iterator_graph)) - topo_index = {node_id: i for i, node_id in enumerate(topo)} - 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_activation.py b/invokeai/app/services/shared/graph_if_activation.py index 80220ee7dda..e9916db102a 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,9 @@ 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, ...]], ...]] = {} + 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) @@ -46,11 +49,51 @@ 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 _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) + if cached is not None: + return cached + + source_graph = self._state._get_source_graph_flat() + 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 + 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 +102,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 +127,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_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 c2ea5d9fa9a..ecee72c78ef 100644 --- a/invokeai/app/services/shared/graph_if_runtime.py +++ b/invokeai/app/services/shared/graph_if_runtime.py @@ -5,73 +5,169 @@ 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 ( + 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 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]) + 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) + branch_source = state.graph.get_node(edge.source.node_id) + if ( + condition_iterators + and condition_iterators < branch_iterators + and branch_iterators != broadcast_axes + 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( state: "GraphExecutionState", edge: 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: _IteratorAncestry, + effective_paths_by_edge: Optional[dict[tuple[str, str], set[tuple[int, ...]]]] = 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, + 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 = 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: + 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, 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( + state, + current_source_id, + visited_source_ids, + get_prepared_nodes_for_source, + get_prepared_edge_iteration_path, + iterator_ancestry, + ) ) - ) + continue - 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, - ) - for input_edge in input_edges - ) - ) + input_edges = state.graph._get_input_edges(current_source_id) + if not input_edges: + cache[cache_key] = {()} + continue + pending.append((current_edge, True)) + pending.extend((input_edge, False) for input_edge in reversed(input_edges)) + + return cache[(edge.source.node_id, edge.source.field)] -def _get_if_iteration_paths( +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, ...]], - source_graph: Optional["nx.DiGraph"] = None, + iterator_ancestry: _IteratorAncestry, ) -> list[tuple[int, ...]]: - """Infer the selected branch's effective iteration paths for an If output.""" + """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} @@ -112,7 +208,8 @@ def _get_if_iteration_paths( next_visited_source_ids, get_prepared_nodes_for_source, get_prepared_edge_iteration_path, - source_graph, + iterator_ancestry, + effective_paths_by_edge, ) for edge in selected_edges ) @@ -128,18 +225,36 @@ def _get_if_iteration_paths( 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 [()] 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)) @@ -150,50 +265,144 @@ 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()): - 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 = () state._tx_set_mapping(state._if_activation_dependencies_by_exec, exec_node_id, dependencies) + state._index_activation_dependencies(exec_node_id, 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 not dependencies: + return dependencies + + 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) + 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 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( + state: "GraphExecutionState", if_node_id: str, branch_field: Optional[str] +) -> tuple[tuple[int, ...], ...]: + """Return only resolved If frames that selected a Collect-backed branch.""" + 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( 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) +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) - 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) + + 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 ) @@ -216,12 +425,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( @@ -314,6 +528,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 new file mode 100644 index 00000000000..3ecab6a1b2c --- /dev/null +++ b/invokeai/app/services/shared/graph_iterator_ancestry.py @@ -0,0 +1,110 @@ +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_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: + """Own the ordered open For/Iterate axes for one source-graph traversal.""" + + 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 = iterator_graph + self._has_iterator_ancestor: dict[str, bool] | None = None + 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: + 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( + has_iterator_ancestor[ancestor_id] + or isinstance(self._state.graph.get_node(ancestor_id), (ForInvocation, IterateInvocation)) + for ancestor_id in iterator_graph.predecessors(source_node_id) + ) + self._iterator_graph = iterator_graph + 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] + + 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 26a6c34c91c..0f6ca9cbece 100644 --- a/invokeai/app/services/shared/graph_materializer.py +++ b/invokeai/app/services/shared/graph_materializer.py @@ -14,6 +14,8 @@ 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 ( COLLECTION_FIELD, @@ -32,6 +34,10 @@ from invokeai.app.services.shared.graph import GraphExecutionState +class UnsupportedSharedIfCollectBranchesError(ValueError): + """Raised when one sibling Collect-backed Iterate must serve multiple If owners.""" + + class _ExecutionNodeBuilder: """Builds concrete execution-graph nodes for a selected runtime planner. @@ -42,7 +48,22 @@ 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, 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. + self._flat_iterator_graph: Optional[tuple["nx.DiGraph", "nx.DiGraph"]] = None + 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) @@ -51,11 +72,51 @@ 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])) + 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 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, @@ -85,14 +146,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))) @@ -167,6 +246,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): @@ -178,6 +350,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: @@ -309,7 +492,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( @@ -1064,8 +1250,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 @@ -1076,8 +1264,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) @@ -1129,10 +1331,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 @@ -1276,8 +1480,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 @@ -1313,13 +1520,64 @@ 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 + branch_path = self._state._get_iteration_path(prepared_iterators[index]) + 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()) + ) + 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_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: return iter(self._get_parent_iteration_mappings_without_iterators(next_node_id, input_edges)) @@ -1328,7 +1586,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( @@ -1341,6 +1608,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( @@ -1411,21 +1680,23 @@ def create_execution_node( return new_nodes - def _has_admitted_source_mapping(self, node_id: str, graph: "nx.DiGraph") -> bool: - if not any(isinstance(node, IfInvocation) for node in self._state.graph.nodes.values()): + def _has_admitted_source_mapping( + self, node_id: str, graph: "nx.DiGraph", iterator_ancestry: _IteratorAncestry + ) -> bool: + 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") 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, ) ) mappings = self._get_if_condition_iteration_mappings(node_id, graph) @@ -1460,7 +1731,15 @@ 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_axes = self._get_ordered_iterator_axes_for_edges(condition_edges, iterator_graph) + condition_iterator_ids = set(condition_axes) + if condition_iterator_ids: + iterator_nodes = condition_axes + 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: @@ -1471,19 +1750,22 @@ 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) - - # 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) + return self._get_parent_iteration_mappings( + node_id, graph, input_edges=condition_edges, iterator_node_ids=iterator_nodes ) - for branch_field in ("true_input", "false_input") - ): + + # 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([[]]) + 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 = { + axis_id for edge in branch_edges for axis_id in iterator_ancestry.get_iterator_axes_for_edge(edge) + } + 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([]) @@ -1509,15 +1791,124 @@ 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_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, 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 + + graph = self._state._get_source_graph_flat() + iterator_graph = self.iterator_graph(graph) + 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): + 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 = 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) + + 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 + 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 { + "true_input", + "false_input", + }: + continue + + condition_edges = self._state.graph._get_input_edges(destination.id, "condition") + condition_iterator_ids = set(self._get_ordered_iterator_axes_for_edges(condition_edges, iterator_graph)) + if not condition_iterator_ids.intersection(iterator_ids): + continue + + for iterate_node_id in iterator_ids - condition_iterator_ids: + if any( + 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 + branch = (destination.id, output_edge.destination.field) + existing_branch = branches.get(iterate_node_id) + 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 nodes: '{existing_branch[0]}' and " + f"'{branch[0]}.{branch[1]}'" + ) + 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. + 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)) + } + self._sibling_iterate_if_collect_branches = branches + return branches + + 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, @@ -1548,17 +1939,29 @@ 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 - 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 - ] - 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_ids.append(candidates[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), ) ) @@ -1585,13 +1988,33 @@ 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() - collectors = ( - n for n in self._state.graph.nodes if isinstance(self._state.graph.get_node(n), CollectInvocation) - ) + """Gets a DiGraph with collector edges removed for iterator ancestry searches. + + 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)) + 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) + if n in base + ] for c in collectors: + collection_sources = { + edge.source.node_id + for edge in self._state.graph._get_input_edges(c, COLLECTION_FIELD) + if base.has_edge(edge.source.node_id, c) + } 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 ( @@ -1603,12 +2026,11 @@ 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() - return [ - n - for n in nx.ancestors(g, node_id) - if isinstance(self._state.graph.get_node(n), (ForInvocation, IterateInvocation)) - ] + 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 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 { @@ -1703,9 +2125,261 @@ 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 _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): + 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 @@ -1717,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 @@ -1734,7 +2409,8 @@ 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 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)) @@ -1742,7 +2418,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, @@ -1762,13 +2442,13 @@ 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") if condition_edges: @@ -1797,13 +2477,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) @@ -1832,13 +2512,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. @@ -1846,7 +2553,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/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/invokeai/app/services/shared/graph_scheduler.py b/invokeai/app/services/shared/graph_scheduler.py index df937bb698a..f910a1e3e90 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,10 @@ 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) + 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() if enqueue: self._project_ready_nodes() @@ -443,10 +446,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 +540,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 +651,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 ce89287c311..1e40a62a898 100644 --- a/tests/app/services/shared/test_execution_engine_differential.py +++ b/tests/app/services/shared/test_execution_engine_differential.py @@ -1,6 +1,10 @@ from __future__ import annotations +import copy +import inspect import json +import random +import sys from itertools import product from pathlib import Path from typing import Any, Callable @@ -10,7 +14,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, @@ -21,6 +25,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, @@ -37,6 +42,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 ( @@ -1337,13 +1343,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: @@ -1352,6 +1362,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 @@ -2530,9 +2542,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) @@ -2549,13 +2572,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: @@ -6580,9 +6603,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 @@ -6654,7 +6678,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 @@ -6808,8 +6838,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 = { @@ -7141,6 +7173,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() @@ -7500,7 +7598,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 @@ -7548,7 +7652,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() @@ -7690,14 +7800,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() @@ -8057,7 +8173,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() @@ -8365,7 +8481,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 @@ -8489,3 +8611,2107 @@ 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() + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +@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( + force_compatibility_scheduler: bool, 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_graph( + GraphExecutionState(graph=graph), force_compatibility_scheduler=force_compatibility_scheduler + ) + + 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() + 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_graph(state) + + after_exec_ids = sorted(state.source_prepared_mapping["after"], key=state._get_iteration_path) + assert trace.count("add") == 2 + 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() + + +@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"], +) +@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() + 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"), + ("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)) + + state = GraphExecutionState(graph=graph) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(state) + trace, state = _run_graph(state) + + 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_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"]])) + 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) + + trace, state = _run_graph(state) + + assert state.is_complete() + assert state.results[next(iter(state.source_prepared_mapping["result_collect"]))].collection == [ + ["a", "b"], + ["c"], + ] + + +@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) + + trace, state = _run_graph(state) + + 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 + ] == [((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) + + trace, state = _run_graph(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"), + ): + 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}" + 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], + ] + + +@pytest.mark.parametrize("force_compatibility_scheduler", [False, True], ids=["generic", "compatibility"]) +@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() + 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"), + ): + graph.add_edge(create_edge(source, source_field, destination, destination_field)) + + state = GraphExecutionState(graph=graph) + 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=( + r"If node 'if' cannot select branch input 'true_input' from unrelated open iterator axes: " + r"condition \['condition_iterate'\], branch \['branch_iterate'\]" + ), + ): + state.next() + + +@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 nodes"): + _run_graph(state) + + +@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")) + 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) + original_recursion_limit = sys.getrecursionlimit() + try: + sys.setrecursionlimit(len(inspect.stack()) + 250) + trace, state = _run_graph(state) + finally: + sys.setrecursionlimit(original_recursion_limit) + + 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: + graph = Graph() + 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")) + 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) + 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")) + + trace, state = _run_graph(state) + + 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() + + +@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) + if force_compatibility_scheduler: + state._execution_scheduler = _ExecutionScheduler(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" + ) + + +@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"]) +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( + 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()) + + +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 + + +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())) + + 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 + + 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_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 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) + 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"), + [ + ([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: + """Deeper open branch axes are explicitly rejected before any node executes.""" + state = GraphExecutionState(graph=build_graph()) + 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 == [] + + +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"), + [ + (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 == [expected_value, expected_value] + assert trace.count("if") == 2 + + +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"] 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 d58b2c0f5a3..74b8f59cc55 100644 --- a/tests/app/services/shared/test_for_scheduler_performance.py +++ b/tests/app/services/shared/test_for_scheduler_performance.py @@ -1,28 +1,31 @@ -"""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 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: @@ -99,6 +102,7 @@ def _run_trivial_loop( mapping_counter: _MappingTraversalCounter | None = None, include_if: bool = False, external_condition: bool = False, + condition_value: bool = True, ) -> float: graph = Graph() graph.add_node(RangeInvocation(id="range", start=0, stop=count)) @@ -107,9 +111,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")) @@ -146,6 +152,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.platform_independent @pytest.mark.parametrize( ("loop_type", "include_if", "external_condition"), @@ -161,6 +468,7 @@ def _run_trivial_loop( 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): @@ -217,10 +525,125 @@ def test_prepared_execution_count_scans_each_source_once() -> None: @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, *, 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")) + 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) + 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 + 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.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 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 " + 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)) @@ -258,3 +681,256 @@ 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 _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[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: 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/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_activation.py b/tests/app/services/shared/test_graph_if_activation.py index 90827e48521..150bc985b93 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")], @@ -124,3 +126,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/app/services/shared/test_graph_if_runtime.py b/tests/app/services/shared/test_graph_if_runtime.py index 7a8c72abbe7..47ae5d400e1 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 @@ -61,6 +62,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") @@ -155,21 +160,81 @@ 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_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 = [ + 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) + assert ( + graph_if_runtime._get_if_iteration_paths_with_ancestry( + state, previous, set(), lambda _: set(), lambda *_: (), _IteratorAncestry(state) + ) + == [] + ) + + +def test_resolved_literal_if_uses_root_frame_path() -> None: + state = GraphExecutionState(graph=Graph(nodes={"if": IfInvocation(id="if", condition=True, true_input=42)})) + 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) -> 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") + + 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_iterator_ancestry.py b/tests/app/services/shared/test_graph_iterator_ancestry.py new file mode 100644 index 00000000000..f5fba803895 --- /dev/null +++ b/tests/app/services/shared/test_graph_iterator_ancestry.py @@ -0,0 +1,116 @@ +import random + +import networkx as nx +import pytest + +from invokeai.app.invocations.logic import IfInvocation +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") + + +def test_supplied_graph_view_keeps_excluded_edges_out_of_collect_ancestry() -> None: + nodes = { + "iterator": IterateInvocation(id="iterator"), + "collect": CollectInvocation(id="collect"), + "if": IfInvocation(id="if", condition=True), + } + edges = [ + Edge( + source=EdgeConnection(node_id="iterator", field="item"), + destination=EdgeConnection(node_id="collect", field="collection"), + ), + Edge( + source=EdgeConnection(node_id="collect", field="collection"), + destination=EdgeConnection(node_id="if", field="true_input"), + ), + ] + state = GraphExecutionState(graph=Graph()) + state.graph = Graph(nodes=nodes, edges=edges) + source_graph = state._get_source_graph_flat() + constrained_graph = nx.subgraph_view( + source_graph, filter_edge=lambda source, destination: (source, destination) != ("iterator", "collect") + ) + + assert not _IteratorAncestry(state, constrained_graph).has_iterator_ancestor("if") diff --git a/tests/app/services/shared/test_graph_materializer.py b/tests/app/services/shared/test_graph_materializer.py index f5c28118e86..886ac490a78 100644 --- a/tests/app/services/shared/test_graph_materializer.py +++ b/tests/app/services/shared/test_graph_materializer.py @@ -69,3 +69,46 @@ 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_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 + + 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-")) + 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 diff --git a/tests/test_graph_execution_state.py b/tests/test_graph_execution_state.py index c2fab0de9f8..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): @@ -4241,6 +4257,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=[])) @@ -5286,8 +5323,7 @@ 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"} def test_if_graph_lowers_legacy_branch_choice_to_frame_scoped_activation_token(): diff --git a/tests/test_node_graph.py b/tests/test_node_graph.py index 6efc2fcfcb8..712198babdf 100644 --- a/tests/test_node_graph.py +++ b/tests/test_node_graph.py @@ -1449,6 +1449,29 @@ 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_add_collection_input_revalidates_iterator_consumers(): integer_source = IntegerInvocation(id="integer", value=1) float_collection = FloatCollectionInvocation(id="float_collection", collection=[2.5])