feat: add dynamic per-image list queue with Start and End nodes
This commit is contained in:
@@ -27,6 +27,11 @@
|
||||
- Intelligent distribution
|
||||
- Handles single images and videos
|
||||
|
||||
#### Dynamic Image-List Processing
|
||||
- Schedule one image at a time to the next available direct worker; collect results in input order
|
||||
- Use **Distributed Image List Start / End** around a per-image processing region
|
||||
- [Wiring, supported inputs, and first-version limits](/docs/image-list-queue.md)
|
||||
|
||||
#### Ease of Use
|
||||
- Auto-setup local workers; easily add remote/cloud ones
|
||||
- Convert any workflow to distributed with 2 nodes
|
||||
|
||||
@@ -4,3 +4,4 @@ from . import tunnel_routes # noqa: F401
|
||||
from . import worker_routes # noqa: F401
|
||||
from . import job_routes # noqa: F401
|
||||
from . import usdu_routes # noqa: F401
|
||||
from . import image_list_routes # noqa: F401
|
||||
|
||||
@@ -0,0 +1,89 @@
|
||||
"""Bounded, typed HTTP transport for dynamic image-list worker tasks."""
|
||||
import json
|
||||
|
||||
from aiohttp import web
|
||||
import server
|
||||
|
||||
from ..image_list.runtime import (
|
||||
MAX_PAYLOAD_BYTES,
|
||||
decode_result,
|
||||
delete_worker_item,
|
||||
set_master_result,
|
||||
store_worker_item,
|
||||
)
|
||||
|
||||
prompt_server = server.PromptServer.instance
|
||||
MAX_REQUEST_BYTES = MAX_PAYLOAD_BYTES * 2 + 1024 * 1024
|
||||
|
||||
|
||||
async def _json_body(request):
|
||||
length = request.content_length
|
||||
if length is not None and length > MAX_REQUEST_BYTES:
|
||||
raise ValueError(f"Request exceeds the {MAX_REQUEST_BYTES}-byte image-list limit")
|
||||
raw = bytearray()
|
||||
async for chunk in request.content.iter_chunked(64 * 1024):
|
||||
raw.extend(chunk)
|
||||
if len(raw) > MAX_REQUEST_BYTES:
|
||||
raise ValueError(f"Request exceeds the {MAX_REQUEST_BYTES}-byte image-list limit")
|
||||
try:
|
||||
return json.loads(raw.decode("utf-8"))
|
||||
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
|
||||
raise ValueError(f"Invalid image-list JSON body: {exc}") from exc
|
||||
|
||||
|
||||
@prompt_server.routes.get("/distributed/image-list/capabilities")
|
||||
async def image_list_capabilities_endpoint(_request):
|
||||
return web.json_response({
|
||||
"supported": True,
|
||||
"protocol_version": 1,
|
||||
"node_ids": ["DistributedImageListStart", "DistributedImageListEnd"],
|
||||
"transport": "typed-float-tensor-json-v1",
|
||||
})
|
||||
|
||||
|
||||
@prompt_server.routes.post("/distributed/image-list/item")
|
||||
async def image_list_item_endpoint(request):
|
||||
try:
|
||||
data = await _json_body(request)
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError("Expected an image-list item object")
|
||||
context_id = data.get("context_id")
|
||||
if not isinstance(context_id, str) or not context_id:
|
||||
raise ValueError("context_id must be a non-empty string")
|
||||
item = data.get("item")
|
||||
store_worker_item(data.get("token"), item)
|
||||
return web.json_response({"status": "stored", "item_index": item["index"]})
|
||||
except (ValueError, TypeError) as exc:
|
||||
return web.json_response({"error": str(exc)}, status=400)
|
||||
except Exception as exc:
|
||||
return web.json_response({"error": str(exc)}, status=500)
|
||||
|
||||
|
||||
@prompt_server.routes.delete("/distributed/image-list/item/{token}")
|
||||
async def image_list_delete_item_endpoint(request):
|
||||
token = request.match_info.get("token", "")
|
||||
return web.json_response({"status": "deleted", "removed": delete_worker_item(token)})
|
||||
|
||||
|
||||
@prompt_server.routes.post("/distributed/image-list/result")
|
||||
async def image_list_result_endpoint(request):
|
||||
try:
|
||||
data = await _json_body(request)
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError("Expected an image-list result object")
|
||||
context_id = data.get("context_id")
|
||||
token = data.get("token")
|
||||
item_index = data.get("item_index")
|
||||
if not isinstance(context_id, str) or not context_id:
|
||||
raise ValueError("context_id must be a non-empty string")
|
||||
if not isinstance(token, str) or not token:
|
||||
raise ValueError("token must be a non-empty string")
|
||||
if isinstance(item_index, bool) or not isinstance(item_index, int) or item_index < 0:
|
||||
raise ValueError("item_index must be a non-negative integer")
|
||||
result = decode_result(data.get("result"))
|
||||
set_master_result(context_id, item_index, token, result)
|
||||
return web.json_response({"status": "accepted", "item_index": item_index})
|
||||
except (ValueError, TypeError) as exc:
|
||||
return web.json_response({"error": str(exc)}, status=400)
|
||||
except Exception as exc:
|
||||
return web.json_response({"error": str(exc)}, status=500)
|
||||
@@ -85,7 +85,7 @@ async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0):
|
||||
data = json.loads(msg.data or "{}")
|
||||
if data.get("type") == "dispatch_ack" and data.get("request_id") == request_id:
|
||||
if data.get("ok"):
|
||||
return
|
||||
return data.get("prompt_id")
|
||||
error_text = data.get("error") or "Worker rejected websocket dispatch."
|
||||
validation_error = data.get("validation_error")
|
||||
node_errors = data.get("node_errors")
|
||||
@@ -107,6 +107,7 @@ async def dispatch_worker_prompt(
|
||||
client_id=None,
|
||||
use_websocket=False,
|
||||
trace_execution_id=None,
|
||||
return_prompt_id=False,
|
||||
):
|
||||
"""Send the prepared prompt to a worker ComfyUI instance."""
|
||||
worker_url = build_worker_url(worker)
|
||||
@@ -120,7 +121,7 @@ async def dispatch_worker_prompt(
|
||||
|
||||
if use_websocket:
|
||||
try:
|
||||
await _dispatch_via_websocket(
|
||||
prompt_id = await _dispatch_via_websocket(
|
||||
worker_url,
|
||||
{
|
||||
"prompt": prompt_obj,
|
||||
@@ -128,7 +129,9 @@ async def dispatch_worker_prompt(
|
||||
},
|
||||
client_id,
|
||||
)
|
||||
return
|
||||
if return_prompt_id and (not isinstance(prompt_id, str) or not prompt_id):
|
||||
raise RuntimeError("Worker websocket acknowledgement omitted prompt_id")
|
||||
return prompt_id if return_prompt_id else None
|
||||
except Exception as exc:
|
||||
worker_id = worker.get("id")
|
||||
if trace_execution_id:
|
||||
@@ -144,6 +147,73 @@ async def dispatch_worker_prompt(
|
||||
timeout=aiohttp.ClientTimeout(total=60),
|
||||
) as resp:
|
||||
resp.raise_for_status()
|
||||
if not return_prompt_id:
|
||||
return None
|
||||
try:
|
||||
response_data = await resp.json()
|
||||
except Exception as exc:
|
||||
raise RuntimeError("Worker /prompt response did not contain JSON acknowledgement") from exc
|
||||
prompt_id = response_data.get("prompt_id") if isinstance(response_data, dict) else None
|
||||
if not isinstance(prompt_id, str) or not prompt_id:
|
||||
raise RuntimeError("Worker /prompt response omitted prompt_id")
|
||||
return prompt_id if return_prompt_id else None
|
||||
|
||||
|
||||
async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0):
|
||||
"""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
|
||||
serverless provider. Unsupported worker releases fail before master inputs run.
|
||||
"""
|
||||
if not workers:
|
||||
raise RuntimeError("Distributed Image List requires at least one selected direct worker")
|
||||
session = await get_client_session()
|
||||
required_nodes = {"DistributedImageListStart", "DistributedImageListEnd"}
|
||||
semaphore = asyncio.Semaphore(8)
|
||||
|
||||
async def inspect(worker):
|
||||
worker_id = str(worker.get("id") or "unknown")
|
||||
async with semaphore:
|
||||
cap_url = build_worker_url(worker, "/distributed/image-list/capabilities")
|
||||
try:
|
||||
async with session.get(cap_url, timeout=aiohttp.ClientTimeout(total=timeout)) as response:
|
||||
if response.status != 200:
|
||||
raise RuntimeError(
|
||||
f"Image-list worker {worker_id} does not expose the capability endpoint "
|
||||
f"(HTTP {response.status}); upgrade that worker before running this feature"
|
||||
)
|
||||
capability = await response.json()
|
||||
except asyncio.TimeoutError as exc:
|
||||
raise RuntimeError(f"Image-list worker {worker_id} capability check timed out") from exc
|
||||
except aiohttp.ClientError as exc:
|
||||
raise RuntimeError(f"Image-list worker {worker_id} capability check failed: {exc}") from exc
|
||||
nodes = set(capability.get("node_ids", [])) if isinstance(capability, dict) else set()
|
||||
if (not isinstance(capability, dict) or capability.get("supported") is not True
|
||||
or capability.get("protocol_version") != 1 or not required_nodes.issubset(nodes)):
|
||||
raise RuntimeError(
|
||||
f"Image-list worker {worker_id} is incompatible: requires protocol v1 and nodes "
|
||||
f"{sorted(required_nodes)}; no work was submitted"
|
||||
)
|
||||
queue_url = build_worker_url(worker, "/queue")
|
||||
try:
|
||||
async with session.get(queue_url, timeout=aiohttp.ClientTimeout(total=timeout)) as response:
|
||||
if response.status != 200:
|
||||
raise RuntimeError(f"Image-list worker {worker_id} queue preflight returned HTTP {response.status}")
|
||||
queue = await response.json()
|
||||
except asyncio.TimeoutError as exc:
|
||||
raise RuntimeError(f"Image-list worker {worker_id} queue preflight timed out") from exc
|
||||
except aiohttp.ClientError as exc:
|
||||
raise RuntimeError(f"Image-list worker {worker_id} queue preflight failed: {exc}") from exc
|
||||
if not isinstance(queue, dict) or not isinstance(queue.get("queue_running"), list) or not isinstance(queue.get("queue_pending"), list):
|
||||
raise RuntimeError(f"Image-list worker {worker_id} returned malformed /queue state")
|
||||
idle = not queue["queue_running"] and not queue["queue_pending"]
|
||||
return worker_id, idle
|
||||
|
||||
statuses = await asyncio.gather(*(inspect(worker) for worker in workers))
|
||||
result = {worker_id: idle for worker_id, idle in statuses}
|
||||
if require_idle and not any(result.values()):
|
||||
raise RuntimeError("No selected feature-compatible direct worker is idle; unrelated worker prompts were left untouched")
|
||||
return {worker_id: True for worker_id, idle in result.items() if idle}
|
||||
|
||||
|
||||
async def select_active_workers(
|
||||
|
||||
@@ -441,11 +441,39 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids):
|
||||
return pruned_prompt
|
||||
|
||||
|
||||
def _resolve_direct_workers_for_image_list(config, requested_ids):
|
||||
"""Resolve only enabled direct workers already selected by orchestration."""
|
||||
requested = None if requested_ids is None else {str(worker_id) for worker_id in requested_ids}
|
||||
direct = []
|
||||
for worker in (config or {}).get("workers", []):
|
||||
worker_id = str(worker.get("id") or "").strip()
|
||||
if not worker_id or not worker.get("enabled", False):
|
||||
continue
|
||||
if requested is not None and worker_id not in requested:
|
||||
continue
|
||||
if str(worker.get("type") or "local").strip().lower() == "serverless":
|
||||
continue
|
||||
try:
|
||||
port = int(worker.get("port", worker.get("listen_port", 8188)) or 8188)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise RuntimeError(f"Invalid port for image-list worker {worker_id}") from exc
|
||||
direct.append({
|
||||
"id": worker_id,
|
||||
"name": str(worker.get("name") or worker_id),
|
||||
"host": worker.get("host"),
|
||||
"port": port,
|
||||
"type": worker.get("type", "local"),
|
||||
})
|
||||
return direct
|
||||
|
||||
|
||||
def generate_job_id_map(prompt_index, prefix):
|
||||
"""Create stable per-node job IDs for distributed nodes."""
|
||||
job_map = {}
|
||||
distributed_nodes = prompt_index.nodes_for_class("DistributedCollector") + prompt_index.nodes_for_class(
|
||||
"UltimateSDUpscaleDistributed"
|
||||
distributed_nodes = (
|
||||
prompt_index.nodes_for_class("DistributedCollector")
|
||||
+ prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
|
||||
+ prompt_index.nodes_for_class("DistributedImageListStart")
|
||||
)
|
||||
for node_id in distributed_nodes:
|
||||
job_map[node_id] = f"{prefix}_{node_id}"
|
||||
@@ -540,6 +568,21 @@ def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id,
|
||||
inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}"
|
||||
|
||||
|
||||
def _override_image_list_nodes(prompt_copy, prompt_index, participant_id, worker_ids,
|
||||
job_id_map, master_url, workflow_meta=None):
|
||||
"""Attach protocol context to Start without changing its visible V3 sockets."""
|
||||
for node_id in prompt_index.nodes_for_class("DistributedImageListStart"):
|
||||
node = prompt_copy.get(node_id)
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
inputs = node.setdefault("inputs", {})
|
||||
inputs["image_list_enabled_worker_ids"] = json.dumps(list(worker_ids))
|
||||
inputs["image_list_master_callback"] = str(master_url or "")
|
||||
inputs["image_list_execution_id"] = job_id_map.get(node_id, node_id)
|
||||
if workflow_meta is not None:
|
||||
inputs["image_list_workflow_meta"] = workflow_meta
|
||||
|
||||
|
||||
def apply_participant_overrides(
|
||||
prompt_copy,
|
||||
participant_id,
|
||||
@@ -575,5 +618,13 @@ def apply_participant_overrides(
|
||||
master_url,
|
||||
enabled_json,
|
||||
)
|
||||
_override_image_list_nodes(
|
||||
prompt_copy,
|
||||
prompt_index,
|
||||
participant_id,
|
||||
enabled_worker_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
)
|
||||
|
||||
return prompt_copy
|
||||
|
||||
@@ -29,7 +29,9 @@ from .orchestration.prompt_transform import (
|
||||
generate_job_id_map,
|
||||
prepare_delegate_master_prompt,
|
||||
prune_prompt_for_worker,
|
||||
_resolve_direct_workers_for_image_list,
|
||||
)
|
||||
from ..image_list.runtime import validate_image_list_regions
|
||||
|
||||
|
||||
prompt_server = server.PromptServer.instance
|
||||
@@ -246,8 +248,42 @@ 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
|
||||
workers = _resolve_enabled_workers(config, requested_ids)
|
||||
prompt_index = PromptIndex(prompt_obj)
|
||||
image_list_start_ids = prompt_index.nodes_for_class("DistributedImageListStart")
|
||||
image_list_mode = bool(image_list_start_ids)
|
||||
if image_list_mode:
|
||||
# Validate before probing or queueing anything so malformed boundaries and
|
||||
# old worker releases fail before any upstream image is materialized.
|
||||
validate_image_list_regions(prompt_obj)
|
||||
other_distributed = (
|
||||
prompt_index.nodes_for_class("DistributedCollector")
|
||||
+ prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
|
||||
)
|
||||
if other_distributed:
|
||||
raise ValueError(
|
||||
"Distributed Image List v1 cannot be mixed with Distributed Collector or "
|
||||
"Ultimate SD Upscale nodes in one prompt"
|
||||
)
|
||||
configured_serverless = [
|
||||
str(worker.get("id")) for worker in config.get("workers", [])
|
||||
if worker.get("enabled", False)
|
||||
and str(worker.get("type") or "local").lower() == "serverless"
|
||||
and (requested_ids is None or str(worker.get("id")) in {str(x) for x in requested_ids})
|
||||
]
|
||||
if configured_serverless:
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
"Image List v1 explicitly excludes unsupported serverless pools: "
|
||||
+ ", ".join(configured_serverless),
|
||||
)
|
||||
workers = _resolve_direct_workers_for_image_list(config, requested_ids)
|
||||
if not workers:
|
||||
raise RuntimeError(
|
||||
"Distributed Image List v1 requires an enabled selected direct local/remote worker; "
|
||||
"serverless pools and master-only execution are unsupported"
|
||||
)
|
||||
else:
|
||||
workers = _resolve_enabled_workers(config, requested_ids)
|
||||
load_balance_requested = _prompt_requests_load_balance(prompt_index)
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
@@ -266,6 +302,10 @@ async def orchestrate_distributed_execution(
|
||||
# Respect master delegate-only configuration
|
||||
if delegate_master is None:
|
||||
delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False))
|
||||
if image_list_mode:
|
||||
# The master must execute each upstream dependency once to obtain the
|
||||
# full typed lists; only the enclosed branch is delegated.
|
||||
delegate_master = False
|
||||
|
||||
if not workers and delegate_master:
|
||||
trace_debug(
|
||||
@@ -281,8 +321,19 @@ async def orchestrate_distributed_execution(
|
||||
trace_execution_id=execution_trace_id,
|
||||
probe_concurrency=worker_probe_concurrency,
|
||||
)
|
||||
if image_list_mode:
|
||||
from .orchestration.dispatch import preflight_image_list_workers
|
||||
if not active_workers:
|
||||
raise RuntimeError(
|
||||
"No selected direct worker is reachable for Distributed Image List; "
|
||||
"the master will not duplicate per-image processing"
|
||||
)
|
||||
idle_ids = await preflight_image_list_workers(active_workers, require_idle=True)
|
||||
active_workers = [worker for worker in active_workers if worker["id"] in idle_ids]
|
||||
if not active_workers:
|
||||
raise RuntimeError("No feature-compatible direct worker is idle; no master job was queued")
|
||||
|
||||
if load_balance_requested:
|
||||
if load_balance_requested and not image_list_mode:
|
||||
candidate_workers = list(active_workers)
|
||||
if not delegate_master:
|
||||
# Include master in load balancing only when master participation is enabled.
|
||||
@@ -385,7 +436,7 @@ async def orchestrate_distributed_execution(
|
||||
+ ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers),
|
||||
)
|
||||
worker_payloads = []
|
||||
if active_workers:
|
||||
if active_workers and not image_list_mode:
|
||||
worker_prep_semaphore = asyncio.Semaphore(worker_prep_concurrency)
|
||||
media_sync_semaphore = asyncio.Semaphore(media_sync_concurrency)
|
||||
worker_payloads = await asyncio.gather(
|
||||
@@ -416,8 +467,9 @@ async def orchestrate_distributed_execution(
|
||||
manager.ensure_capacity(len(serverless_jobs))
|
||||
|
||||
# Do not leave collector queues behind when offline compilation/media validation fails.
|
||||
for job_id in job_id_map.values():
|
||||
await _ensure_distributed_queue(job_id)
|
||||
if not image_list_mode:
|
||||
for job_id in job_id_map.values():
|
||||
await _ensure_distributed_queue(job_id)
|
||||
|
||||
if regular_payloads:
|
||||
await asyncio.gather(
|
||||
@@ -468,4 +520,4 @@ async def orchestrate_distributed_execution(
|
||||
execution_trace_id,
|
||||
f"Orchestration complete: prompt_id={prompt_id}, dispatched_workers={len(worker_payloads)}, delegate_master={delegate_master}",
|
||||
)
|
||||
return prompt_id, prompt_number, len(worker_payloads), node_errors
|
||||
return prompt_id, prompt_number, (len(active_workers) if image_list_mode else len(worker_payloads)), node_errors
|
||||
|
||||
@@ -0,0 +1,72 @@
|
||||
# Dynamic image-list queue
|
||||
|
||||
Process a list of images across **direct local or remote ComfyUI workers** without assigning fixed chunks to each worker.
|
||||
|
||||
The master evaluates the input list once. Each selected idle worker receives one image job; when it returns a result, it gets the next unassigned image. Faster workers can therefore process more images. The master restores input order before running downstream nodes.
|
||||
|
||||
## Wiring
|
||||
|
||||
```text
|
||||
Image list or IMAGE batch
|
||||
|
|
||||
Distributed Image List Start -------- context ---------+
|
||||
| image |
|
||||
v |
|
||||
Per-image processing |
|
||||
| IMAGE |
|
||||
v v
|
||||
Distributed Image List End
|
||||
| image_list
|
||||
v
|
||||
Preview / Save / next stage
|
||||
```
|
||||
|
||||
1. Install the same feature-capable version of Distributed on the master and every selected worker. Restart those instances and refresh the frontend after installation.
|
||||
2. Enable Distributed and select at least one reachable, idle **direct** worker in the panel. The master must be reachable from remote workers at its configured callback address.
|
||||
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.
|
||||
6. Put preview/save nodes **after End**, then queue normally. API callers should submit through `/distributed/queue`, not directly through `/prompt`.
|
||||
|
||||
ComfyUI represents these as ordinary `IMAGE` sockets with list semantics, not a separate `IMAGE_LIST` socket type. Outside the region there is a list; on a worker, Start provides one frame and its aligned values. End returns ordered lists, not a concatenated tensor batch, so differently sized images can be preserved.
|
||||
|
||||
## Inputs and outputs
|
||||
|
||||
Start accepts:
|
||||
|
||||
- `images`: one image, an IMAGE batch, or a ComfyUI image list. A batch is split into individual frames.
|
||||
- Optional `reference_image`, `mask`, `text`, `strength`, and `seed`.
|
||||
- A `per_item` / `shared` mode for each optional input. `per_item` requires exactly one value per image; `shared` requires exactly one value, reused for all images. Inputs are never silently padded by repeating the last element.
|
||||
|
||||
Reference images and masks default to `per_item`. Text, strength and seed default to `shared`; strength defaults to `1.0`, seed to `0`.
|
||||
|
||||
Start outputs, in order: `image`, `reference_image`, `mask`, `text`, `strength`, `seed`, `item_index` (zero-based), `context`. Connect only optional values that your processing region actually provides and uses. Seeds are tied to input values, not worker assignments; use a per-item seed list for deterministic variation.
|
||||
|
||||
End accepts `images`, optional `masks` / `texts`, and `context`. It returns `image_list`, `masks`, and `texts` in original item order. Each job must produce exactly one image, at most one mask, and at most one text value. Unconnected optional outputs contain `None` entries.
|
||||
|
||||
Models are **loaded and cached on workers**, not serialized from the master. Connect loader outputs directly to the processing nodes. Each worker needs the region's node packs and model files. Per-image data must pass through Start; do not bypass it with a second image list connected directly into the processing region.
|
||||
|
||||
## Scope and failure behavior
|
||||
|
||||
- This version supports direct workers only, 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 workers are excluded. Workers joining later are not added to an already-running region.
|
||||
- Only the enclosed processing is distributed. List creation and downstream saving run on the master, even when master-delegate-only is enabled. There is no master fallback when workers are unavailable.
|
||||
- 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. Transport preserves floating-point tensor dtype and values rather than encoding to lossy images. Each tensor is limited to 50 MiB, and the complete JSON/base64 HTTP request to 101 MiB; image/reference/mask combinations can hit the aggregate request limit first.
|
||||
- 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 and the configured worker timeout fail the region. User interruption stops further scheduling and issues targeted queue deletion/interruption for this region's acknowledged prompts. An in-progress submission is allowed to resolve so its prompt ID can be cleaned up. Unreachable workers or a lost submission acknowledgement can require checking that worker manually; cancellation is not a guarantee of remote termination through a broken network.
|
||||
- State is in memory. Restarting the master does not resume an interrupted region. A region also has a one-hour overall deadline.
|
||||
|
||||
Use trusted workers and trusted workflows, as with the existing ComfyUI API. Capability tokens protect assignment/result matching; they are not a replacement for network access control.
|
||||
|
||||
## 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 lightweight queue/history facade is test code; no managed ComfyUI service or model is started.
|
||||
|
||||
```bash
|
||||
COMFYUI_SOURCE_ROOT=/path/to/ComfyUI /path/to/ComfyUI/.venv/bin/python -m pytest tests -q
|
||||
bash scripts/test-web.sh
|
||||
```
|
||||
|
||||
Without `COMFYUI_SOURCE_ROOT`, real-framework tests are explicitly skipped. This coverage is not a substitute for checking your model workflow and connectivity on the actual GPU worker machines.
|
||||
@@ -0,0 +1 @@
|
||||
"""Dynamic per-image image-list execution internals."""
|
||||
@@ -0,0 +1,798 @@
|
||||
"""Runtime primitives for the V3 dynamic image-list feature.
|
||||
|
||||
The network protocol transports typed float tensors as bounded JSON/base64 payloads;
|
||||
no Python object deserialization or arbitrary-path access is used.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import math
|
||||
import secrets
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
PROTOCOL_VERSION = 1
|
||||
START_CLASS = "DistributedImageListStart"
|
||||
END_CLASS = "DistributedImageListEnd"
|
||||
START_NODE_ID = "DistributedImageListStart"
|
||||
END_NODE_ID = "DistributedImageListEnd"
|
||||
# Start output indexes are a stable internal prompt protocol.
|
||||
IMAGE_OUT, REFERENCE_OUT, MASK_OUT, TEXT_OUT, STRENGTH_OUT, SEED_OUT, INDEX_OUT, CONTEXT_OUT = range(8)
|
||||
MAX_PAYLOAD_BYTES = 50 * 1024 * 1024
|
||||
MAX_ITEMS = 4096
|
||||
_ALLOWED_FLOAT_DTYPES = {"float16", "bfloat16", "float32", "float64"}
|
||||
|
||||
_LOCK = threading.RLock()
|
||||
_MASTER_JOBS: dict[str, "ImageListJob"] = {}
|
||||
_WORKER_ITEMS: dict[str, tuple[float, dict[str, Any]]] = {}
|
||||
_WORKER_ITEM_TTL_SECONDS = 3600.0
|
||||
|
||||
|
||||
class ImageListCancelledError(RuntimeError):
|
||||
"""Raised after owned worker prompts have been cleaned up on cancellation."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class ImageListJob:
|
||||
context_id: str
|
||||
start_id: str
|
||||
end_id: str
|
||||
original_prompt: dict[str, Any]
|
||||
items: list[dict[str, Any]]
|
||||
worker_ids: list[str]
|
||||
master_callback_url: str
|
||||
workflow_meta: dict[str, Any] | None = None
|
||||
results: dict[int, dict[str, Any]] = field(default_factory=dict)
|
||||
received_indexes: set[int] = field(default_factory=set)
|
||||
expected_tokens: dict[int, str] = field(default_factory=dict)
|
||||
expected_workers: dict[int, str] = field(default_factory=dict)
|
||||
event: asyncio.Event | None = None
|
||||
created_at: float = field(default_factory=time.monotonic)
|
||||
failed: str | None = None
|
||||
|
||||
|
||||
def _torch():
|
||||
import torch
|
||||
return torch
|
||||
|
||||
|
||||
def unwrap_single(value: Any) -> Any:
|
||||
"""Unwrap the one-element wrappers V3/INPUT_IS_LIST uses for scalar inputs."""
|
||||
while isinstance(value, (list, tuple)) and len(value) == 1:
|
||||
value = value[0]
|
||||
return value
|
||||
|
||||
|
||||
def _image_frames(value: Any, name: str, *, mask: bool = False) -> list[Any]:
|
||||
torch = _torch()
|
||||
if value is None:
|
||||
return []
|
||||
if isinstance(value, (list, tuple)):
|
||||
frames: list[Any] = []
|
||||
for item in value:
|
||||
frames.extend(_image_frames(item, name, mask=mask))
|
||||
return frames
|
||||
if not isinstance(value, torch.Tensor):
|
||||
raise TypeError(f"{name} must be a ComfyUI tensor or a list of tensors")
|
||||
if not value.is_floating_point():
|
||||
raise TypeError(f"{name} tensor must have a floating-point dtype")
|
||||
if mask:
|
||||
if value.ndim == 2:
|
||||
value = value.unsqueeze(0)
|
||||
elif value.ndim == 4 and value.shape[-1] == 1:
|
||||
value = value[..., 0]
|
||||
if value.ndim != 3:
|
||||
raise ValueError(f"{name} MASK must have shape [B,H,W] or [B,H,W,1]")
|
||||
if min(value.shape) < 1:
|
||||
raise ValueError(f"{name} MASK dimensions must be positive")
|
||||
return [value[i:i + 1] for i in range(value.shape[0])]
|
||||
if value.ndim == 3:
|
||||
value = value.unsqueeze(0)
|
||||
if value.ndim != 4:
|
||||
raise ValueError(f"{name} IMAGE must have shape [B,H,W,C] or [H,W,C]")
|
||||
if min(value.shape[:3]) < 1 or value.shape[-1] not in (3, 4):
|
||||
raise ValueError(f"{name} IMAGE must have positive dimensions and 3 or 4 channels")
|
||||
return [value[i:i + 1] for i in range(value.shape[0])]
|
||||
|
||||
|
||||
def _scalar_items(value: Any, name: str) -> list[Any]:
|
||||
if value is None:
|
||||
return []
|
||||
if isinstance(value, (list, tuple)):
|
||||
return list(value)
|
||||
return [value]
|
||||
|
||||
|
||||
def _normalize_aligned(value: Any, mode: str, count: int, name: str, kind: str) -> list[Any]:
|
||||
if value is None:
|
||||
return [None] * count
|
||||
if mode not in ("per_item", "shared"):
|
||||
raise ValueError(f"{name}_mode must be 'per_item' or 'shared'")
|
||||
if kind == "image":
|
||||
items = _image_frames(value, name)
|
||||
elif kind == "mask":
|
||||
items = _image_frames(value, name, mask=True)
|
||||
else:
|
||||
items = _scalar_items(value, name)
|
||||
if mode == "shared":
|
||||
if len(items) != 1:
|
||||
raise ValueError(f"{name} is shared but has {len(items)} values; shared inputs require exactly one value")
|
||||
items = [items[0] for _ in range(count)]
|
||||
if len(items) != count:
|
||||
raise ValueError(f"{name} has {len(items)} per-item values; expected exactly {count}")
|
||||
if kind == "text" and any(item is not None and not isinstance(item, str) for item in items):
|
||||
raise TypeError(f"{name} per-item values must be strings")
|
||||
if kind == "int" and any(item is not None and (isinstance(item, bool) or not isinstance(item, int)) for item in items):
|
||||
raise TypeError(f"{name} per-item values must be integers")
|
||||
if kind == "float" and any(item is not None and (isinstance(item, bool) or not isinstance(item, (int, float)) or not math.isfinite(float(item))) for item in items):
|
||||
raise TypeError(f"{name} per-item values must be finite numbers")
|
||||
return items
|
||||
|
||||
|
||||
def normalize_image_list_inputs(images: Any, reference_image: Any = None, mask: Any = None,
|
||||
text: Any = None, strength: Any = None, seed: Any = None,
|
||||
reference_image_mode: Any = "per_item", mask_mode: Any = "per_item",
|
||||
text_mode: Any = "per_item", strength_mode: Any = "per_item",
|
||||
seed_mode: Any = "per_item") -> list[dict[str, Any]]:
|
||||
"""Normalize batches and explicitly shared/per-item optional inputs."""
|
||||
frames = _image_frames(images, "images")
|
||||
if not frames:
|
||||
raise ValueError("Distributed Image List Start requires at least one image")
|
||||
if len(frames) > MAX_ITEMS:
|
||||
raise ValueError(f"Image list contains {len(frames)} items; maximum is {MAX_ITEMS}")
|
||||
count = len(frames)
|
||||
refs = _normalize_aligned(reference_image, str(unwrap_single(reference_image_mode)), count,
|
||||
"reference_image", "image")
|
||||
masks = _normalize_aligned(mask, str(unwrap_single(mask_mode)), count, "mask", "mask")
|
||||
texts = _normalize_aligned(text, str(unwrap_single(text_mode)), count, "text", "text")
|
||||
strengths = _normalize_aligned(strength, str(unwrap_single(strength_mode)), count, "strength", "float")
|
||||
seeds = _normalize_aligned(seed, str(unwrap_single(seed_mode)), count, "seed", "int")
|
||||
return [
|
||||
{"image": frames[i], "reference_image": refs[i], "mask": masks[i],
|
||||
"text": texts[i], "strength": strengths[i], "seed": seeds[i], "index": i}
|
||||
for i in range(count)
|
||||
]
|
||||
|
||||
|
||||
def _tensor_dtype_name(tensor: Any) -> str:
|
||||
name = str(tensor.dtype).removeprefix("torch.")
|
||||
if name not in _ALLOWED_FLOAT_DTYPES:
|
||||
raise TypeError(f"Tensor dtype {name!r} is unsupported; use float16, bfloat16, float32, or float64")
|
||||
return name
|
||||
|
||||
|
||||
def encode_value(value: Any, kind: str) -> dict[str, Any] | None:
|
||||
"""Encode only the explicit IMAGE/MASK/STRING/FLOAT/INT protocol types."""
|
||||
if value is None:
|
||||
return None
|
||||
if kind in ("IMAGE", "MASK"):
|
||||
torch = _torch()
|
||||
if not isinstance(value, torch.Tensor) or not value.is_floating_point():
|
||||
raise TypeError(f"{kind} result must be a floating-point tensor")
|
||||
if kind == "IMAGE":
|
||||
if value.ndim != 4 or value.shape[0] != 1 or value.shape[-1] not in (3, 4):
|
||||
raise ValueError("IMAGE item must have shape [1,H,W,3|4]")
|
||||
else:
|
||||
if value.ndim == 2:
|
||||
value = value.unsqueeze(0)
|
||||
if value.ndim == 4 and value.shape[-1] == 1:
|
||||
value = value[..., 0]
|
||||
if value.ndim != 3 or value.shape[0] != 1:
|
||||
raise ValueError("MASK item must have shape [1,H,W] or [1,H,W,1]")
|
||||
if value.numel() * value.element_size() > MAX_PAYLOAD_BYTES:
|
||||
raise ValueError(f"{kind} item exceeds the {MAX_PAYLOAD_BYTES}-byte transport limit")
|
||||
cpu = value.detach().to("cpu").contiguous()
|
||||
raw = cpu.view(torch.uint8).numpy().tobytes()
|
||||
if len(raw) > MAX_PAYLOAD_BYTES:
|
||||
raise ValueError(f"{kind} item exceeds the {MAX_PAYLOAD_BYTES}-byte transport limit")
|
||||
return {"type": kind, "dtype": _tensor_dtype_name(cpu),
|
||||
"shape": list(cpu.shape), "data": base64.b64encode(raw).decode("ascii")}
|
||||
if kind == "STRING":
|
||||
if not isinstance(value, str):
|
||||
raise TypeError("STRING transport value must be a string")
|
||||
return {"type": kind, "value": value}
|
||||
if kind == "INT":
|
||||
if isinstance(value, bool) or not isinstance(value, int):
|
||||
raise TypeError("INT transport value must be an integer")
|
||||
return {"type": kind, "value": value}
|
||||
if kind == "FLOAT":
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(float(value)):
|
||||
raise TypeError("FLOAT transport value must be finite")
|
||||
return {"type": kind, "value": float(value)}
|
||||
raise ValueError(f"Unsupported image-list transport type {kind!r}")
|
||||
|
||||
|
||||
def decode_value(data: Any, expected_kind: str) -> Any:
|
||||
if data is None:
|
||||
return None
|
||||
if not isinstance(data, dict) or data.get("type") != expected_kind:
|
||||
raise ValueError(f"Expected {expected_kind} transport value")
|
||||
if expected_kind in ("IMAGE", "MASK"):
|
||||
torch = _torch()
|
||||
dtype_name = data.get("dtype")
|
||||
if dtype_name not in _ALLOWED_FLOAT_DTYPES:
|
||||
raise ValueError(f"Unsupported tensor dtype {dtype_name!r}")
|
||||
shape = data.get("shape")
|
||||
if not isinstance(shape, list) or not shape or len(shape) > 4:
|
||||
raise ValueError("Tensor shape must be a list of 1 to 4 positive dimensions")
|
||||
if any(isinstance(dim, bool) or not isinstance(dim, int) or dim < 1 for dim in shape):
|
||||
raise ValueError("Tensor shape dimensions must be positive integers")
|
||||
expected_count = math.prod(shape)
|
||||
element_size = {"float16": 2, "bfloat16": 2, "float32": 4, "float64": 8}[dtype_name]
|
||||
expected_bytes = expected_count * element_size
|
||||
if expected_bytes > MAX_PAYLOAD_BYTES:
|
||||
raise ValueError("Tensor exceeds the image-list transport limit")
|
||||
encoded = data.get("data")
|
||||
if not isinstance(encoded, str) or len(encoded) > ((MAX_PAYLOAD_BYTES + 2) // 3) * 4:
|
||||
raise ValueError("Tensor data is missing or exceeds the transport limit")
|
||||
try:
|
||||
raw = base64.b64decode(encoded, validate=True)
|
||||
except Exception as exc:
|
||||
raise ValueError("Tensor data must be valid base64") from exc
|
||||
if len(raw) != expected_bytes:
|
||||
raise ValueError("Tensor byte length does not match shape and dtype")
|
||||
dtype = getattr(torch, dtype_name)
|
||||
return torch.frombuffer(bytearray(raw), dtype=torch.uint8).view(dtype).reshape(shape).clone()
|
||||
value = data.get("value")
|
||||
return encode_value(value, expected_kind) and value
|
||||
|
||||
|
||||
def encode_item(item: dict[str, Any]) -> dict[str, Any]:
|
||||
return {"image": encode_value(item["image"], "IMAGE"),
|
||||
"reference_image": encode_value(item.get("reference_image"), "IMAGE"),
|
||||
"mask": encode_value(item.get("mask"), "MASK"),
|
||||
"text": encode_value(item.get("text"), "STRING"),
|
||||
"strength": encode_value(item.get("strength"), "FLOAT"),
|
||||
"seed": encode_value(item.get("seed"), "INT"),
|
||||
"index": int(item["index"])}
|
||||
|
||||
|
||||
def decode_item(data: Any) -> dict[str, Any]:
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError("item must be an object")
|
||||
index = data.get("index")
|
||||
if isinstance(index, bool) or not isinstance(index, int) or index < 0:
|
||||
raise ValueError("item.index must be a non-negative integer")
|
||||
image = decode_value(data.get("image"), "IMAGE")
|
||||
if image is None or image.ndim != 4 or image.shape[0] != 1 or image.shape[-1] not in (3, 4):
|
||||
raise ValueError("item.image must contain exactly one RGB/RGBA image frame")
|
||||
reference = decode_value(data.get("reference_image"), "IMAGE")
|
||||
if reference is not None and (reference.ndim != 4 or reference.shape[0] != 1 or reference.shape[-1] not in (3, 4)):
|
||||
raise ValueError("item.reference_image must contain exactly one RGB/RGBA image frame")
|
||||
mask = decode_value(data.get("mask"), "MASK")
|
||||
if mask is not None and (mask.ndim != 3 or mask.shape[0] != 1):
|
||||
raise ValueError("item.mask must contain exactly one MASK frame")
|
||||
return {"image": image,
|
||||
"reference_image": reference,
|
||||
"mask": mask,
|
||||
"text": decode_value(data.get("text"), "STRING"),
|
||||
"strength": decode_value(data.get("strength"), "FLOAT"),
|
||||
"seed": decode_value(data.get("seed"), "INT"), "index": index}
|
||||
|
||||
|
||||
def _sweep_worker_items() -> None:
|
||||
now = time.monotonic()
|
||||
with _LOCK:
|
||||
for token, (expires, _) in list(_WORKER_ITEMS.items()):
|
||||
if expires <= now:
|
||||
del _WORKER_ITEMS[token]
|
||||
|
||||
|
||||
def store_worker_item(token: str, item: dict[str, Any]) -> None:
|
||||
if not isinstance(token, str) or len(token) < 32 or len(token) > 128:
|
||||
raise ValueError("token must be a 32-128 character capability")
|
||||
clean = decode_item(item)
|
||||
_sweep_worker_items()
|
||||
with _LOCK:
|
||||
if token in _WORKER_ITEMS:
|
||||
raise ValueError("duplicate item token")
|
||||
_WORKER_ITEMS[token] = (time.monotonic() + _WORKER_ITEM_TTL_SECONDS, clean)
|
||||
|
||||
|
||||
def consume_worker_item(token: str) -> dict[str, Any]:
|
||||
if not isinstance(token, str) or not token:
|
||||
raise ValueError("worker item token is required")
|
||||
_sweep_worker_items()
|
||||
with _LOCK:
|
||||
entry = _WORKER_ITEMS.pop(token, None)
|
||||
if entry is None:
|
||||
raise RuntimeError("Image-list item is missing, expired, or already consumed")
|
||||
return entry[1]
|
||||
|
||||
|
||||
def delete_worker_item(token: str) -> bool:
|
||||
with _LOCK:
|
||||
return _WORKER_ITEMS.pop(str(token), None) is not None
|
||||
|
||||
|
||||
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) -> str:
|
||||
if not isinstance(prompt, dict):
|
||||
raise ValueError("Image List Start requires the original ComfyUI prompt for region extraction")
|
||||
if not worker_ids:
|
||||
raise RuntimeError("Distributed Image List requires at least one preflighted direct worker; master fallback is disabled")
|
||||
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,
|
||||
workflow_meta)
|
||||
with _LOCK:
|
||||
_MASTER_JOBS[context_id] = state
|
||||
return context_id
|
||||
|
||||
|
||||
def get_master_job(context_id: str) -> ImageListJob | None:
|
||||
with _LOCK:
|
||||
return _MASTER_JOBS.get(context_id)
|
||||
|
||||
|
||||
def remove_master_job(context_id: str) -> None:
|
||||
with _LOCK:
|
||||
_MASTER_JOBS.pop(context_id, None)
|
||||
|
||||
|
||||
def _prompt_links(node: Any, prompt: dict[str, Any]):
|
||||
if not isinstance(node, dict):
|
||||
return
|
||||
inputs = node.get("inputs", {})
|
||||
if not isinstance(inputs, dict):
|
||||
return
|
||||
for input_name, value in inputs.items():
|
||||
if isinstance(value, (list, tuple)) and len(value) == 2 and str(value[0]) in prompt:
|
||||
yield input_name, str(value[0]), value[1]
|
||||
|
||||
|
||||
def _direct_link(value: Any, prompt: dict[str, Any], node_id: str, input_name: str) -> tuple[str, int]:
|
||||
if not (isinstance(value, (list, tuple)) and len(value) == 2 and str(value[0]) in prompt):
|
||||
raise ValueError(f"{node_id}.{input_name} must be connected to an in-graph node")
|
||||
try:
|
||||
return str(value[0]), int(value[1])
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError(f"{node_id}.{input_name} has an invalid link") from exc
|
||||
|
||||
|
||||
def _ancestors_for_inputs(prompt: dict[str, Any], end_id: str, start_id: str,
|
||||
input_names: tuple[str, ...]) -> set[str]:
|
||||
result = {end_id}
|
||||
pending = []
|
||||
end_inputs = prompt[end_id].get("inputs", {})
|
||||
for name in input_names:
|
||||
if name in end_inputs:
|
||||
source, _ = _direct_link(end_inputs[name], prompt, end_id, name)
|
||||
pending.append(source)
|
||||
while pending:
|
||||
node_id = pending.pop()
|
||||
if node_id in result:
|
||||
continue
|
||||
result.add(node_id)
|
||||
if node_id == start_id:
|
||||
continue
|
||||
node = prompt[node_id]
|
||||
for _, source, _ in _prompt_links(node, prompt):
|
||||
if source not in result:
|
||||
pending.append(source)
|
||||
return result
|
||||
|
||||
|
||||
def _downstream_until(prompt: dict[str, Any], start_id: str, end_id: str) -> set[str]:
|
||||
downstream: dict[str, set[str]] = {str(node_id): set() for node_id in prompt}
|
||||
for node_id, node in prompt.items():
|
||||
for _, source, _ in _prompt_links(node, prompt):
|
||||
downstream[source].add(str(node_id))
|
||||
visited = set()
|
||||
pending = [start_id]
|
||||
while pending:
|
||||
node_id = pending.pop()
|
||||
for target in downstream.get(node_id, ()):
|
||||
if target == end_id or target in visited:
|
||||
continue
|
||||
visited.add(target)
|
||||
pending.append(target)
|
||||
return visited
|
||||
|
||||
|
||||
def validate_image_list_regions(prompt: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
"""Validate pairing, non-nesting, and closed per-image graph boundaries."""
|
||||
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 = []
|
||||
for start_id in starts:
|
||||
candidates = []
|
||||
for end_id in ends:
|
||||
context = prompt[end_id].get("inputs", {}).get("context")
|
||||
if isinstance(context, (list, tuple)) and len(context) == 2 and str(context[0]) == start_id and context[1] == CONTEXT_OUT:
|
||||
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]
|
||||
if "images" not in prompt[end_id].get("inputs", {}):
|
||||
raise ValueError(f"End node {end_id}.images must be connected to the per-item result branch")
|
||||
end_inputs = prompt[end_id].get("inputs", {})
|
||||
for name in ("images", "masks", "texts"):
|
||||
if name in end_inputs:
|
||||
_direct_link(end_inputs[name], prompt, end_id, name)
|
||||
region_nodes = _ancestors_for_inputs(prompt, end_id, start_id, ("images", "masks", "texts"))
|
||||
if start_id not in region_nodes:
|
||||
raise ValueError(f"Start node {start_id} is not an ancestor of End node {end_id}")
|
||||
forbidden = {n for n in region_nodes if n != start_id and n != end_id and
|
||||
isinstance(prompt[n], dict) and prompt[n].get("class_type") in
|
||||
{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", {})
|
||||
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")
|
||||
downstream = _downstream_until(prompt, start_id, end_id)
|
||||
outside = sorted(downstream - region_nodes)
|
||||
if outside:
|
||||
raise ValueError(f"Image-list region {start_id}->{end_id} has unsafe cross-boundary nodes: {outside}")
|
||||
# Model loaders and constants can be shared; item-dependent processing cannot.
|
||||
regions.append({"start_id": start_id, "end_id": end_id,
|
||||
"node_ids": region_nodes, "branch_ids": (region_nodes & downstream) - {start_id}})
|
||||
regions.sort(key=lambda region: (0, int(region["start_id"])) if region["start_id"].isdigit()
|
||||
else (1, region["start_id"]))
|
||||
internal_sets = [r["branch_ids"] for r in regions]
|
||||
for i in range(len(regions)):
|
||||
for j in range(i + 1, len(regions)):
|
||||
overlap = internal_sets[i] & internal_sets[j]
|
||||
if overlap:
|
||||
raise ValueError(f"Image-list regions share ambiguous processing nodes: {sorted(overlap)}")
|
||||
return regions
|
||||
|
||||
|
||||
def build_worker_prompt(prompt: dict[str, Any], region: dict[str, Any], item_token: str,
|
||||
context_id: str, item_index: int, callback_url: str) -> dict[str, Any]:
|
||||
"""Copy only the per-image branch and its worker-side dependencies."""
|
||||
start_id, end_id = region["start_id"], region["end_id"]
|
||||
_direct_link(prompt[end_id].get("inputs", {}).get("context"), prompt, end_id, "context")
|
||||
kept = {node_id: json.loads(json.dumps(prompt[node_id])) for node_id in region["node_ids"]}
|
||||
start_inputs = kept[start_id].setdefault("inputs", {})
|
||||
for name in ("images", "reference_image", "mask", "text", "strength", "seed"):
|
||||
start_inputs.pop(name, None)
|
||||
start_inputs.update({"worker_item_token": item_token, "worker_context_id": context_id})
|
||||
end_inputs = kept[end_id].setdefault("inputs", {})
|
||||
end_inputs.update({"worker_result_token": item_token,
|
||||
"worker_item_index": int(item_index),
|
||||
"master_callback_url": callback_url})
|
||||
return kept
|
||||
|
||||
|
||||
def create_worker_task_token() -> str:
|
||||
return secrets.token_urlsafe(36)
|
||||
|
||||
|
||||
async def dynamic_schedule(item_count: int, workers: list[str], dispatch, wait_next,
|
||||
cleanup, interrupted=None):
|
||||
"""One in-flight job per worker; completed workers immediately take next item."""
|
||||
if item_count <= 0:
|
||||
raise ValueError("Dynamic image-list queue has no items")
|
||||
if not workers:
|
||||
raise RuntimeError("No free, reachable, feature-compatible direct workers are available")
|
||||
pending_index = 0
|
||||
inflight: dict[str, Any] = {}
|
||||
results: dict[int, Any] = {}
|
||||
|
||||
def check_interrupted():
|
||||
if interrupted is not None and interrupted():
|
||||
raise ImageListCancelledError("Image-list job interrupted by user")
|
||||
|
||||
async def next_result():
|
||||
# Keep one waiter alive while checking ComfyUI's thread-safe interrupt flag.
|
||||
# Restarting wait_next on a timeout can lose a result it already consumed.
|
||||
waiter = asyncio.create_task(wait_next(inflight))
|
||||
try:
|
||||
while True:
|
||||
check_interrupted()
|
||||
done, _ = await asyncio.wait({waiter}, timeout=0.1)
|
||||
check_interrupted()
|
||||
if done:
|
||||
return waiter.result()
|
||||
finally:
|
||||
if not waiter.done():
|
||||
waiter.cancel()
|
||||
await asyncio.gather(waiter, return_exceptions=True)
|
||||
|
||||
async def submit(worker, index):
|
||||
# An outer timeout must not discard an acknowledgement for a job we own.
|
||||
task = asyncio.create_task(dispatch(worker, index))
|
||||
try:
|
||||
return await asyncio.shield(task)
|
||||
except asyncio.CancelledError:
|
||||
try:
|
||||
inflight[worker] = await task
|
||||
except Exception:
|
||||
pass
|
||||
raise
|
||||
|
||||
try:
|
||||
for worker in workers:
|
||||
check_interrupted()
|
||||
if pending_index >= item_count:
|
||||
break
|
||||
inflight[worker] = await submit(worker, pending_index)
|
||||
pending_index += 1
|
||||
while inflight:
|
||||
check_interrupted()
|
||||
worker, index, result = await next_result()
|
||||
if worker not in inflight:
|
||||
raise RuntimeError(f"Unexpected completion from unassigned worker {worker}")
|
||||
if isinstance(index, bool) or not isinstance(index, int) or index != inflight[worker].get("index"):
|
||||
raise RuntimeError(f"Worker {worker} returned an unexpected image index {index!r}")
|
||||
if index in results:
|
||||
raise RuntimeError(f"Duplicate image-list result for item {index}")
|
||||
results[index] = result
|
||||
del inflight[worker]
|
||||
if pending_index < item_count:
|
||||
inflight[worker] = await submit(worker, pending_index)
|
||||
pending_index += 1
|
||||
missing = sorted(set(range(item_count)) - set(results))
|
||||
if missing:
|
||||
raise RuntimeError(f"Image-list results are incomplete; missing indexes {missing}")
|
||||
return [results[i] for i in range(item_count)]
|
||||
except BaseException:
|
||||
if inflight:
|
||||
await asyncio.gather(*(cleanup(worker, handle) for worker, handle in inflight.items()), return_exceptions=True)
|
||||
raise
|
||||
|
||||
|
||||
def set_master_result(context_id: str, item_index: int, token: str, result: dict[str, Any]) -> None:
|
||||
with _LOCK:
|
||||
state = _MASTER_JOBS.get(context_id)
|
||||
if state is None:
|
||||
raise ValueError("Unknown or expired image-list context")
|
||||
expected = state.expected_tokens.get(item_index)
|
||||
if expected is None or not secrets.compare_digest(expected, token):
|
||||
raise ValueError("Image-list result token does not match an assigned item")
|
||||
if item_index in state.received_indexes:
|
||||
raise ValueError(f"Duplicate image-list result for item {item_index}")
|
||||
state.received_indexes.add(item_index)
|
||||
state.results[item_index] = result
|
||||
event = state.event
|
||||
if event is not None:
|
||||
event.set()
|
||||
|
||||
|
||||
def register_assignment(context_id: str, item_index: int, worker_id: str, token: str) -> None:
|
||||
state = get_master_job(context_id)
|
||||
if state is None:
|
||||
raise RuntimeError("Image-list context expired before worker dispatch")
|
||||
with _LOCK:
|
||||
if item_index in state.expected_tokens:
|
||||
raise RuntimeError(f"Image-list item {item_index} was already assigned")
|
||||
state.expected_tokens[item_index] = token
|
||||
state.expected_workers[item_index] = worker_id
|
||||
|
||||
|
||||
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")
|
||||
if image is None or image.ndim != 4 or image.shape[0] != 1 or image.shape[-1] not in (3, 4):
|
||||
raise ValueError("result.image must contain exactly one RGB/RGBA image")
|
||||
mask = decode_value(result.get("mask"), "MASK")
|
||||
if mask is not None and (mask.ndim != 3 or mask.shape[0] != 1):
|
||||
raise ValueError("result.mask must contain exactly one MASK frame")
|
||||
text = decode_value(result.get("text"), "STRING")
|
||||
return {"image": image, "mask": mask, "text": text}
|
||||
|
||||
|
||||
def normalize_worker_result(images: Any, masks: Any = None, texts: Any = None) -> dict[str, Any]:
|
||||
def one(value, name, kind):
|
||||
if value is None:
|
||||
return None
|
||||
if kind in ("IMAGE", "MASK"):
|
||||
values = _image_frames(value, name, mask=(kind == "MASK"))
|
||||
if len(values) != 1:
|
||||
raise ValueError(f"Worker End expects exactly one {kind} result, received {len(values)}")
|
||||
return values[0]
|
||||
values = _scalar_items(value, name)
|
||||
values = [v for v in values if v is not None]
|
||||
if not values:
|
||||
return None
|
||||
if len(values) != 1:
|
||||
raise ValueError(f"Worker End expects at most one {kind} result, received {len(values)}")
|
||||
return values[0]
|
||||
image = one(images, "images", "IMAGE")
|
||||
if image is None:
|
||||
raise ValueError("Worker End received no image result")
|
||||
mask = one(masks, "masks", "MASK")
|
||||
text = one(texts, "texts", "STRING")
|
||||
if text is not None and not isinstance(text, str):
|
||||
raise TypeError("Worker End text result must be a string")
|
||||
return {"image": image, "mask": mask, "text": text}
|
||||
|
||||
|
||||
def encode_result(result: dict[str, Any]) -> dict[str, Any]:
|
||||
normalized = normalize_worker_result(result.get("image"), result.get("mask"), result.get("text"))
|
||||
return {"image": encode_value(normalized["image"], "IMAGE"),
|
||||
"mask": encode_value(normalized["mask"], "MASK"),
|
||||
"text": encode_value(normalized["text"], "STRING")}
|
||||
|
||||
|
||||
def _collect_output_lists(results: list[dict[str, Any]]) -> tuple[list[Any], list[Any], list[Any]]:
|
||||
if not results:
|
||||
raise RuntimeError("Image-list execution returned no results")
|
||||
images, masks, texts = [], [], []
|
||||
for expected_index, result in enumerate(results):
|
||||
if not isinstance(result, dict) or "image" not in result:
|
||||
raise RuntimeError(f"Image-list result {expected_index} is missing its image")
|
||||
images.append(result["image"])
|
||||
masks.append(result.get("mask"))
|
||||
texts.append(result.get("text"))
|
||||
return images, masks, texts
|
||||
|
||||
|
||||
async def run_image_list_job(context_id: str) -> tuple[list[Any], list[Any], list[Any]]:
|
||||
"""Resolve compatible idle direct workers and schedule one image at a time."""
|
||||
state = get_master_job(context_id)
|
||||
if state is None:
|
||||
raise RuntimeError("Distributed Image List context is missing or expired")
|
||||
try:
|
||||
return await _run_image_list_job(state)
|
||||
finally:
|
||||
remove_master_job(context_id)
|
||||
|
||||
|
||||
async def _run_image_list_job(state: ImageListJob):
|
||||
context_id = state.context_id
|
||||
state.event = asyncio.Event()
|
||||
from ..api.orchestration.dispatch import (
|
||||
preflight_image_list_workers, dispatch_worker_prompt,
|
||||
)
|
||||
from ..api.orchestration.prompt_transform import _resolve_direct_workers_for_image_list
|
||||
from ..utils.config import load_config, get_worker_timeout_seconds
|
||||
from ..utils.network import build_master_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:
|
||||
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=True)
|
||||
workers = [worker for worker in workers if worker["id"] in worker_statuses]
|
||||
if not workers:
|
||||
raise RuntimeError("No free, reachable, feature-compatible direct workers are available; master execution is disabled")
|
||||
session = await get_client_session()
|
||||
callback_url = state.master_callback_url or build_master_url(config=config)
|
||||
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_ids: set[str] = set()
|
||||
worker_path_separators: dict[str, str | None] = {}
|
||||
|
||||
async def dispatch(worker_id: str, index: int) -> dict[str, Any]:
|
||||
worker = next((entry for entry in workers if entry["id"] == worker_id), None)
|
||||
if worker is None:
|
||||
raise RuntimeError(f"Worker {worker_id} was removed before 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])}
|
||||
input_url = build_worker_url(worker, "/distributed/image-list/item")
|
||||
async with session.post(input_url, json={"token": token, **item_payload},
|
||||
timeout=aiohttp.ClientTimeout(total=60)) as response:
|
||||
if response.status >= 400:
|
||||
raise RuntimeError(f"Worker {worker_id} rejected image-list item {index}: HTTP {response.status}")
|
||||
worker_prompt = build_worker_prompt(state.original_prompt, region, token, context_id,
|
||||
index, callback_url)
|
||||
if worker_id not in worker_path_separators:
|
||||
worker_path_separators[worker_id] = None
|
||||
if worker.get("host") and str(worker.get("type") or "local").lower() != "local":
|
||||
from ..api.orchestration.media_sync import fetch_worker_path_separator
|
||||
worker_path_separators[worker_id] = await fetch_worker_path_separator(worker)
|
||||
separator = worker_path_separators[worker_id]
|
||||
if separator:
|
||||
from ..api.orchestration.media_sync import convert_paths_for_platform
|
||||
worker_prompt = convert_paths_for_platform(worker_prompt, separator)
|
||||
if worker_id not in prepared_worker_ids:
|
||||
from ..api.orchestration.media_sync import sync_worker_media
|
||||
await asyncio.wait_for(sync_worker_media(worker, worker_prompt), timeout=60)
|
||||
prepared_worker_ids.add(worker_id)
|
||||
prompt_id = await dispatch_worker_prompt(
|
||||
worker, worker_prompt, state.workflow_meta,
|
||||
client_id=f"image-list-{context_id[:20]}", use_websocket=use_websocket,
|
||||
return_prompt_id=True,
|
||||
)
|
||||
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")
|
||||
return {"prompt_id": prompt_id, "token": token, "index": index,
|
||||
"started_at": time.monotonic(), "worker": worker}
|
||||
|
||||
async def wait_next(inflight: dict[str, Any]):
|
||||
while True:
|
||||
if state.failed:
|
||||
raise RuntimeError(state.failed)
|
||||
now = time.monotonic()
|
||||
expired = [f"{wid}/item {handle['index']}" for wid, handle in inflight.items()
|
||||
if now - handle["started_at"] > timeout_seconds]
|
||||
if expired:
|
||||
raise TimeoutError(f"Image-list worker timeout after {timeout_seconds:.1f}s: {', '.join(expired)}")
|
||||
# Treat worker-side execution errors as failures even when End never posts a result.
|
||||
async def history_status(worker_id, handle):
|
||||
url = build_worker_url(handle["worker"], f"/history/{handle['prompt_id']}")
|
||||
try:
|
||||
async with session.get(url, timeout=aiohttp.ClientTimeout(total=3)) as response:
|
||||
if response.status >= 400:
|
||||
return None
|
||||
data = await response.json()
|
||||
entry = data.get(handle["prompt_id"])
|
||||
if not isinstance(entry, dict):
|
||||
return None
|
||||
status = entry.get("status", {}) or {}
|
||||
if status.get("status_str") == "error":
|
||||
return f"Worker {worker_id} failed image-list item {handle['index']}: {status.get('messages', status)}"
|
||||
if status.get("completed") and status.get("status_str") not in (None, "success"):
|
||||
return f"Worker {worker_id} failed image-list item {handle['index']}: {status}"
|
||||
if status.get("completed"):
|
||||
with _LOCK:
|
||||
result = state.results.get(handle["index"])
|
||||
if result is None:
|
||||
return f"Worker {worker_id} completed item {handle['index']} without posting a result"
|
||||
return worker_id, handle["index"], result
|
||||
except (asyncio.TimeoutError, aiohttp.ClientError):
|
||||
return None
|
||||
return None
|
||||
errors = await asyncio.gather(*(history_status(wid, handle) for wid, handle in inflight.items()))
|
||||
failure = next((error for error in errors if isinstance(error, str)), None)
|
||||
if failure:
|
||||
raise RuntimeError(failure)
|
||||
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
|
||||
# successful history before releasing the slot / next region.
|
||||
with _LOCK:
|
||||
state.results.pop(completed[1], None)
|
||||
return completed
|
||||
await asyncio.sleep(0.25)
|
||||
|
||||
async def cleanup(worker_id: str, handle: dict[str, Any]) -> None:
|
||||
worker = handle["worker"]
|
||||
try:
|
||||
prompt_id = handle["prompt_id"]
|
||||
# Delete if still queued, then interrupt only this prompt if it began
|
||||
# running meanwhile. Never send a global interrupt or clear a queue.
|
||||
async with session.post(build_worker_url(worker, "/queue"), json={"delete": [prompt_id]},
|
||||
timeout=aiohttp.ClientTimeout(total=5)) as response:
|
||||
response.raise_for_status()
|
||||
async with session.post(build_worker_url(worker, "/interrupt"), json={"prompt_id": prompt_id},
|
||||
timeout=aiohttp.ClientTimeout(total=5)) as response:
|
||||
response.raise_for_status()
|
||||
except Exception as exc:
|
||||
log(f"[Distributed Image List] Owned prompt cleanup failed for worker {worker_id}: {exc}")
|
||||
try:
|
||||
async with session.delete(build_worker_url(worker, f"/distributed/image-list/item/{handle['token']}"),
|
||||
timeout=aiohttp.ClientTimeout(total=5)):
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def interrupted() -> bool:
|
||||
try:
|
||||
import comfy.model_management
|
||||
return bool(comfy.model_management.processing_interrupted())
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
results = await dynamic_schedule(len(state.items), [worker["id"] for worker in workers],
|
||||
dispatch, wait_next, cleanup, interrupted)
|
||||
return _collect_output_lists(results)
|
||||
@@ -0,0 +1,226 @@
|
||||
"""Public V3 Start/End nodes for dynamic per-image image-list execution."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
from comfy_api.v0_0_2 import io
|
||||
|
||||
from ..utils.async_helpers import run_async_in_server_loop
|
||||
from ..image_list.runtime import (
|
||||
create_master_job,
|
||||
normalize_image_list_inputs,
|
||||
normalize_worker_result,
|
||||
run_image_list_job,
|
||||
encode_result,
|
||||
unwrap_single,
|
||||
)
|
||||
|
||||
|
||||
class DistributedImageListStart(io.ComfyNode):
|
||||
"""Materialize a batch once, then expose ordered per-item typed values."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="DistributedImageListStart",
|
||||
display_name="Distributed Image List Start",
|
||||
category="image/distributed",
|
||||
description=("Starts one per-image processing region. Optional inputs explicitly choose shared "
|
||||
"(exactly one value) or per-item (exactly one value per image) alignment."),
|
||||
inputs=[
|
||||
io.Image.Input("images", optional=True, tooltip="Image list or image batch; each frame becomes one worker job."),
|
||||
io.Image.Input("reference_image", optional=True),
|
||||
io.Combo.Input("reference_image_mode", options=["per_item", "shared"], default="per_item"),
|
||||
io.Mask.Input("mask", optional=True),
|
||||
io.Combo.Input("mask_mode", options=["per_item", "shared"], default="per_item"),
|
||||
io.String.Input("text", optional=True),
|
||||
io.Combo.Input("text_mode", options=["per_item", "shared"], default="shared"),
|
||||
io.Float.Input("strength", optional=True, default=1.0),
|
||||
io.Combo.Input("strength_mode", options=["per_item", "shared"], default="shared"),
|
||||
io.Int.Input("seed", optional=True, default=0, min=0, max=0xffffffffffffffff),
|
||||
io.Combo.Input("seed_mode", options=["per_item", "shared"], default="shared"),
|
||||
],
|
||||
outputs=[
|
||||
io.Image.Output(display_name="image", is_output_list=True),
|
||||
io.Image.Output(display_name="reference_image", is_output_list=True),
|
||||
io.Mask.Output(display_name="mask", is_output_list=True),
|
||||
io.String.Output(display_name="text", is_output_list=True),
|
||||
io.Float.Output(display_name="strength", is_output_list=True),
|
||||
io.Int.Output(display_name="seed", is_output_list=True),
|
||||
io.Int.Output(display_name="item_index", is_output_list=True),
|
||||
io.Custom("DISTRIBUTED_IMAGE_LIST_CONTEXT").Output(display_name="context", is_output_list=True),
|
||||
],
|
||||
is_input_list=True,
|
||||
accept_all_inputs=True,
|
||||
not_idempotent=True,
|
||||
hidden=[io.Hidden.prompt, io.Hidden.unique_id],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, images=None, reference_image=None, reference_image_mode="per_item",
|
||||
mask=None, mask_mode="per_item", text=None, text_mode="shared",
|
||||
strength=None, strength_mode="shared", seed=None, seed_mode="shared",
|
||||
worker_item_token="", worker_context_id="", image_list_enabled_worker_ids="[]",
|
||||
image_list_master_callback="", image_list_execution_id="",
|
||||
image_list_workflow_meta=None):
|
||||
worker_item_token = unwrap_single(worker_item_token)
|
||||
worker_context_id = unwrap_single(worker_context_id)
|
||||
image_list_workflow_meta = unwrap_single(image_list_workflow_meta)
|
||||
if worker_item_token:
|
||||
item = run_async_in_server_loop(_consume_worker_item(worker_item_token), timeout=10)
|
||||
context_id = worker_context_id
|
||||
if not isinstance(context_id, str) or not context_id:
|
||||
raise RuntimeError("Worker image-list prompt is missing its paired context")
|
||||
return io.NodeOutput(
|
||||
[item["image"]], [item.get("reference_image")], [item.get("mask")],
|
||||
[item.get("text")], [item.get("strength")], [item.get("seed")],
|
||||
[item["index"]], [context_id],
|
||||
)
|
||||
|
||||
items = normalize_image_list_inputs(
|
||||
images, reference_image, mask, text, strength, seed,
|
||||
reference_image_mode, mask_mode, text_mode, strength_mode, seed_mode,
|
||||
)
|
||||
prompt = getattr(cls.hidden, "prompt", None)
|
||||
start_id = str(getattr(cls.hidden, "unique_id", ""))
|
||||
if not prompt or not start_id:
|
||||
raise RuntimeError("Distributed Image List Start needs a submitted ComfyUI prompt context")
|
||||
try:
|
||||
worker_ids = json.loads(str(unwrap_single(image_list_enabled_worker_ids or "[]")))
|
||||
except (TypeError, ValueError) as exc:
|
||||
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 ..image_list.runtime import validate_image_list_regions
|
||||
regions = validate_image_list_regions(prompt)
|
||||
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}")
|
||||
context_id = create_master_job(
|
||||
start_id, region["end_id"], prompt, items, worker_ids,
|
||||
str(unwrap_single(image_list_master_callback or "")),
|
||||
image_list_workflow_meta if isinstance(image_list_workflow_meta, dict) else None,
|
||||
)
|
||||
return io.NodeOutput(
|
||||
[item["image"] for item in items],
|
||||
[item.get("reference_image") for item in items],
|
||||
[item.get("mask") for item in items],
|
||||
[item.get("text") for item in items],
|
||||
[item.get("strength") for item in items],
|
||||
[item.get("seed") for item in items],
|
||||
[item["index"] for item in items],
|
||||
[context_id for _ in items],
|
||||
)
|
||||
|
||||
|
||||
async def _consume_worker_item(token):
|
||||
# Item state is local to each worker process; source nodes consume it exactly once.
|
||||
from ..image_list.runtime import consume_worker_item
|
||||
return consume_worker_item(token)
|
||||
|
||||
|
||||
class DistributedImageListEnd(io.ComfyNode):
|
||||
"""Collect branch results in source order without evaluating the branch on master."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="DistributedImageListEnd",
|
||||
display_name="Distributed Image List End",
|
||||
category="image/distributed",
|
||||
description="Closes a paired image-list region and returns ordered image, mask, and text lists.",
|
||||
inputs=[
|
||||
io.Image.Input("images", lazy=True),
|
||||
io.Mask.Input("masks", optional=True, lazy=True),
|
||||
io.String.Input("texts", optional=True, lazy=True),
|
||||
io.Custom("DISTRIBUTED_IMAGE_LIST_CONTEXT").Input("context"),
|
||||
],
|
||||
outputs=[
|
||||
io.Image.Output(display_name="image_list", is_output_list=True),
|
||||
io.Mask.Output(display_name="masks", is_output_list=True),
|
||||
io.String.Output(display_name="texts", is_output_list=True),
|
||||
],
|
||||
is_input_list=True,
|
||||
is_output_node=True,
|
||||
accept_all_inputs=True,
|
||||
not_idempotent=True,
|
||||
hidden=[io.Hidden.prompt, io.Hidden.unique_id],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def check_lazy_status(cls, images=None, masks=None, texts=None, context=None,
|
||||
worker_result_token="", **kwargs):
|
||||
worker_result_token = unwrap_single(worker_result_token)
|
||||
if worker_result_token:
|
||||
# Decide which optional sockets are connected from the actual worker prompt;
|
||||
# absent optional inputs must not be demanded from Comfy's lazy executor.
|
||||
prompt = getattr(cls.hidden, "prompt", {}) or {}
|
||||
end_inputs = next((node.get("inputs", {}) for node in prompt.values()
|
||||
if isinstance(node, dict)
|
||||
and node.get("class_type") == "DistributedImageListEnd"
|
||||
and node.get("inputs", {}).get("worker_result_token") == worker_result_token), {})
|
||||
required = []
|
||||
for name, value in (("images", images), ("masks", masks), ("texts", texts)):
|
||||
connected = name == "images" or name in end_inputs
|
||||
missing = value is None or (isinstance(value, (list, tuple)) and all(v is None for v in value))
|
||||
if connected and missing:
|
||||
required.append(name)
|
||||
return required
|
||||
# Master only needs Start's context. Never request branch outputs here.
|
||||
if context is None or (isinstance(context, (tuple, list)) and all(item is None for item in context)):
|
||||
return ["context"]
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def execute(cls, images=None, masks=None, texts=None, context=None,
|
||||
worker_result_token="", worker_item_index=-1, master_callback_url=""):
|
||||
worker_result_token = unwrap_single(worker_result_token)
|
||||
master_callback_url = unwrap_single(master_callback_url)
|
||||
if worker_result_token:
|
||||
context_values = [item for item in (context if isinstance(context, (list, tuple)) else [context])
|
||||
if isinstance(item, str) and item]
|
||||
if not context_values or len(set(context_values)) != 1:
|
||||
raise RuntimeError("Worker Image List End is missing its paired context")
|
||||
result = normalize_worker_result(images, masks, texts)
|
||||
payload = encode_result(result)
|
||||
run_async_in_server_loop(
|
||||
submit_worker_result(str(master_callback_url), str(worker_result_token),
|
||||
context_values[0], int(unwrap_single(worker_item_index)), payload), timeout=60,
|
||||
)
|
||||
return io.NodeOutput([result["image"]], [result.get("mask")], [result.get("text")])
|
||||
|
||||
context_values = [item for item in (context if isinstance(context, (list, tuple)) else [context])
|
||||
if isinstance(item, str) and item]
|
||||
if not context_values or len(set(context_values)) != 1:
|
||||
raise RuntimeError("Distributed Image List End requires the paired Start context wire")
|
||||
context_id = context_values[0]
|
||||
from ..image_list.runtime import ImageListCancelledError
|
||||
try:
|
||||
# Put the deadline inside the coroutine so timeout cancels owned remote work.
|
||||
import asyncio
|
||||
result_lists = run_async_in_server_loop(asyncio.wait_for(run_image_list_job(context_id), 3600))
|
||||
except ImageListCancelledError:
|
||||
import comfy.model_management
|
||||
raise comfy.model_management.InterruptProcessingException()
|
||||
images_out, masks_out, texts_out = result_lists
|
||||
if len(images_out) != len(masks_out) or len(images_out) != len(texts_out):
|
||||
raise RuntimeError("Image-list collector received mismatched image, mask, or text result counts")
|
||||
return io.NodeOutput(images_out, masks_out, texts_out)
|
||||
|
||||
|
||||
async def submit_worker_result(callback_url, item_token, context_id, item_index, result):
|
||||
if not callback_url.startswith(("http://", "https://")):
|
||||
raise RuntimeError("Worker image-list callback URL is missing or invalid")
|
||||
import aiohttp
|
||||
from ..utils.network import get_client_session
|
||||
session = await get_client_session()
|
||||
payload = {"context_id": context_id, "token": item_token,
|
||||
"item_index": item_index, "result": result}
|
||||
async with session.post(
|
||||
f"{callback_url.rstrip('/')}/distributed/image-list/result",
|
||||
json=payload,
|
||||
timeout=aiohttp.ClientTimeout(total=60),
|
||||
) as response:
|
||||
if response.status >= 400:
|
||||
body = (await response.text())[:500]
|
||||
raise RuntimeError(f"Master rejected image-list result (HTTP {response.status}): {body}")
|
||||
+3
-1
@@ -9,6 +9,7 @@ from comfy_api.v0_0_2 import io
|
||||
from . import utilities as _utilities
|
||||
from .collector import DistributedCollectorNode as _CollectorRuntime
|
||||
from .distributed_upscale import UltimateSDUpscaleDistributed as _UpscaleRuntime
|
||||
from .image_list import DistributedImageListEnd, DistributedImageListStart
|
||||
|
||||
|
||||
class DistributedSeed(io.ComfyNode):
|
||||
@@ -178,4 +179,5 @@ class UltimateSDUpscaleDistributed(io.ComfyNode):
|
||||
|
||||
|
||||
NODES = [DistributedCollector, DistributedSeed, DistributedModelName, DistributedValue,
|
||||
ImageBatchDivider, AudioBatchDivider, DistributedEmptyImage, UltimateSDUpscaleDistributed]
|
||||
ImageBatchDivider, AudioBatchDivider, DistributedEmptyImage, UltimateSDUpscaleDistributed,
|
||||
DistributedImageListStart, DistributedImageListEnd]
|
||||
|
||||
@@ -0,0 +1,395 @@
|
||||
"""Model-free multi-process acceptance using real ComfyUI executors and loopback HTTP.
|
||||
|
||||
Only the transport test facade (/prompt, /queue, /history, /interrupt) is synthetic.
|
||||
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 importlib
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
import uuid
|
||||
from unittest.mock import patch
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[2]
|
||||
COMFY_ROOT = Path(sys.argv[1]).resolve()
|
||||
WORKER = '--worker' in sys.argv
|
||||
DELAY = float(sys.argv[3]) if WORKER else 0
|
||||
sys.path.insert(0, str(COMFY_ROOT))
|
||||
os.environ['COMFYUI_IS_WORKER'] = '1'
|
||||
import comfy.cli_args
|
||||
comfy.cli_args.args.cpu = True
|
||||
comfy.cli_args.args.disable_assets = True
|
||||
from app.assets.manager import default_asset_manager
|
||||
import comfy.model_management
|
||||
import execution
|
||||
import nodes
|
||||
import server
|
||||
import torch
|
||||
from aiohttp import web
|
||||
import aiohttp
|
||||
|
||||
torch.set_num_threads(1)
|
||||
SOURCE_CALLS, BRANCH_CALLS, CAPTURED = [], [], []
|
||||
COUNT = 7
|
||||
|
||||
|
||||
def source_rows():
|
||||
return [{'image': torch.full((1, 5 + i, 7, 3 if i % 2 == 0 else 4), i / 16, dtype=torch.float32),
|
||||
'reference_image': torch.full((1, 2, 3, 3), (i + 1) / 16),
|
||||
'mask': torch.full((1, 5 + i, 7), (i + 1) / 16),
|
||||
'text': f'item-{i}', 'strength': (i + 1) / 16, 'seed': 100 + i, 'index': i}
|
||||
for i in range(COUNT)]
|
||||
|
||||
|
||||
class TestImageSource:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {}}
|
||||
RETURN_TYPES = ('IMAGE', 'IMAGE', 'MASK', 'STRING', 'FLOAT', 'INT')
|
||||
OUTPUT_IS_LIST = (True,) * 6
|
||||
FUNCTION = 'run'
|
||||
def run(self):
|
||||
assert not WORKER, 'upstream image source leaked into worker prompt'
|
||||
SOURCE_CALLS.append(True)
|
||||
rows = source_rows()
|
||||
return tuple([row[key] for row in rows] for key in
|
||||
('image', 'reference_image', 'mask', 'text', 'strength', 'seed'))
|
||||
|
||||
|
||||
class TestImageOperation:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'image': ('IMAGE',), 'reference_image': ('IMAGE',), 'mask': ('MASK',),
|
||||
'text': ('STRING', {'forceInput': True}),
|
||||
'strength': ('FLOAT', {'forceInput': True}),
|
||||
'seed': ('INT', {'forceInput': True}), 'index': ('INT', {'forceInput': True})}}
|
||||
RETURN_TYPES = ('IMAGE', 'MASK', 'STRING')
|
||||
FUNCTION = 'run'
|
||||
def run(self, image, reference_image, mask, text, strength, seed, index):
|
||||
assert WORKER, 'per-image branch executed on master'
|
||||
BRANCH_CALLS.append(index)
|
||||
# Sleep in short chunks so the real Comfy interrupt signal can be tested.
|
||||
until = time.monotonic() + DELAY
|
||||
while time.monotonic() < until:
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
time.sleep(0.01)
|
||||
assert seed == 100 + index
|
||||
assert text == f'item-{index}'
|
||||
assert strength == (index + 1) / 16
|
||||
assert torch.all(reference_image == strength)
|
||||
return image + strength, mask * 0.5, f'{text}:seed-{seed}'
|
||||
|
||||
|
||||
class TestCapture:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'images': ('IMAGE',), 'masks': ('MASK',), 'texts': ('STRING', {'forceInput': True})}}
|
||||
INPUT_IS_LIST = True
|
||||
OUTPUT_NODE = True
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = 'run'
|
||||
def run(self, images, masks, texts):
|
||||
CAPTURED.append((images, masks, texts))
|
||||
return ()
|
||||
|
||||
|
||||
async def setup():
|
||||
assets = default_asset_manager()
|
||||
ps = server.PromptServer(asyncio.get_running_loop(), assets)
|
||||
assert await nodes.load_custom_node(str(ROOT))
|
||||
module = sys.modules[str(ROOT).replace('.', '_x_')]
|
||||
nodes.NODE_CLASS_MAPPINGS.update(TestImageSource=TestImageSource,
|
||||
TestImageOperation=TestImageOperation, TestCapture=TestCapture)
|
||||
app = web.Application()
|
||||
# Register the actual feature routes, not test replacements.
|
||||
app.add_routes([route for route in ps.routes if route.path.startswith('/distributed/image-list/')])
|
||||
return ps, assets, module, app
|
||||
|
||||
|
||||
async def listen(app):
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, '127.0.0.1', 0)
|
||||
await site.start()
|
||||
return runner, site._server.sockets[0].getsockname()[1]
|
||||
|
||||
|
||||
async def worker_main():
|
||||
ps, assets, module, app = await setup()
|
||||
pending, running, history = [], {}, {}
|
||||
wake = asyncio.Event()
|
||||
interrupted_ids = []
|
||||
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 {})
|
||||
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({})
|
||||
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})
|
||||
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)
|
||||
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.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)
|
||||
consumer = asyncio.create_task(consume())
|
||||
print(f'READY:{port}', flush=True)
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
finally:
|
||||
consumer.cancel()
|
||||
await asyncio.gather(consumer, return_exceptions=True)
|
||||
await runner.cleanup()
|
||||
|
||||
|
||||
def graph(callback, worker_ids):
|
||||
return {
|
||||
'1': {'class_type': 'TestImageSource', 'inputs': {}},
|
||||
'2': {'class_type': 'DistributedImageListStart', 'inputs': {
|
||||
'images': ['1', 0], 'reference_image': ['1', 1], 'mask': ['1', 2], 'text': ['1', 3],
|
||||
'strength': ['1', 4], 'seed': ['1', 5],
|
||||
'reference_image_mode': 'per_item', 'mask_mode': 'per_item', 'text_mode': 'per_item',
|
||||
'strength_mode': 'per_item', 'seed_mode': 'per_item',
|
||||
'image_list_enabled_worker_ids': json.dumps(worker_ids), 'image_list_master_callback': callback}},
|
||||
'3': {'class_type': 'TestImageOperation', 'inputs': {
|
||||
key: ['2', index] for key, index in [('image', 0), ('reference_image', 1), ('mask', 2),
|
||||
('text', 3), ('strength', 4), ('seed', 5), ('index', 6)]}},
|
||||
'4': {'class_type': 'DistributedImageListEnd', 'inputs': {
|
||||
'images': ['3', 0], 'masks': ['3', 1], 'texts': ['3', 2], 'context': ['2', 7]}},
|
||||
'5': {'class_type': 'TestCapture', 'inputs': {'images': ['4', 0], 'masks': ['4', 1], 'texts': ['4', 2]}},
|
||||
}
|
||||
|
||||
|
||||
async def coordinator():
|
||||
ps, assets, module, app = await setup()
|
||||
runtime = importlib.import_module(module.__name__ + '.image_list.runtime')
|
||||
config_mod = importlib.import_module(module.__name__ + '.utils.config')
|
||||
orchestration = importlib.import_module(module.__name__ + '.api.queue_orchestration')
|
||||
runner, port = await listen(app)
|
||||
processes, stderr_files = [], []
|
||||
# Old worker capability refusal is read-only and precedes input evaluation.
|
||||
legacy_app = web.Application()
|
||||
legacy_requests = []
|
||||
async def legacy(request):
|
||||
legacy_requests.append(request.method)
|
||||
return web.json_response({'error': 'not found'}, status=404)
|
||||
legacy_app.router.add_route('*', '/{path:.*}', legacy)
|
||||
legacy_runner, legacy_port = await listen(legacy_app)
|
||||
dispatch_module = importlib.import_module(module.__name__ + '.api.orchestration.dispatch')
|
||||
try:
|
||||
try:
|
||||
await dispatch_module.preflight_image_list_workers([
|
||||
{'id': 'legacy', 'type': 'local', 'host': '127.0.0.1', 'port': legacy_port}])
|
||||
raise AssertionError('legacy worker accepted')
|
||||
except RuntimeError as exc:
|
||||
assert 'upgrade' in str(exc), str(exc)
|
||||
assert legacy_requests == ['GET'] and not SOURCE_CALLS
|
||||
finally:
|
||||
await legacy_runner.cleanup()
|
||||
print('IMAGE_LIST_PREFLIGHT_OK legacy worker rejected before materialization or writes')
|
||||
scratch = Path(os.environ.get('TMPDIR', Path.home() / '.hermes/cache/scratch'))
|
||||
scratch.mkdir(parents=True, exist_ok=True)
|
||||
try:
|
||||
with tempfile.TemporaryDirectory(prefix='image-list-network-', dir=scratch) as tmp:
|
||||
workers = []
|
||||
for name, delay in [('fast', 0.02), ('slow', 0.5)]:
|
||||
err = open(Path(tmp) / f'{name}.log', 'w+')
|
||||
stderr_files.append(err)
|
||||
proc = await asyncio.create_subprocess_exec(sys.executable, str(Path(__file__).resolve()),
|
||||
str(COMFY_ROOT), '--worker', str(delay), stdout=asyncio.subprocess.PIPE, stderr=err,
|
||||
env={**os.environ, 'TMPDIR': tmp})
|
||||
processes.append(proc)
|
||||
async def ready():
|
||||
while True:
|
||||
line = await proc.stdout.readline()
|
||||
if not line:
|
||||
err.seek(0)
|
||||
raise AssertionError(f'Worker exited: {err.read()}')
|
||||
if line.startswith(b'READY:'):
|
||||
return int(line.split(b':')[1])
|
||||
worker_port = await asyncio.wait_for(ready(), 60)
|
||||
workers.append({'id': name, 'type': 'local', 'host': '127.0.0.1',
|
||||
'port': worker_port, 'enabled': True})
|
||||
config = {'workers': workers, 'settings': {'websocket_orchestration': False},
|
||||
'master': {'host': f'http://127.0.0.1:{port}'}}
|
||||
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
|
||||
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')
|
||||
async def enqueue(prepared, workflow, client_id, **kwargs):
|
||||
checked = await execution.validate_prompt('network-acceptance', prepared, None)
|
||||
assert checked[0], checked
|
||||
await asyncio.to_thread(executor.execute, prepared, 'network-acceptance', {}, checked[2])
|
||||
return {'prompt_id': 'network-acceptance', 'number': 1, 'node_errors': {}}
|
||||
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', [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'
|
||||
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(CAPTURED) == 1
|
||||
images, masks, texts = CAPTURED[0]
|
||||
assert len(images) == len(masks) == len(texts) == COUNT
|
||||
for i, row in enumerate(source_rows()):
|
||||
assert torch.equal(images[i], row['image'] + row['strength']), (i, images[i])
|
||||
assert torch.equal(masks[i], row['mask'] * 0.5), i
|
||||
assert texts[i] == f'item-{i}:seed-{100 + i}', texts
|
||||
async with aiohttp.ClientSession() as session:
|
||||
stats = []
|
||||
for worker in workers:
|
||||
async with session.get(f"http://127.0.0.1:{worker['port']}/test-stats") as response:
|
||||
stats.append(await response.json())
|
||||
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))
|
||||
assert not runtime._MASTER_JOBS
|
||||
# Three sequential regions on one worker, including optional outputs
|
||||
# left unconnected. Region N+1 must not race the prior worker End.
|
||||
chained = {'1': {'class_type': 'EmptyImage', 'inputs': {
|
||||
'height': 8, 'width': 8, 'batch_size': 2, 'color': 0}}}
|
||||
previous = '1'
|
||||
for start, end in [('2', '3'), ('4', '5'), ('6', '7')]:
|
||||
chained[start] = {'class_type': 'DistributedImageListStart', 'inputs': {
|
||||
'images': [previous, 0], 'reference_image_mode': 'per_item', 'mask_mode': 'per_item',
|
||||
'text_mode': 'shared', 'strength_mode': 'shared', 'seed_mode': 'shared'}}
|
||||
chained[end] = {'class_type': 'DistributedImageListEnd', 'inputs': {
|
||||
'images': [start, 0], 'context': [start, 7]}}
|
||||
previous = end
|
||||
chained['8'] = {'class_type': 'TestCapture', 'inputs': {
|
||||
'images': ['7', 0], 'masks': ['7', 1], 'texts': ['7', 2]}}
|
||||
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 executor.success, executor.status_messages
|
||||
assert len(CAPTURED) == 2 and len(CAPTURED[1][0]) == 2
|
||||
assert all(torch.equal(image, torch.zeros((1, 8, 8, 3))) for image in CAPTURED[1][0])
|
||||
assert CAPTURED[1][1:] == ([None, None], [None, None])
|
||||
assert not runtime._MASTER_JOBS
|
||||
print('IMAGE_LIST_SEQUENTIAL_OK three regions; one worker; optional outputs omitted')
|
||||
# Real Comfy user-interrupt flag while a real worker executor is active.
|
||||
slow = workers[1]
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async def slow_stats():
|
||||
async with session.get(f"http://127.0.0.1:{slow['port']}/test-stats") as response:
|
||||
return await response.json()
|
||||
async def until(predicate):
|
||||
for _ in range(150):
|
||||
current = await slow_stats()
|
||||
if predicate(current):
|
||||
return current
|
||||
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'],
|
||||
f'http://127.0.0.1:{port}')
|
||||
with patch.object(config_mod, 'load_config', return_value=config), \
|
||||
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
|
||||
task = asyncio.create_task(runtime.run_image_list_job(context))
|
||||
active = await until(lambda state: bool(state['running']))
|
||||
owned_id = active['running'][0]
|
||||
comfy.model_management.interrupt_current_processing(True)
|
||||
try:
|
||||
await asyncio.wait_for(task, 5)
|
||||
raise AssertionError('interrupt was ignored')
|
||||
except runtime.ImageListCancelledError:
|
||||
pass
|
||||
finally:
|
||||
comfy.model_management.interrupt_current_processing(False)
|
||||
after = await until(lambda state: not state['running'] and not state['pending'])
|
||||
assert after['interrupted'] == [owned_id], after
|
||||
assert len(after['calls']) <= len(before['calls']) + 1
|
||||
assert set(before['history']) <= set(after['history'])
|
||||
assert not runtime._MASTER_JOBS
|
||||
print('IMAGE_LIST_CANCEL_OK owned prompt stopped; no further dispatch; context removed')
|
||||
# A worker exception before End must fail, not hang or return a partial list.
|
||||
bad_rows = source_rows()
|
||||
bad_rows[0]['seed'] = -1
|
||||
context = runtime.create_master_job('2', '4', prompt, bad_rows, ['fast'],
|
||||
f'http://127.0.0.1:{port}')
|
||||
with patch.object(config_mod, 'load_config', return_value=config), \
|
||||
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
|
||||
try:
|
||||
await asyncio.wait_for(runtime.run_image_list_job(context), 5)
|
||||
raise AssertionError('worker exception was ignored')
|
||||
except RuntimeError as exc:
|
||||
assert 'failed image-list item' in str(exc), str(exc)
|
||||
assert not runtime._MASTER_JOBS
|
||||
print('IMAGE_LIST_FAILURE_OK worker exception surfaced; context removed')
|
||||
print('IMAGE_LIST_NETWORK_OK ' + json.dumps({'workers': 2, 'items': COUNT,
|
||||
'assignments': dict(zip(['fast', 'slow'], counts)), 'ordered_exact_outputs': True,
|
||||
'master_branch_calls': len(BRANCH_CALLS), 'source_calls': len(SOURCE_CALLS)}))
|
||||
finally:
|
||||
for proc in processes:
|
||||
if proc.returncode is None:
|
||||
proc.terminate()
|
||||
await asyncio.gather(*(proc.wait() for proc in processes))
|
||||
for err in stderr_files:
|
||||
err.close()
|
||||
await runner.cleanup()
|
||||
network = importlib.import_module(module.__name__ + '.utils.network')
|
||||
session = await network.get_client_session()
|
||||
await session.close()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
asyncio.run(worker_main() if WORKER else coordinator())
|
||||
@@ -196,6 +196,70 @@ async def check_execution(mapping, module, prompt_server, asset_manager):
|
||||
print('V3_EXECUTION_OK eight nodes; upscale GPU boundary mocked; preview artifact verified')
|
||||
|
||||
|
||||
async def check_image_list_lazy_executor(mapping, module, prompt_server, asset_manager):
|
||||
"""Execute real native V3 list/lazy flow; master must skip the processing branch."""
|
||||
branch_calls = []
|
||||
class ImageListProbe:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'image': ('IMAGE',)}}
|
||||
RETURN_TYPES = ('IMAGE',)
|
||||
FUNCTION = 'run'
|
||||
CATEGORY = 'tests'
|
||||
def run(self, image):
|
||||
branch_calls.append(image)
|
||||
return (image,)
|
||||
|
||||
nodes.NODE_CLASS_MAPPINGS['ImageListTestProbe'] = ImageListProbe
|
||||
runtime = sys.modules[module.__name__ + '.image_list.runtime']
|
||||
original_contexts = set(runtime._MASTER_JOBS)
|
||||
image_list_nodes = sys.modules[module.__name__ + '.nodes.image_list']
|
||||
async def fake_worker_results(_context_id):
|
||||
return (
|
||||
[torch.full((1, 8, 8, 3), value, dtype=torch.float32) for value in (0.125, 0.5, 0.875)],
|
||||
[None, None, None], [None, None, None],
|
||||
)
|
||||
|
||||
graph = {
|
||||
'1': {'class_type': 'EmptyImage', 'inputs': {'height': 8, 'width': 8, 'batch_size': 3, 'color': 0}},
|
||||
'2': {'class_type': 'DistributedImageListStart', 'inputs': {
|
||||
'images': ['1', 0], 'reference_image_mode': 'per_item', 'mask_mode': 'per_item',
|
||||
'text_mode': 'shared', 'strength': 1.0, 'strength_mode': 'shared',
|
||||
'seed': 17, 'seed_mode': 'shared', 'image_list_enabled_worker_ids': '["test-worker"]',
|
||||
'image_list_master_callback': 'http://127.0.0.1:1',
|
||||
}},
|
||||
'3': {'class_type': 'ImageListTestProbe', 'inputs': {'image': ['2', 0]}},
|
||||
'4': {'class_type': 'DistributedImageListEnd', 'inputs': {
|
||||
'images': ['3', 0], 'context': ['2', 7],
|
||||
}},
|
||||
'5': {'class_type': 'PreviewImage', 'inputs': {'images': ['4', 0]}},
|
||||
}
|
||||
valid = await execution.validate_prompt('image-list-lazy', graph, None)
|
||||
assert valid[0], valid
|
||||
import folder_paths
|
||||
scratch = Path(os.environ.get('TMPDIR', Path.home() / '.hermes/cache/scratch'))
|
||||
try:
|
||||
with tempfile.TemporaryDirectory(prefix='image-list-lazy-', dir=scratch) as temp:
|
||||
with patch.object(folder_paths, 'temp_directory', temp), \
|
||||
patch.object(image_list_nodes, 'run_image_list_job', fake_worker_results):
|
||||
executor = execution.PromptExecutor(
|
||||
prompt_server, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=asset_manager)
|
||||
await asyncio.to_thread(executor.execute, graph, 'image-list-lazy', {}, valid[2])
|
||||
assert executor.success, executor.status_messages
|
||||
assert branch_calls == [], 'master evaluated the lazy per-image branch'
|
||||
records = executor.history_result['outputs']['5']['images']
|
||||
assert len(records) == 3, records
|
||||
for record in records:
|
||||
preview = Path(temp) / record.get('subfolder', '') / record['filename']
|
||||
with Image.open(preview) as image:
|
||||
assert image.size == (8, 8) and image.mode == 'RGB'
|
||||
finally:
|
||||
nodes.NODE_CLASS_MAPPINGS.pop('ImageListTestProbe', None)
|
||||
for context_id in set(runtime._MASTER_JOBS) - original_contexts:
|
||||
runtime.remove_master_job(context_id)
|
||||
print('IMAGE_LIST_LAZY_EXECUTION_OK real PromptExecutor; branch skipped on master; 3 previews verified')
|
||||
|
||||
|
||||
async def main():
|
||||
asset_manager = default_asset_manager()
|
||||
prompt_server = server.PromptServer(asyncio.get_running_loop(), asset_manager)
|
||||
@@ -206,9 +270,10 @@ async def main():
|
||||
extension = await module.comfy_entrypoint()
|
||||
classes = await extension.get_node_list()
|
||||
mapping = {cls.GET_SCHEMA().node_id: cls for cls in classes}
|
||||
assert len(classes) == len(mapping) == len(BASELINE) == 8
|
||||
assert set(mapping) == set(BASELINE)
|
||||
for node_id, cls in mapping.items():
|
||||
assert len(classes) == len(mapping) == len(BASELINE) + 2 == 10
|
||||
assert set(mapping) == set(BASELINE) | {'DistributedImageListStart', 'DistributedImageListEnd'}
|
||||
for node_id in BASELINE:
|
||||
cls = mapping[node_id]
|
||||
assert issubclass(cls, io.ComfyNode)
|
||||
assert nodes.NODE_CLASS_MAPPINGS[node_id] is cls
|
||||
old = dict(BASELINE[node_id])
|
||||
@@ -243,9 +308,24 @@ async def main():
|
||||
if old['output_node']:
|
||||
expected_hidden.extend(name for name in ['prompt', 'extra_pnginfo'] if name not in expected_hidden)
|
||||
assert list(new['input'].get('hidden', {})) == expected_hidden, (node_id, new['input'].get('hidden'))
|
||||
start_schema = mapping['DistributedImageListStart'].GET_SCHEMA()
|
||||
end_schema = mapping['DistributedImageListEnd'].GET_SCHEMA()
|
||||
assert start_schema.node_id == 'DistributedImageListStart'
|
||||
assert start_schema.display_name == 'Distributed Image List Start'
|
||||
assert end_schema.node_id == 'DistributedImageListEnd'
|
||||
assert end_schema.display_name == 'Distributed Image List End'
|
||||
assert [out.io_type for out in start_schema.outputs] == [
|
||||
'IMAGE', 'IMAGE', 'MASK', 'STRING', 'FLOAT', 'INT', 'INT', 'DISTRIBUTED_IMAGE_LIST_CONTEXT']
|
||||
assert all(out.is_output_list for out in start_schema.outputs)
|
||||
assert start_schema.is_input_list and start_schema.accept_all_inputs
|
||||
assert [out.io_type for out in end_schema.outputs] == ['IMAGE', 'MASK', 'STRING']
|
||||
assert all(out.is_output_list for out in end_schema.outputs)
|
||||
assert end_schema.is_input_list and end_schema.is_output_node
|
||||
print('IMAGE_LIST_SCHEMA_OK exact IDs, names, typed list outputs')
|
||||
print('SCHEMA_PARITY_OK', len(mapping))
|
||||
check_saved_workflows(mapping)
|
||||
await check_execution(mapping, module, prompt_server, asset_manager)
|
||||
await check_image_list_lazy_executor(mapping, module, prompt_server, asset_manager)
|
||||
print('V3_ACCEPTANCE_OK')
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,248 @@
|
||||
"""Regression tests for typed dynamic image-list normalization and graph scheduling."""
|
||||
import asyncio
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from image_list.runtime import (
|
||||
END_CLASS,
|
||||
IMAGE_OUT,
|
||||
CONTEXT_OUT,
|
||||
START_CLASS,
|
||||
ImageListCancelledError,
|
||||
build_worker_prompt,
|
||||
create_master_job,
|
||||
decode_item,
|
||||
decode_result,
|
||||
dynamic_schedule,
|
||||
encode_item,
|
||||
encode_result,
|
||||
normalize_image_list_inputs,
|
||||
normalize_worker_result,
|
||||
remove_master_job,
|
||||
set_master_result,
|
||||
validate_image_list_regions,
|
||||
)
|
||||
|
||||
|
||||
def start_node(image_link):
|
||||
return {"class_type": START_CLASS, "inputs": {
|
||||
"images": image_link,
|
||||
"reference_image_mode": "per_item", "mask_mode": "per_item",
|
||||
"text_mode": "shared", "strength_mode": "shared", "seed_mode": "shared",
|
||||
}}
|
||||
|
||||
|
||||
def graph_one_region():
|
||||
return {
|
||||
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
|
||||
"2": {"class_type": "EmptyImage", "inputs": {"height": 8, "width": 8, "batch_size": 2}},
|
||||
"3": start_node(["2", 0]),
|
||||
"4": {"class_type": "GenericFrameOperation", "inputs": {
|
||||
"image": ["3", IMAGE_OUT], "model": ["1", 0], "seed": ["3", 5],
|
||||
}},
|
||||
"5": {"class_type": END_CLASS, "inputs": {
|
||||
"images": ["4", 0], "context": ["3", CONTEXT_OUT],
|
||||
}},
|
||||
"6": {"class_type": "SaveImage", "inputs": {"images": ["5", 0]}},
|
||||
}
|
||||
|
||||
|
||||
def graph_two_sequential_regions():
|
||||
return {
|
||||
"1": {"class_type": "EmptyImage", "inputs": {"height": 4, "width": 4, "batch_size": 2}},
|
||||
"2": start_node(["1", 0]),
|
||||
"3": {"class_type": "FirstFrameOperation", "inputs": {"image": ["2", IMAGE_OUT]}},
|
||||
"4": {"class_type": END_CLASS, "inputs": {"images": ["3", 0], "context": ["2", CONTEXT_OUT]}},
|
||||
"5": start_node(["4", 0]),
|
||||
"6": {"class_type": "SecondFrameOperation", "inputs": {"image": ["5", IMAGE_OUT]}},
|
||||
"7": {"class_type": END_CLASS, "inputs": {"images": ["6", 0], "context": ["5", CONTEXT_OUT]}},
|
||||
"8": {"class_type": "SaveImage", "inputs": {"images": ["7", 0]}},
|
||||
}
|
||||
|
||||
|
||||
class ImageListNormalizationTests(unittest.TestCase):
|
||||
def test_batch_frame_normalization_preserves_sizes_channels_and_float_values(self):
|
||||
images = [
|
||||
torch.tensor([[[[0.125, 0.25, 0.5]]]], dtype=torch.float64),
|
||||
torch.arange(2 * 3 * 4, dtype=torch.float32).reshape(1, 2, 3, 4) / 17,
|
||||
]
|
||||
refs = torch.ones((2, 5, 7, 3), dtype=torch.float16) * 0.333
|
||||
masks = [torch.full((1, 2, 3), 0.125, dtype=torch.float64),
|
||||
torch.full((1, 8, 9), 0.875, dtype=torch.float32)]
|
||||
rows = normalize_image_list_inputs(
|
||||
images, refs, masks, ["first", "second"], [0.2], [37],
|
||||
"per_item", "per_item", "per_item", "shared", "shared",
|
||||
)
|
||||
self.assertEqual(len(rows), 2)
|
||||
self.assertEqual([tuple(row["image"].shape) for row in rows], [(1, 1, 1, 3), (1, 2, 3, 4)])
|
||||
self.assertEqual([row["image"].dtype for row in rows], [torch.float64, torch.float32])
|
||||
self.assertTrue(torch.equal(rows[0]["reference_image"], refs[:1]))
|
||||
self.assertEqual(rows[1]["mask"].dtype, torch.float32)
|
||||
self.assertEqual([row["index"] for row in rows], [0, 1])
|
||||
self.assertEqual([row["seed"] for row in rows], [37, 37])
|
||||
self.assertEqual([row["strength"] for row in rows], [0.2, 0.2])
|
||||
|
||||
def test_explicit_per_item_and_shared_scalar_contracts(self):
|
||||
image = torch.zeros((3, 2, 4, 3), dtype=torch.float32)
|
||||
rows = normalize_image_list_inputs(image, text="same", strength=0.4, seed=9,
|
||||
text_mode="shared", strength_mode="shared", seed_mode="shared")
|
||||
self.assertEqual([row["text"] for row in rows], ["same"] * 3)
|
||||
self.assertEqual([row["strength"] for row in rows], [0.4] * 3)
|
||||
self.assertEqual([row["seed"] for row in rows], [9] * 3)
|
||||
with self.assertRaisesRegex(ValueError, "expected exactly 3"):
|
||||
normalize_image_list_inputs(image, text=["one", "two"], text_mode="per_item")
|
||||
with self.assertRaisesRegex(ValueError, "shared inputs require exactly one"):
|
||||
normalize_image_list_inputs(image, seed=[1, 2], seed_mode="shared")
|
||||
with self.assertRaisesRegex(ValueError, "expected exactly 3"):
|
||||
normalize_image_list_inputs(image, mask=torch.zeros((2, 2, 4)), mask_mode="per_item")
|
||||
|
||||
def test_no_implicit_tail_repetition_and_no_empty_primary_batch(self):
|
||||
images = torch.zeros((2, 2, 2, 3), dtype=torch.float32)
|
||||
with self.assertRaisesRegex(ValueError, "expected exactly 2"):
|
||||
normalize_image_list_inputs(images, reference_image=torch.zeros((1, 2, 2, 3)),
|
||||
reference_image_mode="per_item")
|
||||
with self.assertRaisesRegex(ValueError, "at least one image"):
|
||||
normalize_image_list_inputs(None)
|
||||
|
||||
def test_typed_transport_round_trips_float_images_masks_and_scalar_values(self):
|
||||
item = {"image": torch.arange(24, dtype=torch.float64).reshape(1, 2, 4, 3) / 7,
|
||||
"reference_image": torch.ones((1, 1, 3, 4), dtype=torch.float16),
|
||||
"mask": torch.full((1, 5, 6), 0.375, dtype=torch.float32),
|
||||
"text": "safe text", "strength": 0.625, "seed": 42, "index": 3}
|
||||
restored = decode_item(encode_item(item))
|
||||
self.assertTrue(torch.equal(restored["image"], item["image"]))
|
||||
self.assertEqual(restored["image"].dtype, torch.float64)
|
||||
self.assertTrue(torch.equal(restored["reference_image"], item["reference_image"]))
|
||||
self.assertTrue(torch.equal(restored["mask"], item["mask"]))
|
||||
self.assertEqual(restored["text"], "safe text")
|
||||
self.assertEqual(restored["seed"], 42)
|
||||
with self.assertRaisesRegex(TypeError, "floating-point tensor"):
|
||||
encode_item({**item, "image": torch.ones((1, 2, 4, 3), dtype=torch.int32)})
|
||||
with self.assertRaisesRegex(ValueError, "Unsupported image-list transport type"):
|
||||
from image_list.runtime import encode_value
|
||||
encode_value("/untrusted/path", "PATH")
|
||||
|
||||
def test_result_transport_preserves_ragged_result_shapes(self):
|
||||
result = normalize_worker_result(
|
||||
[torch.ones((1, 3, 5, 4), dtype=torch.float32)],
|
||||
[torch.ones((1, 3, 5), dtype=torch.float16)],
|
||||
["done"],
|
||||
)
|
||||
encoded = encode_result(result)
|
||||
decoded = decode_result(encoded)
|
||||
self.assertEqual(tuple(decoded["image"].shape), (1, 3, 5, 4))
|
||||
self.assertEqual(decoded["image"].dtype, torch.float32)
|
||||
self.assertEqual(decoded["mask"].dtype, torch.float16)
|
||||
self.assertEqual(decoded["text"], "done")
|
||||
with self.assertRaisesRegex(ValueError, "exactly one IMAGE"):
|
||||
normalize_worker_result(torch.zeros((2, 3, 4, 3)))
|
||||
|
||||
|
||||
class ImageListGraphBoundaryTests(unittest.TestCase):
|
||||
def test_prunes_upstream_item_source_but_keeps_external_model_loaders(self):
|
||||
prompt = graph_one_region()
|
||||
region = validate_image_list_regions(prompt)[0]
|
||||
worker = build_worker_prompt(prompt, region, "token_12345678901234567890123456789012",
|
||||
"ctx", 1, "http://master.invalid")
|
||||
self.assertEqual(set(worker), {"1", "3", "4", "5"})
|
||||
self.assertEqual(worker["3"]["inputs"]["worker_item_token"], "token_12345678901234567890123456789012")
|
||||
self.assertNotIn("images", worker["3"]["inputs"])
|
||||
self.assertEqual(worker["4"]["inputs"]["model"], ["1", 0])
|
||||
self.assertNotIn("6", worker)
|
||||
|
||||
def test_two_sequential_regions_are_valid_and_isolated(self):
|
||||
regions = validate_image_list_regions(graph_two_sequential_regions())
|
||||
self.assertEqual([(r["start_id"], r["end_id"]) for r in regions], [("2", "4"), ("5", "7")])
|
||||
p = graph_two_sequential_regions()
|
||||
second = build_worker_prompt(p, regions[1], "item_token_123456789012345678901234", "ctx", 0,
|
||||
"http://master.invalid")
|
||||
self.assertEqual(set(second), {"5", "6", "7"})
|
||||
|
||||
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)
|
||||
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)
|
||||
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)
|
||||
|
||||
def test_rejects_duplicate_or_missing_results(self):
|
||||
prompt = graph_one_region()
|
||||
prompt["5"]["inputs"].pop("images")
|
||||
with self.assertRaisesRegex(ValueError, "images must be connected"):
|
||||
validate_image_list_regions(prompt)
|
||||
|
||||
|
||||
class DynamicImageSchedulingTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_fast_workers_get_more_items_and_order_is_stable(self):
|
||||
assignments = {"fast": [], "slow": []}
|
||||
async def dispatch(worker, index):
|
||||
assignments[worker].append(index)
|
||||
delay = 0.002 if worker == "fast" else 0.025
|
||||
task = asyncio.create_task(asyncio.sleep(delay, result=(index, f"result-{index}")))
|
||||
return {"index": index, "task": task}
|
||||
async def wait_next(inflight):
|
||||
done, _ = await asyncio.wait([h["task"] for h in inflight.values()], return_when=asyncio.FIRST_COMPLETED)
|
||||
task = next(iter(done))
|
||||
worker = next(w for w, handle in inflight.items() if handle["task"] is task)
|
||||
index, result = task.result()
|
||||
return worker, index, result
|
||||
async def cleanup(_worker, handle):
|
||||
handle["task"].cancel()
|
||||
ordered = await dynamic_schedule(12, ["fast", "slow"], dispatch, wait_next, cleanup)
|
||||
self.assertEqual(ordered, [f"result-{i}" for i in range(12)])
|
||||
self.assertGreater(len(assignments["fast"]), len(assignments["slow"]))
|
||||
self.assertEqual(sorted(assignments["fast"] + assignments["slow"]), list(range(12)))
|
||||
|
||||
async def test_failure_timeout_and_cancel_clean_only_current_assignments(self):
|
||||
async def dispatch(worker, index):
|
||||
return {"index": index, "prompt_id": f"{worker}-{index}"}
|
||||
cancelled = []
|
||||
async def cleanup(worker, handle):
|
||||
cancelled.append((worker, handle["prompt_id"]))
|
||||
async def timeout(_inflight):
|
||||
raise TimeoutError("worker timed out")
|
||||
with self.assertRaisesRegex(TimeoutError, "timed out"):
|
||||
await dynamic_schedule(4, ["one", "two"], dispatch, timeout, cleanup)
|
||||
self.assertCountEqual(cancelled, [("one", "one-0"), ("two", "two-1")])
|
||||
|
||||
cancelled.clear()
|
||||
async def wrong_index(inflight):
|
||||
worker, handle = next(iter(inflight.items()))
|
||||
return worker, handle["index"] + 99, "bad"
|
||||
with self.assertRaisesRegex(RuntimeError, "unexpected image index"):
|
||||
await dynamic_schedule(2, ["one"], dispatch, wrong_index, cleanup)
|
||||
self.assertEqual(len(cancelled), 1)
|
||||
|
||||
cancelled.clear()
|
||||
with self.assertRaises(ImageListCancelledError):
|
||||
await dynamic_schedule(2, ["one", "two"], dispatch, timeout, cleanup,
|
||||
interrupted=lambda: True)
|
||||
self.assertEqual(len(cancelled), 0)
|
||||
|
||||
async def test_no_workers_and_duplicate_results_fail_closed(self):
|
||||
async def unused(*_args):
|
||||
self.fail("must not dispatch without workers")
|
||||
with self.assertRaisesRegex(RuntimeError, "No free"):
|
||||
await dynamic_schedule(2, [], unused, unused, unused)
|
||||
async def dispatch(worker, index):
|
||||
return {"index": index}
|
||||
async def duplicate(inflight):
|
||||
worker, _handle = next(iter(inflight.items()))
|
||||
return worker, 0, "same"
|
||||
with self.assertRaisesRegex(RuntimeError, "unexpected image index"):
|
||||
await dynamic_schedule(2, ["one"], dispatch, duplicate, lambda *_: asyncio.sleep(0))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,178 @@
|
||||
"""Regression coverage at the scheduler and real HTTP transport boundaries."""
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import json
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
from aiohttp import web
|
||||
from aiohttp.test_utils import TestClient, TestServer
|
||||
|
||||
from image_list import runtime
|
||||
from test_image_list import graph_two_sequential_regions, start_node
|
||||
|
||||
|
||||
def load_routes():
|
||||
root = Path(__file__).resolve().parents[1]
|
||||
name = '_image_list_transport_tests'
|
||||
routes = web.RouteTableDef()
|
||||
modules = {
|
||||
name: types.ModuleType(name),
|
||||
name + '.api': types.ModuleType(name + '.api'),
|
||||
name + '.image_list': types.ModuleType(name + '.image_list'),
|
||||
name + '.image_list.runtime': runtime,
|
||||
'server': types.SimpleNamespace(PromptServer=types.SimpleNamespace(
|
||||
instance=types.SimpleNamespace(routes=routes))),
|
||||
}
|
||||
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)
|
||||
with patch.dict(sys.modules, modules):
|
||||
spec.loader.exec_module(module)
|
||||
return module, routes
|
||||
|
||||
|
||||
class TransportTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
self.module, routes = load_routes()
|
||||
app = web.Application()
|
||||
app.add_routes(routes)
|
||||
self.client = TestClient(TestServer(app))
|
||||
await self.client.start_server()
|
||||
|
||||
async def asyncTearDown(self):
|
||||
await self.client.close()
|
||||
runtime._WORKER_ITEMS.clear()
|
||||
|
||||
async def test_fragmented_chunked_item_upload_is_consumable(self):
|
||||
token = runtime.create_worker_task_token()
|
||||
item = {'image': torch.full((1, 2, 3, 3), 0.25), 'index': 2}
|
||||
data = json.dumps({'token': token, 'context_id': 'context', 'item': runtime.encode_item(item)}).encode()
|
||||
async def chunks():
|
||||
yield data[:8]
|
||||
await asyncio.sleep(0.02)
|
||||
yield data[8:]
|
||||
response = await self.client.post('/distributed/image-list/item', data=chunks())
|
||||
self.assertEqual(response.status, 200, await response.text())
|
||||
result = runtime.consume_worker_item(token)
|
||||
self.assertTrue(torch.equal(result['image'], item['image']))
|
||||
self.assertEqual(result['index'], 2)
|
||||
|
||||
async def test_complete_item_upload_and_replay_rejection(self):
|
||||
token = runtime.create_worker_task_token()
|
||||
data = {'token': token, 'context_id': 'context', 'item': runtime.encode_item(
|
||||
{'image': torch.zeros((1, 2, 3, 3)), 'index': 0})}
|
||||
response = await self.client.post('/distributed/image-list/item', json=data)
|
||||
self.assertEqual(response.status, 200, await response.text())
|
||||
response = await self.client.post('/distributed/image-list/item', json=data)
|
||||
self.assertEqual(response.status, 400)
|
||||
runtime.consume_worker_item(token)
|
||||
with self.assertRaisesRegex(RuntimeError, 'already consumed'):
|
||||
runtime.consume_worker_item(token)
|
||||
|
||||
async def test_chunked_body_limit_is_aggregate(self):
|
||||
async def chunks():
|
||||
yield b'{"context_id":"'
|
||||
await asyncio.sleep(0.01)
|
||||
yield b'x' * 60
|
||||
yield b'"}'
|
||||
with patch.object(self.module, 'MAX_REQUEST_BYTES', 64):
|
||||
response = await self.client.post('/distributed/image-list/item', data=chunks())
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn('exceeds', await response.text())
|
||||
|
||||
async def test_invalid_result_shapes_are_rejected(self):
|
||||
for kind, shape in [('image', (1, 2, 3)), ('mask', (1, 2, 3, 1))]:
|
||||
result = runtime.encode_result({'image': torch.zeros((1, 2, 3, 3)),
|
||||
'mask': torch.zeros((1, 2, 3))})
|
||||
result[kind]['shape'] = list(shape)
|
||||
with self.assertRaises(ValueError):
|
||||
runtime.decode_result(result)
|
||||
|
||||
|
||||
class SchedulingBoundaryTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_user_interrupt_during_result_wait_cleans_up_promptly(self):
|
||||
entered = asyncio.Event()
|
||||
cancelled = asyncio.Event()
|
||||
wait_cancelled = asyncio.Event()
|
||||
dispatches, cleanups = [], []
|
||||
async def dispatch(worker, index):
|
||||
dispatches.append(index)
|
||||
return {'index': index}
|
||||
async def wait_next(_inflight):
|
||||
entered.set()
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
finally:
|
||||
wait_cancelled.set()
|
||||
async def cleanup(worker, handle):
|
||||
cleanups.append((worker, handle['index']))
|
||||
task = asyncio.create_task(runtime.dynamic_schedule(
|
||||
3, ['one'], dispatch, wait_next, cleanup, cancelled.is_set))
|
||||
await entered.wait()
|
||||
cancelled.set()
|
||||
with self.assertRaises(runtime.ImageListCancelledError):
|
||||
await asyncio.wait_for(task, timeout=1)
|
||||
self.assertTrue(wait_cancelled.is_set())
|
||||
self.assertEqual(dispatches, [0])
|
||||
self.assertEqual(cleanups, [('one', 0)])
|
||||
|
||||
async def test_already_interrupted_does_not_submit_any_work(self):
|
||||
async def unused(*args):
|
||||
self.fail('pre-cancelled job must not dispatch or wait')
|
||||
with self.assertRaises(runtime.ImageListCancelledError):
|
||||
await runtime.dynamic_schedule(2, ['one'], unused, unused, unused, lambda: True)
|
||||
|
||||
async def test_task_cancel_during_submit_waits_for_owned_handle(self):
|
||||
entered, release = asyncio.Event(), asyncio.Event()
|
||||
cleanups = []
|
||||
async def dispatch(worker, index):
|
||||
entered.set()
|
||||
await release.wait()
|
||||
return {'index': index, 'prompt_id': 'owned'}
|
||||
async def unused(*args):
|
||||
self.fail('cancelled dispatch must not enter result wait')
|
||||
async def cleanup(worker, handle):
|
||||
cleanups.append(handle['prompt_id'])
|
||||
task = asyncio.create_task(runtime.dynamic_schedule(2, ['one'], dispatch, unused, cleanup))
|
||||
await entered.wait()
|
||||
task.cancel()
|
||||
release.set()
|
||||
with self.assertRaises(asyncio.CancelledError):
|
||||
await task
|
||||
self.assertEqual(cleanups, ['owned'])
|
||||
|
||||
def test_output_nodes_inside_region_are_rejected(self):
|
||||
prompt = graph_two_sequential_regions()
|
||||
output_class = type('OutputOperation', (), {'OUTPUT_NODE': True})
|
||||
registry = types.SimpleNamespace(NODE_CLASS_MAPPINGS={'FirstFrameOperation': output_class})
|
||||
with patch.dict(sys.modules, {'nodes': registry}):
|
||||
with self.assertRaisesRegex(ValueError, 'must be placed after'):
|
||||
runtime.validate_image_list_regions(prompt)
|
||||
|
||||
def test_shared_scalars_are_validated_before_dispatch(self):
|
||||
for kwargs in ({'text': 4, 'text_mode': 'shared'},
|
||||
{'strength': float('nan'), 'strength_mode': 'shared'},
|
||||
{'seed': True, 'seed_mode': 'shared'}):
|
||||
with self.assertRaises(TypeError):
|
||||
runtime.normalize_image_list_inputs(torch.zeros((1, 2, 2, 3)), **kwargs)
|
||||
|
||||
def test_three_sequential_regions_can_share_worker_loader(self):
|
||||
prompt = graph_two_sequential_regions()
|
||||
prompt.pop('8')
|
||||
prompt.update({
|
||||
'8': start_node(['7', 0]),
|
||||
'9': {'class_type': 'ThirdOperation', 'inputs': {'image': ['8', 0]}},
|
||||
'10': {'class_type': runtime.END_CLASS,
|
||||
'inputs': {'images': ['9', 0], 'context': ['8', runtime.CONTEXT_OUT]}},
|
||||
'11': {'class_type': 'GenericModelLoader', 'inputs': {}},
|
||||
})
|
||||
for node_id in ['3', '6', '9']:
|
||||
prompt[node_id]['inputs']['model'] = ['11', 0]
|
||||
regions = runtime.validate_image_list_regions(prompt)
|
||||
self.assertEqual(len(regions), 3)
|
||||
for region in regions:
|
||||
self.assertIn('11', region['node_ids'])
|
||||
@@ -19,3 +19,18 @@ def test_native_v3_runtime():
|
||||
)
|
||||
assert result.returncode == 0, result.stdout + '\n' + result.stderr
|
||||
assert 'V3_ACCEPTANCE_OK' in result.stdout
|
||||
|
||||
|
||||
def test_image_list_network_runtime():
|
||||
comfy_root = os.environ.get('COMFYUI_SOURCE_ROOT')
|
||||
if not comfy_root:
|
||||
pytest.skip('Set COMFYUI_SOURCE_ROOT for model-free multi-process HTTP acceptance')
|
||||
helper = Path(__file__).parent / 'helpers' / 'image_list_network_check.py'
|
||||
result = subprocess.run(
|
||||
[sys.executable, str(helper), str(Path(comfy_root).resolve())],
|
||||
capture_output=True, text=True, timeout=180,
|
||||
env={**os.environ, 'COMFYUI_IS_WORKER': '1'},
|
||||
)
|
||||
assert result.returncode == 0, result.stdout + '\n' + result.stderr
|
||||
assert 'IMAGE_LIST_NETWORK_OK' in result.stdout
|
||||
assert 'IMAGE_LIST_CANCEL_OK' in result.stdout
|
||||
|
||||
@@ -225,6 +225,7 @@ export const ENDPOINTS = {
|
||||
|
||||
export const NODE_CLASSES = {
|
||||
DISTRIBUTED_COLLECTOR: 'DistributedCollector',
|
||||
IMAGE_LIST_START: 'DistributedImageListStart',
|
||||
DISTRIBUTED_SEED: 'DistributedSeed',
|
||||
DISTRIBUTED_EMPTY_IMAGE: 'DistributedEmptyImage',
|
||||
UPSCALE_DISTRIBUTED: 'UltimateSDUpscaleDistributed',
|
||||
|
||||
+12
-2
@@ -9,7 +9,9 @@ export function setupInterceptor(extension) {
|
||||
const hasCollector = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_COLLECTOR).length > 0;
|
||||
const hasDistUpscale = findNodesByClass(prompt.output, NODE_CLASSES.UPSCALE_DISTRIBUTED).length > 0;
|
||||
|
||||
if (hasCollector || hasDistUpscale) {
|
||||
const hasImageList = findNodesByClass(prompt.output, NODE_CLASSES.IMAGE_LIST_START).length > 0;
|
||||
|
||||
if (hasCollector || hasDistUpscale || hasImageList) {
|
||||
const result = await executeParallelDistributed(extension, prompt);
|
||||
// Immediate status check for instant feedback
|
||||
checkAllWorkerStatuses(extension);
|
||||
@@ -25,12 +27,20 @@ export function setupInterceptor(extension) {
|
||||
export async function executeParallelDistributed(extension, promptWrapper) {
|
||||
const traceExecutionId = `exec_${Date.now()}_${generateUUID().slice(0, 6)}`;
|
||||
try {
|
||||
const enabledWorkers = extension.enabledWorkers;
|
||||
const hasImageList = findNodesByClass(promptWrapper.output, NODE_CLASSES.IMAGE_LIST_START).length > 0;
|
||||
// Image-list v1 uses direct workers only; do not contact paid provider pools.
|
||||
const enabledWorkers = hasImageList
|
||||
? extension.enabledWorkers.filter(worker => worker.type !== "serverless")
|
||||
: extension.enabledWorkers;
|
||||
extension.log(`[exec:${traceExecutionId}] Starting distributed execution`, "debug");
|
||||
|
||||
// Pre-flight health check on all enabled workers
|
||||
const activeWorkers = await performPreflightCheck(extension, enabledWorkers);
|
||||
|
||||
if (hasImageList && activeWorkers.length === 0) {
|
||||
throw new Error("Distributed Image List requires an active direct worker; master fallback is disabled.");
|
||||
}
|
||||
|
||||
// Case: Enabled workers but all offline
|
||||
if (activeWorkers.length === 0 && enabledWorkers.length > 0) {
|
||||
extension.log("No active workers found. All enabled workers are offline.");
|
||||
|
||||
@@ -0,0 +1,56 @@
|
||||
import { afterEach, describe, expect, it, vi } from 'vitest';
|
||||
vi.mock('../../scripts/api.js', () => ({ api: { clientId: 'test-client' } }), { virtual: true });
|
||||
vi.mock('../workerLifecycle.js', () => ({
|
||||
checkAllWorkerStatuses: vi.fn(), getWorkerUrl: () => 'http://worker.invalid',
|
||||
}));
|
||||
vi.mock('../workerUtils.js', () => ({
|
||||
applyProbeResultToWorkerDot: vi.fn(),
|
||||
findNodesByClass: (prompt, cls) => Object.values(prompt).filter(node => node.class_type === cls),
|
||||
}));
|
||||
import { api } from '../../scripts/api.js';
|
||||
import { setupInterceptor } from '../executionUtils.js';
|
||||
|
||||
const prompt = { output: { '1': { class_type: 'DistributedImageListStart', inputs: {} } } };
|
||||
function extension(workers, online = true) {
|
||||
return {
|
||||
isEnabled: true, enabledWorkers: workers, log: vi.fn(), config: { settings: {} },
|
||||
originalQueuePrompt: vi.fn().mockResolvedValue({ prompt_id: 'local' }),
|
||||
api: { probeWorker: vi.fn().mockResolvedValue({ ok: online }),
|
||||
getServerlessStatus: vi.fn(),
|
||||
queueDistributed: vi.fn().mockResolvedValue({ prompt_id: 'distributed' }) },
|
||||
ui: { showToast: vi.fn() },
|
||||
};
|
||||
}
|
||||
afterEach(() => vi.useRealTimers());
|
||||
describe('image-list queue interception', () => {
|
||||
it('routes image-list graphs through backend orchestration', async () => {
|
||||
vi.useFakeTimers();
|
||||
const ext = extension([{ id: 'one', type: 'local' }]);
|
||||
setupInterceptor(ext);
|
||||
expect(await api.queuePrompt(0, prompt)).toEqual({ prompt_id: 'distributed' });
|
||||
expect(ext.api.queueDistributed).toHaveBeenCalledWith(expect.objectContaining({ enabled_worker_ids: ['one'] }));
|
||||
expect(ext.originalQueuePrompt).not.toHaveBeenCalled();
|
||||
});
|
||||
it('never falls back to master when all image-list workers are offline', async () => {
|
||||
vi.useFakeTimers();
|
||||
const ext = extension([{ id: 'one', type: 'local' }], false);
|
||||
setupInterceptor(ext);
|
||||
expect(await api.queuePrompt(0, prompt)).toBeNull();
|
||||
expect(ext.originalQueuePrompt).not.toHaveBeenCalled();
|
||||
expect(ext.api.queueDistributed).not.toHaveBeenCalled();
|
||||
});
|
||||
it('does not probe serverless providers for an image-list job', async () => {
|
||||
vi.useFakeTimers();
|
||||
const ext = extension([{ id: 'pool', type: 'serverless' }]);
|
||||
setupInterceptor(ext);
|
||||
expect(await api.queuePrompt(0, prompt)).toBeNull();
|
||||
expect(ext.api.getServerlessStatus).not.toHaveBeenCalled();
|
||||
expect(ext.originalQueuePrompt).not.toHaveBeenCalled();
|
||||
});
|
||||
it('preserves ordinary prompt handling', async () => {
|
||||
const ext = extension([]);
|
||||
setupInterceptor(ext);
|
||||
expect(await api.queuePrompt(0, { output: {} })).toEqual({ prompt_id: 'local' });
|
||||
expect(ext.api.queueDistributed).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user