Continue independent branches when a node fails during execution

tests/execution/test_execution.py

@@ -1,3 +1,4 @@
+from collections import Counter
 from io import BytesIO
 import numpy
 from PIL import Image
@@ -28,6 +29,9 @@ class RunResult:
         self.runs: Dict[str,bool] = {}
         self.cached: Dict[str,bool] = {}
         self.prompt_id: str = prompt_id
+        self.node_errors = []
+        self.execution_success = None
+        self.run_counts: Dict[str, int] = {}
 
     def get_output(self, node: Node):
         return self.outputs.get(node.id, None)
@@ -152,8 +156,15 @@ class ComfyClient:
                     if data['node'] is None:
                         break
                     result.runs[data['node']] = True
+                    result.run_counts[data['node']] = result.run_counts.get(data['node'], 0) + 1
                 elif message['type'] == 'execution_error':
                     raise Exception(message['data'])
+                elif message['type'] == 'execution_node_error':
+                    if message['data']['prompt_id'] == prompt_id:
+                        result.node_errors.append(message['data'])
+                elif message['type'] == 'execution_success':
+                    if message['data']['prompt_id'] == prompt_id:
+                        result.execution_success = message['data']
                 elif message['type'] == 'execution_cached':
                     if message['data']['prompt_id'] == prompt_id:
                         cached_nodes = message['data'].get('nodes', [])
@@ -305,6 +316,270 @@ class TestExecution:
         except Exception as e:
             assert 'prompt_id' in e.args[0], f"Did not get back a proper error message: {e}"
 
+    def test_partial_failure_after_error(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
+        error_node = g.node("TestSyncError", value=image.out(0))
+        blocked_output = g.node("PreviewImage", images=error_node.out(0))
+        successful_output = g.node("SaveImage", images=image.out(0))
+
+        result = client.run(g)
+
+        assert result.did_run(error_node)
+        assert result.was_executed(successful_output)
+        assert len(result.get_images(successful_output)) == 1
+        assert len(result.node_errors) == 1
+        assert result.node_errors[0]['node_id'] == error_node.id
+        assert result.execution_success['completion_status'] == 'partial_success'
+        assert result.execution_success['has_errors'] is True
+        assert result.execution_success['execution_error_count'] == 1
+        assert blocked_output.id in result.execution_success['blocked_output_node_ids']
+        assert successful_output.id in result.execution_success['successful_output_node_ids']
+
+        retry = client.run(g)
+        assert retry.did_run(error_node), "Failed nodes must be retried on a new prompt"
+        assert len(retry.node_errors) == 1
+
+    def test_partial_failure_reuses_failed_node_for_late_lazy_link(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
+        error_node = g.node("TestSyncError", value=image.out(0))
+        g.node("PreviewImage", images=error_node.out(0))
+
+        mask = g.node("StubMask", value=0.0, height=32, width=32, batch_size=1)
+        unused_image = g.node("StubImage", content="WHITE", height=32, width=32, batch_size=1)
+        lazy_mix = g.node("TestLazyMixImages", image1=error_node.out(0), image2=unused_image.out(0), mask=mask.out(0))
+        g.node("PreviewImage", images=lazy_mix.out(0))
+        successful_output = g.node("SaveImage", images=image.out(0))
+
+        result = client.run(g)
+
+        assert result.run_counts[error_node.id] == 1
+        assert lazy_mix.id in result.execution_success['blocked_node_ids']
+        assert successful_output.id in result.execution_success['successful_output_node_ids']
+        assert result.execution_success['completion_status'] == 'partial_success'
+
+    def test_partial_failure_handles_mixed_dynamic_results(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        zero = g.node("StubInt", value=0)
+        one = g.node("StubInt", value=1)
+        values = g.node("TestMakeListNode", value1=zero.out(0), value2=one.out(0))
+        mixed = g.node("TestMixedExpansionFailure", value=values.out(0))
+        output = g.node("PreviewImage", images=mixed.out(0))
+
+        result = client.run(g)
+
+        assert len(result.node_errors) == 1
+        assert result.node_errors[0]['node_id'] == mixed.id
+        assert len(result.get_images(output)) == 1
+        assert output.id in result.execution_success['successful_output_node_ids']
+        assert output.id not in result.execution_success['blocked_output_node_ids']
+        assert result.execution_success['completion_status'] == 'partial_success'
+
+        retry = client.run(g)
+        assert retry.did_run(mixed), "Failure-tainted dynamic parents must not be reused from cache"
+        assert len(retry.node_errors) == 1
+
+    def test_partial_failure_keeps_oom_terminal(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
+        oom_node = g.node("TestOOMError", value=image.out(0))
+        g.node("PreviewImage", images=oom_node.out(0))
+        g.node("SaveImage", images=image.out(0))
+
+        with pytest.raises(Exception) as exc_info:
+            client.run(g)
+
+        assert exc_info.value.args[0]['node_id'] == oom_node.id
+        assert exc_info.value.args[0]['exception_type'] == 'torch.OutOfMemoryError'
+
+    def test_partial_failure_accepts_cached_output(self, client: ComfyClient, builder: GraphBuilder, server):
+        g = builder
+        image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
+        error_node = g.node("TestSyncError", value=image.out(0))
+        g.node("PreviewImage", images=error_node.out(0))
+        successful_output = g.node("SaveImage", images=image.out(0))
+
+        client.run(g, partial_execution_targets=[successful_output.id])
+        result = client.run(g)
+
+        if server["should_cache_results"]:
+            assert result.was_cached(successful_output)
+        assert len(result.node_errors) == 1
+        assert result.execution_success['completion_status'] == 'partial_success'
+        assert successful_output.id in result.execution_success['successful_output_node_ids']
+
+    def test_partial_failure_keeps_executor_errors_terminal(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
+        bad = g.node("TestMalformedExpansion", value=image.out(0))
+        g.node("PreviewImage", images=bad.out(0))
+        g.node("SaveImage", images=image.out(0))
+
+        with pytest.raises(Exception) as exc_info:
+            client.run(g)
+
+        assert exc_info.value.args[0]['node_id'] == bad.id
+        assert exc_info.value.args[0]['exception_type'] == 'KeyError'
+
+    def test_partial_failure_malformed_result_terminal(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
+        bad = g.node("TestMalformedResult", value=image.out(0))
+        g.node("PreviewImage", images=bad.out(0))
+        g.node("SaveImage", images=image.out(0))
+
+        with pytest.raises(Exception) as exc_info:
+            client.run(g)
+
+        assert exc_info.value.args[0]['node_id'] == bad.id
+        assert exc_info.value.args[0]['exception_type'] == 'TypeError'
+
+    def test_partial_failure_cyclic_expansion_reports_cycle(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
+        cyclic = g.node("TestCyclicExpansion", value=image.out(0))
+        g.node("PreviewImage", images=cyclic.out(0))
+        g.node("SaveImage", images=image.out(0))
+
+        with pytest.raises(Exception) as exc_info:
+            client.run(g)
+
+        assert exc_info.value.args[0]['exception_type'] == 'graph.DependencyCycleError'
+
+    def test_partial_failure_async_output_partial_failure(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        zero = g.node("StubInt", value=0)
+        one = g.node("StubInt", value=1)
+        values = g.node("TestMakeListNode", value1=zero.out(0), value2=one.out(0))
+        mixed = g.node("TestMixedExpansionFailure", value=values.out(0))
+        async_output = g.node("TestAsyncOutput", value=mixed.out(0), seconds=0.1)
+
+        result = client.run(g)
+
+        assert len(result.node_errors) == 1
+        assert result.execution_success['completion_status'] == 'partial_success'
+        assert async_output.id in result.execution_success['successful_output_node_ids']
+        assert async_output.id not in result.execution_success['blocked_node_ids']
+
+        retry = client.run(g)
+        assert retry.did_run(async_output), "Async outputs with failure-blocked invocations must be retried"
+
+    def test_partial_failure_blocker_beats_user_blocker(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
+        error_node = g.node("TestSyncError", value=image.out(0))
+        user_blocker = g.node("TestExecutionBlocker", input=image.out(0), block=True, verbose=False)
+        combo = g.node("TestMakeListNode", value1=user_blocker.out(0), value2=error_node.out(0))
+        g.node("PreviewImage", images=combo.out(0))
+        successful_output = g.node("SaveImage", images=image.out(0))
+
+        result = client.run(g)
+
+        assert result.execution_success['completion_status'] == 'partial_success'
+        assert combo.id in result.execution_success['blocked_node_ids']
+        assert successful_output.id in result.execution_success['successful_output_node_ids']
+
+        retry = client.run(g)
+        assert retry.did_run(combo), "Nodes blocked by a failure must be retried even when also user-blocked"
+
+    def test_partial_failure_retries_failed_expansion_side_branch(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
+        expander = g.node("TestExpansionWithFailingOutput", image=image.out(0))
+        output = g.node("PreviewImage", images=expander.out(0))
+
+        result = client.run(g)
+
+        assert len(result.node_errors) == 1
+        assert result.node_errors[0]['node_id'] == expander.id
+        assert result.execution_success['completion_status'] == 'partial_success'
+        assert len(result.get_images(output)) == 1
+
+        retry = client.run(g)
+        assert retry.did_run(expander), "Expansion parents with failed side branches must be retried"
+        assert len(retry.node_errors) == 1
+
+    def test_partial_failure_with_partial_targets(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
+        error_node = g.node("TestSyncError", value=image.out(0))
+        failing_output = g.node("PreviewImage", images=error_node.out(0))
+        successful_output = g.node("SaveImage", images=image.out(0))
+        unselected_output = g.node("SaveImage", images=image.out(0))
+
+        result = client.run(
+            g,
+            partial_execution_targets=[failing_output.id, successful_output.id],
+        )
+
+        assert len(result.node_errors) == 1
+        assert result.execution_success['completion_status'] == 'partial_success'
+        assert successful_output.id in result.execution_success['successful_output_node_ids']
+        assert failing_output.id in result.execution_success['blocked_output_node_ids']
+        assert not result.was_executed(unselected_output), "Unselected outputs must not execute"
+
+    def test_partial_failure_after_sibling_output(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
+        fast_output = g.node("TestAsyncOutput", value=image.out(0), seconds=0.05)
+        slow_error = g.node("TestAsyncError", value=image.out(0), error_after=0.5)
+        g.node("PreviewImage", images=slow_error.out(0))
+
+        result = client.run(g)
+
+        assert len(result.node_errors) == 1
+        assert result.node_errors[0]['node_id'] == slow_error.id
+        assert result.execution_success['completion_status'] == 'partial_success'
+        assert fast_output.id in result.execution_success['successful_output_node_ids']
+
+    def test_partial_failure_never_stores_failed_outputs_externally(self, client: ComfyClient, builder: GraphBuilder, server):
+        def record_stored_counts(prefix):
+            record_graph = GraphBuilder(prefix=prefix)
+            record = record_graph.node("TestCacheProviderRecord")
+            record_result = client.run(record_graph)
+            return Counter(record_result.get_output(record)['stored_class_types'])
+
+        baseline = record_stored_counts("cache_baseline")
+
+        g = builder
+        # Unique 31x31 signature so StubImage executes fresh instead of hitting the cache
+        image = g.node("StubImage", content="BLACK", height=31, width=31, batch_size=1)
+        error_node = g.node("TestSyncError", value=image.out(0))
+        g.node("PreviewImage", images=error_node.out(0))
+        g.node("SaveImage", images=image.out(0))
+
+        result = client.run(g)
+        assert result.execution_success['completion_status'] == 'partial_success'
+
+        delta = record_stored_counts("cache_record") - baseline
+        if server["should_cache_results"]:
+            assert delta["StubImage"] >= 1, "Successful outputs should reach external cache providers"
+        assert delta["TestSyncError"] == 0, "Failed node outputs must never reach external cache providers"
+        assert delta["PreviewImage"] == 0, "Failure-blocked outputs must never reach external cache providers"
+
+    def test_history_status_omits_summary_by_default(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
+        g.node("PreviewImage", images=image.out(0))
+
+        result = client.run(g)
+
+        history = client.get_history(result.prompt_id)[result.prompt_id]
+        assert history['status']['status_str'] == 'success'
+        assert 'execution_summary' not in history['status']
+
+    def test_partial_failure_fails_when_no_output_survives(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
+        error_node = g.node("TestSyncError", value=image.out(0))
+        g.node("PreviewImage", images=error_node.out(0))
+
+        with pytest.raises(Exception) as exc_info:
+            client.run(g)
+
+        assert exc_info.value.args[0]['node_id'] == error_node.id
+
     @pytest.mark.parametrize("test_value, expect_error", [
         (5, True),
         ("foo", True),