Compare commits

..
Author SHA1 Message Date
Robert Wojciechowski fef70b7d94 Update pyproject.toml 2026-10-14 21:51:03 +11:00
Robert Wojciechowski c64338e566 Update pyproject.toml 2026-10-14 21:50:38 +11:00
Robert Wojciechowski a399556ee0 Merge pull request #92 from robertvoy/fix/issue-91-preserve-image-alpha
fix: preserve alpha and validate distributed image batches
2026-10-14 03:47:36 -07:00
Robert Wojciechowski 5159cc7166 fix: integrate alpha preservation with current collector transports
Merge main while retaining exact transfer, serverless cleanup and image assembly error propagation. Update RGB fixtures and verify mixed RGB/RGBA exact HTTP transport without value loss.
2026-10-14 10:46:16 +00:00
Robert Wojciechowski f149cbdf4d fix: preserve image-list target proof after cache fingerprinting 2026-10-14 04:37:56 +00:00
Robert Wojciechowski efa5dd624c Merge pull request #102 from robertvoy/feat/image-list-serverless
feat: schedule image-list items on bounded serverless slots
2026-10-13 21:18:54 -07:00
Robert Wojciechowski 0cc1631eec Merge pull request #101 from robertvoy/feat/serverless-multimodal
feat: add versioned lossless serverless media transport
2026-10-13 21:18:39 -07:00
Robert Wojciechowski 1940208035 Merge pull request #100 from robertvoy/feat/image-list-execution-targets
feat: preserve image-list execution targets and context values
2026-10-13 21:18:23 -07:00
Robert Wojciechowski 0ec9f5d4bd Merge pull request #104 from robertvoy/fix/image-list-native-cache-reuse
fix: preserve native image-list processing cache reuse
2026-10-13 21:18:09 -07:00
Robert Wojciechowski 98987db131 feat: integrate native image-list caching and exact serverless value overrides 2026-10-14 04:08:48 +00:00
Robert Wojciechowski 2df73e4b9f feat: add image-list return format and quality controls 2026-10-14 00:44:45 +00:00
Robert Wojciechowski b14c51c365 Merge branch 'feat/image-list-serverless-base' into feat/image-list-serverless 2026-10-14 00:37:48 +00:00
Robert Wojciechowski 3699477f75 Merge branch 'feat/serverless-multimodal' into feat/image-list-serverless-base 2026-10-14 00:37:48 +00:00
Robert Wojciechowski e9c95bf110 feat: schedule image-list items on bounded serverless slots 2026-10-14 00:16:54 +00:00
Robert Wojciechowski 1fb4d4a284 fix: preserve mixed-mode preflight across execution planner integration 2026-10-13 23:56:35 +00:00
Robert Wojciechowski 4e823397b7 feat: preserve required output utilities in image-list execution 2026-10-13 23:55:26 +00:00
Robert Wojciechowski 213d491f06 fix: preserve image alpha in distributed collection 2026-10-10 05:59:52 +00:00
37 changed files with 3359 additions and 178 deletions
+2
View File
@@ -5,6 +5,7 @@ from aiohttp import web
import server
from ..image_list.runtime import (
IMAGE_TRANSFER_FORMATS,
MAX_REQUEST_BYTES,
PROTOCOL_VERSION,
decode_result,
@@ -40,6 +41,7 @@ async def image_list_capabilities_endpoint(_request):
"protocol_version": PROTOCOL_VERSION,
"node_ids": ["DistributedImageListStart", "DistributedImageListEnd", "DistributedImageListGetValue"],
"transport": "typed-float-tensor-json-v1",
"image_transfer_formats": list(IMAGE_TRANSFER_FORMATS),
})
+3 -1
View File
@@ -136,7 +136,9 @@ def _decode_canonical_png_tensor(image_payload):
try:
with Image.open(io.BytesIO(png_bytes)) as img:
img = img.convert("RGB")
# PNG callbacks must retain alpha just like local IMAGE tensors.
has_alpha = "A" in img.getbands() or "transparency" in img.info
img = img.convert("RGBA" if has_alpha else "RGB")
tensor = pil_to_tensor(img)
return ensure_contiguous(tensor)
except Exception as exc:
+20 -1
View File
@@ -188,7 +188,8 @@ 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, *, require_value_results=False):
async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0, *,
required_transfer_formats=(), 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
@@ -196,6 +197,14 @@ async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0,
"""
if not workers:
raise RuntimeError("Distributed Image List requires at least one selected direct worker")
if (not isinstance(required_transfer_formats, (list, tuple, set, frozenset))
or any(type(fmt) is not str for fmt in required_transfer_formats)):
raise ValueError("required_transfer_formats must contain image transfer format names")
required_formats = set(required_transfer_formats) - {"exact"}
if required_formats:
from ...serverless_bridge.image_transfer import validate_options
for fmt in required_formats:
validate_options(fmt)
session = await get_client_session()
required_nodes = {"DistributedImageListStart", "DistributedImageListEnd"}
semaphore = asyncio.Semaphore(8)
@@ -242,6 +251,16 @@ async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0,
"upgrade that worker before running this feature; no work was submitted"
)
if required_formats:
advertised = capability.get("image_transfer_formats")
if (not isinstance(advertised, list)
or any(type(fmt) is not str for fmt in advertised)
or not required_formats.issubset(advertised)):
raise RuntimeError(
f"Image-list worker {worker_id} lacks required image_transfer_formats "
f"{sorted(required_formats)}; upgrade that worker before compressed IMAGE returns; "
"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; "
+65 -43
View File
@@ -13,7 +13,7 @@ import hashlib
import json
from graphlib import CycleError, TopologicalSorter
from ...serverless_bridge.contract import SOCKET_KINDS, identifier, validate_bindings
from ...serverless_bridge.contract import SOCKET_KINDS, VALUE_SOCKETS, identifier, validate_bindings
from ...serverless_bridge.image_transfer import DEFAULT_FORMAT, validate_options
from .serverless_prompt import (
_PRIMITIVE_NODES,
@@ -54,7 +54,7 @@ def _validate_identifier(value, context):
raise ValueError(f"{context} must be a valid nonempty media identifier") from None
def _validate_prompt(prepared_prompt):
def _validate_prompt(prepared_prompt, *, require_collectors=True):
"""Validate the entire source graph, including branches later discarded."""
if not isinstance(prepared_prompt, dict) or not prepared_prompt:
raise ValueError("prepared prompt must be a nonempty API graph mapping")
@@ -77,7 +77,7 @@ def _validate_prompt(prepared_prompt):
if any(not isinstance(name, str) or not name for name in inputs):
raise ValueError(f"prompt node {node_id!r} input names must be nonempty strings")
if not collectors:
if require_collectors and not collectors:
raise ValueError("prepared prompt contains no DistributedCollector nodes")
dependencies = {}
@@ -140,52 +140,74 @@ def _validate_prompt(prepared_prompt):
return sorted(collectors)
def _build_exports(source_nodes, collector_ids, request_id):
bindings = []
exports = {}
def _build_explicit_exports(source_nodes, specs, request_id):
bindings, exports = [], {}
reserved = set(source_nodes)
for spec in specs:
if not isinstance(spec, dict) or set(spec) != {
"collector_id", "job_id", "participant_id", "socket", "source",
"transfer_format", "transfer_quality"}:
raise ValueError("Invalid explicit media export specification")
collector_id, socket = spec["collector_id"], spec["socket"]
kind = "VALUE" if socket in VALUE_SOCKETS else SOCKET_KINDS.get(socket)
if kind is None:
raise ValueError("Unsupported explicit media export socket")
link = spec["source"]
if (not _is_link(link) or link[0] not in source_nodes
or type(link[1]) is not int or link[1] < 0):
raise ValueError("Explicit media export requires a valid source link")
fmt, quality = validate_options(spec["transfer_format"], spec["transfer_quality"])
if kind != "IMAGE" and fmt != DEFAULT_FORMAT:
raise ValueError("Image transfer controls apply only to IMAGE exports")
# Same stable identities as Collector compilation, including collisions.
identity = json.dumps([collector_id, socket], separators=(",", ":"))
digest = hashlib.sha256(identity.encode("utf-8")).hexdigest()
base_id = f"serverless:media:{digest}"
node_id, suffix = base_id, 0
while node_id in reserved:
suffix += 1
node_id = f"{base_id}:{suffix}"
reserved.add(node_id)
binding = {"binding_id": f"b:{digest}", "collector_id": collector_id,
"job_id": spec["job_id"], "participant_id": spec["participant_id"],
"node_id": node_id, "socket": socket, "kind": kind}
bindings.append(binding)
exports[node_id] = {"class_type": "MediaExportV1", "inputs": {
"request_id": request_id, "binding_json": json.dumps(binding),
socket: copy.deepcopy(link)}}
if fmt != DEFAULT_FORMAT:
exports[node_id]["inputs"].update(transfer_format=fmt, transfer_quality=quality)
validate_bindings(bindings)
return bindings, exports
def _build_exports(source_nodes, collector_ids, request_id):
specs = []
for collector_id in collector_ids:
inputs = source_nodes[collector_id].get("inputs", {})
for socket in COLLECTOR_SOCKETS:
kind = SOCKET_KINDS[socket]
if socket not in inputs:
continue
# Hashing avoids separator ambiguity and keeps long valid source
# identities within the bridge's identifier length bound. Identity
# is stable across request IDs and dictionary insertion orders.
identity = json.dumps([collector_id, socket], separators=(",", ":"))
digest = hashlib.sha256(identity.encode("utf-8")).hexdigest()
base_id = f"serverless:media:{digest}"
node_id = base_id
suffix = 0
while node_id in reserved:
suffix += 1
node_id = f"{base_id}:{suffix}"
reserved.add(node_id)
binding = {
"binding_id": f"b:{digest}",
"collector_id": collector_id,
"job_id": inputs["multi_job_id"],
"participant_id": inputs["worker_id"],
"node_id": node_id,
"socket": socket,
"kind": kind,
}
bindings.append(binding)
exports[node_id] = {
"class_type": "MediaExportV1",
"inputs": {
"request_id": request_id,
"binding_json": json.dumps(binding),
socket: copy.deepcopy(inputs[socket]),
},
}
if kind == "IMAGE":
fmt, quality = collector_image_policy(inputs)
if fmt != DEFAULT_FORMAT:
exports[node_id]["inputs"].update(transfer_format=fmt, transfer_quality=quality)
validate_bindings(bindings)
return bindings, exports
fmt, quality = collector_image_policy(inputs) if socket == "images" else (DEFAULT_FORMAT, 90)
specs.append({"collector_id": collector_id, "job_id": inputs["multi_job_id"],
"participant_id": inputs["worker_id"], "socket": socket,
"source": inputs[socket], "transfer_format": fmt, "transfer_quality": quality})
return _build_explicit_exports(source_nodes, specs, request_id)
def compile_export_prompt(prepared_prompt, request_id, specs):
"""Compile explicit internal exports, without fabricating a Collector node.
Callers supply linked source sockets and literal binding identities; only
exporters and required ancestors survive. This does not add Collector
sockets: its public compiler remains images/audio only.
"""
_validate_identifier(request_id, "request_id")
if _validate_prompt(prepared_prompt, require_collectors=False):
raise ValueError("Explicit exports cannot contain Collector nodes")
bindings, exports = _build_explicit_exports(prepared_prompt, specs, request_id)
compiled = _prune_graph(prepared_prompt, bindings, exports, _make_scalar_resolver(prepared_prompt))
return {node_id: compiled[node_id] for node_id in sorted(compiled)}, bindings
def compile_media_prompt(prepared_prompt: dict, request_id: str) -> tuple[dict, list[dict]]:
+2
View File
@@ -171,6 +171,8 @@ class ServerlessManager:
for rid in done[:-200]:
self.tasks.pop(rid, None)
self.runs.pop(rid, None)
# Additive ownership handle for item schedulers; legacy callers ignore it.
return [record_id for record_id, _run, _semaphore in new_runs]
def is_pending(self, collector_job_id, participant_id):
return any(not task.done() and run["worker"]["id"] == participant_id and
+26 -24
View File
@@ -38,7 +38,7 @@ from .orchestration.prompt_transform import (
WORKER_SINK_CLASS,
_resolve_direct_workers_for_image_list,
)
from ..image_list.runtime import validate_image_list_regions
from ..image_list.runtime import validate_image_list_regions, _has_selected_serverless_workers
from ..image_list.execution_plan import prepare_image_list_execution
@@ -337,25 +337,25 @@ async def orchestrate_distributed_execution(
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, explicit_targets=True)
configured_serverless = [
str(worker.get("id")) for worker in config.get("workers", [])
if worker.get("enabled", False)
and str(worker.get("type") or "local").lower() == "serverless"
and (requested_ids is None or str(worker.get("id")) in {str(x) for x in requested_ids})
]
if configured_serverless:
trace_debug(
execution_trace_id,
"Image List v1 explicitly excludes unsupported serverless pools: "
+ ", ".join(configured_serverless),
)
from ..serverless_bridge.image_transfer import DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options
required_transfer_formats = set()
for end_id in prompt_index.nodes_for_class("DistributedImageListEnd"):
end_inputs = prompt_obj[end_id].get("inputs", {})
fmt, _ = validate_options(end_inputs.get("transfer_format", DEFAULT_FORMAT),
end_inputs.get("transfer_quality", DEFAULT_QUALITY))
if fmt != DEFAULT_FORMAT:
required_transfer_formats.add(fmt)
workers = _resolve_direct_workers_for_image_list(config, requested_ids)
if _has_selected_serverless_workers(config, requested_ids):
from ..image_list.serverless import resolve_serverless_workers
# Pool identity is retained once, never expanded by Images per run.
# Validate profiles before reachability can hide an invalid selection.
workers += resolve_serverless_workers(config, requested_ids)
if not workers and delegate_master:
raise RuntimeError(
"Distributed Image List in orchestrator-only mode requires a selected direct worker; "
"serverless pools are unsupported"
"Distributed Image List in orchestrator-only mode requires a selected direct worker "
"or comfy-media-v1 serverless pool"
)
else:
workers = _resolve_enabled_workers(config, requested_ids)
@@ -401,18 +401,20 @@ async def orchestrate_distributed_execution(
from .orchestration.dispatch import preflight_image_list_workers
if not active_workers and delegate_master:
raise RuntimeError(
"No selected direct worker is reachable for Distributed Image List in orchestrator-only mode"
"No selected worker is reachable for Distributed Image List 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)
direct_workers = [worker for worker in active_workers if worker.get("type") != "serverless"]
preflight_kwargs = ({"required_transfer_formats": sorted(required_transfer_formats)}
if required_transfer_formats else {})
if requires_values:
preflight_kwargs["require_value_results"] = True
# Preserve busy direct candidates for queued runs; runtime rechecks idle
# availability before assignment. Never direct-probe serverless pools.
if direct_workers:
await preflight_image_list_workers(direct_workers, require_idle=False, **preflight_kwargs)
else:
delegate_master = selected_delegate_master
+94 -14
View File
@@ -1,8 +1,8 @@
# Dynamic image-list queue
Process a list of images across the master and direct local or remote ComfyUI workers without assigning fixed chunks to each participant.
Process a list of images across the master, direct local or remote ComfyUI workers, and compatible Runpod serverless pools without assigning fixed chunks to each participant.
The master evaluates the input list once. By default, the master and each selected idle worker receive one image job at a time; when a participant finishes, it gets the next unassigned image. Faster participants can therefore process more images. The master restores input order before running downstream nodes.
The master evaluates the input list once. By default, the master and each selected direct worker receive one image job at a time; a serverless pool contributes bounded concurrent slots. All slots share the same queue: when a participant finishes, it gets the next unassigned image. Faster participants can therefore process more images. Each source frame is assigned once, and the master restores input order before running downstream nodes.
The existing **Orchestrator Only** setting (`settings.master_delegate_only`) excludes the master from per-image processing. The master still creates the input list and runs downstream nodes. API callers can override the setting for a run with `delegate_master`. There is no second image-list-specific toggle.
@@ -23,8 +23,8 @@ Distributed Image List Start -------- context ---------+
Preview / Save / next stage
```
1. Install the same feature-capable version of Distributed on the master and every selected worker. Workers on an older image-list protocol fail preflight before any work is sent. Restart those instances and refresh the frontend after installation.
2. Enable Distributed and select the **direct** workers to use. The master also participates unless **Orchestrator Only** is enabled. With the master participating, an empty worker selection runs locally through the same image-list scheduler; orchestrator-only mode requires at least one reachable, idle direct worker. The master must be reachable from remote workers at its configured callback address.
1. Install the same feature-capable version of Distributed on the master and every selected direct worker. Direct workers on an older image-list protocol fail preflight before any work is sent. Restart those instances and refresh the frontend after installation. Serverless endpoints instead need the compatible media bridge and supervision described below.
2. Enable Distributed and select direct workers and/or **Typed media v1** serverless pools. The master also participates unless **Orchestrator Only** is enabled. With the master participating, an empty worker selection runs locally through the same image-list scheduler; orchestrator-only mode requires at least one eligible direct worker or compatible serverless pool. Serverless-only selection is supported. The master must be reachable from direct remote workers at its configured callback address; serverless jobs do not use this callback.
3. Feed your image list or tensor batch into **Distributed Image List Start → images**.
4. Connect its `image` output through the per-image processing nodes to **Distributed Image List End → images**.
5. Connect Start's **context** directly to End's **context**. This is a dedicated `DISTRIBUTED_IMAGE_LIST_CONTEXT` socket, not editable text.
@@ -84,7 +84,80 @@ Per-image data must pass through Start; do not bypass it with upstream images, l
End's optional `texts` is a connection-only socket. For workflows saved before this socket became connection-only, recreate End and reconnect its wires if an old text widget or literal `texts: ""` remains in the compiled prompt. Arbitrary text literals are not accepted at End.
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.
File-based dependencies keep ComfyUI's `[output]` / `[temp]` root annotations and are read on the master. Direct-worker copies use separate content-addressed paths under `input/distributed-media/`; each direct worker's rewritten prompt is reused across items without reusing item tokens. Cloud requests use the PR2 bounded file-input adapters instead.
## Image return format and quality
End adds `transfer_format` (default **exact**) and `transfer_quality` (default **90**).
These apply to returned images from **both direct workers and serverless pools**.
Master-local results bypass encoding. Input images, masks, generic values and
returned masks/text retain their exact existing transport.
- `exact` preserves floating-point dtype and values.
- `legacy_png`, `png`, `ffv1` preserve quantized 8-bit image pixels, not the
original floating-point colors. `legacy_png` is PNG without compression.
- `jpeg`, `webp`, `h264`, `h265`, `av1` are lossy choices; quality is 1–100.
Video codecs encode each assigned image independently, not the whole source list.
- Dimensions, floating dtype and alpha survive; alpha is an exact sidecar.
- Direct workers must advertise the selected codec in `image_transfer_formats`;
older workers fail before submission for non-exact choices. Exact remains
compatible with the existing protocol. Serverless endpoints need the updated
complete bridge; health probing does not prove codec availability.
- The selected software encoder must be installed. Encoding failures stop the
item; there is no silent format fallback. No hardware encoder is selected.
Compression can reduce transfer size, but does **not** bypass decoded tensor,
conservative codec scratch, or complete request limits. For details and codec
limits, see [image return options](serverless-media.md#image-return-format-and-quality).
Mixed local/remote results may have different color fidelity when compression is
selected, because local images are deliberately not quantized.
## Serverless pools and bounded capacity
Select `media_profile: "comfy-media-v1"` for each cloud pool. The endpoint must
include the [standalone PR2 media bridge](../serverless_bridge/README.md), route
typed requests through its executor, and provide the required supervisor. Merely
installing the bridge nodes or choosing the profile is not sufficient. The
endpoint needs the workflow's models and custom nodes, but not the full
Distributed extension. A selected `legacy-image-v1` pool fails before the master
prompt is queued; it is not silently skipped in favor of local execution.
Provider health/authentication probing does not verify bridge/model installation.
- Each pool is selected once by its configured pool ID. **Images per run
(`jobs_per_run`) does not multiply input images or set image-list parallelism.**
An assigned cloud job processes exactly one source item and its paired values.
- Cloud slots are bounded by the existing **16 outstanding jobs per pool** cap,
available manager capacity (**64 pending jobs** across pools), reservations,
and the source item count. Unresolved submissions and unconfirmed cancellations
still consume capacity. Selecting more items does not enqueue all of them as
provider jobs at once. Endpoint-side worker limits can reduce actual concurrency.
- Cloud slots, direct workers and the optional master pull from the same queue;
no region is cloned once per Images-per-run setting. Results are restored to
source order even when participants finish out of order.
- Each cloud request uses the standard media-v1 envelope and existing durable
serverless job lifecycle. The private item graph replaces Start/End with typed
inputs and explicit exporter roots, not independent all-output execution.
IMAGE returns are exact by default, or explicitly encoded as selected at End.
MASK tensor data, input images and text retain exact typed codecs. There is no
image-list-specific endpoint protocol or new visible node socket.
- Connected End `value_1` … `value_10` overrides are computed on the endpoint and returned using exact typed `VALUE` exports. Sparse slot identities are preserved; only explicitly connected overrides are transported. The master overlays them on each source item’s original values. BOOLEAN, INT, FLOAT, STRING, IMAGE and MASK remain distinct; missing/extra slots, multiple values and nonfinite values fail. Override tensors remain exact even when the main image uses a compressed return format. End overrides require the updated complete bridge (including `value_codec.py`), not merely an older healthy media endpoint.
- Cloud `value_N` tensors must have native IMAGE `[B,H,W,C]` or MASK `[B,H,W]`
shapes, not arbitrary rank-one/rank-two tensor payloads. Typed tensors must be
finite; connected masks must match their image grid. Unsupported values fail
rather than being resized, flattened or converted to lossy images.
- Inputs are inline and the **complete serialized Runpod request is capped at
10,000,000 bytes**, including workflow and base64. Default inline result JSON
has the same cap. The PR2 contract's larger artifact ceilings do not remove
these inline limits; larger outputs need an operator-provided artifact publisher.
There is no built-in large-input uploader. See [typed media limits](serverless-media.md#size-and-operational-limits).
- Missing/duplicate/mismatched results or provider failures fail the region.
Cancellation closes only this region's owned cloud jobs and collector queues,
including jobs whose dispatch was cancelled before returning a handle. It does
not clear unrelated pool work. No failed or ambiguous job is retried automatically.
- A provider cancellation or timeout is **not proof that endpoint ComfyUI stopped**.
The PR2 supervisor must quarantine/reconcile or terminate an uncertain endpoint
before reuse. Preserve its retained files until execution and readers are known
to have stopped; see the [bridge uncertainty policy](../serverless_bridge/README.md).
## Native cache reuse
@@ -99,21 +172,21 @@ Repeated items can reuse **ComfyUI's native processing cache** on the same parti
## 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.
- 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.
- This version supports master, direct workers and `comfy-media-v1` serverless pools in one image-list region. Direct capability/idle HTTP probes are never sent to serverless pools; pool probing uses the existing serverless health/authentication path.
- Queue submission checks reachability and feature compatibility without freezing direct-worker selection to that moment's idle snapshot. A selected direct worker busy with an earlier prompt remains a candidate for the queued run. At region start, saved direct candidates are checked again and only idle, compatible workers receive items; transiently unreachable workers are excluded. Malformed capability/queue responses and incompatible versions fail explicitly. Workers excluded by initial reachability or selected later are not added to the queued run, and workers becoming free after region start are not added mid-region.
- **Orchestrator Only** can queue while selected direct workers are busy. At region start it requires an idle compatible direct worker or a compatible selected serverless pool; 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 and compatible serverless pools. Explicit IDs select the named configured workers/pools, including disabled ones; an empty list stays empty and unknown IDs fail. Pool IDs are retained in Start's hidden selection metadata, not expanded participant IDs. Local direct workers use the master's loopback callback; remote direct 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 workers are selected or reachable. **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.
- 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.
- Every selected direct 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. Cloud jobs use explicit media exporter roots. 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.
- The entire input list is materialized before scheduling; this is not streaming input. The limit is 4,096 frames. Direct-worker input transport and default exact returns preserve floating-point tensor dtype and values. Opt-in image return encoding does not relax decoded-size limits. Each direct-worker 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. Serverless inline limits are smaller, as described above. 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.
- Worker errors fail the region. For direct workers, 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. Serverless jobs instead use their configured serverless timeout and durable manager lifecycle; they are excluded from direct-history/liveness polling.
- User interruption is checked during submission and result waiting. Cleanup targets only this region's acknowledged prompt IDs and staged input tokens. Submission may wait briefly for an acknowledgement before cancellation; remote transport and cleanup have finite time limits. Local cleanup stops at node/progress boundaries and waits for its execution thread to finish before the parent prompt can resume. A synchronous node that does not cooperate with cancellation can delay this beyond the configured deadline. Lost acknowledgements are ambiguous, never retried automatically. Unreachable cleanup is logged as unconfirmed and may require checking the worker manually; cancellation does not guarantee remote termination through a broken network.
- State is in memory. Restarting the master does not resume an interrupted region. Abandoned master contexts and expired worker staging are actively swept on the server loop; active/queued master work is protected from age-only expiry. End also removes its context on success, failure or interruption. This releases feature-owned references, not ComfyUI's separate execution cache.
- Image-list ordering/context state is in memory. Restarting the master does not resume an interrupted region. Cloud job IDs and uncertainty state remain in the existing durable serverless ledger for reconciliation, not automatic list resumption. Abandoned master contexts and expired direct-worker staging are actively swept on the server loop; active/queued master work is protected from age-only expiry. End also removes its context on success, failure or interruption. This releases feature-owned references, not ComfyUI's separate execution cache.
### Time and size limits
@@ -129,6 +202,13 @@ Use trusted workers and trusted workflows, as with the existing ComfyUI API. Cap
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.
Offline cloud integration tests additionally cover mixed-role and cloud-only
scheduling, out-of-order image/mask/text pairing, one job per source item regardless
of Images per run, profile/selection rejection before queueing, exact cleanup on
setup/dispatch/result failures, and the real CPU ComfyUI `processing_interrupted`
signal while awaiting a mocked cloud result. Provider calls are mocked: this is
not live Runpod, endpoint-supervisor or GPU-workflow acceptance.
```bash
COMFYUI_SOURCE_ROOT=/path/to/ComfyUI /path/to/ComfyUI/.venv/bin/python -m pytest tests -q
bash scripts/test-web.sh
+19 -3
View File
@@ -14,6 +14,20 @@ DEFAULT_ORPHAN_GRACE_SECONDS = 300.0
_LIFECYCLE_MARKER = object()
def _without_runtime_fingerprints(prompt: dict) -> dict:
"""Compare submitted structure without ComfyUI's per-node cache annotation.
PromptQueue snapshots the graph before IsChangedCache adds ``is_changed``
to executing nodes. Ignore only that top-level field, never inputs or nested
metadata, and leave both the queued and executing graphs untouched.
"""
return {
node_id: {key: value for key, value in node.items() if key != "is_changed"}
if isinstance(node, dict) else node
for node_id, node in prompt.items()
}
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.
@@ -23,8 +37,9 @@ def has_safe_execution_targets(prompt_server: Any, prompt_id: str | None, prompt
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:
execution_prompt = _without_runtime_fingerprints(prompt)
prepared, targets = prepare_image_list_execution(execution_prompt)
if prepared != execution_prompt or not targets:
return False
queue = getattr(prompt_server, "prompt_queue", None)
for name in ("get_current_queue_volatile", "get_current_queue"):
@@ -35,7 +50,8 @@ def has_safe_execution_targets(prompt_server: Any, prompt_id: str | None, prompt
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):
and entry[1] == prompt_id and isinstance(entry[2], dict)
and _without_runtime_fingerprints(entry[2]) == execution_prompt):
roots = entry[4]
return (isinstance(roots, (list, tuple))
and all(isinstance(root, str) for root in roots)
+163 -36
View File
@@ -17,6 +17,15 @@ import time
from dataclasses import dataclass, field
from typing import Any
try:
from ..serverless_bridge.image_transfer import (
FORMATS as IMAGE_TRANSFER_FORMATS, DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options,
)
except ImportError: # Standalone image_list.runtime (without extension bootstrap).
from serverless_bridge.image_transfer import (
FORMATS as IMAGE_TRANSFER_FORMATS, DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options,
)
# v2: Start carries generic value_N slots instead of fixed reference/text/strength/seed fields.
PROTOCOL_VERSION = 2
START_CLASS = "DistributedImageListStart"
@@ -362,6 +371,8 @@ def decode_value(data: Any, expected_kind: str) -> Any:
return None
if not isinstance(data, dict) or data.get("type") != expected_kind:
raise ValueError(f"Expected {expected_kind} transport value")
if "encoding" in data or "descriptor" in data:
raise ValueError("Compressed transport is supported only for the returned IMAGE")
if expected_kind in ("IMAGE", "MASK", "TENSOR"):
torch = _torch()
dtype_name = data.get("dtype")
@@ -526,7 +537,7 @@ def create_master_job(start_id: str, end_id: str, prompt: dict[str, Any], items:
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")
raise RuntimeError("Distributed Image List requires at least one selected 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)
@@ -889,10 +900,48 @@ def register_assignment(context_id: str, item_index: int, worker_id: str, token:
state.expected_workers[item_index] = worker_id
def _image_codec():
# Exact returns and inputs must not import optional media codec dependencies.
if __package__ == "image_list":
from serverless_bridge import image_codec
else:
from ..serverless_bridge import image_codec
return image_codec
def _decode_return_image(data: Any) -> Any:
if (not isinstance(data, dict) or set(data) != {"type", "encoding", "descriptor", "data"}
or data["type"] != "IMAGE" or data["encoding"] != "image-transfer-v1"):
raise ValueError("Invalid compressed IMAGE envelope")
descriptor = data["descriptor"]
if not isinstance(descriptor, dict) or descriptor.get("encoding") != "image-transfer-v1":
raise ValueError("Invalid compressed IMAGE descriptor")
shape = descriptor.get("shape")
if (not isinstance(shape, list) or len(shape) != 4
or any(type(dim) is not int or dim < 1 for dim in shape)
or shape[0] != 1 or shape[-1] not in (3, 4)):
raise ValueError("Compressed IMAGE must have shape [1,H,W,3|4]")
encoded = data["data"]
if not isinstance(encoded, str) or len(encoded) > ((MAX_PAYLOAD_BYTES + 2) // 3) * 4:
raise ValueError("Compressed IMAGE data is missing or exceeds the transport limit")
try:
raw = base64.b64decode(encoded, validate=True)
except (ValueError, TypeError) as exc:
raise ValueError("Compressed IMAGE data must be valid base64") from exc
if len(raw) > MAX_PAYLOAD_BYTES:
raise ValueError("Compressed IMAGE exceeds the transport limit")
return _image_codec().decode_image(descriptor, raw, max_bytes=MAX_PAYLOAD_BYTES)
def decode_result(result: Any) -> dict[str, Any]:
if not isinstance(result, dict):
raise ValueError("result must be an object")
image = decode_value(result.get("image"), "IMAGE")
image_data = result.get("image")
if isinstance(image_data, dict) and ("encoding" in image_data or "descriptor" in image_data):
encode_request_payload(result) # Bound combined JSON/base64 before codec allocation.
image = _decode_return_image(image_data)
else:
image = decode_value(image_data, "IMAGE")
if image is None or image.ndim != 4 or image.shape[0] != 1 or image.shape[-1] not in (3, 4):
raise ValueError("result.image must contain exactly one RGB/RGBA image")
mask = decode_value(result.get("mask"), "MASK")
@@ -947,16 +996,32 @@ def normalize_worker_result(images: Any, masks: Any = None, texts: Any = None,
return result
def encode_result(result: dict[str, Any]) -> dict[str, Any]:
def encode_result(result: dict[str, Any], *, transfer_format=DEFAULT_FORMAT,
transfer_quality=DEFAULT_QUALITY) -> dict[str, Any]:
fmt, quality = validate_options(transfer_format, transfer_quality)
normalized = normalize_worker_result(result.get("image"), result.get("mask"), result.get("text"))
encoded = {"image": encode_value(normalized["image"], "IMAGE"),
"mask": encode_value(normalized["mask"], "MASK"),
if fmt == "exact":
image = encode_value(normalized["image"], "IMAGE")
else:
value = normalized["image"]
# Reject before moving/copying a potentially oversized device tensor.
if value.numel() * value.element_size() > MAX_PAYLOAD_BYTES:
raise ValueError("IMAGE item exceeds the image-list transport limit")
descriptor, raw = _image_codec().encode_image(
value.detach().to("cpu").contiguous(), transfer_format=fmt,
transfer_quality=quality, max_bytes=MAX_PAYLOAD_BYTES,
)
image = {"type": "IMAGE", "encoding": "image-transfer-v1", "descriptor": descriptor,
"data": base64.b64encode(raw).decode("ascii")}
encoded = {"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()}
if fmt != "exact":
encode_request_payload(encoded)
return encoded
@@ -985,7 +1050,7 @@ def _collect_output_lists(results: list[dict[str, Any]], *, state: ImageListJob
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."""
"""Schedule one queue across master, direct peers and bounded cloud slots."""
state = get_master_job(context_id)
if state is None:
raise RuntimeError("Distributed Image List context is missing or expired")
@@ -1011,37 +1076,80 @@ async def run_image_list_job(context_id: str, *, model_callback=None,
remove_master_job(context_id)
def _has_selected_serverless_workers(config, requested_ids):
"""Keep direct/master-only execution free of optional cloud adapter imports.
This is only an import guard, not validation. The adapter validates original
pool IDs/configuration without repairing them, before any cloud work starts.
"""
requested = None if requested_ids is None else {str(value) for value in requested_ids}
return any(
str(worker.get("type") or "local").strip().lower() == "serverless"
and (worker.get("enabled", False) if requested is None
else str(worker.get("id") or "").strip() in requested)
for worker in config.get("workers", [])
)
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()
from ..api.orchestration.dispatch import (
preflight_image_list_workers, dispatch_worker_prompt,
)
from ..api.orchestration.prompt_transform import _resolve_direct_workers_for_image_list
from ..utils.config import load_config, get_worker_timeout_seconds
from ..utils.network import build_master_callback_url, build_worker_url, get_client_session
from ..utils.logging import log
import aiohttp
from ..utils.config import load_config
config = load_config()
workers = _resolve_direct_workers_for_image_list(config, state.worker_ids)
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.
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
# Both item paths below select End explicitly, never all OUTPUT_NODE roots.
# Every item path selects 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")
controller = None
try:
if _has_selected_serverless_workers(config, state.worker_ids):
from .serverless import create_serverless_slots
controller = await create_serverless_slots(state, region, config, state.event)
return await _run_image_list_slots(state, config, workers, region, controller,
include_values=include_values)
finally:
# Own cloud jobs even if direct preflight/local setup fails, or dispatch
# is cancelled before returning an acknowledged handle to the scheduler.
if controller is not None:
await controller.close()
async def _run_image_list_slots(state, config, workers, region, controller, *, include_values=False):
context_id = state.context_id
from ..api.orchestration.dispatch import (
preflight_image_list_workers, dispatch_worker_prompt,
)
from ..utils.config import get_worker_timeout_seconds
from ..utils.network import build_master_callback_url, build_worker_url, get_client_session
from ..utils.logging import log
import aiohttp
serverless_slots = list(controller.slots) if controller is not None else []
end_inputs = state.original_prompt[state.end_id]["inputs"]
transfer_format, _ = validate_options(
end_inputs.get("transfer_format", DEFAULT_FORMAT),
end_inputs.get("transfer_quality", DEFAULT_QUALITY),
)
transfer_requirements = ({"required_transfer_formats": (transfer_format,)}
if transfer_format != "exact" else {})
override_names = end_override_names(end_inputs)
value_preflight = {"require_value_results": True} if override_names else {}
# Direct capability/idle probes must never be sent to a serverless pool.
# Busy peers are filtered in either role; they are fatal only when neither
# the master nor an explicitly selected cloud slot can process the queue.
worker_statuses = await preflight_image_list_workers(
workers, require_idle=not state.master_participates and not serverless_slots,
**transfer_requirements, **value_preflight,
) if workers else {}
workers = [worker for worker in workers if worker["id"] in worker_statuses]
if not workers and not serverless_slots and not state.master_participates:
raise RuntimeError("No reachable, feature-compatible image-list workers are available; master execution is disabled")
session = await get_client_session() if workers else None
timeout_seconds = max(5.0, float(get_worker_timeout_seconds()))
use_websocket = bool((config.get("settings", {}) or {}).get("websocket_orchestration", False))
prepared_worker_prompts: dict[str, dict[str, Any]] = {}
@@ -1056,6 +1164,8 @@ async def _run_image_list_job(state: ImageListJob, *, include_values: bool = Fal
local_executor = LocalExecutionSlot(model_callback=state.local_model_callback)
async def dispatch(worker_id: Any, index: int) -> dict[str, Any]:
if worker_id in serverless_slots:
return await controller.dispatch(worker_id, index)
if worker_id is master_slot:
assert local_executor is not None
token = create_worker_task_token()
@@ -1076,13 +1186,15 @@ async def _run_image_list_job(state: ImageListJob, *, include_values: bool = Fal
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")
if transfer_requirements or value_preflight:
# A peer may restart/downgrade between items. Recheck both return
# contracts before staging or POSTing each assignment.
statuses = await preflight_image_list_workers(
[worker], require_idle=bool(transfer_requirements),
**transfer_requirements, **value_preflight,
)
if worker_id not in statuses:
raise RuntimeError(f"Image-list worker {worker_id} unavailable before return dispatch")
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])}
@@ -1138,6 +1250,10 @@ async def _run_image_list_job(state: ImageListJob, *, include_values: bool = Fal
state.event.clear()
if state.failed:
raise RuntimeError(state.failed)
if controller is not None:
completed = controller.poll(inflight)
if completed is not None:
return completed
local = inflight.get(master_slot)
if local is not None and local["future"].done():
# A posted result does not release this slot until execution has
@@ -1187,10 +1303,15 @@ async def _run_image_list_job(state: ImageListJob, *, include_values: bool = Fal
return None
return None
errors = await asyncio.gather(*(history_status(wid, handle) for wid, handle in inflight.items()
if wid is not master_slot))
if wid is not master_slot and wid not in serverless_slots))
failure = next((error for error in errors if isinstance(error, str)), None)
if failure:
raise RuntimeError(failure)
# A cloud manager task may finish during direct HTTP history waits.
if controller is not None:
completed = controller.poll(inflight)
if completed is not None:
return completed
completed = next((entry for entry in errors if isinstance(entry, tuple)), None)
if completed is not None:
# The callback precedes the worker prompt finishing. Wait for its
@@ -1201,7 +1322,8 @@ async def _run_image_list_job(state: ImageListJob, *, include_values: bool = Fal
return completed
now = time.monotonic()
expired = [f"{wid}/item {handle['index']}" for wid, handle in inflight.items()
if wid is not master_slot and now - handle["last_seen"] > timeout_seconds]
if wid is not master_slot and wid not in serverless_slots
and now - handle["last_seen"] > timeout_seconds]
if expired:
raise TimeoutError(f"Image-list owned-prompt liveness timeout after {timeout_seconds:.1f}s: {', '.join(expired)}")
try:
@@ -1210,6 +1332,10 @@ async def _run_image_list_job(state: ImageListJob, *, include_values: bool = Fal
pass
async def cleanup(worker_id: Any, handle: dict[str, Any]) -> None:
if worker_id in serverless_slots:
# The controller owns exact manager records, including dispatches
# that never returned a handle. Close it once in the outer finally.
return
if worker_id is master_slot:
# The outer finally owns the non-abandonable stop/join. The generic
# scheduler's bounded remote cleanup must never detach local work.
@@ -1250,7 +1376,8 @@ async def _run_image_list_job(state: ImageListJob, *, include_values: bool = Fal
return False
try:
slots = ([master_slot] if local_executor is not None else []) + [worker["id"] for worker in workers]
slots = (([master_slot] if local_executor is not None else [])
+ [worker["id"] for worker in workers] + serverless_slots)
results = await dynamic_schedule(len(state.items), slots,
dispatch, wait_next, cleanup, interrupted)
return _collect_output_lists(results, state=state, include_values=include_values)
+435
View File
@@ -0,0 +1,435 @@
"""Image-list items on explicitly selected comfy-media-v1 pools.
Only graph/media adaptation lives here. The shared manager owns admission,
provider submission, durable reservations, cancellation and atomic decoding.
A provider terminal state is not proof that an arbitrary endpoint stopped its
local Comfy prompt: the selected profile must honour its supervision contract.
"""
from __future__ import annotations
import asyncio
import copy
import hashlib
import uuid
from . import runtime
from ..api.orchestration.runpod_client import serialize_request
from ..api.orchestration.serverless_media_prompt import compile_export_prompt
from ..api.orchestration.serverless_runtime import PER_WORKER_JOB_LIMIT, get_manager
from ..api.orchestration.serverless_typed_inputs import prepare_media_inputs, _artifact_id
from ..serverless_bridge import codec, contract
from ..utils.serverless_config import validate_serverless_worker
def _identity(value):
contract.identifier(value)
if any(part in ("", ".", "..") for part in value.split("/")):
raise ValueError("Unsafe image-list media identifier")
return value
def resolve_serverless_workers(config, requested_ids):
"""Return pool configurations, never jobs_per_run-expanded participants.
None means enabled defaults; an explicit selection includes disabled pools.
Direct worker IDs may coexist in that selection and are left to their own
resolver. Unknown, duplicate configured and malformed identities fail closed.
"""
workers = (config or {}).get("workers", [])
indexed = {}
for worker in workers:
wid = worker.get("id")
if not isinstance(wid, str) or not wid or wid != wid.strip():
raise ValueError("Invalid configured image-list worker ID")
if wid in indexed:
raise ValueError("Duplicate configured image-list worker ID")
indexed[wid] = worker
requested = None
if requested_ids is not None:
if not isinstance(requested_ids, (list, tuple, set)):
raise ValueError("Image-list worker selection must be a list of exact IDs")
if any(not isinstance(wid, str) or not wid or wid != wid.strip() for wid in requested_ids):
raise ValueError("Invalid selected image-list worker ID")
requested = set(requested_ids)
if requested - indexed.keys():
raise ValueError("Unknown explicitly selected image-list worker ID(s)")
result = []
for wid, worker in indexed.items():
if requested is None:
if worker.get("enabled", False) is not True:
continue
elif wid not in requested:
continue
if str(worker.get("type") or "").strip().lower() != "serverless":
continue
clean = validate_serverless_worker(worker)
_identity(wid)
if clean.get("media_profile") != contract.PROFILE:
raise ValueError(f"Image-list serverless pool {wid!r} requires comfy-media-v1")
# Pool IDs, not virtual legacy participant IDs, own capacity and secrets.
if "pool_id" in clean and clean["pool_id"] != wid:
raise ValueError("Image-list selection requires a pool, not an expanded participant")
clean["pool_id"] = wid
result.append(clean)
return result
def _link(value, prompt):
if not isinstance(value, (list, tuple)):
return None
if (len(value) != 2 or not isinstance(value[0], str)
or type(value[1]) is not int or value[1] < 0):
raise ValueError("Invalid image-list graph link")
_identity(value[0])
if value[0] not in prompt:
raise ValueError("Image-list graph link source is missing")
return value[0], value[1]
def _new_node_id(graph, label):
base = "image-list:input:" + label
candidate, suffix = base, 0
while candidate in graph:
suffix += 1
candidate = f"{base}:{suffix}"
return candidate
def _merge_artifacts(*groups):
indexed = {}
for group in groups:
for artifact in group:
contract.validate_artifact(artifact)
aid = artifact["id"]
if aid in indexed and indexed[aid] != artifact:
raise ValueError("Conflicting image-list input artifact identity")
indexed[aid] = artifact
if (len(indexed) > contract.MAX_ARTIFACTS or
sum(a["byte_length"] for a in indexed.values()) > contract.MAX_TOTAL_BYTES):
raise ValueError("Image-list media inputs exceed artifact limits")
return [indexed[aid] for aid in sorted(indexed)]
async def prepare_serverless_item(prompt, region, item, worker, *, request_id, job_id):
"""Compile one immutable Start/End region to an endpoint-native media job.
Tensor Start outputs use exact codec artifacts and references to existing
consumer inputs. PrimitiveString placeholders keep exporter sockets linked
until the endpoint handler injects MediaImportV1 and prunes unused ancestors.
Scalar values remain typed literals, except exported text stays linked.
"""
_identity(request_id)
_identity(job_id)
worker = validate_serverless_worker(worker)
_identity(worker["id"])
if worker.get("media_profile") != contract.PROFILE:
raise ValueError("Image-list serverless execution requires comfy-media-v1")
if not isinstance(prompt, dict) or not prompt:
raise ValueError("Image-list prompt must be an API graph")
# Check original tokens before looking them up or pruning any nodes. Do not
# let int/bool IDs or slots become valid through string/int normalization.
for node_id, node in prompt.items():
_identity(node_id)
if not isinstance(node, dict) or not isinstance(node.get("inputs", {}), dict):
raise ValueError("Invalid image-list prompt node")
for name, value in node.get("inputs", {}).items():
_identity(name)
_link(value, prompt)
start, end = region["start_id"], region["end_id"]
_identity(start)
_identity(end)
if (prompt[start].get("class_type") != runtime.START_CLASS or
prompt[end].get("class_type") != runtime.END_CLASS):
raise ValueError("Image-list region has mismatched Start/End identities")
if _link(prompt[end].get("inputs", {}).get("context"), prompt) != (start, runtime.CONTEXT_OUT):
raise ValueError("Image-list End context must reference its exact Start")
ids = region["node_ids"]
for node_id in ids:
_identity(node_id)
if not {start, end} <= set(ids):
raise ValueError("Image-list region is missing its boundary nodes")
graph = {nid: copy.deepcopy(prompt[nid]) for nid in sorted(ids)}
end_inputs = graph[end].get("inputs", {})
from ..serverless_bridge.image_transfer import DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options
transfer_format, transfer_quality = validate_options(
end_inputs.get("transfer_format", DEFAULT_FORMAT),
end_inputs.get("transfer_quality", DEFAULT_QUALITY))
if "images" not in end_inputs:
raise ValueError("Image-list End requires images")
override_names = runtime.end_override_names(end_inputs)
graph[end]["inputs"] = {key: end_inputs[key] for key in
("images", "masks", "texts") + override_names if key in end_inputs}
del graph[start]
if type(item.get("index")) is not int or item["index"] < 0:
raise ValueError("Image-list item index must be a nonnegative integer")
image = item.get("image")
if (not isinstance(image, runtime._torch().Tensor) or image.ndim != 4
or image.shape[0] != 1 or image.shape[-1] not in (3, 4)):
raise ValueError("Image-list item must contain exactly one RGB/RGBA image frame")
values = item.get("values", {})
if not isinstance(values, dict) or any(name not in runtime.VALUE_NAMES for name in values):
raise ValueError("Invalid image-list value_N slots")
materialized, artifacts = {}, {}
placeholders = {}
def replacement(slot):
if slot in materialized:
return copy.deepcopy(materialized[slot])
if slot == runtime.IMAGE_OUT:
value, kind = image, "IMAGE"
elif slot == runtime.MASK_OUT:
value, kind = item.get("mask"), "MASK"
elif slot == runtime.INDEX_OUT:
value, kind = item["index"], "INT"
elif runtime.VALUE_OUT_BASE <= slot < runtime.VALUE_OUT_BASE + runtime.MAX_VALUE_SLOTS:
name = runtime.VALUE_NAMES[slot - runtime.VALUE_OUT_BASE]
value = values.get(name)
kind = runtime.value_kind(value, name)
if kind == "TENSOR":
if value.ndim == 4:
kind = "IMAGE"
elif value.ndim == 3:
kind = "MASK"
else:
raise ValueError(f"{name} serverless tensor requires IMAGE [B,H,W,C] or MASK [B,H,W]")
else:
raise ValueError("Unsupported Start context/control output in endpoint computation")
if value is None:
raise ValueError("Connected Start output is empty; None cannot be submitted")
if kind in ("IMAGE", "MASK"):
if kind == "MASK" and slot == runtime.MASK_OUT:
if value.ndim != 3 or value.shape[0] != 1 or value.shape[1:3] != image.shape[1:3]:
raise ValueError("Image-list item mask must be paired with its image grid")
if value.numel() * value.element_size() > contract.MAX_ARTIFACT_BYTES:
raise ValueError("Image-list tensor exceeds media artifact byte limit")
descriptor, payload = codec.encode_payload(kind, value.detach().to("cpu"),
max_bytes=contract.MAX_ARTIFACT_BYTES)
aid = _artifact_id(kind, descriptor, payload)
if aid not in artifacts and (len(artifacts) >= contract.MAX_ARTIFACTS or
len(payload) + sum(a["byte_length"] for a in artifacts.values()) > contract.MAX_TOTAL_BYTES):
raise ValueError("Image-list media inputs exceed artifact limits")
artifacts[aid] = contract.artifact_from_payload(aid, kind, descriptor, payload)
placeholder = _new_node_id(graph, str(slot))
graph[placeholder] = {"class_type": "PrimitiveString", "inputs": {"value": ""}}
placeholders[placeholder] = aid
output = [placeholder, 0]
else:
# value_kind rejects objects/None/nonfinite values and bool remains bool.
runtime.value_kind(value, "Start output")
output = value
materialized[slot] = output
return copy.deepcopy(output)
for node_id in list(graph):
inputs = graph[node_id].setdefault("inputs", {})
for name, value in list(inputs.items()):
link = _link(value, prompt)
if link and link[0] == start:
inputs[name] = replacement(link[1])
elif link:
inputs[name] = list(link)
# Start text/index/generic scalar substitutions may have removed End's link.
text = graph[end]["inputs"].get("texts")
if "texts" in graph[end]["inputs"] and not isinstance(text, list):
if not isinstance(text, str):
raise ValueError("Image-list End texts must be STRING")
nid = _new_node_id(graph, "text")
graph[nid] = {"class_type": "PrimitiveString", "inputs": {"value": text}}
graph[end]["inputs"]["texts"] = [nid, 0]
# Literal/scalar overrides also use the exact bridge codec, never numeric
# primitive widgets (Comfy may coerce their types). The generic import slot
# keeps the export linked until native injection, just like tensor inputs.
for name in override_names:
value = graph[end]["inputs"][name]
if not isinstance(value, list):
descriptor, payload = codec.encode_payload("VALUE", value,
max_bytes=contract.MAX_ARTIFACT_BYTES)
aid = _artifact_id("VALUE", descriptor, payload)
artifacts[aid] = contract.artifact_from_payload(aid, "VALUE", descriptor, payload)
nid = _new_node_id(graph, name)
graph[nid] = {"class_type": "PrimitiveString", "inputs": {"value": ""}}
placeholders[nid] = aid
graph[end]["inputs"][name] = [nid, 0]
outputs = graph.pop(end)["inputs"]
specs = [{"collector_id": end, "job_id": job_id, "participant_id": worker["id"],
"socket": name, "source": value,
"transfer_format": transfer_format if name == "images" else DEFAULT_FORMAT,
"transfer_quality": transfer_quality if name == "images" else DEFAULT_QUALITY}
for name, value in outputs.items()]
workflow, manifest = compile_export_prompt(graph, request_id, specs)
refs = []
for nid, node in workflow.items():
for name, value in node["inputs"].items():
if isinstance(value, list) and value[0] in placeholders:
refs.append({"node_id": nid, "input_name": name, "artifact_id": placeholders[value[0]]})
# Only actual surviving consumer slots count as uses of an input artifact.
used = {ref["artifact_id"] for ref in refs}
images, fixed_artifacts, fixed_refs = await prepare_media_inputs(workflow)
merged = _merge_artifacts([artifacts[aid] for aid in used], fixed_artifacts)
refs.extend(fixed_refs)
media = {"version": contract.VERSION, "profile": contract.PROFILE,
"request_id": request_id, "artifacts": merged,
"inputs": sorted(refs, key=lambda ref: (ref["node_id"], ref["input_name"])),
"outputs": manifest}
# Includes handler identity/link/ref validation and aggregate + exact JSON caps.
serialize_request(workflow, images, worker["timeout_seconds"], media=media)
return {"worker": worker, "workflow": workflow, "images": images,
"manifest": manifest, "media": media}
class ServerlessSlots:
"""Nonblocking slot adapter; completion means the manager task has drained."""
def __init__(self, state, region, workers, manager, event):
self.state, self.region, self.manager, self.event = state, region, manager, event
self.workers = {worker["id"]: worker for worker in workers}
# The manager remains the atomic authority at start(). Bound this local
# scheduler too, including other callers' currently pending tasks.
available = max(0, 64 - sum(not task.done() for task in manager.tasks.values())) if manager else 0
self.slots = [("serverless", worker["id"], index)
for index in range(PER_WORKER_JOB_LIMIT) for worker in workers][:available]
if workers and not self.slots:
raise ValueError("Distributed already has too many pending serverless jobs (limit 64)")
digest = hashlib.sha256(state.context_id.encode()).hexdigest()
self._scope = f"image-list:{digest}:{uuid.uuid4().hex}"
self._active = {}
self._queues = {}
self._executions = set()
self._assigned = set()
self._submissions = set()
self._closed = False
self._close_task = None
async def dispatch(self, slot, index):
if self._closed:
raise RuntimeError("Image-list serverless controller is closed")
if slot not in self.slots or slot in self._active:
raise RuntimeError("Unknown or occupied serverless image-list slot")
if type(index) is not int or not 0 <= index < len(self.state.items) or index in self._assigned:
raise ValueError("Invalid or duplicate serverless image-list index")
if self.state.items[index].get("index") != index:
raise ValueError("Serverless image-list item index mismatch")
self._assigned.add(index)
identity = f"{self._scope}:{index}"
handle = {"index": index, "execution_id": identity, "job_id": identity,
"worker_id": slot[1], "task": None}
self._active[slot] = handle
# Own execution identity BEFORE start, even if its caller loses the handle.
self._executions.add(identity)
task = asyncio.create_task(self._submit(handle))
self._submissions.add(task)
try:
await asyncio.shield(task)
return handle
except BaseException:
if not task.done():
task.cancel()
await asyncio.gather(task, return_exceptions=True)
self._remove_queue(handle["job_id"])
await self.manager.cancel(execution_id=identity)
raise
finally:
self._submissions.discard(task)
async def _submit(self, handle):
job = await prepare_serverless_item(
self.state.original_prompt, self.region, self.state.items[handle["index"]],
self.workers[handle["worker_id"]], request_id=handle["execution_id"], job_id=handle["job_id"])
if self._closed:
raise RuntimeError("Image-list serverless controller closed during preparation")
self.manager.ensure_capacity(1)
queue = asyncio.Queue()
async with self.manager.server.distributed_jobs_lock:
if self._closed:
raise RuntimeError("Image-list serverless controller is closed")
if handle["job_id"] in self.manager.server.distributed_pending_jobs:
raise RuntimeError("Duplicate serverless image-list queue identity")
self._queues[handle["job_id"]] = queue
self.manager.server.distributed_pending_jobs[handle["job_id"]] = queue
handle["sockets"] = {binding["socket"] for binding in job["manifest"]}
record_ids = await self.manager.start(handle["execution_id"], self.state.prompt_id, [job])
if len(record_ids) != 1 or record_ids[0] not in self.manager.tasks:
raise RuntimeError("Missing serverless image-list lifecycle task")
handle["task"] = self.manager.tasks[record_ids[0]]
handle["task"].add_done_callback(lambda _: self.event.set())
def _remove_queue(self, job_id):
queue = self._queues.pop(job_id, None)
# Same event-loop turn, without an await: never delete a replacement queue.
if queue is not None and self.manager.server.distributed_pending_jobs.get(job_id) is queue:
del self.manager.server.distributed_pending_jobs[job_id]
def poll(self, inflight):
if self._closed:
raise RuntimeError("Image-list serverless controller is closed")
for slot, handle in inflight.items():
if slot not in self.slots:
continue # Direct/master slots belong to the surrounding scheduler.
if self._active.get(slot) is not handle:
raise RuntimeError("Mismatched or duplicate serverless image-list completion handle")
task = handle["task"]
if task is None or not task.done():
continue
if task.cancelled():
raise RuntimeError("Serverless image-list task was cancelled")
if task.exception() is not None:
raise RuntimeError("Serverless image-list lifecycle task failed") from task.exception()
queue = self._queues.get(handle["job_id"])
if queue is None or self.manager.server.distributed_pending_jobs.get(handle["job_id"]) is not queue:
raise RuntimeError("Missing serverless image-list result queue")
if queue.qsize() != 1:
raise RuntimeError("Missing or duplicate serverless image-list terminal result")
message = queue.get_nowait()
if (not isinstance(message, dict) or message.get("worker_id") != handle["worker_id"]
or message.get("is_last") is not True):
raise RuntimeError("Mismatched serverless image-list result identity")
if "serverless_error" in message:
raise RuntimeError(f"Serverless image-list item {handle['index']} failed: {message['serverless_error']}")
media = message.get("media")
if not isinstance(media, dict) or set(media) != handle["sockets"]:
raise RuntimeError("Missing or unexpected serverless image-list result sockets")
if any(not isinstance(values, list) or len(values) != 1 or values[0] is None
for values in media.values()):
raise RuntimeError("Serverless image-list result requires exactly one value per socket")
overrides = {name: media[name] for name in runtime.VALUE_NAMES if name in media}
result = runtime.normalize_worker_result(media["images"], media.get("masks"), media.get("texts"),
values=overrides)
runtime._checked_overrides(result, runtime.end_override_names(
self.state.original_prompt[self.region["end_id"]].get("inputs", {})))
if result["mask"] is not None and result["mask"].shape[1:3] != result["image"].shape[1:3]:
raise ValueError("Serverless image-list result mask is not paired with the image grid")
self._remove_queue(handle["job_id"])
del self._active[slot]
return slot, handle["index"], result
return None
async def _close(self):
# Stop preparation/start before cancellation discovery, then retire exact
# destinations before manager callbacks can publish cancellation errors.
submissions = tuple(self._submissions)
for task in submissions:
if not task.done():
task.cancel()
await asyncio.gather(*submissions, return_exceptions=True)
for job_id in list(self._queues):
self._remove_queue(job_id)
await asyncio.gather(*(self.manager.cancel(execution_id=eid) for eid in self._executions))
self._active.clear()
async def close(self):
if self._close_task is None:
self._closed = True
self._close_task = asyncio.create_task(self._close())
cancelled = False
while not self._close_task.done():
try:
await asyncio.shield(self._close_task)
except asyncio.CancelledError:
cancelled = True
self._close_task.result()
if cancelled:
raise asyncio.CancelledError
async def create_serverless_slots(state, region, config, event):
workers = resolve_serverless_workers(config, state.worker_ids)
return ServerlessSlots(state, region, workers, get_manager() if workers else None, event)
+52 -5
View File
@@ -121,6 +121,46 @@ class DistributedCollectorNode:
return value[0]
return value
@staticmethod
def _combine_image_tensors(tensors, sources):
"""Validate BHWC batches, retaining alpha and deterministic input order."""
nonempty = []
reference_shape = None
reference_source = None
for tensor, source in zip(tensors, sources):
if not isinstance(tensor, torch.Tensor):
raise ValueError(f"Collector expected an IMAGE tensor from {source}")
shape = tuple(tensor.shape)
if tensor.ndim != 4 or shape[-1] not in (1, 3, 4) or min(shape[1:3]) <= 0:
raise ValueError(
f"Collector received invalid IMAGE shape {shape} from {source}; "
"expected [B, H, W, C] with 1, 3, or 4 channels"
)
# Empty master placeholders must not determine the worker image grid.
if shape[0] == 0:
continue
if reference_shape is None:
reference_shape, reference_source = shape, source
elif shape[1:3] != reference_shape[1:3]:
raise ValueError(
"Collector cannot combine images with different spatial dimensions: "
f"{reference_source} has shape {reference_shape}; {source} has shape {shape}"
)
nonempty.append(tensor)
if not nonempty:
raise ValueError("Collector cannot return an empty image batch")
channels = max(tensor.shape[-1] for tensor in nonempty)
normalized = []
for tensor in nonempty:
if tensor.shape[-1] == 1 and channels > 1:
tensor = tensor.expand(-1, -1, -1, 3)
if tensor.shape[-1] == 3 and channels == 4:
# RGB is opaque; never discard alpha from another participant.
tensor = torch.cat((tensor, torch.ones_like(tensor[..., :1])), dim=-1)
normalized.append(ensure_contiguous(tensor))
return ensure_contiguous(torch.cat(normalized, dim=0))
def _normalize_images_input(self, images):
"""Collapse ComfyUI list IMAGE inputs into a normal batched IMAGE tensor."""
if isinstance(images, (list, tuple)):
@@ -130,7 +170,9 @@ class DistributedCollectorNode:
raise TypeError("Collector expected IMAGE list items to be torch.Tensor instances")
if len(images) == 1:
return ensure_contiguous(images[0])
return ensure_contiguous(torch.cat([ensure_contiguous(image) for image in images], dim=0))
return self._combine_image_tensors(
images, [f"IMAGE list item {index}" for index in range(len(images))]
)
return ensure_contiguous(images)
def _normalize_audio_input(self, audio):
@@ -624,9 +666,11 @@ class DistributedCollectorNode:
):
"""Assemble final tensor, or return None when the job contains only audio."""
ordered_tensors = []
sources = []
if not delegate_mode and images_on_cpu is not None:
for i in range(master_batch_size):
ordered_tensors.append(images_on_cpu[i:i+1])
sources.append(f"master image {i}")
ordered_worker_ids = [str(worker_id) for worker_id in (worker_order or [])]
seen = set()
@@ -636,6 +680,7 @@ class DistributedCollectorNode:
continue
for idx in sorted(worker_images[worker_id_str].keys()):
ordered_tensors.append(worker_images[worker_id_str][idx])
sources.append(f"worker {worker_id_str!r} image {idx}")
# Append any unexpected worker ids deterministically.
for worker_id_str in sorted(worker_images.keys()):
@@ -643,6 +688,7 @@ class DistributedCollectorNode:
continue
for idx in sorted(worker_images[worker_id_str].keys()):
ordered_tensors.append(worker_images[worker_id_str][idx])
sources.append(f"worker {worker_id_str!r} image {idx}")
cpu_tensors = []
for t in ordered_tensors:
@@ -652,9 +698,9 @@ class DistributedCollectorNode:
cpu_tensors.append(t)
if cpu_tensors:
return ensure_contiguous(torch.cat(cpu_tensors, dim=0))
return self._combine_image_tensors(cpu_tensors, sources)
if fallback_images is not None:
return ensure_contiguous(fallback_images)
return self._combine_image_tensors([fallback_images], ["master fallback"])
return None
async def execute(self, images, audio, load_balance=False, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False, transfer_format=transfer_codec.DEFAULT_TRANSFER_FORMAT, transfer_quality=transfer_codec.DEFAULT_TRANSFER_QUALITY):
@@ -944,8 +990,9 @@ class DistributedCollectorNode:
f"{sorted(worker_images)}: {e}"
) from e
log(f"Master - Error combining images: {e}")
# Preserve collected audio for regular workers when image assembly fails.
return (images, combined_audio)
# Do not silently discard worker images or pass an empty
# fallback batch downstream to SaveImage.
raise
except BaseException:
# Remove the queue first, preventing cancelled work from delivering into
# a collector that is already unwinding.
+12 -2
View File
@@ -6,6 +6,8 @@ import json
from comfy_api.v0_0_2 import io
from ..utils.async_helpers import run_async_in_server_loop
from ..serverless_bridge.image_transfer import (FORMATS, DEFAULT_FORMAT, DEFAULT_QUALITY,
FORMAT_TOOLTIP, QUALITY_TOOLTIP, validate_options)
from ..image_list.runtime import (
MAX_VALUE_SLOTS,
VALUE_NAMES,
@@ -204,6 +206,10 @@ class DistributedImageListEnd(io.ComfyNode):
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.Combo.Input("transfer_format", options=list(FORMATS), default=DEFAULT_FORMAT,
optional=True, tooltip=FORMAT_TOOLTIP),
io.Int.Input("transfer_quality", default=DEFAULT_QUALITY, min=1, max=100,
optional=True, tooltip=QUALITY_TOOLTIP),
*[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],
@@ -251,7 +257,10 @@ class DistributedImageListEnd(io.ComfyNode):
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, worker_context_id="",
worker_item_identity=None, **values):
worker_item_identity=None, transfer_format=DEFAULT_FORMAT,
transfer_quality=DEFAULT_QUALITY, **values):
transfer_format, transfer_quality = validate_options(unwrap_single(transfer_format),
unwrap_single(transfer_quality))
worker_result_token = unwrap_single(worker_result_token)
master_callback_url = unwrap_single(master_callback_url)
worker_context_id = unwrap_single(worker_context_id)
@@ -300,7 +309,8 @@ class DistributedImageListEnd(io.ComfyNode):
set_master_result(context_id, int(unwrap_single(worker_item_index)),
str(worker_result_token), result)
else:
payload = encode_result(result)
payload = encode_result(result, transfer_format=transfer_format,
transfer_quality=transfer_quality)
run_async_in_server_loop(
submit_worker_result(str(master_callback_url), str(worker_result_token),
context_id, int(unwrap_single(worker_item_index)), payload), timeout=60,
+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "ComfyUI-Distributed"
description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
version = "1.5.0"
version = "1.6.0"
license = {file = "LICENSE"}
dependencies = []
+18 -2
View File
@@ -132,7 +132,7 @@ For typed inputs, `media.artifacts` contains records with exactly:
| Field | Value |
| --- | --- |
| `id`, `kind` | Unique artifact ID and one of the five native types |
| `id`, `kind` | Unique artifact ID and a native media type or internal `VALUE` |
| `descriptor` | Kind-specific codec metadata; use the shared codec, not pickle |
| `byte_length`, `sha256` | Exact payload byte count and lowercase SHA-256 hex |
| `transport` | `{"type": "inline", "data_base64": "..."}` or the URL transport below |
@@ -149,7 +149,7 @@ to owned staging paths. They are not an arbitrary filename-upload interface.
The returned `media` object has exactly `version`, `profile`, `request_id`,
`bindings`, `artifacts`, and `items`:
- `bindings`: every declared binding, unchanged, plus `count` (1–64).
- `bindings`: every declared binding, unchanged, plus `count` (1–64 for media; exactly 1 for End value overrides).
- `artifacts`: the same artifact schema above with inline or URL transport.
- `items`: `{ "binding_id": "...", "item_index": 0, "artifact_id": "..." }`
records covering every index exactly once, ordered by `item_index` when decoded.
@@ -160,6 +160,22 @@ only and must not be returned to the master. The executor reads validated
`media_exports` history records and bounded `/view` payloads before constructing
the external envelope. [contract.py](contract.py) is the schema authority.
## Internal image-list values
Image-list End uses sparse `value_1` … `value_10` bindings of kind `VALUE`.
These are internal exports, not new Collector sockets. `value-v1` descriptors
preserve BOOLEAN, arbitrary-size INT, finite binary64 FLOAT, STRING, IMAGE and
MASK distinctly; tensors/text reuse exact native codecs. Each binding returns
exactly one value. Only explicit End overrides are exported; the master rebuilds
unchanged values from its source items. Compression of the main IMAGE never
changes value-tensor fidelity.
Deploy the complete updated bridge, including `value_codec.py`, before using
End overrides. `MediaImportV1` retains its original five output positions and
appends the internal wildcard VALUE output at slot 5. Older bridge installations
do not support these bindings; ordinary provider health checks do not verify
this support. Request/result size limits still apply.
## Optional image return encoding
IMAGE `MediaExportV1` nodes may also contain literal `transfer_format` and
+7 -1
View File
@@ -24,7 +24,7 @@ _LAYOUTS = {"IMAGE": "BHWC", "MASK": "BHW", "AUDIO": "BCT"}
def _validate_kind(kind):
if type(kind) is not str or kind not in ("IMAGE", "MASK", "STRING", "AUDIO"):
if type(kind) is not str or kind not in ("IMAGE", "MASK", "STRING", "AUDIO", "VALUE"):
raise ValueError("Unsupported payload kind")
@@ -92,6 +92,9 @@ def encode_payload(kind: str, value, *, max_bytes: int = MAX_PAYLOAD_BYTES) -> t
"""
_validate_kind(kind)
limit = _byte_limit(max_bytes)
if kind == "VALUE":
from .value_codec import encode_value
return encode_value(value, max_bytes=limit)
if kind == "STRING":
return _encode_string(value, limit)
@@ -146,6 +149,9 @@ def decode_payload(kind: str, descriptor: dict, payload: bytes, *, max_bytes: in
raise ValueError("Payload must be bytes")
if len(payload) > limit:
raise ValueError("Payload exceeds byte limit")
if kind == "VALUE":
from .value_codec import decode_value
return decode_value(descriptor, payload, max_bytes=limit)
if kind == "IMAGE" and descriptor.get("encoding") == "image-transfer-v1":
# Lazy and IMAGE-only: exact tensors and other kinds stay unchanged.
from .image_codec import decode_image
+8 -3
View File
@@ -19,6 +19,9 @@ MAX_ARTIFACTS = 128
MAX_ITEMS = 64
MAX_DESCRIPTOR_BYTES = 128 * 1024
SOCKET_KINDS = {"images": "IMAGE", "audio": "AUDIO", "masks": "MASK", "texts": "STRING", "video": "VIDEO"}
# Internal End slots do not extend the public Collector or legacy media sockets.
VALUE_SOCKETS = tuple(f"value_{i}" for i in range(1, 11))
ARTIFACT_KINDS = frozenset(SOCKET_KINDS.values()) | {"VALUE"}
BINDING_FIELDS = {"binding_id", "collector_id", "job_id", "participant_id", "node_id", "socket", "kind"}
_ID = re.compile(r"[A-Za-z0-9_.:/-]{1,200}\Z")
_HASH = re.compile(r"[a-f0-9]{64}\Z")
@@ -68,7 +71,8 @@ def validate_binding(binding):
raise ValueError("Invalid media binding")
for field in BINDING_FIELDS - {"socket", "kind"}:
identifier(binding[field])
if SOCKET_KINDS.get(binding["socket"]) != binding["kind"]:
kind = "VALUE" if binding["socket"] in VALUE_SOCKETS else SOCKET_KINDS.get(binding["socket"])
if kind is None or kind != binding["kind"]:
raise ValueError("Media socket/type mismatch")
return binding
@@ -113,7 +117,7 @@ def validate_artifact(artifact, *, max_bytes=MAX_ARTIFACT_BYTES):
if not isinstance(artifact, dict) or set(artifact) != fields:
raise ValueError("Invalid media artifact")
identifier(artifact["id"])
if artifact["kind"] not in SOCKET_KINDS.values():
if artifact["kind"] not in ARTIFACT_KINDS:
raise ValueError("Unsupported media kind")
if not isinstance(artifact["descriptor"], dict):
raise ValueError("Invalid media descriptor")
@@ -191,7 +195,8 @@ def validate_response(response, *, request_id, bindings, max_total_bytes=MAX_TOT
if bid in counts or expected.get(bid) != binding:
raise ValueError("Media binding identity mismatch")
# Connected outputs must produce at least one value. Empty STRING is a value.
counts[bid] = bounded_int(row["count"], MAX_ITEMS, minimum=1)
maximum = 1 if binding["socket"] in VALUE_SOCKETS else MAX_ITEMS
counts[bid] = bounded_int(row["count"], maximum, minimum=1)
artifacts = response["artifacts"]
if not isinstance(artifacts, list) or not 1 <= len(artifacts) <= MAX_ARTIFACTS:
raise ValueError("Invalid media artifacts count")
+7 -3
View File
@@ -40,7 +40,7 @@ DRAIN_GRACE_SECONDS = 10.0
HTTP_TIMEOUT_SECONDS = 15.0
_UNCERTAIN_WORK = set()
GET_ATTEMPTS = 3
_IMPORT_SLOTS = {"IMAGE": 0, "AUDIO": 1, "MASK": 2, "STRING": 3, "VIDEO": 4}
_IMPORT_SLOTS = {"IMAGE": 0, "AUDIO": 1, "MASK": 2, "STRING": 3, "VIDEO": 4, "VALUE": 5}
# Used ONLY to run the shared manifest validator before /view. Never fetched.
_VALIDATION_TRANSPORT = {"type": "url", "url": "https://media-validation.invalid/artifact",
"expires_at": "9999-01-01T00:00:00Z"}
@@ -194,8 +194,12 @@ def _validate_request(input_data):
or aid not in indexed or (node_id, name) in destinations):
raise ValueError("Unknown or duplicate media input reference")
binding = exporters.get(node_id)
if binding is not None and (name != binding["socket"] or indexed[aid]["kind"] != binding["kind"]):
raise ValueError("Exporter input must match its declared media socket and kind")
if binding is not None:
# VALUE exports accept the six explicit runtime types. IMAGE/MASK
# imports can feed both native consumers and a sparse override.
allowed = {"VALUE", "IMAGE", "MASK", "STRING"} if binding["kind"] == "VALUE" else {binding["kind"]}
if name != binding["socket"] or indexed[aid]["kind"] not in allowed:
raise ValueError("Exporter input must match its declared media socket and kind")
destinations.add((node_id, name))
used.add(aid)
if used != set(indexed):
+15 -6
View File
@@ -11,7 +11,7 @@ from .codec import decode_payload, encode_payload
from .image_transfer import (FORMATS, DEFAULT_FORMAT, DEFAULT_QUALITY,
FORMAT_TOOLTIP, QUALITY_TOOLTIP, validate_options)
from .contract import (
MAX_ARTIFACT_BYTES, MAX_ITEMS, MAX_TOTAL_BYTES, SOCKET_KINDS,
MAX_ARTIFACT_BYTES, MAX_ITEMS, MAX_TOTAL_BYTES, SOCKET_KINDS, VALUE_SOCKETS, ARTIFACT_KINDS,
bounded_int, identifier, validate_binding, validate_json,
)
@@ -38,6 +38,9 @@ def _json(text):
def _cpu(kind, value):
# Runtime transfer only, never sampling or device discovery. Pure codecs
# deliberately reject non-CPU storage rather than making hidden transfers.
if kind == "VALUE":
from .value_codec import value_kind
kind = value_kind(value)
if kind in ("IMAGE", "MASK"):
return value.detach().cpu()
if kind == "AUDIO":
@@ -72,6 +75,8 @@ def export_values(request_id, binding, values, *, transfer_format=DEFAULT_FORMAT
raise ValueError("Image transfer controls apply only to IMAGE exports")
if not isinstance(values, list) or not 1 <= len(values) <= MAX_ITEMS:
raise ValueError("Connected media output has no values or exceeds item limit")
if binding["socket"] in VALUE_SOCKETS and len(values) != 1:
raise ValueError("End override requires exactly one value")
request_key = hashlib.sha256(request_id.encode()).hexdigest()
subfolder = f"distributed-media-v1/{request_key}"
root = Path(folder_paths.get_temp_directory()).resolve()
@@ -114,7 +119,7 @@ def import_value(artifact, filename):
if not isinstance(artifact, dict) or set(artifact) != fields:
raise ValueError("Invalid staged media metadata")
identifier(artifact["id"])
if artifact["kind"] not in SOCKET_KINDS.values():
if artifact["kind"] not in ARTIFACT_KINDS:
raise ValueError("Invalid staged media kind")
size = bounded_int(artifact["byte_length"], MAX_ARTIFACT_BYTES)
if not isinstance(filename, str):
@@ -138,13 +143,14 @@ class MediaImportV1(io.ComfyNode):
def define_schema(cls):
return io.Schema(node_id="MediaImportV1", category="distributed/internal", is_deprecated=True,
inputs=[io.String.Input("artifact_json"), io.String.Input("filename")],
outputs=[io.Image.Output(), io.Audio.Output(), io.Mask.Output(), io.String.Output(), io.Video.Output()])
outputs=[io.Image.Output(), io.Audio.Output(), io.Mask.Output(), io.String.Output(), io.Video.Output(),
io.AnyType.Output()])
@classmethod
def execute(cls, artifact_json, filename):
artifact = _json(artifact_json)
value = import_value(artifact, filename)
return io.NodeOutput(*(value if kind == artifact["kind"] else None for kind in SOCKET_KINDS.values()))
return io.NodeOutput(*(value if kind == artifact["kind"] else None for kind in (*SOCKET_KINDS.values(), "VALUE")))
class MediaExportV1(io.ComfyNode):
@@ -155,6 +161,7 @@ class MediaExportV1(io.ComfyNode):
io.Image.Input("images", optional=True), io.Audio.Input("audio", optional=True),
io.Mask.Input("masks", optional=True), io.String.Input("texts", optional=True, force_input=True),
io.Video.Input("video", optional=True),
*[io.AnyType.Input(name, optional=True) for name in VALUE_SOCKETS],
io.Combo.Input("transfer_format", options=list(FORMATS), default=DEFAULT_FORMAT,
optional=True, tooltip=FORMAT_TOOLTIP),
io.Int.Input("transfer_quality", default=DEFAULT_QUALITY, min=1, max=100,
@@ -163,10 +170,12 @@ class MediaExportV1(io.ComfyNode):
@classmethod
def execute(cls, request_id, binding_json, images=None, audio=None, masks=None, texts=None, video=None,
transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY):
transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY, **overrides):
binding = _json(_scalar(binding_json))
validate_binding(binding)
provided = dict(images=images, audio=audio, masks=masks, texts=texts, video=video)
if set(overrides) - set(VALUE_SOCKETS):
raise ValueError("Unknown End value slot")
provided = dict(images=images, audio=audio, masks=masks, texts=texts, video=video, **overrides)
if {name for name, value in provided.items() if value is not None} != {binding["socket"]}:
raise ValueError("Media export must receive exactly its declared typed socket")
exported = export_values(_scalar(request_id), binding, provided[binding["socket"]],
+85
View File
@@ -0,0 +1,85 @@
"""Exact six-type values for internal sparse End exports; no pack imports.
Tensor/text payloads reuse the existing media codec. Scalars have distinct
encodings: BOOLEAN is one byte, INT is minimal signed little-endian, FLOAT is
IEEE binary64. No scalar passes through float() or a lossy JSON number parser.
"""
import math
import struct
import torch
KINDS = frozenset(("IMAGE", "MASK", "STRING", "BOOLEAN", "INT", "FLOAT"))
_SCALAR_ENCODINGS = {"BOOLEAN": "bool-v1", "INT": "int-le-v1", "FLOAT": "float64-le-v1"}
def value_kind(value):
if isinstance(value, torch.Tensor):
if value.ndim == 4 and value.shape[-1] in (3, 4):
return "IMAGE"
if value.ndim == 3:
return "MASK"
raise ValueError("VALUE tensor requires IMAGE [B,H,W,3|4] or MASK [B,H,W]")
kind = {bool: "BOOLEAN", int: "INT", float: "FLOAT", str: "STRING"}.get(type(value))
if kind is None:
raise ValueError("VALUE requires IMAGE, MASK, STRING, INT, FLOAT or BOOLEAN")
if kind == "FLOAT" and not math.isfinite(value):
raise ValueError("VALUE numbers must be finite")
return kind
def encode_value(value, *, max_bytes):
from .codec import encode_payload
kind = value_kind(value)
if kind in ("IMAGE", "MASK", "STRING"):
descriptor, payload = encode_payload(kind, value, max_bytes=max_bytes)
else:
descriptor = {"encoding": _SCALAR_ENCODINGS[kind]}
size = ((value if value >= 0 else ~value).bit_length() + 8) // 8 if kind == "INT" else (8 if kind == "FLOAT" else 1)
if size > max_bytes:
raise ValueError("VALUE payload exceeds byte limit")
if kind == "BOOLEAN":
payload = b"\x01" if value else b"\x00"
elif kind == "INT":
payload = value.to_bytes(size, "little", signed=True)
else:
payload = struct.pack("<d", value)
return {"encoding": "value-v1", "type": kind, "descriptor": descriptor}, payload
def decode_value(descriptor, payload, *, max_bytes):
from .codec import decode_payload
if (set(descriptor) != {"encoding", "type", "descriptor"}
or descriptor["encoding"] != "value-v1"
or type(descriptor["type"]) is not str or descriptor["type"] not in KINDS):
raise ValueError("Invalid VALUE descriptor")
kind, inner = descriptor["type"], descriptor["descriptor"]
if kind in ("IMAGE", "MASK", "STRING"):
# Generic values always stay exact even if End.images is compressed.
if not isinstance(inner, dict) or inner.get("encoding") != (
"utf8-v1" if kind == "STRING" else "tensor-le-v1"):
raise ValueError("VALUE requires exact typed payloads")
value = decode_payload(kind, inner, payload, max_bytes=max_bytes)
if value_kind(value) != kind:
raise ValueError("VALUE tensor/type mismatch")
return value
if inner != {"encoding": _SCALAR_ENCODINGS[kind]}:
raise ValueError("Invalid VALUE scalar encoding")
if kind == "BOOLEAN":
if payload not in (b"\x00", b"\x01"):
raise ValueError("Invalid BOOLEAN payload")
return payload == b"\x01"
if kind == "FLOAT":
if len(payload) != 8:
raise ValueError("Invalid FLOAT payload length")
value = struct.unpack("<d", payload)[0]
if not math.isfinite(value):
raise ValueError("VALUE numbers must be finite")
return value
if not payload:
raise ValueError("Invalid INT payload length")
value = int.from_bytes(payload, "little", signed=True)
if ((value if value >= 0 else ~value).bit_length() + 8) // 8 != len(payload):
raise ValueError("Invalid INT payload length")
return value
@@ -0,0 +1,60 @@
"""CPU-only proof against real PromptQueue and IsChangedCache; no services."""
import asyncio
from pathlib import Path
import sys
from types import SimpleNamespace
core = Path(sys.argv[1]).resolve()
root = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(core))
import comfy.cli_args
comfy.cli_args.args.cpu = True
import execution
import nodes
from comfy_execution.graph import DynamicPrompt
sys.path.insert(0, str(root))
from image_list.lifecycle import has_safe_execution_targets
class Source:
RETURN_TYPES = ('IMAGE',)
@classmethod
def INPUT_TYPES(cls):
return {'required': {'seed': ('INT',)}}
@classmethod
def IS_CHANGED(cls, seed):
return f'fingerprint-{seed}'
nodes.NODE_CLASS_MAPPINGS['ProofSource'] = Source
nodes.NODE_CLASS_MAPPINGS['ProofUtility'] = type('Utility', (), {'OUTPUT_NODE': True})
prompt = {
'source': {'class_type': 'ProofSource', 'inputs': {'seed': 7}},
'start': {'class_type': 'DistributedImageListStart', 'inputs': {'images': ['source', 0]}},
'op': {'class_type': 'ProofUtility', 'inputs': {'image': ['start', 0]}},
'end': {'class_type': 'DistributedImageListEnd', 'inputs': {
'images': ['op', 0], 'context': ['start', 3]}},
}
server = SimpleNamespace(queue_updated=lambda: None)
queue = execution.PromptQueue(server)
server.prompt_queue = queue
queue.put((0, 'owned', prompt, {}, ['end']))
item, _ = queue.get(timeout=0)
assert item[2] is prompt
snapshot = queue.get_current_queue_volatile()[0][0][2]
assert snapshot is not prompt
assert has_safe_execution_targets(server, 'owned', prompt)
cache = execution.IsChangedCache('owned', DynamicPrompt(prompt), None)
assert asyncio.run(cache.get('source')) == ['fingerprint-7']
assert prompt['source']['is_changed'] == ['fingerprint-7']
assert 'is_changed' not in snapshot['source']
assert has_safe_execution_targets(server, 'owned', prompt), 'core fingerprint invalidated safe roots'
prompt['source']['inputs']['seed'] = 8
assert not has_safe_execution_targets(server, 'owned', prompt), 'changed input accepted'
prompt['source']['inputs']['seed'] = 7
assert not has_safe_execution_targets(server, 'not-owned', prompt)
queue.get_current_queue_volatile()[0][0][4].append('op')
assert not has_safe_execution_targets(server, 'owned', prompt), 'unsafe output root accepted'
print('FINGERPRINT_PROOF_OK')
+38 -16
View File
@@ -198,8 +198,8 @@ async def setup():
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)
def checked_encode(result, **options):
encoded = encode_result(result, **options)
if 'values' in encoded:
assert set(encoded['values']) == {f'value_{n}' for n in (2, 3, 4, 5, 7, 8, 9, 10)}
return encoded
@@ -246,7 +246,10 @@ async def worker_main():
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'),
delivery = next(node['inputs'] for node in entry[2].values()
if node['class_type'] == 'DistributedImageListEnd')
submissions.append({'prompt_id': entry[1], 'item_index': delivery['worker_item_index'],
'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']})
@@ -514,7 +517,14 @@ async def coordinator():
ALLOW_MASTER = True
config['settings']['master_delegate_only'] = False
participation_counts = {}
for mode, selected in [('mixed', [w['id'] for w in workers]), ('master-only', [])]:
for mode, selected, fmt in [('mixed', [w['id'] for w in workers], 'exact'),
('master-only', [], 'exact'),
('mixed-png', [w['id'] for w in workers], 'png'),
('master-only-png', [], 'png')]:
case_prompt = copy.deepcopy(prompt)
if fmt != 'exact':
case_prompt['4']['inputs'].update(transfer_format=fmt, transfer_quality=73)
saved_case = copy.deepcopy(case_prompt)
# 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)
@@ -530,8 +540,8 @@ async def coordinator():
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')
case_prompt, None, 'test', selected), 45)
assert case_prompt == saved_case and 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'
@@ -540,7 +550,7 @@ async def coordinator():
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 = []
remote_assignments = []
async with aiohttp.ClientSession() as session:
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:
@@ -550,27 +560,38 @@ async def coordinator():
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})
assigned = [submission['item_index'] for submission in submissions]
# A repeated item can hit native processing cache even
# when End changes return format. Count fresh deliveries,
# not recomputation, for exact-once scheduling evidence.
assert Counter(calls) <= Counter(assigned), (mode, worker['id'], calls, assigned)
assert len(calls) == len({event['prompt_id'] for event in events}), (mode, events)
assert {event['prompt_id'] for event in events} <= {
submission['prompt_id'] for submission in submissions}
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)
remote_assignments.extend(assigned)
assert not current['running'] and not current['pending'], current
if mode == 'mixed':
assert remote_calls, 'mixed execution failed to use direct workers'
if mode.startswith('mixed'):
assert remote_assignments, '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)
assert not remote_assignments and len(local_calls) == COUNT
assert Counter(local_calls + remote_assignments) == Counter(range(COUNT)), (local_calls, remote_assignments)
check_carried_values()
images, masks, texts = CAPTURED[-1]
assert len(images) == len(masks) == len(texts) == COUNT
for i, row in enumerate(source_rows()):
assert torch.equal(images[i], row['image'] + row['strength']), (mode, i)
expected = row['image'] + row['strength']
if fmt == 'png' and i in remote_assignments:
expected = expected.clone()
expected[..., :3] = (expected[..., :3].clamp(0, 1) * 255).to(torch.uint8).float() / 255
assert torch.equal(images[i], expected), (mode, i)
assert torch.equal(masks[i], row['mask'] * 0.5), (mode, i)
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)}
participation_counts[mode] = {'master': len(local_calls), 'workers': len(remote_assignments)}
# 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)
@@ -599,8 +620,9 @@ async def coordinator():
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))
assert len(outer_executions) == 3 + len(participation_counts) and all(row['captures'] == 1 for row in outer_executions)
print('IMAGE_LIST_TRANSFER_HTTP_OK mixed-png and master-only-png; masks/text/alpha/local images exact')
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'],
@@ -0,0 +1,169 @@
"""Real CPU PromptExecutor + standalone bridge proof, no provider or service."""
import asyncio
import importlib
import json
from pathlib import Path
import shutil
import struct
import sys
import types
from media_end_to_end_check import bootstrap
ROOT = Path(__file__).resolve().parents[2]
async def check(base):
import torch
import execution
import folder_paths
import nodes
import server
from app.assets.manager import default_asset_manager
from comfy_api.v0_0_2 import io
torch.set_num_threads(1)
assert not torch.cuda.is_initialized()
assert execution.comfy.model_management.get_torch_device().type == 'cpu'
# Deploy ONLY the bridge folder to a separate import root. Import it before
# master code: it cannot rely on the Distributed image_list package.
shutil.copytree(ROOT / 'serverless_bridge', base / 'endpoint_bridge')
sys.path.insert(0, str(base))
previously_loaded = set(sys.modules)
bridge = importlib.import_module('endpoint_bridge')
handler = importlib.import_module('endpoint_bridge.handler')
extension = await bridge.comfy_entrypoint()
classes = await extension.get_node_list()
assert not any('.image_list' in name for name in set(sys.modules) - previously_loaded)
calls = []
class CalculateOverrides(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(node_id='CalculateOverrides', inputs=[io.Image.Input('image'), io.Mask.Input('mask'),
*[io.AnyType.Input(name) for name in ('integer', 'boolean', 'floating', 'text', 'reference', 'reference_mask')]],
outputs=[io.Image.Output(), io.Mask.Output(), io.String.Output(), io.Int.Output(),
io.Float.Output(), io.Boolean.Output()])
@classmethod
def execute(cls, image, mask, integer, boolean, floating, text, reference, reference_mask):
assert type(integer) is int and type(boolean) is bool and type(floating) is float
calls.append(integer)
return io.NodeOutput(reference + image, reference_mask + mask,
text + ':calculated', integer + 7, floating / 2, not boolean)
class MustNotExecute(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(node_id='MustNotExecute', inputs=[io.Image.Input('images')], outputs=[], is_output_node=True)
@classmethod
def execute(cls, images):
raise AssertionError('unrelated output was not pruned')
for cls in [*classes, CalculateOverrides, MustNotExecute]:
cls.GET_NODE_INFO_V1()
nodes.NODE_CLASS_MAPPINGS[cls.GET_SCHEMA().node_id] = cls
package = '_serverless_value_master'
for suffix in ('', '.api', '.api.orchestration', '.image_list'):
module = types.ModuleType(package + suffix)
module.__path__ = [str(ROOT.joinpath(*suffix.lstrip('.').split('.')))] if suffix else [str(ROOT)]
sys.modules[module.__name__] = module
adapter = importlib.import_module(package + '.image_list.serverless')
decode = importlib.import_module(package + '.api.orchestration.serverless_typed_results').collect_typed_results
rt = adapter.runtime
assets = default_asset_manager()
ps = server.PromptServer(asyncio.get_running_loop(), assets)
def node(kind, **inputs):
return {'class_type': kind, 'inputs': inputs}
image = torch.arange(24, dtype=torch.float64).reshape(1, 2, 3, 4) / 997
mask = torch.arange(6, dtype=torch.bfloat16).reshape(1, 2, 3) / 16
values = {'value_1': 2**63 + 123, 'value_2': False, 'value_3': -0.0, 'value_4': ' 雪\r\n',
'value_5': image.repeat(2, 1, 1, 1), 'value_6': mask.repeat(2, 1, 1),
'value_7': 'master-only unchanged'}
item = {'image': image, 'mask': mask, 'values': values, 'index': 0}
graph = {
'start': node(rt.START_CLASS),
'calculate': node('CalculateOverrides', image=['start', 0], mask=['start', 1],
integer=['start', 4], boolean=['start', 5], floating=['start', 6], text=['start', 7],
reference=['start', 8], reference_mask=['start', 9]),
'end': node(rt.END_CLASS, context=['start', 3], images=['start', 0]),
'preview': node('MustNotExecute', images=['start', 0]),
}
names = ('value_1', 'value_2', 'value_4', 'value_6', 'value_8', 'value_10')
expected = dict(zip(names, (values['value_5'] + image, values['value_6'] + mask,
values['value_4'] + ':calculated', values['value_1'] + 7, -0.0, True)))
graph['end']['inputs'].update({name: ['calculate', index] for index, name in enumerate(names)})
region = {'start_id': 'start', 'end_id': 'end', 'node_ids': set(graph)}
worker = {'id': 'pool', 'type': 'serverless', 'endpoint_id': 'ep', 'timeout_seconds': 30,
'jobs_per_run': 16, 'enabled': True, 'media_profile': 'comfy-media-v1'}
async def run(prompt, request_id):
job = await adapter.prepare_serverless_item(prompt, region, item, worker,
request_id=request_id, job_id='job')
assert not any('Distributed' in n['class_type'] for n in job['workflow'].values())
workflow, media, legacy, targets = handler._validate_request(
{key: job[key] for key in ('workflow', 'images', 'media')})
stage = handler._Staging(folder_paths.get_input_directory(), request_id)
async def submit(native):
assert 'preview' not in native
valid = await execution.validate_prompt(request_id, native, targets)
assert valid[0], valid
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
await asyncio.to_thread(executor.execute, native, request_id, {}, targets)
assert executor.success, executor.status_messages
return executor.history_result
async def fetch(artifact):
transport = artifact['transport']
path = Path(folder_paths.get_temp_directory()) / transport['subfolder'] / transport['filename']
return path.read_bytes()
execution_adapter = types.SimpleNamespace(media=media, targets=targets, stage=stage,
remaining=lambda: 30, submit_and_wait=submit, fetch=fetch)
try:
response = await handler._run(execution_adapter, workflow, legacy, None, None)
grouped = await decode(response, media, job['manifest'])
finally:
stage.close(cleanup=True)
result = rt.normalize_worker_result(grouped['job']['images'], values={
name: grouped['job'][name] for name in rt.VALUE_NAMES if name in grouped['job']})
state = rt.ImageListJob('ctx', 'start', 'end', prompt, [item], ['pool'], '')
return job, rt._collect_output_lists([result], state=state, include_values=True)[3][0]
job, actual = await run(graph, 'calculated')
assert set(b['socket'] for b in job['manifest']) == {'images', *names}
assert calls == [values['value_1']]
for name, wanted in {**values, **expected}.items():
got = actual[name]
assert type(got) is type(wanted), name
if isinstance(wanted, torch.Tensor):
assert got.dtype == wanted.dtype and got.shape == wanted.shape
assert torch.equal(got.contiguous().view(torch.uint8), wanted.contiguous().view(torch.uint8)), name
elif isinstance(wanted, float):
assert struct.pack('!d', got) == struct.pack('!d', wanted), name
else:
assert got == wanted, name
# Direct overrides exercise exact VALUE import injection for bool/int/float/
# string, native tensor import slots, and no execution of the calculation.
for name in names:
graph['end']['inputs'].pop(name)
graph['end']['inputs'].update({f'value_{i + 1}': ['start', i + 4] for i in range(6)})
_, direct = await run(graph, 'passthrough')
assert calls == [values['value_1']]
assert type(direct['value_1']) is int and direct['value_1'] == values['value_1']
assert direct['value_2'] is False
assert struct.pack('!d', direct['value_3']) == struct.pack('!d', -0.0)
assert direct['value_4'] == values['value_4']
assert torch.equal(direct['value_5'], values['value_5'])
assert torch.equal(direct['value_6'], values['value_6'])
assert not torch.cuda.is_initialized()
print('IMAGE_LIST_SERVERLESS_VALUES_CPU_OK calculated six types; sparse master overlay; standalone bridge; real PromptExecutor')
if __name__ == '__main__':
asyncio.run(check(bootstrap()))
+7 -2
View File
@@ -432,12 +432,17 @@ async def main():
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 [input.id for input in end_schema.inputs] == [
'images', 'masks', 'texts', 'context', 'transfer_format', 'transfer_quality'] + list(runtime.VALUE_NAMES)
format_input, quality_input = end_schema.inputs[4:6]
assert format_input.options == list(runtime.IMAGE_TRANSFER_FORMATS)
assert format_input.default == 'exact' and quality_input.default == 90
assert format_input.optional and quality_input.optional
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:]:
for override in end_schema.inputs[6:]:
assert override.optional and override.lazy
assert override.get_io_type() == 'IMAGE,MASK,STRING,INT,FLOAT,BOOLEAN'
getter = mapping['DistributedImageListGetValue']
+30 -2
View File
@@ -405,8 +405,9 @@ def test_exact_collector_rejects_old_worker_8bit_result(transport, monkeypatch,
def test_exact_collector_shape_mismatch_raises_instead_of_returning_local_images(transport, monkeypatch):
async def scenario():
async with endpoint(transport, monkeypatch) as ctx:
await send(ctx, frames(count=1, channels=4))
with pytest.raises(RuntimeError, match="shape|combin|mismatch"):
# Channel mixing is supported; incompatible spatial sizes still fail.
await send(ctx, frames(count=1, channels=4, width=10))
with pytest.raises(RuntimeError, match="spatial dimensions"):
await ctx.node.execute(
frames(count=1, channels=3), None, multi_job_id=JOB_ID,
enabled_worker_ids=json.dumps([WORKER_ID]), transfer_format="exact",
@@ -416,6 +417,33 @@ def test_exact_collector_shape_mismatch_raises_instead_of_returning_local_images
asyncio.run(scenario())
@pytest.mark.parametrize("local_channels,worker_channels", [(3, 4), (4, 3)])
def test_exact_collector_mixed_rgb_rgba_preserves_every_value(transport, monkeypatch, local_channels, worker_channels):
local = frames(count=1, channels=local_channels)
worker = frames(count=1, channels=worker_channels)
async def scenario():
async with endpoint(transport, monkeypatch) as ctx:
await send(ctx, worker)
result, _audio = await ctx.node.execute(
local, None, multi_job_id=JOB_ID,
enabled_worker_ids=json.dumps([WORKER_ID]), transfer_format="exact",
)
assert result.shape == (2, 7, 9, 4)
assert result.is_contiguous()
# RGB views into RGBA have a four-channel stride; materialize the
# sliced comparison without changing the output-contiguity check.
assert_exact(result[:1, ..., :local_channels].contiguous(), local)
assert_exact(result[1:, ..., :worker_channels].contiguous(), worker)
if local_channels == 3:
assert torch.all(result[:1, ..., 3] == 1)
if worker_channels == 3:
assert torch.all(result[1:, ..., 3] == 1)
assert JOB_ID not in transport.server.distributed_pending_jobs
asyncio.run(scenario())
def test_exact_incomplete_worker_raises_and_cleans_queue(transport, monkeypatch):
"""A timed-out worker's first frame is not a successful partial result."""
envelope, parts = exact_chunk(transport, count=1)
+225
View File
@@ -0,0 +1,225 @@
"""Real PNG/tensor regressions; only server/config/HTTP plumbing is stubbed."""
import asyncio
import importlib.util
import io
import base64
import sys
import types
from pathlib import Path
from unittest.mock import patch
import numpy as np
import pytest
import torch
from aiohttp import web
from PIL import Image
ROOT = Path(__file__).resolve().parents[1]
def _load_file(name, path):
spec = importlib.util.spec_from_file_location(name, path)
module = importlib.util.module_from_spec(spec)
sys.modules[name] = module
spec.loader.exec_module(module)
return module
@pytest.fixture
def media_modules():
# Existing harnesses provide ComfyUI/config doubles. Restore their module
# replacements so this file can coexist with the rest of the test suite.
with patch.dict(sys.modules):
image_utils = _load_file("real_image_utils", ROOT / "utils/image.py")
collector_tests = _load_file("collector_harness", ROOT / "tests/test_collector_list_inputs.py")
collector = collector_tests._load_collector_module()
sys.modules["aiohttp"].web = web
route_tests = _load_file("route_harness", ROOT / "tests/api/test_distributed_queue.py")
routes = route_tests._load_job_routes_module()
for module in (collector, routes):
module.pil_to_tensor = image_utils.pil_to_tensor
module.ensure_contiguous = image_utils.ensure_contiguous
collector.tensor_to_pil = image_utils.tensor_to_pil
routes.prompt_server = collector.prompt_server
yield collector, routes
def _png_payload(image):
buffer = io.BytesIO()
image.save(buffer, format="PNG")
return "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode("ascii")
@pytest.mark.parametrize("mode", ["RGB", "RGBA", "L", "LA", "palette", "palette_alpha"])
def test_png_decoder_preserves_alpha_and_normalizes_opaque_modes(media_modules, mode):
_, routes = media_modules
if mode.startswith("palette"):
image = Image.new("P", (12, 8), 0)
image.putpalette([16, 32, 64] + [0] * 765)
if mode == "palette_alpha":
image.info["transparency"] = 0
else:
color = {"RGB": (16, 32, 64), "RGBA": (16, 32, 64, 96), "L": 32, "LA": (32, 96)}[mode]
image = Image.new(mode, (12, 8), color)
alpha = mode in {"RGBA", "LA", "palette_alpha"}
expected = np.array(image.convert("RGBA" if alpha else "RGB"), dtype=np.float32) / 255
decoded = routes._decode_canonical_png_tensor(_png_payload(image))
assert decoded.shape == (1, 8, 12, 4 if alpha else 3)
assert decoded.is_contiguous()
np.testing.assert_array_equal(decoded[0].numpy(), expected)
async def _master_collect(module, local, worker_images, delegate=False):
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
queue = asyncio.Queue()
for index, tensor in enumerate(worker_images):
queue.put_nowait({"worker_id": "worker-a", "image_index": index,
"tensor": tensor, "is_last": index == len(worker_images) - 1})
module.prompt_server.distributed_pending_jobs = {"image-test": queue}
return await module.DistributedCollectorNode().execute(
local, None, multi_job_id="image-test", enabled_worker_ids='["worker-a"]',
delegate_only=delegate,
)
@pytest.mark.parametrize("local_channels,worker_channels", [(3, 3), (4, 4), (3, 4), (4, 3), (None, 4)])
def test_worker_png_callback_and_collector_preserve_pixels_and_alpha(
media_modules, local_channels, worker_channels
):
module, routes = media_modules
local = None if local_channels is None else torch.full((1, 8, 12, local_channels), 0.5)
worker = torch.full((2, 8, 12, worker_channels), 0.25)
if worker_channels == 4:
worker[..., 3] = torch.linspace(0, 1, 12).reshape(1, 1, 12)
responses = []
async def scenario():
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
module.prompt_server.distributed_pending_jobs = {"image-test": asyncio.Queue()}
class Response:
def __init__(self, payload):
self.payload = payload
async def __aenter__(self):
async def request_json():
return self.payload
response = await routes.job_complete_endpoint(types.SimpleNamespace(json=request_json))
responses.append(response.status)
assert response.status == 200, response.text
return self
async def __aexit__(self, *_args):
return False
def raise_for_status(self):
pass
class Session:
def post(self, _url, json, timeout):
return Response(json)
async def get_session():
return Session()
module.get_client_session = get_session
module.encode_audio_payload = lambda _audio: None
collector = module.DistributedCollectorNode()
await collector.send_batch_to_master(worker, None, "image-test", "http://master", "worker-a")
result, _ = await collector.execute(
local, None, multi_job_id="image-test", enabled_worker_ids='["worker-a"]',
delegate_only=local is None,
)
assert "image-test" not in module.prompt_server.distributed_pending_jobs
return result
result = asyncio.run(scenario())
channels = max(local_channels or 3, worker_channels)
offset = 0 if local is None else 1
assert responses == [200, 200]
assert result.shape == (offset + 2, 8, 12, channels)
assert result.is_contiguous()
if local is not None:
assert torch.equal(result[:1, ..., :local_channels], local)
if local_channels == 3 and channels == 4:
assert torch.all(result[:1, ..., 3] == 1)
assert torch.allclose(result[offset:, ..., :worker_channels], worker, atol=1 / 255)
if worker_channels == 3 and channels == 4:
assert torch.all(result[offset:, ..., 3] == 1)
def test_collector_mixed_rgb_rgba_list_promotes_rgb_to_opaque_rgba(media_modules):
module, _ = media_modules
rgb = torch.zeros(1, 8, 12, 3)
rgba = torch.full((1, 8, 12, 4), 0.5)
result, _ = module.DistributedCollectorNode().run(images=[rgb, rgba])
assert result.shape == (2, 8, 12, 4)
assert torch.equal(result[:1, ..., :3], rgb)
assert torch.all(result[:1, ..., 3] == 1)
assert torch.equal(result[1:], rgba)
@pytest.mark.parametrize("other_channels", [1, 3, 4])
def test_collector_preserves_grayscale_or_promotes_it_for_color_batches(media_modules, other_channels):
module, _ = media_modules
gray = torch.full((1, 8, 12, 1), 0.25)
other = torch.full((1, 8, 12, other_channels), 0.5)
result, _ = module.DistributedCollectorNode().run(images=[gray, other])
assert result.shape == (2, 8, 12, other_channels)
assert torch.all(result[:1, ..., :min(other_channels, 3)] == 0.25)
if other_channels == 4:
assert torch.all(result[:1, ..., 3] == 1)
assert torch.equal(result[1:], other)
def test_mixed_channel_merge_keeps_configured_worker_and_image_order(media_modules):
module, _ = media_modules
collector = module.DistributedCollectorNode()
master = torch.zeros(1, 8, 12, 3)
worker_a = torch.full((1, 8, 12, 4), 0.25)
worker_b = torch.full((1, 8, 12, 3), 0.5)
unexpected = torch.full((1, 8, 12, 3), 0.75)
result = collector._reorder_and_combine_tensors(
{"worker-a": {0: worker_a}, "worker-b": {2: worker_b, 0: master},
"unexpected": {0: unexpected}},
["worker-b", "worker-a"], 1, master, False, master,
)
assert result.shape == (5, 8, 12, 4)
assert result[:, 0, 0, 0].tolist() == [0, 0, 0.5, 0.25, 0.75]
assert result[:, 0, 0, 3].tolist() == [1, 1, 1, 0.25, 1]
def test_collector_spatial_mismatch_raises_with_shapes_not_successful_fallback(media_modules):
module, _ = media_modules
local = torch.zeros(1, 8, 12, 3)
worker = torch.zeros(1, 8, 13, 3)
with pytest.raises(ValueError, match="spatial dimensions") as error:
asyncio.run(_master_collect(module, local, [worker]))
message = str(error.value)
assert "master image 0" in message
assert "worker 'worker-a' image 0" in message
assert "(1, 8, 12, 3)" in message and "(1, 8, 13, 3)" in message
assert "image-test" not in module.prompt_server.distributed_pending_jobs
@pytest.mark.parametrize("shape", [(1, 8, 12), (1, 8, 12, 5)])
def test_collector_rejects_invalid_worker_image_shape(media_modules, shape):
module, _ = media_modules
with pytest.raises(ValueError, match="worker 'worker-a' image 0"):
asyncio.run(_master_collect(module, None, [torch.zeros(shape)], delegate=True))
def test_delegate_master_ignores_empty_placeholder_shape(media_modules):
module, _ = media_modules
placeholder = torch.empty(0, 64, 64, 3)
worker = torch.full((1, 8, 12, 4), 0.5)
result, _ = asyncio.run(_master_collect(module, placeholder, [worker], delegate=True))
assert torch.equal(result, worker)
def test_empty_image_results_raise_instead_of_reaching_save_image(media_modules):
module, _ = media_modules
with pytest.raises(ValueError, match="empty image batch"):
asyncio.run(_master_collect(module, torch.empty(0, 64, 64, 3),
[torch.empty(0, 8, 12, 3)], delegate=True))
-1
View File
@@ -491,7 +491,6 @@ class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
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)})
@@ -0,0 +1,81 @@
"""Execution-target proof must survive core cache annotations, not graph edits."""
import copy
import os
from pathlib import Path
import subprocess
import sys
from types import SimpleNamespace
import pytest
from image_list.lifecycle import has_safe_execution_targets
from test_image_list_master_runtime import graph
ROOT = Path(__file__).resolve().parents[1]
@pytest.fixture
def proof(monkeypatch):
prompt = graph()
queued = copy.deepcopy(prompt)
entry = [0, 'owned', queued, {}, ['end']]
server = SimpleNamespace(prompt_queue=SimpleNamespace(
get_current_queue_volatile=lambda: ([entry], [])))
monkeypatch.setitem(sys.modules, 'nodes', SimpleNamespace(NODE_CLASS_MAPPINGS={
'Operation': type('Flagged', (), {'OUTPUT_NODE': True})}))
return prompt, queued, entry, server
@pytest.mark.parametrize('value', [['digest'], [None], float('nan'), {'nested': [1, 2]}])
@pytest.mark.parametrize('side', ['executing', 'queued', 'both'])
def test_top_level_fingerprints_do_not_invalidate_proof(proof, value, side):
prompt, queued, _, server = proof
if side in ('executing', 'both'):
prompt['source']['is_changed'] = value
if side in ('queued', 'both'):
queued['source']['is_changed'] = ['different queued fingerprint']
assert has_safe_execution_targets(server, 'owned', prompt)
# Comparisons must not strip or rewrite the caller's data.
if side in ('executing', 'both'):
assert prompt['source']['is_changed'] is value
if side in ('queued', 'both'):
assert queued['source']['is_changed'] == ['different queued fingerprint']
@pytest.mark.parametrize('change', ['input', 'link', 'class', 'node_id', 'nested_is_changed', 'other_metadata'])
def test_fingerprint_tolerance_does_not_hide_graph_changes(proof, change):
prompt, _, _, server = proof
prompt['source']['is_changed'] = ['digest']
if change == 'input':
prompt['source']['inputs']['seed'] = 2
elif change == 'link':
prompt['op']['inputs']['image'] = ['start', 1]
elif change == 'class':
prompt['source']['class_type'] = 'OtherSource'
elif change == 'node_id':
prompt['new_source'] = prompt.pop('source')
prompt['start']['inputs']['images'] = ['new_source', 0]
elif change == 'nested_is_changed':
prompt['source']['inputs']['is_changed'] = True
else:
prompt['source']['_meta'] = {'title': 'changed'}
assert not has_safe_execution_targets(server, 'owned', prompt)
@pytest.mark.parametrize('roots', [['op', 'end'], [], ['end', 'end'], None, [True]])
def test_fingerprints_never_grant_unscoped_or_invalid_roots(proof, roots):
prompt, _, entry, server = proof
prompt['source']['is_changed'] = ['digest']
entry[4] = roots
assert not has_safe_execution_targets(server, 'owned', prompt)
def test_real_core_queue_snapshot_and_fingerprint():
core = Path(os.environ.get('COMFYUI_SOURCE_ROOT', str(Path(sys.executable).parents[2])))
if not (core / 'execution.py').exists():
pytest.skip('real ComfyUI source not available beside this interpreter')
result = subprocess.run(
[sys.executable, str(ROOT / 'tests/helpers/image_list_fingerprint_proof.py'), str(core)],
text=True, capture_output=True, timeout=90)
assert result.returncode == 0, result.stdout + result.stderr
assert 'FINGERPRINT_PROOF_OK' in result.stdout
+44
View File
@@ -191,6 +191,50 @@ class MasterRuntimeTests(unittest.IsolatedAsyncioTestCase):
self.assertEqual(result[2], [str(i) for i in range(9)])
self.preflight.assert_awaited_once_with(self.workers, require_idle=False)
async def test_busy_direct_peer_does_not_block_selected_cloud_slots(self):
self.workers.append({'id': 'worker', 'type': 'remote'})
context = self.create(participates=False, count=3)
self.config['workers'] = self.workers + [{'id': 'pool', 'type': 'serverless'}]
self.state.worker_ids.append('pool')
self.state.original_prompt['end']['inputs']['value_10'] = ['start', 2]
fixture = self
cloud_calls = []
async def busy_preflight(workers, **kwargs):
if kwargs['require_idle']:
raise RuntimeError('No selected feature-compatible direct worker is idle')
return {}
self.preflight.side_effect = busy_preflight
class Controller:
slots = [('serverless', 'pool', 0)]
close = AsyncMock()
async def dispatch(self, slot, index):
cloud_calls.append(index)
return {'index': index}
def poll(self, inflight):
slot = self.slots[0]
if slot in inflight:
index = inflight[slot]['index']
return slot, index, {'image': fixture.items[index]['image'],
'values': {'value_10': index}}
return None
controller = Controller()
create_slots = AsyncMock(return_value=controller)
with patch.dict(sys.modules, {'_master_slot_tests.image_list.serverless':
types.SimpleNamespace(create_serverless_slots=create_slots)}):
result = await self.rt.run_image_list_job(context, include_values=True)
self.assertEqual(cloud_calls, [0, 1, 2])
self.assertEqual(result[3], [{'value_1': str(i), 'value_10': i} for i in range(3)])
self.preflight.assert_awaited_once_with(self.workers, require_idle=False, require_value_results=True)
self.network.get_client_session.assert_not_awaited()
controller.close.assert_awaited_once()
self.assertFalse(self.local_calls or self.dispatched or self.rt._MASTER_JOBS)
async def test_orchestrator_mode_never_instantiates_local_slot(self):
self.workers.append({'id': 'worker', 'type': 'remote'})
await self.rt.run_image_list_job(self.create(participates=False))
+19 -3
View File
@@ -272,12 +272,28 @@ class MasterSelectionTests(unittest.IsolatedAsyncioTestCase):
await self.queue()
self.q.queue_prompt_payload.assert_not_awaited()
async def test_serverless_only_uses_master_without_provider_dispatch(self):
self.config["workers"] = [{"id": "pool", "type": "serverless", "enabled": True}]
async def test_serverless_pool_is_retained_with_participating_master(self):
self.config["workers"] = [{
"id": "pool", "name": "Cloud", "type": "serverless", "enabled": True,
"endpoint_id": "endpoint", "media_profile": "comfy-media-v1",
}]
result = await self.queue()
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
self.assertEqual(result[2], 0)
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["pool"])
self.assertEqual(result[2], 1)
self.preflight.assert_not_awaited()
self.q.dispatch_worker_prompt.assert_not_awaited()
self.assertEqual(self.ps.distributed_pending_jobs, {})
async def test_legacy_serverless_pool_is_rejected_not_silently_ignored(self):
self.config["workers"] = [{
"id": "pool", "name": "Cloud", "type": "serverless", "enabled": True,
"endpoint_id": "endpoint",
}]
with self.assertRaisesRegex(ValueError, "comfy-media-v1"):
await self.queue()
self.q.queue_prompt_payload.assert_not_awaited()
self.preflight.assert_not_awaited()
async def test_incompatible_worker_is_not_silently_hidden_by_master(self):
self.preflight.side_effect = RuntimeError("upgrade worker")
+11 -2
View File
@@ -291,16 +291,25 @@ def test_legacy_worker_path_still_consumes_and_delivers_job_context(boundaries):
assert rt.get_master_job(context).received_indexes == {0}
def test_remote_cache_hit_submits_fresh_callback_metadata(boundaries, monkeypatch):
@pytest.mark.parametrize('fmt', ['exact', 'png'])
def test_remote_cache_hit_submits_fresh_callback_metadata(boundaries, monkeypatch, fmt):
data = item()
token = stage(data)
prompt = worker_prompt(token, 'fresh-context', data, local=False)
prompt['end']['inputs'].update(transfer_format=fmt, transfer_quality=73)
submit = AsyncMock()
monkeypatch.setattr(boundaries, 'submit_worker_result', submit)
cached = ([data['image']], [data['mask']], [0], [rt.item_content_identity(data)])
end(boundaries, prompt, cached)
result = end(boundaries, prompt, cached)
assert len(result) == 4
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}
payload = submit.await_args.args[4]
if fmt == 'png':
assert payload['image']['descriptor']['format'] == fmt
assert payload['image']['descriptor']['quality'] == 73
else:
assert 'encoding' not in payload['image']
assert token not in rt._WORKER_ITEMS
+728
View File
@@ -0,0 +1,728 @@
"""CPU-only image-list compilation and lifecycle tests; providers are mocked."""
import asyncio
import base64
import copy
import importlib
import json
import sys
import tempfile
import types
import unittest
from pathlib import Path
from unittest.mock import AsyncMock, patch
import pytest
import torch
ROOT = Path(__file__).resolve().parents[1]
PACKAGE = "_image_list_serverless_tests"
for suffix in ("", ".api", ".api.orchestration", ".image_list", ".serverless_bridge"):
module = types.ModuleType(PACKAGE + suffix)
module.__path__ = [str(ROOT.joinpath(*suffix.strip(".").split(".")))] if suffix else [str(ROOT)]
sys.modules[module.__name__] = module
adapter = importlib.import_module(PACKAGE + ".image_list.serverless")
runtime = adapter.runtime
manager_module = importlib.import_module(PACKAGE + ".api.orchestration.serverless_runtime")
handler = importlib.import_module(PACKAGE + ".serverless_bridge.handler")
codec, contract = adapter.codec, adapter.contract
def node(kind, **inputs):
return {"class_type": kind, "inputs": inputs}
def pool(wid="pool", jobs=1, enabled=True, **overrides):
return dict(id=wid, type="serverless", endpoint_id="ep", timeout_seconds=30,
jobs_per_run=jobs, enabled=enabled, media_profile=contract.PROFILE, **overrides)
def source():
return {
"upstream": node("LoadImage", image="never-read.png"),
"ns/start": node(runtime.START_CLASS, images=["upstream", 0]),
"model": node("CheckpointLoaderSimple", ckpt_name="fixed.safetensors"),
"process": node("Process", image=["ns/start", 0], mask=["ns/start", 1],
index=["ns/start", 2], text=["ns/start", 4], integer=["ns/start", 5],
floating=["ns/start", 6], boolean=["ns/start", 7], model=["model", 0]),
"ns/end": node(runtime.END_CLASS, context=["ns/start", 3], images=["process", 0],
masks=["process", 1], texts=["ns/start", 4], worker_result_token="drop-me"),
"downstream": node("SaveImage", images=["ns/end", 0]),
}
def region(prompt=None):
return {"start_id": "ns/start", "end_id": "ns/end",
"node_ids": set(prompt or source()) - {"upstream", "downstream"}}
def item(index=0, dtype=torch.float32):
return {"index": index, "image": torch.arange(24, dtype=dtype).reshape(1, 2, 3, 4) / 24,
"mask": torch.arange(6, dtype=dtype).reshape(1, 2, 3) / 6,
"values": {"value_1": " \n雪", "value_2": 123, "value_3": 0.25, "value_4": False}}
async def prepare(prompt=None, value=None, worker=None):
prompt = prompt or source()
return await adapter.prepare_serverless_item(prompt, region(prompt), value or item(), worker or pool(),
request_id="request:1", job_id="queue:1")
async def decoded_inputs(job):
return {artifact["id"]: codec.decode_payload(artifact["kind"], artifact["descriptor"],
await contract.resolve_artifact(artifact)) for artifact in job["media"]["artifacts"]}
def test_selection_is_exact_no_expansion_or_profile_fallback():
a, b = pool("a", jobs=16), pool("b", enabled=False)
direct = {"id": "direct", "type": "remote", "enabled": True}
config = {"workers": [a, b, direct]}
assert [w["id"] for w in adapter.resolve_serverless_workers(config, None)] == ["a"]
resolved = adapter.resolve_serverless_workers(config, ["b", "direct"])
assert [w["id"] for w in resolved] == ["b"]
assert resolved[0]["pool_id"] == "b"
resolved[0]["jobs_per_run"] = 2
assert b["jobs_per_run"] == 1
assert adapter.resolve_serverless_workers(config, []) == []
legacy = pool("legacy", enabled=False)
legacy.pop("media_profile")
config["workers"].append(legacy)
assert [w["id"] for w in adapter.resolve_serverless_workers(config, None)] == ["a"]
with pytest.raises(ValueError, match="comfy-media-v1"):
adapter.resolve_serverless_workers(config, ["legacy"])
with pytest.raises(ValueError, match="Unknown"):
adapter.resolve_serverless_workers(config, ["a:1"])
with pytest.raises(ValueError, match="Invalid"):
adapter.resolve_serverless_workers(config, [" a"])
with pytest.raises(ValueError, match="Duplicate"):
adapter.resolve_serverless_workers({"workers": [a, a]}, ["a"])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
def test_deterministic_graph_exact_codecs_and_source_immutable(dtype):
async def check():
prompt, value = source(), item(dtype=dtype)
original, saved_image = copy.deepcopy(prompt), value["image"].clone()
one = await prepare(prompt, value)
two = await prepare(dict(reversed(list(prompt.items()))), value)
assert one == two
assert prompt == original and torch.equal(value["image"], saved_image)
assert "upstream" not in one["workflow"] and "downstream" not in one["workflow"]
assert one["workflow"]["model"] == prompt["model"]
assert not any("distributed" in n["class_type"].lower() for n in one["workflow"].values())
proc = one["workflow"]["process"]["inputs"]
for name, expected in [("index", 0), ("integer", 123), ("floating", 0.25), ("boolean", False), ("text", " \n雪")]:
assert proc[name] == expected and type(proc[name]) is type(expected)
output = next(b for b in one["manifest"] if b["socket"] == "texts")
text_source = one["workflow"][output["node_id"]]["inputs"]["texts"][0]
assert one["workflow"][text_source] == node("PrimitiveString", value=" \n雪")
values = await decoded_inputs(one)
for ref in one["media"]["inputs"]:
result = values[ref["artifact_id"]]
assert result.dtype == dtype
assert torch.equal(result, value[ref["input_name"]])
assert len(values) == 2
handler._validate_request({key: one[key] for key in ("workflow", "images", "media")})
asyncio.run(check())
def test_direct_passthrough_exports_and_tensor_generic_values_deduplicate():
async def check():
prompt, value = source(), item()
prompt["process"]["inputs"]["reference"] = ["ns/start", 8]
prompt["process"]["inputs"]["reference_mask"] = ["ns/start", 9]
value["values"].update(value_5=value["image"], value_6=value["mask"])
prompt["ns/end"]["inputs"]["images"] = ["ns/start", 0]
job = await prepare(prompt, value)
assert "process" in job["workflow"] # mask still depends on processing
assert len(job["media"]["artifacts"]) == 2
refs = job["media"]["inputs"]
assert len(refs) == 5
image_binding = next(b for b in job["manifest"] if b["socket"] == "images")
assert any(r["node_id"] == image_binding["node_id"] and r["input_name"] == "images" for r in refs)
decoded = await decoded_inputs(job)
for ref in refs:
tensor = decoded[ref["artifact_id"]]
assert torch.equal(tensor, value["mask"] if "mask" in ref["input_name"] else value["image"])
asyncio.run(check())
@pytest.mark.parametrize("bad", [None, object(), torch.ones(2), torch.ones(2, 3)])
def test_connected_unsupported_generic_values_fail(bad):
value = item()
value["values"]["value_1"] = bad
with pytest.raises((ValueError, TypeError)):
asyncio.run(prepare(value=value))
@pytest.mark.parametrize("bad", [True, -1, "0", 0.0])
def test_start_slot_identity_not_coerced(bad):
prompt = source()
prompt["process"]["inputs"]["image"] = ["ns/start", bad]
with pytest.raises(ValueError, match="link"):
asyncio.run(prepare(prompt))
@pytest.mark.parametrize("bad", ["../bad", "ns//bad", "ns/./bad", "bad ", 42])
def test_node_identity_not_normalized_even_on_pruned_branch(bad):
prompt = source()
prompt[bad] = node("PrimitiveString", value="unused")
with pytest.raises(ValueError):
asyncio.run(prepare(prompt))
def test_placeholder_namespace_collision_and_empty_text():
async def check():
prompt, value = source(), item()
prompt["image-list:input:0"] = node("PrimitiveInt", value=5)
prompt["process"]["inputs"]["fixed"] = ["image-list:input:0", 0]
value["values"]["value_1"] = ""
job = await prepare(prompt, value)
assert job["workflow"]["image-list:input:0"] == prompt["image-list:input:0"]
assert job["workflow"]["process"]["inputs"]["image"] == ["image-list:input:0:1", 0]
assert any(n == node("PrimitiveString", value="") for n in job["workflow"].values())
asyncio.run(check())
def test_no_connected_none_and_payload_budget_before_submission():
value = item()
value["mask"] = None
with pytest.raises(ValueError, match="empty"):
asyncio.run(prepare(value=value))
client = importlib.import_module(PACKAGE + ".api.orchestration.runpod_client")
with patch.object(client, "MAX_REQUEST_BYTES", 10):
with pytest.raises(ValueError, match="Serialized"):
asyncio.run(prepare())
def test_fixed_file_inputs_are_prepared_on_private_copy_and_merge(tmp_path, monkeypatch):
input_dir = tmp_path / "input"
input_dir.mkdir()
(tmp_path / "output").mkdir()
(tmp_path / "temp").mkdir()
payload = b"bounded original file bytes"
(input_dir / "fixed.png").write_bytes(payload)
folders = types.SimpleNamespace(get_input_directory=lambda: str(input_dir),
get_output_directory=lambda: str(tmp_path / "output"), get_temp_directory=lambda: str(tmp_path / "temp"),
get_annotated_filepath=lambda name: str(input_dir / name))
monkeypatch.setitem(sys.modules, "folder_paths", folders)
prompt = source()
prompt["fixed"] = node("LoadImage", image="fixed.png")
prompt["process"]["inputs"]["fixed"] = ["fixed", 0]
original = copy.deepcopy(prompt)
job = asyncio.run(prepare(prompt))
assert prompt == original
assert len(job["images"]) == 1
assert base64.b64decode(job["images"][0]["image"]) == payload
assert job["workflow"]["fixed"]["inputs"]["image"] == job["images"][0]["name"]
assert len(job["media"]["artifacts"]) == 2
def test_merge_rejects_conflicting_artifact_ids_and_combined_caps():
async def check():
job = await prepare()
artifact = job["media"]["artifacts"][0]
assert adapter._merge_artifacts([artifact], [copy.deepcopy(artifact)]) == [artifact]
bad = copy.deepcopy(artifact)
bad["sha256"] = "0" * 64
with pytest.raises(ValueError, match="identity"):
adapter._merge_artifacts([artifact], [bad])
with patch.object(contract, "MAX_TOTAL_BYTES", 1):
with pytest.raises(ValueError, match="limits"):
adapter._merge_artifacts([artifact])
asyncio.run(check())
def test_handler_injects_exact_native_values_and_prunes_placeholders():
async def check():
value = item(dtype=torch.float64)
prompt = source()
prompt["ns/end"]["inputs"].update(images=["ns/start", 0], masks=["ns/start", 1])
job = await prepare(prompt, value)
workflow, media, legacy, targets = handler._validate_request(
{key: job[key] for key in ("workflow", "images", "media")})
buffers = {}
def write(payload):
name = f"owned-{len(buffers)}.bin"
buffers[name] = payload
return name
class Submitted(Exception):
pass
async def submit(native_graph):
assert all(n["class_type"] in {"MediaExportV1", "MediaImportV1", "PrimitiveString"}
for n in native_graph.values())
for binding in job["manifest"]:
link = native_graph[binding["node_id"]]["inputs"][binding["socket"]]
native = native_graph[link[0]]
if binding["kind"] == "STRING":
assert native == node("PrimitiveString", value=value["values"]["value_1"])
continue
assert native["class_type"] == "MediaImportV1"
metadata = json.loads(native["inputs"]["artifact_json"])
decoded = codec.decode_payload(metadata["kind"], metadata["descriptor"],
buffers[native["inputs"]["filename"]])
assert link[1] == {"IMAGE": 0, "MASK": 2}[binding["kind"]]
assert torch.equal(decoded, value["image" if binding["kind"] == "IMAGE" else "mask"])
# Stop at the mocked HTTP boundary; no endpoint/process/provider exists.
raise Submitted()
execution = types.SimpleNamespace(media=media, targets=targets, remaining=lambda: 30,
stage=types.SimpleNamespace(write=write), submit_and_wait=submit)
with pytest.raises(Submitted):
await handler._run(execution, workflow, legacy, None, None)
asyncio.run(check())
def test_fixed_typed_inputs_merge_with_start_refs_without_losing_consumers():
async def check():
prompt = source()
prompt["process"]["inputs"]["fixed_audio"] = 0
audio = {"waveform": torch.ones(1, 2, 3), "sample_rate": 48000}
descriptor, payload = codec.encode_payload("AUDIO", audio)
artifact = contract.artifact_from_payload("fixed-audio", "AUDIO", descriptor, payload)
fixed = AsyncMock(return_value=([], [artifact], [
{"node_id": "process", "input_name": "fixed_audio", "artifact_id": "fixed-audio"}]))
with patch.object(adapter, "prepare_media_inputs", fixed):
job = await prepare(prompt)
assert len(job["media"]["artifacts"]) == 3
assert len(job["media"]["inputs"]) == 3
assert any(ref["input_name"] == "fixed_audio" for ref in job["media"]["inputs"])
decoded = await decoded_inputs(job)
assert decoded["fixed-audio"]["sample_rate"] == 48000
assert torch.equal(decoded["fixed-audio"]["waveform"], audio["waveform"])
asyncio.run(check())
class LifecycleTests(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
self.temp = tempfile.TemporaryDirectory()
self.server = types.SimpleNamespace(distributed_pending_jobs={}, distributed_jobs_lock=asyncio.Lock())
self.manager = manager_module.ServerlessManager(self.server,
types.SimpleNamespace(get=lambda _wid: "fake-key"), Path(self.temp.name) / "jobs.json", AsyncMock())
self.manager.poll_seconds = 0
self.state = runtime.ImageListJob("context", "ns/start", "ns/end", source(),
[item(i) for i in range(40)], ["pool"], "", prompt_id="master-prompt")
self.gates, self.jobs, self.mutations = {}, {}, {}
self.started = asyncio.Event()
async def submit(workflow, images, timeout, *, media):
jid = "provider-" + str(len(self.jobs))
self.jobs[jid] = dict(workflow=workflow, images=images, media=media)
self.gates[jid] = asyncio.Event()
self.started.set()
return jid
async def status(jid):
await self.gates[jid].wait()
output = self.output(self.jobs[jid])
if jid in self.mutations:
self.mutations[jid](output)
return {"status": "COMPLETED", "output": output}
self.api = types.SimpleNamespace(submit=AsyncMock(side_effect=submit), status=AsyncMock(side_effect=status),
cancel=AsyncMock(return_value={"status": "CANCELLED"}))
self.patches = [patch.object(manager_module, "RunpodClient", return_value=self.api),
patch.object(adapter, "get_manager", return_value=self.manager)]
for p in self.patches:
p.start()
self.event = asyncio.Event()
self.controllers = []
self.controller = await self.controller_for()
async def controller_for(self, workers=None, state=None):
ctl = await adapter.create_serverless_slots(state or self.state, region(),
{"workers": workers or [pool()]}, self.event)
self.controllers.append(ctl)
return ctl
async def asyncTearDown(self):
for ctl in self.controllers:
await ctl.close()
await self.manager.cancel()
for p in reversed(self.patches):
p.stop()
self.temp.cleanup()
def output(self, job):
media = job["media"]
output = dict(version=1, profile=contract.PROFILE, request_id=media["request_id"],
bindings=[], artifacts=[], items=[])
index = int(media["request_id"].rsplit(":", 1)[1])
data = item(index)
values = {"images": data["image"], "masks": data["mask"], "texts": str(index)}
values.update(getattr(self, "override_values", {}))
for n, binding in enumerate(media["outputs"]):
descriptor, payload = codec.encode_payload(binding["kind"], values[binding["socket"]])
aid = f"result:{n}"
output["bindings"].append(dict(binding, count=1))
output["artifacts"].append(contract.artifact_from_payload(aid, binding["kind"], descriptor, payload))
output["items"].insert(0, dict(binding_id=binding["binding_id"], item_index=0, artifact_id=aid))
return {"media": output}
async def dispatch(self, index=0, ctl=None, slot=None):
ctl = ctl or self.controller
slot = slot or ctl.slots[index]
handle = await ctl.dispatch(slot, index)
return slot, handle
async def finish(self, handle):
# Submission has run by the time dispatch's shield completes, but yield
# explicitly rather than assuming provider status has already been entered.
while len(self.jobs) <= handle["index"]:
await asyncio.sleep(0)
jid = next(j for j, job in self.jobs.items()
if job["media"]["request_id"] == handle["execution_id"])
self.gates[jid].set()
await handle["task"]
async def test_no_pool_does_not_initialize_manager_or_credentials(self):
state = copy.copy(self.state)
state.worker_ids = []
with patch.object(adapter, "get_manager", side_effect=AssertionError("must stay lazy")):
ctl = await adapter.create_serverless_slots(state, region(), {"workers": [pool()]}, self.event)
assert ctl.slots == []
await ctl.close()
async def test_slots_are_bounded_not_jobs_per_run_and_cannot_collide(self):
assert len(self.controller.slots) == manager_module.PER_WORKER_JOB_LIMIT == 16
ctl = await self.controller_for([pool(jobs=16)])
assert ctl.slots == self.controller.slots
assert all(isinstance(slot, tuple) and len(slot) == 3 for slot in ctl.slots)
assert len(set(ctl.slots)) == 16 and "pool" not in ctl.slots
many = [pool(f"p{i}") for i in range(8)]
state = copy.copy(self.state)
state.worker_ids = [w["id"] for w in many]
ctl = await self.controller_for(many, state)
assert len(ctl.slots) == 64
assert all(sum(s[1] == w["id"] for s in ctl.slots) <= 16 for w in many)
async def test_reverse_completion_order_and_slot_reuse(self):
first, second = await self.dispatch(0), await self.dispatch(1)
inflight = dict([first, second])
assert self.controller.poll(inflight) is None
await self.finish(second[1])
await asyncio.sleep(0)
assert self.event.is_set()
slot, index, result = self.controller.poll(inflight)
assert (slot, index) == (second[0], 1) and result["text"] == "1"
assert torch.equal(result["image"], item(1)["image"])
assert second[1]["job_id"] not in self.server.distributed_pending_jobs
del inflight[slot]
third = await self.dispatch(2, slot=slot)
assert third[1]["execution_id"] != second[1]["execution_id"]
await self.finish(first[1])
assert self.controller.poll(inflight)[1] == 0
await self.finish(third[1])
assert self.controller.poll(dict([third]))[1] == 2
async def test_many_items_submit_once_each_through_bounded_reusable_slots(self):
ctl = await self.controller_for([pool(jobs=16)])
first_wave = asyncio.Event()
active = 0
peak = 0
original_submit = self.api.submit.side_effect
async def submit(*args, **kwargs):
jid = await original_submit(*args, **kwargs)
if len(self.jobs) == 16:
first_wave.set()
return jid
async def status(jid):
nonlocal active, peak
active += 1
peak = max(peak, active)
try:
await first_wave.wait()
await asyncio.sleep(0.001)
return {"status": "COMPLETED", "output": self.output(self.jobs[jid])}
finally:
active -= 1
self.api.submit.side_effect = submit
self.api.status.side_effect = status
async def next_result(inflight):
while True:
self.event.clear()
completed = ctl.poll(inflight)
if completed is not None:
return completed
await self.event.wait()
async def cleanup(*_):
pass
results = await asyncio.wait_for(runtime.dynamic_schedule(
len(self.state.items), ctl.slots, ctl.dispatch, next_result, cleanup), 10)
assert self.api.submit.await_count == len(self.state.items)
assert [row["text"] for row in results] == [str(i) for i in range(len(self.state.items))]
assert peak == 16
assert not self.server.distributed_pending_jobs
assert all(task.done() for task in self.manager.tasks.values())
async def test_waits_for_task_drain_not_provider_terminal_or_queue_message(self):
entered, release = asyncio.Event(), asyncio.Event()
original = self.manager._notify_capacity
async def delayed():
entered.set()
await release.wait()
await original()
with patch.object(self.manager, "_notify_capacity", delayed):
slot, handle = await self.dispatch()
self.gates["provider-0"].set()
await entered.wait()
assert self.server.distributed_pending_jobs[handle["job_id"]].qsize() == 1
assert self.controller.poll({slot: handle}) is None
with pytest.raises(RuntimeError, match="occupied"):
await self.controller.dispatch(slot, 1)
release.set()
await handle["task"]
assert self.controller.poll({slot: handle})[1] == 0
async def test_unknown_reservations_survive_and_prevent_provider_admission(self):
for i in range(16):
self.manager.ledger.update(f"unknown-{i}", worker_id="pool", status="UNKNOWN", job_id=None)
slot, handle = await self.dispatch()
await asyncio.sleep(0)
self.api.submit.assert_not_awaited()
assert self.controller.poll({slot: handle}) is None
await self.controller.close()
assert sum(r["status"] == "UNKNOWN" for r in self.manager.ledger.all()) == 16
assert not self.server.distributed_pending_jobs
assert self.manager.pool_is_active("pool")
async def test_local_pending_cap_checked_before_journal(self):
waiting = [asyncio.create_task(asyncio.Event().wait()) for _ in range(64)]
self.manager.tasks.update({f"unrelated-{i}": t for i, t in enumerate(waiting)})
self.manager.runs.update({f"unrelated-{i}": {"execution_id": "unrelated"} for i in range(64)})
try:
with pytest.raises(ValueError, match="limit 64"):
await self.dispatch()
assert not self.manager.ledger.all()
self.api.submit.assert_not_awaited()
assert not self.server.distributed_pending_jobs
finally:
for i, task in enumerate(waiting):
self.manager.tasks.pop(f"unrelated-{i}")
self.manager.runs.pop(f"unrelated-{i}")
task.cancel()
await asyncio.gather(*waiting, return_exceptions=True)
async def test_missing_corrupt_or_mismatched_results_are_atomic_errors(self):
mutations = [lambda o: o["media"]["items"].pop(),
lambda o: o["media"]["artifacts"][-1].update(sha256="0" * 64),
lambda o: o["media"].update(request_id="wrong-request"),
lambda o: o["media"]["bindings"][-1].update(participant_id="wrong-worker")]
for index, mutation in enumerate(mutations):
slot, handle = await self.dispatch(index)
self.mutations[f"provider-{index}"] = mutation
await self.finish(handle)
queue = self.server.distributed_pending_jobs[handle["job_id"]]
assert queue.qsize() == 1
with pytest.raises(RuntimeError, match="failed"):
self.controller.poll({slot: handle})
async def test_terminal_queue_identity_completeness_and_multiplicity(self):
for index, kind in enumerate(("missing", "duplicate", "worker", "mask", "text", "multiple", "empty")):
slot, handle = await self.dispatch(index)
await self.finish(handle)
queue = self.server.distributed_pending_jobs[handle["job_id"]]
message = queue.get_nowait()
if kind == "duplicate":
queue.put_nowait(message)
if kind == "worker":
message["worker_id"] = "unrelated"
if kind in ("mask", "text"):
message["media"].pop("masks" if kind == "mask" else "texts")
if kind == "multiple":
message["media"]["images"].append(item()["image"])
if kind == "empty":
message["media"]["texts"] = [None]
if kind != "missing":
queue.put_nowait(message)
with pytest.raises(RuntimeError):
self.controller.poll({slot: handle})
async def test_wrong_frame_count_channels_and_unpaired_mask_fail(self):
for index, kind in enumerate(("batch", "channels", "grid")):
slot, handle = await self.dispatch(index)
await self.finish(handle)
queue = self.server.distributed_pending_jobs[handle["job_id"]]
message = queue.get_nowait()
if kind == "batch":
message["media"]["images"] = [torch.ones(2, 2, 3, 4)]
elif kind == "channels":
message["media"]["images"] = [torch.ones(1, 2, 3, 1)]
else:
message["media"]["masks"] = [torch.ones(1, 5, 5)]
queue.put_nowait(message)
with pytest.raises(ValueError):
self.controller.poll({slot: handle})
async def test_close_cancels_only_owned_execution_and_exact_queues(self):
state = copy.copy(self.state)
state.context_id = "unrelated"
unrelated = await self.controller_for(state=state)
other_slot, other = await self.dispatch(0, unrelated)
slot, handle = await self.dispatch(0)
queue = self.server.distributed_pending_jobs[handle["job_id"]]
replacement = asyncio.Queue()
self.server.distributed_pending_jobs[handle["job_id"]] = replacement
await self.controller.close()
assert other["task"].done() is False
assert self.server.distributed_pending_jobs[handle["job_id"]] is replacement
assert other["job_id"] in self.server.distributed_pending_jobs
assert handle["task"].done()
assert len(self.api.cancel.await_args_list) == 1
assert queue.empty()
self.server.distributed_pending_jobs.pop(handle["job_id"])
# A late manager callback skips a retired destination instead of recreating it.
run = next(r for r in self.manager.runs.values() if r["execution_id"] == handle["execution_id"])
await self.manager._deliver(run, error="late")
assert handle["job_id"] not in self.server.distributed_pending_jobs
async def test_cancel_dispatch_after_start_loses_handle_but_not_ownership(self):
entered = asyncio.Event()
original = self.manager.start
async def delayed(*args):
ids = await original(*args)
entered.set()
await asyncio.Event().wait()
return ids
with patch.object(self.manager, "start", delayed):
task = asyncio.create_task(self.controller.dispatch(self.controller.slots[0], 0))
await entered.wait()
await self.started.wait()
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
await self.controller.close()
assert all(t.done() for t in self.manager.tasks.values())
assert not self.server.distributed_pending_jobs
self.api.cancel.assert_awaited_once_with("provider-0")
async def test_close_during_preparation_prevents_late_start(self):
entered = asyncio.Event()
prepared = await prepare()
async def resistant(*args, **kwargs):
entered.set()
try:
await asyncio.Event().wait()
except asyncio.CancelledError:
return prepared
with patch.object(adapter, "prepare_serverless_item", resistant):
task = asyncio.create_task(self.controller.dispatch(self.controller.slots[0], 0))
await entered.wait()
await self.controller.close()
with pytest.raises(RuntimeError, match="closed"):
await task
assert not self.manager.ledger.all()
assert not self.server.distributed_pending_jobs
self.api.submit.assert_not_awaited()
async def test_cancel_during_provider_submission_preserves_unknown(self):
entered = asyncio.Event()
async def ambiguous(*args, **kwargs):
entered.set()
await asyncio.Event().wait()
self.api.submit.side_effect = ambiguous
slot, handle = await self.dispatch()
await entered.wait()
await self.controller.close()
assert next(iter(self.manager.ledger.all()))["status"] == "UNKNOWN"
assert self.manager.pool_is_active("pool")
self.api.cancel.assert_not_awaited()
assert not self.server.distributed_pending_jobs
async def test_cancelled_wait_closes_owned_work(self):
entered = asyncio.Event()
async def job():
try:
await self.dispatch()
entered.set()
await asyncio.Event().wait()
finally:
await self.controller.close()
task = asyncio.create_task(job())
await entered.wait()
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
assert all(t.done() for t in self.manager.tasks.values())
assert not self.server.distributed_pending_jobs
async def test_repeated_cancellation_does_not_abandon_close(self):
slot, handle = await self.dispatch()
entered, release = asyncio.Event(), asyncio.Event()
original = self.manager.cancel
async def delayed(**kwargs):
entered.set()
await release.wait()
return await original(**kwargs)
with patch.object(self.manager, "cancel", delayed):
closing = asyncio.create_task(self.controller.close())
await entered.wait()
closing.cancel()
await asyncio.sleep(0)
closing.cancel()
await asyncio.sleep(0)
assert not closing.done()
assert not self.server.distributed_pending_jobs
release.set()
with pytest.raises(asyncio.CancelledError):
await closing
assert handle["task"].done()
assert self.controller._close_task.done()
async def test_serialization_failure_never_journals_or_submits(self):
with patch.object(adapter, "serialize_request", side_effect=ValueError("budget")):
with pytest.raises(ValueError, match="budget"):
await self.dispatch()
assert not self.manager.ledger.all()
assert not self.server.distributed_pending_jobs
self.api.submit.assert_not_awaited()
async def test_sparse_overrides_overlay_master_values_without_forwarding(self):
from test_image_list_serverless_values import assert_exact
self.override_values = {"value_2": 2**63 + 333, "value_10": False}
self.state.original_prompt["ns/end"]["inputs"].update(
value_2=["process", 2], value_10=["process", 3])
slot, handle = await self.dispatch()
await self.finish(handle)
result = self.controller.poll({slot: handle})[2]
assert set(result["values"]) == set(self.override_values)
for name, expected in self.override_values.items():
assert_exact(result["values"][name], expected)
state = copy.copy(self.state)
state.items = [self.state.items[0]]
columns = runtime._collect_output_lists([result], state=state, include_values=True)
assert columns[3][0] == {**state.items[0]["values"], **self.override_values}
assert self.state.items[0]["values"]["value_2"] == 123
async def test_override_missing_extra_wrong_nonfinite_multiple_results_fail(self):
self.override_values = {"value_2": 42}
self.state.original_prompt["ns/end"]["inputs"]["value_2"] = ["process", 2]
for index, bad in enumerate(("missing", "extra", "wrong", "nonfinite", "multiple", "none")):
slot, handle = await self.dispatch(index)
await self.finish(handle)
queue = self.server.distributed_pending_jobs[handle["job_id"]]
message = queue.get_nowait()
media = message["media"]
if bad == "missing":
media.pop("value_2")
elif bad == "extra":
media["value_3"] = [1]
else:
media["value_2"] = {"wrong": [{}], "nonfinite": [float("nan")],
"multiple": [1, 2], "none": [None]}[bad]
queue.put_nowait(message)
with pytest.raises((ValueError, TypeError, RuntimeError)):
self.controller.poll({slot: handle})
async def test_no_override_results_preserve_original_return_shape(self):
slot, handle = await self.dispatch()
await self.finish(handle)
result = self.controller.poll({slot: handle})[2]
assert set(result) == {"image", "mask", "text"}
assert not any(binding["kind"] == "VALUE" for binding in self.jobs["provider-0"]["media"]["outputs"])
async def test_duplicate_index_and_mismatched_handle_fail(self):
slot, handle = await self.dispatch()
with pytest.raises(ValueError, match="duplicate"):
await self.controller.dispatch(self.controller.slots[1], 0)
with pytest.raises(RuntimeError, match="Mismatched"):
self.controller.poll({slot: dict(handle)})
await self.finish(handle)
self.controller.poll({slot: handle})
with pytest.raises(RuntimeError, match="duplicate"):
self.controller.poll({slot: handle})
@@ -0,0 +1,359 @@
"""Offline image-list queue/scheduler integration; no provider or GPU calls."""
import asyncio
import copy
import json
import os
from pathlib import Path
import subprocess
import sys
import types
import unittest
from unittest.mock import AsyncMock, patch
import torch
import test_image_list_master_runtime as direct_fixture
from test_serverless_integration import PKG, load_modules
def pool(**updates):
return dict({"id": "pool", "name": "Cloud", "type": "serverless", "enabled": True,
"provider": "runpod", "endpoint_id": "endpoint", "jobs_per_run": 16,
"timeout_seconds": 30, "media_profile": "comfy-media-v1"}, **updates)
class MixedRuntimeTests(unittest.IsolatedAsyncioTestCase):
def setUp(self):
# Reuse only fixture setup, not inherited tests. The transport/slot doubles
# execute real runtime staging, assignment and result ordering code.
self.fixture = direct_fixture.MasterRuntimeTests()
self.fixture.setUp()
# This reused TestCase has no asyncio runner: invoke its patch cleanup
# directly, not doCleanups(), which would silently fail and leak aiohttp.
self.addCleanup(self.fixture.modules.stop)
self.f = f = self.fixture
self.rt = f.rt
self.configured_pool = pool()
f.config["workers"] = [self.configured_pool]
self.cloud_calls = []
self.completed = []
self.closed = 0
self.hold = False
self.missing = False
self.dispatch_error = False
self.poll_entered = asyncio.Event()
self.dispatch_entered = asyncio.Event()
self.hold_dispatch = False
self.slot_count = 2
fixture = self
class Controller:
def __init__(self, state, event):
self.state, self.event = state, event
self.slots = [("serverless", "pool", index) for index in range(fixture.slot_count)]
self.tasks = []
self.closed = False
async def dispatch(self, slot, index):
assert not self.closed
fixture.cloud_calls.append((slot, index))
fixture.dispatch_entered.set()
if fixture.dispatch_error:
raise RuntimeError("serverless dispatch failed")
async def execute():
if fixture.hold:
await asyncio.Event().wait()
# Make a later source item complete before the first local one.
await asyncio.sleep(0.001)
fixture.completed.append(index)
if fixture.missing:
raise RuntimeError("serverless completed without an image result")
item = state.items[index]
return {"image": item["image"], "mask": item["mask"],
"text": item["values"]["value_1"]}
state = self.state
task = asyncio.create_task(execute())
self.tasks.append(task)
task.add_done_callback(lambda _: self.event.set())
# Own the task before any await that might be cancelled.
if fixture.hold_dispatch:
await asyncio.Event().wait()
return {"index": index, "task": task}
def poll(self, inflight):
fixture.poll_entered.set()
for slot, handle in inflight.items():
if slot in self.slots and handle["task"].done():
return slot, handle["index"], handle["task"].result()
return None
async def close(self):
self.closed = True
for task in self.tasks:
if not task.done():
task.cancel()
await asyncio.gather(*self.tasks, return_exceptions=True)
fixture.closed += 1
async def create(state, region, config, event):
self.assertEqual(region["end_id"], "end")
self.assertEqual(config["workers"][-1]["id"], "pool")
self.controller = Controller(state, event)
return self.controller
self.create_controller = AsyncMock(side_effect=create)
sys.modules['_master_slot_tests.image_list.serverless'] = types.SimpleNamespace(
create_serverless_slots=self.create_controller)
self.direct_dispatch = sys.modules['_master_slot_tests.api.orchestration.dispatch'].dispatch_worker_prompt
def create(self, *, master=True, direct=True, count=9):
f = self.f
if direct:
f.workers.append({"id": "peer", "type": "local"})
f.config["workers"] = f.workers + [self.configured_pool]
f.local_delay = 0.04
context = f.create(participates=master, count=count) if (master or direct) else None
if context is None:
# Fixture.create assumes only direct peers; cloud-only uses a pool ID.
f.items = self.rt.normalize_image_list_inputs(torch.zeros((count, 2, 2, 3)),
mask=torch.zeros((count, 2, 2)), values={"value_1": [str(i) for i in range(count)]})
context = self.rt.create_master_job('start', 'end', direct_fixture.graph(), f.items, ['pool'], '')
f.context = context
f.state = self.rt.get_master_job(context)
else:
f.state.worker_ids.append('pool')
return context
async def test_mixed_roles_share_one_queue_and_preserve_pairing(self):
context = self.create()
original = copy.deepcopy(self.f.state.original_prompt)
result = await self.rt.run_image_list_job(context)
indices = self.f.local_calls + self.f.dispatched + [i for _, i in self.cloud_calls]
self.assertEqual(sorted(indices), list(range(9)))
self.assertTrue(self.f.local_calls and self.f.dispatched and self.cloud_calls)
self.assertNotEqual(self.completed[0], 0)
self.assertEqual(result[2], [str(i) for i in range(9)])
for index, (image, mask) in enumerate(zip(*result[:2])):
self.assertTrue(torch.equal(image, self.f.items[index]['image']))
self.assertTrue(torch.equal(mask, self.f.items[index]['mask']))
self.assertEqual(self.f.state.original_prompt, original)
self.f.preflight.assert_awaited_once_with(self.f.workers, require_idle=False)
for call in self.direct_dispatch.await_args_list:
self.assertEqual(call.kwargs['partial_execution_targets'], ['end'])
self.assertEqual(self.closed, 1)
self.assertFalse(self.rt._MASTER_JOBS)
async def test_cloud_only_one_job_per_item_ignores_images_per_run(self):
for jobs_per_run in (1, 16):
with self.subTest(jobs_per_run=jobs_per_run):
self.configured_pool['jobs_per_run'] = jobs_per_run
self.cloud_calls.clear()
result = await self.rt.run_image_list_job(self.create(master=False, direct=False, count=7))
self.assertEqual([i for _, i in self.cloud_calls], list(range(7)))
self.assertEqual(result[2], [str(i) for i in range(7)])
self.assertEqual(len({slot for slot, _ in self.cloud_calls}), self.slot_count)
self.f.preflight.assert_not_awaited()
self.f.network.get_client_session.assert_not_awaited()
self.assertFalse(self.f.slot_instances)
async def test_missing_cloud_result_fails_without_refill(self):
self.missing = True
with self.assertRaisesRegex(RuntimeError, 'without an image result'):
await self.rt.run_image_list_job(self.create(master=False, direct=False))
self.assertEqual(len(self.cloud_calls), self.slot_count)
self.assertEqual(self.closed, 1)
self.assertFalse(self.rt._MASTER_JOBS)
async def test_controller_closed_when_direct_preflight_fails(self):
self.f.preflight.side_effect = RuntimeError('direct preflight failed')
with self.assertRaisesRegex(RuntimeError, 'direct preflight failed'):
await self.rt.run_image_list_job(self.create())
self.assertEqual(self.closed, 1)
self.assertFalse(self.cloud_calls)
async def test_controller_closed_when_local_setup_fails(self):
slot_module = sys.modules['_master_slot_tests.image_list.local_executor']
with patch.object(slot_module, 'LocalExecutionSlot', side_effect=RuntimeError('local setup failed')):
with self.assertRaisesRegex(RuntimeError, 'local setup failed'):
await self.rt.run_image_list_job(self.create())
self.assertEqual(self.closed, 1)
self.assertFalse(self.cloud_calls)
async def test_controller_closed_on_dispatch_failure(self):
self.dispatch_error = True
with self.assertRaisesRegex(RuntimeError, 'serverless dispatch failed'):
await self.rt.run_image_list_job(self.create(master=False, direct=False))
self.assertEqual(len(self.cloud_calls), 1)
self.assertEqual(self.closed, 1)
async def test_cancelled_dispatch_without_returned_handle_is_owned(self):
self.hold = self.hold_dispatch = True
with patch.object(self.rt, 'SUBMISSION_CANCEL_GRACE_SECONDS', 0.01):
task = asyncio.create_task(self.rt.run_image_list_job(self.create(master=False, direct=False)))
await asyncio.wait_for(self.dispatch_entered.wait(), 1)
task.cancel()
with self.assertRaises(self.rt.ImageListCancelledError):
await asyncio.wait_for(task, 1)
self.assertEqual(len(self.cloud_calls), 1)
self.assertEqual(self.closed, 1)
self.assertTrue(all(t.done() for t in self.controller.tasks))
async def interrupt_probe(self, models):
self.hold = True
task = asyncio.create_task(self.rt.run_image_list_job(self.create(master=False, direct=False)))
try:
await asyncio.wait_for(self.poll_entered.wait(), 1)
assigned = list(self.cloud_calls)
models.interrupt_current_processing(True)
with self.assertRaises(self.rt.ImageListCancelledError):
await asyncio.wait_for(task, 1)
self.assertEqual(self.cloud_calls, assigned)
self.assertEqual(self.closed, 1)
self.assertTrue(all(t.done() for t in self.controller.tasks))
self.assertFalse(self.rt._MASTER_JOBS)
finally:
models.interrupt_current_processing(False)
if not task.done():
task.cancel()
await asyncio.gather(task, return_exceptions=True)
class QueueSelectionTests(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
self.q, modules, self.ps = load_modules()
self.modules_patch = patch.dict(sys.modules, modules)
self.modules_patch.start()
self.addCleanup(self.modules_patch.stop)
self.pool = pool()
self.peer = {'id': 'peer', 'name': 'Peer', 'type': 'local', 'enabled': True}
self.config = {'workers': [self.pool], 'settings': {'master_delegate_only': True}}
self.q.load_config = lambda: self.config
self.q.select_active_workers = AsyncMock(side_effect=lambda workers, ws, delegate, **kw: (workers, delegate))
self.preflight = AsyncMock(side_effect=lambda workers, **kw: {w['id']: True for w in workers})
sys.modules[PKG + '.api.orchestration.dispatch'].preflight_image_list_workers = self.preflight
self.q.queue_prompt_payload = AsyncMock(return_value={'prompt_id': 'outer', 'number': 1})
self.prompt = direct_fixture.graph()
async def queue(self, **kwargs):
return await self.q.orchestrate_distributed_execution(self.prompt, None, 'client', **kwargs)
async def test_delegate_pool_selected_once_no_direct_probe_or_source_mutation(self):
original = copy.deepcopy(self.prompt)
result = await self.queue()
selected = self.q.select_active_workers.await_args.args[0]
self.assertEqual([w['id'] for w in selected], ['pool'])
prepared = self.q.queue_prompt_payload.await_args.args[0]
inputs = prepared['start']['inputs']
self.assertEqual(json.loads(inputs['image_list_enabled_worker_ids']), ['pool'])
self.assertIs(inputs['image_list_master_participates'], False)
self.assertEqual(self.q.queue_prompt_payload.await_args.kwargs['partial_execution_targets'], ['end'])
self.assertEqual(result[2], 1)
self.assertEqual(self.prompt, original)
self.preflight.assert_not_awaited()
self.q.dispatch_worker_prompt.assert_not_awaited()
self.assertEqual(self.ps.distributed_pending_jobs, {})
async def test_mixed_selection_preflights_only_direct_peers(self):
self.config['workers'].insert(0, self.peer)
result = await self.queue(delegate_master=False)
self.assertEqual([w['id'] for w in self.preflight.await_args.args[0]], ['peer'])
self.assertFalse(self.preflight.await_args.kwargs['require_idle'])
inputs = self.q.queue_prompt_payload.await_args.args[0]['start']['inputs']
self.assertEqual(json.loads(inputs['image_list_enabled_worker_ids']), ['peer', 'pool'])
self.assertEqual(result[2], 2)
async def test_legacy_pool_rejected_before_probe_or_queue(self):
self.pool['media_profile'] = 'legacy-image-v1'
with self.assertRaisesRegex((ValueError, RuntimeError), 'comfy-media-v1'):
await self.queue()
self.q.select_active_workers.assert_not_awaited()
self.q.queue_prompt_payload.assert_not_awaited()
async def test_unknown_selection_rejected_before_queue(self):
with self.assertRaisesRegex((ValueError, RuntimeError), 'Unknown|unknown'):
await self.queue(enabled_worker_ids=['missing'])
self.q.select_active_workers.assert_not_awaited()
self.q.queue_prompt_payload.assert_not_awaited()
async def test_disabled_pool_can_be_explicitly_selected(self):
self.pool['enabled'] = False
await self.queue(enabled_worker_ids=['pool'])
self.assertEqual([w['id'] for w in self.q.select_active_workers.await_args.args[0]], ['pool'])
async def test_explicit_empty_selection_does_not_import_cloud_adapter(self):
sys.modules[PKG + '.image_list.serverless'] = None
result = await self.queue(enabled_worker_ids=[], delegate_master=False)
self.assertEqual(result[2], 0)
self.preflight.assert_not_awaited()
async def test_cloud_auth_or_capacity_failure_never_queues_master(self):
for error in ('authentication failed', 'capacity exhausted'):
with self.subTest(error=error):
self.q.select_active_workers.side_effect = ValueError(error)
with self.assertRaisesRegex(ValueError, error):
await self.queue()
self.q.queue_prompt_payload.assert_not_awaited()
async def test_real_selection_uses_manager_health_not_direct_http(self):
from test_dispatch_selection import _load_dispatch_module
dispatch = _load_dispatch_module()
health = AsyncMock(return_value={'ok': True, 'queueRemaining': 3})
sys.modules['dist_dispatch_testpkg.api.orchestration.serverless_runtime'] = types.SimpleNamespace(
get_manager=lambda: types.SimpleNamespace(health=health))
self.q.select_active_workers = dispatch.select_active_workers
with patch.object(dispatch, 'get_client_session', AsyncMock(side_effect=AssertionError('no HTTP'))), \
patch.object(dispatch, 'probe_worker', AsyncMock(side_effect=AssertionError('no direct probe'))):
await self.queue()
health.assert_awaited_once()
self.assertEqual(health.await_args.args[0]['id'], 'pool')
self.preflight.assert_not_awaited()
self.q.queue_prompt_payload.reset_mock()
health.side_effect = ValueError('provider authentication failed')
with self.assertRaisesRegex(ValueError, 'authentication failed'):
await self.queue()
self.q.queue_prompt_payload.assert_not_awaited()
async def test_unreachable_pool_never_enables_master_fallback(self):
self.q.select_active_workers.side_effect = None
self.q.select_active_workers.return_value = ([], False)
with self.assertRaisesRegex(RuntimeError, 'orchestrator|Orchestrator'):
await self.queue()
self.q.queue_prompt_payload.assert_not_awaited()
class RealInterruptTests(unittest.TestCase):
@unittest.skipUnless(os.environ.get('COMFYUI_SOURCE_ROOT'), 'set COMFYUI_SOURCE_ROOT for CPU interrupt probe')
def test_processing_interrupted_signal_during_cloud_result_wait(self):
result = subprocess.run([sys.executable, str(Path(__file__).resolve()), '--interrupt-probe',
os.environ['COMFYUI_SOURCE_ROOT']], text=True, capture_output=True, timeout=45)
self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
self.assertIn('SERVERLESS_INTERRUPT_OK', result.stdout)
def real_interrupt_probe(comfy_root):
sys.path.insert(0, str(Path(comfy_root).resolve()))
import comfy.cli_args
comfy.cli_args.args.cpu = True
comfy.cli_args.args.disable_assets = True
import comfy.model_management as models
real_comfy = sys.modules['comfy']
case = MixedRuntimeTests()
case.setUp()
try:
# Keep actual thread-safe ComfyUI interruption state, not a task.cancel
# surrogate or a fake processing_interrupted callback.
sys.modules['comfy'] = real_comfy
sys.modules['comfy.model_management'] = models
asyncio.run(case.interrupt_probe(models))
finally:
case.f.modules.stop()
print('SERVERLESS_INTERRUPT_OK no refill, owned close, actual ComfyUI flag')
if __name__ == '__main__':
if len(sys.argv) > 1 and sys.argv[1] == '--interrupt-probe':
real_interrupt_probe(sys.argv[2])
else:
unittest.main()
+165
View File
@@ -0,0 +1,165 @@
"""Sparse End overrides through the standalone serverless bridge, CPU only."""
import asyncio
import copy
import math
import os
from pathlib import Path
import struct
import subprocess
import sys
import pytest
import torch
from unittest.mock import patch
from test_image_list_serverless import adapter, codec, contract, handler, item, prepare, source
@pytest.mark.skipif(not os.environ.get('COMFYUI_SOURCE_ROOT'), reason='set COMFYUI_SOURCE_ROOT for real CPU executor')
def test_real_cpu_standalone_bridge_calculated_overrides(tmp_path):
env = dict(os.environ, CUDA_VISIBLE_DEVICES='', HIP_VISIBLE_DEVICES='', ROCR_VISIBLE_DEVICES='',
PYTHONDONTWRITEBYTECODE='1')
env.pop('PYTHONPATH', None)
helper = Path(__file__).parent / 'helpers' / 'image_list_serverless_values_check.py'
result = subprocess.run([sys.executable, str(helper), os.environ['COMFYUI_SOURCE_ROOT'], str(tmp_path)],
env=env, text=True, capture_output=True, timeout=90)
assert result.returncode == 0, result.stdout + result.stderr
assert 'IMAGE_LIST_SERVERLESS_VALUES_CPU_OK' in result.stdout
def assert_exact(actual, expected):
assert type(actual) is type(expected)
if isinstance(expected, torch.Tensor):
assert actual.dtype == expected.dtype and actual.shape == expected.shape
assert torch.equal(actual.contiguous().view(torch.uint8), expected.contiguous().view(torch.uint8))
elif isinstance(expected, float):
assert struct.pack('!d', actual) == struct.pack('!d', expected)
else:
assert actual == expected
@pytest.mark.parametrize('value', [False, True, 2**63 + 123, -(2**100), -0.0, 1.0000000000000002,
'', ' \r\n雪 ', torch.ones(2, 2, 3, 4, dtype=torch.float64),
torch.ones(2, 2, 3, dtype=torch.bfloat16)])
def test_value_codec_exact_types(value):
descriptor, payload = codec.encode_payload('VALUE', value)
assert_exact(codec.decode_payload('VALUE', descriptor, payload), value)
@pytest.mark.parametrize('value', [None, {}, [], (1,), object(), float('inf'), float('nan'),
torch.ones(2), torch.ones(1, 2, 3, dtype=torch.int64),
torch.full((1, 2, 3), float('nan'))])
def test_value_codec_rejects_unsupported_or_nonfinite(value):
with pytest.raises((ValueError, TypeError)):
codec.encode_payload('VALUE', value)
def test_sparse_exports_never_create_collector_or_forward_unchanged_values():
async def check():
prompt, value = source(), item()
prompt['ns/end']['inputs'].update(value_2=['ns/start', 5], value_10=['process', 2])
original = copy.deepcopy(prompt)
job = await prepare(prompt, value)
assert prompt == original
bindings = {binding['socket']: binding for binding in job['manifest']}
assert set(bindings) == {'images', 'masks', 'texts', 'value_2', 'value_10'}
assert bindings['value_2']['kind'] == bindings['value_10']['kind'] == 'VALUE'
assert all('Collector' not in node['class_type'] for node in job['workflow'].values())
assert all(node['class_type'] not in (adapter.runtime.START_CLASS, adapter.runtime.END_CLASS)
for node in job['workflow'].values())
handler._validate_request({key: job[key] for key in ('workflow', 'images', 'media')})
asyncio.run(check())
@pytest.mark.parametrize('slot', ['value_0', 'value_11', 'value_01', 'values.value_1'])
def test_unknown_end_override_slot_rejected(slot):
prompt = source()
prompt['ns/end']['inputs'][slot] = ['ns/start', 4]
with pytest.raises(ValueError, match='slot'):
asyncio.run(prepare(prompt))
def test_value_codec_budgets_and_malformed_payload():
with pytest.raises(ValueError, match='byte limit'):
codec.encode_payload('VALUE', '雪', max_bytes=2)
descriptor, payload = codec.encode_payload('VALUE', True)
with pytest.raises(ValueError):
codec.decode_payload('VALUE', dict(descriptor, type='INT'), payload)
descriptor, payload = codec.encode_payload('VALUE', -0.0)
with pytest.raises(ValueError):
codec.decode_payload('VALUE', descriptor, struct.pack('<d', math.inf))
with pytest.raises(ValueError):
codec.decode_payload('VALUE', descriptor, payload + b'0')
with pytest.raises(ValueError, match='byte limit'):
codec.decode_payload('VALUE', descriptor, payload, max_bytes=7)
@pytest.mark.parametrize('fmt', ['exact', 'legacy_png', 'png', 'ffv1', 'jpeg', 'webp', 'h264', 'h265', 'av1'])
def test_end_image_policy_never_compresses_override_tensors(fmt):
prompt = source()
prompt['ns/end']['inputs'].update(value_10=['ns/start', 0], transfer_format=fmt, transfer_quality=73)
job = asyncio.run(prepare(prompt))
for binding in job['manifest']:
controls = job['workflow'][binding['node_id']]['inputs']
if binding['socket'] == 'images' and fmt != 'exact':
assert controls['transfer_format'] == fmt and controls['transfer_quality'] == 73
else:
assert 'transfer_format' not in controls and 'transfer_quality' not in controls
default = asyncio.run(prepare())
assert all('transfer_format' not in n['inputs'] and 'transfer_quality' not in n['inputs']
for n in default['workflow'].values())
def value_response():
binding = dict(binding_id='binding', collector_id='end', job_id='job', participant_id='pool',
node_id='export', socket='value_10', kind='VALUE')
descriptor, payload = codec.encode_payload('VALUE', 2**63 + 7)
artifact = contract.artifact_from_payload('artifact', 'VALUE', descriptor, payload)
return binding, dict(version=1, profile=contract.PROFILE, request_id='request',
bindings=[dict(binding, count=1)], artifacts=[artifact],
items=[dict(binding_id='binding', item_index=0, artifact_id='artifact')])
@pytest.mark.parametrize('bad', ['missing', 'extra', 'wrong_kind', 'multiple', 'extra_artifact', 'nonfinite'])
def test_override_response_rejected_atomically(bad):
binding, response = value_response()
if bad == 'missing':
response['bindings'] = []
elif bad == 'extra':
response['bindings'].append(dict(binding, binding_id='other', socket='value_9', count=1))
elif bad == 'wrong_kind':
response['artifacts'][0]['kind'] = 'STRING'
elif bad == 'multiple':
response['bindings'][0]['count'] = 2
elif bad == 'extra_artifact':
response['artifacts'].append(dict(response['artifacts'][0], id='extra'))
else:
descriptor, _ = codec.encode_payload('VALUE', 1.0)
response['artifacts'][0] = contract.artifact_from_payload(
'artifact', 'VALUE', descriptor, struct.pack('<d', float('nan')))
with pytest.raises(ValueError):
asyncio.run(contract.decode_response(response, request_id='request', bindings=[binding]))
def test_override_request_and_result_budgets():
prompt = source()
prompt['ns/end']['inputs']['value_10'] = '雪' * 100
with patch.object(contract, 'MAX_ARTIFACT_BYTES', 128):
with pytest.raises(ValueError, match='byte limit'):
asyncio.run(prepare(prompt))
binding, response = value_response()
with pytest.raises(ValueError, match='total byte limit'):
asyncio.run(contract.decode_response(response, request_id='request', bindings=[binding], max_total_bytes=1))
def test_collector_compiler_keeps_only_its_public_sockets():
from importlib import import_module
compiler = import_module(adapter.__package__.rsplit('.', 1)[0] + '.api.orchestration.serverless_media_prompt')
assert compiler.COLLECTOR_SOCKETS == ('images', 'audio')
graph = {'source': {'class_type': 'Source', 'inputs': {}},
'collector': {'class_type': 'DistributedCollector', 'inputs': {
'images': ['source', 0], 'value_10': ['source', 1], 'masks': ['source', 2],
'multi_job_id': 'job', 'worker_id': 'pool'}}}
_, bindings = compiler.compile_media_prompt(graph, 'request')
assert [binding['socket'] for binding in bindings] == ['images']
+70
View File
@@ -0,0 +1,70 @@
"""End policy reaches cloud exports and initial worker compatibility preflight."""
import asyncio
import copy
import unittest
import pytest
import torch
import test_image_list_serverless as cloud
import test_image_list_master_selection as selection
FORMATS = ('exact', 'legacy_png', 'png', 'ffv1', 'jpeg', 'webp', 'h264', 'h265', 'av1')
@pytest.mark.parametrize('fmt', FORMATS)
def test_cloud_returns_use_selected_image_format_but_inputs_remain_exact(fmt):
async def check():
prompt = cloud.source()
prompt['ns/end']['inputs'].update(transfer_format=fmt, transfer_quality=73)
original = copy.deepcopy(prompt)
value = cloud.item(dtype=torch.float64)
job = await cloud.prepare(prompt, value)
assert prompt == original
for binding in job['manifest']:
controls = job['workflow'][binding['node_id']]['inputs']
if binding['kind'] == 'IMAGE' and fmt != 'exact':
assert controls['transfer_format'] == fmt
assert controls['transfer_quality'] == 73
else:
assert 'transfer_format' not in controls and 'transfer_quality' not in controls
for artifact in job['media']['artifacts']:
assert artifact['descriptor']['encoding'] == 'tensor-le-v1'
inputs = await cloud.decoded_inputs(job)
for ref in job['media']['inputs']:
assert torch.equal(inputs[ref['artifact_id']], value[ref['input_name']])
cloud.handler._validate_request({k: job[k] for k in ('workflow', 'images', 'media')})
asyncio.run(check())
@pytest.mark.parametrize('fmt,quality', [('JPEG',90),('jpeg',True),('jpeg',101),('jpeg','90')])
def test_cloud_rejects_invalid_policy(fmt, quality):
prompt = cloud.source()
prompt['ns/end']['inputs'].update(transfer_format=fmt, transfer_quality=quality)
with pytest.raises(ValueError):
asyncio.run(cloud.prepare(prompt))
class TransferPreflightTests(unittest.IsolatedAsyncioTestCase):
asyncSetUp = selection.MasterSelectionTests.asyncSetUp
queue = selection.MasterSelectionTests.queue
async def test_queue_requests_image_codec_capability(self):
self.prompt['3']['inputs'].update(transfer_format='jpeg', transfer_quality=73)
await self.queue()
self.assertEqual(self.preflight.call_args.kwargs['required_transfer_formats'], ['jpeg'])
async def test_invalid_policy_fails_before_worker_probe_or_queue(self):
self.prompt['3']['inputs'].update(transfer_format='bad-codec')
with self.assertRaises(ValueError):
await self.queue()
self.q.select_active_workers.assert_not_awaited()
self.preflight.assert_not_awaited()
self.q.queue_prompt_payload.assert_not_awaited()
async def test_missing_capability_cannot_silently_fall_back_to_master(self):
self.prompt['3']['inputs'].update(transfer_format='jpeg')
self.preflight.side_effect = RuntimeError('upgrade worker for image return encoding')
with self.assertRaisesRegex(RuntimeError, 'upgrade'):
await self.queue()
self.q.queue_prompt_payload.assert_not_awaited()
+312
View File
@@ -0,0 +1,312 @@
"""CPU-only direct-worker IMAGE return codecs and capability negotiation."""
import asyncio
import base64
import importlib.util
import json
from pathlib import Path
import sys
import types
from unittest.mock import AsyncMock, Mock, patch
import pytest
import torch
from image_list import runtime as rt
from serverless_bridge import image_codec
from serverless_bridge.image_transfer import FORMATS
ROOT = Path(__file__).resolve().parents[1]
def result(channels=4, dtype=torch.float32):
image = torch.linspace(-0.1, 1.1, 17 * 19 * channels, dtype=dtype).reshape(1, 17, 19, channels)
return {"image": image, "mask": torch.full((1, 17, 19), 0.123456789, dtype=torch.float64),
"text": "exact text: \u00e9\n"}
def legacy_value(value, kind):
if isinstance(value, str):
return {"type": kind, "value": value}
raw = value.contiguous().view(torch.uint8).numpy().tobytes()
return {"type": kind, "dtype": str(value.dtype).removeprefix("torch."),
"shape": list(value.shape), "data": base64.b64encode(raw).decode("ascii")}
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
def test_exact_result_and_inputs_are_byte_for_byte_legacy(dtype):
source = result(dtype=dtype)
expected = {name: legacy_value(source[name], kind)
for name, kind in (("image", "IMAGE"), ("mask", "MASK"), ("text", "STRING"))}
body = rt.encode_request_payload(expected)
with patch.object(image_codec, "encode_image", side_effect=AssertionError("exact used codec")):
assert rt.encode_request_payload(rt.encode_result(source)) == body
assert rt.encode_request_payload(rt.encode_result(source, transfer_format="exact", transfer_quality=90)) == body
decoded = rt.decode_result(expected)
assert torch.equal(decoded["image"].view(torch.uint8), source["image"].view(torch.uint8))
item = {"image": source["image"], "mask": source["mask"], "index": 0, "values": {"value_1": source["image"]}}
encoded = rt.encode_item(item)
assert encoded["image"] == expected["image"]
assert encoded["mask"] == expected["mask"]
assert encoded["values"]["value_1"] == legacy_value(source["image"], "TENSOR")
@pytest.mark.parametrize("fmt", FORMATS[1:])
@pytest.mark.parametrize("channels", [3, 4])
def test_all_compressed_formats_roundtrip_only_main_image(fmt, channels):
source = result(channels)
encoded = rt.encode_result(source, transfer_format=fmt, transfer_quality=90)
image = encoded["image"]
assert set(image) == {"type", "encoding", "descriptor", "data"}
assert image["type"] == "IMAGE"
assert image["encoding"] == "image-transfer-v1"
assert image["descriptor"]["format"] == fmt
assert image["descriptor"]["shape"] == [1, 17, 19, channels]
assert image["descriptor"]["quality"] == 90
assert base64.b64decode(image["data"], validate=True)
decoded = rt.decode_result(json.loads(rt.encode_request_payload(encoded)))
assert decoded["image"].shape == source["image"].shape
assert decoded["image"].dtype == source["image"].dtype
assert decoded["image"].device.type == "cpu"
assert torch.isfinite(decoded["image"]).all()
pixels = source["image"][..., :3].float().mul(255).clamp(0, 255).to(torch.uint8).float().div(255)
if fmt in ("legacy_png", "png", "ffv1"):
assert torch.equal(decoded["image"][..., :3], pixels)
else:
assert (decoded["image"][..., :3] - pixels).abs().mean() < 0.08
if channels == 4:
assert torch.equal(decoded["image"][..., 3], source["image"][..., 3])
assert encoded["mask"] == legacy_value(source["mask"], "MASK")
assert torch.equal(decoded["mask"], source["mask"])
assert decoded["text"] == source["text"]
@pytest.mark.parametrize("fmt,quality", [("nope", 90), (["node", 0], 90), ("png", True), ("png", 0), ("exact", 101)])
def test_invalid_options_fail_closed(fmt, quality):
with pytest.raises(ValueError):
rt.encode_result(result(), transfer_format=fmt, transfer_quality=quality)
def test_compressed_inputs_masks_and_generic_tensors_are_rejected():
source = result()
encoded = rt.encode_result(source, transfer_format="png")
for kind in ("IMAGE", "MASK", "TENSOR", "STRING"):
with pytest.raises(ValueError):
rt.decode_value({**encoded["image"], "type": kind}, kind)
item = rt.encode_item({"image": source["image"], "index": 0})
with pytest.raises(ValueError):
rt.decode_item({**item, "image": encoded["image"]})
for key, kind in (("mask", "MASK"), ("text", "STRING")):
with pytest.raises(ValueError):
rt.decode_result({**encoded, key: {**encoded["image"], "type": kind}})
@pytest.mark.parametrize("shape", [[2, 17, 19, 3], [1, 17, 19, 1], [17, 19, 3], [True, 17, 19, 3], [1, 0, 19, 3]])
def test_invalid_image_shapes_rejected_before_codec_decode(shape):
encoded = rt.encode_result(result(), transfer_format="png")
encoded["image"]["descriptor"]["shape"] = shape
with patch.object(image_codec, "decode_image", side_effect=AssertionError("native codec reached")):
with pytest.raises(ValueError):
rt.decode_result(encoded)
@pytest.mark.parametrize("change", [
{"encoding": "image-transfer-v2"}, {"type": "MASK"}, {"extra": 1}, {"data": "%%%"},
{"descriptor": None}, {"descriptor": {"encoding": "tensor-le-v1", "shape": [1, 17, 19, 4]}},
])
def test_malformed_compressed_envelopes_rejected(change):
encoded = rt.encode_result(result(), transfer_format="png")
encoded["image"].update(change)
with pytest.raises(ValueError):
rt.decode_result(encoded)
def test_single_frame_and_image_list_not_bridge_budgets(monkeypatch):
with pytest.raises(ValueError):
rt.encode_result({"image": torch.zeros(2, 17, 19, 3)}, transfer_format="png")
with patch.object(image_codec, "encode_image", wraps=image_codec.encode_image) as encode:
encoded = rt.encode_result(result(), transfer_format="png")
assert encode.call_args.kwargs["max_bytes"] == rt.MAX_PAYLOAD_BYTES
with patch.object(image_codec, "decode_image", wraps=image_codec.decode_image) as decode:
rt.decode_result(encoded)
assert decode.call_args.kwargs["max_bytes"] == rt.MAX_PAYLOAD_BYTES
monkeypatch.setattr(rt, "MAX_PAYLOAD_BYTES", 64)
with pytest.raises(ValueError, match="limit"):
rt.encode_result({"image": torch.empty(1, 4, 4, 3, device="meta")}, transfer_format="png")
with patch.object(image_codec, "decode_image", side_effect=AssertionError("oversize reached codec")):
with pytest.raises(ValueError, match="limit"):
rt.decode_result(encoded)
def test_compressed_request_combined_budget(monkeypatch):
encoded = rt.encode_result(result(), transfer_format="png")
monkeypatch.setattr(rt, "MAX_REQUEST_BYTES", len(rt.encode_request_payload(encoded)) - 1)
with pytest.raises(ValueError, match="request"):
rt.encode_result(result(), transfer_format="png")
with patch.object(image_codec, "decode_image", side_effect=AssertionError("oversize reached codec")):
with pytest.raises(ValueError, match="request"):
rt.decode_result(encoded)
@pytest.mark.parametrize("fmt", FORMATS)
def test_return_formats_preserve_exact_sparse_override_values(fmt):
from test_image_list_values import assert_values, values
source = result()
source["values"] = values()
source["values"]["value_10"] = source["values"].pop("value_6")
encoded = rt.encode_result(source, transfer_format=fmt, transfer_quality=73)
assert set(encoded["values"]) == set(source["values"])
for name, value in source["values"].items():
assert encoded["values"][name] == rt.encode_value(value, rt.value_kind(value, name))
decoded = rt.decode_result(json.loads(rt.encode_request_payload(encoded)))
assert_values(decoded["values"], source["values"])
assert torch.equal(decoded["mask"], source["mask"])
assert decoded["text"] == source["text"]
def test_compressed_aggregate_budget_includes_override_values(monkeypatch):
source = result()
encoded_without_values = rt.encode_result(source, transfer_format="png")
source["values"] = {"value_10": "x" * 500}
encoded = rt.encode_result(source, transfer_format="png")
monkeypatch.setattr(rt, "MAX_REQUEST_BYTES", len(rt.encode_request_payload(encoded_without_values)) + 100)
with pytest.raises(ValueError, match="request"):
rt.encode_result(source, transfer_format="png")
with patch.object(image_codec, "decode_image", side_effect=AssertionError("oversize reached codec")):
with pytest.raises(ValueError, match="request"):
rt.decode_result(encoded)
@pytest.fixture
def dispatch_module(monkeypatch):
name = "_image_list_transfer_tests"
for suffix in ("", ".api", ".api.orchestration", ".utils"):
module = types.ModuleType(name + suffix)
module.__path__ = [str(ROOT / suffix.lstrip(".").replace(".", "/"))]
monkeypatch.setitem(sys.modules, name + suffix, module)
monkeypatch.setitem(sys.modules, name + ".utils.logging", types.SimpleNamespace(log=lambda *a: None, debug_log=lambda *a: None))
monkeypatch.setitem(sys.modules, name + ".utils.network", types.SimpleNamespace(
build_worker_url=lambda worker, path="": "http://" + worker["id"] + path,
get_client_session=AsyncMock(), probe_worker=AsyncMock()))
spec = importlib.util.spec_from_file_location(name + ".api.orchestration.dispatch", ROOT / "api/orchestration/dispatch.py")
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
def capability():
return {"supported": True, "protocol_version": 2, "execution_targets": True,
"node_ids": ["DistributedImageListStart", "DistributedImageListEnd"]}
class Response:
status = 200
def __init__(self, value):
self.value = value
async def __aenter__(self):
return self
async def __aexit__(self, *args):
pass
async def json(self):
return self.value
@pytest.mark.parametrize("advertised", [None, [], ["exact"], "png", ["png", 3]])
def test_nonexact_requires_advertised_capability_before_queue(dispatch_module, advertised):
cap = capability()
if advertised is not None:
cap["image_transfer_formats"] = advertised
session = types.SimpleNamespace(get=Mock(return_value=Response(cap)), post=Mock())
dispatch_module.get_client_session = AsyncMock(return_value=session)
with pytest.raises(RuntimeError, match="upgrade.*no work was submitted"):
asyncio.run(dispatch_module.preflight_image_list_workers(
[{"id": "worker"}], required_transfer_formats=["png"]))
assert session.get.call_count == 1
session.post.assert_not_called()
@pytest.mark.parametrize("requested", [(), ("exact",), ("png", "jpeg")])
def test_exact_legacy_or_advertised_nonexact_capability_succeeds(dispatch_module, requested):
cap = capability()
if "png" in requested:
cap["image_transfer_formats"] = list(FORMATS)
session = types.SimpleNamespace(get=Mock(side_effect=[Response(cap), Response({"queue_running": [], "queue_pending": []})]))
dispatch_module.get_client_session = AsyncMock(return_value=session)
assert asyncio.run(dispatch_module.preflight_image_list_workers(
[{"id": "worker"}], required_transfer_formats=requested)) == {"worker": True}
def test_capability_endpoint_advertises_shared_formats(monkeypatch):
from aiohttp import web
name = "_transfer_routes"
monkeypatch.setitem(sys.modules, name, types.ModuleType(name))
monkeypatch.setitem(sys.modules, name + ".api", types.ModuleType(name + ".api"))
monkeypatch.setitem(sys.modules, name + ".image_list", types.ModuleType(name + ".image_list"))
monkeypatch.setitem(sys.modules, name + ".image_list.runtime", rt)
monkeypatch.setitem(sys.modules, "server", types.SimpleNamespace(PromptServer=types.SimpleNamespace(
instance=types.SimpleNamespace(routes=web.RouteTableDef()))))
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)
spec.loader.exec_module(module)
response = asyncio.run(module.image_list_capabilities_endpoint(None))
caps = json.loads(response.text)
assert caps["image_transfer_formats"] == list(FORMATS)
assert all(caps[name] is True for name in ("supported", "execution_targets", "value_results"))
assert "DistributedImageListGetValue" in caps["node_ids"]
@pytest.mark.parametrize("reject_call", [None, 1, 2, 3])
@pytest.mark.parametrize("with_values", [False, True])
def test_runtime_initial_and_each_nonexact_assignment_preflight(reject_call, with_values):
from test_image_list_runtime_boundaries import RuntimeBoundaryTests
async def exercise():
fixture = RuntimeBoundaryTests()
await fixture.asyncSetUp()
try:
context = fixture.create(2)
state = fixture.rt.get_master_job(context)
state.original_prompt[state.end_id]["inputs"].update(transfer_format="png", transfer_quality=87)
if with_values:
state.original_prompt[state.end_id]["inputs"]["value_10"] = [state.start_id, 2]
for item in state.items:
item["values"] = {"value_1": "retained"}
original_set = fixture.rt.set_master_result
def set_with_override(context, index, token, result):
return original_set(context, index, token, {**result, "values": {"value_10": index}})
fixture.rt.set_master_result = set_with_override
calls = []
async def preflight(workers, **kwargs):
calls.append(kwargs)
assert kwargs["required_transfer_formats"] == ("png",)
assert kwargs.get("require_value_results", False) is with_values
if len(calls) == reject_call:
raise RuntimeError("upgrade worker; no work was submitted")
return {"worker": True}
sys.modules["_image_list_boundaries.api.orchestration.dispatch"].preflight_image_list_workers = preflight
if reject_call:
with pytest.raises(RuntimeError, match="upgrade"):
await fixture.rt.run_image_list_job(context, include_values=with_values)
assert len(fixture.uploads) == max(0, reject_call - 2)
assert len(fixture.prompts) == max(0, reject_call - 2)
else:
result_lists = await fixture.rt.run_image_list_job(context, include_values=with_values)
assert len(calls) == 3
assert len(fixture.uploads) == 2
if with_values:
assert result_lists[3] == [{"value_1": "retained", "value_10": index} for index in range(2)]
assert not fixture.rt._MASTER_JOBS
finally:
# Do not call doCleanups(): this borrowed IsolatedAsyncioTestCase has
# no runner, but its sys.modules patch still must be restored.
fixture.patch_modules.stop()
asyncio.run(exercise())
+2 -2
View File
@@ -94,7 +94,7 @@ async def check():
assert await extension.get_node_list() == [MediaImportV1, MediaExportV1]
imp = MediaImportV1.GET_NODE_INFO_V1()
exp = MediaExportV1.GET_NODE_INFO_V1()
assert imp['output'] == ['IMAGE','AUDIO','MASK','STRING','VIDEO'], imp
assert imp['output'] == ['IMAGE','AUDIO','MASK','STRING','VIDEO','*'], imp
assert exp['output_node'] and MediaExportV1.define_schema().is_input_list, exp
with tempfile.TemporaryDirectory(prefix='typed-media-', dir=os.environ.get('TMPDIR')) as scratch:
base = Path(scratch)
@@ -152,7 +152,7 @@ async def check():
mapped, missing, hidden = execution.get_input_data({'artifact_json':json.dumps(artifact),'filename':filename},MediaImportV1,'import')
got = await execution._async_map_node_over_list('probe','import',MediaImportV1,mapped,MediaImportV1.FUNCTION,v3_data=hidden)
got = await execution.resolve_map_node_over_list_results(got)
assert len(got[0].result) == 5
assert len(got[0].result) == 6
assert all(x is None for n,x in enumerate(got[0].result) if n != list(SOCKET_KINDS).index(socket))
bad = dict(artifact,sha256='0'*64)
try: import_value(bad,filename)
+5 -5
View File
@@ -56,7 +56,7 @@ def test_serverless_pending_waits_past_normal_worker_timeout():
queue.put_nowait({
"worker_id": participant,
"image_index": 0,
"tensor": torch.ones(1, 2, 2, 2),
"tensor": torch.ones(1, 2, 2, 3),
"is_last": True,
})
@@ -74,7 +74,7 @@ def test_serverless_pending_waits_past_normal_worker_timeout():
images, _audio = asyncio.run(run_collection())
assert tuple(images.shape) == (1, 2, 2, 2)
assert tuple(images.shape) == (1, 2, 2, 3)
assert manager.pending_calls
assert probe_calls == []
assert job_id not in module.prompt_server.distributed_pending_jobs
@@ -162,12 +162,12 @@ def test_serverless_tensor_shape_mismatch_raises_instead_of_silent_fallback():
queue.put_nowait({
"worker_id": "serverless:pool:1",
"image_index": 0,
"tensor": torch.ones(1, 3, 2, 2),
"tensor": torch.ones(1, 3, 2, 3),
"is_last": True,
})
manager = _FakeManager()
collector = _set_up_collector(module, job_id, queue, manager)
local_images = torch.zeros(1, 2, 2, 2)
local_images = torch.zeros(1, 2, 2, 3)
with pytest.raises(RuntimeError, match="serverless.*shape|shape.*serverless"):
asyncio.run(collector.execute(
@@ -188,7 +188,7 @@ def test_regular_worker_timeout_keeps_partial_behavior_without_creating_manager(
collector = _set_up_collector(module, job_id, queue)
module.get_worker_timeout_seconds = lambda: 0.01
module.HEARTBEAT_INTERVAL = 0.1
local_images = torch.zeros(1, 2, 2, 2)
local_images = torch.zeros(1, 2, 2, 3)
collector._get_serverless_manager = lambda *_args, **_kwargs: pytest.fail("regular collection must not create manager")
images, _audio = asyncio.run(collector.execute(