diff --git a/api/image_list_routes.py b/api/image_list_routes.py index 8940c8b..92a303d 100644 --- a/api/image_list_routes.py +++ b/api/image_list_routes.py @@ -35,6 +35,7 @@ async def _json_body(request): async def image_list_capabilities_endpoint(_request): return web.json_response({ "supported": True, + "execution_targets": True, "protocol_version": PROTOCOL_VERSION, "node_ids": ["DistributedImageListStart", "DistributedImageListEnd"], "transport": "typed-float-tensor-json-v1", diff --git a/api/orchestration/dispatch.py b/api/orchestration/dispatch.py index 39a17b4..4213f31 100644 --- a/api/orchestration/dispatch.py +++ b/api/orchestration/dispatch.py @@ -77,6 +77,8 @@ async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0): "workflow": payload.get("workflow"), "client_id": client_id, } + if "partial_execution_targets" in payload: + ws_payload["partial_execution_targets"] = payload["partial_execution_targets"] ws_url = worker_url.replace("http://", "ws://").replace("https://", "wss://") ws_url = f"{ws_url}/distributed/worker_ws" session = await get_client_session() @@ -128,11 +130,17 @@ async def dispatch_worker_prompt( use_websocket=False, trace_execution_id=None, return_prompt_id=False, + *, + partial_execution_targets=None, ): """Send the prepared prompt to a worker ComfyUI instance.""" worker_url = build_worker_url(worker) url = build_worker_url(worker, "/prompt") payload = {"prompt": prompt_obj} + target_payload = {} + if partial_execution_targets is not None: + target_payload["partial_execution_targets"] = list(partial_execution_targets) + payload.update(target_payload) extra_data = {} if workflow_meta: extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta @@ -146,6 +154,7 @@ async def dispatch_worker_prompt( { "prompt": prompt_obj, "workflow": workflow_meta, + **target_payload, }, client_id, ) @@ -227,6 +236,12 @@ async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0): f"{sorted(required_nodes)}; no work was submitted" ) + if capability.get("execution_targets") is not True: + raise RuntimeError( + f"Image-list worker {worker_id} lacks explicit execution_targets support; " + "upgrade that worker before running this feature; no work was submitted" + ) + check = "queue" queue_url = build_worker_url(worker, "/queue") async with session.get(queue_url, timeout=aiohttp.ClientTimeout(total=timeout)) as response: diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index d10872a..709682f 100644 --- a/api/queue_orchestration.py +++ b/api/queue_orchestration.py @@ -38,6 +38,7 @@ from .orchestration.prompt_transform import ( _resolve_direct_workers_for_image_list, ) from ..image_list.runtime import validate_image_list_regions +from ..image_list.execution_plan import prepare_image_list_execution prompt_server = server.PromptServer.instance @@ -265,13 +266,17 @@ async def orchestrate_distributed_execution( media_sync_timeout_seconds, ) = _resolve_orchestration_limits(config) requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None + image_list_mode = any(isinstance(node, dict) and node.get("class_type") == "DistributedImageListStart" + for node in prompt_obj.values()) + target_kwargs = {} + if image_list_mode: + prompt_obj, outer_targets = prepare_image_list_execution(prompt_obj) + target_kwargs["partial_execution_targets"] = outer_targets prompt_index = PromptIndex(prompt_obj) - image_list_start_ids = prompt_index.nodes_for_class("DistributedImageListStart") - image_list_mode = bool(image_list_start_ids) if image_list_mode: # Validate before probing or queueing anything so malformed boundaries and # old worker releases fail before any upstream image is materialized. - validate_image_list_regions(prompt_obj) + validate_image_list_regions(prompt_obj, explicit_targets=True) other_distributed = ( prompt_index.nodes_for_class("DistributedCollector") + prompt_index.nodes_for_class("UltimateSDUpscaleDistributed") @@ -409,6 +414,7 @@ async def orchestrate_distributed_execution( workflow_meta, client_id, include_queue_metadata=True, + **target_kwargs, ) return ( queue_result["prompt_id"], @@ -506,6 +512,7 @@ async def orchestrate_distributed_execution( workflow_meta, client_id, include_queue_metadata=True, + **target_kwargs, ) except Exception: if serverless_jobs: diff --git a/api/worker_routes.py b/api/worker_routes.py index a92b541..3f6b33f 100644 --- a/api/worker_routes.py +++ b/api/worker_routes.py @@ -61,6 +61,13 @@ async def worker_ws_endpoint(request): }) continue + if not isinstance(data, dict): + await ws.send_json({ + "type": "dispatch_ack", "request_id": None, "ok": False, + "error": "Expected a websocket message object.", + }) + continue + if data.get("type") != "dispatch_prompt": await ws.send_json({ "type": "dispatch_ack", @@ -81,10 +88,16 @@ async def worker_ws_endpoint(request): continue try: + target_kwargs = {} + if "partial_execution_targets" in data: + from ..utils.async_helpers import validate_execution_targets + target_kwargs["partial_execution_targets"] = validate_execution_targets( + prompt, data["partial_execution_targets"]) prompt_id = await queue_prompt_payload( prompt, workflow_meta=data.get("workflow"), client_id=data.get("client_id"), + **target_kwargs, ) await ws.send_json({ "type": "dispatch_ack", diff --git a/docs/image-list-queue.md b/docs/image-list-queue.md index 5295a06..5e09a75 100644 --- a/docs/image-list-queue.md +++ b/docs/image-list-queue.md @@ -70,7 +70,9 @@ File-based dependencies keep ComfyUI's `[output]` / `[temp]` root annotations an - Only the enclosed processing is distributed. List creation and downstream saving run on the master in both modes. With participation enabled, the master is a normal scheduler participant and can run alone when no direct workers are available. **Orchestrator-only mode never silently falls back to master processing.** Failed items are not reassigned automatically. - Local items execute while the enclosing master prompt waits; they are not queued behind that prompt on the master's `/prompt` queue. The enclosed branch is not also evaluated as a full list by the outer prompt. - Start/End regions may be chained sequentially and share model-loader dependencies. Nested/overlapping per-image regions and data escaping a region without End are rejected. Collector and distributed-upscaler nodes cannot be mixed into the same prompt in this first version. -- Output nodes inside a region are rejected because ComfyUI would execute them independently on the master. Avoid other side-effecting processing nodes as well; save after End. +- Queue preparation works on an execution copy, leaving the saved canvas unchanged. Unused per-item preview/save branches are pruned by dependency reachability, not node names. Required computational utilities marked `OUTPUT_NODE` stay in the graph; explicit `partial_execution_targets` prevent them from also becoming independent outer-master roots. End and legitimate post-End or unrelated outputs are retained. Bypass/rejoin branches still fail boundary validation rather than being silently removed. +- Every selected worker must advertise `execution_targets: true` as well as the existing image-list protocol. Older workers fail preflight before dispatch. Both HTTP and WebSocket item submissions target End only; the participating master uses its existing End-only local executor. Submit through `/distributed/queue`: an ordinary all-output `/prompt` request does not receive this relaxation. +- **Target selection is not a sandbox.** A required utility still runs its own save, preview or other side effects when it executes for an item. Keep intentional final saving after End. - The entire input list is materialized before scheduling; this is not streaming input. The limit is 4,096 frames. Remote transport preserves floating-point tensor dtype and values rather than encoding to lossy images. Each transported tensor is limited to 50 MiB, and the complete JSON/base64 HTTP request to 101 MiB; image/mask/value combinations can hit the aggregate request limit first. In-process master items use tensors directly and are not subject to those HTTP limits. - Results are accepted once per item using assignment tokens. Missing, duplicate, malformed and wrong-index results fail rather than silently returning a partial list. No failed or ambiguous submission is retried automatically. - Worker errors fail the region. The worker timeout measures loss of evidence for the owned prompt ID, not maximum render duration. Pending/running evidence refreshes liveness; unrelated busy jobs and failed polls do not. Healthy cold loads and long renders can continue. @@ -89,7 +91,7 @@ Use trusted workers and trusted workflows, as with the existing ComfyUI API. Cap ## Verification -The generic acceptance helper starts two isolated loopback worker test servers and uses actual CPU ComfyUI `PromptExecutor` instances. Production queue preparation, list nodes, graph extraction, HTTP item/result routes, dynamic scheduling, float transport, ordering and cancellation are exercised. The harness checks orchestrator-only, mixed master/worker and master-only execution, including exact-once item assignment, source-once materialization and cached local model-loader dependencies. The lightweight queue/history facade is test code; no managed ComfyUI service or model is started. +The generic acceptance helper starts two isolated loopback workers using real ComfyUI HTTP handlers, `PromptQueue` records and CPU `PromptExecutor` instances. Production queue preparation, list nodes, graph extraction, HTTP item/result routes, dynamic scheduling, float transport, ordering and cancellation are exercised. The harness checks orchestrator-only, mixed master/worker and master-only execution, including exact-once required output utilities, discarded previews, unchanged source graphs, End-only worker roots, post-End capture, source-once materialization and cached local model-loader dependencies. WebSocket target propagation is covered by focused route/dispatch tests. No managed ComfyUI service or model is started. ```bash COMFYUI_SOURCE_ROOT=/path/to/ComfyUI /path/to/ComfyUI/.venv/bin/python -m pytest tests -q diff --git a/image_list/execution_plan.py b/image_list/execution_plan.py new file mode 100644 index 0000000..9d196d7 --- /dev/null +++ b/image_list/execution_plan.py @@ -0,0 +1,70 @@ +"""Prepare image-list execution without changing the saved/API prompt or registry.""" +from __future__ import annotations + +import copy +import sys + +from .runtime import ( + END_CLASS, + _ancestors_for_inputs, + _downstream_until, + _image_list_pairs, + _is_output_node, + _prompt_links, + validate_image_list_regions, +) + + +def _reachable(roots: set[str], edges: dict[str, set[str]]) -> set[str]: + visited = set() + pending = list(roots) + while pending: + node_id = pending.pop() + if node_id in visited: + continue + visited.add(node_id) + pending.extend(edges[node_id] - visited) + return visited + + +def prepare_image_list_execution(prompt: dict) -> tuple[dict, list[str]]: + """Return an execution copy and its outer-master output targets. + + Disposable per-item side branches are removed before boundary validation. + Required output-node utilities remain normal dependencies, but must not be + independent outer roots. Callers must use these targets (and End-only targets + for item execution); this does not suppress side effects inside required nodes. + """ + pairs = _image_list_pairs(prompt) + prepared = copy.deepcopy(prompt) + upstream = {node_id: set() for node_id in prepared} + downstream = {node_id: set() for node_id in prepared} + for node_id, node in prepared.items(): + for _, source, _ in _prompt_links(node, prepared): + upstream[node_id].add(source) + downstream[source].add(node_id) + + # Never erase another region or its dependencies to make a bad graph pass. + boundaries = {node_id for pair in pairs for node_id in pair} + required = _reachable(boundaries, upstream) + discarded = set() + for start_id, end_id in pairs: + region_nodes = _ancestors_for_inputs(prepared, end_id, start_id, ("images", "masks", "texts")) + outside = _downstream_until(prepared, start_id, end_id) - region_nodes + # Preserve bypass joins and their whole chain into post-End computation. + # They must fail boundary validation, not disappear with debug previews. + post_end = _reachable({end_id}, downstream) + protected = required | _reachable(post_end, upstream) + discarded.update(outside - protected) + for node_id in discarded: + del prepared[node_id] + + regions = validate_image_list_regions(prepared, explicit_targets=True) + per_item = set().union(*(region["branch_ids"] for region in regions)) + registry = getattr(sys.modules.get("nodes"), "NODE_CLASS_MAPPINGS", {}) + targets = [ + node_id for node_id, node in prepared.items() + if node_id not in per_item and isinstance(node, dict) + and (node.get("class_type") == END_CLASS or _is_output_node(node, registry)) + ] + return prepared, targets diff --git a/image_list/lifecycle.py b/image_list/lifecycle.py index fe4c0c1..3c3d5d7 100644 --- a/image_list/lifecycle.py +++ b/image_list/lifecycle.py @@ -14,6 +14,38 @@ DEFAULT_ORPHAN_GRACE_SECONDS = 300.0 _LIFECYCLE_MARKER = object() +def has_safe_execution_targets(prompt_server: Any, prompt_id: str | None, prompt: dict) -> bool: + """Prove outer roots from the running queue, not a forgeable prompt marker. + + Plain /prompt remains strict when ComfyUI selected all outputs, including a + flagged per-item dependency. No saved workflow field grants this exemption. + """ + if not prompt_id: + return False + from .execution_plan import prepare_image_list_execution + prepared, targets = prepare_image_list_execution(prompt) + if prepared != prompt or not targets: + return False + queue = getattr(prompt_server, "prompt_queue", None) + for name in ("get_current_queue_volatile", "get_current_queue"): + reader = getattr(queue, name, None) + if not callable(reader): + continue + try: + running, _pending = reader() + for entry in running: + if (isinstance(entry, (tuple, list)) and len(entry) > 4 + and entry[1] == prompt_id and entry[2] == prompt): + roots = entry[4] + return (isinstance(roots, (list, tuple)) + and all(isinstance(root, str) for root in roots) + and len(roots) == len(targets) and set(roots) == set(targets)) + return False + except Exception: + continue + return False + + def _queue_prompt_ids(prompt_server: Any) -> set[str]: """Read running/pending prompt IDs, preferring ComfyUI's volatile view.""" queue = getattr(prompt_server, "prompt_queue", None) diff --git a/image_list/runtime.py b/image_list/runtime.py index c8f6848..f63f295 100644 --- a/image_list/runtime.py +++ b/image_list/runtime.py @@ -376,13 +376,19 @@ def delete_worker_item(token: str) -> bool: def create_master_job(start_id: str, end_id: str, prompt: dict[str, Any], items: list[dict[str, Any]], worker_ids: list[str], master_callback_url: str, workflow_meta: dict[str, Any] | None = None, - prompt_id: str | None = None, *, master_participates: bool = False) -> str: + prompt_id: str | None = None, *, master_participates: bool = False, + explicit_targets: bool = False) -> str: if not isinstance(prompt, dict): raise ValueError("Image List Start requires the original ComfyUI prompt for region extraction") if not isinstance(master_participates, bool): raise ValueError("master_participates must be a boolean") if not worker_ids and not master_participates: raise RuntimeError("Distributed Image List requires at least one preflighted direct worker; master fallback is disabled") + # Only the target-controlled Start caller may relax the outer prompt guard. + regions = validate_image_list_regions(prompt, explicit_targets=explicit_targets is True) + if not any(region["start_id"] == str(start_id) and region["end_id"] == str(end_id) + for region in regions): + raise ValueError("Image-list region does not match the requested Start/End pair") context_id = secrets.token_urlsafe(36) state = ImageListJob(context_id, str(start_id), str(end_id), prompt, items, [str(worker_id) for worker_id in worker_ids], master_callback_url, @@ -497,15 +503,23 @@ def _safe_shared_dependency(node: dict[str, Any], prompt: dict[str, Any], regist return types <= scalar_types and not any(_prompt_links(node, prompt)) -def validate_image_list_regions(prompt: dict[str, Any]) -> list[dict[str, Any]]: - """Validate pairing, non-nesting, and closed per-image graph boundaries.""" +def _is_output_node(node: dict[str, Any], registry) -> bool: + node_class = registry.get(node.get("class_type")) + if node_class is None: + return False + info = node_class.GET_NODE_INFO_V1() if hasattr(node_class, "GET_NODE_INFO_V1") else {} + return bool(info.get("output_node", getattr(node_class, "OUTPUT_NODE", False))) + + +def _image_list_pairs(prompt: dict[str, Any]) -> list[tuple[str, str]]: + """Pair boundaries before execution-copy pruning or full validation.""" if not isinstance(prompt, dict): raise ValueError("Prompt must be an object") starts = sorted(str(k) for k, v in prompt.items() if isinstance(v, dict) and v.get("class_type") == START_CLASS) ends = sorted(str(k) for k, v in prompt.items() if isinstance(v, dict) and v.get("class_type") == END_CLASS) if len(starts) != len(ends) or not starts: raise ValueError("Distributed Image List needs one paired Start and End per region") - regions = [] + pairs = [] for start_id in starts: candidates = [] for end_id in ends: @@ -514,7 +528,19 @@ def validate_image_list_regions(prompt: dict[str, Any]) -> list[dict[str, Any]]: candidates.append(end_id) if len(candidates) != 1: raise ValueError(f"Start node {start_id} must connect its typed context output directly to exactly one End") - end_id = candidates[0] + pairs.append((start_id, candidates[0])) + return pairs + + +def validate_image_list_regions(prompt: dict[str, Any], *, explicit_targets: bool = False) -> list[dict[str, Any]]: + """Validate pairing, non-nesting, and closed per-image graph boundaries. + + Required output-node utilities are allowed only when callers explicitly + select execution roots, rather than letting every OUTPUT_NODE run independently. + All data-boundary checks still apply in that mode. + """ + regions = [] + for start_id, end_id in _image_list_pairs(prompt): if "images" not in prompt[end_id].get("inputs", {}): raise ValueError(f"End node {end_id}.images must be connected to the per-item result branch") end_inputs = prompt[end_id].get("inputs", {}) @@ -529,8 +555,6 @@ def validate_image_list_regions(prompt: dict[str, Any]) -> list[dict[str, Any]]: {START_CLASS, END_CLASS, "DistributedCollector", "UltimateSDUpscaleDistributed"}} if forbidden: raise ValueError("Nested/ambiguous regions and other distributed nodes are unsupported in an image-list region: " + ", ".join(sorted(forbidden))) - # Output nodes execute independently of the lazy End, so they would - # evaluate the processing branch on the master as well as the worker. import sys registry = getattr(sys.modules.get("nodes"), "NODE_CLASS_MAPPINGS", {}) shared = region_nodes & _start_ancestors(prompt, start_id) @@ -545,13 +569,10 @@ def validate_image_list_regions(prompt: dict[str, Any]) -> list[dict[str, Any]]: f"Image-list region {start_id}->{end_id} bypasses Start through upstream nodes " f"{sorted(unsafe)} ({'; '.join(links)}); route per-item data through Start" ) - for node_id in region_nodes - {start_id, end_id}: - node_class = registry.get(prompt[node_id].get("class_type")) - if node_class is None: - continue - info = node_class.GET_NODE_INFO_V1() if hasattr(node_class, "GET_NODE_INFO_V1") else {} - if info.get("output_node", getattr(node_class, "OUTPUT_NODE", False)): - raise ValueError(f"Output node {node_id} must be placed after Distributed Image List End") + if not explicit_targets: + for node_id in region_nodes - {start_id, end_id}: + if _is_output_node(prompt[node_id], registry): + raise ValueError(f"Output node {node_id} must be placed after Distributed Image List End") downstream = _downstream_until(prompt, start_id, end_id) outside = sorted(downstream - region_nodes) if outside: @@ -822,7 +843,8 @@ async def _run_image_list_job(state: ImageListJob): if not workers and not state.master_participates: raise RuntimeError("No free, reachable, feature-compatible direct workers are available; master execution is disabled") session = await get_client_session() if workers else None - regions = validate_image_list_regions(state.original_prompt) + # Both item paths below select End explicitly, never all OUTPUT_NODE roots. + regions = validate_image_list_regions(state.original_prompt, explicit_targets=True) region = next((entry for entry in regions if entry["start_id"] == state.start_id and entry["end_id"] == state.end_id), None) if region is None: raise RuntimeError("Image-list region no longer matches the submitted prompt") @@ -892,6 +914,7 @@ async def _run_image_list_job(state: ImageListJob): worker, worker_prompt, state.workflow_meta, client_id=f"image-list-{context_id[:20]}", use_websocket=use_websocket, return_prompt_id=True, + partial_execution_targets=[state.end_id], ) if not isinstance(prompt_id, str) or not prompt_id: raise RuntimeError(f"Worker {worker_id} did not return a prompt_id; submission status is ambiguous") diff --git a/nodes/image_list.py b/nodes/image_list.py index 3c0bb62..1057b3c 100644 --- a/nodes/image_list.py +++ b/nodes/image_list.py @@ -99,13 +99,18 @@ class DistributedImageListStart(io.ComfyNode): raise RuntimeError("Invalid preflighted image-list worker list") from exc if not isinstance(worker_ids, list) or not all(isinstance(item, str) for item in worker_ids): raise RuntimeError("Invalid preflighted image-list worker list") + from comfy_execution.utils import get_executing_context + from ..image_list.lifecycle import has_safe_execution_targets + import server + executing = get_executing_context() + explicit_targets = has_safe_execution_targets( + getattr(server.PromptServer, "instance", None), + executing.prompt_id if executing is not None else None, prompt) from ..image_list.runtime import validate_image_list_regions - regions = validate_image_list_regions(prompt) + regions = validate_image_list_regions(prompt, explicit_targets=explicit_targets) region = next((row for row in regions if row["start_id"] == start_id), None) if region is None: raise RuntimeError(f"No safe paired Image List End was found for Start node {start_id}") - from comfy_execution.utils import get_executing_context - executing = get_executing_context() master_participates = unwrap_single(image_list_master_participates) if master_participates is None: from ..utils.config import load_config @@ -118,6 +123,7 @@ class DistributedImageListStart(io.ComfyNode): image_list_workflow_meta if isinstance(image_list_workflow_meta, dict) else None, prompt_id=executing.prompt_id if executing is not None else None, master_participates=master_participates, + explicit_targets=explicit_targets, ) return io.NodeOutput( [item["image"] for item in items], diff --git a/tests/api/test_worker_routes.py b/tests/api/test_worker_routes.py index f6d0b2e..c66ebe6 100644 --- a/tests/api/test_worker_routes.py +++ b/tests/api/test_worker_routes.py @@ -271,6 +271,74 @@ def _load_worker_routes_module(): worker_routes = _load_worker_routes_module() +class WebsocketExecutionTargetTests(unittest.IsolatedAsyncioTestCase): + def setUp(self): + from test_async_helpers import _load_async_helpers_module + self.modules = patch.dict(sys.modules) + self.modules.start() + self.addCleanup(self.modules.stop) + self.helper, self.server = _load_async_helpers_module() + sys.modules['dist_api_worker_testpkg.utils.async_helpers'] = self.helper + self.prompt = {'9': {'class_type': 'Output'}, '10': {'class_type': 'Output'}} + + async def receive(self, payload): + import json + sent = [] + class Socket: + async def prepare(self, request): + pass + async def send_json(self, value): + sent.append(value) + def __aiter__(self): + async def messages(): + yield types.SimpleNamespace(type=worker_routes.aiohttp.WSMsgType.TEXT, + data=json.dumps(payload)) + return messages() + with patch.object(worker_routes, 'web', types.SimpleNamespace(WebSocketResponse=lambda **kw: Socket())), \ + patch.object(worker_routes, 'queue_prompt_payload', self.helper.queue_prompt_payload), \ + patch.object(worker_routes, 'PromptValidationError', self.helper.PromptValidationError): + await worker_routes.worker_ws_endpoint(_FakeRequest()) + return sent[0] + + def payload(self, targets): + return {'type': 'dispatch_prompt', 'request_id': 'req', 'prompt': self.prompt, + 'partial_execution_targets': targets} + + async def test_explicit_targets_reach_queue_and_acknowledgement(self): + ack = await self.receive(self.payload(['9'])) + self.assertTrue(ack['ok'], ack) + self.assertEqual(ack['request_id'], 'req') + self.assertEqual(self.server.prompt_queue.items[0][4], ['9']) + + async def test_malformed_targets_and_messages_are_rejected(self): + for targets in (None, [], '9', [True], [9], ['absent'], ['9', '9']): + with self.subTest(targets=targets): + ack = await self.receive(self.payload(targets)) + self.assertFalse(ack['ok']) + self.assertIn('partial_execution_targets', ack['error']) + self.assertFalse((await self.receive([]))['ok']) + self.assertEqual(self.server.prompt_queue.items, []) + + async def test_validation_cannot_silently_queue_a_partial_subset(self): + ack = await self.receive(self.payload(['9', '10'])) + self.assertFalse(ack['ok']) + self.assertEqual(ack['validation_error']['type'], 'partial_execution_targets_failed_validation') + self.assertEqual(self.server.prompt_queue.items, []) + + async def test_missing_target_field_keeps_legacy_queue_call(self): + from unittest.mock import AsyncMock + queued = AsyncMock(return_value='legacy') + original = self.helper.queue_prompt_payload + self.helper.queue_prompt_payload = queued + try: + payload = self.payload(['9']) + payload.pop('partial_execution_targets') + self.assertTrue((await self.receive(payload))['ok']) + self.assertNotIn('partial_execution_targets', queued.call_args.kwargs) + finally: + self.helper.queue_prompt_payload = original + + class WorkerRoutesTests(unittest.IsolatedAsyncioTestCase): async def test_network_info_reports_actual_master_and_available_ports(self): with patch.object(worker_routes, "_get_cuda_info", return_value=(0, 1, 4)), \ diff --git a/tests/helpers/image_list_network_check.py b/tests/helpers/image_list_network_check.py index 95fb384..1f58ecf 100644 --- a/tests/helpers/image_list_network_check.py +++ b/tests/helpers/image_list_network_check.py @@ -1,11 +1,12 @@ """Model-free multi-process acceptance using real ComfyUI executors and loopback HTTP. -Only the transport test facade (/prompt, /queue, /history, /interrupt) is synthetic. +Isolated listeners use real core /prompt, /queue, /history and /interrupt handlers. Start/End, graph extraction, scheduling, payload codecs and callbacks are production code. No managed ComfyUI service, user workflow, model, or remote worker is touched. """ import asyncio from collections import Counter +import copy import importlib import json import os @@ -36,6 +37,7 @@ import aiohttp torch.set_num_threads(1) SOURCE_CALLS, BRANCH_CALLS, CAPTURED, MODEL_CALLS = [], [], [], [] +UTILITY_CALLS, ITEM_EXECUTIONS = [], [] ALLOW_MASTER = False COUNT = 7 @@ -85,6 +87,51 @@ class TestCachedModel: return (object(),) +def assert_item_execution(prompt, index): + """Prove the operation belongs to an item, not merely a participating host.""" + from comfy_execution.utils import get_executing_context + context = get_executing_context() + assert context is not None + assert prompt['2']['inputs'].get('worker_item_token'), 'per-item node ran in outer graph' + running, _ = server.PromptServer.instance.prompt_queue.get_current_queue_volatile() + assert len(running) == 1, running + entry = running[0] + if WORKER: + assert entry[1] == context.prompt_id and entry[2] == prompt + assert entry[4] == ['4'], entry[4] + else: + assert ALLOW_MASTER, 'per-item node ran on delegate-only master' + assert entry[1] != context.prompt_id, 'per-item node ran in outer executor' + assert set(entry[4]) == {'4', '5'}, entry[4] + return {'index': index, 'prompt_id': context.prompt_id, + 'role': 'worker' if WORKER else 'master-item', 'queue_prompt_id': entry[1]} + + +class TestRequiredOutputUtility: + @classmethod + def INPUT_TYPES(cls): + return {'required': {'image': ('IMAGE',), 'index': ('INT', {'forceInput': True})}, + 'hidden': {'prompt': 'PROMPT'}} + OUTPUT_NODE = True + RETURN_TYPES = ('IMAGE',) + FUNCTION = 'run' + def run(self, image, index, prompt): + assert_item_execution(prompt, index) + UTILITY_CALLS.append(index) + return (image,) + + +class TestDiscardedPreview: + @classmethod + def INPUT_TYPES(cls): + return {'required': {'image': ('IMAGE',)}} + OUTPUT_NODE = True + RETURN_TYPES = () + FUNCTION = 'run' + def run(self, image): + raise AssertionError('discarded per-item output preview executed') + + class TestImageOperation: @classmethod def INPUT_TYPES(cls): @@ -92,11 +139,13 @@ class TestImageOperation: 'text': ('STRING', {'forceInput': True}), 'strength': ('FLOAT', {'forceInput': True}), 'seed': ('INT', {'forceInput': True}), 'index': ('INT', {'forceInput': True}), - 'flag': ('BOOLEAN', {'forceInput': True}), 'model': ('MODEL',)}} + 'flag': ('BOOLEAN', {'forceInput': True}), 'model': ('MODEL',)}, + 'hidden': {'prompt': 'PROMPT'}} RETURN_TYPES = ('IMAGE', 'MASK', 'STRING') FUNCTION = 'run' - def run(self, image, reference_image, mask, text, strength, seed, index, flag, model): + def run(self, image, reference_image, mask, text, strength, seed, index, flag, model, prompt): assert WORKER or ALLOW_MASTER, 'per-image branch executed in orchestrator-only mode' + ITEM_EXECUTIONS.append(assert_item_execution(prompt, index)) BRANCH_CALLS.append(index) # Sleep in short chunks so the real Comfy interrupt signal can be tested. until = time.monotonic() + DELAY @@ -120,6 +169,7 @@ class TestCapture: RETURN_TYPES = () FUNCTION = 'run' def run(self, images, masks, texts): + assert not WORKER, 'post-End capture leaked into a worker prompt' CAPTURED.append((images, masks, texts)) return () @@ -131,7 +181,9 @@ async def setup(): module = sys.modules[str(ROOT).replace('.', '_x_')] nodes.NODE_CLASS_MAPPINGS.update(TestImageSource=TestImageSource, TestImageOperation=TestImageOperation, TestCapture=TestCapture, - TestCachedModel=TestCachedModel, TestSharedFlag=TestSharedFlag) + TestCachedModel=TestCachedModel, TestSharedFlag=TestSharedFlag, + TestRequiredOutputUtility=TestRequiredOutputUtility, + TestDiscardedPreview=TestDiscardedPreview) # Match the CORS-enabled direct-worker launch contract for canvas probes. app = web.Application(middlewares=[server.create_cors_middleware('*')]) # Register the actual feature routes, not test replacements. @@ -147,61 +199,56 @@ async def listen(app): return runner, site._server.sockets[0].getsockname()[1] +def finish_queue_item(ps, item_id, executor): + ps.prompt_queue.task_done(item_id, executor.history_result, + status=execution.PromptQueue.ExecutionStatus( + status_str='success' if executor.success else 'error', + completed=executor.success, messages=executor.status_messages), + process_item=lambda item: item[:5] + item[6:]) + + async def worker_main(): ps, assets, module, app = await setup() - pending, running, history = [], {}, {} wake = asyncio.Event() - interrupted_ids = [] + interrupted_ids, submissions = [], [] + core_handlers = {(route.method, route.path): route.handler for route in ps.routes} async def prompt(request): data = await request.json() - prompt_id = str(uuid.uuid4()) - graph = data['prompt'] - valid = await execution.validate_prompt(prompt_id, graph, None) - if not valid[0]: - return web.json_response({'error': valid[1], 'node_errors': valid[3]}, status=400) - pending.append((prompt_id, graph, valid[2])) - wake.set() - return web.json_response({'prompt_id': prompt_id}) - async def queue(request): - if request.method == 'POST': - deleted = (await request.json()).get('delete', []) - pending[:] = [entry for entry in pending if entry[0] not in deleted] - return web.json_response({}) - return web.json_response({'queue_running': [[0, key] for key in running], - 'queue_pending': [[0, entry[0]] for entry in pending]}) - async def get_history(request): - pid = request.match_info['prompt_id'] - return web.json_response({pid: history[pid]} if pid in history else {}) + # Use core's actual validation/enqueue/response schema, including + # partial_execution_targets; only observe it and wake our CPU consumer. + response = await core_handlers['POST', '/prompt'](request) + if response.status == 200: + accepted = json.loads(response.body) + _, pending = ps.prompt_queue.get_current_queue_volatile() + entry = next(item for item in pending if item[1] == accepted['prompt_id']) + submissions.append({'prompt_id': entry[1], 'requested': data.get('partial_execution_targets'), + 'queued': entry[4], 'node_ids': sorted(entry[2]), + 'end_ids': [key for key, node in entry[2].items() + if node['class_type'] == 'DistributedImageListEnd']}) + wake.set() + return response async def interrupt(request): - pid = (await request.json())['prompt_id'] - interrupted_ids.append(pid) - if pid in running: - comfy.model_management.interrupt_current_processing(True) - return web.json_response({}) + interrupted_ids.append((await request.json())['prompt_id']) + return await core_handlers['POST', '/interrupt'](request) async def stats(request): - return web.json_response({'calls': BRANCH_CALLS, 'running': list(running), - 'pending': [e[0] for e in pending], 'interrupted': interrupted_ids, - 'history': history}) + running, pending = ps.prompt_queue.get_current_queue_volatile() + return web.json_response({'calls': BRANCH_CALLS, 'utility_calls': UTILITY_CALLS, + 'item_executions': ITEM_EXECUTIONS, 'submissions': submissions, + 'running': [e[1] for e in running], 'pending': [e[1] for e in pending], + 'interrupted': interrupted_ids, 'history': ps.prompt_queue.get_history()}) async def consume(): executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets) while True: await wake.wait() - while pending: - pid, graph, outputs = pending.pop(0) - running[pid] = True - await asyncio.to_thread(executor.execute, graph, pid, {}, outputs) - history[pid] = {'status': {'completed': True, - 'status_str': 'success' if executor.success else 'error', - 'messages': str(executor.status_messages)}} - running.pop(pid, None) + while (queued := ps.prompt_queue.get(timeout=0)) is not None: + entry, item_id = queued + await asyncio.to_thread(executor.execute, entry[2], entry[1], entry[3], entry[4]) + finish_queue_item(ps, item_id, executor) wake.clear() - async def probe(_request): - return web.json_response({'exec_info': {'queue_remaining': len(pending) + len(running)}}) - app.router.add_get('/prompt', probe) + app.add_routes([route for route in ps.routes if + (route.path in ('/queue', '/history/{prompt_id}') or + (route.path == '/prompt' and route.method == 'GET'))]) app.router.add_post('/prompt', prompt) - app.router.add_get('/queue', queue) - app.router.add_post('/queue', queue) - app.router.add_get('/history/{prompt_id}', get_history) app.router.add_post('/interrupt', interrupt) app.router.add_get('/test-stats', stats) runner, port = await listen(app) @@ -228,12 +275,14 @@ def graph(callback, worker_ids): '3': {'class_type': 'TestImageOperation', 'inputs': { key: ['2', index] for key, index in [('image', 0), ('mask', 1), ('index', 2), ('reference_image', 4), ('text', 5), ('strength', 6), ('seed', 7), ('flag', 8)]} - | {'model': ['6', 0]}}, + | {'image': ['8', 0], 'model': ['6', 0]}}, '4': {'class_type': 'DistributedImageListEnd', 'inputs': { 'images': ['3', 0], 'masks': ['3', 1], 'texts': ['3', 2], 'context': ['2', 3]}}, '5': {'class_type': 'TestCapture', 'inputs': {'images': ['4', 0], 'masks': ['4', 1], 'texts': ['4', 2]}}, '6': {'class_type': 'TestCachedModel', 'inputs': {}}, '7': {'class_type': 'TestSharedFlag', 'inputs': {}}, + '8': {'class_type': 'TestRequiredOutputUtility', 'inputs': {'image': ['2', 0], 'index': ['2', 2]}}, + '9': {'class_type': 'TestDiscardedPreview', 'inputs': {'image': ['8', 0]}}, } @@ -241,6 +290,8 @@ async def coordinator(): global ALLOW_MASTER ps, assets, module, app = await setup() runtime = importlib.import_module(module.__name__ + '.image_list.runtime') + execution_plan = importlib.import_module(module.__name__ + '.image_list.execution_plan') + lifecycle = importlib.import_module(module.__name__ + '.image_list.lifecycle') config_mod = importlib.import_module(module.__name__ + '.utils.config') orchestration = importlib.import_module(module.__name__ + '.api.queue_orchestration') runner, port = await listen(app) @@ -294,16 +345,49 @@ async def coordinator(): prompt = graph(f'http://127.0.0.1:{port}', [w['id'] for w in workers]) valid = await execution.validate_prompt('network-acceptance', prompt, None) assert valid[0], valid + assert {'4', '5', '8', '9'} == set(valid[2]), valid executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets) # Exercise production queue preparation; only master enqueue is redirected # to this isolated executor, not the managed server's live queue. prompt['2']['inputs'].pop('image_list_enabled_worker_ids') prompt['2']['inputs'].pop('image_list_master_callback') + saved_prompt = copy.deepcopy(prompt) + outer_executions = [] async def enqueue(prepared, workflow, client_id, **kwargs): - checked = await execution.validate_prompt('network-acceptance', prepared, None) + targets = kwargs['partial_execution_targets'] + planned, planned_targets = execution_plan.prepare_image_list_execution(prepared) + assert planned == prepared and targets == planned_targets + if '9' in saved_prompt and 'TestImageOperation' in { + node['class_type'] for node in prepared.values()}: + assert set(prepared) == set(saved_prompt) - {'9'}, prepared + assert prepared['8'] == saved_prompt['8'], 'required utility changed' + assert prepared['5'] == saved_prompt['5'], 'post-End output changed' + assert set(targets) == {'4', '5'}, targets + assert nodes.NODE_CLASS_MAPPINGS['TestRequiredOutputUtility'].OUTPUT_NODE is True + assert nodes.NODE_CLASS_MAPPINGS['TestDiscardedPreview'].OUTPUT_NODE is True + pid = 'network-acceptance' if not outer_executions else str(uuid.uuid4()) + checked = await execution.validate_prompt(pid, prepared, targets) assert checked[0], checked - await asyncio.to_thread(executor.execute, prepared, 'network-acceptance', {}, checked[2]) - return {'prompt_id': 'network-acceptance', 'number': 1, 'node_errors': {}} + assert not checked[3] and Counter(checked[2]) == Counter(targets), checked + assert not lifecycle.has_safe_execution_targets(ps, pid, prepared) + # Exactly the core queue tuple/put/get/task_done lifecycle. Start + # must prove its roots from this real running record, not a mock. + number = len(outer_executions) + ps.prompt_queue.put((number, pid, prepared, {'client_id': client_id}, checked[2], {})) + entry, item_id = ps.prompt_queue.get(timeout=0) + assert lifecycle.has_safe_execution_targets(ps, pid, prepared) + before_captures = len(CAPTURED) + await asyncio.to_thread(executor.execute, entry[2], entry[1], entry[3], entry[4]) + finish_queue_item(ps, item_id, executor) + assert ps.prompt_queue.get_current_queue_volatile() == ([], []) + history = ps.prompt_queue.get_history(pid)[pid] + assert Counter(history['prompt'][4]) == Counter(targets) + assert not lifecycle.has_safe_execution_targets(ps, pid, prepared) + outer_executions.append({'prompt_id': pid, 'targets': targets, + 'captures': len(CAPTURED) - before_captures}) + if executor.success: + assert len(CAPTURED) == before_captures + 1, 'post-End Capture must run exactly once' + return {'prompt_id': pid, 'number': number, 'node_errors': checked[3]} with patch.object(config_mod, 'load_config', return_value=config), \ patch.object(orchestration, 'load_config', return_value=config), \ patch.object(orchestration, 'queue_prompt_payload', enqueue), \ @@ -312,11 +396,12 @@ async def coordinator(): prompt, None, 'test', [w['id'] for w in workers]), 45) assert queued[0] == 'network-acceptance' and queued[2] == 2, queued assert not ps.distributed_pending_jobs, 'image-list created unused Collector queues' + assert prompt == saved_prompt, 'queue preparation changed the saved/source graph' if not executor.success: for err in stderr_files: err.flush(); err.seek(0); print(err.read(), file=sys.stderr) assert executor.success, executor.status_messages - assert len(SOURCE_CALLS) == 1 and not BRANCH_CALLS + assert len(SOURCE_CALLS) == 1 and not BRANCH_CALLS and not UTILITY_CALLS assert len(CAPTURED) == 1 images, masks, texts = CAPTURED[0] assert len(images) == len(masks) == len(texts) == COUNT @@ -332,6 +417,14 @@ async def coordinator(): counts = [len(row['calls']) for row in stats] assert counts[0] > counts[1] > 0, counts assert Counter(i for row in stats for i in row['calls']) == Counter(range(COUNT)) + for row in stats: + assert Counter(row['utility_calls']) == Counter(row['calls']), row + assert len(row['submissions']) == len(row['calls']) + assert len({event['prompt_id'] for event in row['item_executions']}) == len(row['calls']) + for submission in row['submissions']: + assert submission['requested'] == submission['queued'] == ['4'], submission + assert '8' in submission['node_ids'] and '9' not in submission['node_ids'] + assert '5' not in submission['node_ids'] assert not runtime._MASTER_JOBS # Three sequential regions on one worker, including optional outputs # left unconnected. Region N+1 must not race the prior worker End. @@ -346,12 +439,14 @@ async def coordinator(): previous = end chained['8'] = {'class_type': 'TestCapture', 'inputs': { 'images': ['7', 0], 'masks': ['7', 1], 'texts': ['7', 2]}} + saved_chained = copy.deepcopy(chained) with patch.object(config_mod, 'load_config', return_value=config), \ patch.object(orchestration, 'load_config', return_value=config), \ patch.object(orchestration, 'queue_prompt_payload', enqueue), \ patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15): await asyncio.wait_for(orchestration.orchestrate_distributed_execution( chained, None, 'test', ['slow']), 45) + assert chained == saved_chained assert executor.success, executor.status_messages assert len(CAPTURED) == 2 and len(CAPTURED[1][0]) == 2 assert all(torch.equal(image, torch.zeros((1, 8, 8, 3))) for image in CAPTURED[1][0]) @@ -366,29 +461,43 @@ async def coordinator(): for mode, selected in [('mixed', [w['id'] for w in workers]), ('master-only', [])]: executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets) before_sources, before_local, before_models = len(SOURCE_CALLS), len(BRANCH_CALLS), len(MODEL_CALLS) + before_utilities, before_events = len(UTILITY_CALLS), len(ITEM_EXECUTIONS) async with aiohttp.ClientSession() as session: - before_worker_calls = [] + before_worker_stats = [] for worker in workers: async with session.get(f"http://127.0.0.1:{worker['port']}/test-stats") as response: - before_worker_calls.append(len((await response.json())['calls'])) + before_worker_stats.append(await response.json()) with patch.object(config_mod, 'load_config', return_value=config), \ patch.object(orchestration, 'load_config', return_value=config), \ patch.object(orchestration, 'queue_prompt_payload', enqueue), \ patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15): queued = await asyncio.wait_for(orchestration.orchestrate_distributed_execution( prompt, None, 'test', selected), 45) + assert prompt == saved_prompt, (mode, 'saved/source graph changed') assert executor.success, (mode, executor.status_messages) assert queued[2] == len(selected) assert len(SOURCE_CALLS) == before_sources + 1, 'source was reevaluated per local item' local_calls = BRANCH_CALLS[before_local:] + assert Counter(UTILITY_CALLS[before_utilities:]) == Counter(local_calls), mode + assert len({event['prompt_id'] for event in ITEM_EXECUTIONS[before_events:]}) == len(local_calls) assert local_calls, f'{mode}: master did not process any items' assert len(MODEL_CALLS) == before_models + 1, 'local model dependency did not stay cached' remote_calls = [] async with aiohttp.ClientSession() as session: - for worker, offset in zip(workers, before_worker_calls): + for worker, before_stats in zip(workers, before_worker_stats): async with session.get(f"http://127.0.0.1:{worker['port']}/test-stats") as response: current = await response.json() - remote_calls.extend(current['calls'][offset:]) + calls = current['calls'][len(before_stats['calls']):] + utilities = current['utility_calls'][len(before_stats['utility_calls']):] + submissions = current['submissions'][len(before_stats['submissions']):] + events = current['item_executions'][len(before_stats['item_executions']):] + assert Counter(utilities) == Counter(calls), (mode, current) + assert len(submissions) == len(calls) == len({event['prompt_id'] for event in events}) + for submission in submissions: + assert submission['requested'] == submission['queued'] == ['4'], submission + assert '8' in submission['node_ids'] and '9' not in submission['node_ids'] + assert '5' not in submission['node_ids'] + remote_calls.extend(calls) assert not current['running'] and not current['pending'], current if mode == 'mixed': assert remote_calls, 'mixed execution failed to use direct workers' @@ -404,8 +513,19 @@ async def coordinator(): assert not runtime._MASTER_JOBS and not runtime._WORKER_ITEMS participation_counts[mode] = {'master': len(local_calls), 'workers': len(remote_calls)} print('IMAGE_LIST_MASTER_OK ' + json.dumps(participation_counts)) + assert len(outer_executions) == 4 and all(row['captures'] == 1 for row in outer_executions) + print('IMAGE_LIST_TARGETS_OK ' + json.dumps({ + 'modes': ['delegate-only', 'mixed', 'master-only'], + 'outer_targets': outer_executions[0]['targets'], 'worker_targets': ['4'], + 'required_utility_once_per_item': True, 'preview_discarded': True, + 'saved_prompt_unchanged': prompt == saved_prompt, 'capture_once_per_outer_prompt': True})) ALLOW_MASTER = False config['settings']['master_delegate_only'] = True + # Preserve the original direct-runtime cancellation/failure fixture: + # it has no outer queue and must not claim target-control evidence. + direct_prompt = copy.deepcopy(prompt) + del direct_prompt['8'], direct_prompt['9'] + direct_prompt['3']['inputs']['image'] = ['2', 0] # Real Comfy user-interrupt flag while a real worker executor is active. slow = workers[1] async with aiohttp.ClientSession() as session: @@ -420,7 +540,7 @@ async def coordinator(): await asyncio.sleep(0.02) raise AssertionError('worker state transition timed out') before = await until(lambda state: not state['running'] and not state['pending']) - context = runtime.create_master_job('2', '4', prompt, source_rows(), ['slow'], + context = runtime.create_master_job('2', '4', direct_prompt, source_rows(), ['slow'], f'http://127.0.0.1:{port}') with patch.object(config_mod, 'load_config', return_value=config), \ patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15): @@ -444,7 +564,7 @@ async def coordinator(): # A worker exception before End must fail, not hang or return a partial list. bad_rows = source_rows() bad_rows[0]['seed'] = -1 - context = runtime.create_master_job('2', '4', prompt, bad_rows, ['fast'], + context = runtime.create_master_job('2', '4', direct_prompt, bad_rows, ['fast'], f'http://127.0.0.1:{port}') with patch.object(config_mod, 'load_config', return_value=config), \ patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15): diff --git a/tests/test_async_helpers.py b/tests/test_async_helpers.py index 72b61cc..c6e66b2 100644 --- a/tests/test_async_helpers.py +++ b/tests/test_async_helpers.py @@ -3,6 +3,7 @@ import sys import types import unittest from pathlib import Path +from unittest.mock import AsyncMock, patch class _PromptQueue: @@ -87,5 +88,59 @@ class QueuePromptPayloadTests(unittest.IsolatedAsyncioTestCase): self.assertEqual(extra_data["extra_pnginfo"]["workflow"], {"id": "workflow-1"}) +class ExplicitTargetTests(unittest.IsolatedAsyncioTestCase): + def setUp(self): + self.module, self.server = _load_async_helpers_module() + self.prompt = {"9": {"class_type": "Output"}, "10": {"class_type": "Output"}} + self.validate = AsyncMock(return_value=(True, None, ["9"], {})) + self.module.execution.validate_prompt = self.validate + + async def test_targets_survive_callback_omission_and_are_queued_exactly(self): + seen = [] + def callback(payload): + seen.append(payload["partial_execution_targets"][:]) + payload["partial_execution_targets"].append("10") + return {"prompt": payload["prompt"]} + self.server.trigger_on_prompt = callback + targets = ["9"] + await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=targets) + self.assertEqual(seen, [["9"]]) + self.assertEqual(targets, ["9"]) + self.assertEqual(self.validate.call_args.args[2], ["9"]) + self.assertEqual(self.server.prompt_queue.items[0][4], ["9"]) + + async def test_invalid_requests_never_reach_validation_or_queue(self): + for targets in ([], "9", [9], [True], ["9", "9"], ["missing"], {}): + with self.subTest(targets=targets): + with self.assertRaises(self.module.PromptValidationError): + await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=targets) + self.validate.assert_not_awaited() + self.assertEqual(self.server.prompt_queue.items, []) + + async def test_removed_target_after_callback_is_rejected(self): + self.server.trigger_on_prompt = lambda payload: {"prompt": {"10": self.prompt["10"]}} + with self.assertRaises(self.module.PromptValidationError): + await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=["9"]) + self.validate.assert_not_awaited() + self.assertEqual(self.server.prompt_queue.items, []) + + async def test_partial_subset_extra_roots_and_target_errors_fail_closed(self): + for roots, errors in ((["9"], {}), (["9", "10", "extra"], {}), + (["9", "10"], {"9": {"errors": [{"message": "bad"}]}}), + (["9", "10"], {"dep": {"dependent_outputs": ["10"], "errors": []}})): + with self.subTest(roots=roots, errors=errors): + self.validate.return_value = (True, None, roots, errors) + with self.assertRaises(self.module.PromptValidationError): + await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=["9", "10"]) + self.assertEqual(self.server.prompt_queue.items, []) + + async def test_none_preserves_legacy_payload_and_validation(self): + seen = [] + self.server.trigger_on_prompt = lambda payload: seen.append(payload.copy()) or payload + await self.module.queue_prompt_payload(self.prompt) + self.assertNotIn("partial_execution_targets", seen[0]) + self.assertIsNone(self.validate.call_args.args[2]) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_dispatch_selection.py b/tests/test_dispatch_selection.py index 401121b..8527ce8 100644 --- a/tests/test_dispatch_selection.py +++ b/tests/test_dispatch_selection.py @@ -4,7 +4,7 @@ import sys import types import unittest from pathlib import Path -from unittest.mock import patch +from unittest.mock import AsyncMock, patch def _load_dispatch_module(): @@ -168,6 +168,7 @@ class _FakeSession: def _valid_capability(): return { "supported": True, + "execution_targets": True, "protocol_version": 2, "node_ids": ["DistributedImageListStart", "DistributedImageListEnd"], } @@ -177,6 +178,43 @@ def _valid_queue(running=None, pending=None): return {"queue_running": running or [], "queue_pending": pending or []} +class ExplicitTargetTransportTests(unittest.IsolatedAsyncioTestCase): + async def test_http_preserves_targets_and_omits_default_field(self): + response = types.SimpleNamespace(raise_for_status=lambda: None) + class Response: + async def __aenter__(self): return response + async def __aexit__(self, *args): pass + calls = [] + def post(url, **kwargs): + calls.append(kwargs['json']) + return Response() + session = types.SimpleNamespace(post=post) + with patch.object(dispatch, 'get_client_session', AsyncMock(return_value=session)): + await dispatch.dispatch_worker_prompt({'id': 'worker'}, {'9': {}}, None, + partial_execution_targets=['9']) + await dispatch.dispatch_worker_prompt({'id': 'worker'}, {'9': {}}, None) + self.assertEqual(calls[0]['partial_execution_targets'], ['9']) + self.assertNotIn('partial_execution_targets', calls[1]) + + async def test_websocket_sends_targets_all_the_way_to_wire(self): + import json + sent = [] + async def receive(): + return types.SimpleNamespace(type=dispatch.aiohttp.WSMsgType.TEXT, data=json.dumps({ + 'type': 'dispatch_ack', 'request_id': sent[-1]['request_id'], + 'ok': True, 'prompt_id': 'worker-prompt'})) + ws = types.SimpleNamespace(send_json=AsyncMock(side_effect=sent.append), + receive=receive, close=AsyncMock()) + session = types.SimpleNamespace(ws_connect=AsyncMock(return_value=ws)) + with patch.object(dispatch, 'get_client_session', AsyncMock(return_value=session)): + result = await dispatch.dispatch_worker_prompt({'id': 'worker'}, {'9': {}}, None, + use_websocket=True, return_prompt_id=True, partial_execution_targets=['9']) + await dispatch.dispatch_worker_prompt({'id': 'worker'}, {'9': {}}, None, use_websocket=True) + self.assertEqual(result, 'worker-prompt') + self.assertEqual(sent[0]['partial_execution_targets'], ['9']) + self.assertNotIn('partial_execution_targets', sent[1]) + + class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase): async def test_select_active_workers_filters_offline(self): workers = [ @@ -417,6 +455,23 @@ class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase): with self.assertRaisesRegex(RuntimeError, "No selected feature-compatible direct worker is idle"): await dispatch.preflight_image_list_workers(workers, require_idle=True) + async def test_image_list_preflight_requires_exact_execution_targets_support(self): + for value in (None, False, 1, "true", []): + with self.subTest(value=value): + capability = _valid_capability() + if value is None: + capability.pop("execution_targets") + else: + capability["execution_targets"] = value + session = _FakeSession({("worker", "capabilities"): (200, capability)}) + async def fake_session(): + return session + with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object( + dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}" + ): + with self.assertRaisesRegex(RuntimeError, "execution_targets.*upgrade"): + await dispatch.preflight_image_list_workers([{"id": "worker"}]) + async def test_image_list_preflight_fails_closed_for_malformed_capabilities(self): malformed = {"supported": True, "protocol_version": 2, "node_ids": None} session = _FakeSession({("worker", "capabilities"): (200, malformed)}) diff --git a/tests/test_image_list.py b/tests/test_image_list.py index c0c7f0f..c54b96c 100644 --- a/tests/test_image_list.py +++ b/tests/test_image_list.py @@ -202,8 +202,10 @@ class ImageListGraphBoundaryTests(unittest.TestCase): prompt["2"] = {"class_type": "VAEDecode", "inputs": {"samples": ["7", 0]}} prompt["7"] = {"class_type": "KSampler", "inputs": {"model": ["1", 0]}} prompt["4"]["inputs"]["reference"] = ["2", 0] - with self.assertRaisesRegex(ValueError, "bypasses Start.*2"): - validate_image_list_regions(prompt) + for explicit_targets in (False, True): + with self.subTest(explicit_targets=explicit_targets): + with self.assertRaisesRegex(ValueError, "bypasses Start.*2"): + validate_image_list_regions(prompt, explicit_targets=explicit_targets) def test_shared_typed_loader_and_leaf_constant_remain_valid(self): import sys @@ -247,19 +249,25 @@ class ImageListGraphBoundaryTests(unittest.TestCase): def test_rejects_missing_context_pair_nested_region_and_cross_boundary_branch(self): prompt = graph_one_region() prompt["5"]["inputs"]["context"] = ["2", 0] - with self.assertRaisesRegex(ValueError, "typed context output"): - validate_image_list_regions(prompt) + for explicit_targets in (False, True): + with self.subTest(explicit_targets=explicit_targets): + with self.assertRaisesRegex(ValueError, "typed context output"): + validate_image_list_regions(prompt, explicit_targets=explicit_targets) nested = graph_one_region() nested["7"] = start_node(["4", 0]) nested["8"] = {"class_type": "NestedOperation", "inputs": {"image": ["7", IMAGE_OUT]}} nested["5"]["inputs"]["images"] = ["8", 0] nested["9"] = {"class_type": END_CLASS, "inputs": {"images": ["8", 0], "context": ["7", CONTEXT_OUT]}} - with self.assertRaisesRegex(ValueError, "Nested/ambiguous"): - validate_image_list_regions(nested) + for explicit_targets in (False, True): + with self.subTest(explicit_targets=explicit_targets): + with self.assertRaisesRegex(ValueError, "Nested/ambiguous"): + validate_image_list_regions(nested, explicit_targets=explicit_targets) outside = graph_one_region() outside["7"] = {"class_type": "PreviewImage", "inputs": {"images": ["3", IMAGE_OUT]}} - with self.assertRaisesRegex(ValueError, "cross-boundary"): - validate_image_list_regions(outside) + for explicit_targets in (False, True): + with self.subTest(explicit_targets=explicit_targets): + with self.assertRaisesRegex(ValueError, "cross-boundary"): + validate_image_list_regions(outside, explicit_targets=explicit_targets) def test_rejects_duplicate_or_missing_results(self): prompt = graph_one_region() diff --git a/tests/test_image_list_execution_plan.py b/tests/test_image_list_execution_plan.py new file mode 100644 index 0000000..74a69e2 --- /dev/null +++ b/tests/test_image_list_execution_plan.py @@ -0,0 +1,245 @@ +"""CPU-only execution-copy planning; no ComfyUI imports or prompt execution.""" +import copy +import sys +import types + +import pytest + +from image_list.execution_plan import prepare_image_list_execution +from image_list.runtime import ( + CONTEXT_OUT, + END_CLASS, + IMAGE_OUT, + START_CLASS, + build_worker_prompt, + validate_image_list_regions, +) + + +class OutputNode: + OUTPUT_NODE = True + + +class SchemaOutputNode: + OUTPUT_NODE = False # V3's schema, not this fallback, is authoritative. + + @classmethod + def GET_NODE_INFO_V1(cls): + return {"output_node": True, "output": ("IMAGE",)} + + +class ModelLoader: + RETURN_TYPES = ("MODEL", "CLIP", "VAE") + + +@pytest.fixture +def registry(monkeypatch): + classes = { + END_CLASS: SchemaOutputNode, + "SaveImage": OutputNode, + "PreviewImage": OutputNode, + "OutputUtility": OutputNode, + "SchemaUtility": SchemaOutputNode, + "ModelLoader": ModelLoader, + } + monkeypatch.setitem(sys.modules, "nodes", types.SimpleNamespace(NODE_CLASS_MAPPINGS=classes)) + return classes + + +def node(class_type, **inputs): + return {"class_type": class_type, "inputs": inputs} + + +def graph(): + return { + "loader": node("ModelLoader"), + "source": node("ImageSource"), + "start": node(START_CLASS, images=["source", 0]), + "process": node("FrameOperation", image=["start", IMAGE_OUT], model=["loader", 0]), + "end": node(END_CLASS, images=["process", 0], context=["start", CONTEXT_OUT]), + "save": node("SaveImage", images=["end", 0]), + } + + +def test_discards_multistep_previews_but_preserves_other_output_roots(registry): + prompt = graph() + prompt.update({ + "debug-transform": node("DebugTransform", images=["start", IMAGE_OUT]), + "debug-preview": node("PreviewImage", images=["debug-transform", 0]), + "process-preview": node("PreviewImage", images=["process", 0]), + "unrelated-source": node("ImageSource"), + "unrelated-preview": node("PreviewImage", images=["unrelated-source", 0]), + "upstream-preview": node("PreviewImage", images=["source", 0]), + "post-transform": node("FrameOperation", images=["end", 0]), + "post-preview": node("PreviewImage", images=["post-transform", 0]), + }) + before = copy.deepcopy(prompt) + with pytest.raises(ValueError, match="cross-boundary"): + validate_image_list_regions(prompt) + prepared, targets = prepare_image_list_execution(prompt) + assert set(prepared) == set(prompt) - {"debug-transform", "debug-preview", "process-preview"} + assert targets == ["end", "save", "unrelated-preview", "upstream-preview", "post-preview"] + validate_image_list_regions(prepared, explicit_targets=True) + assert prompt == before + + +@pytest.mark.parametrize("class_type", ["OutputUtility", "SchemaUtility"]) +def test_required_output_utility_keeps_complete_dependency_chain(registry, class_type): + prompt = graph() + prompt.update({ + "utility-config": node("Constant", value=7), + "utility-first": node(class_type, images=["process", 0], config=["utility-config", 0]), + "utility-second": node(class_type, images=["utility-first", 0]), + "result": node("FrameOperation", image=["utility-second", 0]), + "discarded-preview": node("PreviewImage", images=["utility-first", 0]), + }) + prompt["end"]["inputs"]["images"] = ["result", 0] + before = copy.deepcopy(prompt) + registered = dict(registry) + prepared, targets = prepare_image_list_execution(prompt) + assert set(prepared) == set(prompt) - {"discarded-preview"} + assert targets == ["end", "save"] + with pytest.raises(ValueError, match="Output node"): + validate_image_list_regions(prepared) + region = validate_image_list_regions(prepared, explicit_targets=True)[0] + worker = build_worker_prompt(prepared, region, "token", "context", 0, "http://master.invalid") + assert set(worker) == set(prepared) - {"source", "save"} + assert worker["utility-second"]["inputs"] == {"images": ["utility-first", 0]} + assert prompt == before + assert registry == registered + assert OutputNode.OUTPUT_NODE is True + assert SchemaOutputNode.OUTPUT_NODE is False + assert SchemaOutputNode.GET_NODE_INFO_V1()["output_node"] is True + + +@pytest.mark.parametrize("source", ["start", "process"]) +@pytest.mark.parametrize("direct", [False, True]) +def test_internal_bypass_into_post_end_computation_is_rejected(registry, source, direct): + prompt = graph() + prompt["post-process"] = node("Composite", images=["end", 0]) + if not direct: + prompt["bypass"] = node("FrameOperation", image=[source, 0]) + source = "bypass" + prompt["post-process"]["inputs"]["reference"] = [source, 0] + prompt["save"]["inputs"]["images"] = ["post-process", 0] + before = copy.deepcopy(prompt) + with pytest.raises(ValueError, match="cross-boundary"): + prepare_image_list_execution(prompt) + assert prompt == before + + +def test_sequential_namespaced_regions_and_intermediate_output_are_preserved(registry): + prompt = graph() + prompt.update({ + "next-start": node(START_CLASS, images=["end", 0]), + "next-process": node("SchemaUtility", images=["next-start", IMAGE_OUT], model=["loader", 0]), + "next-end": node(END_CLASS, images=["next-process", 0], context=["next-start", CONTEXT_OUT]), + "next-save": node("SaveImage", images=["next-end", 0]), + "first-debug": node("PreviewImage", images=["process", 0]), + "second-debug": node("PreviewImage", images=["next-process", 0]), + }) + renamed = {key: f"group:7/{key}" for key in prompt} + for spec in prompt.values(): + for name, value in spec["inputs"].items(): + if isinstance(value, list): + spec["inputs"][name] = [renamed[value[0]], value[1]] + prompt = {renamed[key]: spec for key, spec in prompt.items()} + prepared, targets = prepare_image_list_execution(prompt) + assert set(prepared) == set(prompt) - {renamed["first-debug"], renamed["second-debug"]} + assert targets == [renamed[name] for name in ("end", "save", "next-end", "next-save")] + regions = validate_image_list_regions(prepared, explicit_targets=True) + assert {(r["start_id"], r["end_id"]) for r in regions} == { + (renamed["start"], renamed["end"]), (renamed["next-start"], renamed["next-end"]), + } + second = next(r for r in regions if r["start_id"] == renamed["next-start"]) + worker = build_worker_prompt(prepared, second, "token", "context", 0, "http://master.invalid") + assert set(worker) == {renamed[name] for name in ("loader", "next-start", "next-process", "next-end")} + + +def test_shared_loader_and_its_independent_output_root_are_preserved(registry): + registry["ModelLoader"] = type("OutputLoader", (ModelLoader, OutputNode), {}) + prompt = graph() + prompt["source"]["inputs"]["model"] = ["loader", 0] + prompt["independent-save"] = node("SaveImage", model=["loader", 0]) + prepared, targets = prepare_image_list_execution(prompt) + assert prepared == prompt + assert targets == ["loader", "end", "save", "independent-save"] + region = validate_image_list_regions(prepared, explicit_targets=True)[0] + assert "loader" in region["node_ids"] + assert "loader" not in region["branch_ids"] + + +@pytest.mark.parametrize("unsafe_class", [None, type("ListLoader", (), { + "RETURN_TYPES": ("MODEL",), "OUTPUT_IS_LIST": (True,), +})]) +def test_unknown_or_list_producing_shared_loader_still_fails_closed(registry, unsafe_class): + if unsafe_class is None: + registry.pop("ModelLoader") + else: + registry["ModelLoader"] = unsafe_class + prompt = graph() + prompt["source"]["inputs"]["model"] = ["loader", 0] + prompt["debug"] = node("PreviewImage", images=["process", 0]) + with pytest.raises(ValueError, match="bypasses Start.*loader"): + prepare_image_list_execution(prompt) + + +def test_upstream_per_item_bypass_is_not_hidden_by_preview_pruning(registry): + prompt = graph() + prompt["process"]["inputs"]["reference"] = ["source", 0] + prompt["debug"] = node("PreviewImage", images=["process", 0]) + with pytest.raises(ValueError, match="bypasses Start.*source"): + prepare_image_list_execution(prompt) + + +@pytest.mark.parametrize("nested", [False, True]) +def test_overlapping_or_nested_regions_cannot_be_discarded_as_side_branches(registry, nested): + prompt = graph() + prompt["other-start"] = node(START_CLASS, images=["process" if nested else "source", 0]) + prompt["other-process"] = node("FrameOperation", image=["other-start", IMAGE_OUT], reference=["process", 0]) + prompt["other-end"] = node(END_CLASS, images=["other-process", 0], context=["other-start", CONTEXT_OUT]) + with pytest.raises(ValueError, match="Nested/ambiguous|cross-boundary|ambiguous processing"): + prepare_image_list_execution(prompt) + + +@pytest.mark.parametrize("invalid", ["missing-end", "wrong-context", "missing-images", "unrelated-images"]) +def test_pairing_and_result_validation_cannot_be_pruned_away(registry, invalid): + prompt = graph() + if invalid == "missing-end": + del prompt["end"] + elif invalid == "wrong-context": + prompt["end"]["inputs"]["context"] = ["source", 0] + elif invalid == "missing-images": + del prompt["end"]["inputs"]["images"] + else: + prompt["end"]["inputs"]["images"] = ["source", 0] + before = copy.deepcopy(prompt) + with pytest.raises(ValueError): + prepare_image_list_execution(prompt) + assert prompt == before + + +def test_execution_copy_is_deep_and_repeated_preparation_is_stable(registry): + prompt = graph() + prompt["process"]["_meta"] = {"title": "Original", "tags": ["saved"]} + before = copy.deepcopy(prompt) + prepared, targets = prepare_image_list_execution(prompt) + prepared_again, targets_again = prepare_image_list_execution(prepared) + assert prepared_again == prepared + assert targets_again == targets + prepared["start"]["inputs"]["images"][0] = "changed" + prepared["process"]["_meta"]["tags"].append("execution-only") + assert prompt == before + assert prepared_again == before + + +def test_v3_false_output_flag_overrides_legacy_true(registry): + registry["NotOutput"] = type("NotOutput", (), { + "OUTPUT_NODE": True, + "GET_NODE_INFO_V1": classmethod(lambda cls: {"output_node": False}), + }) + prompt = graph() + prompt["not-output"] = node("NotOutput", images=["end", 0]) + prepared, targets = prepare_image_list_execution(prompt) + assert prepared == prompt + assert targets == ["end", "save"] diff --git a/tests/test_image_list_integration_gaps.py b/tests/test_image_list_integration_gaps.py index b11ea66..1d261f7 100644 --- a/tests/test_image_list_integration_gaps.py +++ b/tests/test_image_list_integration_gaps.py @@ -35,6 +35,15 @@ def load_routes(): return module, routes +class CapabilityTests(unittest.IsolatedAsyncioTestCase): + async def test_capability_advertises_target_control_without_tensor_protocol_bump(self): + module, _routes = load_routes() + response = await module.image_list_capabilities_endpoint(None) + payload = json.loads(response.text) + self.assertIs(payload['execution_targets'], True) + self.assertEqual(payload['protocol_version'], 2) + + class TransportTests(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self): self.module, routes = load_routes() diff --git a/tests/test_image_list_lifecycle.py b/tests/test_image_list_lifecycle.py index 998a617..cb34c66 100644 --- a/tests/test_image_list_lifecycle.py +++ b/tests/test_image_list_lifecycle.py @@ -11,6 +11,83 @@ from aiohttp import web from image_list import lifecycle, runtime +class ExecutionTargetProofTests(unittest.TestCase): + def setUp(self): + from test_image_list_master_runtime import graph + self.prompt = graph() + self.registry = SimpleNamespace(NODE_CLASS_MAPPINGS={ + 'Operation': type('Flagged', (), {'OUTPUT_NODE': True})}) + self.modules = patch.dict('sys.modules', {'nodes': self.registry}) + self.modules.start() + self.addCleanup(self.modules.stop) + + def server(self, roots, *, prompt=None, running=True): + item = (0, 'owned', self.prompt if prompt is None else prompt, {}, roots, {}) + return SimpleNamespace(prompt_queue=SimpleNamespace( + get_current_queue_volatile=lambda: ([item], []) if running else ([], [item]))) + + def test_outer_roots_in_running_queue_are_authoritative(self): + self.assertTrue(lifecycle.has_safe_execution_targets(self.server(['end']), 'owned', self.prompt)) + for roots in (['op', 'end'], [], ['end', 'end'], None, [True]): + with self.subTest(roots=roots): + self.assertFalse(lifecycle.has_safe_execution_targets(self.server(roots), 'owned', self.prompt)) + self.assertFalse(lifecycle.has_safe_execution_targets(self.server(['end']), 'other', self.prompt)) + self.assertFalse(lifecycle.has_safe_execution_targets(self.server(['end'], running=False), 'owned', self.prompt)) + self.assertFalse(lifecycle.has_safe_execution_targets(self.server(['end'], prompt={}), 'owned', self.prompt)) + + def test_prompt_marker_is_not_an_execution_guarantee(self): + self.prompt['start']['inputs']['image_list_explicit_targets'] = True + self.assertFalse(lifecycle.has_safe_execution_targets(self.server(['op', 'end']), 'owned', self.prompt)) + with self.assertRaisesRegex(ValueError, 'must be placed after'): + runtime.create_master_job('start', 'end', self.prompt, [], [], '', master_participates=True) + self.assertEqual(runtime._MASTER_JOBS, {}) + + def test_start_node_uses_running_queue_proof(self): + import importlib.util + import sys + from pathlib import Path + from types import ModuleType + import torch + + root = Path(__file__).resolve().parents[1] + package = '_image_list_start_targets' + modules = {} + for suffix in ('', '.nodes', '.image_list', '.utils'): + module = ModuleType(package + suffix) + module.__path__ = [str(root / suffix.lstrip('.'))] + modules[package + suffix] = module + modules[package + '.image_list.runtime'] = runtime + modules[package + '.image_list.lifecycle'] = lifecycle + modules[package + '.utils.async_helpers'] = SimpleNamespace(run_async_in_server_loop=None) + io = SimpleNamespace(ComfyNode=type('ComfyNode', (), {}), NodeOutput=lambda *args: args, + **{name: object() for name in ('Image', 'Mask', 'String', 'Int', 'Float', 'Boolean')}) + modules['comfy_api.v0_0_2'] = SimpleNamespace(io=io) + modules['comfy_execution.utils'] = SimpleNamespace( + get_executing_context=lambda: SimpleNamespace(prompt_id='owned')) + server = SimpleNamespace(PromptServer=SimpleNamespace(instance=self.server(['end']))) + modules['server'] = server + spec = importlib.util.spec_from_file_location(package + '.nodes.image_list', root / 'nodes/image_list.py') + node_module = importlib.util.module_from_spec(spec) + with patch.dict(sys.modules, modules): + spec.loader.exec_module(node_module) + start = node_module.DistributedImageListStart + start.hidden = SimpleNamespace(prompt=self.prompt, unique_id='start') + result = start.execute(images=torch.zeros((1, 2, 2, 3)), image_list_master_participates=True) + context = result[3][0] + self.assertIsNotNone(runtime.get_master_job(context)) + runtime.remove_master_job(context) + server.PromptServer.instance = self.server(['op', 'end']) + with self.assertRaisesRegex(ValueError, 'must be placed after'): + start.execute(images=torch.zeros((1, 2, 2, 3)), image_list_master_participates=True) + self.assertEqual(runtime._MASTER_JOBS, {}) + + def test_owned_target_controlled_start_can_create_job(self): + context = runtime.create_master_job('start', 'end', self.prompt, [], [], '', + master_participates=True, explicit_targets=True) + self.addCleanup(runtime.remove_master_job, context) + self.assertIsNotNone(runtime.get_master_job(context)) + + class ImageListLifecycleTests(unittest.IsolatedAsyncioTestCase): def setUp(self): with runtime._LOCK: diff --git a/tests/test_image_list_master_selection.py b/tests/test_image_list_master_selection.py index 58e2e5a..ef56b61 100644 --- a/tests/test_image_list_master_selection.py +++ b/tests/test_image_list_master_selection.py @@ -34,6 +34,23 @@ class MasterSelectionTests(unittest.IsolatedAsyncioTestCase): def prepared_inputs(self): return self.q.queue_prompt_payload.call_args.args[0]["2"]["inputs"] + async def test_prepared_master_prunes_side_outputs_without_mutating_input(self): + import copy + self.prompt['op'] = {'class_type': 'Flagged', 'inputs': {'image': ['2', 0]}} + self.prompt['3']['inputs']['images'] = ['op', 0] + self.prompt['preview'] = {'class_type': 'Flagged', 'inputs': {'image': ['op', 0]}} + self.prompt['save'] = {'class_type': 'Flagged', 'inputs': {'image': ['3', 0]}} + original = copy.deepcopy(self.prompt) + registry = type('Registry', (), {'NODE_CLASS_MAPPINGS': { + 'Flagged': type('Flagged', (), {'OUTPUT_NODE': True})}}) + with patch.dict(sys.modules, {'nodes': registry}): + await self.queue() + queued = self.q.queue_prompt_payload.call_args + self.assertEqual(queued.kwargs['partial_execution_targets'], ['3', 'save']) + self.assertIn('op', queued.args[0]) + self.assertNotIn('preview', queued.args[0]) + self.assertEqual(self.prompt, original) + async def test_default_includes_master_and_direct_workers(self): result = await self.queue() self.assertIs(self.prepared_inputs()["image_list_master_participates"], True) diff --git a/tests/test_image_list_runtime_boundaries.py b/tests/test_image_list_runtime_boundaries.py index 341a97f..2129e1e 100644 --- a/tests/test_image_list_runtime_boundaries.py +++ b/tests/test_image_list_runtime_boundaries.py @@ -84,7 +84,7 @@ class RuntimeBoundaryTests(unittest.IsolatedAsyncioTestCase): # Other legacy test modules replace aiohttp during collection. This # controlled provider fixture needs only its explicit timeout/error seam. sys.modules['aiohttp'] = types.SimpleNamespace( - ClientTimeout=lambda **kwargs: kwargs, ClientError=ConnectionError) + ClientTimeout=lambda **kwargs: kwargs, ClientError=ConnectionError, ClientConnectionError=ConnectionError) comfy = types.ModuleType('comfy') comfy.model_management = types.SimpleNamespace(processing_interrupted=lambda: self.interrupted) sys.modules['comfy'] = comfy @@ -94,6 +94,7 @@ class RuntimeBoundaryTests(unittest.IsolatedAsyncioTestCase): prompt['7']['inputs']['image'] = 'shared/ref.png' async def dispatch_prompt(self, worker, prompt, meta, **kwargs): + self.assertEqual(kwargs.get('partial_execution_targets'), ['5']) if self.reject: raise RuntimeError('prompt rejected') self.prompts.append((json.loads(json.dumps(prompt)), meta)) @@ -136,6 +137,23 @@ class RuntimeBoundaryTests(unittest.IsolatedAsyncioTestCase): ['worker'], 'http://public-master', self.metadata) return self.context + async def test_runtime_preflight_rejects_old_target_capability_before_upload(self): + from test_dispatch_selection import _load_dispatch_module, _FakeSession, _valid_capability + real_dispatch = _load_dispatch_module() + capability = _valid_capability() + capability.pop('execution_targets') + session = _FakeSession({('worker', 'capabilities'): (200, capability)}) + module = sys.modules['_image_list_boundaries.api.orchestration.dispatch'] + module.preflight_image_list_workers = real_dispatch.preflight_image_list_workers + with patch.object(real_dispatch, 'get_client_session', AsyncMock(return_value=session)), patch.object( + real_dispatch, 'build_worker_url', side_effect=lambda worker, path: f"http://{worker['id']}{path}" + ): + with self.assertRaisesRegex(RuntimeError, 'execution_targets.*upgrade'): + await self.rt.run_image_list_job(self.create(1)) + self.assertFalse(self.uploads) + self.dispatch.assert_not_awaited() + self.assertFalse(self.rt._MASTER_JOBS) + async def test_each_assignment_reuses_media_paths_with_fresh_tokens_and_metadata(self): context = self.create() await self.rt.run_image_list_job(context) diff --git a/tests/test_v3_migration.py b/tests/test_v3_migration.py index 626437d..452163c 100644 --- a/tests/test_v3_migration.py +++ b/tests/test_v3_migration.py @@ -1,4 +1,5 @@ """Real-framework acceptance; opt in with COMFYUI_SOURCE_ROOT.""" +import json import os from pathlib import Path import subprocess @@ -35,3 +36,15 @@ def test_image_list_network_runtime(): assert 'IMAGE_LIST_NETWORK_OK' in result.stdout assert 'IMAGE_LIST_CANCEL_OK' in result.stdout assert 'IMAGE_LIST_MASTER_OK' in result.stdout + assert 'IMAGE_LIST_PREFLIGHT_OK' in result.stdout + assert 'IMAGE_LIST_SEQUENTIAL_OK' in result.stdout + assert 'IMAGE_LIST_FAILURE_OK' in result.stdout + reports = [json.loads(line.removeprefix('IMAGE_LIST_TARGETS_OK ')) + for line in result.stdout.splitlines() if line.startswith('IMAGE_LIST_TARGETS_OK ')] + assert len(reports) == 1, result.stdout + assert reports[0] == { + 'modes': ['delegate-only', 'mixed', 'master-only'], + 'outer_targets': ['4', '5'], 'worker_targets': ['4'], + 'required_utility_once_per_item': True, 'preview_discarded': True, + 'saved_prompt_unchanged': True, 'capture_once_per_outer_prompt': True, + } diff --git a/utils/async_helpers.py b/utils/async_helpers.py index eb34d15..01207b1 100644 --- a/utils/async_helpers.py +++ b/utils/async_helpers.py @@ -105,24 +105,63 @@ class PromptValidationError(RuntimeError): super().__init__(f"Invalid prompt: {merged}") +def validate_execution_targets(prompt, targets): + """Reject malformed explicit roots instead of falling back to all outputs.""" + if (not isinstance(targets, list) or not targets + or any(not isinstance(node_id, str) or not node_id or node_id not in prompt + for node_id in targets) + or len(set(targets)) != len(targets)): + raise PromptValidationError({ + "type": "invalid_partial_execution_targets", + "message": "partial_execution_targets must be a non-empty list of unique existing node IDs", + "details": "", "extra_info": {}, + }) + return list(targets) + + async def queue_prompt_payload( prompt_obj, workflow_meta=None, client_id=None, include_queue_metadata=False, + *, + partial_execution_targets=None, ): - """Validate and queue a prompt via ComfyUI's prompt queue.""" + """Validate and queue a prompt, preserving any explicitly requested roots.""" + targets = (validate_execution_targets(prompt_obj, partial_execution_targets) + if partial_execution_targets is not None else None) payload = {"prompt": prompt_obj} + if targets is not None: + payload["partial_execution_targets"] = list(targets) payload = prompt_server.trigger_on_prompt(payload) prompt = payload["prompt"] + if targets is not None: + # Hooks may replace the envelope or mutate its target list. The caller's + # requested roots remain authoritative, but must still exist afterwards. + validate_execution_targets(prompt, targets) prompt_id = str(uuid.uuid4()) - valid = await execution.validate_prompt(prompt_id, prompt, None) + valid = await execution.validate_prompt(prompt_id, prompt, targets) if not valid[0]: error_payload = valid[1] if len(valid) > 1 else "Prompt outputs failed validation" node_errors = valid[3] if len(valid) > 3 else {} raise PromptValidationError(error_payload, node_errors) + if targets is not None: + node_errors = valid[3] if len(valid) > 3 else {} + target_errors = { + node_id: entry for node_id, entry in node_errors.items() + if node_id in targets or (isinstance(entry, dict) + and set(entry.get("dependent_outputs", [])) & set(targets)) + } + if set(valid[2]) != set(targets) or target_errors: + raise PromptValidationError({ + "type": "partial_execution_targets_failed_validation", + "message": "Not all requested execution targets passed validation", + "details": "Refusing to queue a subset or additional output roots", + "extra_info": {"requested_targets": targets}, + }, node_errors) + extra_data = {"create_time": int(time.time() * 1000)} if workflow_meta: extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta