feat(Core): support partial graph execution

comfy_execution/graph.py

@@ -3,11 +3,12 @@ from typing import Type, Literal
 import nodes
 import asyncio
 import inspect
-from comfy_execution.graph_utils import is_link, ExecutionBlocker
+from comfy_execution.graph_utils import is_link, ExecutionBlocker, ExecutionFailureBlocker
 from comfy.comfy_types.node_typing import ComfyNodeABC, InputTypeDict, InputTypeOptions
 
 # NOTE: ExecutionBlocker code got moved to graph_utils.py to prevent torch being imported too soon during unit tests
 ExecutionBlocker = ExecutionBlocker
+ExecutionFailureBlocker = ExecutionFailureBlocker
 
 class DependencyCycleError(Exception):
     pass
@@ -202,14 +203,21 @@ class ExecutionList(TopologicalSort):
         self.staged_node_id = None
         self.execution_cache = {}
         self.execution_cache_listeners = {}
+        self.transient_cache = {}
+        self.failure_tainted_parents = set()
 
     def is_cached(self, node_id):
-        return self.output_cache.get_local(node_id) is not None
+        return node_id in self.transient_cache or self.output_cache.get_local(node_id) is not None
+
+    def _get_cache_value(self, node_id):
+        if node_id in self.transient_cache:
+            return self.transient_cache[node_id]
+        return self.output_cache.get_local(node_id)
 
     def cache_link(self, from_node_id, to_node_id, from_socket=None):
         if to_node_id not in self.execution_cache:
             self.execution_cache[to_node_id] = {}
-        value = self.output_cache.get_local(from_node_id)
+        value = self._get_cache_value(from_node_id)
         self.execution_cache[to_node_id][from_node_id] = value
         if from_node_id not in self.execution_cache_listeners:
             self.execution_cache_listeners[from_node_id] = set()
@@ -218,6 +226,8 @@ class ExecutionList(TopologicalSort):
             self.output_link_callback(value.outputs[from_socket])
 
     def get_cache(self, from_node_id, to_node_id):
+        if from_node_id in self.transient_cache:
+            return self.transient_cache[from_node_id]
         if to_node_id not in self.execution_cache:
             return None
         value = self.execution_cache[to_node_id].get(from_node_id)
@@ -227,7 +237,9 @@ class ExecutionList(TopologicalSort):
         self.output_cache.set_local(from_node_id, value)
         return value
 
-    def cache_update(self, node_id, value):
+    def cache_update(self, node_id, value, transient=False):
+        if transient:
+            self.transient_cache[node_id] = value
         if node_id in self.execution_cache_listeners:
             for to_node_id, from_socket in self.execution_cache_listeners[node_id]:
                 if to_node_id in self.execution_cache:
@@ -239,6 +251,25 @@ class ExecutionList(TopologicalSort):
         super().add_strong_link(from_node_id, from_socket, to_node_id)
         self.cache_link(from_node_id, to_node_id, from_socket)
 
+    def add_completion_link(self, from_node_id, to_node_id):
+        # Block to_node_id until from_node_id finishes, without consuming any of its output sockets.
+        if not self.is_cached(from_node_id):
+            self.add_node(from_node_id)
+            if to_node_id not in self.blocking[from_node_id]:
+                self.blocking[from_node_id][to_node_id] = {}
+                self.blockCount[to_node_id] += 1
+
+    def mark_failure_tainted(self, node_id):
+        # Taint all ephemeral ancestors of a failed or failure-blocked node so dynamically-expanded parents are
+        # never cached as reusable when part of their expansion did not complete.
+        parent_id = self.dynprompt.get_parent_node_id(node_id)
+        while parent_id is not None and parent_id not in self.failure_tainted_parents:
+            self.failure_tainted_parents.add(parent_id)
+            parent_id = self.dynprompt.get_parent_node_id(parent_id)
+
+    def is_failure_tainted(self, node_id):
+        return node_id in self.failure_tainted_parents
+
     async def stage_node_execution(self):
         assert self.staged_node_id is None
         if self.is_empty():
@@ -329,7 +360,9 @@ class ExecutionList(TopologicalSort):
         blocked_by = { node_id: {} for node_id in self.pendingNodes }
         for from_node_id in self.blocking:
             for to_node_id in self.blocking[from_node_id]:
-                if True in self.blocking[from_node_id][to_node_id].values():
+                # Strong links have a True socket entry; completion links have no socket entries at all.
+                sockets = self.blocking[from_node_id][to_node_id]
+                if len(sockets) == 0 or True in sockets.values():
                     blocked_by[to_node_id][from_node_id] = True
         to_remove = [node_id for node_id in blocked_by if len(blocked_by[node_id]) == 0]
         while len(to_remove) > 0:

comfy_execution/graph_utils.py

@@ -153,3 +153,9 @@ class ExecutionBlocker:
     """
     def __init__(self, message):
         self.message = message
+
+
+class ExecutionFailureBlocker(ExecutionBlocker):
+    def __init__(self, node_id):
+        super().__init__(None)
+        self.node_id = node_id

comfy_execution/jobs.py

@@ -217,10 +217,14 @@ def normalize_history_item(prompt_id: str, history_item: dict, include_outputs:
     outputs_count, preview_output = get_outputs_summary(outputs)
 
     execution_error = None
+    execution_errors = []
     execution_start_time = None
     execution_end_time = None
+    execution_success = None
     was_interrupted = False
+    execution_summary = {}
     if status_info:
+        execution_summary = status_info.get('execution_summary') or {}
         messages = status_info.get('messages', [])
         for entry in messages:
             if isinstance(entry, (list, tuple)) and len(entry) >= 2:
@@ -230,10 +234,22 @@ def normalize_history_item(prompt_id: str, history_item: dict, include_outputs:
                         execution_start_time = event_data.get('timestamp')
                     elif event_name in ('execution_success', 'execution_error', 'execution_interrupted'):
                         execution_end_time = event_data.get('timestamp')
-                        if event_name == 'execution_error':
+                        if event_name == 'execution_success':
+                            execution_success = event_data
+                        elif event_name == 'execution_error':
                             execution_error = event_data
                         elif event_name == 'execution_interrupted':
                             was_interrupted = True
+                    elif event_name == 'execution_node_error':
+                        execution_errors.append(event_data)
+
+    completion_status = execution_summary.get('completion_status')
+    if completion_status is None and execution_success is not None:
+        completion_status = execution_success.get('completion_status', 'success')
+    if completion_status is None and status_str == 'success':
+        completion_status = 'success'
+    has_errors = execution_summary.get('has_errors', bool(execution_errors)) if completion_status is not None else None
+    execution_error_count = execution_summary.get('execution_error_count', len(execution_errors)) if completion_status is not None else None
 
     if status_str == 'success':
         status = JobStatus.COMPLETED
@@ -250,6 +266,9 @@ def normalize_history_item(prompt_id: str, history_item: dict, include_outputs:
         'execution_start_time': execution_start_time,
         'execution_end_time': execution_end_time,
         'execution_error': execution_error,
+        'completion_status': completion_status,
+        'has_errors': has_errors,
+        'execution_error_count': execution_error_count,
         'outputs_count': outputs_count,
         'preview_output': preview_output,
         'workflow_id': workflow_id,
@@ -258,6 +277,7 @@ def normalize_history_item(prompt_id: str, history_item: dict, include_outputs:
     if include_outputs:
         job['outputs'] = normalize_outputs(outputs)
         job['execution_status'] = status_info
+        job['execution_errors'] = execution_errors
         job['workflow'] = {
             'prompt': prompt,
             'extra_data': extra_data,

comfy_execution/progress.py

@@ -17,6 +17,7 @@ class NodeState(Enum):
     Running = "running"
     Finished = "finished"
     Error = "error"
+    Blocked = "blocked"
 
 
 class NodeProgressState(TypedDict):
@@ -301,17 +302,25 @@ class ProgressRegistry:
                     node_id, value, max_value, entry, self.prompt_id, image
                 )
 
-    def finish_progress(self, node_id: str) -> None:
-        """Finish progress tracking for a node"""
+    def _finish_progress(self, node_id: str, state: NodeState) -> None:
         entry = self.ensure_entry(node_id)
-        entry["state"] = NodeState.Finished
+        entry["state"] = state
         entry["value"] = entry["max"]
 
-        # Notify all enabled handlers
         for handler in self.handlers.values():
             if handler.enabled:
                 handler.finish_handler(node_id, entry, self.prompt_id)
 
+    def finish_progress(self, node_id: str) -> None:
+        """Finish progress tracking for a node"""
+        self._finish_progress(node_id, NodeState.Finished)
+
+    def error_progress(self, node_id: str) -> None:
+        self._finish_progress(node_id, NodeState.Error)
+
+    def block_progress(self, node_id: str) -> None:
+        self._finish_progress(node_id, NodeState.Blocked)
+
     def reset_handlers(self) -> None:
         """Reset all handlers"""
         for handler in self.handlers.values():