From 67b8d23c2f1da768f31c950ad2c52c82dcdda98a Mon Sep 17 00:00:00 2001 From: bigcat88 Date: Mon, 13 Jul 2026 12:39:33 +0300 Subject: [PATCH] feat(Core): support partial graph execution --- comfy_api/feature_flags.py | 1 + comfy_execution/graph.py | 43 ++- comfy_execution/graph_utils.py | 6 + comfy_execution/jobs.py | 22 +- comfy_execution/progress.py | 17 +- execution.py | 232 ++++++++++-- main.py | 3 +- openapi.yaml | 7 + server.py | 17 + tests/execution/test_async_nodes.py | 35 ++ tests/execution/test_execution.py | 339 +++++++++++++++++- tests/execution/test_jobs.py | 70 ++++ .../testing_nodes/testing-pack/__init__.py | 3 + .../testing-pack/async_test_nodes.py | 156 ++++++++ .../testing-pack/cache_provider_test_nodes.py | 51 +++ 15 files changed, 967 insertions(+), 35 deletions(-) create mode 100644 tests/execution/testing_nodes/testing-pack/cache_provider_test_nodes.py diff --git a/comfy_api/feature_flags.py b/comfy_api/feature_flags.py index cb14a5be0..40d29cebe 100644 --- a/comfy_api/feature_flags.py +++ b/comfy_api/feature_flags.py @@ -100,6 +100,7 @@ def _parse_cli_feature_flags() -> dict[str, Any]: # Default server capabilities _CORE_FEATURE_FLAGS: dict[str, Any] = { "supports_preview_metadata": True, + "supports_node_failure_policy": True, "supports_model_type_tags": True, "max_upload_size": args.max_upload_size * 1024 * 1024, # Convert MB to bytes "extension": {"manager": {"supports_v4": True}}, diff --git a/comfy_execution/graph.py b/comfy_execution/graph.py index 479ee8a53..73041e00f 100644 --- a/comfy_execution/graph.py +++ b/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 @@ -201,19 +202,28 @@ 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): if to_node_id not in self.execution_cache: self.execution_cache[to_node_id] = {} - self.execution_cache[to_node_id][from_node_id] = self.output_cache.get_local(from_node_id) + self.execution_cache[to_node_id][from_node_id] = self._get_cache_value(from_node_id) if from_node_id not in self.execution_cache_listeners: self.execution_cache_listeners[from_node_id] = set() self.execution_cache_listeners[from_node_id].add(to_node_id) 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) @@ -223,7 +233,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 in self.execution_cache_listeners[node_id]: if to_node_id in self.execution_cache: @@ -233,6 +245,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) + 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(): @@ -323,7 +354,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: diff --git a/comfy_execution/graph_utils.py b/comfy_execution/graph_utils.py index 496d2c634..7b62a74c3 100644 --- a/comfy_execution/graph_utils.py +++ b/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 diff --git a/comfy_execution/jobs.py b/comfy_execution/jobs.py index f0ad59f86..160dbff13 100644 --- a/comfy_execution/jobs.py +++ b/comfy_execution/jobs.py @@ -204,10 +204,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: @@ -217,10 +221,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 @@ -237,6 +253,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, @@ -245,6 +264,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, diff --git a/comfy_execution/progress.py b/comfy_execution/progress.py index 731b8dc66..218f317e0 100644 --- a/comfy_execution/progress.py +++ b/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(): diff --git a/execution.py b/execution.py index 19b8cdd68..56f63198b 100644 --- a/execution.py +++ b/execution.py @@ -32,9 +32,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 @@ -51,6 +55,16 @@ class ExecutionResult(Enum): SUCCESS = 0 FAILURE = 1 PENDING = 2 + BLOCKED = 3 + + +NODE_FAILURE_POLICY_FAIL_FAST = "fail_fast" +NODE_FAILURE_POLICY_CONTINUE_INDEPENDENT = "continue_independent" +NODE_FAILURE_POLICIES = frozenset({ + NODE_FAILURE_POLICY_FAIL_FAST, + NODE_FAILURE_POLICY_CONTINUE_INDEPENDENT, +}) +NODE_FAILURE_POLICY_EXTRA_DATA_KEY = "_node_failure_policy" class DuplicateNodeError(Exception): pass @@ -104,6 +118,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 @@ -152,7 +207,7 @@ class CacheSet: } return result -SENSITIVE_EXTRA_DATA_KEYS = ("auth_token_comfy_org", "api_key_comfy_org") +SENSITIVE_EXTRA_DATA_KEYS = ("auth_token_comfy_org", "api_key_comfy_org", NODE_FAILURE_POLICY_EXTRA_DATA_KEY) def get_input_data(inputs, class_def, unique_id, execution_list=None, dynprompt=None, extra_data={}): is_v3 = issubclass(class_def, _ComfyNodeInternal) @@ -255,16 +310,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) @@ -290,7 +351,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) @@ -301,7 +366,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) @@ -448,11 +517,16 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed, execution_list.cache_update(unique_id, cached) return (ExecutionResult.SUCCESS, None, None) + continue_on_failure = extra_data.get(NODE_FAILURE_POLICY_EXTRA_DATA_KEY) == NODE_FAILURE_POLICY_CONTINUE_INDEPENDENT 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()) @@ -465,7 +539,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: @@ -518,6 +593,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, @@ -536,6 +614,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) @@ -550,7 +630,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)] @@ -605,14 +685,26 @@ 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) + if continue_on_failure: + # 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) + if continue_on_failure: + has_failure_blocker, has_normal_output = _failure_blocker_state(output_data) + else: + has_failure_blocker, has_normal_output = False, True + 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") @@ -649,11 +741,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) @@ -670,6 +774,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 = { @@ -708,6 +813,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 @@ -728,6 +848,8 @@ class PromptExecutor: set_preview_method(extra_data.get("preview_method")) nodes.interrupt_processing(False) + self.success = True + self.execution_summary = None if "client_id" in extra_data: self.server.client_id = extra_data["client_id"] @@ -772,24 +894,60 @@ class PromptExecutor: executed = set() execution_list = ExecutionList(dynamic_prompt, self.caches.outputs) 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 = [] + continue_independent = extra_data.get( + NODE_FAILURE_POLICY_EXTRA_DATA_KEY, + NODE_FAILURE_POLICY_FAIL_FAST, + ) == NODE_FAILURE_POLICY_CONTINUE_INDEPENDENT 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 continue_independent and 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: @@ -817,7 +975,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 = {} @@ -1275,6 +1462,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): @@ -1286,6 +1474,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) diff --git a/main.py b/main.py index 580074b19..17797032e 100644 --- a/main.py +++ b/main.py @@ -366,7 +366,8 @@ def prompt_worker(q, server_instance): status=execution.PromptQueue.ExecutionStatus( status_str='success' if e.success else 'error', completed=e.success, - messages=e.status_messages), process_item=remove_sensitive) + messages=e.status_messages, + execution_summary=e.execution_summary), process_item=remove_sensitive) if server_instance.client_id is not None: server_instance.send_sync("executing", {"node": None, "prompt_id": prompt_id}, server_instance.client_id) diff --git a/openapi.yaml b/openapi.yaml index e00643bad..0b897ca43 100644 --- a/openapi.yaml +++ b/openapi.yaml @@ -922,6 +922,13 @@ components: number: description: Priority number for the queue (lower numbers have higher priority) type: number + node_failure_policy: + default: fail_fast + description: Controls whether a runtime node failure terminates the prompt or only blocks dependent nodes + enum: + - fail_fast + - continue_independent + type: string partial_execution_targets: description: List of node names to execute items: diff --git a/server.py b/server.py index e28fe2d22..43459f880 100644 --- a/server.py +++ b/server.py @@ -1097,6 +1097,19 @@ class PromptServer(): if "partial_execution_targets" in json_data: partial_execution_targets = json_data["partial_execution_targets"] + node_failure_policy = json_data.get( + "node_failure_policy", + execution.NODE_FAILURE_POLICY_FAIL_FAST, + ) + if not isinstance(node_failure_policy, str) or node_failure_policy not in execution.NODE_FAILURE_POLICIES: + error = { + "type": "invalid_node_failure_policy", + "message": "node_failure_policy must be 'fail_fast' or 'continue_independent'", + "details": f"Invalid node_failure_policy: {node_failure_policy!r}", + "extra_info": {}, + } + return web.json_response({"error": error, "node_errors": {}}, status=400) + self.node_replace_manager.apply_replacements(prompt) valid = await execution.validate_prompt(prompt_id, prompt, partial_execution_targets) @@ -1104,6 +1117,10 @@ class PromptServer(): if "extra_data" in json_data: extra_data = json_data["extra_data"] + extra_data.pop(execution.NODE_FAILURE_POLICY_EXTRA_DATA_KEY, None) + if node_failure_policy == execution.NODE_FAILURE_POLICY_CONTINUE_INDEPENDENT: + extra_data[execution.NODE_FAILURE_POLICY_EXTRA_DATA_KEY] = node_failure_policy + if "client_id" in json_data: extra_data["client_id"] = json_data["client_id"] diff --git a/tests/execution/test_async_nodes.py b/tests/execution/test_async_nodes.py index 54660c112..178f4e360 100644 --- a/tests/execution/test_async_nodes.py +++ b/tests/execution/test_async_nodes.py @@ -281,6 +281,41 @@ class TestAsyncNodes: # Verify the sync error was caught even though async was running assert 'prompt_id' in e.args[0] + def test_async_sibling_completes_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("TestAsyncError", value=image.out(0), error_after=0.05) + sleep_node = g.node("TestSleep", value=image.out(0), seconds=0.1) + g.node("PreviewImage", images=error_node.out(0)) + successful_output = g.node("SaveImage", images=sleep_node.out(0)) + + result = client.run(g, node_failure_policy="continue_independent") + + assert result.did_run(error_node) + assert result.did_run(sleep_node) + assert result.was_executed(successful_output) + assert len(result.get_images(successful_output)) == 1 + assert result.execution_success['completion_status'] == 'partial_success' + assert result.node_errors[0]['node_id'] == error_node.id + + def test_async_sibling_completes_after_multiple_errors(self, client: ComfyClient, builder: GraphBuilder): + g = builder + image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1) + error1 = g.node("TestAsyncError", value=image.out(0), error_after=0.02) + error2 = g.node("TestAsyncError", value=image.out(0), error_after=0.04) + sleep_node = g.node("TestSleep", value=image.out(0), seconds=0.06) + g.node("PreviewImage", images=error1.out(0)) + g.node("PreviewImage", images=error2.out(0)) + successful_output = g.node("SaveImage", images=sleep_node.out(0)) + + result = client.run(g, node_failure_policy="continue_independent") + + assert {error['node_id'] for error in result.node_errors} == {error1.id, error2.id} + assert result.did_run(sleep_node) + assert result.was_executed(successful_output) + assert result.execution_success['completion_status'] == 'partial_success' + assert result.execution_success['execution_error_count'] == 2 + # Edge Cases def test_async_with_execution_blocker(self, client: ComfyClient, builder: GraphBuilder): diff --git a/tests/execution/test_execution.py b/tests/execution/test_execution.py index 15e2304fc..0ea76b27b 100644 --- a/tests/execution/test_execution.py +++ b/tests/execution/test_execution.py @@ -1,3 +1,4 @@ +from collections import Counter from io import BytesIO import numpy from PIL import Image @@ -13,8 +14,35 @@ import uuid import urllib.request import urllib.parse import urllib.error +from comfy_execution.graph import DynamicPrompt, ExecutionList from comfy_execution.graph_utils import GraphBuilder, Node + +def test_execution_list_transient_cache_is_prompt_scoped(): + class OutputCache: + def __init__(self): + self.values = {} + self.set_calls = [] + + def get_local(self, node_id): + return self.values.get(node_id) + + def set_local(self, node_id, value): + self.set_calls.append((node_id, value)) + self.values[node_id] = value + + output_cache = OutputCache() + execution_list = ExecutionList(DynamicPrompt({}), output_cache) + failure_entry = object() + + execution_list.cache_update("failed", failure_entry, transient=True) + execution_list.cache_link("failed", "late_consumer") + + assert execution_list.is_cached("failed") + assert execution_list.get_cache("failed", "late_consumer") is failure_entry + assert output_cache.set_calls == [] + assert not ExecutionList(DynamicPrompt({}), output_cache).is_cached("failed") + def run_warmup(client, prefix="warmup"): """Run a simple workflow to warm up the server.""" warmup_g = GraphBuilder(prefix=prefix) @@ -28,6 +56,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) @@ -66,10 +97,12 @@ class ComfyClient: ws.connect("ws://{}/ws?clientId={}".format(self.server_address, self.client_id)) self.ws = ws - def queue_prompt(self, prompt, partial_execution_targets=None): + def queue_prompt(self, prompt, partial_execution_targets=None, node_failure_policy=None): p = {"prompt": prompt, "client_id": self.client_id} if partial_execution_targets is not None: p["partial_execution_targets"] = partial_execution_targets + if node_failure_policy is not None: + p["node_failure_policy"] = node_failure_policy data = json.dumps(p).encode('utf-8') req = urllib.request.Request("http://{}/prompt".format(self.server_address), data=data) return json.loads(urllib.request.urlopen(req).read()) @@ -133,13 +166,13 @@ class ComfyClient: def set_test_name(self, name): self.test_name = name - def run(self, graph, partial_execution_targets=None): + def run(self, graph, partial_execution_targets=None, node_failure_policy=None): prompt = graph.finalize() for node in graph.nodes.values(): if node.class_type == 'SaveImage': node.inputs['filename_prefix'] = self.test_name - prompt_id = self.queue_prompt(prompt, partial_execution_targets)['prompt_id'] + prompt_id = self.queue_prompt(prompt, partial_execution_targets, node_failure_policy)['prompt_id'] result = RunResult(prompt_id) while True: out = self.ws.recv() @@ -152,8 +185,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 +345,299 @@ 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_continue_independent_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, node_failure_policy="continue_independent") + + 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'] + history = client.get_history(result.prompt_id)[result.prompt_id] + assert '_node_failure_policy' not in history['prompt'][3] + + retry = client.run(g, node_failure_policy="continue_independent") + assert retry.did_run(error_node), "Failed nodes must be retried on a new prompt" + assert len(retry.node_errors) == 1 + + def test_continue_independent_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, node_failure_policy="continue_independent") + + 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_continue_independent_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, node_failure_policy="continue_independent") + + 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, node_failure_policy="continue_independent") + assert retry.did_run(mixed), "Failure-tainted dynamic parents must not be reused from cache" + assert len(retry.node_errors) == 1 + + def test_continue_independent_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, node_failure_policy="continue_independent") + + assert exc_info.value.args[0]['node_id'] == oom_node.id + assert exc_info.value.args[0]['exception_type'] == 'torch.OutOfMemoryError' + + def test_continue_independent_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, node_failure_policy="continue_independent") + + 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_continue_independent_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, node_failure_policy="continue_independent") + + assert exc_info.value.args[0]['node_id'] == bad.id + assert exc_info.value.args[0]['exception_type'] == 'KeyError' + + def test_continue_independent_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, node_failure_policy="continue_independent") + + assert exc_info.value.args[0]['node_id'] == bad.id + assert exc_info.value.args[0]['exception_type'] == 'TypeError' + + def test_continue_independent_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, node_failure_policy="continue_independent") + + assert exc_info.value.args[0]['exception_type'] == 'graph.DependencyCycleError' + + def test_continue_independent_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, node_failure_policy="continue_independent") + + 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, node_failure_policy="continue_independent") + assert retry.did_run(async_output), "Async outputs with failure-blocked invocations must be retried" + + def test_continue_independent_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, node_failure_policy="continue_independent") + + 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, node_failure_policy="continue_independent") + assert retry.did_run(combo), "Nodes blocked by a failure must be retried even when also user-blocked" + + def test_continue_independent_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, node_failure_policy="continue_independent") + + 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, node_failure_policy="continue_independent") + assert retry.did_run(expander), "Expansion parents with failed side branches must be retried" + assert len(retry.node_errors) == 1 + + def test_continue_independent_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], + node_failure_policy="continue_independent", + ) + + 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_explicit_fail_fast_policy_matches_default(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)) + g.node("SaveImage", images=image.out(0)) + + with pytest.raises(Exception) as exc_info: + client.run(g, node_failure_policy="fail_fast") + + assert exc_info.value.args[0]['node_id'] == error_node.id + history = client.get_history(exc_info.value.args[0]['prompt_id']) + entry = next(iter(history.values())) + assert entry['status']['status_str'] == 'error' + assert 'execution_summary' not in entry['status'] + + def test_continue_independent_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, node_failure_policy="continue_independent") + + 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_continue_independent_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, node_failure_policy="continue_independent") + 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_continue_independent_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, node_failure_policy="continue_independent") + + assert exc_info.value.args[0]['node_id'] == error_node.id + + def test_invalid_node_failure_policy(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)) + + with pytest.raises(urllib.error.HTTPError) as exc_info: + client.queue_prompt(g.finalize(), node_failure_policy="continue_everything") + + assert exc_info.value.code == 400 + @pytest.mark.parametrize("test_value, expect_error", [ (5, True), ("foo", True), diff --git a/tests/execution/test_jobs.py b/tests/execution/test_jobs.py index f7cb612e4..f41617419 100644 --- a/tests/execution/test_jobs.py +++ b/tests/execution/test_jobs.py @@ -500,6 +500,76 @@ class TestNormalizeHistoryItem: 'extra_data': {'create_time': 1234567890, 'client_id': 'abc'}, } + def test_missing_status(self): + history_item = { + 'prompt': ( + 5, + 'prompt-without-status', + {'nodes': {}}, + {'create_time': 100}, + [], + ), + 'status': None, + 'outputs': {}, + } + + job = normalize_history_item('prompt-without-status', history_item) + + assert job['status'] == 'completed' + assert 'completion_status' not in job + + def test_partial_success_metadata_and_errors(self): + node_error = { + 'prompt_id': 'prompt-partial', + 'node_id': '2', + 'node_type': 'TestSyncError', + 'exception_message': 'failed', + 'exception_type': 'RuntimeError', + 'traceback': [], + 'current_inputs': {}, + 'current_outputs': [], + 'timestamp': 200, + } + history_item = { + 'prompt': ( + 5, + 'prompt-partial', + {'nodes': {}}, + {'create_time': 100}, + ['3', '4'], + ), + 'status': { + 'status_str': 'success', + 'completed': True, + 'execution_summary': { + 'completion_status': 'partial_success', + 'has_errors': True, + 'execution_error_count': 1, + }, + 'messages': [ + ('execution_start', {'prompt_id': 'prompt-partial', 'timestamp': 150}), + ('execution_node_error', node_error), + ('execution_success', { + 'prompt_id': 'prompt-partial', + 'completion_status': 'partial_success', + 'has_errors': True, + 'execution_error_count': 1, + 'timestamp': 300, + }), + ], + }, + 'outputs': {'4': {'images': [{'filename': 'survived.png'}]}}, + } + + job = normalize_history_item('prompt-partial', history_item, include_outputs=True) + + assert job['status'] == 'completed' + assert job['completion_status'] == 'partial_success' + assert job['has_errors'] is True + assert job['execution_error_count'] == 1 + assert job['execution_errors'] == [node_error] + assert job['outputs']['4']['images'] == [{'filename': 'survived.png'}] + def test_include_outputs_normalizes_3d_strings(self): """Detail view should transform string 3D filenames into file output dicts.""" history_item = { diff --git a/tests/execution/testing_nodes/testing-pack/__init__.py b/tests/execution/testing_nodes/testing-pack/__init__.py index 3d5ac8a94..f0254e9b6 100644 --- a/tests/execution/testing_nodes/testing-pack/__init__.py +++ b/tests/execution/testing_nodes/testing-pack/__init__.py @@ -5,6 +5,7 @@ from .conditions import CONDITION_NODE_CLASS_MAPPINGS, CONDITION_NODE_DISPLAY_NA from .stubs import TEST_STUB_NODE_CLASS_MAPPINGS, TEST_STUB_NODE_DISPLAY_NAME_MAPPINGS from .async_test_nodes import ASYNC_TEST_NODE_CLASS_MAPPINGS, ASYNC_TEST_NODE_DISPLAY_NAME_MAPPINGS from .api_test_nodes import API_TEST_NODE_CLASS_MAPPINGS, API_TEST_NODE_DISPLAY_NAME_MAPPINGS +from .cache_provider_test_nodes import CACHE_PROVIDER_TEST_NODE_CLASS_MAPPINGS, CACHE_PROVIDER_TEST_NODE_DISPLAY_NAME_MAPPINGS # NODE_CLASS_MAPPINGS = GENERAL_NODE_CLASS_MAPPINGS.update(COMPONENT_NODE_CLASS_MAPPINGS) # NODE_DISPLAY_NAME_MAPPINGS = GENERAL_NODE_DISPLAY_NAME_MAPPINGS.update(COMPONENT_NODE_DISPLAY_NAME_MAPPINGS) @@ -17,6 +18,7 @@ NODE_CLASS_MAPPINGS.update(CONDITION_NODE_CLASS_MAPPINGS) NODE_CLASS_MAPPINGS.update(TEST_STUB_NODE_CLASS_MAPPINGS) NODE_CLASS_MAPPINGS.update(ASYNC_TEST_NODE_CLASS_MAPPINGS) NODE_CLASS_MAPPINGS.update(API_TEST_NODE_CLASS_MAPPINGS) +NODE_CLASS_MAPPINGS.update(CACHE_PROVIDER_TEST_NODE_CLASS_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS = {} NODE_DISPLAY_NAME_MAPPINGS.update(TEST_NODE_DISPLAY_NAME_MAPPINGS) @@ -26,3 +28,4 @@ NODE_DISPLAY_NAME_MAPPINGS.update(CONDITION_NODE_DISPLAY_NAME_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(TEST_STUB_NODE_DISPLAY_NAME_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(ASYNC_TEST_NODE_DISPLAY_NAME_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(API_TEST_NODE_DISPLAY_NAME_MAPPINGS) +NODE_DISPLAY_NAME_MAPPINGS.update(CACHE_PROVIDER_TEST_NODE_DISPLAY_NAME_MAPPINGS) diff --git a/tests/execution/testing_nodes/testing-pack/async_test_nodes.py b/tests/execution/testing_nodes/testing-pack/async_test_nodes.py index 589dabf17..03cf2a9b8 100644 --- a/tests/execution/testing_nodes/testing-pack/async_test_nodes.py +++ b/tests/execution/testing_nodes/testing-pack/async_test_nodes.py @@ -135,6 +135,148 @@ class TestSyncError(ComfyNodeABC): raise RuntimeError("Intentional sync execution error for testing") +class TestOOMError(ComfyNodeABC): + @classmethod + def INPUT_TYPES(cls): + return {"required": {"value": (IO.ANY, {})}} + + RETURN_TYPES = (IO.ANY,) + FUNCTION = "oom_error" + CATEGORY = "experimental/async" + + def oom_error(self, value): + raise torch.OutOfMemoryError("Intentional out of memory error for testing") + + +class TestMixedExpansionFailure(ComfyNodeABC): + @classmethod + def INPUT_TYPES(cls): + return {"required": {"value": ("INT", {})}} + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "expand" + CATEGORY = "experimental/async" + + def expand(self, value): + image = torch.zeros([1, 32, 32, 3]) + if value == 0: + return (image,) + + graph = GraphBuilder() + error = graph.node("TestSyncError", value=image) + return { + "result": (error.out(0),), + "expand": graph.finalize(), + } + + +class TestMalformedExpansion(ComfyNodeABC): + """Expands to a graph referencing a missing node class, so the failure + happens in the executor after the node function has returned.""" + + @classmethod + def INPUT_TYPES(cls): + return {"required": {"value": (IO.ANY, {})}} + + RETURN_TYPES = (IO.ANY,) + FUNCTION = "expand" + CATEGORY = "experimental/async" + + def expand(self, value): + graph = GraphBuilder() + missing = graph.node("TestNodeClassThatDoesNotExist", value=value) + return { + "result": (missing.out(0),), + "expand": graph.finalize(), + } + + +class TestMalformedResult(ComfyNodeABC): + """Returns a non-tuple result so the failure happens while the executor + merges results, after the node function has returned.""" + + @classmethod + def INPUT_TYPES(cls): + return {"required": {"value": (IO.ANY, {})}} + + RETURN_TYPES = (IO.ANY,) + FUNCTION = "run" + CATEGORY = "experimental/async" + + def run(self, value): + return 5 + + +class TestCyclicExpansion(ComfyNodeABC): + """Expands to an output node that consumes this node's own pending output, + forming a cycle through the expansion completion link.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": {"value": (IO.ANY, {})}, + "hidden": {"unique_id": "UNIQUE_ID"}, + } + + RETURN_TYPES = (IO.ANY,) + FUNCTION = "expand" + CATEGORY = "experimental/async" + + def expand(self, value, unique_id): + graph = GraphBuilder() + graph.node("TestAsyncOutput", value=[unique_id, 0], seconds=0.0) + return { + "result": (value,), + "expand": graph.finalize(), + } + + +class TestExpansionWithFailingOutput(ComfyNodeABC): + """Expands to a subgraph whose result succeeds while a side branch ending + in an output node fails.""" + + @classmethod + def INPUT_TYPES(cls): + return {"required": {"image": (IO.IMAGE, {})}} + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "expand" + CATEGORY = "experimental/async" + + def expand(self, image): + graph = GraphBuilder() + error = graph.node("TestSyncError", value=image) + graph.node("PreviewImage", images=error.out(0)) + passthrough = graph.node("StubImage", content="WHITE", height=32, width=32, batch_size=1) + return { + "result": (passthrough.out(0),), + "expand": graph.finalize(), + } + + +class TestAsyncOutput(ComfyNodeABC): + """Async output node with no return sockets, used to test partial failure + handling across pending async invocations.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "value": (IO.ANY, {}), + "seconds": (IO.FLOAT, {"default": 0.1}), + }, + } + + RETURN_TYPES = () + OUTPUT_NODE = True + FUNCTION = "run" + CATEGORY = "experimental/async" + + async def run(self, value, seconds=0.1): + await asyncio.sleep(seconds) + return {"ui": {"values": [1]}} + + class TestAsyncLazyCheck(ComfyNodeABC): """Test node with async check_lazy_status.""" @@ -322,6 +464,13 @@ ASYNC_TEST_NODE_CLASS_MAPPINGS = { "TestAsyncValidationError": TestAsyncValidationError, "TestAsyncTimeout": TestAsyncTimeout, "TestSyncError": TestSyncError, + "TestOOMError": TestOOMError, + "TestMixedExpansionFailure": TestMixedExpansionFailure, + "TestMalformedExpansion": TestMalformedExpansion, + "TestMalformedResult": TestMalformedResult, + "TestCyclicExpansion": TestCyclicExpansion, + "TestExpansionWithFailingOutput": TestExpansionWithFailingOutput, + "TestAsyncOutput": TestAsyncOutput, "TestAsyncLazyCheck": TestAsyncLazyCheck, "TestDynamicAsyncGeneration": TestDynamicAsyncGeneration, "TestAsyncResourceUser": TestAsyncResourceUser, @@ -335,6 +484,13 @@ ASYNC_TEST_NODE_DISPLAY_NAME_MAPPINGS = { "TestAsyncValidationError": "Test Async Validation Error", "TestAsyncTimeout": "Test Async Timeout", "TestSyncError": "Test Sync Error", + "TestOOMError": "Test OOM Error", + "TestMixedExpansionFailure": "Test Mixed Expansion Failure", + "TestMalformedExpansion": "Test Malformed Expansion", + "TestMalformedResult": "Test Malformed Result", + "TestCyclicExpansion": "Test Cyclic Expansion", + "TestExpansionWithFailingOutput": "Test Expansion With Failing Output", + "TestAsyncOutput": "Test Async Output", "TestAsyncLazyCheck": "Test Async Lazy Check", "TestDynamicAsyncGeneration": "Test Dynamic Async Generation", "TestAsyncResourceUser": "Test Async Resource User", diff --git a/tests/execution/testing_nodes/testing-pack/cache_provider_test_nodes.py b/tests/execution/testing_nodes/testing-pack/cache_provider_test_nodes.py new file mode 100644 index 000000000..01f5293ca --- /dev/null +++ b/tests/execution/testing_nodes/testing-pack/cache_provider_test_nodes.py @@ -0,0 +1,51 @@ +from comfy.comfy_types.node_typing import ComfyNodeABC +from comfy_api.latest._caching import CacheProvider +from comfy_execution.cache_provider import register_cache_provider + + +class _RecordingCacheProvider(CacheProvider): + """Records the class types of every externally stored cache entry so tests + can assert that failed or failure-blocked outputs never leave the process.""" + + def __init__(self): + self.stored_class_types = [] + + async def on_lookup(self, context): + return None + + async def on_store(self, context, value): + self.stored_class_types.append(context.class_type) + + +RECORDING_CACHE_PROVIDER = _RecordingCacheProvider() +register_cache_provider(RECORDING_CACHE_PROVIDER) + + +class TestCacheProviderRecord(ComfyNodeABC): + """Reports which node class types have been stored through the external + cache provider interface since the server started.""" + + @classmethod + def INPUT_TYPES(cls): + return {"required": {}} + + @classmethod + def IS_CHANGED(cls): + return float("NaN") + + RETURN_TYPES = () + OUTPUT_NODE = True + FUNCTION = "report" + CATEGORY = "Testing/Nodes" + + def report(self): + return {"ui": {"stored_class_types": list(RECORDING_CACHE_PROVIDER.stored_class_types)}} + + +CACHE_PROVIDER_TEST_NODE_CLASS_MAPPINGS = { + "TestCacheProviderRecord": TestCacheProviderRecord, +} + +CACHE_PROVIDER_TEST_NODE_DISPLAY_NAME_MAPPINGS = { + "TestCacheProviderRecord": "Test Cache Provider Record", +}