Polish lazy If runtime diagnostics and idempotency

Author: JPPhoto

invokeai/app/services/shared/graph.py

@@ -253,7 +253,8 @@ def is_deferred_by_unresolved_if(self, exec_node_id: str) -> bool:
         return False
 
     def mark_exec_node_skipped(self, exec_node_id: str) -> None:
-        if self._state._get_prepared_exec_metadata(exec_node_id).state == "executed":
+        state = self._state._get_prepared_exec_metadata(exec_node_id).state
+        if state in ("executed", "skipped"):
             return
 
         self._state._remove_from_ready_queues(exec_node_id)
@@ -488,6 +489,9 @@ def get_iteration_node(
         prepared_iterator_nodes: list[str],
     ) -> Optional[str]:
         prepared_nodes = self._get_prepared_nodes_for_source(source_node_id)
+        if len(prepared_nodes) == 1 and not prepared_iterator_nodes:
+            return next(iter(prepared_nodes))
+
         parent_iterators = self._get_parent_iterator_exec_nodes(source_node_id, graph, prepared_iterator_nodes)
         if len(prepared_nodes) == 1:
             prepared_node_id = next(iter(prepared_nodes))
@@ -793,8 +797,10 @@ def _prepare_if_inputs(self, node: IfInvocation, input_edges: list[Edge]) -> Non
 
             found_value, copied_value = self._try_get_copied_result_value(edge)
             if not found_value:
+                iteration_path = self._state._get_iteration_path(node.id)
                 raise RuntimeError(
-                    "IfInvocation selected input edge points at an exec node with no stored result output"
+                    "IfInvocation selected input edge points at an exec node with no stored result output: "
+                    f"if_exec_id={node.id}, source_exec_id={edge.source.node_id}, iteration_path={iteration_path}"
                 )
 
             setattr(node, edge.destination.field, copied_value)

tests/test_graph_execution_state.py

@@ -23,6 +23,7 @@
 # This import must happen before other invoke imports or test in other files(!!) break
 from tests.test_nodes import (
     AnyTypeTestInvocation,
+    AnyTypeTestInvocationOutput,
     PromptCollectionTestInvocation,
     PromptTestInvocation,
     TextToImageTestInvocation,
@@ -847,9 +848,14 @@ def test_prepare_if_inputs_raises_when_selected_branch_source_has_no_result():
     g._resolved_if_exec_branches[if_exec_id] = "true_input"
 
     if_node = g.execution_graph.get_node(if_exec_id)
-    with pytest.raises(RuntimeError, match="selected input edge"):
+    with pytest.raises(RuntimeError) as exc_info:
         g._prepare_inputs(if_node)
 
+    message = str(exc_info.value)
+    assert if_exec_id in message
+    assert true_value_exec_id in message
+    assert "iteration_path=()" in message
+
 
 def test_get_collect_iteration_mappings_ignores_skipped_prepared_exec_nodes():
     graph = Graph()
@@ -957,7 +963,7 @@ def test_mark_exec_node_skipped_does_not_hide_already_executed_results():
     g = GraphExecutionState(graph=graph)
 
     exec_id = g._create_execution_node("value", [])[0]
-    g.results[exec_id] = AnyTypeTestInvocation(id="result", value="value").invoke(Mock(InvocationContext))
+    g.results[exec_id] = AnyTypeTestInvocationOutput(value="value")
     g.executed.add(exec_id)
     g._set_prepared_exec_state(exec_id, "executed")
 
@@ -967,6 +973,21 @@ def test_mark_exec_node_skipped_does_not_hide_already_executed_results():
     assert g.results[exec_id].value == "value"
 
 
+def test_mark_exec_node_skipped_is_idempotent_for_skipped_state():
+    graph = Graph()
+    graph.add_node(AnyTypeTestInvocation(id="value", value="value"))
+
+    g = GraphExecutionState(graph=graph)
+
+    exec_id = g._create_execution_node("value", [])[0]
+
+    g._if_scheduler().mark_exec_node_skipped(exec_id)
+    g._if_scheduler().mark_exec_node_skipped(exec_id)
+
+    assert g._get_prepared_exec_metadata(exec_id).state == "skipped"
+    assert g.executed_history.count("value") == 1
+
+
 def test_are_connection_types_compatible_accepts_subclass_to_base():
     """A subclass output should be connectable to a base-class input.