Compare commits

..
Author SHA1 Message Date
Robert Wojciechowski 98987db131 feat: integrate native image-list caching and exact serverless value overrides 2026-10-14 04:08:48 +00:00
Robert Wojciechowski 593047badf Merge Collector transport with image-list values and native cache reuse 2026-10-14 03:53:29 +00:00
Robert Wojciechowski a9ef36f951 feat: unify Collector media transport and add direct exact returns
Use the existing images/audio Collector for typed serverless pools. Preserve saved legacy_png defaults, map them to serverless png, and retain every transfer format with exact opt-in. Remove the separate Media Collector.

Send direct exact frames through bounded raw tensor multipart transport without 8-bit fallback. Reject incompatible legacy pools, invalid payloads and incomplete exact results. Cover real CPU HTTP round trips, native ComfyUI execution, mixed participant ordering and failure paths.
2026-10-14 03:41:05 +00:00
Robert Wojciechowski 6db5692d78 fix: preserve native image-list processing cache reuse 2026-10-14 03:15:49 +00:00
Robert Wojciechowski 591ffc7396 fix: preserve image-list worker selection across queued runs (#103) 2026-10-13 19:52:29 -07:00
Robert Wojciechowski 5f9a2f0fc6 fix: carry image-list values through End context
Preserve source values without extra wires, collect explicit per-item overrides, and expose ordered bundles through Get Value. Keep legacy sockets and no-override transport compatible.
2026-10-14 01:51:13 +00:00
Robert Wojciechowski 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 b636fd5202 feat: add explicit image return codecs and quality controls 2026-10-14 00:37:42 +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 f0927f1166 feat: add versioned lossless serverless media transport 2026-10-13 23:53:58 +00:00
Robert Wojciechowski c30036ae31 feat: preserve required output utilities in image-list execution 2026-10-13 23:26:25 +00:00
Robert Wojciechowski 41caf02481 docs: tighten image-list and serverless wording 2026-10-13 22:38:10 +00:00
Robert Wojciechowski 494d5506a0 README: add Distributed Image List Start/End to nodes table 2026-10-13 22:27:07 +00:00
Robert Wojciechowski 7ed60b93ca feat: dynamically distribute image lists across direct workers (#98)
* feat: add Runpod serverless image collector pools

* fix: prevent worker settings clipping and remove preview copy

* fix: initialize worker settings before mounting the form

* fix: validate complete serverless requests against 10 MB limit

* fix: simplify serverless cards and keep status dots circular

* fix: match serverless endpoint box to master card

* fix: preserve serverless capacity and recovery across job failures

Retain durable reservations until provider termination or explicit recovery, protect unresolved pools and credentials, and add scoped reconciliation. Consolidate lifecycle and secret policies and split prompt/media stages with regression coverage.

* feat: add dynamic per-image list queue with Start and End nodes

* fix(image-list): reject upstream data bypasses around Start

* fix(image-list): make optional End text a connection-only input

* fix(image-list): isolate unavailable workers during preflight

* fix(image-list): bound acknowledgements and configure run deadline

* fix(media-sync): preserve output and temp source roots

* fix(image-list): own job lifecycle and preserve worker preparation

* feat(image-list): let master participate unless orchestrator-only

* feat(image-list)!: replace fixed Start extras with growable value_N slots

Distributed Image List Start no longer has reference_image, text,
strength, seed or per-slot mode dropdowns. It now takes images, mask and
up to 10 generic value_N inputs (IMAGE, MASK, STRING, INT, FLOAT,
BOOLEAN) that grow via native V3 Autogrow, each with a matching output.

- One value is shared by every image; one per image pairs in order; any
  other count fails before dispatch. Only images is split into frames.
- Outputs are declared up front (image, mask, item_index, context,
  value_1..value_10). web/imageListStart.js hides unused value outputs,
  mirrors input types, and follows Autogrow compaction so downstream
  links keep their value.
- Image-list protocol bumped to v2; older workers fail preflight.

Breaking: saved workflows using the old Start inputs must recreate Start.

* fix(serverless): disable Clear saved API key while the worker is enabled

The checkbox read "Clear saved API key (disable pool first)", but the
panel never defines "pool" and saving with it ticked on an enabled worker
was rejected by the backend.

- Disable the checkbox while the worker is enabled, with a tooltip saying
  to turn the worker off first. Also disable it when no key is saved.
- Use "worker" instead of "pool" in user-facing serverless errors, the
  serverless guide and the README entry.

* feat(serverless): replace Concurrent jobs with a single Images per run setting

- Remove the separate Concurrent jobs field; rename Jobs per run to
  Images per run (1-16) with a tooltip explaining that Runpod's max
  workers decides how many GPUs run them.
- Backend uses a fixed per-worker limit of 16 outstanding jobs, so all
  images of a run are submitted at once. Unresolved jobs still count
  against it. The limit is constant, so it never has to change while
  jobs are unresolved.
- Legacy 'concurrency' in saved configs and older browser tabs is
  accepted and ignored/dropped; it no longer triggers the busy guard.
- Revert the enable toggle and show the reason when the backend rejects it.
- Trim a pasted api.runpod.ai/v2/<id> URL to the endpoint ID; plainer
  endpoint-ID error otherwise.
- Confirm before a type change deletes a saved Runpod API key.
- Update docs and tests.
2026-10-13 15:20:05 -07:00
88 changed files with 15778 additions and 328 deletions
+4 -2
View File
@@ -28,7 +28,7 @@
- Handles single images and videos
#### Dynamic Image-List Processing
- Schedule one image at a time to the next available direct worker; collect results in input order
- Process a whole image list in parallel: the master and every worker take the next image when they free up, and results come back in input order
- Use **Distributed Image List Start / End** around a per-image processing region
- [Wiring, supported inputs, and first-version limits](/docs/image-list-queue.md)
@@ -48,7 +48,7 @@ ComfyUI Distributed supports four types of workers:
- **Local Workers** - Additional GPUs on the same machine (auto-configured on first launch)
- **Remote Workers** - GPUs on other computers in your network
- **Cloud Workers** - GPUs hosted on a cloud service like Runpod, accessible via secure tunnels
- **Serverless (Runpod)** - Queue-based Runpod endpoints for image Collector jobs; no Distributed install or inbound tunnel on workers. [Setup, endpoint contract, and preview limitations](/docs/serverless-workers.md).
- **Serverless (Runpod)** - Queue-based endpoints for image Collector jobs; no Distributed install or inbound tunnel on the endpoint. [Setup, endpoint contract, and preview limitations](/docs/serverless-workers.md).
> For detailed setup instructions, see the [setup guide](/docs/worker-setup-guides.md)
@@ -184,6 +184,8 @@ Use **Distributed Value** when you want per-worker overrides (for example, diffe
| **Audio Segment Divider** | Splits an audio waveform into up to ten sequential time segments |
| **Distributed Model Name** | Passes model paths to workers, enabling workflows to use models not present on the master in orchestrator-only mode |
| **Distributed Empty Image** | Produces an empty IMAGE batch used when the master delegates all work |
| **Distributed Image List Start** | Turns an image list or batch into one job per image, scheduled to the next available participant; mask and value inputs are shared across images or paired one per image |
| **Distributed Image List End** | Closes the paired region and returns images, masks, and texts as lists in the original input order |
---
+19 -5
View File
@@ -97,10 +97,11 @@ _SECRET_FIELDS = {"api_key", "clear_api_key"}
_SERVERLESS_LOCAL_FIELDS = {"host", "port", "listen_port", "cuda_device", "extra_args"}
_SERVERLESS_ONLY_FIELDS = {
"provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds", "api_key_configured",
"media_profile",
}
_SERVERLESS_CONFIG_FIELDS = {
"worker_id", "type", "provider", "endpoint_id", "jobs_per_run", "concurrency",
"timeout_seconds", "enabled", "name", "api_key", "clear_api_key",
"timeout_seconds", "enabled", "name", "api_key", "clear_api_key", "media_profile",
}
@@ -199,7 +200,13 @@ def _clean_bulk_type_conversions(old_workers, new_workers):
if not isinstance(worker, dict) or worker.get("id") is None:
continue
old = old_by_id.get(str(worker["id"]))
if old is None or old.get("type") != "serverless" or worker.get("type") == "serverless":
if old is None or old.get("type") != "serverless":
continue
if worker.get("type") == "serverless":
# Older clients do not know this field. Only an explicit selection
# may change the endpoint contract, including on bulk saves.
if "media_profile" not in worker and "media_profile" in old:
worker["media_profile"] = old["media_profile"]
continue
converted = dict(worker)
for field in _SERVERLESS_ONLY_FIELDS:
@@ -218,7 +225,10 @@ def _bulk_serverless_changes(old_workers, new_workers):
if new is None or new.get("type") != "serverless":
_ensure_pool_idle(worker_id)
secret_ids.add(worker_id)
elif old.get("endpoint_id") != new.get("endpoint_id"):
elif (
old.get("endpoint_id") != new.get("endpoint_id")
or old.get("media_profile", "legacy-image-v1") != new.get("media_profile", "legacy-image-v1")
):
_ensure_pool_idle(worker_id)
elif new is not None and new.get("type") == "serverless":
secret_ids.add(worker_id)
@@ -401,7 +411,7 @@ async def _update_worker_endpoint(request, secret_changes):
candidate.setdefault("provider", "runpod")
candidate.setdefault("enabled", False)
# Legacy "concurrency" is accepted (see _SERVERLESS_CONFIG_FIELDS) but ignored.
for field in ("provider", "endpoint_id", "jobs_per_run", "timeout_seconds", "enabled", "name"):
for field in ("provider", "endpoint_id", "jobs_per_run", "timeout_seconds", "enabled", "name", "media_profile"):
if field in data:
candidate[field] = data[field]
candidate = _clean_serverless_worker(candidate)
@@ -424,8 +434,12 @@ async def _update_worker_endpoint(request, secret_changes):
current_type == "serverless"
and existing.get("endpoint_id") != validated.get("endpoint_id")
)
media_profile_changed = (
current_type == "serverless"
and existing.get("media_profile", "legacy-image-v1") != validated.get("media_profile", "legacy-image-v1")
)
if current_type == "serverless" and (
endpoint_changed or next_secret != prior_secret
endpoint_changed or media_profile_changed or next_secret != prior_secret
):
_ensure_pool_idle(worker_id)
+8 -3
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,
@@ -35,9 +36,12 @@ async def _json_body(request):
async def image_list_capabilities_endpoint(_request):
return web.json_response({
"supported": True,
"execution_targets": True,
"value_results": True,
"protocol_version": PROTOCOL_VERSION,
"node_ids": ["DistributedImageListStart", "DistributedImageListEnd"],
"node_ids": ["DistributedImageListStart", "DistributedImageListEnd", "DistributedImageListGetValue"],
"transport": "typed-float-tensor-json-v1",
"image_transfer_formats": list(IMAGE_TRANSFER_FORMATS),
})
@@ -51,8 +55,9 @@ async def image_list_item_endpoint(request):
if not isinstance(context_id, str) or not context_id:
raise ValueError("context_id must be a non-empty string")
item = data.get("item")
store_worker_item(data.get("token"), item)
return web.json_response({"status": "stored", "item_index": item["index"]})
identity = store_worker_item(data.get("token"), item)
return web.json_response({"status": "stored", "item_index": item["index"],
"item_identity": identity})
except (ValueError, TypeError) as exc:
return web.json_response({"error": str(exc)}, status=400)
except Exception as exc:
+35
View File
@@ -37,6 +37,8 @@ prompt_server = server.PromptServer.instance
# envelope: JSON { "job_id", "worker_id", "format", "start_index", "count", "width", "height",
# "is_last", "audio"?: {...} }
# media: one part per frame for still formats, or one video file for video formats
# format "exact" adds envelope "descriptors" (one tensor-le-v1 descriptor per
# frame) and sends one raw tensor part per frame; values are not quantized.
def _decode_image_sync(image_path):
@@ -299,6 +301,8 @@ async def _enqueue_worker_items(multi_job_id, items):
def _decode_compressed_chunk_sync(envelope, parts):
"""Decode a compressed chunk into per-frame IMAGE tensors (threadpool)."""
if envelope.get("format") == transfer_codec.EXACT_FORMAT:
return _decode_exact_chunk_sync(envelope, parts)
frames = transfer_codec.decode_chunk(envelope, parts)
return [
ensure_contiguous(torch.from_numpy(transfer_codec.uint8_frame_to_tensor_array(frame)))
@@ -306,6 +310,35 @@ def _decode_compressed_chunk_sync(envelope, parts):
]
def _decode_exact_chunk_sync(envelope, parts):
"""Decode raw exact IMAGE frames, checking each against the envelope geometry."""
from ..serverless_bridge import codec as bridge_codec
count, width, height = envelope.get("count"), envelope.get("width"), envelope.get("height")
for name, value in (("count", count), ("width", width), ("height", height)):
if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
raise transfer_codec.TransferCodecError(f"envelope.{name}: expected positive integer")
descriptors = envelope.get("descriptors")
if not isinstance(descriptors, list) or len(descriptors) != count or len(parts) != count:
raise transfer_codec.TransferCodecError(
f"exact chunk expected {count} descriptors and frames, got "
f"{len(descriptors) if isinstance(descriptors, list) else 'no'} descriptors and {len(parts)} frames"
)
tensors = []
for descriptor, payload in zip(descriptors, parts):
if not isinstance(descriptor, dict) or descriptor.get("encoding") != "tensor-le-v1":
raise transfer_codec.TransferCodecError("exact frames require tensor-le-v1 encoding")
shape = descriptor.get("shape") if isinstance(descriptor, dict) else None
if not isinstance(shape, list) or len(shape) != 4 or shape[:3] != [1, height, width]:
raise transfer_codec.TransferCodecError(
f"exact frame shape {shape} does not match the envelope's 1x{height}x{width} frames"
)
try:
tensors.append(bridge_codec.decode_payload("IMAGE", descriptor, payload))
except ValueError as exc:
raise transfer_codec.TransferCodecError(f"exact frame: {exc}") from None
return tensors
def _validate_compressed_envelope(envelope):
errors = []
if not isinstance(envelope, dict):
@@ -384,6 +417,7 @@ async def _job_complete_compressed(request):
"image_index": start_index + pos,
"is_last": bool(is_last and pos == last_pos),
"audio": decoded_audio if pos == last_pos else None,
"format": envelope["format"],
}
for pos, tensor in enumerate(tensors)
]
@@ -462,6 +496,7 @@ async def job_complete_endpoint(request):
"image_index": int(batch_idx),
"is_last": is_last,
"audio": decoded_audio,
"format": transfer_codec.LEGACY_FORMAT,
}
],
)
+41 -1
View File
@@ -77,6 +77,8 @@ async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0):
"workflow": payload.get("workflow"),
"client_id": client_id,
}
if "partial_execution_targets" in payload:
ws_payload["partial_execution_targets"] = payload["partial_execution_targets"]
ws_url = worker_url.replace("http://", "ws://").replace("https://", "wss://")
ws_url = f"{ws_url}/distributed/worker_ws"
session = await get_client_session()
@@ -128,11 +130,17 @@ async def dispatch_worker_prompt(
use_websocket=False,
trace_execution_id=None,
return_prompt_id=False,
*,
partial_execution_targets=None,
):
"""Send the prepared prompt to a worker ComfyUI instance."""
worker_url = build_worker_url(worker)
url = build_worker_url(worker, "/prompt")
payload = {"prompt": prompt_obj}
target_payload = {}
if partial_execution_targets is not None:
target_payload["partial_execution_targets"] = list(partial_execution_targets)
payload.update(target_payload)
extra_data = {}
if workflow_meta:
extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
@@ -146,6 +154,7 @@ async def dispatch_worker_prompt(
{
"prompt": prompt_obj,
"workflow": workflow_meta,
**target_payload,
},
client_id,
)
@@ -179,7 +188,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):
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
@@ -187,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)
@@ -227,6 +245,28 @@ async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0):
f"{sorted(required_nodes)}; no work was submitted"
)
if capability.get("execution_targets") is not True:
raise RuntimeError(
f"Image-list worker {worker_id} lacks explicit execution_targets support; "
"upgrade that worker before running this feature; no work was submitted"
)
if 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; "
"upgrade that worker before using End overrides; no work was submitted"
)
check = "queue"
queue_url = build_worker_url(worker, "/queue")
async with session.get(queue_url, timeout=aiohttp.ClientTimeout(total=timeout)) as response:
+7 -3
View File
@@ -18,7 +18,7 @@ class AmbiguousSubmission(RunpodError):
"""Submission may have been accepted: never retry it automatically."""
def serialize_request(workflow, images, timeout_seconds):
def serialize_request(workflow, images, timeout_seconds, *, media=None):
"""Serialize and size-check the exact JSON body used by Runpod ``/run``."""
payload = {
"input": {"workflow": workflow, "images": images},
@@ -27,6 +27,10 @@ def serialize_request(workflow, images, timeout_seconds):
"ttl": timeout_seconds * 1000,
},
}
if media is not None:
from ...serverless_bridge.handler import _validate_request
payload["input"]["media"] = media
_validate_request(payload["input"])
body = json.dumps(
payload,
allow_nan=False,
@@ -108,8 +112,8 @@ class RunpodClient:
async def health(self):
return await self._request("GET", "health")
async def submit(self, workflow, images, timeout_seconds):
body = serialize_request(workflow, images, timeout_seconds)
async def submit(self, workflow, images, timeout_seconds, *, media=None):
body = serialize_request(workflow, images, timeout_seconds, media=media)
result = await self._request("POST", "run", request_body=body)
try:
return checked_id(result.get("id"))
@@ -0,0 +1,239 @@
"""Pure compilation of prepared Distributed Collectors for comfy-media-v1 pools.
The legacy image compiler remains separate. This compiler emits one MediaExportV1
execution target per connected Collector socket (``images`` and/or ``audio``),
retaining only its upstream computation. It never imports ComfyUI or
inspects/executes runtime node classes.
"""
from __future__ import annotations
import copy
import hashlib
import json
from graphlib import CycleError, TopologicalSorter
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,
_REMOVED_SCALAR_NODES,
_is_link,
_is_unsupported_distributed,
_make_scalar_resolver,
_prune_graph,
)
_COLLECTOR = "DistributedCollector"
# Collector inputs, in output-slot order, that a typed pool returns.
COLLECTOR_SOCKETS = ("images", "audio")
# Collector widget defaults (the saved-workflow contract is legacy_png / 90).
COLLECTOR_DEFAULT_FORMAT = "legacy_png"
COLLECTOR_DEFAULT_QUALITY = 90
def collector_image_policy(inputs):
"""Return the endpoint IMAGE (format, quality) for one Collector.
legacy_png is the Collector's saved default and names a JSON/base64 framing
that only applies to direct workers; serverless returns the same 8-bit
pixels as compressed png. exact stays opt-in. No other format is altered.
"""
fmt, quality = validate_options(inputs.get("transfer_format", COLLECTOR_DEFAULT_FORMAT),
inputs.get("transfer_quality", COLLECTOR_DEFAULT_QUALITY))
return ("png" if fmt == "legacy_png" else fmt), quality
def _validate_identifier(value, context):
# Validate the original token before any lookup. In particular, neither
# integer node IDs nor whitespace in prepared identities is normalized.
try:
identifier(value)
except ValueError:
raise ValueError(f"{context} must be a valid nonempty media identifier") from None
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")
collectors = []
for node_id, node in prepared_prompt.items():
_validate_identifier(node_id, "prompt node ID")
if not isinstance(node, dict):
raise ValueError(f"prompt node {node_id!r} must be an object")
class_type = node.get("class_type")
if not isinstance(class_type, str) or not class_type.strip():
raise ValueError(f"prompt node {node_id!r} must have a nonempty class_type")
if class_type == _COLLECTOR:
collectors.append(node_id)
elif _is_unsupported_distributed(class_type):
raise ValueError(f"unsupported distributed or batch node: {class_type}")
inputs = node.get("inputs", {})
if not isinstance(inputs, dict):
raise ValueError(f"prompt node {node_id!r} inputs must be an object")
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 require_collectors and not collectors:
raise ValueError("prepared prompt contains no DistributedCollector nodes")
dependencies = {}
for node_id, node in prepared_prompt.items():
parents = set()
for input_name, value in node.get("inputs", {}).items():
# Ordinary inputs may contain literal lists. As in the legacy
# compiler, exactly two elements denotes an API link, not media.
if not _is_link(value):
continue
source_id, slot = value
context = f"link at node {node_id!r} input {input_name!r}"
_validate_identifier(source_id, context + " source ID")
if source_id not in prepared_prompt:
raise ValueError(f"invalid {context}: source {source_id!r} does not exist")
if type(slot) is not int or slot < 0:
raise ValueError(f"invalid {context}: output slot must be a nonnegative integer")
source_class = prepared_prompt[source_id]["class_type"]
if source_class in _REMOVED_SCALAR_NODES or source_class in _PRIMITIVE_NODES:
if slot != 0:
raise ValueError(f"invalid output slot {slot} for scalar node {source_id!r}")
elif source_class == _COLLECTOR and slot >= len(COLLECTOR_SOCKETS):
raise ValueError(f"invalid output slot {slot} for collector {source_id!r}")
parents.add(source_id)
dependencies[node_id] = parents
for collector_id in collectors:
inputs = prepared_prompt[collector_id].get("inputs", {})
collector_image_policy(inputs)
for field in ("multi_job_id", "worker_id"):
_validate_identifier(inputs.get(field), f"collector {collector_id!r} {field}")
connected = [socket for socket in COLLECTOR_SOCKETS if socket in inputs]
if not connected:
raise ValueError(f"collector {collector_id!r} requires at least one connected media socket")
for socket in connected:
value = inputs[socket]
if not _is_link(value):
raise ValueError(f"collector {collector_id!r} {socket} must be linked; use PrimitiveString for literal text")
if prepared_prompt[value[0]]["class_type"] in _REMOVED_SCALAR_NODES:
# Export sockets must remain actual links after scalar folding.
# A PrimitiveString wrapper supports distributed text values.
raise ValueError(f"collector {collector_id!r} {socket} cannot link directly to a distributed scalar; use a typed node or PrimitiveString")
# Iterative topological sorting checks ordinary and unused cycles too, not
# only the primitive chains that the shared scalar resolver traverses.
try:
ordered = list(TopologicalSorter(dependencies).static_order())
except CycleError:
raise ValueError("cycle in prepared media prompt") from None
collector_set = set(collectors)
has_collector_ancestor = {}
for node_id in ordered:
has_collector_ancestor[node_id] = any(
parent in collector_set or has_collector_ancestor[parent]
for parent in dependencies[node_id]
)
if node_id in collector_set and has_collector_ancestor[node_id]:
raise ValueError("collector-to-collector dependence is unsupported")
return sorted(collectors)
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:
if socket not in inputs:
continue
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]]:
"""Return a worker graph and canonical typed-media export bindings.
Each DistributedCollector must have prepared literal ``multi_job_id`` and
``worker_id`` identities and at least one linked ``images``/``audio`` input.
Present media inputs must be API links. Collector transport controls become
the IMAGE export policy (``legacy_png`` is sent as ``png``); other controls
are not exported.
Only MediaExportV1 targets and their ancestors survive. Required ordinary
output/utility nodes are preserved regardless of their runtime OUTPUT_NODE
flag; callers must dispatch exactly the returned binding ``node_id`` values
as partial_execution_targets. Distributed scalar dependencies use the same
resolver as the legacy compiler. Direct scalar-to-media links need a typed
wrapper so that each export input remains a link.
Validation and compilation never modify the caller's source graph. This
structural compiler does not infer socket types or unknown nodes' output
counts; ComfyUI performs runtime type validation at the endpoint.
"""
_validate_identifier(request_id, "request_id")
collector_ids = _validate_prompt(prepared_prompt)
bindings, exports = _build_exports(prepared_prompt, collector_ids, 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
+38 -5
View File
@@ -146,6 +146,10 @@ class ServerlessManager:
async def start(self, execution_id, prompt_id, prepared_jobs):
self.ensure_capacity(len(prepared_jobs))
from .serverless_typed_results import validate_prepared_job
# Validate the entire batch before creating even the first local slot.
for prepared in prepared_jobs:
validate_prepared_job(prepared)
new_runs = []
# Journal every local slot before starting any billed submission.
for prepared in prepared_jobs:
@@ -167,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
@@ -177,14 +183,30 @@ class ServerlessManager:
return any(run["worker"]["id"] == participant_id and
any(m["job_id"] == collector_job_id for m in run["manifest"]) for run in self.runs.values())
async def _deliver(self, run, images=None, error=None):
async with self.server.distributed_jobs_lock:
async def _deliver(self, run, images=None, error=None, *, media=None):
if media is not None and not error:
# Validate all destinations before exposing any partial success.
for binding in run["manifest"]:
if binding["participant_id"] != run["worker"]["id"]:
raise ValueError("Serverless media participant identity mismatch")
values = media[binding["job_id"]][binding["socket"]]
if not isinstance(values, list) or not values:
raise ValueError("Missing serverless media values")
async with self.server.distributed_jobs_lock:
delivered = set()
for binding in run["manifest"]:
job_id = binding["job_id"]
if error or media is not None:
if job_id in delivered:
continue
delivered.add(job_id)
queue = self.server.distributed_pending_jobs.get(binding["job_id"])
if queue is None:
continue # A cancelled/finished collector must never be resurrected.
if error:
queue.put_nowait({"worker_id": run["worker"]["id"], "serverless_error": error, "is_last": True})
elif media is not None:
queue.put_nowait({"worker_id": run["worker"]["id"], "media": media[job_id], "is_last": True})
else:
batch = images[binding["job_id"]]
for index, tensor in enumerate(batch):
@@ -218,13 +240,18 @@ class ServerlessManager:
submitted = False
provider_terminal = False
try:
from .serverless_typed_results import validate_prepared_job, collect_typed_results
typed = validate_prepared_job(run)
async with asyncio.timeout(max(0, run["deadline"] - time.monotonic())):
async with semaphore:
await self._reserve_provider_capacity(run)
client = await self._client(run["worker"])
run["client"] = client
submitted = True
job_id = await client.submit(run["workflow"], run["images"], run["worker"]["timeout_seconds"])
if typed:
job_id = await client.submit(run["workflow"], run["images"], run["worker"]["timeout_seconds"], media=run["media"])
else:
job_id = await client.submit(run["workflow"], run["images"], run["worker"]["timeout_seconds"])
run["provider_job_id"] = job_id
# Persist immediately: all subsequent actions use this exact ID.
self.ledger.update(rid, job_id=job_id, status="IN_QUEUE", updated_at=time.time())
@@ -240,8 +267,12 @@ class ServerlessManager:
if not isinstance(output, dict) or output.get("error") or output.get("errors"):
raise RunpodError("Serverless workflow reported an output error; inspect the provider job")
self.ledger.update(rid, status="DOWNLOADING", updated_at=time.time())
images = await decode_results(output, run["manifest"])
await self._deliver(run, images=images)
if typed:
media = await collect_typed_results(output, run["media"], run["manifest"])
await self._deliver(run, media=media)
else:
images = await decode_results(output, run["manifest"])
await self._deliver(run, images=images)
self.ledger.update(rid, status="COMPLETED", updated_at=time.time())
return
self.ledger.update(rid, status=status, updated_at=time.time())
@@ -276,6 +307,7 @@ class ServerlessManager:
# base64 inputs, workflow metadata, or a client holding credentials.
run.pop("workflow", None)
run.pop("images", None)
run.pop("media", None)
run["client"] = None
# A released local semaphore alone cannot release a durable reservation.
await self._notify_capacity()
@@ -359,6 +391,7 @@ class ServerlessManager:
self.ledger.update(rid, status="CANCELLED", updated_at=time.time())
run.pop("workflow", None)
run.pop("images", None)
run.pop("media", None)
await self._deliver(run, error="Serverless job cancelled before submission")
continue
# Do not send two cancellation requests for the same live job in
@@ -0,0 +1,225 @@
"""Bounded local native inputs for comfy-media-v1 requests.
Only stock LoadAudio/LoadVideo schemas are adapted. Original nodes and links
remain in the private request graph for endpoint validation; only the existing
image preparer rewrites filenames. No upload, transcode, resample or provider
operation belongs here. The caller still enforces the complete request-body cap.
"""
import hashlib
import io
import json
from pathlib import Path
from ...serverless_bridge import codec, contract, video_codec
from . import serverless_media
_LOADERS = {"LoadAudio": ("audio", "AUDIO"), "LoadVideo": ("file", "VIDEO")}
Error = serverless_media.ServerlessMediaError
def _loader_references(workflow):
if not isinstance(workflow, dict):
raise Error("Serverless workflow must be an API graph")
loaders = {}
for node_id, node in workflow.items():
if not isinstance(node, dict):
continue
class_type = node.get("class_type")
if not isinstance(class_type, str) or class_type not in _LOADERS:
continue
contract.identifier(node_id)
key, kind = _LOADERS[class_type]
inputs = node.get("inputs")
if not isinstance(inputs, dict) or set(inputs) != {key}:
raise Error(f"{class_type} inputs must contain only its native filename; extra controls are unsupported")
if not isinstance(inputs[key], str):
raise Error(f"{class_type} requires a literal filename; linked inputs are unsupported")
normalized, _suffix = serverless_media._normalize_media_reference(inputs[key])
loaders[node_id] = (kind, normalized)
references = []
used = set()
for node_id, node in workflow.items():
if not isinstance(node, dict) or not isinstance(node.get("inputs"), dict):
continue
for key, value in node["inputs"].items():
if not isinstance(value, (list, tuple)) or not value or not isinstance(value[0], str):
continue
if value[0] not in loaders:
continue
if type(value) is not list or len(value) != 2 or type(value[1]) is not int or value[1] != 0:
raise Error("Native media loader links require output slot 0")
contract.identifier(node_id)
contract.identifier(key)
references.append((node_id, key, value[0]))
used.add(value[0])
if used != set(loaders):
raise Error("Native media loader has no consumer")
return loaders, references
def _source_path(value, folder_paths):
"""Resolve native annotations; authorization/opening is owned by _read_path."""
annotation = serverless_media._MEDIA_ANNOTATION.search(value)
unannotated = value[:annotation.start()] if annotation else value
if Path(unannotated).is_absolute():
return unannotated
try:
return folder_paths.get_annotated_filepath(value)
except Exception:
raise Error("Native media input path is missing or invalid") from None
def _decode_audio(data, limit):
"""Match core nodes_audio.load, bounding each frame *before* accumulation.
Float32 output bytes are charged even for narrow integer PCM. Native float64
remains float64, as in f32_pcm (there is no silent dtype conversion). Raw
retained arrays and each live AV plane allocation have a separate <= limit
guard, so wider/padded raw frames cannot bypass the output budget. PyAV owns
the one frame being decoded; it is checked before to_ndarray or torch.cat.
Concatenation/conversion use bounded temporary copies, not an RSS guarantee.
"""
if limit <= 0:
raise Error("AUDIO decoded input exceeds byte limit")
try:
import av
import torch
from comfy_extras.nodes_audio import f32_pcm
except ImportError:
raise Error("Native AUDIO decoding is unavailable") from None
try:
with av.open(io.BytesIO(data)) as container:
if not container.streams.audio:
raise Error("Native AUDIO input has no audio stream")
stream = container.streams.audio[0]
rate = stream.codec_context.sample_rate
channels = stream.channels
if type(rate) is not int or rate <= 0 or type(channels) is not int or channels <= 0:
raise Error("Native AUDIO stream shape or sample rate is invalid")
frames = []
output_bytes = raw_bytes = 0
dtype = None
for frame in container.decode(streams=stream.index):
samples = frame.samples
if (type(samples) is not int or samples <= 0 or frame.sample_rate != rate
or len(frame.layout.channels) != channels):
raise Error("Native AUDIO frame shape or sample rate changed")
if frame.format.name not in {"s16", "s16p", "s32", "s32p", "flt", "fltp", "dbl", "dblp"}:
raise Error("Unsupported native AUDIO PCM format")
itemsize = frame.format.bytes
count = samples * channels
# Native f32_pcm promotes integer PCM to float32, not float64.
decoded_bytes = count * max(4, itemsize)
frame_bytes = count * itemsize
plane_bytes = sum(plane.buffer_size for plane in frame.planes)
if (decoded_bytes > limit - output_bytes or frame_bytes > limit - raw_bytes
or plane_bytes > limit - raw_bytes):
raise Error("AUDIO decoded input or raw frame exceeds byte limit")
output_bytes += decoded_bytes
raw_bytes += frame_bytes
array = frame.to_ndarray()
if array.ndim != 2 or array.size != count or array.nbytes != frame_bytes:
raise Error("Native AUDIO frame shape is invalid")
buf = torch.from_numpy(array)
if buf.dtype not in (torch.int16, torch.int32, torch.float32, torch.float64):
raise Error("Unsupported native AUDIO PCM dtype")
if dtype is not None and buf.dtype != dtype:
raise Error("Native AUDIO frame dtype changed")
dtype = buf.dtype
if buf.shape[0] != channels:
buf = buf.view(-1, channels).t()
if tuple(buf.shape) != (channels, samples):
raise Error("Native AUDIO frame shape is invalid")
frames.append(buf)
if not frames:
raise Error("Native AUDIO input has no decoded frames")
waveform = f32_pcm(torch.cat(frames, dim=1)).unsqueeze(0)
return {"waveform": waveform, "sample_rate": rate}
except Error:
raise
except Exception:
raise Error("Native AUDIO input could not be decoded safely") from None
def _artifact_id(kind, descriptor, payload):
digest = hashlib.sha256()
digest.update(kind.encode("ascii") + b"\0")
# Include interpretation (not just waveform bytes), notably sample_rate.
digest.update(json.dumps(descriptor, sort_keys=True, separators=(",", ":"), allow_nan=False).encode())
digest.update(b"\0")
digest.update(payload)
return "input-" + kind.lower() + "-" + digest.hexdigest()
async def prepare_media_inputs(workflow: dict) -> tuple[list, list, list]:
"""Return legacy images, inline typed artifacts, and every consumer reference.
Paths are read once into bounded bytes through the video adapter's no-follow
directory-fd reader. Decoders receive BytesIO, never a second path open.
Source files and unique decoded artifacts have independent aggregate caps.
Repeated content of the same kind is decoded once and charged once.
"""
loaders, references = _loader_references(workflow)
image_view = {node_id: node for node_id, node in workflow.items() if node_id not in loaders}
# Retain the strict scan for every other loader before expensive decoding.
serverless_media._iter_image_references(image_view)
if not loaders:
return await serverless_media.prepare_images(image_view), [], []
try:
import folder_paths
except ImportError:
raise Error("ComfyUI media directories are unavailable") from None
roots = list(serverless_media._folder_roots(folder_paths).values())
artifacts = {}
source_artifacts = {}
source_digests = set()
filenames = {}
loader_artifacts = {}
file_total = decoded_total = 0
for node_id, (kind, filename) in loaders.items():
# Avoid retaining file buffers merely to deduplicate repeat references.
if (kind, filename) in filenames:
loader_artifacts[node_id] = filenames[kind, filename]
continue
path = _source_path(filename, folder_paths)
data = video_codec._read_path(path, contract.MAX_ARTIFACT_BYTES, roots)
source_digest = hashlib.sha256(data).hexdigest()
if source_digest not in source_digests:
if len(data) > contract.MAX_TOTAL_BYTES - file_total:
raise Error("Native media source files exceed total byte limit")
file_total += len(data)
source_digests.add(source_digest)
source_key = (kind, source_digest)
if source_key not in source_artifacts:
if len(artifacts) >= contract.MAX_ARTIFACTS:
raise Error("Native media input artifact count exceeds limit")
limit = min(contract.MAX_ARTIFACT_BYTES, contract.MAX_TOTAL_BYTES - decoded_total)
if kind == "AUDIO":
value = _decode_audio(data, limit)
descriptor, payload = codec.encode_payload(kind, value, max_bytes=limit)
del value
else:
from comfy_api.latest import InputImpl
value = InputImpl.VideoFromFile(io.BytesIO(data))
descriptor, payload = video_codec.encode_video(value, max_bytes=limit, allowed_roots=roots)
del value
aid = _artifact_id(kind, descriptor, payload)
if aid not in artifacts:
if len(payload) > limit:
raise Error("Native media decoded inputs exceed total byte limit")
artifacts[aid] = contract.artifact_from_payload(
aid, kind, descriptor, payload, max_bytes=contract.MAX_ARTIFACT_BYTES)
decoded_total += len(payload)
source_artifacts[source_key] = aid
del payload
aid = source_artifacts[source_key]
loader_artifacts[node_id] = filenames[kind, filename] = aid
del data
refs = [{"node_id": node_id, "input_name": key, "artifact_id": loader_artifacts[source]}
for node_id, key, source in references]
# Shared node dictionaries deliberately propagate only legacy image rewrites.
images = await serverless_media.prepare_images(image_view)
return images, list(artifacts.values()), refs
@@ -0,0 +1,79 @@
"""Explicit media-v1 preflight and atomic, lossless result materialization."""
from urllib.parse import urljoin
import aiohttp
from ...serverless_bridge import contract
from ...serverless_bridge.contract import MAX_ARTIFACT_BYTES, PROFILE, validate_bindings
from .serverless_media import (
MAX_REDIRECTS, _CHUNK_SIZE, _DOWNLOAD_TIMEOUT_SECONDS,
_PublicResolver, _validate_artifact_url,
)
def validate_prepared_job(prepared):
"""Fail before journalling/submission; legacy jobs keep their original path."""
profile = prepared["worker"].get("media_profile", "legacy-image-v1")
media = prepared.get("media")
if profile == "legacy-image-v1" and media is None:
return False
if profile != PROFILE or media is None:
raise ValueError("Serverless media profile/request mismatch")
from .runpod_client import serialize_request
# The shared endpoint validator also checks exporter controls and input refs.
serialize_request(prepared["workflow"], prepared["images"],
prepared["worker"]["timeout_seconds"], media=media)
expected = validate_bindings(prepared["manifest"])
if expected != validate_bindings(media["outputs"]):
raise ValueError("Serverless media manifest identity mismatch")
if any(binding["participant_id"] != prepared["worker"]["id"] for binding in expected.values()):
raise ValueError("Serverless media participant identity mismatch")
return True
async def download_public_artifact(url, *, max_bytes=MAX_ARTIFACT_BYTES):
"""Contract-sized public HTTPS download, never using a provider session/key."""
if type(max_bytes) is not int or not 0 <= max_bytes <= MAX_ARTIFACT_BYTES:
raise ValueError("Media artifact exceeds the permitted size")
current_url = _validate_artifact_url(url)
connector = aiohttp.TCPConnector(resolver=_PublicResolver(), use_dns_cache=False, limit=4)
timeout = aiohttp.ClientTimeout(total=_DOWNLOAD_TIMEOUT_SECONDS, connect=10, sock_read=15)
try:
async with aiohttp.ClientSession(connector=connector, timeout=timeout, trust_env=False,
cookie_jar=aiohttp.DummyCookieJar()) as session:
for redirect_count in range(MAX_REDIRECTS + 1):
async with session.get(current_url, allow_redirects=False) as response:
if response.status in {301, 302, 303, 307, 308}:
location = response.headers.get("Location")
if not location or redirect_count >= MAX_REDIRECTS:
raise ValueError("Media artifact redirect is invalid")
current_url = _validate_artifact_url(urljoin(current_url, location))
continue
if response.status != 200:
raise ValueError("Media artifact download failed")
payload = bytearray()
async for chunk in response.content.iter_chunked(_CHUNK_SIZE):
if not isinstance(chunk, bytes) or len(payload) + len(chunk) > max_bytes:
raise ValueError("Media artifact exceeds the permitted size")
payload.extend(chunk)
return bytes(payload)
except Exception:
# Signed URLs and transport exception details must not reach the ledger.
raise ValueError("Media artifact download failed or exceeded its byte limit") from None
raise ValueError("Media artifact download failed")
async def collect_typed_results(output, media, manifest, *, download=None):
"""Decode every binding before returning any collector's native value lists."""
if not isinstance(output, dict) or output.get("error") or output.get("errors"):
raise ValueError("Invalid serverless media output")
decoded = await contract.decode_response(output.get("media"), request_id=media["request_id"],
bindings=manifest,
download=download if download is not None else download_public_artifact)
grouped = {}
for binding in manifest:
values = decoded[binding["binding_id"]]
if not isinstance(values, list) or not values:
raise ValueError("Missing serverless media values")
grouped.setdefault(binding["job_id"], {}).setdefault(binding["socket"], []).extend(values)
return grouped
+120 -35
View File
@@ -1,4 +1,5 @@
import asyncio
import copy
import time
import uuid
@@ -37,7 +38,8 @@ 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
prompt_server = server.PromptServer.instance
@@ -85,6 +87,10 @@ def _resolve_enabled_workers(config, requested_ids=None):
if worker.get("type") == "serverless":
from ..utils.serverless_config import validate_serverless_worker
pool = validate_serverless_worker(worker)
if pool.get("media_profile") == "comfy-media-v1":
from ..serverless_bridge.contract import identifier
# Do not repair malformed media identities by stripping them.
identifier(pool["id"])
for slot in range(pool["jobs_per_run"]):
workers.append(dict(pool, pool_id=worker_id, id=f"serverless:{worker_id}:{slot + 1}",
name=f"{pool.get('name', worker_id)} / job {slot + 1}"))
@@ -154,6 +160,32 @@ def _prompt_requests_load_balance(prompt_index):
return False
def _validate_media_workers(workers, prompt_index=None):
"""Reject unknown serverless media profiles before probing or dispatching.
A Distributed Collector accepts direct workers, legacy-image-v1 pools and
comfy-media-v1 pools in the same run; each pool's profile selects its
endpoint contract. Unknown profiles fail rather than guessing a contract.
legacy-image-v1 endpoints return 8-bit files, so a Collector set to exact
cannot use them.
"""
exact = prompt_index is not None and any(
str(prompt_index.inputs_by_node.get(node_id, {}).get("transfer_format", "")).strip().lower() == "exact"
for node_id in prompt_index.nodes_for_class("DistributedCollector")
)
for worker in workers:
if worker.get("type") != "serverless":
continue
profile = worker.get("media_profile", "legacy-image-v1")
if profile not in ("legacy-image-v1", "comfy-media-v1"):
raise ValueError("Unsupported serverless media_profile")
if exact and profile == "legacy-image-v1":
raise ValueError(
"transfer_format exact requires comfy-media-v1 serverless pools; "
"legacy-image-v1 endpoints return 8-bit images"
)
async def _prepare_worker_payload(
worker,
prompt_index,
@@ -169,11 +201,33 @@ async def _prepare_worker_payload(
):
"""Prepare one worker prompt payload with bounded concurrency and media-sync timeout."""
async with worker_prep_semaphore:
_validate_media_workers([worker], prompt_index)
worker_prompt = prompt_index.copy_prompt()
if worker.get("type") == "serverless":
from .orchestration.runpod_client import serialize_request
if worker.get("media_profile", "legacy-image-v1") == "comfy-media-v1":
from ..serverless_bridge.contract import PROFILE, VERSION
from .orchestration.serverless_media_prompt import compile_media_prompt
from .orchestration.serverless_typed_inputs import prepare_media_inputs
# The typed compiler validates the complete private graph and
# prunes it itself, retaining every connected Collector branch.
worker_prompt = apply_participant_overrides(
worker_prompt, worker["id"], enabled_ids, job_id_map,
"", delegate_master, prompt_index,
)
request_id = str(uuid.uuid4())
workflow, manifest = compile_media_prompt(worker_prompt, request_id)
images, artifacts, refs = await prepare_media_inputs(workflow)
media = {
"version": VERSION, "profile": PROFILE, "request_id": request_id,
"artifacts": artifacts, "inputs": refs, "outputs": copy.deepcopy(manifest),
}
serialize_request(workflow, images, worker["timeout_seconds"], media=media)
return worker, {"worker": worker, "workflow": workflow, "manifest": manifest,
"images": images, "media": media}
from .orchestration.serverless_prompt import compile_serverless_prompt
from .orchestration.serverless_media import prepare_images
from .orchestration.runpod_client import serialize_request
worker_prompt = apply_participant_overrides(
prune_prompt_for_worker(worker_prompt), worker["id"], enabled_ids,
job_id_map, "", delegate_master, prompt_index,
@@ -265,42 +319,53 @@ async def orchestrate_distributed_execution(
media_sync_timeout_seconds,
) = _resolve_orchestration_limits(config)
requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None
prompt_index = PromptIndex(prompt_obj)
image_list_start_ids = prompt_index.nodes_for_class("DistributedImageListStart")
image_list_mode = bool(image_list_start_ids)
image_list_mode = any(isinstance(node, dict) and node.get("class_type") == "DistributedImageListStart"
for node in prompt_obj.values())
target_kwargs = {}
if image_list_mode:
# 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)
other_distributed = (
prompt_index.nodes_for_class("DistributedCollector")
+ prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
)
if other_distributed:
# Check the source graph before execution-root pruning can remove an
# incompatible mode (or mask it with an unrelated boundary error).
if any(isinstance(node, dict) and node.get("class_type") in (
"DistributedCollector", "UltimateSDUpscaleDistributed") for node in prompt_obj.values()):
raise ValueError(
"Distributed Image List v1 cannot be mixed with Distributed Collector or "
"Ultimate SD Upscale nodes in one prompt"
)
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),
)
prompt_obj, outer_targets = prepare_image_list_execution(prompt_obj)
target_kwargs["partial_execution_targets"] = outer_targets
prompt_index = PromptIndex(prompt_obj)
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)
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)
_validate_media_workers(workers, prompt_index)
if any(worker.get("media_profile") == "comfy-media-v1" for worker in workers):
from ..serverless_bridge.contract import identifier
# Validate the source, not PromptIndex's JSON copy (which coerces
# integer keys to strings), before probing or dispatching workers.
for node_id in prompt_obj:
identifier(node_id)
load_balance_requested = _prompt_requests_load_balance(prompt_index)
trace_debug(
execution_trace_id,
@@ -336,13 +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"
)
idle_ids = (await preflight_image_list_workers(active_workers, require_idle=delegate_master)
if active_workers else {})
active_workers = [worker for worker in active_workers if worker["id"] in idle_ids]
if not active_workers and delegate_master:
raise RuntimeError("No feature-compatible direct worker is idle in orchestrator-only mode")
from ..image_list.runtime import END_CLASS, end_override_names
requires_values = any(end_override_names(node.get("inputs", {}))
for node in prompt_obj.values() if node.get("class_type") == END_CLASS)
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
@@ -409,6 +481,7 @@ async def orchestrate_distributed_execution(
workflow_meta,
client_id,
include_queue_metadata=True,
**target_kwargs,
)
return (
queue_result["prompt_id"],
@@ -477,6 +550,15 @@ async def orchestrate_distributed_execution(
manager = None
if serverless_jobs:
from .orchestration.serverless_runtime import get_manager
# Repeat the manager's whole-batch identity/profile preflight *before*
# any direct dispatch, master queueing or collector queue allocation.
# Pure legacy jobs are a no-op in that validator; keep their path free
# of typed imports (and compatible with existing legacy media adapters).
for prepared in serverless_jobs:
if (prepared["worker"].get("media_profile", "legacy-image-v1") != "legacy-image-v1"
or prepared.get("media") is not None):
from .orchestration.serverless_typed_results import validate_prepared_job
validate_prepared_job(prepared)
manager = get_manager()
manager.ensure_capacity(len(serverless_jobs))
@@ -506,6 +588,7 @@ async def orchestrate_distributed_execution(
workflow_meta,
client_id,
include_queue_metadata=True,
**target_kwargs,
)
except Exception:
if serverless_jobs:
@@ -522,8 +605,10 @@ async def orchestrate_distributed_execution(
await manager.cancel(execution_id=execution_trace_id)
async with prompt_server.distributed_jobs_lock:
for prepared in serverless_jobs:
for binding in prepared["manifest"]:
q = prompt_server.distributed_pending_jobs.get(binding["job_id"])
# Typed manifests have several sockets per collector, but
# each participant gets only one terminal collector event.
for job_id in {binding["job_id"] for binding in prepared["manifest"]}:
q = prompt_server.distributed_pending_jobs.get(job_id)
if q is not None:
q.put_nowait({"worker_id": prepared["worker"]["id"], "is_last": True,
"serverless_error": "Serverless job setup failed; no automatic retry was attempted"})
+13
View File
@@ -61,6 +61,13 @@ async def worker_ws_endpoint(request):
})
continue
if not isinstance(data, dict):
await ws.send_json({
"type": "dispatch_ack", "request_id": None, "ok": False,
"error": "Expected a websocket message object.",
})
continue
if data.get("type") != "dispatch_prompt":
await ws.send_json({
"type": "dispatch_ack",
@@ -81,10 +88,16 @@ async def worker_ws_endpoint(request):
continue
try:
target_kwargs = {}
if "partial_execution_targets" in data:
from ..utils.async_helpers import validate_execution_targets
target_kwargs["partial_execution_targets"] = validate_execution_targets(
prompt, data["partial_execution_targets"])
prompt_id = await queue_prompt_payload(
prompt,
workflow_meta=data.get("workflow"),
client_id=data.get("client_id"),
**target_kwargs,
)
await ws.send_json({
"type": "dispatch_ack",
+29 -4
View File
@@ -141,13 +141,16 @@ Also, for collector-based flows:
## Endpoint: `POST /distributed/job_complete`
Submit one completed worker image back to the master collector queue.
Submit direct-worker image/audio results to the master Collector queue. The route
accepts canonical JSON completions or multipart image chunks, including opt-in
exact floating-point IMAGE transport. Serverless typed results use their separate
bridge/provider contract.
### URL
- `http://<master-host>:<master-port>/distributed/job_complete`
### Request Body
### JSON request body
```json
{
@@ -164,13 +167,35 @@ Submit one completed worker image back to the master collector queue.
- `job_id` (string, required)
- `worker_id` (string, required)
- `batch_idx` (integer >= 0, required)
- `image` (string, required)
- `image` (string, optional if `audio` is present)
- PNG payload as either:
- data URL: `data:image/png;base64,...`
- raw base64 PNG bytes
- `is_last` (boolean, required)
Legacy multipart/tensor payload formats are no longer accepted on this endpoint.
At least one of `image` or `audio` must be present. Optional `audio` is an object
with `sample_rate` (positive integer), `shape` (`[batch, channels, samples]`),
`dtype: "float32"`, and `data` (base64 waveform bytes). Audio-only completions
still include `batch_idx`; the Collector uses `0` and `is_last: true`.
### Multipart image chunks
Send `multipart/form-data` with a JSON `envelope` field and repeated binary
`media` file fields. The envelope contains `job_id`, `worker_id`, `format`,
`start_index` (nonnegative integer), `count`, `width`, `height`, `is_last`, and
optional canonical `audio`.
- `png`, `jpeg`, `webp`: one `media` part per frame.
- `ffv1`, `h264`, `h265` (experimental), `av1`: one video part per chunk.
- `exact`: one raw tensor part per frame and an envelope `descriptors` list with
one `tensor-le-v1` IMAGE descriptor per part. Each descriptor shape must be
`[1, height, width, channels]`; decoding preserves supported float dtype and
values instead of quantizing to 8-bit.
`legacy_png` uses JSON, not multipart. Exact transport requires support on both
sides and never retries as an 8-bit format when an older master rejects it.
See [result transfer](compressed-result-transfer.md) for format controls,
non-exact compatibility fallbacks and request limits.
### CORS note
+36 -25
View File
@@ -1,38 +1,49 @@
# Compressed Result Transfer
# Compressed and Exact Result Transfer
**Distributed Collector** has two optional widgets that control how workers send their images back to the master. This matters most when workers run in the cloud and you pay for data leaving it (egress).
**Distributed Collector** (`DistributedCollector`) collects images and/or audio. Its `transfer_format` and `transfer_quality` widgets control **remote image returns only**; they do not re-encode master-local images or change audio transport. The saved defaults remain `legacy_png` and quality `90`.
| `transfer_format` | Type | Notes |
| `transfer_format` | Fidelity | Direct-worker transport |
|---|---|---|
| `legacy_png` (default) | lossless | Original behaviour: one JSON request per frame, uncompressed PNG in base64. Works with any master version. |
| `png` | lossless | Compressed PNG, binary, several frames per request. Same tensors as `legacy_png`. |
| `ffv1` | lossless | Whole batch as one FFV1 video. Same tensors as `legacy_png`; best lossless choice for frame sequences. |
| `jpeg`, `webp` | lossy | Per-frame stills. Good for image batches. |
| `h264`, `h265`, `av1` | lossy | Whole batch as one video. By far the smallest for video frame batches. |
| `legacy_png` (default) | 8-bit pixels | Original behaviour: one JSON request per frame, uncompressed PNG in base64. Compatible with legacy masters. |
| `png` | Lossless 8-bit pixels | Compressed PNG, binary, several frames per request. |
| `ffv1` | Lossless 8-bit pixels | FFV1 video chunks; useful for frame sequences. |
| `jpeg`, `webp` | Lossy 8-bit images | Per-frame stills. |
| `h264`, `h265` (**experimental**), `av1` | Lossy 8-bit images | Video chunks; useful for similar consecutive frames. |
| `exact` (opt-in) | Original floating-point tensors | Binary multipart chunks with one raw tensor payload and `tensor-le-v1` descriptor per frame. No 8-bit quantization. |
`transfer_quality` (1–100, default 90) applies to the lossy formats. JPEG/WebP use it directly; video codecs map it to CRF (90 → CRF 13 for h264/h265, CRF 16 for av1; 100 → CRF 0).
“Lossless” PNG/FFV1 means lossless **after 8-bit quantization**, not preservation of the source float tensor. Direct non-exact returns decode to RGB float32. `exact` preserves supported floating-point dtype, shape, channels and values: float16, bfloat16, float32 or float64, finite dense BHWC tensors with 1, 3 or 4 channels. It generally costs substantially more bandwidth.
Measured on a 121-frame 1280×720 batch (one worker's result):
`transfer_quality` (1–100, default 90) applies to lossy formats, not `exact`, PNG or FFV1. JPEG/WebP use it directly; video codecs map it to CRF (90 → CRF 13 for h264/h265, CRF 16 for av1; 100 → CRF 0). A quality of 100 is not a substitute for exact float transport.
## Serverless profiles
- **`comfy-media-v1`:** the same Collector supplies the endpoint's image-return policy. The saved `legacy_png` default is mapped to compressed **`png`** with the same 8-bit pixel precision; normal serverless returns are not exact by default. All other formats retain their selection, including opt-in `exact` and experimental `h265`. Missing codecs or rejected formats fail instead of selecting another codec. Typed image returns preserve channel count and carry alpha separately where needed. See [serverless images and audio](serverless-media.md).
- **`legacy-image-v1`:** retains its image-file/`SaveImage` contract, not the typed return-format controls. Selecting a legacy pool when any Collector requests `exact` is rejected before dispatch; there is no 8-bit fallback. Audio requires direct workers or `comfy-media-v1`.
## Direct-worker compatibility and limits
- Choose the format on the master; the value is sent to workers with the workflow.
- **`exact` requires support on both master and worker.** If an older master rejects an exact chunk, the worker raises an error and never resends it as `legacy_png` or another 8-bit format. An exact Collector also rejects direct-worker image returns labelled with a non-exact format. Update both sides or explicitly select a non-exact format.
- For **non-exact compressed** formats only, the specific first-request older-master rejection (`400` with `Invalid JSON payload`) causes a retry using `legacy_png`. Other rejections, including HTTP 413, remain errors.
- For direct video returns only, an unavailable encoder falls back to `h264` for lossy video if available, otherwise `png`; unavailable `ffv1` falls back to `png`. The worker logs this. This compatibility fallback does **not** apply to `exact` or typed serverless returns.
- Multipart requests use a 32 MiB chunk target (`COMFYUI_DISTRIBUTED_TRANSFER_CHUNK_BYTES`) and a 96 MiB request budget (`COMFYUI_DISTRIBUTED_TRANSFER_MAX_REQUEST_BYTES`), reserving space for the envelope and accounting separately for audio. A single large frame may exceed the chunk target but must fit the request budget. Keep the request budget below the master's `--max-upload-size` (default 100 MB).
- Non-exact compressed chunks are all encoded before sending. Exact transfer instead plans frame sizes from shape/dtype before sending and encodes one chunk at a time. An oversized individual frame is rejected before any chunks are sent; choose another format or raise the worker and master limits together.
- Audio is attached to the final image chunk when it fits, otherwise sent separately. Audio-only completions use the canonical JSON envelope. Direct audio uses float32 waveform bytes regardless of the image format; `exact` is an IMAGE option, not a new audio codec.
- Each worker logs the format and bytes sent. These request limits are for direct multipart transfers; serverless request/artifact limits are documented separately.
## Previously documented size example
The following measurements are retained from a 121-frame 1280×720 slow-pan batch (one worker's result); they are not new validation of the current implementation or exact transport.
| Format | Sent | vs legacy | PSNR |
|---|---|---|---|
| `legacy_png` | 446 MB | 1× | lossless |
| `png` | 163 MB | 2.7× | lossless |
| `ffv1` | 128 MB | 3.5× | lossless |
| `legacy_png` | 446 MB | 1× | lossless at 8-bit precision |
| `png` | 163 MB | 2.7× | lossless at 8-bit precision |
| `ffv1` | 128 MB | 3.5× | lossless at 8-bit precision |
| `jpeg` q90 | 27 MB | 17× | 42.7 dB |
| `webp` q90 | 13 MB | 33× | 40.6 dB |
| `h264` q90 | 2.4 MB | 186× | 40.2 dB |
| `h265` q90 | 1.0 MB | 434× | 40.1 dB |
| `h265` q90 (experimental) | 1.0 MB | 434× | 40.1 dB |
| `av1` q90 | 0.4 MB | 1031× | 40.5 dB |
The test clip is a slow pan, which video codecs compress especially well; expect smaller (still large) gains on content with more motion.
Notes:
- Choose the format on the master; the value is sent to workers with the workflow.
- The master must run this version to accept compressed results. A worker sending to an older master gets a 400 on its first request and automatically resends with `legacy_png`.
- If a worker's PyAV build lacks the requested video encoder, it falls back to `h264` (lossy formats) or `png` (`ffv1`), and logs it.
- Lossy formats only affect frames produced by workers; the master's own frames are never re-encoded.
- Request size is bounded on the final encoded bytes. Frames are batched into requests of up to 32 MiB (`COMFYUI_DISTRIBUTED_TRANSFER_CHUNK_BYTES`); a single large frame may go alone in a request of up to 96 MiB (`COMFYUI_DISTRIBUTED_TRANSFER_MAX_REQUEST_BYTES`). Keep the latter below the master's `--max-upload-size` (default 100 MB) if you change either.
- All chunks are encoded before anything is sent. If a single frame cannot fit within the request limit, the worker fails with an error and sends nothing; use a lossy format or raise both limits.
- A rejection other than the older-master case above (for example the master's size limit, HTTP 413) is reported as an error rather than retried as `legacy_png`.
- Each worker logs what it sent, e.g. `Sent 121 image(s) as h264 in 1 request(s): 2.40 MB (raw 334.49 MB, 139.4x smaller)`.
Video codecs compress this slow pan especially well; gains depend on content and motion.
+137 -19
View File
@@ -1,10 +1,10 @@
# 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. It 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.
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.
## Wiring
@@ -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.
@@ -52,30 +52,141 @@ Start outputs, in order: `image`, `mask`, `item_index` (zero-based), `context`,
Values that cannot be transported exactly (models, latents, conditioning and other objects) are rejected. Load models inside the region from a loader connected directly to the processing nodes, as described below.
End accepts `images`, optional `masks` / `texts`, and `context`. It returns `image_list`, `masks`, and `texts` in original item order. Each job must produce exactly one image, at most one mask, and at most one text value. Unconnected optional outputs contain `None` entries.
End accepts `images`, optional `masks` / `texts`, and `context`. Its first three outputs remain `image_list`, `masks`, and `texts` in original item order. Each job must produce exactly one image, at most one mask, and at most one text value. Unconnected optional mask/text outputs contain `None` entries.
Models are **loaded and cached on each processing participant**, not serialized between machines. Connect loader outputs directly to the processing nodes. The master needs the region's model files when participating, as does each selected worker. **The master still validates the prompt:** it needs the relevant node definitions and any model-name/file availability required by their validation (including combo lists). This is not a model-free-master guarantee.
### Keep values through End without extra wires
Start's existing **context** connection now also lets End recover each item's original `value_1` … `value_10`. You do not need to wire those values individually into End. Context remains a lightweight execution reference, not a general-purpose data pipe.
End appends a **values** output (`DISTRIBUTED_IMAGE_LIST_VALUES`): one value bundle per image, in the same source order. Connect it to **Distributed Image List Get Value**, then choose `value_name` to extract the field you need. The getter preserves its type and ComfyUI maps it across the bundles, producing a list aligned with End's images. Missing fields fail clearly instead of producing empty placeholders.
For example, to preserve filenames:
```text
filenames → Start.value_1
Start.context ─────────────→ End.context
Start.image → processing → End.images
├─ image_list → Save.images
└─ values → Get Value [value_1] → Save.filename_prefix
```
**New or changed per-item values:** connect the calculated result to End's optional `value_N` input. It replaces the original field with that number, or adds it if Start did not supply it. Every connected override must produce exactly one supported value for each item. A tensor batch remains one value, just as it does at Start. Image, mask, string, integer, float and boolean values are supported; arbitrary objects/models/latents and non-finite numbers are not.
End shows override inputs progressively. Use **Add value override** to expose a later slot without connecting the earlier ones. Disconnecting an override never renumbers the other End fields: `value_3` always means `value_3`. Unlike Start's native Autogrow inputs, End overrides must keep these identities stable. Existing image/mask/text/context sockets and output indices are unchanged. If the frontend extension is unavailable, Python still declares all optional override inputs and the values output.
Original values are retained on the master; workers only return explicitly connected overrides, avoiding a redundant return trip for unchanged reference images or other values. Overrides use the existing exact value transport and size limits. Selected workers must advertise `value_results: true` when overrides are requested, checked before execution and again at dispatch. Original-value passthrough alone does not require that additional capability.
Values leaving the region must come from **End**, not directly from Start. To chain regions, extract an End value and connect it to the next Start's corresponding input. Values remain available downstream after End cleans up the execution context; they do not depend on a later lookup of expired job state.
Models are loaded and cached on each processing participant, not serialized between machines. Connect loader outputs directly to the processing nodes. The master needs the region's model files when participating, as does each selected worker. The master still validates the prompt: it needs the relevant node definitions and any model-name/file availability required by their validation (including combo lists). This is not a model-free-master guarantee.
Per-image data must pass through Start; do not bypass it with upstream images, latents or generated batches connected directly into the processing region. Shared dependencies across Start are restricted to registered model outputs or scalar constants, not arbitrary image-producing leaf nodes.
End's optional `texts` is a connection-only socket. For workflows saved before this fix, 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.
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 while being 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
Repeated items can reuse **ComfyUI's native processing cache** on the same participant. Start uses a deterministic identity of the exact image, optional mask, all supplied value slots and source index; tensor dtype, shape and bytes are preserved. Fresh job IDs, assignment tokens and callback addresses live at End, so they no longer invalidate otherwise unchanged upstream processing. End still validates the current assignment and delivers results in source order on every run.
- Processing settings, seeds, links and node fingerprints still participate in ComfyUI's normal signatures. Changed content invalidates reuse; nodes that declare themselves noncacheable still run.
- The master's private executor retains one native LRU output cache between successful slots, capped at **128 entries when retained**. An active graph may temporarily exceed that target; oversized generations are not retained. Failed or cancelled local slots do not seed reuse, and `--cache-none` disables the local output cache. Direct workers use their own configured native caches.
- Cache hits are best-effort, not guaranteed: eviction, restarts, changed ordering/index, changed settings and assignment to a different worker can cause misses. There is no worker affinity, disk cache or separate completed-result store.
- Reuse retains tensor/model references and can increase RAM/VRAM use; the entry cap is **not a byte limit**. Exact identity calculation also adds CPU hashing and, for GPU-resident inputs, a CPU copy.
- Remote items are still uploaded. Each upload returns a worker-computed identity; if a compatible older worker does not return a valid identity, that assignment uses the legacy graph without stable-input reuse. A worker downgrade between items also falls back safely. Assignment expiry and current-item validation apply even when processing hits the cache.
- Keep side-effect nodes outside the region when they must run on every submission. A cached processing node follows ordinary ComfyUI semantics and does not repeat its save/preview side effects; End's delivery and downstream saving still run.
## Scope and failure behavior
- This version supports the master and direct workers, not Runpod serverless pools. Unsupported pools are excluded without contacting the provider for this job.
- Worker eligibility is determined before execution and checked again before dispatch. Busy or transiently unreachable workers are excluded; healthy candidates can continue. Malformed capability/queue responses and incompatible versions fail explicitly. Workers joining later are not added to an already-running region.
- 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.
- Output nodes inside a region are rejected because ComfyUI would execute them independently on the master. Avoid other side-effecting processing nodes as well; save after End.
- 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 do not incur those HTTP limits.
- 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 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. 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.
- 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.
- 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.
- 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
@@ -89,7 +200,14 @@ Use trusted workers and trusted workflows, as with the existing ComfyUI API. Cap
## Verification
The generic acceptance helper starts two isolated loopback worker test servers and uses actual CPU ComfyUI `PromptExecutor` instances. Production queue preparation, list nodes, graph extraction, HTTP item/result routes, dynamic scheduling, float transport, ordering and cancellation are exercised. The harness checks orchestrator-only, mixed master/worker and master-only execution, including exact-once item assignment, source-once materialization and cached local model-loader dependencies. The lightweight queue/history facade is test code; no managed ComfyUI service or model is started.
The generic acceptance helper starts two isolated loopback workers using real ComfyUI HTTP handlers, `PromptQueue` records and CPU `PromptExecutor` instances. Production queue preparation, list nodes, graph extraction, HTTP item/result routes, dynamic scheduling, float transport, ordering and cancellation are exercised. The harness checks orchestrator-only, mixed master/worker and master-only execution, including exact-once required output utilities, discarded previews, unchanged source graphs, End-only worker roots, post-End capture, source-once materialization and cached local model-loader dependencies. WebSocket target propagation is covered by focused route/dispatch tests. No managed ComfyUI service or model is started.
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
+187
View File
@@ -0,0 +1,187 @@
# Serverless images and audio
**Distributed Collector** (`DistributedCollector`) collects `IMAGE` and/or `AUDIO`
from direct workers and serverless pools. There is no separate
`DistributedMediaCollector`; replace that removed node with `DistributedCollector`
for image/audio workflows. The Collector has no MASK, STRING, or native VIDEO
output sockets.
`comfy-media-v1` is an **explicit endpoint profile**, not an automatically detected
capability. It returns typed image/audio items to the Collector. The standalone
bridge's generic contract also supports `MASK`, `STRING`, and native ComfyUI
`VIDEO`, but that does not add those outputs to the Collector. Omitting
`media_profile` still selects the image-only `legacy-image-v1` contract.
See [legacy serverless setup](serverless-workers.md) for credentials and recovery.
## Use the Collector with a typed endpoint
1. Prepare an endpoint with the [standalone media bridge](../serverless_bridge/README.md)
and the workflow's models/custom nodes. The endpoint needs only the bridge,
**not the full ComfyUI-Distributed extension**.
2. In the serverless worker settings, set **Endpoint media profile** to
**Typed media v1 (requires endpoint bridge)** (`media_profile: "comfy-media-v1"`).
Start with **Images per run = 1**. This existing setting counts remote jobs,
not the number of typed items each job returns.
3. Use **Distributed Collector** (`DistributedCollector`), connect `images`,
`audio`, or both, and connect the corresponding outputs downstream. Image-only
and audio-only workflows use the same node.
4. Queue through Distributed. **Delegate Only** excludes local upstream
generation; normal collection includes the master's local values.
**Test Saved Endpoint** checks provider authentication/health, not bridge
installation, model availability, or typed execution. Selecting a profile does
not install or upgrade anything on the endpoint.
The Collector accepts optional `images` and `audio` inputs, including ComfyUI list
inputs, and returns two ordinary outputs—not list outputs:
| Socket / output index | Native type | Aggregation |
| --- | --- | --- |
| `images` / 0 | `IMAGE` | Concatenate BHWC tensors along the batch dimension |
| `audio` / 1 | `AUDIO` | Concatenate waveforms along the sample dimension |
For each connected socket, order is **master first, then prepared participant
order, then that participant's item order**—never completion order. Delegate-only
mode omits the master. Image dimensions/channels and audio batch/channel dimensions
must be compatible for concatenation. Use a common audio sample rate; the Collector
does not resample or align streams. Audio-only collection returns `None` for images;
when no audio is supplied it returns a silent placeholder, not an empty list.
The typed wire contract carries 1–64 items per connected socket. The entire remote
response is validated and decoded before any of its results are delivered;
missing, duplicate, mismatched, or corrupt typed results fail collection. Validated
wire lists are then combined into the Collector's ordinary outputs.
Direct workers, `legacy-image-v1` pools and `comfy-media-v1` pools can share an
image-only Collector run. Audio requires direct workers or `comfy-media-v1` pools.
If any Collector requests `exact`, a selected `legacy-image-v1` pool is rejected
before dispatch instead of silently returning 8-bit images.
## Image return format and quality
The Collector's saved `transfer_format` default remains **`legacy_png`**, with
`transfer_quality` **90**. For `comfy-media-v1`, the compiler maps `legacy_png` to
**`png`**, so normal serverless returns use compressed 8-bit PNG. `exact` is opt-in.
These controls affect **remote IMAGE returns only**, not master-local images,
uploaded inputs or AUDIO. The bridge's generic MASK, STRING and native VIDEO
codecs are also unaffected.
- `exact`: preserve supported floating-point dtype and values, with no image
quantization. Available for direct workers and `comfy-media-v1` pools; it is
never downgraded to an 8-bit format. An older master that rejects exact transfer
causes an error, not a legacy retry. Update both sides before selecting it.
- `legacy_png`, `png`, `ffv1`: lossless **8-bit image pixels**, not exact float
tensors. For typed serverless returns, image values are clipped to [0,1] and
quantized; `legacy_png` is sent as `png`. Direct workers retain the original
JSON/base64 framing for `legacy_png`.
- `jpeg`, `webp`: lossy still-image compression.
- `h264`, `h265` (**experimental**), `av1`: lossy video encoding of each IMAGE
tensor's frame batch. This is an image-return encoding, not a native VIDEO output.
- Quality 1–100 affects lossy formats only, using the existing Collector's
quality-to-CRF mapping for video. Dimensions, frame order and channel count
are preserved for typed serverless returns; alpha is carried separately where
needed, not silently dropped.
Both master and endpoint must have this bridge revision for compressed image
returns. Pillow handles still formats and CPU/software PyAV encoders handle video;
unavailable codecs fail clearly rather than silently changing the selection.
Compression can reduce result egress and inline response size, but does not relax
encoded/decoded byte budgets or the input request limit. Compressed IMAGE batches
are capped at 240 frames and 8192 pixels per dimension, with a conservative
256 MiB aggregate decode budget (including codec scratch and alpha). Direct
workers have different limits and non-exact compatibility fallbacks; see
[result transfer](compressed-result-transfer.md). The bridge's internal
`MediaExportV1` default is still `exact` for generic callers that omit controls;
the Collector compiler explicitly supplies `png` for the normal saved default.
## Inputs and fidelity
Automatic local-file transfer supports stock `LoadImage`, plus explicit stock
`LoadAudio` (`audio` filename) and `LoadVideo` (`file` filename) adapters. Audio and
video loaders require a literal filename, no extra controls, and output slot 0.
Files must be under ComfyUI's input/output/temp roots. Custom file loaders are not
supported; detected unsupported file references fail rather than silently
becoming endpoint-local paths. Models and custom nodes are never transferred.
- `LoadImage` uses the existing bounded file/base64 image staging path.
- `LoadAudio` is decoded on the master using native PCM semantics with bounded
accumulation, then transferred as a waveform and sample rate. It does not
resample; integer PCM conversion follows ComfyUI's `f32_pcm` behavior.
- `LoadVideo` transfers opaque file bytes as native `VideoFromFile`, without
decoding frames or transcoding the container.
- Typed tensor transport is `tensor-le-v1`: contiguous little-endian float16,
bfloat16, float32, or float64 bytes with exact dtype/shape/layout metadata.
Only dense, finite tensors with positive dimensions are supported. IMAGE has
1, 3, or 4 channels. Bridge export explicitly detaches/stages tensors to CPU;
the pure codec itself rejects non-CPU storage. Remote decoding returns
independently owned CPU values. STRING transport is exact UTF-8 (`utf8-v1`).
### Bridge-level native VIDEO support
The generic bridge can import/export native VIDEO, and the typed input adapter
can stage `LoadVideo` for upstream processing. This is not a Collector VIDEO
output or a generic video converter.
The adapter supports the pinned native classes **exactly**, not arbitrary
subclasses or objects exposing `save_to()` / `get_components()`:
- `VideoFromComponents`: images, optional matching BHW or BHW1 alpha, optional
audio, positive rational frame rate, bounded JSON metadata, positive integer
bit depth, and `sRGB`, `HDR`, or `HDR PQ` color space are preserved.
- `VideoFromFile`: regular local files within allowed media roots or an exact
`BytesIO`; opaque bytes, finite signed start/duration, and optional four-integer
crop are preserved without normalization.
- `VideoFromList`: 1–64 flat exact `VideoFromFile` segments, plus optional
`complete_audio`, in order.
Unknown classes, incompatible native state, nested/component segment lists,
unsupported tensors, invalid metadata, and oversized payloads are rejected.
There is no fallback transcode, frame extraction, color conversion, or container
repair/validation. A downstream native consumer still needs to read the original
container. Decoded file-backed VIDEO values own independent `BytesIO` storage;
component tensors also own their storage, so staging-file cleanup cannot
invalidate a returned VIDEO.
## Size and operational limits
- The typed contract caps each artifact at **256 MiB**, aggregate artifact bytes
at **512 MiB**, and each artifact list at **128** entries. These are validation
ceilings, not a promise that inline requests can carry that much data.
- The built-in master packs inputs inline and enforces **10,000,000 bytes for the
complete serialized Runpod `/run` request**, including workflow, policy and
base64. Legacy image reads also retain their 10,000,000-byte aggregate guard.
There is **no built-in master-side large-input uploader** in PR2.
- The endpoint defaults to inline results and caps its result JSON at
**10,000,000 bytes**. Larger results require an operator-supplied publisher
returning expiring public HTTPS artifact URLs. See the bridge guide for
bounded download/upload responsibilities and retention.
- The generic five-kind typed envelope is a serverless bridge contract. Direct
workers support Collector images/audio, including opt-in exact IMAGE transport,
not generic MASK/STRING/VIDEO Collector outputs. The typed compiler rejects
unsupported distributed/batch nodes and collector-to-collector dependencies.
Downstream processing stays on the master. [Image-list scheduling](image-list-queue.md)
is a separate execution mode.
## Timeout is not proof that Comfy stopped
Provider cancellation, a handler timeout, or cancellation of its Python task
**does not establish that the local Comfy prompt stopped**. The bridge never
posts `/interrupt`: the pinned core route checks a queue snapshot and then
interrupts globally, so even a prompt ID does not make it race-free.
If submission/termination or callback teardown cannot be confirmed, the bridge
raises `MediaExecutionUncertain` and retains owned input/exporter files—even
after terminal Comfy history if a downloader/publisher ignores cancellation.
Callback teardown has a separate 10-second grace; late callback exit does not
clean retained files. Callbacks must be cooperative async code; synchronous
event-loop blocking cannot be hard-timed-out.
The endpoint operator must **quarantine the owning worker/environment, reconcile
the local prompt and readers or terminate that environment, and establish that
they have stopped before cleanup or worker reuse**. Do not automatically retry
the request or turn every handler exception into a safely reusable terminal
provider job. Master-side provider status/reconciliation alone is not proof
that a surviving local Comfy process or callback has stopped.
These docs describe the implementation, not a test result or live Runpod/provider
or GPU end-to-end validation.
+23 -12
View File
@@ -1,9 +1,12 @@
# Runpod serverless workers (preview)
A **Serverless (Runpod)** worker is a queue-based endpoint, not a permanent
ComfyUI machine. Runpod chooses the actual GPU workers. Distributed submits one
normal ComfyUI API workflow per requested image and collects the returned images.
The endpoint workers **do not need ComfyUI-Distributed installed**.
ComfyUI machine. Runpod chooses the actual GPU workers. This page describes the
default **`legacy-image-v1`** profile: Distributed submits normal ComfyUI API
workflows and collects returned image files. For image/audio collection and
opt-in exact IMAGE returns, use [the `comfy-media-v1` profile](serverless-media.md)
with the same **Distributed Collector** (`DistributedCollector`). The endpoint
workers do not need the full ComfyUI-Distributed extension installed.
## Setup
@@ -14,7 +17,8 @@ The endpoint workers **do not need ComfyUI-Distributed installed**.
4. Save, then choose **Test Saved Endpoint**. This checks authentication and
endpoint health only; it does not start GPU inference or prove model/node compatibility.
5. Enable the worker when you want queued Distributed workflows to use it.
6. Use an image **Distributed Collector** and queue through Distributed as usual.
6. Connect the **Distributed Collector**'s `images` input and queue through
Distributed as usual. Leave `audio` unconnected for this legacy profile.
Enable **Delegate Only** if the master should collect without generating locally.
Keys are write-only. Leave the key field blank to retain a saved key, or select
@@ -93,7 +97,7 @@ incomplete serverless batch.
- A zero-worker endpoint can be healthy: Runpod scales from zero on submission.
- Interrupting a collector or using the panel cancellation action requests
cancellation only for jobs submitted by this master. It never purges the endpoint.
- A network failure during submission is **not retried automatically**: the job
- A network failure during submission is not retried automatically: the job
may already be billable. Check the Runpod job list before queuing again.
## Persistence and recovery
@@ -107,7 +111,7 @@ available on Windows.
`distributed/serverless/jobs.json` records provider job IDs and sanitized states,
not image data, workflow payloads, or API keys. A restarted master cannot resume
an interrupted collector. It does **not** automatically resubmit jobs. Retained
an interrupted collector. It does not automatically resubmit jobs. Retained
IDs are available through `/distributed/serverless/status/<worker_id>`; the worker's
cancel action can request cancellation of known pre-restart jobs. If no job ID
was received, check Runpod manually. Cancellation is best-effort and may race
@@ -127,7 +131,7 @@ Recovery uses the same trusted-master API as the panel; it never resubmits a job
every unresolved record with its local `record_id`, provider `job_id` (if
received), and state. These records remain visible if endpoint health fails.
2. Request cancellation from the panel, or send
`POST /distributed/serverless/reconcile` with the body below to **read** the
`POST /distributed/serverless/reconcile` with the body below to read the
retained provider IDs. Only a confirmed terminal state releases protection.
Network errors, expired IDs, or ambiguous responses do not.
@@ -143,7 +147,7 @@ Recovery uses the same trusted-master API as the panel; it never resubmits a job
{"worker_id": "YOUR_WORKER_ID", "record_id": "LOCAL_RECORD_ID", "provider_checked": true}
```
This records `RESOLVED`: an explicit operator acknowledgment, **not** a
This records `RESOLVED`: an explicit operator acknowledgment, not a
provider-confirmed cancellation. It cannot resolve a live local task and does
not stop remote compute. It releases the record's reservation and may let
already-waiting jobs proceed. Never automate this acknowledgment as a retry
@@ -152,9 +156,16 @@ Recovery uses the same trusted-master API as the panel; it never resubmits a job
Worker setting changes become available once all its local tasks and unresolved records
are finished. No credentials or ledger files need to be edited by hand.
## First-version limits
## Legacy profile limits
- Image collectors only. Audio/video transfer and audio collectors are rejected.
- Only the Collector's image branch is supported by `legacy-image-v1`; connected
audio is rejected. `comfy-media-v1` adds audio through the same Collector and
supports the configured image-return formats.
- `exact` is opt-in and requires direct workers or `comfy-media-v1` pools. If any
Collector requests it, a selected legacy pool is rejected before dispatch,
not used as an 8-bit fallback. Other Collector transfer-format controls do not
change this legacy profile's `SaveImage`/image-file contract. With
`comfy-media-v1`, the saved `legacy_png` default is mapped to compressed `png`.
- Coordinated `UltimateSDUpscaleDistributed` workloads and unsupported Distributed
node types are rejected before submission.
- Downstream post-collector processing remains local; the remote graph contains
@@ -168,6 +179,6 @@ are finished. No credentials or ledger files need to be edited by hand.
- No remote GPU discovery, endpoint provisioning, model synchronization, automatic
scaling changes, or inbound master callback/tunnel is needed.
This feature is verified with offline provider/transport tests. A successful
health test is not a paid end-to-end generation test; test your deployed handler
These docs describe the implementation, not a test result. A successful health
check is not a paid end-to-end generation test; validate your deployed handler
and model inventory with one image per run before increasing it.
+71
View File
@@ -0,0 +1,71 @@
"""Prepare image-list execution without changing the saved/API prompt or registry."""
from __future__ import annotations
import copy
import sys
from .runtime import (
END_CLASS,
END_RESULT_INPUTS,
_ancestors_for_inputs,
_downstream_until,
_image_list_pairs,
_is_output_node,
_prompt_links,
validate_image_list_regions,
)
def _reachable(roots: set[str], edges: dict[str, set[str]]) -> set[str]:
visited = set()
pending = list(roots)
while pending:
node_id = pending.pop()
if node_id in visited:
continue
visited.add(node_id)
pending.extend(edges[node_id] - visited)
return visited
def prepare_image_list_execution(prompt: dict) -> tuple[dict, list[str]]:
"""Return an execution copy and its outer-master output targets.
Disposable per-item side branches are removed before boundary validation.
Required output-node utilities remain normal dependencies, but must not be
independent outer roots. Callers must use these targets (and End-only targets
for item execution); this does not suppress side effects inside required nodes.
"""
pairs = _image_list_pairs(prompt)
prepared = copy.deepcopy(prompt)
upstream = {node_id: set() for node_id in prepared}
downstream = {node_id: set() for node_id in prepared}
for node_id, node in prepared.items():
for _, source, _ in _prompt_links(node, prepared):
upstream[node_id].add(source)
downstream[source].add(node_id)
# Never erase another region or its dependencies to make a bad graph pass.
boundaries = {node_id for pair in pairs for node_id in pair}
required = _reachable(boundaries, upstream)
discarded = set()
for start_id, end_id in pairs:
region_nodes = _ancestors_for_inputs(prepared, end_id, start_id, END_RESULT_INPUTS)
outside = _downstream_until(prepared, start_id, end_id) - region_nodes
# Preserve bypass joins and their whole chain into post-End computation.
# They must fail boundary validation, not disappear with debug previews.
post_end = _reachable({end_id}, downstream)
protected = required | _reachable(post_end, upstream)
discarded.update(outside - protected)
for node_id in discarded:
del prepared[node_id]
regions = validate_image_list_regions(prepared, explicit_targets=True)
per_item = set().union(*(region["branch_ids"] for region in regions))
registry = getattr(sys.modules.get("nodes"), "NODE_CLASS_MAPPINGS", {})
targets = [
node_id for node_id, node in prepared.items()
if node_id not in per_item and isinstance(node, dict)
and (node.get("class_type") == END_CLASS or _is_output_node(node, registry))
]
return prepared, targets
+32
View File
@@ -14,6 +14,38 @@ DEFAULT_ORPHAN_GRACE_SECONDS = 300.0
_LIFECYCLE_MARKER = object()
def has_safe_execution_targets(prompt_server: Any, prompt_id: str | None, prompt: dict) -> bool:
"""Prove outer roots from the running queue, not a forgeable prompt marker.
Plain /prompt remains strict when ComfyUI selected all outputs, including a
flagged per-item dependency. No saved workflow field grants this exemption.
"""
if not prompt_id:
return False
from .execution_plan import prepare_image_list_execution
prepared, targets = prepare_image_list_execution(prompt)
if prepared != prompt or not targets:
return False
queue = getattr(prompt_server, "prompt_queue", None)
for name in ("get_current_queue_volatile", "get_current_queue"):
reader = getattr(queue, name, None)
if not callable(reader):
continue
try:
running, _pending = reader()
for entry in running:
if (isinstance(entry, (tuple, list)) and len(entry) > 4
and entry[1] == prompt_id and entry[2] == prompt):
roots = entry[4]
return (isinstance(roots, (list, tuple))
and all(isinstance(root, str) for root in roots)
and len(roots) == len(targets) and set(roots) == set(targets))
return False
except Exception:
continue
return False
def _queue_prompt_ids(prompt_server: Any) -> set[str]:
"""Read running/pending prompt IDs, preferring ComfyUI's volatile view."""
queue = getattr(prompt_server, "prompt_queue", None)
+40 -1
View File
@@ -24,6 +24,34 @@ from .runtime import ImageListCancelledError
_EXECUTION_LOCK = threading.Lock()
_CACHE_LOCK = threading.Lock()
# One process-local native output cache, not a second completed-result cache.
# Core's LRU can temporarily exceed its target for the active graph; oversized
# generations are discarded rather than retained after an item/run finishes.
NATIVE_CACHE_MAX_ENTRIES = 128
_RETAINED_OUTPUTS = None
def _take_outputs(execution):
global _RETAINED_OUTPUTS
with _CACHE_LOCK:
outputs, _RETAINED_OUTPUTS = _RETAINED_OUTPUTS, None
if getattr(execution.args, "cache_none", False):
return execution.NullCache()
return outputs if outputs is not None else execution.LRUCache(
execution.CacheKeySetInputSignature, max_size=NATIVE_CACHE_MAX_ENTRIES)
def _retain_outputs(outputs):
global _RETAINED_OUTPUTS
# Only the bounded native LRU is eligible. NullCache and test doubles are not.
if getattr(outputs, "max_size", None) != NATIVE_CACHE_MAX_ENTRIES:
return
outputs.clean_unused()
if len(outputs.cache) > NATIVE_CACHE_MAX_ENTRIES:
return
with _CACHE_LOCK:
_RETAINED_OUTPUTS = outputs
class LocalExecutionSlot:
@@ -32,6 +60,7 @@ class LocalExecutionSlot:
self._stop = threading.Event()
self._future = None
self._caches = None
self._failed = False
self._model_callback = model_callback
self._closed = False
self._loop = None
@@ -55,6 +84,10 @@ class LocalExecutionSlot:
self._loop.call_soon_threadsafe(self._task.cancel)
async def close(self):
if self._closed:
return
# A stop requested before close indicates cancellation, not normal drain.
cancelled = self._stop.is_set()
self.request_stop()
try:
if self._future is not None:
@@ -72,6 +105,8 @@ class LocalExecutionSlot:
waiter.exception()
finally:
self._pool.shutdown(wait=True, cancel_futures=True)
if not cancelled and not self._failed and self._caches is not None:
_retain_outputs(self._caches.outputs)
self._caches = None
self._closed = True
@@ -84,7 +119,11 @@ class LocalExecutionSlot:
try:
asyncio.run(self._execute(prompt, prompt_id, end_id, workflow_meta))
except asyncio.CancelledError as exc:
self._failed = True
raise ImageListCancelledError("Local image-list execution cancelled") from exc
except BaseException:
self._failed = True
raise
finally:
with self._state_lock:
self._loop = self._task = None
@@ -125,7 +164,7 @@ class LocalExecutionSlot:
try:
with torch.inference_mode():
if self._caches is None:
outputs = execution.HierarchicalCache(execution.CacheKeySetInputSignature)
outputs = _take_outputs(execution)
objects = execution.HierarchicalCache(execution.CacheKeySetID)
self._caches = SimpleNamespace(outputs=outputs, objects=objects, all=[outputs, objects])
caches = self._caches
+427 -65
View File
@@ -7,14 +7,25 @@ from __future__ import annotations
import asyncio
import base64
import hashlib
import json
import math
import secrets
import struct
import threading
import time
from dataclasses import dataclass, field
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"
@@ -29,6 +40,7 @@ VALUE_OUT_BASE = 4
MAX_VALUE_SLOTS = 10
VALUE_NAMES = tuple(f"value_{i}" for i in range(1, MAX_VALUE_SLOTS + 1))
VALUES_INPUT = "values"
END_RESULT_INPUTS = ("images", "masks", "texts") + VALUE_NAMES
MAX_PAYLOAD_BYTES = 50 * 1024 * 1024
MAX_REQUEST_BYTES = MAX_PAYLOAD_BYTES * 2 + 1024 * 1024
MAX_ITEMS = 4096
@@ -179,6 +191,50 @@ def _normalize_values(values: Any, count: int) -> dict[str, list[Any]]:
return aligned
def end_override_names(inputs: dict[str, Any]) -> tuple[str, ...]:
"""Fixed keys stay stable even when an intermediate socket disconnects."""
for name in inputs:
if (name.startswith("value_") or name.startswith("values.")) and name not in VALUE_NAMES:
raise ValueError(f"Unknown End value slot {name!r}")
return tuple(name for name in VALUE_NAMES if name in inputs)
def _validate_result_values(values: Any) -> dict[str, Any]:
"""Validate flat override values; tensors are values, never frame lists."""
if not isinstance(values, dict):
raise ValueError("End result values must be a mapping of value_N slots")
for name, value in values.items():
if name not in VALUE_NAMES:
raise ValueError(f"Unknown End value slot {name!r}")
kind = value_kind(value, name)
if kind == "TENSOR":
_tensor_dtype_name(value)
if not _torch().isfinite(value).all().item():
raise ValueError(f"End {name} tensor values must be finite")
return values
def normalize_value_overrides(values: Any) -> dict[str, Any]:
if not isinstance(values, dict):
raise ValueError("End overrides must be a mapping of value_N slots")
normalized = {}
for name, raw in values.items():
items = list(raw) if isinstance(raw, (list, tuple)) else [raw]
if len(items) != 1:
raise ValueError(f"Worker End override {name} expects exactly one value, received {len(items)}")
normalized[name] = items[0]
return _validate_result_values(normalized)
def _checked_overrides(result: dict[str, Any], expected: tuple[str, ...]) -> dict[str, Any]:
values = _validate_result_values(result.get("values", {}))
missing, extra = set(expected) - values.keys(), values.keys() - set(expected)
if missing or extra:
raise ValueError(f"End result overrides do not match connected slots: "
f"missing {sorted(missing)}, unexpected {sorted(extra)}")
return values
def normalize_image_list_inputs(images: Any, mask: Any = None, values: Any = None) -> list[dict[str, Any]]:
"""Split images into items and pair masks/generic values by the 1-or-N rule."""
frames = _image_frames(images, "images")
@@ -197,6 +253,64 @@ def normalize_image_list_inputs(images: Any, mask: Any = None, values: Any = Non
]
def is_valid_item_identity(identity: Any) -> bool:
return (isinstance(identity, str) and len(identity) == 71
and identity.startswith("sha256:")
and all(char in "0123456789abcdef" for char in identity[7:]))
def item_content_identity(item: dict[str, Any]) -> str:
"""Hash normalized content, not transport bytes or dispatch metadata.
Length-framed, typed fields distinguish scalars, absent masks, slot names,
shapes and exact tensor bits (including bfloat16). CPU copying changes neither
dtype nor values; local items deliberately have no HTTP transport-size limit.
"""
digest = hashlib.sha256(b"distributed-image-list-item-v1\0")
def field(tag, raw):
digest.update(tag.encode("ascii") + b"\0")
digest.update(len(raw).to_bytes(8, "big"))
digest.update(raw)
def value(raw, name):
if raw is None:
field("NONE", b"")
return
kind = value_kind(raw, name)
if kind == "TENSOR":
cpu = raw.detach().to("cpu").contiguous()
field("DTYPE", str(cpu.dtype).encode("ascii"))
field("SHAPE", json.dumps(list(cpu.shape), separators=(",", ":")).encode("ascii"))
field("TENSOR", memoryview(cpu.view(_torch().uint8).numpy()).cast("B"))
elif kind == "FLOAT":
field(kind, struct.pack("!d", raw))
elif kind == "BOOLEAN":
field(kind, b"1" if raw else b"0")
elif kind == "INT":
field(kind, str(raw).encode("ascii"))
else:
field(kind, raw.encode("utf-8", errors="surrogatepass"))
index = item["index"]
if isinstance(index, bool) or not isinstance(index, int) or index < 0:
raise ValueError("item.index must be a non-negative integer")
field("INDEX", str(index).encode("ascii"))
field("IMAGE", b"")
value(item["image"], "image")
field("MASK", b"")
value(item.get("mask"), "mask")
values = item.get("values") or {}
if not isinstance(values, dict) or any(name not in VALUE_NAMES for name in values):
raise ValueError("item.values must be a mapping of value_N slots")
# Canonical socket order, not the producer's mapping insertion order.
for name in VALUE_NAMES:
if name in values:
field("SLOT", name.encode("ascii"))
value(values[name], name)
return "sha256:" + digest.hexdigest()
def _tensor_dtype_name(tensor: Any) -> str:
name = str(tensor.dtype).removeprefix("torch.")
if name not in _ALLOWED_FLOAT_DTYPES:
@@ -257,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")
@@ -346,26 +462,64 @@ def _sweep_worker_items() -> None:
del _WORKER_ITEMS[token]
def store_worker_item(token: str, item: dict[str, Any]) -> None:
def store_worker_item(token: str, item: dict[str, Any]) -> str:
if not isinstance(token, str) or len(token) < 32 or len(token) > 128:
raise ValueError("token must be a 32-128 character capability")
clean = decode_item(item)
# Never trust a master-supplied digest; hash the decoded worker-owned item.
identity = item_content_identity(clean)
_sweep_worker_items()
with _LOCK:
if token in _WORKER_ITEMS:
raise ValueError("duplicate item token")
_WORKER_ITEMS[token] = (time.monotonic() + _WORKER_ITEM_TTL_SECONDS, clean)
return identity
def consume_worker_item(token: str) -> dict[str, Any]:
def _pop_worker_item(token: str) -> tuple[float, dict[str, Any]]:
if not isinstance(token, str) or not token:
raise ValueError("worker item token is required")
_sweep_worker_items()
with _LOCK:
entry = _WORKER_ITEMS.pop(token, None)
if entry is None:
if entry is None or entry[0] <= time.monotonic():
raise RuntimeError("Image-list item is missing, expired, or already consumed")
return entry[1]
return entry
def _check_worker_item_identity(item: dict[str, Any], identity: str, index: int) -> None:
if not is_valid_item_identity(identity):
raise ValueError("Invalid worker item identity")
actual = item_content_identity(item) if "image" in item else item.get("item_identity")
if actual != identity:
raise ValueError("Worker item identity does not match the staged item")
if isinstance(index, bool) or not isinstance(index, int) or item["index"] != index:
raise ValueError("Worker item index does not match the staged item")
def consume_worker_item(token: str, *, item_identity: str | None = None,
item_index: int | None = None) -> dict[str, Any]:
with _LOCK:
expires, item = _pop_worker_item(token)
if "image" not in item:
raise RuntimeError("Image-list item is already consumed")
if item_identity is not None:
_check_worker_item_identity(item, item_identity, item_index)
if expires <= time.monotonic():
raise RuntimeError("Image-list item is expired")
# Start misses release staged tensors, but retain a tiny per-token
# receipt until End. This is NOT a completed-result/content cache.
# Keeping the original deadline prevents cache hits bypassing TTL.
_WORKER_ITEMS[token] = (expires, {"item_identity": item_identity, "index": item["index"]})
return item
def finish_worker_item(token: str, item_identity: str, item_index: int) -> None:
"""Validate and release this dispatch, whether Start ran or hit native cache."""
expires, item = _pop_worker_item(token)
_check_worker_item_identity(item, item_identity, item_index)
if expires <= time.monotonic():
raise RuntimeError("Image-list item is expired")
def delete_worker_item(token: str) -> bool:
@@ -376,13 +530,19 @@ def delete_worker_item(token: str) -> bool:
def create_master_job(start_id: str, end_id: str, prompt: dict[str, Any], items: list[dict[str, Any]],
worker_ids: list[str], master_callback_url: str,
workflow_meta: dict[str, Any] | None = None,
prompt_id: str | None = None, *, master_participates: bool = False) -> str:
prompt_id: str | None = None, *, master_participates: bool = False,
explicit_targets: bool = False) -> str:
if not isinstance(prompt, dict):
raise ValueError("Image List Start requires the original ComfyUI prompt for region extraction")
if not isinstance(master_participates, bool):
raise ValueError("master_participates must be a boolean")
if not worker_ids and not master_participates:
raise RuntimeError("Distributed Image List requires at least one preflighted direct worker; master fallback is disabled")
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)
for region in regions):
raise ValueError("Image-list region does not match the requested Start/End pair")
context_id = secrets.token_urlsafe(36)
state = ImageListJob(context_id, str(start_id), str(end_id), prompt, items,
[str(worker_id) for worker_id in worker_ids], master_callback_url,
@@ -497,15 +657,23 @@ def _safe_shared_dependency(node: dict[str, Any], prompt: dict[str, Any], regist
return types <= scalar_types and not any(_prompt_links(node, prompt))
def validate_image_list_regions(prompt: dict[str, Any]) -> list[dict[str, Any]]:
"""Validate pairing, non-nesting, and closed per-image graph boundaries."""
def _is_output_node(node: dict[str, Any], registry) -> bool:
node_class = registry.get(node.get("class_type"))
if node_class is None:
return False
info = node_class.GET_NODE_INFO_V1() if hasattr(node_class, "GET_NODE_INFO_V1") else {}
return bool(info.get("output_node", getattr(node_class, "OUTPUT_NODE", False)))
def _image_list_pairs(prompt: dict[str, Any]) -> list[tuple[str, str]]:
"""Pair boundaries before execution-copy pruning or full validation."""
if not isinstance(prompt, dict):
raise ValueError("Prompt must be an object")
starts = sorted(str(k) for k, v in prompt.items() if isinstance(v, dict) and v.get("class_type") == START_CLASS)
ends = sorted(str(k) for k, v in prompt.items() if isinstance(v, dict) and v.get("class_type") == END_CLASS)
if len(starts) != len(ends) or not starts:
raise ValueError("Distributed Image List needs one paired Start and End per region")
regions = []
pairs = []
for start_id in starts:
candidates = []
for end_id in ends:
@@ -514,14 +682,27 @@ def validate_image_list_regions(prompt: dict[str, Any]) -> list[dict[str, Any]]:
candidates.append(end_id)
if len(candidates) != 1:
raise ValueError(f"Start node {start_id} must connect its typed context output directly to exactly one End")
end_id = candidates[0]
pairs.append((start_id, candidates[0]))
return pairs
def validate_image_list_regions(prompt: dict[str, Any], *, explicit_targets: bool = False) -> list[dict[str, Any]]:
"""Validate pairing, non-nesting, and closed per-image graph boundaries.
Required output-node utilities are allowed only when callers explicitly
select execution roots, rather than letting every OUTPUT_NODE run independently.
All data-boundary checks still apply in that mode.
"""
regions = []
for start_id, end_id in _image_list_pairs(prompt):
if "images" not in prompt[end_id].get("inputs", {}):
raise ValueError(f"End node {end_id}.images must be connected to the per-item result branch")
end_inputs = prompt[end_id].get("inputs", {})
for name in ("images", "masks", "texts"):
end_override_names(end_inputs)
for name in END_RESULT_INPUTS:
if name in end_inputs:
_direct_link(end_inputs[name], prompt, end_id, name)
region_nodes = _ancestors_for_inputs(prompt, end_id, start_id, ("images", "masks", "texts"))
region_nodes = _ancestors_for_inputs(prompt, end_id, start_id, END_RESULT_INPUTS)
if start_id not in region_nodes:
raise ValueError(f"Start node {start_id} is not an ancestor of End node {end_id}")
forbidden = {n for n in region_nodes if n != start_id and n != end_id and
@@ -529,8 +710,6 @@ def validate_image_list_regions(prompt: dict[str, Any]) -> list[dict[str, Any]]:
{START_CLASS, END_CLASS, "DistributedCollector", "UltimateSDUpscaleDistributed"}}
if forbidden:
raise ValueError("Nested/ambiguous regions and other distributed nodes are unsupported in an image-list region: " + ", ".join(sorted(forbidden)))
# Output nodes execute independently of the lazy End, so they would
# evaluate the processing branch on the master as well as the worker.
import sys
registry = getattr(sys.modules.get("nodes"), "NODE_CLASS_MAPPINGS", {})
shared = region_nodes & _start_ancestors(prompt, start_id)
@@ -545,13 +724,10 @@ def validate_image_list_regions(prompt: dict[str, Any]) -> list[dict[str, Any]]:
f"Image-list region {start_id}->{end_id} bypasses Start through upstream nodes "
f"{sorted(unsafe)} ({'; '.join(links)}); route per-item data through Start"
)
for node_id in region_nodes - {start_id, end_id}:
node_class = registry.get(prompt[node_id].get("class_type"))
if node_class is None:
continue
info = node_class.GET_NODE_INFO_V1() if hasattr(node_class, "GET_NODE_INFO_V1") else {}
if info.get("output_node", getattr(node_class, "OUTPUT_NODE", False)):
raise ValueError(f"Output node {node_id} must be placed after Distributed Image List End")
if not explicit_targets:
for node_id in region_nodes - {start_id, end_id}:
if _is_output_node(prompt[node_id], registry):
raise ValueError(f"Output node {node_id} must be placed after Distributed Image List End")
downstream = _downstream_until(prompt, start_id, end_id)
outside = sorted(downstream - region_nodes)
if outside:
@@ -571,18 +747,28 @@ def validate_image_list_regions(prompt: dict[str, Any]) -> list[dict[str, Any]]:
def build_worker_prompt(prompt: dict[str, Any], region: dict[str, Any], item_token: str,
context_id: str, item_index: int, callback_url: str) -> dict[str, Any]:
"""Copy only the per-image branch and its worker-side dependencies."""
context_id: str, item_index: int, callback_url: str, *,
item_identity: str | None = None) -> dict[str, Any]:
"""Copy the branch; opt into stable ancestry only after a worker handshake."""
start_id, end_id = region["start_id"], region["end_id"]
_direct_link(prompt[end_id].get("inputs", {}).get("context"), prompt, end_id, "context")
kept = {node_id: json.loads(json.dumps(prompt[node_id])) for node_id in region["node_ids"]}
start_inputs = kept[start_id].setdefault("inputs", {})
# Per-item data reaches workers only through the staged item token.
for name in list(start_inputs):
if name in ("images", "mask") or name.startswith(VALUES_INPUT + "."):
start_inputs.pop(name)
start_inputs.update({"worker_item_token": item_token, "worker_context_id": context_id})
end_inputs = kept[end_id].setdefault("inputs", {})
if item_identity is not None:
if not is_valid_item_identity(item_identity):
raise ValueError("Invalid worker item identity")
# Native ancestry signatures include EVERY literal, not just a custom
# fingerprint. Remove inherited execution IDs, callbacks and all source
# inputs wholesale; transient delivery state belongs downstream at End.
kept[start_id]["inputs"] = {"worker_item_identity": item_identity}
end_inputs.update(worker_context_id=context_id, worker_item_identity=item_identity)
else:
# Preserve old raw /prompt and older-worker behavior byte-for-byte.
start_inputs = kept[start_id].setdefault("inputs", {})
for name in list(start_inputs):
if name in ("images", "mask") or name.startswith(VALUES_INPUT + "."):
start_inputs.pop(name)
start_inputs.update({"worker_item_token": item_token, "worker_context_id": context_id})
end_inputs.update({"worker_result_token": item_token,
"worker_item_index": int(item_index),
"master_callback_url": callback_url})
@@ -692,6 +878,7 @@ def set_master_result(context_id: str, item_index: int, token: str, result: dict
raise ValueError("Image-list result token does not match an assigned item")
if item_index in state.received_indexes:
raise ValueError(f"Duplicate image-list result for item {item_index}")
_checked_overrides(result, end_override_names(state.original_prompt[state.end_id].get("inputs", {})))
state.received_indexes.add(item_index)
state.results[item_index] = result
event = state.event
@@ -713,20 +900,72 @@ 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")
if mask is not None and (mask.ndim != 3 or mask.shape[0] != 1):
raise ValueError("result.mask must contain exactly one MASK frame")
text = decode_value(result.get("text"), "STRING")
return {"image": image, "mask": mask, "text": text}
decoded = {"image": image, "mask": mask, "text": text}
if "values" in result:
raw_values = result["values"]
if not isinstance(raw_values, dict):
raise ValueError("result.values must be an object")
values = {}
for name, encoded in raw_values.items():
if name not in VALUE_NAMES:
raise ValueError(f"result.values has unknown slot {name!r}")
if not isinstance(encoded, dict) or encoded.get("type") not in _VALUE_KINDS:
raise ValueError(f"result.values.{name} has an unsupported type")
values[name] = decode_value(encoded, encoded["type"])
decoded["values"] = _validate_result_values(values)
return decoded
def normalize_worker_result(images: Any, masks: Any = None, texts: Any = None) -> dict[str, Any]:
def normalize_worker_result(images: Any, masks: Any = None, texts: Any = None,
values: Any = None) -> dict[str, Any]:
def one(value, name, kind):
if value is None:
return None
@@ -749,17 +988,45 @@ def normalize_worker_result(images: Any, masks: Any = None, texts: Any = None) -
text = one(texts, "texts", "STRING")
if text is not None and not isinstance(text, str):
raise TypeError("Worker End text result must be a string")
return {"image": image, "mask": mask, "text": text}
result = {"image": image, "mask": mask, "text": text}
if values is not None:
overrides = normalize_value_overrides(values)
if overrides:
result["values"] = overrides
return result
def encode_result(result: dict[str, Any]) -> dict[str, Any]:
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"))
return {"image": encode_value(normalized["image"], "IMAGE"),
"mask": encode_value(normalized["mask"], "MASK"),
"text": encode_value(normalized["text"], "STRING")}
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
def _collect_output_lists(results: list[dict[str, Any]]) -> tuple[list[Any], list[Any], list[Any]]:
def _collect_output_lists(results: list[dict[str, Any]], *, state: ImageListJob | None = None,
include_values: bool = False) -> tuple[list[Any], ...]:
if not results:
raise RuntimeError("Image-list execution returned no results")
images, masks, texts = [], [], []
@@ -769,11 +1036,21 @@ def _collect_output_lists(results: list[dict[str, Any]]) -> tuple[list[Any], lis
images.append(result["image"])
masks.append(result.get("mask"))
texts.append(result.get("text"))
return images, masks, texts
if not include_values:
return images, masks, texts
if state is None or len(results) != len(state.items):
raise RuntimeError("Image-list values require exactly one result per source item")
expected = end_override_names(state.original_prompt[state.end_id].get("inputs", {}))
# Reconstruct unchanged values from the master source, not worker transport.
# Do this before the context is removed, and only after all items succeeded.
values = [{**item.get("values", {}), **_checked_overrides(result, expected)}
for item, result in zip(state.items, results)]
return images, masks, texts, values
async def run_image_list_job(context_id: str, *, model_callback=None) -> tuple[list[Any], list[Any], list[Any]]:
"""Schedule a shared queue across the optional master and direct workers."""
async def run_image_list_job(context_id: str, *, model_callback=None,
include_values: bool = False) -> tuple[list[Any], ...]:
"""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")
@@ -786,7 +1063,7 @@ async def run_image_list_job(context_id: str, *, model_callback=None) -> tuple[l
timer = asyncio.timeout(deadline if deadline > 0 else None)
try:
async with timer:
return await _run_image_list_job(state)
return await _run_image_list_job(state, include_values=include_values)
except TimeoutError:
if timer.expired():
raise TimeoutError(f"Image-list total deadline exceeded ({deadline:g}s); cleanup attempted") from None
@@ -799,33 +1076,80 @@ async def run_image_list_job(context_id: str, *, model_callback=None) -> tuple[l
remove_master_job(context_id)
async def _run_image_list_job(state: ImageListJob):
context_id = state.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):
state.event = asyncio.Event()
state.event_loop = asyncio.get_running_loop()
from ..api.orchestration.prompt_transform import _resolve_direct_workers_for_image_list
from ..utils.config import load_config
config = load_config()
workers = _resolve_direct_workers_for_image_list(config, state.worker_ids)
# 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 ..api.orchestration.prompt_transform import _resolve_direct_workers_for_image_list
from ..utils.config import load_config, get_worker_timeout_seconds
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
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.
worker_statuses = await preflight_image_list_workers(workers, require_idle=not state.master_participates) if workers else {}
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 state.master_participates:
raise RuntimeError("No free, reachable, feature-compatible direct workers are available; master execution is disabled")
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
regions = validate_image_list_regions(state.original_prompt)
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")
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]] = {}
@@ -840,6 +1164,8 @@ async def _run_image_list_job(state: ImageListJob):
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()
@@ -849,7 +1175,8 @@ async def _run_image_list_job(state: ImageListJob):
# /prompt queue, or repeated execution of the original source.
with _LOCK:
_WORKER_ITEMS[token] = (time.monotonic() + _WORKER_ITEM_TTL_SECONDS, state.items[index])
prompt = build_worker_prompt(state.original_prompt, region, token, context_id, index, "")
prompt = build_worker_prompt(state.original_prompt, region, token, context_id, index, "",
item_identity=item_content_identity(state.items[index]))
prompt[state.start_id]["inputs"]["image_list_local_execution"] = True
prompt[state.end_id]["inputs"]["image_list_local_execution"] = True
future = local_executor.submit(prompt, f"{context_id}:{index}", state.end_id, state.workflow_meta)
@@ -859,6 +1186,15 @@ async def _run_image_list_job(state: ImageListJob):
if worker is None:
raise RuntimeError(f"Worker {worker_id} was removed before dispatch")
assert session is not None
if 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])}
@@ -872,6 +1208,14 @@ async def _run_image_list_job(state: ImageListJob):
timeout=aiohttp.ClientTimeout(total=60)) as response:
if response.status >= 400:
raise RuntimeError(f"Worker {worker_id} rejected image-list item {index}: HTTP {response.status}")
try:
stored = await response.json()
except (ValueError, aiohttp.ClientError):
stored = None
# A per-upload acknowledgement gates the new graph. Older workers may
# return no identity; do not change global protocol/preflight contracts.
identity = stored.get("item_identity") if isinstance(stored, dict) else None
identity_options = {"item_identity": identity} if is_valid_item_identity(identity) else {}
callback_url = build_master_callback_url(worker, config=config)
if worker_id not in prepared_worker_prompts:
worker_prompt = build_worker_prompt(state.original_prompt, region, token, context_id,
@@ -884,14 +1228,17 @@ async def _run_image_list_job(state: ImageListJob):
from ..api.orchestration.media_sync import sync_worker_media
await asyncio.wait_for(sync_worker_media(worker, worker_prompt), timeout=60)
prepared_worker_prompts[worker_id] = worker_prompt
# Reuse media rewrites, never assignment identity or a mutable submitted graph.
# Reuse legacy media rewrites, never assignment identity or a mutable
# submitted graph. Keeping this template legacy also permits a worker to
# downgrade between items without inheriting a previous stable identity.
worker_prompt = build_worker_prompt(prepared_worker_prompts[worker_id], region, token,
context_id, index, callback_url)
context_id, index, callback_url, **identity_options)
handle["submission_pending"] = True
prompt_id = await dispatch_worker_prompt(
worker, worker_prompt, state.workflow_meta,
client_id=f"image-list-{context_id[:20]}", use_websocket=use_websocket,
return_prompt_id=True,
partial_execution_targets=[state.end_id],
)
if not isinstance(prompt_id, str) or not prompt_id:
raise RuntimeError(f"Worker {worker_id} did not return a prompt_id; submission status is ambiguous")
@@ -903,6 +1250,10 @@ async def _run_image_list_job(state: ImageListJob):
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
@@ -952,10 +1303,15 @@ async def _run_image_list_job(state: ImageListJob):
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
@@ -966,7 +1322,8 @@ async def _run_image_list_job(state: ImageListJob):
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:
@@ -975,6 +1332,10 @@ async def _run_image_list_job(state: ImageListJob):
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.
@@ -1015,10 +1376,11 @@ async def _run_image_list_job(state: ImageListJob):
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)
return _collect_output_lists(results, state=state, include_values=include_values)
finally:
if local_executor is not None:
# No timeout here: Python cannot safely kill a synchronous node.
+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)
+188 -3
View File
@@ -27,6 +27,39 @@ def _encode_all_chunks(frames, fmt, quality, max_chunk_bytes, max_frame_bytes):
return list(transfer_codec.encode_chunks(frames, fmt, quality, max_chunk_bytes, max_frame_bytes))
def _exact_frame_bytes(image_batch):
"""Wire size of one exact frame, computed from shape/dtype before encoding."""
from ..serverless_bridge import codec as bridge_codec
if not isinstance(image_batch, torch.Tensor) or image_batch.layout != torch.strided or image_batch.is_nested:
raise transfer_codec.TransferCodecError("exact transfer requires a dense IMAGE tensor")
if image_batch.dtype not in bridge_codec._DTYPE_NAMES:
raise transfer_codec.TransferCodecError(f"exact transfer does not support dtype {image_batch.dtype}")
if image_batch.ndim != 4:
raise transfer_codec.TransferCodecError(
f"Expected IMAGE tensor [B,H,W,C], got shape {tuple(image_batch.shape)}")
batch, height, width, channels = (int(value) for value in image_batch.shape)
if min(batch, height, width) <= 0 or channels not in (1, 3, 4):
raise transfer_codec.TransferCodecError("exact IMAGE dimensions must be positive with 1, 3 or 4 channels")
size = height * width * channels * image_batch.element_size()
if size > bridge_codec.MAX_PAYLOAD_BYTES:
raise transfer_codec.FrameTooLargeError("exact frame exceeds the tensor codec byte limit")
return size
def _encode_exact_frames(frames):
"""Encode [1,H,W,C] CPU slices losslessly (runs in an executor)."""
from ..serverless_bridge import codec as bridge_codec
descriptors, parts = [], []
for frame in frames:
try:
descriptor, payload = bridge_codec.encode_payload("IMAGE", frame)
except ValueError as exc:
raise transfer_codec.TransferCodecError(f"exact transfer cannot encode this IMAGE: {exc}") from None
descriptors.append(descriptor)
parts.append(payload)
return descriptors, parts
class DistributedCollectorNode:
INPUT_IS_LIST = True
EMPTY_AUDIO = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
@@ -200,6 +233,10 @@ class DistributedCollectorNode:
"""Send an image batch, optionally with audio, or an audio-only completion."""
batch_size = 0 if image_batch is None else image_batch.shape[0]
fmt = self._resolve_transfer_format(transfer_format) if batch_size else transfer_codec.LEGACY_FORMAT
if fmt == transfer_codec.EXACT_FORMAT:
# Never downgrade exact to 8-bit: an older master fails the job clearly.
await self._send_exact_batch(image_batch, audio, multi_job_id, master_url, worker_id)
return
if fmt != transfer_codec.LEGACY_FORMAT:
sent = await self._send_compressed_batch(
image_batch, audio, multi_job_id, master_url, worker_id, fmt, transfer_quality
@@ -329,6 +366,100 @@ class DistributedCollectorNode:
)
return True
async def _send_exact_batch(self, image_batch, audio, multi_job_id, master_url, worker_id):
"""Send original floating-point frames as binary multipart chunks.
Frame sizes are known from shape and dtype, so every request is planned
(and an oversized frame rejected) before anything is sent. Frames are
encoded one chunk at a time to avoid holding a second copy of the batch.
Any master rejection raises: there is no lossy or 8-bit fallback.
"""
encoded_audio = encode_audio_payload(audio)
session = await get_client_session()
url = f"{master_url}/distributed/job_complete"
loop = asyncio.get_running_loop()
media_limit = max(1, TRANSFER_MAX_REQUEST_BYTES - self.TRANSFER_ENVELOPE_RESERVE_BYTES)
chunk_limit = min(TRANSFER_MAX_CHUNK_BYTES, media_limit)
frame_bytes = _exact_frame_bytes(image_batch)
total = int(image_batch.shape[0])
if frame_bytes > media_limit:
raise transfer_codec.FrameTooLargeError(
f"Frame 0 is {frame_bytes} bytes as exact, above the {media_limit}-byte request limit. "
"Use png or a lossy format, or raise the master's --max-upload-size together with "
"COMFYUI_DISTRIBUTED_TRANSFER_MAX_REQUEST_BYTES."
)
# Bound descriptor/part-header overhead even for tiny frames.
per_chunk = min(128, max(1, chunk_limit // frame_bytes))
ranges = [(start, min(per_chunk, total - start)) for start in range(0, total, per_chunk)]
audio_bytes = len(json.dumps(encoded_audio)) if encoded_audio is not None else 0
last_media = ranges[-1][1] * frame_bytes
audio_inline = encoded_audio is not None and last_media + audio_bytes <= media_limit
audio_separate = encoded_audio is not None and not audio_inline
height, width = int(image_batch.shape[1]), int(image_batch.shape[2])
if audio_separate and audio_bytes > media_limit:
raise transfer_codec.FrameTooLargeError("exact completion audio exceeds the request limit")
cpu_batch = ensure_contiguous(image_batch.detach().cpu())
started = time.time()
sent_bytes = 0
for index, (start, count) in enumerate(ranges):
frames = [cpu_batch[i:i + 1] for i in range(start, start + count)]
descriptors, parts = await loop.run_in_executor(None, _encode_exact_frames, frames)
is_final_chunk = index == len(ranges) - 1
envelope = {
"job_id": str(multi_job_id),
"worker_id": str(worker_id),
"format": transfer_codec.EXACT_FORMAT,
"start_index": start,
"count": count,
"width": width,
"height": height,
"descriptors": descriptors,
"is_last": bool(is_final_chunk and not audio_separate),
}
if is_final_chunk and audio_inline:
envelope["audio"] = encoded_audio
chunk_bytes = sum(len(p) for p in parts)
body = self._build_multipart(envelope, parts)()
if body.size is None or body.size > TRANSFER_MAX_REQUEST_BYTES:
raise transfer_codec.FrameTooLargeError("exact multipart body exceeds the request limit")
timeout_seconds = min(600, 60 + chunk_bytes // (256 * 1024))
try:
async with session.post(
url,
data=body,
timeout=aiohttp.ClientTimeout(total=timeout_seconds),
) as response:
if response.status >= 400:
body = await response.text()
raise RuntimeError(
f"Master rejected exact chunk {index + 1}/{len(ranges)} ({chunk_bytes} bytes) "
f"with HTTP {response.status}: {body[:300]}. exact needs a master with exact "
"transfer support and requests within its --max-upload-size; it is never "
"resent as an 8-bit format. Update the master or choose png."
)
except Exception as e:
log(f"Worker - Failed to send exact chunk to master: {e}")
debug_log(f"Worker - Full error details: URL={url}")
raise
sent_bytes += chunk_bytes
requests = len(ranges)
if audio_separate:
await self._post_json_completion(session, url, {
"job_id": str(multi_job_id),
"worker_id": str(worker_id),
"batch_idx": 0,
"audio": encoded_audio,
"is_last": True,
})
requests += 1
log(
f"Worker - Sent {total} image(s) as exact in {requests} request(s): "
f"{sent_bytes / 1e6:.2f} MB in {time.time() - started:.1f}s"
)
@staticmethod
async def _post_json_completion(session, url, payload):
timeout_seconds = 60 if "image" in payload else 600
@@ -457,6 +588,31 @@ class DistributedCollectorNode:
worker_images[worker_id][image_index] = tensor
return 1
# Collector sockets a comfy-media-v1 serverless pool may return.
TYPED_MEDIA_SOCKETS = ("images", "audio")
def _store_typed_media(self, worker_images: dict, worker_audio: dict, worker_id: str, media) -> int:
"""Store one atomic typed serverless delivery: ordered IMAGE/AUDIO item lists.
Each IMAGE item is a BHWC batch kept at its list position; AUDIO items are
joined in order, exactly as list inputs are on the master.
"""
if (not isinstance(media, dict) or not media
or any(socket not in self.TYPED_MEDIA_SOCKETS for socket in media)):
raise RuntimeError(f"Invalid serverless media delivery from {worker_id}")
for socket, items in media.items():
if not isinstance(items, list) or not items or any(item is None for item in items):
raise RuntimeError(f"Serverless participant {worker_id} returned no {socket} items")
stored = 0
for index, tensor in enumerate(media.get("images", [])):
if not isinstance(tensor, torch.Tensor) or tensor.ndim != 4:
raise RuntimeError(f"Serverless participant {worker_id} returned an invalid IMAGE item")
worker_images.setdefault(worker_id, {})[index] = tensor
stored += 1
if "audio" in media:
worker_audio[worker_id] = self._normalize_audio_input(media["audio"])
return stored
def _reorder_and_combine_tensors(
self,
worker_images: dict,
@@ -534,6 +690,10 @@ class DistributedCollectorNode:
serverless_manager = self._get_serverless_manager()
return serverless_manager
try:
exact_requested = transfer_codec.normalize_format(transfer_format) == transfer_codec.EXACT_FORMAT
except transfer_codec.TransferCodecError:
exact_requested = False
num_workers = len(expected_workers)
if num_workers == 0:
return (images, audio if audio is not None else self.EMPTY_AUDIO)
@@ -609,6 +769,27 @@ class DistributedCollectorNode:
raise RuntimeError(f"Serverless participant {participant} failed: {message}")
result_worker_id = str(result["worker_id"])
if "media" in result:
if result_worker_id not in serverless_workers:
raise RuntimeError(f"Unexpected serverless media participant {result_worker_id}")
if result_worker_id in workers_done:
raise RuntimeError(f"Duplicate serverless media delivery from {result_worker_id}")
if result.get("is_last") is not True:
raise RuntimeError("Serverless media delivery must be atomic and final")
collected_count += self._store_typed_media(
worker_images, worker_audio, result_worker_id, result["media"]
)
debug_log(f"Master - Got typed media from serverless participant {result_worker_id}")
mark_worker_done(result_worker_id)
return
if (exact_requested and result.get("tensor") is not None
and result.get("format") != transfer_codec.EXACT_FORMAT):
raise RuntimeError(
f"Worker {result_worker_id} sent {result.get('format')} images, but this Collector "
"requested exact. Update ComfyUI-Distributed on that worker; exact is never "
"replaced by an 8-bit format."
)
is_last = result.get("is_last", False)
collected_count += self._store_worker_result(worker_images, result)
debug_log(
@@ -733,6 +914,10 @@ class DistributedCollectorNode:
consume_result(item)
break
if exact_requested and workers_done != expected_workers:
raise RuntimeError(
f"Exact transfer incomplete: missing completion from {sorted(expected_workers - workers_done)}"
)
total_collected = sum(len(imgs) for imgs in worker_images.values())
# Remove the queue before combining so a late provider callback cannot
@@ -753,10 +938,10 @@ class DistributedCollectorNode:
f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
return (combined, combined_audio)
except Exception as e:
if serverless_results:
if serverless_results or exact_requested:
raise RuntimeError(
"Serverless result batch shape mismatch while combining participants "
f"{sorted(serverless_results)}: {e}"
"Serverless or exact result batch shape mismatch while combining participants "
f"{sorted(worker_images)}: {e}"
) from e
log(f"Master - Error combining images: {e}")
# Preserve collected audio for regular workers when image assembly fails.
+151 -17
View File
@@ -6,10 +6,13 @@ 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,
VALUES_INPUT,
end_override_names,
create_master_job,
normalize_image_list_inputs,
normalize_worker_result,
@@ -17,9 +20,40 @@ from ..image_list.runtime import (
encode_result,
encode_request_payload,
unwrap_single,
is_valid_item_identity,
)
def _stable_worker_binding(prompt, start_id, identity):
"""Resolve transient state only from the exact Start -> End context wire."""
if not is_valid_item_identity(identity):
raise ValueError("Invalid worker item identity")
start = prompt.get(start_id) if isinstance(prompt, dict) else None
if (not isinstance(start, dict) or start.get("class_type") != "DistributedImageListStart"
or start.get("inputs", {}).get("worker_item_identity") != identity):
raise RuntimeError("Worker item identity is missing its paired Start")
candidates = []
for end_id, node in prompt.items():
if not isinstance(node, dict) or node.get("class_type") != "DistributedImageListEnd":
continue
inputs = node.get("inputs", {})
link = inputs.get("context")
if isinstance(link, (list, tuple)) and list(link) == [start_id, 3]:
candidates.append((end_id, inputs))
if len(candidates) != 1:
raise RuntimeError("Worker Start must have exactly one paired End context wire")
end_id, inputs = candidates[0]
if inputs.get("worker_item_identity") != identity:
raise ValueError("Paired End worker item identity does not match Start")
for name in ("worker_result_token", "worker_context_id"):
if not isinstance(inputs.get(name), str) or not inputs[name]:
raise RuntimeError(f"Worker paired End is missing {name}")
index = inputs.get("worker_item_index")
if isinstance(index, bool) or not isinstance(index, int) or index < 0:
raise ValueError("Worker paired End item index is invalid")
return end_id, inputs
# Types the transport can carry exactly. Anything else (models, latents, ...)
# must be loaded inside the region or shared as a model dependency instead.
_VALUE_TYPES = [io.Image, io.Mask, io.String, io.Int, io.Float, io.Boolean]
@@ -69,10 +103,26 @@ class DistributedImageListStart(io.ComfyNode):
worker_item_token="", worker_context_id="", image_list_enabled_worker_ids="[]",
image_list_master_callback="", image_list_execution_id="",
image_list_workflow_meta=None, image_list_master_participates=None,
image_list_local_execution=False):
image_list_local_execution=False, worker_item_identity=None):
worker_item_token = unwrap_single(worker_item_token)
worker_context_id = unwrap_single(worker_context_id)
worker_item_identity = unwrap_single(worker_item_identity)
image_list_workflow_meta = unwrap_single(image_list_workflow_meta)
if worker_item_identity is not None:
prompt = getattr(cls.hidden, "prompt", None)
start_id = str(getattr(cls.hidden, "unique_id", ""))
_, delivery = _stable_worker_binding(prompt, start_id, worker_item_identity)
# Only cache misses execute Start. Never put this token/context in
# Start's literal inputs or returned context: both would poison reuse.
from ..image_list.runtime import consume_worker_item
item = consume_worker_item(delivery["worker_result_token"],
item_identity=worker_item_identity,
item_index=delivery["worker_item_index"])
item_values = item.get("values") or {}
return io.NodeOutput(
[item["image"]], [item.get("mask")], [item["index"]], [worker_item_identity],
*[[item_values.get(name)] for name in VALUE_NAMES],
)
if worker_item_token:
if unwrap_single(image_list_local_execution):
from ..image_list.runtime import consume_worker_item
@@ -99,13 +149,18 @@ class DistributedImageListStart(io.ComfyNode):
raise RuntimeError("Invalid preflighted image-list worker list") from exc
if not isinstance(worker_ids, list) or not all(isinstance(item, str) for item in worker_ids):
raise RuntimeError("Invalid preflighted image-list worker list")
from comfy_execution.utils import get_executing_context
from ..image_list.lifecycle import has_safe_execution_targets
import server
executing = get_executing_context()
explicit_targets = has_safe_execution_targets(
getattr(server.PromptServer, "instance", None),
executing.prompt_id if executing is not None else None, prompt)
from ..image_list.runtime import validate_image_list_regions
regions = validate_image_list_regions(prompt)
regions = validate_image_list_regions(prompt, explicit_targets=explicit_targets)
region = next((row for row in regions if row["start_id"] == start_id), None)
if region is None:
raise RuntimeError(f"No safe paired Image List End was found for Start node {start_id}")
from comfy_execution.utils import get_executing_context
executing = get_executing_context()
master_participates = unwrap_single(image_list_master_participates)
if master_participates is None:
from ..utils.config import load_config
@@ -118,6 +173,7 @@ class DistributedImageListStart(io.ComfyNode):
image_list_workflow_meta if isinstance(image_list_workflow_meta, dict) else None,
prompt_id=executing.prompt_id if executing is not None else None,
master_participates=master_participates,
explicit_targets=explicit_targets,
)
return io.NodeOutput(
[item["image"] for item in items],
@@ -143,17 +199,27 @@ class DistributedImageListEnd(io.ComfyNode):
node_id="DistributedImageListEnd",
display_name="Distributed Image List End",
category="image/distributed",
description="Closes a paired image-list region and returns ordered image, mask, and text lists.",
description=("Closes a paired image-list region and returns ordered image, mask, text and values lists. "
"Start values are retained; connected value_N inputs replace or add that key per item."),
inputs=[
io.Image.Input("images", lazy=True),
io.Mask.Input("masks", optional=True, lazy=True),
io.String.Input("texts", optional=True, lazy=True, force_input=True),
io.Custom("DISTRIBUTED_IMAGE_LIST_CONTEXT").Input("context"),
io.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],
],
outputs=[
io.Image.Output(display_name="image_list", is_output_list=True),
io.Mask.Output(display_name="masks", is_output_list=True),
io.String.Output(display_name="texts", is_output_list=True),
io.Custom("DISTRIBUTED_IMAGE_LIST_VALUES").Output(
id="values", display_name="values", is_output_list=True),
],
is_input_list=True,
is_output_node=True,
@@ -175,7 +241,8 @@ class DistributedImageListEnd(io.ComfyNode):
and node.get("class_type") == "DistributedImageListEnd"
and node.get("inputs", {}).get("worker_result_token") == worker_result_token), {})
required = []
for name, value in (("images", images), ("masks", masks), ("texts", texts)):
for name, value in (("images", images), ("masks", masks), ("texts", texts),
*((name, kwargs.get(name)) for name in end_override_names(end_inputs))):
connected = name == "images" or name in end_inputs
missing = value is None or (isinstance(value, (list, tuple)) and all(v is None for v in value))
if connected and missing:
@@ -189,26 +256,67 @@ class DistributedImageListEnd(io.ComfyNode):
@classmethod
def execute(cls, images=None, masks=None, texts=None, context=None,
worker_result_token="", worker_item_index=-1, master_callback_url="",
image_list_local_execution=False):
image_list_local_execution=False, worker_context_id="",
worker_item_identity=None, 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)
worker_item_identity = unwrap_single(worker_item_identity)
if worker_item_identity is not None and not worker_result_token:
raise RuntimeError("Worker item identity is missing its delivery token")
if worker_result_token:
context_values = [item for item in (context if isinstance(context, (list, tuple)) else [context])
if isinstance(item, str) and item]
if not context_values or len(set(context_values)) != 1:
raise RuntimeError("Worker Image List End is missing its paired context")
result = normalize_worker_result(images, masks, texts)
prompt = getattr(cls.hidden, "prompt", {}) or {}
context_id = context_values[0]
if worker_item_identity is not None:
end_id = str(getattr(cls.hidden, "unique_id", ""))
node = prompt.get(end_id, {})
link = node.get("inputs", {}).get("context")
if (node.get("class_type") != "DistributedImageListEnd"
or not isinstance(link, (list, tuple)) or len(link) != 2):
raise RuntimeError("Worker End is missing its paired Start context wire")
paired_id, end_inputs = _stable_worker_binding(prompt, link[0], worker_item_identity)
if (paired_id != end_id or context_values != [worker_item_identity]
or end_inputs["worker_context_id"] != worker_context_id
or end_inputs["worker_result_token"] != worker_result_token
or end_inputs["worker_item_index"] != unwrap_single(worker_item_index)
or end_inputs.get("master_callback_url", "") != master_callback_url):
raise RuntimeError("Worker End identity/context does not match its current delivery")
from ..image_list.runtime import finish_worker_item
# On a miss this consumes Start's receipt; on a hit it validates
# and drops the current staged tensors. Both paths enforce TTL.
finish_worker_item(worker_result_token, worker_item_identity,
unwrap_single(worker_item_index))
context_id = worker_context_id
else:
end_inputs = next((node.get("inputs", {}) for node in prompt.values()
if isinstance(node, dict) and node.get("class_type") == "DistributedImageListEnd"
and node.get("inputs", {}).get("worker_result_token") == worker_result_token), {})
expected = end_override_names(end_inputs)
# Missing connected lazy values must fail instead of becoming an
# implicit fallback to the original Start value.
supplied = {name: values.get(name) for name in expected}
supplied.update(values)
result = normalize_worker_result(images, masks, texts, supplied)
if unwrap_single(image_list_local_execution):
from ..image_list.runtime import set_master_result
set_master_result(context_values[0], int(unwrap_single(worker_item_index)),
set_master_result(context_id, int(unwrap_single(worker_item_index)),
str(worker_result_token), result)
else:
payload = encode_result(result)
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_values[0], int(unwrap_single(worker_item_index)), payload), timeout=60,
context_id, int(unwrap_single(worker_item_index)), payload), timeout=60,
)
return io.NodeOutput([result["image"]], [result.get("mask")], [result.get("text")])
return io.NodeOutput([result["image"]], [result.get("mask")], [result.get("text")],
[result.get("values", {})])
context_values = [item for item in (context if isinstance(context, (list, tuple)) else [context])
if isinstance(item, str) and item]
@@ -223,14 +331,40 @@ class DistributedImageListEnd(io.ComfyNode):
# Add them to its existing tracker; never start/end another tracker
# that could unpin a model still needed by the parent.
model_callback = getattr(execution_list, "output_link_callback", None)
result_lists = run_async_in_server_loop(run_image_list_job(context_id, model_callback=model_callback))
result_lists = run_async_in_server_loop(run_image_list_job(
context_id, model_callback=model_callback, include_values=True))
except ImageListCancelledError:
import comfy.model_management
raise comfy.model_management.InterruptProcessingException()
images_out, masks_out, texts_out = result_lists
if len(images_out) != len(masks_out) or len(images_out) != len(texts_out):
raise RuntimeError("Image-list collector received mismatched image, mask, or text result counts")
return io.NodeOutput(images_out, masks_out, texts_out)
images_out, masks_out, texts_out, values_out = result_lists
if len({len(output) for output in result_lists}) != 1:
raise RuntimeError("Image-list collector received mismatched image, mask, text or values result counts")
return io.NodeOutput(images_out, masks_out, texts_out, values_out)
class DistributedImageListGetValue(io.ComfyNode):
"""Ordinary scalar execution lets core map over End's ordered values list."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DistributedImageListGetValue",
display_name="Distributed Image List Get Value",
category="image/distributed",
inputs=[io.Custom("DISTRIBUTED_IMAGE_LIST_VALUES").Input("values"),
io.Combo.Input("value_name", options=list(VALUE_NAMES), default="value_1")],
outputs=[io.AnyType.Output(id="value", display_name="value")],
)
@classmethod
def execute(cls, values, value_name="value_1"):
if value_name not in VALUE_NAMES:
raise ValueError(f"Unknown image-list value slot {value_name!r}")
if not isinstance(values, dict):
raise TypeError("Image List Get Value expects one values mapping from Image List End")
if value_name not in values:
raise ValueError(f"Image List Get Value: {value_name} is absent from this item's values")
return io.NodeOutput(values[value_name])
async def submit_worker_result(callback_url, item_token, context_id, item_index, result):
+2 -2
View File
@@ -10,7 +10,7 @@ from . import utilities as _utilities
from .collector import DistributedCollectorNode as _CollectorRuntime
from ..utils import transfer_codec as _transfer_codec
from .distributed_upscale import UltimateSDUpscaleDistributed as _UpscaleRuntime
from .image_list import DistributedImageListEnd, DistributedImageListStart
from .image_list import DistributedImageListEnd, DistributedImageListStart, DistributedImageListGetValue
class DistributedSeed(io.ComfyNode):
@@ -212,4 +212,4 @@ class DistributedWorkerSink(io.ComfyNode):
NODES = [DistributedCollector, DistributedSeed, DistributedModelName, DistributedValue,
ImageBatchDivider, AudioBatchDivider, DistributedEmptyImage, UltimateSDUpscaleDistributed,
DistributedWorkerSink, DistributedImageListStart, DistributedImageListEnd]
DistributedWorkerSink, DistributedImageListStart, DistributedImageListEnd, DistributedImageListGetValue]
+282
View File
@@ -0,0 +1,282 @@
# Standalone `comfy-media-v1` endpoint bridge
This directory is independently installable. It registers two internal Comfy V3
nodes, `MediaImportV1` and `MediaExportV1`, and provides the asynchronous HTTP
executor `execute_media_request`. It does **not** start ComfyUI, provision an
endpoint, install models, or provide a provider worker/supervisor. The endpoint
needs this bridge, not the full ComfyUI-Distributed extension.
The bridge's generic contract supports `IMAGE`, `AUDIO`, `MASK`, `STRING`, and
native ComfyUI `VIDEO`. **Distributed Collector** (`DistributedCollector`) uses
only the image/audio outputs and combines their returned lists into ordinary
IMAGE/AUDIO outputs. It has no MASK/STRING/VIDEO outputs, and the separate
`DistributedMediaCollector` node has been removed. For Collector setup, ordering,
and limits, see [serverless images and audio](../docs/serverless-media.md).
## Install in the endpoint image
1. Copy this complete directory to `<ComfyUI>/custom_nodes/serverless_bridge`.
Use a compatible ComfyUI with `comfy_api.v0_0_2` and the native VIDEO classes
used by [video_codec.py](video_codec.py). These exact-class/private-state
adapters are tied to the supported core implementation, not arbitrary VIDEO
plugins or future core layouts.
2. Restart ComfyUI so `comfy_entrypoint()` registers the two nodes. Keep the
workflow's own models and custom nodes installed separately.
3. Make `<ComfyUI>/custom_nodes` importable by the handler process (for example,
via its `PYTHONPATH`). The handler uses `aiohttp`; the nodes/codecs use
ComfyUI's PyTorch/native media dependencies. The handler module itself does
not import ComfyUI, PyTorch, Runpod, or Distributed bootstrap code.
4. Give the handler filesystem access to **the same input directory** used by
that ComfyUI instance, plus its trusted HTTP URL. These are operator settings,
never request-controlled paths/URLs. Staging/cleanup requires POSIX directory
descriptor and no-follow file operations.
Installing the nodes alone is insufficient: the endpoint must route media-v1
inputs through this executor and implement the uncertainty policy below. A
legacy image handler cannot interpret the typed envelope.
## Callable integration
```python
from serverless_bridge.handler import execute_media_request
async def execute_endpoint_input(
input_data, *, comfy_url, input_dir, temp_dir=None,
download=None, publish=None,
):
# Invoke only under a supervisor implementing the uncertainty policy below.
# Exceptions intentionally propagate; this is NOT a provider worker loop.
return await execute_media_request(
input_data,
comfy_url=comfy_url,
input_dir=input_dir,
timeout_seconds=1800,
session=None,
download=download,
publish=publish,
temp_dir=temp_dir,
)
```
Pass the provider job's **`input` object**, not the whole job. Return the resulting
`{"media": ...}` object as provider output; Runpod exposes it as `output.media`.
Do not add another `output` wrapper inside it. `timeout_seconds` must be finite
and positive. If supplied, `session` is a caller-owned `aiohttp.ClientSession`
for local Comfy HTTP and remains open; otherwise the executor creates/closes its
own session. Input dictionaries are copied before graph rewriting.
**This callable is not a turnkey Runpod handler.** An outer SDK that catches every
exception and immediately reuses the worker is unsafe. The supervisor must
recognize `serverless_bridge.handler.MediaExecutionUncertain`, stop intake, and
quarantine/reconcile or terminate the owning environment before reuse. Merely
returning `{"error": ...}` or marking a provider job `FAILED` is not that policy.
## Request and response
A minimal generic bridge text request builder (no files or model required by this
graph). This illustrates the bridge contract, not a Collector text output:
```python
import json
import uuid
def make_text_request(text="hello"):
request_id = uuid.uuid4().hex
binding = {
"binding_id": "text-binding",
"collector_id": "collector-1",
"job_id": "collector-job-1",
"participant_id": "serverless:example:0",
"node_id": "export-1",
"socket": "texts",
"kind": "STRING",
}
return {
"workflow": {
"text-1": {"class_type": "PrimitiveString", "inputs": {"value": text}},
"export-1": {
"class_type": "MediaExportV1",
"inputs": {
"request_id": request_id,
"binding_json": json.dumps(binding),
"texts": ["text-1", 0],
},
},
},
"images": [],
"media": {
"version": 1,
"profile": "comfy-media-v1",
"request_id": request_id,
"artifacts": [],
"inputs": [],
"outputs": [binding],
},
}
```
`workflow` is an API graph, not an editor workflow. Top-level input fields are
exactly `workflow`, `media`, and optional `images`. `media` has exactly the six
fields shown above. `media.outputs` declares one exporter per connected socket;
its binding must match that node's `binding_json` and request ID exactly. Socket
names/types are `images/IMAGE`, `audio/AUDIO`, `masks/MASK`, `texts/STRING`, and
`video/VIDEO`. Use unique request IDs. Identifiers are 1–200 characters from
`[A-Za-z0-9_.:/-]`; request scopes additionally reject empty, `.` or `..` slash
components. The executor validates the graph and submits **once**, with exactly
the declared exporter IDs as `partial_execution_targets`. Unrelated graph roots,
cycles, undeclared exporters, or caller-supplied `MediaImportV1` nodes fail.
For typed inputs, `media.artifacts` contains records with exactly:
| Field | Value |
| --- | --- |
| `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 |
Each `media.inputs` entry is
`{"node_id": "consumer", "input_name": "socket", "artifact_id": "artifact-id"}`.
It replaces an existing consumer input with a generated import link; every
artifact must be used, and destinations cannot repeat. The executor verifies
bytes/checksums, stages owned files, inserts imports, then prunes disconnected
loader ancestors. Optional legacy `images` entries are `{ "name": "input.png",
"image": "<base64>" }`, must match stock `LoadImage` references, and are rewritten
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 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.
All identity, completeness, type, size, and checksum checks must pass; do not
consume a partial result. The internal history transport `comfy-file` is local
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
`transfer_quality` controls. The **bridge-level defaults** are `exact` and `90`;
generic callers that omit the controls retain that representation. Other media
exports must not carry image controls. Callers should not alter binding identity
metadata.
The **Collector default is different**: its saved widget value remains
`legacy_png` with quality `90`. For a `comfy-media-v1` pool, the master compiler
maps that value to an explicit `png` export, so normal serverless image returns
use compressed 8-bit PNG. `exact` is opt-in at the Collector. Other selected
formats are unchanged. Direct workers also support opt-in exact IMAGE transport;
an older master rejection is an error, not an 8-bit retry. A Collector requesting
`exact` cannot dispatch to a `legacy-image-v1` pool.
Formats: `exact`, `legacy_png`, `png`, `ffv1`, `jpeg`, `webp`, `h264`,
`h265` (**experimental**), `av1`.
`exact` preserves float tensor bits. The image/video codecs quantize RGB/grayscale
to 8-bit; PNG/FFV1 are lossless at that precision, while the other codecs are
lossy. Quality is an integer 1–100 and affects only lossy formats. IMAGE batch
order, shape and channel count are retained, and alpha is not discarded.
Native VIDEO/MASK/AUDIO/STRING bridge values still use their original exact codecs;
the image-format controls do not add those sockets to the Collector.
Copy/update the complete bridge directory on both sides, including
`image_transfer.py` and `image_codec.py`. Pillow and the selected CPU/software
PyAV codec must be available; missing encoders fail without automatic fallback.
Compressed IMAGE artifacts use `image-transfer-v1`; decode through the shared
`codec.decode_payload`, which enforces payload and decoded-shape limits.
See [image transfer choices](../docs/serverless-media.md#image-return-format-and-quality).
## Transport budgets and ownership
- Contract maxima: 256 MiB per artifact, 512 MiB aggregate artifact bytes,
128 artifacts, 64 items per binding. Request input totals include decoded
legacy-image bytes. Decoders validate tensor/VIDEO sizes and layouts before
allocating native storage; these are payload limits, **not peak-RSS guarantees**.
- The built-in master currently emits inline inputs only. Its complete Runpod
request JSON is capped at **10,000,000 bytes**, including base64, graph and
policy. The larger contract limits do not bypass this cap; there is no
built-in large-input/object-storage uploader.
- The executor's result JSON is capped at **10,000,000 bytes**. Inline results
are the default. When needed, `publish` is called for selected larger
artifacts to fit that budget; supplying it does not force all results to URLs.
Without a publisher, an oversized response fails instead of being truncated.
- Optional **async** `download(url, *, max_bytes)` returns `bytes`. The executor
supplies the declared artifact length and checks exact length and SHA-256.
The callback must enforce public HTTPS, DNS/IP restrictions, redirect checks,
and bounded streaming **before allocation**; syntactic contract validation
does not provide SSRF protection. No downloader is installed by default.
- Optional **async** `publish(metadata, payload)` uploads the exact `bytes` and
returns exactly `{"type": "url", "url": "https://...", "expires_at": "..."}`.
`metadata` is the artifact record without `transport`. Expiry must be a future,
timezone-aware ISO timestamp; URLs use HTTPS port 443, no userinfo/fragment,
and at most 8192 characters. Keep objects available long enough for collection.
The callback/operator owns upload credentials, object expiry and cleanup,
including orphan uploads after failure. There is no bridge storage backend.
- The master uses a separate bounded public-HTTPS downloader without provider
credentials. Never send provider credentials to artifact hosts or log signed
URLs. Endpoint callbacks must enforce equivalent protections themselves.
Callbacks must cooperate with async cancellation, finish/abort their own I/O,
and not detach readers/uploads that outlive the callback. Synchronous code that
blocks the event loop cannot be hard-timed-out by this executor.
## Cancellation, quarantine, and file retention
`timeout_seconds` covers staging, submission, polling, downloads and publication.
On failure/cancellation, there is a separate **10-second callback-drain grace**
and, when needed, a **10-second prompt-reconciliation grace**. These are recovery
allowances, not evidence that execution stopped. Ambiguous `/prompt` submission
is never retried. An identified pending prompt may be deleted from `/queue`,
with exact read-back; running prompts are never sent `/interrupt`. The pinned
core route has a snapshot/global-interrupt race even when given a prompt ID.
`MediaExecutionUncertain` means prompt acceptance/termination **or callback
teardown** is unconfirmed. Even successful terminal Comfy history is insufficient
if an async downloader/publisher ignores cancellation. The bridge retains input
and exporter files in that case; late callback exit does not automatically
remove retained files. Before accepting another job on that environment:
1. Quarantine it and stop new intake; do not automatically resubmit the request.
2. Reconcile the local prompt and all readers, or terminate the owning execution
environment when their termination cannot be established.
3. Only then clean retained files and allow reuse/replacement. A provider terminal
job status alone cannot certify a surviving local Comfy process or callback.
Staged inputs use random exclusive files under
`input_dir/distributed-media-v1/<sha256(request_id)>/`. Cleanup removes only files
created by that invocation, after confirmed terminal state (or before submission)
and reader teardown; it never recursively deletes a shared namespace.
`temp_dir=None` leaves exporter retention to the Comfy temp-directory/environment
lifecycle. Optionally pass the **trusted local Comfy temp root** to enable
best-effort cleanup of validated, invocation-owned exporter UUID `.bin` files
under the corresponding hashed namespace, only after terminal history and
collection/publisher teardown. Symlinked, hard-linked, replaced, missing, or
unidentifiable files are not claimed; directories and all unidentified leftovers
remain lifecycle-owned. This is not a temp-directory sweep. Independently decoded
VIDEO objects do not depend on either staging or exporter files.
This guide describes the source contract, not a test result. No live provider/GPU
validation is claimed, and this directory is not a turnkey provider deployment.
+14
View File
@@ -0,0 +1,14 @@
"""Standalone typed-media bridge; copy this directory into endpoint custom_nodes.
Importing codecs does not load ComfyUI, providers or Distributed bootstrap.
"""
async def comfy_entrypoint():
from comfy_api.v0_0_2 import ComfyExtension
from .nodes import MediaImportV1, MediaExportV1
class MediaBridgeExtension(ComfyExtension):
async def get_node_list(self):
return [MediaImportV1, MediaExportV1]
return MediaBridgeExtension()
+195
View File
@@ -0,0 +1,195 @@
"""Lossless, bounded CPU payloads shared with the standalone endpoint bridge.
Descriptors contain only kind-specific metadata. Tensor bytes are contiguous,
little-endian floating-point values, never a pickle or an image/audio container.
The optional byte budget can lower, but never raise, MAX_PAYLOAD_BYTES.
"""
import ctypes
import sys
import torch
MAX_PAYLOAD_BYTES = 256 * 1024 * 1024
_DTYPES = {
"float16": (torch.float16, 2),
"bfloat16": (torch.bfloat16, 2),
"float32": (torch.float32, 4),
"float64": (torch.float64, 8),
}
_DTYPE_NAMES = {dtype: name for name, (dtype, _) in _DTYPES.items()}
_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", "VALUE"):
raise ValueError("Unsupported payload kind")
def _byte_limit(max_bytes):
if type(max_bytes) is not int or max_bytes < 0:
raise ValueError("max_bytes must be a nonnegative integer")
return min(max_bytes, MAX_PAYLOAD_BYTES)
def _sample_rate(value):
if type(value) is not int or value <= 0:
raise ValueError("AUDIO sample_rate must be a positive integer")
return value
def _tensor_size(kind, shape, itemsize, limit):
if type(shape) is not list or len(shape) != len(_LAYOUTS[kind]):
raise ValueError("Tensor shape has an invalid rank")
if any(type(dimension) is not int or dimension <= 0 for dimension in shape):
raise ValueError("Tensor shape dimensions must be positive integers")
if kind == "IMAGE" and shape[-1] not in (1, 3, 4):
raise ValueError("IMAGE shape must have 1, 3 or 4 channels")
size = itemsize
for dimension in shape:
# Divide before multiplying, including for arbitrary-size Python ints.
if dimension > limit // size:
raise ValueError("Tensor payload exceeds byte limit")
size *= dimension
return size
def _native_little_endian(payload, itemsize):
"""Convert native to wire byte order, or vice versa (an involution)."""
if sys.byteorder == "little":
return payload
swapped = bytearray(len(payload))
for offset in range(itemsize):
swapped[offset::itemsize] = payload[itemsize - 1 - offset::itemsize]
return bytes(swapped)
def _encode_string(value, limit):
if not isinstance(value, str):
raise ValueError("STRING value must be text")
# Check the exact UTF-8 size before allocating the encoded payload.
if len(value) > limit:
raise ValueError("STRING payload exceeds byte limit")
size = 0
for character in value:
codepoint = ord(character)
if 0xD800 <= codepoint <= 0xDFFF:
raise ValueError("STRING value must be valid UTF-8 text")
size += 1 + (codepoint > 0x7F) + (codepoint > 0x7FF) + (codepoint > 0xFFFF)
if size > limit:
raise ValueError("STRING payload exceeds byte limit")
return {"encoding": "utf8-v1"}, value.encode("utf-8")
def encode_payload(kind: str, value, *, max_bytes: int = MAX_PAYLOAD_BYTES) -> tuple[dict, bytes]:
"""Encode a typed value without dtype conversion, normalization or GPU use.
IMAGE/MASK accept dense CPU tensors; AUDIO accepts a waveform CPU tensor
plus a positive integer sample_rate. STRING accepts exact Unicode text.
Invalid values and exceeded budgets raise nonsensitive ValueError messages.
"""
_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)
sample_rate = None
if kind == "AUDIO":
if not isinstance(value, dict) or set(value) != {"waveform", "sample_rate"}:
raise ValueError("AUDIO value must contain waveform and sample_rate")
sample_rate = _sample_rate(value["sample_rate"])
value = value["waveform"]
if not isinstance(value, torch.Tensor):
raise ValueError("Tensor value must be a torch.Tensor")
if value.device.type != "cpu":
raise ValueError("Tensor value must be on CPU")
if value.layout != torch.strided or value.is_nested:
raise ValueError("Tensor value must have a dense strided layout")
if value.dtype not in _DTYPE_NAMES:
raise ValueError("Unsupported tensor dtype")
dtype_name = _DTYPE_NAMES[value.dtype]
_, itemsize = _DTYPES[dtype_name]
shape = list(value.shape)
size = _tensor_size(kind, shape, itemsize, limit)
tensor = value.detach().resolve_neg().contiguous()
if not torch.isfinite(tensor).all().item():
raise ValueError("Tensor values must be finite")
# A uint16 view preserves BF16 bits without conversion through float32.
if tensor.dtype == torch.bfloat16:
tensor = tensor.view(torch.uint16)
# The live, contiguous CPU tensor owns this pointer for the entire copy.
payload = _native_little_endian(ctypes.string_at(tensor.data_ptr(), size), itemsize)
descriptor = {
"encoding": "tensor-le-v1", "dtype": dtype_name,
"shape": shape, "layout": _LAYOUTS[kind],
}
if kind == "AUDIO":
descriptor["sample_rate"] = sample_rate
return descriptor, payload
def decode_payload(kind: str, descriptor: dict, payload: bytes, *, max_bytes: int = MAX_PAYLOAD_BYTES):
"""Decode validated wire bytes into independent contiguous CPU tensors/text.
All metadata, shape products and byte lengths are checked before tensor
allocation. No GPU device discovery, transfers or object deserialization
occur. AUDIO returns {'waveform': tensor, 'sample_rate': int}.
"""
_validate_kind(kind)
limit = _byte_limit(max_bytes)
if not isinstance(descriptor, dict):
raise ValueError("Payload descriptor must be a dictionary")
if not isinstance(payload, bytes):
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
return decode_image(descriptor, payload, max_bytes=limit)
if kind == "STRING":
if set(descriptor) != {"encoding"} or descriptor["encoding"] != "utf8-v1":
raise ValueError("Invalid STRING descriptor encoding or fields")
try:
return payload.decode("utf-8")
except UnicodeDecodeError:
raise ValueError("STRING payload must be valid UTF-8 bytes") from None
fields = {"encoding", "dtype", "shape", "layout"}
if kind == "AUDIO":
fields.add("sample_rate")
if set(descriptor) != fields:
raise ValueError("Invalid tensor descriptor fields")
if descriptor["encoding"] != "tensor-le-v1":
raise ValueError("Unsupported tensor descriptor encoding")
if descriptor["layout"] != _LAYOUTS[kind]:
raise ValueError("Invalid tensor descriptor layout")
dtype_name = descriptor["dtype"]
if type(dtype_name) is not str or dtype_name not in _DTYPES:
raise ValueError("Unsupported tensor descriptor dtype")
dtype, itemsize = _DTYPES[dtype_name]
size = _tensor_size(kind, descriptor["shape"], itemsize, limit)
if len(payload) != size:
raise ValueError("Tensor payload length does not match shape and dtype")
sample_rate = _sample_rate(descriptor["sample_rate"]) if kind == "AUDIO" else None
# frombuffer is always CPU. Copy immutable wire bytes into a writable private
# buffer, then clone so the returned tensor owns its storage independently.
buffer = bytearray(_native_little_endian(payload, itemsize))
storage_dtype = torch.uint16 if dtype == torch.bfloat16 else dtype
tensor = torch.frombuffer(buffer, dtype=storage_dtype).view(dtype)
if not torch.isfinite(tensor).all().item():
raise ValueError("Tensor payload values must be finite")
tensor = tensor.clone().reshape(descriptor["shape"])
if kind == "AUDIO":
return {"waveform": tensor, "sample_rate": sample_rate}
return tensor
+254
View File
@@ -0,0 +1,254 @@
"""Versioned, bounded media envelopes shared by master and endpoint bridge.
No provider or ComfyUI bootstrap imports. Network I/O is supplied by the caller,
which must enforce public-HTTPS resolution, redirect and streaming byte limits.
"""
import base64
import binascii
import hashlib
import json
import re
from datetime import datetime, timezone
from urllib.parse import urlsplit
PROFILE = "comfy-media-v1"
VERSION = 1
MAX_ARTIFACT_BYTES = 256 * 1024 * 1024
MAX_TOTAL_BYTES = 512 * 1024 * 1024
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")
def identifier(value):
if not isinstance(value, str) or not _ID.fullmatch(value):
raise ValueError("Invalid media identifier")
return value
def bounded_int(value, maximum, *, minimum=0):
if type(value) is not int or not minimum <= value <= maximum:
raise ValueError("Invalid media count or size")
return value
def validate_json(value, *, max_bytes=MAX_DESCRIPTOR_BYTES):
"""Bound metadata depth and size without accepting non-finite JSON numbers."""
def walk(item, depth=0):
if depth > 12:
raise ValueError("Media metadata nesting exceeds limit")
if isinstance(item, dict):
if len(item) > 4096 or any(not isinstance(k, str) for k in item):
raise ValueError("Invalid media metadata object")
for child in item.values():
walk(child, depth + 1)
elif isinstance(item, list):
if len(item) > 4096:
raise ValueError("Media metadata list exceeds limit")
for child in item:
walk(child, depth + 1)
elif item is not None and type(item) not in (str, int, float, bool):
raise ValueError("Invalid media metadata value")
walk(value)
try:
raw = json.dumps(value, allow_nan=False, separators=(",", ":")).encode("utf-8")
except (ValueError, TypeError, UnicodeError, RecursionError):
raise ValueError("Invalid media JSON") from None
if len(raw) > max_bytes:
raise ValueError("Media JSON exceeds byte limit")
return raw
def validate_binding(binding):
if not isinstance(binding, dict) or set(binding) != BINDING_FIELDS:
raise ValueError("Invalid media binding")
for field in BINDING_FIELDS - {"socket", "kind"}:
identifier(binding[field])
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
def validate_bindings(bindings):
if not isinstance(bindings, list) or not 1 <= len(bindings) <= MAX_ARTIFACTS:
raise ValueError("Invalid media bindings count")
result = {}
destinations = set()
for binding in bindings:
validate_binding(binding)
bid = binding["binding_id"]
destination = (binding["job_id"], binding["participant_id"], binding["collector_id"], binding["socket"])
if bid in result or destination in destinations:
raise ValueError("Duplicate media binding")
result[bid] = binding
destinations.add(destination)
return result
def _validate_url_transport(transport):
if set(transport) != {"type", "url", "expires_at"}:
raise ValueError("Invalid media URL transport")
url = transport["url"]
if not isinstance(url, str) or len(url) > 8192:
raise ValueError("Invalid media URL")
try:
parsed = urlsplit(url)
if parsed.scheme != "https" or not parsed.hostname or parsed.username or parsed.password or parsed.fragment:
raise ValueError()
if parsed.port not in (None, 443):
raise ValueError()
expiry = datetime.fromisoformat(transport["expires_at"].replace("Z", "+00:00"))
if expiry.tzinfo is None or expiry <= datetime.now(timezone.utc):
raise ValueError()
except (ValueError, TypeError, AttributeError):
raise ValueError("Invalid or expired media URL") from None
def validate_artifact(artifact, *, max_bytes=MAX_ARTIFACT_BYTES):
fields = {"id", "kind", "descriptor", "byte_length", "sha256", "transport"}
if not isinstance(artifact, dict) or set(artifact) != fields:
raise ValueError("Invalid media artifact")
identifier(artifact["id"])
if artifact["kind"] not in ARTIFACT_KINDS:
raise ValueError("Unsupported media kind")
if not isinstance(artifact["descriptor"], dict):
raise ValueError("Invalid media descriptor")
validate_json(artifact["descriptor"])
size = bounded_int(artifact["byte_length"], max_bytes)
digest = artifact["sha256"]
if not isinstance(digest, str) or not _HASH.fullmatch(digest):
raise ValueError("Invalid media checksum")
transport = artifact["transport"]
if not isinstance(transport, dict):
raise ValueError("Invalid media transport")
if transport.get("type") == "inline":
if set(transport) != {"type", "data_base64"} or not isinstance(transport["data_base64"], str):
raise ValueError("Invalid inline media transport")
if len(transport["data_base64"]) != 4 * ((size + 2) // 3):
raise ValueError("Inline media length mismatch")
elif transport.get("type") == "url":
_validate_url_transport(transport)
else:
raise ValueError("Unsupported media transport")
return artifact
def artifact_from_payload(artifact_id, kind, descriptor, payload, *, transport=None, max_bytes=MAX_ARTIFACT_BYTES):
if not isinstance(payload, bytes):
raise ValueError("Media payload must be bytes")
bounded_int(len(payload), max_bytes)
result = {"id": artifact_id, "kind": kind, "descriptor": descriptor,
"byte_length": len(payload), "sha256": hashlib.sha256(payload).hexdigest(),
"transport": transport if transport is not None else {
"type": "inline", "data_base64": base64.b64encode(payload).decode("ascii")}}
return validate_artifact(result, max_bytes=max_bytes)
async def resolve_artifact(artifact, *, download=None, max_bytes=MAX_ARTIFACT_BYTES):
"""download(url, *, max_bytes) must return bounded bytes; never log its URL."""
validate_artifact(artifact, max_bytes=max_bytes)
transport = artifact["transport"]
if transport["type"] == "inline":
try:
payload = base64.b64decode(transport["data_base64"], validate=True)
except (ValueError, binascii.Error):
raise ValueError("Invalid media base64") from None
else:
if download is None:
raise ValueError("Media URL download is not configured")
payload = await download(transport["url"], max_bytes=artifact["byte_length"])
if not isinstance(payload, bytes) or len(payload) != artifact["byte_length"]:
raise ValueError("Media byte length mismatch")
if hashlib.sha256(payload).hexdigest() != artifact["sha256"]:
raise ValueError("Media checksum mismatch")
return payload
def validate_response(response, *, request_id, bindings, max_total_bytes=MAX_TOTAL_BYTES):
"""Validate ALL identity/completeness before downloading or delivering anything."""
expected = validate_bindings(bindings)
identifier(request_id)
if not isinstance(response, dict) or set(response) != {"version", "profile", "request_id", "bindings", "artifacts", "items"}:
raise ValueError("Invalid media response envelope")
if type(response["version"]) is not int or response["version"] != VERSION or response["profile"] != PROFILE:
raise ValueError("Unsupported media response version/profile")
if response["request_id"] != request_id:
raise ValueError("Media request identity mismatch")
returned = response["bindings"]
if not isinstance(returned, list) or len(returned) != len(expected):
raise ValueError("Missing or extra media bindings")
counts = {}
for row in returned:
if not isinstance(row, dict) or set(row) != BINDING_FIELDS | {"count"}:
raise ValueError("Invalid result binding")
binding = {key: row[key] for key in BINDING_FIELDS}
validate_binding(binding)
bid = binding["binding_id"]
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.
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")
indexed = {}
total = 0
for artifact in artifacts:
validate_artifact(artifact)
aid = artifact["id"]
if aid in indexed:
raise ValueError("Duplicate media artifact")
total += artifact["byte_length"]
if total > max_total_bytes:
raise ValueError("Media response exceeds total byte limit")
indexed[aid] = artifact
items = response["items"]
if not isinstance(items, list) or len(items) != sum(counts.values()):
raise ValueError("Missing or extra media items")
ordered = {bid: [None] * count for bid, count in counts.items()}
used = set()
for item in items:
if not isinstance(item, dict) or set(item) != {"binding_id", "item_index", "artifact_id"}:
raise ValueError("Invalid media item")
bid, aid = identifier(item["binding_id"]), identifier(item["artifact_id"])
if bid not in counts or aid not in indexed:
raise ValueError("Unknown media item identity")
index = bounded_int(item["item_index"], counts[bid] - 1)
if ordered[bid][index] is not None or indexed[aid]["kind"] != expected[bid]["kind"]:
raise ValueError("Duplicate or mistyped media item")
ordered[bid][index] = indexed[aid]
used.add(aid)
if used != set(indexed) or any(value is None for values in ordered.values() for value in values):
raise ValueError("Incomplete media result manifest")
return ordered
async def decode_response(response, *, request_id, bindings, download=None, max_total_bytes=MAX_TOTAL_BYTES):
ordered = validate_response(response, request_id=request_id, bindings=bindings, max_total_bytes=max_total_bytes)
# Resolve and decode before exposing a single result to the collector queue.
from .codec import decode_payload
cache = {}
decoded = {}
for bid, artifacts in ordered.items():
values = []
for artifact in artifacts:
aid = artifact["id"]
if aid not in cache:
payload = await resolve_artifact(artifact, download=download)
if artifact["kind"] == "VIDEO":
from .video_codec import decode_video
cache[aid] = decode_video(artifact["descriptor"], payload)
else:
cache[aid] = decode_payload(artifact["kind"], artifact["descriptor"], payload)
values.append(cache[aid])
decoded[bid] = values
return decoded
+798
View File
@@ -0,0 +1,798 @@
"""Pure HTTP executor for media-v1 requests against a local ComfyUI instance.
No ComfyUI, tensor, or provider imports. ``download(url, *, max_bytes)`` must
provide public-HTTPS/redirect/DNS protection and bounded streaming. Optional
``publish(metadata, payload)`` returns a contract-valid, expiring HTTPS transport.
Only identified invocation-owned files are deleted. Optional trusted local
``temp_dir`` enables exporter cleanup after terminal history and collection.
Without it, exporter retention belongs to the ComfyUI temp-directory lifecycle.
The pinned native /interrupt route is not atomic: running-prompt cancellation
fails closed and retains files instead of risking an unrelated successor prompt.
"""
import asyncio
import base64
import binascii
import copy
import hashlib
import json
import math
import os
import re
import stat
import uuid
from pathlib import Path
from urllib.parse import quote, urlsplit
import aiohttp
from .contract import (
BINDING_FIELDS, MAX_ARTIFACT_BYTES, MAX_ARTIFACTS, MAX_TOTAL_BYTES,
PROFILE, VERSION, artifact_from_payload, identifier, resolve_artifact,
validate_artifact, validate_bindings, validate_json, validate_response,
)
MAX_RESPONSE_JSON_BYTES = 10_000_000
MAX_HTTP_JSON_BYTES = 32 * 1024 * 1024
POLL_INTERVAL_SECONDS = 0.25
CANCEL_GRACE_SECONDS = 10.0
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, "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"}
class MediaExecutionUncertain(RuntimeError):
"""The local prompt may still execute; staged files were deliberately retained."""
class _HTTPError(RuntimeError):
def __init__(self, status):
super().__init__(f"Local Comfy HTTP request failed (status {status})")
self.status = status
def _scope_id(value):
identifier(value)
if any(part in ("", ".", "..") for part in value.split("/")):
raise ValueError("Unsafe media scope identifier")
return value
def _json_bytes(value):
return json.dumps(value, allow_nan=False).encode("utf-8")
def _link(value, workflow):
if (not isinstance(value, list) or len(value) != 2 or
not isinstance(value[0], str) or value[0] not in workflow or
type(value[1]) is not int or value[1] < 0):
raise ValueError("Invalid workflow socket link")
return value[0]
def _validate_request(input_data):
if (not isinstance(input_data, dict) or not {"workflow", "media"} <= set(input_data)
or set(input_data) - {"workflow", "media", "images"}):
raise ValueError("Invalid media request envelope")
media = input_data["media"]
if (not isinstance(media, dict) or set(media) !=
{"version", "profile", "request_id", "artifacts", "inputs", "outputs"}):
raise ValueError("Invalid media request controls")
if type(media["version"]) is not int or media["version"] != VERSION or media["profile"] != PROFILE:
raise ValueError("Unsupported media request version/profile")
_scope_id(media["request_id"])
bindings = validate_bindings(media["outputs"])
targets = []
for binding in bindings.values():
for field in BINDING_FIELDS - {"socket", "kind"}:
_scope_id(binding[field])
targets.append(binding["node_id"])
if len(set(targets)) != len(targets):
raise ValueError("Duplicate media exporter node")
workflow = input_data["workflow"]
if not isinstance(workflow, dict) or not workflow:
raise ValueError("Invalid API workflow")
validate_json(workflow, max_bytes=16 * 1024 * 1024)
dependencies = {}
for node_id, node in workflow.items():
_scope_id(node_id)
if (not isinstance(node, dict) or not isinstance(node.get("class_type"), str)
or not node["class_type"] or not isinstance(node.get("inputs"), dict)):
raise ValueError("Invalid API workflow node")
if node["class_type"] == "MediaImportV1":
raise ValueError("Request may not supply its own media import filenames")
if node["class_type"] == "MediaExportV1" and node_id not in targets:
raise ValueError("Undeclared media exporter")
dependencies[node_id] = set()
for name, value in node["inputs"].items():
_scope_id(name)
if isinstance(value, list):
dependencies[node_id].add(_link(value, workflow))
for binding in bindings.values():
node = workflow.get(binding["node_id"])
if node is None or node["class_type"] != "MediaExportV1":
raise ValueError("Missing declared MediaExportV1 node")
controls = node["inputs"]
from .image_transfer import DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options
required = {"request_id", "binding_json", binding["socket"]}
optional = {"transfer_format", "transfer_quality"} if binding["kind"] == "IMAGE" else set()
if not required <= set(controls) or set(controls) - required - optional:
raise ValueError("Media exporter must have exactly its declared socket and controls")
validate_options(controls.get("transfer_format", DEFAULT_FORMAT),
controls.get("transfer_quality", DEFAULT_QUALITY))
if controls["request_id"] != media["request_id"]:
raise ValueError("Media exporter request identity mismatch")
text = controls["binding_json"]
if not isinstance(text, str) or len(text) > 128 * 1024:
raise ValueError("Invalid media exporter binding JSON")
try:
declared = json.loads(text)
except (ValueError, RecursionError):
raise ValueError("Invalid media exporter binding JSON") from None
if declared != binding:
raise ValueError("Media exporter binding identity mismatch")
_link(controls[binding["socket"]], workflow)
# The execution roots are exactly the declared exporters, not every class
# with OUTPUT_NODE=True. Required upstream output utilities may remain.
visited = set()
pending = list(targets)
while pending:
node_id = pending.pop()
if node_id not in visited:
visited.add(node_id)
pending.extend(dependencies[node_id])
if visited != set(workflow):
raise ValueError("Workflow contains nodes outside declared output roots")
if any(set(targets) & sources for sources in dependencies.values()):
raise ValueError("A media exporter cannot be an upstream dependency")
# Kahn traversal bounds cycle checking without relying on recursion depth.
counts = {node: len(sources) for node, sources in dependencies.items()}
consumers = {node: [] for node in workflow}
for node, sources in dependencies.items():
for source in sources:
consumers[source].append(node)
pending = [node for node, count in counts.items() if not count]
processed = 0
while pending:
processed += 1
for consumer in consumers[pending.pop()]:
counts[consumer] -= 1
if not counts[consumer]:
pending.append(consumer)
if processed != len(workflow):
raise ValueError("Workflow contains a dependency cycle")
artifacts = media["artifacts"]
if not isinstance(artifacts, list) or len(artifacts) > MAX_ARTIFACTS:
raise ValueError("Invalid input artifacts count")
indexed, total = {}, 0
for artifact in artifacts:
validate_artifact(artifact)
aid = _scope_id(artifact["id"])
if aid in indexed:
raise ValueError("Duplicate input artifact")
indexed[aid] = artifact
total += artifact["byte_length"]
inputs = media["inputs"]
if not isinstance(inputs, list) or len(inputs) > 4096:
raise ValueError("Invalid input references")
destinations, used = set(), set()
exporters = {binding["node_id"]: binding for binding in bindings.values()}
for ref in inputs:
if not isinstance(ref, dict) or set(ref) != {"node_id", "input_name", "artifact_id"}:
raise ValueError("Invalid media input reference")
for value in ref.values():
_scope_id(value)
node_id, name, aid = ref["node_id"], ref["input_name"], ref["artifact_id"]
if (node_id not in workflow or name not in workflow[node_id]["inputs"]
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:
# 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):
raise ValueError("Every input artifact must be used")
images = input_data.get("images", [])
if not isinstance(images, list) or len(images) > MAX_ARTIFACTS:
raise ValueError("Invalid legacy images count")
names = set()
legacy = []
for image in images:
if not isinstance(image, dict) or set(image) != {"name", "image"}:
raise ValueError("Invalid legacy image")
name, encoded = image["name"], image["image"]
if (not isinstance(name, str) or not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_. -]{0,254}", name)
or name in names or not isinstance(encoded, str)):
raise ValueError("Invalid or duplicate legacy image basename")
names.add(name)
if len(encoded) > 4 * ((MAX_ARTIFACT_BYTES + 2) // 3):
raise ValueError("Legacy image exceeds byte limit")
# Exact decoded length, including padding, before allocating bytes.
size = (len(encoded) // 4) * 3 - (len(encoded) - len(encoded.rstrip("=")))
if len(encoded) % 4 or size < 0 or size > MAX_ARTIFACT_BYTES:
raise ValueError("Invalid legacy image base64 length")
total += size
matches = [node_id for node_id, node in workflow.items()
if node["class_type"] == "LoadImage" and node["inputs"].get("image") == name]
if not matches or any((node_id, "image") in destinations for node_id in matches):
raise ValueError("Unused or conflicting legacy image binding")
legacy.append((encoded, size, matches))
if total > MAX_TOTAL_BYTES:
raise ValueError("Media inputs exceed total byte limit")
return copy.deepcopy(workflow), copy.deepcopy(media), legacy, targets
class _Staging:
"""Exclusive random files under a hashed namespace; never delete siblings."""
def __init__(self, input_dir, request_id):
self.root = Path(input_dir).resolve()
self.key = hashlib.sha256(request_id.encode()).hexdigest()
self.relative = f"distributed-media-v1/{self.key}"
self.names = []
self.fd = None
def _open(self):
self.root.mkdir(parents=True, exist_ok=True)
directory = self.root
for component in ("distributed-media-v1", self.key):
directory = directory / component
if directory.is_symlink():
raise ValueError("Unsafe media staging directory symlink")
directory.mkdir(exist_ok=True, mode=0o700)
if not directory.resolve().is_relative_to(self.root):
raise ValueError("Unsafe media staging directory")
self.fd = os.open(directory, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
def write(self, payload):
if self.fd is None:
self._open()
name = uuid.uuid4().hex + ".bin"
fd = os.open(name, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600, dir_fd=self.fd)
self.names.append(name)
with os.fdopen(fd, "wb") as stream:
stream.write(payload)
return f"{self.relative}/{name}"
def close(self, *, cleanup):
if self.fd is not None:
try:
if cleanup:
for name in self.names:
try:
os.unlink(name, dir_fd=self.fd)
except FileNotFoundError:
pass
finally:
os.close(self.fd)
self.fd = None
if cleanup:
# Never recursively remove a shared namespace or another run.
for directory in (self.root / self.relative, self.root / "distributed-media-v1"):
try:
directory.rmdir()
except OSError:
pass
class _OwnedExports:
"""Best-effort retirement of validated terminal-history files, never a sweep.
Walk every directory component with openat/O_NOFOLLOW, then pin the directory
fd and each regular file's identity. Missing, linked, replaced, or unidentifiable
files stay with the endpoint lifecycle. Never use history-supplied path strings.
"""
def __init__(self, temp_dir, request_id):
self.root = Path(temp_dir).absolute() if temp_dir is not None else None
self.key = hashlib.sha256(request_id.encode()).hexdigest()
self.fd = None
self.files = {}
def _open(self):
if self.fd is not None:
return True
if self.root is None or ".." in self.root.parts:
return False
flags = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW
fd = None
try:
fd = os.open(self.root.anchor, flags)
for component in (*self.root.parts[1:], "distributed-media-v1", self.key):
child = os.open(component, flags, dir_fd=fd)
os.close(fd)
fd = child
self.fd, fd = fd, None
return True
except OSError:
return False
finally:
if fd is not None:
os.close(fd)
def _identity(self, name):
info = os.stat(name, dir_fd=self.fd, follow_symlinks=False)
if not stat.S_ISREG(info.st_mode) or info.st_nlink != 1:
return None
return (info.st_dev, info.st_ino, info.st_size, info.st_mtime_ns, info.st_ctime_ns)
def identify(self, entry, media):
if self.root is None or not isinstance(entry.get("outputs"), dict):
return
# A failed prompt may contain valid completed exporters alongside missing
# or malformed ones. Claim only independently validated declared records.
for binding in media["outputs"]:
node_id = binding["node_id"]
try:
_, files = _manifest({"outputs": {node_id: entry["outputs"].get(node_id)}},
dict(media, outputs=[binding]))
except ValueError:
continue
if not files or not self._open():
continue
for artifact in files:
name = artifact["id"] + ".bin" # _manifest enforces UUID hex.
try:
identity = self._identity(name)
except OSError:
continue
if identity is not None:
self.files.setdefault(name, identity)
def close(self, *, cleanup):
if self.fd is None:
return
try:
if cleanup:
for name, identity in self.files.items():
try:
if self._identity(name) == identity:
os.unlink(name, dir_fd=self.fd)
except OSError:
pass # Endpoint lifecycle owns inaccessible leftovers.
finally:
os.close(self.fd)
self.fd = None
# Deliberately retain directories: they may belong to concurrent runs.
async def _body(response, maximum):
length = response.content_length
if length is not None and length > maximum:
raise ValueError("Local Comfy response exceeds declared byte length")
body = bytearray()
async for chunk in response.content.iter_chunked(64 * 1024):
if len(body) + len(chunk) > maximum:
raise ValueError("Local Comfy response exceeds declared byte length")
body.extend(chunk)
return bytes(body)
class _Execution:
def __init__(self, session, url, stage, media, targets, deadline, exports):
self.session, self.url, self.stage = session, url, stage
self.media, self.targets = media, targets
self.deadline = deadline
self.run_id = uuid.uuid4().hex
self.prompt_id = None
self.attempted = False
self.terminal = False
self.exports = exports
self.interruption_blocked = False
def remaining(self):
remaining = self.deadline - asyncio.get_running_loop().time()
if remaining <= 0:
raise asyncio.TimeoutError("Media execution deadline expired")
return remaining
async def json_request(self, method, path, *, body=None, permit_rejection=False):
async with self.session.request(method, self.url + path, json=body,
timeout=aiohttp.ClientTimeout(total=min(HTTP_TIMEOUT_SECONDS, self.remaining())),
allow_redirects=False) as response:
raw = await _body(response, MAX_HTTP_JSON_BYTES)
# Native queue-delete and interrupt routes return an empty 200.
if response.status == 200 and not raw and method == "POST" and path in {"/queue", "/interrupt"}:
return {}
try:
value = json.loads(raw)
except (ValueError, UnicodeError, RecursionError):
if response.status != 200:
raise _HTTPError(response.status) from None
raise ValueError("Invalid local Comfy JSON response") from None
if response.status != 200:
# Native /prompt rejects the entire graph with a 400 error.
# Other failures may hide acceptance and are never retried.
if (permit_rejection and response.status == 400 and isinstance(value, dict)
and value.get("error") and "prompt_id" not in value):
self.terminal = True
raise _HTTPError(response.status)
if not isinstance(value, dict):
raise ValueError("Invalid local Comfy JSON object")
return value
async def get_json(self, path):
for attempt in range(GET_ATTEMPTS):
try:
return await self.json_request("GET", path)
except (_HTTPError, aiohttp.ClientError, asyncio.TimeoutError) as exc:
if isinstance(exc, _HTTPError) and exc.status != 429 and exc.status < 500:
raise
if attempt == GET_ATTEMPTS - 1:
raise
await asyncio.sleep(POLL_INTERVAL_SECONDS * (2 ** attempt))
async def history(self):
response = await self.get_json("/history/" + quote(self.prompt_id, safe=""))
entry = response.get(self.prompt_id)
if entry is None:
return None
if not isinstance(entry, dict) or not isinstance(entry.get("status"), dict):
raise ValueError("Malformed Comfy prompt history")
status = entry["status"]
messages = status.get("messages", [])
if not isinstance(messages, list):
raise ValueError("Malformed Comfy history messages")
failed = status.get("status_str") in {"error", "interrupted"} or any(
isinstance(message, list) and message and message[0] in {"execution_error", "execution_interrupted"}
for message in messages)
if failed or status.get("completed") is True:
self.terminal = True
self.exports.identify(entry, self.media)
if failed or status.get("status_str") != "success":
raise RuntimeError("Local Comfy prompt execution failed")
return entry
return None
async def submit_and_wait(self, workflow):
body = {"prompt": workflow, "client_id": self.run_id,
"extra_data": {"distributed_media_run_id": self.run_id},
"partial_execution_targets": list(self.targets)}
self.remaining() # CPU staging must not spend past the absolute deadline.
self.attempted = True
response = await self.json_request("POST", "/prompt", body=body, permit_rejection=True)
self.prompt_id = _scope_id(response.get("prompt_id"))
if response.get("node_errors"):
raise RuntimeError("Local Comfy accepted an incomplete set of execution targets")
while True:
entry = await self.history()
if entry is not None:
return entry
await asyncio.sleep(POLL_INTERVAL_SECONDS)
async def settle(self):
"""Bounded recovery; caller shields this from task cancellation."""
deleted = False
while True:
if self.prompt_id:
try:
await self.history()
except RuntimeError:
if not self.terminal:
raise
if self.terminal:
return
queue = await self.get_json("/queue")
groups = []
for key in ("queue_running", "queue_pending"):
rows = queue.get(key)
if not isinstance(rows, list):
raise ValueError("Malformed Comfy queue")
group = {}
for row in rows:
if not isinstance(row, list) or len(row) < 4 or not isinstance(row[3], dict):
raise ValueError("Malformed Comfy queue entry")
pid = _scope_id(row[1])
if pid in group:
raise ValueError("Duplicate Comfy queue identity")
group[pid] = row[3]
groups.append(group)
running, pending = groups
if self.prompt_id is None:
own = [pid for group in groups for pid, extra in group.items()
if extra.get("distributed_media_run_id") == self.run_id
and extra.get("client_id") == self.run_id]
if len(own) > 1:
raise ValueError("Ambiguous Comfy queue ownership")
if own:
self.prompt_id = own[0]
if self.prompt_id in pending and not deleted:
await self.json_request("POST", "/queue", body={"delete": [self.prompt_id]})
deleted = True
elif self.prompt_id in running:
# Pinned ComfyUI server.py checks a snapshot then globally
# interrupts, outside the queue mutex. Supplying prompt_id does
# NOT close that race. No safe capability negotiation exists on
# this route; do not invoke it, even after our own queue read.
self.interruption_blocked = True
raise RuntimeError("Native prompt-scoped interruption unavailable")
elif deleted and self.prompt_id not in running and self.prompt_id not in pending:
self.terminal = True # queue deletion confirmed by exact read-back
return
await asyncio.sleep(POLL_INTERVAL_SECONDS)
async def fetch(self, artifact):
transport = artifact["transport"]
params = {"filename": transport["filename"], "subfolder": transport["subfolder"], "type": "temp"}
for attempt in range(GET_ATTEMPTS):
try:
async with self.session.get(self.url + "/view", params=params, allow_redirects=False,
timeout=aiohttp.ClientTimeout(total=min(HTTP_TIMEOUT_SECONDS, self.remaining()))) as response:
if response.status != 200:
raise _HTTPError(response.status)
payload = await _body(response, artifact["byte_length"])
break
except (_HTTPError, aiohttp.ClientError, asyncio.TimeoutError) as exc:
if isinstance(exc, _HTTPError) and exc.status != 429 and exc.status < 500:
raise
if attempt == GET_ATTEMPTS - 1:
raise
await asyncio.sleep(POLL_INTERVAL_SECONDS * (2 ** attempt))
if len(payload) != artifact["byte_length"]:
raise ValueError("Media output byte length mismatch")
if hashlib.sha256(payload).hexdigest() != artifact["sha256"]:
raise ValueError("Media output checksum mismatch")
return payload
def _manifest(entry, media):
outputs = entry.get("outputs")
if not isinstance(outputs, dict):
raise ValueError("Missing media history outputs")
expected_ids = {binding["node_id"] for binding in media["outputs"]}
if any(node_id not in expected_ids and isinstance(value, dict) and "media_exports" in value
for node_id, value in outputs.items()):
raise ValueError("Unexpected media export history")
response = {"version": VERSION, "profile": PROFILE, "request_id": media["request_id"],
"bindings": [], "artifacts": [], "items": []}
files = []
subfolder = "distributed-media-v1/" + hashlib.sha256(media["request_id"].encode()).hexdigest()
for binding in media["outputs"]:
node = outputs.get(binding["node_id"], {})
records = node.get("media_exports") if isinstance(node, dict) else None
if not isinstance(records, list) or len(records) != 1:
raise ValueError("Missing or duplicate media export result")
record = records[0]
if (not isinstance(record, dict) or set(record) != {"request_id", "binding", "artifacts", "items"}
or record["request_id"] != media["request_id"]):
raise ValueError("Media export request identity mismatch")
if not isinstance(record["artifacts"], list) or len(record["artifacts"]) > MAX_ARTIFACTS:
raise ValueError("Invalid exported artifacts count")
converted = []
for artifact in record["artifacts"]:
if not isinstance(artifact, dict) or not isinstance(artifact.get("transport"), dict):
raise ValueError("Invalid exported artifact")
transport = artifact["transport"]
aid = artifact.get("id")
if (not isinstance(aid, str) or not re.fullmatch(r"[a-f0-9]{32}", aid)
or transport != {"type": "comfy-file", "filename": aid + ".bin",
"subfolder": subfolder, "folder_type": "temp"}):
raise ValueError("Unsafe or unowned exported media file")
converted.append(dict(artifact, transport=dict(_VALIDATION_TRANSPORT)))
partial = dict(response, bindings=[record["binding"]], artifacts=converted, items=record["items"])
validate_response(partial, request_id=media["request_id"], bindings=[binding])
response["bindings"].append(copy.deepcopy(record["binding"]))
response["artifacts"].extend(converted)
response["items"].extend(copy.deepcopy(record["items"]))
files.extend(copy.deepcopy(record["artifacts"]))
# Validate the whole manifest BEFORE the first /view or publisher callback.
validate_response(response, request_id=media["request_id"], bindings=media["outputs"])
return response, files
def _publication_plan(response, files, publish):
skeleton = dict(response, artifacts=[dict(a, transport={"type": "inline", "data_base64": ""}) for a in files])
base_size = len(_json_bytes({"media": skeleton}))
sizes = {a["id"]: 4 * ((a["byte_length"] + 2) // 3) for a in files}
predicted = base_size + sum(sizes.values())
published = set()
if predicted > MAX_RESPONSE_JSON_BYTES and publish is None:
raise ValueError("Media response exceeds inline JSON limit; a publisher is required")
# Reserve enough space for a maximum-size escaped URL, including metadata.
# Small remaining inline values preserve their exact bytes and ordering.
for aid in sorted(sizes, key=sizes.get, reverse=True):
if predicted <= MAX_RESPONSE_JSON_BYTES:
break
if sizes[aid] > 65_536:
published.add(aid)
predicted -= sizes[aid] - 65_536
if predicted > MAX_RESPONSE_JSON_BYTES:
raise ValueError("Media response metadata exceeds JSON limit")
return published
async def _run(execution, workflow, legacy, download, publish):
media = execution.media
imports = {}
for artifact in media["artifacts"]:
execution.remaining()
payload = await resolve_artifact(artifact, download=download)
execution.remaining()
filename = execution.stage.write(payload)
node_id = "distributed-media-import:" + uuid.uuid4().hex
while node_id in workflow:
node_id = "distributed-media-import:" + uuid.uuid4().hex
metadata = {key: value for key, value in artifact.items() if key != "transport"}
workflow[node_id] = {"class_type": "MediaImportV1", "inputs": {
"artifact_json": json.dumps(metadata, allow_nan=False), "filename": filename}}
imports[artifact["id"]] = [node_id, _IMPORT_SLOTS[artifact["kind"]]]
for ref in media["inputs"]:
workflow[ref["node_id"]]["inputs"][ref["input_name"]] = list(imports[ref["artifact_id"]])
for encoded, size, matches in legacy:
execution.remaining()
try:
payload = base64.b64decode(encoded, validate=True)
except (ValueError, binascii.Error):
raise ValueError("Invalid legacy image base64") from None
if len(payload) != size:
raise ValueError("Legacy image byte length mismatch")
execution.remaining()
filename = execution.stage.write(payload)
for node_id in matches:
workflow[node_id]["inputs"]["image"] = filename
# Replacements can disconnect file loaders whose original endpoint-local
# paths are unavailable. Keep exactly the post-rewrite exporter closure;
# shared ancestors still needed by another branch must survive.
reachable, pending = set(), list(execution.targets)
while pending:
node_id = pending.pop()
if node_id in reachable:
continue
reachable.add(node_id)
pending.extend(_link(value, workflow) for value in workflow[node_id]["inputs"].values()
if isinstance(value, list))
workflow = {node_id: node for node_id, node in workflow.items() if node_id in reachable}
entry = await execution.submit_and_wait(workflow)
response, files = _manifest(entry, media)
published = _publication_plan(response, files, publish)
artifacts = []
for artifact in files:
payload = await execution.fetch(artifact)
transport = None
if artifact["id"] in published:
metadata = {key: copy.deepcopy(value) for key, value in artifact.items() if key != "transport"}
transport = await publish(metadata, payload)
if not isinstance(transport, dict) or transport.get("type") != "url":
raise ValueError("Publisher must return an expiring HTTPS URL transport")
artifacts.append(artifact_from_payload(artifact["id"], artifact["kind"], artifact["descriptor"],
payload, transport=transport))
response["artifacts"] = artifacts
validate_response(response, request_id=media["request_id"], bindings=media["outputs"])
result = {"media": response}
if len(_json_bytes(result)) > MAX_RESPONSE_JSON_BYTES:
raise ValueError("Media response exceeds JSON byte limit")
execution.remaining()
return result
def _release_uncertain_work(task):
_UNCERTAIN_WORK.discard(task)
try:
task.result()
except BaseException:
pass
async def _drain_work(work):
"""Cancel once; a misbehaving callback cannot make teardown unbounded."""
if not work.done():
work.cancel()
deadline = asyncio.get_running_loop().time() + DRAIN_GRACE_SECONDS
while not work.done():
remaining = deadline - asyncio.get_running_loop().time()
if remaining <= 0:
# Keep a strong reference until it exits, but never delete its files
# asynchronously: uncertain ownership belongs to endpoint recovery.
_UNCERTAIN_WORK.add(work)
work.add_done_callback(_release_uncertain_work)
return False
try:
# Unlike wait_for, wait's timeout does not wait for cancellation.
await asyncio.wait({work}, timeout=remaining)
except asyncio.CancelledError:
continue
_release_uncertain_work(work)
return True
async def execute_media_request(input_data, *, comfy_url, input_dir, temp_dir=None,
timeout_seconds=1800, session=None, download=None, publish=None) -> dict:
"""Validate, stage, submit once, collect exact media, and retire owned inputs.
``timeout_seconds`` covers staging, HTTP, polling, and publishing. A separate
bounded cancellation grace is allowed ONLY to establish the local prompt is
terminal. On uncertain acceptance/termination, raise MediaExecutionUncertain
and retain input files rather than racing an executing prompt. Never retry an
ambiguous submission, even when the exception is a timeout.
Async callback teardown has its own bounded grace; non-cooperative callbacks
also raise MediaExecutionUncertain and retain files, even after terminal history.
``temp_dir`` is a trusted endpoint setting for this ComfyUI's local temp root,
never a request-supplied path. Only validated exporter UUID.bin files beneath
distributed-media-v1/sha256(request_id) are retired, after all collection and
publisher work has unwound. Unidentified files remain lifecycle-owned. With
``temp_dir=None`` all exporter files remain lifecycle-owned as before.
The pinned native /interrupt route has a check/global-interrupt race. Running
prompts therefore fail closed with MediaExecutionUncertain, without posting
/interrupt. A prompt-scoped request body alone is not a safety guarantee.
"""
if (type(timeout_seconds) not in (int, float) or not math.isfinite(timeout_seconds)
or timeout_seconds <= 0):
raise ValueError("Invalid media execution timeout")
if not isinstance(comfy_url, str):
raise ValueError("Invalid local Comfy URL")
try:
parsed = urlsplit(comfy_url)
if (parsed.scheme not in {"http", "https"} or not parsed.hostname or parsed.username
or parsed.password or parsed.query or parsed.fragment):
raise ValueError()
parsed.port
except ValueError:
raise ValueError("Invalid local Comfy URL") from None
deadline = asyncio.get_running_loop().time() + timeout_seconds
workflow, media, legacy, targets = _validate_request(input_data)
stage = _Staging(input_dir, media["request_id"])
exports = _OwnedExports(temp_dir, media["request_id"])
owned_session = session is None
if owned_session:
session = aiohttp.ClientSession()
execution = _Execution(session, comfy_url.rstrip("/"), stage, media, targets, deadline, exports)
work = None
try:
remaining = execution.remaining()
work = asyncio.create_task(_run(execution, workflow, legacy, download, publish))
return await asyncio.wait_for(asyncio.shield(work), remaining)
except BaseException as exc:
# Retire files only after readers unwind. Cooperative async callbacks
# normally drain immediately; an uncooperative one requires quarantine.
if work is not None and not await _drain_work(work):
raise MediaExecutionUncertain(
"Media callback teardown unconfirmed; owned files retained; "
"quarantine this endpoint before reuse; do not resubmit automatically"
) from exc
if execution.attempted and not execution.terminal:
execution.deadline = asyncio.get_running_loop().time() + CANCEL_GRACE_SECONDS
recovery = asyncio.create_task(asyncio.wait_for(execution.settle(), CANCEL_GRACE_SECONDS))
# Shield the bounded recovery even from repeated task.cancel() calls.
while not recovery.done():
try:
await asyncio.shield(recovery)
except asyncio.CancelledError:
continue
except Exception:
break
try:
recovery.result()
except BaseException:
pass
if not execution.terminal:
raise MediaExecutionUncertain(
"Local Comfy prompt terminal state unconfirmed; owned staged files retained "
f"under {stage.relative}; prompt_id={execution.prompt_id or 'unknown'}; "
"do not resubmit this request automatically"
+ ("; native prompt-scoped interruption unavailable (non-atomic /interrupt route)"
if execution.interruption_blocked else "")
) from exc
raise
finally:
try:
try:
exports.close(cleanup=(work is None or work.done()) and execution.terminal)
finally:
stage.close(cleanup=(work is None or work.done()) and
(not execution.attempted or execution.terminal))
finally:
if owned_session:
await session.close()
+634
View File
@@ -0,0 +1,634 @@
"""Bounded optional IMAGE returns, independent of ComfyUI and Distributed.
Exact mode is the existing tensor wire format, unchanged. All other modes clip
and truncate RGB/luma to uint8 (as the Collector does), normalize in float32,
then cast to the source dtype. RGBA always carries an exact tensor alpha
sidecar, including PNG/WebP. ``legacy_png`` means uncompressed binary PNG here,
not the old Collector's JSON/base64 transport.
The compressed wire format concatenates length-delimited stills or one Matroska
video, followed by optional alpha bytes. No archive, path, URL or pickle is
accepted. Limits cover wire bytes, descriptor structure, output tensor bytes,
and aggregate padded pixel buffers, not process RSS/native codec overhead.
Video preflight reads EBML track dimensions and block counts before creating a
decoder. We extract packets ourselves, avoiding native container stream probing
(which can decode before caller checks). FFmpeg max_pixels also bounds
bitstream-driven resolution changes, including native buffer alignment.
Software codecs only, one thread, no fallback to another encoder/format.
"""
import io
import struct
import zlib
from fractions import Fraction
import torch
from . import codec
from .codec import MAX_PAYLOAD_BYTES
from .image_transfer import DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options
MAX_FRAMES = 240
MAX_DIMENSION = 8192
MAX_DESCRIPTOR_BYTES = 64 * 1024
MAX_CONTAINER_ELEMENTS = 16384
_VIDEO_CODECS = {
"ffv1": ("ffv1", None, "bgr0"),
"h264": ("libx264", 51, "yuv420p"),
"h265": ("libx265", 51, "yuv420p"),
"av1": ("libsvtav1", 63, "yuv420p"),
}
_CODEC_IDS = {"ffv1": b"V_FFV1", "h264": b"V_MPEG4/ISO/AVC",
"h265": b"V_MPEGH/ISO/HEVC", "av1": b"V_AV1"}
class _BoundedIO(io.BytesIO):
"""Enforce the final file size on every write, including muxer flush/seek."""
def __init__(self, limit):
super().__init__()
self.limit = limit
def write(self, data):
if self.tell() + len(data) > self.limit:
raise ValueError("Encoded IMAGE exceeds byte limit")
return super().write(data)
def seek(self, offset, whence=0):
if whence == 0:
position = offset
elif whence == 1:
position = self.tell() + offset
elif whence == 2:
position = self.getbuffer().nbytes + offset
else:
raise ValueError("Invalid image buffer seek")
if not 0 <= position <= self.limit:
raise ValueError("Encoded IMAGE exceeds byte limit")
return super().seek(offset, whence)
def _quality_to_crf(quality, maximum):
return int(round(maximum * (1.0 - quality / 100.0) ** 0.6))
def _padded_size(fmt, height, width):
if fmt in ("h264", "h265", "av1"):
height, width = height + height % 2, width + width % 2
# libx265 / SVT-AV1 cannot safely encode arbitrarily small pictures.
if fmt in ("h265", "av1"):
height, width = max(64, height), max(64, width)
return height, width
def _shape_budget(shape, fmt, limit, itemsize=4, alpha_size=0, coded_frames=None):
size = codec._tensor_size("IMAGE", shape, itemsize, limit)
count, height, width, _ = shape
if count > MAX_FRAMES or max(height, width) > MAX_DIMENSION:
raise ValueError("IMAGE frame count or dimensions exceed limit")
ph, pw = _padded_size(fmt, height, width)
# Reserve output and padded RGBA-sized decoded buffers BEFORE allocation.
bh, bw = ((ph + 63) // 64 * 64, (pw + 63) // 64 * 64) if fmt in _VIDEO_CODECS else (ph, pw)
# Slice/OBU counts conservatively bound native frames, including AV1 hidden
# frames and multiple outputs from one packet. Reserve 16 bytes per pixel
# for hostile high-bit-depth streams, not only our usual 8-bit RGB frames.
native_count = max(count, coded_frames or count)
if native_count > MAX_FRAMES * 4:
raise ValueError("IMAGE coded frame count exceeds limit")
pixels = codec._tensor_size("IMAGE", [native_count, bh, bw, 4], 4, limit)
if size + pixels + alpha_size > limit:
raise ValueError("IMAGE decoded aggregate exceeds byte limit")
return ph, pw
def _descriptor_budget(value):
"""Bound the JSON-shaped tree without serializing attacker-sized values."""
remaining = MAX_DESCRIPTOR_BYTES
def walk(obj, depth):
nonlocal remaining
if depth > 5:
raise ValueError("IMAGE descriptor nesting exceeds limit")
remaining -= 8
if remaining < 0:
raise ValueError("IMAGE descriptor exceeds byte limit")
if obj is None:
return
if type(obj) is str:
remaining -= len(obj) * 6
elif type(obj) is int:
if obj.bit_length() > 63:
raise ValueError("IMAGE descriptor integer exceeds limit")
remaining -= 20
elif type(obj) is list:
if len(obj) > MAX_FRAMES:
raise ValueError("IMAGE descriptor list exceeds limit")
for item in obj:
walk(item, depth + 1)
elif type(obj) is dict:
if len(obj) > 10:
raise ValueError("IMAGE descriptor fields exceed limit")
for key, item in obj.items():
if type(key) is not str:
raise ValueError("Invalid IMAGE descriptor key")
walk(key, depth + 1)
walk(item, depth + 1)
else:
raise ValueError("Invalid IMAGE descriptor value")
if remaining < 0:
raise ValueError("IMAGE descriptor exceeds byte limit")
walk(value, 0)
def _alpha_metadata(alpha, shape, limit):
if shape[-1] != 4:
if alpha is not None:
raise ValueError("Unexpected IMAGE alpha sidecar")
return 0, torch.float32
if type(alpha) is not dict or set(alpha) != {"descriptor", "size"}:
raise ValueError("Missing or invalid IMAGE alpha sidecar")
d = alpha["descriptor"]
if type(d) is not dict or set(d) != {"encoding", "dtype", "shape", "layout"}:
raise ValueError("Invalid IMAGE alpha descriptor")
if d["encoding"] != "tensor-le-v1" or d["layout"] != "BHWC" or d["shape"] != shape[:3] + [1]:
raise ValueError("Invalid IMAGE alpha shape or encoding")
if type(d["dtype"]) is not str or d["dtype"] not in codec._DTYPES:
raise ValueError("Invalid IMAGE alpha dtype")
dtype, itemsize = codec._DTYPES[d["dtype"]]
size = codec._tensor_size("IMAGE", d["shape"], itemsize, limit)
if type(alpha["size"]) is not int or alpha["size"] != size:
raise ValueError("Invalid IMAGE alpha size")
return size, torch.float64 if dtype == torch.float64 else torch.float32
def _import_av():
try:
import av
except ImportError:
raise ValueError("PyAV is required for IMAGE video transfer") from None
return av
def _encode_still(frame, fmt, quality, limit):
from PIL import Image
sink = _BoundedIO(limit)
if frame.shape[-1] == 1:
frame = frame[..., 0]
image = Image.fromarray(frame)
if fmt in ("legacy_png", "png"):
image.save(sink, format="PNG", compress_level=0 if fmt == "legacy_png" else 4)
elif fmt == "jpeg":
image.save(sink, format="JPEG", quality=quality, subsampling=0 if quality >= 90 else 2)
else:
image.save(sink, format="WEBP", quality=quality, method=4)
return sink.getvalue()
def _encode_video(frames, fmt, quality, height, width, limit):
import numpy as np
av = _import_av()
encoder, maximum, pix_fmt = _VIDEO_CODECS[fmt]
try:
av.codec.Codec(encoder, "w")
except Exception:
raise ValueError("IMAGE software encoder unavailable: " + encoder) from None
sink = _BoundedIO(limit)
options = {}
if maximum is not None:
options["crf"] = str(_quality_to_crf(quality, maximum))
if fmt == "h264":
options["preset"] = "medium"
elif fmt == "h265":
options.update(preset="fast", **{"x265-params": "log-level=error:pools=none:frame-threads=1"})
elif fmt == "av1":
options.update(preset="8", **{"svtav1-params": "lp=1"})
else:
options.update(level="3", slicecrc="1")
try:
with av.open(sink, mode="w", format="matroska") as container:
stream = container.add_stream(encoder, rate=24)
stream.width, stream.height = width, height
stream.pix_fmt = pix_fmt
stream.codec_context.thread_count = 1
stream.options = options
for index, frame in enumerate(frames):
if frame.shape[-1] == 1:
frame = np.repeat(frame, 3, axis=-1)
h, w = frame.shape[:2]
frame = np.pad(frame, ((0, height - h), (0, width - w), (0, 0)), mode="edge")
vf = av.VideoFrame.from_ndarray(frame, format="rgb24")
vf.pts = index
vf.time_base = Fraction(1, 24)
for packet in stream.encode(vf):
container.mux(packet)
for packet in stream.encode():
container.mux(packet)
data = sink.getvalue()
if len(data) > limit:
raise ValueError("Encoded IMAGE exceeds byte limit")
return data
except av.error.FFmpegError:
raise ValueError("IMAGE software video encode failed (codec, dimensions or byte limit)") from None
except ValueError:
raise
except Exception:
raise ValueError("IMAGE software video encode failed (codec, dimensions or byte limit)") from None
def encode_image(value, *, transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY,
max_bytes=MAX_PAYLOAD_BYTES):
"""Return (descriptor, bytes); reject unsupported inputs, never fall back."""
fmt, quality = validate_options(transfer_format, transfer_quality)
if fmt == "exact":
return codec.encode_payload("IMAGE", value, max_bytes=max_bytes)
limit = codec._byte_limit(max_bytes)
if not isinstance(value, torch.Tensor) or value.device.type != "cpu":
raise ValueError("IMAGE must be a CPU tensor")
if value.layout != torch.strided or value.is_nested:
raise ValueError("IMAGE must have dense strided layout")
if value.dtype not in codec._DTYPE_NAMES:
raise ValueError("Unsupported IMAGE dtype")
shape = list(value.shape)
codec._tensor_size("IMAGE", shape, value.element_size(), limit)
alpha_size = shape[0] * shape[1] * shape[2] * value.element_size() if shape[-1] == 4 else 0
itemsize = value.element_size()
ph, pw = _shape_budget(shape, fmt, limit, itemsize, alpha_size)
if not torch.isfinite(value).all().item():
raise ValueError("IMAGE values must be finite")
value = value.detach().resolve_neg()
alpha, alpha_payload = None, b""
if shape[-1] == 4:
ad, alpha_payload = codec.encode_payload("IMAGE", value[..., 3:4], max_bytes=limit)
alpha = {"descriptor": ad, "size": len(alpha_payload)}
import numpy as np
def frames():
for frame in value:
arr = frame[..., :3].float().numpy()
with np.errstate(over="ignore", invalid="ignore"):
yield np.ascontiguousarray(np.clip(arr * 255.0, 0, 255).astype(np.uint8))
budget = limit - len(alpha_payload)
if fmt in _VIDEO_CODECS:
parts = [_encode_video(frames(), fmt, quality, ph, pw, budget)]
header = _video_header(parts[0], fmt, ph, pw, shape[0])
coded_frames = _coded_frame_bound(parts[0], fmt, header)
_shape_budget(shape, fmt, limit, itemsize, alpha_size, coded_frames)
else:
parts = []
for frame in frames():
part = _encode_still(frame, fmt, quality, budget)
budget -= len(part)
parts.append(part)
descriptor = {"encoding": "image-transfer-v1", "format": fmt, "quality": quality,
"shape": shape, "layout": "BHWC", "dtype": codec._DTYPE_NAMES[value.dtype],
"parts": [len(p) for p in parts], "alpha": alpha}
_descriptor_budget(descriptor)
payload = b"".join(parts) + alpha_payload
if len(payload) > limit:
raise ValueError("Encoded IMAGE exceeds byte limit")
return descriptor, payload
def _still_header(data, fmt):
"""Read dimensions before Pillow (WebP may allocate pixels during open)."""
if fmt in ("png", "legacy_png"):
if data[:8] != b"\x89PNG\r\n\x1a\n":
raise ValueError("Invalid IMAGE PNG header")
offset, dimensions, saw_data, chunks = 8, None, False, 0
while offset < len(data):
chunks += 1
if chunks > MAX_CONTAINER_ELEMENTS:
raise ValueError("IMAGE PNG chunk count exceeds limit")
if offset + 12 > len(data):
raise ValueError("Truncated IMAGE PNG chunk")
length = int.from_bytes(data[offset:offset + 4], "big")
end = offset + 12 + length
if end > len(data):
raise ValueError("Truncated IMAGE PNG chunk")
tag = data[offset + 4:offset + 8]
body = data[offset + 8:end - 4]
if zlib.crc32(data[offset + 4:end - 4]) != int.from_bytes(data[end - 4:end], "big"):
raise ValueError("Invalid IMAGE PNG checksum")
if tag == b"IHDR" and offset == 8 and length == 13:
width, height, depth, color, compression, filtering, interlace = struct.unpack(">IIBBBBB", body)
if depth != 8 or color not in (0, 2) or (compression, filtering, interlace) != (0, 0, 0):
raise ValueError("Unsupported IMAGE PNG layout")
dimensions = (height, width)
elif tag == b"IDAT" and dimensions:
saw_data = True
elif tag == b"IEND" and not length and end == len(data) and saw_data:
return dimensions
else:
raise ValueError("Unsupported IMAGE PNG chunk")
offset = end
raise ValueError("Truncated IMAGE PNG")
if fmt == "webp":
# Our encoder writes a single lossy VP8 still. Reject animation/alpha/
# extended WebP before WebPAnimDecoder can allocate its canvas.
if len(data) < 30 or data[:4] != b"RIFF" or data[8:16] != b"WEBPVP8 ":
raise ValueError("Unsupported IMAGE WebP header")
size = int.from_bytes(data[16:20], "little")
if int.from_bytes(data[4:8], "little") + 8 != len(data) or 20 + size + size % 2 != len(data):
raise ValueError("Invalid IMAGE WebP length")
if data[23:26] != b"\x9d\x01\x2a":
raise ValueError("Invalid IMAGE WebP frame")
width, height = struct.unpack("<HH", data[26:30])
return height & 0x3fff, width & 0x3fff
if data[:2] != b"\xff\xd8" or data[-2:] != b"\xff\xd9":
raise ValueError("Invalid IMAGE JPEG header or trailer")
offset, dimensions = 2, None
while offset + 4 <= len(data):
if data[offset] != 0xff:
raise ValueError("Invalid IMAGE JPEG marker")
marker = data[offset + 1]
size = int.from_bytes(data[offset + 2:offset + 4], "big")
end = offset + 2 + size
if size < 2 or end > len(data):
raise ValueError("Truncated IMAGE JPEG marker")
if marker in (0xc0, 0xc2):
if dimensions or size < 8 or data[offset + 4] != 8 or data[offset + 9] not in (1, 3):
raise ValueError("Unsupported IMAGE JPEG layout")
dimensions = struct.unpack(">HH", data[offset + 5:offset + 9])
elif marker == 0xda:
if dimensions is None:
raise ValueError("Missing IMAGE JPEG dimensions")
return dimensions
offset = end
raise ValueError("Missing IMAGE JPEG pixels")
def _video_header(data, fmt, height, width, count):
"""Strict bounded EBML preflight: one track, no lacing/attachments/encodings.
Native stream probing is not the first dimension check: it may instantiate a
decoder. Validate the Matroska header and count blocks without decompressing.
"""
nodes = 0
def vint(pos, end, identifier=False):
if pos >= end or data[pos] == 0:
raise ValueError("Invalid IMAGE EBML integer")
length = 9 - data[pos].bit_length()
if length > (4 if identifier else 8) or pos + length > end:
raise ValueError("Truncated IMAGE EBML integer")
value = int.from_bytes(data[pos:pos + length], "big")
if not identifier:
value &= (1 << (7 * length)) - 1
if value == (1 << (7 * length)) - 1:
raise ValueError("Unknown-sized IMAGE EBML elements are forbidden")
return value, pos + length
def elements(start, end):
nonlocal nodes
while start < end:
nodes += 1
if nodes > MAX_CONTAINER_ELEMENTS:
raise ValueError("IMAGE container element count exceeds limit")
tag, pos = vint(start, end, True)
length, body = vint(pos, end)
finish = body + length
if finish > end:
raise ValueError("Truncated IMAGE EBML element")
yield tag, body, finish
start = finish
def integer(start, end):
if not 1 <= end - start <= 8:
raise ValueError("Invalid IMAGE EBML numeric field")
return int.from_bytes(data[start:end], "big")
roots = list(elements(0, len(data)))
if len(roots) != 2 or [r[0] for r in roots] != [0x1a45dfa3, 0x18538067]:
raise ValueError("Invalid IMAGE Matroska structure")
tracks, blocks, packets = [], 0, []
for tag, start, end in elements(roots[1][1], roots[1][2]):
if tag == 0x1654ae6b:
for entry, es, ee in elements(start, end):
if entry in (0xbf, 0xec):
continue
if entry != 0xae:
raise ValueError("Unexpected IMAGE track element")
fields, dims = {}, {}
for field, fs, fe in elements(es, ee):
if field in fields:
raise ValueError("Duplicate IMAGE track field")
fields[field] = (fs, fe)
if field == 0x6d80:
raise ValueError("IMAGE track encodings are forbidden")
if field == 0xe0:
for dim, ds, de in elements(fs, fe):
if dim in dims:
raise ValueError("Duplicate IMAGE video field")
dims[dim] = (ds, de)
if any(key not in fields for key in (0xd7, 0x83, 0x86)) or any(key not in dims for key in (0xb0, 0xba)):
raise ValueError("Incomplete IMAGE video track")
if integer(*fields[0xd7]) != 1 or integer(*fields[0x83]) != 1:
raise ValueError("Unsupported IMAGE video track")
if fields[0x86][1] - fields[0x86][0] > 32:
raise ValueError("IMAGE codec identifier exceeds limit")
codec_id = bytes(data[slice(*fields[0x86])])
if codec_id != _CODEC_IDS[fmt] and not (fmt == "ffv1" and codec_id == b"V_MS/VFW/FOURCC"):
raise ValueError("IMAGE video codec mismatch")
if integer(*dims[0xb0]) != width or integer(*dims[0xba]) != height:
raise ValueError("IMAGE video header dimensions mismatch")
if codec_id == b"V_MS/VFW/FOURCC":
cs, ce = fields.get(0x63a2, (0, 0))
if ce - cs < 40 or data[cs + 16:cs + 20] != b"FFV1":
raise ValueError("Invalid IMAGE FFV1 codec header")
if struct.unpack("<ii", data[cs + 4:cs + 12]) != (width, height):
raise ValueError("IMAGE FFV1 header dimensions mismatch")
tracks.append(fields)
elif tag == 0x1f43b675:
def block(bs, be):
nonlocal blocks
track, pos = vint(bs, be)
if track != 1 or pos + 3 >= be or data[pos + 2] & 0x06:
raise ValueError("Invalid or laced IMAGE video block")
blocks += 1
if blocks > count:
raise ValueError("IMAGE video frame count exceeds descriptor")
packets.append((pos + 3, be))
for field, fs, fe in elements(start, end):
if field == 0xa3:
block(fs, fe)
elif field == 0xa0:
group = [e for e in elements(fs, fe) if e[0] == 0xa1]
if len(group) != 1:
raise ValueError("Invalid IMAGE video block group")
block(group[0][1], group[0][2])
elif field not in (0xe7, 0xec, 0xbf):
raise ValueError("Unsupported IMAGE cluster field")
elif tag not in (0x114d9b74, 0x1549a966, 0x1c53bb6b, 0x1254c367, 0xec, 0xbf):
raise ValueError("Unsupported IMAGE Matroska element")
if len(tracks) != 1 or blocks != count:
raise ValueError("IMAGE video track or frame count mismatch")
cs, ce = tracks[0].get(0x63a2, (0, 0))
if not 0 < ce - cs <= MAX_DESCRIPTOR_BYTES:
raise ValueError("IMAGE codec configuration exceeds limit or is missing")
if bytes(data[slice(*tracks[0][0x86])]) == b"V_MS/VFW/FOURCC":
cs += 40 # Skip BITMAPINFOHEADER; FFV1 receives only its codec config.
return bytes(data[cs:ce]), packets
def _coded_frame_bound(data, fmt, header):
"""Conservative per-packet allocation bound before native decode.
FFV1 has one picture per packet. Every AVC/HEVC picture needs at least one
VCL NAL; every AV1 output (including show-existing) needs a frame/header OBU.
Counting slices overestimates pictures, which is safe. Do not trust a
Matroska block count alone: a malicious packet may contain several pictures.
"""
config, packets = header
if fmt == "ffv1":
return len(packets)
total, units = 0, 0
if fmt in ("h264", "h265"):
index = 4 if fmt == "h264" else 21
if len(config) <= index or config[0] != 1:
raise ValueError("Invalid IMAGE video configuration")
length_size = (config[index] & 3) + 1
for start, end in packets:
while start < end:
units += 1
if units > MAX_CONTAINER_ELEMENTS:
raise ValueError("IMAGE coded unit count exceeds limit")
if fmt in ("h264", "h265"):
if start + length_size > end:
raise ValueError("Truncated IMAGE NAL length")
length = int.from_bytes(data[start:start + length_size], "big")
start += length_size
if length < (1 if fmt == "h264" else 2) or length > end - start:
raise ValueError("Truncated IMAGE NAL unit")
kind = (data[start] & 31) if fmt == "h264" else (data[start] >> 1) & 63
total += (1 <= kind <= 5) if fmt == "h264" else (kind <= 31)
else:
tag = data[start]
start += 1
if tag & 0x81 or not tag & 2:
raise ValueError("Unsupported IMAGE AV1 OBU header")
kind = (tag >> 3) & 15
if tag & 4:
start += 1
length = 0
for shift in range(0, 56, 7):
if start >= end:
raise ValueError("Truncated IMAGE AV1 OBU length")
value = data[start]
start += 1
length |= (value & 127) << shift
if not value & 128:
break
else:
raise ValueError("IMAGE AV1 OBU length exceeds limit")
if length > end - start:
raise ValueError("Truncated IMAGE AV1 OBU")
total += kind in (3, 6)
start += length
if total > MAX_FRAMES * 4:
raise ValueError("IMAGE coded frame count exceeds limit")
if total < len(packets):
raise ValueError("IMAGE video has missing coded frames")
return total
def _decode_video(data, fmt, height, width, count, header):
av = _import_av()
maximum = ((height + 63) // 64 * 64) * ((width + 63) // 64 * 64)
ctx = av.CodecContext.create("libdav1d" if fmt == "av1" else "hevc" if fmt == "h265" else fmt, "r")
ctx.width, ctx.height = width, height
ctx.extradata, packets = header
ctx.thread_count = 1
ctx.options = {"max_pixels": str(maximum), "threads": "1", "err_detect": "explode"}
ctx.open()
if "max_pixels" in ctx.options:
raise ValueError("IMAGE decoder does not support max_pixels safety limit")
if (ctx.width, ctx.height) != (width, height):
raise ValueError("IMAGE bitstream dimensions mismatch")
decoded = 0
for bounds in packets + [None]:
packet = av.Packet(data[slice(*bounds)]) if bounds is not None else None
for frame in ctx.decode(packet):
if decoded >= count or (frame.width, frame.height) != (width, height):
raise ValueError("IMAGE decoded dimensions or frame count mismatch")
decoded += 1
yield frame.to_ndarray(format="rgb24")
if decoded != count:
raise ValueError("IMAGE decoded frame count mismatch")
def decode_image(descriptor, payload, *, max_bytes=MAX_PAYLOAD_BYTES):
"""Decode to an independent contiguous CPU tensor after bounded preflight."""
limit = codec._byte_limit(max_bytes)
if not isinstance(descriptor, dict):
raise ValueError("IMAGE descriptor must be a dictionary")
if descriptor.get("encoding") == "tensor-le-v1":
return codec.decode_payload("IMAGE", descriptor, payload, max_bytes=limit)
_descriptor_budget(descriptor)
if set(descriptor) != {"encoding", "format", "quality", "shape", "layout", "dtype", "parts", "alpha"}:
raise ValueError("Invalid IMAGE descriptor fields")
if descriptor["encoding"] != "image-transfer-v1" or descriptor["layout"] != "BHWC":
raise ValueError("Invalid IMAGE descriptor encoding or layout")
fmt, _ = validate_options(descriptor["format"], descriptor["quality"])
if fmt == "exact":
raise ValueError("Exact IMAGE requires the tensor encoding")
if not isinstance(payload, bytes):
raise ValueError("IMAGE payload must be bytes")
if len(payload) > limit:
raise ValueError("IMAGE payload exceeds byte limit")
shape = descriptor["shape"]
if type(descriptor["dtype"]) is not str or descriptor["dtype"] not in codec._DTYPES:
raise ValueError("Unsupported IMAGE descriptor dtype")
dtype, itemsize = codec._DTYPES[descriptor["dtype"]]
codec._tensor_size("IMAGE", shape, itemsize, limit)
alpha_size, _ = _alpha_metadata(descriptor["alpha"], shape, limit)
if alpha_size and descriptor["alpha"]["descriptor"]["dtype"] != descriptor["dtype"]:
raise ValueError("IMAGE alpha dtype mismatch")
ph, pw = _shape_budget(shape, fmt, limit, itemsize, alpha_size)
count, height, width, channels = shape
lengths = descriptor["parts"]
if type(lengths) is not list or len(lengths) != (1 if fmt in _VIDEO_CODECS else count):
raise ValueError("IMAGE part count mismatch")
total = alpha_size
for length in lengths:
if type(length) is not int or not 0 < length <= limit - total:
raise ValueError("IMAGE part length exceeds byte limit")
total += length
if total != len(payload):
raise ValueError("IMAGE payload length does not match parts")
parts, offset = [], 0
for length in lengths:
parts.append(memoryview(payload)[offset:offset + length])
offset += length
try:
if fmt in _VIDEO_CODECS:
video_header = _video_header(parts[0], fmt, ph, pw, count)
coded_frames = _coded_frame_bound(parts[0], fmt, video_header)
_shape_budget(shape, fmt, limit, itemsize, alpha_size, coded_frames)
else:
for part in parts:
if _still_header(part, fmt) != (height, width):
raise ValueError("IMAGE header dimensions mismatch")
# All headers and aggregate budgets are checked before output allocation.
alpha = None
if alpha_size:
alpha = codec.decode_payload("IMAGE", descriptor["alpha"]["descriptor"], payload[offset:], max_bytes=limit)
result = torch.empty(shape, dtype=dtype, device="cpu")
if fmt in _VIDEO_CODECS:
frames = _decode_video(parts[0], fmt, ph, pw, count, video_header)
else:
def stills():
import numpy as np
from PIL import Image
for part in parts:
with Image.open(io.BytesIO(part)) as image:
if image.size != (width, height) or getattr(image, "n_frames", 1) != 1:
raise ValueError("IMAGE still dimensions or frame count mismatch")
if image.mode not in ("L", "RGB"):
raise ValueError("Unsupported IMAGE still pixel mode")
yield np.array(image.convert("RGB"), dtype=np.uint8)
frames = stills()
color = 3 if channels == 4 else channels
for index, frame in enumerate(frames):
# Float32 division followed by the source dtype cast reproduces
# Collector normalized pixels without narrowing exact alpha.
result[index, ..., :color] = torch.from_numpy(frame[:height, :width, :color].copy()).float().div_(255)
if alpha is not None:
result[..., 3:4] = alpha
return result
except ValueError:
raise
except Exception:
raise ValueError("Malformed or unsupported IMAGE compressed payload") from None
+30
View File
@@ -0,0 +1,30 @@
"""Image-return policy shared by the master and standalone endpoint bridge.
This module is deliberately dependency-free so graph compilation does not import
Torch, Pillow, PyAV or ComfyUI. Existing Collector settings stay unchanged.
"""
FORMATS = ("exact", "legacy_png", "png", "ffv1", "jpeg", "webp", "h264", "h265", "av1")
DEFAULT_FORMAT = "exact"
DEFAULT_QUALITY = 90
FORMAT_TOOLTIP = (
"Image returns only; master-local images and other media stay unchanged. "
"exact: original floating-point tensor, no quantization. "
"legacy_png / png / ffv1: lossless 8-bit pixels, not exact float tensors. "
"jpeg / webp: lossy stills; h264 / h265 / av1: lossy frame-batch video. "
"Alpha is preserved separately when the codec cannot carry it. "
"Requires matching transfer support on the worker/endpoint."
)
QUALITY_TOOLTIP = (
"Quality 1–100 for lossy image returns; ignored by exact and lossless modes. "
"JPEG/WebP use it directly; video formats use the existing Collector CRF mapping."
)
def validate_options(transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY):
"""Reject malformed controls rather than repairing or silently substituting."""
if type(transfer_format) is not str or transfer_format not in FORMATS:
raise ValueError("Unsupported image transfer_format")
if type(transfer_quality) is not int or not 1 <= transfer_quality <= 100:
raise ValueError("Image transfer_quality must be an integer from 1 to 100")
return transfer_format, transfer_quality
+184
View File
@@ -0,0 +1,184 @@
"""Internal ComfyUI boundary nodes for the standalone media-v1 handler."""
import hashlib
import json
import uuid
from pathlib import Path
import folder_paths
from comfy_api.v0_0_2 import io
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, VALUE_SOCKETS, ARTIFACT_KINDS,
bounded_int, identifier, validate_binding, validate_json,
)
def _scalar(value):
if isinstance(value, list):
if len(value) != 1:
raise ValueError("Media bridge control must have exactly one value")
return value[0]
return value
def _json(text):
if not isinstance(text, str) or len(text) > 256 * 1024:
raise ValueError("Invalid media bridge JSON")
try:
value = json.loads(text)
except (ValueError, TypeError, RecursionError):
raise ValueError("Invalid media bridge JSON") from None
validate_json(value)
return value
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":
return {"waveform": value["waveform"].detach().cpu(), "sample_rate": value["sample_rate"]}
return value
def encode_value(kind, value, *, transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY):
fmt, quality = validate_options(transfer_format, transfer_quality)
if kind == "IMAGE" and fmt != DEFAULT_FORMAT:
from .image_codec import encode_image
return encode_image(_cpu(kind, value), transfer_format=fmt, transfer_quality=quality)
if kind == "VIDEO":
from .video_codec import encode_video_cpu
return encode_video_cpu(value)
return encode_payload(kind, _cpu(kind, value))
def decode_value(kind, descriptor, payload):
if kind == "VIDEO":
from .video_codec import decode_video
return decode_video(descriptor, payload)
return decode_payload(kind, descriptor, payload)
def export_values(request_id, binding, values, *, transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY):
"""Write bounded owned payloads; history contains metadata, never base64."""
identifier(request_id)
validate_binding(binding)
fmt, quality = validate_options(transfer_format, transfer_quality)
if binding["kind"] != "IMAGE" and fmt != 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()
directory = (root / subfolder).resolve()
if not directory.is_relative_to(root):
raise ValueError("Unsafe media output directory")
directory.mkdir(parents=True, exist_ok=True)
written, artifacts, items = [], [], []
total = 0
try:
for index, value in enumerate(values):
descriptor, payload = encode_value(binding["kind"], value,
transfer_format=fmt, transfer_quality=quality)
validate_json(descriptor)
total += len(payload)
if total > MAX_TOTAL_BYTES:
raise ValueError("Media output exceeds total byte limit")
aid = uuid.uuid4().hex
filename = aid + ".bin"
path = directory / filename
with path.open("xb") as stream:
written.append(path)
stream.write(payload)
artifacts.append({"id": aid, "kind": binding["kind"], "descriptor": descriptor,
"byte_length": len(payload), "sha256": hashlib.sha256(payload).hexdigest(),
"transport": {"type": "comfy-file", "filename": filename,
"subfolder": subfolder, "folder_type": "temp"}})
items.append({"binding_id": binding["binding_id"], "item_index": index, "artifact_id": aid})
except BaseException:
for path in written:
path.unlink(missing_ok=True)
raise
return {"request_id": request_id, "binding": dict(binding, count=len(values)),
"artifacts": artifacts, "items": items}
def import_value(artifact, filename):
"""Read a handler-staged, checksummed file within native media roots only."""
fields = {"id", "kind", "descriptor", "byte_length", "sha256"}
if not isinstance(artifact, dict) or set(artifact) != fields:
raise ValueError("Invalid staged media metadata")
identifier(artifact["id"])
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):
raise ValueError("Invalid staged media filename")
roots = [Path(getter()).resolve() for getter in (
folder_paths.get_input_directory, folder_paths.get_output_directory, folder_paths.get_temp_directory)]
path = Path(folder_paths.get_annotated_filepath(filename)).resolve()
if not any(path.is_relative_to(root) for root in roots) or not path.is_file():
raise ValueError("Staged media file is outside allowed roots or unavailable")
if path.stat().st_size != size:
raise ValueError("Staged media file length mismatch")
with path.open("rb") as stream:
payload = stream.read(size + 1)
if len(payload) != size or hashlib.sha256(payload).hexdigest() != artifact["sha256"]:
raise ValueError("Staged media file failed integrity validation")
return decode_value(artifact["kind"], artifact["descriptor"], payload)
class MediaImportV1(io.ComfyNode):
@classmethod
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(),
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(), "VALUE")))
class MediaExportV1(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(node_id="MediaExportV1", category="distributed/internal", is_deprecated=True,
inputs=[io.String.Input("request_id"), io.String.Input("binding_json"),
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,
optional=True, tooltip=QUALITY_TOOLTIP)],
outputs=[], is_output_node=True, is_input_list=True)
@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, **overrides):
binding = _json(_scalar(binding_json))
validate_binding(binding)
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"]],
transfer_format=_scalar(transfer_format),
transfer_quality=_scalar(transfer_quality))
return io.NodeOutput(ui={"media_exports": [exported]})
+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
+482
View File
@@ -0,0 +1,482 @@
"""Lossless transport of pinned native ComfyUI VIDEO representations.
No save_to/get_components fallback is allowed: an unknown class must not silently
transcode. Tensor components use the shared CPU codec. File bytes are opaque and
retain signed trim/crop parameters; decoded files own independent BytesIOs.
VideoFromList's native constructor flattens its inputs, so wire segments are a
bounded, flat list of exact VideoFromFile objects, never recursive playlists.
Offsets are relative to the containing VIDEO payload. Canonical part order is
images/alpha/audio, or segments/complete_audio. Every byte is covered exactly
once. All descriptor structure and tensor sizes are preflighted before decoding
any media storage. File containers are not parsed until a native consumer reads
them; this adapter does not attempt to validate or repair container codecs.
"""
from fractions import Fraction
import io
import json
import math
import os
from pathlib import Path
import stat
from .codec import MAX_PAYLOAD_BYTES, decode_payload, encode_payload
MAX_SEGMENTS = 64
MAX_METADATA_BYTES = 64 * 1024
MAX_METADATA_DEPTH = 16
_ENCODING = "native-video-v1"
_BASE_FIELDS = {"encoding", "representation", "byte_length"}
_COMPONENT_FIELDS = {"images", "alpha", "audio", "frame_rate", "metadata", "bit_depth", "color_space"}
_FILE_FIELDS = {"data", "start_time", "duration", "crop"}
_SEGMENT_FIELDS = {"segments", "complete_audio"}
_TENSOR_BYTES = {"float16": 2, "bfloat16": 2, "float32": 4, "float64": 8}
def _limit(max_bytes):
if type(max_bytes) is not int or max_bytes < 0:
raise ValueError("VIDEO max_bytes must be a nonnegative integer")
return min(max_bytes, MAX_PAYLOAD_BYTES)
def _native_types():
# Lazy imports keep non-VIDEO users independent of ComfyUI and PyAV.
from comfy_api.latest._input_impl.video_types import (
VideoFromComponents, VideoFromFile, VideoFromList,
)
from comfy_api.latest._util import VideoComponents
return VideoFromComponents, VideoFromFile, VideoFromList, VideoComponents
def _fields(value, fields, label):
if type(value) is not dict or set(value) != fields:
raise ValueError(f"Invalid VIDEO {label} fields")
def _integer(value, *, positive=False):
if type(value) is not int or value < (1 if positive else 0):
raise ValueError("VIDEO length, offset and dimensions must be integers in range")
return value
def _number(value):
if type(value) not in (int, float) or (type(value) is float and not math.isfinite(value)):
raise ValueError("VIDEO trim values must be finite numbers")
return value
def _crop(value, *, wire=False):
if value is None:
return None
if type(value) is not (list if wire else tuple) or len(value) != 4:
raise ValueError("VIDEO crop must contain four integers")
if any(type(part) is not int for part in value):
raise ValueError("VIDEO crop must contain four integers")
# Native crop normalization is a consumer concern, not a transport operation.
return list(value) if not wire else tuple(value)
def _frame_rate(value):
_fields(value, {"numerator", "denominator"}, "frame_rate")
_integer(value["numerator"], positive=True)
_integer(value["denominator"], positive=True)
def _encoding_options(bit_depth, color_space):
_integer(bit_depth, positive=True)
if type(color_space) is not str or color_space not in ("sRGB", "HDR", "HDR PQ"):
raise ValueError("Unsupported native VIDEO color_space")
def _metadata(value):
"""Validate bounded JSON before allocating its serialized form; return a copy."""
if value is None:
return None
if type(value) is not dict:
raise ValueError("VIDEO metadata must be a JSON object or null")
remaining = MAX_METADATA_BYTES
def charge(size):
nonlocal remaining
remaining -= size
if remaining < 0:
raise ValueError("VIDEO metadata exceeds byte limit")
def visit(item, depth):
if depth > MAX_METADATA_DEPTH:
raise ValueError("VIDEO metadata exceeds depth limit")
kind = type(item)
if kind in (dict, list):
# Each member costs at least one JSON byte. Also bounds iteration
# before walking potentially huge containers or cyclic structures.
if len(item) > remaining:
raise ValueError("VIDEO metadata exceeds byte limit")
charge(2 + max(0, len(item) - 1))
if kind is dict:
for key, child in item.items():
if type(key) is not str:
raise ValueError("VIDEO metadata keys must be strings")
visit(key, depth + 1)
charge(1)
visit(child, depth + 1)
else:
for child in item:
visit(child, depth + 1)
elif kind in (str, int, float, bool) or item is None:
if kind is str and len(item) > remaining:
raise ValueError("VIDEO metadata exceeds byte limit")
# Avoid converting an unbounded Python integer to decimal first.
if kind is int and item.bit_length() > MAX_METADATA_BYTES * 4:
raise ValueError("VIDEO metadata exceeds byte limit")
if kind is float and not math.isfinite(item):
raise ValueError("VIDEO metadata numbers must be finite")
try:
charge(len(json.dumps(item, ensure_ascii=True, allow_nan=False)))
except (ValueError, OverflowError):
raise ValueError("Invalid or oversized VIDEO metadata scalar") from None
else:
raise ValueError("VIDEO metadata must contain only JSON values")
visit(value, 0)
return json.loads(json.dumps(value, ensure_ascii=True, allow_nan=False, separators=(",", ":")))
def _attrs(value, names):
attributes = vars(value)
if not all(name in attributes for name in names):
raise ValueError("Native VIDEO attributes do not match the pinned adapter")
return [attributes[name] for name in names]
def _roots(allowed_roots):
if allowed_roots is None:
import folder_paths
allowed_roots = [folder_paths.get_input_directory(), folder_paths.get_output_directory(),
folder_paths.get_temp_directory()]
if type(allowed_roots) not in (list, tuple):
raise ValueError("VIDEO allowed_roots must be a list or tuple")
try:
roots = [Path(root).resolve() for root in allowed_roots]
except (OSError, TypeError, ValueError, RuntimeError):
raise ValueError("Invalid VIDEO allowed roots") from None
return roots
def _read_path(value, limit, allowed_roots):
"""Resolve authorization, reject special files, and bound all reads.
Open each resolved path component with NOFOLLOW using directory fds, so an
intermediate symlink swapped in after resolve() cannot bypass authorization.
NONBLOCK prevents an attacker swapping in a FIFO from hanging the worker.
"""
roots = _roots(allowed_roots)
directory_fd = file_fd = None
try:
path = Path(value).resolve(strict=True)
if not any(path.is_relative_to(root) for root in roots):
raise ValueError("VIDEO file is outside allowed roots")
directory_flags = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW
directory_fd = os.open(path.anchor, directory_flags)
for part in path.parts[1:-1]:
next_fd = os.open(part, directory_flags, dir_fd=directory_fd)
os.close(directory_fd)
directory_fd = next_fd
file_fd = os.open(path.name, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK, dir_fd=directory_fd)
info = os.fstat(file_fd)
if not stat.S_ISREG(info.st_mode):
raise ValueError("VIDEO source must be a regular file")
if info.st_size > limit:
raise ValueError("VIDEO file exceeds byte limit")
# Read at most the stat size plus a sentinel; reject growth or truncation.
chunks = []
count = 0
while count <= info.st_size:
chunk = os.read(file_fd, min(1024 * 1024, info.st_size + 1 - count))
if not chunk:
break
chunks.append(chunk)
count += len(chunk)
if count != info.st_size:
raise ValueError("VIDEO source changed size during bounded read")
return b"".join(chunks)
except (OSError, RuntimeError):
raise ValueError("VIDEO file cannot be read safely") from None
finally:
if file_fd is not None:
os.close(file_fd)
if directory_fd is not None:
os.close(directory_fd)
def _file_bytes(value, limit, allowed_roots):
if type(value) is io.BytesIO:
try:
with value.getbuffer() as view:
if view.nbytes > limit:
raise ValueError("VIDEO file exceeds byte limit")
return bytes(view)
except ValueError:
raise ValueError("VIDEO BytesIO is closed or exceeds byte limit") from None
if isinstance(value, (str, os.PathLike)):
return _read_path(value, limit, allowed_roots)
raise ValueError("VIDEO file must be a local path or exact BytesIO")
def encode_video(value, *, max_bytes=MAX_PAYLOAD_BYTES, allowed_roots=None):
"""Encode an exact native VIDEO without opening/transcoding its container.
Paths are authorized only under allowed_roots (when supplied), otherwise
under ComfyUI input/output/temp roots. Budgets may lower, never raise, the
shared maximum. Unsupported native state raises ValueError, not a fallback.
Non-CPU tensors are rejected; runtime staging belongs to the bridge boundary.
"""
return _encode_video(value, max_bytes=max_bytes, allowed_roots=allowed_roots, stage_cpu=False)
def encode_video_cpu(value, *, max_bytes=MAX_PAYLOAD_BYTES, allowed_roots=None):
"""Bridge-only boundary: detach/stage native tensors, then losslessly encode.
Share the exact-type adapter with encode_video rather than reconstructing
native videos, calling get_components, or duplicating private attributes.
Images, alpha and both audio representations move without dtype conversion;
metadata/options and opaque file bytes/trim/crop take the unchanged path.
"""
return _encode_video(value, max_bytes=max_bytes, allowed_roots=allowed_roots, stage_cpu=True)
def _stage_tensor_cpu(value):
import torch
if not isinstance(value, torch.Tensor):
raise ValueError("Tensor value must be a torch.Tensor")
return value.detach().cpu()
def _encode_video(value, *, max_bytes, allowed_roots, stage_cpu):
limit = _limit(max_bytes)
components_class, file_class, list_class, components_type = _native_types()
descriptor = {"encoding": _ENCODING}
chunks = []
size = 0
def add(payload, nested=None):
nonlocal size
if len(payload) > limit - size:
raise ValueError("VIDEO payload exceeds byte limit")
slot = {"offset": size, "length": len(payload)}
if nested is not None:
slot["descriptor"] = nested
chunks.append(payload)
size += len(payload)
return slot
def tensor(kind, tensor_value):
if stage_cpu:
if kind == "AUDIO":
if not isinstance(tensor_value, dict) or set(tensor_value) != {"waveform", "sample_rate"}:
raise ValueError("AUDIO value must contain waveform and sample_rate")
tensor_value = {**tensor_value, "waveform": _stage_tensor_cpu(tensor_value["waveform"])}
else:
tensor_value = _stage_tensor_cpu(tensor_value)
nested, payload = encode_payload(kind, tensor_value, max_bytes=limit - size)
return add(payload, nested)
if type(value) is components_class:
component, bit_depth, color_space = _attrs(value, [
"_VideoFromComponents__components", "_VideoFromComponents__bit_depth",
"_VideoFromComponents__color_space",
])
if type(component) is not components_type:
raise ValueError("Unsupported native VIDEO components class")
images, fps, audio, metadata, alpha = _attrs(component, [
"images", "frame_rate", "audio", "metadata", "alpha",
])
if type(fps) is not Fraction or fps <= 0:
raise ValueError("VIDEO frame_rate must be a positive Fraction")
_encoding_options(bit_depth, color_space)
descriptor.update(representation="components", bit_depth=bit_depth, color_space=color_space,
frame_rate={"numerator": fps.numerator, "denominator": fps.denominator},
metadata=_metadata(metadata))
descriptor["images"] = tensor("IMAGE", images)
# Native get_components() returns BHW1 alpha; MASK nodes use BHW.
# Select an existing lossless tensor layout without squeezing/converting.
alpha_kind = "IMAGE" if getattr(alpha, "ndim", None) == 4 else "MASK"
descriptor["alpha"] = tensor(alpha_kind, alpha) if alpha is not None else None
if alpha is not None:
_alpha_header(descriptor["alpha"], descriptor["images"])
descriptor["audio"] = tensor("AUDIO", audio) if audio is not None else None
elif type(value) is file_class:
source, start, duration, crop = _attrs(value, [
"_VideoFromFile__file", "_VideoFromFile__start_time", "_VideoFromFile__duration",
"_VideoFromFile__crop",
])
descriptor.update(representation="file", start_time=_number(start), duration=_number(duration),
crop=_crop(crop))
descriptor["data"] = add(_file_bytes(source, limit, allowed_roots))
elif type(value) is list_class:
videos, complete_audio = _attrs(value, ["videos", "complete_audio"])
if type(videos) is not list or not 1 <= len(videos) <= MAX_SEGMENTS:
raise ValueError("VIDEO must have between 1 and 64 segments")
if any(type(segment) is not file_class for segment in videos):
raise ValueError("VIDEO segments must be exact native VideoFromFile objects")
descriptor.update(representation="segments", segments=[])
for segment in videos:
nested, payload = encode_video(segment, max_bytes=limit - size, allowed_roots=allowed_roots)
descriptor["segments"].append(add(payload, nested))
descriptor["complete_audio"] = tensor("AUDIO", complete_audio) if complete_audio is not None else None
else:
raise ValueError("Unsupported native VIDEO class")
descriptor["byte_length"] = size
return descriptor, b"".join(chunks)
def _tensor_header(kind, descriptor, length):
"""Allocation-free preflight of the shared tensor-le-v1 descriptor.
The shared decoder still validates the bytes/finite values. Doing its small
structural check here first prevents allocation for an early part when a
later part's descriptor is already known to be invalid.
"""
fields = {"encoding", "dtype", "shape", "layout"}
if kind == "AUDIO":
fields.add("sample_rate")
_fields(descriptor, fields, "tensor")
layout = {"IMAGE": "BHWC", "MASK": "BHW", "AUDIO": "BCT"}[kind]
if type(descriptor["encoding"]) is not str or descriptor["encoding"] != "tensor-le-v1":
raise ValueError("Invalid VIDEO tensor encoding")
if type(descriptor["layout"]) is not str or descriptor["layout"] != layout:
raise ValueError("Invalid VIDEO tensor layout")
dtype = descriptor["dtype"]
if type(dtype) is not str or dtype not in _TENSOR_BYTES:
raise ValueError("Invalid VIDEO tensor dtype")
shape = descriptor["shape"]
if type(shape) is not list or len(shape) != len(layout):
raise ValueError("Invalid VIDEO tensor rank")
size = _TENSOR_BYTES[dtype]
for dimension in shape:
_integer(dimension, positive=True)
if dimension > length // size:
raise ValueError("VIDEO tensor shape exceeds byte limit")
size *= dimension
if kind == "IMAGE" and shape[-1] not in (1, 3, 4):
raise ValueError("Invalid VIDEO image channels")
if size != length:
raise ValueError("VIDEO tensor length does not match shape")
if kind == "AUDIO":
_integer(descriptor["sample_rate"], positive=True)
def _alpha_header(alpha, images):
_fields(alpha, {"offset", "length", "descriptor"}, "alpha part")
descriptor = alpha["descriptor"]
if type(descriptor) is not dict:
raise ValueError("Invalid VIDEO alpha descriptor")
layout = descriptor.get("layout")
if type(layout) is not str or layout not in ("BHW", "BHWC"):
raise ValueError("Invalid VIDEO alpha layout")
kind = "MASK" if layout == "BHW" else "IMAGE"
_tensor_header(kind, descriptor, _integer(alpha["length"]))
expected = images["descriptor"]["shape"][:3] + ([1] if kind == "IMAGE" else [])
if descriptor["shape"] != expected:
raise ValueError("VIDEO alpha shape must match images")
return kind
def _validate_video(descriptor, size, *, file_only=False):
if type(descriptor) is not dict:
raise ValueError("VIDEO descriptor must be a dictionary")
representation = descriptor.get("representation")
if type(representation) is not str or representation not in ("components", "file", "segments"):
raise ValueError("Unsupported VIDEO representation")
if file_only and representation != "file":
# Depth is at most one: native lists already flatten their segments.
raise ValueError("VIDEO segments cannot contain nested lists or components")
fields = {"components": _COMPONENT_FIELDS, "file": _FILE_FIELDS, "segments": _SEGMENT_FIELDS}[representation]
_fields(descriptor, _BASE_FIELDS | fields, "descriptor")
if type(descriptor["encoding"]) is not str or descriptor["encoding"] != _ENCODING:
raise ValueError("Unsupported VIDEO encoding")
if _integer(descriptor["byte_length"]) != size:
raise ValueError("VIDEO payload length does not match descriptor")
result = dict(descriptor)
cursor = 0
def part(slot, kind=None, video=False):
nonlocal cursor
fields = {"offset", "length", "descriptor"} if kind or video else {"offset", "length"}
_fields(slot, fields, "part")
offset = _integer(slot["offset"])
length = _integer(slot["length"])
if offset != cursor or length > size - cursor:
raise ValueError("VIDEO parts must not overlap, leave gaps or exceed bounds")
cursor += length
if kind:
_tensor_header(kind, slot["descriptor"], length)
if video:
return {**slot, "descriptor": _validate_video(slot["descriptor"], length, file_only=True)}
return slot
if representation == "components":
_frame_rate(descriptor["frame_rate"])
_encoding_options(descriptor["bit_depth"], descriptor["color_space"])
result["metadata"] = _metadata(descriptor["metadata"])
part(descriptor["images"], "IMAGE")
if descriptor["alpha"] is not None:
part(descriptor["alpha"], _alpha_header(descriptor["alpha"], descriptor["images"]))
if descriptor["audio"] is not None:
part(descriptor["audio"], "AUDIO")
elif representation == "file":
_number(descriptor["start_time"])
_number(descriptor["duration"])
_crop(descriptor["crop"], wire=True)
part(descriptor["data"])
else:
segments = descriptor["segments"]
if type(segments) is not list or not 1 <= len(segments) <= MAX_SEGMENTS:
raise ValueError("VIDEO must have between 1 and 64 segments")
result["segments"] = [part(segment, video=True) for segment in segments]
if descriptor["complete_audio"] is not None:
part(descriptor["complete_audio"], "AUDIO")
if cursor != size:
raise ValueError("VIDEO payload contains trailing bytes")
return result
def _decode_validated(descriptor, payload, types):
components_class, file_class, list_class, components_type = types
def data(slot):
return payload[slot["offset"]:slot["offset"] + slot["length"]]
def tensor(kind, slot):
if slot is None:
return None
return decode_payload(kind, slot["descriptor"], data(slot), max_bytes=slot["length"])
if descriptor["representation"] == "file":
# VideoFromFile retains the BytesIO, unlike a temporary filesystem path.
return file_class(io.BytesIO(data(descriptor["data"])), start_time=descriptor["start_time"],
duration=descriptor["duration"], crop=_crop(descriptor["crop"], wire=True))
if descriptor["representation"] == "components":
fps = descriptor["frame_rate"]
return components_class(components_type(
images=tensor("IMAGE", descriptor["images"]),
frame_rate=Fraction(fps["numerator"], fps["denominator"]),
audio=tensor("AUDIO", descriptor["audio"]),
metadata=descriptor["metadata"],
alpha=tensor("IMAGE" if descriptor["alpha"] is not None and
descriptor["alpha"]["descriptor"]["layout"] == "BHWC" else "MASK",
descriptor["alpha"]),
), bit_depth=descriptor["bit_depth"], color_space=descriptor["color_space"])
videos = [_decode_validated(slot["descriptor"], data(slot), types) for slot in descriptor["segments"]]
return list_class(videos, complete_audio=tensor("AUDIO", descriptor["complete_audio"]))
def decode_video(descriptor, payload, *, max_bytes=MAX_PAYLOAD_BYTES):
"""Decode a bounded VIDEO into independently owned native CPU media objects."""
limit = _limit(max_bytes)
if type(payload) is not bytes:
raise ValueError("VIDEO payload must be bytes")
if len(payload) > limit:
raise ValueError("VIDEO payload exceeds byte limit")
validated = _validate_video(descriptor, len(payload))
return _decode_validated(validated, payload, _native_types())
+68
View File
@@ -271,6 +271,74 @@ def _load_worker_routes_module():
worker_routes = _load_worker_routes_module()
class WebsocketExecutionTargetTests(unittest.IsolatedAsyncioTestCase):
def setUp(self):
from test_async_helpers import _load_async_helpers_module
self.modules = patch.dict(sys.modules)
self.modules.start()
self.addCleanup(self.modules.stop)
self.helper, self.server = _load_async_helpers_module()
sys.modules['dist_api_worker_testpkg.utils.async_helpers'] = self.helper
self.prompt = {'9': {'class_type': 'Output'}, '10': {'class_type': 'Output'}}
async def receive(self, payload):
import json
sent = []
class Socket:
async def prepare(self, request):
pass
async def send_json(self, value):
sent.append(value)
def __aiter__(self):
async def messages():
yield types.SimpleNamespace(type=worker_routes.aiohttp.WSMsgType.TEXT,
data=json.dumps(payload))
return messages()
with patch.object(worker_routes, 'web', types.SimpleNamespace(WebSocketResponse=lambda **kw: Socket())), \
patch.object(worker_routes, 'queue_prompt_payload', self.helper.queue_prompt_payload), \
patch.object(worker_routes, 'PromptValidationError', self.helper.PromptValidationError):
await worker_routes.worker_ws_endpoint(_FakeRequest())
return sent[0]
def payload(self, targets):
return {'type': 'dispatch_prompt', 'request_id': 'req', 'prompt': self.prompt,
'partial_execution_targets': targets}
async def test_explicit_targets_reach_queue_and_acknowledgement(self):
ack = await self.receive(self.payload(['9']))
self.assertTrue(ack['ok'], ack)
self.assertEqual(ack['request_id'], 'req')
self.assertEqual(self.server.prompt_queue.items[0][4], ['9'])
async def test_malformed_targets_and_messages_are_rejected(self):
for targets in (None, [], '9', [True], [9], ['absent'], ['9', '9']):
with self.subTest(targets=targets):
ack = await self.receive(self.payload(targets))
self.assertFalse(ack['ok'])
self.assertIn('partial_execution_targets', ack['error'])
self.assertFalse((await self.receive([]))['ok'])
self.assertEqual(self.server.prompt_queue.items, [])
async def test_validation_cannot_silently_queue_a_partial_subset(self):
ack = await self.receive(self.payload(['9', '10']))
self.assertFalse(ack['ok'])
self.assertEqual(ack['validation_error']['type'], 'partial_execution_targets_failed_validation')
self.assertEqual(self.server.prompt_queue.items, [])
async def test_missing_target_field_keeps_legacy_queue_call(self):
from unittest.mock import AsyncMock
queued = AsyncMock(return_value='legacy')
original = self.helper.queue_prompt_payload
self.helper.queue_prompt_payload = queued
try:
payload = self.payload(['9'])
payload.pop('partial_execution_targets')
self.assertTrue((await self.receive(payload))['ok'])
self.assertNotIn('partial_execution_targets', queued.call_args.kwargs)
finally:
self.helper.queue_prompt_payload = original
class WorkerRoutesTests(unittest.IsolatedAsyncioTestCase):
async def test_network_info_reports_actual_master_and_available_ports(self):
with patch.object(worker_routes, "_get_cuda_info", return_value=(0, 1, 4)), \
@@ -0,0 +1,172 @@
"""Offline CPU proof against ComfyUI's real signatures and PromptExecutor.
Exercises one native worker executor and separate private master-slot lifetimes,
including their shared stable Start / fresh End contract.
"""
import asyncio
import copy
import importlib
from pathlib import Path
import sys
import types
ROOT = Path(__file__).resolve().parents[2]
def main(core):
sys.path.insert(0, str(Path(core).resolve()))
import comfy.cli_args
comfy.cli_args.args.cpu = True # Must precede all model-management imports.
comfy.cli_args.args.disable_assets = True
import execution
import nodes
import torch
from comfy_execution.caching import CacheKeySetInputSignature
name = '_native_cache_identity_probe'
for suffix in ('', '.image_list', '.nodes', '.utils'):
package = types.ModuleType(name + suffix)
package.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
sys.modules[name + suffix] = package
def forbidden(*args, **kwargs):
raise AssertionError('offline native proof attempted a server-loop bridge')
sys.modules[name + '.utils.async_helpers'] = types.SimpleNamespace(run_async_in_server_loop=forbidden)
rt = importlib.import_module(name + '.image_list.runtime')
boundary = importlib.import_module(name + '.nodes.image_list')
calls = []
class Operation:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'image': ('IMAGE',), 'mask': ('MASK',),
'text': ('STRING', {'forceInput': True})}}
RETURN_TYPES = ('IMAGE', 'MASK', 'STRING')
FUNCTION = 'run'
def run(self, image, mask, text):
calls.append(text)
return image + 1, mask + 1, text + ':done'
nodes.NODE_CLASS_MAPPINGS.update(DistributedImageListStart=boundary.DistributedImageListStart,
DistributedImageListEnd=boundary.DistributedImageListEnd,
Operation=Operation)
prompt = {
'start': {'class_type': rt.START_CLASS, 'inputs': {'image_list_execution_id': 'outer-A'}},
'op': {'class_type': 'Operation', 'inputs': {
'image': ['start', 0], 'mask': ['start', 1], 'text': ['start', 4]}},
'end': {'class_type': rt.END_CLASS, 'inputs': {
'images': ['op', 0], 'masks': ['op', 1], 'texts': ['op', 2],
'context': ['start', 3], 'value_3': ['start', 2]}},
}
data = rt.normalize_image_list_inputs(torch.zeros(1, 2, 3, 3, dtype=torch.float64),
torch.ones(1, 2, 3, dtype=torch.bfloat16),
{'value_1': 'original', 'value_2': 3})[0]
region = rt.validate_image_list_regions(prompt)[0]
def prepare(data):
token = rt.create_worker_task_token()
context = rt.create_master_job('start', 'end', prompt, [data], ['worker'], '')
rt.register_assignment(context, data['index'], 'worker', token)
identity = rt.store_worker_item(token, rt.encode_item(data))
result = rt.build_worker_prompt(prompt, region, token, context, data['index'], '', item_identity=identity)
result['start']['inputs']['image_list_local_execution'] = True
result['end']['inputs']['image_list_local_execution'] = True
return token, context, result
async def signature(graph, executor):
dynamic = execution.DynamicPrompt(graph)
changed = execution.IsChangedCache('signature-only', dynamic, executor.caches.outputs)
keys = CacheKeySetInputSignature(dynamic, graph.keys(), changed)
await keys.add_keys(graph.keys())
return {node: keys.get_data_key(node) for node in graph}
async def probe():
server = types.SimpleNamespace(client_id=None, last_node_id=None, send_sync=lambda *a, **kw: None)
executor = execution.PromptExecutor(server, cache_type=execution.CacheType.LRU,
cache_args={'lru': 64, 'ram': 0, 'ram_inactive': 0})
ta, ca, a = prepare(data)
prompt['start']['inputs']['image_list_execution_id'] = 'outer-B'
tb, cb, b = prepare(data)
assert ta != tb and ca != cb
ka, kb = await signature(a, executor), await signature(b, executor)
assert ka['start'] == kb['start'] and ka['op'] == kb['op']
assert ka['end'] != kb['end']
assert a['start']['inputs'] == b['start']['inputs']
# Exercise actual V3 list mapping, hidden.prompt/unique_id, lazy inputs,
# native node signatures and cache hits, not a reimplementation of them.
for number, (token, context, graph) in enumerate(((ta, ca, a), (tb, cb, b))):
valid = await execution.validate_prompt(f'native-{number}', graph, ['end'])
assert valid[0], valid
await executor.execute_async(graph, f'native-{number}', {}, ['end'])
assert executor.success, executor.status_messages
assert token not in rt._WORKER_ITEMS
result = rt.get_master_job(context).results[0]
assert torch.equal(result['image'], data['image'] + 1)
assert torch.equal(result['mask'], data['mask'] + 1)
assert result['text'] == 'original:done'
assert result['values'] == {'value_3': 0}
cached = next(entry['nodes'] for event, entry in executor.status_messages if event == 'execution_cached')
if number:
assert 'start' in cached and 'op' in cached and 'end' not in cached, cached
assert calls == ['original'], calls
for failure in ('expired', 'mismatched'):
token, context, graph = prepare(data)
expires, staged = rt._WORKER_ITEMS[token]
if failure == 'expired':
rt._WORKER_ITEMS[token] = (0, staged)
else:
staged['image'] = staged['image'] + 9
await executor.execute_async(graph, f'rejected-{failure}', {}, ['end'])
assert not executor.success, failure
assert not rt.get_master_job(context).received_indexes
assert token not in rt._WORKER_ITEMS
cached = next(entry['nodes'] for event, entry in executor.status_messages if event == 'execution_cached')
assert 'start' in cached and 'op' in cached and 'end' not in cached, cached
assert calls == ['original'], calls
for part in ('image', 'mask', 'value', 'index'):
changed = copy.deepcopy(data)
if part in ('image', 'mask'):
changed[part].reshape(-1)[0] += 1
elif part == 'value':
changed['values']['value_2'] = 4 # Even an unconnected slot is content.
else:
changed['index'] = 1
token, context, graph = prepare(changed)
keys = await signature(graph, executor)
assert keys['start'] != ka['start'] and keys['op'] != ka['op'], part
await executor.execute_async(graph, f'changed-{part}', {}, ['end'])
assert executor.success, executor.status_messages
assert token not in rt._WORKER_ITEMS
assert rt.get_master_job(context).received_indexes == {changed['index']}
assert len(calls) == 5, calls
# Combine the real Start/End identity path with cache transfer between
# private master slots, not just independent signature/retention probes.
local = importlib.import_module(name + '.image_list.local_executor')
before = len(calls)
for number in range(2):
token, context, graph = prepare(data)
slot = local.LocalExecutionSlot()
try:
await asyncio.wrap_future(slot.submit(graph, f'local-{number}', 'end'))
result = rt.get_master_job(context).results[0]
assert torch.equal(result['image'], data['image'] + 1)
assert result['values'] == {'value_3': 0}
assert token not in rt._WORKER_ITEMS
finally:
await slot.close()
rt.remove_master_job(context)
assert local._RETAINED_OUTPUTS is not None
assert len(calls) == before + 1, calls
assert not rt._WORKER_ITEMS
assert not torch.cuda.is_initialized()
print('NATIVE_CACHE_IDENTITY_OK stable Start/op keys; fresh End key; '
'1 operation for 2 same-data jobs; 4 changed-content misses; '
'expired/mismatched cached inputs rejected; current delivery and cleanup')
asyncio.run(probe())
if __name__ == '__main__':
main(sys.argv[1])
+312 -77
View File
@@ -1,11 +1,12 @@
"""Model-free multi-process acceptance using real ComfyUI executors and loopback HTTP.
Only the transport test facade (/prompt, /queue, /history, /interrupt) is synthetic.
Isolated listeners use real core /prompt, /queue, /history and /interrupt handlers.
Start/End, graph extraction, scheduling, payload codecs and callbacks are production code.
No managed ComfyUI service, user workflow, model, or remote worker is touched.
"""
import asyncio
from collections import Counter
import copy
import importlib
import json
import os
@@ -36,13 +37,14 @@ import aiohttp
torch.set_num_threads(1)
SOURCE_CALLS, BRANCH_CALLS, CAPTURED, MODEL_CALLS = [], [], [], []
UTILITY_CALLS, ITEM_EXECUTIONS, VALUE_CAPTURES = [], [], []
ALLOW_MASTER = False
COUNT = 7
def source_rows():
return [{'image': torch.full((1, 5 + i, 7, 3 if i % 2 == 0 else 4), i / 16, dtype=torch.float32),
'reference_image': torch.full((1, 2, 3, 3), (i + 1) / 16),
'reference_image': torch.full((2, 2, 3, 3), (i + 1) / 16),
'mask': torch.full((1, 5 + i, 7), (i + 1) / 16),
'text': f'item-{i}', 'strength': (i + 1) / 16, 'seed': 100 + i, 'index': i}
for i in range(COUNT)]
@@ -85,6 +87,53 @@ class TestCachedModel:
return (object(),)
def assert_item_execution(prompt, index):
"""Prove the operation belongs to an item, not merely a participating host."""
from comfy_execution.utils import get_executing_context
context = get_executing_context()
assert context is not None
assert prompt['4']['inputs'].get('worker_result_token'), 'per-item node ran in outer graph'
assert (prompt['2']['inputs'].get('worker_item_identity')
or prompt['2']['inputs'].get('worker_item_token')), 'per-item source binding missing'
running, _ = server.PromptServer.instance.prompt_queue.get_current_queue_volatile()
assert len(running) == 1, running
entry = running[0]
if WORKER:
assert entry[1] == context.prompt_id and entry[2] == prompt
assert entry[4] == ['4'], entry[4]
else:
assert ALLOW_MASTER, 'per-item node ran on delegate-only master'
assert entry[1] != context.prompt_id, 'per-item node ran in outer executor'
assert set(entry[4]) == {'4', '5'}, entry[4]
return {'index': index, 'prompt_id': context.prompt_id,
'role': 'worker' if WORKER else 'master-item', 'queue_prompt_id': entry[1]}
class TestRequiredOutputUtility:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'image': ('IMAGE',), 'index': ('INT', {'forceInput': True})},
'hidden': {'prompt': 'PROMPT'}}
OUTPUT_NODE = True
RETURN_TYPES = ('IMAGE',)
FUNCTION = 'run'
def run(self, image, index, prompt):
assert_item_execution(prompt, index)
UTILITY_CALLS.append(index)
return (image,)
class TestDiscardedPreview:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'image': ('IMAGE',)}}
OUTPUT_NODE = True
RETURN_TYPES = ()
FUNCTION = 'run'
def run(self, image):
raise AssertionError('discarded per-item output preview executed')
class TestImageOperation:
@classmethod
def INPUT_TYPES(cls):
@@ -92,11 +141,13 @@ class TestImageOperation:
'text': ('STRING', {'forceInput': True}),
'strength': ('FLOAT', {'forceInput': True}),
'seed': ('INT', {'forceInput': True}), 'index': ('INT', {'forceInput': True}),
'flag': ('BOOLEAN', {'forceInput': True}), 'model': ('MODEL',)}}
'flag': ('BOOLEAN', {'forceInput': True}), 'model': ('MODEL',)},
'hidden': {'prompt': 'PROMPT'}}
RETURN_TYPES = ('IMAGE', 'MASK', 'STRING')
FUNCTION = 'run'
def run(self, image, reference_image, mask, text, strength, seed, index, flag, model):
def run(self, image, reference_image, mask, text, strength, seed, index, flag, model, prompt):
assert WORKER or ALLOW_MASTER, 'per-image branch executed in orchestrator-only mode'
ITEM_EXECUTIONS.append(assert_item_execution(prompt, index))
BRANCH_CALLS.append(index)
# Sleep in short chunks so the real Comfy interrupt signal can be tested.
until = time.monotonic() + DELAY
@@ -111,16 +162,32 @@ class TestImageOperation:
return image + strength, mask * 0.5, f'{text}:seed-{seed}'
class TestValueOverride:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'text': ('STRING', {'forceInput': True}),
'index': ('INT', {'forceInput': True})}}
RETURN_TYPES = ('STRING', 'INT')
FUNCTION = 'run'
def run(self, text, index):
assert WORKER or ALLOW_MASTER, 'override branch executed on delegate-only master'
return (text + ':override', 200 + index)
class TestCapture:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'images': ('IMAGE',), 'masks': ('MASK',), 'texts': ('STRING', {'forceInput': True})}}
return {'required': {'images': ('IMAGE',), 'masks': ('MASK',), 'texts': ('STRING', {'forceInput': True})},
'optional': {'values': ('DISTRIBUTED_IMAGE_LIST_VALUES',),
'picked': ('STRING', {'forceInput': True})}}
INPUT_IS_LIST = True
OUTPUT_NODE = True
RETURN_TYPES = ()
FUNCTION = 'run'
def run(self, images, masks, texts):
def run(self, images, masks, texts, values=None, picked=None):
assert not WORKER, 'post-End capture leaked into a worker prompt'
CAPTURED.append((images, masks, texts))
VALUE_CAPTURES.append((values, picked))
return ()
@@ -129,9 +196,19 @@ async def setup():
ps = server.PromptServer(asyncio.get_running_loop(), assets)
assert await nodes.load_custom_node(str(ROOT))
module = sys.modules[str(ROOT).replace('.', '_x_')]
boundary = importlib.import_module(module.__name__ + '.nodes.image_list')
encode_result = boundary.encode_result
def checked_encode(result, **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
boundary.encode_result = checked_encode
nodes.NODE_CLASS_MAPPINGS.update(TestImageSource=TestImageSource,
TestImageOperation=TestImageOperation, TestCapture=TestCapture,
TestCachedModel=TestCachedModel, TestSharedFlag=TestSharedFlag)
TestCachedModel=TestCachedModel, TestSharedFlag=TestSharedFlag,
TestRequiredOutputUtility=TestRequiredOutputUtility,
TestDiscardedPreview=TestDiscardedPreview, TestValueOverride=TestValueOverride)
# Match the CORS-enabled direct-worker launch contract for canvas probes.
app = web.Application(middlewares=[server.create_cors_middleware('*')])
# Register the actual feature routes, not test replacements.
@@ -147,61 +224,59 @@ async def listen(app):
return runner, site._server.sockets[0].getsockname()[1]
def finish_queue_item(ps, item_id, executor):
ps.prompt_queue.task_done(item_id, executor.history_result,
status=execution.PromptQueue.ExecutionStatus(
status_str='success' if executor.success else 'error',
completed=executor.success, messages=executor.status_messages),
process_item=lambda item: item[:5] + item[6:])
async def worker_main():
ps, assets, module, app = await setup()
pending, running, history = [], {}, {}
wake = asyncio.Event()
interrupted_ids = []
interrupted_ids, submissions = [], []
core_handlers = {(route.method, route.path): route.handler for route in ps.routes}
async def prompt(request):
data = await request.json()
prompt_id = str(uuid.uuid4())
graph = data['prompt']
valid = await execution.validate_prompt(prompt_id, graph, None)
if not valid[0]:
return web.json_response({'error': valid[1], 'node_errors': valid[3]}, status=400)
pending.append((prompt_id, graph, valid[2]))
wake.set()
return web.json_response({'prompt_id': prompt_id})
async def queue(request):
if request.method == 'POST':
deleted = (await request.json()).get('delete', [])
pending[:] = [entry for entry in pending if entry[0] not in deleted]
return web.json_response({})
return web.json_response({'queue_running': [[0, key] for key in running],
'queue_pending': [[0, entry[0]] for entry in pending]})
async def get_history(request):
pid = request.match_info['prompt_id']
return web.json_response({pid: history[pid]} if pid in history else {})
# Use core's actual validation/enqueue/response schema, including
# partial_execution_targets; only observe it and wake our CPU consumer.
response = await core_handlers['POST', '/prompt'](request)
if response.status == 200:
accepted = json.loads(response.body)
_, pending = ps.prompt_queue.get_current_queue_volatile()
entry = next(item for item in pending if item[1] == accepted['prompt_id'])
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']})
wake.set()
return response
async def interrupt(request):
pid = (await request.json())['prompt_id']
interrupted_ids.append(pid)
if pid in running:
comfy.model_management.interrupt_current_processing(True)
return web.json_response({})
interrupted_ids.append((await request.json())['prompt_id'])
return await core_handlers['POST', '/interrupt'](request)
async def stats(request):
return web.json_response({'calls': BRANCH_CALLS, 'running': list(running),
'pending': [e[0] for e in pending], 'interrupted': interrupted_ids,
'history': history})
running, pending = ps.prompt_queue.get_current_queue_volatile()
return web.json_response({'calls': BRANCH_CALLS, 'utility_calls': UTILITY_CALLS,
'item_executions': ITEM_EXECUTIONS, 'submissions': submissions,
'running': [e[1] for e in running], 'pending': [e[1] for e in pending],
'interrupted': interrupted_ids, 'history': ps.prompt_queue.get_history()})
async def consume():
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
while True:
await wake.wait()
while pending:
pid, graph, outputs = pending.pop(0)
running[pid] = True
await asyncio.to_thread(executor.execute, graph, pid, {}, outputs)
history[pid] = {'status': {'completed': True,
'status_str': 'success' if executor.success else 'error',
'messages': str(executor.status_messages)}}
running.pop(pid, None)
while (queued := ps.prompt_queue.get(timeout=0)) is not None:
entry, item_id = queued
await asyncio.to_thread(executor.execute, entry[2], entry[1], entry[3], entry[4])
finish_queue_item(ps, item_id, executor)
wake.clear()
async def probe(_request):
return web.json_response({'exec_info': {'queue_remaining': len(pending) + len(running)}})
app.router.add_get('/prompt', probe)
app.add_routes([route for route in ps.routes if
(route.path in ('/queue', '/history/{prompt_id}') or
(route.path == '/prompt' and route.method == 'GET'))])
app.router.add_post('/prompt', prompt)
app.router.add_get('/queue', queue)
app.router.add_post('/queue', queue)
app.router.add_get('/history/{prompt_id}', get_history)
app.router.add_post('/interrupt', interrupt)
app.router.add_get('/test-stats', stats)
runner, port = await listen(app)
@@ -224,23 +299,51 @@ def graph(callback, worker_ids):
'images': ['1', 0], 'mask': ['1', 2],
'values.value_1': ['1', 1], 'values.value_2': ['1', 3],
'values.value_3': ['1', 4], 'values.value_4': ['1', 5], 'values.value_5': ['7', 0],
'values.value_6': ['1', 2],
'image_list_enabled_worker_ids': json.dumps(worker_ids), 'image_list_master_callback': callback}},
'3': {'class_type': 'TestImageOperation', 'inputs': {
key: ['2', index] for key, index in [('image', 0), ('mask', 1), ('index', 2),
('reference_image', 4), ('text', 5), ('strength', 6), ('seed', 7), ('flag', 8)]}
| {'model': ['6', 0]}},
| {'image': ['8', 0], 'model': ['6', 0]}},
'4': {'class_type': 'DistributedImageListEnd', 'inputs': {
'images': ['3', 0], 'masks': ['3', 1], 'texts': ['3', 2], 'context': ['2', 3]}},
'5': {'class_type': 'TestCapture', 'inputs': {'images': ['4', 0], 'masks': ['4', 1], 'texts': ['4', 2]}},
'images': ['3', 0], 'masks': ['3', 1], 'texts': ['3', 2], 'context': ['2', 3],
'value_2': ['10', 0], 'value_3': ['2', 6], 'value_4': ['2', 7], 'value_5': ['2', 8],
'value_7': ['3', 1], 'value_8': ['2', 4], 'value_9': ['2', 9], 'value_10': ['10', 1]}},
'5': {'class_type': 'TestCapture', 'inputs': {'images': ['4', 0], 'masks': ['4', 1], 'texts': ['4', 2],
'values': ['4', 3], 'picked': ['get', 0]}},
'6': {'class_type': 'TestCachedModel', 'inputs': {}},
'7': {'class_type': 'TestSharedFlag', 'inputs': {}},
'8': {'class_type': 'TestRequiredOutputUtility', 'inputs': {'image': ['2', 0], 'index': ['2', 2]}},
'9': {'class_type': 'TestDiscardedPreview', 'inputs': {'image': ['8', 0]}},
'10': {'class_type': 'TestValueOverride', 'inputs': {'text': ['2', 5], 'index': ['2', 2]}},
'get': {'class_type': 'DistributedImageListGetValue', 'inputs': {'values': ['4', 3], 'value_name': 'value_2'}},
}
def check_carried_values():
values, picked = VALUE_CAPTURES[-1]
assert len(values) == len(picked) == COUNT
for index, (actual, source) in enumerate(zip(values, source_rows())):
assert set(actual) == {f'value_{n}' for n in range(1, 11)}, actual.keys()
assert torch.equal(actual['value_1'], source['reference_image'])
assert actual['value_2'] == picked[index] == f'item-{index}:override'
assert actual['value_3'] == source['strength']
assert actual['value_4'] == source['seed']
assert actual['value_5'] is True
assert torch.equal(actual['value_6'], source['mask'])
assert torch.equal(actual['value_7'], source['mask'] * 0.5)
assert torch.equal(actual['value_8'], source['reference_image'])
assert actual['value_8'].shape[0] == 2, 'value tensor batch was split'
assert torch.equal(actual['value_9'], source['mask'])
assert actual['value_10'] == 200 + index
async def coordinator():
global ALLOW_MASTER
ps, assets, module, app = await setup()
runtime = importlib.import_module(module.__name__ + '.image_list.runtime')
execution_plan = importlib.import_module(module.__name__ + '.image_list.execution_plan')
lifecycle = importlib.import_module(module.__name__ + '.image_list.lifecycle')
config_mod = importlib.import_module(module.__name__ + '.utils.config')
orchestration = importlib.import_module(module.__name__ + '.api.queue_orchestration')
runner, port = await listen(app)
@@ -294,16 +397,49 @@ async def coordinator():
prompt = graph(f'http://127.0.0.1:{port}', [w['id'] for w in workers])
valid = await execution.validate_prompt('network-acceptance', prompt, None)
assert valid[0], valid
assert {'4', '5', '8', '9'} == set(valid[2]), valid
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
# Exercise production queue preparation; only master enqueue is redirected
# to this isolated executor, not the managed server's live queue.
prompt['2']['inputs'].pop('image_list_enabled_worker_ids')
prompt['2']['inputs'].pop('image_list_master_callback')
saved_prompt = copy.deepcopy(prompt)
outer_executions = []
async def enqueue(prepared, workflow, client_id, **kwargs):
checked = await execution.validate_prompt('network-acceptance', prepared, None)
targets = kwargs['partial_execution_targets']
planned, planned_targets = execution_plan.prepare_image_list_execution(prepared)
assert planned == prepared and targets == planned_targets
if '9' in saved_prompt and 'TestImageOperation' in {
node['class_type'] for node in prepared.values()}:
assert set(prepared) == set(saved_prompt) - {'9'}, prepared
assert prepared['8'] == saved_prompt['8'], 'required utility changed'
assert prepared['5'] == saved_prompt['5'], 'post-End output changed'
assert set(targets) == {'4', '5'}, targets
assert nodes.NODE_CLASS_MAPPINGS['TestRequiredOutputUtility'].OUTPUT_NODE is True
assert nodes.NODE_CLASS_MAPPINGS['TestDiscardedPreview'].OUTPUT_NODE is True
pid = 'network-acceptance' if not outer_executions else str(uuid.uuid4())
checked = await execution.validate_prompt(pid, prepared, targets)
assert checked[0], checked
await asyncio.to_thread(executor.execute, prepared, 'network-acceptance', {}, checked[2])
return {'prompt_id': 'network-acceptance', 'number': 1, 'node_errors': {}}
assert not checked[3] and Counter(checked[2]) == Counter(targets), checked
assert not lifecycle.has_safe_execution_targets(ps, pid, prepared)
# Exactly the core queue tuple/put/get/task_done lifecycle. Start
# must prove its roots from this real running record, not a mock.
number = len(outer_executions)
ps.prompt_queue.put((number, pid, prepared, {'client_id': client_id}, checked[2], {}))
entry, item_id = ps.prompt_queue.get(timeout=0)
assert lifecycle.has_safe_execution_targets(ps, pid, prepared)
before_captures = len(CAPTURED)
await asyncio.to_thread(executor.execute, entry[2], entry[1], entry[3], entry[4])
finish_queue_item(ps, item_id, executor)
assert ps.prompt_queue.get_current_queue_volatile() == ([], [])
history = ps.prompt_queue.get_history(pid)[pid]
assert Counter(history['prompt'][4]) == Counter(targets)
assert not lifecycle.has_safe_execution_targets(ps, pid, prepared)
outer_executions.append({'prompt_id': pid, 'targets': targets,
'captures': len(CAPTURED) - before_captures})
if executor.success:
assert len(CAPTURED) == before_captures + 1, 'post-End Capture must run exactly once'
return {'prompt_id': pid, 'number': number, 'node_errors': checked[3]}
with patch.object(config_mod, 'load_config', return_value=config), \
patch.object(orchestration, 'load_config', return_value=config), \
patch.object(orchestration, 'queue_prompt_payload', enqueue), \
@@ -312,12 +448,14 @@ async def coordinator():
prompt, None, 'test', [w['id'] for w in workers]), 45)
assert queued[0] == 'network-acceptance' and queued[2] == 2, queued
assert not ps.distributed_pending_jobs, 'image-list created unused Collector queues'
assert prompt == saved_prompt, 'queue preparation changed the saved/source graph'
if not executor.success:
for err in stderr_files:
err.flush(); err.seek(0); print(err.read(), file=sys.stderr)
assert executor.success, executor.status_messages
assert len(SOURCE_CALLS) == 1 and not BRANCH_CALLS
assert len(SOURCE_CALLS) == 1 and not BRANCH_CALLS and not UTILITY_CALLS
assert len(CAPTURED) == 1
check_carried_values()
images, masks, texts = CAPTURED[0]
assert len(images) == len(masks) == len(texts) == COUNT
for i, row in enumerate(source_rows()):
@@ -332,30 +470,46 @@ async def coordinator():
counts = [len(row['calls']) for row in stats]
assert counts[0] > counts[1] > 0, counts
assert Counter(i for row in stats for i in row['calls']) == Counter(range(COUNT))
for row in stats:
assert Counter(row['utility_calls']) == Counter(row['calls']), row
assert len(row['submissions']) == len(row['calls'])
assert len({event['prompt_id'] for event in row['item_executions']}) == len(row['calls'])
for submission in row['submissions']:
assert submission['requested'] == submission['queued'] == ['4'], submission
assert '8' in submission['node_ids'] and '9' not in submission['node_ids']
assert '5' not in submission['node_ids']
assert not runtime._MASTER_JOBS
# Three sequential regions on one worker, including optional outputs
# left unconnected. Region N+1 must not race the prior worker End.
chained = {'1': {'class_type': 'EmptyImage', 'inputs': {
'height': 8, 'width': 8, 'batch_size': 2, 'color': 0}}}
chained['constant'] = {'class_type': 'TestSharedFlag', 'inputs': {}}
previous = '1'
value_source = 'constant'
for start, end in [('2', '3'), ('4', '5'), ('6', '7')]:
chained[start] = {'class_type': 'DistributedImageListStart', 'inputs': {
'images': [previous, 0]}}
'images': [previous, 0], 'values.value_1': [value_source, 0]}}
chained[end] = {'class_type': 'DistributedImageListEnd', 'inputs': {
'images': [start, 0], 'context': [start, 3]}}
value_source = f'get-{end}'
chained[value_source] = {'class_type': 'DistributedImageListGetValue', 'inputs': {
'values': [end, 3], 'value_name': 'value_1'}}
previous = end
chained['8'] = {'class_type': 'TestCapture', 'inputs': {
'images': ['7', 0], 'masks': ['7', 1], 'texts': ['7', 2]}}
'images': ['7', 0], 'masks': ['7', 1], 'texts': ['7', 2], 'values': ['7', 3]}}
saved_chained = copy.deepcopy(chained)
with patch.object(config_mod, 'load_config', return_value=config), \
patch.object(orchestration, 'load_config', return_value=config), \
patch.object(orchestration, 'queue_prompt_payload', enqueue), \
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
await asyncio.wait_for(orchestration.orchestrate_distributed_execution(
chained, None, 'test', ['slow']), 45)
assert chained == saved_chained
assert executor.success, executor.status_messages
assert len(CAPTURED) == 2 and len(CAPTURED[1][0]) == 2
assert all(torch.equal(image, torch.zeros((1, 8, 8, 3))) for image in CAPTURED[1][0])
assert CAPTURED[1][1:] == ([None, None], [None, None])
assert VALUE_CAPTURES[1][0] == [{'value_1': True}, {'value_1': True}]
assert not runtime._MASTER_JOBS
print('IMAGE_LIST_SEQUENTIAL_OK three regions; one worker; optional outputs omitted')
# Same public queue path, now with master participation enabled. The
@@ -363,50 +517,131 @@ 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)
before_sources, before_local, before_models = len(SOURCE_CALLS), len(BRANCH_CALLS), len(MODEL_CALLS)
before_utilities, before_events = len(UTILITY_CALLS), len(ITEM_EXECUTIONS)
async with aiohttp.ClientSession() as session:
before_worker_calls = []
before_worker_stats = []
for worker in workers:
async with session.get(f"http://127.0.0.1:{worker['port']}/test-stats") as response:
before_worker_calls.append(len((await response.json())['calls']))
before_worker_stats.append(await response.json())
with patch.object(config_mod, 'load_config', return_value=config), \
patch.object(orchestration, 'load_config', return_value=config), \
patch.object(orchestration, 'queue_prompt_payload', enqueue), \
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
queued = await asyncio.wait_for(orchestration.orchestrate_distributed_execution(
prompt, None, 'test', selected), 45)
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'
local_calls = BRANCH_CALLS[before_local:]
assert Counter(UTILITY_CALLS[before_utilities:]) == Counter(local_calls), mode
assert len({event['prompt_id'] for event in ITEM_EXECUTIONS[before_events:]}) == len(local_calls)
assert local_calls, f'{mode}: master did not process any items'
assert len(MODEL_CALLS) == before_models + 1, 'local model dependency did not stay cached'
remote_calls = []
remote_assignments = []
async with aiohttp.ClientSession() as session:
for worker, offset in zip(workers, before_worker_calls):
for worker, before_stats in zip(workers, before_worker_stats):
async with session.get(f"http://127.0.0.1:{worker['port']}/test-stats") as response:
current = await response.json()
remote_calls.extend(current['calls'][offset:])
calls = current['calls'][len(before_stats['calls']):]
utilities = current['utility_calls'][len(before_stats['utility_calls']):]
submissions = current['submissions'][len(before_stats['submissions']):]
events = current['item_executions'][len(before_stats['item_executions']):]
assert Counter(utilities) == Counter(calls), (mode, current)
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_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)
before_models, before_utilities = len(MODEL_CALLS), len(UTILITY_CALLS)
before_captures = len(CAPTURED)
previous_outputs = CAPTURED[-1]
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
with patch.object(config_mod, 'load_config', return_value=config), \
patch.object(orchestration, 'load_config', return_value=config), \
patch.object(orchestration, 'queue_prompt_payload', enqueue), \
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
await asyncio.wait_for(orchestration.orchestrate_distributed_execution(
prompt, None, 'test', []), 45)
assert executor.success, executor.status_messages
assert len(SOURCE_CALLS) == before_sources + 1
assert len(BRANCH_CALLS) == before_local
assert len(MODEL_CALLS) == before_models
assert len(UTILITY_CALLS) == before_utilities
assert len(CAPTURED) == before_captures + 1
for before, after in zip(previous_outputs[:2], CAPTURED[-1][:2]):
assert len(before) == len(after) == COUNT
assert all(torch.equal(a, b) for a, b in zip(before, after))
assert previous_outputs[2] == CAPTURED[-1][2]
check_carried_values()
assert not runtime._MASTER_JOBS and not runtime._WORKER_ITEMS
assert prompt == saved_prompt
print('IMAGE_LIST_CACHE_REQUEUE_OK unchanged items reused; fresh ordered delivery and cleanup')
print('IMAGE_LIST_MASTER_OK ' + json.dumps(participation_counts))
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')
print('IMAGE_LIST_TARGETS_OK ' + json.dumps({
'modes': ['delegate-only', 'mixed', 'master-only'],
'outer_targets': outer_executions[0]['targets'], 'worker_targets': ['4'],
'required_utility_once_per_item': True, 'preview_discarded': True,
'saved_prompt_unchanged': prompt == saved_prompt, 'capture_once_per_outer_prompt': True}))
ALLOW_MASTER = False
config['settings']['master_delegate_only'] = True
# Preserve the original direct-runtime cancellation/failure fixture:
# it has no outer queue and must not claim target-control evidence.
direct_prompt = copy.deepcopy(prompt)
del direct_prompt['8'], direct_prompt['9']
direct_prompt['3']['inputs']['image'] = ['2', 0]
# Real Comfy user-interrupt flag while a real worker executor is active.
rows = source_rows()
items = runtime.normalize_image_list_inputs(
[row['image'] for row in rows], [row['mask'] for row in rows],
{f'value_{n}': [row[key] for row in rows]
for n, key in ((1, 'reference_image'), (2, 'text'), (3, 'strength'), (4, 'seed'), (6, 'mask'))}
| {'value_5': True})
slow = workers[1]
async with aiohttp.ClientSession() as session:
async def slow_stats():
@@ -420,11 +655,11 @@ async def coordinator():
await asyncio.sleep(0.02)
raise AssertionError('worker state transition timed out')
before = await until(lambda state: not state['running'] and not state['pending'])
context = runtime.create_master_job('2', '4', prompt, source_rows(), ['slow'],
context = runtime.create_master_job('2', '4', direct_prompt, items, ['slow'],
f'http://127.0.0.1:{port}')
with patch.object(config_mod, 'load_config', return_value=config), \
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
task = asyncio.create_task(runtime.run_image_list_job(context))
task = asyncio.create_task(runtime.run_image_list_job(context, include_values=True))
active = await until(lambda state: bool(state['running']))
owned_id = active['running'][0]
comfy.model_management.interrupt_current_processing(True)
@@ -442,14 +677,14 @@ async def coordinator():
assert not runtime._MASTER_JOBS
print('IMAGE_LIST_CANCEL_OK owned prompt stopped; no further dispatch; context removed')
# A worker exception before End must fail, not hang or return a partial list.
bad_rows = source_rows()
bad_rows[0]['seed'] = -1
context = runtime.create_master_job('2', '4', prompt, bad_rows, ['fast'],
bad_rows = copy.deepcopy(items)
bad_rows[0]['values']['value_4'] = -1
context = runtime.create_master_job('2', '4', direct_prompt, bad_rows, ['fast'],
f'http://127.0.0.1:{port}')
with patch.object(config_mod, 'load_config', return_value=config), \
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
try:
await asyncio.wait_for(runtime.run_image_list_job(context), 5)
await asyncio.wait_for(runtime.run_image_list_job(context, include_values=True), 5)
raise AssertionError('worker exception was ignored')
except RuntimeError as exc:
assert 'failed image-list item' in str(exc), str(exc)
@@ -0,0 +1,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()))
@@ -0,0 +1,114 @@
"""CPU-only native Comfy execution across private-slot lifetimes."""
import asyncio
import importlib
import os
from pathlib import Path
import sys
import types
ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, sys.argv[1])
import comfy.cli_args
comfy.cli_args.args.cpu = True
comfy.cli_args.args.disable_assets = True
import execution
import nodes
pkg = '_local_retention_probe'
for name, path in [(pkg, ROOT), (pkg + '.image_list', ROOT / 'image_list')]:
module = types.ModuleType(name)
module.__path__ = [str(path)]
sys.modules[name] = module
local = importlib.import_module(pkg + '.image_list.local_executor')
CALLS = []
OUTPUTS = []
class Counted:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'value': ('INT',)}}
RETURN_TYPES = ('INT',)
FUNCTION = 'run'
def run(self, value):
if value < 0:
raise ValueError('intentional retention failure')
CALLS.append(value)
return (value,)
class Uncached(Counted):
@classmethod
def IS_CHANGED(cls, **kwargs):
return float('nan')
class Capture:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'value': ('INT',), 'delivery': ('STRING',)}}
RETURN_TYPES = ()
OUTPUT_NODE = True
FUNCTION = 'run'
def run(self, value, delivery):
OUTPUTS.append((value, delivery))
return ()
nodes.NODE_CLASS_MAPPINGS.update(RetentionCounted=Counted, RetentionUncached=Uncached, RetentionCapture=Capture)
async def run(slot, value, delivery, uncached=False):
prompt = {
'source': {'class_type': 'RetentionUncached' if uncached else 'RetentionCounted', 'inputs': {'value': value}},
'end': {'class_type': 'RetentionCapture', 'inputs': {'value': ['source', 0], 'delivery': delivery}},
}
await asyncio.wrap_future(slot.submit(prompt, delivery, 'end'))
async def main():
first = local.LocalExecutionSlot()
await run(first, 1, 'first-1')
await run(first, 2, 'first-2')
await first.close()
assert local._RETAINED_OUTPUTS is not None
second = local.LocalExecutionSlot()
await run(second, 1, 'second-1')
await run(second, 2, 'second-2')
await run(second, 3, 'changed')
await second.close()
assert CALLS == [1, 2, 3], CALLS
assert OUTPUTS == [(1, 'first-1'), (2, 'first-2'), (1, 'second-1'), (2, 'second-2'), (3, 'changed')], OUTPUTS
third = local.LocalExecutionSlot()
await run(third, 4, 'uncached-1', True)
await run(third, 4, 'uncached-2', True)
await third.close()
assert CALLS[-2:] == [4, 4]
comfy.cli_args.args.cache_none = True
disabled = local.LocalExecutionSlot()
await run(disabled, 1, 'disabled-1')
await run(disabled, 1, 'disabled-2')
await disabled.close()
assert CALLS[-2:] == [1, 1]
assert local._RETAINED_OUTPUTS is None
comfy.cli_args.args.cache_none = False
failed = local.LocalExecutionSlot()
await run(failed, 7, 'before-failure')
try:
await run(failed, -1, 'failure')
except RuntimeError:
pass
else:
raise AssertionError('failure fixture unexpectedly succeeded')
await failed.close()
assert local._RETAINED_OUTPUTS is None
cancelled = local.LocalExecutionSlot()
await run(cancelled, 8, 'cancelled')
cancelled.request_stop()
await cancelled.close()
assert local._RETAINED_OUTPUTS is None
oversized = execution.LRUCache(execution.CacheKeySetInputSignature, max_size=local.NATIVE_CACHE_MAX_ENTRIES)
empty = execution.DynamicPrompt({})
await oversized.set_prompt(empty, [], execution.IsChangedCache('bound', empty, oversized))
oversized.min_generation = oversized.generation
oversized.cache = {n: n for n in range(local.NATIVE_CACHE_MAX_ENTRIES + 1)}
local._retain_outputs(oversized)
assert local._RETAINED_OUTPUTS is None
print('PASS: cross-run native cache reuse, fresh delivery, input change, noncacheable nodes, cache_none, cancellation, retention bound')
asyncio.run(main())
+512
View File
@@ -0,0 +1,512 @@
"""Real CPU Comfy V3 + HTTP proof; invoked only by the opt-in parent test.
The only harness pieces are deterministic source/assertion nodes and a bounded
queue worker. /prompt, /history and /view are Comfy's real handlers; compilation,
packing, staging, codecs, bridge execution, response decoding and Collector are
unmodified production code. No normal instance or Runpod client is contacted.
"""
import asyncio
import copy
import importlib
import io as bytes_io
import json
import os
from pathlib import Path
import socket
import sys
import types
import wave
from contextlib import asynccontextmanager
from fractions import Fraction
ROOT = Path(__file__).resolve().parents[2]
def bootstrap():
if len(sys.argv) != 3:
raise SystemExit("Use test_media_end_to_end.py with explicit COMFYUI_SOURCE_ROOT")
core, run_root = (Path(value).resolve() for value in sys.argv[1:])
assert run_root.is_dir() and run_root != core and not run_root.is_relative_to(core)
for name in ("CUDA_VISIBLE_DEVICES", "HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES"):
assert os.environ.get(name) == "", f"CPU isolation missing: {name}"
assert not any(name in sys.modules for name in ("torch", "execution", "folder_paths"))
sys.path.insert(0, str(core))
import comfy.cli_args
args = comfy.cli_args.args
# Set CPU and all writable roots BEFORE any framework/tensor import.
args.cpu = True
args.disable_assets = True
args.disable_all_custom_nodes = True
args.enable_manager = False
args.offline = True
args.base_directory = str(run_root)
args.front_end_root = str(run_root)
args.database_url = "sqlite:///:memory:"
for name in ("input", "output", "temp", "user"):
(run_root / name).mkdir()
setattr(args, name + "_directory", str(run_root / name))
return run_root
async def check(base):
import av
import numpy as np
import torch
import execution
import folder_paths
import nodes
import server
from aiohttp import web
from app.assets.manager import default_asset_manager
from comfy_api.v0_0_2 import io
from comfy_api.latest._input_impl.video_types import (
VideoFromComponents, VideoFromFile, VideoFromList,
)
from comfy_api.latest._util import VideoComponents
from comfy_extras import nodes_audio
assert not torch.cuda.is_initialized()
torch.set_num_threads(1)
assert execution.comfy.model_management.get_torch_device().type == "cpu"
for name in ("input", "output", "temp", "user"):
assert Path(getattr(folder_paths, "get_" + name + "_directory")()).is_relative_to(base)
# Namespace scaffolding avoids the unrelated pack bootstrap (service/config
# side effects). Every imported module below is the real source on disk.
package_name = "_media_e2e_acceptance"
for suffix in ("", ".api", ".api.orchestration", ".nodes"):
package = types.ModuleType(package_name + suffix)
package.__path__ = [str(ROOT.joinpath(*suffix.lstrip(".").split(".")))] if suffix else [str(ROOT)]
sys.modules[package.__name__] = package
def load(name):
return importlib.import_module(package_name + "." + name)
asset_manager = default_asset_manager()
ps = server.PromptServer(asyncio.get_running_loop(), asset_manager)
ps.distributed_pending_jobs = {}
ps.distributed_jobs_lock = asyncio.Lock()
compiler = load("api.orchestration.serverless_media_prompt")
packer = load("api.orchestration.serverless_typed_inputs")
decoder = load("api.orchestration.serverless_typed_results")
serializer = load("api.orchestration.runpod_client").serialize_request
handler = load("serverless_bridge.handler")
contract = load("serverless_bridge.contract")
video_codec = load("serverless_bridge.video_codec")
bridge = load("serverless_bridge")
collector = load("nodes.v3").DistributedCollector
extension = await bridge.comfy_entrypoint()
bridge_classes = await extension.get_node_list()
assert {cls.GET_SCHEMA().node_id for cls in bridge_classes} == {"MediaImportV1", "MediaExportV1"}
sockets = list(contract.SOCKET_KINDS)
collector_sockets = ["images", "audio"]
kinds = list(contract.SOCKET_KINDS.values())
assert kinds == ["IMAGE", "AUDIO", "MASK", "STRING", "VIDEO"]
io_types = [io.Image, io.Audio, io.Mask, io.String, io.Video]
fixtures, expected_sinks, seen_sinks, source_calls = {}, {}, {}, []
class AcceptanceMediaSource(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(node_id="AcceptanceMediaSource", inputs=[io.String.Input("key")],
outputs=[kind.Output(is_output_list=True) for kind in io_types])
@classmethod
def execute(cls, key):
source_calls.append(key)
return io.NodeOutput(*(fixtures[key][socket] for socket in sockets))
def tensor_equal(actual, expected):
assert actual.device.type == "cpu" and actual.dtype == expected.dtype
assert actual.shape == expected.shape
# Byte comparison preserves signed zero and catches dtype narrowing.
assert torch.equal(actual.contiguous().view(torch.uint8), expected.contiguous().view(torch.uint8))
def exact(kind, actual, expected):
if kind in ("IMAGE", "MASK"):
tensor_equal(actual, expected)
elif kind == "AUDIO":
assert actual["sample_rate"] == expected["sample_rate"]
tensor_equal(actual["waveform"], expected["waveform"])
elif kind == "STRING":
assert type(actual) is str and actual.encode("utf-8") == expected.encode("utf-8")
else:
assert video_codec.encode_video(actual) == expected
if type(actual) is VideoFromFile:
buffer = vars(actual)["_VideoFromFile__file"]
assert type(buffer) is bytes_io.BytesIO and not buffer.closed
position = buffer.tell()
buffer.seek(0)
with av.open(buffer) as container:
assert len(list(container.decode(video=0))) == 3
buffer.seek(position)
elif type(actual) is VideoFromList:
for item in actual.videos:
exact("VIDEO", item, video_codec.encode_video(item))
else:
assert type(actual) is VideoFromComponents
components = actual.get_components()
assert type(components.frame_rate) is Fraction
assert components.images.dtype == torch.float64
assert components.alpha.dtype == torch.bfloat16
assert components.audio["waveform"].dtype == torch.float64
assert actual.get_bit_depth() == 10 and actual.get_color_space() == "HDR PQ"
class AcceptanceMediaAssert(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(node_id="AcceptanceMediaAssert", inputs=[io.String.Input("tag"),
*[kind.Input(socket, optional=True, **({"force_input": True} if socket == "texts" else {}))
for socket, kind in zip(sockets, io_types)]], outputs=[],
is_input_list=True, is_output_node=True)
@classmethod
def execute(cls, tag, **media):
assert len(tag) == 1 and tag[0] not in seen_sinks
expected = expected_sinks[tag[0]]
assert set(media) == set(expected)
for socket in expected:
assert len(media[socket]) == len(expected[socket]), (tag, socket)
# Existing Collector outputs are a batched IMAGE and one AUDIO,
# not lists. Native V3 wraps each output once for this list sink.
for actual, wanted in zip(media[socket], expected[socket]):
exact(contract.SOCKET_KINDS[socket], actual, wanted)
counts = {socket: len(media[socket]) for socket in expected}
seen_sinks[tag[0]] = counts
return io.NodeOutput(ui={"accepted": [counts]})
for cls in [*bridge_classes, collector, AcceptanceMediaSource, AcceptanceMediaAssert]:
assert issubclass(cls, io.ComfyNode)
cls.GET_NODE_INFO_V1()
nodes.NODE_CLASS_MAPPINGS[cls.GET_SCHEMA().node_id] = cls
def node(class_type, **inputs):
return {"class_type": class_type, "inputs": inputs}
def add_bridge_exports(workflow, manifest, request_id, participant, job_id, links):
"""Exercise bridge-only kinds without inventing Collector sockets."""
for socket, link in links.items():
export_id = "bridge-export:" + socket
binding = {"binding_id": "bridge-binding:" + socket, "collector_id": "bridge-only",
"job_id": job_id, "participant_id": participant, "node_id": export_id,
"socket": socket, "kind": contract.SOCKET_KINDS[socket]}
assert export_id not in workflow
manifest.append(binding)
workflow[export_id] = node("MediaExportV1", request_id=request_id,
binding_json=json.dumps(binding), **{socket: link})
contract.validate_bindings(manifest)
def tiny_video(variant):
buffer = bytes_io.BytesIO()
with av.open(buffer, "w", format="matroska") as container:
stream = container.add_stream("ffv1", rate=Fraction(24, 1))
stream.width = stream.height = 16
stream.pix_fmt = "yuv420p"
for level in (16 + variant, 64 + variant, 128 + variant):
frame = av.VideoFrame.from_ndarray(np.full((16, 16, 3), level, dtype=np.uint8), format="rgb24")
for packet in stream.encode(frame):
container.mux(packet)
for packet in stream.encode():
container.mux(packet)
return buffer.getvalue()
def native_fixtures(variant):
image = torch.arange(48, dtype=torch.float64).reshape(2, 2, 3, 4) / 997 + variant / 100
image[0, 0, 0, 0] = -0.0
mask = torch.tensor([[[-0.0, 0.125, 0.875]]], dtype=torch.bfloat16)
audio = {"waveform": torch.arange(38, dtype=torch.float64).reshape(1, 2, 19) / 997,
"sample_rate": 44100}
audio["waveform"][0, 0, 0] = -0.0
component = VideoFromComponents(VideoComponents(
images=image[..., :3], frame_rate=Fraction(30000, 1001), alpha=mask.reshape(1, 1, 3, 1).repeat(2, 2, 1, 1),
audio=audio, metadata={"title": "雪", "nested": [True, None, -4, 0.125]}),
bit_depth=10, color_space="HDR PQ")
payload = tiny_video(variant)
file = VideoFromFile(bytes_io.BytesIO(payload), start_time=-0.0, duration=0, crop=(0, 0, 8, 8))
sequence = VideoFromList([VideoFromFile(bytes_io.BytesIO(payload)), file], complete_audio=audio)
return {"images": [image, image[:1].clone()], "audio": [audio], "masks": [mask],
"texts": [f" café 🎥 {variant}\r\n雪 ", ""], "video": [component, sequence]}
for key, variant in (("local", 0), ("worker-a", 1), ("worker-b", 2)):
fixtures[key] = native_fixtures(variant)
native_root = base / "master-input"
native_root.mkdir()
wav = native_root / "stereo.wav"
pcm = np.array([-32768, -12001, -1, 0, 1, 8193, 32767, 17] * 2, dtype="<i2")
with wave.open(str(wav), "wb") as out:
out.setnchannels(2)
out.setsampwidth(2)
out.setframerate(44100)
out.writeframes(pcm.tobytes())
waveform, sample_rate = nodes_audio.load(str(wav))
clip = native_root / "native.mkv"
clip.write_bytes(tiny_video(3))
native_expected = copy.copy(fixtures["worker-a"])
native_expected["audio"] = [{"waveform": waveform.unsqueeze(0), "sample_rate": sample_rate}]
native_expected["video"] = [VideoFromFile(str(clip))]
expectations = {}
for key, values in (("worker-a", native_expected), ("worker-b", fixtures["worker-b"])):
expectations[key] = {socket: [video_codec.encode_video(v, allowed_roots=[native_root]) if socket == "video" else v
for v in values[socket]] for socket in sockets}
# Native routes only, bound to a new ephemeral loopback listener. No fake
# history, permissive validator, bridge/codec mock, or existing service.
allowed_paths = {"/prompt", "/history/{prompt_id}", "/view", "/queue"}
ps.app.add_routes([route for route in ps.routes if route.path in allowed_paths])
completed = []
stop = asyncio.Event()
async def queue_worker():
while not stop.is_set():
item = await asyncio.to_thread(ps.prompt_queue.get, 0.1)
if item is None:
continue
queued, item_id = item
assert len(queued) in (5, 6), "Unsupported core PromptQueue tuple; adapt acceptance worker explicitly"
executor = execution.PromptExecutor(ps, cache_args={"ram": 0, "ram_inactive": 0}, asset_manager=asset_manager)
await asyncio.to_thread(executor.execute, queued[2], queued[1], queued[3], queued[4])
status = execution.PromptQueue.ExecutionStatus(
status_str="success" if executor.success else "error", completed=executor.success,
messages=executor.status_messages)
ps.prompt_queue.task_done(item_id, executor.history_result, status=status)
completed.append((queued, executor.success, executor.status_messages))
@asynccontextmanager
async def serve():
runner = web.AppRunner(ps.app)
await runner.setup()
site = web.TCPSite(runner, "127.0.0.1", 0)
await site.start()
port = site._server.sockets[0].getsockname()[1]
ALLOWED_CONNECTIONS.add(("127.0.0.1", port))
worker = asyncio.create_task(queue_worker())
try:
yield f"http://127.0.0.1:{port}"
finally:
stop.set()
await asyncio.wait_for(worker, 10)
await runner.cleanup()
ALLOWED_CONNECTIONS.clear()
grouped_results = {}
async with serve() as url:
for key in ("worker-a", "worker-b"):
participant = "serverless:" + key
request_id = "acceptance-" + key
graph = {"source": node("AcceptanceMediaSource", key=key),
"collector": node("DistributedCollector", load_balance=False, transfer_format="exact",
multi_job_id="job-main", worker_id=participant,
**{socket: ["source", index] for index, socket in enumerate(collector_sockets)}),
"audio-collector": node("DistributedCollector", load_balance=False,
multi_job_id="job-audio", worker_id=participant, audio=["source", 1])}
if key == "worker-a":
graph.update({"wav": node("LoadAudio", audio=wav.name), "clip": node("LoadVideo", file=clip.name)})
graph["collector"]["inputs"]["audio"] = ["wav", 0]
before = copy.deepcopy(graph)
workflow, manifest = compiler.compile_media_prompt(graph, request_id)
assert graph == before
bridge_links = {socket: ["source", sockets.index(socket)] for socket in sockets[2:]}
if key == "worker-a":
workflow["clip"] = copy.deepcopy(graph["clip"])
bridge_links["video"] = ["clip", 0]
add_bridge_exports(workflow, manifest, request_id, participant, "job-bridge", bridge_links)
assert len(manifest) == 6 and {b["kind"] for b in manifest} == set(kinds)
folder_paths.set_input_directory(str(native_root))
workflow_before = copy.deepcopy(workflow)
images, artifacts, refs = await packer.prepare_media_inputs(workflow)
assert images == [] and workflow == workflow_before
if key == "worker-a":
assert {a["kind"] for a in artifacts} == {"AUDIO", "VIDEO"} and len(artifacts) == 2
assert len(refs) == 2
assert {ref["input_name"] for ref in refs} == {"audio", "video"}
assert all(workflow[ref["node_id"]]["class_type"] == "MediaExportV1" for ref in refs)
# Endpoint cannot fall back to the master's local files/loaders.
wav.unlink()
clip.unlink()
print("NATIVE_INPUTS_AND_DIRECT_EXPORT_REFS_OK WAV + native file VIDEO")
else:
assert not artifacts and not refs
folder_paths.set_input_directory(str(base / "input"))
media = {"version": 1, "profile": contract.PROFILE, "request_id": request_id,
"artifacts": artifacts, "inputs": refs, "outputs": manifest}
body = json.loads(serializer(workflow, images, 30, media=media))
assert decoder.validate_prepared_job({"worker": {"id": participant, "media_profile": contract.PROFILE,
"timeout_seconds": 30}, "workflow": workflow, "images": images,
"media": media, "manifest": manifest}) is True
before = copy.deepcopy(body["input"])
result = await handler.execute_media_request(body["input"], comfy_url=url,
input_dir=base / "input", temp_dir=base / "temp", timeout_seconds=30)
assert body["input"] == before
assert completed and completed[-1][1], completed
queued = completed[-1][0]
assert set(queued[4]) == {b["node_id"] for b in manifest}
submitted = queued[2]
if key == "worker-a":
assert "wav" not in submitted and "clip" not in submitted
imports = [n for n in submitted.values() if n["class_type"] == "MediaImportV1"]
assert len(imports) == 2
assert {json.loads(n["inputs"]["artifact_json"])["kind"] for n in imports} == {"AUDIO", "VIDEO"}
assert not list((base / "input").rglob("*.bin"))
assert not list((base / "temp").rglob("*.bin"))
if key == "worker-b":
# Retire the originals as well as handler-owned staging/export
# files before constructing the result's lazy native videos.
for original in fixtures[key]["video"][1].videos:
vars(original)["_VideoFromFile__file"].close()
# Exercise the wire format, not an in-memory handler-only object.
result = json.loads(json.dumps(result, allow_nan=False))
grouped = await decoder.collect_typed_results(result, media, manifest)
assert set(grouped) == {"job-main", "job-audio", "job-bridge"}
assert set(grouped["job-main"]) == set(collector_sockets)
assert set(grouped["job-bridge"]) == set(sockets[2:])
exact("AUDIO", grouped["job-audio"]["audio"][0], fixtures[key]["audio"][0])
all_media = {**grouped["job-main"], **grouped["job-bridge"]}
for socket, kind in zip(sockets, kinds):
values, expected = all_media[socket], expectations[key][socket]
assert len(values) == len(expected), (key, socket)
for actual, wanted in zip(values, expected):
exact(kind, actual, wanted)
if key == "worker-b":
original = fixtures[key]["video"][0].get_components()
decoded = grouped["job-bridge"]["video"][0].get_components()
for field in ("images", "alpha"):
assert getattr(decoded, field).data_ptr() != getattr(original, field).data_ptr()
assert decoded.audio["waveform"].data_ptr() != original.audio["waveform"].data_ptr()
assert decoded.metadata == original.metadata and decoded.metadata is not original.metadata
grouped_results[participant] = grouped
# Late invalid artifact must reject the entire grouped response.
corrupt = copy.deepcopy(result)
corrupt["media"]["artifacts"][-1]["sha256"] = "0" * 64
try:
await decoder.collect_typed_results(corrupt, media, manifest)
except ValueError:
pass
else:
raise AssertionError("Result decoder accepted a corrupt final artifact")
assert ps.distributed_pending_jobs == {}
# Exercise user-selected image return encoding through the actual
# compiler, HTTP handler, V3 exporter and atomic response decoder.
for fmt in ("legacy_png", "png", "jpeg", "h264"):
key = "encoded-" + fmt
fixtures[key] = native_fixtures(3)
# Keep full alpha precision and use nontrivial image dimensions.
torch.manual_seed(41)
rgba = torch.rand((2, 64, 64, 4), dtype=torch.float64)
fixtures[key]["images"] = [rgba]
request_id = "request-" + key
graph = {"source": node("AcceptanceMediaSource", key=key),
"collector": node("DistributedCollector", transfer_format=fmt,
transfer_quality=73, multi_job_id="job-" + key,
worker_id="serverless:" + key,
**{socket: ["source", i] for i, socket in enumerate(collector_sockets)})}
if fmt == "legacy_png":
# A saved Collector with no newly selected controls still uses
# legacy_png / 90, mapped to PNG by the serverless compiler.
graph["collector"]["inputs"].pop("transfer_format")
graph["collector"]["inputs"].pop("transfer_quality")
original = copy.deepcopy(graph)
workflow, manifest = compiler.compile_media_prompt(graph, request_id)
if fmt == "legacy_png":
binding = next(b for b in manifest if b["kind"] == "IMAGE")
controls = workflow[binding["node_id"]]["inputs"]
assert (controls["transfer_format"], controls["transfer_quality"]) == ("png", 90)
add_bridge_exports(workflow, manifest, request_id, "serverless:" + key, "job-bridge-" + key,
{socket: ["source", sockets.index(socket)] for socket in sockets[2:]})
media = {"version": 1, "profile": contract.PROFILE, "request_id": request_id,
"artifacts": [], "inputs": [], "outputs": manifest}
body = json.loads(serializer(workflow, [], 30, media=media))
result = await handler.execute_media_request(body["input"], comfy_url=url,
input_dir=base / "input", temp_dir=base / "temp", timeout_seconds=30)
assert graph == original
grouped = await decoder.collect_typed_results(result, media, manifest)
actual = {**grouped["job-" + key], **grouped["job-bridge-" + key]}
image = actual["images"][0]
assert image.shape == rgba.shape and image.dtype == rgba.dtype
assert torch.equal(image[..., 3], rgba[..., 3]), "alpha must not be discarded/narrowed"
assert not torch.equal(image[..., :3], rgba[..., :3]), "selection must actually encode"
if fmt in ("legacy_png", "png"):
pixels = (rgba[..., :3].float().clamp(0, 1) * 255).to(torch.uint8)
assert torch.equal(image[..., :3], (pixels.float() / 255).to(rgba.dtype))
for socket, kind in zip(sockets[1:], kinds[1:]):
for got, wanted in zip(actual[socket], fixtures[key][socket]):
exact(kind, got, video_codec.encode_video(wanted) if kind == "VIDEO" else wanted)
assert not list((base / "temp").rglob("*.bin"))
print("IMAGE_TRANSFER_HTTP_OK default legacy_png -> png/90; png jpeg h264; alpha and other media exact")
print("REAL_HTTP_BRIDGE_RESULTS_OK Collector IMAGE/AUDIO + independent MASK/STRING/VIDEO bridge exports; native VIDEO survives cleanup")
print("ATOMIC_DECODE_REJECTION_OK corrupt final artifact rejected before delivery")
participants = list(grouped_results)
def expected_batches(ordered, include_images=True):
audio_items = [item for media in ordered for item in media["audio"]]
expected = {"audio": [{"waveform": torch.cat([item["waveform"] for item in audio_items], dim=-1),
"sample_rate": audio_items[0]["sample_rate"]}]}
if include_images:
expected["images"] = [torch.cat([image for media in ordered for image in media["images"]])]
return expected
for delegate in (False, True):
tag = "delegate" if delegate else "master"
job = "collect-" + tag
queue = asyncio.Queue()
for participant in reversed(participants):
await queue.put({"worker_id": participant, "media": grouped_results[participant]["job-main"], "is_last": True})
ps.distributed_pending_jobs[job] = queue
ordered = ([] if delegate else [fixtures["local"]]) + [grouped_results[p]["job-main"] for p in participants]
expected_sinks[tag] = expected_batches(ordered)
graph = {"collector": node("DistributedCollector", load_balance=False, multi_job_id=job,
enabled_worker_ids=json.dumps(participants), delegate_only=delegate, transfer_format="exact"),
"assert": node("AcceptanceMediaAssert", tag=tag,
**{socket: ["collector", index] for index, socket in enumerate(collector_sockets)})}
if not delegate:
graph["source"] = node("AcceptanceMediaSource", key="local")
graph["collector"]["inputs"].update({socket: ["source", index] for index, socket in enumerate(collector_sockets)})
# A second audio-only Collector consumes its own job group from the same
# atomic response. It must return None for images and batch audio in order.
audio_job, audio_tag = job + "-audio", tag + "-audio"
audio_queue = asyncio.Queue()
for participant in reversed(participants):
await audio_queue.put({"worker_id": participant,
"media": grouped_results[participant]["job-audio"], "is_last": True})
ps.distributed_pending_jobs[audio_job] = audio_queue
audio_ordered = ([] if delegate else [fixtures["local"]]) + [grouped_results[p]["job-audio"] for p in participants]
expected_sinks[audio_tag] = expected_batches(audio_ordered, include_images=False)
graph["audio-collector"] = node("DistributedCollector", load_balance=False, multi_job_id=audio_job,
enabled_worker_ids=json.dumps(participants), delegate_only=delegate)
graph["audio-assert"] = node("AcceptanceMediaAssert", tag=audio_tag, audio=["audio-collector", 1])
if not delegate:
graph["audio-collector"]["inputs"]["audio"] = ["source", 1]
valid = await execution.validate_prompt(job, graph, None)
assert valid[0] and not valid[3], valid
executor = execution.PromptExecutor(ps, cache_args={"ram": 0, "ram_inactive": 0}, asset_manager=asset_manager)
await asyncio.to_thread(executor.execute, graph, job, {}, valid[2])
assert executor.success, executor.status_messages
assert executor.history_result["outputs"]["assert"]["accepted"] == [seen_sinks[tag]]
assert executor.history_result["outputs"]["audio-assert"]["accepted"] == [seen_sinks[audio_tag]]
assert job not in ps.distributed_pending_jobs and audio_job not in ps.distributed_pending_jobs
assert source_calls == ["worker-a", "worker-b", "encoded-legacy_png", "encoded-png", "encoded-jpeg", "encoded-h264", "local"], source_calls
assert set(seen_sinks) == {"master", "delegate", "master-audio", "delegate-audio"}
assert not list((base / "output").rglob("*"))
assert not torch.cuda.is_initialized()
print("MASTER_AND_DELEGATE_ORDER_OK", json.dumps(seen_sinks, sort_keys=True))
print("MEDIA_END_TO_END_CPU_OK real PromptExecutor/HTTP; no queue-orchestration or provider acceptance claimed")
ALLOWED_CONNECTIONS = set()
def network_guard(event, args):
"""Fail closed on provider/normal-instance traffic, including during imports."""
if event == "socket.connect":
sock, address = args
if sock.family in (socket.AF_INET, socket.AF_INET6):
assert address in ALLOWED_CONNECTIONS, "Acceptance forbids non-owned network connections"
elif event == "socket.bind":
sock, address = args
if sock.family in (socket.AF_INET, socket.AF_INET6):
assert address == ("127.0.0.1", 0), "Acceptance binds only ephemeral loopback"
if __name__ == "__main__":
sys.addaudithook(network_guard)
asyncio.run(asyncio.wait_for(check(bootstrap()), timeout=120))
+48 -5
View File
@@ -36,7 +36,7 @@ BASELINE = json.loads((ROOT / 'tests/fixtures/v1_node_contracts.json').read_text
ADDED_OPTIONAL_INPUTS = {
'DistributedCollector': {
'transfer_format': [
['legacy_png', 'png', 'ffv1', 'jpeg', 'webp', 'h264', 'h265', 'av1'],
['legacy_png', 'png', 'ffv1', 'jpeg', 'webp', 'h264', 'h265', 'av1', 'exact'],
{'default': 'legacy_png'},
],
'transfer_quality': ['INT', {'default': 90, 'min': 1, 'max': 100, 'step': 1}],
@@ -45,7 +45,8 @@ ADDED_OPTIONAL_INPUTS = {
# Nodes added after the V1 baseline; not part of saved-workflow contracts.
# DistributedWorkerSink is internal (dev-only) and only inserted by the master.
ADDED_NODES = {'DistributedWorkerSink', 'DistributedImageListStart', 'DistributedImageListEnd'}
ADDED_NODES = {'DistributedWorkerSink', 'DistributedImageListStart', 'DistributedImageListEnd',
'DistributedImageListGetValue'}
def normalized_input(value):
@@ -273,7 +274,7 @@ async def check_image_list_lazy_executor(mapping, module, prompt_server, asset_m
original_contexts = set(runtime._MASTER_JOBS)
image_list_nodes = sys.modules[module.__name__ + '.nodes.image_list']
workflow_meta = {'nodes': [{'id': 2, 'type': 'DistributedImageListStart'}], 'extra': {'acceptance': True}}
async def fake_worker_results(_context_id, *, model_callback=None):
async def fake_worker_results(_context_id, *, model_callback=None, include_values=False):
state = runtime.get_master_job(_context_id)
assert state.prompt_id == 'image-list-lazy', state.prompt_id
assert state.workflow_meta == workflow_meta, state.workflow_meta
@@ -281,6 +282,7 @@ async def check_image_list_lazy_executor(mapping, module, prompt_server, asset_m
return (
[torch.full((1, 8, 8, 3), value, dtype=torch.float32) for value in (0.125, 0.5, 0.875)],
[None, None, None], [None, None, None],
[{}, {}, {}],
)
graph = {
@@ -366,10 +368,16 @@ async def main():
extension = await module.comfy_entrypoint()
classes = await extension.get_node_list()
mapping = {cls.GET_SCHEMA().node_id: cls for cls in classes}
assert len(classes) == len(mapping) == len(BASELINE) + len(ADDED_NODES) == 11
assert len(classes) == len(mapping) == len(BASELINE) + len(ADDED_NODES) == 12
assert set(mapping) == set(BASELINE) | ADDED_NODES
for node_id in ADDED_NODES:
assert nodes.NODE_CLASS_MAPPINGS[node_id] is mapping[node_id]
assert 'DistributedMediaCollector' not in mapping
collector_info = mapping['DistributedCollector'].GET_NODE_INFO_V1()
assert collector_info['output'] == ['IMAGE', 'AUDIO'], collector_info
assert not any(collector_info['output_is_list']), collector_info
assert list(collector_info['input']['optional']) == [
'images', 'audio', 'transfer_format', 'transfer_quality']
sink_info = mapping['DistributedWorkerSink'].GET_NODE_INFO_V1()
assert sink_info['output_node'] and sink_info['dev_only'], sink_info
for node_id, cls in mapping.items():
@@ -423,6 +431,42 @@ async def main():
assert start_schema.display_name == 'Distributed Image List Start'
assert end_schema.node_id == 'DistributedImageListEnd'
assert end_schema.display_name == 'Distributed Image List End'
runtime = sys.modules[module.__name__ + '.image_list.runtime']
assert [input.id for input in end_schema.inputs] == [
'images', 'masks', 'texts', 'context', '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[6:]:
assert override.optional and override.lazy
assert override.get_io_type() == 'IMAGE,MASK,STRING,INT,FLOAT,BOOLEAN'
getter = mapping['DistributedImageListGetValue']
getter_schema = getter.GET_SCHEMA()
assert getter_schema.display_name == 'Distributed Image List Get Value'
assert not getter_schema.is_input_list
assert getter_schema.outputs[0].id == getter_schema.outputs[0].display_name == 'value'
assert getter_schema.outputs[0].io_type == '*'
assert getter_schema.inputs[1].options == list(runtime.VALUE_NAMES)
assert getter_schema.inputs[1].default == 'value_1'
# Real core maps scalar mappings, never arbitrary tensor batches.
tensor = torch.ones(3, 2, 2, 3)
data, missing, hidden = execution.get_input_data({'value_name': 'value_1'}, getter, 'get-value')
data['values'] = [{'value_1': tensor}, {'value_1': 17}, {'value_1': False}]
mapped = await execution._async_map_node_over_list('get-value', 'get-value', getter, data,
getter.FUNCTION, v3_data=hidden)
mapped = await execution.resolve_map_node_over_list_results(mapped)
assert len(mapped) == 3 and mapped[0].result[0] is tensor
assert mapped[1].result == (17,) and mapped[2].result == (False,)
try:
getter.execute({}, 'value_1')
raise AssertionError('missing Get Value key accepted')
except ValueError as exc:
assert 'absent' in str(exc)
texts_input = next(input for input in end_schema.inputs if input.id == 'texts')
texts_schema = texts_input.as_dict()
assert texts_schema.get('forceInput') is True, texts_schema
@@ -466,7 +510,6 @@ async def main():
rows = runtime_mod.normalize_image_list_inputs(prepared['images'], None, seen['values'])
assert [row['values'] for row in rows] == [{'value_1': 'shared', 'value_3': n} for n in (4, 5, 6)]
print('IMAGE_LIST_AUTOGROW_OK real parser; value_N slots shared/per-item')
assert [out.io_type for out in end_schema.outputs] == ['IMAGE', 'MASK', 'STRING']
assert all(out.is_output_list for out in end_schema.outputs)
assert end_schema.is_input_list and end_schema.is_output_node
print('IMAGE_LIST_SCHEMA_OK exact IDs, names, typed list outputs')
+55
View File
@@ -3,6 +3,7 @@ import sys
import types
import unittest
from pathlib import Path
from unittest.mock import AsyncMock, patch
class _PromptQueue:
@@ -87,5 +88,59 @@ class QueuePromptPayloadTests(unittest.IsolatedAsyncioTestCase):
self.assertEqual(extra_data["extra_pnginfo"]["workflow"], {"id": "workflow-1"})
class ExplicitTargetTests(unittest.IsolatedAsyncioTestCase):
def setUp(self):
self.module, self.server = _load_async_helpers_module()
self.prompt = {"9": {"class_type": "Output"}, "10": {"class_type": "Output"}}
self.validate = AsyncMock(return_value=(True, None, ["9"], {}))
self.module.execution.validate_prompt = self.validate
async def test_targets_survive_callback_omission_and_are_queued_exactly(self):
seen = []
def callback(payload):
seen.append(payload["partial_execution_targets"][:])
payload["partial_execution_targets"].append("10")
return {"prompt": payload["prompt"]}
self.server.trigger_on_prompt = callback
targets = ["9"]
await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=targets)
self.assertEqual(seen, [["9"]])
self.assertEqual(targets, ["9"])
self.assertEqual(self.validate.call_args.args[2], ["9"])
self.assertEqual(self.server.prompt_queue.items[0][4], ["9"])
async def test_invalid_requests_never_reach_validation_or_queue(self):
for targets in ([], "9", [9], [True], ["9", "9"], ["missing"], {}):
with self.subTest(targets=targets):
with self.assertRaises(self.module.PromptValidationError):
await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=targets)
self.validate.assert_not_awaited()
self.assertEqual(self.server.prompt_queue.items, [])
async def test_removed_target_after_callback_is_rejected(self):
self.server.trigger_on_prompt = lambda payload: {"prompt": {"10": self.prompt["10"]}}
with self.assertRaises(self.module.PromptValidationError):
await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=["9"])
self.validate.assert_not_awaited()
self.assertEqual(self.server.prompt_queue.items, [])
async def test_partial_subset_extra_roots_and_target_errors_fail_closed(self):
for roots, errors in ((["9"], {}), (["9", "10", "extra"], {}),
(["9", "10"], {"9": {"errors": [{"message": "bad"}]}}),
(["9", "10"], {"dep": {"dependent_outputs": ["10"], "errors": []}})):
with self.subTest(roots=roots, errors=errors):
self.validate.return_value = (True, None, roots, errors)
with self.assertRaises(self.module.PromptValidationError):
await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=["9", "10"])
self.assertEqual(self.server.prompt_queue.items, [])
async def test_none_preserves_legacy_payload_and_validation(self):
seen = []
self.server.trigger_on_prompt = lambda payload: seen.append(payload.copy()) or payload
await self.module.queue_prompt_payload(self.prompt)
self.assertNotIn("partial_execution_targets", seen[0])
self.assertIsNone(self.validate.call_args.args[2])
if __name__ == "__main__":
unittest.main()
+524
View File
@@ -0,0 +1,524 @@
"""Exact Collector transport through real serializers, HTTP, route and queue.
Only ComfyUI's process-wide server/model-management and unrelated orchestration
are replaced. Tensor/audio codecs, multipart bodies and aiohttp are real; no
running ComfyUI service, GPU, provider or external network access is required.
"""
import asyncio
from contextlib import asynccontextmanager
import importlib
import importlib.util
import json
from pathlib import Path
import sys
import types
from unittest.mock import AsyncMock
import pytest
import torch
ROOT = Path(__file__).resolve().parents[1]
PKG = "dist_collector_exact_testpkg"
JOB_ID = "exact-job"
WORKER_ID = "worker-1"
@pytest.fixture
def transport(monkeypatch):
# Other legacy test loaders may leave aiohttp stubs behind. Replace only
# those entries, never unload/reload real extension-backed dependencies.
for name, loaded_module in list(sys.modules.items()):
if (name == "aiohttp" or name.startswith("aiohttp.")) and not getattr(loaded_module, "__file__", None):
monkeypatch.delitem(sys.modules, name)
aiohttp = importlib.import_module("aiohttp")
web = importlib.import_module("aiohttp.web")
def module(name, **attrs):
result = types.ModuleType(name)
result.__dict__.update(attrs)
monkeypatch.setitem(sys.modules, name, result)
return result
def load(relative):
name = f"{PKG}.{relative}"
spec = importlib.util.spec_from_file_location(name, ROOT.joinpath(*relative.split(".")).with_suffix(".py"))
assert spec is not None and spec.loader is not None
result = importlib.util.module_from_spec(spec)
monkeypatch.setitem(sys.modules, name, result)
spec.loader.exec_module(result)
parent, _, child = name.rpartition(".")
setattr(sys.modules[parent], child, result)
return result
module(PKG, __path__=[str(ROOT)])
for package in ("nodes", "api", "utils", "serverless_bridge"):
module(f"{PKG}.{package}", __path__=[str(ROOT / package)])
server = types.SimpleNamespace(routes=web.RouteTableDef(), distributed_pending_jobs={})
module("server", PromptServer=types.SimpleNamespace(instance=server))
mm = module("comfy.model_management", throw_exception_if_processing_interrupted=lambda: None)
progress = module("comfy.utils", ProgressBar=lambda _total: types.SimpleNamespace(update=lambda _n: None))
module("comfy", __path__=[], model_management=mm, utils=progress)
module(f"{PKG}.utils.logging", debug_log=lambda *_a, **_kw: None, log=lambda *_a, **_kw: None)
module(f"{PKG}.utils.config", get_worker_timeout_seconds=lambda: 0.1,
load_config=lambda: {"workers": []}, is_master_delegate_only=lambda: False)
async def handle_api_error(_request, error, status=500):
return web.json_response({"error": str(error)}, status=status)
module(f"{PKG}.utils.network", handle_api_error=handle_api_error,
get_client_session=AsyncMock(), build_worker_url=lambda _worker: "http://unused.invalid",
probe_worker=AsyncMock(return_value=None))
module(f"{PKG}.utils.async_helpers", run_async_in_server_loop=asyncio.run)
module(f"{PKG}.api.queue_orchestration", ensure_distributed_state=lambda: None,
orchestrate_distributed_execution=AsyncMock())
module(f"{PKG}.api.queue_request", parse_queue_request_payload=lambda value: value)
for name in ("utils.constants", "utils.image", "utils.audio_payload", "utils.transfer_codec"):
load(name)
bridge_codec = load("serverless_bridge.codec")
load("serverless_bridge.image_transfer")
image_codec = load("serverless_bridge.image_codec")
collector = load("nodes.collector")
routes = load("api.job_routes")
return types.SimpleNamespace(aiohttp=aiohttp, web=web, server=server,
collector=collector, routes=routes, bridge_codec=bridge_codec,
image_codec=image_codec)
@asynccontextmanager
async def endpoint(transport, monkeypatch, *, request_limit=2 * 1024 * 1024,
chunk_target=1024 * 1024, client_max_size=None, handler=None):
"""A disposable loopback endpoint; record actual HTTP lengths, not estimates."""
queue = asyncio.Queue()
transport.server.distributed_jobs_lock = asyncio.Lock()
transport.server.distributed_pending_jobs = {JOB_ID: queue}
requests = []
web = transport.web
@web.middleware
async def record_request(request, handler):
record = {"bytes": request.content_length, "content_type": request.content_type}
requests.append(record)
response = await handler(request)
record["status"] = response.status
if response.status < 400 and request.content_type.startswith("multipart/"):
form = await request.post()
record["envelope"] = json.loads(form["envelope"])
return response
app = web.Application(client_max_size=client_max_size or request_limit, middlewares=[record_request])
app.router.add_post("/distributed/job_complete", handler or transport.routes.job_complete_endpoint)
runner = web.AppRunner(app)
await runner.setup()
try:
site = web.TCPSite(runner, "127.0.0.1", 0)
await site.start()
port = site._server.sockets[0].getsockname()[1]
async with transport.aiohttp.ClientSession() as session:
monkeypatch.setattr(transport.collector, "get_client_session", AsyncMock(return_value=session))
monkeypatch.setattr(transport.collector, "TRANSFER_MAX_REQUEST_BYTES", request_limit)
monkeypatch.setattr(transport.collector, "TRANSFER_MAX_CHUNK_BYTES", chunk_target)
yield types.SimpleNamespace(
session=session, url=f"http://127.0.0.1:{port}", requests=requests, queue=queue,
node=transport.collector.DistributedCollectorNode(),
)
finally:
await runner.cleanup()
def frames(dtype=torch.float32, channels=3, count=5, height=7, width=9):
# Non-contiguous, out-of-range and non-8-bit-representable values exercise
# detach/contiguous without permitting clipping, alpha loss or quantization.
values = torch.arange(count * height * width * channels, dtype=torch.float64)
values = (values.remainder(41) - 11) / 17
images = values.reshape(count, width, height, channels).transpose(1, 2).to(dtype)
images[0, 0, 0, 0] = -0.0
return images.requires_grad_(True)
def assert_exact(actual, expected):
assert actual.dtype == expected.dtype
assert actual.shape == expected.shape
assert actual.device.type == "cpu"
assert actual.is_contiguous()
# Compare raw bits as well as values, including BF16 and signed zero.
assert torch.equal(actual, expected)
assert torch.equal(actual.view(torch.uint8), expected.detach().contiguous().view(torch.uint8))
def queued_items(queue):
return [queue.get_nowait() for _ in range(queue.qsize())]
async def send(ctx, images, *, audio=None, fmt="exact"):
await ctx.node.send_batch_to_master(images, audio, JOB_ID, ctx.url, WORKER_ID, fmt, 17)
@pytest.mark.parametrize("delegate_only", [False, True])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
@pytest.mark.parametrize("channels", [1, 3, 4])
def test_exact_sender_route_and_collector_preserve_dtype_values_and_channels(
transport, monkeypatch, dtype, channels, delegate_only):
images = frames(dtype, channels)
frame_bytes = images[0].numel() * images.element_size()
reserve = transport.collector.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
async def scenario():
async with endpoint(transport, monkeypatch, request_limit=reserve + 2 * frame_bytes,
chunk_target=2 * frame_bytes) as ctx:
await send(ctx, images)
items = queued_items(ctx.queue)
assert [item["image_index"] for item in items] == list(range(len(images)))
assert [item["is_last"] for item in items] == [False] * 4 + [True]
assert all(item["worker_id"] == WORKER_ID and item["format"] == "exact" for item in items)
for item, expected in zip(items, images.split(1)):
assert_exact(item["tensor"], expected)
await ctx.queue.put(item)
envelopes = [record["envelope"] for record in ctx.requests]
assert [(e["start_index"], e["count"]) for e in envelopes] == [(0, 2), (2, 2), (4, 1)]
assert all(record["bytes"] <= reserve + 2 * frame_bytes for record in ctx.requests)
master_images = None if delegate_only else images[:1].detach().clone().add(0.03125)
actual, _audio = await ctx.node.execute(
master_images, None, multi_job_id=JOB_ID, enabled_worker_ids=json.dumps([WORKER_ID]),
delegate_only=delegate_only, transfer_format="exact",
)
expected = images if delegate_only else torch.cat([master_images, images])
assert_exact(actual, expected)
assert JOB_ID not in transport.server.distributed_pending_jobs
asyncio.run(scenario())
def test_exact_single_frame_can_exceed_chunk_target_but_not_request_limit(transport, monkeypatch):
images = frames(count=3)
frame_bytes = images[0].numel() * images.element_size()
reserve = transport.collector.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
async def scenario():
async with endpoint(transport, monkeypatch, request_limit=reserve + frame_bytes,
chunk_target=frame_bytes // 2) as ctx:
await send(ctx, images)
assert [record["envelope"]["count"] for record in ctx.requests] == [1, 1, 1]
assert all(record["bytes"] <= reserve + frame_bytes for record in ctx.requests)
assert_exact(torch.cat([item["tensor"] for item in queued_items(ctx.queue)]), images)
asyncio.run(scenario())
def test_exact_oversized_frame_is_rejected_before_encoding_or_post(transport, monkeypatch):
images = frames(count=3)
frame_bytes = images[0].numel() * images.element_size()
reserve = transport.collector.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
encode = AsyncMock(side_effect=AssertionError("oversized frame must not be encoded"))
monkeypatch.setattr(transport.collector, "_encode_exact_frames", encode)
async def scenario():
async with endpoint(transport, monkeypatch, request_limit=reserve + frame_bytes - 1) as ctx:
with pytest.raises(transport.collector.transfer_codec.FrameTooLargeError, match="request limit"):
await send(ctx, images)
assert ctx.requests == []
assert ctx.queue.empty()
encode.assert_not_called()
asyncio.run(scenario())
@pytest.mark.parametrize("shape", [(0, 2, 3, 3), (1, 0, 3, 3), (1, 2, 0, 3), (1, 2, 3, 0),
(1, 2, 3, 2), (1, 2, 3, 5), (1, 2, 3), (2, 3)])
def test_exact_sender_rejects_invalid_geometry_before_post(transport, monkeypatch, shape):
image = torch.empty(shape)
with pytest.raises(transport.collector.transfer_codec.TransferCodecError):
transport.collector._exact_frame_bytes(image)
async def scenario():
async with endpoint(transport, monkeypatch) as ctx:
# An empty batch is also a no-media completion error (ValueError).
with pytest.raises(ValueError):
await send(ctx, image)
assert ctx.requests == []
assert ctx.queue.empty()
asyncio.run(scenario())
@pytest.mark.parametrize("dtype", [torch.uint8, torch.int32, torch.bool, torch.complex64])
def test_exact_sender_rejects_unsupported_dtype_before_post(transport, monkeypatch, dtype):
async def scenario():
async with endpoint(transport, monkeypatch) as ctx:
with pytest.raises(transport.collector.transfer_codec.TransferCodecError, match="dtype"):
await send(ctx, torch.zeros(1, 2, 3, 3, dtype=dtype))
assert ctx.requests == []
asyncio.run(scenario())
def test_exact_frame_planning_enforces_bridge_byte_limit(transport, monkeypatch):
monkeypatch.setattr(transport.bridge_codec, "MAX_PAYLOAD_BYTES", 32)
with pytest.raises(transport.collector.transfer_codec.TransferCodecError, match="[Bb]yte|limit"):
transport.collector._exact_frame_bytes(torch.zeros(1, 2, 3, 3))
def test_exact_frame_planning_rejects_non_dense_layout(transport):
image = torch.zeros(1, 2, 3, 3).to_sparse()
with pytest.raises(transport.collector.transfer_codec.TransferCodecError, match="dense|layout|strided"):
transport.collector._exact_frame_bytes(image)
@pytest.mark.parametrize("value", [float("nan"), float("inf"), -float("inf")])
def test_exact_sender_rejects_nonfinite_payload(transport, monkeypatch, value):
images = torch.full((1, 2, 3, 3), value)
async def scenario():
async with endpoint(transport, monkeypatch) as ctx:
with pytest.raises(transport.collector.transfer_codec.TransferCodecError, match="finite"):
await send(ctx, images)
assert ctx.requests == []
assert ctx.queue.empty()
asyncio.run(scenario())
def exact_chunk(transport, count=2):
images = frames(count=count).detach()
encoded = [transport.bridge_codec.encode_payload("IMAGE", frame) for frame in images.split(1)]
return {
"job_id": JOB_ID, "worker_id": WORKER_ID, "format": "exact",
"start_index": 0, "count": count, "width": images.shape[2], "height": images.shape[1],
"is_last": True, "descriptors": [descriptor for descriptor, _ in encoded],
}, [payload for _, payload in encoded]
@pytest.mark.parametrize("case", [
"zero_count", "bool_width", "negative_height", "negative_index", "bool_index",
"missing_descriptor", "extra_descriptor", "missing_part", "extra_part", "not_a_descriptor",
"wrong_height", "wrong_batch", "bad_channels", "bad_dtype", "bad_layout", "bad_encoding",
"extra_descriptor_field", "truncated", "trailing_bytes", "nonfinite", "compressed_spoof",
])
def test_exact_route_rejects_bad_geometry_and_payload_atomically(transport, monkeypatch, case):
envelope, parts = exact_chunk(transport)
descriptor = envelope["descriptors"][-1]
if case == "zero_count":
envelope["count"] = 0
elif case == "bool_width":
envelope["width"] = True
elif case == "negative_height":
envelope["height"] = -1
elif case == "negative_index":
envelope["start_index"] = -1
elif case == "bool_index":
envelope["start_index"] = True
elif case == "missing_descriptor":
envelope["descriptors"].pop()
elif case == "extra_descriptor":
envelope["descriptors"].append(dict(descriptor))
elif case == "missing_part":
parts.pop()
elif case == "extra_part":
parts.append(parts[-1])
elif case == "not_a_descriptor":
envelope["descriptors"][-1] = None
elif case == "wrong_height":
descriptor["shape"][1] += 1
elif case == "wrong_batch":
descriptor["shape"][0] = 2
elif case == "bad_channels":
descriptor["shape"][-1] = 2
elif case == "bad_dtype":
descriptor["dtype"] = "uint8"
elif case == "bad_layout":
descriptor["layout"] = "BCHW"
elif case == "bad_encoding":
descriptor["encoding"] = "pickle"
elif case == "extra_descriptor_field":
descriptor["unknown"] = True
elif case == "truncated":
parts[-1] = parts[-1][:-1]
elif case == "trailing_bytes":
parts[-1] += b"\0"
elif case == "nonfinite":
bad = torch.full(tuple(descriptor["shape"]), float("nan"))
parts[-1] = bad.numpy().tobytes()
elif case == "compressed_spoof":
# A valid serverless image-transfer descriptor must NOT be accepted as
# exact by the direct route, even though the shared bridge can decode it.
envelope["descriptors"][-1], parts[-1] = transport.image_codec.encode_image(
frames(count=1).detach(), transfer_format="png", transfer_quality=90,
)
# Prove the spoof is a valid compressed bridge payload, not simply junk.
decoded = transport.bridge_codec.decode_payload("IMAGE", envelope["descriptors"][-1], parts[-1])
assert list(decoded.shape) == descriptor["shape"]
async def scenario():
async with endpoint(transport, monkeypatch) as ctx:
async with ctx.session.post(
ctx.url + "/distributed/job_complete",
data=ctx.node._build_multipart(envelope, parts),
) as response:
body = await response.text()
assert response.status == 400, body
# A corrupt second frame must not leave the first in the queue.
assert ctx.queue.empty()
asyncio.run(scenario())
@pytest.mark.parametrize("status,body", [(400, "Invalid JSON payload: old master"),
(413, "upload limit"), (500, "unexpected failure")])
def test_exact_master_rejection_never_falls_back(transport, monkeypatch, status, body):
async def reject(_request):
return transport.web.Response(status=status, text=body)
async def scenario():
async with endpoint(transport, monkeypatch, handler=reject, chunk_target=1) as ctx:
legacy = AsyncMock(side_effect=AssertionError("exact must never use legacy transport"))
compressed = AsyncMock(side_effect=AssertionError("exact must never use compressed transport"))
monkeypatch.setattr(ctx.node, "_send_legacy_batch", legacy)
monkeypatch.setattr(ctx.node, "_send_compressed_batch", compressed)
with pytest.raises(RuntimeError, match=f"HTTP {status}") as raised:
await send(ctx, frames(count=3))
assert "never" in str(raised.value) and "8-bit" in str(raised.value)
assert len(ctx.requests) == 1
assert ctx.requests[0]["content_type"].startswith("multipart/")
assert ctx.queue.empty()
legacy.assert_not_called()
compressed.assert_not_called()
asyncio.run(scenario())
@pytest.mark.parametrize("fmt", ["legacy_png", "png"])
def test_exact_collector_rejects_old_worker_8bit_result(transport, monkeypatch, fmt):
async def scenario():
async with endpoint(transport, monkeypatch) as ctx:
await send(ctx, frames(count=1).detach(), fmt=fmt)
assert ctx.queue.qsize() == 1
with pytest.raises(RuntimeError, match="requested exact"):
await ctx.node.execute(
None, None, multi_job_id=JOB_ID, enabled_worker_ids=json.dumps([WORKER_ID]),
delegate_only=True, transfer_format="exact",
)
assert JOB_ID not in transport.server.distributed_pending_jobs
asyncio.run(scenario())
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"):
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",
)
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)
envelope["is_last"] = False
async def scenario():
async with endpoint(transport, monkeypatch) as ctx:
async with ctx.session.post(
ctx.url + "/distributed/job_complete", data=ctx.node._build_multipart(envelope, parts),
) as response:
assert response.status == 200, await response.text()
assert ctx.queue.qsize() == 1
with pytest.raises(RuntimeError, match="Exact transfer incomplete"):
await asyncio.wait_for(ctx.node.execute(
frames(count=1).detach(), None, multi_job_id=JOB_ID,
enabled_worker_ids=json.dumps([WORKER_ID]), transfer_format="exact",
), timeout=2)
assert JOB_ID not in transport.server.distributed_pending_jobs
asyncio.run(scenario())
def test_exact_oversized_multipart_metadata_is_rejected_before_post(transport, monkeypatch):
images = frames(count=1)
reserve = transport.collector.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
limit = reserve + images.numel() * images.element_size()
async def scenario():
async with endpoint(transport, monkeypatch, request_limit=limit) as ctx:
with pytest.raises(transport.collector.transfer_codec.FrameTooLargeError, match="request limit"):
await ctx.node.send_batch_to_master(
images, None, JOB_ID, ctx.url, "w" * limit, "exact", 90,
)
assert ctx.requests == []
assert ctx.queue.empty()
asyncio.run(scenario())
def test_exact_actual_server_upload_rejection_is_413_without_retry(transport, monkeypatch):
async def scenario():
async with endpoint(transport, monkeypatch, client_max_size=1024) as ctx:
with pytest.raises(RuntimeError, match="HTTP 413"):
await send(ctx, frames(count=3, height=64, width=64))
assert len(ctx.requests) == 1
assert ctx.requests[0]["status"] == 413
assert ctx.queue.empty()
asyncio.run(scenario())
def test_exact_oversized_audio_is_rejected_before_any_image_post(transport, monkeypatch):
images = frames(count=3)
audio = {"waveform": torch.zeros(1, 1, 100000), "sample_rate": 48000}
reserve = transport.collector.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
async def scenario():
async with endpoint(transport, monkeypatch, request_limit=reserve + 1024) as ctx:
with pytest.raises(transport.collector.transfer_codec.FrameTooLargeError, match="audio.*request limit"):
await send(ctx, images, audio=audio)
assert ctx.requests == []
assert ctx.queue.empty()
asyncio.run(scenario())
def test_exact_tiny_frames_cannot_overflow_multipart_metadata_budget(transport, monkeypatch):
images = torch.full((2048, 1, 1, 1), 0.125, dtype=torch.float16)
reserve = transport.collector.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
limit = reserve + images.numel() * images.element_size()
async def scenario():
async with endpoint(transport, monkeypatch, request_limit=limit, chunk_target=limit) as ctx:
await send(ctx, images)
assert len(ctx.requests) > 1
assert all(record["bytes"] <= limit for record in ctx.requests)
assert_exact(torch.cat([item["tensor"] for item in queued_items(ctx.queue)]), images)
asyncio.run(scenario())
@pytest.mark.parametrize("separate", [False, True])
def test_exact_audio_is_preserved_and_only_final_completion_is_terminal(transport, monkeypatch, separate):
images = frames(count=3)
audio = {"waveform": torch.linspace(-0.9, 0.9, 60).reshape(1, 1, -1), "sample_rate": 48000}
frame_bytes = images[0].numel() * images.element_size()
reserve = transport.collector.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
media_budget = frame_bytes if separate else frame_bytes + 8192
limit = reserve + media_budget
async def scenario():
async with endpoint(transport, monkeypatch, request_limit=limit, chunk_target=frame_bytes) as ctx:
await send(ctx, images, audio=audio)
items = queued_items(ctx.queue)
assert len(ctx.requests) == 3 + int(separate)
assert [item["is_last"] for item in items] == [False] * (len(items) - 1) + [True]
assert all(item["audio"] is None for item in items[:-1])
assert items[-1]["audio"]["sample_rate"] == audio["sample_rate"]
assert_exact(items[-1]["audio"]["waveform"], audio["waveform"])
assert_exact(torch.cat([item["tensor"] for item in items if item["tensor"] is not None]), images)
assert all(record["bytes"] <= limit for record in ctx.requests)
if separate:
assert items[-1]["tensor"] is None
assert ctx.requests[-1]["content_type"] == "application/json"
asyncio.run(scenario())
+13 -5
View File
@@ -4,6 +4,7 @@ import types
import asyncio
from pathlib import Path
import pytest
import torch
@@ -19,6 +20,13 @@ def _register_real_transfer_codec(package_name, utils_pkg):
def _load_collector_module():
# Restore shared dependencies even when a different test imports this loader.
# Keep real modules imported during loading (notably extension modules) cached.
with pytest.MonkeyPatch.context() as monkeypatch:
return _load_collector_module_with_stubs(monkeypatch)
def _load_collector_module_with_stubs(monkeypatch):
module_path = Path(__file__).resolve().parents[1] / "nodes" / "collector.py"
package_name = "dist_collector_list_testpkg"
@@ -53,7 +61,7 @@ def _load_collector_module():
)
server_module = types.ModuleType("server")
server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
sys.modules["server"] = server_module
monkeypatch.setitem(sys.modules, "server", server_module)
comfy_module = types.ModuleType("comfy")
model_management = types.ModuleType("comfy.model_management")
@@ -77,13 +85,13 @@ def _load_collector_module():
comfy_utils.ProgressBar = ProgressBar
comfy_module.utils = comfy_utils
sys.modules["comfy"] = comfy_module
sys.modules["comfy.model_management"] = model_management
sys.modules["comfy.utils"] = comfy_utils
monkeypatch.setitem(sys.modules, "comfy", comfy_module)
monkeypatch.setitem(sys.modules, "comfy.model_management", model_management)
monkeypatch.setitem(sys.modules, "comfy.utils", comfy_utils)
aiohttp_module = types.ModuleType("aiohttp")
aiohttp_module.ClientTimeout = lambda total: types.SimpleNamespace(total=total)
sys.modules["aiohttp"] = aiohttp_module
monkeypatch.setitem(sys.modules, "aiohttp", aiohttp_module)
logging_module = types.ModuleType(f"{package_name}.utils.logging")
logging_module.debug_log = lambda *_args, **_kwargs: None
+112
View File
@@ -0,0 +1,112 @@
"""Real Collector aggregation of typed bridge deliveries and direct results.
Only Comfy's host services are stubbed by the shared Collector loader. Tensor
normalization, participant ordering, queue consumption and cleanup are real.
"""
import asyncio
import json
import types
from unittest.mock import AsyncMock
import pytest
import torch
from test_collector_list_inputs import _load_collector_module
@pytest.fixture
def collector_module():
module = _load_collector_module()
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
module.prompt_server.distributed_serverless = types.SimpleNamespace(cancel=AsyncMock())
return module
def native_media(value=0):
image = torch.arange(24, dtype=torch.float64).reshape(1, 2, 3, 4) / 997 + value
image[0, 0, 0, 0] = -0.0
audio = {"waveform": torch.arange(6, dtype=torch.float64).reshape(1, 2, 3) / 997 + value,
"sample_rate": 24000}
return {"images": image, "audio": audio}
def assert_exact(actual, expected):
assert actual.dtype == expected.dtype and actual.shape == expected.shape
assert torch.equal(actual.contiguous().view(torch.uint8), expected.contiguous().view(torch.uint8))
def collect_queue(module, messages, participants, *, local=None, delegate=True):
queue = asyncio.Queue()
for message in messages:
queue.put_nowait(message)
module.prompt_server.distributed_pending_jobs["typed-job"] = queue
values = local or {"images": None, "audio": None}
return asyncio.run(asyncio.wait_for(module.DistributedCollectorNode().execute(
**values, multi_job_id="typed-job", enabled_worker_ids=json.dumps(participants),
delegate_only=delegate), timeout=2))
@pytest.mark.parametrize("delegate", [False, True])
@pytest.mark.parametrize("socket_set", [("images",), ("audio",), ("images", "audio")])
def test_typed_and_direct_results_preserve_batch_and_audio_order(collector_module, delegate, socket_set):
module = collector_module
local, typed_a, typed_b, direct = [native_media(n) for n in range(4)]
typed = {socket: [typed_a[socket], typed_b[socket]] for socket in socket_set}
direct_message = {"worker_id": "direct", "is_last": True}
if "images" in socket_set:
direct_message.update(tensor=direct["images"], image_index=0)
if "audio" in socket_set:
direct_message["audio"] = direct["audio"]
messages = [direct_message, {"worker_id": "serverless:pool:1", "media": typed, "is_last": True}]
# Delivery order differs from participant order. Duplicate configured IDs
# must not duplicate a participant's data or leave collection waiting.
participants = ["serverless:pool:1", "direct", "serverless:pool:1"]
local_inputs = {socket: local[socket] if socket in socket_set else None for socket in ("images", "audio")}
images, audio = collect_queue(module, messages, participants, local=local_inputs, delegate=delegate)
expected = ([] if delegate else [local]) + [typed_a, typed_b, direct]
if "images" in socket_set:
assert_exact(images, torch.cat([row["images"] for row in expected]))
else:
assert images is None
if "audio" in socket_set:
assert audio["sample_rate"] == 24000
assert_exact(audio["waveform"], torch.cat([row["audio"]["waveform"] for row in expected], dim=-1))
else:
assert torch.equal(audio["waveform"], module.DistributedCollectorNode.EMPTY_AUDIO["waveform"])
assert module.prompt_server.distributed_pending_jobs == {}
module.prompt_server.distributed_serverless.cancel.assert_not_awaited()
@pytest.mark.parametrize("media", [{}, {"texts": ["not a Collector socket"]}, {"images": []},
{"audio": [None]}, {"images": "not a list"},
{"images": [torch.zeros(2, 2, 3)]}])
def test_invalid_typed_delivery_cancels_and_removes_queue(collector_module, media):
module = collector_module
with pytest.raises(RuntimeError, match="media|items|IMAGE"):
collect_queue(module, [{"worker_id": "serverless:pool:1", "media": media, "is_last": True}],
["serverless:pool:1"])
assert module.prompt_server.distributed_pending_jobs == {}
module.prompt_server.distributed_serverless.cancel.assert_awaited_once_with(collector_job_id="typed-job")
@pytest.mark.parametrize("failure", ["nonfinal", "unexpected", "duplicate", "provider-error", "shape"])
def test_typed_delivery_protocol_failures_are_not_partial_success(collector_module, failure):
media = {"images": [native_media()["images"]]}
first = {"worker_id": "serverless:pool:1", "media": media, "is_last": True}
messages = [first]
participants = ["serverless:pool:1", "serverless:pool:2"]
if failure == "nonfinal":
first["is_last"] = False
elif failure == "unexpected":
first["worker_id"] = "other"
elif failure == "duplicate":
messages.append(dict(first))
elif failure == "provider-error":
messages.append({"worker_id": "serverless:pool:2", "serverless_error": "offline failure", "is_last": True})
else:
messages.append({"worker_id": "serverless:pool:2", "media": {"images": [torch.zeros(1, 4, 5, 3)]},
"is_last": True})
with pytest.raises(RuntimeError, match="atomic|Unexpected|Duplicate|failed|shape mismatch"):
collect_queue(collector_module, messages, participants)
assert collector_module.prompt_server.distributed_pending_jobs == {}
collector_module.prompt_server.distributed_serverless.cancel.assert_awaited_once_with(collector_job_id="typed-job")
+75 -1
View File
@@ -4,7 +4,7 @@ import sys
import types
import unittest
from pathlib import Path
from unittest.mock import patch
from unittest.mock import AsyncMock, patch
def _load_dispatch_module():
@@ -168,6 +168,7 @@ class _FakeSession:
def _valid_capability():
return {
"supported": True,
"execution_targets": True,
"protocol_version": 2,
"node_ids": ["DistributedImageListStart", "DistributedImageListEnd"],
}
@@ -177,6 +178,43 @@ def _valid_queue(running=None, pending=None):
return {"queue_running": running or [], "queue_pending": pending or []}
class ExplicitTargetTransportTests(unittest.IsolatedAsyncioTestCase):
async def test_http_preserves_targets_and_omits_default_field(self):
response = types.SimpleNamespace(raise_for_status=lambda: None)
class Response:
async def __aenter__(self): return response
async def __aexit__(self, *args): pass
calls = []
def post(url, **kwargs):
calls.append(kwargs['json'])
return Response()
session = types.SimpleNamespace(post=post)
with patch.object(dispatch, 'get_client_session', AsyncMock(return_value=session)):
await dispatch.dispatch_worker_prompt({'id': 'worker'}, {'9': {}}, None,
partial_execution_targets=['9'])
await dispatch.dispatch_worker_prompt({'id': 'worker'}, {'9': {}}, None)
self.assertEqual(calls[0]['partial_execution_targets'], ['9'])
self.assertNotIn('partial_execution_targets', calls[1])
async def test_websocket_sends_targets_all_the_way_to_wire(self):
import json
sent = []
async def receive():
return types.SimpleNamespace(type=dispatch.aiohttp.WSMsgType.TEXT, data=json.dumps({
'type': 'dispatch_ack', 'request_id': sent[-1]['request_id'],
'ok': True, 'prompt_id': 'worker-prompt'}))
ws = types.SimpleNamespace(send_json=AsyncMock(side_effect=sent.append),
receive=receive, close=AsyncMock())
session = types.SimpleNamespace(ws_connect=AsyncMock(return_value=ws))
with patch.object(dispatch, 'get_client_session', AsyncMock(return_value=session)):
result = await dispatch.dispatch_worker_prompt({'id': 'worker'}, {'9': {}}, None,
use_websocket=True, return_prompt_id=True, partial_execution_targets=['9'])
await dispatch.dispatch_worker_prompt({'id': 'worker'}, {'9': {}}, None, use_websocket=True)
self.assertEqual(result, 'worker-prompt')
self.assertEqual(sent[0]['partial_execution_targets'], ['9'])
self.assertNotIn('partial_execution_targets', sent[1])
class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
async def test_select_active_workers_filters_offline(self):
workers = [
@@ -417,6 +455,42 @@ class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
with self.assertRaisesRegex(RuntimeError, "No selected feature-compatible direct worker is idle"):
await dispatch.preflight_image_list_workers(workers, require_idle=True)
async def test_image_list_preflight_requires_exact_execution_targets_support(self):
for value in (None, False, 1, "true", []):
with self.subTest(value=value):
capability = _valid_capability()
if value is None:
capability.pop("execution_targets")
else:
capability["execution_targets"] = value
session = _FakeSession({("worker", "capabilities"): (200, capability)})
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
with self.assertRaisesRegex(RuntimeError, "execution_targets.*upgrade"):
await dispatch.preflight_image_list_workers([{"id": "worker"}])
async def test_image_list_value_results_requires_exact_true_only_for_overrides(self):
for value in (None, False, 1, "true", [], True):
capability = _valid_capability()
if value is not None:
capability['value_results'] = value
session = _FakeSession({('worker', 'capabilities'): (200, capability),
('worker', 'queue'): (200, _valid_queue())})
async def fake_session():
return session
with patch.object(dispatch, 'get_client_session', side_effect=fake_session), patch.object(
dispatch, 'build_worker_url', side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
self.assertEqual(await dispatch.preflight_image_list_workers([{'id': 'worker'}]), {'worker': True})
if value is True:
self.assertEqual(await dispatch.preflight_image_list_workers(
[{'id': 'worker'}], require_value_results=True), {'worker': True})
else:
with self.assertRaisesRegex(RuntimeError, 'value_results.*upgrade'):
await dispatch.preflight_image_list_workers([{'id': 'worker'}], require_value_results=True)
async def test_image_list_preflight_fails_closed_for_malformed_capabilities(self):
malformed = {"supported": True, "protocol_version": 2, "node_ids": None}
session = _FakeSession({("worker", "capabilities"): (200, malformed)})
+16 -8
View File
@@ -202,8 +202,10 @@ class ImageListGraphBoundaryTests(unittest.TestCase):
prompt["2"] = {"class_type": "VAEDecode", "inputs": {"samples": ["7", 0]}}
prompt["7"] = {"class_type": "KSampler", "inputs": {"model": ["1", 0]}}
prompt["4"]["inputs"]["reference"] = ["2", 0]
with self.assertRaisesRegex(ValueError, "bypasses Start.*2"):
validate_image_list_regions(prompt)
for explicit_targets in (False, True):
with self.subTest(explicit_targets=explicit_targets):
with self.assertRaisesRegex(ValueError, "bypasses Start.*2"):
validate_image_list_regions(prompt, explicit_targets=explicit_targets)
def test_shared_typed_loader_and_leaf_constant_remain_valid(self):
import sys
@@ -247,19 +249,25 @@ class ImageListGraphBoundaryTests(unittest.TestCase):
def test_rejects_missing_context_pair_nested_region_and_cross_boundary_branch(self):
prompt = graph_one_region()
prompt["5"]["inputs"]["context"] = ["2", 0]
with self.assertRaisesRegex(ValueError, "typed context output"):
validate_image_list_regions(prompt)
for explicit_targets in (False, True):
with self.subTest(explicit_targets=explicit_targets):
with self.assertRaisesRegex(ValueError, "typed context output"):
validate_image_list_regions(prompt, explicit_targets=explicit_targets)
nested = graph_one_region()
nested["7"] = start_node(["4", 0])
nested["8"] = {"class_type": "NestedOperation", "inputs": {"image": ["7", IMAGE_OUT]}}
nested["5"]["inputs"]["images"] = ["8", 0]
nested["9"] = {"class_type": END_CLASS, "inputs": {"images": ["8", 0], "context": ["7", CONTEXT_OUT]}}
with self.assertRaisesRegex(ValueError, "Nested/ambiguous"):
validate_image_list_regions(nested)
for explicit_targets in (False, True):
with self.subTest(explicit_targets=explicit_targets):
with self.assertRaisesRegex(ValueError, "Nested/ambiguous"):
validate_image_list_regions(nested, explicit_targets=explicit_targets)
outside = graph_one_region()
outside["7"] = {"class_type": "PreviewImage", "inputs": {"images": ["3", IMAGE_OUT]}}
with self.assertRaisesRegex(ValueError, "cross-boundary"):
validate_image_list_regions(outside)
for explicit_targets in (False, True):
with self.subTest(explicit_targets=explicit_targets):
with self.assertRaisesRegex(ValueError, "cross-boundary"):
validate_image_list_regions(outside, explicit_targets=explicit_targets)
def test_rejects_duplicate_or_missing_results(self):
prompt = graph_one_region()
+245
View File
@@ -0,0 +1,245 @@
"""CPU-only execution-copy planning; no ComfyUI imports or prompt execution."""
import copy
import sys
import types
import pytest
from image_list.execution_plan import prepare_image_list_execution
from image_list.runtime import (
CONTEXT_OUT,
END_CLASS,
IMAGE_OUT,
START_CLASS,
build_worker_prompt,
validate_image_list_regions,
)
class OutputNode:
OUTPUT_NODE = True
class SchemaOutputNode:
OUTPUT_NODE = False # V3's schema, not this fallback, is authoritative.
@classmethod
def GET_NODE_INFO_V1(cls):
return {"output_node": True, "output": ("IMAGE",)}
class ModelLoader:
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
@pytest.fixture
def registry(monkeypatch):
classes = {
END_CLASS: SchemaOutputNode,
"SaveImage": OutputNode,
"PreviewImage": OutputNode,
"OutputUtility": OutputNode,
"SchemaUtility": SchemaOutputNode,
"ModelLoader": ModelLoader,
}
monkeypatch.setitem(sys.modules, "nodes", types.SimpleNamespace(NODE_CLASS_MAPPINGS=classes))
return classes
def node(class_type, **inputs):
return {"class_type": class_type, "inputs": inputs}
def graph():
return {
"loader": node("ModelLoader"),
"source": node("ImageSource"),
"start": node(START_CLASS, images=["source", 0]),
"process": node("FrameOperation", image=["start", IMAGE_OUT], model=["loader", 0]),
"end": node(END_CLASS, images=["process", 0], context=["start", CONTEXT_OUT]),
"save": node("SaveImage", images=["end", 0]),
}
def test_discards_multistep_previews_but_preserves_other_output_roots(registry):
prompt = graph()
prompt.update({
"debug-transform": node("DebugTransform", images=["start", IMAGE_OUT]),
"debug-preview": node("PreviewImage", images=["debug-transform", 0]),
"process-preview": node("PreviewImage", images=["process", 0]),
"unrelated-source": node("ImageSource"),
"unrelated-preview": node("PreviewImage", images=["unrelated-source", 0]),
"upstream-preview": node("PreviewImage", images=["source", 0]),
"post-transform": node("FrameOperation", images=["end", 0]),
"post-preview": node("PreviewImage", images=["post-transform", 0]),
})
before = copy.deepcopy(prompt)
with pytest.raises(ValueError, match="cross-boundary"):
validate_image_list_regions(prompt)
prepared, targets = prepare_image_list_execution(prompt)
assert set(prepared) == set(prompt) - {"debug-transform", "debug-preview", "process-preview"}
assert targets == ["end", "save", "unrelated-preview", "upstream-preview", "post-preview"]
validate_image_list_regions(prepared, explicit_targets=True)
assert prompt == before
@pytest.mark.parametrize("class_type", ["OutputUtility", "SchemaUtility"])
def test_required_output_utility_keeps_complete_dependency_chain(registry, class_type):
prompt = graph()
prompt.update({
"utility-config": node("Constant", value=7),
"utility-first": node(class_type, images=["process", 0], config=["utility-config", 0]),
"utility-second": node(class_type, images=["utility-first", 0]),
"result": node("FrameOperation", image=["utility-second", 0]),
"discarded-preview": node("PreviewImage", images=["utility-first", 0]),
})
prompt["end"]["inputs"]["images"] = ["result", 0]
before = copy.deepcopy(prompt)
registered = dict(registry)
prepared, targets = prepare_image_list_execution(prompt)
assert set(prepared) == set(prompt) - {"discarded-preview"}
assert targets == ["end", "save"]
with pytest.raises(ValueError, match="Output node"):
validate_image_list_regions(prepared)
region = validate_image_list_regions(prepared, explicit_targets=True)[0]
worker = build_worker_prompt(prepared, region, "token", "context", 0, "http://master.invalid")
assert set(worker) == set(prepared) - {"source", "save"}
assert worker["utility-second"]["inputs"] == {"images": ["utility-first", 0]}
assert prompt == before
assert registry == registered
assert OutputNode.OUTPUT_NODE is True
assert SchemaOutputNode.OUTPUT_NODE is False
assert SchemaOutputNode.GET_NODE_INFO_V1()["output_node"] is True
@pytest.mark.parametrize("source", ["start", "process"])
@pytest.mark.parametrize("direct", [False, True])
def test_internal_bypass_into_post_end_computation_is_rejected(registry, source, direct):
prompt = graph()
prompt["post-process"] = node("Composite", images=["end", 0])
if not direct:
prompt["bypass"] = node("FrameOperation", image=[source, 0])
source = "bypass"
prompt["post-process"]["inputs"]["reference"] = [source, 0]
prompt["save"]["inputs"]["images"] = ["post-process", 0]
before = copy.deepcopy(prompt)
with pytest.raises(ValueError, match="cross-boundary"):
prepare_image_list_execution(prompt)
assert prompt == before
def test_sequential_namespaced_regions_and_intermediate_output_are_preserved(registry):
prompt = graph()
prompt.update({
"next-start": node(START_CLASS, images=["end", 0]),
"next-process": node("SchemaUtility", images=["next-start", IMAGE_OUT], model=["loader", 0]),
"next-end": node(END_CLASS, images=["next-process", 0], context=["next-start", CONTEXT_OUT]),
"next-save": node("SaveImage", images=["next-end", 0]),
"first-debug": node("PreviewImage", images=["process", 0]),
"second-debug": node("PreviewImage", images=["next-process", 0]),
})
renamed = {key: f"group:7/{key}" for key in prompt}
for spec in prompt.values():
for name, value in spec["inputs"].items():
if isinstance(value, list):
spec["inputs"][name] = [renamed[value[0]], value[1]]
prompt = {renamed[key]: spec for key, spec in prompt.items()}
prepared, targets = prepare_image_list_execution(prompt)
assert set(prepared) == set(prompt) - {renamed["first-debug"], renamed["second-debug"]}
assert targets == [renamed[name] for name in ("end", "save", "next-end", "next-save")]
regions = validate_image_list_regions(prepared, explicit_targets=True)
assert {(r["start_id"], r["end_id"]) for r in regions} == {
(renamed["start"], renamed["end"]), (renamed["next-start"], renamed["next-end"]),
}
second = next(r for r in regions if r["start_id"] == renamed["next-start"])
worker = build_worker_prompt(prepared, second, "token", "context", 0, "http://master.invalid")
assert set(worker) == {renamed[name] for name in ("loader", "next-start", "next-process", "next-end")}
def test_shared_loader_and_its_independent_output_root_are_preserved(registry):
registry["ModelLoader"] = type("OutputLoader", (ModelLoader, OutputNode), {})
prompt = graph()
prompt["source"]["inputs"]["model"] = ["loader", 0]
prompt["independent-save"] = node("SaveImage", model=["loader", 0])
prepared, targets = prepare_image_list_execution(prompt)
assert prepared == prompt
assert targets == ["loader", "end", "save", "independent-save"]
region = validate_image_list_regions(prepared, explicit_targets=True)[0]
assert "loader" in region["node_ids"]
assert "loader" not in region["branch_ids"]
@pytest.mark.parametrize("unsafe_class", [None, type("ListLoader", (), {
"RETURN_TYPES": ("MODEL",), "OUTPUT_IS_LIST": (True,),
})])
def test_unknown_or_list_producing_shared_loader_still_fails_closed(registry, unsafe_class):
if unsafe_class is None:
registry.pop("ModelLoader")
else:
registry["ModelLoader"] = unsafe_class
prompt = graph()
prompt["source"]["inputs"]["model"] = ["loader", 0]
prompt["debug"] = node("PreviewImage", images=["process", 0])
with pytest.raises(ValueError, match="bypasses Start.*loader"):
prepare_image_list_execution(prompt)
def test_upstream_per_item_bypass_is_not_hidden_by_preview_pruning(registry):
prompt = graph()
prompt["process"]["inputs"]["reference"] = ["source", 0]
prompt["debug"] = node("PreviewImage", images=["process", 0])
with pytest.raises(ValueError, match="bypasses Start.*source"):
prepare_image_list_execution(prompt)
@pytest.mark.parametrize("nested", [False, True])
def test_overlapping_or_nested_regions_cannot_be_discarded_as_side_branches(registry, nested):
prompt = graph()
prompt["other-start"] = node(START_CLASS, images=["process" if nested else "source", 0])
prompt["other-process"] = node("FrameOperation", image=["other-start", IMAGE_OUT], reference=["process", 0])
prompt["other-end"] = node(END_CLASS, images=["other-process", 0], context=["other-start", CONTEXT_OUT])
with pytest.raises(ValueError, match="Nested/ambiguous|cross-boundary|ambiguous processing"):
prepare_image_list_execution(prompt)
@pytest.mark.parametrize("invalid", ["missing-end", "wrong-context", "missing-images", "unrelated-images"])
def test_pairing_and_result_validation_cannot_be_pruned_away(registry, invalid):
prompt = graph()
if invalid == "missing-end":
del prompt["end"]
elif invalid == "wrong-context":
prompt["end"]["inputs"]["context"] = ["source", 0]
elif invalid == "missing-images":
del prompt["end"]["inputs"]["images"]
else:
prompt["end"]["inputs"]["images"] = ["source", 0]
before = copy.deepcopy(prompt)
with pytest.raises(ValueError):
prepare_image_list_execution(prompt)
assert prompt == before
def test_execution_copy_is_deep_and_repeated_preparation_is_stable(registry):
prompt = graph()
prompt["process"]["_meta"] = {"title": "Original", "tags": ["saved"]}
before = copy.deepcopy(prompt)
prepared, targets = prepare_image_list_execution(prompt)
prepared_again, targets_again = prepare_image_list_execution(prepared)
assert prepared_again == prepared
assert targets_again == targets
prepared["start"]["inputs"]["images"][0] = "changed"
prepared["process"]["_meta"]["tags"].append("execution-only")
assert prompt == before
assert prepared_again == before
def test_v3_false_output_flag_overrides_legacy_true(registry):
registry["NotOutput"] = type("NotOutput", (), {
"OUTPUT_NODE": True,
"GET_NODE_INFO_V1": classmethod(lambda cls: {"output_node": False}),
})
prompt = graph()
prompt["not-output"] = node("NotOutput", images=["end", 0])
prepared, targets = prepare_image_list_execution(prompt)
assert prepared == prompt
assert targets == ["end", "save"]
@@ -35,6 +35,15 @@ def load_routes():
return module, routes
class CapabilityTests(unittest.IsolatedAsyncioTestCase):
async def test_capability_advertises_target_control_without_tensor_protocol_bump(self):
module, _routes = load_routes()
response = await module.image_list_capabilities_endpoint(None)
payload = json.loads(response.text)
self.assertIs(payload['execution_targets'], True)
self.assertEqual(payload['protocol_version'], 2)
class TransportTests(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
self.module, routes = load_routes()
+77
View File
@@ -11,6 +11,83 @@ from aiohttp import web
from image_list import lifecycle, runtime
class ExecutionTargetProofTests(unittest.TestCase):
def setUp(self):
from test_image_list_master_runtime import graph
self.prompt = graph()
self.registry = SimpleNamespace(NODE_CLASS_MAPPINGS={
'Operation': type('Flagged', (), {'OUTPUT_NODE': True})})
self.modules = patch.dict('sys.modules', {'nodes': self.registry})
self.modules.start()
self.addCleanup(self.modules.stop)
def server(self, roots, *, prompt=None, running=True):
item = (0, 'owned', self.prompt if prompt is None else prompt, {}, roots, {})
return SimpleNamespace(prompt_queue=SimpleNamespace(
get_current_queue_volatile=lambda: ([item], []) if running else ([], [item])))
def test_outer_roots_in_running_queue_are_authoritative(self):
self.assertTrue(lifecycle.has_safe_execution_targets(self.server(['end']), 'owned', self.prompt))
for roots in (['op', 'end'], [], ['end', 'end'], None, [True]):
with self.subTest(roots=roots):
self.assertFalse(lifecycle.has_safe_execution_targets(self.server(roots), 'owned', self.prompt))
self.assertFalse(lifecycle.has_safe_execution_targets(self.server(['end']), 'other', self.prompt))
self.assertFalse(lifecycle.has_safe_execution_targets(self.server(['end'], running=False), 'owned', self.prompt))
self.assertFalse(lifecycle.has_safe_execution_targets(self.server(['end'], prompt={}), 'owned', self.prompt))
def test_prompt_marker_is_not_an_execution_guarantee(self):
self.prompt['start']['inputs']['image_list_explicit_targets'] = True
self.assertFalse(lifecycle.has_safe_execution_targets(self.server(['op', 'end']), 'owned', self.prompt))
with self.assertRaisesRegex(ValueError, 'must be placed after'):
runtime.create_master_job('start', 'end', self.prompt, [], [], '', master_participates=True)
self.assertEqual(runtime._MASTER_JOBS, {})
def test_start_node_uses_running_queue_proof(self):
import importlib.util
import sys
from pathlib import Path
from types import ModuleType
import torch
root = Path(__file__).resolve().parents[1]
package = '_image_list_start_targets'
modules = {}
for suffix in ('', '.nodes', '.image_list', '.utils'):
module = ModuleType(package + suffix)
module.__path__ = [str(root / suffix.lstrip('.'))]
modules[package + suffix] = module
modules[package + '.image_list.runtime'] = runtime
modules[package + '.image_list.lifecycle'] = lifecycle
modules[package + '.utils.async_helpers'] = SimpleNamespace(run_async_in_server_loop=None)
io = SimpleNamespace(ComfyNode=type('ComfyNode', (), {}), NodeOutput=lambda *args: args,
**{name: object() for name in ('Image', 'Mask', 'String', 'Int', 'Float', 'Boolean')})
modules['comfy_api.v0_0_2'] = SimpleNamespace(io=io)
modules['comfy_execution.utils'] = SimpleNamespace(
get_executing_context=lambda: SimpleNamespace(prompt_id='owned'))
server = SimpleNamespace(PromptServer=SimpleNamespace(instance=self.server(['end'])))
modules['server'] = server
spec = importlib.util.spec_from_file_location(package + '.nodes.image_list', root / 'nodes/image_list.py')
node_module = importlib.util.module_from_spec(spec)
with patch.dict(sys.modules, modules):
spec.loader.exec_module(node_module)
start = node_module.DistributedImageListStart
start.hidden = SimpleNamespace(prompt=self.prompt, unique_id='start')
result = start.execute(images=torch.zeros((1, 2, 2, 3)), image_list_master_participates=True)
context = result[3][0]
self.assertIsNotNone(runtime.get_master_job(context))
runtime.remove_master_job(context)
server.PromptServer.instance = self.server(['op', 'end'])
with self.assertRaisesRegex(ValueError, 'must be placed after'):
start.execute(images=torch.zeros((1, 2, 2, 3)), image_list_master_participates=True)
self.assertEqual(runtime._MASTER_JOBS, {})
def test_owned_target_controlled_start_can_create_job(self):
context = runtime.create_master_job('start', 'end', self.prompt, [], [], '',
master_participates=True, explicit_targets=True)
self.addCleanup(runtime.remove_master_job, context)
self.assertIsNotNone(runtime.get_master_job(context))
class ImageListLifecycleTests(unittest.IsolatedAsyncioTestCase):
def setUp(self):
with runtime._LOCK:
@@ -0,0 +1,20 @@
"""Run the real framework's CPU-only private-cache retention regression."""
import os
from pathlib import Path
import subprocess
import sys
import pytest
def test_native_local_cache_retention():
core = os.environ.get('COMFYUI_SOURCE_ROOT')
if not core:
pytest.skip('Set COMFYUI_SOURCE_ROOT for CPU-only real-framework validation')
helper = Path(__file__).parent / 'helpers' / 'local_cache_retention_check.py'
env = dict(os.environ, CUDA_VISIBLE_DEVICES='', HIP_VISIBLE_DEVICES='')
env.pop('PYTHONPATH', None)
result = subprocess.run([sys.executable, str(helper), core], env=env,
text=True, capture_output=True, timeout=90)
assert result.returncode == 0, result.stdout + result.stderr
assert 'PASS: cross-run native cache reuse' in result.stdout
+113
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))
@@ -199,6 +243,73 @@ class MasterRuntimeTests(unittest.IsolatedAsyncioTestCase):
with self.assertRaisesRegex(RuntimeError, 'worker|fallback'):
self.rt.create_master_job('start', 'end', graph(), self.items, [], '')
async def test_values_reconstructed_in_source_order_without_transporting_originals(self):
self.workers.append({'id': 'remote', 'type': 'remote'})
context = self.create(count=6)
tensor = torch.ones(4, 2, 2, 3, dtype=torch.float64)
for index, item in enumerate(self.items):
item['values'].update(value_2=tensor, value_3=True, value_4=index / 8)
self.local_delay = 0.1 # Remote results arrive ahead of master item zero.
with patch.object(self.rt, 'encode_result', side_effect=AssertionError('originals must not return')):
result = await self.rt.run_image_list_job(context, include_values=True)
self.assertFalse(self.rt._MASTER_JOBS)
self.assertEqual(len(result), 4)
for index, row in enumerate(result[3]):
self.assertEqual(row['value_1'], str(index))
self.assertIs(row['value_2'], tensor)
self.assertIs(row['value_3'], True)
self.assertEqual(row['value_4'], index / 8)
self.assertTrue(self.dispatched and self.local_calls)
async def test_overrides_merge_add_and_recheck_capability_before_every_item(self):
self.workers.append({'id': 'remote', 'type': 'remote'})
context = self.create(participates=False, count=3)
self.state.original_prompt['end']['inputs'].update(value_1=['start', 4], value_10=['start', 2])
original_set = self.rt.set_master_result
def override(context, index, token, result):
result['values'] = {'value_1': f'new-{index}', 'value_10': index}
return original_set(context, index, token, result)
with patch.object(self.rt, 'set_master_result', side_effect=override):
result = await self.rt.run_image_list_job(context, include_values=True)
self.assertEqual(result[3], [{'value_1': f'new-{index}', 'value_10': index} for index in range(3)])
self.assertEqual(self.preflight.await_count, 4)
self.assertTrue(all(call.kwargs['require_value_results'] is True
for call in self.preflight.await_args_list))
self.assertEqual(self.items[0]['values'], {'value_1': '0'})
async def test_override_capability_downgrade_stops_before_upload(self):
self.workers.append({'id': 'remote', 'type': 'remote'})
context = self.create(participates=False)
self.state.original_prompt['end']['inputs']['value_1'] = ['start', 4]
self.preflight.side_effect = [{'remote': True}, RuntimeError('value_results unsupported')]
with patch.object(self, 'post', wraps=self.post) as upload:
with self.assertRaisesRegex(RuntimeError, 'value_results'):
await self.rt.run_image_list_job(context, include_values=True)
upload.assert_not_called()
self.assertFalse(self.dispatched or self.rt._MASTER_JOBS)
async def test_missing_expected_override_fails_without_partial_output(self):
self.workers.append({'id': 'remote', 'type': 'remote'})
context = self.create(participates=False)
self.state.original_prompt['end']['inputs']['value_1'] = ['start', 4]
with self.assertRaisesRegex(ValueError, 'missing'):
await self.rt.run_image_list_job(context, include_values=True)
self.assertFalse(self.rt._MASTER_JOBS)
async def test_values_failure_and_cancellation_return_no_partial_lists(self):
self.local_error = RuntimeError('branch failed')
with self.assertRaisesRegex(RuntimeError, 'branch failed'):
await self.rt.run_image_list_job(self.create(), include_values=True)
self.local_error = None
self.local_started.clear()
self.local_delay = 0.1
task = asyncio.create_task(self.rt.run_image_list_job(self.create(), include_values=True))
await self.local_started.wait()
task.cancel()
with self.assertRaises(self.rt.ImageListCancelledError):
await task
self.assertFalse(self.rt._MASTER_JOBS)
async def test_local_failure_surfaces_and_reclaims_staged_token(self):
self.local_error = RuntimeError('local branch failed')
with self.assertRaisesRegex(RuntimeError, 'local branch failed'):
@@ -373,6 +484,8 @@ def real_framework_probe(comfy_root):
rt.remove_master_job(context)
for behavior in ('normal', 'expand', 'error', 'cancel', 'interrupt', 'sync', 'already-interrupted'):
# Independent lifecycle cases start cold; retention is tested separately.
local._RETAINED_OUTPUTS = None
mode[0] = behavior
prompt['op']['class_type'] = 'Expanded' if behavior == 'expand' else 'Operation'
loads.clear()
+196 -8
View File
@@ -1,8 +1,9 @@
"""Image-list participation uses the existing master orchestrator setting."""
import json
import sys
import types
import unittest
from unittest.mock import AsyncMock, patch
from unittest.mock import AsyncMock, Mock, patch
from test_serverless_integration import PKG, load_modules
@@ -34,6 +35,23 @@ class MasterSelectionTests(unittest.IsolatedAsyncioTestCase):
def prepared_inputs(self):
return self.q.queue_prompt_payload.call_args.args[0]["2"]["inputs"]
async def test_prepared_master_prunes_side_outputs_without_mutating_input(self):
import copy
self.prompt['op'] = {'class_type': 'Flagged', 'inputs': {'image': ['2', 0]}}
self.prompt['3']['inputs']['images'] = ['op', 0]
self.prompt['preview'] = {'class_type': 'Flagged', 'inputs': {'image': ['op', 0]}}
self.prompt['save'] = {'class_type': 'Flagged', 'inputs': {'image': ['3', 0]}}
original = copy.deepcopy(self.prompt)
registry = type('Registry', (), {'NODE_CLASS_MAPPINGS': {
'Flagged': type('Flagged', (), {'OUTPUT_NODE': True})}})
with patch.dict(sys.modules, {'nodes': registry}):
await self.queue()
queued = self.q.queue_prompt_payload.call_args
self.assertEqual(queued.kwargs['partial_execution_targets'], ['3', 'save'])
self.assertIn('op', queued.args[0])
self.assertNotIn('preview', queued.args[0])
self.assertEqual(self.prompt, original)
async def test_default_includes_master_and_direct_workers(self):
result = await self.queue()
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
@@ -45,11 +63,21 @@ class MasterSelectionTests(unittest.IsolatedAsyncioTestCase):
self.q.dispatch_worker_prompt.assert_not_awaited()
self.assertEqual(self.ps.distributed_pending_jobs, {})
async def test_end_overrides_require_value_capability_before_master_queue(self):
self.prompt['3']['inputs']['value_10'] = ['2', 2]
await self.queue()
self.assertIs(self.preflight.call_args.kwargs['require_value_results'], True)
self.q.queue_prompt_payload.reset_mock()
self.preflight.side_effect = RuntimeError('value_results unsupported')
with self.assertRaisesRegex(RuntimeError, 'value_results'):
await self.queue()
self.q.queue_prompt_payload.assert_not_awaited()
async def test_orchestrator_setting_excludes_master_processing(self):
self.config["settings"]["master_delegate_only"] = True
await self.queue()
self.assertIs(self.prepared_inputs()["image_list_master_participates"], False)
self.assertTrue(self.preflight.call_args.kwargs["require_idle"])
self.assertIs(self.preflight.call_args.kwargs["require_idle"], False)
# Upstream source stays on master even in orchestrator-only mode.
self.assertIn("1", self.q.queue_prompt_payload.call_args.args[0])
@@ -80,10 +108,154 @@ class MasterSelectionTests(unittest.IsolatedAsyncioTestCase):
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), [])
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
async def test_busy_workers_do_not_prevent_enabled_master(self):
self.preflight.return_value = {}
def use_real_preflight(self, queues):
"""Exercise capability/queue parsing with read-only fake HTTP responses."""
from test_dispatch_selection import dispatch, _FakeSession, _valid_capability
responses = {(worker_id, "capabilities"): (200, _valid_capability())
for worker_id in queues}
responses.update({(worker_id, "queue"): (200, queue)
for worker_id, queue in queues.items()})
session = _FakeSession(responses)
for target, value in (
("get_client_session", AsyncMock(return_value=session)),
("build_worker_url", lambda worker, path: f"http://{worker['id']}{path}"),
):
patched = patch.object(dispatch, target, value)
patched.start()
self.addCleanup(patched.stop)
self.preflight = AsyncMock(wraps=dispatch.preflight_image_list_workers)
sys.modules[PKG + ".api.orchestration.dispatch"].preflight_image_list_workers = self.preflight
return responses
async def run_prepared_selection(self):
"""Run real runtime selection up to the scheduler, without any item I/O."""
rt = sys.modules[PKG + ".image_list.runtime"]
inputs = self.prepared_inputs()
image = object()
context = rt.create_master_job(
"2", "3", self.q.queue_prompt_payload.call_args.args[0],
[{"image": image, "index": 0}],
json.loads(inputs["image_list_enabled_worker_ids"]),
inputs["image_list_master_callback"],
master_participates=inputs["image_list_master_participates"],
explicit_targets=True,
)
self.addCleanup(rt.remove_master_job, context)
self.scheduler = AsyncMock(return_value=[{"image": image}])
self.local_slot = Mock(return_value=types.SimpleNamespace(close=AsyncMock()))
self.runtime_session = AsyncMock(return_value=object())
config = sys.modules[PKG + ".utils.config"]
network = sys.modules[PKG + ".utils.network"]
with patch.object(config, "load_config", lambda: self.config), \
patch.object(config, "get_worker_timeout_seconds", lambda: 60, create=True), \
patch.object(config, "get_image_list_job_timeout_seconds", lambda: 0, create=True), \
patch.object(network, "get_client_session", self.runtime_session, create=True), \
patch.object(network, "build_worker_url", Mock(), create=True), \
patch.object(rt, "dynamic_schedule", self.scheduler), \
patch.dict(sys.modules, {PKG + ".image_list.local_executor": types.SimpleNamespace(
LocalExecutionSlot=self.local_slot)}):
await rt.run_image_list_job(context)
self.assertIsNone(rt.get_master_job(context))
return self.scheduler.call_args.args[1]
async def test_busy_worker_at_enqueue_is_available_to_scheduler_when_idle_at_start(self):
responses = self.use_real_preflight({
"worker": {"queue_running": [[1, "prompt-a"]], "queue_pending": []},
})
result = await self.queue()
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker"])
self.assertEqual(result[2], 1) # Candidates, not a promise of execution-time availability.
self.q.dispatch_worker_prompt.assert_not_awaited()
responses[("worker", "queue")] = (200, {"queue_running": [], "queue_pending": []})
# A newly enabled worker is not implicitly added to this queued selection.
self.config["workers"].append({"id": "later", "enabled": True})
slots = await self.run_prepared_selection()
self.assertEqual([slot for slot in slots if isinstance(slot, str)], ["worker"])
self.assertEqual(len(slots), 2)
self.assertEqual(self.preflight.await_count, 2)
self.assertEqual(self.preflight.await_args_list[0].args, self.preflight.await_args_list[1].args)
async def test_worker_still_busy_at_start_is_excluded_from_scheduler(self):
self.use_real_preflight({
"worker": {"queue_running": [], "queue_pending": [[1, "prompt-a"]]},
})
await self.queue()
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), [])
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker"])
slots = await self.run_prepared_selection()
self.assertEqual(len(slots), 1)
self.assertNotIn("worker", slots)
self.assertEqual(self.preflight.await_count, 2)
self.runtime_session.assert_not_awaited()
async def test_explicit_empty_selection_stays_master_only_at_start(self):
self.use_real_preflight({"worker": {"queue_running": [], "queue_pending": []}})
await self.queue(enabled_worker_ids=[])
slots = await self.run_prepared_selection()
self.assertEqual(len(slots), 1)
self.assertNotIn("worker", slots)
self.preflight.assert_not_awaited()
self.runtime_session.assert_not_awaited()
async def test_busy_incompatible_worker_fails_before_master_queue(self):
responses = self.use_real_preflight({
"worker": {"queue_running": [[1, "prompt-a"]], "queue_pending": []},
})
capability = responses[("worker", "capabilities")][1]
for field, invalid, message in (
("protocol_version", 1, "incompatible"),
("execution_targets", False, "execution_targets"),
("value_results", False, "value_results"),
):
with self.subTest(field=field), patch.dict(capability, {field: invalid}):
if field == "value_results":
self.prompt["3"]["inputs"]["value_10"] = ["2", 2]
with self.assertRaisesRegex(RuntimeError, message):
await self.queue()
self.q.queue_prompt_payload.assert_not_awaited()
self.q.dispatch_worker_prompt.assert_not_awaited()
async def test_orchestrator_with_all_workers_busy_queues_for_idle_start(self):
responses = self.use_real_preflight({
"worker": {"queue_running": [[1, "prompt-a"]], "queue_pending": []},
})
result = await self.queue(delegate_master=True)
self.q.queue_prompt_payload.assert_awaited_once()
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker"])
self.assertIs(self.prepared_inputs()["image_list_master_participates"], False)
self.assertEqual(result[2], 1)
self.q.dispatch_worker_prompt.assert_not_awaited()
responses[("worker", "queue")] = (200, {"queue_running": [], "queue_pending": []})
self.assertEqual(await self.run_prepared_selection(), ["worker"])
self.local_slot.assert_not_called()
self.assertEqual(self.preflight.await_count, 2)
self.assertIs(self.preflight.await_args_list[0].kwargs["require_idle"], False)
self.assertIs(self.preflight.await_args_list[1].kwargs["require_idle"], True)
async def test_orchestrator_rechecks_availability_without_master_fallback(self):
responses = self.use_real_preflight({"worker": {"queue_running": [], "queue_pending": []}})
await self.queue(delegate_master=True)
responses[("worker", "queue")] = (200, {"queue_running": [[1, "other"]], "queue_pending": []})
with self.assertRaisesRegex(RuntimeError, "No selected feature-compatible direct worker is idle"):
await self.run_prepared_selection()
self.assertEqual(self.preflight.await_count, 2)
self.scheduler.assert_not_awaited()
self.local_slot.assert_not_called()
async def test_orchestrator_retains_busy_candidates_when_another_worker_is_idle(self):
self.config["workers"].append({"id": "idle", "enabled": True})
responses = self.use_real_preflight({
"worker": {"queue_running": [[1, "prompt-a"]], "queue_pending": []},
"idle": {"queue_running": [], "queue_pending": []},
})
await self.queue(delegate_master=True)
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker", "idle"])
responses[("worker", "queue")] = (200, {"queue_running": [], "queue_pending": []})
responses[("idle", "queue")] = (200, {"queue_running": [[1, "other"]], "queue_pending": []})
self.assertEqual(await self.run_prepared_selection(), ["worker"])
self.local_slot.assert_not_called()
async def test_orchestrator_without_workers_fails_before_queue(self):
self.config["settings"]["master_delegate_only"] = True
@@ -100,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")
@@ -0,0 +1,536 @@
"""Stable content signatures with fresh, expiring delivery capabilities."""
import asyncio
import copy
import importlib.util
import json
import os
from pathlib import Path
import subprocess
import sys
import types
from unittest.mock import AsyncMock
import pytest
import torch
from image_list import runtime as rt
ROOT = Path(__file__).resolve().parents[1]
def item():
return rt.normalize_image_list_inputs(
torch.arange(18, dtype=torch.float64).reshape(1, 2, 3, 3),
torch.zeros(1, 2, 3, dtype=torch.bfloat16),
{'value_1': 'original', 'value_2': True, 'value_3': 7, 'value_4': 0.25,
'value_10': torch.ones(2, 3, dtype=torch.float16)})[0]
def graph():
return {
'source': {'class_type': 'Source', 'inputs': {}},
'start': {'class_type': rt.START_CLASS, 'inputs': {
'images': ['source', 0], 'image_list_execution_id': 'transient',
'image_list_workflow_meta': {'job': 'metadata'},
'image_list_enabled_worker_ids': '["worker"]',
'image_list_master_callback': 'http://old-master',
'values.value_1': 'inherited', 'worker_context_id': 'old-context'}},
'op': {'class_type': 'Operation', 'inputs': {'image': ['start', 0]}},
'end': {'class_type': rt.END_CLASS, 'inputs': {
'images': ['op', 0], 'context': ['start', 3], 'value_3': ['start', 2]}},
}
def worker_prompt(token, context, data, *, local=True):
prompt = graph()
result = rt.build_worker_prompt(prompt, rt.validate_image_list_regions(prompt)[0],
token, context, data['index'], 'http://master',
item_identity=rt.item_content_identity(data))
if local:
result['start']['inputs']['image_list_local_execution'] = True
result['end']['inputs']['image_list_local_execution'] = True
return result
@pytest.fixture(autouse=True)
def clean_state():
yield
rt._WORKER_ITEMS.clear()
rt._MASTER_JOBS.clear()
@pytest.fixture
def boundaries(monkeypatch):
"""Load only the boundary module; restore every replaced module on exit."""
name = '_native_cache_identity_tests'
for suffix in ('', '.nodes', '.utils', '.image_list', '.api'):
module = types.ModuleType(name + suffix)
module.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
monkeypatch.setitem(sys.modules, name + suffix, module)
monkeypatch.setitem(sys.modules, name + '.image_list.runtime', rt)
io = types.SimpleNamespace(ComfyNode=type('ComfyNode', (), {}),
NodeOutput=lambda *args: args)
for cls in ('Image', 'Mask', 'String', 'Int', 'Float', 'Boolean'):
setattr(io, cls, object())
monkeypatch.setitem(sys.modules, 'comfy_api', types.ModuleType('comfy_api'))
monkeypatch.setitem(sys.modules, 'comfy_api.v0_0_2', types.SimpleNamespace(io=io))
monkeypatch.setitem(sys.modules, name + '.utils.async_helpers',
types.SimpleNamespace(run_async_in_server_loop=lambda coro, **kw: asyncio.run(coro)))
spec = importlib.util.spec_from_file_location(name + '.nodes.image_list', ROOT / 'nodes/image_list.py')
module = importlib.util.module_from_spec(spec)
monkeypatch.setitem(sys.modules, spec.name, module)
spec.loader.exec_module(module)
return module
def start(boundaries, prompt):
boundaries.DistributedImageListStart.hidden = types.SimpleNamespace(prompt=prompt, unique_id='start')
return boundaries.DistributedImageListStart.execute(**prompt['start']['inputs'])
def end(boundaries, prompt, cached):
boundaries.DistributedImageListEnd.hidden = types.SimpleNamespace(prompt=prompt, unique_id='end')
inputs = {k: v for k, v in prompt['end']['inputs'].items()
if k not in ('images', 'context', 'value_3')}
return boundaries.DistributedImageListEnd.execute(
images=cached[0], context=cached[3], value_3=cached[2], **inputs)
def stage(data):
token = rt.create_worker_task_token()
rt.store_worker_item(token, rt.encode_item(data))
return token
def master(data, token):
context = rt.create_master_job('start', 'end', graph(), [data], ['worker'], '')
rt.register_assignment(context, data['index'], 'worker', token)
return context
def test_identity_is_exact_typed_order_independent_and_not_transport_bounded(monkeypatch):
original = item()
identity = rt.item_content_identity(original)
assert rt.is_valid_item_identity(identity)
reordered = copy.deepcopy(original)
reordered['values'] = dict(reversed(list(reordered['values'].items())))
assert rt.item_content_identity(reordered) == identity
strided = copy.deepcopy(original)
strided['image'] = original['image'].transpose(1, 2).contiguous().transpose(1, 2)
assert not strided['image'].is_contiguous()
assert rt.item_content_identity(strided) == identity
monkeypatch.setattr(rt, 'MAX_PAYLOAD_BYTES', 1)
with pytest.raises(ValueError, match='transport limit'):
rt.encode_item(original)
assert rt.item_content_identity(original) == identity
scalar_identities = set()
for value in (True, 1, 1.0, '1', 0.0, -0.0):
changed = copy.deepcopy(original)
changed['values']['value_2'] = value
scalar_identities.add(rt.item_content_identity(changed))
assert len(scalar_identities) == 6
@pytest.mark.parametrize('change', ['image', 'mask', 'value', 'index', 'dtype', 'shape', 'slot', 'missing-mask'])
def test_every_item_component_invalidates_identity(change):
original, changed = item(), item()
if change in ('image', 'mask'):
changed[change].reshape(-1)[0] = 1
elif change == 'value':
changed['values']['value_10'][0, 0] = 2
elif change == 'index':
changed['index'] = 1
elif change == 'dtype':
changed['image'] = changed['image'].float()
elif change == 'shape':
changed['values']['value_10'] = changed['values']['value_10'].reshape(3, 2)
elif change == 'slot':
changed['values']['value_9'] = changed['values'].pop('value_10')
else:
changed['mask'] = None
assert rt.item_content_identity(original) != rt.item_content_identity(changed)
def test_float64_low_bits_and_bfloat16_bytes_are_not_quantized():
data, changed = item(), item()
changed['image'].reshape(-1)[0] = torch.nextafter(torch.tensor(0., dtype=torch.float64),
torch.tensor(1., dtype=torch.float64))
assert rt.item_content_identity(data) != rt.item_content_identity(changed)
changed = item()
changed['mask'].view(torch.int16).reshape(-1)[0] = 1
assert rt.item_content_identity(data) != rt.item_content_identity(changed)
def test_stable_start_replaced_wholesale_and_legacy_graph_preserved():
prompt = graph()
frozen = copy.deepcopy(prompt)
region = rt.validate_image_list_regions(prompt)[0]
legacy = rt.build_worker_prompt(prompt, region, 'a', 'ctx-a', 0, 'http://a')
assert legacy['start']['inputs']['worker_item_token'] == 'a'
assert legacy['start']['inputs']['image_list_execution_id'] == 'transient'
identity = rt.item_content_identity(item())
a = rt.build_worker_prompt(prompt, region, 'a', 'ctx-a', 0, 'http://a', item_identity=identity)
b = rt.build_worker_prompt(prompt, region, 'b', 'ctx-b', 0, 'http://b', item_identity=identity)
assert a['start']['inputs'] == b['start']['inputs'] == {'worker_item_identity': identity}
assert a['end']['inputs']['worker_context_id'] == 'ctx-a'
assert b['end']['inputs']['worker_result_token'] == 'b'
assert b['end']['inputs']['worker_item_identity'] == identity
assert b['end']['inputs']['value_3'] == ['start', 2]
assert prompt == frozen
@pytest.mark.parametrize('bad', ['', 'sha256:' + 'z' * 64, 'a' * 64, False, ['bad']])
def test_malformed_identity_rejected(bad):
assert not rt.is_valid_item_identity(bad)
prompt = graph()
with pytest.raises((ValueError, TypeError), match='identity'):
rt.build_worker_prompt(prompt, rt.validate_image_list_regions(prompt)[0],
'token', 'context', 0, '', item_identity=bad)
def test_cache_miss_then_hit_delivers_current_context_token_and_values(boundaries):
data = item()
a = stage(data)
context_a = master(data, a)
prompt_a = worker_prompt(a, context_a, data)
cached = start(boundaries, prompt_a)
assert cached[3] == [rt.item_content_identity(data)]
assert 'image' not in rt._WORKER_ITEMS[a][1] # only a small, expiring receipt remains
end(boundaries, prompt_a, cached)
assert a not in rt._WORKER_ITEMS
assert rt.get_master_job(context_a).results[0]['values'] == {'value_3': 0}
b = stage(data)
context_b = master(data, b)
prompt_b = worker_prompt(b, context_b, data)
assert prompt_a['start'] == prompt_b['start']
# A native hit does NOT execute Start; End must consume this current upload.
end(boundaries, prompt_b, cached)
assert b not in rt._WORKER_ITEMS
state = rt.get_master_job(context_b)
assert state.expected_tokens[0] == b
assert torch.equal(state.results[0]['image'], data['image'])
assert rt._collect_output_lists([state.results[0]], state=state, include_values=True)[3][0]['value_3'] == 0
with pytest.raises(RuntimeError, match='missing|expired|consumed'):
end(boundaries, prompt_b, cached)
@pytest.mark.parametrize('miss', [False, True])
def test_expired_item_cannot_deliver_even_with_cached_start(boundaries, miss):
data = item()
token = stage(data)
context = master(data, token)
prompt = worker_prompt(token, context, data)
cached = start(boundaries, prompt) if miss else ([data['image']], [data['mask']], [0], [rt.item_content_identity(data)])
_, staged = rt._WORKER_ITEMS[token]
rt._WORKER_ITEMS[token] = (0, staged)
with pytest.raises(RuntimeError, match='missing|expired|consumed'):
end(boundaries, prompt, cached)
assert not rt.get_master_job(context).received_indexes
assert token not in rt._WORKER_ITEMS
def test_identity_mismatch_on_start_and_on_cache_hit_fails_closed(boundaries):
data, different = item(), item()
different['image'][0, 0, 0, 0] = 3
for hit in (False, True):
token = stage(different)
context = master(different, token)
prompt = worker_prompt(token, context, data)
cached = ([data['image']], [data['mask']], [0], [rt.item_content_identity(data)])
with pytest.raises((ValueError, RuntimeError), match='identity'):
end(boundaries, prompt, cached) if hit else start(boundaries, prompt)
assert not rt.get_master_job(context).received_indexes
rt.delete_worker_item(token)
@pytest.mark.parametrize('change', ['ambiguous', 'wrong-link', 'wrong-start-id', 'end-identity', 'end-index'])
def test_start_requires_exact_paired_end_binding(boundaries, change):
data = item()
token = stage(data)
prompt = worker_prompt(token, 'context', data)
if change == 'ambiguous':
prompt['second-end'] = copy.deepcopy(prompt['end'])
elif change == 'wrong-link':
prompt['end']['inputs']['context'] = ['start', 2]
elif change == 'wrong-start-id':
prompt['renamed'] = prompt.pop('start')
elif change == 'end-identity':
prompt['end']['inputs']['worker_item_identity'] = 'sha256:' + '0' * 64
else:
prompt['end']['inputs']['worker_item_index'] = 1
boundaries.DistributedImageListStart.hidden = types.SimpleNamespace(prompt=prompt, unique_id='start')
with pytest.raises((RuntimeError, ValueError), match='paired|identity|index'):
boundaries.DistributedImageListStart.execute(worker_item_identity=rt.item_content_identity(data),
image_list_local_execution=True)
def test_end_checks_stable_context_before_delivering(boundaries):
data = item()
token = stage(data)
context = master(data, token)
prompt = worker_prompt(token, context, data)
cached = list(start(boundaries, prompt))
cached[3] = ['sha256:' + '0' * 64]
with pytest.raises((RuntimeError, ValueError), match='identity|context'):
end(boundaries, prompt, cached)
assert not rt.get_master_job(context).received_indexes
def test_legacy_worker_path_still_consumes_and_delivers_job_context(boundaries):
data = item()
token = stage(data)
context = master(data, token)
prompt = graph()
prompt = rt.build_worker_prompt(prompt, rt.validate_image_list_regions(prompt)[0], token, context, 0, '')
prompt['start']['inputs']['image_list_local_execution'] = True
prompt['end']['inputs']['image_list_local_execution'] = True
cached = start(boundaries, prompt)
assert cached[3] == [context]
assert token not in rt._WORKER_ITEMS
end(boundaries, prompt, cached)
assert rt.get_master_job(context).received_indexes == {0}
@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)])
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
def test_item_route_returns_server_computed_identity(boundaries, monkeypatch):
from aiohttp import web
name = '_native_cache_identity_tests'
server = types.SimpleNamespace(PromptServer=types.SimpleNamespace(instance=types.SimpleNamespace(routes=web.RouteTableDef())))
monkeypatch.setitem(sys.modules, 'server', server)
spec = importlib.util.spec_from_file_location(name + '.api.image_list_routes', ROOT / 'api/image_list_routes.py')
module = importlib.util.module_from_spec(spec)
monkeypatch.setitem(sys.modules, spec.name, module)
spec.loader.exec_module(module)
data = item()
token = rt.create_worker_task_token()
monkeypatch.setattr(module, '_json_body', AsyncMock(return_value={
'token': token, 'context_id': 'context', 'item': rt.encode_item(data),
'item_identity': 'sha256:' + '0' * 64}))
response = asyncio.run(module.image_list_item_endpoint(object()))
assert response.status == 200
assert json.loads(response.body)['item_identity'] == rt.item_content_identity(data)
assert len(rt._WORKER_ITEMS[token]) == 2
@pytest.mark.parametrize('responses', [
('valid', 'valid'), ('valid', 'absent'), ('invalid', 'valid'),
('bad-json', 'valid'), ('non-object', 'valid'),
])
def test_runtime_opts_in_per_upload_and_can_downgrade(monkeypatch, responses):
name = '_native_cache_identity_runtime'
for suffix in ('', '.image_list', '.api', '.api.orchestration', '.utils'):
module = types.ModuleType(name + suffix)
module.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
monkeypatch.setitem(sys.modules, name + suffix, module)
spec = importlib.util.spec_from_file_location(name + '.image_list.runtime', ROOT / 'image_list/runtime.py')
runtime = importlib.util.module_from_spec(spec)
monkeypatch.setitem(sys.modules, spec.name, runtime)
spec.loader.exec_module(runtime)
workers = [{'id': 'worker', 'type': 'local'}]
dispatched, uploads = [], []
preflight = AsyncMock(return_value={'worker': True})
sync_media = AsyncMock()
class Response:
status = 200
def __init__(self, data):
self.data = data
async def __aenter__(self):
return self
async def __aexit__(self, *args):
pass
async def json(self):
if self.data == 'bad-json':
raise ValueError('old response is not JSON')
return self.data
def raise_for_status(self):
pass
class Session:
def post(self, url, *, data, **kw):
assert url.endswith('/distributed/image-list/item')
body = json.loads(data)
identity = runtime.store_worker_item(body['token'], body['item'])
mode = responses[len(uploads)]
uploads.append(body)
response = {'valid': {'item_identity': identity}, 'absent': {'status': 'stored'},
'invalid': {'item_identity': 'bad'}, 'bad-json': 'bad-json', 'non-object': []}[mode]
return Response(response)
def get(self, url, **kw):
pid = url.rsplit('/', 1)[-1]
return Response({pid: {'status': {'completed': True, 'status_str': 'success'}}})
def delete(self, url, **kw):
runtime.delete_worker_item(url.rsplit('/', 1)[-1])
return Response({})
async def dispatch(worker, prompt, meta, **kwargs):
dispatched.append(copy.deepcopy(prompt))
delivery = prompt['end']['inputs']
token, index = delivery['worker_result_token'], delivery['worker_item_index']
identity = prompt['start']['inputs'].get('worker_item_identity')
options = {'item_identity': identity, 'item_index': index} if identity else {}
data = runtime.consume_worker_item(token, **options)
if identity:
runtime.finish_worker_item(token, identity, index)
runtime.set_master_result(context, index, token, {
'image': data['image'], 'text': str(index), 'values': {'value_3': index}})
assert kwargs['partial_execution_targets'] == ['end']
return f'worker-{index}'
dependencies = {
'.api.orchestration.dispatch': types.SimpleNamespace(
preflight_image_list_workers=preflight, dispatch_worker_prompt=dispatch),
'.api.orchestration.prompt_transform': types.SimpleNamespace(
_resolve_direct_workers_for_image_list=lambda *args: workers),
'.api.orchestration.media_sync': types.SimpleNamespace(sync_worker_media=sync_media),
'.utils.config': types.SimpleNamespace(load_config=lambda: {'settings': {}},
get_worker_timeout_seconds=lambda: 60),
'.utils.network': types.SimpleNamespace(build_master_callback_url=lambda *a, **kw: 'http://master',
build_worker_url=lambda worker, path: 'http://worker' + path,
get_client_session=AsyncMock(return_value=Session())),
'.utils.logging': types.SimpleNamespace(log=lambda *args: None),
}
for suffix, module in dependencies.items():
monkeypatch.setitem(sys.modules, name + suffix, module)
monkeypatch.setitem(sys.modules, 'aiohttp', types.SimpleNamespace(
ClientTimeout=lambda **kwargs: kwargs, ClientError=ConnectionError))
comfy = types.ModuleType('comfy')
comfy.model_management = types.SimpleNamespace(processing_interrupted=lambda: False)
monkeypatch.setitem(sys.modules, 'comfy', comfy)
monkeypatch.setitem(sys.modules, 'comfy.model_management', comfy.model_management)
items = [item(), item()]
items[1]['index'] = 1
context = runtime.create_master_job('start', 'end', graph(), items, ['worker'], '')
state = runtime.get_master_job(context)
try:
result = asyncio.run(runtime._run_image_list_job(state, include_values=True))
assert result[2] == ['0', '1']
assert [row['value_3'] for row in result[3]] == [0, 1]
assert len(uploads) == len(dispatched) == 2
assert sync_media.await_count == 1
for index, prompt in enumerate(dispatched):
start_inputs = prompt['start']['inputs']
if responses[index] == 'valid':
assert start_inputs == {'worker_item_identity': runtime.item_content_identity(items[index])}
assert prompt['end']['inputs']['worker_context_id'] == context
else:
assert 'worker_item_identity' not in start_inputs
assert start_inputs['worker_item_token'] == uploads[index]['token']
assert start_inputs['worker_context_id'] == context
assert 'worker_item_identity' not in prompt['end']['inputs']
assert prompt['end']['inputs']['worker_result_token'] == uploads[index]['token']
assert not runtime._WORKER_ITEMS
finally:
runtime.remove_master_job(context)
runtime._WORKER_ITEMS.clear()
def test_runtime_local_path_hashes_without_transport(monkeypatch):
name = '_native_cache_identity_local_runtime'
for suffix in ('', '.image_list', '.api', '.api.orchestration', '.utils'):
module = types.ModuleType(name + suffix)
module.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
monkeypatch.setitem(sys.modules, name + suffix, module)
spec = importlib.util.spec_from_file_location(name + '.image_list.runtime', ROOT / 'image_list/runtime.py')
runtime = importlib.util.module_from_spec(spec)
monkeypatch.setitem(sys.modules, spec.name, runtime)
spec.loader.exec_module(runtime)
dispatched = []
class Slot:
def __init__(self, **kw):
pass
def submit(self, prompt, prompt_id, end_id, workflow_meta):
dispatched.append(prompt)
delivery = prompt[end_id]['inputs']
identity = prompt['start']['inputs']['worker_item_identity']
data = runtime.consume_worker_item(delivery['worker_result_token'], item_identity=identity,
item_index=delivery['worker_item_index'])
runtime.finish_worker_item(delivery['worker_result_token'], identity, data['index'])
runtime.set_master_result(delivery['worker_context_id'], data['index'], delivery['worker_result_token'],
{'image': data['image'], 'values': {'value_3': data['index']}})
future = asyncio.get_running_loop().create_future()
future.set_result(None)
return future
def request_stop(self):
pass
async def close(self):
pass
def forbidden(*args, **kwargs):
raise AssertionError('local path attempted HTTP/transport')
dependencies = {
'.image_list.local_executor': types.SimpleNamespace(LocalExecutionSlot=Slot),
'.api.orchestration.dispatch': types.SimpleNamespace(
preflight_image_list_workers=forbidden, dispatch_worker_prompt=forbidden),
'.api.orchestration.prompt_transform': types.SimpleNamespace(
_resolve_direct_workers_for_image_list=lambda *args: []),
'.utils.config': types.SimpleNamespace(load_config=lambda: {'settings': {}},
get_worker_timeout_seconds=lambda: 60),
'.utils.network': types.SimpleNamespace(build_master_callback_url=forbidden,
build_worker_url=forbidden, get_client_session=forbidden),
'.utils.logging': types.SimpleNamespace(log=lambda *args: None),
}
for suffix, module in dependencies.items():
monkeypatch.setitem(sys.modules, name + suffix, module)
monkeypatch.setitem(sys.modules, 'aiohttp', types.SimpleNamespace(
ClientTimeout=lambda **kwargs: kwargs, ClientError=ConnectionError))
comfy = types.ModuleType('comfy')
comfy.model_management = types.SimpleNamespace(processing_interrupted=lambda: False)
monkeypatch.setitem(sys.modules, 'comfy', comfy)
monkeypatch.setitem(sys.modules, 'comfy.model_management', comfy.model_management)
monkeypatch.setattr(runtime, 'MAX_PAYLOAD_BYTES', 1)
monkeypatch.setattr(runtime, 'encode_item', forbidden)
data = item()
context = runtime.create_master_job('start', 'end', graph(), [data], [], '', master_participates=True)
try:
result = asyncio.run(runtime._run_image_list_job(runtime.get_master_job(context), include_values=True))
assert dispatched[0]['start']['inputs'] == {'worker_item_identity': runtime.item_content_identity(data),
'image_list_local_execution': True}
assert result[3][0]['value_3'] == 0
assert not runtime._WORKER_ITEMS
finally:
runtime.remove_master_job(context)
runtime._WORKER_ITEMS.clear()
def test_real_native_signatures_and_cached_delivery():
core = Path(os.environ.get('COMFYUI_SOURCE_ROOT', str(Path(sys.executable).parents[2])))
if not (core / 'comfy_execution/caching.py').exists():
pytest.skip('real ComfyUI source not available beside this interpreter')
result = subprocess.run([sys.executable, str(ROOT / 'tests/helpers/image_list_native_cache_identity.py'), str(core)],
text=True, capture_output=True, timeout=90)
assert result.returncode == 0, result.stdout + result.stderr
assert 'NATIVE_CACHE_IDENTITY_OK' in result.stdout
+19 -1
View File
@@ -84,7 +84,7 @@ class RuntimeBoundaryTests(unittest.IsolatedAsyncioTestCase):
# Other legacy test modules replace aiohttp during collection. This
# controlled provider fixture needs only its explicit timeout/error seam.
sys.modules['aiohttp'] = types.SimpleNamespace(
ClientTimeout=lambda **kwargs: kwargs, ClientError=ConnectionError)
ClientTimeout=lambda **kwargs: kwargs, ClientError=ConnectionError, ClientConnectionError=ConnectionError)
comfy = types.ModuleType('comfy')
comfy.model_management = types.SimpleNamespace(processing_interrupted=lambda: self.interrupted)
sys.modules['comfy'] = comfy
@@ -94,6 +94,7 @@ class RuntimeBoundaryTests(unittest.IsolatedAsyncioTestCase):
prompt['7']['inputs']['image'] = 'shared/ref.png'
async def dispatch_prompt(self, worker, prompt, meta, **kwargs):
self.assertEqual(kwargs.get('partial_execution_targets'), ['5'])
if self.reject:
raise RuntimeError('prompt rejected')
self.prompts.append((json.loads(json.dumps(prompt)), meta))
@@ -136,6 +137,23 @@ class RuntimeBoundaryTests(unittest.IsolatedAsyncioTestCase):
['worker'], 'http://public-master', self.metadata)
return self.context
async def test_runtime_preflight_rejects_old_target_capability_before_upload(self):
from test_dispatch_selection import _load_dispatch_module, _FakeSession, _valid_capability
real_dispatch = _load_dispatch_module()
capability = _valid_capability()
capability.pop('execution_targets')
session = _FakeSession({('worker', 'capabilities'): (200, capability)})
module = sys.modules['_image_list_boundaries.api.orchestration.dispatch']
module.preflight_image_list_workers = real_dispatch.preflight_image_list_workers
with patch.object(real_dispatch, 'get_client_session', AsyncMock(return_value=session)), patch.object(
real_dispatch, 'build_worker_url', side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
with self.assertRaisesRegex(RuntimeError, 'execution_targets.*upgrade'):
await self.rt.run_image_list_job(self.create(1))
self.assertFalse(self.uploads)
self.dispatch.assert_not_awaited()
self.assertFalse(self.rt._MASTER_JOBS)
async def test_each_assignment_reuses_media_paths_with_fresh_tokens_and_metadata(self):
context = self.create()
await self.rt.run_image_list_job(context)
+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())
+165
View File
@@ -0,0 +1,165 @@
"""Context-carried values and sparse, per-item End overrides."""
import copy
from unittest.mock import patch
import pytest
import torch
from image_list import runtime as rt
from image_list.execution_plan import prepare_image_list_execution
def value_graph():
# Use a minimal, closed region with an override-only ancestor.
return {
'source': {'class_type': 'Source', 'inputs': {}},
'start': {'class_type': rt.START_CLASS, 'inputs': {'images': ['source', 0]}},
'override': {'class_type': 'Transform', 'inputs': {'value': ['start', 4]}},
'end': {'class_type': rt.END_CLASS, 'inputs': {
'images': ['start', 0], 'context': ['start', 3], 'value_1': ['override', 0],
'value_10': ['start', 5]}},
}
def values():
return {'value_1': torch.arange(72, dtype=torch.float64).reshape(2, 3, 4, 3),
'value_2': torch.ones(3, 2, 2, dtype=torch.bfloat16), 'value_3': 'name',
'value_4': 42, 'value_5': 0.125, 'value_6': True}
def assert_values(actual, expected):
assert actual.keys() == expected.keys()
for key, value in expected.items():
assert type(actual[key]) is type(value)
if isinstance(value, torch.Tensor):
assert actual[key].dtype == value.dtype
assert torch.equal(actual[key], value)
else:
assert actual[key] == value
def test_normalize_and_roundtrip_all_override_kinds_keeps_batches_whole():
expected = values()
result = rt.normalize_worker_result(torch.zeros(1, 2, 2, 3),
values={key: [value] for key, value in expected.items()})
assert_values(result['values'], expected)
wire = rt.encode_result(result)
assert_values(rt.decode_result(wire)['values'], expected)
assert wire['values']['value_1']['shape'][0] == 2
assert wire['values']['value_2']['shape'][0] == 3
@pytest.mark.parametrize('bad', [None, [], [None], [1, 2], [[1]], object(),
float('nan'), float('inf'), torch.tensor([float('nan')]),
torch.tensor([float('inf')]), torch.ones(1, dtype=torch.int32)])
def test_invalid_override_rejected(bad):
with pytest.raises((ValueError, TypeError)):
rt.normalize_worker_result(torch.zeros(1, 2, 2, 3), values={'value_1': bad})
def test_unknown_override_names_and_wire_types_rejected():
for name in ('value_0', 'value_11', 'values.value_1'):
with pytest.raises(ValueError, match='slot'):
rt.normalize_worker_result(torch.zeros(1, 2, 2, 3), values={name: 1})
base = rt.encode_result({'image': torch.zeros(1, 2, 2, 3)})
for payload in (None, [], {'value_0': {'type': 'INT', 'value': 1}},
{'value_1': {'type': 'LATENT', 'value': {}}},
{'value_1': {'type': 'FLOAT', 'value': float('nan')}},
{'value_1': {'type': 'STRING', 'value': None}},
{'value_1': rt.encode_value(torch.tensor([float('inf')]), 'TENSOR')}):
with pytest.raises((ValueError, TypeError)):
rt.decode_result({**base, 'values': payload})
def test_no_overrides_keeps_legacy_wire_and_does_not_encode_original_tensors():
image = torch.zeros(1, 2, 2, 3)
result = rt.normalize_worker_result(image)
assert set(result) == {'image', 'mask', 'text'}
with patch.object(rt, 'encode_value', wraps=rt.encode_value) as encode:
wire = rt.encode_result(result)
assert set(wire) == {'image', 'mask', 'text'}
assert [call.args[1] for call in encode.call_args_list] == ['IMAGE', 'MASK', 'STRING']
assert set(rt.decode_result(wire)) == {'image', 'mask', 'text'}
def test_override_ancestors_survive_plan_and_worker_prompt():
prompt = value_graph()
frozen = copy.deepcopy(prompt)
prepared, targets = prepare_image_list_execution(prompt)
assert prepared == frozen and prompt == frozen
assert targets == ['end']
region = rt.validate_image_list_regions(prepared)[0]
assert 'override' in region['branch_ids']
worker = rt.build_worker_prompt(prepared, region, 'token', 'context', 0, 'http://master')
assert worker['end']['inputs']['value_1'] == ['override', 0]
assert worker['end']['inputs']['value_10'] == ['start', 5]
assert 'source' not in worker
for bad in ('value_0', 'value_11', 'values.value_1'):
invalid = copy.deepcopy(prompt)
invalid['end']['inputs'][bad] = ['start', 4]
with pytest.raises(ValueError, match='slot'):
rt.validate_image_list_regions(invalid)
def test_missing_or_unexpected_result_override_fails_before_recording():
prompt = value_graph()
items = rt.normalize_image_list_inputs(torch.zeros(1, 2, 2, 3), values={'value_1': 'original'})
context = rt.create_master_job('start', 'end', prompt, items, ['worker'], '')
try:
rt.register_assignment(context, 0, 'worker', 'token')
for result_values in ({}, {'value_1': 'updated'},
{'value_1': 'updated', 'value_10': 10, 'value_9': False}):
with pytest.raises(ValueError, match='override'):
rt.set_master_result(context, 0, 'token', {'image': items[0]['image'], 'values': result_values})
assert rt.get_master_job(context).received_indexes == set()
rt.set_master_result(context, 0, 'token', {'image': items[0]['image'],
'values': {'value_1': 'updated', 'value_10': 10}})
assert rt.get_master_job(context).received_indexes == {0}
finally:
rt.remove_master_job(context)
def test_override_return_request_budget_is_aggregate():
result = rt.normalize_worker_result(torch.zeros(1, 2, 2, 3), values={'value_1': ['x' * 500]})
with patch.object(rt, 'MAX_REQUEST_BYTES', 400):
with pytest.raises(ValueError, match='request'):
rt.encode_request_payload({'result': rt.encode_result(result)})
class TestResultHTTP:
def test_missing_and_invalid_overrides_fail_through_real_http(self):
import asyncio
from aiohttp import web
from aiohttp.test_utils import TestClient, TestServer
from test_image_list_integration_gaps import load_routes
async def check():
_module, routes = load_routes()
app = web.Application()
app.add_routes(routes)
prompt = value_graph()
items = rt.normalize_image_list_inputs(torch.zeros(1, 2, 2, 3), values={'value_1': 'original'})
context = rt.create_master_job('start', 'end', prompt, items, ['worker'], '')
token = rt.create_worker_task_token()
rt.register_assignment(context, 0, 'worker', token)
try:
async with TestClient(TestServer(app)) as client:
caps = await (await client.get('/distributed/image-list/capabilities')).json()
assert caps['value_results'] is True and caps['protocol_version'] == 2
result = rt.encode_result({'image': items[0]['image']})
envelope = {'context_id': context, 'token': token, 'item_index': 0, 'result': result}
response = await client.post('/distributed/image-list/result', json=envelope)
assert response.status == 400
assert 'missing' in (await response.json())['error']
assert not rt.get_master_job(context).received_indexes
result['values'] = {'value_1': {'type': 'INT', 'value': [1, 2]},
'value_10': {'type': 'INT', 'value': 10}}
response = await client.post('/distributed/image-list/result', json=envelope)
assert response.status == 400
assert not rt.get_master_job(context).received_indexes
result['values']['value_1'] = {'type': 'STRING', 'value': 'new'}
response = await client.post('/distributed/image-list/result', json=envelope)
assert response.status == 200
assert rt.get_master_job(context).results[0]['values'] == {'value_1': 'new', 'value_10': 10}
finally:
rt.remove_master_job(context)
asyncio.run(check())
+198
View File
@@ -0,0 +1,198 @@
"""Real ComfyUI V3 bridge mapping, file ownership and exact CPU round trips."""
import os
from pathlib import Path
import subprocess
import sys
import pytest
ROOT = Path(__file__).resolve().parents[1]
def test_real_media_bridge_nodes():
core = os.environ.get("COMFYUI_SOURCE_ROOT")
if core is None:
pytest.skip("Set COMFYUI_SOURCE_ROOT for CPU media bridge acceptance")
assert core and (Path(core) / "comfy_api/latest/_input_impl/video_types.py").is_file(), (
"Configured COMFYUI_SOURCE_ROOT must provide native video classes"
)
probe = r'''
import asyncio, importlib.util, io, json, os, sys, tempfile
from fractions import Fraction
import av
import numpy as np
from pathlib import Path
from unittest.mock import patch
import torch
sys.path.insert(0, sys.argv[1])
import comfy.cli_args
comfy.cli_args.args.cpu = True
comfy.cli_args.args.disable_assets = True
def forbidden(*args, **kwargs):
raise AssertionError('Bridge tests must not query or initialize a GPU')
torch.cuda.is_available = lambda: False
torch.cuda.device_count = lambda: 0
for name in ('init', '_lazy_init', 'current_device'):
setattr(torch.cuda, name, forbidden)
import execution
import folder_paths
from comfy_api.latest._input_impl.video_types import VideoFromComponents, VideoFromFile, VideoFromList
from comfy_api.latest._util import VideoComponents
for name in ('is_available', 'device_count'):
setattr(torch.cuda, name, forbidden)
spec = importlib.util.spec_from_file_location('typed_bridge_probe', Path(sys.argv[2])/'serverless_bridge/__init__.py', submodule_search_locations=[str(Path(sys.argv[2])/'serverless_bridge')])
bridge = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = bridge
spec.loader.exec_module(bridge)
from typed_bridge_probe.nodes import MediaImportV1, MediaExportV1, export_values, import_value
from typed_bridge_probe.contract import SOCKET_KINDS
from typed_bridge_probe.video_codec import encode_video
def tiny_video():
buffer = io.BytesIO()
with av.open(buffer, 'w', format='matroska') as container:
stream = container.add_stream('ffv1', rate=Fraction(24, 1))
stream.width = stream.height = 32
stream.pix_fmt = 'yuv420p'
for level in (16, 64, 128):
frame = av.VideoFrame.from_ndarray(np.full((32, 32, 3), level, dtype=np.uint8), format='rgb24')
for packet in stream.encode(frame): container.mux(packet)
for packet in stream.encode(): container.mux(packet)
return buffer.getvalue()
def exact_tensor(actual, expected):
assert actual.device.type == 'cpu' and actual.dtype == expected.dtype
assert actual.shape == expected.shape
assert torch.equal(actual.contiguous().view(torch.uint8), expected.contiguous().view(torch.uint8))
def exact_video(actual, expected, wire):
assert type(actual) is type(expected)
assert encode_video(actual) == wire
if type(actual) is VideoFromComponents:
a = vars(actual)['_VideoFromComponents__components']
e = vars(expected)['_VideoFromComponents__components']
assert type(a) is VideoComponents and type(a.frame_rate) is Fraction
assert a.frame_rate == e.frame_rate and a.metadata == e.metadata
assert a.metadata is not e.metadata
for name in ('images', 'alpha'): exact_tensor(getattr(a, name), getattr(e, name))
exact_tensor(a.audio['waveform'], e.audio['waveform'])
assert a.audio['sample_rate'] == e.audio['sample_rate']
assert actual.get_bit_depth() == 10 and actual.get_color_space() == 'HDR PQ'
else:
values = actual.videos if type(actual) is VideoFromList else [actual]
for value in values:
buffer = vars(value)['_VideoFromFile__file']
assert type(buffer) is io.BytesIO and not buffer.closed
buffer.seek(0)
with av.open(buffer) as container:
assert len(list(container.decode(video=0))) == 3
if type(actual) is VideoFromList:
exact_tensor(actual.complete_audio['waveform'], expected.complete_audio['waveform'])
async def check():
extension = await bridge.comfy_entrypoint()
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 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)
roots = {key: base/key for key in ('input','output','temp')}
for path in roots.values(): path.mkdir()
def annotated(name):
for key, root in roots.items():
suffix = f' [{key}]'
if name.endswith(suffix): return str(root/name[:-len(suffix)])
return str(roots['input']/name)
with patch.object(folder_paths,'get_input_directory',lambda: str(roots['input'])), patch.object(folder_paths,'get_output_directory',lambda: str(roots['output'])), patch.object(folder_paths,'get_temp_directory',lambda: str(roots['temp'])), patch.object(folder_paths,'get_annotated_filepath',annotated):
examples = {'images': torch.rand(2,2,3,4,dtype=torch.float64),
'masks': torch.tensor([[[0.123456789,0.987654321]]],dtype=torch.float64),
'audio': {'waveform':torch.randn(1,6,31,dtype=torch.float64),'sample_rate':44100},
'texts': ' café\n🎥\r\n'}
media = tiny_video()
sources = [roots['input']/f'source-{n}.mkv' for n in range(3)]
for path in sources: path.write_bytes(media)
images = torch.arange(36, dtype=torch.float64, device='cpu').reshape(2,2,3,3) / 997
images[0,0,0,0] = -0.0
component_video = VideoFromComponents(VideoComponents(
images=images, frame_rate=Fraction(30000,1001),
alpha=torch.full((2,2,3,1),0.123456789,dtype=torch.float64,device='cpu'),
audio=examples['audio'], metadata={'title':'雪', 'nested':[True,None,-4,0.125]}),
bit_depth=10,color_space='HDR PQ')
videos = [component_video,
VideoFromFile(str(sources[0]),start_time=-0.0,duration=0,crop=(0,0,2,2)),
VideoFromList([VideoFromFile(str(sources[1]),start_time=-1.25,duration=-2.5,crop=(-3,5,7,-9)),
VideoFromFile(str(sources[2]),start_time=0,duration=0)],
complete_audio=examples['audio'])]
for socket,value in [*examples.items(), *[('video',v) for v in videos]]:
kind = SOCKET_KINDS[socket]
expected_wire = encode_video(value) if kind == 'VIDEO' else None
binding = dict(binding_id='b-'+socket,collector_id='collector',job_id='job',participant_id='worker',node_id='export-'+socket,socket=socket,kind=kind)
controls = {'request_id':'request','binding_json':json.dumps(binding)}
data, missing, v3 = execution.get_input_data(controls,MediaExportV1,'export-'+socket)
assert not missing, missing
data[socket] = [value]
results = await execution._async_map_node_over_list('probe','export-'+socket,MediaExportV1,data,MediaExportV1.FUNCTION,v3_data=v3)
results = await execution.resolve_map_node_over_list_results(results)
raw = results[0].ui
exported = raw['media_exports'][0]
assert exported['binding']['count'] == 1
artifact = exported['artifacts'][0]
transport = artifact.pop('transport')
filename = transport['subfolder']+'/'+transport['filename']+' [temp]'
decoded = import_value(artifact,filename)
if kind in ('IMAGE','MASK'):
assert decoded.dtype == value.dtype and torch.equal(decoded,value)
elif kind == 'AUDIO':
assert decoded['sample_rate'] == value['sample_rate']
assert torch.equal(decoded['waveform'],value['waveform'])
elif kind == 'VIDEO': exact_video(decoded,value,expected_wire)
else: assert decoded == value
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) == 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)
except ValueError: pass
else: raise AssertionError('checksum accepted')
try: import_value(artifact,str(base/'outside.bin'))
except ValueError: pass
else: raise AssertionError('outside path accepted')
if kind == 'VIDEO':
mapped_video = got[0].result[list(SOCKET_KINDS).index(socket)]
# Remove both the exported artifact and the original input
# containers before checking independently owned BytesIOs.
Path(annotated(filename)).unlink()
originals = value.videos if type(value) is VideoFromList else [value]
for original in originals:
if type(original) is VideoFromFile:
Path(vars(original)['_VideoFromFile__file']).unlink()
exact_video(decoded,value,expected_wire)
exact_video(mapped_video,value,expected_wire)
if type(value) is VideoFromFile:
assert tuple(decoded.get_components().images.shape) == (3,2,2,3)
text_binding = dict(binding_id='empty',collector_id='collector',job_id='job',participant_id='worker',node_id='export-empty',socket='texts',kind='STRING')
out = export_values('request',text_binding,[''])
assert out['artifacts'][0]['byte_length'] == 0
before = set(roots['temp'].rglob('*.bin'))
try: export_values('request',text_binding,['first',None])
except ValueError: pass
else: raise AssertionError('invalid text accepted')
assert set(roots['temp'].rglob('*.bin')) == before
print('MEDIA_BRIDGE_CPU_OK')
asyncio.run(check())
'''
env = dict(os.environ)
env.pop("PYTHONPATH", None)
env.pop("COMFYUI_SOURCE_ROOT", None)
scratch = Path(os.environ.get("TMPDIR", str(Path.home() / ".hermes/cache/scratch")))
scratch.mkdir(parents=True, exist_ok=True)
env.update(CUDA_VISIBLE_DEVICES="", HIP_VISIBLE_DEVICES="", PYTHONDONTWRITEBYTECODE="1",
TMPDIR=str(scratch))
result = subprocess.run([sys.executable, "-c", probe, str(Path(core).resolve()), str(ROOT)],
capture_output=True, text=True, timeout=90, env=env)
assert result.returncode == 0, result.stdout + result.stderr
assert "MEDIA_BRIDGE_CPU_OK" in result.stdout
+328
View File
@@ -0,0 +1,328 @@
import importlib.util
import math
import struct
import types
from pathlib import Path
import pytest
import torch
MODULE_PATH = Path(__file__).resolve().parents[1] / "serverless_bridge" / "codec.py"
spec = importlib.util.spec_from_file_location("dist_media_codec_test", MODULE_PATH)
codec = importlib.util.module_from_spec(spec)
spec.loader.exec_module(codec)
DTYPES = [torch.float16, torch.bfloat16, torch.float32, torch.float64]
SHAPES = [
("IMAGE", (2, 3, 5, 1), "BHWC"),
("IMAGE", (2, 3, 5, 3), "BHWC"),
("IMAGE", (2, 3, 5, 4), "BHWC"),
("MASK", (2, 3, 5), "BHW"),
("AUDIO", (2, 1, 17), "BCT"),
("AUDIO", (2, 6, 17), "BCT"),
]
def audio(tensor, sample_rate=48000):
return {"waveform": tensor, "sample_rate": sample_rate}
def descriptor(_kind="MASK", **changes):
result = {
"encoding": "tensor-le-v1", "dtype": "float32",
"shape": [1, 1, 1], "layout": "BHW",
}
if _kind == "IMAGE":
result.update(shape=[1, 1, 1, 3], layout="BHWC")
elif _kind == "AUDIO":
result.update(layout="BCT", sample_rate=48000)
result.update(changes)
return result
@pytest.fixture(autouse=True)
def forbid_gpu_queries(monkeypatch):
def forbidden(*args, **kwargs):
pytest.fail("codec must not query or initialize a GPU")
for name in ("is_available", "init", "_lazy_init", "device_count"):
monkeypatch.setattr(torch.cuda, name, forbidden)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("kind,shape,layout", SHAPES)
def test_exact_roundtrip_preserves_dtype_shape_channels_and_bits(dtype, kind, shape, layout):
count = math.prod(shape)
tensor = (torch.arange(count, device="cpu", dtype=torch.float64) / 997 - 0.125).to(dtype).reshape(shape)
tensor.flatten()[0] = -0.0
tensor.requires_grad_(True)
value = audio(tensor) if kind == "AUDIO" else tensor
metadata, payload = codec.encode_payload(kind, value)
expected = {
"encoding": "tensor-le-v1", "dtype": str(dtype).removeprefix("torch."),
"shape": list(shape), "layout": layout,
}
if kind == "AUDIO":
expected["sample_rate"] = 48000
assert metadata == expected
assert type(payload) is bytes
assert len(payload) == count * tensor.element_size()
result = codec.decode_payload(kind, metadata, payload)
if kind == "AUDIO":
assert set(result) == {"waveform", "sample_rate"}
assert result["sample_rate"] == 48000
result = result["waveform"]
assert result.device.type == "cpu"
assert result.is_contiguous()
assert not result.requires_grad
assert result.dtype == dtype
assert result.shape == shape
assert torch.equal(result.view(torch.uint8), tensor.detach().view(torch.uint8))
@pytest.mark.parametrize("dtype,fmt,expected_bits", [
(torch.float16, "e", None), (torch.bfloat16, "H", [0x3F80, 0x8000, 0xC020]),
(torch.float32, "f", None), (torch.float64, "d", None),
])
def test_wire_bytes_are_little_endian(dtype, fmt, expected_bits):
values = [1.0, -0.0, -2.5]
tensor = torch.tensor(values, dtype=dtype, device="cpu").reshape(1, 1, 3)
metadata, payload = codec.encode_payload("MASK", tensor)
expected = struct.pack("<3" + fmt, *(expected_bits if expected_bits else values))
assert payload == expected
assert torch.equal(codec.decode_payload("MASK", metadata, expected).view(torch.uint8), tensor.view(torch.uint8))
def test_bfloat16_extremes_and_subnormal_bits_are_lossless():
bits = torch.tensor([0, 0x8000, 1, 0x8001, 0x3F81, 0x7F7F, 0xFF7F], dtype=torch.uint16, device="cpu")
value = bits.view(torch.bfloat16).reshape(1, 1, -1)
metadata, payload = codec.encode_payload("MASK", value)
assert payload == struct.pack("<7H", *bits.tolist())
assert torch.equal(codec.decode_payload("MASK", metadata, payload).view(torch.uint16).flatten(), bits)
def test_fractional_mask_is_not_quantized_clipped_or_normalized():
value = torch.tensor([0.1234567890123, 1.0000000000001, -2.5], dtype=torch.float64, device="cpu").reshape(1, 1, 3)
metadata, payload = codec.encode_payload("MASK", value)
assert torch.equal(codec.decode_payload("MASK", metadata, payload), value)
@pytest.mark.parametrize("text", ["", " ", "\t\n ", "café 雪 🎵", "e\u0301", "\x00\ufeff\r\n"])
def test_strings_preserve_exact_utf8(text):
metadata, payload = codec.encode_payload("STRING", text)
assert metadata == {"encoding": "utf8-v1"}
assert payload == text.encode("utf-8")
assert codec.decode_payload("STRING", metadata, payload) == text
def test_noncontiguous_and_lazy_negative_tensors_roundtrip():
value = torch.arange(24, dtype=torch.float32, device="cpu").reshape(2, 3, 4).transpose(1, 2)
for tensor in (value, torch._neg_view(value)):
metadata, payload = codec.encode_payload("MASK", tensor)
result = codec.decode_payload("MASK", metadata, payload)
assert result.is_contiguous()
assert torch.equal(result, tensor)
def test_decode_is_independent_and_ignores_default_device(monkeypatch):
source = torch.tensor([0.25], device="cpu").reshape(1, 1, 1)
metadata, payload = codec.encode_payload("MASK", source)
buffers = []
original = torch.frombuffer
def capture(buffer, *args, **kwargs):
buffers.append(buffer)
return original(buffer, *args, **kwargs)
monkeypatch.setattr(torch, "frombuffer", capture)
with torch.device("meta"):
first = codec.decode_payload("MASK", metadata, payload)
second = codec.decode_payload("MASK", metadata, payload)
assert first.device.type == second.device.type == "cpu"
first.fill_(9)
source.fill_(7)
for buffer in buffers:
if isinstance(buffer, bytearray):
buffer[:] = b"\x00" * len(buffer)
assert second.item() == 0.25
assert codec.decode_payload("MASK", metadata, payload).item() == 0.25
@pytest.mark.parametrize("kind", ["image", " MASK", "LATENT", "", None, 1, [], {}])
def test_unknown_kinds_raise_value_error(kind):
with pytest.raises(ValueError, match="kind"):
codec.encode_payload(kind, "")
with pytest.raises(ValueError, match="kind"):
codec.decode_payload(kind, {}, b"")
@pytest.mark.parametrize("kind,value", [
("STRING", None), ("STRING", b"text"), ("STRING", 123),
("IMAGE", [[[[0.5]]]]), ("MASK", None),
("AUDIO", torch.ones(1, 1, 1, device="cpu")), ("AUDIO", {}),
("AUDIO", {"waveform": "invalid", "sample_rate": 48000}),
("AUDIO", {"waveform": torch.ones(1, 1, 1, device="cpu")}),
])
def test_unsupported_values_raise_value_error(kind, value):
with pytest.raises(ValueError):
codec.encode_payload(kind, value)
@pytest.mark.parametrize("dtype", [torch.uint8, torch.int64, torch.bool, torch.complex64])
def test_unsupported_tensor_dtype(dtype):
with pytest.raises(ValueError, match="dtype"):
codec.encode_payload("MASK", torch.ones(1, 1, 1, dtype=dtype, device="cpu"))
@pytest.mark.parametrize("kind,shape", [
("IMAGE", (1, 2, 3)), ("IMAGE", (1, 1, 1, 2)), ("IMAGE", (1, 1, 1, 5)),
("MASK", (1, 2)), ("MASK", (1, 2, 3, 1)), ("AUDIO", (2, 4)),
("IMAGE", (1, 0, 2, 3)), ("MASK", (0, 1, 1)), ("AUDIO", (1, 0, 4)),
("AUDIO", (1, 2, 0)),
])
def test_invalid_tensor_shape(kind, shape):
tensor = torch.empty(shape, dtype=torch.float32, device="cpu")
with pytest.raises(ValueError, match="shape|rank|channel|dimension"):
codec.encode_payload(kind, audio(tensor) if kind == "AUDIO" else tensor)
def test_rejects_non_cpu_and_sparse_without_transferring(monkeypatch):
meta = torch.empty((1, 1, 1), device="meta")
sparse = torch.ones((1, 1, 1), device="cpu").to_sparse()
def forbidden(*args, **kwargs):
pytest.fail("codec must not move tensors between devices")
monkeypatch.setattr(torch.Tensor, "cpu", forbidden)
monkeypatch.setattr(torch.Tensor, "to", forbidden)
for tensor in (meta, sparse):
with pytest.raises(ValueError, match="CPU|layout|strided"):
codec.encode_payload("MASK", tensor)
@pytest.mark.parametrize("sample_rate", [True, False, 0, -1, 44100.0, "44100", None])
def test_invalid_sample_rate(sample_rate):
with pytest.raises(ValueError, match="sample_rate"):
codec.encode_payload("AUDIO", audio(torch.ones(1, 1, 1, device="cpu"), sample_rate))
with pytest.raises(ValueError, match="sample_rate"):
codec.decode_payload("AUDIO", descriptor("AUDIO", sample_rate=sample_rate), b"\x00" * 4)
@pytest.mark.parametrize("bad", [None, [], "descriptor", {}, {"encoding": "utf8-v1"}])
def test_malformed_tensor_descriptor(bad):
with pytest.raises(ValueError, match="descriptor"):
codec.decode_payload("MASK", bad, b"\x00" * 4)
@pytest.mark.parametrize("changes", [
{"encoding": "pickle"}, {"encoding": None}, {"encoding": []},
{"dtype": "uint8"}, {"dtype": torch.float32}, {"dtype": []},
{"layout": "BHWC"}, {"layout": None}, {"layout": []},
{"shape": (1, 1, 1)}, {"shape": None}, {"shape": "1,1,1"},
{"shape": [1, 1]}, {"shape": [1, 1, 1, 1]}, {"shape": [1, 0, 1]},
{"shape": [1, -1, 1]}, {"shape": [1, True, 1]}, {"shape": [1, 1.0, 1]},
{"shape": [1, "1", 1]}, {"shape": [1, [], 1]}, {"kind": "MASK"},
])
def test_invalid_tensor_descriptor_fields(changes):
with pytest.raises(ValueError):
codec.decode_payload("MASK", descriptor(**changes), b"\x00" * 4)
@pytest.mark.parametrize("metadata", [{}, {"encoding": "tensor-le-v1"}, {"encoding": "utf8-v1", "kind": "STRING"}, {"encoding": []}])
def test_invalid_string_descriptor(metadata):
with pytest.raises(ValueError):
codec.decode_payload("STRING", metadata, b"")
@pytest.mark.parametrize("payload", [b"", b"\x00" * 3, b"\x00" * 5, b"\x00" * 8])
def test_shape_length_mismatch(payload):
with pytest.raises(ValueError, match="length"):
codec.decode_payload("MASK", descriptor(), payload)
@pytest.mark.parametrize("payload", [None, "text", [0, 0, 0, 0], 4])
def test_nonbyte_payload_rejected(payload):
with pytest.raises(ValueError, match="bytes"):
codec.decode_payload("MASK", descriptor(), payload)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("value", [float("nan"), float("inf"), -float("inf")])
def test_nonfinite_values_fail_on_encode_and_decode(dtype, value):
tensor = torch.tensor([value], dtype=dtype, device="cpu").reshape(1, 1, 1)
with pytest.raises(ValueError, match="finite"):
codec.encode_payload("MASK", tensor)
fmt = {torch.float16: "e", torch.float32: "f", torch.float64: "d"}
if dtype == torch.bfloat16:
payload = struct.pack("<H", 0x7FC0 if math.isnan(value) else (0x7F80 if value > 0 else 0xFF80))
else:
payload = struct.pack("<" + fmt[dtype], value)
metadata = descriptor(dtype=str(dtype).removeprefix("torch."))
with pytest.raises(ValueError, match="finite"):
codec.decode_payload("MASK", metadata, payload)
def test_invalid_utf8_errors_do_not_echo_content():
with pytest.raises(ValueError) as error:
codec.encode_payload("STRING", "private-value\ud800")
assert "private-value" not in str(error.value)
with pytest.raises(ValueError) as error:
codec.decode_payload("STRING", {"encoding": "utf8-v1"}, b"private-value\xff")
assert "private-value" not in str(error.value)
@pytest.mark.parametrize("max_bytes", [True, False, -1, 4.0, "4", None])
def test_invalid_byte_budget(max_bytes):
with pytest.raises(ValueError, match="max_bytes"):
codec.encode_payload("STRING", "", max_bytes=max_bytes)
with pytest.raises(ValueError, match="max_bytes"):
codec.decode_payload("STRING", {"encoding": "utf8-v1"}, b"", max_bytes=max_bytes)
def test_exact_byte_budget_and_zero_budget_for_empty_string():
tensor = torch.ones(1, 1, 1, device="cpu")
metadata, payload = codec.encode_payload("MASK", tensor, max_bytes=4)
assert codec.decode_payload("MASK", metadata, payload, max_bytes=4).item() == 1
with pytest.raises(ValueError, match="limit|budget"):
codec.encode_payload("MASK", tensor, max_bytes=3)
with pytest.raises(ValueError, match="limit|budget"):
codec.decode_payload("MASK", metadata, payload, max_bytes=3)
metadata, payload = codec.encode_payload("STRING", "", max_bytes=0)
assert codec.decode_payload("STRING", metadata, payload, max_bytes=0) == ""
for text in ("a", "é", "🎵"):
byte_count = len(text.encode("utf-8"))
metadata, payload = codec.encode_payload("STRING", text, max_bytes=byte_count)
assert codec.decode_payload("STRING", metadata, payload, max_bytes=byte_count) == text
with pytest.raises(ValueError, match="limit|budget"):
codec.encode_payload("STRING", text, max_bytes=byte_count - 1)
with pytest.raises(ValueError, match="limit|budget"):
codec.decode_payload("STRING", metadata, payload, max_bytes=byte_count - 1)
def test_oversized_shape_rejected_before_tensor_allocation(monkeypatch):
assert codec.MAX_PAYLOAD_BYTES == 256 * 1024 * 1024
def forbidden(*args, **kwargs):
pytest.fail("oversized shape must fail before tensor allocation")
for name in ("frombuffer", "empty", "tensor", "zeros", "isfinite"):
monkeypatch.setattr(torch, name, forbidden)
for shape in ([1, 1, codec.MAX_PAYLOAD_BYTES], [2**100, 2**100, 1]):
with pytest.raises(ValueError, match="limit|budget"):
codec.decode_payload("MASK", descriptor(shape=shape), b"\x00" * 4)
def test_oversized_encode_rejected_before_materialization(monkeypatch):
huge_view = torch.ones(1, device="cpu").expand(1, 1, codec.MAX_PAYLOAD_BYTES)
def forbidden(*args, **kwargs):
pytest.fail("oversized tensor must fail before materialization or scanning")
monkeypatch.setattr(torch.Tensor, "contiguous", forbidden)
monkeypatch.setattr(torch, "isfinite", forbidden)
with pytest.raises(ValueError, match="limit|budget"):
codec.encode_payload("MASK", huge_view)
@pytest.mark.parametrize("itemsize", [2, 4, 8])
def test_native_little_endian_conversion_handles_big_endian(monkeypatch, itemsize):
raw = bytes(range(itemsize * 3))
monkeypatch.setattr(codec, "sys", types.SimpleNamespace(byteorder="little"))
assert codec._native_little_endian(raw, itemsize) == raw
monkeypatch.setattr(codec, "sys", types.SimpleNamespace(byteorder="big"))
expected = b"".join(raw[i:i + itemsize][::-1] for i in range(0, len(raw), itemsize))
assert codec._native_little_endian(raw, itemsize) == expected
assert codec._native_little_endian(expected, itemsize) == raw
+152
View File
@@ -0,0 +1,152 @@
"""Offline contract tests: malformed results never reach transport/delivery."""
import asyncio
import copy
import importlib.util
from pathlib import Path
import pytest
ROOT = Path(__file__).resolve().parents[1]
spec = importlib.util.spec_from_file_location("media_contract_test", ROOT / "serverless_bridge" / "contract.py")
media = importlib.util.module_from_spec(spec)
spec.loader.exec_module(media)
def binding(**changes):
return dict(binding_id="binding-1", collector_id="100", job_id="job-1",
participant_id="serverless:worker:0", node_id="export-1", socket="texts", kind="STRING", **changes)
def response():
artifact = media.artifact_from_payload("a", "STRING", {"encoding": "utf8-v1"}, b" \n")
return {"version": 1, "profile": media.PROFILE, "request_id": "request-1",
"bindings": [dict(binding(), count=1)], "artifacts": [artifact],
"items": [{"binding_id": "binding-1", "item_index": 0, "artifact_id": "a"}]}
def validate(value):
return media.validate_response(value, request_id="request-1", bindings=[binding()])
def test_valid_manifest_and_exact_inline_payload():
value = response()
ordered = validate(value)
assert asyncio.run(media.resolve_artifact(ordered["binding-1"][0])) == b" \n"
def test_order_uses_indices_not_arrival():
value = response()
value["bindings"][0]["count"] = 2
value["artifacts"].append(media.artifact_from_payload("b", "STRING", {"encoding": "utf8-v1"}, b"second"))
value["items"] = [{"binding_id": "binding-1", "item_index": 1, "artifact_id": "b"}, value["items"][0]]
assert [item["id"] for item in validate(value)["binding-1"]] == ["a", "b"]
@pytest.mark.parametrize("field,value", [("version", True), ("version", 2), ("profile", "legacy"), ("request_id", "another")])
def test_envelope_identity(field, value):
obj = response()
obj[field] = value
with pytest.raises(ValueError):
validate(obj)
@pytest.mark.parametrize("field", ["collector_id", "participant_id", "job_id", "node_id", "binding_id"])
def test_binding_identity(field):
value = response()
value["bindings"][0][field] = "another"
with pytest.raises(ValueError):
validate(value)
@pytest.mark.parametrize("count", [0, -1, True, 1.0, "1", 65])
def test_invalid_count(count):
value = response()
value["bindings"][0]["count"] = count
with pytest.raises(ValueError):
validate(value)
@pytest.mark.parametrize("mutation", ["missing_binding", "missing_item", "extra_item", "duplicate_binding", "duplicate_artifact", "unused_artifact", "wrong_type", "wrong_index"])
def test_incomplete_or_ambiguous_manifest(mutation):
value = response()
if mutation == "missing_binding":
value["bindings"] = []
elif mutation == "missing_item":
value["items"] = []
elif mutation == "extra_item":
value["items"] *= 2
elif mutation == "duplicate_binding":
value["bindings"] *= 2
elif mutation == "duplicate_artifact":
value["artifacts"] *= 2
elif mutation == "unused_artifact":
artifact = copy.deepcopy(value["artifacts"][0])
artifact["id"] = "unused"
value["artifacts"].append(artifact)
elif mutation == "wrong_type":
value["artifacts"][0]["kind"] = "MASK"
elif mutation == "wrong_index":
value["items"][0]["item_index"] = True
with pytest.raises(ValueError):
validate(value)
def test_duplicate_item_indices():
value = response()
value["bindings"][0]["count"] = 2
value["items"] *= 2
with pytest.raises(ValueError, match="Duplicate"):
validate(value)
def test_total_budget():
with pytest.raises(ValueError, match="total byte"):
media.validate_response(response(), request_id="request-1", bindings=[binding()], max_total_bytes=1)
def test_checksum_and_base64_validation():
artifact = response()["artifacts"][0]
artifact["sha256"] = "0" * 64
with pytest.raises(ValueError, match="checksum"):
asyncio.run(media.resolve_artifact(artifact))
artifact["transport"]["data_base64"] = "!!!!"
with pytest.raises(ValueError, match="base64"):
asyncio.run(media.resolve_artifact(artifact))
def test_empty_text_is_a_value():
artifact = media.artifact_from_payload("empty", "STRING", {"encoding": "utf8-v1"}, b"")
assert asyncio.run(media.resolve_artifact(artifact)) == b""
@pytest.mark.parametrize("url,expires", [("http://example.com/a", "2099-01-01T00:00:00Z"),
("https://user:secret@example.com/a", "2099-01-01T00:00:00Z"),
("https://example.com/a", "2000-01-01T00:00:00Z"),
("https://example.com/a", "2099-01-01"),
("https://example.com:1234/a", "2099-01-01T00:00:00Z")])
def test_invalid_url_envelopes(url, expires):
with pytest.raises(ValueError):
media.artifact_from_payload("a", "STRING", {"encoding": "utf8-v1"}, b"a",
transport={"type": "url", "url": url, "expires_at": expires})
def test_url_resolution_is_bounded_and_hash_verified():
artifact = media.artifact_from_payload("a", "STRING", {"encoding": "utf8-v1"}, b"hello",
transport={"type": "url", "url": "https://example.com/a", "expires_at": "2099-01-01T00:00:00Z"})
calls = []
async def download(url, *, max_bytes):
calls.append((url, max_bytes))
return b"hello"
assert asyncio.run(media.resolve_artifact(artifact, download=download)) == b"hello"
assert calls == [("https://example.com/a", 5)]
def test_descriptor_depth_and_nonfinite_rejected():
value = {}
cursor = value
for _ in range(15):
cursor["next"] = {}
cursor = cursor["next"]
for descriptor in [value, {"x": float("nan")}, {"x": object()}]:
with pytest.raises(ValueError):
media.artifact_from_payload("a", "STRING", descriptor, b"")
+48
View File
@@ -0,0 +1,48 @@
"""Opt-in CPU acceptance of Collector IMAGE/AUDIO and the five-kind bridge.
COMFYUI_SOURCE_ROOT must name an existing compatible checkout; run with that
checkout's Python environment. No installation, normal instance, or provider is
used. The child owns its loopback server and every writable Comfy directory.
"""
import os
from pathlib import Path
import subprocess
import sys
import tempfile
import pytest
ROOT = Path(__file__).resolve().parents[1]
def test_typed_media_end_to_end_cpu():
configured = os.environ.get("COMFYUI_SOURCE_ROOT")
if configured is None:
pytest.skip("Set COMFYUI_SOURCE_ROOT for real CPU typed-media end-to-end acceptance")
core = Path(configured).expanduser().resolve()
assert configured and all((core / name).is_file() for name in (
"execution.py", "server.py", "comfy/cli_args.py",
"comfy_api/latest/_input_impl/video_types.py",
)), "Configured COMFYUI_SOURCE_ROOT must provide PromptExecutor and native V3 VIDEO"
scratch = Path(os.environ.get("TMPDIR", str(Path.home() / ".hermes/cache/scratch")))
scratch.mkdir(parents=True, exist_ok=True)
env = dict(os.environ)
for name in ("PYTHONPATH", "COMFYUI_SOURCE_ROOT"):
env.pop(name, None)
env.update(CUDA_VISIBLE_DEVICES="", HIP_VISIBLE_DEVICES="", ROCR_VISIBLE_DEVICES="",
PYTHONDONTWRITEBYTECODE="1", PYTHONHASHSEED="0", COMFYUI_IS_WORKER="1",
OMP_NUM_THREADS="1", MKL_NUM_THREADS="1")
with tempfile.TemporaryDirectory(prefix="media-e2e-", dir=scratch) as run_root:
env["TMPDIR"] = run_root
result = subprocess.run(
[sys.executable, str(ROOT / "tests/helpers/media_end_to_end_check.py"),
str(core), run_root], cwd=run_root, env=env,
capture_output=True, text=True, timeout=150,
)
assert result.returncode == 0, result.stdout + result.stderr
for marker in ("NATIVE_INPUTS_AND_DIRECT_EXPORT_REFS_OK", "REAL_HTTP_BRIDGE_RESULTS_OK",
"MASTER_AND_DELEGATE_ORDER_OK", "ATOMIC_DECODE_REJECTION_OK",
"IMAGE_TRANSFER_HTTP_OK",
"MEDIA_END_TO_END_CPU_OK"):
assert marker in result.stdout, result.stdout + result.stderr
print(result.stdout)
+877
View File
@@ -0,0 +1,877 @@
"""CPU-only endpoint adapter tests against an actual loopback Comfy HTTP fake."""
import asyncio
import base64
import copy
import hashlib
import importlib
import json
import sys
import types
from contextlib import asynccontextmanager
from pathlib import Path
import aiohttp
from aiohttp import web
import pytest
ROOT = Path(__file__).resolve().parents[1]
PACKAGE = "media_handler_testpkg"
package = types.ModuleType(PACKAGE)
package.__path__ = [str(ROOT / "serverless_bridge")]
sys.modules[PACKAGE] = package
handler = importlib.import_module(PACKAGE + ".handler")
contract = importlib.import_module(PACKAGE + ".contract")
def request(kind="STRING", *, incoming=True):
socket = {value: key for key, value in contract.SOCKET_KINDS.items()}[kind]
binding = dict(binding_id="binding:1", collector_id="collector:1", job_id="job/42",
participant_id="worker:0", node_id="export:1", socket=socket, kind=kind)
artifact = contract.artifact_from_payload("input:1", kind, {"encoding": "test"}, b"input bytes")
return {"workflow": {
"source:1": {"class_type": "Source", "inputs": {"value": "literal"}},
"export:1": {"class_type": "MediaExportV1", "inputs": {
"request_id": "request:1", "binding_json": json.dumps(binding), socket: ["source:1", 0]}},
}, "media": {"version": 1, "profile": contract.PROFILE, "request_id": "request:1",
"artifacts": [artifact] if incoming else [],
"inputs": [{"node_id": "source:1", "input_name": "value", "artifact_id": "input:1"}] if incoming else [],
"outputs": [binding]}}
class Comfy:
def __init__(self, root, data=None, *, payload=b"result bytes"):
self.root = root
self.data = data or request()
self.payload = payload
self.submitted = []
self.views = []
self.deletes = []
self.interrupts = []
self.polls = 0
self.active = False
self.mode = "success"
self.history_failures = 0
self.record_mutation = None
self.prompt_seen = asyncio.Event()
self.terminal_seen = False
def staged(self):
return list(self.root.glob("distributed-media-v1/*/*.bin"))
def record(self):
binding = self.data["media"]["outputs"][0]
aid = "a" * 32
artifact = contract.artifact_from_payload(aid, binding["kind"], {"encoding": "test"}, self.payload)
artifact["transport"] = {"type": "comfy-file", "filename": aid + ".bin",
"subfolder": "distributed-media-v1/" + hashlib.sha256(b"request:1").hexdigest(),
"folder_type": "temp"}
value = {"request_id": "request:1", "binding": dict(binding, count=1), "artifacts": [artifact],
"items": [{"binding_id": binding["binding_id"], "item_index": 0, "artifact_id": aid}]}
if self.record_mutation:
self.record_mutation(value)
return value
async def submit(self, req):
data = await req.json()
self.submitted.append(data)
self.active = True
self.prompt_seen.set()
if self.data["media"]["artifacts"]:
assert self.staged()
link = data["prompt"]["source:1"]["inputs"]["value"]
imported = data["prompt"][link[0]]
assert imported["class_type"] == "MediaImportV1"
metadata = json.loads(imported["inputs"]["artifact_json"])
assert "transport" not in metadata
assert (self.root / imported["inputs"]["filename"]).read_bytes() == b"input bytes"
if self.mode == "ambiguous":
return web.Response(status=500)
if self.mode == "partial":
return web.json_response({"prompt_id": "prompt:1", "node_errors": {"export:1": {"errors": ["bad"]}}})
return web.json_response({"prompt_id": "prompt:1", "node_errors": {}})
async def history(self, req):
self.polls += 1
if self.history_failures:
self.history_failures -= 1
return web.Response(status=503)
if self.mode in {"queued", "running", "unrelated", "ambiguous"}:
return web.json_response({})
if self.data["media"]["artifacts"]:
assert self.staged(), "staging must survive until terminal confirmation"
self.active = False
self.terminal_seen = True
if self.mode in {"error", "interrupted"}:
status = {"completed": False, "status_str": "error", "messages": [["execution_interrupted", {}]]}
else:
status = {"completed": True, "status_str": "success", "messages": []}
return web.json_response({"prompt:1": {"status": status, "outputs": {} if self.mode == "missing" else {
"export:1": {"media_exports": [self.record()]}}}})
async def view(self, req):
self.views.append(dict(req.query))
assert self.terminal_seen
payload = self.payload if self.mode != "corrupt" else b"X" * len(self.payload)
if self.mode == "oversized":
payload += b"X"
return web.Response(body=payload)
def entry(self, pid="prompt:1"):
extra = dict(self.submitted[0].get("extra_data", {}), client_id=self.submitted[0]["client_id"])
return [1, pid, self.submitted[0]["prompt"], extra, ["export:1"]]
async def queue(self, req):
if req.method == "POST":
body = await req.json()
assert body == {"delete": ["prompt:1"]}
self.deletes.append(body)
assert self.staged()
self.mode = "deleted"
self.active = False
return web.Response(status=200)
running = [self.entry()] if self.mode == "running" else []
if self.mode == "unrelated":
running = [self.entry("not-our-prompt")]
pending = [self.entry()] if self.mode in {"queued", "ambiguous"} else []
return web.json_response({"queue_running": running, "queue_pending": pending})
async def interrupt(self, req):
assert self.mode == "running"
assert self.staged()
self.interrupts.append(await req.json())
self.mode = "interrupted"
return web.Response(status=200)
@asynccontextmanager
async def serve(comfy):
app = web.Application()
app.add_routes([web.post("/prompt", comfy.submit), web.get("/history/{prompt_id}", comfy.history),
web.get("/view", comfy.view), web.get("/queue", comfy.queue),
web.post("/queue", comfy.queue), web.post("/interrupt", comfy.interrupt)])
runner = web.AppRunner(app)
await runner.setup()
site = web.TCPSite(runner, "127.0.0.1", 0)
await site.start()
port = site._server.sockets[0].getsockname()[1]
try:
yield f"http://127.0.0.1:{port}"
finally:
await runner.cleanup()
def run(coro):
return asyncio.run(coro)
@pytest.mark.parametrize("kind,index", [("IMAGE", 0), ("AUDIO", 1), ("MASK", 2), ("STRING", 3), ("VIDEO", 4)])
def test_loopback_submit_poll_view_cleanup_and_caller_ownership(tmp_path, kind, index):
async def check():
data = request(kind)
original = copy.deepcopy(data)
comfy = Comfy(tmp_path, data)
async with serve(comfy) as url, aiohttp.ClientSession() as session:
result = await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path, session=session)
assert not session.closed
contract.validate_response(result["media"], request_id="request:1", bindings=data["media"]["outputs"])
assert await contract.resolve_artifact(result["media"]["artifacts"][0]) == comfy.payload
assert len(comfy.submitted) == 1
assert comfy.submitted[0]["partial_execution_targets"] == ["export:1"]
assert comfy.submitted[0]["prompt"]["source:1"]["inputs"]["value"][1] == index
assert comfy.views == [{"filename": "a" * 32 + ".bin", "subfolder": "distributed-media-v1/" + hashlib.sha256(b"request:1").hexdigest(), "type": "temp"}]
assert data == original
assert not comfy.staged()
run(check())
@pytest.mark.parametrize("mode,match", [("corrupt", "checksum"), ("missing", "Missing"), ("error", "execution"), ("oversized", "length")])
def test_terminal_failures_clean_only_after_terminal(tmp_path, mode, match):
async def check():
comfy = Comfy(tmp_path)
comfy.mode = mode
async with serve(comfy) as url:
with pytest.raises((ValueError, RuntimeError), match=match):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
assert len(comfy.submitted) == 1
assert not comfy.staged()
run(check())
@pytest.mark.parametrize("mutation", [
lambda d: d["media"].update(version=True),
lambda d: d["media"].update(profile="legacy"),
lambda d: d["media"].update(request_id="../escape"),
lambda d: d["media"]["outputs"][0].update(node_id="missing"),
lambda d: d["workflow"]["export:1"].update(class_type="SaveImage"),
lambda d: d["workflow"]["export:1"]["inputs"].update(request_id="other"),
lambda d: d["workflow"]["export:1"]["inputs"].update(binding_json="{}"),
lambda d: d["workflow"]["export:1"]["inputs"].update(images=["source:1", 0]),
lambda d: d["workflow"]["export:1"]["inputs"].update(texts="literal"),
lambda d: d["workflow"].update(unrelated={"class_type": "SaveImage", "inputs": {}}),
lambda d: d["workflow"].update(unsafe={"class_type": "MediaImportV1", "inputs": {}}),
lambda d: d["media"]["inputs"].append(copy.deepcopy(d["media"]["inputs"][0])),
lambda d: d["media"]["inputs"][0].update(node_id="missing"),
lambda d: d["media"]["inputs"][0].update(input_name="missing"),
lambda d: d["media"].update(inputs=[]),
lambda d: d["media"]["artifacts"][0].update(sha256="0" * 64),
])
def test_rejection_before_submit(tmp_path, mutation):
async def check():
data = request()
mutation(data)
comfy = Comfy(tmp_path, data)
async with serve(comfy) as url:
with pytest.raises(ValueError):
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path)
assert not comfy.submitted
assert not comfy.staged()
run(check())
@pytest.mark.parametrize("mutation", [
lambda r: r.update(request_id="other"),
lambda r: r["binding"].update(count=2),
lambda r: r["binding"].update(collector_id="other"),
lambda r: r["items"][0].update(item_index=1),
lambda r: r["artifacts"][0]["transport"].update(filename="../escape"),
lambda r: r["artifacts"][0]["transport"].update(subfolder="distributed-media-v1/other"),
lambda r: r["artifacts"][0]["transport"].update(folder_type="input"),
lambda r: r["artifacts"][0].update(byte_length=contract.MAX_ARTIFACT_BYTES + 1),
])
def test_bad_result_manifest_rejected_before_any_file_read(tmp_path, mutation):
async def check():
comfy = Comfy(tmp_path)
comfy.record_mutation = mutation
async with serve(comfy) as url:
with pytest.raises(ValueError):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
assert not comfy.views
assert not comfy.staged()
run(check())
def test_idempotent_history_retry_without_resubmission(tmp_path):
async def check():
comfy = Comfy(tmp_path)
comfy.history_failures = 2
async with serve(comfy) as url:
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
assert comfy.polls == 3
assert len(comfy.submitted) == 1
run(check())
@pytest.mark.parametrize("mode", ["queued"])
def test_deadline_cancels_only_own_prompt_before_cleanup(tmp_path, monkeypatch, mode):
monkeypatch.setattr(handler, "POLL_INTERVAL_SECONDS", 0.005)
async def check():
comfy = Comfy(tmp_path)
comfy.mode = mode
async with serve(comfy) as url:
with pytest.raises(TimeoutError):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path, timeout_seconds=0.04)
assert len(comfy.submitted) == 1
assert bool(comfy.deletes) == (mode == "queued")
assert bool(comfy.interrupts) == (mode == "running")
assert not comfy.staged()
run(check())
def test_unconfirmed_terminal_retains_files_and_never_interrupts_unrelated(tmp_path, monkeypatch):
monkeypatch.setattr(handler, "POLL_INTERVAL_SECONDS", 0.005)
monkeypatch.setattr(handler, "CANCEL_GRACE_SECONDS", 0.06)
async def check():
comfy = Comfy(tmp_path)
comfy.mode = "unrelated"
async with serve(comfy) as url:
with pytest.raises(RuntimeError, match="retained"):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path, timeout_seconds=0.04)
assert comfy.staged()
assert not comfy.interrupts
assert not comfy.deletes
assert len(comfy.submitted) == 1
run(check())
def test_ambiguous_post_is_not_retried_and_own_queue_entry_is_reconciled(tmp_path):
async def check():
comfy = Comfy(tmp_path)
comfy.mode = "ambiguous"
async with serve(comfy) as url:
with pytest.raises(RuntimeError):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
assert len(comfy.submitted) == 1
assert len(comfy.deletes) == 1
assert not comfy.staged()
run(check())
def test_task_cancellation_fails_closed_and_preserves_caller_session(tmp_path):
async def check():
comfy = Comfy(tmp_path)
comfy.mode = "running"
async with serve(comfy) as url, aiohttp.ClientSession() as session:
task = asyncio.create_task(handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path, session=session))
await comfy.prompt_seen.wait()
await asyncio.sleep(0.01)
task.cancel()
with pytest.raises(handler.MediaExecutionUncertain, match="prompt-scoped interruption unavailable"):
await task
assert not session.closed
assert not comfy.interrupts
assert comfy.staged()
run(check())
def test_legacy_images_are_owned_and_rewritten_without_overwriting_user_files(tmp_path):
async def check():
data = request(incoming=False)
data["workflow"]["source:1"] = {"class_type": "LoadImage", "inputs": {"image": "same.png"}}
data["images"] = [{"name": "same.png", "image": base64.b64encode(b"new image").decode()}]
(tmp_path / "same.png").write_bytes(b"user image")
comfy = Comfy(tmp_path, data)
async with serve(comfy) as url:
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path)
filename = comfy.submitted[0]["prompt"]["source:1"]["inputs"]["image"]
assert filename.startswith("distributed-media-v1/")
assert filename.endswith(".bin")
assert (tmp_path / "same.png").read_bytes() == b"user image"
assert not comfy.staged()
run(check())
def test_url_input_uses_bounded_injected_downloader(tmp_path):
async def check():
data = request()
data["media"]["artifacts"][0]["transport"] = {"type": "url", "url": "https://example.org/a", "expires_at": "2099-01-01T00:00:00Z"}
calls = []
async def download(url, *, max_bytes):
calls.append((url, max_bytes))
return b"input bytes"
comfy = Comfy(tmp_path, data)
async with serve(comfy) as url:
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path, download=download)
assert calls == [("https://example.org/a", len(b"input bytes"))]
run(check())
def test_large_result_requires_publisher_and_validates_returned_transport(tmp_path, monkeypatch):
monkeypatch.setattr(handler, "MAX_RESPONSE_JSON_BYTES", 100_000)
async def check():
comfy = Comfy(tmp_path, payload=b"Z" * 100_000)
async with serve(comfy) as url:
with pytest.raises(ValueError, match="publisher"):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
assert not comfy.views
assert not comfy.staged()
published = []
async def publish(metadata, payload):
assert "transport" not in metadata
assert payload == comfy.payload
published.append(metadata)
return {"type": "url", "url": "https://example.org/result", "expires_at": "2099-01-01T00:00:00Z"}
async with serve(comfy) as url:
result = await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path, publish=publish)
assert len(published) == 1
assert result["media"]["artifacts"][0]["transport"]["type"] == "url"
assert len(json.dumps(result).encode()) <= 100_000
async def bad_publish(metadata, payload):
return {"type": "url", "url": "http://example.org/result", "expires_at": "2099-01-01T00:00:00Z"}
async with serve(comfy) as url:
with pytest.raises(ValueError):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path, publish=bad_publish)
run(check())
@pytest.mark.parametrize("name", ["../photo.png", "folder/photo.png", "folder\\photo.png", ".", "..", "/photo.png"])
def test_legacy_basename_rejected_without_submit(tmp_path, name):
async def check():
data = request(incoming=False)
data["images"] = [{"name": name, "image": "YQ=="}]
comfy = Comfy(tmp_path, data)
async with serve(comfy) as url:
with pytest.raises(ValueError):
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path)
assert not comfy.submitted
run(check())
def test_staging_rejects_symlink_without_writing_or_submit(tmp_path):
async def check():
outside = tmp_path / "unrelated"
outside.mkdir()
(tmp_path / "distributed-media-v1").symlink_to(outside, target_is_directory=True)
comfy = Comfy(tmp_path)
async with serve(comfy) as url:
with pytest.raises(ValueError, match="symlink"):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
assert not comfy.submitted
assert not list(outside.iterdir())
run(check())
def test_cleanup_preserves_other_runs_in_same_request_namespace(tmp_path):
async def check():
directory = tmp_path / "distributed-media-v1" / hashlib.sha256(b"request:1").hexdigest()
directory.mkdir(parents=True)
sibling = directory / ("b" * 32 + ".bin")
sibling.write_bytes(b"other run")
comfy = Comfy(tmp_path)
async with serve(comfy) as url:
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
assert sibling.read_bytes() == b"other run"
assert comfy.staged() == [sibling]
run(check())
def test_input_total_limit_checked_before_downloader_or_submit(tmp_path):
async def check():
data = request()
data["media"]["artifacts"] = []
data["media"]["inputs"] = []
for index in range(3):
artifact = contract.artifact_from_payload(f"input:{index}", "STRING", {}, b"",
transport={"type": "url", "url": "https://example.org/input", "expires_at": "2099-01-01T00:00:00Z"})
artifact["byte_length"] = contract.MAX_ARTIFACT_BYTES
data["media"]["artifacts"].append(artifact)
name = f"value:{index}"
data["workflow"]["source:1"]["inputs"][name] = "literal"
data["media"]["inputs"].append({"node_id": "source:1", "input_name": name, "artifact_id": artifact["id"]})
comfy = Comfy(tmp_path, data)
async def forbidden_download(*args, **kwargs):
pytest.fail("oversized request reached downloader")
async with serve(comfy) as url:
with pytest.raises(ValueError, match="total byte"):
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path, download=forbidden_download)
assert not comfy.submitted
run(check())
def test_output_total_limit_checked_before_fetch(tmp_path):
def excessive(record):
first = record["artifacts"][0]
record["artifacts"] = []
record["items"] = []
record["binding"]["count"] = 3
for index, char in enumerate("abc"):
artifact = copy.deepcopy(first)
artifact["id"] = char * 32
artifact["byte_length"] = contract.MAX_ARTIFACT_BYTES
artifact["transport"]["filename"] = artifact["id"] + ".bin"
record["artifacts"].append(artifact)
record["items"].append({"binding_id": "binding:1", "item_index": index, "artifact_id": artifact["id"]})
async def check():
comfy = Comfy(tmp_path)
comfy.record_mutation = excessive
async with serve(comfy) as url:
with pytest.raises(ValueError, match="total byte"):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
assert not comfy.views
assert not comfy.staged()
run(check())
def test_partial_target_acceptance_is_not_reported_as_success(tmp_path):
async def check():
comfy = Comfy(tmp_path)
comfy.mode = "partial"
async with serve(comfy) as url:
with pytest.raises(RuntimeError, match="incomplete"):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
assert len(comfy.submitted) == 1
assert not comfy.views
assert not comfy.staged()
run(check())
# Native routes are compiled from the pinned source, without importing ComfyUI
# (which would initialize the GPU). All dependencies and HTTP are CPU-only.
@pytest.mark.parametrize("current,requested,race,expected", [
("own", "own", False, ["own"]),
("other", "own", False, []),
("own", "own", True, ["successor"]),
])
def test_pinned_native_interrupt_route_has_snapshot_race(current, requested, race, expected):
import ast
import logging
source = Path("/home/robert/ComfyUI_dev/server.py")
if not source.is_file():
pytest.skip("Pinned ComfyUI source unavailable")
tree = ast.parse(source.read_text())
routes = [node for node in ast.walk(tree) if isinstance(node, ast.AsyncFunctionDef)
and any("/interrupt" in ast.unparse(d) for d in node.decorator_list)]
assert len(routes) == 1
route = routes[0]
route.decorator_list = []
interrupts = []
class Queue:
running = current
def get_current_queue(self):
snapshot = [[1, self.running, {}, {}, []]]
if race:
self.running = "successor"
return snapshot, []
queue = Queue()
scope = {"json": json, "logging": logging, "web": web,
"self": types.SimpleNamespace(prompt_queue=queue),
"nodes": types.SimpleNamespace(interrupt_processing=lambda: interrupts.append(queue.running))}
exec(compile(ast.Module(body=[route], type_ignores=[]), str(source), "exec"), scope)
async def check():
async def body():
return {"prompt_id": requested}
response = await scope[route.name](types.SimpleNamespace(json=body))
assert response.status == 200
run(check())
assert interrupts == expected
def exporter_file(root, *, name="a" * 32 + ".bin", payload=b"result bytes"):
directory = root / "distributed-media-v1" / hashlib.sha256(b"request:1").hexdigest()
directory.mkdir(parents=True, exist_ok=True)
path = directory / name
path.write_bytes(payload)
return path
@pytest.mark.parametrize("mode", ["success", "corrupt", "oversized", "error", "interrupted", "partial"])
def test_owned_export_cleanup_after_terminal_success_or_failure(tmp_path, mode):
async def check():
temp = tmp_path / "temp"
owned = exporter_file(temp)
sibling = exporter_file(temp, name="b" * 32 + ".bin", payload=b"unrelated")
ordinary = exporter_file(temp, name="keep.txt", payload=b"ordinary")
comfy = Comfy(tmp_path / "input")
comfy.mode = mode
view = comfy.view
async def reading(req):
assert owned.exists(), "cleanup raced artifact collection"
return await view(req)
comfy.view = reading
async with serve(comfy) as url:
if mode == "success":
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
else:
with pytest.raises((ValueError, RuntimeError)):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
assert comfy.terminal_seen
assert not owned.exists()
assert sibling.read_bytes() == b"unrelated"
assert ordinary.read_bytes() == b"ordinary"
assert not comfy.staged()
run(check())
def test_omitted_temp_root_leaves_exporter_lifecycle_ownership(tmp_path):
async def check():
owned = exporter_file(tmp_path / "temp")
comfy = Comfy(tmp_path / "input")
async with serve(comfy) as url:
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root)
assert owned.read_bytes() == comfy.payload
run(check())
@pytest.mark.parametrize("mode", ["running", "unrelated"])
def test_uncertain_exports_and_inputs_retained_without_interrupt(tmp_path, monkeypatch, mode):
monkeypatch.setattr(handler, "POLL_INTERVAL_SECONDS", 0.005)
monkeypatch.setattr(handler, "CANCEL_GRACE_SECONDS", 0.04)
async def check():
temp = tmp_path / "temp"
owned = exporter_file(temp)
comfy = Comfy(tmp_path / "input")
comfy.mode = mode
async with serve(comfy) as url:
with pytest.raises(handler.MediaExecutionUncertain):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root,
temp_dir=temp, timeout_seconds=0.03)
assert owned.exists()
assert comfy.staged()
assert not comfy.interrupts
assert not comfy.deletes
run(check())
@pytest.mark.parametrize("component", ["root", "namespace", "request", "file"])
def test_export_cleanup_never_traverses_symlinks(tmp_path, component):
async def check():
temp = tmp_path / "temp"
owned = exporter_file(temp)
selected = {"root": temp, "namespace": owned.parent.parent,
"request": owned.parent, "file": owned}[component]
outside = tmp_path / "outside"
selected.rename(outside)
selected.symlink_to(outside, target_is_directory=component != "file")
comfy = Comfy(tmp_path / "input")
async with serve(comfy) as url:
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
assert selected.is_symlink()
assert owned.read_bytes() == comfy.payload
run(check())
@pytest.mark.parametrize("mutation", [
lambda r: r.update(request_id="other"),
lambda r: r["binding"].update(collector_id="other"),
lambda r: r["artifacts"][0]["transport"].update(filename="../escape"),
lambda r: r["artifacts"][0]["transport"].update(subfolder="distributed-media-v1/other"),
])
def test_unidentified_exports_are_not_deleted(tmp_path, mutation):
async def check():
temp = tmp_path / "temp"
owned = exporter_file(temp)
comfy = Comfy(tmp_path / "input")
comfy.record_mutation = mutation
async with serve(comfy) as url:
with pytest.raises(ValueError):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
assert owned.exists()
assert not comfy.views
run(check())
@pytest.mark.parametrize("failure", ["required", "publish"])
def test_publication_failures_cleanup_identified_exports(tmp_path, monkeypatch, failure):
monkeypatch.setattr(handler, "MAX_RESPONSE_JSON_BYTES", 100_000)
async def check():
temp = tmp_path / "temp"
comfy = Comfy(tmp_path / "input", payload=b"Z" * 100_000)
owned = exporter_file(temp, payload=comfy.payload)
async def publish(metadata, payload):
assert owned.exists()
raise ValueError("publisher failed")
async with serve(comfy) as url:
with pytest.raises(ValueError, match="publisher"):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root,
temp_dir=temp, publish=publish if failure == "publish" else None)
assert not owned.exists()
run(check())
def test_repeated_cancel_drains_collection_before_cleanup(tmp_path, monkeypatch):
monkeypatch.setattr(handler, "MAX_RESPONSE_JSON_BYTES", 100_000)
async def check():
temp = tmp_path / "temp"
comfy = Comfy(tmp_path / "input", payload=b"Z" * 100_000)
owned = exporter_file(temp, payload=comfy.payload)
started, draining, release = asyncio.Event(), asyncio.Event(), asyncio.Event()
async def publish(metadata, payload):
started.set()
try:
await asyncio.Future()
finally:
draining.set()
await release.wait()
assert owned.exists(), "cleanup while publisher still reading"
async with serve(comfy) as url:
task = asyncio.create_task(handler.execute_media_request(comfy.data, comfy_url=url,
input_dir=comfy.root, temp_dir=temp, publish=publish))
await asyncio.wait_for(started.wait(), 1)
task.cancel()
await asyncio.wait_for(draining.wait(), 1)
task.cancel()
await asyncio.sleep(0)
assert owned.exists()
release.set()
with pytest.raises(asyncio.CancelledError):
await task
assert not owned.exists()
run(check())
@pytest.mark.parametrize("cancel_caller", [False, True])
def test_stubborn_publisher_teardown_is_bounded_and_retains_files(tmp_path, monkeypatch, cancel_caller):
monkeypatch.setattr(handler, "MAX_RESPONSE_JSON_BYTES", 100_000)
monkeypatch.setattr(handler, "DRAIN_GRACE_SECONDS", 0.03)
async def check():
temp = tmp_path / "temp"
comfy = Comfy(tmp_path / "input", payload=b"Z" * 100_000)
owned = exporter_file(temp, payload=comfy.payload)
started, draining, release = asyncio.Event(), asyncio.Event(), asyncio.Event()
async def publish(metadata, payload):
started.set()
try:
await asyncio.Future()
finally:
draining.set()
await release.wait()
assert owned.exists()
async with serve(comfy) as url:
task = asyncio.create_task(handler.execute_media_request(comfy.data, comfy_url=url,
input_dir=comfy.root, temp_dir=temp, publish=publish,
timeout_seconds=5 if cancel_caller else 0.15))
try:
await asyncio.wait_for(started.wait(), 1)
if cancel_caller:
task.cancel()
await asyncio.wait_for(draining.wait(), 1)
if cancel_caller:
task.cancel() # must not reset the grace or recancel the reader
done, _ = await asyncio.wait({task}, timeout=0.5)
assert task in done, "callback suppressed cancellation indefinitely"
with pytest.raises(handler.MediaExecutionUncertain, match="teardown"):
task.result()
assert comfy.terminal_seen and owned.exists() and comfy.staged()
assert handler._UNCERTAIN_WORK
pending = list(handler._UNCERTAIN_WORK)
release.set()
await asyncio.gather(*pending, return_exceptions=True)
await asyncio.sleep(0)
assert not handler._UNCERTAIN_WORK
assert owned.exists() and comfy.staged(), "late callback must not retire files"
finally:
release.set()
if not task.done():
task.cancel()
await asyncio.gather(task, return_exceptions=True)
run(check())
def test_export_replaced_during_collection_is_preserved(tmp_path):
async def check():
temp = tmp_path / "temp"
owned = exporter_file(temp)
comfy = Comfy(tmp_path / "input")
view = comfy.view
async def replace(req):
replacement = owned.with_suffix(".replacement")
replacement.write_bytes(b"unrelated replacement")
replacement.replace(owned)
return await view(req)
comfy.view = replace
async with serve(comfy) as url:
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
assert owned.read_bytes() == b"unrelated replacement"
run(check())
@pytest.mark.parametrize("status", ["success", "error"])
def test_partial_history_cleans_only_valid_identified_exporters(tmp_path, status):
async def check():
temp = tmp_path / "temp"
owned = exporter_file(temp)
unidentified = exporter_file(temp, name="c" * 32 + ".bin")
data = request()
second = dict(data["media"]["outputs"][0], binding_id="binding:2", node_id="export:2", collector_id="collector:2")
data["media"]["outputs"].append(second)
data["workflow"]["export:2"] = {"class_type": "MediaExportV1", "inputs": {
"request_id": "request:1", "binding_json": json.dumps(second), "texts": ["source:1", 0]}}
comfy = Comfy(tmp_path / "input", data)
comfy.mode = status
async with serve(comfy) as url:
with pytest.raises((ValueError, RuntimeError)):
await handler.execute_media_request(data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
assert not owned.exists()
assert unidentified.exists()
assert not comfy.views
run(check())
def test_missing_history_exports_remain_lifecycle_owned(tmp_path):
async def check():
temp = tmp_path / "temp"
owned = exporter_file(temp)
comfy = Comfy(tmp_path / "input")
comfy.mode = "missing"
async with serve(comfy) as url:
with pytest.raises(ValueError, match="Missing"):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
assert owned.exists()
assert not comfy.staged()
run(check())
@pytest.mark.parametrize("replacement", ["symlink", "directory", "hardlink"])
def test_nonregular_or_linked_export_files_are_preserved(tmp_path, replacement):
async def check():
temp = tmp_path / "temp"
owned = exporter_file(temp)
outside = tmp_path / "outside"
outside.write_bytes(b"unrelated")
if replacement == "hardlink":
alias = tmp_path / "alias"
alias.hardlink_to(owned)
else:
owned.unlink()
if replacement == "symlink":
owned.symlink_to(outside)
else:
owned.mkdir()
comfy = Comfy(tmp_path / "input")
async with serve(comfy) as url:
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
assert owned.exists()
assert outside.read_bytes() == b"unrelated"
run(check())
@pytest.mark.parametrize("shared", [False, True])
def test_direct_exporter_input_prunes_only_obsolete_loader_ancestors(tmp_path, shared):
async def check():
data = request("VIDEO")
data["workflow"]["source:1"] = {"class_type": "LoadVideo", "inputs": {"file": "unavailable.mp4"}}
data["media"]["inputs"] = [{"node_id": "export:1", "input_name": "video", "artifact_id": "input:1"}]
if shared:
second = dict(data["media"]["outputs"][0], binding_id="binding:2", node_id="export:2", collector_id="collector:2")
data["media"]["outputs"].append(second)
data["workflow"]["export:2"] = {"class_type": "MediaExportV1", "inputs": {
"request_id": "request:1", "binding_json": json.dumps(second), "video": ["source:1", 0]}}
fake_data = copy.deepcopy(data)
fake_data["media"].update(artifacts=[], inputs=[])
comfy = Comfy(tmp_path, fake_data)
history = comfy.history
async def full_history(req):
response = await history(req)
body = json.loads(response.text)
if shared:
record = comfy.record()
record["binding"] = dict(second, count=1)
record["artifacts"][0]["id"] = "b" * 32
record["artifacts"][0]["transport"]["filename"] = "b" * 32 + ".bin"
record["items"] = [{"binding_id": second["binding_id"], "item_index": 0, "artifact_id": "b" * 32}]
body["prompt:1"]["outputs"]["export:2"] = {"media_exports": [record]}
return web.json_response(body)
comfy.history = full_history
async with serve(comfy) as url:
result = await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path)
contract.validate_response(result["media"], request_id="request:1", bindings=data["media"]["outputs"])
submitted = comfy.submitted[0]["prompt"]
imported, slot = submitted["export:1"]["inputs"]["video"]
assert slot == 4
assert submitted[imported]["class_type"] == "MediaImportV1"
assert json.loads(submitted[imported]["inputs"]["artifact_json"])["kind"] == "VIDEO"
assert ("source:1" in submitted) is shared
assert set(submitted) == ({"export:1", imported, "source:1", "export:2"} if shared else {"export:1", imported})
assert not comfy.staged()
run(check())
@pytest.mark.parametrize("target,kind", [("request_id", "VIDEO"), ("binding_json", "VIDEO"), ("video", "STRING")])
def test_direct_exporter_input_rejects_controls_and_wrong_kind(tmp_path, target, kind):
async def check():
data = request("VIDEO")
data["media"]["inputs"] = [{"node_id": "export:1", "input_name": target, "artifact_id": "input:1"}]
data["media"]["artifacts"][0]["kind"] = kind
comfy = Comfy(tmp_path, data)
async with serve(comfy) as url:
with pytest.raises(ValueError):
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path)
assert not comfy.submitted
assert not comfy.staged()
run(check())
def test_slow_download_times_out_before_any_prompt(tmp_path):
async def check():
data = request()
data["media"]["artifacts"][0]["transport"] = {"type": "url", "url": "https://example.org/input", "expires_at": "2099-01-01T00:00:00Z"}
async def download(*args, **kwargs):
await asyncio.sleep(60)
comfy = Comfy(tmp_path, data)
async with serve(comfy) as url:
with pytest.raises(TimeoutError):
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path, download=download, timeout_seconds=0.02)
assert not comfy.submitted
assert not comfy.staged()
run(check())
+330
View File
@@ -0,0 +1,330 @@
"""Real CPU-only IMAGE transfers; no ComfyUI/bootstrap/provider imports."""
import importlib
import sys
import types
from pathlib import Path
import numpy as np
import pytest
import torch
from PIL import Image
# Load a private package, not Distributed's root __init__.
ROOT = Path(__file__).resolve().parents[1] / "serverless_bridge"
PACKAGE = "image_codec_tests_bridge"
pkg = types.ModuleType(PACKAGE)
pkg.__path__ = [str(ROOT)]
sys.modules[PACKAGE] = pkg
codec = importlib.import_module(PACKAGE + ".codec")
images = importlib.import_module(PACKAGE + ".image_codec")
FORMATS = ("legacy_png", "png", "jpeg", "webp", "ffv1", "h264", "h265", "av1")
LOSSLESS = ("legacy_png", "png", "ffv1")
@pytest.fixture(autouse=True)
def no_gpu(monkeypatch):
def forbidden(*a, **kw):
pytest.fail("GPU query or initialization is forbidden")
for name in ("is_available", "init", "_lazy_init", "device_count"):
monkeypatch.setattr(torch.cuda, name, forbidden)
def available(fmt):
if fmt in images._VIDEO_CODECS:
av = pytest.importorskip("av")
name = images._VIDEO_CODECS[fmt][0]
if name not in av.codecs_available:
pytest.skip("Software encoder unavailable: " + name)
def sample(channels=3, dtype=torch.float32, shape=(3, 65, 67)):
g = torch.Generator(device="cpu").manual_seed(42)
return torch.rand((*shape, channels), generator=g, device="cpu", dtype=dtype)
def quantized(value):
return torch.from_numpy(np.clip(value.float().numpy() * 255, 0, 255).astype(np.uint8)).float() / 255
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
def test_exact_unchanged(dtype):
value = sample(4, dtype, (2, 3, 5))
value.flatten()[0] = -0.0
expected = codec.encode_payload("IMAGE", value)
assert images.encode_image(value) == expected
result = images.decode_image(*expected)
assert torch.equal(result.view(torch.uint8), value.view(torch.uint8))
@pytest.mark.parametrize("fmt", FORMATS)
@pytest.mark.parametrize("channels", [1, 3, 4])
def test_real_roundtrips_shape_order_quantization_alpha(fmt, channels):
available(fmt)
value = sample(channels)
# Large inter-frame difference also verifies frame ordering.
value[0] *= 0.2
value[2] = value[2] * 0.2 + 0.8
if channels == 4:
value[..., 3] = torch.linspace(-0.1, 1.1, value[..., 3].numel()).reshape(value[..., 3].shape)
descriptor, payload = images.encode_image(value, transfer_format=fmt)
assert descriptor["encoding"] == "image-transfer-v1"
result = codec.decode_payload("IMAGE", descriptor, payload)
assert result.shape == value.shape
assert result.device.type == "cpu" and result.is_contiguous()
color = 3 if channels == 4 else channels
expected = quantized(value[..., :color])
if fmt in LOSSLESS:
assert torch.equal(result[..., :color], expected)
else:
# Random RGB noise loses chroma detail in 4:2:0 even at high quality.
assert (result[..., :color] - expected).abs().mean() < 0.12
assert result[0, ..., :color].mean() < result[2, ..., :color].mean()
if channels == 4:
assert torch.equal(result[..., 3], value[..., 3])
@pytest.mark.parametrize("fmt", ["png", "ffv1", "jpeg", "webp", "h264", "h265", "av1"])
def test_float64_alpha_is_exact(fmt):
available(fmt)
value = sample(4, torch.float64, (1, 65, 67))
descriptor, payload = images.encode_image(value, transfer_format=fmt)
out = images.decode_image(descriptor, payload)
assert out.dtype == torch.float64
assert torch.equal(out[..., 3].contiguous().view(torch.uint8), value[..., 3].contiguous().view(torch.uint8))
@pytest.mark.parametrize("fmt", ["jpeg", "webp", "h264", "h265", "av1"])
def test_lossy_quality_changes_final_payload_size(fmt):
available(fmt)
value = sample(shape=(2, 96, 96))
low = images.encode_image(value, transfer_format=fmt, transfer_quality=15)[1]
high = images.encode_image(value, transfer_format=fmt, transfer_quality=95)[1]
assert len(high) > len(low)
def test_truncation_clipping_grayscale_noncontiguous():
value = torch.tensor([-0.1, 0.5, 1.1, 0.12345], device="cpu").reshape(1, 2, 2, 1).transpose(1, 2)
out = images.decode_image(*images.encode_image(value, transfer_format="png"))
assert torch.equal(out, quantized(value))
assert out[0, 0, 1, 0] == 1
@pytest.mark.parametrize("value", [None, [], torch.ones(2, 3, 3), torch.ones(1, 2, 3, 2), torch.ones(1, 2, 3, 0), torch.ones(1, 2, 3, 3, dtype=torch.uint8), torch.ones(1, 2, 3, 3, device="meta"), torch.ones(1, 2, 3, 3).to_sparse()])
def test_invalid_values(value):
with pytest.raises(ValueError):
images.encode_image(value, transfer_format="png")
@pytest.mark.parametrize("bad", [float("nan"), float("inf"), -float("inf")])
def test_nonfinite_rejected(bad):
with pytest.raises(ValueError, match="finite"):
images.encode_image(torch.full((1, 1, 1, 3), bad), transfer_format="png")
@pytest.mark.parametrize("options", [{"transfer_format": "PNG"}, {"transfer_format": []}, {"transfer_quality": True}, {"transfer_quality": 0}, {"transfer_quality": 101}, {"transfer_quality": 1.5}, {"max_bytes": True}, {"max_bytes": -1}])
def test_strict_options(options):
with pytest.raises(ValueError):
images.encode_image(sample(shape=(1, 1, 1)), **options)
@pytest.mark.parametrize("change", [{"extra": 1}, {"format": "exact"}, {"quality": True}, {"layout": "BCHW"}, {"shape": [1, 2**60, 1, 3]}, {"shape": [True, 1, 1, 3]}, {"parts": [0]}, {"parts": [True]}, {"parts": [2**60]}, {"alpha": {"extra": 0}}])
def test_bad_descriptors(change):
descriptor, payload = images.encode_image(sample(shape=(1, 3, 5)), transfer_format="png")
descriptor.update(change)
with pytest.raises(ValueError):
images.decode_image(descriptor, payload)
def test_truncated_appended_and_mismatched_parts():
d, p = images.encode_image(sample(shape=(2, 3, 5)), transfer_format="png")
for payload in (p[:-1], p + b"x"):
with pytest.raises(ValueError):
images.decode_image(d, payload)
for change in ({"shape": [1, 3, 5, 3]}, {"parts": [len(p)]}, {"shape": [2, 2, 5, 3]}):
bad = dict(d, **change)
with pytest.raises(ValueError):
images.decode_image(bad, p)
@pytest.mark.parametrize("fmt", ["png", "jpeg", "webp", "ffv1", "h264", "h265", "av1"])
def test_payload_and_decoded_budget(fmt):
available(fmt)
value = sample(shape=(1, 65, 67))
d, p = images.encode_image(value, transfer_format=fmt)
with pytest.raises(ValueError, match="limit|budget"):
images.decode_image(d, p, max_bytes=100)
with pytest.raises(ValueError, match="limit|budget"):
images.encode_image(value, transfer_format=fmt, max_bytes=100)
def test_oversized_shape_before_materialization(monkeypatch):
value = torch.ones(1).expand(1, 8192, 8192, 4)
def forbidden(*a, **kw):
pytest.fail("Materialized an oversized tensor")
monkeypatch.setattr(torch.Tensor, "contiguous", forbidden)
monkeypatch.setattr(torch, "isfinite", forbidden)
with pytest.raises(ValueError):
images.encode_image(value, transfer_format="png")
def test_header_dimensions_checked_before_pixels(monkeypatch):
d, p = images.encode_image(sample(shape=(1, 9, 11)), transfer_format="png")
d["shape"] = [1, 1, 1, 3]
def forbidden(*a, **kw):
pytest.fail("Pixels decoded before header dimensions were checked")
monkeypatch.setattr(Image.Image, "load", forbidden)
with pytest.raises(ValueError):
images.decode_image(d, p)
def test_other_kinds_do_not_dispatch():
d, p = images.encode_image(sample(shape=(1, 3, 5)), transfer_format="png")
for kind in ("MASK", "AUDIO", "STRING"):
with pytest.raises(ValueError):
codec.decode_payload(kind, d, p)
def test_missing_video_encoder_is_explicit(monkeypatch):
monkeypatch.setitem(images._VIDEO_CODECS, "h264", ("missing_software_encoder", 51, "yuv420p"))
with pytest.raises(ValueError, match="encoder|unavailable"):
images.encode_image(sample(shape=(1, 64, 64)), transfer_format="h264")
def test_bounded_sink_catches_final_flush_bytes():
sink = images._BoundedIO(4)
sink.write(b"1234")
with pytest.raises(ValueError, match="limit"):
sink.write(b"5")
assert sink.getvalue() == b"1234"
def test_crf_mapping():
assert images._quality_to_crf(90, 51) == 13
assert images._quality_to_crf(90, 63) == 16
assert images._quality_to_crf(100, 63) == 0
@pytest.mark.parametrize("fmt", FORMATS)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
def test_compressed_preserves_source_dtype_and_exact_alpha(fmt, dtype):
available(fmt)
value = sample(4, dtype, (1, 17, 19))
value[..., 3].flatten()[0] = -0.0
d, p = images.encode_image(value, transfer_format=fmt)
with torch.device("meta"):
result = images.decode_image(d, p)
assert result.device.type == "cpu" and result.dtype == dtype
assert d["dtype"] == str(dtype).removeprefix("torch.")
assert torch.equal(result[..., 3].contiguous().view(torch.uint8), value[..., 3].contiguous().view(torch.uint8))
if fmt in LOSSLESS:
assert torch.equal(result[..., :3], quantized(value[..., :3]).to(dtype))
@pytest.mark.parametrize("fmt", ["ffv1", "h264", "h265", "av1"])
def test_video_count_and_header_fail_before_native_decoder(fmt, monkeypatch):
available(fmt)
d, p = images.encode_image(sample(shape=(2, 65, 67)), transfer_format=fmt)
def forbidden(*a, **kw):
pytest.fail("Native decoder reached before header validation")
monkeypatch.setattr(images, "_decode_video", forbidden)
for shape in ([1, 65, 67, 3], [3, 65, 67, 3], [2, 32, 32, 3]):
bad = dict(d, shape=shape)
with pytest.raises(ValueError):
images.decode_image(bad, p)
for payload in (p[:-1], p[:40], p + b"extra"):
bad = dict(d, parts=[len(payload)])
with pytest.raises(ValueError):
images.decode_image(bad, payload)
@pytest.mark.parametrize("fmt", ["ffv1", "h264", "h265", "av1"])
def test_video_does_not_use_native_container_probe(fmt, monkeypatch):
available(fmt)
d, p = images.encode_image(sample(shape=(2, 65, 67)), transfer_format=fmt)
av = images._import_av()
def forbidden(*a, **kw):
pytest.fail("Unbounded native stream probing is forbidden")
monkeypatch.setattr(av, "open", forbidden)
assert images.decode_image(d, p).shape == (2, 65, 67, 3)
@pytest.mark.parametrize("fmt", ["ffv1", "h264", "h265", "av1"])
def test_real_video_flush_bound(fmt):
available(fmt)
frames = [np.zeros((64, 64, 3), dtype=np.uint8)]
good = images._encode_video(iter(frames), fmt, 90, 64, 64, 100000)
with pytest.raises(ValueError, match="limit"):
images._encode_video(iter(frames), fmt, 90, 64, 64, len(good) - 1)
def test_alpha_invalid_dtype_nonfinite_and_extra_fields():
import copy
d, p = images.encode_image(sample(4, shape=(1, 3, 5)), transfer_format="png")
for field, value in (("dtype", "float64"), ("shape", [1, 3, 6, 1]), ("extra", 1)):
bad = copy.deepcopy(d)
bad["alpha"]["descriptor"][field] = value
with pytest.raises(ValueError):
images.decode_image(bad, p)
import struct
bad = p[:-4] + struct.pack("<f", float("nan"))
with pytest.raises(ValueError, match="finite"):
images.decode_image(d, bad)
def test_descriptor_structure_bound_and_frame_limit():
d, p = images.encode_image(sample(shape=(1, 3, 5)), transfer_format="png")
for bad in (dict(d, quality="x" * 70000), dict(d, parts=[1] * 241), dict(d, shape=[241, 1, 1, 3]), dict(d, dtype="x")):
with pytest.raises(ValueError):
images.decode_image(bad, p)
with pytest.raises(ValueError):
images.encode_image(torch.zeros(241, 1, 1, 3), transfer_format="png")
@pytest.mark.parametrize("fmt", ["png", "jpeg", "webp"])
def test_all_still_headers_checked_before_pillow_open(fmt, monkeypatch):
d, p = images.encode_image(sample(shape=(1, 31, 33)), transfer_format=fmt)
d["shape"] = [1, 2, 2, 3]
def forbidden(*a, **kw):
pytest.fail("Pillow opened a mismatched header")
monkeypatch.setattr(Image, "open", forbidden)
with pytest.raises(ValueError, match="dimensions"):
images.decode_image(d, p)
@pytest.mark.parametrize("fmt", ["ffv1", "h264", "h265", "av1"])
def test_video_longer_batch_order_and_coded_allocation_bound(fmt):
available(fmt)
value = torch.linspace(0, 1, 16).reshape(16, 1, 1, 1).expand(16, 64, 64, 3)
d, p = images.encode_image(value, transfer_format=fmt)
header = images._video_header(p, fmt, 64, 64, 16)
coded = images._coded_frame_bound(p, fmt, header)
assert 16 <= coded <= images.MAX_FRAMES * 4
result = images.decode_image(d, p)
assert torch.all((result.mean(dim=(1, 2, 3)) - value.mean(dim=(1, 2, 3))).abs() < 0.03)
@pytest.mark.parametrize("fmt", ["h264", "h265", "av1"])
def test_packet_unit_bounds_before_native_decode(fmt):
available(fmt)
d, p = images.encode_image(sample(shape=(1, 64, 64)), transfer_format=fmt)
config, bounds = images._video_header(p, fmt, 64, 64, 1)
start, end = bounds[0]
packet = p[start:end]
with pytest.raises(ValueError):
images._coded_frame_bound(packet[:-1], fmt, (config, [(0, len(packet) - 1)]))
# An adversarial single block can contain many coded frames: bound it too.
enormous = packet * (images.MAX_FRAMES * 4 + 1)
with pytest.raises(ValueError, match="limit"):
images._coded_frame_bound(enormous, fmt, (config, [(0, len(enormous))]))
@pytest.mark.parametrize("fmt", ["h264", "h265", "av1"])
def test_bitstream_resolution_cannot_bypass_native_pixel_limit(fmt, monkeypatch):
available(fmt)
d, p = images.encode_image(sample(shape=(1, 256, 256)), transfer_format=fmt)
header = images._video_header(p, fmt, 256, 256, 1)
av = images._import_av()
# A lying container is not enough to disable max_pixels in the decoder.
# The bitstream still contains 256x256 pixels but the budget permits 64x64.
with pytest.raises((ValueError, av.error.FFmpegError)):
list(images._decode_video(p, fmt, 64, 64, 1, header))
+83
View File
@@ -0,0 +1,83 @@
"""Participant identity, typed compilation and legacy Collector pruning."""
import copy
import importlib
import json
import sys
import types
from pathlib import Path
import pytest
ROOT = Path(__file__).resolve().parents[1]
PACKAGE = "media_transform_testpkg"
for suffix in ("", ".api", ".api.orchestration", ".utils", ".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
pt = importlib.import_module(PACKAGE + ".api.orchestration.prompt_transform")
compiler = importlib.import_module(PACKAGE + ".api.orchestration.serverless_media_prompt")
SOCKETS = ("images", "audio")
def prompt(sockets=SOCKETS):
nodes = {
"source": {"class_type": "MediaSource", "inputs": {}},
"collector": {"class_type": "DistributedCollector", "inputs": {
socket: ["source", SOCKETS.index(socket)] for socket in sockets}},
}
for socket in sockets:
nodes["sink:" + socket] = {"class_type": "Capture", "inputs": {"value": ["collector", SOCKETS.index(socket)]}}
return nodes
def prepare(graph, participant, delegate):
index = pt.PromptIndex(graph)
jobs = pt.generate_job_id_map(index, "run")
assert jobs == {"collector": "run_collector"}
return pt.apply_participant_overrides(index.copy_prompt(), participant,
["serverless:pool:1", "serverless:pool:2"], jobs, "http://master", delegate, index)
def test_typed_worker_preparation_compiles_both_bindings_without_source_mutation():
graph = prompt()
before = copy.deepcopy(graph)
worker = prepare(graph, "serverless:pool:2", True)
compiled, bindings = compiler.compile_media_prompt(worker, "request:2")
assert len(bindings) == 2
assert {row["participant_id"] for row in bindings} == {"serverless:pool:2"}
assert {row["job_id"] for row in bindings} == {"run_collector"}
assert {row["socket"] for row in bindings} == set(SOCKETS)
assert set(compiled) == {"source"} | {row["node_id"] for row in bindings}
assert graph == before
@pytest.mark.parametrize("sockets", [("images",), ("audio",), SOCKETS])
def test_delegate_master_keeps_collector_socket_contract_and_legacy_placeholder(sockets):
graph = prompt(sockets)
prepared = prepare(graph, "master", True)
before = copy.deepcopy(prepared)
delegated = pt.prepare_delegate_master_prompt(prepared, ["collector"])
assert "source" not in delegated
assert {"collector"} | {"sink:" + s for s in sockets} <= set(delegated)
inputs = delegated["collector"]["inputs"]
assert "audio" not in inputs and "expected_media" not in inputs
placeholders = [key for key, node in delegated.items() if node['class_type'] == 'DistributedEmptyImage']
if "images" in sockets:
assert len(placeholders) == 1
assert inputs['images'] == [placeholders[0], 0]
else:
assert not placeholders and 'images' not in inputs
assert inputs["delegate_only"] is True
assert json.loads(inputs["enabled_worker_ids"]) == ["serverless:pool:1", "serverless:pool:2"]
assert prepared == before
@pytest.mark.parametrize("socket", SOCKETS)
def test_participating_master_retains_exact_media_links(socket):
graph = prompt()
original = copy.deepcopy(graph)
master = prepare(graph, "master", False)
assert master["collector"]["inputs"][socket] == graph["collector"]["inputs"][socket]
assert master["collector"]["inputs"]["delegate_only"] is False
assert "worker_id" not in master["collector"]["inputs"]
assert graph == original
+552
View File
@@ -0,0 +1,552 @@
"""Offline queue integration: real compilers, packers and wire validators.
Only host dispatch/probing/queueing and the manager lifecycle are mocked. No
provider requests, services or GPU work are permitted by these tests.
"""
import asyncio
import base64
import copy
import importlib
import json
import os
from pathlib import Path
import socket
import subprocess
import sys
import types
import uuid
import wave
from unittest.mock import AsyncMock, Mock, patch
import aiohttp
import pytest
# Preload before the shared integration loader snapshots sys.modules: extension
# modules cannot safely be unloaded and imported again between fixture contexts.
import torch
from PIL import Image
from test_serverless_integration import load_modules, module, PKG
PROFILE = "comfy-media-v1"
def node(class_type, **inputs):
return {"class_type": class_type, "inputs": inputs}
def collector_prompt():
return {
"source": node("EmptyImage", width=32, height=32, batch_size=1, color=0),
"collector": node("DistributedCollector", images=["source", 0], load_balance=False),
"save": node("SaveImage", images=["collector", 0], filename_prefix="local"),
}
def legacy_prompt():
return {
"source": node("EmptyImage", width=32, height=32, batch_size=1, color=0),
"collector": node("DistributedCollector", images=["source", 0]),
"save": node("SaveImage", images=["collector", 0], filename_prefix="local"),
}
@pytest.fixture
def harness(monkeypatch):
q, modules, ps = load_modules()
with patch.dict(sys.modules, modules):
pool = {"id": "pool", "name": "Media", "type": "serverless", "provider": "runpod",
"endpoint_id": "offline-endpoint", "enabled": True, "jobs_per_run": 2,
"timeout_seconds": 30, "media_profile": PROFILE}
config = {"workers": [pool], "settings": {"master_delegate_only": True}}
q.load_config = lambda: config
q.select_active_workers = AsyncMock(side_effect=lambda workers, _ws, delegate, **kw: (workers, delegate))
q.queue_prompt_payload = AsyncMock(return_value={"prompt_id": "master-prompt", "number": 7, "node_errors": {}})
manager = types.SimpleNamespace(start=AsyncMock(), cancel=AsyncMock(), ensure_capacity=Mock())
sys.modules[PKG + ".api.orchestration.serverless_runtime"] = module("runtime", get_manager=lambda: manager)
# These are the actual modules, not permissive serializer/packing stubs.
client = importlib.import_module(PKG + ".api.orchestration.runpod_client")
compiler = importlib.import_module(PKG + ".api.orchestration.serverless_media_prompt")
packer = importlib.import_module(PKG + ".api.orchestration.serverless_typed_inputs")
validator = importlib.import_module(PKG + ".api.orchestration.serverless_typed_results")
handler = importlib.import_module(PKG + ".serverless_bridge.handler")
def forbidden(*args, **kwargs):
raise AssertionError("Offline queue tests must not access a network or GPU")
monkeypatch.setattr(aiohttp, "ClientSession", forbidden)
monkeypatch.setattr(socket, "create_connection", forbidden)
monkeypatch.setattr(socket.socket, "connect", forbidden)
monkeypatch.setattr(socket.socket, "connect_ex", forbidden)
monkeypatch.setattr(torch.cuda, "init", forbidden)
monkeypatch.setattr(torch.cuda, "current_device", forbidden)
h = types.SimpleNamespace(q=q, ps=ps, pool=pool, config=config, manager=manager,
client=client, compiler=compiler, packer=packer,
validator=validator, handler=handler)
yield h
@pytest.fixture
def media_paths(tmp_path, monkeypatch):
paths = types.SimpleNamespace(get_annotated_filepath=lambda name: str(tmp_path / name))
for kind in ("input", "output", "temp"):
setattr(paths, "get_" + kind + "_directory", lambda: str(tmp_path))
monkeypatch.setitem(sys.modules, "folder_paths", paths)
return tmp_path
def run(h, prompt=None, **kwargs):
return asyncio.run(h.q.orchestrate_distributed_execution(
prompt if prompt is not None else collector_prompt(), None, "client", **kwargs))
def assert_no_dispatch(h, *, no_probe=False):
h.q.queue_prompt_payload.assert_not_awaited()
h.q.dispatch_worker_prompt.assert_not_awaited()
h.manager.start.assert_not_awaited()
h.manager.cancel.assert_not_awaited()
assert h.ps.distributed_pending_jobs == {}
if no_probe:
h.q.select_active_workers.assert_not_awaited()
h.q.select_least_busy_worker.assert_not_awaited()
h.q.fetch_worker_has_node.assert_not_awaited()
h.q.sync_worker_media.assert_not_awaited()
@pytest.mark.parametrize("delegate", [True, False])
def test_real_compile_pack_serialize_and_canonical_bindings(harness, media_paths, delegate):
h = harness
image_path = media_paths / "picture.png"
Image.new("RGBA", (2, 3), color=(17, 29, 41, 53)).save(image_path)
image_bytes = image_path.read_bytes()
prompt = collector_prompt()
prompt["image"] = node("LoadImage", image="picture.png")
prompt["audio"] = node("GeneratedAudio")
prompt["collector"]["inputs"].update(images=["image", 0], audio=["audio", 0])
prompt["collector2"] = node("DistributedCollector", audio=["audio", 0])
prompt["unused"] = node("PreviewImage", images=["image", 0])
source = copy.deepcopy(prompt)
with patch.object(h.compiler, "compile_media_prompt", wraps=h.compiler.compile_media_prompt) as compile_call, \
patch.object(h.packer, "prepare_media_inputs", wraps=h.packer.prepare_media_inputs) as pack_call, \
patch.object(h.client, "serialize_request", wraps=h.client.serialize_request) as serialize_call, \
patch.object(h.q, "prune_prompt_for_worker", side_effect=AssertionError("typed graph must bypass legacy pruning")):
result = run(h, prompt, delegate_master=delegate)
assert result == ("master-prompt", 7, 2, {})
assert compile_call.call_count == pack_call.await_count == 2
assert serialize_call.call_count >= 2
assert prompt == source
h.q.dispatch_worker_prompt.assert_not_awaited()
jobs = h.manager.start.call_args.args[2]
assert [j["worker"]["id"] for j in jobs] == ["serverless:pool:1", "serverless:pool:2"]
assert len({j["media"]["request_id"] for j in jobs}) == 2
master = h.q.queue_prompt_payload.call_args.args[0]
for job in jobs:
media = job["media"]
assert str(uuid.UUID(media["request_id"])) == media["request_id"]
assert uuid.UUID(media["request_id"]).version == 4
assert media["profile"] == PROFILE and media["version"] == 1
assert media["outputs"] == job["manifest"]
assert media["outputs"] is not job["manifest"]
assert media["artifacts"] == media["inputs"] == []
assert {b["socket"] for b in job["manifest"]} == {"images", "audio"}
assert len(job["manifest"]) == 3
assert base64.b64decode(job["images"][0]["image"]) == image_bytes
assert job["workflow"]["image"]["inputs"]["image"] == job["images"][0]["name"]
assert not ({"save", "unused", "collector", "collector2"} & set(job["workflow"]))
body = h.client.serialize_request(job["workflow"], job["images"], 30, media=media)
wire = json.loads(body)
assert wire["input"] == {"workflow": job["workflow"], "images": job["images"], "media": media}
assert wire["policy"] == {"executionTimeout": 30000, "ttl": 30000}
_, _, _, targets = h.handler._validate_request(wire["input"])
assert targets == [b["node_id"] for b in job["manifest"]]
assert h.validator.validate_prepared_job(job) is True
for binding in job["manifest"]:
assert binding["participant_id"] == job["worker"]["id"]
assert binding["job_id"] == master[binding["collector_id"]]["inputs"]["multi_job_id"]
controls = job["workflow"][binding["node_id"]]["inputs"]
assert json.loads(controls["binding_json"]) == binding
assert controls["request_id"] == media["request_id"]
assert master["collector"]["inputs"]["delegate_only"] is delegate
assert "expected_media" not in master["collector"]["inputs"]
if delegate:
assert not {"source", "audio"} & set(master)
# LoadImage is an existing explicitly retained lightweight source; the
# Collector is nevertheless detached from it and uses the placeholder.
assert master["image"] == prompt["image"]
assert "audio" not in master["collector"]["inputs"]
placeholder = master["collector"]["inputs"]["images"][0]
assert master[placeholder]["class_type"] == "DistributedEmptyImage"
else:
assert master["collector"]["inputs"]["audio"] == ["audio", 0]
assert "audio" in master and "image" in master
@pytest.mark.parametrize("audio_only", [False, True])
def test_delegate_prunes_image_audio_sources_and_audio_only_needs_no_placeholder(harness, audio_only):
prompt = collector_prompt()
prompt["audio"] = node("GeneratedAudio")
prompt["collector"]["inputs"]["audio"] = ["audio", 0]
if audio_only:
del prompt["collector"]["inputs"]["images"]
prompt["save"] = node("AudioSink", audio=["collector", 1])
original = copy.deepcopy(prompt)
run(harness, prompt)
master = harness.q.queue_prompt_payload.call_args.args[0]
inputs = master["collector"]["inputs"]
assert "audio" not in inputs and "expected_media" not in inputs
assert "source" not in master and "audio" not in master
if audio_only:
assert "images" not in inputs and set(master) == {"collector", "save"}
else:
assert master[inputs["images"][0]]["class_type"] == "DistributedEmptyImage"
for job in harness.manager.start.call_args.args[2]:
assert {b["socket"] for b in job["manifest"]} == ({"audio"} if audio_only else {"images", "audio"})
assert "audio" in job["workflow"]
assert ("source" in job["workflow"]) is not audio_only
assert harness.validator.validate_prepared_job(job) is True
assert prompt == original
@pytest.mark.parametrize("selected", ["master", "worker"])
def test_typed_load_balance_selects_master_or_one_serverless_slot(harness, selected):
h = harness
prompt = collector_prompt()
prompt["collector"]["inputs"]["load_balance"] = "true"
h.q.select_least_busy_worker.side_effect = lambda candidates, **kw: next(
w for w in candidates if (w["id"] == "master") == (selected == "master"))
result = run(h, prompt, delegate_master=False)
h.q.select_least_busy_worker.assert_awaited_once()
master = h.q.queue_prompt_payload.call_args.args[0]
if selected == "master":
assert result[2] == 0
h.manager.start.assert_not_awaited()
assert master["collector"]["inputs"]["delegate_only"] is False
assert "source" in master
else:
assert result[2] == 1
assert len(h.manager.start.call_args.args[2]) == 1
assert master["collector"]["inputs"]["delegate_only"] is True
assert "source" not in master
@pytest.mark.parametrize("profile", [None, "legacy-image-v1"])
def test_exact_rejects_legacy_pool_before_probe(harness, profile):
h = harness
prompt = collector_prompt()
prompt["collector"]["inputs"]["transfer_format"] = "exact"
if profile is None:
h.pool.pop("media_profile")
else:
h.pool["media_profile"] = profile
with pytest.raises(ValueError, match="exact requires comfy-media-v1"):
run(h, prompt)
assert_no_dispatch(h, no_probe=True)
@pytest.mark.parametrize("profile", [None, "auto", "comfy-media-v2", False])
def test_invalid_explicit_profiles_fail_closed(harness, profile):
harness.pool["media_profile"] = profile
with pytest.raises(ValueError, match="media_profile"):
run(harness)
assert_no_dispatch(harness, no_probe=True)
@pytest.mark.parametrize("direct_type", ["local", "remote", "cloud"])
@pytest.mark.parametrize("explicit", [False, True])
def test_direct_workers_participate_alongside_typed_pools(harness, direct_type, explicit):
h = harness
h.config["workers"].append({"id": "direct", "type": direct_type, "host": "offline.invalid",
"enabled": not explicit, "port": 8188})
kwargs = {"enabled_worker_ids": ["pool", "direct"]} if explicit else {}
prompt = collector_prompt()
original = copy.deepcopy(prompt)
assert run(h, prompt, delegate_master=False, **kwargs)[2] == 3
assert prompt == original
h.q.dispatch_worker_prompt.assert_awaited_once()
master = h.q.queue_prompt_payload.call_args.args[0]
assert json.loads(master["collector"]["inputs"]["enabled_worker_ids"]) == [
"serverless:pool:1", "serverless:pool:2", "direct"]
jobs = h.manager.start.call_args.args[2]
assert len(jobs) == 2 and all(j["media"]["profile"] == PROFILE for j in jobs)
if direct_type != "local":
h.q.sync_worker_media.assert_awaited_once()
def test_disabled_or_unselected_workers_do_not_participate(harness):
h = harness
h.config["workers"] += [
{"id": "direct", "type": "local", "enabled": False},
dict(h.pool, id="legacy", media_profile="legacy-image-v1", enabled=False),
]
assert run(h)[2] == 2
h.pool["enabled"] = False
assert run(h, enabled_worker_ids=["pool"])[2] == 2
@pytest.mark.parametrize("workers", [True, False])
def test_image_list_cannot_mix_with_collector_even_master_only(harness, workers):
if not workers:
harness.config["workers"] = []
prompt = collector_prompt()
prompt["incompatible"] = node("DistributedImageListStart")
with pytest.raises(ValueError, match="cannot be mixed"):
run(harness, prompt, delegate_master=False)
assert_no_dispatch(harness, no_probe=True)
def test_mixed_direct_typed_and_legacy_profiles_keep_each_wire_contract(harness):
h = harness
h.config["workers"] += [dict(h.pool, id="legacy", media_profile="legacy-image-v1"),
{"id": "direct", "type": "local", "enabled": True}]
assert run(h, collector_prompt())[2] == 5
h.q.dispatch_worker_prompt.assert_awaited_once()
jobs = h.manager.start.call_args.args[2]
assert len(jobs) == 4
for job in jobs:
if job["worker"]["media_profile"] == PROFILE:
assert job["media"]["outputs"] == job["manifest"]
assert h.validator.validate_prepared_job(job) is True
else:
assert "media" not in job
assert any(n["class_type"] == "SaveImage" for n in job["workflow"].values())
def test_typed_master_only_is_supported_without_serverless_work(harness):
h = harness
h.config["workers"] = []
assert run(h, delegate_master=False)[2] == 0
master = h.q.queue_prompt_payload.call_args.args[0]
assert master["collector"]["inputs"]["enabled_worker_ids"] == "[]"
assert master["collector"]["inputs"]["images"] == ["source", 0]
h.manager.start.assert_not_awaited()
h.q.dispatch_worker_prompt.assert_not_awaited()
@pytest.mark.parametrize("offline", [False, True])
def test_collector_delegate_without_workers_preserves_master_fallback(harness, offline):
h = harness
if offline:
h.q.select_active_workers.side_effect = lambda *a, **kw: ([], False)
else:
h.config["workers"] = []
assert run(h)[2] == 0
master = h.q.queue_prompt_payload.call_args.args[0]
assert master["collector"]["inputs"]["delegate_only"] is False
assert master["collector"]["inputs"]["images"] == ["source", 0]
h.manager.start.assert_not_awaited()
@pytest.mark.parametrize("bad", [" pool ", "pool white", "pool\n"])
def test_pool_media_identity_is_not_repaired(harness, bad):
harness.pool["id"] = bad
with pytest.raises(ValueError, match="identifier"):
run(harness)
assert_no_dispatch(harness, no_probe=True)
def test_integer_source_node_identity_is_not_normalized(harness):
prompt = collector_prompt()
prompt[1] = prompt.pop("source")
prompt["collector"]["inputs"]["images"] = [1, 0]
with pytest.raises(ValueError, match="identifier"):
run(harness, prompt)
assert_no_dispatch(harness, no_probe=True)
@pytest.mark.parametrize("broken", ["dangling-unused", "literal", "missing-file", "bad-loader-controls", "oversize"])
def test_compilation_packing_and_size_errors_precede_all_dispatch(harness, media_paths, broken):
h = harness
prompt = collector_prompt()
if broken == "dangling-unused":
prompt["unused"] = node("PreviewImage", images=["missing", 0])
elif broken == "literal":
prompt["collector"]["inputs"]["images"] = "must use a linked image node"
elif broken in ("missing-file", "bad-loader-controls"):
prompt["audio"] = node("LoadAudio", audio="missing.wav")
prompt["collector"]["inputs"]["audio"] = ["audio", 0]
if broken == "bad-loader-controls":
prompt["audio"]["inputs"]["duration"] = 5
else:
prompt["source"]["inputs"]["label"] = "雪" * 1000
original = copy.deepcopy(prompt)
with patch.object(h.client, "MAX_REQUEST_BYTES", 100 if broken == "oversize" else h.client.MAX_REQUEST_BYTES):
with pytest.raises((ValueError, OSError)):
run(h, prompt)
assert prompt == original
assert_no_dispatch(h)
@pytest.mark.parametrize("corruption", ["manifest", "participant", "request-control", "binding-control", "profile", "input-ref"])
def test_whole_batch_identity_preflight_before_master_and_manager(harness, corruption):
h = harness
real_prepare = h.q._prepare_worker_payload
async def corrupt_later(*args, **kwargs):
worker, prepared = await real_prepare(*args, **kwargs)
if worker["id"].endswith(":2"):
binding = prepared["manifest"][0]
if corruption == "manifest":
binding["job_id"] = "wrong-job"
elif corruption == "participant":
prepared["worker"] = dict(worker, id="different-participant")
elif corruption == "request-control":
prepared["workflow"][binding["node_id"]]["inputs"]["request_id"] = "different-request"
elif corruption == "binding-control":
prepared["workflow"][binding["node_id"]]["inputs"]["binding_json"] = "{}"
elif corruption == "profile":
prepared["worker"] = dict(worker, media_profile="legacy-image-v1")
else:
prepared["media"]["inputs"] = [{"node_id": "missing", "input_name": "texts", "artifact_id": "unknown"}]
return worker, prepared
with patch.object(h.q, "_prepare_worker_payload", side_effect=corrupt_later):
with pytest.raises(ValueError):
run(h)
h.manager.ensure_capacity.assert_not_called()
assert_no_dispatch(h)
def test_batch_preflight_also_precedes_regular_direct_dispatch(harness):
h = harness
h.pool.pop("media_profile")
h.config["workers"].append({"id": "direct", "type": "local", "enabled": True})
real_prepare = h.q._prepare_worker_payload
async def corrupt_later(*args, **kwargs):
worker, prepared = await real_prepare(*args, **kwargs)
if worker["id"] == "serverless:pool:2":
prepared["worker"] = dict(worker, media_profile=PROFILE)
return worker, prepared
with patch.object(h.q, "_prepare_worker_payload", side_effect=corrupt_later):
with pytest.raises(ValueError, match="mismatch"):
run(h, legacy_prompt())
assert_no_dispatch(h)
def test_capacity_preflight_precedes_collector_state_and_master(harness):
h = harness
h.manager.ensure_capacity.side_effect = ValueError("capacity exhausted")
with pytest.raises(ValueError, match="capacity"):
run(h)
assert_no_dispatch(h)
def test_master_queue_failure_cleans_queues_and_never_starts_manager(harness):
h = harness
h.q.queue_prompt_payload.side_effect = ValueError("invalid master graph")
with pytest.raises(ValueError, match="invalid master"):
run(h)
assert h.ps.distributed_pending_jobs == {}
h.manager.start.assert_not_awaited()
h.q.dispatch_worker_prompt.assert_not_awaited()
def test_manager_setup_failure_emits_one_terminal_per_collector_participant(harness):
h = harness
prompt = collector_prompt()
prompt["audio"] = node("GeneratedAudio")
prompt["collector"]["inputs"]["audio"] = ["audio", 0]
h.manager.start.side_effect = ValueError("journal failure")
with pytest.raises(ValueError, match="journal failure"):
run(h, prompt)
h.manager.cancel.assert_awaited_once()
queued = next(iter(h.ps.distributed_pending_jobs.values()))
assert queued.qsize() == 2 # two participants, not four per-socket completions
messages = [queued.get_nowait(), queued.get_nowait()]
assert {m["worker_id"] for m in messages} == {"serverless:pool:1", "serverless:pool:2"}
assert all(m["is_last"] is True and m["serverless_error"] for m in messages)
@pytest.mark.parametrize("profile", [None, "legacy-image-v1"])
def test_legacy_envelope_and_compiler_are_unchanged(harness, profile):
h = harness
if profile is None:
h.pool.pop("media_profile")
else:
h.pool["media_profile"] = profile
prompt = legacy_prompt()
original = copy.deepcopy(prompt)
with patch.object(h.compiler, "compile_media_prompt", side_effect=AssertionError("legacy must not use typed compiler")), \
patch.object(h.packer, "prepare_media_inputs", side_effect=AssertionError("legacy must not use typed packer")):
assert run(h, prompt)[2] == 2
for job in h.manager.start.call_args.args[2]:
assert "media" not in job
assert any(n["class_type"] == "SaveImage" for n in job["workflow"].values())
expected = {"input": {"workflow": job["workflow"], "images": job["images"]},
"policy": {"executionTimeout": 30000, "ttl": 30000}}
assert h.client.serialize_request(job["workflow"], job["images"], 30) == json.dumps(
expected, allow_nan=False, ensure_ascii=False, separators=(",", ":")).encode()
assert prompt == original
def test_legacy_no_worker_delegate_fallback_is_preserved(harness):
h = harness
h.config["workers"] = []
assert run(h, legacy_prompt())[2] == 0
assert h.q.queue_prompt_payload.call_args.args[0]["collector"]["inputs"]["delegate_only"] is False
def test_native_audio_compile_pack_serialized_input_cpu(harness, media_paths):
"""Exercise real PyAV + core PCM conversion, not a mocked packing result."""
if os.environ.get("_MEDIA_QUEUE_CORE_TEST") != "1":
core = os.environ.get("COMFYUI_SOURCE_ROOT")
if not core:
pytest.skip("Set COMFYUI_SOURCE_ROOT for CPU native audio queue integration")
env = dict(os.environ)
env.pop("PYTHONPATH", None)
env.pop("COMFYUI_SOURCE_ROOT", None)
scratch = Path.home() / ".hermes/cache/scratch"
scratch.mkdir(parents=True, exist_ok=True)
env.update(_MEDIA_QUEUE_CORE_TEST="1", CUDA_VISIBLE_DEVICES="", HIP_VISIBLE_DEVICES="",
PYTHONDONTWRITEBYTECODE="1", TMPDIR=str(scratch))
bootstrap = """
import sys
sys.path.insert(0, sys.argv[1])
import comfy.cli_args
comfy.cli_args.args.cpu = True
import torch
def forbidden(*args, **kwargs):
raise AssertionError('CPU audio integration may not initialize a GPU')
torch.cuda.is_available = lambda: False
torch.cuda.device_count = lambda: 0
torch.cuda.init = forbidden
torch.cuda.current_device = forbidden
torch.cuda._lazy_init = forbidden
import pytest
raise SystemExit(pytest.main([sys.argv[2], '-q', '-p', 'no:cacheprovider', '-k',
'test_native_audio_compile_pack_serialized_input_cpu']))
"""
result = subprocess.run([sys.executable, "-c", bootstrap, core, str(Path(__file__).resolve())],
env=env, capture_output=True, text=True, timeout=120)
assert result.returncode == 0, result.stdout + result.stderr
assert "1 passed" in result.stdout
return
h = harness
h.pool["jobs_per_run"] = 1
path = media_paths / "tone.wav"
pcm = torch.tensor([-100, 0, 100, 32767, -32768, 7], dtype=torch.int16)
with wave.open(str(path), "wb") as stream:
stream.setnchannels(1)
stream.setsampwidth(2)
stream.setframerate(8000)
stream.writeframes(pcm.numpy().tobytes())
original_bytes = path.read_bytes()
prompt = collector_prompt()
prompt["audio"] = node("LoadAudio", audio="tone.wav")
prompt["collector"]["inputs"]["audio"] = ["audio", 0]
original = copy.deepcopy(prompt)
run(h, prompt)
job = h.manager.start.call_args.args[2][0]
wire = json.loads(h.client.serialize_request(job["workflow"], job["images"], 30, media=job["media"]))
artifacts = wire["input"]["media"]["artifacts"]
refs = wire["input"]["media"]["inputs"]
assert len(artifacts) == len(refs) == 1
binding = next(b for b in job["manifest"] if b["socket"] == "audio")
assert refs == [{"node_id": binding["node_id"], "input_name": "audio", "artifact_id": artifacts[0]["id"]}]
contract = importlib.import_module(PKG + ".serverless_bridge.contract")
codec = importlib.import_module(PKG + ".serverless_bridge.codec")
payload = asyncio.run(contract.resolve_artifact(artifacts[0]))
value = codec.decode_payload("AUDIO", artifacts[0]["descriptor"], payload)
assert value["sample_rate"] == 8000
assert value["waveform"].shape == (1, 1, 6)
assert torch.equal(value["waveform"], (pcm.to(torch.float32) / 32768).reshape(1, 1, -1))
assert path.read_bytes() == original_bytes
assert prompt == original
assert h.validator.validate_prepared_job(job) is True
+96
View File
@@ -0,0 +1,96 @@
"""Collector controls -> compiler -> endpoint validation, no network or GPU."""
import copy
import json
import pytest
from test_serverless_media_prompt import compiler, _prompt
from test_media_handler import handler
from test_collector_typed_media import collector_module, native_media, assert_exact
FORMATS = ('legacy_png', 'png', 'ffv1', 'jpeg', 'webp', 'h264', 'h265', 'av1', 'exact')
def compile_request(fmt='legacy_png', quality=90):
prompt = _prompt()
prompt['source:1'] = {'class_type': 'Source', 'inputs': {}}
prompt['collector:1']['inputs'].update(transfer_format=fmt, transfer_quality=quality,
audio=['source:1', 1])
original = copy.deepcopy(prompt)
workflow, outputs = compiler.compile_media_prompt(prompt, 'request:1')
assert prompt == original
return {'workflow': workflow, 'images': [], 'media': {
'version': 1, 'profile': 'comfy-media-v1', 'request_id': 'request:1',
'artifacts': [], 'inputs': [], 'outputs': outputs}}
@pytest.mark.parametrize('fmt', FORMATS)
def test_compiler_and_handler_preserve_image_policy_only(fmt):
data = compile_request(fmt, 73)
assert {b['socket'] for b in data['media']['outputs']} == {'images', 'audio'}
for binding in data['media']['outputs']:
inputs = data['workflow'][binding['node_id']]['inputs']
if binding['kind'] == 'IMAGE' and fmt != 'exact':
assert inputs['transfer_format'] == ('png' if fmt == 'legacy_png' else fmt)
assert inputs['transfer_quality'] == 73
else:
assert 'transfer_format' not in inputs and 'transfer_quality' not in inputs
handler._validate_request(data)
def test_saved_collector_defaults_compile_to_png_not_exact():
workflow, bindings = compiler.compile_media_prompt(_prompt(), 'request:1')
inputs = workflow[bindings[0]['node_id']]['inputs']
assert (inputs['transfer_format'], inputs['transfer_quality']) == ('png', 90)
@pytest.mark.parametrize('fmt,quality', [('JPEG',90),('',90),('jpeg',True),('jpeg',0),
('jpeg',101),('jpeg','90'),(['source:1',0],90)])
def test_compiler_rejects_invalid_or_linked_controls(fmt, quality):
with pytest.raises(ValueError):
compile_request(fmt, quality)
@pytest.mark.parametrize('field,value', [('transfer_format','unknown'),('transfer_quality',0),
('transfer_quality',True),('transfer_quality','90')])
def test_handler_rejects_tampered_policy(field, value):
data = compile_request('jpeg')
binding = next(b for b in data['media']['outputs'] if b['kind'] == 'IMAGE')
data['workflow'][binding['node_id']]['inputs'][field] = value
with pytest.raises(ValueError):
handler._validate_request(data)
@pytest.mark.parametrize('socket,kind', [('audio', 'AUDIO'), ('masks', 'MASK'), ('texts', 'STRING'), ('video', 'VIDEO')])
def test_handler_rejects_image_controls_on_nonimage_bridge_exports(socket, kind):
# The bridge still supports all five kinds; these are direct bridge exports,
# not nonexistent sockets on DistributedCollector.
data = compile_request('exact')
binding = next(b for b in data['media']['outputs'] if b['kind'] == 'AUDIO')
binding.update(socket=socket, kind=kind)
inputs = data['workflow'][binding['node_id']]['inputs']
inputs.pop('audio')
inputs.update({socket: ['source:1', 1], 'binding_json': json.dumps(binding)})
handler._validate_request(data)
inputs['transfer_format'] = 'jpeg'
with pytest.raises(ValueError):
handler._validate_request(data)
@pytest.mark.parametrize('fmt', FORMATS)
def test_master_passthrough_stays_native_for_every_format(collector_module, fmt):
media = native_media()
actual = collector_module.DistributedCollectorNode().run(
images=[media['images']], audio=[media['audio']], pass_through=[True],
transfer_format=[fmt], transfer_quality=[73])
assert actual[0] is media['images']
assert actual[1] is media['audio']
assert_exact(actual[0], media['images'])
def test_collector_widget_contract_keeps_saved_defaults_and_appends_exact(collector_module):
optional = collector_module.DistributedCollectorNode.INPUT_TYPES()['optional']
assert list(optional) == ['images', 'audio', 'transfer_format', 'transfer_quality']
assert tuple(optional['transfer_format'][0]) == FORMATS
assert optional['transfer_format'][1]['default'] == 'legacy_png'
assert optional['transfer_quality'][1]['default'] == 90
+562
View File
@@ -0,0 +1,562 @@
"""Native ComfyUI VIDEO tests, isolated from the custom-node import graph.
The outer pytest process starts one bounded CPU-only subprocess. Inside it the
real pinned core is imported *after* cli_args.args.cpu=True; no fake VideoInput
classes, media transcodes, providers, server, or GPU queries are needed.
"""
import os
from pathlib import Path
import subprocess
import sys
import pytest
if os.environ.get("_DISTRIBUTED_VIDEO_CODEC_CORE_TEST") != "1":
def test_native_video_codec_in_cpu_subprocess():
configured = os.environ.get("COMFYUI_SOURCE_ROOT", os.environ.get("COMFYUI_VIDEO_TEST_ROOT"))
if configured is None:
pytest.skip("Set COMFYUI_SOURCE_ROOT for CPU native VIDEO acceptance")
core = Path(configured).expanduser().resolve()
assert configured and (core / "comfy_api/latest/_input_impl/video_types.py").is_file(), (
"Configured COMFYUI_SOURCE_ROOT (or COMFYUI_VIDEO_TEST_ROOT) must provide native video classes"
)
env = dict(os.environ)
scratch = Path(os.environ.get("TMPDIR", str(Path.home() / ".hermes/cache/scratch")))
scratch.mkdir(parents=True, exist_ok=True)
env["TMPDIR"] = str(scratch)
for name in ("PYTHONPATH", "COMFYUI_SOURCE_ROOT"):
env.pop(name, None)
env.update(CUDA_VISIBLE_DEVICES="", HIP_VISIBLE_DEVICES="",
_DISTRIBUTED_VIDEO_CODEC_CORE_TEST="1", PYTHONDONTWRITEBYTECODE="1")
bootstrap = """
import sys
sys.path.insert(0, sys.argv[1])
import comfy.cli_args
comfy.cli_args.args.cpu = True
import torch
def forbidden(*args, **kwargs):
raise AssertionError('VIDEO tests must not query or initialize a GPU')
# Third-party core imports probe availability even in CPU mode: answer without
# executing a GPU query, then forbid even these probes during codec execution.
torch.cuda.is_available = lambda: False
torch.cuda.device_count = lambda: 0
for name in ('init', '_lazy_init', 'current_device'):
setattr(torch.cuda, name, forbidden)
import comfy_api.latest._input_impl.video_types
for name in ('is_available', 'device_count'):
setattr(torch.cuda, name, forbidden)
import pytest
raise SystemExit(pytest.main([sys.argv[2], '-q', '-p', 'no:cacheprovider']))
"""
result = subprocess.run(
[sys.executable, "-c", bootstrap, str(core), str(Path(__file__).resolve())],
env=env, text=True, capture_output=True, timeout=150,
cwd=Path(__file__).resolve().parents[1],
)
print(result.stdout)
assert result.returncode == 0, result.stdout + result.stderr
else:
import copy
from fractions import Fraction
import importlib
import io
import json
import math
import types
import av
import numpy as np
import torch
from comfy_api.latest._input import VideoInput
from comfy_api.latest._input_impl.video_types import (
VideoFromComponents, VideoFromFile, VideoFromList,
)
from comfy_api.latest._util import VideoComponents
# Only this sibling package is synthetic; all native media types are real.
package_name = "_distributed_video_codec_test"
package = types.ModuleType(package_name)
package.__path__ = [str(Path(__file__).resolve().parents[1] / "serverless_bridge")]
sys.modules[package_name] = package
codec = importlib.import_module(package_name + ".video_codec")
def audio(dtype=torch.float64):
return {"waveform": torch.tensor([[[0.123456789, -0.0, -1.5, 2.0]]], dtype=dtype, device="cpu"),
"sample_rate": 48000}
def components(dtype=torch.float64, *, metadata=None, alpha=True, sound=True):
image = torch.arange(36, device="cpu", dtype=torch.float64).reshape(2, 2, 3, 3) / 997
image[0, 0, 0, 0] = -0.0
return VideoFromComponents(VideoComponents(
images=image.to(dtype), frame_rate=Fraction(30000, 1001),
audio=audio(dtype) if sound else None,
alpha=torch.full((2, 2, 3), 0.123456789, dtype=dtype, device="cpu") if alpha else None,
metadata=metadata,
), bit_depth=10, color_space="HDR")
def exact_tensor(actual, expected):
assert actual.device.type == "cpu"
assert actual.dtype == expected.dtype
assert actual.shape == expected.shape
assert torch.equal(actual.contiguous().view(torch.uint8), expected.contiguous().view(torch.uint8))
def file_value(data=b"opaque original media bytes", **kwargs):
return VideoFromFile(io.BytesIO(data), **kwargs)
def forbid(*args, **kwargs):
raise AssertionError("Unexpected materialization/transcoding/allocation")
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
@pytest.mark.parametrize("alpha,sound", [(False, False), (True, False), (False, True), (True, True)])
def test_components_exact_roundtrip(dtype, alpha, sound, monkeypatch):
metadata = {"title": "雪 🎞", "flags": [True, None, -4, 0.125], "nested": {"empty": []}}
source = components(dtype, metadata=metadata, alpha=alpha, sound=sound)
expected = source.get_components()
monkeypatch.setattr(VideoFromComponents, "get_components", forbid)
monkeypatch.setattr(VideoFromComponents, "save_to", forbid)
descriptor, payload = codec.encode_video(source)
assert descriptor["encoding"] == "native-video-v1"
assert descriptor["representation"] == "components"
assert descriptor["frame_rate"] == {"numerator": 30000, "denominator": 1001}
# This is the actual JSON wire descriptor, not a Python-only shortcut.
descriptor = json.loads(json.dumps(descriptor, allow_nan=False))
result = codec.decode_video(descriptor, payload)
assert type(result) is VideoFromComponents
assert isinstance(result, VideoInput)
actual = vars(result)["_VideoFromComponents__components"]
exact_tensor(actual.images, expected.images)
assert actual.frame_rate == expected.frame_rate
assert type(actual.frame_rate) is Fraction
assert actual.metadata == metadata
assert actual.metadata is not metadata
assert vars(result)["_VideoFromComponents__bit_depth"] == 10
assert vars(result)["_VideoFromComponents__color_space"] == "HDR"
if alpha:
exact_tensor(actual.alpha, expected.alpha)
else:
assert actual.alpha is None
if sound:
exact_tensor(actual.audio["waveform"], expected.audio["waveform"])
assert actual.audio["sample_rate"] == expected.audio["sample_rate"]
else:
assert actual.audio is None
actual.images.fill_(99)
assert not torch.equal(actual.images, expected.images)
actual.metadata["nested"]["empty"].append("new")
assert descriptor["metadata"]["nested"]["empty"] == []
assert metadata["nested"]["empty"] == []
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
def test_native_bhw1_alpha_shape_is_preserved(dtype):
source = components(dtype)
original = source.get_components()
vars(source)["_VideoFromComponents__components"].alpha = original.alpha.unsqueeze(-1)
descriptor, payload = codec.encode_video(source)
result = codec.decode_video(descriptor, payload).get_components()
exact_tensor(result.alpha, original.alpha.unsqueeze(-1))
assert descriptor["alpha"]["descriptor"]["shape"] == [2, 2, 3, 1]
@pytest.mark.parametrize("start,duration,crop", [
(-1.25, -2.5, (-3, 5, 7, -9)), (0, 0, None), (-0.0, 0.0, (0, 0, 2, 2)),
])
def test_file_preserves_bytes_cursor_signed_trim_crop_and_ownership(start, duration, crop, monkeypatch):
source = file_value(start_time=start, duration=duration, crop=crop)
stream = vars(source)["_VideoFromFile__file"]
stream.seek(7)
monkeypatch.setattr(VideoFromFile, "get_components", forbid)
monkeypatch.setattr(VideoFromFile, "save_to", forbid)
descriptor, payload = codec.encode_video(source)
assert stream.tell() == 7
assert payload == stream.getvalue()
assert descriptor["start_time"] == start
assert descriptor["duration"] == duration
result = codec.decode_video(json.loads(json.dumps(descriptor)), payload)
second = codec.decode_video(descriptor, payload)
attrs = vars(result)
assert type(result) is VideoFromFile
assert attrs["_VideoFromFile__start_time"] == start
assert type(attrs["_VideoFromFile__start_time"]) is type(start)
assert attrs["_VideoFromFile__duration"] == duration
assert attrs["_VideoFromFile__crop"] == crop
if type(start) is float:
assert math.copysign(1, attrs["_VideoFromFile__start_time"]) == math.copysign(1, start)
decoded = attrs["_VideoFromFile__file"]
assert type(decoded) is io.BytesIO
assert decoded is not stream
stream.close()
assert decoded.getvalue() == payload
decoded.write(b"changed")
assert vars(second)["_VideoFromFile__file"].getvalue() == payload
def test_real_tiny_pyav_video_is_still_readable_after_source_closed():
buffer = io.BytesIO()
with av.open(buffer, "w", format="matroska") as container:
stream = container.add_stream("ffv1", rate=Fraction(24, 1))
# Width 32 avoids a native fillborders filter limitation on tiny
# unaligned YUV frames; this is still a tiny CPU-only fixture.
stream.width = stream.height = 32
stream.pix_fmt = "yuv420p"
for value in (16, 64, 128):
frame = av.VideoFrame.from_ndarray(np.full((32, 32, 3), value, dtype=np.uint8), format="rgb24")
for packet in stream.encode(frame):
container.mux(packet)
for packet in stream.encode():
container.mux(packet)
data = buffer.getvalue()
source = VideoFromFile(buffer, crop=(0, 0, 2, 2))
descriptor, payload = codec.encode_video(source)
expected = source.get_components()
result = codec.decode_video(descriptor, payload)
buffer.close()
assert payload == data
actual = result.get_components()
exact_tensor(actual.images, expected.images)
assert actual.frame_rate == expected.frame_rate
assert tuple(actual.images.shape) == (3, 2, 2, 3)
def test_segments_order_complete_audio_no_transcode_and_independent_buffers(monkeypatch):
values = [file_value(b"first", start_time=-2), file_value(b"second", duration=-4, crop=(1, 2, 3, 4))]
source = VideoFromList([VideoFromList(values[:1]), values[1]], complete_audio=audio())
monkeypatch.setattr(VideoFromFile, "get_components", forbid)
monkeypatch.setattr(VideoFromFile, "save_to", forbid)
monkeypatch.setattr(VideoFromList, "get_components", forbid)
monkeypatch.setattr(VideoFromList, "save_to", forbid)
descriptor, payload = codec.encode_video(source)
assert descriptor["representation"] == "segments"
assert payload.startswith(b"firstsecond")
assert len(descriptor["segments"]) == 2
decoded = codec.decode_video(json.loads(json.dumps(descriptor)), payload)
assert type(decoded) is VideoFromList
assert [vars(v)["_VideoFromFile__file"].getvalue() for v in decoded.videos] == [b"first", b"second"]
exact_tensor(decoded.complete_audio["waveform"], source.complete_audio["waveform"])
for old, new in zip(source.videos, decoded.videos):
assert old is not new
vars(old)["_VideoFromFile__file"].close()
assert not vars(new)["_VideoFromFile__file"].closed
assert vars(decoded.videos[0])["_VideoFromFile__start_time"] == -2
assert vars(decoded.videos[1])["_VideoFromFile__crop"] == (1, 2, 3, 4)
def test_path_roots_symlinks_and_no_path_in_wire(tmp_path, monkeypatch):
allowed = tmp_path / "allowed"
allowed.mkdir()
media = allowed / "private-name.mkv"
media.write_bytes(b"media bytes")
sibling = tmp_path / "allowed-sibling"
sibling.mkdir()
secret = sibling / "secret.mkv"
secret.write_bytes(b"private")
(allowed / "escape.mkv").symlink_to(secret)
(allowed / "inside.mkv").symlink_to(media)
descriptor, payload = codec.encode_video(VideoFromFile(str(media)), allowed_roots=[allowed])
assert payload == b"media bytes"
assert "private-name" not in json.dumps(descriptor)
assert codec.encode_video(VideoFromFile(str(allowed / "inside.mkv")), allowed_roots=[allowed])[1] == payload
for path in (secret, allowed / "escape.mkv", allowed, allowed / "missing"):
with pytest.raises(ValueError):
codec.encode_video(VideoFromFile(str(path)), allowed_roots=[allowed])
with pytest.raises(ValueError):
codec.encode_video(VideoFromFile(str(media)), allowed_roots=[])
import folder_paths
for name in ("get_input_directory", "get_output_directory", "get_temp_directory"):
monkeypatch.setattr(folder_paths, name, lambda: str(allowed))
assert codec.encode_video(VideoFromFile(str(media)))[1] == payload
with pytest.raises(ValueError):
codec.encode_video(VideoFromFile(str(secret)))
result = codec.decode_video(descriptor, payload)
media.unlink()
assert vars(result)["_VideoFromFile__file"].getvalue() == b"media bytes"
def test_path_special_files_and_symlink_swap_are_rejected(tmp_path, monkeypatch):
fifo = tmp_path / "fifo"
os.mkfifo(fifo)
with pytest.raises(ValueError):
codec.encode_video(VideoFromFile(str(fifo)), allowed_roots=[tmp_path])
allowed = tmp_path / "allowed"
allowed.mkdir()
actual = allowed / "subdir"
actual.mkdir()
media = actual / "clip.mkv"
media.write_bytes(b"allowed")
outside = tmp_path / "outside"
outside.mkdir()
(outside / "clip.mkv").write_bytes(b"forbidden")
real_resolve = codec.Path.resolve
swapped = False
def swap_after_resolution(path, *args, **kwargs):
nonlocal swapped
resolved = real_resolve(path, *args, **kwargs)
if path == media and not swapped:
swapped = True
actual.rename(allowed / "old")
actual.symlink_to(outside, target_is_directory=True)
return resolved
monkeypatch.setattr(codec.Path, "resolve", swap_after_resolution)
with pytest.raises(ValueError):
codec.encode_video(VideoFromFile(str(media)), allowed_roots=[allowed])
assert swapped
def test_file_size_budget_checked_before_read(tmp_path, monkeypatch):
path = tmp_path / "large.mkv"
path.write_bytes(b"12345")
monkeypatch.setattr(codec.os, "read", forbid)
with pytest.raises(ValueError, match="limit|budget"):
codec.encode_video(VideoFromFile(str(path)), max_bytes=4, allowed_roots=[tmp_path])
@pytest.mark.parametrize("limit", [True, False, -1, 1.5, None])
def test_invalid_limits(limit):
with pytest.raises(ValueError):
codec.encode_video(file_value(), max_bytes=limit)
with pytest.raises(ValueError):
codec.decode_video({}, b"", max_bytes=limit)
def test_aggregate_budget_and_exact_boundary():
value = VideoFromList([file_value(b"12"), file_value(b"345")])
descriptor, payload = codec.encode_video(value, max_bytes=5)
assert payload == b"12345"
codec.decode_video(descriptor, payload, max_bytes=5)
with pytest.raises(ValueError):
codec.encode_video(value, max_bytes=4)
with pytest.raises(ValueError):
codec.decode_video(descriptor, payload, max_bytes=4)
stream = vars(value.videos[0])["_VideoFromFile__file"]
stream.seek(1)
with pytest.raises(ValueError):
codec.encode_video(value.videos[0], max_bytes=1)
assert stream.tell() == 1
@pytest.mark.parametrize("kind", ["components", "file", "segments"])
def test_unknown_subclasses_never_materialized(kind):
bases = {"components": VideoFromComponents, "file": VideoFromFile, "segments": VideoFromList}
cls = type("UnknownVideo", (bases[kind],), {"get_components": forbid, "save_to": forbid})
value = object.__new__(cls)
with pytest.raises(ValueError, match="class|type|unsupported|Unsupported"):
codec.encode_video(value)
with pytest.raises(ValueError):
codec.encode_video(object())
@pytest.mark.parametrize("kind", ["components", "file", "segments"])
def test_bridge_rejects_subclasses_without_materializing(kind, monkeypatch):
nodes = importlib.import_module(package_name + ".nodes")
bases = {"components": VideoFromComponents, "file": VideoFromFile, "segments": VideoFromList}
cls = type("UnknownVideo", (bases[kind],), {"get_components": forbid, "save_to": forbid})
monkeypatch.setattr(torch.Tensor, "cpu", forbid)
with pytest.raises(ValueError, match="Unsupported native VIDEO class"):
nodes.encode_value("VIDEO", object.__new__(cls))
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
@pytest.mark.parametrize("kind", ["components", "segments", "file"])
def test_bridge_stages_cpu_tensors_without_materialization(dtype, kind, monkeypatch):
nodes = importlib.import_module(package_name + ".nodes")
if kind == "components":
source = components(dtype, metadata={"nested": ["雪", None, True]})
state = vars(source)["_VideoFromComponents__components"]
state.alpha = state.alpha.unsqueeze(-1)
tensors = [state.images, state.alpha, state.audio["waveform"]]
elif kind == "segments":
source = VideoFromList([file_value(start_time=-1.5, duration=-2, crop=(-1, 2, 3, 4))],
complete_audio=audio(dtype))
tensors = [source.complete_audio["waveform"]]
else:
source = file_value(start_time=-0.0, duration=-2, crop=(-1, 2, 3, 4))
tensors = []
for tensor in tensors:
tensor.requires_grad_(True)
expected = codec.encode_video(source)
detached, staged = [], []
real_detach, real_cpu = torch.Tensor.detach, torch.Tensor.cpu
def detach(tensor):
detached.append(id(tensor))
return real_detach(tensor)
def cpu(tensor, *args, **kwargs):
assert not tensor.requires_grad
staged.append(tensor)
return real_cpu(tensor, *args, **kwargs)
for cls in (VideoFromComponents, VideoFromFile, VideoFromList):
monkeypatch.setattr(cls, "get_components", forbid)
monkeypatch.setattr(cls, "save_to", forbid)
monkeypatch.setattr(torch.Tensor, "detach", detach)
monkeypatch.setattr(torch.Tensor, "cpu", cpu)
actual = nodes.encode_value("VIDEO", source)
assert len(staged) == len(tensors)
assert all(id(tensor) in detached for tensor in tensors)
assert all(tensor.requires_grad for tensor in tensors)
assert actual == expected
decoded = nodes.decode_value("VIDEO", *actual)
assert type(decoded) is type(source)
assert codec.encode_video(decoded) == expected
assert len(staged) == len(tensors) # Pure codec never stages implicitly.
@pytest.mark.parametrize("field", ["images", "alpha", "audio", "complete_audio"])
def test_pure_video_rejects_non_cpu_storage_without_transfer(field, monkeypatch):
if field == "complete_audio":
source = VideoFromList([file_value()], complete_audio=audio())
target = source.complete_audio["waveform"]
else:
source = components()
target = getattr(vars(source)["_VideoFromComponents__components"], field)
if field == "audio":
target = target["waveform"]
# A property spy on real CPU tensors exercises the guard without a GPU
# allocation, subclass, device discovery or CUDA initialization.
real_device = torch.Tensor.device
monkeypatch.setattr(torch.Tensor, "device", property(
lambda tensor: types.SimpleNamespace(type="cuda") if tensor is target
else real_device.__get__(tensor)))
monkeypatch.setattr(torch.Tensor, "cpu", forbid)
with pytest.raises(ValueError, match="must be on CPU"):
codec.encode_video(source)
@pytest.mark.parametrize("field,bad", [
("_VideoFromFile__start_time", True), ("_VideoFromFile__start_time", float("nan")),
("_VideoFromFile__duration", float("inf")), ("_VideoFromFile__crop", (1, 2, True, 4)),
("_VideoFromFile__crop", [1, 2, 3, 4]), ("_VideoFromFile__crop", (1, 2, 3)),
("_VideoFromFile__file", b"not a BytesIO"),
])
def test_invalid_native_file_attributes(field, bad):
value = file_value()
setattr(value, field, bad)
with pytest.raises(ValueError):
codec.encode_video(value)
def test_missing_attributes_and_invalid_component_state():
value = file_value()
del vars(value)["_VideoFromFile__duration"]
with pytest.raises(ValueError):
codec.encode_video(value)
for field, bad in [("frame_rate", 24.0), ("frame_rate", Fraction(0)),
("alpha", torch.ones(1, 1, 1, device="cpu"))]:
value = components()
setattr(vars(value)["_VideoFromComponents__components"], field, bad)
with pytest.raises(ValueError):
codec.encode_video(value)
@pytest.mark.parametrize("metadata", [
{"value": float("nan")}, {"value": float("inf")}, {1: "non-string key"},
{"value": (1, 2)}, {"value": object()}, {"large": "x" * 65536},
])
def test_invalid_metadata_rejected_both_directions(metadata):
with pytest.raises(ValueError):
codec.encode_video(components(metadata=metadata))
descriptor, payload = codec.encode_video(components())
descriptor["metadata"] = metadata
with pytest.raises(ValueError):
codec.decode_video(descriptor, payload)
def test_metadata_depth_cycles_and_exact_size():
metadata = {}
current = metadata
for _ in range(40):
current["child"] = {}
current = current["child"]
for value in (metadata,):
with pytest.raises(ValueError):
codec.encode_video(components(metadata=value))
cyclic = {}
cyclic["self"] = cyclic
with pytest.raises(ValueError):
codec.encode_video(components(metadata=cyclic))
exact = {"x": "a" * (65536 - 8)}
assert len(json.dumps(exact, separators=(",", ":")).encode()) == 65536
descriptor, payload = codec.encode_video(components(metadata=exact))
assert codec.decode_video(descriptor, payload).get_components().metadata == exact
@pytest.mark.parametrize("field,bad", [
("encoding", "other"), ("representation", "other"), ("byte_length", True),
("byte_length", -1), ("byte_length", 12345), ("start_time", False),
("duration", float("inf")), ("crop", [0, 0, True, 4]), ("crop", [1, 2, 3]),
])
def test_malformed_file_descriptors(field, bad):
descriptor, payload = codec.encode_video(file_value())
descriptor[field] = bad
with pytest.raises(ValueError):
codec.decode_video(descriptor, payload)
@pytest.mark.parametrize("field,bad", [("offset", True), ("offset", -1), ("offset", 1),
("length", True), ("length", -1), ("length", 2**80)])
def test_invalid_offsets_lengths(field, bad):
descriptor, payload = codec.encode_video(file_value())
descriptor["data"][field] = bad
with pytest.raises(ValueError):
codec.decode_video(descriptor, payload)
def test_truncation_trailing_gaps_overlap_and_unknown_fields():
value = VideoFromList([file_value(b"123"), file_value(b"456")])
original, payload = codec.encode_video(value)
for wire in (payload[:-1], payload + b"x"):
with pytest.raises(ValueError):
codec.decode_video(original, wire)
for offset in (2, 4):
descriptor = copy.deepcopy(original)
descriptor["segments"][1]["offset"] = offset
with pytest.raises(ValueError):
codec.decode_video(descriptor, payload)
descriptor = copy.deepcopy(original)
descriptor["surprise"] = "bad"
with pytest.raises(ValueError):
codec.decode_video(descriptor, payload)
descriptor = copy.deepcopy(original)
descriptor["segments"][0]["descriptor"]["byte_length"] += 1
with pytest.raises(ValueError):
codec.decode_video(descriptor, payload)
descriptor = copy.deepcopy(original)
descriptor["segments"][0]["descriptor"] = descriptor
with pytest.raises(ValueError):
codec.decode_video(descriptor, payload)
def test_all_descriptors_validated_before_tensor_allocation(monkeypatch):
original, payload = codec.encode_video(components())
monkeypatch.setattr(codec, "decode_payload", forbid)
cases = []
for key, value in [("shape", [2**80, 2, 3]), ("shape", [True, 2, 3]),
("dtype", "pickle"), ("layout", "BHWC")]:
descriptor = copy.deepcopy(original)
descriptor["alpha"]["descriptor"][key] = value
cases.append(descriptor)
descriptor = copy.deepcopy(original)
descriptor["audio"]["descriptor"]["sample_rate"] = True
cases.append(descriptor)
descriptor = copy.deepcopy(original)
descriptor["alpha"]["offset"] += 1
cases.append(descriptor)
for descriptor in cases:
with pytest.raises(ValueError):
codec.decode_video(descriptor, payload)
@pytest.mark.parametrize("field,bad", [("numerator", True), ("denominator", True),
("denominator", 0), ("numerator", -3),
("denominator", -1001)])
def test_invalid_fraction_fields(field, bad):
descriptor, payload = codec.encode_video(components())
descriptor["frame_rate"][field] = bad
with pytest.raises(ValueError):
codec.decode_video(descriptor, payload)
def test_segment_limit_and_exact_maximum(monkeypatch):
value = VideoFromList([file_value(b"x") for _ in range(64)])
descriptor, payload = codec.encode_video(value)
result = codec.decode_video(descriptor, payload)
assert len(result.videos) == 64
value.videos.append(file_value(b"x"))
with pytest.raises(ValueError):
codec.encode_video(value)
descriptor["segments"].append(copy.deepcopy(descriptor["segments"][-1]))
monkeypatch.setattr(codec, "decode_payload", forbid)
with pytest.raises(ValueError):
codec.decode_video(descriptor, payload)
value.videos = []
with pytest.raises(ValueError):
codec.encode_video(value)
value.videos = [components()]
with pytest.raises(ValueError):
codec.encode_video(value)
@pytest.mark.parametrize("descriptor,payload", [(None, b""), ([], b""), ({}, bytearray()), ({}, "")])
def test_invalid_top_level_types(descriptor, payload):
with pytest.raises(ValueError):
codec.decode_video(descriptor, payload)
+26
View File
@@ -60,6 +60,29 @@ class ValidateServerlessWorkerTests(unittest.TestCase):
self.assertIsNot(result, worker)
self.assertEqual(worker, {"id": "worker-1", "endpoint_id": "ep_123-abc", "name": "GPU"})
def test_omitted_media_profile_remains_absent_and_means_legacy(self):
result = serverless_config.validate_serverless_worker(
{"id": "worker-1", "endpoint_id": "ep-123"}
)
self.assertNotIn("media_profile", result)
self.assertEqual(result.get("media_profile", "legacy-image-v1"), "legacy-image-v1")
def test_explicit_media_profiles_are_retained(self):
for profile in ("legacy-image-v1", "comfy-media-v1"):
with self.subTest(profile=profile):
result = serverless_config.validate_serverless_worker(
{"id": "worker-1", "endpoint_id": "ep-123", "media_profile": profile}
)
self.assertEqual(result["media_profile"], profile)
def test_media_profile_rejects_unknown_nonstring_and_null_values(self):
for profile in ("", "auto", "comfy-media-v2", " comfy-media-v1", "COMFY-MEDIA-V1", None, True, 1, [], {}):
with self.subTest(profile=profile):
with self.assertRaisesRegex(ValueError, "media_profile"):
serverless_config.validate_serverless_worker(
{"id": "worker-1", "endpoint_id": "ep-123", "media_profile": profile}
)
def test_accepts_inclusive_numeric_bounds(self):
for jobs, timeout in ((1, 30), (16, 86400)):
with self.subTest(jobs=jobs, timeout=timeout):
@@ -307,6 +330,7 @@ class PublicConfigTests(unittest.TestCase):
"type": "serverless",
"provider": "runpod",
"endpoint_id": "ep-123",
"media_profile": "comfy-media-v1",
"api_key": "inline-secret",
"clear_api_key": "another-secret",
"credential": "credential-secret",
@@ -334,6 +358,8 @@ class PublicConfigTests(unittest.TestCase):
self.assertNotIn("api_key", local)
self.assertNotIn("credentials_note", result["settings"])
self.assertEqual(first["name"], "remote GPU")
self.assertEqual(first["media_profile"], "comfy-media-v1")
self.assertNotIn("media_profile", second)
self.assertEqual(config, original)
self.assertIsNot(result, config)
self.assertIsNot(result["workers"], config["workers"])
+426
View File
@@ -0,0 +1,426 @@
"""CPU-only graph/compiler tests; no ComfyUI package bootstrap is needed."""
import copy
import importlib
import json
import subprocess
import sys
import types
from pathlib import Path
import pytest
ROOT = Path(__file__).resolve().parents[1]
PACKAGE = "serverless_media_prompt_testpkg"
# Real paths let sibling imports use production code, without running the
# extension's ComfyUI-dependent __init__.py or API registration side effects.
for suffix in ("", ".api", ".api.orchestration", ".serverless_bridge"):
package = types.ModuleType(PACKAGE + suffix)
package.__path__ = [str(ROOT.joinpath(*suffix.strip(".").split(".")))] if suffix else [str(ROOT)]
sys.modules[package.__name__] = package
compiler = importlib.import_module(f"{PACKAGE}.api.orchestration.serverless_media_prompt")
legacy = importlib.import_module(f"{PACKAGE}.api.orchestration.serverless_prompt")
contract = importlib.import_module(f"{PACKAGE}.serverless_bridge.contract")
compile_media_prompt = compiler.compile_media_prompt
SOCKETS = {"images": "IMAGE", "audio": "AUDIO"}
def _node(class_type, **inputs):
return {"class_type": class_type, "inputs": inputs}
def _collector(**sockets):
return _node(
"DistributedCollector", multi_job_id="job/42", worker_id="worker_0", **sockets
)
def _prompt():
return {
"source:1": _node("LoadImage", image="example.png"),
"collector:1": _collector(images=["source:1", 0]),
}
def _run(prompt):
return compile_media_prompt(prompt, "request:1")
@pytest.mark.parametrize("socket,kind", SOCKETS.items())
def test_one_export_for_each_typed_socket_with_exact_canonical_binding(socket, kind):
prompt = {"source": _node("Source"), "collector": _collector(**{socket: ["source", 0]})}
graph, bindings = _run(prompt)
assert len(bindings) == 1
binding = bindings[0]
assert set(binding) == contract.BINDING_FIELDS
assert binding == {
"binding_id": binding["binding_id"],
"collector_id": "collector",
"job_id": "job/42",
"participant_id": "worker_0",
"node_id": binding["node_id"],
"socket": socket,
"kind": kind,
}
assert all(isinstance(value, str) for value in binding.values())
assert contract.validate_bindings(bindings) == {binding["binding_id"]: binding}
assert graph[binding["node_id"]] == {
"class_type": "MediaExportV1",
"inputs": {
"request_id": "request:1",
"binding_json": json.dumps(binding),
socket: ["source", 0],
**({"transfer_format": "png", "transfer_quality": 90} if kind == "IMAGE" else {}),
},
}
assert set(graph) == {"source", binding["node_id"]}
def test_multiple_collectors_have_independent_image_audio_targets():
prompt = {
"shared": _node("MultiMediaSource"),
"collector:z": _collector(audio=["shared", 1]),
"collector:a": _collector(images=["shared", 0], audio=["shared", 1]),
}
prompt["collector:z"]["inputs"].update(multi_job_id="job:z", worker_id="participant:z")
graph, bindings = _run(prompt)
assert [(b["collector_id"], b["socket"]) for b in bindings] == [
("collector:a", "images"), ("collector:a", "audio"), ("collector:z", "audio"),
]
assert len({b["binding_id"] for b in bindings}) == 3
targets = {b["node_id"] for b in bindings}
assert len(targets) == 3
assert set(graph) == {"shared"} | targets
for binding in bindings:
export = graph[binding["node_id"]]
inputs = prompt[binding["collector_id"]]["inputs"]
assert export["inputs"][binding["socket"]] == inputs[binding["socket"]]
assert binding["job_id"] == inputs["multi_job_id"]
assert binding["participant_id"] == inputs["worker_id"]
def test_prunes_descendants_and_unrelated_outputs_preserving_required_output_utility():
prompt = _prompt()
prompt.update({
"required-output": _node("RequiredOutputUtility", images=["source:1", 0]),
"consumer": _node("ImageTransform", image=["required-output", 0]),
"descendant": _node("ImageTransform", image=["collector:1", 0]),
"after": _node("SaveImage", images=["descendant", 0]),
"unused-output": _node("SaveImage", images=["source:1", 0]),
"unrelated": _node("LoadAudio"),
})
prompt["required-output"]["_meta"] = {"title": "Required OUTPUT_NODE utility"}
prompt["collector:1"]["inputs"].update(images=["consumer", 0], load_balance=True, pass_through=False)
graph, bindings = _run(prompt)
assert set(graph) == {"source:1", "required-output", "consumer", bindings[0]["node_id"]}
assert graph["required-output"] == prompt["required-output"]
assert not {"load_balance", "pass_through"} & graph[bindings[0]["node_id"]]["inputs"].keys()
def test_upstream_primitive_string_preserves_empty_and_link_looking_text():
for text in ("", '["source:1", 0]', "hello\nworld"):
prompt = {"text": _node("PrimitiveString", value=text),
"source": _node("TextToImage", text=["text", 0]),
"collector": _collector(images=["source", 0])}
graph, bindings = _run(prompt)
assert graph["text"] == prompt["text"]
assert graph["source"]["inputs"]["text"] == ["text", 0]
assert graph[bindings[0]["node_id"]]["inputs"]["images"] == ["source", 0]
@pytest.mark.parametrize("class_type,inputs", [
("DistributedSeed", {}),
("DistributedSeed", {"seed": 10, "is_worker": True, "worker_id": "worker_1"}),
("DistributedSeed", {"seed": 10, "is_worker": False, "worker_id": "worker_1"}),
("DistributedSeed", {"seed": 10, "is_worker": True, "worker_id": "bad"}),
("DistributedValue", {"default_value": "3.9", "worker_values": '{"_type":"INT","1":"8.7"}', "is_worker": True, "worker_id": "worker_0"}),
("DistributedValue", {"default_value": "3.9", "worker_values": '{"_type":"INT","2":""}', "is_worker": True, "worker_id": "worker_1"}),
("DistributedValue", {"default_value": "fallback", "worker_values": "malformed"}),
("DistributedValue", {"default_value": 1.25, "worker_values": '{"_type":"FLOAT","1":"bad"}', "is_worker": True, "worker_id": "0"}),
("DistributedModelName", {"text": "model.safetensors"}),
("DistributedModelName", {"text": {"name": "checkpoint"}}),
("DistributedModelName", {"text": ["a", "b", "c"]}),
])
def test_upstream_scalar_values_match_legacy_compiler(class_type, inputs):
prompt = _prompt()
prompt["scalar"] = _node(class_type, **inputs)
prompt["source:1"]["inputs"]["parameter"] = ["scalar", 0]
old_prompt = copy.deepcopy(prompt)
old_prompt["collector:1"]["class_type"] = "DistributedCollector"
expected, _ = legacy.compile_serverless_prompt(old_prompt)
graph, _ = _run(prompt)
assert graph["source:1"] == expected["source:1"]
assert "scalar" not in graph
def test_imports_same_scalar_resolver_and_resolves_primitive_chains():
assert compiler._make_scalar_resolver is legacy._make_scalar_resolver
prompt = _prompt()
prompt.update({
"int": _node("PrimitiveInt", value=25),
"int-alias": _node("PrimitiveInt", value=["int", 0]),
"float": _node("PrimitiveFloat", value=1.25),
"bool": _node("PrimitiveBoolean", value=True),
"string": _node("PrimitiveString", value="model.safetensors"),
"seed": _node("DistributedSeed", seed=["int-alias", 0], is_worker=True, worker_id="worker_0"),
"value": _node("DistributedValue", default_value=["float", 0], worker_values='{"_type":"FLOAT","1":"2.75"}', is_worker=True, worker_id="worker_0"),
"enabled": _node("DistributedValue", default_value=["bool", 0]),
"model": _node("DistributedModelName", text=["string", 0]),
})
prompt["source:1"]["inputs"].update(seed=["seed", 0], cfg=["value", 0], enabled=["enabled", 0], model=["model", 0])
graph, bindings = _run(prompt)
assert graph["source:1"]["inputs"] == {"image": "example.png", "seed": 26, "cfg": 2.75, "enabled": True, "model": "model.safetensors"}
assert set(graph) == {"source:1", bindings[0]["node_id"]}
def test_distributed_text_value_can_be_wrapped_by_primitive_string():
prompt = {
"value": _node("DistributedModelName", text="model.safetensors"),
"text": _node("PrimitiveString", value=["value", 0]),
"source": _node("TextToImage", text=["text", 0]),
"collector": _collector(images=["source", 0]),
}
graph, bindings = _run(prompt)
assert graph["text"]["inputs"]["value"] == "model.safetensors"
assert graph["source"]["inputs"]["text"] == ["text", 0]
assert graph[bindings[0]["node_id"]]["inputs"]["images"] == ["source", 0]
assert "value" not in graph
@pytest.mark.parametrize("class_type,input_name", [
("DistributedSeed", "seed"), ("DistributedValue", "default_value"), ("DistributedModelName", "text"),
])
def test_dynamic_scalar_resolution_stays_fail_closed(class_type, input_name):
prompt = _prompt()
prompt["dynamic"] = _node("ArbitraryValue")
prompt["scalar"] = _node(class_type, **{input_name: ["dynamic", 0]})
prompt["source:1"]["inputs"]["parameter"] = ["scalar", 0]
with pytest.raises(ValueError, match="literal or supported Primitive"):
_run(prompt)
def test_namespaced_export_ids_avoid_existing_ids_even_on_pruned_nodes():
prompt = _prompt()
_, first = _run(prompt)
original_id = first[0]["node_id"]
assert ":" in original_id and not original_id.isdigit()
prompt[original_id] = _node("ExistingUtility")
_, second = _run(prompt)
assert second[0]["node_id"] not in prompt
assert second[0]["binding_id"] == first[0]["binding_id"]
prompt[second[0]["node_id"]] = _node("AnotherUtility")
graph, third = _run(prompt)
assert third[0]["node_id"] not in prompt
assert third == _run(dict(reversed(list(prompt.items()))))[1]
assert set(graph) == {"source:1", third[0]["node_id"]}
def test_long_valid_identifiers_remain_exact_and_generated_ids_fit_contract():
collector_id = "namespace/" + "c" * 190
prompt = {"source:1": _node("Source"), collector_id: _collector(images=["source:1", 0])}
prompt[collector_id]["inputs"].update(multi_job_id="j" * 200, worker_id="p" * 200)
_, bindings = _run(prompt)
assert bindings[0]["collector_id"] == collector_id
assert bindings[0]["job_id"] == "j" * 200
assert bindings[0]["participant_id"] == "p" * 200
contract.validate_bindings(bindings)
def test_deterministic_independent_copies_and_request_identity():
prompt = _prompt()
prompt["source:1"]["_meta"] = {"nested": {"values": [1, 2, 3]}}
original = copy.deepcopy(prompt)
graph, bindings = _run(prompt)
assert prompt == original
assert (graph, bindings) == _run(prompt)
assert (graph, bindings) == _run(dict(reversed(list(prompt.items()))))
other_graph, other_bindings = compile_media_prompt(prompt, "request:2")
assert other_bindings == bindings
target = bindings[0]["node_id"]
assert other_graph[target]["inputs"]["request_id"] == "request:2"
graph["source:1"]["_meta"]["nested"]["values"].append(4)
graph[target]["inputs"]["images"][0] = "changed"
bindings[0]["collector_id"] = "changed"
assert prompt == original
assert json.loads(graph[target]["inputs"]["binding_json"])["collector_id"] == "collector:1"
assert other_graph["source:1"] == original["source:1"]
@pytest.mark.parametrize("bad", [None, [], {}, {"x": None}, {"x": []}, {"x": {}},
{"x": {"class_type": None}}, {"x": {"class_type": ""}},
{"x": {"class_type": "Source", "inputs": []}},
{"x": {"class_type": "Source", "inputs": None}},
{"x": {"class_type": "Source", "inputs": {1: "literal"}}},
])
def test_rejects_malformed_prompt_or_missing_collectors(bad):
with pytest.raises(ValueError):
_run(bad)
@pytest.mark.parametrize("bad", [None, 1, True, "", " bad", "bad ", "bad\\id", "bad\nid", "x" * 201])
def test_rejects_malformed_ids_without_normalizing(bad):
prompt = _prompt()
prompt[bad] = _node("Source")
with pytest.raises(ValueError):
_run(prompt)
@pytest.mark.parametrize("bad", [None, 1, True, "", " bad", "bad ", "bad\\id", "x" * 201])
@pytest.mark.parametrize("field", ["multi_job_id", "worker_id", "request_id"])
def test_rejects_malformed_binding_and_request_identity(field, bad):
prompt = _prompt()
request_id = "request:1"
if field == "request_id":
request_id = bad
else:
prompt["collector:1"]["inputs"][field] = bad
with pytest.raises(ValueError):
compile_media_prompt(prompt, request_id)
@pytest.mark.parametrize("field", ["multi_job_id", "worker_id"])
def test_requires_prepared_literal_hidden_identity(field):
prompt = _prompt()
del prompt["collector:1"]["inputs"][field]
with pytest.raises(ValueError):
_run(prompt)
prompt["collector:1"]["inputs"][field] = ["source:1", 0]
with pytest.raises(ValueError):
_run(prompt)
@pytest.mark.parametrize("socket", SOCKETS)
@pytest.mark.parametrize("value", [None, "literal", 5, True, [], ["source:1"], ["source:1", 0, 1], ("source:1", 0)])
def test_every_present_media_socket_requires_an_api_link(socket, value):
prompt = _prompt()
prompt["collector:1"]["inputs"][socket] = value
with pytest.raises(ValueError, match="link"):
_run(prompt)
def test_each_collector_requires_at_least_one_media_socket():
prompt = _prompt()
prompt["empty"] = _collector(load_balance=True, pass_through=False)
with pytest.raises(ValueError, match="connected|linked"):
_run(prompt)
@pytest.mark.parametrize("link", [["missing", 0], ["source:1", -1], ["source:1", 0.0],
["source:1", "0"], ["source:1", True], ["source:1", False], [1, 0],
[True, 0], [None, 0], [[], 0], [{}, 0], [" source:1", 0],
])
@pytest.mark.parametrize("location", ["collector", "ordinary", "unrelated"])
def test_rejects_invalid_links_everywhere_before_pruning(location, link):
prompt = _prompt()
prompt["1"] = _node("Source") # An integer/bool source must not be stringified to this ID.
if location == "collector":
prompt["collector:1"]["inputs"]["images"] = link
elif location == "ordinary":
prompt["source:1"]["inputs"]["parameter"] = link
else:
prompt["unrelated"] = _node("Utility", value=link)
with pytest.raises(ValueError):
_run(prompt)
@pytest.mark.parametrize("class_type", [
"DistributedMediaCollector", "DistributedEmptyImage", "UltimateSDUpscaleDistributed",
"DistributedImageListStart", "DistributedImageListEnd", "DistributedWorkerSink",
"ImageBatchDivider", "AudioBatchDivider", "CustomDistributedNode", "custombatchdivider",
])
def test_rejects_unsupported_distributed_families_even_if_unused(class_type):
prompt = _prompt()
prompt["unsupported"] = _node(class_type)
with pytest.raises(ValueError, match="unsupported"):
_run(prompt)
@pytest.mark.parametrize("class_type", ["DistributedSeed", "DistributedValue", "DistributedModelName", "PrimitiveInt", "PrimitiveString"])
def test_rejects_invalid_known_scalar_output_slots(class_type):
prompt = _prompt()
prompt["scalar"] = _node(class_type)
prompt["source:1"]["inputs"]["parameter"] = ["scalar", 1]
with pytest.raises(ValueError, match="slot"):
_run(prompt)
@pytest.mark.parametrize("socket", SOCKETS)
def test_rejects_direct_distributed_scalar_media_links(socket):
prompt = {"scalar": _node("DistributedValue", default_value="text"), "collector": _collector(**{socket: ["scalar", 0]})}
with pytest.raises(ValueError, match="scalar"):
_run(prompt)
@pytest.mark.parametrize("indirect", [False, True])
def test_rejects_collector_chains_across_different_media(indirect):
prompt = _prompt()
source = "collector:1"
if indirect:
prompt["processing"] = _node("ImageToAudio", images=[source, 0])
source = "processing"
prompt["collector:2"] = _collector(audio=[source, 0])
with pytest.raises(ValueError, match="collector-to-collector"):
_run(prompt)
@pytest.mark.parametrize("mode", ["ordinary", "self", "unrelated", "primitive", "collector"])
def test_rejects_cycles_in_entire_prompt(mode):
prompt = _prompt()
if mode == "ordinary":
prompt["source:1"]["inputs"]["other"] = ["other", 0]
prompt["other"] = _node("Utility", value=["source:1", 0])
elif mode == "self":
prompt["source:1"]["inputs"]["other"] = ["source:1", 0]
elif mode == "unrelated":
prompt["unused"] = _node("Utility", value=["unused", 0])
elif mode == "primitive":
prompt["p1"] = _node("PrimitiveInt", value=["p2", 0])
prompt["p2"] = _node("PrimitiveInt", value=["p1", 0])
else:
prompt["source:1"]["inputs"]["other"] = ["collector:1", 0]
with pytest.raises(ValueError, match="cycle"):
_run(prompt)
def test_deep_acyclic_graph_does_not_depend_on_python_recursion_limit():
prompt = {"n:0": _node("Source")}
for index in range(1, 1500):
prompt[f"n:{index}"] = _node("Utility", value=[f"n:{index - 1}", 0])
prompt["collector"] = _collector(images=["n:1499", 0])
graph, bindings = _run(prompt)
assert len(graph) == 1501
assert graph[bindings[0]["node_id"]]["inputs"]["images"] == ["n:1499", 0]
def test_binding_count_respects_bridge_contract():
prompt = {"source": _node("Source")}
for index in range(contract.MAX_ARTIFACTS + 1):
prompt[f"collector:{index}"] = _collector(images=["source", 0])
with pytest.raises(ValueError, match="count"):
_run(prompt)
def test_import_and_compile_do_not_import_comfy_or_tensor_libraries():
script = '''
import builtins, importlib, pathlib, sys, types
root = pathlib.Path(sys.argv[1])
name = "isolated_media_compile"
for suffix in ("", ".api", ".api.orchestration", ".serverless_bridge"):
package = types.ModuleType(name + suffix)
package.__path__ = [str(root.joinpath(*suffix.strip(".").split(".")))] if suffix else [str(root)]
sys.modules[package.__name__] = package
original_import = builtins.__import__
def guarded_import(name, *args, **kwargs):
if name.split(".")[0] in {"comfy", "comfy_api", "nodes", "server", "torch", "numpy"}:
raise AssertionError("unexpected runtime import: " + name)
return original_import(name, *args, **kwargs)
builtins.__import__ = guarded_import
module = importlib.import_module(name + ".api.orchestration.serverless_media_prompt")
prompt = {"s": {"class_type": "Source", "inputs": {}}, "c": {"class_type": "DistributedCollector", "inputs": {"images": ["s", 0], "multi_job_id": "job", "worker_id": "worker"}}}
graph, bindings = module.compile_media_prompt(prompt, "request")
assert graph[bindings[0]["node_id"]]["class_type"] == "MediaExportV1"
'''
result = subprocess.run([sys.executable, "-c", script, str(ROOT)], text=True, capture_output=True)
assert result.returncode == 0, result.stderr
+100 -4
View File
@@ -217,6 +217,78 @@ class ServerlessConfigRouteTests(_ConfigRouteFixture):
self.assertNotIn(obsolete, saved)
self.assertEqual(self.secrets.get("pool-1"), secret)
async def test_media_profile_persists_and_is_public_without_changing_credentials(self):
for profile in ("legacy-image-v1", "comfy-media-v1"):
with self.subTest(profile=profile):
response = await config_routes.update_worker_endpoint(_FakeRequest({
"worker_id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
"media_profile": profile, "api_key": "saved-key", "enabled": False,
}))
self.assertEqual(response.status, 200)
self.assertEqual(self.read_config_file()["workers"][0]["media_profile"], profile)
config_module.invalidate_config_cache()
public = await config_routes.get_config_endpoint(_FakeRequest())
self.assertEqual(public.status, 200)
self.assertEqual(public.payload["workers"][0]["media_profile"], profile)
self.assertTrue(public.payload["workers"][0]["api_key_configured"])
self.assertNotIn("saved-key", json.dumps(public.payload))
self.assertEqual(self.secrets.get("pool-1"), "saved-key")
async def test_old_clients_omit_media_profile_without_downgrading_saved_profile(self):
for stored in ({}, {"media_profile": "comfy-media-v1"}):
with self.subTest(stored=stored):
worker = {"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", **stored}
self.config["workers"] = [worker]
self.assertTrue(config_module.save_config(self.config))
response = await config_routes.update_worker_endpoint(_FakeRequest({
"worker_id": "pool-1", "name": "Old client save", "jobs_per_run": 2,
}))
self.assertEqual(response.status, 200)
saved = self.read_config_file()["workers"][0]
self.assertEqual({k: v for k, v in saved.items() if k == "media_profile"}, stored)
# The bulk route is also used by older clients with no profile field.
response = await config_routes.update_config_endpoint(_FakeRequest({"workers": [{
"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "jobs_per_run": 3,
}]}))
self.assertEqual(response.status, 200)
saved = self.read_config_file()["workers"][0]
self.assertEqual({k: v for k, v in saved.items() if k == "media_profile"}, stored)
async def test_invalid_media_profile_is_rejected_by_worker_and_bulk_routes(self):
for profile in ("auto", "comfy-media-v2", None, False, 1, [], {}):
with self.subTest(profile=profile):
response = await config_routes.update_worker_endpoint(_FakeRequest({
"worker_id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
"media_profile": profile, "api_key": "never-save",
}))
self.assertEqual(response.status, 400)
response = await config_routes.update_config_endpoint(_FakeRequest({"workers": [{
"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "media_profile": profile,
}]}))
self.assertEqual(response.status, 400)
self.assertEqual(self.read_config_file()["workers"], [])
self.assertIsNone(self.secrets.get("pool-1"))
async def test_active_pool_blocks_media_profile_changes_in_both_directions(self):
for old_profile, new_profile in ((None, "comfy-media-v1"), ("comfy-media-v1", "legacy-image-v1")):
with self.subTest(old_profile=old_profile):
worker = {"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123"}
if old_profile is not None:
worker["media_profile"] = old_profile
self.config["workers"] = [worker]
self.assertTrue(config_module.save_config(self.config))
original = self.read_config_file()
with patch.object(config_routes, "_serverless_pool_is_active", return_value=True):
response = await config_routes.update_worker_endpoint(_FakeRequest({
"worker_id": "pool-1", "media_profile": new_profile,
}))
self.assertEqual(response.status, 409)
response = await config_routes.update_config_endpoint(_FakeRequest({
"workers": [dict(worker, media_profile=new_profile)],
}))
self.assertEqual(response.status, 409)
self.assertEqual(self.read_config_file(), original)
async def test_enable_only_update_preserves_serverless_settings_and_saved_key(self):
self.secrets.set("pool-1", "saved-secret")
self.config["workers"] = [{
@@ -302,13 +374,14 @@ class ServerlessConfigRouteTests(_ConfigRouteFixture):
"workers": [{
"id": "pool-1", "type": "local", "port": 8188,
"provider": "runpod", "endpoint_id": "legacy-ep", "api_key_configured": True,
"media_profile": "comfy-media-v1",
}]
}))
self.assertEqual(response.status, 200)
self.assertIsNone(self.secrets.get("pool-1"))
saved = self.read_config_file()["workers"][0]
for field in ("provider", "endpoint_id", "api_key_configured"):
for field in ("provider", "endpoint_id", "api_key_configured", "media_profile"):
self.assertNotIn(field, saved)
async def test_bulk_invalid_serverless_record_does_not_replace_configuration(self):
@@ -341,6 +414,7 @@ class ServerlessConfigRouteTests(_ConfigRouteFixture):
self.config["workers"] = [{
"id": "pool-1", "name": "Pool", "type": "serverless", "provider": "runpod",
"endpoint_id": "ep-123", "jobs_per_run": 2, "concurrency": 2, "timeout_seconds": 600,
"media_profile": "comfy-media-v1",
}]
self.assertTrue(config_module.save_config(self.config))
@@ -353,7 +427,7 @@ class ServerlessConfigRouteTests(_ConfigRouteFixture):
self.assertIsNone(self.secrets.get("pool-1"))
worker = self.read_config_file()["workers"][0]
self.assertEqual(worker["type"], "local")
for serverless_field in ("provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds"):
for serverless_field in ("provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds", "media_profile"):
self.assertNotIn(serverless_field, worker)
async def test_enabled_toggle_allowed_during_active_run_but_endpoint_change_is_rejected(self):
@@ -575,19 +649,41 @@ class PersistedPoolGuardTests(_ConfigRouteFixture):
runtime_module.manager = self.manager
else:
self.manager.tasks["retained"] = types.SimpleNamespace(done=lambda: True)
for fields in ({"api_key": "new-key"}, {"clear_api_key": True}, {"endpoint_id": "ep-other"}, {"type": "local", "port": 8188}):
for fields in ({"api_key": "new-key"}, {"clear_api_key": True}, {"endpoint_id": "ep-other"}, {"type": "local", "port": 8188}, {"media_profile": "comfy-media-v1"}):
response = await config_routes.update_worker_endpoint(_FakeRequest(dict(fields, worker_id="pool-1")))
self.assertEqual(response.status, 409, fields)
response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"}))
self.assertEqual(response.status, 409)
worker = original["workers"][0]
for workers in ([], [dict(worker, type="local")], [dict(worker, endpoint_id="ep-other")]):
for workers in ([], [dict(worker, type="local")], [dict(worker, endpoint_id="ep-other")], [dict(worker, media_profile="comfy-media-v1")]):
response = await config_routes.update_config_endpoint(_FakeRequest({"workers": workers}))
self.assertEqual(response.status, 409, workers)
self.assertEqual(self.read_config_file(), original)
self.assertEqual(self.secrets.get("pool-1"), "saved-key")
self.session_factory.assert_not_awaited()
async def test_retained_jobs_allow_same_or_omitted_media_profile(self):
self.manager.ledger.update("retained", worker_id="pool-1", endpoint_id="ep-123", status="UNKNOWN")
for profile in ("legacy-image-v1", "comfy-media-v1"):
with self.subTest(profile=profile):
worker = {"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123"}
# An absent legacy field and an explicit legacy selection are equivalent.
if profile != "legacy-image-v1":
worker["media_profile"] = profile
self.config["workers"] = [worker]
self.assertTrue(config_module.save_config(self.config))
for fields in ({"media_profile": profile}, {}):
response = await config_routes.update_worker_endpoint(_FakeRequest({
"worker_id": "pool-1", "name": "Same profile", **fields,
}))
self.assertEqual(response.status, 200)
response = await config_routes.update_config_endpoint(_FakeRequest({"workers": [{
"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", **fields,
}]}))
self.assertEqual(response.status, 200)
self.assertEqual(self.read_config_file()["workers"][0]["media_profile"], profile)
self.session_factory.assert_not_awaited()
async def test_retained_jobs_allow_images_per_run_and_legacy_concurrency_saves(self):
# Upgrading from a config with concurrency=1 must not trip the busy guard,
# and the image count stays editable while old jobs are unresolved.
+330
View File
@@ -0,0 +1,330 @@
"""Bounded input packing; native acceptance is opt-in and CPU-only."""
import asyncio
import base64
import copy
import importlib
import io
import os
from pathlib import Path
import subprocess
import sys
import types
import wave
from unittest.mock import Mock
import pytest
ROOT = Path(__file__).resolve().parents[1]
PACKAGE = "_typed_inputs_tests"
for suffix in ("", ".api", ".api.orchestration", ".serverless_bridge"):
package = types.ModuleType(PACKAGE + suffix)
package.__path__ = [str(ROOT.joinpath(*suffix.strip(".").split(".")))] if suffix else [str(ROOT)]
sys.modules[package.__name__] = package
packer = importlib.import_module(PACKAGE + ".api.orchestration.serverless_typed_inputs")
codec = importlib.import_module(PACKAGE + ".serverless_bridge.codec")
def node(kind, **inputs):
return {"class_type": kind, "inputs": inputs}
def graph(kind="LoadAudio", filename="tone.wav"):
key = "audio" if kind == "LoadAudio" else "file"
socket = "audio" if kind == "LoadAudio" else "video"
return {"source": node(kind, **{key: filename}),
"consumer": node("Consumer", **{socket: ["source", 0]})}
def run(workflow):
return asyncio.run(packer.prepare_media_inputs(workflow))
@pytest.fixture
def paths(tmp_path, monkeypatch):
roots = {name: tmp_path / name for name in ("input", "output", "temp")}
for path in roots.values():
path.mkdir()
def resolve(value):
for name in roots:
marker = " [" + name + "]"
if value.endswith(marker):
return str(roots[name] / value[:-len(marker)])
return str(roots["input"] / value)
folders = types.SimpleNamespace(get_annotated_filepath=resolve)
for name, path in roots.items():
setattr(folders, "get_" + name + "_directory", lambda path=path: str(path))
monkeypatch.setitem(sys.modules, "folder_paths", folders)
return roots
@pytest.mark.parametrize("kind,key", [("LoadAudio", "audio"), ("LoadVideo", "file")])
@pytest.mark.parametrize("bad", [{}, None, [], {"filename": "a.wav"},
{"audio": ["filename", 0]}, {"file": ["filename", 0]}])
def test_rejects_non_native_or_linked_inputs_before_read(kind, key, bad, monkeypatch):
workflow = graph(kind)
workflow["source"]["inputs"] = bad
read = Mock(side_effect=AssertionError("must not read"))
monkeypatch.setattr(packer.video_codec, "_read_path", read)
with pytest.raises(ValueError):
run(workflow)
read.assert_not_called()
@pytest.mark.parametrize("kind,key", [("LoadAudio", "audio"), ("LoadVideo", "file")])
@pytest.mark.parametrize("extra", ["start_time", "duration", "sample_rate", "crop", "upload"])
def test_rejects_extra_loader_controls(kind, key, extra):
workflow = graph(kind)
workflow["source"]["inputs"][extra] = 0
with pytest.raises(ValueError, match="controls|inputs"):
run(workflow)
@pytest.mark.parametrize("slot", [1, -1, True, 0.0, "0", None])
def test_rejects_nonzero_or_noninteger_output_slot(slot):
workflow = graph()
workflow["consumer"]["inputs"]["audio"] = ["source", slot]
with pytest.raises(ValueError, match="output|slot"):
run(workflow)
def test_rejects_unused_loader_and_malformed_link():
workflow = graph()
workflow.pop("consumer")
with pytest.raises(ValueError, match="consumer"):
run(workflow)
workflow = graph()
workflow["consumer"]["inputs"]["audio"].append(0)
with pytest.raises(ValueError):
run(workflow)
@pytest.mark.parametrize("key,value", [("audio", "private.wav"), ("file", "private.mp4"), ("video", "private.mov")])
def test_unknown_media_loaders_still_use_strict_scan(key, value):
with pytest.raises(ValueError):
run({"source": node("CustomLoader", **{key: value})})
def test_images_still_rewrite_only_the_private_graph(paths):
payload = b"legacy staging preserves the original file bytes"
(paths["input"] / "picture.png").write_bytes(payload)
source = {"image": node("LoadImage", image="picture.png"),
"export": node("MediaExportV1", images=["image", 0])}
workflow = copy.deepcopy(source)
images, artifacts, refs = run(workflow)
assert artifacts == refs == []
assert len(images) == 1
assert base64.b64decode(images[0]["image"]) == payload
assert workflow["image"]["inputs"]["image"] == images[0]["name"]
assert workflow["export"] == source["export"]
assert source["image"]["inputs"]["image"] == "picture.png"
if os.environ.get("_DISTRIBUTED_TYPED_INPUT_CORE_TEST") != "1":
def test_native_input_acceptance_cpu_subprocess():
configured = os.environ.get("COMFYUI_SOURCE_ROOT")
if configured is None:
pytest.skip("Set COMFYUI_SOURCE_ROOT for CPU native input acceptance")
core = Path(configured).expanduser().resolve()
assert configured and (core / "comfy_extras/nodes_audio.py").is_file()
env = dict(os.environ)
scratch = Path(os.environ.get("TMPDIR", str(Path.home() / ".hermes/cache/scratch")))
scratch.mkdir(parents=True, exist_ok=True)
env.pop("PYTHONPATH", None)
env.pop("COMFYUI_SOURCE_ROOT", None)
env.update(TMPDIR=str(scratch), CUDA_VISIBLE_DEVICES="", HIP_VISIBLE_DEVICES="",
PYTHONDONTWRITEBYTECODE="1", _DISTRIBUTED_TYPED_INPUT_CORE_TEST="1")
bootstrap = """
import sys
sys.path.insert(0, sys.argv[1])
import comfy.cli_args
comfy.cli_args.args.cpu = True
import torch
def forbidden(*args, **kwargs):
raise AssertionError('Input packing must not query or initialize a GPU')
torch.cuda.is_available = lambda: False
torch.cuda.device_count = lambda: 0
for name in ('init', '_lazy_init', 'current_device'):
setattr(torch.cuda, name, forbidden)
import comfy_extras.nodes_audio
import comfy_api.latest._input_impl.video_types
for name in ('is_available', 'device_count'):
setattr(torch.cuda, name, forbidden)
import pytest
raise SystemExit(pytest.main([sys.argv[2], '-q', '-p', 'no:cacheprovider']))
"""
result = subprocess.run([sys.executable, "-c", bootstrap, str(core), str(Path(__file__).resolve())],
cwd=ROOT, env=env, text=True, capture_output=True, timeout=150)
print(result.stdout)
assert result.returncode == 0, result.stdout + result.stderr
else:
import av
import numpy as np
import torch
from comfy_extras import nodes_audio
from comfy_api.latest._input_impl.video_types import VideoFromFile
def wav_bytes(channels=2, sample_width=2, rate=22050, variant=0):
values = np.array([-32768, -12001 + variant, -1, 0, 1, 8193, 32767] * channels,
dtype="<i2" if sample_width == 2 else "<i4")
buffer = io.BytesIO()
with wave.open(buffer, "wb") as out:
out.setnchannels(channels)
out.setsampwidth(sample_width)
out.setframerate(rate)
out.writeframes(values.tobytes())
return buffer.getvalue()
def unpack(artifact):
payload = base64.b64decode(artifact["transport"]["data_base64"])
assert artifact["byte_length"] == len(payload)
return payload
@pytest.mark.parametrize("root", ["input", "output", "temp"])
@pytest.mark.parametrize("channels,width", [(1, 2), (2, 2), (2, 4)])
def test_exact_native_audio_and_every_consumer_ref(paths, root, channels, width):
source = paths[root] / "tone.wav"
source.write_bytes(wav_bytes(channels, width))
workflow = graph(filename="tone.wav [" + root + "]")
workflow["export"] = node("MediaExportV1", audio=["source", 0])
workflow["consumer"]["inputs"]["second_audio"] = ["source", 0]
before = copy.deepcopy(workflow)
images, artifacts, refs = run(workflow)
assert images == [] and len(artifacts) == 1
artifact = artifacts[0]
expected, rate = nodes_audio.load(str(source))
actual = codec.decode_payload("AUDIO", artifact["descriptor"], unpack(artifact))
assert actual["sample_rate"] == rate == 22050
assert actual["waveform"].dtype == expected.dtype
assert torch.equal(actual["waveform"].view(torch.uint8), expected.unsqueeze(0).contiguous().view(torch.uint8))
assert refs == [{"node_id": "consumer", "input_name": "audio", "artifact_id": artifact["id"]},
{"node_id": "consumer", "input_name": "second_audio", "artifact_id": artifact["id"]},
{"node_id": "export", "input_name": "audio", "artifact_id": artifact["id"]}]
assert workflow == before
@pytest.mark.parametrize("root", ["input", "output", "temp"])
def test_exact_video_bytes_without_container_decode(paths, root, monkeypatch):
payload = b"opaque native video bytes\x00\xff"
(paths[root] / "clip.mp4").write_bytes(payload)
workflow = graph("LoadVideo", "clip.mp4 [" + root + "]")
workflow["export"] = node("MediaExportV1", video=["source", 0])
before = copy.deepcopy(workflow)
def forbidden(*args, **kwargs):
raise AssertionError("must not open, decode or transcode video")
monkeypatch.setattr(av, "open", forbidden)
monkeypatch.setattr(VideoFromFile, "get_components", forbidden)
images, artifacts, refs = run(workflow)
assert images == [] and len(artifacts) == 1
artifact = artifacts[0]
assert unpack(artifact) == payload
assert artifact["descriptor"] == {"encoding": "native-video-v1", "representation": "file",
"byte_length": len(payload), "start_time": 0, "duration": 0, "crop": None,
"data": {"offset": 0, "length": len(payload)}}
assert refs == [{"node_id": name, "input_name": "video", "artifact_id": artifact["id"]}
for name in ("consumer", "export")]
native = packer.video_codec.decode_video(artifact["descriptor"], unpack(artifact))
assert vars(native)["_VideoFromFile__file"].getvalue() == payload
assert workflow == before
def test_content_dedupe_kind_separation_and_mixed_images(paths):
data = wav_bytes()
(paths["input"] / "a.wav").write_bytes(data)
(paths["output"] / "b.wav").write_bytes(data)
(paths["input"] / "image.png").write_bytes(b"image bytes")
workflow = graph(filename="a.wav")
workflow["same"] = node("LoadAudio", audio="b.wav [output]")
workflow["video"] = node("LoadVideo", file="a.wav")
workflow["image"] = node("LoadImage", image="image.png")
workflow["export"] = node("MediaExportV1", audio=["same", 0], video=["video", 0], images=["image", 0])
before = copy.deepcopy(workflow)
images, artifacts, refs = run(workflow)
assert len(images) == 1 and len(artifacts) == 2 and len(refs) == 3
assert refs[0]["artifact_id"] == refs[1]["artifact_id"] != refs[2]["artifact_id"]
assert {a["kind"] for a in artifacts} == {"AUDIO", "VIDEO"}
before["image"]["inputs"]["image"] = images[0]["name"]
assert workflow == before
second = graph(filename="b.wav [output]")
assert run(second)[1][0]["id"] == artifacts[0]["id"]
@pytest.mark.parametrize("kind", ["LoadAudio", "LoadVideo"])
def test_path_rejections_and_per_file_bounds(paths, tmp_path, monkeypatch, kind):
outside = tmp_path / "outside.wav"
outside.write_bytes(wav_bytes())
(paths["input"] / "escape.wav").symlink_to(outside)
(paths["input"] / "dirlink").symlink_to(tmp_path, target_is_directory=True)
os.mkfifo(paths["input"] / "fifo.wav")
values = ["../outside.wav", str(outside), "escape.wav", "dirlink/outside.wav", "fifo.wav",
"absent.wav", "tone.wav [unknown]", "https://example.com/a.wav", "input [temp]", "C:\\a.wav"]
for value in values:
with pytest.raises(ValueError):
run(graph(kind, value))
(paths["input"] / "large.wav").write_bytes(wav_bytes())
monkeypatch.setattr(packer.contract, "MAX_ARTIFACT_BYTES", 8)
with pytest.raises(ValueError, match="limit|size"):
run(graph(kind, "large.wav"))
def test_aggregate_decoded_budget_and_duplicate_not_double_charged(paths, monkeypatch):
payload = wav_bytes()
(paths["input"] / "a.wav").write_bytes(payload)
(paths["input"] / "b.wav").write_bytes(wav_bytes(variant=1))
workflow = graph(filename="a.wav")
workflow["other"] = node("LoadAudio", audio="a.wav")
workflow["export"] = node("MediaExportV1", audio=["other", 0])
# File source bytes have a separate aggregate bound; use a longer WAV
# whose float output dominates its tiny header.
with wave.open(io.BytesIO(payload), "rb") as wav:
frames = wav.readframes(wav.getnframes())
buffer = io.BytesIO()
with wave.open(buffer, "wb") as wav:
wav.setparams((2, 2, 22050, 0, "NONE", "not compressed"))
wav.writeframes(frames * 100)
(paths["input"] / "a.wav").write_bytes(buffer.getvalue())
artifact = run(graph(filename="a.wav"))[1][0]
monkeypatch.setattr(packer.contract, "MAX_TOTAL_BYTES", artifact["byte_length"])
assert len(run(workflow)[1]) == 1
workflow["other"]["inputs"]["audio"] = "b.wav"
with pytest.raises(ValueError, match="limit|size|budget"):
run(workflow)
def test_audio_charged_before_ndarray_cat_or_conversion(monkeypatch):
frame = types.SimpleNamespace(samples=100, sample_rate=22050,
layout=types.SimpleNamespace(channels=[1, 2]),
format=types.SimpleNamespace(name="s16", bytes=2),
planes=[types.SimpleNamespace(buffer_size=400)], to_ndarray=Mock())
stream = types.SimpleNamespace(index=0, channels=2,
codec_context=types.SimpleNamespace(sample_rate=22050))
container = Mock()
container.__enter__ = Mock(return_value=container)
container.__exit__ = Mock(return_value=False)
container.streams = types.SimpleNamespace(audio=[stream])
container.decode = Mock(return_value=iter([frame]))
monkeypatch.setattr(av, "open", Mock(return_value=container))
cat = Mock(side_effect=AssertionError("no accumulation allowed"))
monkeypatch.setattr(torch, "cat", cat)
with pytest.raises(ValueError, match="limit|size|budget"):
packer._decode_audio(b"input", 799)
frame.to_ndarray.assert_not_called()
cat.assert_not_called()
frame.samples = 1
frame.planes[0].buffer_size = 800
container.decode.return_value = iter([frame])
with pytest.raises(ValueError, match="limit|size|budget"):
packer._decode_audio(b"input", 799)
frame.to_ndarray.assert_not_called()
def test_audio_rejects_empty_invalid_and_channel_or_rate_drift(paths, monkeypatch):
(paths["input"] / "bad.wav").write_bytes(b"not audio")
with pytest.raises(ValueError):
run(graph(filename="bad.wav"))
stream = types.SimpleNamespace(index=0, channels=2, codec_context=types.SimpleNamespace(sample_rate=22050))
for frames in ([], [types.SimpleNamespace(samples=1, sample_rate=44100,
layout=types.SimpleNamespace(channels=[1, 2]))],
[types.SimpleNamespace(samples=1, sample_rate=22050,
layout=types.SimpleNamespace(channels=[1]))]):
container = Mock()
container.__enter__ = Mock(return_value=container)
container.__exit__ = Mock(return_value=False)
container.streams = types.SimpleNamespace(audio=[stream])
container.decode = Mock(return_value=iter(frames))
monkeypatch.setattr(av, "open", Mock(return_value=container))
with pytest.raises(ValueError):
packer._decode_audio(b"input", 1000)
+319
View File
@@ -0,0 +1,319 @@
"""Offline typed submission/delivery regressions; no provider or GPU access."""
import asyncio
import copy
import importlib
import json
import socket
import tempfile
import types
import unittest
from pathlib import Path
from unittest.mock import AsyncMock, patch
# Load once before integration fixtures snapshot/restore sys.modules. PyTorch
# cannot be re-imported after those fixtures unload its extension-backed modules.
import torch
from test_runpod_client import Response, Session
from test_serverless_runtime import runtime, Secrets
client = importlib.import_module(runtime.__package__ + ".runpod_client")
bridge = runtime.__package__.split(".")[0] + ".serverless_bridge"
contract = importlib.import_module(bridge + ".contract")
def prepared(sockets=("texts",)):
worker = {"id": "pool:1", "pool_id": "pool", "endpoint_id": "ep1",
"timeout_seconds": 30, "type": "serverless", "media_profile": contract.PROFILE}
bindings = [dict(binding_id=f"binding-{i}", collector_id="collector", job_id="collector-job",
participant_id=worker["id"], node_id=f"export-{i}", socket=name,
kind=contract.SOCKET_KINDS[name]) for i, name in enumerate(sockets)]
media = dict(version=1, profile=contract.PROFILE, request_id="request-1",
artifacts=[], inputs=[], outputs=copy.deepcopy(bindings))
graph = {"source": {"class_type": "Source", "inputs": {"value": "café雪"}}}
for binding in bindings:
graph[binding["node_id"]] = {"class_type": "MediaExportV1", "inputs": {
"request_id": media["request_id"], "binding_json": json.dumps(binding),
binding["socket"]: ["source", 0]}}
return dict(worker=worker, workflow=graph, images=[], manifest=bindings, media=media)
def response(job):
envelope = dict(version=1, profile=contract.PROFILE, request_id=job["media"]["request_id"],
bindings=[], artifacts=[], items=[])
for index, binding in enumerate(job["manifest"]):
envelope["bindings"].append(dict(binding, count=2))
for item_index, text in enumerate(("", " \n雪")):
aid = f"artifact-{index}-{item_index}"
descriptor, payload = {"encoding": "utf8-v1"}, text.encode()
if binding["kind"] != "STRING":
codec = importlib.import_module(bridge + ".codec")
values = {"IMAGE": torch.arange(8, dtype=torch.float64).reshape(1, 1, 2, 4),
"MASK": torch.arange(2, dtype=torch.float16).reshape(1, 1, 2),
"AUDIO": {"waveform": torch.arange(4, dtype=torch.float32).reshape(1, 2, 2), "sample_rate": 48000}}
descriptor, payload = codec.encode_payload(binding["kind"], values[binding["kind"]])
artifact = contract.artifact_from_payload(aid, binding["kind"], descriptor, payload)
envelope["artifacts"].append(artifact)
envelope["items"].insert(0, dict(binding_id=binding["binding_id"], item_index=item_index, artifact_id=aid))
return {"media": envelope}
class TypedClientTests(unittest.IsolatedAsyncioTestCase):
async def test_exact_typed_bytes_and_legacy_unchanged(self):
job = prepared()
session = Session([Response(payload={"id": "job-1"})])
api = client.RunpodClient(session, "ep1", "test-only")
await api.submit(job["workflow"], [], 30, media=job["media"])
expected = {"input": {"workflow": job["workflow"], "images": [], "media": job["media"]},
"policy": {"executionTimeout": 30000, "ttl": 30000}}
self.assertEqual(session.calls[0][2]["data"], json.dumps(expected, allow_nan=False,
ensure_ascii=False, separators=(",", ":")).encode())
self.assertEqual(len(session.calls), 1)
expected["input"].pop("media")
legacy = json.dumps(expected, allow_nan=False, ensure_ascii=False, separators=(",", ":")).encode()
self.assertEqual(client.serialize_request(job["workflow"], [], 30), legacy)
self.assertEqual(client.serialize_request(job["workflow"], [], 30, media=None), legacy)
async def test_invalid_envelopes_never_post(self):
for field, value in [("version", True), ("profile", "legacy-image-v1"),
("request_id", "../bad"), ("outputs", []), ("inputs", [{}])]:
with self.subTest(field=field):
job = prepared()
job["media"][field] = value
session = Session([])
with self.assertRaises(ValueError):
await client.RunpodClient(session, "ep1", "test-only").submit(
job["workflow"], [], 30, media=job["media"])
self.assertFalse(session.calls)
session = Session([])
with self.assertRaises(ValueError):
await client.RunpodClient(session, "ep1", "test-only").submit({}, [], 30, media={})
self.assertFalse(session.calls)
async def test_typed_size_boundary_uses_exact_body_before_post(self):
job = prepared()
size = len(client.serialize_request(job["workflow"], [], 30, media=job["media"]))
session = Session([Response(payload={"id": "boundary"})])
api = client.RunpodClient(session, "ep1", "test-only")
with patch.object(client, "MAX_REQUEST_BYTES", size):
await api.submit(job["workflow"], [], 30, media=job["media"])
job["workflow"]["source"]["inputs"]["value"] += "x"
with self.assertRaisesRegex(ValueError, "Serialized serverless request body"):
await api.submit(job["workflow"], [], 30, media=job["media"])
self.assertEqual(len(session.calls), 1)
class TypedRuntimeTests(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
self.temp = tempfile.TemporaryDirectory()
self.server = types.SimpleNamespace(distributed_pending_jobs={"collector-job": asyncio.Queue()},
distributed_jobs_lock=asyncio.Lock())
self.manager = runtime.ServerlessManager(self.server, Secrets(), Path(self.temp.name) / "jobs.json",
AsyncMock(return_value=object()))
self.manager.poll_seconds = 0
self.job = prepared()
self.api = types.SimpleNamespace(submit=AsyncMock(return_value="provider-job"),
status=AsyncMock(return_value={"status": "COMPLETED", "output": response(self.job)}),
cancel=AsyncMock(return_value={"status": "CANCELLED"}))
self.client_patch = patch.object(runtime, "RunpodClient", return_value=self.api)
self.client_patch.start()
async def asyncTearDown(self):
await self.manager.cancel()
self.client_patch.stop()
self.temp.cleanup()
async def run_job(self):
await self.manager.start("execution", "prompt", [self.job])
await self.manager.wait_for_execution("execution")
async def test_preflight_all_jobs_before_any_journal_or_task(self):
for mutation in ("profile", "missing_media", "manifest", "participant", "exporter", "empty_media"):
with self.subTest(mutation=mutation):
bad = prepared()
if mutation == "profile":
bad["worker"]["media_profile"] = "legacy-image-v1"
elif mutation == "missing_media":
bad.pop("media")
elif mutation == "manifest":
bad["manifest"][0]["job_id"] = "other"
elif mutation == "participant":
bad["worker"]["id"] = "other-participant"
elif mutation == "exporter":
bad["workflow"]["export-0"]["inputs"]["request_id"] = "other"
else:
bad["media"] = {}
with self.assertRaises(ValueError):
await self.manager.start("execution", "prompt", [self.job, bad])
self.assertFalse(self.manager.ledger.all())
self.assertFalse(self.manager.tasks)
self.assertFalse(self.manager.runs)
self.api.submit.assert_not_awaited()
async def test_text_roundtrip_id_journal_and_payload_cleanup(self):
async def status(_job):
self.assertEqual(self.manager.ledger.all()[0]["job_id"], "provider-job")
return {"status": "COMPLETED", "output": response(self.job)}
self.api.status.side_effect = status
await self.run_job()
self.api.submit.assert_awaited_once_with(self.job["workflow"], [], 30, media=self.job["media"])
queue = self.server.distributed_pending_jobs["collector-job"]
self.assertEqual(queue.get_nowait(), {"worker_id": "pool:1", "media": {"texts": ["", " \n雪"]}, "is_last": True})
self.assertTrue(queue.empty())
self.assertEqual(self.manager.ledger.all()[0]["status"], "COMPLETED")
run = next(iter(self.manager.runs.values()))
for field in ("workflow", "images", "media"):
self.assertNotIn(field, run)
self.assertNotIn("data_base64", self.manager.ledger.path.read_text())
self.api.cancel.assert_not_awaited()
async def test_all_sockets_grouped_without_native_value_conversion(self):
self.job = prepared(tuple(contract.SOCKET_KINDS))
native = {binding["binding_id"]: [object(), {"native": object()}] for binding in self.job["manifest"]}
with patch.object(contract, "decode_response", AsyncMock(return_value=native)) as decode:
await self.run_job()
decode.assert_awaited_once()
self.assertEqual(decode.call_args.kwargs["request_id"], "request-1")
self.assertEqual(decode.call_args.kwargs["bindings"], self.job["manifest"])
queue = self.server.distributed_pending_jobs["collector-job"]
item = queue.get_nowait()
self.assertEqual(set(item), {"worker_id", "media", "is_last"})
self.assertEqual(set(item["media"]), set(contract.SOCKET_KINDS))
for binding in self.job["manifest"]:
for actual, expected in zip(item["media"][binding["socket"]], native[binding["binding_id"]]):
self.assertIs(actual, expected)
self.assertTrue(queue.empty())
async def test_identity_and_late_decode_failure_deliver_only_one_error(self):
self.job = prepared(("texts", "images"))
for mutation in ("request", "binding", "checksum", "descriptor"):
with self.subTest(mutation=mutation):
output = response(self.job)
if mutation == "request":
output["media"]["request_id"] = "other"
elif mutation == "binding":
output["media"]["bindings"][-1]["participant_id"] = "other"
elif mutation == "checksum":
output["media"]["artifacts"][-1]["sha256"] = "0" * 64
else:
output["media"]["artifacts"][-1]["descriptor"] = {}
self.api.status.return_value = {"status": "COMPLETED", "output": output}
await self.run_job()
queue = self.server.distributed_pending_jobs["collector-job"]
item = queue.get_nowait()
self.assertIn("serverless_error", item)
self.assertNotIn("media", item)
self.assertTrue(queue.empty())
self.api.cancel.assert_not_awaited()
async def test_closed_queue_is_not_resurrected(self):
self.server.distributed_pending_jobs.clear()
await self.run_job()
self.assertFalse(self.server.distributed_pending_jobs)
self.assertEqual(self.manager.ledger.all()[0]["status"], "COMPLETED")
async def test_real_tensor_audio_and_mask_codecs_preserve_dtype_and_shape(self):
self.job = prepared(("texts", "images", "audio", "masks"))
self.api.status.return_value = {"status": "COMPLETED", "output": response(self.job)}
await self.run_job()
result = self.server.distributed_pending_jobs["collector-job"].get_nowait()["media"]
self.assertEqual(result["texts"], ["", " \n雪"])
self.assertEqual(result["images"][0].dtype, torch.float64)
self.assertEqual(result["images"][0].shape, (1, 1, 2, 4))
self.assertEqual(result["masks"][0].dtype, torch.float16)
self.assertEqual(result["audio"][0]["sample_rate"], 48000)
self.assertTrue(torch.equal(result["audio"][0]["waveform"], torch.arange(4, dtype=torch.float32).reshape(1, 2, 2)))
self.assertTrue(all(len(values) == 2 for values in result.values()))
async def test_run_rechecks_profile_before_submission(self):
await self.manager.start("execution", "prompt", [self.job])
self.job["worker"]["media_profile"] = "legacy-image-v1"
await self.manager.wait_for_execution("execution")
self.api.submit.assert_not_awaited()
self.assertEqual(self.manager.ledger.all()[0]["status"], "FAILED")
async def test_multiple_collectors_receive_one_message_each(self):
self.job = prepared(("texts", "images"))
self.job["manifest"][-1]["job_id"] = "second-job"
self.job["media"]["outputs"] = copy.deepcopy(self.job["manifest"])
self.job["workflow"]["export-1"]["inputs"]["binding_json"] = json.dumps(self.job["manifest"][-1])
self.server.distributed_pending_jobs["second-job"] = asyncio.Queue()
self.api.status.return_value = {"status": "COMPLETED", "output": response(self.job)}
await self.run_job()
for job_id, name in (("collector-job", "texts"), ("second-job", "images")):
queue = self.server.distributed_pending_jobs[job_id]
self.assertEqual(queue.qsize(), 1)
item = queue.get_nowait()
self.assertEqual(set(item["media"]), {name})
self.assertTrue(item["is_last"])
async def test_prestart_cancel_releases_media_and_deduplicates_errors(self):
self.job = prepared(("texts", "images"))
await self.manager.start("execution", "prompt", [self.job])
await self.manager.cancel(execution_id="execution")
self.api.submit.assert_not_awaited()
run = next(iter(self.manager.runs.values()))
self.assertNotIn("media", run)
self.assertEqual(self.server.distributed_pending_jobs["collector-job"].qsize(), 1)
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED")
async def test_ambiguous_submit_keeps_reservation_without_media(self):
self.api.submit.side_effect = runtime.AmbiguousSubmission("Unknown submission")
await self.run_job()
self.api.submit.assert_awaited_once()
self.api.cancel.assert_not_awaited()
self.assertEqual(self.manager.ledger.all()[0]["status"], "UNKNOWN")
self.assertTrue(self.manager.pool_is_active("pool"))
self.assertNotIn("media", next(iter(self.manager.runs.values())))
async def test_legacy_still_uses_tensor_messages_and_no_media_kwarg(self):
self.job.pop("media")
self.job["worker"]["media_profile"] = "legacy-image-v1"
self.job["manifest"] = [{"job_id": "collector-job", "prefix": "dist_x", "kind": "image"}]
with patch.object(runtime, "decode_results", AsyncMock(return_value={"collector-job": ["a", "b"]})):
await self.run_job()
self.api.submit.assert_awaited_once_with(self.job["workflow"], [], 30)
queue = self.server.distributed_pending_jobs["collector-job"]
self.assertEqual([queue.get_nowait(), queue.get_nowait()], [
{"worker_id": "pool:1", "tensor": "a", "image_index": 0, "is_last": False},
{"worker_id": "pool:1", "tensor": "b", "image_index": 1, "is_last": True}])
class TypedDownloadTests(unittest.IsolatedAsyncioTestCase):
async def test_stream_bound_cap_and_redirect_policy_without_credentials(self):
typed = importlib.import_module(runtime.__package__ + ".serverless_typed_results")
class DownloadResponse(Response):
headers = {}
async def chunks(self, _size):
yield b"ab"
yield b"cd"
@property
def content(self):
return types.SimpleNamespace(iter_chunked=self.chunks)
class DownloadSession(Response):
def get(self, url, **kwargs):
self.calls.append((url, kwargs))
return self.response
session = DownloadSession()
session.calls = []
session.response = DownloadResponse()
with patch.object(typed.aiohttp, "TCPConnector", return_value=object()), \
patch.object(typed.aiohttp, "ClientSession", return_value=session) as factory:
self.assertEqual(await typed.download_public_artifact("https://example.com/a", max_bytes=contract.MAX_ARTIFACT_BYTES), b"abcd")
self.assertFalse(factory.call_args.kwargs["trust_env"])
self.assertNotIn("headers", factory.call_args.kwargs)
self.assertFalse(session.calls[0][1]["allow_redirects"])
with self.assertRaises(ValueError):
await typed.download_public_artifact("https://example.com/a", max_bytes=3)
session.response.status = 302
session.response.headers = {"Location": "https://127.0.0.1/private"}
with self.assertRaises(ValueError):
await typed.download_public_artifact("https://example.com/a", max_bytes=4)
for url in ("http://example.com/a", "https://127.0.0.1/a", "https://user:pass@example.com/a"):
with self.assertRaises(ValueError):
await typed.download_public_artifact(url, max_bytes=4)
resolver = typed._PublicResolver()
with patch.object(resolver, "_getaddrinfo", AsyncMock(return_value=[
(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("127.0.0.1", 443))])):
with self.assertRaises(OSError):
await resolver.resolve("example.com", 443)
+17
View File
@@ -133,6 +133,7 @@ def test_quality_to_crf_mapping():
def test_normalize_format_rejects_unknown_values():
assert tc.normalize_format(" H264 ") == "h264"
assert tc.normalize_format(" EXACT ") == "exact"
assert tc.normalize_format(None) == tc.DEFAULT_TRANSFER_FORMAT
with pytest.raises(tc.TransferCodecError):
tc.normalize_format("gif")
@@ -158,6 +159,22 @@ def test_decode_chunk_validates_envelope_against_payload():
tc.decode_chunk({**video_env, "count": 5}, video["parts"])
def test_exact_is_appended_without_changing_saved_format_defaults():
assert tc.TRANSFER_FORMATS == (
"legacy_png", "png", "ffv1", "jpeg", "webp", "h264", "h265", "av1", "exact",
)
assert tc.DEFAULT_TRANSFER_FORMAT == "legacy_png"
assert tc.EXACT_FORMAT not in tc.IMAGE_FORMATS + tc.VIDEO_FORMATS
def test_exact_does_not_enter_uint8_chunk_codec():
frames = tc.tensor_batch_to_uint8(_smooth_frames(count=1))
with pytest.raises(tc.TransferCodecError, match="exact"):
list(tc.encode_chunks(frames, "exact", 90, 10**9))
with pytest.raises(tc.TransferCodecError, match="exact"):
tc.decode_chunk({"format": "exact", "count": 1, "width": 91, "height": 67}, [b"raw tensor"])
def test_legacy_format_is_not_chunk_encoded():
frames = tc.tensor_batch_to_uint8(_smooth_frames(count=1))
with pytest.raises(tc.TransferCodecError):
+7 -4
View File
@@ -69,7 +69,8 @@ def v3(monkeypatch):
module(f"{PKG}.nodes.utilities")
module(f"{PKG}.nodes.distributed_upscale", UltimateSDUpscaleDistributed=object)
module(f"{PKG}.nodes.collector", DistributedCollectorNode=CollectorRuntime)
module(f"{PKG}.nodes.image_list", DistributedImageListStart=object, DistributedImageListEnd=object)
module(f"{PKG}.nodes.image_list", DistributedImageListStart=object, DistributedImageListEnd=object,
DistributedImageListGetValue=object)
module("comfy", __path__=[])
module("comfy.samplers", KSampler=types.SimpleNamespace(SAMPLERS=["euler"], SCHEDULERS=["normal"]))
module("comfy_api", __path__=[])
@@ -97,6 +98,7 @@ def test_registered_collector_schema_exposes_optional_transfer_inputs(v3):
fmt = inputs["transfer_format"]
assert fmt.kind == "Combo"
assert fmt.kwargs["options"] == list(codec.TRANSFER_FORMATS)
assert fmt.kwargs["options"][-1] == "exact"
assert fmt.kwargs["default"] == "legacy_png"
assert fmt.kwargs["optional"] is True
quality = inputs["transfer_quality"]
@@ -106,13 +108,14 @@ def test_registered_collector_schema_exposes_optional_transfer_inputs(v3):
assert schema.is_input_list is True
def test_registered_collector_forwards_transfer_options(v3):
@pytest.mark.parametrize("fmt", ["av1", "exact"])
def test_registered_collector_forwards_transfer_options(v3, fmt):
mod, _, calls = v3
mod.DistributedCollector.execute(images=["img"], audio=[None], transfer_format=["av1"], transfer_quality=[70],
mod.DistributedCollector.execute(images=["img"], audio=[None], transfer_format=[fmt], transfer_quality=[70],
multi_job_id=["job"], is_worker=[True])
assert calls[-1]["transfer_format"] == ["av1"]
assert calls[-1]["transfer_format"] == [fmt]
assert calls[-1]["transfer_quality"] == [70]
assert calls[-1]["multi_job_id"] == ["job"]
+14
View File
@@ -1,4 +1,5 @@
"""Real-framework acceptance; opt in with COMFYUI_SOURCE_ROOT."""
import json
import os
from pathlib import Path
import subprocess
@@ -33,5 +34,18 @@ def test_image_list_network_runtime():
)
assert result.returncode == 0, result.stdout + '\n' + result.stderr
assert 'IMAGE_LIST_NETWORK_OK' in result.stdout
assert 'IMAGE_LIST_VALUES_OK' in result.stdout
assert 'IMAGE_LIST_CANCEL_OK' in result.stdout
assert 'IMAGE_LIST_MASTER_OK' in result.stdout
assert 'IMAGE_LIST_PREFLIGHT_OK' in result.stdout
assert 'IMAGE_LIST_SEQUENTIAL_OK' in result.stdout
assert 'IMAGE_LIST_FAILURE_OK' in result.stdout
reports = [json.loads(line.removeprefix('IMAGE_LIST_TARGETS_OK '))
for line in result.stdout.splitlines() if line.startswith('IMAGE_LIST_TARGETS_OK ')]
assert len(reports) == 1, result.stdout
assert reports[0] == {
'modes': ['delegate-only', 'mixed', 'master-only'],
'outer_targets': ['4', '5'], 'worker_targets': ['4'],
'required_utility_once_per_item': True, 'preview_discarded': True,
'saved_prompt_unchanged': True, 'capture_once_per_outer_prompt': True,
}
+41 -2
View File
@@ -105,24 +105,63 @@ class PromptValidationError(RuntimeError):
super().__init__(f"Invalid prompt: {merged}")
def validate_execution_targets(prompt, targets):
"""Reject malformed explicit roots instead of falling back to all outputs."""
if (not isinstance(targets, list) or not targets
or any(not isinstance(node_id, str) or not node_id or node_id not in prompt
for node_id in targets)
or len(set(targets)) != len(targets)):
raise PromptValidationError({
"type": "invalid_partial_execution_targets",
"message": "partial_execution_targets must be a non-empty list of unique existing node IDs",
"details": "", "extra_info": {},
})
return list(targets)
async def queue_prompt_payload(
prompt_obj,
workflow_meta=None,
client_id=None,
include_queue_metadata=False,
*,
partial_execution_targets=None,
):
"""Validate and queue a prompt via ComfyUI's prompt queue."""
"""Validate and queue a prompt, preserving any explicitly requested roots."""
targets = (validate_execution_targets(prompt_obj, partial_execution_targets)
if partial_execution_targets is not None else None)
payload = {"prompt": prompt_obj}
if targets is not None:
payload["partial_execution_targets"] = list(targets)
payload = prompt_server.trigger_on_prompt(payload)
prompt = payload["prompt"]
if targets is not None:
# Hooks may replace the envelope or mutate its target list. The caller's
# requested roots remain authoritative, but must still exist afterwards.
validate_execution_targets(prompt, targets)
prompt_id = str(uuid.uuid4())
valid = await execution.validate_prompt(prompt_id, prompt, None)
valid = await execution.validate_prompt(prompt_id, prompt, targets)
if not valid[0]:
error_payload = valid[1] if len(valid) > 1 else "Prompt outputs failed validation"
node_errors = valid[3] if len(valid) > 3 else {}
raise PromptValidationError(error_payload, node_errors)
if targets is not None:
node_errors = valid[3] if len(valid) > 3 else {}
target_errors = {
node_id: entry for node_id, entry in node_errors.items()
if node_id in targets or (isinstance(entry, dict)
and set(entry.get("dependent_outputs", [])) & set(targets))
}
if set(valid[2]) != set(targets) or target_errors:
raise PromptValidationError({
"type": "partial_execution_targets_failed_validation",
"message": "Not all requested execution targets passed validation",
"details": "Refusing to queue a subset or additional output roots",
"extra_info": {"requested_targets": targets},
}, node_errors)
extra_data = {"create_time": int(time.time() * 1000)}
if workflow_meta:
extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
+8
View File
@@ -84,6 +84,14 @@ def validate_serverless_worker(worker: dict) -> dict:
if "enabled" in result and type(result["enabled"]) is not bool:
raise ServerlessConfigError("enabled must be a boolean")
# Omission preserves old saved records and means legacy-image-v1. Selecting
# comfy-media-v1 declares an installed endpoint bridge, not health detection;
# its schema/representations/limits are fixed by serverless_bridge/contract.py.
if "media_profile" in result:
profile = result["media_profile"]
if not isinstance(profile, str) or profile not in ("legacy-image-v1", "comfy-media-v1"):
raise ServerlessConfigError("media_profile must be legacy-image-v1 or comfy-media-v1")
for field, (minimum, maximum, default) in _NUMERIC_FIELDS.items():
value = result.get(field, default)
if type(value) is not int or not minimum <= value <= maximum:
+18 -5
View File
@@ -17,9 +17,13 @@ jpeg, webp Lossy per-frame stills.
h264, h265, av1
Lossy video for the whole batch. Smallest by far for video
frame batches, where consecutive frames are similar.
exact Original floating-point tensors (raw little-endian bytes, dtype,
channels and values preserved). Largest; opt-in. Encoded by the
Collector/job_complete route with serverless_bridge.codec, never by
encode_chunks/decode_chunk below.
All formats quantize to 8-bit exactly like the legacy path did, so the
lossless formats reproduce the legacy master tensors bit-for-bit.
All formats except exact quantize to 8-bit exactly like the legacy path did,
so the lossless formats reproduce the legacy master tensors bit-for-bit.
"""
import io
import math
@@ -29,20 +33,25 @@ import numpy as np
from PIL import Image
LEGACY_FORMAT = "legacy_png"
EXACT_FORMAT = "exact"
IMAGE_FORMATS = ("png", "jpeg", "webp")
VIDEO_FORMATS = ("ffv1", "h264", "h265", "av1")
LOSSLESS_FORMATS = (LEGACY_FORMAT, "png", "ffv1")
TRANSFER_FORMATS = (LEGACY_FORMAT, "png", "ffv1", "jpeg", "webp", "h264", "h265", "av1")
# Appended, so the pre-existing option order (and saved widget values) is unchanged.
TRANSFER_FORMATS = (LEGACY_FORMAT, "png", "ffv1", "jpeg", "webp", "h264", "h265", "av1", EXACT_FORMAT)
DEFAULT_TRANSFER_FORMAT = LEGACY_FORMAT
DEFAULT_TRANSFER_QUALITY = 90
TRANSFER_FORMAT_TOOLTIP = (
"How workers send images back to the master. "
"How remote workers send images back to the master; the master's own images are never re-encoded. "
"legacy_png: original uncompressed PNG (largest). "
"png / ffv1: lossless, identical result. "
"jpeg / webp: lossy stills. "
"h264 / h265 / av1: lossy video of the whole batch, smallest for frame sequences."
"h264 / h265 / av1: lossy video of the whole batch, smallest for frame sequences "
"(h265 is experimental). "
"exact: original floating-point values, no 8-bit rounding; much larger, opt-in. "
"Serverless media pools send legacy_png as png (same pixels) and never substitute another format."
)
TRANSFER_QUALITY_TOOLTIP = (
"Quality for lossy transfer formats (ignored by lossless ones). "
@@ -298,6 +307,8 @@ def encode_chunks(frames, fmt, quality, max_chunk_bytes, max_frame_bytes=None):
fmt = normalize_format(fmt)
if fmt == LEGACY_FORMAT:
raise TransferCodecError("legacy_png is sent through the JSON path, not encode_chunks")
if fmt == EXACT_FORMAT:
raise TransferCodecError("exact tensors are not encoded by encode_chunks")
total = int(frames.shape[0])
if total == 0:
return
@@ -378,6 +389,8 @@ def decode_chunk(envelope, parts):
fmt = normalize_format(envelope.get("format"))
if fmt == LEGACY_FORMAT:
raise TransferCodecError("legacy_png chunks are not sent as multipart")
if fmt == EXACT_FORMAT:
raise TransferCodecError("exact chunks are decoded as tensors, not frames")
count = envelope.get("count")
width = envelope.get("width")
height = envelope.get("height")
+174
View File
@@ -0,0 +1,174 @@
import { app } from "/scripts/app.js";
// End's Python schema owns all ten optional override slots. Only presentation
// grows here: unlike Start's Autogrow inputs, End keys must never compact, since
// value_3 always overrides source value_3, even if value_2 is disconnected.
const START = "DistributedImageListStart";
const END = "DistributedImageListEnd";
const GET = "DistributedImageListGetValue";
const PIPE = "DISTRIBUTED_IMAGE_LIST_VALUES";
const VALUE_TYPES = "IMAGE,MASK,STRING,INT,FLOAT,BOOLEAN";
const MAX_VALUES = 10;
const pendingGraphs = new WeakSet();
function ordinal(name) {
const match = /^value_([1-9]|10)$/.exec(name ?? "");
return match ? Number(match[1]) : null;
}
function getLink(graph, id) {
if (id == null) return undefined;
if (typeof graph?.getLink === "function") return graph.getLink(id);
return graph?.links?.get?.(id) ?? graph?.links?.[id];
}
function sourceOf(node, inputName) {
const input = node.inputs?.find(slot => slot.name === inputName);
const link = getLink(node.graph, input?.link);
const source = link && node.graph?.getNodeById?.(link.origin_id);
return source ? { node: source, slot: link.origin_slot, link } : null;
}
function isReroute(node) {
return [node?.comfyClass, node?.type].includes("Reroute");
}
function throughReroutes(source) {
const seen = new Set();
while (source && isReroute(source.node)) {
if (seen.has(source.node)) return null;
seen.add(source.node);
source = sourceOf(source.node, source.node.inputs?.[0]?.name);
}
return source;
}
function sourceType(source, seen) {
source = throughReroutes(source);
if (!source) return "*";
const { node, slot, link } = source;
if (node.comfyClass === GET) return getterType(node, seen);
const type = node.outputs?.[slot]?.type;
return type && type !== "*" ? type : (link.type || "*");
}
function getterType(node, seen = new Set()) {
if (seen.has(node)) return "*";
seen.add(node);
const name = node.widgets?.find(widget => widget.name === "value_name")?.value;
if (!ordinal(name)) return "*";
const endSource = throughReroutes(sourceOf(node, "values"));
if (endSource?.node.comfyClass !== END || endSource.slot !== 3) return "*";
const end = endSource.node;
const override = sourceOf(end, name);
if (override) return sourceType(override, seen);
const startSource = throughReroutes(sourceOf(end, "context"));
if (startSource?.node.comfyClass !== START || startSource.slot !== 3) return "*";
return sourceType(sourceOf(startSource.node, `values.${name}`), seen);
}
function compatible(outputType, inputType) {
if (!outputType || !inputType || outputType === "*" || inputType === "*") return true;
return String(outputType).split(",").some(type => String(inputType).split(",").includes(type));
}
function syncGetter(node) {
const output = node.outputs?.[0];
if (!output) return;
const type = getterType(node);
if (output.type === type) return;
output.type = type;
for (const id of [...(output.links ?? [])]) {
const link = getLink(node.graph, id);
if (!link) continue;
const target = node.graph?.getNodeById?.(link.target_id);
if (target && !compatible(type, target.inputs?.[link.target_slot]?.type)) {
target.disconnectInput(link.target_slot);
} else {
link.type = type;
}
}
node.setDirtyCanvas?.(true, true);
}
function syncEnd(node) {
if (!Array.isArray(node.inputs) || !Array.isArray(node.outputs)) return;
// Old saved End nodes have only the three original outputs. Append rather
// than replacing any socket, preserving every existing output index/link.
if (node.outputs.length === 3) {
const shape = node.outputs[0]?.shape;
node.addOutput("values", PIPE, shape === undefined ? undefined : { shape });
}
const connected = node.inputs.filter(input => ordinal(input.name) && input.link != null);
const highest = Math.max(0, ...connected.map(input => ordinal(input.name)));
const visible = Math.min(MAX_VALUES, Math.max(highest + 1, node._distributedVisibleOverrides ?? 1));
for (let i = node.inputs.length - 1; i >= 0; i--) {
const input = node.inputs[i];
if (ordinal(input.name) > visible && input.link == null) node.removeInput(i);
}
for (let i = 1; i <= visible; i++) {
if (!node.inputs.some(input => input.name === `value_${i}`)) {
node.addInput(`value_${i}`, VALUE_TYPES);
}
}
node.setDirtyCanvas?.(true, true);
}
function schedule(graph) {
if (!graph || pendingGraphs.has(graph)) return;
pendingGraphs.add(graph);
const run = () => {
// Keep the guard while syncing, since disconnectInput can re-enter the
// connection hook. A fresh user event schedules a fresh traversal.
try {
for (const node of graph._nodes ?? []) if (node.comfyClass === END) syncEnd(node);
for (const node of graph._nodes ?? []) if (node.comfyClass === GET) syncGetter(node);
} catch (error) {
console.error("[Distributed] Image List values presentation failed:", error);
} finally {
pendingGraphs.delete(graph);
}
};
// Let native Start Autogrow and its existing output-sync finish first.
if (typeof requestAnimationFrame === "function") requestAnimationFrame(() => setTimeout(run, 0));
else setTimeout(run, 0);
}
app.registerExtension({
name: "Distributed.ImageListEndValues",
nodeCreated(node) {
if (![START, END, GET].includes(node.comfyClass) && !isReroute(node)) return;
if (node.comfyClass === END) {
node.addWidget("button", "Add value override", null, () => {
const shown = Math.max(0, ...(node.inputs ?? []).map(input => ordinal(input.name) ?? 0));
node._distributedVisibleOverrides = Math.min(MAX_VALUES, shown + 1);
schedule(node.graph);
}, { serialize: false });
}
for (const hook of ["onConnectionsChange", "onConfigure", "onAdded", "onWidgetChanged"]) {
const original = node[hook];
node[hook] = function (...args) {
const result = original?.apply(this, args);
schedule(this.graph);
return result;
};
}
// Widget callbacks are also used by frontend versions that do not call
// the node's onWidgetChanged hook for a combo selection.
if (node.comfyClass === GET) {
const widget = node.widgets?.find(widget => widget.name === "value_name");
if (widget) {
const callback = widget.callback;
widget.callback = function (...args) {
const result = callback?.apply(this, args);
schedule(node.graph);
return result;
};
}
}
schedule(node.graph);
},
loadedGraphNode(node) {
schedule(node.graph);
},
});
+190
View File
@@ -0,0 +1,190 @@
import { readFileSync } from "node:fs";
import vm from "node:vm";
import { describe, expect, it } from "vitest";
const source = readFileSync(new URL("../imageListValues.js", import.meta.url), "utf8")
.replace(/^import .*;\r?\n/gm, "");
const PIPE = "DISTRIBUTED_IMAGE_LIST_VALUES";
function setup() {
let extension;
let nextId = 1;
const tasks = [];
const graph = {
_nodes: [], links: new Map(),
getNodeById(id) { return this._nodes.find(node => node.id === id); },
getLink(id) { return this.links.get(id); },
};
vm.runInNewContext(source, {
app: { registerExtension: ext => { extension = ext; } },
setTimeout: callback => tasks.push(callback), console,
});
function node(comfyClass, inputs, outputs, widgets = []) {
const n = {
id: nextId++, comfyClass, graph, widgets,
inputs: inputs.map(name => ({ name, type: "*", link: null })),
outputs: outputs.map(([name, type]) => ({ name, type, links: [], shape: 6 })),
addWidget(type, name, value, callback, options) { this.widgets.push({ type, name, value, callback, options }); },
addInput(name, type) { this.inputs.push({ name, type, link: null }); },
addOutput(name, type, extra) { this.outputs.push({ name, type, links: [], ...extra }); },
removeInput(slot) {
if (this.inputs[slot].link != null) throw new Error("Removed wired input");
this.inputs.splice(slot, 1);
this.inputs.forEach((input, index) => {
if (input.link != null) graph.links.get(input.link).target_slot = index;
});
},
disconnectInput(slot) {
const input = this.inputs[slot];
const link = graph.links.get(input.link);
if (link) {
const out = graph.getNodeById(link.origin_id).outputs[link.origin_slot];
out.links = out.links.filter(id => id !== link.id);
graph.links.delete(link.id);
}
input.link = null;
this.onConnectionsChange?.();
},
setDirtyCanvas() {},
};
graph._nodes.push(n);
extension.nodeCreated(n);
return n;
}
const flush = () => { while (tasks.length) tasks.shift()(); };
const connect = (from, slot, to, name) => {
const targetSlot = to.inputs.findIndex(input => input.name === name);
if (targetSlot < 0) throw new Error(`Missing input ${name}`);
to.disconnectInput(targetSlot);
const id = nextId++;
graph.links.set(id, { id, origin_id: from.id, origin_slot: slot, target_id: to.id,
target_slot: targetSlot, type: from.outputs[slot].type });
from.outputs[slot].links.push(id);
to.inputs[targetSlot].link = id;
from.onConnectionsChange?.();
to.onConnectionsChange?.();
return id;
};
const start = node("DistributedImageListStart", ["images", "mask", "values.value_1", "values.value_2", "values.value_3"],
[["image", "IMAGE"], ["mask", "MASK"], ["item_index", "INT"], ["context", "DISTRIBUTED_IMAGE_LIST_CONTEXT"]]);
const end = node("DistributedImageListEnd", ["images", "masks", "texts", "context", ...Array.from({ length: 10 }, (_, i) => `value_${i + 1}`)],
[["image_list", "IMAGE"], ["masks", "MASK"], ["texts", "STRING"], ["values", PIPE]]);
const get = node("DistributedImageListGetValue", ["values"], [["value", "*"]], [{ name: "value_name", value: "value_1" }]);
connect(start, 3, end, "context");
connect(end, 3, get, "values");
flush();
return { node, graph, start, end, get, flush, connect, extension };
}
const valueNames = node => node.inputs.filter(input => input.name.startsWith("value_")).map(input => input.name);
const addOverride = h => { h.end.widgets.find(widget => widget.name === "Add value override").callback(); h.flush(); };
describe("Image List End context values", () => {
it("keeps original sockets and shows only the first optional override initially", () => {
const h = setup();
expect(h.end.inputs.map(input => input.name)).toEqual(["images", "masks", "texts", "context", "value_1"]);
expect(h.end.outputs.map(output => output.name)).toEqual(["image_list", "masks", "texts", "values"]);
expect(h.end.widgets[0].options.serialize).toBe(false);
});
it("appends the values output when loading an old three-output End", () => {
const h = setup();
const old = h.node("DistributedImageListEnd", ["images", "masks", "texts", "context"],
[["image_list", "IMAGE"], ["masks", "MASK"], ["texts", "STRING"]]);
h.flush();
expect(old.outputs[3]).toMatchObject({ name: "values", type: PIPE, shape: 6 });
expect(valueNames(old)).toEqual(["value_1"]);
});
it("grows overrides without compacting field identities or existing links", () => {
const h = setup();
addOverride(h);
addOverride(h);
const src = h.node("PrimitiveInt", [], [["INT", "INT"]]);
const thirdLink = h.connect(src, 0, h.end, "value_3");
h.connect(src, 0, h.end, "value_2");
h.flush();
h.end.disconnectInput(h.end.inputs.findIndex(input => input.name === "value_2"));
h.flush();
expect(valueNames(h.end)).toEqual(["value_1", "value_2", "value_3", "value_4"]);
const link = h.graph.links.get(thirdLink);
expect(h.end.inputs[link.target_slot].name).toBe("value_3");
for (let i = 0; i < 20; i++) addOverride(h);
expect(valueNames(h.end)).toHaveLength(10);
});
it.each(["IMAGE", "MASK", "STRING", "INT", "FLOAT", "BOOLEAN"])("infers %s from the original Start value through context", type => {
const h = setup();
const source = h.node("Source", [], [["value", type]]);
h.connect(source, 0, h.start, "values.value_1");
h.flush();
expect(h.get.outputs[0].type).toBe(type);
});
it("prefers the End override type, then returns to the original after disconnect", () => {
const h = setup();
h.connect(h.node("Source", [], [["value", "INT"]]), 0, h.start, "values.value_1");
h.connect(h.node("Source", [], [["value", "STRING"]]), 0, h.end, "value_1");
h.flush();
expect(h.get.outputs[0].type).toBe("STRING");
h.end.disconnectInput(h.end.inputs.findIndex(input => input.name === "value_1"));
h.flush();
expect(h.get.outputs[0].type).toBe("INT");
});
it("updates a selected field and disconnects incompatible consumers", () => {
const h = setup();
h.connect(h.node("Source", [], [["value", "STRING"]]), 0, h.start, "values.value_1");
h.connect(h.node("Source", [], [["value", "INT"]]), 0, h.start, "values.value_2");
h.flush();
const consumer = h.node("Save", ["filename"], []);
consumer.inputs[0].type = "STRING";
h.connect(h.get, 0, consumer, "filename");
h.get.widgets[0].value = "value_2";
h.get.widgets[0].callback();
h.flush();
expect(h.get.outputs[0].type).toBe("INT");
expect(consumer.inputs[0].link).toBe(null);
});
it("keeps compatible consumer links and uses wildcard for unavailable graph metadata", () => {
const h = setup();
h.connect(h.node("Source", [], [["value", "INT"]]), 0, h.start, "values.value_1");
h.flush();
const consumer = h.node("Consumer", ["value"], []);
consumer.inputs[0].type = "INT,FLOAT";
const id = h.connect(h.get, 0, consumer, "value");
h.get.disconnectInput(0);
h.flush();
expect(h.get.outputs[0].type).toBe("*");
expect(h.graph.links.has(id)).toBe(true);
});
it("follows reroutes for context and values, and terminates on a cycle", () => {
const h = setup();
h.connect(h.node("Source", [], [["value", "FLOAT"]]), 0, h.start, "values.value_1");
const reroute = h.node("Reroute", ["input"], [["output", "*"]]);
h.connect(h.end, 3, reroute, "input");
h.connect(reroute, 0, h.get, "values");
h.flush();
expect(h.get.outputs[0].type).toBe("FLOAT");
h.connect(reroute, 0, reroute, "input");
h.flush();
expect(h.get.outputs[0].type).toBe("*");
});
it("recomputes types across sequential regions", () => {
const h = setup();
const secondStart = h.node("DistributedImageListStart", ["values.value_1"],
[["image", "IMAGE"], ["mask", "MASK"], ["item_index", "INT"], ["context", "DISTRIBUTED_IMAGE_LIST_CONTEXT"]]);
const secondEnd = h.node("DistributedImageListEnd", ["images", "masks", "texts", "context"],
[["image_list", "IMAGE"], ["masks", "MASK"], ["texts", "STRING"], ["values", PIPE]]);
const secondGet = h.node("DistributedImageListGetValue", ["values"], [["value", "*"]], [{ name: "value_name", value: "value_1" }]);
h.connect(h.node("Source", [], [["value", "BOOLEAN"]]), 0, h.start, "values.value_1");
h.connect(h.get, 0, secondStart, "values.value_1");
h.connect(secondStart, 3, secondEnd, "context");
h.connect(secondEnd, 3, secondGet, "values");
h.flush();
expect(secondGet.outputs[0].type).toBe("BOOLEAN");
});
});
+94 -6
View File
@@ -7,7 +7,7 @@ import { createWorkerSettingsForm } from "../ui/settingsForm.js";
import { DistributedUI } from "../ui.js";
import { checkWorkerStatus, formatServerlessStatus, getWorkerUrl, launchWorker, stopWorker, viewWorkerLog } from "../workerLifecycle.js";
import { handleClearMemory, handleInterruptWorkers } from "../workerUtils.js";
import { isRemoteWorker, normalizeRunpodEndpointId, saveWorkerSettings, testServerlessWorker } from "../workerSettings.js";
import { cancelWorkerSettings, isRemoteWorker, normalizeRunpodEndpointId, saveWorkerSettings, testServerlessWorker } from "../workerSettings.js";
import { performPreflightCheck } from "../executionUtils.js";
function createDocumentMock() {
@@ -44,6 +44,7 @@ function createDocumentMock() {
document: {
createElement: makeElement,
getElementById: (id) => elements.get(id) || null,
querySelector: () => null,
},
};
}
@@ -125,19 +126,94 @@ describe("Runpod serverless worker panel", () => {
expect(elements.get(`${field}-group-saved`).style.display).toBe("none");
}
expect(elements.get("runpod-text-saved").style.display).toBe("none");
for (const field of ["endpoint-id", "api-key", "clear-api-key", "jobs-per-run", "timeout-seconds", "serverless-test"]) {
for (const field of ["endpoint-id", "media-profile", "api-key", "clear-api-key", "jobs-per-run", "timeout-seconds", "serverless-test"]) {
expect(elements.get(`${field}-group-saved`).style.display).toBe("flex");
}
// The separate concurrency setting was removed; a legacy value is ignored.
expect(elements.has("concurrency-saved")).toBe(false);
expect(elements.has("concurrency-group-saved")).toBe(false);
expect(elements.get("endpoint-id-saved").value).toBe("ep_123");
expect(elements.get("media-profile-saved").value).toBe("legacy-image-v1");
expect(elements.get("jobs-per-run-saved").value).toBe("2");
expect(elements.get("timeout-seconds-saved").value).toBe("900");
expect(elements.get("api-key-saved").value).toBe("");
expect(elements.get("api-key-saved").placeholder).toMatch(/leave blank/i);
});
it("renders explicit endpoint profiles with installation and no-auto-detection guidance", () => {
createWorkerSettingsForm(createFormUi(), {}, {
id: "profile", name: "Runpod", type: "serverless", media_profile: "comfy-media-v1",
});
const select = elements.get("media-profile-profile");
const group = elements.get("media-profile-group-profile");
expect(group.children[0].textContent).toBe("Endpoint media profile");
expect(group.children[0].htmlFor).toBe(select.id);
expect(select.children.map((option) => [option.value, option.textContent])).toEqual([
["legacy-image-v1", "Legacy images (existing handler)"],
["comfy-media-v1", "Typed media v1 (requires endpoint bridge)"],
]);
expect(select.value).toBe("comfy-media-v1");
expect(select.title).toMatch(/after installing.*endpoint bridge/i);
expect(select.title).toMatch(/health checks do not auto-detect/i);
expect(select.title).toMatch(/schema.*formats and limits are fixed/i);
const type = elements.get("worker-type-profile");
type.value = "remote";
type.dispatchEvent();
expect(group.style.display).toBe("none");
type.value = "serverless";
type.dispatchEvent();
expect(group.style.display).toBe("flex");
expect(select.value).toBe("comfy-media-v1");
});
it.each(["legacy-image-v1", "comfy-media-v1"])("round-trips an explicit %s selection through the save payload and reload", async (profile) => {
const worker = { id: "roundtrip", name: "Runpod", type: "serverless", endpoint_id: "ep_123", enabled: false };
let persisted;
const extension = makeSaveExtension(worker, vi.fn(async (id, payload) => {
persisted = JSON.parse(JSON.stringify({ id, ...payload }));
return {};
}));
const ui = createFormUi();
const form = createWorkerSettingsForm(ui, extension, worker);
elements.get("media-profile-roundtrip").value = profile;
const saveButton = form.children.at(-1).children.find((button) => button.textContent === "Save");
await saveButton.onclick();
expect(persisted.media_profile).toBe(profile);
expect(worker.media_profile).toBe(profile);
expect(persisted).not.toHaveProperty("api_key");
expect(persisted).not.toHaveProperty("concurrency");
createWorkerSettingsForm(ui, extension, persisted);
expect(elements.get("media-profile-roundtrip").value).toBe(profile);
});
it.each([undefined, "comfy-media-v1"])("restores the saved profile on cancel (%s)", (profile) => {
const worker = { id: "cancel", name: "Runpod", type: "serverless", media_profile: profile };
const extension = makeSaveExtension(worker);
createWorkerSettingsForm(createFormUi(), extension, worker);
elements.get("media-profile-cancel").value = profile ? "legacy-image-v1" : "comfy-media-v1";
cancelWorkerSettings(extension, worker.id);
expect(elements.get("media-profile-cancel").value).toBe(profile ?? "legacy-image-v1");
expect(extension.api.updateWorker).not.toHaveBeenCalled();
});
it("keeps the saved profile after a busy-guard error and rejects invalid selections", async () => {
const worker = { id: "busy", name: "Runpod", type: "serverless", endpoint_id: "ep_123", media_profile: "comfy-media-v1" };
const extension = makeSaveExtension(worker, vi.fn().mockRejectedValue(new Error("Finish or reconcile unresolved serverless jobs")));
createWorkerSettingsForm(createFormUi(), extension, worker);
elements.get("media-profile-busy").value = "legacy-image-v1";
await saveWorkerSettings(extension, worker.id);
expect(worker.media_profile).toBe("comfy-media-v1");
expect(extension.app.extensionManager.toast.add).toHaveBeenLastCalledWith(
expect.objectContaining({ summary: "Save Failed", detail: expect.stringMatching(/unresolved/) }),
);
elements.get("media-profile-busy").value = "auto";
await saveWorkerSettings(extension, worker.id);
expect(extension.api.updateWorker).toHaveBeenCalledTimes(1);
expect(extension.app.extensionManager.toast.add).toHaveBeenLastCalledWith(
expect.objectContaining({ summary: "Validation Error", detail: expect.stringMatching(/media profile/) }),
);
});
it("preserves a saved cloud port on first render and applies defaults only on type changes", () => {
vi.spyOn(document, "getElementById").mockReturnValue(null);
createWorkerSettingsForm(createFormUi(), {}, {
@@ -218,6 +294,7 @@ describe("Runpod serverless worker panel", () => {
elements.set("name-w2", { value: "Runpod worker" });
elements.set("worker-type-w2", { value: "serverless" });
elements.set("endpoint-id-w2", { value: "ep_123" });
elements.set("media-profile-w2", { value: "legacy-image-v1" });
elements.set("api-key-w2", { value: "secret-token" });
elements.set("clear-api-key-w2", { checked: false });
elements.set("jobs-per-run-w2", { value: "4" });
@@ -227,6 +304,7 @@ describe("Runpod serverless worker panel", () => {
expect(extension.api.updateWorker).toHaveBeenCalledWith("w2", {
name: "Runpod worker", type: "serverless", provider: "runpod", endpoint_id: "ep_123",
media_profile: "legacy-image-v1",
api_key: "secret-token", jobs_per_run: 4, timeout_seconds: 2400, enabled: false,
});
expect(worker).toMatchObject({ type: "serverless", provider: "runpod", enabled: false, api_key_configured: true });
@@ -244,6 +322,7 @@ describe("Runpod serverless worker panel", () => {
elements.set("name-legacy", { value: "Runpod" });
elements.set("worker-type-legacy", { value: "serverless" });
elements.set("endpoint-id-legacy", { value: "ep_123" });
elements.set("media-profile-legacy", { value: "legacy-image-v1" });
elements.set("api-key-legacy", { value: "" });
elements.set("clear-api-key-legacy", { checked: false });
elements.set("jobs-per-run-legacy", { value: "2" });
@@ -278,6 +357,7 @@ describe("Runpod serverless worker panel", () => {
elements.set("name-url", { value: "Runpod" });
elements.set("worker-type-url", { value: "serverless" });
elements.set("endpoint-id-url", { value: "https://api.runpod.ai/v2/abc123xyz/runsync" });
elements.set("media-profile-url", { value: "legacy-image-v1" });
elements.set("api-key-url", { value: "" });
elements.set("clear-api-key-url", { checked: false });
elements.set("jobs-per-run-url", { value: "1" });
@@ -294,7 +374,7 @@ describe("Runpod serverless worker panel", () => {
});
it("asks before a type change deletes the saved API key, and cancels cleanly", async () => {
const worker = { id: "conv", name: "Runpod", type: "serverless", provider: "runpod", enabled: false, api_key_configured: true };
const worker = { id: "conv", name: "Runpod", type: "serverless", provider: "runpod", enabled: false, api_key_configured: true, media_profile: "comfy-media-v1" };
const extension = makeSaveExtension(worker);
elements.set("name-conv", { value: "Runpod" });
elements.set("worker-type-conv", { value: "remote" });
@@ -307,11 +387,13 @@ describe("Runpod serverless worker panel", () => {
expect(confirmSpy).toHaveBeenCalledWith(expect.stringMatching(/saved Runpod API key will be deleted/));
expect(extension.api.updateWorker).not.toHaveBeenCalled();
expect(worker.type).toBe("serverless");
expect(worker.media_profile).toBe("comfy-media-v1");
confirmSpy.mockReturnValue(true);
await saveWorkerSettings(extension, "conv");
expect(extension.api.updateWorker).toHaveBeenCalledTimes(1);
expect(extension.api.updateWorker.mock.calls[0][1].type).toBe("remote");
expect(worker).not.toHaveProperty("media_profile");
} finally {
delete globalThis.confirm;
}
@@ -341,6 +423,7 @@ describe("Runpod serverless worker panel", () => {
elements.set("name-w-blank", { value: "Runpod" });
elements.set("worker-type-w-blank", { value: "serverless" });
elements.set("endpoint-id-w-blank", { value: "ep_123" });
elements.set("media-profile-w-blank", { value: "legacy-image-v1" });
elements.set("api-key-w-blank", { value: "" });
elements.set("clear-api-key-w-blank", { checked: false });
elements.set("jobs-per-run-w-blank", { value: "1" });
@@ -361,6 +444,7 @@ describe("Runpod serverless worker panel", () => {
elements.set("name-w3", { value: "Runpod" });
elements.set("worker-type-w3", { value: "serverless" });
elements.set("endpoint-id-w3", { value: "ep_123" });
elements.set("media-profile-w3", { value: "legacy-image-v1" });
elements.set("api-key-w3", { value: "" });
elements.set("clear-api-key-w3", { checked: true });
elements.set("jobs-per-run-w3", { value: "1" });
@@ -370,6 +454,7 @@ describe("Runpod serverless worker panel", () => {
expect(extension.api.updateWorker).toHaveBeenCalledWith("w3", {
name: "Runpod", type: "serverless", provider: "runpod", endpoint_id: "ep_123",
media_profile: "legacy-image-v1",
clear_api_key: true, jobs_per_run: 1, timeout_seconds: 1800, enabled: true,
});
expect(worker.api_key_configured).toBe(false);
@@ -527,9 +612,6 @@ describe("Runpod serverless worker panel", () => {
});
it("tests only the saved Runpod endpoint and reports the backend message", async () => {
const button = document.createElement("button");
button.id = "test-serverless-w6";
button.textContent = "Test Saved Endpoint";
const extension = {
api: {
testServerlessWorker: vi.fn().mockResolvedValue({ ok: true, message: "Endpoint reachable" }),
@@ -537,9 +619,15 @@ describe("Runpod serverless worker panel", () => {
},
app: { extensionManager: { toast: { add: vi.fn() } } },
};
const worker = { id: "w6", name: "Runpod", type: "serverless" };
extension.config = { workers: [worker] };
createWorkerSettingsForm(createFormUi(), extension, worker);
const button = elements.get("test-serverless-w6");
await expect(testServerlessWorker(extension, "w6")).resolves.toEqual({ ok: true, message: "Endpoint reachable" });
expect(worker).not.toHaveProperty("media_profile");
expect(elements.get("media-profile-w6").value).toBe("legacy-image-v1");
expect(extension.api.testServerlessWorker).toHaveBeenCalledWith("w6");
expect(extension.api.queueDistributed).not.toHaveBeenCalled();
expect(extension.app.extensionManager.toast.add).toHaveBeenCalledWith(expect.objectContaining({
+30
View File
@@ -4,6 +4,9 @@ import { cancelWorkerSettings, deleteWorker, isRemoteWorker, saveWorkerSettings,
export const IMAGES_PER_RUN_TOOLTIP =
"Separate Runpod jobs submitted each run, all at once. Runpod's max workers decides how many GPUs actually run them.";
const MEDIA_PROFILE_TOOLTIP =
"Select Typed media v1 only after installing the ComfyUI-Distributed endpoint bridge. This declares bridge support; endpoint health checks do not auto-detect it. The v1 schema, media formats and limits are fixed.";
export function createWorkerSettingsForm(ui, extension, worker) {
const form = document.createElement("div");
form.style.cssText = "display: flex; flex-direction: column; gap: 8px;";
@@ -59,6 +62,7 @@ export function createWorkerSettingsForm(ui, extension, worker) {
const portInput = portGroup.input;
const serverlessGroups = [
endpointGroup.group,
mediaProfileGroup,
apiKeyGroup.group,
clearApiKeyGroup,
jobsPerRunGroup.group,
@@ -137,6 +141,32 @@ export function createWorkerSettingsForm(ui, extension, worker) {
endpointGroup.group.style.display = "none";
form.appendChild(endpointGroup.group);
const mediaProfileGroup = document.createElement("div");
mediaProfileGroup.id = `media-profile-group-${worker.id}`;
mediaProfileGroup.style.cssText = "display: none; flex-direction: column; gap: 4px;";
mediaProfileGroup.title = MEDIA_PROFILE_TOOLTIP;
const mediaProfileLabel = document.createElement("label");
mediaProfileLabel.htmlFor = `media-profile-${worker.id}`;
mediaProfileLabel.textContent = "Endpoint media profile";
mediaProfileLabel.style.cssText = typeLabel.style.cssText;
const mediaProfileSelect = document.createElement("select");
mediaProfileSelect.id = mediaProfileLabel.htmlFor;
mediaProfileSelect.style.cssText = typeSelect.style.cssText;
mediaProfileSelect.title = MEDIA_PROFILE_TOOLTIP;
for (const [value, label] of [
["legacy-image-v1", "Legacy images (existing handler)"],
["comfy-media-v1", "Typed media v1 (requires endpoint bridge)"],
]) {
const option = document.createElement("option");
option.value = value;
option.textContent = label;
mediaProfileSelect.appendChild(option);
}
mediaProfileSelect.value = worker.media_profile ?? "legacy-image-v1";
mediaProfileGroup.appendChild(mediaProfileLabel);
mediaProfileGroup.appendChild(mediaProfileSelect);
form.appendChild(mediaProfileGroup);
const apiKeyGroup = ui.createFormGroup("Runpod API Key (write-only):", "", `api-key-${worker.id}`, "password", worker.api_key_configured ? "Leave blank to keep saved API key" : "Runpod API key");
apiKeyGroup.group.id = `api-key-group-${worker.id}`;
apiKeyGroup.input.autocomplete = "new-password";
+10
View File
@@ -61,6 +61,7 @@ function showServerlessValidationError(extension, detail) {
async function saveServerlessWorkerSettings(extension, worker, workerId, name) {
const endpointId = normalizeRunpodEndpointId(document.getElementById(`endpoint-id-${workerId}`).value);
const mediaProfile = document.getElementById(`media-profile-${workerId}`).value;
const apiKey = document.getElementById(`api-key-${workerId}`).value || "";
const clearApiKey = Boolean(document.getElementById(`clear-api-key-${workerId}`).checked);
const jobsPerRun = Number(document.getElementById(`jobs-per-run-${workerId}`).value);
@@ -81,6 +82,10 @@ async function saveServerlessWorkerSettings(extension, worker, workerId, name) {
);
return;
}
if (!["legacy-image-v1", "comfy-media-v1"].includes(mediaProfile)) {
showServerlessValidationError(extension, "Select a supported endpoint media profile");
return;
}
if (!Number.isInteger(jobsPerRun) || jobsPerRun < 1 || jobsPerRun > 16) {
showServerlessValidationError(extension, "Images per run must be a whole number from 1 to 16");
return;
@@ -95,6 +100,7 @@ async function saveServerlessWorkerSettings(extension, worker, workerId, name) {
type: "serverless",
provider: "runpod",
endpoint_id: endpointId,
media_profile: mediaProfile,
...(apiKey.length > 0 && !clearApiKey ? { api_key: apiKey } : {}),
...(clearApiKey ? { clear_api_key: true } : {}),
jobs_per_run: jobsPerRun,
@@ -111,6 +117,7 @@ async function saveServerlessWorkerSettings(extension, worker, workerId, name) {
worker.type = "serverless";
worker.provider = "runpod";
worker.endpoint_id = endpointId;
worker.media_profile = mediaProfile;
worker.jobs_per_run = jobsPerRun;
// Replaced by a fixed backend limit; drop any value loaded from an older config.
delete worker.concurrency;
@@ -310,6 +317,7 @@ export async function saveWorkerSettings(extension, workerId) {
// Update local config
worker.name = name.trim();
worker.type = workerType;
delete worker.media_profile;
if (isRemote) {
worker.host = hostTrimmed;
delete worker.cuda_device;
@@ -366,6 +374,8 @@ export function cancelWorkerSettings(extension, workerId) {
document.getElementById(`args-${workerId}`).value = worker.extra_args || "";
const endpointInput = document.getElementById(`endpoint-id-${workerId}`);
if (endpointInput) endpointInput.value = worker.endpoint_id || "";
const mediaProfileInput = document.getElementById(`media-profile-${workerId}`);
if (mediaProfileInput) mediaProfileInput.value = worker.media_profile ?? "legacy-image-v1";
const apiKeyInput = document.getElementById(`api-key-${workerId}`);
if (apiKeyInput) {
apiKeyInput.value = "";