Compare commits
20
Commits
@@ -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,
|
||||
@@ -36,9 +37,11 @@ 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),
|
||||
})
|
||||
|
||||
|
||||
@@ -52,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:
|
||||
|
||||
+38
-1
@@ -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):
|
||||
@@ -134,7 +136,9 @@ def _decode_canonical_png_tensor(image_payload):
|
||||
|
||||
try:
|
||||
with Image.open(io.BytesIO(png_bytes)) as img:
|
||||
img = img.convert("RGB")
|
||||
# PNG callbacks must retain alpha just like local IMAGE tensors.
|
||||
has_alpha = "A" in img.getbands() or "transparency" in img.info
|
||||
img = img.convert("RGBA" if has_alpha else "RGB")
|
||||
tensor = pil_to_tensor(img)
|
||||
return ensure_contiguous(tensor)
|
||||
except Exception as exc:
|
||||
@@ -299,6 +303,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 +312,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 +419,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 +498,7 @@ async def job_complete_endpoint(request):
|
||||
"image_index": int(batch_idx),
|
||||
"is_last": is_last,
|
||||
"audio": decoded_audio,
|
||||
"format": transfer_codec.LEGACY_FORMAT,
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
@@ -188,7 +188,8 @@ async def dispatch_worker_prompt(
|
||||
return prompt_id if return_prompt_id else None
|
||||
|
||||
|
||||
async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0):
|
||||
async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0, *,
|
||||
required_transfer_formats=(), require_value_results=False):
|
||||
"""Require protocol/node support and report which direct workers are idle.
|
||||
|
||||
This performs read-only GETs only; it never submits a prompt or contacts a
|
||||
@@ -196,6 +197,14 @@ async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0):
|
||||
"""
|
||||
if not workers:
|
||||
raise RuntimeError("Distributed Image List requires at least one selected direct worker")
|
||||
if (not isinstance(required_transfer_formats, (list, tuple, set, frozenset))
|
||||
or any(type(fmt) is not str for fmt in required_transfer_formats)):
|
||||
raise ValueError("required_transfer_formats must contain image transfer format names")
|
||||
required_formats = set(required_transfer_formats) - {"exact"}
|
||||
if required_formats:
|
||||
from ...serverless_bridge.image_transfer import validate_options
|
||||
for fmt in required_formats:
|
||||
validate_options(fmt)
|
||||
session = await get_client_session()
|
||||
required_nodes = {"DistributedImageListStart", "DistributedImageListEnd"}
|
||||
semaphore = asyncio.Semaphore(8)
|
||||
@@ -242,6 +251,22 @@ async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0):
|
||||
"upgrade that worker before running this feature; no work was submitted"
|
||||
)
|
||||
|
||||
if required_formats:
|
||||
advertised = capability.get("image_transfer_formats")
|
||||
if (not isinstance(advertised, list)
|
||||
or any(type(fmt) is not str for fmt in advertised)
|
||||
or not required_formats.issubset(advertised)):
|
||||
raise RuntimeError(
|
||||
f"Image-list worker {worker_id} lacks required image_transfer_formats "
|
||||
f"{sorted(required_formats)}; upgrade that worker before compressed IMAGE returns; "
|
||||
"no work was submitted"
|
||||
)
|
||||
if require_value_results and capability.get("value_results") is not True:
|
||||
raise RuntimeError(
|
||||
f"Image-list worker {worker_id} lacks explicit value_results support; "
|
||||
"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:
|
||||
|
||||
@@ -401,20 +401,6 @@ def prune_prompt_for_worker(prompt_obj, use_worker_sink=False):
|
||||
|
||||
def prepare_delegate_master_prompt(prompt_obj, collector_ids):
|
||||
"""Prune master prompt so it only executes post-collector nodes in delegate mode."""
|
||||
# Typed collectors need no fabricated IMAGE placeholder. Sever all media
|
||||
# links, including STRING scalar branches, before retained-config analysis.
|
||||
# Their expected_media control was frozen by participant overrides.
|
||||
if any((prompt_obj.get(cid) or {}).get("class_type") == "DistributedMediaCollector"
|
||||
for cid in collector_ids):
|
||||
from ...serverless_bridge.contract import SOCKET_KINDS
|
||||
prompt_obj = json.loads(json.dumps(prompt_obj))
|
||||
for cid in collector_ids:
|
||||
node = prompt_obj.get(cid) or {}
|
||||
if node.get("class_type") == "DistributedMediaCollector":
|
||||
inputs = node.setdefault("inputs", {})
|
||||
inputs.setdefault("expected_media", json.dumps([name for name in SOCKET_KINDS if name in inputs]))
|
||||
for name in SOCKET_KINDS:
|
||||
inputs.pop(name, None)
|
||||
downstream = _find_downstream_nodes(prompt_obj, collector_ids)
|
||||
nodes_to_keep = set(collector_ids)
|
||||
nodes_to_keep.update(downstream)
|
||||
@@ -519,7 +505,6 @@ def generate_job_id_map(prompt_index, prefix):
|
||||
job_map = {}
|
||||
distributed_nodes = (
|
||||
prompt_index.nodes_for_class("DistributedCollector")
|
||||
+ prompt_index.nodes_for_class("DistributedMediaCollector")
|
||||
+ prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
|
||||
+ prompt_index.nodes_for_class("DistributedImageListStart")
|
||||
)
|
||||
@@ -552,10 +537,8 @@ def _override_collector_nodes(
|
||||
enabled_json,
|
||||
delegate_master,
|
||||
):
|
||||
"""Configure legacy and typed collectors without changing their media."""
|
||||
collector_ids = (prompt_index.nodes_for_class("DistributedCollector")
|
||||
+ prompt_index.nodes_for_class("DistributedMediaCollector"))
|
||||
for node_id in collector_ids:
|
||||
"""Configure DistributedCollector nodes for master or worker role."""
|
||||
for node_id in prompt_index.nodes_for_class("DistributedCollector"):
|
||||
node = prompt_copy.get(node_id)
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
@@ -568,9 +551,6 @@ def _override_collector_nodes(
|
||||
inputs["multi_job_id"] = job_id_map.get(node_id, node_id)
|
||||
inputs["is_worker"] = not is_master
|
||||
inputs["enabled_worker_ids"] = enabled_json
|
||||
if node.get("class_type") == "DistributedMediaCollector":
|
||||
from ...serverless_bridge.contract import SOCKET_KINDS
|
||||
inputs["expected_media"] = json.dumps([socket for socket in SOCKET_KINDS if socket in inputs])
|
||||
if is_master:
|
||||
inputs["delegate_only"] = bool(delegate_master)
|
||||
inputs.pop("master_url", None)
|
||||
|
||||
@@ -1,8 +1,9 @@
|
||||
"""Pure compilation of prepared typed-media collectors for serverless workers.
|
||||
"""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 socket, retaining only its upstream computation.
|
||||
It never imports ComfyUI or inspects/executes runtime node classes.
|
||||
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
|
||||
@@ -12,8 +13,8 @@ import hashlib
|
||||
import json
|
||||
from graphlib import CycleError, TopologicalSorter
|
||||
|
||||
from ...serverless_bridge.contract import SOCKET_KINDS, identifier, validate_bindings
|
||||
from ...serverless_bridge.image_transfer import DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options
|
||||
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,
|
||||
@@ -24,7 +25,24 @@ from .serverless_prompt import (
|
||||
)
|
||||
|
||||
|
||||
_COLLECTOR = "DistributedMediaCollector"
|
||||
_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):
|
||||
@@ -36,7 +54,7 @@ def _validate_identifier(value, context):
|
||||
raise ValueError(f"{context} must be a valid nonempty media identifier") from None
|
||||
|
||||
|
||||
def _validate_prompt(prepared_prompt):
|
||||
def _validate_prompt(prepared_prompt, *, require_collectors=True):
|
||||
"""Validate the entire source graph, including branches later discarded."""
|
||||
if not isinstance(prepared_prompt, dict) or not prepared_prompt:
|
||||
raise ValueError("prepared prompt must be a nonempty API graph mapping")
|
||||
@@ -51,7 +69,7 @@ def _validate_prompt(prepared_prompt):
|
||||
raise ValueError(f"prompt node {node_id!r} must have a nonempty class_type")
|
||||
if class_type == _COLLECTOR:
|
||||
collectors.append(node_id)
|
||||
elif class_type == "DistributedCollector" or _is_unsupported_distributed(class_type):
|
||||
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):
|
||||
@@ -59,8 +77,8 @@ def _validate_prompt(prepared_prompt):
|
||||
if any(not isinstance(name, str) or not name for name in inputs):
|
||||
raise ValueError(f"prompt node {node_id!r} input names must be nonempty strings")
|
||||
|
||||
if not collectors:
|
||||
raise ValueError("prepared prompt contains no DistributedMediaCollector nodes")
|
||||
if require_collectors and not collectors:
|
||||
raise ValueError("prepared prompt contains no DistributedCollector nodes")
|
||||
|
||||
dependencies = {}
|
||||
for node_id, node in prepared_prompt.items():
|
||||
@@ -81,18 +99,17 @@ def _validate_prompt(prepared_prompt):
|
||||
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(SOCKET_KINDS):
|
||||
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", {})
|
||||
validate_options(inputs.get("transfer_format", DEFAULT_FORMAT),
|
||||
inputs.get("transfer_quality", DEFAULT_QUALITY))
|
||||
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 SOCKET_KINDS if socket in inputs]
|
||||
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:
|
||||
@@ -123,61 +140,84 @@ def _validate_prompt(prepared_prompt):
|
||||
return sorted(collectors)
|
||||
|
||||
|
||||
def _build_exports(source_nodes, collector_ids, request_id):
|
||||
bindings = []
|
||||
exports = {}
|
||||
def _build_explicit_exports(source_nodes, specs, request_id):
|
||||
bindings, exports = [], {}
|
||||
reserved = set(source_nodes)
|
||||
for collector_id in collector_ids:
|
||||
inputs = source_nodes[collector_id].get("inputs", {})
|
||||
for socket, kind in SOCKET_KINDS.items():
|
||||
if socket not in inputs:
|
||||
continue
|
||||
# Hashing avoids separator ambiguity and keeps long valid source
|
||||
# identities within the bridge's identifier length bound. Identity
|
||||
# is stable across request IDs and dictionary insertion orders.
|
||||
identity = json.dumps([collector_id, socket], separators=(",", ":"))
|
||||
digest = hashlib.sha256(identity.encode("utf-8")).hexdigest()
|
||||
base_id = f"serverless:media:{digest}"
|
||||
node_id = base_id
|
||||
suffix = 0
|
||||
while node_id in reserved:
|
||||
suffix += 1
|
||||
node_id = f"{base_id}:{suffix}"
|
||||
reserved.add(node_id)
|
||||
binding = {
|
||||
"binding_id": f"b:{digest}",
|
||||
"collector_id": collector_id,
|
||||
"job_id": inputs["multi_job_id"],
|
||||
"participant_id": inputs["worker_id"],
|
||||
"node_id": node_id,
|
||||
"socket": socket,
|
||||
"kind": kind,
|
||||
}
|
||||
bindings.append(binding)
|
||||
exports[node_id] = {
|
||||
"class_type": "MediaExportV1",
|
||||
"inputs": {
|
||||
"request_id": request_id,
|
||||
"binding_json": json.dumps(binding),
|
||||
socket: copy.deepcopy(inputs[socket]),
|
||||
},
|
||||
}
|
||||
if kind == "IMAGE":
|
||||
fmt, quality = validate_options(inputs.get("transfer_format", DEFAULT_FORMAT),
|
||||
inputs.get("transfer_quality", DEFAULT_QUALITY))
|
||||
if fmt != DEFAULT_FORMAT:
|
||||
exports[node_id]["inputs"].update(transfer_format=fmt, transfer_quality=quality)
|
||||
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 DistributedMediaCollector must have prepared literal ``multi_job_id``
|
||||
and ``worker_id`` identities and at least one linked media input. Present
|
||||
media inputs (including texts) must be API links; literal text belongs in a
|
||||
PrimitiveString node. Nonmedia collector controls are not exported.
|
||||
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
|
||||
|
||||
@@ -171,6 +171,8 @@ class ServerlessManager:
|
||||
for rid in done[:-200]:
|
||||
self.tasks.pop(rid, None)
|
||||
self.runs.pop(rid, None)
|
||||
# Additive ownership handle for item schedulers; legacy callers ignore it.
|
||||
return [record_id for record_id, _run, _semaphore in new_runs]
|
||||
|
||||
def is_pending(self, collector_job_id, participant_id):
|
||||
return any(not task.done() and run["worker"]["id"] == participant_id and
|
||||
|
||||
+64
-82
@@ -38,7 +38,7 @@ from .orchestration.prompt_transform import (
|
||||
WORKER_SINK_CLASS,
|
||||
_resolve_direct_workers_for_image_list,
|
||||
)
|
||||
from ..image_list.runtime import validate_image_list_regions
|
||||
from ..image_list.runtime import validate_image_list_regions, _has_selected_serverless_workers
|
||||
from ..image_list.execution_plan import prepare_image_list_execution
|
||||
|
||||
|
||||
@@ -153,44 +153,36 @@ def _is_load_balance_enabled(value):
|
||||
|
||||
|
||||
def _prompt_requests_load_balance(prompt_index):
|
||||
collector_ids = (prompt_index.nodes_for_class("DistributedCollector")
|
||||
+ prompt_index.nodes_for_class("DistributedMediaCollector"))
|
||||
for node_id in collector_ids:
|
||||
for node_id in prompt_index.nodes_for_class("DistributedCollector"):
|
||||
inputs = prompt_index.inputs_by_node.get(node_id, {})
|
||||
if _is_load_balance_enabled(inputs.get("load_balance", False)):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _validate_media_workers(prompt_index, workers):
|
||||
"""Reject incompatible selections before probing, pruning or dispatching.
|
||||
def _validate_media_workers(workers, prompt_index=None):
|
||||
"""Reject unknown serverless media profiles before probing or dispatching.
|
||||
|
||||
Typed collectors currently have only a master/serverless transport. Check
|
||||
all selected workers, not just reachable or load-balanced participants, so
|
||||
a bad selection cannot silently turn into a different execution mode.
|
||||
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.
|
||||
"""
|
||||
typed = bool(prompt_index.nodes_for_class("DistributedMediaCollector"))
|
||||
if typed and (prompt_index.nodes_for_class("DistributedCollector")
|
||||
or prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")):
|
||||
raise ValueError(
|
||||
"DistributedMediaCollector cannot be mixed with DistributedCollector or "
|
||||
"UltimateSDUpscaleDistributed in one prompt"
|
||||
)
|
||||
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":
|
||||
if typed:
|
||||
raise ValueError(
|
||||
"DistributedMediaCollector does not support selected direct workers; "
|
||||
"use master execution or comfy-media-v1 serverless pools"
|
||||
)
|
||||
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 typed != (profile == "comfy-media-v1"):
|
||||
if exact and profile == "legacy-image-v1":
|
||||
raise ValueError(
|
||||
"DistributedMediaCollector requires comfy-media-v1 serverless pools; "
|
||||
"legacy collectors require legacy-image-v1"
|
||||
"transfer_format exact requires comfy-media-v1 serverless pools; "
|
||||
"legacy-image-v1 endpoints return 8-bit images"
|
||||
)
|
||||
|
||||
|
||||
@@ -209,7 +201,7 @@ 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(prompt_index, [worker])
|
||||
_validate_media_workers([worker], prompt_index)
|
||||
worker_prompt = prompt_index.copy_prompt()
|
||||
if worker.get("type") == "serverless":
|
||||
from .orchestration.runpod_client import serialize_request
|
||||
@@ -217,9 +209,8 @@ async def _prepare_worker_payload(
|
||||
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 legacy pruner only recognizes Collector/Upscale roots.
|
||||
# The typed compiler validates the complete private graph and
|
||||
# prunes it itself, retaining every connected typed branch.
|
||||
# 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,
|
||||
@@ -332,62 +323,49 @@ async def orchestrate_distributed_execution(
|
||||
for node in prompt_obj.values())
|
||||
target_kwargs = {}
|
||||
if image_list_mode:
|
||||
if any(isinstance(node, dict) and node.get("class_type") in {
|
||||
"DistributedCollector", "DistributedMediaCollector", "UltimateSDUpscaleDistributed"
|
||||
} for node in prompt_obj.values()):
|
||||
raise ValueError("Distributed Image List cannot be mixed with other distributed collectors")
|
||||
# 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"
|
||||
)
|
||||
prompt_obj, outer_targets = prepare_image_list_execution(prompt_obj)
|
||||
target_kwargs["partial_execution_targets"] = outer_targets
|
||||
prompt_index = PromptIndex(prompt_obj)
|
||||
typed_media_mode = bool(prompt_index.nodes_for_class("DistributedMediaCollector"))
|
||||
if typed_media_mode:
|
||||
from ..serverless_bridge.contract import identifier
|
||||
# PromptIndex's JSON copy would otherwise turn integer keys into strings.
|
||||
for node_id in prompt_obj:
|
||||
identifier(node_id)
|
||||
image_list_start_ids = prompt_index.nodes_for_class("DistributedImageListStart")
|
||||
image_list_mode = bool(image_list_start_ids)
|
||||
if image_list_mode:
|
||||
# Validate before probing or queueing anything so malformed boundaries and
|
||||
# old worker releases fail before any upstream image is materialized.
|
||||
other_distributed = (
|
||||
prompt_index.nodes_for_class("DistributedCollector")
|
||||
+ prompt_index.nodes_for_class("DistributedMediaCollector")
|
||||
+ prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
|
||||
)
|
||||
if other_distributed:
|
||||
raise ValueError(
|
||||
"Distributed Image List v1 cannot be mixed with Distributed Collector, "
|
||||
"DistributedMediaCollector or Ultimate SD Upscale nodes in one prompt"
|
||||
)
|
||||
validate_image_list_regions(prompt_obj, explicit_targets=True)
|
||||
configured_serverless = [
|
||||
str(worker.get("id")) for worker in config.get("workers", [])
|
||||
if worker.get("enabled", False)
|
||||
and str(worker.get("type") or "local").lower() == "serverless"
|
||||
and (requested_ids is None or str(worker.get("id")) in {str(x) for x in requested_ids})
|
||||
]
|
||||
if configured_serverless:
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
"Image List v1 explicitly excludes unsupported serverless pools: "
|
||||
+ ", ".join(configured_serverless),
|
||||
)
|
||||
from ..serverless_bridge.image_transfer import DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options
|
||||
required_transfer_formats = set()
|
||||
for end_id in prompt_index.nodes_for_class("DistributedImageListEnd"):
|
||||
end_inputs = prompt_obj[end_id].get("inputs", {})
|
||||
fmt, _ = validate_options(end_inputs.get("transfer_format", DEFAULT_FORMAT),
|
||||
end_inputs.get("transfer_quality", DEFAULT_QUALITY))
|
||||
if fmt != DEFAULT_FORMAT:
|
||||
required_transfer_formats.add(fmt)
|
||||
workers = _resolve_direct_workers_for_image_list(config, requested_ids)
|
||||
if _has_selected_serverless_workers(config, requested_ids):
|
||||
from ..image_list.serverless import resolve_serverless_workers
|
||||
# Pool identity is retained once, never expanded by Images per run.
|
||||
# Validate profiles before reachability can hide an invalid selection.
|
||||
workers += resolve_serverless_workers(config, requested_ids)
|
||||
if not workers and delegate_master:
|
||||
raise RuntimeError(
|
||||
"Distributed Image List in orchestrator-only mode requires a selected direct worker; "
|
||||
"serverless pools are unsupported"
|
||||
"Distributed Image List in orchestrator-only mode requires a selected direct worker "
|
||||
"or comfy-media-v1 serverless pool"
|
||||
)
|
||||
else:
|
||||
if typed_media_mode and requested_ids is not None:
|
||||
configured_ids = {worker.get("id") for worker in config.get("workers", [])}
|
||||
if any(worker_id not in configured_ids for worker_id in requested_ids):
|
||||
raise ValueError("Unknown explicitly selected media worker ID")
|
||||
workers = _resolve_enabled_workers(config, requested_ids)
|
||||
_validate_media_workers(prompt_index, workers)
|
||||
if typed_media_mode and delegate_master and not workers:
|
||||
raise RuntimeError("Delegate-only DistributedMediaCollector requires a selected comfy-media-v1 pool")
|
||||
_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,
|
||||
@@ -403,8 +381,8 @@ async def orchestrate_distributed_execution(
|
||||
),
|
||||
)
|
||||
|
||||
# Keep the legacy Collector fallback. Image-list and typed-media delegate
|
||||
# modes must not silently enable master participation.
|
||||
# Keep the legacy Collector fallback, but never override image-list
|
||||
# orchestrator-only mode merely because its direct workers are unavailable.
|
||||
if not image_list_mode and not workers and delegate_master:
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
@@ -419,19 +397,24 @@ async def orchestrate_distributed_execution(
|
||||
trace_execution_id=execution_trace_id,
|
||||
probe_concurrency=worker_probe_concurrency,
|
||||
)
|
||||
if typed_media_mode and delegate_master and not active_workers:
|
||||
raise RuntimeError("No selected comfy-media-v1 worker is reachable for delegate-only DistributedMediaCollector")
|
||||
if image_list_mode:
|
||||
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
|
||||
|
||||
@@ -520,8 +503,7 @@ async def orchestrate_distributed_execution(
|
||||
)
|
||||
|
||||
if delegate_master and not image_list_mode:
|
||||
collector_ids = (find_nodes_by_class(master_prompt, "DistributedCollector")
|
||||
+ find_nodes_by_class(master_prompt, "DistributedMediaCollector"))
|
||||
collector_ids = find_nodes_by_class(master_prompt, "DistributedCollector")
|
||||
upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
|
||||
if upscale_nodes:
|
||||
debug_log(
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
+130
-14
@@ -1,8 +1,8 @@
|
||||
# Dynamic image-list queue
|
||||
|
||||
Process a list of images across the master and direct local or remote ComfyUI workers without assigning fixed chunks to each participant.
|
||||
Process a list of images across the master, direct local or remote ComfyUI workers, and compatible Runpod serverless pools without assigning fixed chunks to each participant.
|
||||
|
||||
The master evaluates the input list once. By default, the master and each selected idle worker receive one image job at a time; when a participant finishes, it gets the next unassigned image. Faster participants can therefore process more images. The master restores input order before running downstream nodes.
|
||||
The master evaluates the input list once. By default, the master and each selected direct worker receive one image job at a time; a serverless pool contributes bounded concurrent slots. All slots share the same queue: when a participant finishes, it gets the next unassigned image. Faster participants can therefore process more images. Each source frame is assigned once, and the master restores input order before running downstream nodes.
|
||||
|
||||
The existing **Orchestrator Only** setting (`settings.master_delegate_only`) excludes the master from per-image processing. The master still creates the input list and runs downstream nodes. API callers can override the setting for a run with `delegate_master`. There is no second image-list-specific toggle.
|
||||
|
||||
@@ -23,8 +23,8 @@ Distributed Image List Start -------- context ---------+
|
||||
Preview / Save / next stage
|
||||
```
|
||||
|
||||
1. Install the same feature-capable version of Distributed on the master and every selected worker. Workers on an older image-list protocol fail preflight before any work is sent. Restart those instances and refresh the frontend after installation.
|
||||
2. Enable Distributed and select the **direct** workers to use. The master also participates unless **Orchestrator Only** is enabled. With the master participating, an empty worker selection runs locally through the same image-list scheduler; orchestrator-only mode requires at least one reachable, idle direct worker. The master must be reachable from remote workers at its configured callback address.
|
||||
1. Install the same feature-capable version of Distributed on the master and every selected direct worker. Direct workers on an older image-list protocol fail preflight before any work is sent. Restart those instances and refresh the frontend after installation. Serverless endpoints instead need the compatible media bridge and supervision described below.
|
||||
2. Enable Distributed and select direct workers and/or **Typed media v1** serverless pools. The master also participates unless **Orchestrator Only** is enabled. With the master participating, an empty worker selection runs locally through the same image-list scheduler; orchestrator-only mode requires at least one eligible direct worker or compatible serverless pool. Serverless-only selection is supported. The master must be reachable from direct remote workers at its configured callback address; serverless jobs do not use this callback.
|
||||
3. Feed your image list or tensor batch into **Distributed Image List Start → images**.
|
||||
4. Connect its `image` output through the per-image processing nodes to **Distributed Image List End → images**.
|
||||
5. Connect Start's **context** directly to End's **context**. This is a dedicated `DISTRIBUTED_IMAGE_LIST_CONTEXT` socket, not editable text.
|
||||
@@ -52,7 +52,31 @@ Start outputs, in order: `image`, `mask`, `item_index` (zero-based), `context`,
|
||||
|
||||
Values that cannot be transported exactly (models, latents, conditioning and other objects) are rejected. Load models inside the region from a loader connected directly to the processing nodes, as described below.
|
||||
|
||||
End accepts `images`, optional `masks` / `texts`, and `context`. It returns `image_list`, `masks`, and `texts` in original item order. Each job must produce exactly one image, at most one mask, and at most one text value. Unconnected optional outputs contain `None` entries.
|
||||
End accepts `images`, optional `masks` / `texts`, and `context`. Its first three outputs remain `image_list`, `masks`, and `texts` in original item order. Each job must produce exactly one image, at most one mask, and at most one text value. Unconnected optional mask/text outputs contain `None` entries.
|
||||
|
||||
### Keep values through End without extra wires
|
||||
|
||||
Start's existing **context** connection now also lets End recover each item's original `value_1` … `value_10`. You do not need to wire those values individually into End. Context remains a lightweight execution reference, not a general-purpose data pipe.
|
||||
|
||||
End appends a **values** output (`DISTRIBUTED_IMAGE_LIST_VALUES`): one value bundle per image, in the same source order. Connect it to **Distributed Image List Get Value**, then choose `value_name` to extract the field you need. The getter preserves its type and ComfyUI maps it across the bundles, producing a list aligned with End's images. Missing fields fail clearly instead of producing empty placeholders.
|
||||
|
||||
For example, to preserve filenames:
|
||||
|
||||
```text
|
||||
filenames → Start.value_1
|
||||
Start.context ─────────────→ End.context
|
||||
Start.image → processing → End.images
|
||||
├─ image_list → Save.images
|
||||
└─ values → Get Value [value_1] → Save.filename_prefix
|
||||
```
|
||||
|
||||
**New or changed per-item values:** connect the calculated result to End's optional `value_N` input. It replaces the original field with that number, or adds it if Start did not supply it. Every connected override must produce exactly one supported value for each item. A tensor batch remains one value, just as it does at Start. Image, mask, string, integer, float and boolean values are supported; arbitrary objects/models/latents and non-finite numbers are not.
|
||||
|
||||
End shows override inputs progressively. Use **Add value override** to expose a later slot without connecting the earlier ones. Disconnecting an override never renumbers the other End fields: `value_3` always means `value_3`. Unlike Start's native Autogrow inputs, End overrides must keep these identities stable. Existing image/mask/text/context sockets and output indices are unchanged. If the frontend extension is unavailable, Python still declares all optional override inputs and the values output.
|
||||
|
||||
Original values are retained on the master; workers only return explicitly connected overrides, avoiding a redundant return trip for unchanged reference images or other values. Overrides use the existing exact value transport and size limits. Selected workers must advertise `value_results: true` when overrides are requested, checked before execution and again at dispatch. Original-value passthrough alone does not require that additional capability.
|
||||
|
||||
Values leaving the region must come from **End**, not directly from Start. To chain regions, extract an End value and connect it to the next Start's corresponding input. Values remain available downstream after End cleans up the execution context; they do not depend on a later lookup of expired job state.
|
||||
|
||||
Models are loaded and cached on each processing participant, not serialized between machines. Connect loader outputs directly to the processing nodes. The master needs the region's model files when participating, as does each selected worker. The master still validates the prompt: it needs the relevant node definitions and any model-name/file availability required by their validation (including combo lists). This is not a model-free-master guarantee.
|
||||
|
||||
@@ -60,24 +84,109 @@ Per-image data must pass through Start; do not bypass it with upstream images, l
|
||||
|
||||
End's optional `texts` is a connection-only socket. For workflows saved before this socket became connection-only, recreate End and reconnect its wires if an old text widget or literal `texts: ""` remains in the compiled prompt. Arbitrary text literals are not accepted at End.
|
||||
|
||||
File-based dependencies keep ComfyUI's `[output]` / `[temp]` root annotations and are read on the master. Their worker copies use separate content-addressed paths under `input/distributed-media/`; each worker's rewritten prompt is reused across items without reusing item tokens.
|
||||
File-based dependencies keep ComfyUI's `[output]` / `[temp]` root annotations and are read on the master. Direct-worker copies use separate content-addressed paths under `input/distributed-media/`; each direct worker's rewritten prompt is reused across items without reusing item tokens. Cloud requests use the PR2 bounded file-input adapters instead.
|
||||
|
||||
## Image return format and quality
|
||||
|
||||
End adds `transfer_format` (default **exact**) and `transfer_quality` (default **90**).
|
||||
These apply to returned images from **both direct workers and serverless pools**.
|
||||
Master-local results bypass encoding. Input images, masks, generic values and
|
||||
returned masks/text retain their exact existing transport.
|
||||
|
||||
- `exact` preserves floating-point dtype and values.
|
||||
- `legacy_png`, `png`, `ffv1` preserve quantized 8-bit image pixels, not the
|
||||
original floating-point colors. `legacy_png` is PNG without compression.
|
||||
- `jpeg`, `webp`, `h264`, `h265`, `av1` are lossy choices; quality is 1–100.
|
||||
Video codecs encode each assigned image independently, not the whole source list.
|
||||
- Dimensions, floating dtype and alpha survive; alpha is an exact sidecar.
|
||||
- Direct workers must advertise the selected codec in `image_transfer_formats`;
|
||||
older workers fail before submission for non-exact choices. Exact remains
|
||||
compatible with the existing protocol. Serverless endpoints need the updated
|
||||
complete bridge; health probing does not prove codec availability.
|
||||
- The selected software encoder must be installed. Encoding failures stop the
|
||||
item; there is no silent format fallback. No hardware encoder is selected.
|
||||
|
||||
Compression can reduce transfer size, but does **not** bypass decoded tensor,
|
||||
conservative codec scratch, or complete request limits. For details and codec
|
||||
limits, see [image return options](serverless-media.md#image-return-format-and-quality).
|
||||
Mixed local/remote results may have different color fidelity when compression is
|
||||
selected, because local images are deliberately not quantized.
|
||||
|
||||
## Serverless pools and bounded capacity
|
||||
|
||||
Select `media_profile: "comfy-media-v1"` for each cloud pool. The endpoint must
|
||||
include the [standalone PR2 media bridge](../serverless_bridge/README.md), route
|
||||
typed requests through its executor, and provide the required supervisor. Merely
|
||||
installing the bridge nodes or choosing the profile is not sufficient. The
|
||||
endpoint needs the workflow's models and custom nodes, but not the full
|
||||
Distributed extension. A selected `legacy-image-v1` pool fails before the master
|
||||
prompt is queued; it is not silently skipped in favor of local execution.
|
||||
Provider health/authentication probing does not verify bridge/model installation.
|
||||
|
||||
- Each pool is selected once by its configured pool ID. **Images per run
|
||||
(`jobs_per_run`) does not multiply input images or set image-list parallelism.**
|
||||
An assigned cloud job processes exactly one source item and its paired values.
|
||||
- Cloud slots are bounded by the existing **16 outstanding jobs per pool** cap,
|
||||
available manager capacity (**64 pending jobs** across pools), reservations,
|
||||
and the source item count. Unresolved submissions and unconfirmed cancellations
|
||||
still consume capacity. Selecting more items does not enqueue all of them as
|
||||
provider jobs at once. Endpoint-side worker limits can reduce actual concurrency.
|
||||
- Cloud slots, direct workers and the optional master pull from the same queue;
|
||||
no region is cloned once per Images-per-run setting. Results are restored to
|
||||
source order even when participants finish out of order.
|
||||
- Each cloud request uses the standard media-v1 envelope and existing durable
|
||||
serverless job lifecycle. The private item graph replaces Start/End with typed
|
||||
inputs and explicit exporter roots, not independent all-output execution.
|
||||
IMAGE returns are exact by default, or explicitly encoded as selected at End.
|
||||
MASK tensor data, input images and text retain exact typed codecs. There is no
|
||||
image-list-specific endpoint protocol or new visible node socket.
|
||||
- Connected End `value_1` … `value_10` overrides are computed on the endpoint and returned using exact typed `VALUE` exports. Sparse slot identities are preserved; only explicitly connected overrides are transported. The master overlays them on each source item’s original values. BOOLEAN, INT, FLOAT, STRING, IMAGE and MASK remain distinct; missing/extra slots, multiple values and nonfinite values fail. Override tensors remain exact even when the main image uses a compressed return format. End overrides require the updated complete bridge (including `value_codec.py`), not merely an older healthy media endpoint.
|
||||
- Cloud `value_N` tensors must have native IMAGE `[B,H,W,C]` or MASK `[B,H,W]`
|
||||
shapes, not arbitrary rank-one/rank-two tensor payloads. Typed tensors must be
|
||||
finite; connected masks must match their image grid. Unsupported values fail
|
||||
rather than being resized, flattened or converted to lossy images.
|
||||
- Inputs are inline and the **complete serialized Runpod request is capped at
|
||||
10,000,000 bytes**, including workflow and base64. Default inline result JSON
|
||||
has the same cap. The PR2 contract's larger artifact ceilings do not remove
|
||||
these inline limits; larger outputs need an operator-provided artifact publisher.
|
||||
There is no built-in large-input uploader. See [typed media limits](serverless-media.md#size-and-operational-limits).
|
||||
- Missing/duplicate/mismatched results or provider failures fail the region.
|
||||
Cancellation closes only this region's owned cloud jobs and collector queues,
|
||||
including jobs whose dispatch was cancelled before returning a handle. It does
|
||||
not clear unrelated pool work. No failed or ambiguous job is retried automatically.
|
||||
- A provider cancellation or timeout is **not proof that endpoint ComfyUI stopped**.
|
||||
The PR2 supervisor must quarantine/reconcile or terminate an uncertain endpoint
|
||||
before reuse. Preserve its retained files until execution and readers are known
|
||||
to have stopped; see the [bridge uncertainty policy](../serverless_bridge/README.md).
|
||||
|
||||
## Native cache reuse
|
||||
|
||||
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.
|
||||
- Queue preparation works on an execution copy, leaving the saved canvas unchanged. Unused per-item preview/save branches are pruned by dependency reachability, not node names. Required computational utilities marked `OUTPUT_NODE` stay in the graph; explicit `partial_execution_targets` prevent them from also becoming independent outer-master roots. End and legitimate post-End or unrelated outputs are retained. Bypass/rejoin branches still fail boundary validation rather than being silently removed.
|
||||
- Every selected worker must advertise `execution_targets: true` as well as the existing image-list protocol. Older workers fail preflight before dispatch. Both HTTP and WebSocket item submissions target End only; the participating master uses its existing End-only local executor. Submit through `/distributed/queue`: an ordinary all-output `/prompt` request does not receive this relaxation.
|
||||
- Every selected direct worker must advertise `execution_targets: true` as well as the existing image-list protocol. Older workers fail preflight before dispatch. Both HTTP and WebSocket item submissions target End only; the participating master uses its existing End-only local executor. Cloud jobs use explicit media exporter roots. Submit through `/distributed/queue`: an ordinary all-output `/prompt` request does not receive this relaxation.
|
||||
- **Target selection is not a sandbox.** A required utility still runs its own save, preview or other side effects when it executes for an item. Keep intentional final saving after End.
|
||||
- The entire input list is materialized before scheduling; this is not streaming input. The limit is 4,096 frames. Remote transport preserves floating-point tensor dtype and values rather than encoding to lossy images. Each transported tensor is limited to 50 MiB, and the complete JSON/base64 HTTP request to 101 MiB; image/mask/value combinations can hit the aggregate request limit first. In-process master items use tensors directly and are not subject to those HTTP limits.
|
||||
- The entire input list is materialized before scheduling; this is not streaming input. The limit is 4,096 frames. Direct-worker input transport and default exact returns preserve floating-point tensor dtype and values. Opt-in image return encoding does not relax decoded-size limits. Each direct-worker tensor is limited to 50 MiB, and the complete JSON/base64 HTTP request to 101 MiB; image/mask/value combinations can hit the aggregate request limit first. Serverless inline limits are smaller, as described above. In-process master items use tensors directly and are not subject to those HTTP limits.
|
||||
- Results are accepted once per item using assignment tokens. Missing, duplicate, malformed and wrong-index results fail rather than silently returning a partial list. No failed or ambiguous submission is retried automatically.
|
||||
- Worker errors fail the region. The worker timeout measures loss of evidence for the owned prompt ID, not maximum render duration. Pending/running evidence refreshes liveness; unrelated busy jobs and failed polls do not. Healthy cold loads and long renders can continue.
|
||||
- Worker errors fail the region. For direct workers, the worker timeout measures loss of evidence for the owned prompt ID, not maximum render duration. Pending/running evidence refreshes liveness; unrelated busy jobs and failed polls do not. Healthy cold loads and long renders can continue. Serverless jobs instead use their configured serverless timeout and durable manager lifecycle; they are excluded from direct-history/liveness polling.
|
||||
- User interruption is checked during submission and result waiting. Cleanup targets only this region's acknowledged prompt IDs and staged input tokens. Submission may wait briefly for an acknowledgement before cancellation; remote transport and cleanup have finite time limits. Local cleanup stops at node/progress boundaries and waits for its execution thread to finish before the parent prompt can resume. A synchronous node that does not cooperate with cancellation can delay this beyond the configured deadline. Lost acknowledgements are ambiguous, never retried automatically. Unreachable cleanup is logged as unconfirmed and may require checking the worker manually; cancellation does not guarantee remote termination through a broken network.
|
||||
- State is in memory. Restarting the master does not resume an interrupted region. Abandoned master contexts and expired worker staging are actively swept on the server loop; active/queued master work is protected from age-only expiry. End also removes its context on success, failure or interruption. This releases feature-owned references, not ComfyUI's separate execution cache.
|
||||
- Image-list ordering/context state is in memory. Restarting the master does not resume an interrupted region. Cloud job IDs and uncertainty state remain in the existing durable serverless ledger for reconciliation, not automatic list resumption. Abandoned master contexts and expired direct-worker staging are actively swept on the server loop; active/queued master work is protected from age-only expiry. End also removes its context on success, failure or interruption. This releases feature-owned references, not ComfyUI's separate execution cache.
|
||||
|
||||
### Time and size limits
|
||||
|
||||
@@ -93,6 +202,13 @@ Use trusted workers and trusted workflows, as with the existing ComfyUI API. Cap
|
||||
|
||||
The generic acceptance helper starts two isolated loopback workers using real ComfyUI HTTP handlers, `PromptQueue` records and CPU `PromptExecutor` instances. Production queue preparation, list nodes, graph extraction, HTTP item/result routes, dynamic scheduling, float transport, ordering and cancellation are exercised. The harness checks orchestrator-only, mixed master/worker and master-only execution, including exact-once required output utilities, discarded previews, unchanged source graphs, End-only worker roots, post-End capture, source-once materialization and cached local model-loader dependencies. WebSocket target propagation is covered by focused route/dispatch tests. No managed ComfyUI service or model is started.
|
||||
|
||||
Offline cloud integration tests additionally cover mixed-role and cloud-only
|
||||
scheduling, out-of-order image/mask/text pairing, one job per source item regardless
|
||||
of Images per run, profile/selection rejection before queueing, exact cleanup on
|
||||
setup/dispatch/result failures, and the real CPU ComfyUI `processing_interrupted`
|
||||
signal while awaiting a mocked cloud result. Provider calls are mocked: this is
|
||||
not live Runpod, endpoint-supervisor or GPU-workflow acceptance.
|
||||
|
||||
```bash
|
||||
COMFYUI_SOURCE_ROOT=/path/to/ComfyUI /path/to/ComfyUI/.venv/bin/python -m pytest tests -q
|
||||
bash scripts/test-web.sh
|
||||
|
||||
+70
-41
@@ -1,12 +1,19 @@
|
||||
# Typed serverless media (PR2)
|
||||
# Serverless images and audio
|
||||
|
||||
`comfy-media-v1` adds typed list results, exact by default, for `IMAGE`, `AUDIO`, `MASK`,
|
||||
`STRING`, and native ComfyUI `VIDEO`. It is an **explicit endpoint profile**, not
|
||||
an automatically detected capability. Omitting `media_profile` still selects
|
||||
`legacy-image-v1`; existing image collectors and handlers keep their old contract.
|
||||
**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 a typed collector
|
||||
## 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,
|
||||
@@ -15,9 +22,9 @@ See [legacy serverless setup](serverless-workers.md) for credentials and recover
|
||||
**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 Media Collector** (`DistributedMediaCollector`), connect at
|
||||
least one typed input, and connect its outputs to the appropriate downstream
|
||||
nodes. For literal text, connect a `PrimitiveString` node to `texts`.
|
||||
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.
|
||||
|
||||
@@ -25,41 +32,55 @@ See [legacy serverless setup](serverless-workers.md) for credentials and recover
|
||||
installation, model availability, or typed execution. Selecting a profile does
|
||||
not install or upgrade anything on the endpoint.
|
||||
|
||||
The collector has five optional inputs and five **list outputs**, in this order:
|
||||
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 | Each list item |
|
||||
| Socket / output index | Native type | Aggregation |
|
||||
| --- | --- | --- |
|
||||
| `images` / 0 | `IMAGE` | One BHWC tensor, including its batch dimension |
|
||||
| `audio` / 1 | `AUDIO` | `{ "waveform": BCT_tensor, "sample_rate": positive_int }` |
|
||||
| `masks` / 2 | `MASK` | One BHW tensor |
|
||||
| `texts` / 3 | `STRING` | Exact text, including whitespace or an empty string |
|
||||
| `video` / 4 | `VIDEO` | One supported native video object |
|
||||
| `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 the prepared participant
|
||||
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. An unconnected socket returns `[]`. A connected remote
|
||||
socket must return 1–64 items; an empty string is an item, an empty list is not.
|
||||
Lists are not concatenated into one tensor, resampled, resized, or aligned across
|
||||
sockets. Different items may have different valid shapes or audio sample rates.
|
||||
The entire remote response is validated and decoded before any of its results
|
||||
are delivered; missing, duplicate, mismatched, or corrupt results fail collection.
|
||||
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
|
||||
|
||||
`transfer_format` defaults to **`exact`** and `transfer_quality` to **90**.
|
||||
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, MASK, AUDIO, STRING or native VIDEO values.
|
||||
uploaded inputs or AUDIO. The bridge's generic MASK, STRING and native VIDEO
|
||||
codecs are also unaffected.
|
||||
|
||||
- `exact`: preserve floating-point dtype and values, with no image quantization.
|
||||
- `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. Image values are clipped to [0,1] and quantized as in the existing
|
||||
Collector. `legacy_png` uses PNG without compression; `png` compresses it.
|
||||
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`, `av1`: lossy video encoding of each IMAGE tensor's frame batch.
|
||||
This does not transcode the native `video` socket or combine separate list items.
|
||||
- `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; alpha is carried separately where needed, not silently dropped.
|
||||
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;
|
||||
@@ -67,8 +88,11 @@ 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). The original Collector's
|
||||
settings, default and fallback behavior are unchanged.
|
||||
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
|
||||
|
||||
@@ -92,7 +116,11 @@ becoming endpoint-local paths. Models and custom nodes are never transferred.
|
||||
the pure codec itself rejects non-CPU storage. Remote decoding returns
|
||||
independently owned CPU values. STRING transport is exact UTF-8 (`utf8-v1`).
|
||||
|
||||
### Native VIDEO is not a generic video converter
|
||||
### 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()`:
|
||||
@@ -127,11 +155,12 @@ invalidate a returned VIDEO.
|
||||
**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.
|
||||
- Remote typed transport is **serverless only**. Direct-worker typed transport
|
||||
is not implemented. Do not mix a typed collector with legacy image collectors,
|
||||
unsupported distributed/batch nodes, or collector-to-collector dependencies.
|
||||
Downstream processing stays on the master. Image-list scheduling is separate
|
||||
PR3 work, not part of this contract.
|
||||
- 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
|
||||
|
||||
@@ -154,5 +183,5 @@ 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 and its offline tests; they do not claim
|
||||
live Runpod/provider or GPU end-to-end validation.
|
||||
These docs describe the implementation, not a test result or live Runpod/provider
|
||||
or GPU end-to-end validation.
|
||||
|
||||
@@ -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
|
||||
@@ -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.
|
||||
|
||||
@@ -6,6 +6,7 @@ import sys
|
||||
|
||||
from .runtime import (
|
||||
END_CLASS,
|
||||
END_RESULT_INPUTS,
|
||||
_ancestors_for_inputs,
|
||||
_downstream_until,
|
||||
_image_list_pairs,
|
||||
@@ -49,7 +50,7 @@ def prepare_image_list_execution(prompt: dict) -> tuple[dict, list[str]]:
|
||||
required = _reachable(boundaries, upstream)
|
||||
discarded = set()
|
||||
for start_id, end_id in pairs:
|
||||
region_nodes = _ancestors_for_inputs(prepared, end_id, start_id, ("images", "masks", "texts"))
|
||||
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.
|
||||
|
||||
+19
-3
@@ -14,6 +14,20 @@ DEFAULT_ORPHAN_GRACE_SECONDS = 300.0
|
||||
_LIFECYCLE_MARKER = object()
|
||||
|
||||
|
||||
def _without_runtime_fingerprints(prompt: dict) -> dict:
|
||||
"""Compare submitted structure without ComfyUI's per-node cache annotation.
|
||||
|
||||
PromptQueue snapshots the graph before IsChangedCache adds ``is_changed``
|
||||
to executing nodes. Ignore only that top-level field, never inputs or nested
|
||||
metadata, and leave both the queued and executing graphs untouched.
|
||||
"""
|
||||
return {
|
||||
node_id: {key: value for key, value in node.items() if key != "is_changed"}
|
||||
if isinstance(node, dict) else node
|
||||
for node_id, node in prompt.items()
|
||||
}
|
||||
|
||||
|
||||
def has_safe_execution_targets(prompt_server: Any, prompt_id: str | None, prompt: dict) -> bool:
|
||||
"""Prove outer roots from the running queue, not a forgeable prompt marker.
|
||||
|
||||
@@ -23,8 +37,9 @@ def has_safe_execution_targets(prompt_server: Any, prompt_id: str | None, prompt
|
||||
if not prompt_id:
|
||||
return False
|
||||
from .execution_plan import prepare_image_list_execution
|
||||
prepared, targets = prepare_image_list_execution(prompt)
|
||||
if prepared != prompt or not targets:
|
||||
execution_prompt = _without_runtime_fingerprints(prompt)
|
||||
prepared, targets = prepare_image_list_execution(execution_prompt)
|
||||
if prepared != execution_prompt or not targets:
|
||||
return False
|
||||
queue = getattr(prompt_server, "prompt_queue", None)
|
||||
for name in ("get_current_queue_volatile", "get_current_queue"):
|
||||
@@ -35,7 +50,8 @@ def has_safe_execution_targets(prompt_server: Any, prompt_id: str | None, prompt
|
||||
running, _pending = reader()
|
||||
for entry in running:
|
||||
if (isinstance(entry, (tuple, list)) and len(entry) > 4
|
||||
and entry[1] == prompt_id and entry[2] == prompt):
|
||||
and entry[1] == prompt_id and isinstance(entry[2], dict)
|
||||
and _without_runtime_fingerprints(entry[2]) == execution_prompt):
|
||||
roots = entry[4]
|
||||
return (isinstance(roots, (list, tuple))
|
||||
and all(isinstance(root, str) for root in roots)
|
||||
|
||||
@@ -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
|
||||
|
||||
+392
-53
@@ -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:
|
||||
@@ -383,7 +537,7 @@ def create_master_job(start_id: str, end_id: str, prompt: dict[str, Any], items:
|
||||
if not isinstance(master_participates, bool):
|
||||
raise ValueError("master_participates must be a boolean")
|
||||
if not worker_ids and not master_participates:
|
||||
raise RuntimeError("Distributed Image List requires at least one preflighted direct worker; master fallback is disabled")
|
||||
raise RuntimeError("Distributed Image List requires at least one selected worker; master fallback is disabled")
|
||||
# Only the target-controlled Start caller may relax the outer prompt guard.
|
||||
regions = validate_image_list_regions(prompt, explicit_targets=explicit_targets is True)
|
||||
if not any(region["start_id"] == str(start_id) and region["end_id"] == str(end_id)
|
||||
@@ -544,10 +698,11 @@ def validate_image_list_regions(prompt: dict[str, Any], *, explicit_targets: boo
|
||||
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
|
||||
@@ -592,18 +747,28 @@ def validate_image_list_regions(prompt: dict[str, Any], *, explicit_targets: boo
|
||||
|
||||
|
||||
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})
|
||||
@@ -713,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
|
||||
@@ -734,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
|
||||
@@ -770,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 = [], [], []
|
||||
@@ -790,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")
|
||||
@@ -807,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
|
||||
@@ -820,34 +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.dispatch import (
|
||||
preflight_image_list_workers, dispatch_worker_prompt,
|
||||
)
|
||||
from ..api.orchestration.prompt_transform import _resolve_direct_workers_for_image_list
|
||||
from ..utils.config import load_config, get_worker_timeout_seconds
|
||||
from ..utils.network import build_master_callback_url, build_worker_url, get_client_session
|
||||
from ..utils.logging import log
|
||||
import aiohttp
|
||||
from ..utils.config import load_config
|
||||
|
||||
config = load_config()
|
||||
workers = _resolve_direct_workers_for_image_list(config, state.worker_ids)
|
||||
if not workers and not state.master_participates:
|
||||
raise RuntimeError("Distributed Image List requires enabled direct workers; serverless pools are unsupported")
|
||||
# Recheck worker capability and idle state immediately before the first item.
|
||||
worker_statuses = await preflight_image_list_workers(workers, require_idle=not state.master_participates) if workers else {}
|
||||
workers = [worker for worker in workers if worker["id"] in worker_statuses]
|
||||
if not workers and not state.master_participates:
|
||||
raise RuntimeError("No free, reachable, feature-compatible direct workers are available; master execution is disabled")
|
||||
session = await get_client_session() if workers else None
|
||||
# Both item paths below select End explicitly, never all OUTPUT_NODE roots.
|
||||
# Every item path selects End explicitly, never all OUTPUT_NODE roots.
|
||||
regions = validate_image_list_regions(state.original_prompt, explicit_targets=True)
|
||||
region = next((entry for entry in regions if entry["start_id"] == state.start_id and entry["end_id"] == state.end_id), None)
|
||||
if region is None:
|
||||
raise RuntimeError("Image-list region no longer matches the submitted prompt")
|
||||
controller = None
|
||||
try:
|
||||
if _has_selected_serverless_workers(config, state.worker_ids):
|
||||
from .serverless import create_serverless_slots
|
||||
controller = await create_serverless_slots(state, region, config, state.event)
|
||||
return await _run_image_list_slots(state, config, workers, region, controller,
|
||||
include_values=include_values)
|
||||
finally:
|
||||
# Own cloud jobs even if direct preflight/local setup fails, or dispatch
|
||||
# is cancelled before returning an acknowledged handle to the scheduler.
|
||||
if controller is not None:
|
||||
await controller.close()
|
||||
|
||||
|
||||
async def _run_image_list_slots(state, config, workers, region, controller, *, include_values=False):
|
||||
context_id = state.context_id
|
||||
from ..api.orchestration.dispatch import (
|
||||
preflight_image_list_workers, dispatch_worker_prompt,
|
||||
)
|
||||
from ..utils.config import get_worker_timeout_seconds
|
||||
from ..utils.network import build_master_callback_url, build_worker_url, get_client_session
|
||||
from ..utils.logging import log
|
||||
import aiohttp
|
||||
|
||||
serverless_slots = list(controller.slots) if controller is not None else []
|
||||
end_inputs = state.original_prompt[state.end_id]["inputs"]
|
||||
transfer_format, _ = validate_options(
|
||||
end_inputs.get("transfer_format", DEFAULT_FORMAT),
|
||||
end_inputs.get("transfer_quality", DEFAULT_QUALITY),
|
||||
)
|
||||
transfer_requirements = ({"required_transfer_formats": (transfer_format,)}
|
||||
if transfer_format != "exact" else {})
|
||||
override_names = end_override_names(end_inputs)
|
||||
value_preflight = {"require_value_results": True} if override_names else {}
|
||||
# Direct capability/idle probes must never be sent to a serverless pool.
|
||||
# Busy peers are filtered in either role; they are fatal only when neither
|
||||
# the master nor an explicitly selected cloud slot can process the queue.
|
||||
worker_statuses = await preflight_image_list_workers(
|
||||
workers, require_idle=not state.master_participates and not serverless_slots,
|
||||
**transfer_requirements, **value_preflight,
|
||||
) if workers else {}
|
||||
workers = [worker for worker in workers if worker["id"] in worker_statuses]
|
||||
if not workers and not serverless_slots and not state.master_participates:
|
||||
raise RuntimeError("No reachable, feature-compatible image-list workers are available; master execution is disabled")
|
||||
session = await get_client_session() if workers else None
|
||||
|
||||
timeout_seconds = max(5.0, float(get_worker_timeout_seconds()))
|
||||
use_websocket = bool((config.get("settings", {}) or {}).get("websocket_orchestration", False))
|
||||
prepared_worker_prompts: dict[str, dict[str, Any]] = {}
|
||||
@@ -862,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()
|
||||
@@ -871,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)
|
||||
@@ -881,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])}
|
||||
@@ -894,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,
|
||||
@@ -906,9 +1228,11 @@ 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,
|
||||
@@ -926,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
|
||||
@@ -975,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
|
||||
@@ -989,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:
|
||||
@@ -998,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.
|
||||
@@ -1038,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.
|
||||
|
||||
@@ -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)
|
||||
+240
-8
@@ -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}
|
||||
@@ -88,6 +121,46 @@ class DistributedCollectorNode:
|
||||
return value[0]
|
||||
return value
|
||||
|
||||
@staticmethod
|
||||
def _combine_image_tensors(tensors, sources):
|
||||
"""Validate BHWC batches, retaining alpha and deterministic input order."""
|
||||
nonempty = []
|
||||
reference_shape = None
|
||||
reference_source = None
|
||||
for tensor, source in zip(tensors, sources):
|
||||
if not isinstance(tensor, torch.Tensor):
|
||||
raise ValueError(f"Collector expected an IMAGE tensor from {source}")
|
||||
shape = tuple(tensor.shape)
|
||||
if tensor.ndim != 4 or shape[-1] not in (1, 3, 4) or min(shape[1:3]) <= 0:
|
||||
raise ValueError(
|
||||
f"Collector received invalid IMAGE shape {shape} from {source}; "
|
||||
"expected [B, H, W, C] with 1, 3, or 4 channels"
|
||||
)
|
||||
# Empty master placeholders must not determine the worker image grid.
|
||||
if shape[0] == 0:
|
||||
continue
|
||||
if reference_shape is None:
|
||||
reference_shape, reference_source = shape, source
|
||||
elif shape[1:3] != reference_shape[1:3]:
|
||||
raise ValueError(
|
||||
"Collector cannot combine images with different spatial dimensions: "
|
||||
f"{reference_source} has shape {reference_shape}; {source} has shape {shape}"
|
||||
)
|
||||
nonempty.append(tensor)
|
||||
if not nonempty:
|
||||
raise ValueError("Collector cannot return an empty image batch")
|
||||
|
||||
channels = max(tensor.shape[-1] for tensor in nonempty)
|
||||
normalized = []
|
||||
for tensor in nonempty:
|
||||
if tensor.shape[-1] == 1 and channels > 1:
|
||||
tensor = tensor.expand(-1, -1, -1, 3)
|
||||
if tensor.shape[-1] == 3 and channels == 4:
|
||||
# RGB is opaque; never discard alpha from another participant.
|
||||
tensor = torch.cat((tensor, torch.ones_like(tensor[..., :1])), dim=-1)
|
||||
normalized.append(ensure_contiguous(tensor))
|
||||
return ensure_contiguous(torch.cat(normalized, dim=0))
|
||||
|
||||
def _normalize_images_input(self, images):
|
||||
"""Collapse ComfyUI list IMAGE inputs into a normal batched IMAGE tensor."""
|
||||
if isinstance(images, (list, tuple)):
|
||||
@@ -97,7 +170,9 @@ class DistributedCollectorNode:
|
||||
raise TypeError("Collector expected IMAGE list items to be torch.Tensor instances")
|
||||
if len(images) == 1:
|
||||
return ensure_contiguous(images[0])
|
||||
return ensure_contiguous(torch.cat([ensure_contiguous(image) for image in images], dim=0))
|
||||
return self._combine_image_tensors(
|
||||
images, [f"IMAGE list item {index}" for index in range(len(images))]
|
||||
)
|
||||
return ensure_contiguous(images)
|
||||
|
||||
def _normalize_audio_input(self, audio):
|
||||
@@ -200,6 +275,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 +408,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 +630,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,
|
||||
@@ -468,9 +666,11 @@ class DistributedCollectorNode:
|
||||
):
|
||||
"""Assemble final tensor, or return None when the job contains only audio."""
|
||||
ordered_tensors = []
|
||||
sources = []
|
||||
if not delegate_mode and images_on_cpu is not None:
|
||||
for i in range(master_batch_size):
|
||||
ordered_tensors.append(images_on_cpu[i:i+1])
|
||||
sources.append(f"master image {i}")
|
||||
|
||||
ordered_worker_ids = [str(worker_id) for worker_id in (worker_order or [])]
|
||||
seen = set()
|
||||
@@ -480,6 +680,7 @@ class DistributedCollectorNode:
|
||||
continue
|
||||
for idx in sorted(worker_images[worker_id_str].keys()):
|
||||
ordered_tensors.append(worker_images[worker_id_str][idx])
|
||||
sources.append(f"worker {worker_id_str!r} image {idx}")
|
||||
|
||||
# Append any unexpected worker ids deterministically.
|
||||
for worker_id_str in sorted(worker_images.keys()):
|
||||
@@ -487,6 +688,7 @@ class DistributedCollectorNode:
|
||||
continue
|
||||
for idx in sorted(worker_images[worker_id_str].keys()):
|
||||
ordered_tensors.append(worker_images[worker_id_str][idx])
|
||||
sources.append(f"worker {worker_id_str!r} image {idx}")
|
||||
|
||||
cpu_tensors = []
|
||||
for t in ordered_tensors:
|
||||
@@ -496,9 +698,9 @@ class DistributedCollectorNode:
|
||||
cpu_tensors.append(t)
|
||||
|
||||
if cpu_tensors:
|
||||
return ensure_contiguous(torch.cat(cpu_tensors, dim=0))
|
||||
return self._combine_image_tensors(cpu_tensors, sources)
|
||||
if fallback_images is not None:
|
||||
return ensure_contiguous(fallback_images)
|
||||
return self._combine_image_tensors([fallback_images], ["master fallback"])
|
||||
return None
|
||||
|
||||
async def execute(self, images, audio, load_balance=False, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False, transfer_format=transfer_codec.DEFAULT_TRANSFER_FORMAT, transfer_quality=transfer_codec.DEFAULT_TRANSFER_QUALITY):
|
||||
@@ -534,6 +736,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 +815,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 +960,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,14 +984,15 @@ 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.
|
||||
return (images, combined_audio)
|
||||
# Do not silently discard worker images or pass an empty
|
||||
# fallback batch downstream to SaveImage.
|
||||
raise
|
||||
except BaseException:
|
||||
# Remove the queue first, preventing cancelled work from delivering into
|
||||
# a collector that is already unwinding.
|
||||
|
||||
+142
-14
@@ -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
|
||||
@@ -149,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,
|
||||
@@ -181,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:
|
||||
@@ -195,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]
|
||||
@@ -229,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):
|
||||
|
||||
@@ -1,229 +0,0 @@
|
||||
"""Typed, lossless list collection for serverless media participants.
|
||||
|
||||
The master compiler supplies expected_media before pruning delegate inputs.
|
||||
The provider manager owns remote deadlines and delivers one atomic native-media
|
||||
message per participant. This node never serializes, concatenates or resamples.
|
||||
"""
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
|
||||
import comfy.model_management
|
||||
from comfy_api.v0_0_2 import io
|
||||
import server
|
||||
|
||||
from ..serverless_bridge.contract import MAX_ITEMS, SOCKET_KINDS, identifier
|
||||
from ..serverless_bridge.image_transfer import (FORMATS, DEFAULT_FORMAT, DEFAULT_QUALITY,
|
||||
FORMAT_TOOLTIP, QUALITY_TOOLTIP, validate_options)
|
||||
from ..utils.async_helpers import run_async_in_server_loop
|
||||
|
||||
POLL_SECONDS = 0.1
|
||||
START_GRACE_SECONDS = 2.0
|
||||
|
||||
|
||||
def _control(value, name, kind):
|
||||
# Only execution controls are scalars wrapped by ComfyUI's list mapper.
|
||||
if isinstance(value, list):
|
||||
if len(value) != 1:
|
||||
raise ValueError(f'{name} must be a singleton control')
|
||||
value = value[0]
|
||||
if type(value) is not kind:
|
||||
raise ValueError(f'Invalid {name} control')
|
||||
return value
|
||||
|
||||
|
||||
def _names(value, name):
|
||||
try:
|
||||
result = json.loads(_control(value, name, str))
|
||||
except (ValueError, TypeError):
|
||||
raise ValueError(f'{name} must be a JSON string list') from None
|
||||
if (not isinstance(result, list) or any(not isinstance(item, str) for item in result)
|
||||
or len(result) != len(set(result))):
|
||||
raise ValueError(f'{name} must contain unique strings')
|
||||
return result
|
||||
|
||||
|
||||
def _validate_media(media, expected):
|
||||
if not isinstance(media, dict) or set(media) != set(expected):
|
||||
raise ValueError('Media sockets do not match expected_media')
|
||||
for socket, items in media.items():
|
||||
if (socket not in SOCKET_KINDS or not isinstance(items, list)
|
||||
or not 1 <= len(items) <= MAX_ITEMS or any(item is None for item in items)):
|
||||
raise ValueError('Connected media sockets require nonempty bounded item lists')
|
||||
|
||||
|
||||
def _prepare(media, controls):
|
||||
worker = _control(controls['is_worker'], 'is_worker', bool)
|
||||
delegate = _control(controls['delegate_only'], 'delegate_only', bool)
|
||||
passthrough = _control(controls['pass_through'], 'pass_through', bool)
|
||||
_control(controls['load_balance'], 'load_balance', bool) # Routing is master-owned.
|
||||
if worker:
|
||||
raise ValueError('DistributedMediaCollector does not support direct-worker execution')
|
||||
participants = _names(controls['enabled_worker_ids'], 'enabled_worker_ids')
|
||||
for participant in participants:
|
||||
identifier(participant)
|
||||
if not participant.startswith('serverless:') or not participant[len('serverless:'):]:
|
||||
raise ValueError('DistributedMediaCollector requires serverless participant IDs')
|
||||
expected = _names(controls['expected_media'], 'expected_media')
|
||||
if any(socket not in SOCKET_KINDS for socket in expected):
|
||||
raise ValueError('Unknown expected_media socket')
|
||||
local = {socket: items for socket, items in media.items() if items is not None}
|
||||
_validate_media(local, local)
|
||||
expected = expected or list(local)
|
||||
if not expected:
|
||||
raise ValueError('At least one connected media socket is required')
|
||||
if not delegate or passthrough:
|
||||
_validate_media(local, expected)
|
||||
if delegate and not passthrough and not participants:
|
||||
raise ValueError('Delegate-only collection requires serverless participants')
|
||||
return local, participants, expected, delegate, passthrough
|
||||
|
||||
|
||||
def _outputs(media):
|
||||
return tuple(media.get(socket, []) for socket in SOCKET_KINDS)
|
||||
|
||||
|
||||
def _consume(message, participants, expected, received):
|
||||
if not isinstance(message, dict):
|
||||
raise RuntimeError('Invalid serverless media delivery')
|
||||
participant = message.get('worker_id')
|
||||
if not isinstance(participant, str) or participant not in participants:
|
||||
raise RuntimeError('Unexpected serverless participant')
|
||||
if participant in received:
|
||||
raise RuntimeError('Duplicate serverless participant delivery')
|
||||
if message.get('is_last') is not True:
|
||||
raise RuntimeError('Serverless media delivery must be atomic and final')
|
||||
if 'serverless_error' in message:
|
||||
if (set(message) != {'worker_id', 'serverless_error', 'is_last'}
|
||||
or not isinstance(message['serverless_error'], str) or not message['serverless_error']):
|
||||
raise RuntimeError('Invalid serverless error delivery')
|
||||
# Do not echo provider-controlled text, credentials or signed URLs.
|
||||
raise RuntimeError(f'Serverless participant {participant} failed to produce media')
|
||||
if set(message) != {'worker_id', 'media', 'is_last'}:
|
||||
raise RuntimeError('Invalid serverless media delivery fields')
|
||||
_validate_media(message['media'], expected)
|
||||
received[participant] = message['media']
|
||||
|
||||
|
||||
async def _collect(job, media, controls):
|
||||
ps = server.PromptServer.instance
|
||||
queue = None
|
||||
try:
|
||||
# Reuse the master's precreated queue, including early completions.
|
||||
async with ps.distributed_jobs_lock:
|
||||
queue = ps.distributed_pending_jobs.setdefault(job, asyncio.Queue())
|
||||
local, participants, expected, delegate, passthrough = _prepare(media, controls)
|
||||
if passthrough or not participants:
|
||||
return _outputs(local)
|
||||
received = {}
|
||||
started = time.monotonic()
|
||||
while len(received) < len(participants):
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
async with ps.distributed_jobs_lock:
|
||||
if ps.distributed_pending_jobs.get(job) is not queue:
|
||||
raise RuntimeError('Serverless collector queue was removed or replaced')
|
||||
try:
|
||||
message = await asyncio.wait_for(queue.get(), timeout=POLL_SECONDS)
|
||||
except asyncio.TimeoutError:
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
# wait_for cancellation may race an atomic manager delivery.
|
||||
if not queue.empty():
|
||||
continue
|
||||
manager = getattr(ps, 'distributed_serverless', None)
|
||||
for participant in participants:
|
||||
if participant in received:
|
||||
continue
|
||||
if manager is not None and manager.is_pending(job, participant):
|
||||
continue # The manager enforces its remote absolute deadline.
|
||||
if manager is not None and manager.owns(job, participant):
|
||||
raise RuntimeError('Serverless participant stopped without delivering media')
|
||||
# The master queues its prompt before manager.start journals slots.
|
||||
if time.monotonic() - started >= START_GRACE_SECONDS:
|
||||
raise RuntimeError('Serverless manager did not start the expected participant')
|
||||
continue
|
||||
_consume(message, participants, expected, received)
|
||||
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
# Drain and close atomically with the manager's delivery lock. A queued
|
||||
# duplicate after the final expected message must not yield success.
|
||||
async with ps.distributed_jobs_lock:
|
||||
if ps.distributed_pending_jobs.get(job) is not queue:
|
||||
raise RuntimeError('Serverless collector queue was removed or replaced')
|
||||
while not queue.empty():
|
||||
_consume(queue.get_nowait(), participants, expected, received)
|
||||
del ps.distributed_pending_jobs[job]
|
||||
ordered = [] if delegate else [local]
|
||||
ordered.extend(received[participant] for participant in participants)
|
||||
return tuple([item for source in ordered for item in source.get(socket, [])]
|
||||
for socket in SOCKET_KINDS)
|
||||
except BaseException:
|
||||
# Remove only our queue before cancellation can trigger final deliveries.
|
||||
async with ps.distributed_jobs_lock:
|
||||
if queue is not None and ps.distributed_pending_jobs.get(job) is queue:
|
||||
del ps.distributed_pending_jobs[job]
|
||||
manager = getattr(ps, 'distributed_serverless', None)
|
||||
if manager is not None:
|
||||
try:
|
||||
await manager.cancel(collector_job_id=job)
|
||||
except BaseException:
|
||||
pass # Preserve the original failure/interruption, not cleanup's.
|
||||
raise
|
||||
finally:
|
||||
async with ps.distributed_jobs_lock:
|
||||
if queue is not None and ps.distributed_pending_jobs.get(job) is queue:
|
||||
del ps.distributed_pending_jobs[job]
|
||||
|
||||
|
||||
async def _bridge_result(coro):
|
||||
# The legacy sync bridge catches Exception, not asyncio.CancelledError.
|
||||
# Carry BaseException as data so cancellation is re-raised on the caller.
|
||||
try:
|
||||
return await coro, None
|
||||
except BaseException as error:
|
||||
return None, error
|
||||
|
||||
|
||||
class DistributedMediaCollector(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
types = (io.Image, io.Audio, io.Mask, io.String, io.Video)
|
||||
return io.Schema(
|
||||
node_id='DistributedMediaCollector', display_name='Distributed Media Collector',
|
||||
category='utils',
|
||||
inputs=[io.Boolean.Input('load_balance', default=False),
|
||||
# Native non-widget inputs expose forceInput via extra_dict.
|
||||
*[kind.Input(socket, optional=True,
|
||||
**({'force_input': True} if socket == 'texts'
|
||||
else {'extra_dict': {'forceInput': True}}))
|
||||
for socket, kind in zip(SOCKET_KINDS, types)],
|
||||
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=[kind.Output(display_name=socket, is_output_list=True)
|
||||
for socket, kind in zip(SOCKET_KINDS, types)],
|
||||
is_input_list=True, accept_all_inputs=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, images=None, audio=None, masks=None, texts=None, video=None,
|
||||
load_balance=False, multi_job_id='', is_worker=False,
|
||||
enabled_worker_ids='[]', delegate_only=False, pass_through=False,
|
||||
expected_media='[]', transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY):
|
||||
validate_options(_control(transfer_format, 'transfer_format', str),
|
||||
_control(transfer_quality, 'transfer_quality', int))
|
||||
media = dict(images=images, audio=audio, masks=masks, texts=texts, video=video)
|
||||
controls = dict(load_balance=load_balance, is_worker=is_worker,
|
||||
enabled_worker_ids=enabled_worker_ids, delegate_only=delegate_only,
|
||||
pass_through=pass_through, expected_media=expected_media)
|
||||
job = _control(multi_job_id, 'multi_job_id', str)
|
||||
if not job:
|
||||
local, participants, _, _, passthrough = _prepare(media, controls)
|
||||
if participants and not passthrough:
|
||||
raise ValueError('Serverless collection requires multi_job_id')
|
||||
return io.NodeOutput(*_outputs(local))
|
||||
identifier(job)
|
||||
result, error = run_async_in_server_loop(_bridge_result(_collect(job, media, controls)))
|
||||
if error is not None:
|
||||
raise error
|
||||
return io.NodeOutput(*result)
|
||||
+2
-3
@@ -10,8 +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 .media_collector import DistributedMediaCollector
|
||||
from .image_list import DistributedImageListEnd, DistributedImageListStart, DistributedImageListGetValue
|
||||
|
||||
|
||||
class DistributedSeed(io.ComfyNode):
|
||||
@@ -213,4 +212,4 @@ class DistributedWorkerSink(io.ComfyNode):
|
||||
|
||||
NODES = [DistributedCollector, DistributedSeed, DistributedModelName, DistributedValue,
|
||||
ImageBatchDivider, AudioBatchDivider, DistributedEmptyImage, UltimateSDUpscaleDistributed,
|
||||
DistributedWorkerSink, DistributedImageListStart, DistributedImageListEnd, DistributedMediaCollector]
|
||||
DistributedWorkerSink, DistributedImageListStart, DistributedImageListEnd, DistributedImageListGetValue]
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "ComfyUI-Distributed"
|
||||
description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
|
||||
version = "1.5.0"
|
||||
version = "1.6.0"
|
||||
license = {file = "LICENSE"}
|
||||
dependencies = []
|
||||
|
||||
|
||||
+44
-13
@@ -6,8 +6,12 @@ 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.
|
||||
|
||||
For collector setup, list ordering, supported native VIDEO states, and PR2
|
||||
limitations, see [typed serverless media](../docs/serverless-media.md).
|
||||
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
|
||||
|
||||
@@ -70,7 +74,8 @@ returning `{"error": ...}` or marking a provider job `FAILED` is not that policy
|
||||
|
||||
## Request and response
|
||||
|
||||
A minimal text request builder (no files or model required by this graph):
|
||||
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
|
||||
@@ -127,7 +132,7 @@ For typed inputs, `media.artifacts` contains records with exactly:
|
||||
|
||||
| Field | Value |
|
||||
| --- | --- |
|
||||
| `id`, `kind` | Unique artifact ID and one of the five native types |
|
||||
| `id`, `kind` | Unique artifact ID and a native media type or internal `VALUE` |
|
||||
| `descriptor` | Kind-specific codec metadata; use the shared codec, not pickle |
|
||||
| `byte_length`, `sha256` | Exact payload byte count and lowercase SHA-256 hex |
|
||||
| `transport` | `{"type": "inline", "data_base64": "..."}` or the URL transport below |
|
||||
@@ -144,7 +149,7 @@ to owned staging paths. They are not an arbitrary filename-upload interface.
|
||||
The returned `media` object has exactly `version`, `profile`, `request_id`,
|
||||
`bindings`, `artifacts`, and `items`:
|
||||
|
||||
- `bindings`: every declared binding, unchanged, plus `count` (1–64).
|
||||
- `bindings`: every declared binding, unchanged, plus `count` (1–64 for media; exactly 1 for End value overrides).
|
||||
- `artifacts`: the same artifact schema above with inline or URL transport.
|
||||
- `items`: `{ "binding_id": "...", "item_index": 0, "artifact_id": "..." }`
|
||||
records covering every index exactly once, ordered by `item_index` when decoded.
|
||||
@@ -155,20 +160,46 @@ 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. Defaults are `exact` and `90`; old exact requests
|
||||
omit the controls and retain their original representation. Other media exports
|
||||
must not carry image controls. The master compiler supplies them from the typed
|
||||
Collector's widgets; callers should not alter binding identity metadata.
|
||||
`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.
|
||||
|
||||
Formats: `exact`, `legacy_png`, `png`, `ffv1`, `jpeg`, `webp`, `h264`, `h265`, `av1`.
|
||||
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 still use their original exact codecs.
|
||||
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
|
||||
@@ -247,5 +278,5 @@ 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.
|
||||
|
||||
Verification scope: source and offline tests. No live provider/GPU validation is
|
||||
claimed, and this directory is not a tested turnkey provider deployment.
|
||||
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.
|
||||
|
||||
@@ -24,7 +24,7 @@ _LAYOUTS = {"IMAGE": "BHWC", "MASK": "BHW", "AUDIO": "BCT"}
|
||||
|
||||
|
||||
def _validate_kind(kind):
|
||||
if type(kind) is not str or kind not in ("IMAGE", "MASK", "STRING", "AUDIO"):
|
||||
if type(kind) is not str or kind not in ("IMAGE", "MASK", "STRING", "AUDIO", "VALUE"):
|
||||
raise ValueError("Unsupported payload kind")
|
||||
|
||||
|
||||
@@ -92,6 +92,9 @@ def encode_payload(kind: str, value, *, max_bytes: int = MAX_PAYLOAD_BYTES) -> t
|
||||
"""
|
||||
_validate_kind(kind)
|
||||
limit = _byte_limit(max_bytes)
|
||||
if kind == "VALUE":
|
||||
from .value_codec import encode_value
|
||||
return encode_value(value, max_bytes=limit)
|
||||
if kind == "STRING":
|
||||
return _encode_string(value, limit)
|
||||
|
||||
@@ -146,6 +149,9 @@ def decode_payload(kind: str, descriptor: dict, payload: bytes, *, max_bytes: in
|
||||
raise ValueError("Payload must be bytes")
|
||||
if len(payload) > limit:
|
||||
raise ValueError("Payload exceeds byte limit")
|
||||
if kind == "VALUE":
|
||||
from .value_codec import decode_value
|
||||
return decode_value(descriptor, payload, max_bytes=limit)
|
||||
if kind == "IMAGE" and descriptor.get("encoding") == "image-transfer-v1":
|
||||
# Lazy and IMAGE-only: exact tensors and other kinds stay unchanged.
|
||||
from .image_codec import decode_image
|
||||
|
||||
@@ -19,6 +19,9 @@ MAX_ARTIFACTS = 128
|
||||
MAX_ITEMS = 64
|
||||
MAX_DESCRIPTOR_BYTES = 128 * 1024
|
||||
SOCKET_KINDS = {"images": "IMAGE", "audio": "AUDIO", "masks": "MASK", "texts": "STRING", "video": "VIDEO"}
|
||||
# Internal End slots do not extend the public Collector or legacy media sockets.
|
||||
VALUE_SOCKETS = tuple(f"value_{i}" for i in range(1, 11))
|
||||
ARTIFACT_KINDS = frozenset(SOCKET_KINDS.values()) | {"VALUE"}
|
||||
BINDING_FIELDS = {"binding_id", "collector_id", "job_id", "participant_id", "node_id", "socket", "kind"}
|
||||
_ID = re.compile(r"[A-Za-z0-9_.:/-]{1,200}\Z")
|
||||
_HASH = re.compile(r"[a-f0-9]{64}\Z")
|
||||
@@ -68,7 +71,8 @@ def validate_binding(binding):
|
||||
raise ValueError("Invalid media binding")
|
||||
for field in BINDING_FIELDS - {"socket", "kind"}:
|
||||
identifier(binding[field])
|
||||
if SOCKET_KINDS.get(binding["socket"]) != binding["kind"]:
|
||||
kind = "VALUE" if binding["socket"] in VALUE_SOCKETS else SOCKET_KINDS.get(binding["socket"])
|
||||
if kind is None or kind != binding["kind"]:
|
||||
raise ValueError("Media socket/type mismatch")
|
||||
return binding
|
||||
|
||||
@@ -113,7 +117,7 @@ def validate_artifact(artifact, *, max_bytes=MAX_ARTIFACT_BYTES):
|
||||
if not isinstance(artifact, dict) or set(artifact) != fields:
|
||||
raise ValueError("Invalid media artifact")
|
||||
identifier(artifact["id"])
|
||||
if artifact["kind"] not in SOCKET_KINDS.values():
|
||||
if artifact["kind"] not in ARTIFACT_KINDS:
|
||||
raise ValueError("Unsupported media kind")
|
||||
if not isinstance(artifact["descriptor"], dict):
|
||||
raise ValueError("Invalid media descriptor")
|
||||
@@ -191,7 +195,8 @@ def validate_response(response, *, request_id, bindings, max_total_bytes=MAX_TOT
|
||||
if bid in counts or expected.get(bid) != binding:
|
||||
raise ValueError("Media binding identity mismatch")
|
||||
# Connected outputs must produce at least one value. Empty STRING is a value.
|
||||
counts[bid] = bounded_int(row["count"], MAX_ITEMS, minimum=1)
|
||||
maximum = 1 if binding["socket"] in VALUE_SOCKETS else MAX_ITEMS
|
||||
counts[bid] = bounded_int(row["count"], maximum, minimum=1)
|
||||
artifacts = response["artifacts"]
|
||||
if not isinstance(artifacts, list) or not 1 <= len(artifacts) <= MAX_ARTIFACTS:
|
||||
raise ValueError("Invalid media artifacts count")
|
||||
|
||||
@@ -40,7 +40,7 @@ DRAIN_GRACE_SECONDS = 10.0
|
||||
HTTP_TIMEOUT_SECONDS = 15.0
|
||||
_UNCERTAIN_WORK = set()
|
||||
GET_ATTEMPTS = 3
|
||||
_IMPORT_SLOTS = {"IMAGE": 0, "AUDIO": 1, "MASK": 2, "STRING": 3, "VIDEO": 4}
|
||||
_IMPORT_SLOTS = {"IMAGE": 0, "AUDIO": 1, "MASK": 2, "STRING": 3, "VIDEO": 4, "VALUE": 5}
|
||||
# Used ONLY to run the shared manifest validator before /view. Never fetched.
|
||||
_VALIDATION_TRANSPORT = {"type": "url", "url": "https://media-validation.invalid/artifact",
|
||||
"expires_at": "9999-01-01T00:00:00Z"}
|
||||
@@ -194,8 +194,12 @@ def _validate_request(input_data):
|
||||
or aid not in indexed or (node_id, name) in destinations):
|
||||
raise ValueError("Unknown or duplicate media input reference")
|
||||
binding = exporters.get(node_id)
|
||||
if binding is not None and (name != binding["socket"] or indexed[aid]["kind"] != binding["kind"]):
|
||||
raise ValueError("Exporter input must match its declared media socket and kind")
|
||||
if binding is not None:
|
||||
# VALUE exports accept the six explicit runtime types. IMAGE/MASK
|
||||
# imports can feed both native consumers and a sparse override.
|
||||
allowed = {"VALUE", "IMAGE", "MASK", "STRING"} if binding["kind"] == "VALUE" else {binding["kind"]}
|
||||
if name != binding["socket"] or indexed[aid]["kind"] not in allowed:
|
||||
raise ValueError("Exporter input must match its declared media socket and kind")
|
||||
destinations.add((node_id, name))
|
||||
used.add(aid)
|
||||
if used != set(indexed):
|
||||
|
||||
@@ -11,7 +11,7 @@ from .codec import decode_payload, encode_payload
|
||||
from .image_transfer import (FORMATS, DEFAULT_FORMAT, DEFAULT_QUALITY,
|
||||
FORMAT_TOOLTIP, QUALITY_TOOLTIP, validate_options)
|
||||
from .contract import (
|
||||
MAX_ARTIFACT_BYTES, MAX_ITEMS, MAX_TOTAL_BYTES, SOCKET_KINDS,
|
||||
MAX_ARTIFACT_BYTES, MAX_ITEMS, MAX_TOTAL_BYTES, SOCKET_KINDS, VALUE_SOCKETS, ARTIFACT_KINDS,
|
||||
bounded_int, identifier, validate_binding, validate_json,
|
||||
)
|
||||
|
||||
@@ -38,6 +38,9 @@ def _json(text):
|
||||
def _cpu(kind, value):
|
||||
# Runtime transfer only, never sampling or device discovery. Pure codecs
|
||||
# deliberately reject non-CPU storage rather than making hidden transfers.
|
||||
if kind == "VALUE":
|
||||
from .value_codec import value_kind
|
||||
kind = value_kind(value)
|
||||
if kind in ("IMAGE", "MASK"):
|
||||
return value.detach().cpu()
|
||||
if kind == "AUDIO":
|
||||
@@ -72,6 +75,8 @@ def export_values(request_id, binding, values, *, transfer_format=DEFAULT_FORMAT
|
||||
raise ValueError("Image transfer controls apply only to IMAGE exports")
|
||||
if not isinstance(values, list) or not 1 <= len(values) <= MAX_ITEMS:
|
||||
raise ValueError("Connected media output has no values or exceeds item limit")
|
||||
if binding["socket"] in VALUE_SOCKETS and len(values) != 1:
|
||||
raise ValueError("End override requires exactly one value")
|
||||
request_key = hashlib.sha256(request_id.encode()).hexdigest()
|
||||
subfolder = f"distributed-media-v1/{request_key}"
|
||||
root = Path(folder_paths.get_temp_directory()).resolve()
|
||||
@@ -114,7 +119,7 @@ def import_value(artifact, filename):
|
||||
if not isinstance(artifact, dict) or set(artifact) != fields:
|
||||
raise ValueError("Invalid staged media metadata")
|
||||
identifier(artifact["id"])
|
||||
if artifact["kind"] not in SOCKET_KINDS.values():
|
||||
if artifact["kind"] not in ARTIFACT_KINDS:
|
||||
raise ValueError("Invalid staged media kind")
|
||||
size = bounded_int(artifact["byte_length"], MAX_ARTIFACT_BYTES)
|
||||
if not isinstance(filename, str):
|
||||
@@ -138,13 +143,14 @@ class MediaImportV1(io.ComfyNode):
|
||||
def define_schema(cls):
|
||||
return io.Schema(node_id="MediaImportV1", category="distributed/internal", is_deprecated=True,
|
||||
inputs=[io.String.Input("artifact_json"), io.String.Input("filename")],
|
||||
outputs=[io.Image.Output(), io.Audio.Output(), io.Mask.Output(), io.String.Output(), io.Video.Output()])
|
||||
outputs=[io.Image.Output(), io.Audio.Output(), io.Mask.Output(), io.String.Output(), io.Video.Output(),
|
||||
io.AnyType.Output()])
|
||||
|
||||
@classmethod
|
||||
def execute(cls, artifact_json, filename):
|
||||
artifact = _json(artifact_json)
|
||||
value = import_value(artifact, filename)
|
||||
return io.NodeOutput(*(value if kind == artifact["kind"] else None for kind in SOCKET_KINDS.values()))
|
||||
return io.NodeOutput(*(value if kind == artifact["kind"] else None for kind in (*SOCKET_KINDS.values(), "VALUE")))
|
||||
|
||||
|
||||
class MediaExportV1(io.ComfyNode):
|
||||
@@ -155,6 +161,7 @@ class MediaExportV1(io.ComfyNode):
|
||||
io.Image.Input("images", optional=True), io.Audio.Input("audio", optional=True),
|
||||
io.Mask.Input("masks", optional=True), io.String.Input("texts", optional=True, force_input=True),
|
||||
io.Video.Input("video", optional=True),
|
||||
*[io.AnyType.Input(name, optional=True) for name in VALUE_SOCKETS],
|
||||
io.Combo.Input("transfer_format", options=list(FORMATS), default=DEFAULT_FORMAT,
|
||||
optional=True, tooltip=FORMAT_TOOLTIP),
|
||||
io.Int.Input("transfer_quality", default=DEFAULT_QUALITY, min=1, max=100,
|
||||
@@ -163,10 +170,12 @@ class MediaExportV1(io.ComfyNode):
|
||||
|
||||
@classmethod
|
||||
def execute(cls, request_id, binding_json, images=None, audio=None, masks=None, texts=None, video=None,
|
||||
transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY):
|
||||
transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY, **overrides):
|
||||
binding = _json(_scalar(binding_json))
|
||||
validate_binding(binding)
|
||||
provided = dict(images=images, audio=audio, masks=masks, texts=texts, video=video)
|
||||
if set(overrides) - set(VALUE_SOCKETS):
|
||||
raise ValueError("Unknown End value slot")
|
||||
provided = dict(images=images, audio=audio, masks=masks, texts=texts, video=video, **overrides)
|
||||
if {name for name, value in provided.items() if value is not None} != {binding["socket"]}:
|
||||
raise ValueError("Media export must receive exactly its declared typed socket")
|
||||
exported = export_values(_scalar(request_id), binding, provided[binding["socket"]],
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
"""Exact six-type values for internal sparse End exports; no pack imports.
|
||||
|
||||
Tensor/text payloads reuse the existing media codec. Scalars have distinct
|
||||
encodings: BOOLEAN is one byte, INT is minimal signed little-endian, FLOAT is
|
||||
IEEE binary64. No scalar passes through float() or a lossy JSON number parser.
|
||||
"""
|
||||
import math
|
||||
import struct
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
KINDS = frozenset(("IMAGE", "MASK", "STRING", "BOOLEAN", "INT", "FLOAT"))
|
||||
_SCALAR_ENCODINGS = {"BOOLEAN": "bool-v1", "INT": "int-le-v1", "FLOAT": "float64-le-v1"}
|
||||
|
||||
|
||||
def value_kind(value):
|
||||
if isinstance(value, torch.Tensor):
|
||||
if value.ndim == 4 and value.shape[-1] in (3, 4):
|
||||
return "IMAGE"
|
||||
if value.ndim == 3:
|
||||
return "MASK"
|
||||
raise ValueError("VALUE tensor requires IMAGE [B,H,W,3|4] or MASK [B,H,W]")
|
||||
kind = {bool: "BOOLEAN", int: "INT", float: "FLOAT", str: "STRING"}.get(type(value))
|
||||
if kind is None:
|
||||
raise ValueError("VALUE requires IMAGE, MASK, STRING, INT, FLOAT or BOOLEAN")
|
||||
if kind == "FLOAT" and not math.isfinite(value):
|
||||
raise ValueError("VALUE numbers must be finite")
|
||||
return kind
|
||||
|
||||
|
||||
def encode_value(value, *, max_bytes):
|
||||
from .codec import encode_payload
|
||||
kind = value_kind(value)
|
||||
if kind in ("IMAGE", "MASK", "STRING"):
|
||||
descriptor, payload = encode_payload(kind, value, max_bytes=max_bytes)
|
||||
else:
|
||||
descriptor = {"encoding": _SCALAR_ENCODINGS[kind]}
|
||||
size = ((value if value >= 0 else ~value).bit_length() + 8) // 8 if kind == "INT" else (8 if kind == "FLOAT" else 1)
|
||||
if size > max_bytes:
|
||||
raise ValueError("VALUE payload exceeds byte limit")
|
||||
if kind == "BOOLEAN":
|
||||
payload = b"\x01" if value else b"\x00"
|
||||
elif kind == "INT":
|
||||
payload = value.to_bytes(size, "little", signed=True)
|
||||
else:
|
||||
payload = struct.pack("<d", value)
|
||||
return {"encoding": "value-v1", "type": kind, "descriptor": descriptor}, payload
|
||||
|
||||
|
||||
def decode_value(descriptor, payload, *, max_bytes):
|
||||
from .codec import decode_payload
|
||||
if (set(descriptor) != {"encoding", "type", "descriptor"}
|
||||
or descriptor["encoding"] != "value-v1"
|
||||
or type(descriptor["type"]) is not str or descriptor["type"] not in KINDS):
|
||||
raise ValueError("Invalid VALUE descriptor")
|
||||
kind, inner = descriptor["type"], descriptor["descriptor"]
|
||||
if kind in ("IMAGE", "MASK", "STRING"):
|
||||
# Generic values always stay exact even if End.images is compressed.
|
||||
if not isinstance(inner, dict) or inner.get("encoding") != (
|
||||
"utf8-v1" if kind == "STRING" else "tensor-le-v1"):
|
||||
raise ValueError("VALUE requires exact typed payloads")
|
||||
value = decode_payload(kind, inner, payload, max_bytes=max_bytes)
|
||||
if value_kind(value) != kind:
|
||||
raise ValueError("VALUE tensor/type mismatch")
|
||||
return value
|
||||
if inner != {"encoding": _SCALAR_ENCODINGS[kind]}:
|
||||
raise ValueError("Invalid VALUE scalar encoding")
|
||||
if kind == "BOOLEAN":
|
||||
if payload not in (b"\x00", b"\x01"):
|
||||
raise ValueError("Invalid BOOLEAN payload")
|
||||
return payload == b"\x01"
|
||||
if kind == "FLOAT":
|
||||
if len(payload) != 8:
|
||||
raise ValueError("Invalid FLOAT payload length")
|
||||
value = struct.unpack("<d", payload)[0]
|
||||
if not math.isfinite(value):
|
||||
raise ValueError("VALUE numbers must be finite")
|
||||
return value
|
||||
if not payload:
|
||||
raise ValueError("Invalid INT payload length")
|
||||
value = int.from_bytes(payload, "little", signed=True)
|
||||
if ((value if value >= 0 else ~value).bit_length() + 8) // 8 != len(payload):
|
||||
raise ValueError("Invalid INT payload length")
|
||||
return value
|
||||
@@ -0,0 +1,60 @@
|
||||
"""CPU-only proof against real PromptQueue and IsChangedCache; no services."""
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
core = Path(sys.argv[1]).resolve()
|
||||
root = Path(__file__).resolve().parents[2]
|
||||
sys.path.insert(0, str(core))
|
||||
import comfy.cli_args
|
||||
comfy.cli_args.args.cpu = True
|
||||
import execution
|
||||
import nodes
|
||||
from comfy_execution.graph import DynamicPrompt
|
||||
sys.path.insert(0, str(root))
|
||||
from image_list.lifecycle import has_safe_execution_targets
|
||||
|
||||
|
||||
class Source:
|
||||
RETURN_TYPES = ('IMAGE',)
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'seed': ('INT',)}}
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, seed):
|
||||
return f'fingerprint-{seed}'
|
||||
|
||||
|
||||
nodes.NODE_CLASS_MAPPINGS['ProofSource'] = Source
|
||||
nodes.NODE_CLASS_MAPPINGS['ProofUtility'] = type('Utility', (), {'OUTPUT_NODE': True})
|
||||
prompt = {
|
||||
'source': {'class_type': 'ProofSource', 'inputs': {'seed': 7}},
|
||||
'start': {'class_type': 'DistributedImageListStart', 'inputs': {'images': ['source', 0]}},
|
||||
'op': {'class_type': 'ProofUtility', 'inputs': {'image': ['start', 0]}},
|
||||
'end': {'class_type': 'DistributedImageListEnd', 'inputs': {
|
||||
'images': ['op', 0], 'context': ['start', 3]}},
|
||||
}
|
||||
server = SimpleNamespace(queue_updated=lambda: None)
|
||||
queue = execution.PromptQueue(server)
|
||||
server.prompt_queue = queue
|
||||
queue.put((0, 'owned', prompt, {}, ['end']))
|
||||
item, _ = queue.get(timeout=0)
|
||||
assert item[2] is prompt
|
||||
snapshot = queue.get_current_queue_volatile()[0][0][2]
|
||||
assert snapshot is not prompt
|
||||
assert has_safe_execution_targets(server, 'owned', prompt)
|
||||
cache = execution.IsChangedCache('owned', DynamicPrompt(prompt), None)
|
||||
assert asyncio.run(cache.get('source')) == ['fingerprint-7']
|
||||
assert prompt['source']['is_changed'] == ['fingerprint-7']
|
||||
assert 'is_changed' not in snapshot['source']
|
||||
assert has_safe_execution_targets(server, 'owned', prompt), 'core fingerprint invalidated safe roots'
|
||||
prompt['source']['inputs']['seed'] = 8
|
||||
assert not has_safe_execution_targets(server, 'owned', prompt), 'changed input accepted'
|
||||
prompt['source']['inputs']['seed'] = 7
|
||||
assert not has_safe_execution_targets(server, 'not-owned', prompt)
|
||||
queue.get_current_queue_volatile()[0][0][4].append('op')
|
||||
assert not has_safe_execution_targets(server, 'owned', prompt), 'unsafe output root accepted'
|
||||
print('FINGERPRINT_PROOF_OK')
|
||||
@@ -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])
|
||||
@@ -37,14 +37,14 @@ import aiohttp
|
||||
|
||||
torch.set_num_threads(1)
|
||||
SOURCE_CALLS, BRANCH_CALLS, CAPTURED, MODEL_CALLS = [], [], [], []
|
||||
UTILITY_CALLS, ITEM_EXECUTIONS = [], []
|
||||
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)]
|
||||
@@ -92,7 +92,9 @@ def assert_item_execution(prompt, index):
|
||||
from comfy_execution.utils import get_executing_context
|
||||
context = get_executing_context()
|
||||
assert context is not None
|
||||
assert prompt['2']['inputs'].get('worker_item_token'), 'per-item node ran in outer graph'
|
||||
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]
|
||||
@@ -160,17 +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 ()
|
||||
|
||||
|
||||
@@ -179,11 +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,
|
||||
TestRequiredOutputUtility=TestRequiredOutputUtility,
|
||||
TestDiscardedPreview=TestDiscardedPreview)
|
||||
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.
|
||||
@@ -221,7 +246,10 @@ async def worker_main():
|
||||
accepted = json.loads(response.body)
|
||||
_, pending = ps.prompt_queue.get_current_queue_volatile()
|
||||
entry = next(item for item in pending if item[1] == accepted['prompt_id'])
|
||||
submissions.append({'prompt_id': entry[1], 'requested': data.get('partial_execution_targets'),
|
||||
delivery = next(node['inputs'] for node in entry[2].values()
|
||||
if node['class_type'] == 'DistributedImageListEnd')
|
||||
submissions.append({'prompt_id': entry[1], 'item_index': delivery['worker_item_index'],
|
||||
'requested': data.get('partial_execution_targets'),
|
||||
'queued': entry[4], 'node_ids': sorted(entry[2]),
|
||||
'end_ids': [key for key, node in entry[2].items()
|
||||
if node['class_type'] == 'DistributedImageListEnd']})
|
||||
@@ -271,21 +299,45 @@ 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)]}
|
||||
| {'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()
|
||||
@@ -403,6 +455,7 @@ async def coordinator():
|
||||
assert executor.success, executor.status_messages
|
||||
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()):
|
||||
@@ -430,15 +483,20 @@ async def coordinator():
|
||||
# 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), \
|
||||
@@ -451,6 +509,7 @@ async def coordinator():
|
||||
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
|
||||
@@ -458,7 +517,16 @@ 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)
|
||||
@@ -472,8 +540,8 @@ async def coordinator():
|
||||
patch.object(orchestration, 'queue_prompt_payload', enqueue), \
|
||||
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
|
||||
queued = await asyncio.wait_for(orchestration.orchestrate_distributed_execution(
|
||||
prompt, None, 'test', selected), 45)
|
||||
assert prompt == saved_prompt, (mode, 'saved/source graph changed')
|
||||
case_prompt, None, 'test', selected), 45)
|
||||
assert case_prompt == saved_case and prompt == saved_prompt, (mode, 'saved/source graph changed')
|
||||
assert executor.success, (mode, executor.status_messages)
|
||||
assert queued[2] == len(selected)
|
||||
assert len(SOURCE_CALLS) == before_sources + 1, 'source was reevaluated per local item'
|
||||
@@ -482,7 +550,7 @@ async def coordinator():
|
||||
assert len({event['prompt_id'] for event in ITEM_EXECUTIONS[before_events:]}) == len(local_calls)
|
||||
assert local_calls, f'{mode}: master did not process any items'
|
||||
assert len(MODEL_CALLS) == before_models + 1, 'local model dependency did not stay cached'
|
||||
remote_calls = []
|
||||
remote_assignments = []
|
||||
async with aiohttp.ClientSession() as session:
|
||||
for worker, before_stats in zip(workers, before_worker_stats):
|
||||
async with session.get(f"http://127.0.0.1:{worker['port']}/test-stats") as response:
|
||||
@@ -492,28 +560,69 @@ async def coordinator():
|
||||
submissions = current['submissions'][len(before_stats['submissions']):]
|
||||
events = current['item_executions'][len(before_stats['item_executions']):]
|
||||
assert Counter(utilities) == Counter(calls), (mode, current)
|
||||
assert len(submissions) == len(calls) == len({event['prompt_id'] for event in events})
|
||||
assigned = [submission['item_index'] for submission in submissions]
|
||||
# A repeated item can hit native processing cache even
|
||||
# when End changes return format. Count fresh deliveries,
|
||||
# not recomputation, for exact-once scheduling evidence.
|
||||
assert Counter(calls) <= Counter(assigned), (mode, worker['id'], calls, assigned)
|
||||
assert len(calls) == len({event['prompt_id'] for event in events}), (mode, events)
|
||||
assert {event['prompt_id'] for event in events} <= {
|
||||
submission['prompt_id'] for submission in submissions}
|
||||
for submission in submissions:
|
||||
assert submission['requested'] == submission['queued'] == ['4'], submission
|
||||
assert '8' in submission['node_ids'] and '9' not in submission['node_ids']
|
||||
assert '5' not in submission['node_ids']
|
||||
remote_calls.extend(calls)
|
||||
remote_assignments.extend(assigned)
|
||||
assert not current['running'] and not current['pending'], current
|
||||
if mode == 'mixed':
|
||||
assert remote_calls, 'mixed execution failed to use direct workers'
|
||||
if mode.startswith('mixed'):
|
||||
assert remote_assignments, 'mixed execution failed to use direct workers'
|
||||
else:
|
||||
assert not remote_calls and len(local_calls) == COUNT
|
||||
assert Counter(local_calls + remote_calls) == Counter(range(COUNT)), (local_calls, remote_calls)
|
||||
assert not remote_assignments and len(local_calls) == COUNT
|
||||
assert Counter(local_calls + remote_assignments) == Counter(range(COUNT)), (local_calls, remote_assignments)
|
||||
check_carried_values()
|
||||
images, masks, texts = CAPTURED[-1]
|
||||
assert len(images) == len(masks) == len(texts) == COUNT
|
||||
for i, row in enumerate(source_rows()):
|
||||
assert torch.equal(images[i], row['image'] + row['strength']), (mode, i)
|
||||
expected = row['image'] + row['strength']
|
||||
if fmt == 'png' and i in remote_assignments:
|
||||
expected = expected.clone()
|
||||
expected[..., :3] = (expected[..., :3].clamp(0, 1) * 255).to(torch.uint8).float() / 255
|
||||
assert torch.equal(images[i], expected), (mode, i)
|
||||
assert torch.equal(masks[i], row['mask'] * 0.5), (mode, i)
|
||||
assert texts[i] == f'item-{i}:seed-{100 + i}', (mode, texts)
|
||||
assert not runtime._MASTER_JOBS and not runtime._WORKER_ITEMS
|
||||
participation_counts[mode] = {'master': len(local_calls), 'workers': len(remote_calls)}
|
||||
participation_counts[mode] = {'master': len(local_calls), 'workers': len(remote_assignments)}
|
||||
# Repeat the completed master-only queue with a fresh outer executor.
|
||||
# Every item must be delivered again without recomputing its branch.
|
||||
before_sources, before_local = len(SOURCE_CALLS), len(BRANCH_CALLS)
|
||||
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) == 4 and all(row['captures'] == 1 for row in outer_executions)
|
||||
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'],
|
||||
@@ -527,6 +636,12 @@ async def coordinator():
|
||||
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():
|
||||
@@ -540,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', direct_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)
|
||||
@@ -562,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
|
||||
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())
|
||||
@@ -93,11 +93,12 @@ async def check(base):
|
||||
contract = load("serverless_bridge.contract")
|
||||
video_codec = load("serverless_bridge.video_codec")
|
||||
bridge = load("serverless_bridge")
|
||||
collector = load("nodes.media_collector").DistributedMediaCollector
|
||||
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]
|
||||
@@ -165,9 +166,10 @@ async def check(base):
|
||||
assert set(media) == set(expected)
|
||||
for socket in expected:
|
||||
assert len(media[socket]) == len(expected[socket]), (tag, socket)
|
||||
# Collector must preserve native objects and ordered item lists,
|
||||
# not merge batches, resample, convert types, or encode again.
|
||||
assert all(a is e for a, e in zip(media[socket], 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]})
|
||||
@@ -180,6 +182,19 @@ async def check(base):
|
||||
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:
|
||||
@@ -199,7 +214,7 @@ async def check(base):
|
||||
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 + variant}
|
||||
"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),
|
||||
@@ -220,7 +235,7 @@ async def check(base):
|
||||
with wave.open(str(wav), "wb") as out:
|
||||
out.setnchannels(2)
|
||||
out.setsampwidth(2)
|
||||
out.setframerate(22050)
|
||||
out.setframerate(44100)
|
||||
out.writeframes(pcm.tobytes())
|
||||
waveform, sample_rate = nodes_audio.load(str(wav))
|
||||
clip = native_root / "native.mkv"
|
||||
@@ -278,17 +293,22 @@ async def check(base):
|
||||
participant = "serverless:" + key
|
||||
request_id = "acceptance-" + key
|
||||
graph = {"source": node("AcceptanceMediaSource", key=key),
|
||||
"collector": node("DistributedMediaCollector", load_balance=False,
|
||||
"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(sockets)}),
|
||||
"text-collector": node("DistributedMediaCollector", load_balance=False,
|
||||
multi_job_id="job-text", worker_id=participant, texts=["source", 3])}
|
||||
**{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"].update(audio=["wav", 0], video=["clip", 0])
|
||||
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)
|
||||
@@ -335,17 +355,19 @@ async def check(base):
|
||||
# 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-text"}
|
||||
assert set(grouped["job-main"]) == set(sockets)
|
||||
assert grouped["job-text"]["texts"] == expectations[key]["texts"]
|
||||
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 = grouped["job-main"][socket], expectations[key][socket]
|
||||
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-main"]["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()
|
||||
@@ -363,7 +385,7 @@ async def check(base):
|
||||
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 ("png", "jpeg", "h264"):
|
||||
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.
|
||||
@@ -372,12 +394,23 @@ async def check(base):
|
||||
fixtures[key]["images"] = [rgba]
|
||||
request_id = "request-" + key
|
||||
graph = {"source": node("AcceptanceMediaSource", key=key),
|
||||
"collector": node("DistributedMediaCollector", transfer_format=fmt,
|
||||
"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(sockets)})}
|
||||
**{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))
|
||||
@@ -385,23 +418,32 @@ async def check(base):
|
||||
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]
|
||||
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 == "png":
|
||||
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 png jpeg h264; alpha and other media exact")
|
||||
print("REAL_HTTP_BRIDGE_RESULTS_OK all five kinds; 6 bindings/request; native VIDEO remains readable after cleanup")
|
||||
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
|
||||
@@ -410,39 +452,40 @@ async def check(base):
|
||||
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] = {socket: [item for values in ordered for item in values[socket]] for socket in sockets}
|
||||
graph = {"collector": node("DistributedMediaCollector", load_balance=False, multi_job_id=job,
|
||||
enabled_worker_ids=json.dumps(participants), expected_media=json.dumps(sockets), delegate_only=delegate),
|
||||
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(sockets)})}
|
||||
**{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(sockets)})
|
||||
# The second group from the very same complete response must reach its
|
||||
# own Collector, rather than being folded into the first collector job.
|
||||
text_job, text_tag = job + "-text", tag + "-text"
|
||||
text_queue = asyncio.Queue()
|
||||
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 text_queue.put({"worker_id": participant,
|
||||
"media": grouped_results[participant]["job-text"], "is_last": True})
|
||||
ps.distributed_pending_jobs[text_job] = text_queue
|
||||
text_ordered = ([] if delegate else [fixtures["local"]]) + [grouped_results[p]["job-text"] for p in participants]
|
||||
expected_sinks[text_tag] = {"texts": [item for values in text_ordered for item in values["texts"]]}
|
||||
graph["text-collector"] = node("DistributedMediaCollector", load_balance=False, multi_job_id=text_job,
|
||||
enabled_worker_ids=json.dumps(participants), expected_media='["texts"]', delegate_only=delegate)
|
||||
graph["text-assert"] = node("AcceptanceMediaAssert", tag=text_tag, texts=["text-collector", 3])
|
||||
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["text-collector"]["inputs"]["texts"] = ["source", 3]
|
||||
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"]["text-assert"]["accepted"] == [seen_sinks[text_tag]]
|
||||
assert job not in ps.distributed_pending_jobs and text_job not in ps.distributed_pending_jobs
|
||||
assert source_calls == ["worker-a", "worker-b", "encoded-png", "encoded-jpeg", "encoded-h264", "local"], source_calls
|
||||
assert set(seen_sinks) == {"master", "delegate", "master-text", "delegate-text"}
|
||||
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))
|
||||
|
||||
@@ -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', 'DistributedMediaCollector'}
|
||||
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,14 +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)
|
||||
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]
|
||||
media_info = mapping['DistributedMediaCollector'].GET_NODE_INFO_V1()
|
||||
assert media_info['output'] == ['IMAGE', 'AUDIO', 'MASK', 'STRING', 'VIDEO'], media_info
|
||||
assert all(media_info['output_is_list']), media_info
|
||||
assert media_info['input']['optional']['texts'][1]['forceInput'] is True
|
||||
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():
|
||||
@@ -427,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
|
||||
@@ -470,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')
|
||||
|
||||
@@ -0,0 +1,552 @@
|
||||
"""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:
|
||||
# Channel mixing is supported; incompatible spatial sizes still fail.
|
||||
await send(ctx, frames(count=1, channels=4, width=10))
|
||||
with pytest.raises(RuntimeError, match="spatial dimensions"):
|
||||
await ctx.node.execute(
|
||||
frames(count=1, channels=3), None, multi_job_id=JOB_ID,
|
||||
enabled_worker_ids=json.dumps([WORKER_ID]), transfer_format="exact",
|
||||
)
|
||||
assert JOB_ID not in transport.server.distributed_pending_jobs
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("local_channels,worker_channels", [(3, 4), (4, 3)])
|
||||
def test_exact_collector_mixed_rgb_rgba_preserves_every_value(transport, monkeypatch, local_channels, worker_channels):
|
||||
local = frames(count=1, channels=local_channels)
|
||||
worker = frames(count=1, channels=worker_channels)
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch) as ctx:
|
||||
await send(ctx, worker)
|
||||
result, _audio = await ctx.node.execute(
|
||||
local, None, multi_job_id=JOB_ID,
|
||||
enabled_worker_ids=json.dumps([WORKER_ID]), transfer_format="exact",
|
||||
)
|
||||
assert result.shape == (2, 7, 9, 4)
|
||||
assert result.is_contiguous()
|
||||
# RGB views into RGBA have a four-channel stride; materialize the
|
||||
# sliced comparison without changing the output-contiguity check.
|
||||
assert_exact(result[:1, ..., :local_channels].contiguous(), local)
|
||||
assert_exact(result[1:, ..., :worker_channels].contiguous(), worker)
|
||||
if local_channels == 3:
|
||||
assert torch.all(result[:1, ..., 3] == 1)
|
||||
if worker_channels == 3:
|
||||
assert torch.all(result[1:, ..., 3] == 1)
|
||||
assert JOB_ID not in transport.server.distributed_pending_jobs
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_exact_incomplete_worker_raises_and_cleans_queue(transport, monkeypatch):
|
||||
"""A timed-out worker's first frame is not a successful partial result."""
|
||||
envelope, parts = exact_chunk(transport, count=1)
|
||||
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())
|
||||
@@ -0,0 +1,225 @@
|
||||
"""Real PNG/tensor regressions; only server/config/HTTP plumbing is stubbed."""
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import io
|
||||
import base64
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from aiohttp import web
|
||||
from PIL import Image
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def _load_file(name, path):
|
||||
spec = importlib.util.spec_from_file_location(name, path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def media_modules():
|
||||
# Existing harnesses provide ComfyUI/config doubles. Restore their module
|
||||
# replacements so this file can coexist with the rest of the test suite.
|
||||
with patch.dict(sys.modules):
|
||||
image_utils = _load_file("real_image_utils", ROOT / "utils/image.py")
|
||||
collector_tests = _load_file("collector_harness", ROOT / "tests/test_collector_list_inputs.py")
|
||||
collector = collector_tests._load_collector_module()
|
||||
sys.modules["aiohttp"].web = web
|
||||
route_tests = _load_file("route_harness", ROOT / "tests/api/test_distributed_queue.py")
|
||||
routes = route_tests._load_job_routes_module()
|
||||
for module in (collector, routes):
|
||||
module.pil_to_tensor = image_utils.pil_to_tensor
|
||||
module.ensure_contiguous = image_utils.ensure_contiguous
|
||||
collector.tensor_to_pil = image_utils.tensor_to_pil
|
||||
routes.prompt_server = collector.prompt_server
|
||||
yield collector, routes
|
||||
|
||||
|
||||
def _png_payload(image):
|
||||
buffer = io.BytesIO()
|
||||
image.save(buffer, format="PNG")
|
||||
return "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode("ascii")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["RGB", "RGBA", "L", "LA", "palette", "palette_alpha"])
|
||||
def test_png_decoder_preserves_alpha_and_normalizes_opaque_modes(media_modules, mode):
|
||||
_, routes = media_modules
|
||||
if mode.startswith("palette"):
|
||||
image = Image.new("P", (12, 8), 0)
|
||||
image.putpalette([16, 32, 64] + [0] * 765)
|
||||
if mode == "palette_alpha":
|
||||
image.info["transparency"] = 0
|
||||
else:
|
||||
color = {"RGB": (16, 32, 64), "RGBA": (16, 32, 64, 96), "L": 32, "LA": (32, 96)}[mode]
|
||||
image = Image.new(mode, (12, 8), color)
|
||||
alpha = mode in {"RGBA", "LA", "palette_alpha"}
|
||||
expected = np.array(image.convert("RGBA" if alpha else "RGB"), dtype=np.float32) / 255
|
||||
decoded = routes._decode_canonical_png_tensor(_png_payload(image))
|
||||
assert decoded.shape == (1, 8, 12, 4 if alpha else 3)
|
||||
assert decoded.is_contiguous()
|
||||
np.testing.assert_array_equal(decoded[0].numpy(), expected)
|
||||
|
||||
|
||||
async def _master_collect(module, local, worker_images, delegate=False):
|
||||
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
queue = asyncio.Queue()
|
||||
for index, tensor in enumerate(worker_images):
|
||||
queue.put_nowait({"worker_id": "worker-a", "image_index": index,
|
||||
"tensor": tensor, "is_last": index == len(worker_images) - 1})
|
||||
module.prompt_server.distributed_pending_jobs = {"image-test": queue}
|
||||
return await module.DistributedCollectorNode().execute(
|
||||
local, None, multi_job_id="image-test", enabled_worker_ids='["worker-a"]',
|
||||
delegate_only=delegate,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("local_channels,worker_channels", [(3, 3), (4, 4), (3, 4), (4, 3), (None, 4)])
|
||||
def test_worker_png_callback_and_collector_preserve_pixels_and_alpha(
|
||||
media_modules, local_channels, worker_channels
|
||||
):
|
||||
module, routes = media_modules
|
||||
local = None if local_channels is None else torch.full((1, 8, 12, local_channels), 0.5)
|
||||
worker = torch.full((2, 8, 12, worker_channels), 0.25)
|
||||
if worker_channels == 4:
|
||||
worker[..., 3] = torch.linspace(0, 1, 12).reshape(1, 1, 12)
|
||||
responses = []
|
||||
|
||||
async def scenario():
|
||||
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
module.prompt_server.distributed_pending_jobs = {"image-test": asyncio.Queue()}
|
||||
|
||||
class Response:
|
||||
def __init__(self, payload):
|
||||
self.payload = payload
|
||||
|
||||
async def __aenter__(self):
|
||||
async def request_json():
|
||||
return self.payload
|
||||
response = await routes.job_complete_endpoint(types.SimpleNamespace(json=request_json))
|
||||
responses.append(response.status)
|
||||
assert response.status == 200, response.text
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *_args):
|
||||
return False
|
||||
|
||||
def raise_for_status(self):
|
||||
pass
|
||||
|
||||
class Session:
|
||||
def post(self, _url, json, timeout):
|
||||
return Response(json)
|
||||
|
||||
async def get_session():
|
||||
return Session()
|
||||
|
||||
module.get_client_session = get_session
|
||||
module.encode_audio_payload = lambda _audio: None
|
||||
collector = module.DistributedCollectorNode()
|
||||
await collector.send_batch_to_master(worker, None, "image-test", "http://master", "worker-a")
|
||||
result, _ = await collector.execute(
|
||||
local, None, multi_job_id="image-test", enabled_worker_ids='["worker-a"]',
|
||||
delegate_only=local is None,
|
||||
)
|
||||
assert "image-test" not in module.prompt_server.distributed_pending_jobs
|
||||
return result
|
||||
|
||||
result = asyncio.run(scenario())
|
||||
channels = max(local_channels or 3, worker_channels)
|
||||
offset = 0 if local is None else 1
|
||||
assert responses == [200, 200]
|
||||
assert result.shape == (offset + 2, 8, 12, channels)
|
||||
assert result.is_contiguous()
|
||||
if local is not None:
|
||||
assert torch.equal(result[:1, ..., :local_channels], local)
|
||||
if local_channels == 3 and channels == 4:
|
||||
assert torch.all(result[:1, ..., 3] == 1)
|
||||
assert torch.allclose(result[offset:, ..., :worker_channels], worker, atol=1 / 255)
|
||||
if worker_channels == 3 and channels == 4:
|
||||
assert torch.all(result[offset:, ..., 3] == 1)
|
||||
|
||||
|
||||
def test_collector_mixed_rgb_rgba_list_promotes_rgb_to_opaque_rgba(media_modules):
|
||||
module, _ = media_modules
|
||||
rgb = torch.zeros(1, 8, 12, 3)
|
||||
rgba = torch.full((1, 8, 12, 4), 0.5)
|
||||
result, _ = module.DistributedCollectorNode().run(images=[rgb, rgba])
|
||||
assert result.shape == (2, 8, 12, 4)
|
||||
assert torch.equal(result[:1, ..., :3], rgb)
|
||||
assert torch.all(result[:1, ..., 3] == 1)
|
||||
assert torch.equal(result[1:], rgba)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("other_channels", [1, 3, 4])
|
||||
def test_collector_preserves_grayscale_or_promotes_it_for_color_batches(media_modules, other_channels):
|
||||
module, _ = media_modules
|
||||
gray = torch.full((1, 8, 12, 1), 0.25)
|
||||
other = torch.full((1, 8, 12, other_channels), 0.5)
|
||||
result, _ = module.DistributedCollectorNode().run(images=[gray, other])
|
||||
assert result.shape == (2, 8, 12, other_channels)
|
||||
assert torch.all(result[:1, ..., :min(other_channels, 3)] == 0.25)
|
||||
if other_channels == 4:
|
||||
assert torch.all(result[:1, ..., 3] == 1)
|
||||
assert torch.equal(result[1:], other)
|
||||
|
||||
|
||||
def test_mixed_channel_merge_keeps_configured_worker_and_image_order(media_modules):
|
||||
module, _ = media_modules
|
||||
collector = module.DistributedCollectorNode()
|
||||
master = torch.zeros(1, 8, 12, 3)
|
||||
worker_a = torch.full((1, 8, 12, 4), 0.25)
|
||||
worker_b = torch.full((1, 8, 12, 3), 0.5)
|
||||
unexpected = torch.full((1, 8, 12, 3), 0.75)
|
||||
result = collector._reorder_and_combine_tensors(
|
||||
{"worker-a": {0: worker_a}, "worker-b": {2: worker_b, 0: master},
|
||||
"unexpected": {0: unexpected}},
|
||||
["worker-b", "worker-a"], 1, master, False, master,
|
||||
)
|
||||
assert result.shape == (5, 8, 12, 4)
|
||||
assert result[:, 0, 0, 0].tolist() == [0, 0, 0.5, 0.25, 0.75]
|
||||
assert result[:, 0, 0, 3].tolist() == [1, 1, 1, 0.25, 1]
|
||||
|
||||
|
||||
def test_collector_spatial_mismatch_raises_with_shapes_not_successful_fallback(media_modules):
|
||||
module, _ = media_modules
|
||||
local = torch.zeros(1, 8, 12, 3)
|
||||
worker = torch.zeros(1, 8, 13, 3)
|
||||
with pytest.raises(ValueError, match="spatial dimensions") as error:
|
||||
asyncio.run(_master_collect(module, local, [worker]))
|
||||
message = str(error.value)
|
||||
assert "master image 0" in message
|
||||
assert "worker 'worker-a' image 0" in message
|
||||
assert "(1, 8, 12, 3)" in message and "(1, 8, 13, 3)" in message
|
||||
assert "image-test" not in module.prompt_server.distributed_pending_jobs
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", [(1, 8, 12), (1, 8, 12, 5)])
|
||||
def test_collector_rejects_invalid_worker_image_shape(media_modules, shape):
|
||||
module, _ = media_modules
|
||||
with pytest.raises(ValueError, match="worker 'worker-a' image 0"):
|
||||
asyncio.run(_master_collect(module, None, [torch.zeros(shape)], delegate=True))
|
||||
|
||||
|
||||
def test_delegate_master_ignores_empty_placeholder_shape(media_modules):
|
||||
module, _ = media_modules
|
||||
placeholder = torch.empty(0, 64, 64, 3)
|
||||
worker = torch.full((1, 8, 12, 4), 0.5)
|
||||
result, _ = asyncio.run(_master_collect(module, placeholder, [worker], delegate=True))
|
||||
assert torch.equal(result, worker)
|
||||
|
||||
|
||||
def test_empty_image_results_raise_instead_of_reaching_save_image(media_modules):
|
||||
module, _ = media_modules
|
||||
with pytest.raises(ValueError, match="empty image batch"):
|
||||
asyncio.run(_master_collect(module, torch.empty(0, 64, 64, 3),
|
||||
[torch.empty(0, 8, 12, 3)], delegate=True))
|
||||
@@ -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")
|
||||
@@ -472,6 +472,25 @@ class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
|
||||
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)})
|
||||
|
||||
@@ -0,0 +1,81 @@
|
||||
"""Execution-target proof must survive core cache annotations, not graph edits."""
|
||||
import copy
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from image_list.lifecycle import has_safe_execution_targets
|
||||
from test_image_list_master_runtime import graph
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def proof(monkeypatch):
|
||||
prompt = graph()
|
||||
queued = copy.deepcopy(prompt)
|
||||
entry = [0, 'owned', queued, {}, ['end']]
|
||||
server = SimpleNamespace(prompt_queue=SimpleNamespace(
|
||||
get_current_queue_volatile=lambda: ([entry], [])))
|
||||
monkeypatch.setitem(sys.modules, 'nodes', SimpleNamespace(NODE_CLASS_MAPPINGS={
|
||||
'Operation': type('Flagged', (), {'OUTPUT_NODE': True})}))
|
||||
return prompt, queued, entry, server
|
||||
|
||||
|
||||
@pytest.mark.parametrize('value', [['digest'], [None], float('nan'), {'nested': [1, 2]}])
|
||||
@pytest.mark.parametrize('side', ['executing', 'queued', 'both'])
|
||||
def test_top_level_fingerprints_do_not_invalidate_proof(proof, value, side):
|
||||
prompt, queued, _, server = proof
|
||||
if side in ('executing', 'both'):
|
||||
prompt['source']['is_changed'] = value
|
||||
if side in ('queued', 'both'):
|
||||
queued['source']['is_changed'] = ['different queued fingerprint']
|
||||
assert has_safe_execution_targets(server, 'owned', prompt)
|
||||
# Comparisons must not strip or rewrite the caller's data.
|
||||
if side in ('executing', 'both'):
|
||||
assert prompt['source']['is_changed'] is value
|
||||
if side in ('queued', 'both'):
|
||||
assert queued['source']['is_changed'] == ['different queued fingerprint']
|
||||
|
||||
|
||||
@pytest.mark.parametrize('change', ['input', 'link', 'class', 'node_id', 'nested_is_changed', 'other_metadata'])
|
||||
def test_fingerprint_tolerance_does_not_hide_graph_changes(proof, change):
|
||||
prompt, _, _, server = proof
|
||||
prompt['source']['is_changed'] = ['digest']
|
||||
if change == 'input':
|
||||
prompt['source']['inputs']['seed'] = 2
|
||||
elif change == 'link':
|
||||
prompt['op']['inputs']['image'] = ['start', 1]
|
||||
elif change == 'class':
|
||||
prompt['source']['class_type'] = 'OtherSource'
|
||||
elif change == 'node_id':
|
||||
prompt['new_source'] = prompt.pop('source')
|
||||
prompt['start']['inputs']['images'] = ['new_source', 0]
|
||||
elif change == 'nested_is_changed':
|
||||
prompt['source']['inputs']['is_changed'] = True
|
||||
else:
|
||||
prompt['source']['_meta'] = {'title': 'changed'}
|
||||
assert not has_safe_execution_targets(server, 'owned', prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('roots', [['op', 'end'], [], ['end', 'end'], None, [True]])
|
||||
def test_fingerprints_never_grant_unscoped_or_invalid_roots(proof, roots):
|
||||
prompt, _, entry, server = proof
|
||||
prompt['source']['is_changed'] = ['digest']
|
||||
entry[4] = roots
|
||||
assert not has_safe_execution_targets(server, 'owned', prompt)
|
||||
|
||||
|
||||
def test_real_core_queue_snapshot_and_fingerprint():
|
||||
core = Path(os.environ.get('COMFYUI_SOURCE_ROOT', str(Path(sys.executable).parents[2])))
|
||||
if not (core / 'execution.py').exists():
|
||||
pytest.skip('real ComfyUI source not available beside this interpreter')
|
||||
result = subprocess.run(
|
||||
[sys.executable, str(ROOT / 'tests/helpers/image_list_fingerprint_proof.py'), str(core)],
|
||||
text=True, capture_output=True, timeout=90)
|
||||
assert result.returncode == 0, result.stdout + result.stderr
|
||||
assert 'FINGERPRINT_PROOF_OK' in result.stdout
|
||||
@@ -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
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -62,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])
|
||||
|
||||
@@ -97,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
|
||||
@@ -117,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
|
||||
@@ -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()
|
||||
@@ -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']
|
||||
@@ -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()
|
||||
@@ -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())
|
||||
@@ -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())
|
||||
@@ -94,7 +94,7 @@ async def check():
|
||||
assert await extension.get_node_list() == [MediaImportV1, MediaExportV1]
|
||||
imp = MediaImportV1.GET_NODE_INFO_V1()
|
||||
exp = MediaExportV1.GET_NODE_INFO_V1()
|
||||
assert imp['output'] == ['IMAGE','AUDIO','MASK','STRING','VIDEO'], imp
|
||||
assert imp['output'] == ['IMAGE','AUDIO','MASK','STRING','VIDEO','*'], imp
|
||||
assert exp['output_node'] and MediaExportV1.define_schema().is_input_list, exp
|
||||
with tempfile.TemporaryDirectory(prefix='typed-media-', dir=os.environ.get('TMPDIR')) as scratch:
|
||||
base = Path(scratch)
|
||||
@@ -152,7 +152,7 @@ async def check():
|
||||
mapped, missing, hidden = execution.get_input_data({'artifact_json':json.dumps(artifact),'filename':filename},MediaImportV1,'import')
|
||||
got = await execution._async_map_node_over_list('probe','import',MediaImportV1,mapped,MediaImportV1.FUNCTION,v3_data=hidden)
|
||||
got = await execution.resolve_map_node_over_list_results(got)
|
||||
assert len(got[0].result) == 5
|
||||
assert len(got[0].result) == 6
|
||||
assert all(x is None for n,x in enumerate(got[0].result) if n != list(SOCKET_KINDS).index(socket))
|
||||
bad = dict(artifact,sha256='0'*64)
|
||||
try: import_value(bad,filename)
|
||||
|
||||
@@ -1,314 +0,0 @@
|
||||
"""CPU-only typed collection and optional native V3 mapping acceptance."""
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
SOCKETS = ('images', 'audio', 'masks', 'texts', 'video')
|
||||
WORKERS = ['serverless:pool:1', 'serverless:pool:2']
|
||||
|
||||
|
||||
class Manager:
|
||||
def __init__(self, *, pending=False, owned=True):
|
||||
self.pending, self.owned = pending, owned
|
||||
self.cancelled = []
|
||||
|
||||
def owns(self, job, worker):
|
||||
return self.owned
|
||||
|
||||
def is_pending(self, job, worker):
|
||||
return self.pending
|
||||
|
||||
async def cancel(self, **kwargs):
|
||||
self.cancelled.append(kwargs)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def module(monkeypatch):
|
||||
package = '_media_collector_test'
|
||||
for name, path in ((package, ROOT), (package + '.nodes', ROOT / 'nodes'),
|
||||
(package + '.utils', ROOT / 'utils'),
|
||||
(package + '.serverless_bridge', ROOT / 'serverless_bridge')):
|
||||
stub = types.ModuleType(name)
|
||||
stub.__path__ = [str(path)]
|
||||
monkeypatch.setitem(sys.modules, name, stub)
|
||||
io = types.SimpleNamespace(ComfyNode=object, NodeOutput=lambda *items: types.SimpleNamespace(result=items))
|
||||
api = types.ModuleType('comfy_api.v0_0_2')
|
||||
api.io = io
|
||||
monkeypatch.setitem(sys.modules, 'comfy_api', types.ModuleType('comfy_api'))
|
||||
monkeypatch.setitem(sys.modules, 'comfy_api.v0_0_2', api)
|
||||
management = types.ModuleType('comfy.model_management')
|
||||
management.throw_exception_if_processing_interrupted = lambda: None
|
||||
comfy = types.ModuleType('comfy')
|
||||
comfy.model_management = management
|
||||
monkeypatch.setitem(sys.modules, 'comfy', comfy)
|
||||
monkeypatch.setitem(sys.modules, 'comfy.model_management', management)
|
||||
ps = types.SimpleNamespace(distributed_pending_jobs={}, distributed_jobs_lock=asyncio.Lock())
|
||||
server = types.ModuleType('server')
|
||||
server.PromptServer = types.SimpleNamespace(instance=ps)
|
||||
monkeypatch.setitem(sys.modules, 'server', server)
|
||||
helper = types.ModuleType(package + '.utils.async_helpers')
|
||||
helper.run_async_in_server_loop = asyncio.run
|
||||
monkeypatch.setitem(sys.modules, helper.__name__, helper)
|
||||
spec = importlib.util.spec_from_file_location(package + '.nodes.media_collector', ROOT / 'nodes/media_collector.py')
|
||||
loaded = importlib.util.module_from_spec(spec)
|
||||
monkeypatch.setitem(sys.modules, spec.name, loaded)
|
||||
spec.loader.exec_module(loaded)
|
||||
monkeypatch.setattr(loaded, 'POLL_SECONDS', 0.005)
|
||||
monkeypatch.setattr(loaded, 'START_GRACE_SECONDS', 0.025)
|
||||
loaded.test_server = ps
|
||||
loaded.test_management = management
|
||||
return loaded
|
||||
|
||||
|
||||
def native_media(offset=0):
|
||||
return {
|
||||
'images': [torch.full((1, 2 + offset, 3, 4), offset, dtype=torch.float64),
|
||||
torch.zeros((2, 1, 1, 3), dtype=torch.float16)],
|
||||
'audio': [{'waveform': torch.zeros((1, 1 + offset, 8), dtype=torch.float64),
|
||||
'sample_rate': 44100 + offset}],
|
||||
'masks': [torch.zeros((1, 2 + offset, 3), dtype=torch.float64)],
|
||||
'texts': ['', ' café\n🎥\r\n'],
|
||||
'video': [object()],
|
||||
}
|
||||
|
||||
|
||||
def setup(module, messages=(), *, manager=None):
|
||||
queue = asyncio.Queue()
|
||||
for message in messages:
|
||||
queue.put_nowait(message)
|
||||
ps = module.test_server
|
||||
ps.distributed_pending_jobs = {'job': queue, 'unrelated': asyncio.Queue()}
|
||||
ps.distributed_serverless = manager if manager is not None else Manager()
|
||||
return ps, queue
|
||||
|
||||
|
||||
def message(worker, media):
|
||||
return {'worker_id': worker, 'media': media, 'is_last': True}
|
||||
|
||||
|
||||
def collect(module, **kwargs):
|
||||
return module.DistributedMediaCollector.execute(**kwargs).result
|
||||
|
||||
|
||||
def remote(module, **kwargs):
|
||||
controls = dict(multi_job_id=['job'], enabled_worker_ids=[json.dumps(WORKERS)],
|
||||
expected_media=[json.dumps(['texts'])], delegate_only=[True])
|
||||
controls.update(kwargs)
|
||||
return collect(module, **controls)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('controls', [{}, {'pass_through': [True]}, {'enabled_worker_ids': ['[]']}])
|
||||
def test_all_five_passthrough_preserves_lists_and_native_values(module, controls):
|
||||
values = native_media()
|
||||
result = collect(module, **values, **controls)
|
||||
assert len(result) == 5
|
||||
for index, socket in enumerate(SOCKETS):
|
||||
assert result[index] is values[socket]
|
||||
|
||||
|
||||
def test_reversed_arrivals_master_then_workers_then_items(module):
|
||||
master, first, second = native_media(), native_media(1), native_media(2)
|
||||
ps, _ = setup(module, [message(WORKERS[1], second), message(WORKERS[0], first)])
|
||||
result = remote(module, **master, delegate_only=[False], expected_media=[json.dumps(list(SOCKETS))])
|
||||
for index, socket in enumerate(SOCKETS):
|
||||
expected = master[socket] + first[socket] + second[socket]
|
||||
assert len(result[index]) == len(expected)
|
||||
assert all(actual is original for actual, original in zip(result[index], expected))
|
||||
assert list(ps.distributed_pending_jobs) == ['unrelated']
|
||||
assert ps.distributed_serverless.cancelled == []
|
||||
|
||||
|
||||
def test_delegate_only_and_empty_string(module):
|
||||
ps, _ = setup(module, [message(WORKERS[1], {'texts': ['b']}), message(WORKERS[0], {'texts': ['']})])
|
||||
result = remote(module)
|
||||
assert result == ([], [], [], ['', 'b'], [])
|
||||
assert 'job' not in ps.distributed_pending_jobs
|
||||
|
||||
|
||||
@pytest.mark.parametrize('kwargs', [{}, {'texts': []}, {'texts': ''}, {'texts': ['x'], 'is_worker': [True]},
|
||||
{'texts': ['x'], 'enabled_worker_ids': ['["direct-worker"]']},
|
||||
{'texts': ['x'], 'expected_media': ['["images"]']},
|
||||
{'texts': ['x'], 'pass_through': [True, False]}, {'delegate_only': True},
|
||||
{'texts': ['x'], 'enabled_worker_ids': ['["serverless:pool:1","serverless:pool:1"]']},
|
||||
{'texts': ['x'], 'load_balance': ['false']}])
|
||||
def test_rejects_missing_media_or_invalid_controls(module, kwargs):
|
||||
with pytest.raises((ValueError, RuntimeError)):
|
||||
collect(module, **kwargs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('bad', [
|
||||
{'worker_id': 'serverless:other:1', 'media': {'texts': ['x']}, 'is_last': True},
|
||||
{'worker_id': WORKERS[0], 'media': {'texts': ['x']}, 'is_last': False},
|
||||
{'worker_id': WORKERS[0], 'media': {'images': [object()]}, 'is_last': True},
|
||||
{'worker_id': WORKERS[0], 'media': {'texts': []}, 'is_last': True},
|
||||
{'worker_id': WORKERS[0], 'media': {'texts': 'x'}, 'is_last': True},
|
||||
{'worker_id': WORKERS[0], 'media': {'texts': [None]}, 'is_last': True},
|
||||
{'worker_id': WORKERS[0], 'media': {'texts': ['x'] * 65}, 'is_last': True},
|
||||
{'worker_id': WORKERS[0], 'media': {'texts': ['x'], 'unknown': ['x']}, 'is_last': True},
|
||||
{'worker_id': WORKERS[0], 'serverless_error': 'sanitized failure', 'is_last': True},
|
||||
])
|
||||
def test_rejects_bad_delivery_without_partial_output_and_cleans(module, bad):
|
||||
ps, _ = setup(module, [bad])
|
||||
with pytest.raises((ValueError, RuntimeError)):
|
||||
remote(module)
|
||||
assert list(ps.distributed_pending_jobs) == ['unrelated']
|
||||
assert ps.distributed_serverless.cancelled == [{'collector_job_id': 'job'}]
|
||||
|
||||
|
||||
@pytest.mark.parametrize('after_complete', [False, True])
|
||||
def test_duplicate_delivery_rejected_including_final_queue_drain(module, after_complete):
|
||||
first = message(WORKERS[0], {'texts': ['a']})
|
||||
second = message(WORKERS[1], {'texts': ['b']})
|
||||
messages = [first, second, first] if after_complete else [first, first, second]
|
||||
ps, _ = setup(module, messages)
|
||||
with pytest.raises(RuntimeError, match='[Dd]uplicate'):
|
||||
remote(module)
|
||||
assert 'job' not in ps.distributed_pending_jobs
|
||||
assert ps.distributed_serverless.cancelled == [{'collector_job_id': 'job'}]
|
||||
|
||||
|
||||
@pytest.mark.parametrize('manager', [None, Manager(owned=False), Manager(pending=False)])
|
||||
def test_missing_manager_or_finished_participant_fails_bounded(module, manager):
|
||||
ps, _ = setup(module)
|
||||
ps.distributed_serverless = manager
|
||||
with pytest.raises(RuntimeError, match='[Ss]erverless'):
|
||||
remote(module)
|
||||
assert 'job' not in ps.distributed_pending_jobs
|
||||
|
||||
|
||||
def test_control_error_also_cleans_exact_job(module):
|
||||
ps, _ = setup(module)
|
||||
with pytest.raises(ValueError):
|
||||
remote(module, expected_media='["unknown"]')
|
||||
assert list(ps.distributed_pending_jobs) == ['unrelated']
|
||||
assert ps.distributed_serverless.cancelled == [{'collector_job_id': 'job'}]
|
||||
|
||||
|
||||
def test_launch_race_and_pending_deadline_owned_by_manager(module, monkeypatch):
|
||||
ps, queue = setup(module)
|
||||
ps.distributed_serverless = None
|
||||
async def runner(coro):
|
||||
async def produce():
|
||||
await asyncio.sleep(0.01)
|
||||
ps.distributed_serverless = Manager(pending=True)
|
||||
await asyncio.sleep(0.035)
|
||||
queue.put_nowait(message(WORKERS[1], {'texts': ['b']}))
|
||||
queue.put_nowait(message(WORKERS[0], {'texts': ['a']}))
|
||||
task = asyncio.create_task(produce())
|
||||
result = await coro
|
||||
await task
|
||||
return result
|
||||
monkeypatch.setattr(module, 'run_async_in_server_loop', lambda coro: asyncio.run(runner(coro)))
|
||||
assert remote(module)[3] == ['a', 'b']
|
||||
assert ps.distributed_serverless.cancelled == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize('cancel_task', [False, True])
|
||||
def test_cancellation_while_blocked_cleans_and_propagates(module, monkeypatch, cancel_task):
|
||||
ps, _ = setup(module, manager=Manager(pending=True))
|
||||
class Interrupted(Exception):
|
||||
pass
|
||||
async def runner(coro):
|
||||
task = asyncio.create_task(coro)
|
||||
await asyncio.sleep(0.012)
|
||||
if cancel_task:
|
||||
task.cancel()
|
||||
else:
|
||||
def interrupt():
|
||||
raise Interrupted('stop')
|
||||
module.test_management.throw_exception_if_processing_interrupted = interrupt
|
||||
return await task
|
||||
monkeypatch.setattr(module, 'run_async_in_server_loop', lambda coro: asyncio.run(runner(coro)))
|
||||
with pytest.raises(asyncio.CancelledError if cancel_task else Interrupted):
|
||||
remote(module)
|
||||
assert list(ps.distributed_pending_jobs) == ['unrelated']
|
||||
assert ps.distributed_serverless.cancelled == [{'collector_job_id': 'job'}]
|
||||
|
||||
|
||||
def test_replaced_queue_not_removed(module, monkeypatch):
|
||||
ps, _ = setup(module, manager=Manager(pending=True))
|
||||
replacement = asyncio.Queue()
|
||||
async def runner(coro):
|
||||
task = asyncio.create_task(coro)
|
||||
await asyncio.sleep(0.012)
|
||||
ps.distributed_pending_jobs['job'] = replacement
|
||||
return await task
|
||||
monkeypatch.setattr(module, 'run_async_in_server_loop', lambda coro: asyncio.run(runner(coro)))
|
||||
with pytest.raises(RuntimeError, match='queue'):
|
||||
remote(module)
|
||||
assert ps.distributed_pending_jobs['job'] is replacement
|
||||
|
||||
|
||||
def test_native_v3_schema_and_list_mapping_cpu():
|
||||
core = os.environ.get('COMFYUI_SOURCE_ROOT')
|
||||
if not core:
|
||||
pytest.skip('Set COMFYUI_SOURCE_ROOT for CPU-only native V3 mapping')
|
||||
probe = r'''
|
||||
import asyncio, importlib.util, sys, types
|
||||
from pathlib import Path
|
||||
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 torch
|
||||
import server
|
||||
# No HTTP listener or PromptServer bootstrap is needed for schema/list mapping.
|
||||
server.PromptServer.instance = types.SimpleNamespace()
|
||||
root = Path(sys.argv[2])
|
||||
for name, path in [('typed_collector', root), ('typed_collector.nodes', root/'nodes'),
|
||||
('typed_collector.utils', root/'utils'),
|
||||
('typed_collector.serverless_bridge', root/'serverless_bridge')]:
|
||||
package = types.ModuleType(name)
|
||||
package.__path__ = [str(path)]
|
||||
sys.modules[name] = package
|
||||
spec = importlib.util.spec_from_file_location('typed_collector.nodes.media_collector', root/'nodes/media_collector.py')
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
cls = module.DistributedMediaCollector
|
||||
info = cls.GET_NODE_INFO_V1()
|
||||
assert info['output'] == ['IMAGE','AUDIO','MASK','STRING','VIDEO'], info
|
||||
assert info['output_is_list'] == [True]*5, info
|
||||
assert cls.define_schema().is_input_list and cls.define_schema().accept_all_inputs
|
||||
for name, kind in zip(['images','audio','masks','texts','video'], info['output']):
|
||||
field = info['input']['optional'][name]
|
||||
assert field[0] == kind and field[1]['forceInput'] is True, field
|
||||
assert info['input']['required']['load_balance'][1]['default'] is False
|
||||
async def check():
|
||||
values = {'load_balance':False, 'pass_through':True}
|
||||
prepared, missing, hidden = execution.get_input_data(values, cls, unique_id='probe')
|
||||
assert not missing, missing
|
||||
image = torch.zeros((1,2,3,4), dtype=torch.float64)
|
||||
audio = {'waveform':torch.zeros((1,2,7)), 'sample_rate':48000}
|
||||
mask = torch.ones((1,2,3), dtype=torch.float64)
|
||||
from comfy_api.input_impl import VideoFromComponents
|
||||
from comfy_api.util import VideoComponents
|
||||
from fractions import Fraction
|
||||
video = VideoFromComponents(VideoComponents(images=image[:,:,:,:3].float(), frame_rate=Fraction(24,1)))
|
||||
native = {'images':[image], 'audio':[audio], 'masks':[mask], 'texts':[''], 'video':[video]}
|
||||
prepared.update(native)
|
||||
results = await execution._async_map_node_over_list('cpu','probe',cls,prepared,cls.FUNCTION,v3_data=hidden)
|
||||
results = await execution.resolve_map_node_over_list_results(results)
|
||||
assert len(results) == 1
|
||||
for index, name in enumerate(native):
|
||||
assert results[0].result[index] is native[name]
|
||||
assert results[0].result[4][0] is video
|
||||
print('MEDIA_COLLECTOR_CPU_OK')
|
||||
asyncio.run(check())
|
||||
'''
|
||||
env = {**os.environ, 'CUDA_VISIBLE_DEVICES': '', 'HIP_VISIBLE_DEVICES': '', 'PYTHONDONTWRITEBYTECODE': '1'}
|
||||
env.pop('PYTHONPATH', None)
|
||||
result = subprocess.run([sys.executable, '-c', probe, core, str(ROOT)], env=env,
|
||||
capture_output=True, text=True, timeout=90)
|
||||
assert result.returncode == 0, result.stdout + result.stderr
|
||||
assert 'MEDIA_COLLECTOR_CPU_OK' in result.stdout
|
||||
@@ -1,4 +1,4 @@
|
||||
"""Opt-in, isolated CPU acceptance of the complete typed-media data path.
|
||||
"""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
|
||||
@@ -42,6 +42,7 @@ def test_typed_media_end_to_end_cpu():
|
||||
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)
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
"""Pure integration of participant identity, typed compilation and delegate pruning."""
|
||||
"""Participant identity, typed compilation and legacy Collector pruning."""
|
||||
import copy
|
||||
import importlib
|
||||
import json
|
||||
@@ -16,17 +16,17 @@ for suffix in ("", ".api", ".api.orchestration", ".utils", ".serverless_bridge")
|
||||
sys.modules[module.__name__] = module
|
||||
pt = importlib.import_module(PACKAGE + ".api.orchestration.prompt_transform")
|
||||
compiler = importlib.import_module(PACKAGE + ".api.orchestration.serverless_media_prompt")
|
||||
contract = importlib.import_module(PACKAGE + ".serverless_bridge.contract")
|
||||
SOCKETS = ("images", "audio")
|
||||
|
||||
|
||||
def prompt():
|
||||
def prompt(sockets=SOCKETS):
|
||||
nodes = {
|
||||
"source": {"class_type": "MediaSource", "inputs": {}},
|
||||
"collector": {"class_type": "DistributedMediaCollector", "inputs": {
|
||||
socket: ["source", index] for index, socket in enumerate(contract.SOCKET_KINDS)}},
|
||||
"collector": {"class_type": "DistributedCollector", "inputs": {
|
||||
socket: ["source", SOCKETS.index(socket)] for socket in sockets}},
|
||||
}
|
||||
for index, socket in enumerate(contract.SOCKET_KINDS):
|
||||
nodes["sink:" + socket] = {"class_type": "Capture", "inputs": {"value": ["collector", index]}}
|
||||
for socket in sockets:
|
||||
nodes["sink:" + socket] = {"class_type": "Capture", "inputs": {"value": ["collector", SOCKETS.index(socket)]}}
|
||||
return nodes
|
||||
|
||||
|
||||
@@ -38,47 +38,41 @@ def prepare(graph, participant, delegate):
|
||||
["serverless:pool:1", "serverless:pool:2"], jobs, "http://master", delegate, index)
|
||||
|
||||
|
||||
def test_typed_worker_preparation_compiles_all_bindings_without_source_mutation():
|
||||
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) == 5
|
||||
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(contract.SOCKET_KINDS)
|
||||
assert {row["socket"] for row in bindings} == set(SOCKETS)
|
||||
assert set(compiled) == {"source"} | {row["node_id"] for row in bindings}
|
||||
assert graph == before
|
||||
|
||||
|
||||
def test_typed_delegate_master_keeps_outputs_without_sources_or_placeholders():
|
||||
graph = prompt()
|
||||
@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 set(delegated) == {"collector"} | {"sink:" + s for s in contract.SOCKET_KINDS}
|
||||
assert "source" not in delegated
|
||||
assert {"collector"} | {"sink:" + s for s in sockets} <= set(delegated)
|
||||
inputs = delegated["collector"]["inputs"]
|
||||
assert not set(contract.SOCKET_KINDS) & set(inputs)
|
||||
assert json.loads(inputs["expected_media"]) == list(contract.SOCKET_KINDS)
|
||||
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
|
||||
|
||||
|
||||
def test_text_source_is_not_retained_as_lightweight_delegate_scalar(monkeypatch):
|
||||
graph = {
|
||||
"text": {"class_type": "PrimitiveString", "inputs": {"value": "local text"}},
|
||||
"collector": {"class_type": "DistributedMediaCollector", "inputs": {"texts": ["text", 0]}},
|
||||
"sink": {"class_type": "Capture", "inputs": {"value": ["collector", 3]}},
|
||||
}
|
||||
monkeypatch.setattr(pt, "_delegate_master_input_is_safe_scalar", lambda *args: True)
|
||||
monkeypatch.setattr(pt, "_is_delegate_master_retained_upstream_node", lambda *args: True)
|
||||
delegated = pt.prepare_delegate_master_prompt(prepare(graph, "master", True), ["collector"])
|
||||
assert "text" not in delegated
|
||||
assert "texts" not in delegated["collector"]["inputs"]
|
||||
assert json.loads(delegated["collector"]["inputs"]["expected_media"]) == ["texts"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("socket", contract.SOCKET_KINDS)
|
||||
@pytest.mark.parametrize("socket", SOCKETS)
|
||||
def test_participating_master_retains_exact_media_links(socket):
|
||||
graph = prompt()
|
||||
original = copy.deepcopy(graph)
|
||||
|
||||
@@ -35,11 +35,11 @@ def node(class_type, **inputs):
|
||||
return {"class_type": class_type, "inputs": inputs}
|
||||
|
||||
|
||||
def text_prompt():
|
||||
def collector_prompt():
|
||||
return {
|
||||
"text": node("PrimitiveString", value=" café雪\n"),
|
||||
"collector": node("DistributedMediaCollector", texts=["text", 0], load_balance=False),
|
||||
"save": node("TextSink", texts=["collector", 3]),
|
||||
"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"),
|
||||
}
|
||||
|
||||
|
||||
@@ -95,7 +95,7 @@ def media_paths(tmp_path, monkeypatch):
|
||||
|
||||
def run(h, prompt=None, **kwargs):
|
||||
return asyncio.run(h.q.orchestrate_distributed_execution(
|
||||
prompt if prompt is not None else text_prompt(), None, "client", **kwargs))
|
||||
prompt if prompt is not None else collector_prompt(), None, "client", **kwargs))
|
||||
|
||||
|
||||
def assert_no_dispatch(h, *, no_probe=False):
|
||||
@@ -117,10 +117,11 @@ def test_real_compile_pack_serialize_and_canonical_bindings(harness, media_paths
|
||||
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 = text_prompt()
|
||||
prompt = collector_prompt()
|
||||
prompt["image"] = node("LoadImage", image="picture.png")
|
||||
prompt["collector"]["inputs"].update(images=["image", 0], masks=["image", 1])
|
||||
prompt["collector2"] = node("DistributedMediaCollector", texts=["text", 0])
|
||||
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, \
|
||||
@@ -145,8 +146,8 @@ def test_real_compile_pack_serialize_and_canonical_bindings(harness, media_paths
|
||||
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"]} == {"texts", "images", "masks"}
|
||||
assert len(job["manifest"]) == 4
|
||||
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"]))
|
||||
@@ -164,40 +165,50 @@ def test_real_compile_pack_serialize_and_canonical_bindings(harness, media_paths
|
||||
assert json.loads(controls["binding_json"]) == binding
|
||||
assert controls["request_id"] == media["request_id"]
|
||||
assert master["collector"]["inputs"]["delegate_only"] is delegate
|
||||
assert json.loads(master["collector"]["inputs"]["expected_media"]) == ["images", "masks", "texts"]
|
||||
assert "expected_media" not in master["collector"]["inputs"]
|
||||
if delegate:
|
||||
assert "text" not in master and "image" not in master
|
||||
assert not any(name in master["collector"]["inputs"] for name in ("images", "audio", "masks", "texts", "video"))
|
||||
assert not any(n["class_type"] == "DistributedEmptyImage" for n in master.values())
|
||||
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"]["texts"] == ["text", 0]
|
||||
assert "text" in master and "image" in master
|
||||
assert master["collector"]["inputs"]["audio"] == ["audio", 0]
|
||||
assert "audio" in master and "image" in master
|
||||
|
||||
|
||||
def test_delegate_disconnects_every_media_socket(harness):
|
||||
prompt = text_prompt()
|
||||
for index, name in enumerate(("images", "audio", "masks", "video")):
|
||||
source_id = "source-" + name
|
||||
prompt[source_id] = node("GeneratedMedia", seed=index)
|
||||
prompt["collector"]["inputs"][name] = [source_id, 0]
|
||||
@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"]
|
||||
sockets = {"images", "audio", "masks", "texts", "video"}
|
||||
assert set(json.loads(inputs["expected_media"])) == sockets
|
||||
assert not sockets.intersection(inputs)
|
||||
assert set(master) == {"collector", "save"}
|
||||
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"]} == sockets
|
||||
assert {"text", "source-images", "source-audio", "source-masks", "source-video"} <= set(job["workflow"])
|
||||
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 = text_prompt()
|
||||
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"))
|
||||
@@ -208,23 +219,24 @@ def test_typed_load_balance_selects_master_or_one_serverless_slot(harness, selec
|
||||
assert result[2] == 0
|
||||
h.manager.start.assert_not_awaited()
|
||||
assert master["collector"]["inputs"]["delegate_only"] is False
|
||||
assert "text" in master
|
||||
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 "text" not in master
|
||||
assert "source" not in master
|
||||
|
||||
|
||||
@pytest.mark.parametrize("profile", [None, "legacy-image-v1", PROFILE])
|
||||
def test_explicit_profile_compatibility_rejected_before_probe(harness, profile):
|
||||
@pytest.mark.parametrize("profile", [None, "legacy-image-v1"])
|
||||
def test_exact_rejects_legacy_pool_before_probe(harness, profile):
|
||||
h = harness
|
||||
prompt = legacy_prompt() if profile == PROFILE else text_prompt()
|
||||
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="requires|require"):
|
||||
with pytest.raises(ValueError, match="exact requires comfy-media-v1"):
|
||||
run(h, prompt)
|
||||
assert_no_dispatch(h, no_probe=True)
|
||||
|
||||
@@ -239,19 +251,26 @@ def test_invalid_explicit_profiles_fail_closed(harness, profile):
|
||||
|
||||
@pytest.mark.parametrize("direct_type", ["local", "remote", "cloud"])
|
||||
@pytest.mark.parametrize("explicit", [False, True])
|
||||
def test_direct_workers_never_silently_skipped_or_fallback(harness, direct_type, explicit):
|
||||
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})
|
||||
prompt = text_prompt()
|
||||
prompt["collector"]["inputs"]["load_balance"] = True
|
||||
kwargs = {"enabled_worker_ids": ["pool", "direct"]} if explicit else {}
|
||||
with pytest.raises(ValueError, match="direct workers"):
|
||||
run(h, prompt, delegate_master=False, **kwargs)
|
||||
assert_no_dispatch(h, no_probe=True)
|
||||
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_incompatible_workers_do_not_block(harness):
|
||||
def test_disabled_or_unselected_workers_do_not_participate(harness):
|
||||
h = harness
|
||||
h.config["workers"] += [
|
||||
{"id": "direct", "type": "local", "enabled": False},
|
||||
@@ -262,27 +281,32 @@ def test_disabled_or_unselected_incompatible_workers_do_not_block(harness):
|
||||
assert run(h, enabled_worker_ids=["pool"])[2] == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("other", ["DistributedCollector", "UltimateSDUpscaleDistributed", "DistributedImageListStart"])
|
||||
@pytest.mark.parametrize("workers", [True, False])
|
||||
def test_mixed_collectors_or_image_list_rejected_even_master_only(harness, other, workers):
|
||||
h = harness
|
||||
def test_image_list_cannot_mix_with_collector_even_master_only(harness, workers):
|
||||
if not workers:
|
||||
h.config["workers"] = []
|
||||
prompt = text_prompt()
|
||||
prompt["incompatible"] = node(other)
|
||||
harness.config["workers"] = []
|
||||
prompt = collector_prompt()
|
||||
prompt["incompatible"] = node("DistributedImageListStart")
|
||||
with pytest.raises(ValueError, match="cannot be mixed"):
|
||||
run(h, prompt, delegate_master=False)
|
||||
assert_no_dispatch(h, no_probe=True)
|
||||
run(harness, prompt, delegate_master=False)
|
||||
assert_no_dispatch(harness, no_probe=True)
|
||||
|
||||
|
||||
def test_mixed_pool_profiles_rejected_before_load_balancing(harness):
|
||||
def test_mixed_direct_typed_and_legacy_profiles_keep_each_wire_contract(harness):
|
||||
h = harness
|
||||
h.config["workers"].append(dict(h.pool, id="legacy", media_profile="legacy-image-v1"))
|
||||
prompt = text_prompt()
|
||||
prompt["collector"]["inputs"]["load_balance"] = True
|
||||
with pytest.raises(ValueError, match="requires"):
|
||||
run(h, prompt, delegate_master=False)
|
||||
assert_no_dispatch(h, no_probe=True)
|
||||
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):
|
||||
@@ -291,27 +315,23 @@ def test_typed_master_only_is_supported_without_serverless_work(harness):
|
||||
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"]["texts"] == ["text", 0]
|
||||
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_typed_delegate_requires_available_worker_without_master_fallback(harness, offline):
|
||||
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"] = []
|
||||
with pytest.raises(RuntimeError, match="comfy-media-v1"):
|
||||
run(h)
|
||||
assert_no_dispatch(h)
|
||||
|
||||
|
||||
def test_unknown_explicit_typed_worker_is_rejected(harness):
|
||||
with pytest.raises(ValueError, match="Unknown explicitly selected"):
|
||||
run(harness, enabled_worker_ids=["unknown"], delegate_master=False)
|
||||
assert_no_dispatch(harness, no_probe=True)
|
||||
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"])
|
||||
@@ -323,9 +343,9 @@ def test_pool_media_identity_is_not_repaired(harness, bad):
|
||||
|
||||
|
||||
def test_integer_source_node_identity_is_not_normalized(harness):
|
||||
prompt = text_prompt()
|
||||
prompt[1] = prompt.pop("text")
|
||||
prompt["collector"]["inputs"]["texts"] = [1, 0]
|
||||
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)
|
||||
@@ -334,18 +354,18 @@ def test_integer_source_node_identity_is_not_normalized(harness):
|
||||
@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 = text_prompt()
|
||||
prompt = collector_prompt()
|
||||
if broken == "dangling-unused":
|
||||
prompt["unused"] = node("PreviewImage", images=["missing", 0])
|
||||
elif broken == "literal":
|
||||
prompt["collector"]["inputs"]["texts"] = "must use a linked string node"
|
||||
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["text"]["inputs"]["value"] = "雪" * 1000
|
||||
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)):
|
||||
@@ -418,9 +438,9 @@ def test_master_queue_failure_cleans_queues_and_never_starts_manager(harness):
|
||||
|
||||
def test_manager_setup_failure_emits_one_terminal_per_collector_participant(harness):
|
||||
h = harness
|
||||
prompt = text_prompt()
|
||||
prompt["image"] = node("EmptyImage", width=2, height=2, batch_size=1, color=0)
|
||||
prompt["collector"]["inputs"]["images"] = ["image", 0]
|
||||
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)
|
||||
@@ -481,6 +501,7 @@ 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
|
||||
@@ -507,7 +528,7 @@ raise SystemExit(pytest.main([sys.argv[2], '-q', '-p', 'no:cacheprovider', '-k',
|
||||
stream.setframerate(8000)
|
||||
stream.writeframes(pcm.numpy().tobytes())
|
||||
original_bytes = path.read_bytes()
|
||||
prompt = text_prompt()
|
||||
prompt = collector_prompt()
|
||||
prompt["audio"] = node("LoadAudio", audio="tone.wav")
|
||||
prompt["collector"]["inputs"]["audio"] = ["audio", 0]
|
||||
original = copy.deepcopy(prompt)
|
||||
|
||||
@@ -6,16 +6,16 @@ import pytest
|
||||
|
||||
from test_serverless_media_prompt import compiler, _prompt
|
||||
from test_media_handler import handler
|
||||
from test_media_collector import module, collect, native_media
|
||||
from test_collector_typed_media import collector_module, native_media, assert_exact
|
||||
|
||||
FORMATS = ('exact', 'legacy_png', 'png', 'ffv1', 'jpeg', 'webp', 'h264', 'h265', 'av1')
|
||||
FORMATS = ('legacy_png', 'png', 'ffv1', 'jpeg', 'webp', 'h264', 'h265', 'av1', 'exact')
|
||||
|
||||
|
||||
def compile_request(fmt='exact', quality=90):
|
||||
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,
|
||||
masks=['source:1', 1], texts=['source:1', 2])
|
||||
audio=['source:1', 1])
|
||||
original = copy.deepcopy(prompt)
|
||||
workflow, outputs = compiler.compile_media_prompt(prompt, 'request:1')
|
||||
assert prompt == original
|
||||
@@ -27,16 +27,23 @@ def compile_request(fmt='exact', quality=90):
|
||||
@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'] == fmt
|
||||
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):
|
||||
@@ -54,23 +61,36 @@ def test_handler_rejects_tampered_policy(field, value):
|
||||
handler._validate_request(data)
|
||||
|
||||
|
||||
def test_handler_rejects_image_controls_on_mask_export():
|
||||
data = compile_request()
|
||||
binding = next(b for b in data['media']['outputs'] if b['kind'] == 'MASK')
|
||||
data['workflow'][binding['node_id']]['inputs']['transfer_format'] = 'jpeg'
|
||||
@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(module, fmt):
|
||||
def test_master_passthrough_stays_native_for_every_format(collector_module, fmt):
|
||||
media = native_media()
|
||||
actual = collect(module, **media, transfer_format=[fmt], transfer_quality=[73])
|
||||
for values, original in zip(actual, media.values()):
|
||||
assert values is original
|
||||
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'])
|
||||
|
||||
|
||||
@pytest.mark.parametrize('quality', [True, 0, 101, '90', [90,91]])
|
||||
def test_node_rejects_bad_quality_even_without_remote_participants(module, quality):
|
||||
with pytest.raises(ValueError):
|
||||
collect(module, texts=['hello'], transfer_quality=quality)
|
||||
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
|
||||
|
||||
@@ -56,7 +56,7 @@ def test_serverless_pending_waits_past_normal_worker_timeout():
|
||||
queue.put_nowait({
|
||||
"worker_id": participant,
|
||||
"image_index": 0,
|
||||
"tensor": torch.ones(1, 2, 2, 2),
|
||||
"tensor": torch.ones(1, 2, 2, 3),
|
||||
"is_last": True,
|
||||
})
|
||||
|
||||
@@ -74,7 +74,7 @@ def test_serverless_pending_waits_past_normal_worker_timeout():
|
||||
|
||||
images, _audio = asyncio.run(run_collection())
|
||||
|
||||
assert tuple(images.shape) == (1, 2, 2, 2)
|
||||
assert tuple(images.shape) == (1, 2, 2, 3)
|
||||
assert manager.pending_calls
|
||||
assert probe_calls == []
|
||||
assert job_id not in module.prompt_server.distributed_pending_jobs
|
||||
@@ -162,12 +162,12 @@ def test_serverless_tensor_shape_mismatch_raises_instead_of_silent_fallback():
|
||||
queue.put_nowait({
|
||||
"worker_id": "serverless:pool:1",
|
||||
"image_index": 0,
|
||||
"tensor": torch.ones(1, 3, 2, 2),
|
||||
"tensor": torch.ones(1, 3, 2, 3),
|
||||
"is_last": True,
|
||||
})
|
||||
manager = _FakeManager()
|
||||
collector = _set_up_collector(module, job_id, queue, manager)
|
||||
local_images = torch.zeros(1, 2, 2, 2)
|
||||
local_images = torch.zeros(1, 2, 2, 3)
|
||||
|
||||
with pytest.raises(RuntimeError, match="serverless.*shape|shape.*serverless"):
|
||||
asyncio.run(collector.execute(
|
||||
@@ -188,7 +188,7 @@ def test_regular_worker_timeout_keeps_partial_behavior_without_creating_manager(
|
||||
collector = _set_up_collector(module, job_id, queue)
|
||||
module.get_worker_timeout_seconds = lambda: 0.01
|
||||
module.HEARTBEAT_INTERVAL = 0.1
|
||||
local_images = torch.zeros(1, 2, 2, 2)
|
||||
local_images = torch.zeros(1, 2, 2, 3)
|
||||
collector._get_serverless_manager = lambda *_args, **_kwargs: pytest.fail("regular collection must not create manager")
|
||||
|
||||
images, _audio = asyncio.run(collector.execute(
|
||||
|
||||
@@ -24,7 +24,7 @@ compiler = importlib.import_module(f"{PACKAGE}.api.orchestration.serverless_medi
|
||||
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", "masks": "MASK", "texts": "STRING", "video": "VIDEO"}
|
||||
SOCKETS = {"images": "IMAGE", "audio": "AUDIO"}
|
||||
|
||||
|
||||
def _node(class_type, **inputs):
|
||||
@@ -33,7 +33,7 @@ def _node(class_type, **inputs):
|
||||
|
||||
def _collector(**sockets):
|
||||
return _node(
|
||||
"DistributedMediaCollector", multi_job_id="job/42", worker_id="worker_0", **sockets
|
||||
"DistributedCollector", multi_job_id="job/42", worker_id="worker_0", **sockets
|
||||
)
|
||||
|
||||
|
||||
@@ -72,26 +72,26 @@ def test_one_export_for_each_typed_socket_with_exact_canonical_binding(socket, k
|
||||
"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_mixed_collectors_all_five_kinds_have_independent_dispatch_targets():
|
||||
def test_multiple_collectors_have_independent_image_audio_targets():
|
||||
prompt = {
|
||||
"shared": _node("MultiMediaSource"),
|
||||
"collector:z": _collector(audio=["shared", 1], video=["shared", 4]),
|
||||
"collector:a": _collector(images=["shared", 0], masks=["shared", 2], texts=["shared", 3]),
|
||||
"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", "masks"), ("collector:a", "texts"),
|
||||
("collector:z", "audio"), ("collector:z", "video"),
|
||||
("collector:a", "images"), ("collector:a", "audio"), ("collector:z", "audio"),
|
||||
]
|
||||
assert len({b["binding_id"] for b in bindings}) == 5
|
||||
assert len({b["binding_id"] for b in bindings}) == 3
|
||||
targets = {b["node_id"] for b in bindings}
|
||||
assert len(targets) == 5
|
||||
assert len(targets) == 3
|
||||
assert set(graph) == {"shared"} | targets
|
||||
for binding in bindings:
|
||||
export = graph[binding["node_id"]]
|
||||
@@ -119,12 +119,15 @@ def test_prunes_descendants_and_unrelated_outputs_preserving_required_output_uti
|
||||
assert not {"load_balance", "pass_through"} & graph[bindings[0]["node_id"]]["inputs"].keys()
|
||||
|
||||
|
||||
def test_stock_primitive_string_preserves_empty_and_link_looking_text():
|
||||
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), "collector": _collector(texts=["text", 0])}
|
||||
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[bindings[0]["node_id"]]["inputs"]["texts"] == ["text", 0]
|
||||
assert graph["source"]["inputs"]["text"] == ["text", 0]
|
||||
assert graph[bindings[0]["node_id"]]["inputs"]["images"] == ["source", 0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("class_type,inputs", [
|
||||
@@ -176,11 +179,13 @@ def test_distributed_text_value_can_be_wrapped_by_primitive_string():
|
||||
prompt = {
|
||||
"value": _node("DistributedModelName", text="model.safetensors"),
|
||||
"text": _node("PrimitiveString", value=["value", 0]),
|
||||
"collector": _collector(texts=["text", 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[bindings[0]["node_id"]]["inputs"]["texts"] == ["text", 0]
|
||||
assert graph["source"]["inputs"]["text"] == ["text", 0]
|
||||
assert graph[bindings[0]["node_id"]]["inputs"]["images"] == ["source", 0]
|
||||
assert "value" not in graph
|
||||
|
||||
|
||||
@@ -321,7 +326,7 @@ def test_rejects_invalid_links_everywhere_before_pruning(location, link):
|
||||
|
||||
|
||||
@pytest.mark.parametrize("class_type", [
|
||||
"DistributedCollector", "DistributedEmptyImage", "UltimateSDUpscaleDistributed",
|
||||
"DistributedMediaCollector", "DistributedEmptyImage", "UltimateSDUpscaleDistributed",
|
||||
"DistributedImageListStart", "DistributedImageListEnd", "DistributedWorkerSink",
|
||||
"ImageBatchDivider", "AudioBatchDivider", "CustomDistributedNode", "custombatchdivider",
|
||||
])
|
||||
@@ -353,9 +358,9 @@ def test_rejects_collector_chains_across_different_media(indirect):
|
||||
prompt = _prompt()
|
||||
source = "collector:1"
|
||||
if indirect:
|
||||
prompt["processing"] = _node("ImageToVideo", images=[source, 0])
|
||||
prompt["processing"] = _node("ImageToAudio", images=[source, 0])
|
||||
source = "processing"
|
||||
prompt["collector:2"] = _collector(video=[source, 0])
|
||||
prompt["collector:2"] = _collector(audio=[source, 0])
|
||||
with pytest.raises(ValueError, match="collector-to-collector"):
|
||||
_run(prompt)
|
||||
|
||||
@@ -413,7 +418,7 @@ def guarded_import(name, *args, **kwargs):
|
||||
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": "DistributedMediaCollector", "inputs": {"texts": ["s", 0], "multi_job_id": "job", "worker_id": "worker"}}}
|
||||
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"
|
||||
'''
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -69,8 +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.media_collector", DistributedMediaCollector=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__=[])
|
||||
@@ -98,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"]
|
||||
@@ -107,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"]
|
||||
|
||||
|
||||
@@ -34,6 +34,7 @@ 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
|
||||
|
||||
+18
-5
@@ -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")
|
||||
|
||||
@@ -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);
|
||||
},
|
||||
});
|
||||
@@ -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");
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user