Compare commits

...
Author SHA1 Message Date
Robert Wojciechowski 6db5692d78 fix: preserve native image-list processing cache reuse 2026-10-14 03:15:49 +00:00
Robert Wojciechowski 591ffc7396 fix: preserve image-list worker selection across queued runs (#103) 2026-10-13 19:52:29 -07:00
Robert Wojciechowski 5f9a2f0fc6 fix: carry image-list values through End context
Preserve source values without extra wires, collect explicit per-item overrides, and expose ordered bundles through Get Value. Keep legacy sockets and no-override transport compatible.
2026-10-14 01:51:13 +00:00
Robert Wojciechowski c30036ae31 feat: preserve required output utilities in image-list execution 2026-10-13 23:26:25 +00:00
33 changed files with 3215 additions and 170 deletions
+6 -3
View File
@@ -35,8 +35,10 @@ async def _json_body(request):
async def image_list_capabilities_endpoint(_request):
return web.json_response({
"supported": True,
"execution_targets": True,
"value_results": True,
"protocol_version": PROTOCOL_VERSION,
"node_ids": ["DistributedImageListStart", "DistributedImageListEnd"],
"node_ids": ["DistributedImageListStart", "DistributedImageListEnd", "DistributedImageListGetValue"],
"transport": "typed-float-tensor-json-v1",
})
@@ -51,8 +53,9 @@ async def image_list_item_endpoint(request):
if not isinstance(context_id, str) or not context_id:
raise ValueError("context_id must be a non-empty string")
item = data.get("item")
store_worker_item(data.get("token"), item)
return web.json_response({"status": "stored", "item_index": item["index"]})
identity = store_worker_item(data.get("token"), item)
return web.json_response({"status": "stored", "item_index": item["index"],
"item_identity": identity})
except (ValueError, TypeError) as exc:
return web.json_response({"error": str(exc)}, status=400)
except Exception as exc:
+22 -1
View File
@@ -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,
)
@@ -179,7 +188,7 @@ async def dispatch_worker_prompt(
return prompt_id if return_prompt_id else None
async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0):
async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0, *, require_value_results=False):
"""Require protocol/node support and report which direct workers are idle.
This performs read-only GETs only; it never submits a prompt or contacts a
@@ -227,6 +236,18 @@ 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"
)
if require_value_results and capability.get("value_results") is not True:
raise RuntimeError(
f"Image-list worker {worker_id} lacks explicit value_results support; "
"upgrade that worker before using End overrides; 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:
+20 -8
View File
@@ -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")
@@ -338,11 +343,16 @@ async def orchestrate_distributed_execution(
raise RuntimeError(
"No selected direct worker is reachable for Distributed Image List in orchestrator-only mode"
)
idle_ids = (await preflight_image_list_workers(active_workers, require_idle=delegate_master)
if active_workers else {})
active_workers = [worker for worker in active_workers if worker["id"] in idle_ids]
if not active_workers and delegate_master:
raise RuntimeError("No feature-compatible direct worker is idle in orchestrator-only mode")
from ..image_list.runtime import END_CLASS, end_override_names
requires_values = any(end_override_names(node.get("inputs", {}))
for node in prompt_obj.values() if node.get("class_type") == END_CLASS)
value_preflight = {"require_value_results": True} if requires_values else {}
# Persist candidates, not this submission-time idle snapshot: a preceding
# master prompt may still be using them. The image-list runtime rechecks
# capability and availability before assigning the first item.
if active_workers:
await preflight_image_list_workers(active_workers, require_idle=False,
**value_preflight)
else:
delegate_master = selected_delegate_master
@@ -409,6 +419,7 @@ async def orchestrate_distributed_execution(
workflow_meta,
client_id,
include_queue_metadata=True,
**target_kwargs,
)
return (
queue_result["prompt_id"],
@@ -506,6 +517,7 @@ async def orchestrate_distributed_execution(
workflow_meta,
client_id,
include_queue_metadata=True,
**target_kwargs,
)
except Exception:
if serverless_jobs:
+13
View File
@@ -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",
+42 -4
View File
@@ -52,7 +52,31 @@ Start outputs, in order: `image`, `mask`, `item_index` (zero-based), `context`,
Values that cannot be transported exactly (models, latents, conditioning and other objects) are rejected. Load models inside the region from a loader connected directly to the processing nodes, as described below.
End accepts `images`, optional `masks` / `texts`, and `context`. It returns `image_list`, `masks`, and `texts` in original item order. Each job must produce exactly one image, at most one mask, and at most one text value. Unconnected optional outputs contain `None` entries.
End accepts `images`, optional `masks` / `texts`, and `context`. Its first three outputs remain `image_list`, `masks`, and `texts` in original item order. Each job must produce exactly one image, at most one mask, and at most one text value. Unconnected optional mask/text outputs contain `None` entries.
### Keep values through End without extra wires
Start's existing **context** connection now also lets End recover each item's original `value_1` … `value_10`. You do not need to wire those values individually into End. Context remains a lightweight execution reference, not a general-purpose data pipe.
End appends a **values** output (`DISTRIBUTED_IMAGE_LIST_VALUES`): one value bundle per image, in the same source order. Connect it to **Distributed Image List Get Value**, then choose `value_name` to extract the field you need. The getter preserves its type and ComfyUI maps it across the bundles, producing a list aligned with End's images. Missing fields fail clearly instead of producing empty placeholders.
For example, to preserve filenames:
```text
filenames → Start.value_1
Start.context ─────────────→ End.context
Start.image → processing → End.images
├─ image_list → Save.images
└─ values → Get Value [value_1] → Save.filename_prefix
```
**New or changed per-item values:** connect the calculated result to End's optional `value_N` input. It replaces the original field with that number, or adds it if Start did not supply it. Every connected override must produce exactly one supported value for each item. A tensor batch remains one value, just as it does at Start. Image, mask, string, integer, float and boolean values are supported; arbitrary objects/models/latents and non-finite numbers are not.
End shows override inputs progressively. Use **Add value override** to expose a later slot without connecting the earlier ones. Disconnecting an override never renumbers the other End fields: `value_3` always means `value_3`. Unlike Start's native Autogrow inputs, End overrides must keep these identities stable. Existing image/mask/text/context sockets and output indices are unchanged. If the frontend extension is unavailable, Python still declares all optional override inputs and the values output.
Original values are retained on the master; workers only return explicitly connected overrides, avoiding a redundant return trip for unchanged reference images or other values. Overrides use the existing exact value transport and size limits. Selected workers must advertise `value_results: true` when overrides are requested, checked before execution and again at dispatch. Original-value passthrough alone does not require that additional capability.
Values leaving the region must come from **End**, not directly from Start. To chain regions, extract an End value and connect it to the next Start's corresponding input. Values remain available downstream after End cleans up the execution context; they do not depend on a later lookup of expired job state.
Models are loaded and cached on each processing participant, not serialized between machines. Connect loader outputs directly to the processing nodes. The master needs the region's model files when participating, as does each selected worker. The master still validates the prompt: it needs the relevant node definitions and any model-name/file availability required by their validation (including combo lists). This is not a model-free-master guarantee.
@@ -62,15 +86,29 @@ End's optional `texts` is a connection-only socket. For workflows saved before t
File-based dependencies keep ComfyUI's `[output]` / `[temp]` root annotations and are read on the master. Their worker copies use separate content-addressed paths under `input/distributed-media/`; each worker's rewritten prompt is reused across items without reusing item tokens.
## Native cache reuse
Repeated items can reuse **ComfyUI's native processing cache** on the same participant. Start uses a deterministic identity of the exact image, optional mask, all supplied value slots and source index; tensor dtype, shape and bytes are preserved. Fresh job IDs, assignment tokens and callback addresses live at End, so they no longer invalidate otherwise unchanged upstream processing. End still validates the current assignment and delivers results in source order on every run.
- Processing settings, seeds, links and node fingerprints still participate in ComfyUI's normal signatures. Changed content invalidates reuse; nodes that declare themselves noncacheable still run.
- The master's private executor retains one native LRU output cache between successful slots, capped at **128 entries when retained**. An active graph may temporarily exceed that target; oversized generations are not retained. Failed or cancelled local slots do not seed reuse, and `--cache-none` disables the local output cache. Direct workers use their own configured native caches.
- Cache hits are best-effort, not guaranteed: eviction, restarts, changed ordering/index, changed settings and assignment to a different worker can cause misses. There is no worker affinity, disk cache or separate completed-result store.
- Reuse retains tensor/model references and can increase RAM/VRAM use; the entry cap is **not a byte limit**. Exact identity calculation also adds CPU hashing and, for GPU-resident inputs, a CPU copy.
- Remote items are still uploaded. Each upload returns a worker-computed identity; if a compatible older worker does not return a valid identity, that assignment uses the legacy graph without stable-input reuse. A worker downgrade between items also falls back safely. Assignment expiry and current-item validation apply even when processing hits the cache.
- Keep side-effect nodes outside the region when they must run on every submission. A cached processing node follows ordinary ComfyUI semantics and does not repeat its save/preview side effects; End's delivery and downstream saving still run.
## Scope and failure behavior
- This version supports the master and direct workers, not Runpod serverless pools. Unsupported pools are excluded without contacting the provider for this job.
- Worker eligibility is determined before execution and checked again before dispatch. Busy or transiently unreachable workers are excluded; healthy candidates can continue. Malformed capability/queue responses and incompatible versions fail explicitly. Workers joining later are not added to an already-running region.
- Queue submission checks selected direct workers for reachability and feature compatibility, but does not freeze the worker list to that moment's idle snapshot. A selected worker busy with an earlier prompt remains a candidate for the queued run. Immediately before a region starts assigning images, its saved candidates are checked again: workers still busy or transiently unreachable are excluded; healthy idle candidates can continue. Malformed capability/queue responses and incompatible versions fail explicitly. Workers excluded by the initial reachability probe are not retained, and newly selected workers are not added to an already-queued run. Workers becoming free after the region starts are not added mid-region.
- **Orchestrator Only** requires a selected, reachable, feature-compatible direct worker at submission, but can queue while all selected workers are busy. At region start it requires at least one idle compatible direct worker and fails if none are available; it never falls back to master processing. The queue response's `worker_count` counts saved remote candidates, not guaranteed execution-time participants.
- Omitted worker IDs select enabled direct workers. Explicit IDs select the named configured direct workers, including disabled ones; an empty list stays empty, unknown IDs fail, and serverless workers are excluded. Local workers use the master's loopback callback; remote workers use its configured public address.
- 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 +127,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
+71
View File
@@ -0,0 +1,71 @@
"""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,
END_RESULT_INPUTS,
_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, END_RESULT_INPUTS)
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
+32
View File
@@ -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)
+40 -1
View File
@@ -24,6 +24,34 @@ from .runtime import ImageListCancelledError
_EXECUTION_LOCK = threading.Lock()
_CACHE_LOCK = threading.Lock()
# One process-local native output cache, not a second completed-result cache.
# Core's LRU can temporarily exceed its target for the active graph; oversized
# generations are discarded rather than retained after an item/run finishes.
NATIVE_CACHE_MAX_ENTRIES = 128
_RETAINED_OUTPUTS = None
def _take_outputs(execution):
global _RETAINED_OUTPUTS
with _CACHE_LOCK:
outputs, _RETAINED_OUTPUTS = _RETAINED_OUTPUTS, None
if getattr(execution.args, "cache_none", False):
return execution.NullCache()
return outputs if outputs is not None else execution.LRUCache(
execution.CacheKeySetInputSignature, max_size=NATIVE_CACHE_MAX_ENTRIES)
def _retain_outputs(outputs):
global _RETAINED_OUTPUTS
# Only the bounded native LRU is eligible. NullCache and test doubles are not.
if getattr(outputs, "max_size", None) != NATIVE_CACHE_MAX_ENTRIES:
return
outputs.clean_unused()
if len(outputs.cache) > NATIVE_CACHE_MAX_ENTRIES:
return
with _CACHE_LOCK:
_RETAINED_OUTPUTS = outputs
class LocalExecutionSlot:
@@ -32,6 +60,7 @@ class LocalExecutionSlot:
self._stop = threading.Event()
self._future = None
self._caches = None
self._failed = False
self._model_callback = model_callback
self._closed = False
self._loop = None
@@ -55,6 +84,10 @@ class LocalExecutionSlot:
self._loop.call_soon_threadsafe(self._task.cancel)
async def close(self):
if self._closed:
return
# A stop requested before close indicates cancellation, not normal drain.
cancelled = self._stop.is_set()
self.request_stop()
try:
if self._future is not None:
@@ -72,6 +105,8 @@ class LocalExecutionSlot:
waiter.exception()
finally:
self._pool.shutdown(wait=True, cancel_futures=True)
if not cancelled and not self._failed and self._caches is not None:
_retain_outputs(self._caches.outputs)
self._caches = None
self._closed = True
@@ -84,7 +119,11 @@ class LocalExecutionSlot:
try:
asyncio.run(self._execute(prompt, prompt_id, end_id, workflow_meta))
except asyncio.CancelledError as exc:
self._failed = True
raise ImageListCancelledError("Local image-list execution cancelled") from exc
except BaseException:
self._failed = True
raise
finally:
with self._state_lock:
self._loop = self._task = None
@@ -125,7 +164,7 @@ class LocalExecutionSlot:
try:
with torch.inference_mode():
if self._caches is None:
outputs = execution.HierarchicalCache(execution.CacheKeySetInputSignature)
outputs = _take_outputs(execution)
objects = execution.HierarchicalCache(execution.CacheKeySetID)
self._caches = SimpleNamespace(outputs=outputs, objects=objects, all=[outputs, objects])
caches = self._caches
+280 -45
View File
@@ -7,9 +7,11 @@ from __future__ import annotations
import asyncio
import base64
import hashlib
import json
import math
import secrets
import struct
import threading
import time
from dataclasses import dataclass, field
@@ -29,6 +31,7 @@ VALUE_OUT_BASE = 4
MAX_VALUE_SLOTS = 10
VALUE_NAMES = tuple(f"value_{i}" for i in range(1, MAX_VALUE_SLOTS + 1))
VALUES_INPUT = "values"
END_RESULT_INPUTS = ("images", "masks", "texts") + VALUE_NAMES
MAX_PAYLOAD_BYTES = 50 * 1024 * 1024
MAX_REQUEST_BYTES = MAX_PAYLOAD_BYTES * 2 + 1024 * 1024
MAX_ITEMS = 4096
@@ -179,6 +182,50 @@ def _normalize_values(values: Any, count: int) -> dict[str, list[Any]]:
return aligned
def end_override_names(inputs: dict[str, Any]) -> tuple[str, ...]:
"""Fixed keys stay stable even when an intermediate socket disconnects."""
for name in inputs:
if (name.startswith("value_") or name.startswith("values.")) and name not in VALUE_NAMES:
raise ValueError(f"Unknown End value slot {name!r}")
return tuple(name for name in VALUE_NAMES if name in inputs)
def _validate_result_values(values: Any) -> dict[str, Any]:
"""Validate flat override values; tensors are values, never frame lists."""
if not isinstance(values, dict):
raise ValueError("End result values must be a mapping of value_N slots")
for name, value in values.items():
if name not in VALUE_NAMES:
raise ValueError(f"Unknown End value slot {name!r}")
kind = value_kind(value, name)
if kind == "TENSOR":
_tensor_dtype_name(value)
if not _torch().isfinite(value).all().item():
raise ValueError(f"End {name} tensor values must be finite")
return values
def normalize_value_overrides(values: Any) -> dict[str, Any]:
if not isinstance(values, dict):
raise ValueError("End overrides must be a mapping of value_N slots")
normalized = {}
for name, raw in values.items():
items = list(raw) if isinstance(raw, (list, tuple)) else [raw]
if len(items) != 1:
raise ValueError(f"Worker End override {name} expects exactly one value, received {len(items)}")
normalized[name] = items[0]
return _validate_result_values(normalized)
def _checked_overrides(result: dict[str, Any], expected: tuple[str, ...]) -> dict[str, Any]:
values = _validate_result_values(result.get("values", {}))
missing, extra = set(expected) - values.keys(), values.keys() - set(expected)
if missing or extra:
raise ValueError(f"End result overrides do not match connected slots: "
f"missing {sorted(missing)}, unexpected {sorted(extra)}")
return values
def normalize_image_list_inputs(images: Any, mask: Any = None, values: Any = None) -> list[dict[str, Any]]:
"""Split images into items and pair masks/generic values by the 1-or-N rule."""
frames = _image_frames(images, "images")
@@ -197,6 +244,64 @@ def normalize_image_list_inputs(images: Any, mask: Any = None, values: Any = Non
]
def is_valid_item_identity(identity: Any) -> bool:
return (isinstance(identity, str) and len(identity) == 71
and identity.startswith("sha256:")
and all(char in "0123456789abcdef" for char in identity[7:]))
def item_content_identity(item: dict[str, Any]) -> str:
"""Hash normalized content, not transport bytes or dispatch metadata.
Length-framed, typed fields distinguish scalars, absent masks, slot names,
shapes and exact tensor bits (including bfloat16). CPU copying changes neither
dtype nor values; local items deliberately have no HTTP transport-size limit.
"""
digest = hashlib.sha256(b"distributed-image-list-item-v1\0")
def field(tag, raw):
digest.update(tag.encode("ascii") + b"\0")
digest.update(len(raw).to_bytes(8, "big"))
digest.update(raw)
def value(raw, name):
if raw is None:
field("NONE", b"")
return
kind = value_kind(raw, name)
if kind == "TENSOR":
cpu = raw.detach().to("cpu").contiguous()
field("DTYPE", str(cpu.dtype).encode("ascii"))
field("SHAPE", json.dumps(list(cpu.shape), separators=(",", ":")).encode("ascii"))
field("TENSOR", memoryview(cpu.view(_torch().uint8).numpy()).cast("B"))
elif kind == "FLOAT":
field(kind, struct.pack("!d", raw))
elif kind == "BOOLEAN":
field(kind, b"1" if raw else b"0")
elif kind == "INT":
field(kind, str(raw).encode("ascii"))
else:
field(kind, raw.encode("utf-8", errors="surrogatepass"))
index = item["index"]
if isinstance(index, bool) or not isinstance(index, int) or index < 0:
raise ValueError("item.index must be a non-negative integer")
field("INDEX", str(index).encode("ascii"))
field("IMAGE", b"")
value(item["image"], "image")
field("MASK", b"")
value(item.get("mask"), "mask")
values = item.get("values") or {}
if not isinstance(values, dict) or any(name not in VALUE_NAMES for name in values):
raise ValueError("item.values must be a mapping of value_N slots")
# Canonical socket order, not the producer's mapping insertion order.
for name in VALUE_NAMES:
if name in values:
field("SLOT", name.encode("ascii"))
value(values[name], name)
return "sha256:" + digest.hexdigest()
def _tensor_dtype_name(tensor: Any) -> str:
name = str(tensor.dtype).removeprefix("torch.")
if name not in _ALLOWED_FLOAT_DTYPES:
@@ -346,26 +451,64 @@ def _sweep_worker_items() -> None:
del _WORKER_ITEMS[token]
def store_worker_item(token: str, item: dict[str, Any]) -> None:
def store_worker_item(token: str, item: dict[str, Any]) -> str:
if not isinstance(token, str) or len(token) < 32 or len(token) > 128:
raise ValueError("token must be a 32-128 character capability")
clean = decode_item(item)
# Never trust a master-supplied digest; hash the decoded worker-owned item.
identity = item_content_identity(clean)
_sweep_worker_items()
with _LOCK:
if token in _WORKER_ITEMS:
raise ValueError("duplicate item token")
_WORKER_ITEMS[token] = (time.monotonic() + _WORKER_ITEM_TTL_SECONDS, clean)
return identity
def consume_worker_item(token: str) -> dict[str, Any]:
def _pop_worker_item(token: str) -> tuple[float, dict[str, Any]]:
if not isinstance(token, str) or not token:
raise ValueError("worker item token is required")
_sweep_worker_items()
with _LOCK:
entry = _WORKER_ITEMS.pop(token, None)
if entry is None:
if entry is None or entry[0] <= time.monotonic():
raise RuntimeError("Image-list item is missing, expired, or already consumed")
return entry[1]
return entry
def _check_worker_item_identity(item: dict[str, Any], identity: str, index: int) -> None:
if not is_valid_item_identity(identity):
raise ValueError("Invalid worker item identity")
actual = item_content_identity(item) if "image" in item else item.get("item_identity")
if actual != identity:
raise ValueError("Worker item identity does not match the staged item")
if isinstance(index, bool) or not isinstance(index, int) or item["index"] != index:
raise ValueError("Worker item index does not match the staged item")
def consume_worker_item(token: str, *, item_identity: str | None = None,
item_index: int | None = None) -> dict[str, Any]:
with _LOCK:
expires, item = _pop_worker_item(token)
if "image" not in item:
raise RuntimeError("Image-list item is already consumed")
if item_identity is not None:
_check_worker_item_identity(item, item_identity, item_index)
if expires <= time.monotonic():
raise RuntimeError("Image-list item is expired")
# Start misses release staged tensors, but retain a tiny per-token
# receipt until End. This is NOT a completed-result/content cache.
# Keeping the original deadline prevents cache hits bypassing TTL.
_WORKER_ITEMS[token] = (expires, {"item_identity": item_identity, "index": item["index"]})
return item
def finish_worker_item(token: str, item_identity: str, item_index: int) -> None:
"""Validate and release this dispatch, whether Start ran or hit native cache."""
expires, item = _pop_worker_item(token)
_check_worker_item_identity(item, item_identity, item_index)
if expires <= time.monotonic():
raise RuntimeError("Image-list item is expired")
def delete_worker_item(token: str) -> bool:
@@ -376,13 +519,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 +646,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,14 +671,27 @@ 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", {})
for name in ("images", "masks", "texts"):
end_override_names(end_inputs)
for name in END_RESULT_INPUTS:
if name in end_inputs:
_direct_link(end_inputs[name], prompt, end_id, name)
region_nodes = _ancestors_for_inputs(prompt, end_id, start_id, ("images", "masks", "texts"))
region_nodes = _ancestors_for_inputs(prompt, end_id, start_id, END_RESULT_INPUTS)
if start_id not in region_nodes:
raise ValueError(f"Start node {start_id} is not an ancestor of End node {end_id}")
forbidden = {n for n in region_nodes if n != start_id and n != end_id and
@@ -529,8 +699,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 +713,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:
@@ -571,18 +736,28 @@ def validate_image_list_regions(prompt: dict[str, Any]) -> list[dict[str, Any]]:
def build_worker_prompt(prompt: dict[str, Any], region: dict[str, Any], item_token: str,
context_id: str, item_index: int, callback_url: str) -> dict[str, Any]:
"""Copy only the per-image branch and its worker-side dependencies."""
context_id: str, item_index: int, callback_url: str, *,
item_identity: str | None = None) -> dict[str, Any]:
"""Copy the branch; opt into stable ancestry only after a worker handshake."""
start_id, end_id = region["start_id"], region["end_id"]
_direct_link(prompt[end_id].get("inputs", {}).get("context"), prompt, end_id, "context")
kept = {node_id: json.loads(json.dumps(prompt[node_id])) for node_id in region["node_ids"]}
start_inputs = kept[start_id].setdefault("inputs", {})
# Per-item data reaches workers only through the staged item token.
for name in list(start_inputs):
if name in ("images", "mask") or name.startswith(VALUES_INPUT + "."):
start_inputs.pop(name)
start_inputs.update({"worker_item_token": item_token, "worker_context_id": context_id})
end_inputs = kept[end_id].setdefault("inputs", {})
if item_identity is not None:
if not is_valid_item_identity(item_identity):
raise ValueError("Invalid worker item identity")
# Native ancestry signatures include EVERY literal, not just a custom
# fingerprint. Remove inherited execution IDs, callbacks and all source
# inputs wholesale; transient delivery state belongs downstream at End.
kept[start_id]["inputs"] = {"worker_item_identity": item_identity}
end_inputs.update(worker_context_id=context_id, worker_item_identity=item_identity)
else:
# Preserve old raw /prompt and older-worker behavior byte-for-byte.
start_inputs = kept[start_id].setdefault("inputs", {})
for name in list(start_inputs):
if name in ("images", "mask") or name.startswith(VALUES_INPUT + "."):
start_inputs.pop(name)
start_inputs.update({"worker_item_token": item_token, "worker_context_id": context_id})
end_inputs.update({"worker_result_token": item_token,
"worker_item_index": int(item_index),
"master_callback_url": callback_url})
@@ -692,6 +867,7 @@ def set_master_result(context_id: str, item_index: int, token: str, result: dict
raise ValueError("Image-list result token does not match an assigned item")
if item_index in state.received_indexes:
raise ValueError(f"Duplicate image-list result for item {item_index}")
_checked_overrides(result, end_override_names(state.original_prompt[state.end_id].get("inputs", {})))
state.received_indexes.add(item_index)
state.results[item_index] = result
event = state.event
@@ -723,10 +899,24 @@ def decode_result(result: Any) -> dict[str, Any]:
if mask is not None and (mask.ndim != 3 or mask.shape[0] != 1):
raise ValueError("result.mask must contain exactly one MASK frame")
text = decode_value(result.get("text"), "STRING")
return {"image": image, "mask": mask, "text": text}
decoded = {"image": image, "mask": mask, "text": text}
if "values" in result:
raw_values = result["values"]
if not isinstance(raw_values, dict):
raise ValueError("result.values must be an object")
values = {}
for name, encoded in raw_values.items():
if name not in VALUE_NAMES:
raise ValueError(f"result.values has unknown slot {name!r}")
if not isinstance(encoded, dict) or encoded.get("type") not in _VALUE_KINDS:
raise ValueError(f"result.values.{name} has an unsupported type")
values[name] = decode_value(encoded, encoded["type"])
decoded["values"] = _validate_result_values(values)
return decoded
def normalize_worker_result(images: Any, masks: Any = None, texts: Any = None) -> dict[str, Any]:
def normalize_worker_result(images: Any, masks: Any = None, texts: Any = None,
values: Any = None) -> dict[str, Any]:
def one(value, name, kind):
if value is None:
return None
@@ -749,17 +939,29 @@ def normalize_worker_result(images: Any, masks: Any = None, texts: Any = None) -
text = one(texts, "texts", "STRING")
if text is not None and not isinstance(text, str):
raise TypeError("Worker End text result must be a string")
return {"image": image, "mask": mask, "text": text}
result = {"image": image, "mask": mask, "text": text}
if values is not None:
overrides = normalize_value_overrides(values)
if overrides:
result["values"] = overrides
return result
def encode_result(result: dict[str, Any]) -> dict[str, Any]:
normalized = normalize_worker_result(result.get("image"), result.get("mask"), result.get("text"))
return {"image": encode_value(normalized["image"], "IMAGE"),
"mask": encode_value(normalized["mask"], "MASK"),
"text": encode_value(normalized["text"], "STRING")}
encoded = {"image": encode_value(normalized["image"], "IMAGE"),
"mask": encode_value(normalized["mask"], "MASK"),
"text": encode_value(normalized["text"], "STRING")}
if "values" in result:
values = _validate_result_values(result["values"])
if values:
encoded["values"] = {name: encode_value(value, value_kind(value, name))
for name, value in values.items()}
return encoded
def _collect_output_lists(results: list[dict[str, Any]]) -> tuple[list[Any], list[Any], list[Any]]:
def _collect_output_lists(results: list[dict[str, Any]], *, state: ImageListJob | None = None,
include_values: bool = False) -> tuple[list[Any], ...]:
if not results:
raise RuntimeError("Image-list execution returned no results")
images, masks, texts = [], [], []
@@ -769,10 +971,20 @@ def _collect_output_lists(results: list[dict[str, Any]]) -> tuple[list[Any], lis
images.append(result["image"])
masks.append(result.get("mask"))
texts.append(result.get("text"))
return images, masks, texts
if not include_values:
return images, masks, texts
if state is None or len(results) != len(state.items):
raise RuntimeError("Image-list values require exactly one result per source item")
expected = end_override_names(state.original_prompt[state.end_id].get("inputs", {}))
# Reconstruct unchanged values from the master source, not worker transport.
# Do this before the context is removed, and only after all items succeeded.
values = [{**item.get("values", {}), **_checked_overrides(result, expected)}
for item, result in zip(state.items, results)]
return images, masks, texts, values
async def run_image_list_job(context_id: str, *, model_callback=None) -> tuple[list[Any], list[Any], list[Any]]:
async def run_image_list_job(context_id: str, *, model_callback=None,
include_values: bool = False) -> tuple[list[Any], ...]:
"""Schedule a shared queue across the optional master and direct workers."""
state = get_master_job(context_id)
if state is None:
@@ -786,7 +998,7 @@ async def run_image_list_job(context_id: str, *, model_callback=None) -> tuple[l
timer = asyncio.timeout(deadline if deadline > 0 else None)
try:
async with timer:
return await _run_image_list_job(state)
return await _run_image_list_job(state, include_values=include_values)
except TimeoutError:
if timer.expired():
raise TimeoutError(f"Image-list total deadline exceeded ({deadline:g}s); cleanup attempted") from None
@@ -799,7 +1011,7 @@ async def run_image_list_job(context_id: str, *, model_callback=None) -> tuple[l
remove_master_job(context_id)
async def _run_image_list_job(state: ImageListJob):
async def _run_image_list_job(state: ImageListJob, *, include_values: bool = False):
context_id = state.context_id
state.event = asyncio.Event()
state.event_loop = asyncio.get_running_loop()
@@ -817,12 +1029,16 @@ async def _run_image_list_job(state: ImageListJob):
if not workers and not state.master_participates:
raise RuntimeError("Distributed Image List requires enabled direct workers; serverless pools are unsupported")
# Recheck worker capability and idle state immediately before the first item.
worker_statuses = await preflight_image_list_workers(workers, require_idle=not state.master_participates) if workers else {}
override_names = end_override_names(state.original_prompt[state.end_id].get("inputs", {}))
value_preflight = {"require_value_results": True} if override_names else {}
worker_statuses = (await preflight_image_list_workers(
workers, require_idle=not state.master_participates, **value_preflight) if workers else {})
workers = [worker for worker in workers if worker["id"] in worker_statuses]
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")
@@ -849,7 +1065,8 @@ async def _run_image_list_job(state: ImageListJob):
# /prompt queue, or repeated execution of the original source.
with _LOCK:
_WORKER_ITEMS[token] = (time.monotonic() + _WORKER_ITEM_TTL_SECONDS, state.items[index])
prompt = build_worker_prompt(state.original_prompt, region, token, context_id, index, "")
prompt = build_worker_prompt(state.original_prompt, region, token, context_id, index, "",
item_identity=item_content_identity(state.items[index]))
prompt[state.start_id]["inputs"]["image_list_local_execution"] = True
prompt[state.end_id]["inputs"]["image_list_local_execution"] = True
future = local_executor.submit(prompt, f"{context_id}:{index}", state.end_id, state.workflow_meta)
@@ -859,6 +1076,13 @@ async def _run_image_list_job(state: ImageListJob):
if worker is None:
raise RuntimeError(f"Worker {worker_id} was removed before dispatch")
assert session is not None
if override_names:
# A worker can restart/downgrade between items. Check before upload,
# not merely once at queue preparation or at the start of the run.
ready = await preflight_image_list_workers([worker], require_idle=False,
require_value_results=True)
if worker_id not in ready:
raise RuntimeError(f"Image-list worker {worker_id} is unavailable for value overrides")
token = create_worker_task_token()
register_assignment(context_id, index, worker_id, token)
item_payload = {"context_id": context_id, "item": encode_item(state.items[index])}
@@ -872,6 +1096,14 @@ async def _run_image_list_job(state: ImageListJob):
timeout=aiohttp.ClientTimeout(total=60)) as response:
if response.status >= 400:
raise RuntimeError(f"Worker {worker_id} rejected image-list item {index}: HTTP {response.status}")
try:
stored = await response.json()
except (ValueError, aiohttp.ClientError):
stored = None
# A per-upload acknowledgement gates the new graph. Older workers may
# return no identity; do not change global protocol/preflight contracts.
identity = stored.get("item_identity") if isinstance(stored, dict) else None
identity_options = {"item_identity": identity} if is_valid_item_identity(identity) else {}
callback_url = build_master_callback_url(worker, config=config)
if worker_id not in prepared_worker_prompts:
worker_prompt = build_worker_prompt(state.original_prompt, region, token, context_id,
@@ -884,14 +1116,17 @@ async def _run_image_list_job(state: ImageListJob):
from ..api.orchestration.media_sync import sync_worker_media
await asyncio.wait_for(sync_worker_media(worker, worker_prompt), timeout=60)
prepared_worker_prompts[worker_id] = worker_prompt
# Reuse media rewrites, never assignment identity or a mutable submitted graph.
# Reuse legacy media rewrites, never assignment identity or a mutable
# submitted graph. Keeping this template legacy also permits a worker to
# downgrade between items without inheriting a previous stable identity.
worker_prompt = build_worker_prompt(prepared_worker_prompts[worker_id], region, token,
context_id, index, callback_url)
context_id, index, callback_url, **identity_options)
handle["submission_pending"] = True
prompt_id = await dispatch_worker_prompt(
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")
@@ -1018,7 +1253,7 @@ async def _run_image_list_job(state: ImageListJob):
slots = ([master_slot] if local_executor is not None else []) + [worker["id"] for worker in workers]
results = await dynamic_schedule(len(state.items), slots,
dispatch, wait_next, cleanup, interrupted)
return _collect_output_lists(results)
return _collect_output_lists(results, state=state, include_values=include_values)
finally:
if local_executor is not None:
# No timeout here: Python cannot safely kill a synchronous node.
+140 -16
View File
@@ -10,6 +10,7 @@ from ..image_list.runtime import (
MAX_VALUE_SLOTS,
VALUE_NAMES,
VALUES_INPUT,
end_override_names,
create_master_job,
normalize_image_list_inputs,
normalize_worker_result,
@@ -17,9 +18,40 @@ from ..image_list.runtime import (
encode_result,
encode_request_payload,
unwrap_single,
is_valid_item_identity,
)
def _stable_worker_binding(prompt, start_id, identity):
"""Resolve transient state only from the exact Start -> End context wire."""
if not is_valid_item_identity(identity):
raise ValueError("Invalid worker item identity")
start = prompt.get(start_id) if isinstance(prompt, dict) else None
if (not isinstance(start, dict) or start.get("class_type") != "DistributedImageListStart"
or start.get("inputs", {}).get("worker_item_identity") != identity):
raise RuntimeError("Worker item identity is missing its paired Start")
candidates = []
for end_id, node in prompt.items():
if not isinstance(node, dict) or node.get("class_type") != "DistributedImageListEnd":
continue
inputs = node.get("inputs", {})
link = inputs.get("context")
if isinstance(link, (list, tuple)) and list(link) == [start_id, 3]:
candidates.append((end_id, inputs))
if len(candidates) != 1:
raise RuntimeError("Worker Start must have exactly one paired End context wire")
end_id, inputs = candidates[0]
if inputs.get("worker_item_identity") != identity:
raise ValueError("Paired End worker item identity does not match Start")
for name in ("worker_result_token", "worker_context_id"):
if not isinstance(inputs.get(name), str) or not inputs[name]:
raise RuntimeError(f"Worker paired End is missing {name}")
index = inputs.get("worker_item_index")
if isinstance(index, bool) or not isinstance(index, int) or index < 0:
raise ValueError("Worker paired End item index is invalid")
return end_id, inputs
# Types the transport can carry exactly. Anything else (models, latents, ...)
# must be loaded inside the region or shared as a model dependency instead.
_VALUE_TYPES = [io.Image, io.Mask, io.String, io.Int, io.Float, io.Boolean]
@@ -69,10 +101,26 @@ class DistributedImageListStart(io.ComfyNode):
worker_item_token="", worker_context_id="", image_list_enabled_worker_ids="[]",
image_list_master_callback="", image_list_execution_id="",
image_list_workflow_meta=None, image_list_master_participates=None,
image_list_local_execution=False):
image_list_local_execution=False, worker_item_identity=None):
worker_item_token = unwrap_single(worker_item_token)
worker_context_id = unwrap_single(worker_context_id)
worker_item_identity = unwrap_single(worker_item_identity)
image_list_workflow_meta = unwrap_single(image_list_workflow_meta)
if worker_item_identity is not None:
prompt = getattr(cls.hidden, "prompt", None)
start_id = str(getattr(cls.hidden, "unique_id", ""))
_, delivery = _stable_worker_binding(prompt, start_id, worker_item_identity)
# Only cache misses execute Start. Never put this token/context in
# Start's literal inputs or returned context: both would poison reuse.
from ..image_list.runtime import consume_worker_item
item = consume_worker_item(delivery["worker_result_token"],
item_identity=worker_item_identity,
item_index=delivery["worker_item_index"])
item_values = item.get("values") or {}
return io.NodeOutput(
[item["image"]], [item.get("mask")], [item["index"]], [worker_item_identity],
*[[item_values.get(name)] for name in VALUE_NAMES],
)
if worker_item_token:
if unwrap_single(image_list_local_execution):
from ..image_list.runtime import consume_worker_item
@@ -99,13 +147,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 +171,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],
@@ -143,17 +197,23 @@ class DistributedImageListEnd(io.ComfyNode):
node_id="DistributedImageListEnd",
display_name="Distributed Image List End",
category="image/distributed",
description="Closes a paired image-list region and returns ordered image, mask, and text lists.",
description=("Closes a paired image-list region and returns ordered image, mask, text and values lists. "
"Start values are retained; connected value_N inputs replace or add that key per item."),
inputs=[
io.Image.Input("images", lazy=True),
io.Mask.Input("masks", optional=True, lazy=True),
io.String.Input("texts", optional=True, lazy=True, force_input=True),
io.Custom("DISTRIBUTED_IMAGE_LIST_CONTEXT").Input("context"),
*[io.MultiType.Input(name, types=_VALUE_TYPES, optional=True, lazy=True,
tooltip="Replace or add this value for each item; exactly one value per item.")
for name in VALUE_NAMES],
],
outputs=[
io.Image.Output(display_name="image_list", is_output_list=True),
io.Mask.Output(display_name="masks", is_output_list=True),
io.String.Output(display_name="texts", is_output_list=True),
io.Custom("DISTRIBUTED_IMAGE_LIST_VALUES").Output(
id="values", display_name="values", is_output_list=True),
],
is_input_list=True,
is_output_node=True,
@@ -175,7 +235,8 @@ class DistributedImageListEnd(io.ComfyNode):
and node.get("class_type") == "DistributedImageListEnd"
and node.get("inputs", {}).get("worker_result_token") == worker_result_token), {})
required = []
for name, value in (("images", images), ("masks", masks), ("texts", texts)):
for name, value in (("images", images), ("masks", masks), ("texts", texts),
*((name, kwargs.get(name)) for name in end_override_names(end_inputs))):
connected = name == "images" or name in end_inputs
missing = value is None or (isinstance(value, (list, tuple)) and all(v is None for v in value))
if connected and missing:
@@ -189,26 +250,63 @@ class DistributedImageListEnd(io.ComfyNode):
@classmethod
def execute(cls, images=None, masks=None, texts=None, context=None,
worker_result_token="", worker_item_index=-1, master_callback_url="",
image_list_local_execution=False):
image_list_local_execution=False, worker_context_id="",
worker_item_identity=None, **values):
worker_result_token = unwrap_single(worker_result_token)
master_callback_url = unwrap_single(master_callback_url)
worker_context_id = unwrap_single(worker_context_id)
worker_item_identity = unwrap_single(worker_item_identity)
if worker_item_identity is not None and not worker_result_token:
raise RuntimeError("Worker item identity is missing its delivery token")
if worker_result_token:
context_values = [item for item in (context if isinstance(context, (list, tuple)) else [context])
if isinstance(item, str) and item]
if not context_values or len(set(context_values)) != 1:
raise RuntimeError("Worker Image List End is missing its paired context")
result = normalize_worker_result(images, masks, texts)
prompt = getattr(cls.hidden, "prompt", {}) or {}
context_id = context_values[0]
if worker_item_identity is not None:
end_id = str(getattr(cls.hidden, "unique_id", ""))
node = prompt.get(end_id, {})
link = node.get("inputs", {}).get("context")
if (node.get("class_type") != "DistributedImageListEnd"
or not isinstance(link, (list, tuple)) or len(link) != 2):
raise RuntimeError("Worker End is missing its paired Start context wire")
paired_id, end_inputs = _stable_worker_binding(prompt, link[0], worker_item_identity)
if (paired_id != end_id or context_values != [worker_item_identity]
or end_inputs["worker_context_id"] != worker_context_id
or end_inputs["worker_result_token"] != worker_result_token
or end_inputs["worker_item_index"] != unwrap_single(worker_item_index)
or end_inputs.get("master_callback_url", "") != master_callback_url):
raise RuntimeError("Worker End identity/context does not match its current delivery")
from ..image_list.runtime import finish_worker_item
# On a miss this consumes Start's receipt; on a hit it validates
# and drops the current staged tensors. Both paths enforce TTL.
finish_worker_item(worker_result_token, worker_item_identity,
unwrap_single(worker_item_index))
context_id = worker_context_id
else:
end_inputs = next((node.get("inputs", {}) for node in prompt.values()
if isinstance(node, dict) and node.get("class_type") == "DistributedImageListEnd"
and node.get("inputs", {}).get("worker_result_token") == worker_result_token), {})
expected = end_override_names(end_inputs)
# Missing connected lazy values must fail instead of becoming an
# implicit fallback to the original Start value.
supplied = {name: values.get(name) for name in expected}
supplied.update(values)
result = normalize_worker_result(images, masks, texts, supplied)
if unwrap_single(image_list_local_execution):
from ..image_list.runtime import set_master_result
set_master_result(context_values[0], int(unwrap_single(worker_item_index)),
set_master_result(context_id, int(unwrap_single(worker_item_index)),
str(worker_result_token), result)
else:
payload = encode_result(result)
run_async_in_server_loop(
submit_worker_result(str(master_callback_url), str(worker_result_token),
context_values[0], int(unwrap_single(worker_item_index)), payload), timeout=60,
context_id, int(unwrap_single(worker_item_index)), payload), timeout=60,
)
return io.NodeOutput([result["image"]], [result.get("mask")], [result.get("text")])
return io.NodeOutput([result["image"]], [result.get("mask")], [result.get("text")],
[result.get("values", {})])
context_values = [item for item in (context if isinstance(context, (list, tuple)) else [context])
if isinstance(item, str) and item]
@@ -223,14 +321,40 @@ class DistributedImageListEnd(io.ComfyNode):
# Add them to its existing tracker; never start/end another tracker
# that could unpin a model still needed by the parent.
model_callback = getattr(execution_list, "output_link_callback", None)
result_lists = run_async_in_server_loop(run_image_list_job(context_id, model_callback=model_callback))
result_lists = run_async_in_server_loop(run_image_list_job(
context_id, model_callback=model_callback, include_values=True))
except ImageListCancelledError:
import comfy.model_management
raise comfy.model_management.InterruptProcessingException()
images_out, masks_out, texts_out = result_lists
if len(images_out) != len(masks_out) or len(images_out) != len(texts_out):
raise RuntimeError("Image-list collector received mismatched image, mask, or text result counts")
return io.NodeOutput(images_out, masks_out, texts_out)
images_out, masks_out, texts_out, values_out = result_lists
if len({len(output) for output in result_lists}) != 1:
raise RuntimeError("Image-list collector received mismatched image, mask, text or values result counts")
return io.NodeOutput(images_out, masks_out, texts_out, values_out)
class DistributedImageListGetValue(io.ComfyNode):
"""Ordinary scalar execution lets core map over End's ordered values list."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DistributedImageListGetValue",
display_name="Distributed Image List Get Value",
category="image/distributed",
inputs=[io.Custom("DISTRIBUTED_IMAGE_LIST_VALUES").Input("values"),
io.Combo.Input("value_name", options=list(VALUE_NAMES), default="value_1")],
outputs=[io.AnyType.Output(id="value", display_name="value")],
)
@classmethod
def execute(cls, values, value_name="value_1"):
if value_name not in VALUE_NAMES:
raise ValueError(f"Unknown image-list value slot {value_name!r}")
if not isinstance(values, dict):
raise TypeError("Image List Get Value expects one values mapping from Image List End")
if value_name not in values:
raise ValueError(f"Image List Get Value: {value_name} is absent from this item's values")
return io.NodeOutput(values[value_name])
async def submit_worker_result(callback_url, item_token, context_id, item_index, result):
+2 -2
View File
@@ -10,7 +10,7 @@ from . import utilities as _utilities
from .collector import DistributedCollectorNode as _CollectorRuntime
from ..utils import transfer_codec as _transfer_codec
from .distributed_upscale import UltimateSDUpscaleDistributed as _UpscaleRuntime
from .image_list import DistributedImageListEnd, DistributedImageListStart
from .image_list import DistributedImageListEnd, DistributedImageListStart, DistributedImageListGetValue
class DistributedSeed(io.ComfyNode):
@@ -212,4 +212,4 @@ class DistributedWorkerSink(io.ComfyNode):
NODES = [DistributedCollector, DistributedSeed, DistributedModelName, DistributedValue,
ImageBatchDivider, AudioBatchDivider, DistributedEmptyImage, UltimateSDUpscaleDistributed,
DistributedWorkerSink, DistributedImageListStart, DistributedImageListEnd]
DistributedWorkerSink, DistributedImageListStart, DistributedImageListEnd, DistributedImageListGetValue]
+68
View File
@@ -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)), \
@@ -0,0 +1,172 @@
"""Offline CPU proof against ComfyUI's real signatures and PromptExecutor.
Exercises one native worker executor and separate private master-slot lifetimes,
including their shared stable Start / fresh End contract.
"""
import asyncio
import copy
import importlib
from pathlib import Path
import sys
import types
ROOT = Path(__file__).resolve().parents[2]
def main(core):
sys.path.insert(0, str(Path(core).resolve()))
import comfy.cli_args
comfy.cli_args.args.cpu = True # Must precede all model-management imports.
comfy.cli_args.args.disable_assets = True
import execution
import nodes
import torch
from comfy_execution.caching import CacheKeySetInputSignature
name = '_native_cache_identity_probe'
for suffix in ('', '.image_list', '.nodes', '.utils'):
package = types.ModuleType(name + suffix)
package.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
sys.modules[name + suffix] = package
def forbidden(*args, **kwargs):
raise AssertionError('offline native proof attempted a server-loop bridge')
sys.modules[name + '.utils.async_helpers'] = types.SimpleNamespace(run_async_in_server_loop=forbidden)
rt = importlib.import_module(name + '.image_list.runtime')
boundary = importlib.import_module(name + '.nodes.image_list')
calls = []
class Operation:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'image': ('IMAGE',), 'mask': ('MASK',),
'text': ('STRING', {'forceInput': True})}}
RETURN_TYPES = ('IMAGE', 'MASK', 'STRING')
FUNCTION = 'run'
def run(self, image, mask, text):
calls.append(text)
return image + 1, mask + 1, text + ':done'
nodes.NODE_CLASS_MAPPINGS.update(DistributedImageListStart=boundary.DistributedImageListStart,
DistributedImageListEnd=boundary.DistributedImageListEnd,
Operation=Operation)
prompt = {
'start': {'class_type': rt.START_CLASS, 'inputs': {'image_list_execution_id': 'outer-A'}},
'op': {'class_type': 'Operation', 'inputs': {
'image': ['start', 0], 'mask': ['start', 1], 'text': ['start', 4]}},
'end': {'class_type': rt.END_CLASS, 'inputs': {
'images': ['op', 0], 'masks': ['op', 1], 'texts': ['op', 2],
'context': ['start', 3], 'value_3': ['start', 2]}},
}
data = rt.normalize_image_list_inputs(torch.zeros(1, 2, 3, 3, dtype=torch.float64),
torch.ones(1, 2, 3, dtype=torch.bfloat16),
{'value_1': 'original', 'value_2': 3})[0]
region = rt.validate_image_list_regions(prompt)[0]
def prepare(data):
token = rt.create_worker_task_token()
context = rt.create_master_job('start', 'end', prompt, [data], ['worker'], '')
rt.register_assignment(context, data['index'], 'worker', token)
identity = rt.store_worker_item(token, rt.encode_item(data))
result = rt.build_worker_prompt(prompt, region, token, context, data['index'], '', item_identity=identity)
result['start']['inputs']['image_list_local_execution'] = True
result['end']['inputs']['image_list_local_execution'] = True
return token, context, result
async def signature(graph, executor):
dynamic = execution.DynamicPrompt(graph)
changed = execution.IsChangedCache('signature-only', dynamic, executor.caches.outputs)
keys = CacheKeySetInputSignature(dynamic, graph.keys(), changed)
await keys.add_keys(graph.keys())
return {node: keys.get_data_key(node) for node in graph}
async def probe():
server = types.SimpleNamespace(client_id=None, last_node_id=None, send_sync=lambda *a, **kw: None)
executor = execution.PromptExecutor(server, cache_type=execution.CacheType.LRU,
cache_args={'lru': 64, 'ram': 0, 'ram_inactive': 0})
ta, ca, a = prepare(data)
prompt['start']['inputs']['image_list_execution_id'] = 'outer-B'
tb, cb, b = prepare(data)
assert ta != tb and ca != cb
ka, kb = await signature(a, executor), await signature(b, executor)
assert ka['start'] == kb['start'] and ka['op'] == kb['op']
assert ka['end'] != kb['end']
assert a['start']['inputs'] == b['start']['inputs']
# Exercise actual V3 list mapping, hidden.prompt/unique_id, lazy inputs,
# native node signatures and cache hits, not a reimplementation of them.
for number, (token, context, graph) in enumerate(((ta, ca, a), (tb, cb, b))):
valid = await execution.validate_prompt(f'native-{number}', graph, ['end'])
assert valid[0], valid
await executor.execute_async(graph, f'native-{number}', {}, ['end'])
assert executor.success, executor.status_messages
assert token not in rt._WORKER_ITEMS
result = rt.get_master_job(context).results[0]
assert torch.equal(result['image'], data['image'] + 1)
assert torch.equal(result['mask'], data['mask'] + 1)
assert result['text'] == 'original:done'
assert result['values'] == {'value_3': 0}
cached = next(entry['nodes'] for event, entry in executor.status_messages if event == 'execution_cached')
if number:
assert 'start' in cached and 'op' in cached and 'end' not in cached, cached
assert calls == ['original'], calls
for failure in ('expired', 'mismatched'):
token, context, graph = prepare(data)
expires, staged = rt._WORKER_ITEMS[token]
if failure == 'expired':
rt._WORKER_ITEMS[token] = (0, staged)
else:
staged['image'] = staged['image'] + 9
await executor.execute_async(graph, f'rejected-{failure}', {}, ['end'])
assert not executor.success, failure
assert not rt.get_master_job(context).received_indexes
assert token not in rt._WORKER_ITEMS
cached = next(entry['nodes'] for event, entry in executor.status_messages if event == 'execution_cached')
assert 'start' in cached and 'op' in cached and 'end' not in cached, cached
assert calls == ['original'], calls
for part in ('image', 'mask', 'value', 'index'):
changed = copy.deepcopy(data)
if part in ('image', 'mask'):
changed[part].reshape(-1)[0] += 1
elif part == 'value':
changed['values']['value_2'] = 4 # Even an unconnected slot is content.
else:
changed['index'] = 1
token, context, graph = prepare(changed)
keys = await signature(graph, executor)
assert keys['start'] != ka['start'] and keys['op'] != ka['op'], part
await executor.execute_async(graph, f'changed-{part}', {}, ['end'])
assert executor.success, executor.status_messages
assert token not in rt._WORKER_ITEMS
assert rt.get_master_job(context).received_indexes == {changed['index']}
assert len(calls) == 5, calls
# Combine the real Start/End identity path with cache transfer between
# private master slots, not just independent signature/retention probes.
local = importlib.import_module(name + '.image_list.local_executor')
before = len(calls)
for number in range(2):
token, context, graph = prepare(data)
slot = local.LocalExecutionSlot()
try:
await asyncio.wrap_future(slot.submit(graph, f'local-{number}', 'end'))
result = rt.get_master_job(context).results[0]
assert torch.equal(result['image'], data['image'] + 1)
assert result['values'] == {'value_3': 0}
assert token not in rt._WORKER_ITEMS
finally:
await slot.close()
rt.remove_master_job(context)
assert local._RETAINED_OUTPUTS is not None
assert len(calls) == before + 1, calls
assert not rt._WORKER_ITEMS
assert not torch.cuda.is_initialized()
print('NATIVE_CACHE_IDENTITY_OK stable Start/op keys; fresh End key; '
'1 operation for 2 same-data jobs; 4 changed-content misses; '
'expired/mismatched cached inputs rejected; current delivery and cleanup')
asyncio.run(probe())
if __name__ == '__main__':
main(sys.argv[1])
+281 -68
View File
@@ -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,13 +37,14 @@ import aiohttp
torch.set_num_threads(1)
SOURCE_CALLS, BRANCH_CALLS, CAPTURED, MODEL_CALLS = [], [], [], []
UTILITY_CALLS, ITEM_EXECUTIONS, VALUE_CAPTURES = [], [], []
ALLOW_MASTER = False
COUNT = 7
def source_rows():
return [{'image': torch.full((1, 5 + i, 7, 3 if i % 2 == 0 else 4), i / 16, dtype=torch.float32),
'reference_image': torch.full((1, 2, 3, 3), (i + 1) / 16),
'reference_image': torch.full((2, 2, 3, 3), (i + 1) / 16),
'mask': torch.full((1, 5 + i, 7), (i + 1) / 16),
'text': f'item-{i}', 'strength': (i + 1) / 16, 'seed': 100 + i, 'index': i}
for i in range(COUNT)]
@@ -85,6 +87,53 @@ 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['4']['inputs'].get('worker_result_token'), 'per-item node ran in outer graph'
assert (prompt['2']['inputs'].get('worker_item_identity')
or prompt['2']['inputs'].get('worker_item_token')), 'per-item source binding missing'
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 +141,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
@@ -111,16 +162,32 @@ class TestImageOperation:
return image + strength, mask * 0.5, f'{text}:seed-{seed}'
class TestValueOverride:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'text': ('STRING', {'forceInput': True}),
'index': ('INT', {'forceInput': True})}}
RETURN_TYPES = ('STRING', 'INT')
FUNCTION = 'run'
def run(self, text, index):
assert WORKER or ALLOW_MASTER, 'override branch executed on delegate-only master'
return (text + ':override', 200 + index)
class TestCapture:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'images': ('IMAGE',), 'masks': ('MASK',), 'texts': ('STRING', {'forceInput': True})}}
return {'required': {'images': ('IMAGE',), 'masks': ('MASK',), 'texts': ('STRING', {'forceInput': True})},
'optional': {'values': ('DISTRIBUTED_IMAGE_LIST_VALUES',),
'picked': ('STRING', {'forceInput': True})}}
INPUT_IS_LIST = True
OUTPUT_NODE = True
RETURN_TYPES = ()
FUNCTION = 'run'
def run(self, images, masks, texts):
def run(self, images, masks, texts, values=None, picked=None):
assert not WORKER, 'post-End capture leaked into a worker prompt'
CAPTURED.append((images, masks, texts))
VALUE_CAPTURES.append((values, picked))
return ()
@@ -129,9 +196,19 @@ async def setup():
ps = server.PromptServer(asyncio.get_running_loop(), assets)
assert await nodes.load_custom_node(str(ROOT))
module = sys.modules[str(ROOT).replace('.', '_x_')]
boundary = importlib.import_module(module.__name__ + '.nodes.image_list')
encode_result = boundary.encode_result
def checked_encode(result):
encoded = encode_result(result)
if 'values' in encoded:
assert set(encoded['values']) == {f'value_{n}' for n in (2, 3, 4, 5, 7, 8, 9, 10)}
return encoded
boundary.encode_result = checked_encode
nodes.NODE_CLASS_MAPPINGS.update(TestImageSource=TestImageSource,
TestImageOperation=TestImageOperation, TestCapture=TestCapture,
TestCachedModel=TestCachedModel, TestSharedFlag=TestSharedFlag)
TestCachedModel=TestCachedModel, TestSharedFlag=TestSharedFlag,
TestRequiredOutputUtility=TestRequiredOutputUtility,
TestDiscardedPreview=TestDiscardedPreview, TestValueOverride=TestValueOverride)
# 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 +224,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)
@@ -224,23 +296,51 @@ def graph(callback, worker_ids):
'images': ['1', 0], 'mask': ['1', 2],
'values.value_1': ['1', 1], 'values.value_2': ['1', 3],
'values.value_3': ['1', 4], 'values.value_4': ['1', 5], 'values.value_5': ['7', 0],
'values.value_6': ['1', 2],
'image_list_enabled_worker_ids': json.dumps(worker_ids), 'image_list_master_callback': callback}},
'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]}},
'images': ['3', 0], 'masks': ['3', 1], 'texts': ['3', 2], 'context': ['2', 3],
'value_2': ['10', 0], 'value_3': ['2', 6], 'value_4': ['2', 7], 'value_5': ['2', 8],
'value_7': ['3', 1], 'value_8': ['2', 4], 'value_9': ['2', 9], 'value_10': ['10', 1]}},
'5': {'class_type': 'TestCapture', 'inputs': {'images': ['4', 0], 'masks': ['4', 1], 'texts': ['4', 2],
'values': ['4', 3], 'picked': ['get', 0]}},
'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]}},
'10': {'class_type': 'TestValueOverride', 'inputs': {'text': ['2', 5], 'index': ['2', 2]}},
'get': {'class_type': 'DistributedImageListGetValue', 'inputs': {'values': ['4', 3], 'value_name': 'value_2'}},
}
def check_carried_values():
values, picked = VALUE_CAPTURES[-1]
assert len(values) == len(picked) == COUNT
for index, (actual, source) in enumerate(zip(values, source_rows())):
assert set(actual) == {f'value_{n}' for n in range(1, 11)}, actual.keys()
assert torch.equal(actual['value_1'], source['reference_image'])
assert actual['value_2'] == picked[index] == f'item-{index}:override'
assert actual['value_3'] == source['strength']
assert actual['value_4'] == source['seed']
assert actual['value_5'] is True
assert torch.equal(actual['value_6'], source['mask'])
assert torch.equal(actual['value_7'], source['mask'] * 0.5)
assert torch.equal(actual['value_8'], source['reference_image'])
assert actual['value_8'].shape[0] == 2, 'value tensor batch was split'
assert torch.equal(actual['value_9'], source['mask'])
assert actual['value_10'] == 200 + index
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 +394,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,12 +445,14 @@ 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
check_carried_values()
images, masks, texts = CAPTURED[0]
assert len(images) == len(masks) == len(texts) == COUNT
for i, row in enumerate(source_rows()):
@@ -332,30 +467,46 @@ 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.
chained = {'1': {'class_type': 'EmptyImage', 'inputs': {
'height': 8, 'width': 8, 'batch_size': 2, 'color': 0}}}
chained['constant'] = {'class_type': 'TestSharedFlag', 'inputs': {}}
previous = '1'
value_source = 'constant'
for start, end in [('2', '3'), ('4', '5'), ('6', '7')]:
chained[start] = {'class_type': 'DistributedImageListStart', 'inputs': {
'images': [previous, 0]}}
'images': [previous, 0], 'values.value_1': [value_source, 0]}}
chained[end] = {'class_type': 'DistributedImageListEnd', 'inputs': {
'images': [start, 0], 'context': [start, 3]}}
value_source = f'get-{end}'
chained[value_source] = {'class_type': 'DistributedImageListGetValue', 'inputs': {
'values': [end, 3], 'value_name': 'value_1'}}
previous = end
chained['8'] = {'class_type': 'TestCapture', 'inputs': {
'images': ['7', 0], 'masks': ['7', 1], 'texts': ['7', 2]}}
'images': ['7', 0], 'masks': ['7', 1], 'texts': ['7', 2], 'values': ['7', 3]}}
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])
assert CAPTURED[1][1:] == ([None, None], [None, None])
assert VALUE_CAPTURES[1][0] == [{'value_1': True}, {'value_1': True}]
assert not runtime._MASTER_JOBS
print('IMAGE_LIST_SEQUENTIAL_OK three regions; one worker; optional outputs omitted')
# Same public queue path, now with master participation enabled. The
@@ -364,37 +515,54 @@ async def coordinator():
config['settings']['master_delegate_only'] = False
participation_counts = {}
for mode, selected in [('mixed', [w['id'] for w in workers]), ('master-only', [])]:
# Keep cold-run exact-once assertions independent of other modes.
importlib.import_module(module.__name__ + '.image_list.local_executor')._RETAINED_OUTPUTS = None
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'
else:
assert not remote_calls and len(local_calls) == COUNT
assert Counter(local_calls + remote_calls) == Counter(range(COUNT)), (local_calls, remote_calls)
check_carried_values()
images, masks, texts = CAPTURED[-1]
assert len(images) == len(masks) == len(texts) == COUNT
for i, row in enumerate(source_rows()):
@@ -403,10 +571,55 @@ async def coordinator():
assert texts[i] == f'item-{i}:seed-{100 + i}', (mode, texts)
assert not runtime._MASTER_JOBS and not runtime._WORKER_ITEMS
participation_counts[mode] = {'master': len(local_calls), 'workers': len(remote_calls)}
# Repeat the completed master-only queue with a fresh outer executor.
# Every item must be delivered again without recomputing its branch.
before_sources, before_local = len(SOURCE_CALLS), len(BRANCH_CALLS)
before_models, before_utilities = len(MODEL_CALLS), len(UTILITY_CALLS)
before_captures = len(CAPTURED)
previous_outputs = CAPTURED[-1]
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
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(
prompt, None, 'test', []), 45)
assert executor.success, executor.status_messages
assert len(SOURCE_CALLS) == before_sources + 1
assert len(BRANCH_CALLS) == before_local
assert len(MODEL_CALLS) == before_models
assert len(UTILITY_CALLS) == before_utilities
assert len(CAPTURED) == before_captures + 1
for before, after in zip(previous_outputs[:2], CAPTURED[-1][:2]):
assert len(before) == len(after) == COUNT
assert all(torch.equal(a, b) for a, b in zip(before, after))
assert previous_outputs[2] == CAPTURED[-1][2]
check_carried_values()
assert not runtime._MASTER_JOBS and not runtime._WORKER_ITEMS
assert prompt == saved_prompt
print('IMAGE_LIST_CACHE_REQUEUE_OK unchanged items reused; fresh ordered delivery and cleanup')
print('IMAGE_LIST_MASTER_OK ' + json.dumps(participation_counts))
print('IMAGE_LIST_VALUES_OK passthrough all six types; sparse override/add; Get; sequential Get-to-Start; direct/mixed/master-only')
assert len(outer_executions) == 5 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.
rows = source_rows()
items = runtime.normalize_image_list_inputs(
[row['image'] for row in rows], [row['mask'] for row in rows],
{f'value_{n}': [row[key] for row in rows]
for n, key in ((1, 'reference_image'), (2, 'text'), (3, 'strength'), (4, 'seed'), (6, 'mask'))}
| {'value_5': True})
slow = workers[1]
async with aiohttp.ClientSession() as session:
async def slow_stats():
@@ -420,11 +633,11 @@ 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, items, ['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):
task = asyncio.create_task(runtime.run_image_list_job(context))
task = asyncio.create_task(runtime.run_image_list_job(context, include_values=True))
active = await until(lambda state: bool(state['running']))
owned_id = active['running'][0]
comfy.model_management.interrupt_current_processing(True)
@@ -442,14 +655,14 @@ async def coordinator():
assert not runtime._MASTER_JOBS
print('IMAGE_LIST_CANCEL_OK owned prompt stopped; no further dispatch; context removed')
# 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'],
bad_rows = copy.deepcopy(items)
bad_rows[0]['values']['value_4'] = -1
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):
try:
await asyncio.wait_for(runtime.run_image_list_job(context), 5)
await asyncio.wait_for(runtime.run_image_list_job(context, include_values=True), 5)
raise AssertionError('worker exception was ignored')
except RuntimeError as exc:
assert 'failed image-list item' in str(exc), str(exc)
@@ -0,0 +1,114 @@
"""CPU-only native Comfy execution across private-slot lifetimes."""
import asyncio
import importlib
import os
from pathlib import Path
import sys
import types
ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, sys.argv[1])
import comfy.cli_args
comfy.cli_args.args.cpu = True
comfy.cli_args.args.disable_assets = True
import execution
import nodes
pkg = '_local_retention_probe'
for name, path in [(pkg, ROOT), (pkg + '.image_list', ROOT / 'image_list')]:
module = types.ModuleType(name)
module.__path__ = [str(path)]
sys.modules[name] = module
local = importlib.import_module(pkg + '.image_list.local_executor')
CALLS = []
OUTPUTS = []
class Counted:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'value': ('INT',)}}
RETURN_TYPES = ('INT',)
FUNCTION = 'run'
def run(self, value):
if value < 0:
raise ValueError('intentional retention failure')
CALLS.append(value)
return (value,)
class Uncached(Counted):
@classmethod
def IS_CHANGED(cls, **kwargs):
return float('nan')
class Capture:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'value': ('INT',), 'delivery': ('STRING',)}}
RETURN_TYPES = ()
OUTPUT_NODE = True
FUNCTION = 'run'
def run(self, value, delivery):
OUTPUTS.append((value, delivery))
return ()
nodes.NODE_CLASS_MAPPINGS.update(RetentionCounted=Counted, RetentionUncached=Uncached, RetentionCapture=Capture)
async def run(slot, value, delivery, uncached=False):
prompt = {
'source': {'class_type': 'RetentionUncached' if uncached else 'RetentionCounted', 'inputs': {'value': value}},
'end': {'class_type': 'RetentionCapture', 'inputs': {'value': ['source', 0], 'delivery': delivery}},
}
await asyncio.wrap_future(slot.submit(prompt, delivery, 'end'))
async def main():
first = local.LocalExecutionSlot()
await run(first, 1, 'first-1')
await run(first, 2, 'first-2')
await first.close()
assert local._RETAINED_OUTPUTS is not None
second = local.LocalExecutionSlot()
await run(second, 1, 'second-1')
await run(second, 2, 'second-2')
await run(second, 3, 'changed')
await second.close()
assert CALLS == [1, 2, 3], CALLS
assert OUTPUTS == [(1, 'first-1'), (2, 'first-2'), (1, 'second-1'), (2, 'second-2'), (3, 'changed')], OUTPUTS
third = local.LocalExecutionSlot()
await run(third, 4, 'uncached-1', True)
await run(third, 4, 'uncached-2', True)
await third.close()
assert CALLS[-2:] == [4, 4]
comfy.cli_args.args.cache_none = True
disabled = local.LocalExecutionSlot()
await run(disabled, 1, 'disabled-1')
await run(disabled, 1, 'disabled-2')
await disabled.close()
assert CALLS[-2:] == [1, 1]
assert local._RETAINED_OUTPUTS is None
comfy.cli_args.args.cache_none = False
failed = local.LocalExecutionSlot()
await run(failed, 7, 'before-failure')
try:
await run(failed, -1, 'failure')
except RuntimeError:
pass
else:
raise AssertionError('failure fixture unexpectedly succeeded')
await failed.close()
assert local._RETAINED_OUTPUTS is None
cancelled = local.LocalExecutionSlot()
await run(cancelled, 8, 'cancelled')
cancelled.request_stop()
await cancelled.close()
assert local._RETAINED_OUTPUTS is None
oversized = execution.LRUCache(execution.CacheKeySetInputSignature, max_size=local.NATIVE_CACHE_MAX_ENTRIES)
empty = execution.DynamicPrompt({})
await oversized.set_prompt(empty, [], execution.IsChangedCache('bound', empty, oversized))
oversized.min_generation = oversized.generation
oversized.cache = {n: n for n in range(local.NATIVE_CACHE_MAX_ENTRIES + 1)}
local._retain_outputs(oversized)
assert local._RETAINED_OUTPUTS is None
print('PASS: cross-run native cache reuse, fresh delivery, input change, noncacheable nodes, cache_none, cancellation, retention bound')
asyncio.run(main())
+36 -4
View File
@@ -45,7 +45,8 @@ ADDED_OPTIONAL_INPUTS = {
# Nodes added after the V1 baseline; not part of saved-workflow contracts.
# DistributedWorkerSink is internal (dev-only) and only inserted by the master.
ADDED_NODES = {'DistributedWorkerSink', 'DistributedImageListStart', 'DistributedImageListEnd'}
ADDED_NODES = {'DistributedWorkerSink', 'DistributedImageListStart', 'DistributedImageListEnd',
'DistributedImageListGetValue'}
def normalized_input(value):
@@ -273,7 +274,7 @@ async def check_image_list_lazy_executor(mapping, module, prompt_server, asset_m
original_contexts = set(runtime._MASTER_JOBS)
image_list_nodes = sys.modules[module.__name__ + '.nodes.image_list']
workflow_meta = {'nodes': [{'id': 2, 'type': 'DistributedImageListStart'}], 'extra': {'acceptance': True}}
async def fake_worker_results(_context_id, *, model_callback=None):
async def fake_worker_results(_context_id, *, model_callback=None, include_values=False):
state = runtime.get_master_job(_context_id)
assert state.prompt_id == 'image-list-lazy', state.prompt_id
assert state.workflow_meta == workflow_meta, state.workflow_meta
@@ -281,6 +282,7 @@ async def check_image_list_lazy_executor(mapping, module, prompt_server, asset_m
return (
[torch.full((1, 8, 8, 3), value, dtype=torch.float32) for value in (0.125, 0.5, 0.875)],
[None, None, None], [None, None, None],
[{}, {}, {}],
)
graph = {
@@ -366,7 +368,7 @@ async def main():
extension = await module.comfy_entrypoint()
classes = await extension.get_node_list()
mapping = {cls.GET_SCHEMA().node_id: cls for cls in classes}
assert len(classes) == len(mapping) == len(BASELINE) + len(ADDED_NODES) == 11
assert len(classes) == len(mapping) == len(BASELINE) + len(ADDED_NODES) == 12
assert set(mapping) == set(BASELINE) | ADDED_NODES
for node_id in ADDED_NODES:
assert nodes.NODE_CLASS_MAPPINGS[node_id] is mapping[node_id]
@@ -423,6 +425,37 @@ async def main():
assert start_schema.display_name == 'Distributed Image List Start'
assert end_schema.node_id == 'DistributedImageListEnd'
assert end_schema.display_name == 'Distributed Image List End'
runtime = sys.modules[module.__name__ + '.image_list.runtime']
assert [input.id for input in end_schema.inputs] == ['images', 'masks', 'texts', 'context'] + list(runtime.VALUE_NAMES)
assert [output.io_type for output in end_schema.outputs] == [
'IMAGE', 'MASK', 'STRING', 'DISTRIBUTED_IMAGE_LIST_VALUES']
assert end_schema.outputs[3].id == end_schema.outputs[3].display_name == 'values'
assert all(output.is_output_list for output in end_schema.outputs)
for override in end_schema.inputs[4:]:
assert override.optional and override.lazy
assert override.get_io_type() == 'IMAGE,MASK,STRING,INT,FLOAT,BOOLEAN'
getter = mapping['DistributedImageListGetValue']
getter_schema = getter.GET_SCHEMA()
assert getter_schema.display_name == 'Distributed Image List Get Value'
assert not getter_schema.is_input_list
assert getter_schema.outputs[0].id == getter_schema.outputs[0].display_name == 'value'
assert getter_schema.outputs[0].io_type == '*'
assert getter_schema.inputs[1].options == list(runtime.VALUE_NAMES)
assert getter_schema.inputs[1].default == 'value_1'
# Real core maps scalar mappings, never arbitrary tensor batches.
tensor = torch.ones(3, 2, 2, 3)
data, missing, hidden = execution.get_input_data({'value_name': 'value_1'}, getter, 'get-value')
data['values'] = [{'value_1': tensor}, {'value_1': 17}, {'value_1': False}]
mapped = await execution._async_map_node_over_list('get-value', 'get-value', getter, data,
getter.FUNCTION, v3_data=hidden)
mapped = await execution.resolve_map_node_over_list_results(mapped)
assert len(mapped) == 3 and mapped[0].result[0] is tensor
assert mapped[1].result == (17,) and mapped[2].result == (False,)
try:
getter.execute({}, 'value_1')
raise AssertionError('missing Get Value key accepted')
except ValueError as exc:
assert 'absent' in str(exc)
texts_input = next(input for input in end_schema.inputs if input.id == 'texts')
texts_schema = texts_input.as_dict()
assert texts_schema.get('forceInput') is True, texts_schema
@@ -466,7 +499,6 @@ async def main():
rows = runtime_mod.normalize_image_list_inputs(prepared['images'], None, seen['values'])
assert [row['values'] for row in rows] == [{'value_1': 'shared', 'value_3': n} for n in (4, 5, 6)]
print('IMAGE_LIST_AUTOGROW_OK real parser; value_N slots shared/per-item')
assert [out.io_type for out in end_schema.outputs] == ['IMAGE', 'MASK', 'STRING']
assert all(out.is_output_list for out in end_schema.outputs)
assert end_schema.is_input_list and end_schema.is_output_node
print('IMAGE_LIST_SCHEMA_OK exact IDs, names, typed list outputs')
+55
View File
@@ -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()
+76 -1
View File
@@ -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,43 @@ 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_value_results_requires_exact_true_only_for_overrides(self):
for value in (None, False, 1, "true", [], True):
capability = _valid_capability()
if value is not None:
capability['value_results'] = value
session = _FakeSession({('worker', 'capabilities'): (200, capability),
('worker', 'queue'): (200, _valid_queue())})
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}"
):
self.assertEqual(await dispatch.preflight_image_list_workers([{'id': 'worker'}]), {'worker': True})
if value is True:
self.assertEqual(await dispatch.preflight_image_list_workers(
[{'id': 'worker'}], require_value_results=True), {'worker': True})
else:
with self.assertRaisesRegex(RuntimeError, 'value_results.*upgrade'):
await dispatch.preflight_image_list_workers([{'id': 'worker'}], require_value_results=True)
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)})
+16 -8
View File
@@ -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()
+245
View File
@@ -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()
+77
View File
@@ -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:
@@ -0,0 +1,20 @@
"""Run the real framework's CPU-only private-cache retention regression."""
import os
from pathlib import Path
import subprocess
import sys
import pytest
def test_native_local_cache_retention():
core = os.environ.get('COMFYUI_SOURCE_ROOT')
if not core:
pytest.skip('Set COMFYUI_SOURCE_ROOT for CPU-only real-framework validation')
helper = Path(__file__).parent / 'helpers' / 'local_cache_retention_check.py'
env = dict(os.environ, CUDA_VISIBLE_DEVICES='', HIP_VISIBLE_DEVICES='')
env.pop('PYTHONPATH', None)
result = subprocess.run([sys.executable, str(helper), core], env=env,
text=True, capture_output=True, timeout=90)
assert result.returncode == 0, result.stdout + result.stderr
assert 'PASS: cross-run native cache reuse' in result.stdout
+69
View File
@@ -199,6 +199,73 @@ class MasterRuntimeTests(unittest.IsolatedAsyncioTestCase):
with self.assertRaisesRegex(RuntimeError, 'worker|fallback'):
self.rt.create_master_job('start', 'end', graph(), self.items, [], '')
async def test_values_reconstructed_in_source_order_without_transporting_originals(self):
self.workers.append({'id': 'remote', 'type': 'remote'})
context = self.create(count=6)
tensor = torch.ones(4, 2, 2, 3, dtype=torch.float64)
for index, item in enumerate(self.items):
item['values'].update(value_2=tensor, value_3=True, value_4=index / 8)
self.local_delay = 0.1 # Remote results arrive ahead of master item zero.
with patch.object(self.rt, 'encode_result', side_effect=AssertionError('originals must not return')):
result = await self.rt.run_image_list_job(context, include_values=True)
self.assertFalse(self.rt._MASTER_JOBS)
self.assertEqual(len(result), 4)
for index, row in enumerate(result[3]):
self.assertEqual(row['value_1'], str(index))
self.assertIs(row['value_2'], tensor)
self.assertIs(row['value_3'], True)
self.assertEqual(row['value_4'], index / 8)
self.assertTrue(self.dispatched and self.local_calls)
async def test_overrides_merge_add_and_recheck_capability_before_every_item(self):
self.workers.append({'id': 'remote', 'type': 'remote'})
context = self.create(participates=False, count=3)
self.state.original_prompt['end']['inputs'].update(value_1=['start', 4], value_10=['start', 2])
original_set = self.rt.set_master_result
def override(context, index, token, result):
result['values'] = {'value_1': f'new-{index}', 'value_10': index}
return original_set(context, index, token, result)
with patch.object(self.rt, 'set_master_result', side_effect=override):
result = await self.rt.run_image_list_job(context, include_values=True)
self.assertEqual(result[3], [{'value_1': f'new-{index}', 'value_10': index} for index in range(3)])
self.assertEqual(self.preflight.await_count, 4)
self.assertTrue(all(call.kwargs['require_value_results'] is True
for call in self.preflight.await_args_list))
self.assertEqual(self.items[0]['values'], {'value_1': '0'})
async def test_override_capability_downgrade_stops_before_upload(self):
self.workers.append({'id': 'remote', 'type': 'remote'})
context = self.create(participates=False)
self.state.original_prompt['end']['inputs']['value_1'] = ['start', 4]
self.preflight.side_effect = [{'remote': True}, RuntimeError('value_results unsupported')]
with patch.object(self, 'post', wraps=self.post) as upload:
with self.assertRaisesRegex(RuntimeError, 'value_results'):
await self.rt.run_image_list_job(context, include_values=True)
upload.assert_not_called()
self.assertFalse(self.dispatched or self.rt._MASTER_JOBS)
async def test_missing_expected_override_fails_without_partial_output(self):
self.workers.append({'id': 'remote', 'type': 'remote'})
context = self.create(participates=False)
self.state.original_prompt['end']['inputs']['value_1'] = ['start', 4]
with self.assertRaisesRegex(ValueError, 'missing'):
await self.rt.run_image_list_job(context, include_values=True)
self.assertFalse(self.rt._MASTER_JOBS)
async def test_values_failure_and_cancellation_return_no_partial_lists(self):
self.local_error = RuntimeError('branch failed')
with self.assertRaisesRegex(RuntimeError, 'branch failed'):
await self.rt.run_image_list_job(self.create(), include_values=True)
self.local_error = None
self.local_started.clear()
self.local_delay = 0.1
task = asyncio.create_task(self.rt.run_image_list_job(self.create(), include_values=True))
await self.local_started.wait()
task.cancel()
with self.assertRaises(self.rt.ImageListCancelledError):
await task
self.assertFalse(self.rt._MASTER_JOBS)
async def test_local_failure_surfaces_and_reclaims_staged_token(self):
self.local_error = RuntimeError('local branch failed')
with self.assertRaisesRegex(RuntimeError, 'local branch failed'):
@@ -373,6 +440,8 @@ def real_framework_probe(comfy_root):
rt.remove_master_job(context)
for behavior in ('normal', 'expand', 'error', 'cancel', 'interrupt', 'sync', 'already-interrupted'):
# Independent lifecycle cases start cold; retention is tested separately.
local._RETAINED_OUTPUTS = None
mode[0] = behavior
prompt['op']['class_type'] = 'Expanded' if behavior == 'expand' else 'Operation'
loads.clear()
+177 -5
View File
@@ -1,8 +1,9 @@
"""Image-list participation uses the existing master orchestrator setting."""
import json
import sys
import types
import unittest
from unittest.mock import AsyncMock, patch
from unittest.mock import AsyncMock, Mock, patch
from test_serverless_integration import PKG, load_modules
@@ -34,6 +35,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)
@@ -45,11 +63,21 @@ class MasterSelectionTests(unittest.IsolatedAsyncioTestCase):
self.q.dispatch_worker_prompt.assert_not_awaited()
self.assertEqual(self.ps.distributed_pending_jobs, {})
async def test_end_overrides_require_value_capability_before_master_queue(self):
self.prompt['3']['inputs']['value_10'] = ['2', 2]
await self.queue()
self.assertIs(self.preflight.call_args.kwargs['require_value_results'], True)
self.q.queue_prompt_payload.reset_mock()
self.preflight.side_effect = RuntimeError('value_results unsupported')
with self.assertRaisesRegex(RuntimeError, 'value_results'):
await self.queue()
self.q.queue_prompt_payload.assert_not_awaited()
async def test_orchestrator_setting_excludes_master_processing(self):
self.config["settings"]["master_delegate_only"] = True
await self.queue()
self.assertIs(self.prepared_inputs()["image_list_master_participates"], False)
self.assertTrue(self.preflight.call_args.kwargs["require_idle"])
self.assertIs(self.preflight.call_args.kwargs["require_idle"], False)
# Upstream source stays on master even in orchestrator-only mode.
self.assertIn("1", self.q.queue_prompt_payload.call_args.args[0])
@@ -80,10 +108,154 @@ class MasterSelectionTests(unittest.IsolatedAsyncioTestCase):
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), [])
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
async def test_busy_workers_do_not_prevent_enabled_master(self):
self.preflight.return_value = {}
def use_real_preflight(self, queues):
"""Exercise capability/queue parsing with read-only fake HTTP responses."""
from test_dispatch_selection import dispatch, _FakeSession, _valid_capability
responses = {(worker_id, "capabilities"): (200, _valid_capability())
for worker_id in queues}
responses.update({(worker_id, "queue"): (200, queue)
for worker_id, queue in queues.items()})
session = _FakeSession(responses)
for target, value in (
("get_client_session", AsyncMock(return_value=session)),
("build_worker_url", lambda worker, path: f"http://{worker['id']}{path}"),
):
patched = patch.object(dispatch, target, value)
patched.start()
self.addCleanup(patched.stop)
self.preflight = AsyncMock(wraps=dispatch.preflight_image_list_workers)
sys.modules[PKG + ".api.orchestration.dispatch"].preflight_image_list_workers = self.preflight
return responses
async def run_prepared_selection(self):
"""Run real runtime selection up to the scheduler, without any item I/O."""
rt = sys.modules[PKG + ".image_list.runtime"]
inputs = self.prepared_inputs()
image = object()
context = rt.create_master_job(
"2", "3", self.q.queue_prompt_payload.call_args.args[0],
[{"image": image, "index": 0}],
json.loads(inputs["image_list_enabled_worker_ids"]),
inputs["image_list_master_callback"],
master_participates=inputs["image_list_master_participates"],
explicit_targets=True,
)
self.addCleanup(rt.remove_master_job, context)
self.scheduler = AsyncMock(return_value=[{"image": image}])
self.local_slot = Mock(return_value=types.SimpleNamespace(close=AsyncMock()))
self.runtime_session = AsyncMock(return_value=object())
config = sys.modules[PKG + ".utils.config"]
network = sys.modules[PKG + ".utils.network"]
with patch.object(config, "load_config", lambda: self.config), \
patch.object(config, "get_worker_timeout_seconds", lambda: 60, create=True), \
patch.object(config, "get_image_list_job_timeout_seconds", lambda: 0, create=True), \
patch.object(network, "get_client_session", self.runtime_session, create=True), \
patch.object(network, "build_worker_url", Mock(), create=True), \
patch.object(rt, "dynamic_schedule", self.scheduler), \
patch.dict(sys.modules, {PKG + ".image_list.local_executor": types.SimpleNamespace(
LocalExecutionSlot=self.local_slot)}):
await rt.run_image_list_job(context)
self.assertIsNone(rt.get_master_job(context))
return self.scheduler.call_args.args[1]
async def test_busy_worker_at_enqueue_is_available_to_scheduler_when_idle_at_start(self):
responses = self.use_real_preflight({
"worker": {"queue_running": [[1, "prompt-a"]], "queue_pending": []},
})
result = await self.queue()
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker"])
self.assertEqual(result[2], 1) # Candidates, not a promise of execution-time availability.
self.q.dispatch_worker_prompt.assert_not_awaited()
responses[("worker", "queue")] = (200, {"queue_running": [], "queue_pending": []})
# A newly enabled worker is not implicitly added to this queued selection.
self.config["workers"].append({"id": "later", "enabled": True})
slots = await self.run_prepared_selection()
self.assertEqual([slot for slot in slots if isinstance(slot, str)], ["worker"])
self.assertEqual(len(slots), 2)
self.assertEqual(self.preflight.await_count, 2)
self.assertEqual(self.preflight.await_args_list[0].args, self.preflight.await_args_list[1].args)
async def test_worker_still_busy_at_start_is_excluded_from_scheduler(self):
self.use_real_preflight({
"worker": {"queue_running": [], "queue_pending": [[1, "prompt-a"]]},
})
await self.queue()
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), [])
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker"])
slots = await self.run_prepared_selection()
self.assertEqual(len(slots), 1)
self.assertNotIn("worker", slots)
self.assertEqual(self.preflight.await_count, 2)
self.runtime_session.assert_not_awaited()
async def test_explicit_empty_selection_stays_master_only_at_start(self):
self.use_real_preflight({"worker": {"queue_running": [], "queue_pending": []}})
await self.queue(enabled_worker_ids=[])
slots = await self.run_prepared_selection()
self.assertEqual(len(slots), 1)
self.assertNotIn("worker", slots)
self.preflight.assert_not_awaited()
self.runtime_session.assert_not_awaited()
async def test_busy_incompatible_worker_fails_before_master_queue(self):
responses = self.use_real_preflight({
"worker": {"queue_running": [[1, "prompt-a"]], "queue_pending": []},
})
capability = responses[("worker", "capabilities")][1]
for field, invalid, message in (
("protocol_version", 1, "incompatible"),
("execution_targets", False, "execution_targets"),
("value_results", False, "value_results"),
):
with self.subTest(field=field), patch.dict(capability, {field: invalid}):
if field == "value_results":
self.prompt["3"]["inputs"]["value_10"] = ["2", 2]
with self.assertRaisesRegex(RuntimeError, message):
await self.queue()
self.q.queue_prompt_payload.assert_not_awaited()
self.q.dispatch_worker_prompt.assert_not_awaited()
async def test_orchestrator_with_all_workers_busy_queues_for_idle_start(self):
responses = self.use_real_preflight({
"worker": {"queue_running": [[1, "prompt-a"]], "queue_pending": []},
})
result = await self.queue(delegate_master=True)
self.q.queue_prompt_payload.assert_awaited_once()
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker"])
self.assertIs(self.prepared_inputs()["image_list_master_participates"], False)
self.assertEqual(result[2], 1)
self.q.dispatch_worker_prompt.assert_not_awaited()
responses[("worker", "queue")] = (200, {"queue_running": [], "queue_pending": []})
self.assertEqual(await self.run_prepared_selection(), ["worker"])
self.local_slot.assert_not_called()
self.assertEqual(self.preflight.await_count, 2)
self.assertIs(self.preflight.await_args_list[0].kwargs["require_idle"], False)
self.assertIs(self.preflight.await_args_list[1].kwargs["require_idle"], True)
async def test_orchestrator_rechecks_availability_without_master_fallback(self):
responses = self.use_real_preflight({"worker": {"queue_running": [], "queue_pending": []}})
await self.queue(delegate_master=True)
responses[("worker", "queue")] = (200, {"queue_running": [[1, "other"]], "queue_pending": []})
with self.assertRaisesRegex(RuntimeError, "No selected feature-compatible direct worker is idle"):
await self.run_prepared_selection()
self.assertEqual(self.preflight.await_count, 2)
self.scheduler.assert_not_awaited()
self.local_slot.assert_not_called()
async def test_orchestrator_retains_busy_candidates_when_another_worker_is_idle(self):
self.config["workers"].append({"id": "idle", "enabled": True})
responses = self.use_real_preflight({
"worker": {"queue_running": [[1, "prompt-a"]], "queue_pending": []},
"idle": {"queue_running": [], "queue_pending": []},
})
await self.queue(delegate_master=True)
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker", "idle"])
responses[("worker", "queue")] = (200, {"queue_running": [], "queue_pending": []})
responses[("idle", "queue")] = (200, {"queue_running": [[1, "other"]], "queue_pending": []})
self.assertEqual(await self.run_prepared_selection(), ["worker"])
self.local_slot.assert_not_called()
async def test_orchestrator_without_workers_fails_before_queue(self):
self.config["settings"]["master_delegate_only"] = True
@@ -0,0 +1,527 @@
"""Stable content signatures with fresh, expiring delivery capabilities."""
import asyncio
import copy
import importlib.util
import json
import os
from pathlib import Path
import subprocess
import sys
import types
from unittest.mock import AsyncMock
import pytest
import torch
from image_list import runtime as rt
ROOT = Path(__file__).resolve().parents[1]
def item():
return rt.normalize_image_list_inputs(
torch.arange(18, dtype=torch.float64).reshape(1, 2, 3, 3),
torch.zeros(1, 2, 3, dtype=torch.bfloat16),
{'value_1': 'original', 'value_2': True, 'value_3': 7, 'value_4': 0.25,
'value_10': torch.ones(2, 3, dtype=torch.float16)})[0]
def graph():
return {
'source': {'class_type': 'Source', 'inputs': {}},
'start': {'class_type': rt.START_CLASS, 'inputs': {
'images': ['source', 0], 'image_list_execution_id': 'transient',
'image_list_workflow_meta': {'job': 'metadata'},
'image_list_enabled_worker_ids': '["worker"]',
'image_list_master_callback': 'http://old-master',
'values.value_1': 'inherited', 'worker_context_id': 'old-context'}},
'op': {'class_type': 'Operation', 'inputs': {'image': ['start', 0]}},
'end': {'class_type': rt.END_CLASS, 'inputs': {
'images': ['op', 0], 'context': ['start', 3], 'value_3': ['start', 2]}},
}
def worker_prompt(token, context, data, *, local=True):
prompt = graph()
result = rt.build_worker_prompt(prompt, rt.validate_image_list_regions(prompt)[0],
token, context, data['index'], 'http://master',
item_identity=rt.item_content_identity(data))
if local:
result['start']['inputs']['image_list_local_execution'] = True
result['end']['inputs']['image_list_local_execution'] = True
return result
@pytest.fixture(autouse=True)
def clean_state():
yield
rt._WORKER_ITEMS.clear()
rt._MASTER_JOBS.clear()
@pytest.fixture
def boundaries(monkeypatch):
"""Load only the boundary module; restore every replaced module on exit."""
name = '_native_cache_identity_tests'
for suffix in ('', '.nodes', '.utils', '.image_list', '.api'):
module = types.ModuleType(name + suffix)
module.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
monkeypatch.setitem(sys.modules, name + suffix, module)
monkeypatch.setitem(sys.modules, name + '.image_list.runtime', rt)
io = types.SimpleNamespace(ComfyNode=type('ComfyNode', (), {}),
NodeOutput=lambda *args: args)
for cls in ('Image', 'Mask', 'String', 'Int', 'Float', 'Boolean'):
setattr(io, cls, object())
monkeypatch.setitem(sys.modules, 'comfy_api', types.ModuleType('comfy_api'))
monkeypatch.setitem(sys.modules, 'comfy_api.v0_0_2', types.SimpleNamespace(io=io))
monkeypatch.setitem(sys.modules, name + '.utils.async_helpers',
types.SimpleNamespace(run_async_in_server_loop=lambda coro, **kw: asyncio.run(coro)))
spec = importlib.util.spec_from_file_location(name + '.nodes.image_list', ROOT / 'nodes/image_list.py')
module = importlib.util.module_from_spec(spec)
monkeypatch.setitem(sys.modules, spec.name, module)
spec.loader.exec_module(module)
return module
def start(boundaries, prompt):
boundaries.DistributedImageListStart.hidden = types.SimpleNamespace(prompt=prompt, unique_id='start')
return boundaries.DistributedImageListStart.execute(**prompt['start']['inputs'])
def end(boundaries, prompt, cached):
boundaries.DistributedImageListEnd.hidden = types.SimpleNamespace(prompt=prompt, unique_id='end')
inputs = {k: v for k, v in prompt['end']['inputs'].items()
if k not in ('images', 'context', 'value_3')}
return boundaries.DistributedImageListEnd.execute(
images=cached[0], context=cached[3], value_3=cached[2], **inputs)
def stage(data):
token = rt.create_worker_task_token()
rt.store_worker_item(token, rt.encode_item(data))
return token
def master(data, token):
context = rt.create_master_job('start', 'end', graph(), [data], ['worker'], '')
rt.register_assignment(context, data['index'], 'worker', token)
return context
def test_identity_is_exact_typed_order_independent_and_not_transport_bounded(monkeypatch):
original = item()
identity = rt.item_content_identity(original)
assert rt.is_valid_item_identity(identity)
reordered = copy.deepcopy(original)
reordered['values'] = dict(reversed(list(reordered['values'].items())))
assert rt.item_content_identity(reordered) == identity
strided = copy.deepcopy(original)
strided['image'] = original['image'].transpose(1, 2).contiguous().transpose(1, 2)
assert not strided['image'].is_contiguous()
assert rt.item_content_identity(strided) == identity
monkeypatch.setattr(rt, 'MAX_PAYLOAD_BYTES', 1)
with pytest.raises(ValueError, match='transport limit'):
rt.encode_item(original)
assert rt.item_content_identity(original) == identity
scalar_identities = set()
for value in (True, 1, 1.0, '1', 0.0, -0.0):
changed = copy.deepcopy(original)
changed['values']['value_2'] = value
scalar_identities.add(rt.item_content_identity(changed))
assert len(scalar_identities) == 6
@pytest.mark.parametrize('change', ['image', 'mask', 'value', 'index', 'dtype', 'shape', 'slot', 'missing-mask'])
def test_every_item_component_invalidates_identity(change):
original, changed = item(), item()
if change in ('image', 'mask'):
changed[change].reshape(-1)[0] = 1
elif change == 'value':
changed['values']['value_10'][0, 0] = 2
elif change == 'index':
changed['index'] = 1
elif change == 'dtype':
changed['image'] = changed['image'].float()
elif change == 'shape':
changed['values']['value_10'] = changed['values']['value_10'].reshape(3, 2)
elif change == 'slot':
changed['values']['value_9'] = changed['values'].pop('value_10')
else:
changed['mask'] = None
assert rt.item_content_identity(original) != rt.item_content_identity(changed)
def test_float64_low_bits_and_bfloat16_bytes_are_not_quantized():
data, changed = item(), item()
changed['image'].reshape(-1)[0] = torch.nextafter(torch.tensor(0., dtype=torch.float64),
torch.tensor(1., dtype=torch.float64))
assert rt.item_content_identity(data) != rt.item_content_identity(changed)
changed = item()
changed['mask'].view(torch.int16).reshape(-1)[0] = 1
assert rt.item_content_identity(data) != rt.item_content_identity(changed)
def test_stable_start_replaced_wholesale_and_legacy_graph_preserved():
prompt = graph()
frozen = copy.deepcopy(prompt)
region = rt.validate_image_list_regions(prompt)[0]
legacy = rt.build_worker_prompt(prompt, region, 'a', 'ctx-a', 0, 'http://a')
assert legacy['start']['inputs']['worker_item_token'] == 'a'
assert legacy['start']['inputs']['image_list_execution_id'] == 'transient'
identity = rt.item_content_identity(item())
a = rt.build_worker_prompt(prompt, region, 'a', 'ctx-a', 0, 'http://a', item_identity=identity)
b = rt.build_worker_prompt(prompt, region, 'b', 'ctx-b', 0, 'http://b', item_identity=identity)
assert a['start']['inputs'] == b['start']['inputs'] == {'worker_item_identity': identity}
assert a['end']['inputs']['worker_context_id'] == 'ctx-a'
assert b['end']['inputs']['worker_result_token'] == 'b'
assert b['end']['inputs']['worker_item_identity'] == identity
assert b['end']['inputs']['value_3'] == ['start', 2]
assert prompt == frozen
@pytest.mark.parametrize('bad', ['', 'sha256:' + 'z' * 64, 'a' * 64, False, ['bad']])
def test_malformed_identity_rejected(bad):
assert not rt.is_valid_item_identity(bad)
prompt = graph()
with pytest.raises((ValueError, TypeError), match='identity'):
rt.build_worker_prompt(prompt, rt.validate_image_list_regions(prompt)[0],
'token', 'context', 0, '', item_identity=bad)
def test_cache_miss_then_hit_delivers_current_context_token_and_values(boundaries):
data = item()
a = stage(data)
context_a = master(data, a)
prompt_a = worker_prompt(a, context_a, data)
cached = start(boundaries, prompt_a)
assert cached[3] == [rt.item_content_identity(data)]
assert 'image' not in rt._WORKER_ITEMS[a][1] # only a small, expiring receipt remains
end(boundaries, prompt_a, cached)
assert a not in rt._WORKER_ITEMS
assert rt.get_master_job(context_a).results[0]['values'] == {'value_3': 0}
b = stage(data)
context_b = master(data, b)
prompt_b = worker_prompt(b, context_b, data)
assert prompt_a['start'] == prompt_b['start']
# A native hit does NOT execute Start; End must consume this current upload.
end(boundaries, prompt_b, cached)
assert b not in rt._WORKER_ITEMS
state = rt.get_master_job(context_b)
assert state.expected_tokens[0] == b
assert torch.equal(state.results[0]['image'], data['image'])
assert rt._collect_output_lists([state.results[0]], state=state, include_values=True)[3][0]['value_3'] == 0
with pytest.raises(RuntimeError, match='missing|expired|consumed'):
end(boundaries, prompt_b, cached)
@pytest.mark.parametrize('miss', [False, True])
def test_expired_item_cannot_deliver_even_with_cached_start(boundaries, miss):
data = item()
token = stage(data)
context = master(data, token)
prompt = worker_prompt(token, context, data)
cached = start(boundaries, prompt) if miss else ([data['image']], [data['mask']], [0], [rt.item_content_identity(data)])
_, staged = rt._WORKER_ITEMS[token]
rt._WORKER_ITEMS[token] = (0, staged)
with pytest.raises(RuntimeError, match='missing|expired|consumed'):
end(boundaries, prompt, cached)
assert not rt.get_master_job(context).received_indexes
assert token not in rt._WORKER_ITEMS
def test_identity_mismatch_on_start_and_on_cache_hit_fails_closed(boundaries):
data, different = item(), item()
different['image'][0, 0, 0, 0] = 3
for hit in (False, True):
token = stage(different)
context = master(different, token)
prompt = worker_prompt(token, context, data)
cached = ([data['image']], [data['mask']], [0], [rt.item_content_identity(data)])
with pytest.raises((ValueError, RuntimeError), match='identity'):
end(boundaries, prompt, cached) if hit else start(boundaries, prompt)
assert not rt.get_master_job(context).received_indexes
rt.delete_worker_item(token)
@pytest.mark.parametrize('change', ['ambiguous', 'wrong-link', 'wrong-start-id', 'end-identity', 'end-index'])
def test_start_requires_exact_paired_end_binding(boundaries, change):
data = item()
token = stage(data)
prompt = worker_prompt(token, 'context', data)
if change == 'ambiguous':
prompt['second-end'] = copy.deepcopy(prompt['end'])
elif change == 'wrong-link':
prompt['end']['inputs']['context'] = ['start', 2]
elif change == 'wrong-start-id':
prompt['renamed'] = prompt.pop('start')
elif change == 'end-identity':
prompt['end']['inputs']['worker_item_identity'] = 'sha256:' + '0' * 64
else:
prompt['end']['inputs']['worker_item_index'] = 1
boundaries.DistributedImageListStart.hidden = types.SimpleNamespace(prompt=prompt, unique_id='start')
with pytest.raises((RuntimeError, ValueError), match='paired|identity|index'):
boundaries.DistributedImageListStart.execute(worker_item_identity=rt.item_content_identity(data),
image_list_local_execution=True)
def test_end_checks_stable_context_before_delivering(boundaries):
data = item()
token = stage(data)
context = master(data, token)
prompt = worker_prompt(token, context, data)
cached = list(start(boundaries, prompt))
cached[3] = ['sha256:' + '0' * 64]
with pytest.raises((RuntimeError, ValueError), match='identity|context'):
end(boundaries, prompt, cached)
assert not rt.get_master_job(context).received_indexes
def test_legacy_worker_path_still_consumes_and_delivers_job_context(boundaries):
data = item()
token = stage(data)
context = master(data, token)
prompt = graph()
prompt = rt.build_worker_prompt(prompt, rt.validate_image_list_regions(prompt)[0], token, context, 0, '')
prompt['start']['inputs']['image_list_local_execution'] = True
prompt['end']['inputs']['image_list_local_execution'] = True
cached = start(boundaries, prompt)
assert cached[3] == [context]
assert token not in rt._WORKER_ITEMS
end(boundaries, prompt, cached)
assert rt.get_master_job(context).received_indexes == {0}
def test_remote_cache_hit_submits_fresh_callback_metadata(boundaries, monkeypatch):
data = item()
token = stage(data)
prompt = worker_prompt(token, 'fresh-context', data, local=False)
submit = AsyncMock()
monkeypatch.setattr(boundaries, 'submit_worker_result', submit)
cached = ([data['image']], [data['mask']], [0], [rt.item_content_identity(data)])
end(boundaries, prompt, cached)
assert submit.await_args.args[:4] == ('http://master', token, 'fresh-context', 0)
assert submit.await_args.args[4]['values']['value_3'] == {'type': 'INT', 'value': 0}
assert token not in rt._WORKER_ITEMS
def test_item_route_returns_server_computed_identity(boundaries, monkeypatch):
from aiohttp import web
name = '_native_cache_identity_tests'
server = types.SimpleNamespace(PromptServer=types.SimpleNamespace(instance=types.SimpleNamespace(routes=web.RouteTableDef())))
monkeypatch.setitem(sys.modules, 'server', server)
spec = importlib.util.spec_from_file_location(name + '.api.image_list_routes', ROOT / 'api/image_list_routes.py')
module = importlib.util.module_from_spec(spec)
monkeypatch.setitem(sys.modules, spec.name, module)
spec.loader.exec_module(module)
data = item()
token = rt.create_worker_task_token()
monkeypatch.setattr(module, '_json_body', AsyncMock(return_value={
'token': token, 'context_id': 'context', 'item': rt.encode_item(data),
'item_identity': 'sha256:' + '0' * 64}))
response = asyncio.run(module.image_list_item_endpoint(object()))
assert response.status == 200
assert json.loads(response.body)['item_identity'] == rt.item_content_identity(data)
assert len(rt._WORKER_ITEMS[token]) == 2
@pytest.mark.parametrize('responses', [
('valid', 'valid'), ('valid', 'absent'), ('invalid', 'valid'),
('bad-json', 'valid'), ('non-object', 'valid'),
])
def test_runtime_opts_in_per_upload_and_can_downgrade(monkeypatch, responses):
name = '_native_cache_identity_runtime'
for suffix in ('', '.image_list', '.api', '.api.orchestration', '.utils'):
module = types.ModuleType(name + suffix)
module.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
monkeypatch.setitem(sys.modules, name + suffix, module)
spec = importlib.util.spec_from_file_location(name + '.image_list.runtime', ROOT / 'image_list/runtime.py')
runtime = importlib.util.module_from_spec(spec)
monkeypatch.setitem(sys.modules, spec.name, runtime)
spec.loader.exec_module(runtime)
workers = [{'id': 'worker', 'type': 'local'}]
dispatched, uploads = [], []
preflight = AsyncMock(return_value={'worker': True})
sync_media = AsyncMock()
class Response:
status = 200
def __init__(self, data):
self.data = data
async def __aenter__(self):
return self
async def __aexit__(self, *args):
pass
async def json(self):
if self.data == 'bad-json':
raise ValueError('old response is not JSON')
return self.data
def raise_for_status(self):
pass
class Session:
def post(self, url, *, data, **kw):
assert url.endswith('/distributed/image-list/item')
body = json.loads(data)
identity = runtime.store_worker_item(body['token'], body['item'])
mode = responses[len(uploads)]
uploads.append(body)
response = {'valid': {'item_identity': identity}, 'absent': {'status': 'stored'},
'invalid': {'item_identity': 'bad'}, 'bad-json': 'bad-json', 'non-object': []}[mode]
return Response(response)
def get(self, url, **kw):
pid = url.rsplit('/', 1)[-1]
return Response({pid: {'status': {'completed': True, 'status_str': 'success'}}})
def delete(self, url, **kw):
runtime.delete_worker_item(url.rsplit('/', 1)[-1])
return Response({})
async def dispatch(worker, prompt, meta, **kwargs):
dispatched.append(copy.deepcopy(prompt))
delivery = prompt['end']['inputs']
token, index = delivery['worker_result_token'], delivery['worker_item_index']
identity = prompt['start']['inputs'].get('worker_item_identity')
options = {'item_identity': identity, 'item_index': index} if identity else {}
data = runtime.consume_worker_item(token, **options)
if identity:
runtime.finish_worker_item(token, identity, index)
runtime.set_master_result(context, index, token, {
'image': data['image'], 'text': str(index), 'values': {'value_3': index}})
assert kwargs['partial_execution_targets'] == ['end']
return f'worker-{index}'
dependencies = {
'.api.orchestration.dispatch': types.SimpleNamespace(
preflight_image_list_workers=preflight, dispatch_worker_prompt=dispatch),
'.api.orchestration.prompt_transform': types.SimpleNamespace(
_resolve_direct_workers_for_image_list=lambda *args: workers),
'.api.orchestration.media_sync': types.SimpleNamespace(sync_worker_media=sync_media),
'.utils.config': types.SimpleNamespace(load_config=lambda: {'settings': {}},
get_worker_timeout_seconds=lambda: 60),
'.utils.network': types.SimpleNamespace(build_master_callback_url=lambda *a, **kw: 'http://master',
build_worker_url=lambda worker, path: 'http://worker' + path,
get_client_session=AsyncMock(return_value=Session())),
'.utils.logging': types.SimpleNamespace(log=lambda *args: None),
}
for suffix, module in dependencies.items():
monkeypatch.setitem(sys.modules, name + suffix, module)
monkeypatch.setitem(sys.modules, 'aiohttp', types.SimpleNamespace(
ClientTimeout=lambda **kwargs: kwargs, ClientError=ConnectionError))
comfy = types.ModuleType('comfy')
comfy.model_management = types.SimpleNamespace(processing_interrupted=lambda: False)
monkeypatch.setitem(sys.modules, 'comfy', comfy)
monkeypatch.setitem(sys.modules, 'comfy.model_management', comfy.model_management)
items = [item(), item()]
items[1]['index'] = 1
context = runtime.create_master_job('start', 'end', graph(), items, ['worker'], '')
state = runtime.get_master_job(context)
try:
result = asyncio.run(runtime._run_image_list_job(state, include_values=True))
assert result[2] == ['0', '1']
assert [row['value_3'] for row in result[3]] == [0, 1]
assert len(uploads) == len(dispatched) == 2
assert sync_media.await_count == 1
for index, prompt in enumerate(dispatched):
start_inputs = prompt['start']['inputs']
if responses[index] == 'valid':
assert start_inputs == {'worker_item_identity': runtime.item_content_identity(items[index])}
assert prompt['end']['inputs']['worker_context_id'] == context
else:
assert 'worker_item_identity' not in start_inputs
assert start_inputs['worker_item_token'] == uploads[index]['token']
assert start_inputs['worker_context_id'] == context
assert 'worker_item_identity' not in prompt['end']['inputs']
assert prompt['end']['inputs']['worker_result_token'] == uploads[index]['token']
assert not runtime._WORKER_ITEMS
finally:
runtime.remove_master_job(context)
runtime._WORKER_ITEMS.clear()
def test_runtime_local_path_hashes_without_transport(monkeypatch):
name = '_native_cache_identity_local_runtime'
for suffix in ('', '.image_list', '.api', '.api.orchestration', '.utils'):
module = types.ModuleType(name + suffix)
module.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
monkeypatch.setitem(sys.modules, name + suffix, module)
spec = importlib.util.spec_from_file_location(name + '.image_list.runtime', ROOT / 'image_list/runtime.py')
runtime = importlib.util.module_from_spec(spec)
monkeypatch.setitem(sys.modules, spec.name, runtime)
spec.loader.exec_module(runtime)
dispatched = []
class Slot:
def __init__(self, **kw):
pass
def submit(self, prompt, prompt_id, end_id, workflow_meta):
dispatched.append(prompt)
delivery = prompt[end_id]['inputs']
identity = prompt['start']['inputs']['worker_item_identity']
data = runtime.consume_worker_item(delivery['worker_result_token'], item_identity=identity,
item_index=delivery['worker_item_index'])
runtime.finish_worker_item(delivery['worker_result_token'], identity, data['index'])
runtime.set_master_result(delivery['worker_context_id'], data['index'], delivery['worker_result_token'],
{'image': data['image'], 'values': {'value_3': data['index']}})
future = asyncio.get_running_loop().create_future()
future.set_result(None)
return future
def request_stop(self):
pass
async def close(self):
pass
def forbidden(*args, **kwargs):
raise AssertionError('local path attempted HTTP/transport')
dependencies = {
'.image_list.local_executor': types.SimpleNamespace(LocalExecutionSlot=Slot),
'.api.orchestration.dispatch': types.SimpleNamespace(
preflight_image_list_workers=forbidden, dispatch_worker_prompt=forbidden),
'.api.orchestration.prompt_transform': types.SimpleNamespace(
_resolve_direct_workers_for_image_list=lambda *args: []),
'.utils.config': types.SimpleNamespace(load_config=lambda: {'settings': {}},
get_worker_timeout_seconds=lambda: 60),
'.utils.network': types.SimpleNamespace(build_master_callback_url=forbidden,
build_worker_url=forbidden, get_client_session=forbidden),
'.utils.logging': types.SimpleNamespace(log=lambda *args: None),
}
for suffix, module in dependencies.items():
monkeypatch.setitem(sys.modules, name + suffix, module)
monkeypatch.setitem(sys.modules, 'aiohttp', types.SimpleNamespace(
ClientTimeout=lambda **kwargs: kwargs, ClientError=ConnectionError))
comfy = types.ModuleType('comfy')
comfy.model_management = types.SimpleNamespace(processing_interrupted=lambda: False)
monkeypatch.setitem(sys.modules, 'comfy', comfy)
monkeypatch.setitem(sys.modules, 'comfy.model_management', comfy.model_management)
monkeypatch.setattr(runtime, 'MAX_PAYLOAD_BYTES', 1)
monkeypatch.setattr(runtime, 'encode_item', forbidden)
data = item()
context = runtime.create_master_job('start', 'end', graph(), [data], [], '', master_participates=True)
try:
result = asyncio.run(runtime._run_image_list_job(runtime.get_master_job(context), include_values=True))
assert dispatched[0]['start']['inputs'] == {'worker_item_identity': runtime.item_content_identity(data),
'image_list_local_execution': True}
assert result[3][0]['value_3'] == 0
assert not runtime._WORKER_ITEMS
finally:
runtime.remove_master_job(context)
runtime._WORKER_ITEMS.clear()
def test_real_native_signatures_and_cached_delivery():
core = Path(os.environ.get('COMFYUI_SOURCE_ROOT', str(Path(sys.executable).parents[2])))
if not (core / 'comfy_execution/caching.py').exists():
pytest.skip('real ComfyUI source not available beside this interpreter')
result = subprocess.run([sys.executable, str(ROOT / 'tests/helpers/image_list_native_cache_identity.py'), str(core)],
text=True, capture_output=True, timeout=90)
assert result.returncode == 0, result.stdout + result.stderr
assert 'NATIVE_CACHE_IDENTITY_OK' in result.stdout
+19 -1
View File
@@ -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)
+165
View File
@@ -0,0 +1,165 @@
"""Context-carried values and sparse, per-item End overrides."""
import copy
from unittest.mock import patch
import pytest
import torch
from image_list import runtime as rt
from image_list.execution_plan import prepare_image_list_execution
def value_graph():
# Use a minimal, closed region with an override-only ancestor.
return {
'source': {'class_type': 'Source', 'inputs': {}},
'start': {'class_type': rt.START_CLASS, 'inputs': {'images': ['source', 0]}},
'override': {'class_type': 'Transform', 'inputs': {'value': ['start', 4]}},
'end': {'class_type': rt.END_CLASS, 'inputs': {
'images': ['start', 0], 'context': ['start', 3], 'value_1': ['override', 0],
'value_10': ['start', 5]}},
}
def values():
return {'value_1': torch.arange(72, dtype=torch.float64).reshape(2, 3, 4, 3),
'value_2': torch.ones(3, 2, 2, dtype=torch.bfloat16), 'value_3': 'name',
'value_4': 42, 'value_5': 0.125, 'value_6': True}
def assert_values(actual, expected):
assert actual.keys() == expected.keys()
for key, value in expected.items():
assert type(actual[key]) is type(value)
if isinstance(value, torch.Tensor):
assert actual[key].dtype == value.dtype
assert torch.equal(actual[key], value)
else:
assert actual[key] == value
def test_normalize_and_roundtrip_all_override_kinds_keeps_batches_whole():
expected = values()
result = rt.normalize_worker_result(torch.zeros(1, 2, 2, 3),
values={key: [value] for key, value in expected.items()})
assert_values(result['values'], expected)
wire = rt.encode_result(result)
assert_values(rt.decode_result(wire)['values'], expected)
assert wire['values']['value_1']['shape'][0] == 2
assert wire['values']['value_2']['shape'][0] == 3
@pytest.mark.parametrize('bad', [None, [], [None], [1, 2], [[1]], object(),
float('nan'), float('inf'), torch.tensor([float('nan')]),
torch.tensor([float('inf')]), torch.ones(1, dtype=torch.int32)])
def test_invalid_override_rejected(bad):
with pytest.raises((ValueError, TypeError)):
rt.normalize_worker_result(torch.zeros(1, 2, 2, 3), values={'value_1': bad})
def test_unknown_override_names_and_wire_types_rejected():
for name in ('value_0', 'value_11', 'values.value_1'):
with pytest.raises(ValueError, match='slot'):
rt.normalize_worker_result(torch.zeros(1, 2, 2, 3), values={name: 1})
base = rt.encode_result({'image': torch.zeros(1, 2, 2, 3)})
for payload in (None, [], {'value_0': {'type': 'INT', 'value': 1}},
{'value_1': {'type': 'LATENT', 'value': {}}},
{'value_1': {'type': 'FLOAT', 'value': float('nan')}},
{'value_1': {'type': 'STRING', 'value': None}},
{'value_1': rt.encode_value(torch.tensor([float('inf')]), 'TENSOR')}):
with pytest.raises((ValueError, TypeError)):
rt.decode_result({**base, 'values': payload})
def test_no_overrides_keeps_legacy_wire_and_does_not_encode_original_tensors():
image = torch.zeros(1, 2, 2, 3)
result = rt.normalize_worker_result(image)
assert set(result) == {'image', 'mask', 'text'}
with patch.object(rt, 'encode_value', wraps=rt.encode_value) as encode:
wire = rt.encode_result(result)
assert set(wire) == {'image', 'mask', 'text'}
assert [call.args[1] for call in encode.call_args_list] == ['IMAGE', 'MASK', 'STRING']
assert set(rt.decode_result(wire)) == {'image', 'mask', 'text'}
def test_override_ancestors_survive_plan_and_worker_prompt():
prompt = value_graph()
frozen = copy.deepcopy(prompt)
prepared, targets = prepare_image_list_execution(prompt)
assert prepared == frozen and prompt == frozen
assert targets == ['end']
region = rt.validate_image_list_regions(prepared)[0]
assert 'override' in region['branch_ids']
worker = rt.build_worker_prompt(prepared, region, 'token', 'context', 0, 'http://master')
assert worker['end']['inputs']['value_1'] == ['override', 0]
assert worker['end']['inputs']['value_10'] == ['start', 5]
assert 'source' not in worker
for bad in ('value_0', 'value_11', 'values.value_1'):
invalid = copy.deepcopy(prompt)
invalid['end']['inputs'][bad] = ['start', 4]
with pytest.raises(ValueError, match='slot'):
rt.validate_image_list_regions(invalid)
def test_missing_or_unexpected_result_override_fails_before_recording():
prompt = value_graph()
items = rt.normalize_image_list_inputs(torch.zeros(1, 2, 2, 3), values={'value_1': 'original'})
context = rt.create_master_job('start', 'end', prompt, items, ['worker'], '')
try:
rt.register_assignment(context, 0, 'worker', 'token')
for result_values in ({}, {'value_1': 'updated'},
{'value_1': 'updated', 'value_10': 10, 'value_9': False}):
with pytest.raises(ValueError, match='override'):
rt.set_master_result(context, 0, 'token', {'image': items[0]['image'], 'values': result_values})
assert rt.get_master_job(context).received_indexes == set()
rt.set_master_result(context, 0, 'token', {'image': items[0]['image'],
'values': {'value_1': 'updated', 'value_10': 10}})
assert rt.get_master_job(context).received_indexes == {0}
finally:
rt.remove_master_job(context)
def test_override_return_request_budget_is_aggregate():
result = rt.normalize_worker_result(torch.zeros(1, 2, 2, 3), values={'value_1': ['x' * 500]})
with patch.object(rt, 'MAX_REQUEST_BYTES', 400):
with pytest.raises(ValueError, match='request'):
rt.encode_request_payload({'result': rt.encode_result(result)})
class TestResultHTTP:
def test_missing_and_invalid_overrides_fail_through_real_http(self):
import asyncio
from aiohttp import web
from aiohttp.test_utils import TestClient, TestServer
from test_image_list_integration_gaps import load_routes
async def check():
_module, routes = load_routes()
app = web.Application()
app.add_routes(routes)
prompt = value_graph()
items = rt.normalize_image_list_inputs(torch.zeros(1, 2, 2, 3), values={'value_1': 'original'})
context = rt.create_master_job('start', 'end', prompt, items, ['worker'], '')
token = rt.create_worker_task_token()
rt.register_assignment(context, 0, 'worker', token)
try:
async with TestClient(TestServer(app)) as client:
caps = await (await client.get('/distributed/image-list/capabilities')).json()
assert caps['value_results'] is True and caps['protocol_version'] == 2
result = rt.encode_result({'image': items[0]['image']})
envelope = {'context_id': context, 'token': token, 'item_index': 0, 'result': result}
response = await client.post('/distributed/image-list/result', json=envelope)
assert response.status == 400
assert 'missing' in (await response.json())['error']
assert not rt.get_master_job(context).received_indexes
result['values'] = {'value_1': {'type': 'INT', 'value': [1, 2]},
'value_10': {'type': 'INT', 'value': 10}}
response = await client.post('/distributed/image-list/result', json=envelope)
assert response.status == 400
assert not rt.get_master_job(context).received_indexes
result['values']['value_1'] = {'type': 'STRING', 'value': 'new'}
response = await client.post('/distributed/image-list/result', json=envelope)
assert response.status == 200
assert rt.get_master_job(context).results[0]['values'] == {'value_1': 'new', 'value_10': 10}
finally:
rt.remove_master_job(context)
asyncio.run(check())
+2 -1
View File
@@ -69,7 +69,8 @@ def v3(monkeypatch):
module(f"{PKG}.nodes.utilities")
module(f"{PKG}.nodes.distributed_upscale", UltimateSDUpscaleDistributed=object)
module(f"{PKG}.nodes.collector", DistributedCollectorNode=CollectorRuntime)
module(f"{PKG}.nodes.image_list", DistributedImageListStart=object, DistributedImageListEnd=object)
module(f"{PKG}.nodes.image_list", DistributedImageListStart=object, DistributedImageListEnd=object,
DistributedImageListGetValue=object)
module("comfy", __path__=[])
module("comfy.samplers", KSampler=types.SimpleNamespace(SAMPLERS=["euler"], SCHEDULERS=["normal"]))
module("comfy_api", __path__=[])
+14
View File
@@ -1,4 +1,5 @@
"""Real-framework acceptance; opt in with COMFYUI_SOURCE_ROOT."""
import json
import os
from pathlib import Path
import subprocess
@@ -33,5 +34,18 @@ def test_image_list_network_runtime():
)
assert result.returncode == 0, result.stdout + '\n' + result.stderr
assert 'IMAGE_LIST_NETWORK_OK' in result.stdout
assert 'IMAGE_LIST_VALUES_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
View File
@@ -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
+174
View File
@@ -0,0 +1,174 @@
import { app } from "/scripts/app.js";
// End's Python schema owns all ten optional override slots. Only presentation
// grows here: unlike Start's Autogrow inputs, End keys must never compact, since
// value_3 always overrides source value_3, even if value_2 is disconnected.
const START = "DistributedImageListStart";
const END = "DistributedImageListEnd";
const GET = "DistributedImageListGetValue";
const PIPE = "DISTRIBUTED_IMAGE_LIST_VALUES";
const VALUE_TYPES = "IMAGE,MASK,STRING,INT,FLOAT,BOOLEAN";
const MAX_VALUES = 10;
const pendingGraphs = new WeakSet();
function ordinal(name) {
const match = /^value_([1-9]|10)$/.exec(name ?? "");
return match ? Number(match[1]) : null;
}
function getLink(graph, id) {
if (id == null) return undefined;
if (typeof graph?.getLink === "function") return graph.getLink(id);
return graph?.links?.get?.(id) ?? graph?.links?.[id];
}
function sourceOf(node, inputName) {
const input = node.inputs?.find(slot => slot.name === inputName);
const link = getLink(node.graph, input?.link);
const source = link && node.graph?.getNodeById?.(link.origin_id);
return source ? { node: source, slot: link.origin_slot, link } : null;
}
function isReroute(node) {
return [node?.comfyClass, node?.type].includes("Reroute");
}
function throughReroutes(source) {
const seen = new Set();
while (source && isReroute(source.node)) {
if (seen.has(source.node)) return null;
seen.add(source.node);
source = sourceOf(source.node, source.node.inputs?.[0]?.name);
}
return source;
}
function sourceType(source, seen) {
source = throughReroutes(source);
if (!source) return "*";
const { node, slot, link } = source;
if (node.comfyClass === GET) return getterType(node, seen);
const type = node.outputs?.[slot]?.type;
return type && type !== "*" ? type : (link.type || "*");
}
function getterType(node, seen = new Set()) {
if (seen.has(node)) return "*";
seen.add(node);
const name = node.widgets?.find(widget => widget.name === "value_name")?.value;
if (!ordinal(name)) return "*";
const endSource = throughReroutes(sourceOf(node, "values"));
if (endSource?.node.comfyClass !== END || endSource.slot !== 3) return "*";
const end = endSource.node;
const override = sourceOf(end, name);
if (override) return sourceType(override, seen);
const startSource = throughReroutes(sourceOf(end, "context"));
if (startSource?.node.comfyClass !== START || startSource.slot !== 3) return "*";
return sourceType(sourceOf(startSource.node, `values.${name}`), seen);
}
function compatible(outputType, inputType) {
if (!outputType || !inputType || outputType === "*" || inputType === "*") return true;
return String(outputType).split(",").some(type => String(inputType).split(",").includes(type));
}
function syncGetter(node) {
const output = node.outputs?.[0];
if (!output) return;
const type = getterType(node);
if (output.type === type) return;
output.type = type;
for (const id of [...(output.links ?? [])]) {
const link = getLink(node.graph, id);
if (!link) continue;
const target = node.graph?.getNodeById?.(link.target_id);
if (target && !compatible(type, target.inputs?.[link.target_slot]?.type)) {
target.disconnectInput(link.target_slot);
} else {
link.type = type;
}
}
node.setDirtyCanvas?.(true, true);
}
function syncEnd(node) {
if (!Array.isArray(node.inputs) || !Array.isArray(node.outputs)) return;
// Old saved End nodes have only the three original outputs. Append rather
// than replacing any socket, preserving every existing output index/link.
if (node.outputs.length === 3) {
const shape = node.outputs[0]?.shape;
node.addOutput("values", PIPE, shape === undefined ? undefined : { shape });
}
const connected = node.inputs.filter(input => ordinal(input.name) && input.link != null);
const highest = Math.max(0, ...connected.map(input => ordinal(input.name)));
const visible = Math.min(MAX_VALUES, Math.max(highest + 1, node._distributedVisibleOverrides ?? 1));
for (let i = node.inputs.length - 1; i >= 0; i--) {
const input = node.inputs[i];
if (ordinal(input.name) > visible && input.link == null) node.removeInput(i);
}
for (let i = 1; i <= visible; i++) {
if (!node.inputs.some(input => input.name === `value_${i}`)) {
node.addInput(`value_${i}`, VALUE_TYPES);
}
}
node.setDirtyCanvas?.(true, true);
}
function schedule(graph) {
if (!graph || pendingGraphs.has(graph)) return;
pendingGraphs.add(graph);
const run = () => {
// Keep the guard while syncing, since disconnectInput can re-enter the
// connection hook. A fresh user event schedules a fresh traversal.
try {
for (const node of graph._nodes ?? []) if (node.comfyClass === END) syncEnd(node);
for (const node of graph._nodes ?? []) if (node.comfyClass === GET) syncGetter(node);
} catch (error) {
console.error("[Distributed] Image List values presentation failed:", error);
} finally {
pendingGraphs.delete(graph);
}
};
// Let native Start Autogrow and its existing output-sync finish first.
if (typeof requestAnimationFrame === "function") requestAnimationFrame(() => setTimeout(run, 0));
else setTimeout(run, 0);
}
app.registerExtension({
name: "Distributed.ImageListEndValues",
nodeCreated(node) {
if (![START, END, GET].includes(node.comfyClass) && !isReroute(node)) return;
if (node.comfyClass === END) {
node.addWidget("button", "Add value override", null, () => {
const shown = Math.max(0, ...(node.inputs ?? []).map(input => ordinal(input.name) ?? 0));
node._distributedVisibleOverrides = Math.min(MAX_VALUES, shown + 1);
schedule(node.graph);
}, { serialize: false });
}
for (const hook of ["onConnectionsChange", "onConfigure", "onAdded", "onWidgetChanged"]) {
const original = node[hook];
node[hook] = function (...args) {
const result = original?.apply(this, args);
schedule(this.graph);
return result;
};
}
// Widget callbacks are also used by frontend versions that do not call
// the node's onWidgetChanged hook for a combo selection.
if (node.comfyClass === GET) {
const widget = node.widgets?.find(widget => widget.name === "value_name");
if (widget) {
const callback = widget.callback;
widget.callback = function (...args) {
const result = callback?.apply(this, args);
schedule(node.graph);
return result;
};
}
}
schedule(node.graph);
},
loadedGraphNode(node) {
schedule(node.graph);
},
});
+190
View File
@@ -0,0 +1,190 @@
import { readFileSync } from "node:fs";
import vm from "node:vm";
import { describe, expect, it } from "vitest";
const source = readFileSync(new URL("../imageListValues.js", import.meta.url), "utf8")
.replace(/^import .*;\r?\n/gm, "");
const PIPE = "DISTRIBUTED_IMAGE_LIST_VALUES";
function setup() {
let extension;
let nextId = 1;
const tasks = [];
const graph = {
_nodes: [], links: new Map(),
getNodeById(id) { return this._nodes.find(node => node.id === id); },
getLink(id) { return this.links.get(id); },
};
vm.runInNewContext(source, {
app: { registerExtension: ext => { extension = ext; } },
setTimeout: callback => tasks.push(callback), console,
});
function node(comfyClass, inputs, outputs, widgets = []) {
const n = {
id: nextId++, comfyClass, graph, widgets,
inputs: inputs.map(name => ({ name, type: "*", link: null })),
outputs: outputs.map(([name, type]) => ({ name, type, links: [], shape: 6 })),
addWidget(type, name, value, callback, options) { this.widgets.push({ type, name, value, callback, options }); },
addInput(name, type) { this.inputs.push({ name, type, link: null }); },
addOutput(name, type, extra) { this.outputs.push({ name, type, links: [], ...extra }); },
removeInput(slot) {
if (this.inputs[slot].link != null) throw new Error("Removed wired input");
this.inputs.splice(slot, 1);
this.inputs.forEach((input, index) => {
if (input.link != null) graph.links.get(input.link).target_slot = index;
});
},
disconnectInput(slot) {
const input = this.inputs[slot];
const link = graph.links.get(input.link);
if (link) {
const out = graph.getNodeById(link.origin_id).outputs[link.origin_slot];
out.links = out.links.filter(id => id !== link.id);
graph.links.delete(link.id);
}
input.link = null;
this.onConnectionsChange?.();
},
setDirtyCanvas() {},
};
graph._nodes.push(n);
extension.nodeCreated(n);
return n;
}
const flush = () => { while (tasks.length) tasks.shift()(); };
const connect = (from, slot, to, name) => {
const targetSlot = to.inputs.findIndex(input => input.name === name);
if (targetSlot < 0) throw new Error(`Missing input ${name}`);
to.disconnectInput(targetSlot);
const id = nextId++;
graph.links.set(id, { id, origin_id: from.id, origin_slot: slot, target_id: to.id,
target_slot: targetSlot, type: from.outputs[slot].type });
from.outputs[slot].links.push(id);
to.inputs[targetSlot].link = id;
from.onConnectionsChange?.();
to.onConnectionsChange?.();
return id;
};
const start = node("DistributedImageListStart", ["images", "mask", "values.value_1", "values.value_2", "values.value_3"],
[["image", "IMAGE"], ["mask", "MASK"], ["item_index", "INT"], ["context", "DISTRIBUTED_IMAGE_LIST_CONTEXT"]]);
const end = node("DistributedImageListEnd", ["images", "masks", "texts", "context", ...Array.from({ length: 10 }, (_, i) => `value_${i + 1}`)],
[["image_list", "IMAGE"], ["masks", "MASK"], ["texts", "STRING"], ["values", PIPE]]);
const get = node("DistributedImageListGetValue", ["values"], [["value", "*"]], [{ name: "value_name", value: "value_1" }]);
connect(start, 3, end, "context");
connect(end, 3, get, "values");
flush();
return { node, graph, start, end, get, flush, connect, extension };
}
const valueNames = node => node.inputs.filter(input => input.name.startsWith("value_")).map(input => input.name);
const addOverride = h => { h.end.widgets.find(widget => widget.name === "Add value override").callback(); h.flush(); };
describe("Image List End context values", () => {
it("keeps original sockets and shows only the first optional override initially", () => {
const h = setup();
expect(h.end.inputs.map(input => input.name)).toEqual(["images", "masks", "texts", "context", "value_1"]);
expect(h.end.outputs.map(output => output.name)).toEqual(["image_list", "masks", "texts", "values"]);
expect(h.end.widgets[0].options.serialize).toBe(false);
});
it("appends the values output when loading an old three-output End", () => {
const h = setup();
const old = h.node("DistributedImageListEnd", ["images", "masks", "texts", "context"],
[["image_list", "IMAGE"], ["masks", "MASK"], ["texts", "STRING"]]);
h.flush();
expect(old.outputs[3]).toMatchObject({ name: "values", type: PIPE, shape: 6 });
expect(valueNames(old)).toEqual(["value_1"]);
});
it("grows overrides without compacting field identities or existing links", () => {
const h = setup();
addOverride(h);
addOverride(h);
const src = h.node("PrimitiveInt", [], [["INT", "INT"]]);
const thirdLink = h.connect(src, 0, h.end, "value_3");
h.connect(src, 0, h.end, "value_2");
h.flush();
h.end.disconnectInput(h.end.inputs.findIndex(input => input.name === "value_2"));
h.flush();
expect(valueNames(h.end)).toEqual(["value_1", "value_2", "value_3", "value_4"]);
const link = h.graph.links.get(thirdLink);
expect(h.end.inputs[link.target_slot].name).toBe("value_3");
for (let i = 0; i < 20; i++) addOverride(h);
expect(valueNames(h.end)).toHaveLength(10);
});
it.each(["IMAGE", "MASK", "STRING", "INT", "FLOAT", "BOOLEAN"])("infers %s from the original Start value through context", type => {
const h = setup();
const source = h.node("Source", [], [["value", type]]);
h.connect(source, 0, h.start, "values.value_1");
h.flush();
expect(h.get.outputs[0].type).toBe(type);
});
it("prefers the End override type, then returns to the original after disconnect", () => {
const h = setup();
h.connect(h.node("Source", [], [["value", "INT"]]), 0, h.start, "values.value_1");
h.connect(h.node("Source", [], [["value", "STRING"]]), 0, h.end, "value_1");
h.flush();
expect(h.get.outputs[0].type).toBe("STRING");
h.end.disconnectInput(h.end.inputs.findIndex(input => input.name === "value_1"));
h.flush();
expect(h.get.outputs[0].type).toBe("INT");
});
it("updates a selected field and disconnects incompatible consumers", () => {
const h = setup();
h.connect(h.node("Source", [], [["value", "STRING"]]), 0, h.start, "values.value_1");
h.connect(h.node("Source", [], [["value", "INT"]]), 0, h.start, "values.value_2");
h.flush();
const consumer = h.node("Save", ["filename"], []);
consumer.inputs[0].type = "STRING";
h.connect(h.get, 0, consumer, "filename");
h.get.widgets[0].value = "value_2";
h.get.widgets[0].callback();
h.flush();
expect(h.get.outputs[0].type).toBe("INT");
expect(consumer.inputs[0].link).toBe(null);
});
it("keeps compatible consumer links and uses wildcard for unavailable graph metadata", () => {
const h = setup();
h.connect(h.node("Source", [], [["value", "INT"]]), 0, h.start, "values.value_1");
h.flush();
const consumer = h.node("Consumer", ["value"], []);
consumer.inputs[0].type = "INT,FLOAT";
const id = h.connect(h.get, 0, consumer, "value");
h.get.disconnectInput(0);
h.flush();
expect(h.get.outputs[0].type).toBe("*");
expect(h.graph.links.has(id)).toBe(true);
});
it("follows reroutes for context and values, and terminates on a cycle", () => {
const h = setup();
h.connect(h.node("Source", [], [["value", "FLOAT"]]), 0, h.start, "values.value_1");
const reroute = h.node("Reroute", ["input"], [["output", "*"]]);
h.connect(h.end, 3, reroute, "input");
h.connect(reroute, 0, h.get, "values");
h.flush();
expect(h.get.outputs[0].type).toBe("FLOAT");
h.connect(reroute, 0, reroute, "input");
h.flush();
expect(h.get.outputs[0].type).toBe("*");
});
it("recomputes types across sequential regions", () => {
const h = setup();
const secondStart = h.node("DistributedImageListStart", ["values.value_1"],
[["image", "IMAGE"], ["mask", "MASK"], ["item_index", "INT"], ["context", "DISTRIBUTED_IMAGE_LIST_CONTEXT"]]);
const secondEnd = h.node("DistributedImageListEnd", ["images", "masks", "texts", "context"],
[["image_list", "IMAGE"], ["masks", "MASK"], ["texts", "STRING"], ["values", PIPE]]);
const secondGet = h.node("DistributedImageListGetValue", ["values"], [["value", "*"]], [{ name: "value_name", value: "value_1" }]);
h.connect(h.node("Source", [], [["value", "BOOLEAN"]]), 0, h.start, "values.value_1");
h.connect(h.get, 0, secondStart, "values.value_1");
h.connect(secondStart, 3, secondEnd, "context");
h.connect(secondEnd, 3, secondGet, "values");
h.flush();
expect(secondGet.outputs[0].type).toBe("BOOLEAN");
});
});