Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
3699477f75 | ||
|
|
1fb4d4a284 | ||
|
|
4e823397b7 |
@@ -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",
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -39,6 +39,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
|
||||
@@ -327,6 +328,16 @@ 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:
|
||||
if any(isinstance(node, dict) and node.get("class_type") in {
|
||||
"DistributedCollector", "DistributedMediaCollector", "UltimateSDUpscaleDistributed"
|
||||
} for node in prompt_obj.values()):
|
||||
raise ValueError("Distributed Image List cannot be mixed with other distributed collectors")
|
||||
prompt_obj, outer_targets = prepare_image_list_execution(prompt_obj)
|
||||
target_kwargs["partial_execution_targets"] = outer_targets
|
||||
prompt_index = PromptIndex(prompt_obj)
|
||||
typed_media_mode = bool(prompt_index.nodes_for_class("DistributedMediaCollector"))
|
||||
if typed_media_mode:
|
||||
@@ -349,7 +360,7 @@ async def orchestrate_distributed_execution(
|
||||
"Distributed Image List v1 cannot be mixed with Distributed Collector, "
|
||||
"DistributedMediaCollector or Ultimate SD Upscale nodes in one prompt"
|
||||
)
|
||||
validate_image_list_regions(prompt_obj)
|
||||
validate_image_list_regions(prompt_obj, explicit_targets=True)
|
||||
configured_serverless = [
|
||||
str(worker.get("id")) for worker in config.get("workers", [])
|
||||
if worker.get("enabled", False)
|
||||
@@ -487,6 +498,7 @@ async def orchestrate_distributed_execution(
|
||||
workflow_meta,
|
||||
client_id,
|
||||
include_queue_metadata=True,
|
||||
**target_kwargs,
|
||||
)
|
||||
return (
|
||||
queue_result["prompt_id"],
|
||||
@@ -594,6 +606,7 @@ async def orchestrate_distributed_execution(
|
||||
workflow_meta,
|
||||
client_id,
|
||||
include_queue_metadata=True,
|
||||
**target_kwargs,
|
||||
)
|
||||
except Exception:
|
||||
if serverless_jobs:
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
|
||||
+38
-15
@@ -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")
|
||||
|
||||
+9
-3
@@ -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],
|
||||
|
||||
@@ -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)), \
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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)})
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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"]
|
||||
@@ -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()
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
+41
-2
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user