Continue independent branches when a node fails during execution

2026-08-16 22:46:38 +08:00 · 2026-08-05 22:41:20 +03:00
parent 6f7cd7fcea
commit 676e686f1a
12 changed files with 872 additions and 38 deletions
--- a/execution.py
+++ b/execution.py
@@ -34,9 +34,13 @@ from comfy_execution.caching import (
    RAM_CACHE_LARGE_INTERMEDIATE,
 )
 from comfy_execution.graph import (
+    DependencyCycleError,
    DynamicPrompt,
    ExecutionBlocker,
+    ExecutionFailureBlocker,
    ExecutionList,
+    NodeInputError,
+    NodeNotFoundError,
    get_input_info,
 )
 from comfy_execution.graph_utils import GraphBuilder, is_link
@@ -53,6 +57,8 @@ class ExecutionResult(Enum):
    SUCCESS = 0
    FAILURE = 1
    PENDING = 2
+    BLOCKED = 3
+

 class DuplicateNodeError(Exception):
    pass
@@ -106,6 +112,47 @@ class CacheEntry(NamedTuple):
    outputs: list


+def _failure_blocker_state(value):
+    if isinstance(value, ExecutionFailureBlocker):
+        return True, False
+    if isinstance(value, dict):
+        values = value.values()
+    elif isinstance(value, (list, tuple)):
+        values = value
+    else:
+        return False, True
+
+    has_failure_blocker = False
+    has_normal_value = False
+    for child in values:
+        child_failure, child_normal = _failure_blocker_state(child)
+        has_failure_blocker = has_failure_blocker or child_failure
+        has_normal_value = has_normal_value or child_normal
+        if has_failure_blocker and has_normal_value:
+            break
+    return has_failure_blocker, has_normal_value
+
+
+def _tag_node_raised(ex):
+    # Marks an exception as raised by node code (vs execution machinery). Best effort: an
+    # exception class rejecting attribute assignment must not mask the original error.
+    try:
+        ex._node_raised = True
+    except Exception:
+        pass
+
+
+def _is_recoverable_node_failure(ex):
+    if isinstance(ex, (
+        comfy.model_management.InterruptProcessingException,
+        DependencyCycleError,
+        NodeInputError,
+        NodeNotFoundError,
+    )):
+        return False
+    return not comfy.model_management.is_oom(ex)
+
+
 class CacheType(Enum):
    CLASSIC = 0
    LRU = 1
@@ -257,16 +304,22 @@ async def _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, f
    async def process_inputs(inputs, index=None, input_is_list=False):
        if allow_interrupt:
            nodes.before_node_execution()
-        execution_block = None
+        # Prefer a runtime-failure blocker over a plain user blocker so failure
+        # taint and retry semantics are not masked by input ordering.
+        blocked_value = None
        for k, v in inputs.items():
-            if input_is_list:
-                for e in v:
-                    if isinstance(e, ExecutionBlocker):
-                        v = e
+            values = v if input_is_list else (v,)
+            for e in values:
+                if isinstance(e, ExecutionBlocker):
+                    if blocked_value is None or isinstance(e, ExecutionFailureBlocker):
+                        blocked_value = e
+                    if isinstance(blocked_value, ExecutionFailureBlocker):
                        break
-            if isinstance(v, ExecutionBlocker):
-                execution_block = execution_block_cb(v) if execution_block_cb else v
+            if isinstance(blocked_value, ExecutionFailureBlocker):
                break
+        execution_block = None
+        if blocked_value is not None:
+            execution_block = execution_block_cb(blocked_value) if execution_block_cb else blocked_value
        if execution_block is None:
            if pre_execute_cb is not None and index is not None:
                pre_execute_cb(index)
@@ -292,7 +345,11 @@ async def _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, f
            if inspect.iscoroutinefunction(f):
                async def async_wrapper(f, prompt_id, unique_id, list_index, args):
                    with CurrentNodeContext(prompt_id, unique_id, list_index):
-                        return await f(**args)
+                        try:
+                            return await f(**args)
+                        except Exception as ex:
+                            _tag_node_raised(ex)
+                            raise
                task = asyncio.create_task(async_wrapper(f, prompt_id, unique_id, index, args=inputs))
                # Give the task a chance to execute without yielding
                await asyncio.sleep(0)
@@ -303,7 +360,11 @@ async def _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, f
                    results.append(task)
            else:
                with CurrentNodeContext(prompt_id, unique_id, index):
-                    result = f(**inputs)
+                    try:
+                        result = f(**inputs)
+                    except Exception as ex:
+                        _tag_node_raised(ex)
+                        raise
                results.append(result)
        else:
            results.append(execution_block)
@@ -451,10 +512,14 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
        return (ExecutionResult.SUCCESS, None, None)

    input_data_all = None
+    failure_blocked_invocations = 0
+    successful_invocations = 0
+    resumed_subgraph = False
    try:
        if unique_id in pending_async_nodes:
+            pending_results, failure_blocked_invocations, successful_invocations = pending_async_nodes[unique_id]
            results = []
-            for r in pending_async_nodes[unique_id]:
+            for r in pending_results:
                if isinstance(r, asyncio.Task):
                    try:
                        results.append(r.result())
@@ -467,7 +532,8 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
            del pending_async_nodes[unique_id]
            output_data, output_ui, has_subgraph = get_output_from_returns(results, class_def)
        elif unique_id in pending_subgraph_results:
-            cached_results = pending_subgraph_results[unique_id]
+            cached_results, failure_blocked_invocations = pending_subgraph_results[unique_id]
+            resumed_subgraph = True
            resolved_outputs = []
            for is_subgraph, result in cached_results:
                if not is_subgraph:
@@ -520,6 +586,9 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
                    return (ExecutionResult.PENDING, None, None)

            def execution_block_cb(block):
+                nonlocal failure_blocked_invocations
+                if isinstance(block, ExecutionFailureBlocker):
+                    failure_blocked_invocations += 1
                if block.message is not None:
                    mes = {
                        "prompt_id": prompt_id,
@@ -538,6 +607,8 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
                else:
                    return block
            def pre_execute_cb(call_index):
+                nonlocal successful_invocations
+                successful_invocations += 1
                # TODO - How to handle this with async functions without contextvars (which requires Python 3.12)?
                GraphBuilder.set_default_prefix(unique_id, call_index, 0)

@@ -552,7 +623,7 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
                    comfy_aimdo.model_vbar.vbars_reset_watermark_limits()

            if has_pending_tasks:
-                pending_async_nodes[unique_id] = output_data
+                pending_async_nodes[unique_id] = (output_data, failure_blocked_invocations, successful_invocations)
                unblock = execution_list.add_external_block(unique_id)
                async def await_completion():
                    tasks = [x for x in output_data if isinstance(x, asyncio.Task)]
@@ -607,14 +678,22 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
            for node_id in new_output_ids:
                execution_list.add_node(node_id)
                execution_list.cache_link(node_id, unique_id)
+                # Wait for expanded output nodes before finishing, so a failure inside the expansion taints
+                # this node before its outputs can be cached as reusable.
+                execution_list.add_completion_link(node_id, unique_id)
            for link in new_output_links:
                execution_list.add_strong_link(link[0], link[1], unique_id)
-            pending_subgraph_results[unique_id] = cached_outputs
+            pending_subgraph_results[unique_id] = (cached_outputs, failure_blocked_invocations)
            return (ExecutionResult.PENDING, None, None)

        cache_entry = CacheEntry(ui=ui_outputs.get(unique_id), outputs=output_data)
-        execution_list.cache_update(unique_id, cache_entry)
-        await caches.outputs.set(unique_id, cache_entry)
+        has_failure_blocker, has_normal_output = _failure_blocker_state(output_data)
+        failure_tainted = has_failure_blocker or failure_blocked_invocations > 0 or execution_list.is_failure_tainted(unique_id)
+        if failure_tainted:
+            execution_list.mark_failure_tainted(unique_id)
+        execution_list.cache_update(unique_id, cache_entry, transient=failure_tainted)
+        if not failure_tainted:
+            await caches.outputs.set(unique_id, cache_entry)

    except comfy.model_management.InterruptProcessingException as iex:
        logging.info("Processing interrupted")
@@ -651,11 +730,23 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
            "exception_message": "{}\n{}".format(ex, tips),
            "exception_type": exception_type,
            "traceback": traceback.format_tb(tb),
-            "current_inputs": input_data_formatted
+            "current_inputs": input_data_formatted,
+            "node_raised": getattr(ex, "_node_raised", False),
        }

        return (ExecutionResult.FAILURE, error_details, ex)

+    if resumed_subgraph:
+        fully_failure_blocked = has_failure_blocker and not has_normal_output
+    else:
+        fully_failure_blocked = successful_invocations == 0 and (
+            failure_blocked_invocations > 0
+            or (has_failure_blocker and not has_normal_output)
+        )
+    if fully_failure_blocked:
+        get_progress_state().block_progress(unique_id)
+        return (ExecutionResult.BLOCKED, None, None)
+
    get_progress_state().finish_progress(unique_id)
    executed.add(unique_id)

@@ -673,6 +764,7 @@ class PromptExecutor:
        self.caches = CacheSet(cache_type=self.cache_type, cache_args=self.cache_args)
        self.status_messages = []
        self.success = True
+        self.execution_summary = None

    def add_message(self, event, data: dict, broadcast: bool):
        data = {
@@ -711,6 +803,21 @@ class PromptExecutor:
            }
            self.add_message("execution_error", mes, broadcast=False)

+    def handle_node_execution_error(self, prompt_id, prompt, current_outputs, executed, error):
+        node_id = error["node_id"]
+        mes = {
+            "prompt_id": prompt_id,
+            "node_id": node_id,
+            "node_type": prompt[node_id]["class_type"],
+            "executed": list(executed),
+            "exception_message": error["exception_message"],
+            "exception_type": error["exception_type"],
+            "traceback": error["traceback"],
+            "current_inputs": error["current_inputs"],
+            "current_outputs": list(current_outputs),
+        }
+        self.add_message("execution_node_error", mes, broadcast=False)
+
    def _notify_prompt_lifecycle(self, event: str, prompt_id: str):
        if not _has_cache_providers():
            return
@@ -732,6 +839,8 @@ class PromptExecutor:

        nodes.interrupt_processing(False)
        self.prompt_model_tracker.start()
+        self.success = True
+        self.execution_summary = None

        if "client_id" in extra_data:
            self.server.client_id = extra_data["client_id"]
@@ -776,24 +885,56 @@ class PromptExecutor:
                executed = set()
                execution_list = ExecutionList(dynamic_prompt, self.caches.outputs, self.prompt_model_tracker.add)
                current_outputs = self.caches.outputs.all_node_ids()
+                output_targets = set(execute_outputs)
+                failed_node_ids = set()
+                blocked_node_ids = set()
+                blocked_output_node_ids = set()
+                successful_output_node_ids = set()
+                node_failures = []
                for node_id in list(execute_outputs):
                    execution_list.add_node(node_id)

                while not execution_list.is_empty():
                    node_id, error, ex = await execution_list.stage_node_execution()
                    if error is not None:
+                        self.success = False
                        self.handle_execution_error(prompt_id, dynamic_prompt.original_prompt, current_outputs, executed, error, ex)
                        break

                    assert node_id is not None, "Node ID should not be None at this point"
                    result, error, ex = await execute(self.server, dynamic_prompt, self.caches, node_id, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_node_outputs)
-                    self.success = result != ExecutionResult.FAILURE
                    if result == ExecutionResult.FAILURE:
-                        self.handle_execution_error(prompt_id, dynamic_prompt.original_prompt, current_outputs, executed, error, ex)
-                        break
+                        if error.get("node_raised") and _is_recoverable_node_failure(ex):
+                            self.handle_node_execution_error(prompt_id, dynamic_prompt.original_prompt, current_outputs, executed, error)
+                            real_node_id = error["node_id"]
+                            failed_node_ids.add(real_node_id)
+                            node_failures.append((error, ex))
+                            blocker = ExecutionFailureBlocker(real_node_id)
+                            class_type = dynamic_prompt.get_node(node_id)["class_type"]
+                            class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
+                            cache_entry = CacheEntry(
+                                ui=None,
+                                outputs=[[blocker] for _ in class_def.RETURN_TYPES],
+                            )
+                            execution_list.cache_update(node_id, cache_entry, transient=True)
+                            execution_list.mark_failure_tainted(node_id)
+                            get_progress_state().error_progress(node_id)
+                            execution_list.complete_node_execution()
+                        else:
+                            self.success = False
+                            self.handle_execution_error(prompt_id, dynamic_prompt.original_prompt, current_outputs, executed, error, ex)
+                            break
                    elif result == ExecutionResult.PENDING:
                        execution_list.unstage_node_execution()
+                    elif result == ExecutionResult.BLOCKED:
+                        real_node_id = dynamic_prompt.get_real_node_id(node_id)
+                        blocked_node_ids.add(real_node_id)
+                        if node_id in output_targets:
+                            blocked_output_node_ids.add(real_node_id)
+                        execution_list.complete_node_execution()
                    else: # result == ExecutionResult.SUCCESS:
+                        if node_id in output_targets:
+                            successful_output_node_ids.add(dynamic_prompt.get_real_node_id(node_id))
                        execution_list.complete_node_execution()

                    if self.cache_type == CacheType.RAM_PRESSURE:
@@ -821,7 +962,36 @@ class PromptExecutor:
                        if cached is not None:
                            display_node_id = dynamic_prompt.get_display_node_id(node_id)
                            _send_cached_ui(self.server, node_id, display_node_id, cached, prompt_id, ui_node_outputs)
-                    self.add_message("execution_success", { "prompt_id": prompt_id }, broadcast=False)
+
+                    if node_failures:
+                        self.execution_summary = {
+                            "has_errors": True,
+                            "execution_error_count": len(node_failures),
+                            "failed_node_ids": sorted(failed_node_ids)[:100],
+                            "blocked_node_ids": sorted(blocked_node_ids)[:100],
+                            "blocked_output_node_ids": sorted(blocked_output_node_ids)[:100],
+                            "successful_output_node_ids": sorted(successful_output_node_ids)[:100],
+                        }
+                        if successful_output_node_ids:
+                            self.execution_summary["completion_status"] = "partial_success"
+                            self.add_message(
+                                "execution_success",
+                                {"prompt_id": prompt_id, **self.execution_summary},
+                                broadcast=False,
+                            )
+                        else:
+                            self.success = False
+                            last_error, last_ex = node_failures[-1]
+                            self.handle_execution_error(
+                                prompt_id,
+                                dynamic_prompt.original_prompt,
+                                current_outputs,
+                                executed,
+                                last_error,
+                                last_ex,
+                            )
+                    else:
+                        self.add_message("execution_success", { "prompt_id": prompt_id }, broadcast=False)

                ui_outputs = {}
                meta_outputs = {}
@@ -1282,6 +1452,7 @@ class PromptQueue:
        status_str: Literal['success', 'error']
        completed: bool
        messages: List[str]
+        execution_summary: Optional[dict] = None

    def task_done(self, item_id, history_result,
                  status: Optional['PromptQueue.ExecutionStatus'], process_item=None):
@@ -1293,6 +1464,8 @@ class PromptQueue:
            status_dict: Optional[dict] = None
            if status is not None:
                status_dict = copy.deepcopy(status._asdict())
+                if status_dict.get("execution_summary") is None:
+                    del status_dict["execution_summary"]

            if process_item is not None:
                prompt = process_item(prompt)