Compare commits
6
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2df73e4b9f | ||
|
|
b14c51c365 | ||
|
|
3699477f75 | ||
|
|
e9c95bf110 | ||
|
|
1fb4d4a284 | ||
|
|
4e823397b7 |
@@ -5,6 +5,7 @@ from aiohttp import web
|
||||
import server
|
||||
|
||||
from ..image_list.runtime import (
|
||||
IMAGE_TRANSFER_FORMATS,
|
||||
MAX_REQUEST_BYTES,
|
||||
PROTOCOL_VERSION,
|
||||
decode_result,
|
||||
@@ -35,9 +36,11 @@ async def _json_body(request):
|
||||
async def image_list_capabilities_endpoint(_request):
|
||||
return web.json_response({
|
||||
"supported": True,
|
||||
"execution_targets": True,
|
||||
"protocol_version": PROTOCOL_VERSION,
|
||||
"node_ids": ["DistributedImageListStart", "DistributedImageListEnd"],
|
||||
"transport": "typed-float-tensor-json-v1",
|
||||
"image_transfer_formats": list(IMAGE_TRANSFER_FORMATS),
|
||||
})
|
||||
|
||||
|
||||
|
||||
@@ -77,6 +77,8 @@ async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0):
|
||||
"workflow": payload.get("workflow"),
|
||||
"client_id": client_id,
|
||||
}
|
||||
if "partial_execution_targets" in payload:
|
||||
ws_payload["partial_execution_targets"] = payload["partial_execution_targets"]
|
||||
ws_url = worker_url.replace("http://", "ws://").replace("https://", "wss://")
|
||||
ws_url = f"{ws_url}/distributed/worker_ws"
|
||||
session = await get_client_session()
|
||||
@@ -128,11 +130,17 @@ async def dispatch_worker_prompt(
|
||||
use_websocket=False,
|
||||
trace_execution_id=None,
|
||||
return_prompt_id=False,
|
||||
*,
|
||||
partial_execution_targets=None,
|
||||
):
|
||||
"""Send the prepared prompt to a worker ComfyUI instance."""
|
||||
worker_url = build_worker_url(worker)
|
||||
url = build_worker_url(worker, "/prompt")
|
||||
payload = {"prompt": prompt_obj}
|
||||
target_payload = {}
|
||||
if partial_execution_targets is not None:
|
||||
target_payload["partial_execution_targets"] = list(partial_execution_targets)
|
||||
payload.update(target_payload)
|
||||
extra_data = {}
|
||||
if workflow_meta:
|
||||
extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
|
||||
@@ -146,6 +154,7 @@ async def dispatch_worker_prompt(
|
||||
{
|
||||
"prompt": prompt_obj,
|
||||
"workflow": workflow_meta,
|
||||
**target_payload,
|
||||
},
|
||||
client_id,
|
||||
)
|
||||
@@ -179,7 +188,8 @@ async def dispatch_worker_prompt(
|
||||
return prompt_id if return_prompt_id else None
|
||||
|
||||
|
||||
async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0):
|
||||
async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0, *,
|
||||
required_transfer_formats=()):
|
||||
"""Require protocol/node support and report which direct workers are idle.
|
||||
|
||||
This performs read-only GETs only; it never submits a prompt or contacts a
|
||||
@@ -187,6 +197,14 @@ async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0):
|
||||
"""
|
||||
if not workers:
|
||||
raise RuntimeError("Distributed Image List requires at least one selected direct worker")
|
||||
if (not isinstance(required_transfer_formats, (list, tuple, set, frozenset))
|
||||
or any(type(fmt) is not str for fmt in required_transfer_formats)):
|
||||
raise ValueError("required_transfer_formats must contain image transfer format names")
|
||||
required_formats = set(required_transfer_formats) - {"exact"}
|
||||
if required_formats:
|
||||
from ...serverless_bridge.image_transfer import validate_options
|
||||
for fmt in required_formats:
|
||||
validate_options(fmt)
|
||||
session = await get_client_session()
|
||||
required_nodes = {"DistributedImageListStart", "DistributedImageListEnd"}
|
||||
semaphore = asyncio.Semaphore(8)
|
||||
@@ -227,6 +245,23 @@ async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0):
|
||||
f"{sorted(required_nodes)}; no work was submitted"
|
||||
)
|
||||
|
||||
if capability.get("execution_targets") is not True:
|
||||
raise RuntimeError(
|
||||
f"Image-list worker {worker_id} lacks explicit execution_targets support; "
|
||||
"upgrade that worker before running this feature; no work was submitted"
|
||||
)
|
||||
|
||||
if required_formats:
|
||||
advertised = capability.get("image_transfer_formats")
|
||||
if (not isinstance(advertised, list)
|
||||
or any(type(fmt) is not str for fmt in advertised)
|
||||
or not required_formats.issubset(advertised)):
|
||||
raise RuntimeError(
|
||||
f"Image-list worker {worker_id} lacks required image_transfer_formats "
|
||||
f"{sorted(required_formats)}; upgrade that worker before compressed IMAGE returns; "
|
||||
"no work was submitted"
|
||||
)
|
||||
|
||||
check = "queue"
|
||||
queue_url = build_worker_url(worker, "/queue")
|
||||
async with session.get(queue_url, timeout=aiohttp.ClientTimeout(total=timeout)) as response:
|
||||
|
||||
@@ -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
|
||||
|
||||
+39
-21
@@ -38,7 +38,8 @@ from .orchestration.prompt_transform import (
|
||||
WORKER_SINK_CLASS,
|
||||
_resolve_direct_workers_for_image_list,
|
||||
)
|
||||
from ..image_list.runtime import validate_image_list_regions
|
||||
from ..image_list.runtime import validate_image_list_regions, _has_selected_serverless_workers
|
||||
from ..image_list.execution_plan import prepare_image_list_execution
|
||||
|
||||
|
||||
prompt_server = server.PromptServer.instance
|
||||
@@ -327,6 +328,16 @@ async def orchestrate_distributed_execution(
|
||||
media_sync_timeout_seconds,
|
||||
) = _resolve_orchestration_limits(config)
|
||||
requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None
|
||||
image_list_mode = any(isinstance(node, dict) and node.get("class_type") == "DistributedImageListStart"
|
||||
for node in prompt_obj.values())
|
||||
target_kwargs = {}
|
||||
if image_list_mode:
|
||||
if any(isinstance(node, dict) and node.get("class_type") in {
|
||||
"DistributedCollector", "DistributedMediaCollector", "UltimateSDUpscaleDistributed"
|
||||
} for node in prompt_obj.values()):
|
||||
raise ValueError("Distributed Image List cannot be mixed with other distributed collectors")
|
||||
prompt_obj, outer_targets = prepare_image_list_execution(prompt_obj)
|
||||
target_kwargs["partial_execution_targets"] = outer_targets
|
||||
prompt_index = PromptIndex(prompt_obj)
|
||||
typed_media_mode = bool(prompt_index.nodes_for_class("DistributedMediaCollector"))
|
||||
if typed_media_mode:
|
||||
@@ -349,24 +360,25 @@ async def orchestrate_distributed_execution(
|
||||
"Distributed Image List v1 cannot be mixed with Distributed Collector, "
|
||||
"DistributedMediaCollector or Ultimate SD Upscale nodes in one prompt"
|
||||
)
|
||||
validate_image_list_regions(prompt_obj)
|
||||
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),
|
||||
)
|
||||
validate_image_list_regions(prompt_obj, explicit_targets=True)
|
||||
from ..serverless_bridge.image_transfer import DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options
|
||||
required_transfer_formats = set()
|
||||
for end_id in prompt_index.nodes_for_class("DistributedImageListEnd"):
|
||||
end_inputs = prompt_obj[end_id].get("inputs", {})
|
||||
fmt, _ = validate_options(end_inputs.get("transfer_format", DEFAULT_FORMAT),
|
||||
end_inputs.get("transfer_quality", DEFAULT_QUALITY))
|
||||
if fmt != DEFAULT_FORMAT:
|
||||
required_transfer_formats.add(fmt)
|
||||
workers = _resolve_direct_workers_for_image_list(config, requested_ids)
|
||||
if _has_selected_serverless_workers(config, requested_ids):
|
||||
from ..image_list.serverless import resolve_serverless_workers
|
||||
# Pool identity is retained once, never expanded by Images per run.
|
||||
# Validate profiles before reachability can hide an invalid selection.
|
||||
workers += resolve_serverless_workers(config, requested_ids)
|
||||
if not workers and delegate_master:
|
||||
raise RuntimeError(
|
||||
"Distributed Image List in orchestrator-only mode requires a selected direct worker; "
|
||||
"serverless pools are unsupported"
|
||||
"Distributed Image List in orchestrator-only mode requires a selected direct worker "
|
||||
"or comfy-media-v1 serverless pool"
|
||||
)
|
||||
else:
|
||||
if typed_media_mode and requested_ids is not None:
|
||||
@@ -414,13 +426,17 @@ async def orchestrate_distributed_execution(
|
||||
from .orchestration.dispatch import preflight_image_list_workers
|
||||
if not active_workers and delegate_master:
|
||||
raise RuntimeError(
|
||||
"No selected direct worker is reachable for Distributed Image List in orchestrator-only mode"
|
||||
"No selected worker is reachable for Distributed Image List in orchestrator-only mode"
|
||||
)
|
||||
idle_ids = (await preflight_image_list_workers(active_workers, require_idle=delegate_master)
|
||||
if active_workers else {})
|
||||
active_workers = [worker for worker in active_workers if worker["id"] in idle_ids]
|
||||
direct_workers = [worker for worker in active_workers if worker.get("type") != "serverless"]
|
||||
transfer_kwargs = ({"required_transfer_formats": sorted(required_transfer_formats)}
|
||||
if required_transfer_formats else {})
|
||||
idle_ids = (await preflight_image_list_workers(direct_workers, require_idle=delegate_master, **transfer_kwargs)
|
||||
if direct_workers else {})
|
||||
active_workers = [worker for worker in active_workers
|
||||
if worker.get("type") == "serverless" or 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")
|
||||
raise RuntimeError("No feature-compatible worker is available in orchestrator-only mode")
|
||||
else:
|
||||
delegate_master = selected_delegate_master
|
||||
|
||||
@@ -487,6 +503,7 @@ async def orchestrate_distributed_execution(
|
||||
workflow_meta,
|
||||
client_id,
|
||||
include_queue_metadata=True,
|
||||
**target_kwargs,
|
||||
)
|
||||
return (
|
||||
queue_result["prompt_id"],
|
||||
@@ -594,6 +611,7 @@ async def orchestrate_distributed_execution(
|
||||
workflow_meta,
|
||||
client_id,
|
||||
include_queue_metadata=True,
|
||||
**target_kwargs,
|
||||
)
|
||||
except Exception:
|
||||
if serverless_jobs:
|
||||
|
||||
@@ -61,6 +61,13 @@ async def worker_ws_endpoint(request):
|
||||
})
|
||||
continue
|
||||
|
||||
if not isinstance(data, dict):
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack", "request_id": None, "ok": False,
|
||||
"error": "Expected a websocket message object.",
|
||||
})
|
||||
continue
|
||||
|
||||
if data.get("type") != "dispatch_prompt":
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
@@ -81,10 +88,16 @@ async def worker_ws_endpoint(request):
|
||||
continue
|
||||
|
||||
try:
|
||||
target_kwargs = {}
|
||||
if "partial_execution_targets" in data:
|
||||
from ..utils.async_helpers import validate_execution_targets
|
||||
target_kwargs["partial_execution_targets"] = validate_execution_targets(
|
||||
prompt, data["partial_execution_targets"])
|
||||
prompt_id = await queue_prompt_payload(
|
||||
prompt,
|
||||
workflow_meta=data.get("workflow"),
|
||||
client_id=data.get("client_id"),
|
||||
**target_kwargs,
|
||||
)
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
|
||||
+95
-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.
|
||||
@@ -60,22 +60,96 @@ 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.
|
||||
- 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).
|
||||
|
||||
## 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.
|
||||
- Worker eligibility is determined before execution and checked again before dispatch. Direct workers must be idle in orchestrator-only mode; master-participating mode can use reachable compatible peers without requiring an empty queue. 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 and compatible serverless pools. Explicit IDs select the named configured workers/pools, including disabled ones; an empty list stays empty and unknown IDs fail. Pool IDs are retained in Start's hidden selection metadata, not expanded participant IDs. Local direct workers use the master's loopback callback; remote direct workers use its configured public address.
|
||||
- Only the enclosed processing is distributed. List creation and downstream saving run on the master in both modes. With participation enabled, the master is a normal scheduler participant and can run alone when no workers are selected or reachable. **Orchestrator-only mode never silently falls back to master processing.** Failed items are not reassigned automatically.
|
||||
- Local items execute while the enclosing master prompt waits; they are not queued behind that prompt on the master's `/prompt` queue. The enclosed branch is not also evaluated as a full list by the outer prompt.
|
||||
- Start/End regions may be chained sequentially and share model-loader dependencies. Nested/overlapping per-image regions and data escaping a region without End are rejected. Collector and distributed-upscaler nodes cannot be mixed into the same prompt in this first version.
|
||||
- Output nodes inside a region are rejected because ComfyUI would execute them independently on the master. Avoid other side-effecting processing nodes as well; save after End.
|
||||
- The entire input list is materialized before scheduling; this is not streaming input. The limit is 4,096 frames. Remote transport preserves floating-point tensor dtype and values rather than encoding to lossy images. Each transported tensor is limited to 50 MiB, and the complete JSON/base64 HTTP request to 101 MiB; image/mask/value combinations can hit the aggregate request limit first. In-process master items use tensors directly and are not subject to those HTTP limits.
|
||||
- Queue preparation works on an execution copy, leaving the saved canvas unchanged. Unused per-item preview/save branches are pruned by dependency reachability, not node names. Required computational utilities marked `OUTPUT_NODE` stay in the graph; explicit `partial_execution_targets` prevent them from also becoming independent outer-master roots. End and legitimate post-End or unrelated outputs are retained. Bypass/rejoin branches still fail boundary validation rather than being silently removed.
|
||||
- Every selected direct worker must advertise `execution_targets: true` as well as the existing image-list protocol. Older workers fail preflight before dispatch. Both HTTP and WebSocket item submissions target End only; the participating master uses its existing End-only local executor. Cloud jobs use explicit media exporter roots. Submit through `/distributed/queue`: an ordinary all-output `/prompt` request does not receive this relaxation.
|
||||
- **Target selection is not a sandbox.** A required utility still runs its own save, preview or other side effects when it executes for an item. Keep intentional final saving after End.
|
||||
- The entire input list is materialized before scheduling; this is not streaming input. The limit is 4,096 frames. Direct-worker input transport and default exact returns preserve floating-point tensor dtype and values. Opt-in image return encoding does not relax decoded-size limits. Each direct-worker tensor is limited to 50 MiB, and the complete JSON/base64 HTTP request to 101 MiB; image/mask/value combinations can hit the aggregate request limit first. Serverless inline limits are smaller, as described above. In-process master items use tensors directly and are not subject to those HTTP limits.
|
||||
- Results are accepted once per item using assignment tokens. Missing, duplicate, malformed and wrong-index results fail rather than silently returning a partial list. No failed or ambiguous submission is retried automatically.
|
||||
- Worker errors fail the region. The worker timeout measures loss of evidence for the owned prompt ID, not maximum render duration. Pending/running evidence refreshes liveness; unrelated busy jobs and failed polls do not. Healthy cold loads and long renders can continue.
|
||||
- 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
|
||||
|
||||
@@ -89,7 +163,14 @@ Use trusted workers and trusted workflows, as with the existing ComfyUI API. Cap
|
||||
|
||||
## Verification
|
||||
|
||||
The generic acceptance helper starts two isolated loopback worker test servers and uses actual CPU ComfyUI `PromptExecutor` instances. Production queue preparation, list nodes, graph extraction, HTTP item/result routes, dynamic scheduling, float transport, ordering and cancellation are exercised. The harness checks orchestrator-only, mixed master/worker and master-only execution, including exact-once item assignment, source-once materialization and cached local model-loader dependencies. The lightweight queue/history facade is test code; no managed ComfyUI service or model is started.
|
||||
The generic acceptance helper starts two isolated loopback workers using real ComfyUI HTTP handlers, `PromptQueue` records and CPU `PromptExecutor` instances. Production queue preparation, list nodes, graph extraction, HTTP item/result routes, dynamic scheduling, float transport, ordering and cancellation are exercised. The harness checks orchestrator-only, mixed master/worker and master-only execution, including exact-once required output utilities, discarded previews, unchanged source graphs, End-only worker roots, post-End capture, source-once materialization and cached local model-loader dependencies. WebSocket target propagation is covered by focused route/dispatch tests. No managed ComfyUI service or model is started.
|
||||
|
||||
Offline cloud integration tests additionally cover mixed-role and cloud-only
|
||||
scheduling, out-of-order image/mask/text pairing, one job per source item regardless
|
||||
of Images per run, profile/selection rejection before queueing, exact cleanup on
|
||||
setup/dispatch/result failures, and the real CPU ComfyUI `processing_interrupted`
|
||||
signal while awaiting a mocked cloud result. Provider calls are mocked: this is
|
||||
not live Runpod, endpoint-supervisor or GPU-workflow acceptance.
|
||||
|
||||
```bash
|
||||
COMFYUI_SOURCE_ROOT=/path/to/ComfyUI /path/to/ComfyUI/.venv/bin/python -m pytest tests -q
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
"""Prepare image-list execution without changing the saved/API prompt or registry."""
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import sys
|
||||
|
||||
from .runtime import (
|
||||
END_CLASS,
|
||||
_ancestors_for_inputs,
|
||||
_downstream_until,
|
||||
_image_list_pairs,
|
||||
_is_output_node,
|
||||
_prompt_links,
|
||||
validate_image_list_regions,
|
||||
)
|
||||
|
||||
|
||||
def _reachable(roots: set[str], edges: dict[str, set[str]]) -> set[str]:
|
||||
visited = set()
|
||||
pending = list(roots)
|
||||
while pending:
|
||||
node_id = pending.pop()
|
||||
if node_id in visited:
|
||||
continue
|
||||
visited.add(node_id)
|
||||
pending.extend(edges[node_id] - visited)
|
||||
return visited
|
||||
|
||||
|
||||
def prepare_image_list_execution(prompt: dict) -> tuple[dict, list[str]]:
|
||||
"""Return an execution copy and its outer-master output targets.
|
||||
|
||||
Disposable per-item side branches are removed before boundary validation.
|
||||
Required output-node utilities remain normal dependencies, but must not be
|
||||
independent outer roots. Callers must use these targets (and End-only targets
|
||||
for item execution); this does not suppress side effects inside required nodes.
|
||||
"""
|
||||
pairs = _image_list_pairs(prompt)
|
||||
prepared = copy.deepcopy(prompt)
|
||||
upstream = {node_id: set() for node_id in prepared}
|
||||
downstream = {node_id: set() for node_id in prepared}
|
||||
for node_id, node in prepared.items():
|
||||
for _, source, _ in _prompt_links(node, prepared):
|
||||
upstream[node_id].add(source)
|
||||
downstream[source].add(node_id)
|
||||
|
||||
# Never erase another region or its dependencies to make a bad graph pass.
|
||||
boundaries = {node_id for pair in pairs for node_id in pair}
|
||||
required = _reachable(boundaries, upstream)
|
||||
discarded = set()
|
||||
for start_id, end_id in pairs:
|
||||
region_nodes = _ancestors_for_inputs(prepared, end_id, start_id, ("images", "masks", "texts"))
|
||||
outside = _downstream_until(prepared, start_id, end_id) - region_nodes
|
||||
# Preserve bypass joins and their whole chain into post-End computation.
|
||||
# They must fail boundary validation, not disappear with debug previews.
|
||||
post_end = _reachable({end_id}, downstream)
|
||||
protected = required | _reachable(post_end, upstream)
|
||||
discarded.update(outside - protected)
|
||||
for node_id in discarded:
|
||||
del prepared[node_id]
|
||||
|
||||
regions = validate_image_list_regions(prepared, explicit_targets=True)
|
||||
per_item = set().union(*(region["branch_ids"] for region in regions))
|
||||
registry = getattr(sys.modules.get("nodes"), "NODE_CLASS_MAPPINGS", {})
|
||||
targets = [
|
||||
node_id for node_id, node in prepared.items()
|
||||
if node_id not in per_item and isinstance(node, dict)
|
||||
and (node.get("class_type") == END_CLASS or _is_output_node(node, registry))
|
||||
]
|
||||
return prepared, targets
|
||||
@@ -14,6 +14,38 @@ DEFAULT_ORPHAN_GRACE_SECONDS = 300.0
|
||||
_LIFECYCLE_MARKER = object()
|
||||
|
||||
|
||||
def has_safe_execution_targets(prompt_server: Any, prompt_id: str | None, prompt: dict) -> bool:
|
||||
"""Prove outer roots from the running queue, not a forgeable prompt marker.
|
||||
|
||||
Plain /prompt remains strict when ComfyUI selected all outputs, including a
|
||||
flagged per-item dependency. No saved workflow field grants this exemption.
|
||||
"""
|
||||
if not prompt_id:
|
||||
return False
|
||||
from .execution_plan import prepare_image_list_execution
|
||||
prepared, targets = prepare_image_list_execution(prompt)
|
||||
if prepared != prompt or not targets:
|
||||
return False
|
||||
queue = getattr(prompt_server, "prompt_queue", None)
|
||||
for name in ("get_current_queue_volatile", "get_current_queue"):
|
||||
reader = getattr(queue, name, None)
|
||||
if not callable(reader):
|
||||
continue
|
||||
try:
|
||||
running, _pending = reader()
|
||||
for entry in running:
|
||||
if (isinstance(entry, (tuple, list)) and len(entry) > 4
|
||||
and entry[1] == prompt_id and entry[2] == prompt):
|
||||
roots = entry[4]
|
||||
return (isinstance(roots, (list, tuple))
|
||||
and all(isinstance(root, str) for root in roots)
|
||||
and len(roots) == len(targets) and set(roots) == set(targets))
|
||||
return False
|
||||
except Exception:
|
||||
continue
|
||||
return False
|
||||
|
||||
|
||||
def _queue_prompt_ids(prompt_server: Any) -> set[str]:
|
||||
"""Read running/pending prompt IDs, preferring ComfyUI's volatile view."""
|
||||
queue = getattr(prompt_server, "prompt_queue", None)
|
||||
|
||||
+191
-39
@@ -15,6 +15,15 @@ import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
from ..serverless_bridge.image_transfer import (
|
||||
FORMATS as IMAGE_TRANSFER_FORMATS, DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options,
|
||||
)
|
||||
except ImportError: # Standalone image_list.runtime (without extension bootstrap).
|
||||
from serverless_bridge.image_transfer import (
|
||||
FORMATS as IMAGE_TRANSFER_FORMATS, DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options,
|
||||
)
|
||||
|
||||
# v2: Start carries generic value_N slots instead of fixed reference/text/strength/seed fields.
|
||||
PROTOCOL_VERSION = 2
|
||||
START_CLASS = "DistributedImageListStart"
|
||||
@@ -257,6 +266,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")
|
||||
@@ -376,13 +387,19 @@ def delete_worker_item(token: str) -> bool:
|
||||
def create_master_job(start_id: str, end_id: str, prompt: dict[str, Any], items: list[dict[str, Any]],
|
||||
worker_ids: list[str], master_callback_url: str,
|
||||
workflow_meta: dict[str, Any] | None = None,
|
||||
prompt_id: str | None = None, *, master_participates: bool = False) -> str:
|
||||
prompt_id: str | None = None, *, master_participates: bool = False,
|
||||
explicit_targets: bool = False) -> str:
|
||||
if not isinstance(prompt, dict):
|
||||
raise ValueError("Image List Start requires the original ComfyUI prompt for region extraction")
|
||||
if not isinstance(master_participates, bool):
|
||||
raise ValueError("master_participates must be a boolean")
|
||||
if not worker_ids and not master_participates:
|
||||
raise RuntimeError("Distributed Image List requires at least one preflighted direct worker; master fallback is disabled")
|
||||
raise RuntimeError("Distributed Image List requires at least one selected worker; master fallback is disabled")
|
||||
# Only the target-controlled Start caller may relax the outer prompt guard.
|
||||
regions = validate_image_list_regions(prompt, explicit_targets=explicit_targets is True)
|
||||
if not any(region["start_id"] == str(start_id) and region["end_id"] == str(end_id)
|
||||
for region in regions):
|
||||
raise ValueError("Image-list region does not match the requested Start/End pair")
|
||||
context_id = secrets.token_urlsafe(36)
|
||||
state = ImageListJob(context_id, str(start_id), str(end_id), prompt, items,
|
||||
[str(worker_id) for worker_id in worker_ids], master_callback_url,
|
||||
@@ -497,15 +514,23 @@ def _safe_shared_dependency(node: dict[str, Any], prompt: dict[str, Any], regist
|
||||
return types <= scalar_types and not any(_prompt_links(node, prompt))
|
||||
|
||||
|
||||
def validate_image_list_regions(prompt: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
"""Validate pairing, non-nesting, and closed per-image graph boundaries."""
|
||||
def _is_output_node(node: dict[str, Any], registry) -> bool:
|
||||
node_class = registry.get(node.get("class_type"))
|
||||
if node_class is None:
|
||||
return False
|
||||
info = node_class.GET_NODE_INFO_V1() if hasattr(node_class, "GET_NODE_INFO_V1") else {}
|
||||
return bool(info.get("output_node", getattr(node_class, "OUTPUT_NODE", False)))
|
||||
|
||||
|
||||
def _image_list_pairs(prompt: dict[str, Any]) -> list[tuple[str, str]]:
|
||||
"""Pair boundaries before execution-copy pruning or full validation."""
|
||||
if not isinstance(prompt, dict):
|
||||
raise ValueError("Prompt must be an object")
|
||||
starts = sorted(str(k) for k, v in prompt.items() if isinstance(v, dict) and v.get("class_type") == START_CLASS)
|
||||
ends = sorted(str(k) for k, v in prompt.items() if isinstance(v, dict) and v.get("class_type") == END_CLASS)
|
||||
if len(starts) != len(ends) or not starts:
|
||||
raise ValueError("Distributed Image List needs one paired Start and End per region")
|
||||
regions = []
|
||||
pairs = []
|
||||
for start_id in starts:
|
||||
candidates = []
|
||||
for end_id in ends:
|
||||
@@ -514,7 +539,19 @@ def validate_image_list_regions(prompt: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
candidates.append(end_id)
|
||||
if len(candidates) != 1:
|
||||
raise ValueError(f"Start node {start_id} must connect its typed context output directly to exactly one End")
|
||||
end_id = candidates[0]
|
||||
pairs.append((start_id, candidates[0]))
|
||||
return pairs
|
||||
|
||||
|
||||
def validate_image_list_regions(prompt: dict[str, Any], *, explicit_targets: bool = False) -> list[dict[str, Any]]:
|
||||
"""Validate pairing, non-nesting, and closed per-image graph boundaries.
|
||||
|
||||
Required output-node utilities are allowed only when callers explicitly
|
||||
select execution roots, rather than letting every OUTPUT_NODE run independently.
|
||||
All data-boundary checks still apply in that mode.
|
||||
"""
|
||||
regions = []
|
||||
for start_id, end_id in _image_list_pairs(prompt):
|
||||
if "images" not in prompt[end_id].get("inputs", {}):
|
||||
raise ValueError(f"End node {end_id}.images must be connected to the per-item result branch")
|
||||
end_inputs = prompt[end_id].get("inputs", {})
|
||||
@@ -529,8 +566,6 @@ def validate_image_list_regions(prompt: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
{START_CLASS, END_CLASS, "DistributedCollector", "UltimateSDUpscaleDistributed"}}
|
||||
if forbidden:
|
||||
raise ValueError("Nested/ambiguous regions and other distributed nodes are unsupported in an image-list region: " + ", ".join(sorted(forbidden)))
|
||||
# Output nodes execute independently of the lazy End, so they would
|
||||
# evaluate the processing branch on the master as well as the worker.
|
||||
import sys
|
||||
registry = getattr(sys.modules.get("nodes"), "NODE_CLASS_MAPPINGS", {})
|
||||
shared = region_nodes & _start_ancestors(prompt, start_id)
|
||||
@@ -545,13 +580,10 @@ def validate_image_list_regions(prompt: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
f"Image-list region {start_id}->{end_id} bypasses Start through upstream nodes "
|
||||
f"{sorted(unsafe)} ({'; '.join(links)}); route per-item data through Start"
|
||||
)
|
||||
for node_id in region_nodes - {start_id, end_id}:
|
||||
node_class = registry.get(prompt[node_id].get("class_type"))
|
||||
if node_class is None:
|
||||
continue
|
||||
info = node_class.GET_NODE_INFO_V1() if hasattr(node_class, "GET_NODE_INFO_V1") else {}
|
||||
if info.get("output_node", getattr(node_class, "OUTPUT_NODE", False)):
|
||||
raise ValueError(f"Output node {node_id} must be placed after Distributed Image List End")
|
||||
if not explicit_targets:
|
||||
for node_id in region_nodes - {start_id, end_id}:
|
||||
if _is_output_node(prompt[node_id], registry):
|
||||
raise ValueError(f"Output node {node_id} must be placed after Distributed Image List End")
|
||||
downstream = _downstream_until(prompt, start_id, end_id)
|
||||
outside = sorted(downstream - region_nodes)
|
||||
if outside:
|
||||
@@ -713,10 +745,48 @@ def register_assignment(context_id: str, item_index: int, worker_id: str, token:
|
||||
state.expected_workers[item_index] = worker_id
|
||||
|
||||
|
||||
def _image_codec():
|
||||
# Exact returns and inputs must not import optional media codec dependencies.
|
||||
if __package__ == "image_list":
|
||||
from serverless_bridge import image_codec
|
||||
else:
|
||||
from ..serverless_bridge import image_codec
|
||||
return image_codec
|
||||
|
||||
|
||||
def _decode_return_image(data: Any) -> Any:
|
||||
if (not isinstance(data, dict) or set(data) != {"type", "encoding", "descriptor", "data"}
|
||||
or data["type"] != "IMAGE" or data["encoding"] != "image-transfer-v1"):
|
||||
raise ValueError("Invalid compressed IMAGE envelope")
|
||||
descriptor = data["descriptor"]
|
||||
if not isinstance(descriptor, dict) or descriptor.get("encoding") != "image-transfer-v1":
|
||||
raise ValueError("Invalid compressed IMAGE descriptor")
|
||||
shape = descriptor.get("shape")
|
||||
if (not isinstance(shape, list) or len(shape) != 4
|
||||
or any(type(dim) is not int or dim < 1 for dim in shape)
|
||||
or shape[0] != 1 or shape[-1] not in (3, 4)):
|
||||
raise ValueError("Compressed IMAGE must have shape [1,H,W,3|4]")
|
||||
encoded = data["data"]
|
||||
if not isinstance(encoded, str) or len(encoded) > ((MAX_PAYLOAD_BYTES + 2) // 3) * 4:
|
||||
raise ValueError("Compressed IMAGE data is missing or exceeds the transport limit")
|
||||
try:
|
||||
raw = base64.b64decode(encoded, validate=True)
|
||||
except (ValueError, TypeError) as exc:
|
||||
raise ValueError("Compressed IMAGE data must be valid base64") from exc
|
||||
if len(raw) > MAX_PAYLOAD_BYTES:
|
||||
raise ValueError("Compressed IMAGE exceeds the transport limit")
|
||||
return _image_codec().decode_image(descriptor, raw, max_bytes=MAX_PAYLOAD_BYTES)
|
||||
|
||||
|
||||
def decode_result(result: Any) -> dict[str, Any]:
|
||||
if not isinstance(result, dict):
|
||||
raise ValueError("result must be an object")
|
||||
image = decode_value(result.get("image"), "IMAGE")
|
||||
image_data = result.get("image")
|
||||
if isinstance(image_data, dict) and ("encoding" in image_data or "descriptor" in image_data):
|
||||
encode_request_payload(result) # Bound combined JSON/base64 before codec allocation.
|
||||
image = _decode_return_image(image_data)
|
||||
else:
|
||||
image = decode_value(image_data, "IMAGE")
|
||||
if image is None or image.ndim != 4 or image.shape[0] != 1 or image.shape[-1] not in (3, 4):
|
||||
raise ValueError("result.image must contain exactly one RGB/RGBA image")
|
||||
mask = decode_value(result.get("mask"), "MASK")
|
||||
@@ -752,11 +822,28 @@ def normalize_worker_result(images: Any, masks: Any = None, texts: Any = None) -
|
||||
return {"image": image, "mask": mask, "text": text}
|
||||
|
||||
|
||||
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 fmt != "exact":
|
||||
encode_request_payload(encoded)
|
||||
return encoded
|
||||
|
||||
|
||||
def _collect_output_lists(results: list[dict[str, Any]]) -> tuple[list[Any], list[Any], list[Any]]:
|
||||
@@ -773,7 +860,7 @@ def _collect_output_lists(results: list[dict[str, Any]]) -> tuple[list[Any], lis
|
||||
|
||||
|
||||
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."""
|
||||
"""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")
|
||||
@@ -799,33 +886,73 @@ async def run_image_list_job(context_id: str, *, model_callback=None) -> tuple[l
|
||||
remove_master_job(context_id)
|
||||
|
||||
|
||||
def _has_selected_serverless_workers(config, requested_ids):
|
||||
"""Keep direct/master-only execution free of optional cloud adapter imports.
|
||||
|
||||
This is only an import guard, not validation. The adapter validates original
|
||||
pool IDs/configuration without repairing them, before any cloud work starts.
|
||||
"""
|
||||
requested = None if requested_ids is None else {str(value) for value in requested_ids}
|
||||
return any(
|
||||
str(worker.get("type") or "local").strip().lower() == "serverless"
|
||||
and (worker.get("enabled", False) if requested is None
|
||||
else str(worker.get("id") or "").strip() in requested)
|
||||
for worker in config.get("workers", [])
|
||||
)
|
||||
|
||||
|
||||
async def _run_image_list_job(state: ImageListJob):
|
||||
context_id = state.context_id
|
||||
state.event = asyncio.Event()
|
||||
state.event_loop = asyncio.get_running_loop()
|
||||
from ..api.orchestration.prompt_transform import _resolve_direct_workers_for_image_list
|
||||
from ..utils.config import load_config
|
||||
|
||||
config = load_config()
|
||||
workers = _resolve_direct_workers_for_image_list(config, state.worker_ids)
|
||||
# Every item path selects End explicitly, never all OUTPUT_NODE roots.
|
||||
regions = validate_image_list_regions(state.original_prompt, explicit_targets=True)
|
||||
region = next((entry for entry in regions if entry["start_id"] == state.start_id and entry["end_id"] == state.end_id), None)
|
||||
if region is None:
|
||||
raise RuntimeError("Image-list region no longer matches the submitted prompt")
|
||||
controller = None
|
||||
try:
|
||||
if _has_selected_serverless_workers(config, state.worker_ids):
|
||||
from .serverless import create_serverless_slots
|
||||
controller = await create_serverless_slots(state, region, config, state.event)
|
||||
return await _run_image_list_slots(state, config, workers, region, controller)
|
||||
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):
|
||||
context_id = state.context_id
|
||||
from ..api.orchestration.dispatch import (
|
||||
preflight_image_list_workers, dispatch_worker_prompt,
|
||||
)
|
||||
from ..api.orchestration.prompt_transform import _resolve_direct_workers_for_image_list
|
||||
from ..utils.config import load_config, get_worker_timeout_seconds
|
||||
from ..utils.config import get_worker_timeout_seconds
|
||||
from ..utils.network import build_master_callback_url, build_worker_url, get_client_session
|
||||
from ..utils.logging import log
|
||||
import aiohttp
|
||||
|
||||
config = load_config()
|
||||
workers = _resolve_direct_workers_for_image_list(config, state.worker_ids)
|
||||
if not workers and not state.master_participates:
|
||||
raise RuntimeError("Distributed Image List requires enabled direct workers; serverless pools are unsupported")
|
||||
# Recheck worker capability and idle state immediately before the first item.
|
||||
worker_statuses = await preflight_image_list_workers(workers, require_idle=not state.master_participates) if workers else {}
|
||||
serverless_slots = list(controller.slots) if controller is not None else []
|
||||
end_inputs = state.original_prompt[state.end_id]["inputs"]
|
||||
transfer_format, _ = validate_options(
|
||||
end_inputs.get("transfer_format", DEFAULT_FORMAT),
|
||||
end_inputs.get("transfer_quality", DEFAULT_QUALITY),
|
||||
)
|
||||
transfer_requirements = ({"required_transfer_formats": (transfer_format,)}
|
||||
if transfer_format != "exact" else {})
|
||||
# Direct capability/idle probes must never be sent to a serverless pool.
|
||||
worker_statuses = await preflight_image_list_workers(
|
||||
workers, require_idle=not state.master_participates, **transfer_requirements,
|
||||
) if workers else {}
|
||||
workers = [worker for worker in workers if worker["id"] in worker_statuses]
|
||||
if not workers and not state.master_participates:
|
||||
raise RuntimeError("No free, reachable, feature-compatible direct workers are available; master execution is disabled")
|
||||
if not workers and not serverless_slots and not state.master_participates:
|
||||
raise RuntimeError("No reachable, feature-compatible image-list workers are available; master execution is disabled")
|
||||
session = await get_client_session() if workers else None
|
||||
regions = validate_image_list_regions(state.original_prompt)
|
||||
region = next((entry for entry in regions if entry["start_id"] == state.start_id and entry["end_id"] == state.end_id), None)
|
||||
if region is None:
|
||||
raise RuntimeError("Image-list region no longer matches the submitted prompt")
|
||||
timeout_seconds = max(5.0, float(get_worker_timeout_seconds()))
|
||||
use_websocket = bool((config.get("settings", {}) or {}).get("websocket_orchestration", False))
|
||||
prepared_worker_prompts: dict[str, dict[str, Any]] = {}
|
||||
@@ -840,6 +967,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()
|
||||
@@ -859,6 +988,13 @@ async def _run_image_list_job(state: ImageListJob):
|
||||
if worker is None:
|
||||
raise RuntimeError(f"Worker {worker_id} was removed before dispatch")
|
||||
assert session is not None
|
||||
if transfer_requirements:
|
||||
# Recheck on every compressed assignment, before staging or POSTing.
|
||||
statuses = await preflight_image_list_workers(
|
||||
[worker], require_idle=True, **transfer_requirements,
|
||||
)
|
||||
if worker_id not in statuses:
|
||||
raise RuntimeError(f"Image-list worker {worker_id} unavailable before compressed 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])}
|
||||
@@ -892,6 +1028,7 @@ async def _run_image_list_job(state: ImageListJob):
|
||||
worker, worker_prompt, state.workflow_meta,
|
||||
client_id=f"image-list-{context_id[:20]}", use_websocket=use_websocket,
|
||||
return_prompt_id=True,
|
||||
partial_execution_targets=[state.end_id],
|
||||
)
|
||||
if not isinstance(prompt_id, str) or not prompt_id:
|
||||
raise RuntimeError(f"Worker {worker_id} did not return a prompt_id; submission status is ambiguous")
|
||||
@@ -903,6 +1040,10 @@ async def _run_image_list_job(state: ImageListJob):
|
||||
state.event.clear()
|
||||
if state.failed:
|
||||
raise RuntimeError(state.failed)
|
||||
if controller is not None:
|
||||
completed = controller.poll(inflight)
|
||||
if completed is not None:
|
||||
return completed
|
||||
local = inflight.get(master_slot)
|
||||
if local is not None and local["future"].done():
|
||||
# A posted result does not release this slot until execution has
|
||||
@@ -952,10 +1093,15 @@ async def _run_image_list_job(state: ImageListJob):
|
||||
return None
|
||||
return None
|
||||
errors = await asyncio.gather(*(history_status(wid, handle) for wid, handle in inflight.items()
|
||||
if wid is not master_slot))
|
||||
if wid is not master_slot and wid not in serverless_slots))
|
||||
failure = next((error for error in errors if isinstance(error, str)), None)
|
||||
if failure:
|
||||
raise RuntimeError(failure)
|
||||
# A cloud manager task may finish during direct HTTP history waits.
|
||||
if controller is not None:
|
||||
completed = controller.poll(inflight)
|
||||
if completed is not None:
|
||||
return completed
|
||||
completed = next((entry for entry in errors if isinstance(entry, tuple)), None)
|
||||
if completed is not None:
|
||||
# The callback precedes the worker prompt finishing. Wait for its
|
||||
@@ -966,7 +1112,8 @@ async def _run_image_list_job(state: ImageListJob):
|
||||
return completed
|
||||
now = time.monotonic()
|
||||
expired = [f"{wid}/item {handle['index']}" for wid, handle in inflight.items()
|
||||
if wid is not master_slot and now - handle["last_seen"] > timeout_seconds]
|
||||
if wid is not master_slot and wid not in serverless_slots
|
||||
and now - handle["last_seen"] > timeout_seconds]
|
||||
if expired:
|
||||
raise TimeoutError(f"Image-list owned-prompt liveness timeout after {timeout_seconds:.1f}s: {', '.join(expired)}")
|
||||
try:
|
||||
@@ -975,6 +1122,10 @@ async def _run_image_list_job(state: ImageListJob):
|
||||
pass
|
||||
|
||||
async def cleanup(worker_id: Any, handle: dict[str, Any]) -> None:
|
||||
if worker_id in serverless_slots:
|
||||
# The controller owns exact manager records, including dispatches
|
||||
# that never returned a handle. Close it once in the outer finally.
|
||||
return
|
||||
if worker_id is master_slot:
|
||||
# The outer finally owns the non-abandonable stop/join. The generic
|
||||
# scheduler's bounded remote cleanup must never detach local work.
|
||||
@@ -1015,7 +1166,8 @@ 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)
|
||||
|
||||
@@ -0,0 +1,413 @@
|
||||
"""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_media_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")
|
||||
graph[end] = {"class_type": "DistributedMediaCollector", "inputs": {
|
||||
**{key: end_inputs[key] for key in ("images", "masks", "texts") if key in end_inputs},
|
||||
"multi_job_id": job_id, "worker_id": worker["id"]}}
|
||||
if transfer_format != DEFAULT_FORMAT:
|
||||
graph[end]["inputs"].update(transfer_format=transfer_format, transfer_quality=transfer_quality)
|
||||
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]
|
||||
workflow, manifest = compile_media_prompt(graph, request_id)
|
||||
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")
|
||||
result = runtime.normalize_worker_result(media["images"], media.get("masks"), media.get("texts"))
|
||||
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)
|
||||
+21
-5
@@ -6,6 +6,8 @@ import json
|
||||
from comfy_api.v0_0_2 import io
|
||||
|
||||
from ..utils.async_helpers import run_async_in_server_loop
|
||||
from ..serverless_bridge.image_transfer import (FORMATS, DEFAULT_FORMAT, DEFAULT_QUALITY,
|
||||
FORMAT_TOOLTIP, QUALITY_TOOLTIP, validate_options)
|
||||
from ..image_list.runtime import (
|
||||
MAX_VALUE_SLOTS,
|
||||
VALUE_NAMES,
|
||||
@@ -99,13 +101,18 @@ class DistributedImageListStart(io.ComfyNode):
|
||||
raise RuntimeError("Invalid preflighted image-list worker list") from exc
|
||||
if not isinstance(worker_ids, list) or not all(isinstance(item, str) for item in worker_ids):
|
||||
raise RuntimeError("Invalid preflighted image-list worker list")
|
||||
from comfy_execution.utils import get_executing_context
|
||||
from ..image_list.lifecycle import has_safe_execution_targets
|
||||
import server
|
||||
executing = get_executing_context()
|
||||
explicit_targets = has_safe_execution_targets(
|
||||
getattr(server.PromptServer, "instance", None),
|
||||
executing.prompt_id if executing is not None else None, prompt)
|
||||
from ..image_list.runtime import validate_image_list_regions
|
||||
regions = validate_image_list_regions(prompt)
|
||||
regions = validate_image_list_regions(prompt, explicit_targets=explicit_targets)
|
||||
region = next((row for row in regions if row["start_id"] == start_id), None)
|
||||
if region is None:
|
||||
raise RuntimeError(f"No safe paired Image List End was found for Start node {start_id}")
|
||||
from comfy_execution.utils import get_executing_context
|
||||
executing = get_executing_context()
|
||||
master_participates = unwrap_single(image_list_master_participates)
|
||||
if master_participates is None:
|
||||
from ..utils.config import load_config
|
||||
@@ -118,6 +125,7 @@ class DistributedImageListStart(io.ComfyNode):
|
||||
image_list_workflow_meta if isinstance(image_list_workflow_meta, dict) else None,
|
||||
prompt_id=executing.prompt_id if executing is not None else None,
|
||||
master_participates=master_participates,
|
||||
explicit_targets=explicit_targets,
|
||||
)
|
||||
return io.NodeOutput(
|
||||
[item["image"] for item in items],
|
||||
@@ -149,6 +157,10 @@ class DistributedImageListEnd(io.ComfyNode):
|
||||
io.Mask.Input("masks", optional=True, lazy=True),
|
||||
io.String.Input("texts", optional=True, lazy=True, force_input=True),
|
||||
io.Custom("DISTRIBUTED_IMAGE_LIST_CONTEXT").Input("context"),
|
||||
io.Combo.Input("transfer_format", options=list(FORMATS), default=DEFAULT_FORMAT,
|
||||
optional=True, tooltip=FORMAT_TOOLTIP),
|
||||
io.Int.Input("transfer_quality", default=DEFAULT_QUALITY, min=1, max=100,
|
||||
optional=True, tooltip=QUALITY_TOOLTIP),
|
||||
],
|
||||
outputs=[
|
||||
io.Image.Output(display_name="image_list", is_output_list=True),
|
||||
@@ -189,7 +201,10 @@ 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, transfer_format=DEFAULT_FORMAT,
|
||||
transfer_quality=DEFAULT_QUALITY):
|
||||
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)
|
||||
if worker_result_token:
|
||||
@@ -203,7 +218,8 @@ class DistributedImageListEnd(io.ComfyNode):
|
||||
set_master_result(context_values[0], 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,
|
||||
|
||||
@@ -271,6 +271,74 @@ def _load_worker_routes_module():
|
||||
worker_routes = _load_worker_routes_module()
|
||||
|
||||
|
||||
class WebsocketExecutionTargetTests(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self):
|
||||
from test_async_helpers import _load_async_helpers_module
|
||||
self.modules = patch.dict(sys.modules)
|
||||
self.modules.start()
|
||||
self.addCleanup(self.modules.stop)
|
||||
self.helper, self.server = _load_async_helpers_module()
|
||||
sys.modules['dist_api_worker_testpkg.utils.async_helpers'] = self.helper
|
||||
self.prompt = {'9': {'class_type': 'Output'}, '10': {'class_type': 'Output'}}
|
||||
|
||||
async def receive(self, payload):
|
||||
import json
|
||||
sent = []
|
||||
class Socket:
|
||||
async def prepare(self, request):
|
||||
pass
|
||||
async def send_json(self, value):
|
||||
sent.append(value)
|
||||
def __aiter__(self):
|
||||
async def messages():
|
||||
yield types.SimpleNamespace(type=worker_routes.aiohttp.WSMsgType.TEXT,
|
||||
data=json.dumps(payload))
|
||||
return messages()
|
||||
with patch.object(worker_routes, 'web', types.SimpleNamespace(WebSocketResponse=lambda **kw: Socket())), \
|
||||
patch.object(worker_routes, 'queue_prompt_payload', self.helper.queue_prompt_payload), \
|
||||
patch.object(worker_routes, 'PromptValidationError', self.helper.PromptValidationError):
|
||||
await worker_routes.worker_ws_endpoint(_FakeRequest())
|
||||
return sent[0]
|
||||
|
||||
def payload(self, targets):
|
||||
return {'type': 'dispatch_prompt', 'request_id': 'req', 'prompt': self.prompt,
|
||||
'partial_execution_targets': targets}
|
||||
|
||||
async def test_explicit_targets_reach_queue_and_acknowledgement(self):
|
||||
ack = await self.receive(self.payload(['9']))
|
||||
self.assertTrue(ack['ok'], ack)
|
||||
self.assertEqual(ack['request_id'], 'req')
|
||||
self.assertEqual(self.server.prompt_queue.items[0][4], ['9'])
|
||||
|
||||
async def test_malformed_targets_and_messages_are_rejected(self):
|
||||
for targets in (None, [], '9', [True], [9], ['absent'], ['9', '9']):
|
||||
with self.subTest(targets=targets):
|
||||
ack = await self.receive(self.payload(targets))
|
||||
self.assertFalse(ack['ok'])
|
||||
self.assertIn('partial_execution_targets', ack['error'])
|
||||
self.assertFalse((await self.receive([]))['ok'])
|
||||
self.assertEqual(self.server.prompt_queue.items, [])
|
||||
|
||||
async def test_validation_cannot_silently_queue_a_partial_subset(self):
|
||||
ack = await self.receive(self.payload(['9', '10']))
|
||||
self.assertFalse(ack['ok'])
|
||||
self.assertEqual(ack['validation_error']['type'], 'partial_execution_targets_failed_validation')
|
||||
self.assertEqual(self.server.prompt_queue.items, [])
|
||||
|
||||
async def test_missing_target_field_keeps_legacy_queue_call(self):
|
||||
from unittest.mock import AsyncMock
|
||||
queued = AsyncMock(return_value='legacy')
|
||||
original = self.helper.queue_prompt_payload
|
||||
self.helper.queue_prompt_payload = queued
|
||||
try:
|
||||
payload = self.payload(['9'])
|
||||
payload.pop('partial_execution_targets')
|
||||
self.assertTrue((await self.receive(payload))['ok'])
|
||||
self.assertNotIn('partial_execution_targets', queued.call_args.kwargs)
|
||||
finally:
|
||||
self.helper.queue_prompt_payload = original
|
||||
|
||||
|
||||
class WorkerRoutesTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_network_info_reports_actual_master_and_available_ports(self):
|
||||
with patch.object(worker_routes, "_get_cuda_info", return_value=(0, 1, 4)), \
|
||||
|
||||
@@ -1,11 +1,12 @@
|
||||
"""Model-free multi-process acceptance using real ComfyUI executors and loopback HTTP.
|
||||
|
||||
Only the transport test facade (/prompt, /queue, /history, /interrupt) is synthetic.
|
||||
Isolated listeners use real core /prompt, /queue, /history and /interrupt handlers.
|
||||
Start/End, graph extraction, scheduling, payload codecs and callbacks are production code.
|
||||
No managed ComfyUI service, user workflow, model, or remote worker is touched.
|
||||
"""
|
||||
import asyncio
|
||||
from collections import Counter
|
||||
import copy
|
||||
import importlib
|
||||
import json
|
||||
import os
|
||||
@@ -36,6 +37,7 @@ import aiohttp
|
||||
|
||||
torch.set_num_threads(1)
|
||||
SOURCE_CALLS, BRANCH_CALLS, CAPTURED, MODEL_CALLS = [], [], [], []
|
||||
UTILITY_CALLS, ITEM_EXECUTIONS = [], []
|
||||
ALLOW_MASTER = False
|
||||
COUNT = 7
|
||||
|
||||
@@ -85,6 +87,51 @@ class TestCachedModel:
|
||||
return (object(),)
|
||||
|
||||
|
||||
def assert_item_execution(prompt, index):
|
||||
"""Prove the operation belongs to an item, not merely a participating host."""
|
||||
from comfy_execution.utils import get_executing_context
|
||||
context = get_executing_context()
|
||||
assert context is not None
|
||||
assert prompt['2']['inputs'].get('worker_item_token'), 'per-item node ran in outer graph'
|
||||
running, _ = server.PromptServer.instance.prompt_queue.get_current_queue_volatile()
|
||||
assert len(running) == 1, running
|
||||
entry = running[0]
|
||||
if WORKER:
|
||||
assert entry[1] == context.prompt_id and entry[2] == prompt
|
||||
assert entry[4] == ['4'], entry[4]
|
||||
else:
|
||||
assert ALLOW_MASTER, 'per-item node ran on delegate-only master'
|
||||
assert entry[1] != context.prompt_id, 'per-item node ran in outer executor'
|
||||
assert set(entry[4]) == {'4', '5'}, entry[4]
|
||||
return {'index': index, 'prompt_id': context.prompt_id,
|
||||
'role': 'worker' if WORKER else 'master-item', 'queue_prompt_id': entry[1]}
|
||||
|
||||
|
||||
class TestRequiredOutputUtility:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'image': ('IMAGE',), 'index': ('INT', {'forceInput': True})},
|
||||
'hidden': {'prompt': 'PROMPT'}}
|
||||
OUTPUT_NODE = True
|
||||
RETURN_TYPES = ('IMAGE',)
|
||||
FUNCTION = 'run'
|
||||
def run(self, image, index, prompt):
|
||||
assert_item_execution(prompt, index)
|
||||
UTILITY_CALLS.append(index)
|
||||
return (image,)
|
||||
|
||||
|
||||
class TestDiscardedPreview:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'image': ('IMAGE',)}}
|
||||
OUTPUT_NODE = True
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = 'run'
|
||||
def run(self, image):
|
||||
raise AssertionError('discarded per-item output preview executed')
|
||||
|
||||
|
||||
class TestImageOperation:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
@@ -92,11 +139,13 @@ class TestImageOperation:
|
||||
'text': ('STRING', {'forceInput': True}),
|
||||
'strength': ('FLOAT', {'forceInput': True}),
|
||||
'seed': ('INT', {'forceInput': True}), 'index': ('INT', {'forceInput': True}),
|
||||
'flag': ('BOOLEAN', {'forceInput': True}), 'model': ('MODEL',)}}
|
||||
'flag': ('BOOLEAN', {'forceInput': True}), 'model': ('MODEL',)},
|
||||
'hidden': {'prompt': 'PROMPT'}}
|
||||
RETURN_TYPES = ('IMAGE', 'MASK', 'STRING')
|
||||
FUNCTION = 'run'
|
||||
def run(self, image, reference_image, mask, text, strength, seed, index, flag, model):
|
||||
def run(self, image, reference_image, mask, text, strength, seed, index, flag, model, prompt):
|
||||
assert WORKER or ALLOW_MASTER, 'per-image branch executed in orchestrator-only mode'
|
||||
ITEM_EXECUTIONS.append(assert_item_execution(prompt, index))
|
||||
BRANCH_CALLS.append(index)
|
||||
# Sleep in short chunks so the real Comfy interrupt signal can be tested.
|
||||
until = time.monotonic() + DELAY
|
||||
@@ -120,6 +169,7 @@ class TestCapture:
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = 'run'
|
||||
def run(self, images, masks, texts):
|
||||
assert not WORKER, 'post-End capture leaked into a worker prompt'
|
||||
CAPTURED.append((images, masks, texts))
|
||||
return ()
|
||||
|
||||
@@ -131,7 +181,9 @@ async def setup():
|
||||
module = sys.modules[str(ROOT).replace('.', '_x_')]
|
||||
nodes.NODE_CLASS_MAPPINGS.update(TestImageSource=TestImageSource,
|
||||
TestImageOperation=TestImageOperation, TestCapture=TestCapture,
|
||||
TestCachedModel=TestCachedModel, TestSharedFlag=TestSharedFlag)
|
||||
TestCachedModel=TestCachedModel, TestSharedFlag=TestSharedFlag,
|
||||
TestRequiredOutputUtility=TestRequiredOutputUtility,
|
||||
TestDiscardedPreview=TestDiscardedPreview)
|
||||
# Match the CORS-enabled direct-worker launch contract for canvas probes.
|
||||
app = web.Application(middlewares=[server.create_cors_middleware('*')])
|
||||
# Register the actual feature routes, not test replacements.
|
||||
@@ -147,61 +199,56 @@ async def listen(app):
|
||||
return runner, site._server.sockets[0].getsockname()[1]
|
||||
|
||||
|
||||
def finish_queue_item(ps, item_id, executor):
|
||||
ps.prompt_queue.task_done(item_id, executor.history_result,
|
||||
status=execution.PromptQueue.ExecutionStatus(
|
||||
status_str='success' if executor.success else 'error',
|
||||
completed=executor.success, messages=executor.status_messages),
|
||||
process_item=lambda item: item[:5] + item[6:])
|
||||
|
||||
|
||||
async def worker_main():
|
||||
ps, assets, module, app = await setup()
|
||||
pending, running, history = [], {}, {}
|
||||
wake = asyncio.Event()
|
||||
interrupted_ids = []
|
||||
interrupted_ids, submissions = [], []
|
||||
core_handlers = {(route.method, route.path): route.handler for route in ps.routes}
|
||||
async def prompt(request):
|
||||
data = await request.json()
|
||||
prompt_id = str(uuid.uuid4())
|
||||
graph = data['prompt']
|
||||
valid = await execution.validate_prompt(prompt_id, graph, None)
|
||||
if not valid[0]:
|
||||
return web.json_response({'error': valid[1], 'node_errors': valid[3]}, status=400)
|
||||
pending.append((prompt_id, graph, valid[2]))
|
||||
wake.set()
|
||||
return web.json_response({'prompt_id': prompt_id})
|
||||
async def queue(request):
|
||||
if request.method == 'POST':
|
||||
deleted = (await request.json()).get('delete', [])
|
||||
pending[:] = [entry for entry in pending if entry[0] not in deleted]
|
||||
return web.json_response({})
|
||||
return web.json_response({'queue_running': [[0, key] for key in running],
|
||||
'queue_pending': [[0, entry[0]] for entry in pending]})
|
||||
async def get_history(request):
|
||||
pid = request.match_info['prompt_id']
|
||||
return web.json_response({pid: history[pid]} if pid in history else {})
|
||||
# Use core's actual validation/enqueue/response schema, including
|
||||
# partial_execution_targets; only observe it and wake our CPU consumer.
|
||||
response = await core_handlers['POST', '/prompt'](request)
|
||||
if response.status == 200:
|
||||
accepted = json.loads(response.body)
|
||||
_, pending = ps.prompt_queue.get_current_queue_volatile()
|
||||
entry = next(item for item in pending if item[1] == accepted['prompt_id'])
|
||||
submissions.append({'prompt_id': entry[1], 'requested': data.get('partial_execution_targets'),
|
||||
'queued': entry[4], 'node_ids': sorted(entry[2]),
|
||||
'end_ids': [key for key, node in entry[2].items()
|
||||
if node['class_type'] == 'DistributedImageListEnd']})
|
||||
wake.set()
|
||||
return response
|
||||
async def interrupt(request):
|
||||
pid = (await request.json())['prompt_id']
|
||||
interrupted_ids.append(pid)
|
||||
if pid in running:
|
||||
comfy.model_management.interrupt_current_processing(True)
|
||||
return web.json_response({})
|
||||
interrupted_ids.append((await request.json())['prompt_id'])
|
||||
return await core_handlers['POST', '/interrupt'](request)
|
||||
async def stats(request):
|
||||
return web.json_response({'calls': BRANCH_CALLS, 'running': list(running),
|
||||
'pending': [e[0] for e in pending], 'interrupted': interrupted_ids,
|
||||
'history': history})
|
||||
running, pending = ps.prompt_queue.get_current_queue_volatile()
|
||||
return web.json_response({'calls': BRANCH_CALLS, 'utility_calls': UTILITY_CALLS,
|
||||
'item_executions': ITEM_EXECUTIONS, 'submissions': submissions,
|
||||
'running': [e[1] for e in running], 'pending': [e[1] for e in pending],
|
||||
'interrupted': interrupted_ids, 'history': ps.prompt_queue.get_history()})
|
||||
async def consume():
|
||||
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
|
||||
while True:
|
||||
await wake.wait()
|
||||
while pending:
|
||||
pid, graph, outputs = pending.pop(0)
|
||||
running[pid] = True
|
||||
await asyncio.to_thread(executor.execute, graph, pid, {}, outputs)
|
||||
history[pid] = {'status': {'completed': True,
|
||||
'status_str': 'success' if executor.success else 'error',
|
||||
'messages': str(executor.status_messages)}}
|
||||
running.pop(pid, None)
|
||||
while (queued := ps.prompt_queue.get(timeout=0)) is not None:
|
||||
entry, item_id = queued
|
||||
await asyncio.to_thread(executor.execute, entry[2], entry[1], entry[3], entry[4])
|
||||
finish_queue_item(ps, item_id, executor)
|
||||
wake.clear()
|
||||
async def probe(_request):
|
||||
return web.json_response({'exec_info': {'queue_remaining': len(pending) + len(running)}})
|
||||
app.router.add_get('/prompt', probe)
|
||||
app.add_routes([route for route in ps.routes if
|
||||
(route.path in ('/queue', '/history/{prompt_id}') or
|
||||
(route.path == '/prompt' and route.method == 'GET'))])
|
||||
app.router.add_post('/prompt', prompt)
|
||||
app.router.add_get('/queue', queue)
|
||||
app.router.add_post('/queue', queue)
|
||||
app.router.add_get('/history/{prompt_id}', get_history)
|
||||
app.router.add_post('/interrupt', interrupt)
|
||||
app.router.add_get('/test-stats', stats)
|
||||
runner, port = await listen(app)
|
||||
@@ -228,12 +275,14 @@ def graph(callback, worker_ids):
|
||||
'3': {'class_type': 'TestImageOperation', 'inputs': {
|
||||
key: ['2', index] for key, index in [('image', 0), ('mask', 1), ('index', 2),
|
||||
('reference_image', 4), ('text', 5), ('strength', 6), ('seed', 7), ('flag', 8)]}
|
||||
| {'model': ['6', 0]}},
|
||||
| {'image': ['8', 0], 'model': ['6', 0]}},
|
||||
'4': {'class_type': 'DistributedImageListEnd', 'inputs': {
|
||||
'images': ['3', 0], 'masks': ['3', 1], 'texts': ['3', 2], 'context': ['2', 3]}},
|
||||
'5': {'class_type': 'TestCapture', 'inputs': {'images': ['4', 0], 'masks': ['4', 1], 'texts': ['4', 2]}},
|
||||
'6': {'class_type': 'TestCachedModel', 'inputs': {}},
|
||||
'7': {'class_type': 'TestSharedFlag', 'inputs': {}},
|
||||
'8': {'class_type': 'TestRequiredOutputUtility', 'inputs': {'image': ['2', 0], 'index': ['2', 2]}},
|
||||
'9': {'class_type': 'TestDiscardedPreview', 'inputs': {'image': ['8', 0]}},
|
||||
}
|
||||
|
||||
|
||||
@@ -241,6 +290,8 @@ async def coordinator():
|
||||
global ALLOW_MASTER
|
||||
ps, assets, module, app = await setup()
|
||||
runtime = importlib.import_module(module.__name__ + '.image_list.runtime')
|
||||
execution_plan = importlib.import_module(module.__name__ + '.image_list.execution_plan')
|
||||
lifecycle = importlib.import_module(module.__name__ + '.image_list.lifecycle')
|
||||
config_mod = importlib.import_module(module.__name__ + '.utils.config')
|
||||
orchestration = importlib.import_module(module.__name__ + '.api.queue_orchestration')
|
||||
runner, port = await listen(app)
|
||||
@@ -294,16 +345,49 @@ async def coordinator():
|
||||
prompt = graph(f'http://127.0.0.1:{port}', [w['id'] for w in workers])
|
||||
valid = await execution.validate_prompt('network-acceptance', prompt, None)
|
||||
assert valid[0], valid
|
||||
assert {'4', '5', '8', '9'} == set(valid[2]), valid
|
||||
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
|
||||
# Exercise production queue preparation; only master enqueue is redirected
|
||||
# to this isolated executor, not the managed server's live queue.
|
||||
prompt['2']['inputs'].pop('image_list_enabled_worker_ids')
|
||||
prompt['2']['inputs'].pop('image_list_master_callback')
|
||||
saved_prompt = copy.deepcopy(prompt)
|
||||
outer_executions = []
|
||||
async def enqueue(prepared, workflow, client_id, **kwargs):
|
||||
checked = await execution.validate_prompt('network-acceptance', prepared, None)
|
||||
targets = kwargs['partial_execution_targets']
|
||||
planned, planned_targets = execution_plan.prepare_image_list_execution(prepared)
|
||||
assert planned == prepared and targets == planned_targets
|
||||
if '9' in saved_prompt and 'TestImageOperation' in {
|
||||
node['class_type'] for node in prepared.values()}:
|
||||
assert set(prepared) == set(saved_prompt) - {'9'}, prepared
|
||||
assert prepared['8'] == saved_prompt['8'], 'required utility changed'
|
||||
assert prepared['5'] == saved_prompt['5'], 'post-End output changed'
|
||||
assert set(targets) == {'4', '5'}, targets
|
||||
assert nodes.NODE_CLASS_MAPPINGS['TestRequiredOutputUtility'].OUTPUT_NODE is True
|
||||
assert nodes.NODE_CLASS_MAPPINGS['TestDiscardedPreview'].OUTPUT_NODE is True
|
||||
pid = 'network-acceptance' if not outer_executions else str(uuid.uuid4())
|
||||
checked = await execution.validate_prompt(pid, prepared, targets)
|
||||
assert checked[0], checked
|
||||
await asyncio.to_thread(executor.execute, prepared, 'network-acceptance', {}, checked[2])
|
||||
return {'prompt_id': 'network-acceptance', 'number': 1, 'node_errors': {}}
|
||||
assert not checked[3] and Counter(checked[2]) == Counter(targets), checked
|
||||
assert not lifecycle.has_safe_execution_targets(ps, pid, prepared)
|
||||
# Exactly the core queue tuple/put/get/task_done lifecycle. Start
|
||||
# must prove its roots from this real running record, not a mock.
|
||||
number = len(outer_executions)
|
||||
ps.prompt_queue.put((number, pid, prepared, {'client_id': client_id}, checked[2], {}))
|
||||
entry, item_id = ps.prompt_queue.get(timeout=0)
|
||||
assert lifecycle.has_safe_execution_targets(ps, pid, prepared)
|
||||
before_captures = len(CAPTURED)
|
||||
await asyncio.to_thread(executor.execute, entry[2], entry[1], entry[3], entry[4])
|
||||
finish_queue_item(ps, item_id, executor)
|
||||
assert ps.prompt_queue.get_current_queue_volatile() == ([], [])
|
||||
history = ps.prompt_queue.get_history(pid)[pid]
|
||||
assert Counter(history['prompt'][4]) == Counter(targets)
|
||||
assert not lifecycle.has_safe_execution_targets(ps, pid, prepared)
|
||||
outer_executions.append({'prompt_id': pid, 'targets': targets,
|
||||
'captures': len(CAPTURED) - before_captures})
|
||||
if executor.success:
|
||||
assert len(CAPTURED) == before_captures + 1, 'post-End Capture must run exactly once'
|
||||
return {'prompt_id': pid, 'number': number, 'node_errors': checked[3]}
|
||||
with patch.object(config_mod, 'load_config', return_value=config), \
|
||||
patch.object(orchestration, 'load_config', return_value=config), \
|
||||
patch.object(orchestration, 'queue_prompt_payload', enqueue), \
|
||||
@@ -312,11 +396,12 @@ async def coordinator():
|
||||
prompt, None, 'test', [w['id'] for w in workers]), 45)
|
||||
assert queued[0] == 'network-acceptance' and queued[2] == 2, queued
|
||||
assert not ps.distributed_pending_jobs, 'image-list created unused Collector queues'
|
||||
assert prompt == saved_prompt, 'queue preparation changed the saved/source graph'
|
||||
if not executor.success:
|
||||
for err in stderr_files:
|
||||
err.flush(); err.seek(0); print(err.read(), file=sys.stderr)
|
||||
assert executor.success, executor.status_messages
|
||||
assert len(SOURCE_CALLS) == 1 and not BRANCH_CALLS
|
||||
assert len(SOURCE_CALLS) == 1 and not BRANCH_CALLS and not UTILITY_CALLS
|
||||
assert len(CAPTURED) == 1
|
||||
images, masks, texts = CAPTURED[0]
|
||||
assert len(images) == len(masks) == len(texts) == COUNT
|
||||
@@ -332,6 +417,14 @@ async def coordinator():
|
||||
counts = [len(row['calls']) for row in stats]
|
||||
assert counts[0] > counts[1] > 0, counts
|
||||
assert Counter(i for row in stats for i in row['calls']) == Counter(range(COUNT))
|
||||
for row in stats:
|
||||
assert Counter(row['utility_calls']) == Counter(row['calls']), row
|
||||
assert len(row['submissions']) == len(row['calls'])
|
||||
assert len({event['prompt_id'] for event in row['item_executions']}) == len(row['calls'])
|
||||
for submission in row['submissions']:
|
||||
assert submission['requested'] == submission['queued'] == ['4'], submission
|
||||
assert '8' in submission['node_ids'] and '9' not in submission['node_ids']
|
||||
assert '5' not in submission['node_ids']
|
||||
assert not runtime._MASTER_JOBS
|
||||
# Three sequential regions on one worker, including optional outputs
|
||||
# left unconnected. Region N+1 must not race the prior worker End.
|
||||
@@ -346,12 +439,14 @@ async def coordinator():
|
||||
previous = end
|
||||
chained['8'] = {'class_type': 'TestCapture', 'inputs': {
|
||||
'images': ['7', 0], 'masks': ['7', 1], 'texts': ['7', 2]}}
|
||||
saved_chained = copy.deepcopy(chained)
|
||||
with patch.object(config_mod, 'load_config', return_value=config), \
|
||||
patch.object(orchestration, 'load_config', return_value=config), \
|
||||
patch.object(orchestration, 'queue_prompt_payload', enqueue), \
|
||||
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
|
||||
await asyncio.wait_for(orchestration.orchestrate_distributed_execution(
|
||||
chained, None, 'test', ['slow']), 45)
|
||||
assert chained == saved_chained
|
||||
assert executor.success, executor.status_messages
|
||||
assert len(CAPTURED) == 2 and len(CAPTURED[1][0]) == 2
|
||||
assert all(torch.equal(image, torch.zeros((1, 8, 8, 3))) for image in CAPTURED[1][0])
|
||||
@@ -363,34 +458,55 @@ 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)
|
||||
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
|
||||
before_sources, before_local, before_models = len(SOURCE_CALLS), len(BRANCH_CALLS), len(MODEL_CALLS)
|
||||
before_utilities, before_events = len(UTILITY_CALLS), len(ITEM_EXECUTIONS)
|
||||
async with aiohttp.ClientSession() as session:
|
||||
before_worker_calls = []
|
||||
before_worker_stats = []
|
||||
for worker in workers:
|
||||
async with session.get(f"http://127.0.0.1:{worker['port']}/test-stats") as response:
|
||||
before_worker_calls.append(len((await response.json())['calls']))
|
||||
before_worker_stats.append(await response.json())
|
||||
with patch.object(config_mod, 'load_config', return_value=config), \
|
||||
patch.object(orchestration, 'load_config', return_value=config), \
|
||||
patch.object(orchestration, 'queue_prompt_payload', enqueue), \
|
||||
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
|
||||
queued = await asyncio.wait_for(orchestration.orchestrate_distributed_execution(
|
||||
prompt, None, 'test', selected), 45)
|
||||
case_prompt, None, 'test', selected), 45)
|
||||
assert case_prompt == saved_case and prompt == saved_prompt, (mode, 'saved/source graph changed')
|
||||
assert executor.success, (mode, executor.status_messages)
|
||||
assert queued[2] == len(selected)
|
||||
assert len(SOURCE_CALLS) == before_sources + 1, 'source was reevaluated per local item'
|
||||
local_calls = BRANCH_CALLS[before_local:]
|
||||
assert Counter(UTILITY_CALLS[before_utilities:]) == Counter(local_calls), mode
|
||||
assert len({event['prompt_id'] for event in ITEM_EXECUTIONS[before_events:]}) == len(local_calls)
|
||||
assert local_calls, f'{mode}: master did not process any items'
|
||||
assert len(MODEL_CALLS) == before_models + 1, 'local model dependency did not stay cached'
|
||||
remote_calls = []
|
||||
async with aiohttp.ClientSession() as session:
|
||||
for worker, offset in zip(workers, before_worker_calls):
|
||||
for worker, before_stats in zip(workers, before_worker_stats):
|
||||
async with session.get(f"http://127.0.0.1:{worker['port']}/test-stats") as response:
|
||||
current = await response.json()
|
||||
remote_calls.extend(current['calls'][offset:])
|
||||
calls = current['calls'][len(before_stats['calls']):]
|
||||
utilities = current['utility_calls'][len(before_stats['utility_calls']):]
|
||||
submissions = current['submissions'][len(before_stats['submissions']):]
|
||||
events = current['item_executions'][len(before_stats['item_executions']):]
|
||||
assert Counter(utilities) == Counter(calls), (mode, current)
|
||||
assert len(submissions) == len(calls) == len({event['prompt_id'] for event in events})
|
||||
for submission in submissions:
|
||||
assert submission['requested'] == submission['queued'] == ['4'], submission
|
||||
assert '8' in submission['node_ids'] and '9' not in submission['node_ids']
|
||||
assert '5' not in submission['node_ids']
|
||||
remote_calls.extend(calls)
|
||||
assert not current['running'] and not current['pending'], current
|
||||
if mode == 'mixed':
|
||||
if mode.startswith('mixed'):
|
||||
assert remote_calls, 'mixed execution failed to use direct workers'
|
||||
else:
|
||||
assert not remote_calls and len(local_calls) == COUNT
|
||||
@@ -398,14 +514,30 @@ async def coordinator():
|
||||
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_calls:
|
||||
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)}
|
||||
print('IMAGE_LIST_MASTER_OK ' + json.dumps(participation_counts))
|
||||
assert len(outer_executions) == 2 + 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_TARGETS_OK ' + json.dumps({
|
||||
'modes': ['delegate-only', 'mixed', 'master-only'],
|
||||
'outer_targets': outer_executions[0]['targets'], 'worker_targets': ['4'],
|
||||
'required_utility_once_per_item': True, 'preview_discarded': True,
|
||||
'saved_prompt_unchanged': prompt == saved_prompt, 'capture_once_per_outer_prompt': True}))
|
||||
ALLOW_MASTER = False
|
||||
config['settings']['master_delegate_only'] = True
|
||||
# Preserve the original direct-runtime cancellation/failure fixture:
|
||||
# it has no outer queue and must not claim target-control evidence.
|
||||
direct_prompt = copy.deepcopy(prompt)
|
||||
del direct_prompt['8'], direct_prompt['9']
|
||||
direct_prompt['3']['inputs']['image'] = ['2', 0]
|
||||
# Real Comfy user-interrupt flag while a real worker executor is active.
|
||||
slow = workers[1]
|
||||
async with aiohttp.ClientSession() as session:
|
||||
@@ -420,7 +552,7 @@ async def coordinator():
|
||||
await asyncio.sleep(0.02)
|
||||
raise AssertionError('worker state transition timed out')
|
||||
before = await until(lambda state: not state['running'] and not state['pending'])
|
||||
context = runtime.create_master_job('2', '4', prompt, source_rows(), ['slow'],
|
||||
context = runtime.create_master_job('2', '4', direct_prompt, source_rows(), ['slow'],
|
||||
f'http://127.0.0.1:{port}')
|
||||
with patch.object(config_mod, 'load_config', return_value=config), \
|
||||
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
|
||||
@@ -444,7 +576,7 @@ async def coordinator():
|
||||
# A worker exception before End must fail, not hang or return a partial list.
|
||||
bad_rows = source_rows()
|
||||
bad_rows[0]['seed'] = -1
|
||||
context = runtime.create_master_job('2', '4', prompt, bad_rows, ['fast'],
|
||||
context = runtime.create_master_job('2', '4', direct_prompt, bad_rows, ['fast'],
|
||||
f'http://127.0.0.1:{port}')
|
||||
with patch.object(config_mod, 'load_config', return_value=config), \
|
||||
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
|
||||
|
||||
@@ -3,6 +3,7 @@ import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
|
||||
class _PromptQueue:
|
||||
@@ -87,5 +88,59 @@ class QueuePromptPayloadTests(unittest.IsolatedAsyncioTestCase):
|
||||
self.assertEqual(extra_data["extra_pnginfo"]["workflow"], {"id": "workflow-1"})
|
||||
|
||||
|
||||
class ExplicitTargetTests(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self):
|
||||
self.module, self.server = _load_async_helpers_module()
|
||||
self.prompt = {"9": {"class_type": "Output"}, "10": {"class_type": "Output"}}
|
||||
self.validate = AsyncMock(return_value=(True, None, ["9"], {}))
|
||||
self.module.execution.validate_prompt = self.validate
|
||||
|
||||
async def test_targets_survive_callback_omission_and_are_queued_exactly(self):
|
||||
seen = []
|
||||
def callback(payload):
|
||||
seen.append(payload["partial_execution_targets"][:])
|
||||
payload["partial_execution_targets"].append("10")
|
||||
return {"prompt": payload["prompt"]}
|
||||
self.server.trigger_on_prompt = callback
|
||||
targets = ["9"]
|
||||
await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=targets)
|
||||
self.assertEqual(seen, [["9"]])
|
||||
self.assertEqual(targets, ["9"])
|
||||
self.assertEqual(self.validate.call_args.args[2], ["9"])
|
||||
self.assertEqual(self.server.prompt_queue.items[0][4], ["9"])
|
||||
|
||||
async def test_invalid_requests_never_reach_validation_or_queue(self):
|
||||
for targets in ([], "9", [9], [True], ["9", "9"], ["missing"], {}):
|
||||
with self.subTest(targets=targets):
|
||||
with self.assertRaises(self.module.PromptValidationError):
|
||||
await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=targets)
|
||||
self.validate.assert_not_awaited()
|
||||
self.assertEqual(self.server.prompt_queue.items, [])
|
||||
|
||||
async def test_removed_target_after_callback_is_rejected(self):
|
||||
self.server.trigger_on_prompt = lambda payload: {"prompt": {"10": self.prompt["10"]}}
|
||||
with self.assertRaises(self.module.PromptValidationError):
|
||||
await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=["9"])
|
||||
self.validate.assert_not_awaited()
|
||||
self.assertEqual(self.server.prompt_queue.items, [])
|
||||
|
||||
async def test_partial_subset_extra_roots_and_target_errors_fail_closed(self):
|
||||
for roots, errors in ((["9"], {}), (["9", "10", "extra"], {}),
|
||||
(["9", "10"], {"9": {"errors": [{"message": "bad"}]}}),
|
||||
(["9", "10"], {"dep": {"dependent_outputs": ["10"], "errors": []}})):
|
||||
with self.subTest(roots=roots, errors=errors):
|
||||
self.validate.return_value = (True, None, roots, errors)
|
||||
with self.assertRaises(self.module.PromptValidationError):
|
||||
await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=["9", "10"])
|
||||
self.assertEqual(self.server.prompt_queue.items, [])
|
||||
|
||||
async def test_none_preserves_legacy_payload_and_validation(self):
|
||||
seen = []
|
||||
self.server.trigger_on_prompt = lambda payload: seen.append(payload.copy()) or payload
|
||||
await self.module.queue_prompt_payload(self.prompt)
|
||||
self.assertNotIn("partial_execution_targets", seen[0])
|
||||
self.assertIsNone(self.validate.call_args.args[2])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -4,7 +4,7 @@ import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
|
||||
def _load_dispatch_module():
|
||||
@@ -168,6 +168,7 @@ class _FakeSession:
|
||||
def _valid_capability():
|
||||
return {
|
||||
"supported": True,
|
||||
"execution_targets": True,
|
||||
"protocol_version": 2,
|
||||
"node_ids": ["DistributedImageListStart", "DistributedImageListEnd"],
|
||||
}
|
||||
@@ -177,6 +178,43 @@ def _valid_queue(running=None, pending=None):
|
||||
return {"queue_running": running or [], "queue_pending": pending or []}
|
||||
|
||||
|
||||
class ExplicitTargetTransportTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_http_preserves_targets_and_omits_default_field(self):
|
||||
response = types.SimpleNamespace(raise_for_status=lambda: None)
|
||||
class Response:
|
||||
async def __aenter__(self): return response
|
||||
async def __aexit__(self, *args): pass
|
||||
calls = []
|
||||
def post(url, **kwargs):
|
||||
calls.append(kwargs['json'])
|
||||
return Response()
|
||||
session = types.SimpleNamespace(post=post)
|
||||
with patch.object(dispatch, 'get_client_session', AsyncMock(return_value=session)):
|
||||
await dispatch.dispatch_worker_prompt({'id': 'worker'}, {'9': {}}, None,
|
||||
partial_execution_targets=['9'])
|
||||
await dispatch.dispatch_worker_prompt({'id': 'worker'}, {'9': {}}, None)
|
||||
self.assertEqual(calls[0]['partial_execution_targets'], ['9'])
|
||||
self.assertNotIn('partial_execution_targets', calls[1])
|
||||
|
||||
async def test_websocket_sends_targets_all_the_way_to_wire(self):
|
||||
import json
|
||||
sent = []
|
||||
async def receive():
|
||||
return types.SimpleNamespace(type=dispatch.aiohttp.WSMsgType.TEXT, data=json.dumps({
|
||||
'type': 'dispatch_ack', 'request_id': sent[-1]['request_id'],
|
||||
'ok': True, 'prompt_id': 'worker-prompt'}))
|
||||
ws = types.SimpleNamespace(send_json=AsyncMock(side_effect=sent.append),
|
||||
receive=receive, close=AsyncMock())
|
||||
session = types.SimpleNamespace(ws_connect=AsyncMock(return_value=ws))
|
||||
with patch.object(dispatch, 'get_client_session', AsyncMock(return_value=session)):
|
||||
result = await dispatch.dispatch_worker_prompt({'id': 'worker'}, {'9': {}}, None,
|
||||
use_websocket=True, return_prompt_id=True, partial_execution_targets=['9'])
|
||||
await dispatch.dispatch_worker_prompt({'id': 'worker'}, {'9': {}}, None, use_websocket=True)
|
||||
self.assertEqual(result, 'worker-prompt')
|
||||
self.assertEqual(sent[0]['partial_execution_targets'], ['9'])
|
||||
self.assertNotIn('partial_execution_targets', sent[1])
|
||||
|
||||
|
||||
class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_select_active_workers_filters_offline(self):
|
||||
workers = [
|
||||
@@ -417,6 +455,23 @@ class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
|
||||
with self.assertRaisesRegex(RuntimeError, "No selected feature-compatible direct worker is idle"):
|
||||
await dispatch.preflight_image_list_workers(workers, require_idle=True)
|
||||
|
||||
async def test_image_list_preflight_requires_exact_execution_targets_support(self):
|
||||
for value in (None, False, 1, "true", []):
|
||||
with self.subTest(value=value):
|
||||
capability = _valid_capability()
|
||||
if value is None:
|
||||
capability.pop("execution_targets")
|
||||
else:
|
||||
capability["execution_targets"] = value
|
||||
session = _FakeSession({("worker", "capabilities"): (200, capability)})
|
||||
async def fake_session():
|
||||
return session
|
||||
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
|
||||
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
|
||||
):
|
||||
with self.assertRaisesRegex(RuntimeError, "execution_targets.*upgrade"):
|
||||
await dispatch.preflight_image_list_workers([{"id": "worker"}])
|
||||
|
||||
async def test_image_list_preflight_fails_closed_for_malformed_capabilities(self):
|
||||
malformed = {"supported": True, "protocol_version": 2, "node_ids": None}
|
||||
session = _FakeSession({("worker", "capabilities"): (200, malformed)})
|
||||
|
||||
@@ -202,8 +202,10 @@ class ImageListGraphBoundaryTests(unittest.TestCase):
|
||||
prompt["2"] = {"class_type": "VAEDecode", "inputs": {"samples": ["7", 0]}}
|
||||
prompt["7"] = {"class_type": "KSampler", "inputs": {"model": ["1", 0]}}
|
||||
prompt["4"]["inputs"]["reference"] = ["2", 0]
|
||||
with self.assertRaisesRegex(ValueError, "bypasses Start.*2"):
|
||||
validate_image_list_regions(prompt)
|
||||
for explicit_targets in (False, True):
|
||||
with self.subTest(explicit_targets=explicit_targets):
|
||||
with self.assertRaisesRegex(ValueError, "bypasses Start.*2"):
|
||||
validate_image_list_regions(prompt, explicit_targets=explicit_targets)
|
||||
|
||||
def test_shared_typed_loader_and_leaf_constant_remain_valid(self):
|
||||
import sys
|
||||
@@ -247,19 +249,25 @@ class ImageListGraphBoundaryTests(unittest.TestCase):
|
||||
def test_rejects_missing_context_pair_nested_region_and_cross_boundary_branch(self):
|
||||
prompt = graph_one_region()
|
||||
prompt["5"]["inputs"]["context"] = ["2", 0]
|
||||
with self.assertRaisesRegex(ValueError, "typed context output"):
|
||||
validate_image_list_regions(prompt)
|
||||
for explicit_targets in (False, True):
|
||||
with self.subTest(explicit_targets=explicit_targets):
|
||||
with self.assertRaisesRegex(ValueError, "typed context output"):
|
||||
validate_image_list_regions(prompt, explicit_targets=explicit_targets)
|
||||
nested = graph_one_region()
|
||||
nested["7"] = start_node(["4", 0])
|
||||
nested["8"] = {"class_type": "NestedOperation", "inputs": {"image": ["7", IMAGE_OUT]}}
|
||||
nested["5"]["inputs"]["images"] = ["8", 0]
|
||||
nested["9"] = {"class_type": END_CLASS, "inputs": {"images": ["8", 0], "context": ["7", CONTEXT_OUT]}}
|
||||
with self.assertRaisesRegex(ValueError, "Nested/ambiguous"):
|
||||
validate_image_list_regions(nested)
|
||||
for explicit_targets in (False, True):
|
||||
with self.subTest(explicit_targets=explicit_targets):
|
||||
with self.assertRaisesRegex(ValueError, "Nested/ambiguous"):
|
||||
validate_image_list_regions(nested, explicit_targets=explicit_targets)
|
||||
outside = graph_one_region()
|
||||
outside["7"] = {"class_type": "PreviewImage", "inputs": {"images": ["3", IMAGE_OUT]}}
|
||||
with self.assertRaisesRegex(ValueError, "cross-boundary"):
|
||||
validate_image_list_regions(outside)
|
||||
for explicit_targets in (False, True):
|
||||
with self.subTest(explicit_targets=explicit_targets):
|
||||
with self.assertRaisesRegex(ValueError, "cross-boundary"):
|
||||
validate_image_list_regions(outside, explicit_targets=explicit_targets)
|
||||
|
||||
def test_rejects_duplicate_or_missing_results(self):
|
||||
prompt = graph_one_region()
|
||||
|
||||
@@ -0,0 +1,245 @@
|
||||
"""CPU-only execution-copy planning; no ComfyUI imports or prompt execution."""
|
||||
import copy
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
|
||||
from image_list.execution_plan import prepare_image_list_execution
|
||||
from image_list.runtime import (
|
||||
CONTEXT_OUT,
|
||||
END_CLASS,
|
||||
IMAGE_OUT,
|
||||
START_CLASS,
|
||||
build_worker_prompt,
|
||||
validate_image_list_regions,
|
||||
)
|
||||
|
||||
|
||||
class OutputNode:
|
||||
OUTPUT_NODE = True
|
||||
|
||||
|
||||
class SchemaOutputNode:
|
||||
OUTPUT_NODE = False # V3's schema, not this fallback, is authoritative.
|
||||
|
||||
@classmethod
|
||||
def GET_NODE_INFO_V1(cls):
|
||||
return {"output_node": True, "output": ("IMAGE",)}
|
||||
|
||||
|
||||
class ModelLoader:
|
||||
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def registry(monkeypatch):
|
||||
classes = {
|
||||
END_CLASS: SchemaOutputNode,
|
||||
"SaveImage": OutputNode,
|
||||
"PreviewImage": OutputNode,
|
||||
"OutputUtility": OutputNode,
|
||||
"SchemaUtility": SchemaOutputNode,
|
||||
"ModelLoader": ModelLoader,
|
||||
}
|
||||
monkeypatch.setitem(sys.modules, "nodes", types.SimpleNamespace(NODE_CLASS_MAPPINGS=classes))
|
||||
return classes
|
||||
|
||||
|
||||
def node(class_type, **inputs):
|
||||
return {"class_type": class_type, "inputs": inputs}
|
||||
|
||||
|
||||
def graph():
|
||||
return {
|
||||
"loader": node("ModelLoader"),
|
||||
"source": node("ImageSource"),
|
||||
"start": node(START_CLASS, images=["source", 0]),
|
||||
"process": node("FrameOperation", image=["start", IMAGE_OUT], model=["loader", 0]),
|
||||
"end": node(END_CLASS, images=["process", 0], context=["start", CONTEXT_OUT]),
|
||||
"save": node("SaveImage", images=["end", 0]),
|
||||
}
|
||||
|
||||
|
||||
def test_discards_multistep_previews_but_preserves_other_output_roots(registry):
|
||||
prompt = graph()
|
||||
prompt.update({
|
||||
"debug-transform": node("DebugTransform", images=["start", IMAGE_OUT]),
|
||||
"debug-preview": node("PreviewImage", images=["debug-transform", 0]),
|
||||
"process-preview": node("PreviewImage", images=["process", 0]),
|
||||
"unrelated-source": node("ImageSource"),
|
||||
"unrelated-preview": node("PreviewImage", images=["unrelated-source", 0]),
|
||||
"upstream-preview": node("PreviewImage", images=["source", 0]),
|
||||
"post-transform": node("FrameOperation", images=["end", 0]),
|
||||
"post-preview": node("PreviewImage", images=["post-transform", 0]),
|
||||
})
|
||||
before = copy.deepcopy(prompt)
|
||||
with pytest.raises(ValueError, match="cross-boundary"):
|
||||
validate_image_list_regions(prompt)
|
||||
prepared, targets = prepare_image_list_execution(prompt)
|
||||
assert set(prepared) == set(prompt) - {"debug-transform", "debug-preview", "process-preview"}
|
||||
assert targets == ["end", "save", "unrelated-preview", "upstream-preview", "post-preview"]
|
||||
validate_image_list_regions(prepared, explicit_targets=True)
|
||||
assert prompt == before
|
||||
|
||||
|
||||
@pytest.mark.parametrize("class_type", ["OutputUtility", "SchemaUtility"])
|
||||
def test_required_output_utility_keeps_complete_dependency_chain(registry, class_type):
|
||||
prompt = graph()
|
||||
prompt.update({
|
||||
"utility-config": node("Constant", value=7),
|
||||
"utility-first": node(class_type, images=["process", 0], config=["utility-config", 0]),
|
||||
"utility-second": node(class_type, images=["utility-first", 0]),
|
||||
"result": node("FrameOperation", image=["utility-second", 0]),
|
||||
"discarded-preview": node("PreviewImage", images=["utility-first", 0]),
|
||||
})
|
||||
prompt["end"]["inputs"]["images"] = ["result", 0]
|
||||
before = copy.deepcopy(prompt)
|
||||
registered = dict(registry)
|
||||
prepared, targets = prepare_image_list_execution(prompt)
|
||||
assert set(prepared) == set(prompt) - {"discarded-preview"}
|
||||
assert targets == ["end", "save"]
|
||||
with pytest.raises(ValueError, match="Output node"):
|
||||
validate_image_list_regions(prepared)
|
||||
region = validate_image_list_regions(prepared, explicit_targets=True)[0]
|
||||
worker = build_worker_prompt(prepared, region, "token", "context", 0, "http://master.invalid")
|
||||
assert set(worker) == set(prepared) - {"source", "save"}
|
||||
assert worker["utility-second"]["inputs"] == {"images": ["utility-first", 0]}
|
||||
assert prompt == before
|
||||
assert registry == registered
|
||||
assert OutputNode.OUTPUT_NODE is True
|
||||
assert SchemaOutputNode.OUTPUT_NODE is False
|
||||
assert SchemaOutputNode.GET_NODE_INFO_V1()["output_node"] is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("source", ["start", "process"])
|
||||
@pytest.mark.parametrize("direct", [False, True])
|
||||
def test_internal_bypass_into_post_end_computation_is_rejected(registry, source, direct):
|
||||
prompt = graph()
|
||||
prompt["post-process"] = node("Composite", images=["end", 0])
|
||||
if not direct:
|
||||
prompt["bypass"] = node("FrameOperation", image=[source, 0])
|
||||
source = "bypass"
|
||||
prompt["post-process"]["inputs"]["reference"] = [source, 0]
|
||||
prompt["save"]["inputs"]["images"] = ["post-process", 0]
|
||||
before = copy.deepcopy(prompt)
|
||||
with pytest.raises(ValueError, match="cross-boundary"):
|
||||
prepare_image_list_execution(prompt)
|
||||
assert prompt == before
|
||||
|
||||
|
||||
def test_sequential_namespaced_regions_and_intermediate_output_are_preserved(registry):
|
||||
prompt = graph()
|
||||
prompt.update({
|
||||
"next-start": node(START_CLASS, images=["end", 0]),
|
||||
"next-process": node("SchemaUtility", images=["next-start", IMAGE_OUT], model=["loader", 0]),
|
||||
"next-end": node(END_CLASS, images=["next-process", 0], context=["next-start", CONTEXT_OUT]),
|
||||
"next-save": node("SaveImage", images=["next-end", 0]),
|
||||
"first-debug": node("PreviewImage", images=["process", 0]),
|
||||
"second-debug": node("PreviewImage", images=["next-process", 0]),
|
||||
})
|
||||
renamed = {key: f"group:7/{key}" for key in prompt}
|
||||
for spec in prompt.values():
|
||||
for name, value in spec["inputs"].items():
|
||||
if isinstance(value, list):
|
||||
spec["inputs"][name] = [renamed[value[0]], value[1]]
|
||||
prompt = {renamed[key]: spec for key, spec in prompt.items()}
|
||||
prepared, targets = prepare_image_list_execution(prompt)
|
||||
assert set(prepared) == set(prompt) - {renamed["first-debug"], renamed["second-debug"]}
|
||||
assert targets == [renamed[name] for name in ("end", "save", "next-end", "next-save")]
|
||||
regions = validate_image_list_regions(prepared, explicit_targets=True)
|
||||
assert {(r["start_id"], r["end_id"]) for r in regions} == {
|
||||
(renamed["start"], renamed["end"]), (renamed["next-start"], renamed["next-end"]),
|
||||
}
|
||||
second = next(r for r in regions if r["start_id"] == renamed["next-start"])
|
||||
worker = build_worker_prompt(prepared, second, "token", "context", 0, "http://master.invalid")
|
||||
assert set(worker) == {renamed[name] for name in ("loader", "next-start", "next-process", "next-end")}
|
||||
|
||||
|
||||
def test_shared_loader_and_its_independent_output_root_are_preserved(registry):
|
||||
registry["ModelLoader"] = type("OutputLoader", (ModelLoader, OutputNode), {})
|
||||
prompt = graph()
|
||||
prompt["source"]["inputs"]["model"] = ["loader", 0]
|
||||
prompt["independent-save"] = node("SaveImage", model=["loader", 0])
|
||||
prepared, targets = prepare_image_list_execution(prompt)
|
||||
assert prepared == prompt
|
||||
assert targets == ["loader", "end", "save", "independent-save"]
|
||||
region = validate_image_list_regions(prepared, explicit_targets=True)[0]
|
||||
assert "loader" in region["node_ids"]
|
||||
assert "loader" not in region["branch_ids"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("unsafe_class", [None, type("ListLoader", (), {
|
||||
"RETURN_TYPES": ("MODEL",), "OUTPUT_IS_LIST": (True,),
|
||||
})])
|
||||
def test_unknown_or_list_producing_shared_loader_still_fails_closed(registry, unsafe_class):
|
||||
if unsafe_class is None:
|
||||
registry.pop("ModelLoader")
|
||||
else:
|
||||
registry["ModelLoader"] = unsafe_class
|
||||
prompt = graph()
|
||||
prompt["source"]["inputs"]["model"] = ["loader", 0]
|
||||
prompt["debug"] = node("PreviewImage", images=["process", 0])
|
||||
with pytest.raises(ValueError, match="bypasses Start.*loader"):
|
||||
prepare_image_list_execution(prompt)
|
||||
|
||||
|
||||
def test_upstream_per_item_bypass_is_not_hidden_by_preview_pruning(registry):
|
||||
prompt = graph()
|
||||
prompt["process"]["inputs"]["reference"] = ["source", 0]
|
||||
prompt["debug"] = node("PreviewImage", images=["process", 0])
|
||||
with pytest.raises(ValueError, match="bypasses Start.*source"):
|
||||
prepare_image_list_execution(prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("nested", [False, True])
|
||||
def test_overlapping_or_nested_regions_cannot_be_discarded_as_side_branches(registry, nested):
|
||||
prompt = graph()
|
||||
prompt["other-start"] = node(START_CLASS, images=["process" if nested else "source", 0])
|
||||
prompt["other-process"] = node("FrameOperation", image=["other-start", IMAGE_OUT], reference=["process", 0])
|
||||
prompt["other-end"] = node(END_CLASS, images=["other-process", 0], context=["other-start", CONTEXT_OUT])
|
||||
with pytest.raises(ValueError, match="Nested/ambiguous|cross-boundary|ambiguous processing"):
|
||||
prepare_image_list_execution(prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("invalid", ["missing-end", "wrong-context", "missing-images", "unrelated-images"])
|
||||
def test_pairing_and_result_validation_cannot_be_pruned_away(registry, invalid):
|
||||
prompt = graph()
|
||||
if invalid == "missing-end":
|
||||
del prompt["end"]
|
||||
elif invalid == "wrong-context":
|
||||
prompt["end"]["inputs"]["context"] = ["source", 0]
|
||||
elif invalid == "missing-images":
|
||||
del prompt["end"]["inputs"]["images"]
|
||||
else:
|
||||
prompt["end"]["inputs"]["images"] = ["source", 0]
|
||||
before = copy.deepcopy(prompt)
|
||||
with pytest.raises(ValueError):
|
||||
prepare_image_list_execution(prompt)
|
||||
assert prompt == before
|
||||
|
||||
|
||||
def test_execution_copy_is_deep_and_repeated_preparation_is_stable(registry):
|
||||
prompt = graph()
|
||||
prompt["process"]["_meta"] = {"title": "Original", "tags": ["saved"]}
|
||||
before = copy.deepcopy(prompt)
|
||||
prepared, targets = prepare_image_list_execution(prompt)
|
||||
prepared_again, targets_again = prepare_image_list_execution(prepared)
|
||||
assert prepared_again == prepared
|
||||
assert targets_again == targets
|
||||
prepared["start"]["inputs"]["images"][0] = "changed"
|
||||
prepared["process"]["_meta"]["tags"].append("execution-only")
|
||||
assert prompt == before
|
||||
assert prepared_again == before
|
||||
|
||||
|
||||
def test_v3_false_output_flag_overrides_legacy_true(registry):
|
||||
registry["NotOutput"] = type("NotOutput", (), {
|
||||
"OUTPUT_NODE": True,
|
||||
"GET_NODE_INFO_V1": classmethod(lambda cls: {"output_node": False}),
|
||||
})
|
||||
prompt = graph()
|
||||
prompt["not-output"] = node("NotOutput", images=["end", 0])
|
||||
prepared, targets = prepare_image_list_execution(prompt)
|
||||
assert prepared == prompt
|
||||
assert targets == ["end", "save"]
|
||||
@@ -35,6 +35,15 @@ def load_routes():
|
||||
return module, routes
|
||||
|
||||
|
||||
class CapabilityTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_capability_advertises_target_control_without_tensor_protocol_bump(self):
|
||||
module, _routes = load_routes()
|
||||
response = await module.image_list_capabilities_endpoint(None)
|
||||
payload = json.loads(response.text)
|
||||
self.assertIs(payload['execution_targets'], True)
|
||||
self.assertEqual(payload['protocol_version'], 2)
|
||||
|
||||
|
||||
class TransportTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
self.module, routes = load_routes()
|
||||
|
||||
@@ -11,6 +11,83 @@ from aiohttp import web
|
||||
from image_list import lifecycle, runtime
|
||||
|
||||
|
||||
class ExecutionTargetProofTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
from test_image_list_master_runtime import graph
|
||||
self.prompt = graph()
|
||||
self.registry = SimpleNamespace(NODE_CLASS_MAPPINGS={
|
||||
'Operation': type('Flagged', (), {'OUTPUT_NODE': True})})
|
||||
self.modules = patch.dict('sys.modules', {'nodes': self.registry})
|
||||
self.modules.start()
|
||||
self.addCleanup(self.modules.stop)
|
||||
|
||||
def server(self, roots, *, prompt=None, running=True):
|
||||
item = (0, 'owned', self.prompt if prompt is None else prompt, {}, roots, {})
|
||||
return SimpleNamespace(prompt_queue=SimpleNamespace(
|
||||
get_current_queue_volatile=lambda: ([item], []) if running else ([], [item])))
|
||||
|
||||
def test_outer_roots_in_running_queue_are_authoritative(self):
|
||||
self.assertTrue(lifecycle.has_safe_execution_targets(self.server(['end']), 'owned', self.prompt))
|
||||
for roots in (['op', 'end'], [], ['end', 'end'], None, [True]):
|
||||
with self.subTest(roots=roots):
|
||||
self.assertFalse(lifecycle.has_safe_execution_targets(self.server(roots), 'owned', self.prompt))
|
||||
self.assertFalse(lifecycle.has_safe_execution_targets(self.server(['end']), 'other', self.prompt))
|
||||
self.assertFalse(lifecycle.has_safe_execution_targets(self.server(['end'], running=False), 'owned', self.prompt))
|
||||
self.assertFalse(lifecycle.has_safe_execution_targets(self.server(['end'], prompt={}), 'owned', self.prompt))
|
||||
|
||||
def test_prompt_marker_is_not_an_execution_guarantee(self):
|
||||
self.prompt['start']['inputs']['image_list_explicit_targets'] = True
|
||||
self.assertFalse(lifecycle.has_safe_execution_targets(self.server(['op', 'end']), 'owned', self.prompt))
|
||||
with self.assertRaisesRegex(ValueError, 'must be placed after'):
|
||||
runtime.create_master_job('start', 'end', self.prompt, [], [], '', master_participates=True)
|
||||
self.assertEqual(runtime._MASTER_JOBS, {})
|
||||
|
||||
def test_start_node_uses_running_queue_proof(self):
|
||||
import importlib.util
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from types import ModuleType
|
||||
import torch
|
||||
|
||||
root = Path(__file__).resolve().parents[1]
|
||||
package = '_image_list_start_targets'
|
||||
modules = {}
|
||||
for suffix in ('', '.nodes', '.image_list', '.utils'):
|
||||
module = ModuleType(package + suffix)
|
||||
module.__path__ = [str(root / suffix.lstrip('.'))]
|
||||
modules[package + suffix] = module
|
||||
modules[package + '.image_list.runtime'] = runtime
|
||||
modules[package + '.image_list.lifecycle'] = lifecycle
|
||||
modules[package + '.utils.async_helpers'] = SimpleNamespace(run_async_in_server_loop=None)
|
||||
io = SimpleNamespace(ComfyNode=type('ComfyNode', (), {}), NodeOutput=lambda *args: args,
|
||||
**{name: object() for name in ('Image', 'Mask', 'String', 'Int', 'Float', 'Boolean')})
|
||||
modules['comfy_api.v0_0_2'] = SimpleNamespace(io=io)
|
||||
modules['comfy_execution.utils'] = SimpleNamespace(
|
||||
get_executing_context=lambda: SimpleNamespace(prompt_id='owned'))
|
||||
server = SimpleNamespace(PromptServer=SimpleNamespace(instance=self.server(['end'])))
|
||||
modules['server'] = server
|
||||
spec = importlib.util.spec_from_file_location(package + '.nodes.image_list', root / 'nodes/image_list.py')
|
||||
node_module = importlib.util.module_from_spec(spec)
|
||||
with patch.dict(sys.modules, modules):
|
||||
spec.loader.exec_module(node_module)
|
||||
start = node_module.DistributedImageListStart
|
||||
start.hidden = SimpleNamespace(prompt=self.prompt, unique_id='start')
|
||||
result = start.execute(images=torch.zeros((1, 2, 2, 3)), image_list_master_participates=True)
|
||||
context = result[3][0]
|
||||
self.assertIsNotNone(runtime.get_master_job(context))
|
||||
runtime.remove_master_job(context)
|
||||
server.PromptServer.instance = self.server(['op', 'end'])
|
||||
with self.assertRaisesRegex(ValueError, 'must be placed after'):
|
||||
start.execute(images=torch.zeros((1, 2, 2, 3)), image_list_master_participates=True)
|
||||
self.assertEqual(runtime._MASTER_JOBS, {})
|
||||
|
||||
def test_owned_target_controlled_start_can_create_job(self):
|
||||
context = runtime.create_master_job('start', 'end', self.prompt, [], [], '',
|
||||
master_participates=True, explicit_targets=True)
|
||||
self.addCleanup(runtime.remove_master_job, context)
|
||||
self.assertIsNotNone(runtime.get_master_job(context))
|
||||
|
||||
|
||||
class ImageListLifecycleTests(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self):
|
||||
with runtime._LOCK:
|
||||
|
||||
@@ -34,6 +34,23 @@ class MasterSelectionTests(unittest.IsolatedAsyncioTestCase):
|
||||
def prepared_inputs(self):
|
||||
return self.q.queue_prompt_payload.call_args.args[0]["2"]["inputs"]
|
||||
|
||||
async def test_prepared_master_prunes_side_outputs_without_mutating_input(self):
|
||||
import copy
|
||||
self.prompt['op'] = {'class_type': 'Flagged', 'inputs': {'image': ['2', 0]}}
|
||||
self.prompt['3']['inputs']['images'] = ['op', 0]
|
||||
self.prompt['preview'] = {'class_type': 'Flagged', 'inputs': {'image': ['op', 0]}}
|
||||
self.prompt['save'] = {'class_type': 'Flagged', 'inputs': {'image': ['3', 0]}}
|
||||
original = copy.deepcopy(self.prompt)
|
||||
registry = type('Registry', (), {'NODE_CLASS_MAPPINGS': {
|
||||
'Flagged': type('Flagged', (), {'OUTPUT_NODE': True})}})
|
||||
with patch.dict(sys.modules, {'nodes': registry}):
|
||||
await self.queue()
|
||||
queued = self.q.queue_prompt_payload.call_args
|
||||
self.assertEqual(queued.kwargs['partial_execution_targets'], ['3', 'save'])
|
||||
self.assertIn('op', queued.args[0])
|
||||
self.assertNotIn('preview', queued.args[0])
|
||||
self.assertEqual(self.prompt, original)
|
||||
|
||||
async def test_default_includes_master_and_direct_workers(self):
|
||||
result = await self.queue()
|
||||
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
|
||||
@@ -100,12 +117,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")
|
||||
|
||||
@@ -84,7 +84,7 @@ class RuntimeBoundaryTests(unittest.IsolatedAsyncioTestCase):
|
||||
# Other legacy test modules replace aiohttp during collection. This
|
||||
# controlled provider fixture needs only its explicit timeout/error seam.
|
||||
sys.modules['aiohttp'] = types.SimpleNamespace(
|
||||
ClientTimeout=lambda **kwargs: kwargs, ClientError=ConnectionError)
|
||||
ClientTimeout=lambda **kwargs: kwargs, ClientError=ConnectionError, ClientConnectionError=ConnectionError)
|
||||
comfy = types.ModuleType('comfy')
|
||||
comfy.model_management = types.SimpleNamespace(processing_interrupted=lambda: self.interrupted)
|
||||
sys.modules['comfy'] = comfy
|
||||
@@ -94,6 +94,7 @@ class RuntimeBoundaryTests(unittest.IsolatedAsyncioTestCase):
|
||||
prompt['7']['inputs']['image'] = 'shared/ref.png'
|
||||
|
||||
async def dispatch_prompt(self, worker, prompt, meta, **kwargs):
|
||||
self.assertEqual(kwargs.get('partial_execution_targets'), ['5'])
|
||||
if self.reject:
|
||||
raise RuntimeError('prompt rejected')
|
||||
self.prompts.append((json.loads(json.dumps(prompt)), meta))
|
||||
@@ -136,6 +137,23 @@ class RuntimeBoundaryTests(unittest.IsolatedAsyncioTestCase):
|
||||
['worker'], 'http://public-master', self.metadata)
|
||||
return self.context
|
||||
|
||||
async def test_runtime_preflight_rejects_old_target_capability_before_upload(self):
|
||||
from test_dispatch_selection import _load_dispatch_module, _FakeSession, _valid_capability
|
||||
real_dispatch = _load_dispatch_module()
|
||||
capability = _valid_capability()
|
||||
capability.pop('execution_targets')
|
||||
session = _FakeSession({('worker', 'capabilities'): (200, capability)})
|
||||
module = sys.modules['_image_list_boundaries.api.orchestration.dispatch']
|
||||
module.preflight_image_list_workers = real_dispatch.preflight_image_list_workers
|
||||
with patch.object(real_dispatch, 'get_client_session', AsyncMock(return_value=session)), patch.object(
|
||||
real_dispatch, 'build_worker_url', side_effect=lambda worker, path: f"http://{worker['id']}{path}"
|
||||
):
|
||||
with self.assertRaisesRegex(RuntimeError, 'execution_targets.*upgrade'):
|
||||
await self.rt.run_image_list_job(self.create(1))
|
||||
self.assertFalse(self.uploads)
|
||||
self.dispatch.assert_not_awaited()
|
||||
self.assertFalse(self.rt._MASTER_JOBS)
|
||||
|
||||
async def test_each_assignment_reuses_media_paths_with_fresh_tokens_and_metadata(self):
|
||||
context = self.create()
|
||||
await self.rt.run_image_list_job(context)
|
||||
|
||||
@@ -0,0 +1,683 @@
|
||||
"""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)}
|
||||
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_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,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,265 @@
|
||||
"""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.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))
|
||||
assert json.loads(response.text)["image_transfer_formats"] == list(FORMATS)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reject_call", [None, 1, 2, 3])
|
||||
def test_runtime_initial_and_each_nonexact_assignment_preflight(reject_call):
|
||||
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)
|
||||
calls = []
|
||||
|
||||
async def preflight(workers, **kwargs):
|
||||
calls.append(kwargs)
|
||||
assert kwargs["required_transfer_formats"] == ("png",)
|
||||
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)
|
||||
assert len(fixture.uploads) == max(0, reject_call - 2)
|
||||
assert len(fixture.prompts) == max(0, reject_call - 2)
|
||||
else:
|
||||
await fixture.rt.run_image_list_job(context)
|
||||
assert len(calls) == 3
|
||||
assert len(fixture.uploads) == 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())
|
||||
@@ -1,4 +1,5 @@
|
||||
"""Real-framework acceptance; opt in with COMFYUI_SOURCE_ROOT."""
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
@@ -35,3 +36,15 @@ def test_image_list_network_runtime():
|
||||
assert 'IMAGE_LIST_NETWORK_OK' in result.stdout
|
||||
assert 'IMAGE_LIST_CANCEL_OK' in result.stdout
|
||||
assert 'IMAGE_LIST_MASTER_OK' in result.stdout
|
||||
assert 'IMAGE_LIST_PREFLIGHT_OK' in result.stdout
|
||||
assert 'IMAGE_LIST_SEQUENTIAL_OK' in result.stdout
|
||||
assert 'IMAGE_LIST_FAILURE_OK' in result.stdout
|
||||
reports = [json.loads(line.removeprefix('IMAGE_LIST_TARGETS_OK '))
|
||||
for line in result.stdout.splitlines() if line.startswith('IMAGE_LIST_TARGETS_OK ')]
|
||||
assert len(reports) == 1, result.stdout
|
||||
assert reports[0] == {
|
||||
'modes': ['delegate-only', 'mixed', 'master-only'],
|
||||
'outer_targets': ['4', '5'], 'worker_targets': ['4'],
|
||||
'required_utility_once_per_item': True, 'preview_discarded': True,
|
||||
'saved_prompt_unchanged': True, 'capture_once_per_outer_prompt': True,
|
||||
}
|
||||
|
||||
+41
-2
@@ -105,24 +105,63 @@ class PromptValidationError(RuntimeError):
|
||||
super().__init__(f"Invalid prompt: {merged}")
|
||||
|
||||
|
||||
def validate_execution_targets(prompt, targets):
|
||||
"""Reject malformed explicit roots instead of falling back to all outputs."""
|
||||
if (not isinstance(targets, list) or not targets
|
||||
or any(not isinstance(node_id, str) or not node_id or node_id not in prompt
|
||||
for node_id in targets)
|
||||
or len(set(targets)) != len(targets)):
|
||||
raise PromptValidationError({
|
||||
"type": "invalid_partial_execution_targets",
|
||||
"message": "partial_execution_targets must be a non-empty list of unique existing node IDs",
|
||||
"details": "", "extra_info": {},
|
||||
})
|
||||
return list(targets)
|
||||
|
||||
|
||||
async def queue_prompt_payload(
|
||||
prompt_obj,
|
||||
workflow_meta=None,
|
||||
client_id=None,
|
||||
include_queue_metadata=False,
|
||||
*,
|
||||
partial_execution_targets=None,
|
||||
):
|
||||
"""Validate and queue a prompt via ComfyUI's prompt queue."""
|
||||
"""Validate and queue a prompt, preserving any explicitly requested roots."""
|
||||
targets = (validate_execution_targets(prompt_obj, partial_execution_targets)
|
||||
if partial_execution_targets is not None else None)
|
||||
payload = {"prompt": prompt_obj}
|
||||
if targets is not None:
|
||||
payload["partial_execution_targets"] = list(targets)
|
||||
payload = prompt_server.trigger_on_prompt(payload)
|
||||
prompt = payload["prompt"]
|
||||
if targets is not None:
|
||||
# Hooks may replace the envelope or mutate its target list. The caller's
|
||||
# requested roots remain authoritative, but must still exist afterwards.
|
||||
validate_execution_targets(prompt, targets)
|
||||
|
||||
prompt_id = str(uuid.uuid4())
|
||||
valid = await execution.validate_prompt(prompt_id, prompt, None)
|
||||
valid = await execution.validate_prompt(prompt_id, prompt, targets)
|
||||
if not valid[0]:
|
||||
error_payload = valid[1] if len(valid) > 1 else "Prompt outputs failed validation"
|
||||
node_errors = valid[3] if len(valid) > 3 else {}
|
||||
raise PromptValidationError(error_payload, node_errors)
|
||||
|
||||
if targets is not None:
|
||||
node_errors = valid[3] if len(valid) > 3 else {}
|
||||
target_errors = {
|
||||
node_id: entry for node_id, entry in node_errors.items()
|
||||
if node_id in targets or (isinstance(entry, dict)
|
||||
and set(entry.get("dependent_outputs", [])) & set(targets))
|
||||
}
|
||||
if set(valid[2]) != set(targets) or target_errors:
|
||||
raise PromptValidationError({
|
||||
"type": "partial_execution_targets_failed_validation",
|
||||
"message": "Not all requested execution targets passed validation",
|
||||
"details": "Refusing to queue a subset or additional output roots",
|
||||
"extra_info": {"requested_targets": targets},
|
||||
}, node_errors)
|
||||
|
||||
extra_data = {"create_time": int(time.time() * 1000)}
|
||||
if workflow_meta:
|
||||
extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
|
||||
|
||||
Reference in New Issue
Block a user