feat(image-list): let master participate unless orchestrator-only
This commit is contained in:
@@ -585,7 +585,7 @@ def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id,
|
||||
|
||||
|
||||
def _override_image_list_nodes(prompt_copy, prompt_index, participant_id, worker_ids,
|
||||
job_id_map, master_url, workflow_meta=None):
|
||||
job_id_map, master_url, delegate_master, 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)
|
||||
@@ -593,6 +593,7 @@ def _override_image_list_nodes(prompt_copy, prompt_index, participant_id, worker
|
||||
continue
|
||||
inputs = node.setdefault("inputs", {})
|
||||
inputs["image_list_enabled_worker_ids"] = json.dumps(list(worker_ids))
|
||||
inputs["image_list_master_participates"] = participant_id == "master" and not delegate_master
|
||||
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:
|
||||
@@ -642,6 +643,7 @@ def apply_participant_overrides(
|
||||
enabled_worker_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
delegate_master,
|
||||
workflow_meta=workflow_meta,
|
||||
)
|
||||
|
||||
|
||||
+19
-21
@@ -239,6 +239,8 @@ async def orchestrate_distributed_execution(
|
||||
execution_trace_id = trace_execution_id or _generate_execution_trace_id()
|
||||
|
||||
config = load_config()
|
||||
if delegate_master is None:
|
||||
delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False))
|
||||
use_websocket = bool(config.get("settings", {}).get("websocket_orchestration", False))
|
||||
master_url = build_master_url(config=config, prompt_server_instance=prompt_server)
|
||||
(
|
||||
@@ -277,10 +279,10 @@ async def orchestrate_distributed_execution(
|
||||
+ ", ".join(configured_serverless),
|
||||
)
|
||||
workers = _resolve_direct_workers_for_image_list(config, requested_ids)
|
||||
if not workers:
|
||||
if not workers and delegate_master:
|
||||
raise RuntimeError(
|
||||
"Distributed Image List v1 requires an enabled selected direct local/remote worker; "
|
||||
"serverless pools and master-only execution are unsupported"
|
||||
"Distributed Image List in orchestrator-only mode requires a selected direct worker; "
|
||||
"serverless pools are unsupported"
|
||||
)
|
||||
else:
|
||||
workers = _resolve_enabled_workers(config, requested_ids)
|
||||
@@ -299,39 +301,35 @@ 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:
|
||||
# Keep the legacy Collector fallback, but never override image-list
|
||||
# orchestrator-only mode merely because its direct workers are unavailable.
|
||||
if not image_list_mode and not workers and delegate_master:
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
"Delegate-only requested but no workers are enabled. Falling back to master execution.",
|
||||
)
|
||||
delegate_master = False
|
||||
|
||||
active_workers, delegate_master = await select_active_workers(
|
||||
active_workers, selected_delegate_master = await select_active_workers(
|
||||
workers,
|
||||
use_websocket,
|
||||
delegate_master,
|
||||
False if image_list_mode else delegate_master,
|
||||
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:
|
||||
if not active_workers and delegate_master:
|
||||
raise RuntimeError(
|
||||
"No selected direct worker is reachable for Distributed Image List; "
|
||||
"the master will not duplicate per-image processing"
|
||||
"No selected direct worker is reachable for Distributed Image List in orchestrator-only mode"
|
||||
)
|
||||
idle_ids = await preflight_image_list_workers(active_workers, require_idle=True)
|
||||
idle_ids = (await preflight_image_list_workers(active_workers, require_idle=delegate_master)
|
||||
if active_workers else {})
|
||||
active_workers = [worker for worker in active_workers if worker["id"] in idle_ids]
|
||||
if not active_workers:
|
||||
raise RuntimeError("No feature-compatible direct worker is idle; no master job was queued")
|
||||
if not active_workers and delegate_master:
|
||||
raise RuntimeError("No feature-compatible direct worker is idle in orchestrator-only mode")
|
||||
else:
|
||||
delegate_master = selected_delegate_master
|
||||
|
||||
if load_balance_requested and not image_list_mode:
|
||||
candidate_workers = list(active_workers)
|
||||
@@ -416,7 +414,7 @@ async def orchestrate_distributed_execution(
|
||||
workflow_meta=workflow_meta,
|
||||
)
|
||||
|
||||
if delegate_master:
|
||||
if delegate_master and not image_list_mode:
|
||||
collector_ids = find_nodes_by_class(master_prompt, "DistributedCollector")
|
||||
upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
|
||||
if upscale_nodes:
|
||||
|
||||
+14
-11
@@ -1,8 +1,10 @@
|
||||
# Dynamic image-list queue
|
||||
|
||||
Process a list of images across **direct local or remote ComfyUI workers** without assigning fixed chunks to each worker.
|
||||
Process a list of images across **the master and direct local or remote ComfyUI workers** without assigning fixed chunks to each participant.
|
||||
|
||||
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.
|
||||
The master evaluates the input list once. By default, the master and each selected idle worker receive one image job at a time; when a participant finishes, it gets the next unassigned image. Faster participants can therefore process more images. The master restores input order before running downstream nodes.
|
||||
|
||||
The existing **Orchestrator Only** setting (`settings.master_delegate_only`) excludes the master from per-image processing. It still creates the input list and runs downstream nodes. API callers can override the setting for a run with `delegate_master`. There is no second image-list-specific toggle.
|
||||
|
||||
## Wiring
|
||||
|
||||
@@ -22,13 +24,13 @@ Distributed Image List Start -------- context ---------+
|
||||
```
|
||||
|
||||
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.
|
||||
2. Enable Distributed and select the **direct** workers to use. The master also participates unless **Orchestrator Only** is enabled. With the master participating, an empty worker selection runs locally through the same image-list scheduler; orchestrator-only mode requires at least one reachable, idle direct worker. The master must be reachable from remote workers at its configured callback address.
|
||||
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.
|
||||
ComfyUI represents these as ordinary `IMAGE` sockets with list semantics, not a separate `IMAGE_LIST` socket type. Outside the region there is a list; for each assigned item, 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
|
||||
|
||||
@@ -44,7 +46,7 @@ Start outputs, in order: `image`, `reference_image`, `mask`, `text`, `strength`,
|
||||
|
||||
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. **The master still validates the prompt:** it needs the relevant node definitions and any model-name/file availability required by their validation (including combo lists). This is not a model-free-master guarantee.
|
||||
Models are **loaded and cached on each processing participant**, not serialized between machines. Connect loader outputs directly to the processing nodes. The master needs the region's model files when participating, as does each selected worker. **The master still validates the prompt:** it needs the relevant node definitions and any model-name/file availability required by their validation (including combo lists). This is not a model-free-master guarantee.
|
||||
|
||||
Per-image data must pass through Start; do not bypass it with upstream images, latents or generated batches connected directly into the processing region. Shared dependencies across Start are restricted to registered model outputs or scalar constants, not arbitrary image-producing leaf nodes.
|
||||
|
||||
@@ -54,16 +56,17 @@ File-based dependencies keep ComfyUI's `[output]` / `[temp]` root annotations wh
|
||||
|
||||
## 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.
|
||||
- This version supports the master and direct workers, not Runpod serverless pools. Unsupported pools are excluded without contacting the provider for this job.
|
||||
- Worker eligibility is determined before execution and checked again before dispatch. Busy or transiently unreachable workers are excluded; healthy candidates can continue. Malformed capability/queue responses and incompatible versions fail explicitly. Workers joining later are not added to an already-running region.
|
||||
- Omitted worker IDs select enabled direct workers. Explicit IDs select the named configured direct workers, including disabled ones; an empty list stays empty, unknown IDs fail, and serverless workers are excluded. Local workers use the master's loopback callback; remote workers use its configured public address.
|
||||
- Only the enclosed processing is distributed. List creation and downstream saving run on the master, even when master-delegate-only is enabled. There is no master fallback when workers are unavailable.
|
||||
- Only the enclosed processing is distributed. List creation and downstream saving run on the master in both modes. With participation enabled, the master is a normal scheduler participant and can run alone when no direct workers are available. **Orchestrator-only mode never silently falls back to master processing.** Failed items are not reassigned automatically.
|
||||
- Local items execute while the enclosing master prompt waits; they are not queued behind that prompt on the master's `/prompt` queue. The enclosed branch is not also evaluated as a full list by the outer prompt.
|
||||
- Start/End regions may be chained sequentially and share model-loader dependencies. Nested/overlapping per-image regions and data escaping a region without End are rejected. Collector and distributed-upscaler nodes cannot be mixed into the same prompt in this first version.
|
||||
- Output nodes inside a region are rejected because ComfyUI would execute them independently on the master. Avoid other side-effecting processing nodes as well; save after End.
|
||||
- The entire input list is materialized before scheduling; this is not streaming input. The limit is 4,096 frames. 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.
|
||||
- The entire input list is materialized before scheduling; this is not streaming input. The limit is 4,096 frames. Remote transport preserves floating-point tensor dtype and values rather than encoding to lossy images. Each transported tensor is limited to 50 MiB, and the complete JSON/base64 HTTP request to 101 MiB; image/reference/mask combinations can hit the aggregate request limit first. In-process master items use tensors directly and do not incur those HTTP limits.
|
||||
- Results are accepted once per item using assignment tokens. Missing, duplicate, malformed and wrong-index results fail rather than silently returning a partial list. No failed or ambiguous submission is retried automatically.
|
||||
- Worker errors fail the region. The worker timeout measures loss of evidence for the **owned prompt ID**, not maximum render duration. Pending/running evidence refreshes liveness; unrelated busy jobs and failed polls do not. Healthy cold loads and long renders can continue.
|
||||
- User interruption is checked during submission and result waiting. Cleanup targets only this region's acknowledged prompt IDs and staged input tokens. Submission may wait briefly for an acknowledgement before cancellation; transport and cleanup have finite time limits. Lost acknowledgements are **ambiguous**, never retried automatically. Unreachable cleanup is logged as unconfirmed and may require checking the worker manually; cancellation does not guarantee remote termination through a broken network.
|
||||
- User interruption is checked during submission and result waiting. Cleanup targets only this region's acknowledged prompt IDs and staged input tokens. Submission may wait briefly for an acknowledgement before cancellation; remote transport and cleanup have finite time limits. Local cleanup stops at node/progress boundaries and waits for its execution thread to finish before the parent prompt can resume. A synchronous node that does not cooperate with cancellation can delay this beyond the configured deadline. Lost acknowledgements are **ambiguous**, never retried automatically. Unreachable cleanup is logged as unconfirmed and may require checking the worker manually; cancellation does not guarantee remote termination through a broken network.
|
||||
- State is in memory. Restarting the master does not resume an interrupted region. Abandoned master contexts and expired worker staging are actively swept on the server loop; active/queued master work is protected from age-only expiry. End also removes its context on success, failure or interruption. This releases feature-owned references, not ComfyUI's separate execution cache.
|
||||
|
||||
### Time and size limits
|
||||
@@ -72,13 +75,13 @@ File-based dependencies keep ComfyUI's `[output]` / `[temp]` root annotations wh
|
||||
|
||||
Uploads, media preparation and submission acknowledgements are independently bounded. WebSocket acknowledgement waiting includes connect/send/receive and is not extended by heartbeats or unrelated messages. A user interrupt allows up to one second for a late acknowledgement; cleanup requests use five-second bounds each.
|
||||
|
||||
**4K float32 RGB outputs are not supported by these transport limits.** A 3840×2160 RGB float32 tensor exceeds 50 MiB before base64 encoding. 1080p-to-4K upscaling therefore cannot use this configuration. Dtypes/values are never silently reduced to fit. Complete encoded input requests are checked before upload, including image/reference/mask combinations; output limits are checked when the worker has computed the result. Reverse proxies may impose lower HTTP limits.
|
||||
**4K float32 RGB outputs cannot be returned by direct workers under these transport limits.** A 3840×2160 RGB float32 tensor exceeds 50 MiB before base64 encoding. Mixed master/worker 1080p-to-4K upscaling therefore cannot use this configuration; master-only execution does not use the HTTP transport. Dtypes/values are never silently reduced to fit. Complete encoded input requests are checked before upload, including image/reference/mask combinations; output limits are checked when the worker has computed the result. Reverse proxies may impose lower HTTP limits.
|
||||
|
||||
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.
|
||||
The generic acceptance helper starts two isolated loopback worker test servers and uses actual CPU ComfyUI `PromptExecutor` instances. Production queue preparation, list nodes, graph extraction, HTTP item/result routes, dynamic scheduling, float transport, ordering and cancellation are exercised. The harness checks orchestrator-only, mixed master/worker and master-only execution, including exact-once item assignment, source-once materialization and cached local model-loader dependencies. The lightweight queue/history facade is test code; no managed ComfyUI service or model is started.
|
||||
|
||||
```bash
|
||||
COMFYUI_SOURCE_ROOT=/path/to/ComfyUI /path/to/ComfyUI/.venv/bin/python -m pytest tests -q
|
||||
|
||||
@@ -0,0 +1,184 @@
|
||||
"""A queue-free Comfy graph slot, only while the synchronous outer End is blocked.
|
||||
|
||||
Do NOT use PromptExecutor.execute[_async] here. Its prompt lifecycle clears the
|
||||
process interrupt, resets progress handlers, changes RAM-cache callbacks and
|
||||
starts/ends a model tracker. Those belong to the suspended parent executor.
|
||||
Instead use core's ExecutionList/execute (including lazy/async/expanded nodes)
|
||||
with private provider-free caches and a private server facade. Model outputs are
|
||||
registered with the existing parent tracker, whose lifetime includes this work.
|
||||
|
||||
Core progress and GraphBuilder defaults are process-global. Borrow them only
|
||||
inside the owned execution thread, restore them in finally, and join that thread
|
||||
before End returns. Cancellation never clears or raises the global interrupt:
|
||||
it stops locally at node/progress boundaries. A non-cooperative synchronous node
|
||||
must finish before cleanup can return; abandoning its thread is not safe.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
import threading
|
||||
from types import SimpleNamespace
|
||||
|
||||
from .runtime import ImageListCancelledError
|
||||
|
||||
|
||||
_EXECUTION_LOCK = threading.Lock()
|
||||
|
||||
|
||||
class LocalExecutionSlot:
|
||||
def __init__(self, *, model_callback=None):
|
||||
self._pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="image-list-master")
|
||||
self._stop = threading.Event()
|
||||
self._future = None
|
||||
self._caches = None
|
||||
self._model_callback = model_callback
|
||||
self._closed = False
|
||||
self._loop = None
|
||||
self._task = None
|
||||
self._state_lock = threading.Lock()
|
||||
|
||||
def submit(self, prompt, prompt_id, end_id, workflow_meta=None):
|
||||
if self._closed or self._stop.is_set():
|
||||
raise ImageListCancelledError("Local image-list slot has stopped")
|
||||
if self._future is not None and not self._future.done():
|
||||
raise RuntimeError("Local image-list slot already has an in-flight item")
|
||||
self._future = self._pool.submit(self._run, prompt, prompt_id, end_id, workflow_meta)
|
||||
return self._future
|
||||
|
||||
def request_stop(self):
|
||||
self._stop.set()
|
||||
# Wakes a graph waiting on an async node as well as stopping synchronous
|
||||
# nodes at their next progress/between-node checkpoint.
|
||||
with self._state_lock:
|
||||
if self._loop is not None and self._task is not None:
|
||||
self._loop.call_soon_threadsafe(self._task.cancel)
|
||||
|
||||
async def close(self):
|
||||
self.request_stop()
|
||||
try:
|
||||
if self._future is not None:
|
||||
# Do not cancel this Future: cancelling an asyncio wrapper cannot
|
||||
# kill its running thread, and falsely makes cleanup look done.
|
||||
waiter = asyncio.wrap_future(self._future)
|
||||
while not waiter.done():
|
||||
try:
|
||||
await asyncio.shield(waiter)
|
||||
except asyncio.CancelledError:
|
||||
continue
|
||||
except Exception:
|
||||
break # The scheduler propagates execution errors.
|
||||
if waiter.done() and not waiter.cancelled():
|
||||
waiter.exception()
|
||||
finally:
|
||||
self._pool.shutdown(wait=True, cancel_futures=True)
|
||||
self._caches = None
|
||||
self._closed = True
|
||||
|
||||
def _run(self, prompt, prompt_id, end_id, workflow_meta):
|
||||
# The normal Comfy prompt worker is single-threaded and blocked in End.
|
||||
# Fail closed if another local slot is nevertheless executing, rather
|
||||
# than interleave writes to core's process-global execution state.
|
||||
if not _EXECUTION_LOCK.acquire(blocking=False):
|
||||
raise RuntimeError("Another local image-list executor is already active")
|
||||
try:
|
||||
asyncio.run(self._execute(prompt, prompt_id, end_id, workflow_meta))
|
||||
except asyncio.CancelledError as exc:
|
||||
raise ImageListCancelledError("Local image-list execution cancelled") from exc
|
||||
finally:
|
||||
with self._state_lock:
|
||||
self._loop = self._task = None
|
||||
_EXECUTION_LOCK.release()
|
||||
|
||||
def _check_stopped(self):
|
||||
import comfy.model_management
|
||||
if self._stop.is_set() or comfy.model_management.processing_interrupted():
|
||||
raise ImageListCancelledError("Local image-list execution interrupted")
|
||||
|
||||
async def _execute(self, prompt, prompt_id, end_id, workflow_meta):
|
||||
import execution
|
||||
import torch
|
||||
import comfy.model_management
|
||||
import comfy.utils
|
||||
from comfy_execution import progress
|
||||
from comfy_execution.graph_utils import GraphBuilder
|
||||
|
||||
self._check_stopped()
|
||||
with self._state_lock:
|
||||
self._loop, self._task = asyncio.get_running_loop(), asyncio.current_task()
|
||||
dynamic_prompt = execution.DynamicPrompt(prompt)
|
||||
private_progress = progress.ProgressRegistry(prompt_id, dynamic_prompt)
|
||||
parent_progress = progress.global_progress_registry
|
||||
parent_hook = comfy.utils.PROGRESS_BAR_HOOK
|
||||
prefix_keys = ("_default_prefix_root", "_default_prefix_call_index", "_default_prefix_graph_index")
|
||||
parent_prefix = {key: getattr(GraphBuilder, key) for key in prefix_keys}
|
||||
progress.global_progress_registry = private_progress
|
||||
|
||||
def progress_hook(value, total, preview=None, prompt_id=None, node_id=None):
|
||||
self._check_stopped()
|
||||
# Never route a child progress/preview event through the parent's
|
||||
# client or reset its handlers. Keep the private progress accurate.
|
||||
if node_id is not None:
|
||||
private_progress.update_progress(node_id, value, total, preview)
|
||||
|
||||
comfy.utils.PROGRESS_BAR_HOOK = progress_hook
|
||||
try:
|
||||
with torch.inference_mode():
|
||||
if self._caches is None:
|
||||
outputs = execution.HierarchicalCache(execution.CacheKeySetInputSignature)
|
||||
objects = execution.HierarchicalCache(execution.CacheKeySetID)
|
||||
self._caches = SimpleNamespace(outputs=outputs, objects=objects, all=[outputs, objects])
|
||||
caches = self._caches
|
||||
changed = execution.IsChangedCache(prompt_id, dynamic_prompt, caches.outputs)
|
||||
for cache in caches.all:
|
||||
await cache.set_prompt(dynamic_prompt, prompt.keys(), changed)
|
||||
cache.clean_unused()
|
||||
execution_list = execution.ExecutionList(dynamic_prompt, caches.outputs, self._model_callback)
|
||||
execution_list.add_node(end_id)
|
||||
server = SimpleNamespace(client_id=None, last_node_id=None, send_sync=lambda *a, **kw: None)
|
||||
assets = execution.NoAssets(execution.args)
|
||||
extra = {"extra_pnginfo": {"workflow": workflow_meta}} if workflow_meta is not None else {}
|
||||
executed, pending_subgraphs, pending_async, ui = set(), {}, {}, {}
|
||||
while not execution_list.is_empty():
|
||||
self._check_stopped()
|
||||
node_id, error, exc = await execution_list.stage_node_execution()
|
||||
if error is not None:
|
||||
raise RuntimeError(f"Local image-list graph failed: {error}") from exc
|
||||
self._check_stopped()
|
||||
result, error, exc = await execution.execute(
|
||||
server, dynamic_prompt, caches, node_id, extra, executed, prompt_id,
|
||||
execution_list, pending_subgraphs, pending_async, ui, assets)
|
||||
if result == execution.ExecutionResult.FAILURE:
|
||||
if isinstance(exc, (comfy.model_management.InterruptProcessingException,
|
||||
ImageListCancelledError)):
|
||||
raise ImageListCancelledError("Local image-list execution interrupted") from exc
|
||||
raise RuntimeError(f"Local image-list node {node_id} failed: {error}") from exc
|
||||
if result == execution.ExecutionResult.PENDING:
|
||||
execution_list.unstage_node_execution()
|
||||
else:
|
||||
execution_list.complete_node_execution()
|
||||
self._check_stopped()
|
||||
finally:
|
||||
# Core execute creates async-node and unblock tasks. This is a private
|
||||
# event loop, so every other task is owned here. Drain/cancel them
|
||||
# before restoring globals, including when the graph failed early.
|
||||
with self._state_lock:
|
||||
self._task = None # Stop requests must not interrupt this drain.
|
||||
try:
|
||||
while True:
|
||||
pending = [task for task in asyncio.all_tasks() if task is not asyncio.current_task()]
|
||||
if not pending:
|
||||
break
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
drain = asyncio.gather(*pending, return_exceptions=True)
|
||||
while not drain.done():
|
||||
try:
|
||||
await asyncio.shield(drain)
|
||||
except asyncio.CancelledError:
|
||||
continue
|
||||
finally:
|
||||
progress.global_progress_registry = parent_progress
|
||||
comfy.utils.PROGRESS_BAR_HOOK = parent_hook
|
||||
for key, value in parent_prefix.items():
|
||||
setattr(GraphBuilder, key, value)
|
||||
+98
-19
@@ -59,6 +59,9 @@ class ImageListJob:
|
||||
failed: str | None = None
|
||||
prompt_id: str | None = None
|
||||
active: bool = False
|
||||
master_participates: bool = False
|
||||
event_loop: asyncio.AbstractEventLoop | None = None
|
||||
local_model_callback: Any = None
|
||||
|
||||
|
||||
def _torch():
|
||||
@@ -329,15 +332,17 @@ def delete_worker_item(token: str) -> bool:
|
||||
def create_master_job(start_id: str, end_id: str, prompt: dict[str, Any], items: list[dict[str, Any]],
|
||||
worker_ids: list[str], master_callback_url: str,
|
||||
workflow_meta: dict[str, Any] | None = None,
|
||||
prompt_id: str | None = None) -> str:
|
||||
prompt_id: str | None = None, *, master_participates: bool = False) -> str:
|
||||
if not isinstance(prompt, dict):
|
||||
raise ValueError("Image List Start requires the original ComfyUI prompt for region extraction")
|
||||
if not worker_ids:
|
||||
if not isinstance(master_participates, bool):
|
||||
raise ValueError("master_participates must be a boolean")
|
||||
if not worker_ids and not master_participates:
|
||||
raise RuntimeError("Distributed Image List requires at least one preflighted direct worker; master fallback is disabled")
|
||||
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, prompt_id=prompt_id)
|
||||
workflow_meta, prompt_id=prompt_id, master_participates=master_participates)
|
||||
with _LOCK:
|
||||
_MASTER_JOBS[context_id] = state
|
||||
return context_id
|
||||
@@ -542,7 +547,7 @@ def create_worker_task_token() -> str:
|
||||
return secrets.token_urlsafe(36)
|
||||
|
||||
|
||||
async def dynamic_schedule(item_count: int, workers: list[str], dispatch, wait_next,
|
||||
async def dynamic_schedule(item_count: int, workers: list[Any], dispatch, wait_next,
|
||||
cleanup, interrupted=None):
|
||||
"""One in-flight job per worker; completed workers immediately take next item."""
|
||||
if item_count <= 0:
|
||||
@@ -550,7 +555,7 @@ async def dynamic_schedule(item_count: int, workers: list[str], dispatch, wait_n
|
||||
if not workers:
|
||||
raise RuntimeError("No free, reachable, feature-compatible direct workers are available")
|
||||
pending_index = 0
|
||||
inflight: dict[str, Any] = {}
|
||||
inflight: dict[Any, Any] = {}
|
||||
results: dict[int, Any] = {}
|
||||
|
||||
def check_interrupted():
|
||||
@@ -645,7 +650,10 @@ def set_master_result(context_id: str, item_index: int, token: str, result: dict
|
||||
state.results[item_index] = result
|
||||
event = state.event
|
||||
if event is not None:
|
||||
event.set()
|
||||
if state.event_loop is not None:
|
||||
state.event_loop.call_soon_threadsafe(event.set)
|
||||
else:
|
||||
event.set()
|
||||
|
||||
|
||||
def register_assignment(context_id: str, item_index: int, worker_id: str, token: str) -> None:
|
||||
@@ -718,13 +726,14 @@ def _collect_output_lists(results: list[dict[str, Any]]) -> tuple[list[Any], lis
|
||||
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."""
|
||||
async def run_image_list_job(context_id: str, *, model_callback=None) -> tuple[list[Any], list[Any], list[Any]]:
|
||||
"""Schedule a shared queue across the optional master and direct workers."""
|
||||
state = get_master_job(context_id)
|
||||
if state is None:
|
||||
raise RuntimeError("Distributed Image List context is missing or expired")
|
||||
with _LOCK:
|
||||
state.active = True
|
||||
state.local_model_callback = model_callback
|
||||
try:
|
||||
from ..utils.config import get_image_list_job_timeout_seconds
|
||||
deadline = get_image_list_job_timeout_seconds()
|
||||
@@ -736,6 +745,10 @@ async def run_image_list_job(context_id: str) -> tuple[list[Any], list[Any], lis
|
||||
if timer.expired():
|
||||
raise TimeoutError(f"Image-list total deadline exceeded ({deadline:g}s); cleanup attempted") from None
|
||||
raise
|
||||
except asyncio.CancelledError as exc:
|
||||
# The synchronous End bridge catches Exception, not BaseException.
|
||||
# Cleanup (including the local thread join) has finished at this point.
|
||||
raise ImageListCancelledError("Image-list job cancelled") from exc
|
||||
finally:
|
||||
remove_master_job(context_id)
|
||||
|
||||
@@ -743,6 +756,7 @@ async def run_image_list_job(context_id: str) -> tuple[list[Any], list[Any], lis
|
||||
async def _run_image_list_job(state: ImageListJob):
|
||||
context_id = state.context_id
|
||||
state.event = asyncio.Event()
|
||||
state.event_loop = asyncio.get_running_loop()
|
||||
from ..api.orchestration.dispatch import (
|
||||
preflight_image_list_workers, dispatch_worker_prompt,
|
||||
)
|
||||
@@ -754,14 +768,14 @@ async def _run_image_list_job(state: ImageListJob):
|
||||
|
||||
config = load_config()
|
||||
workers = _resolve_direct_workers_for_image_list(config, state.worker_ids)
|
||||
if not workers:
|
||||
if not workers and not state.master_participates:
|
||||
raise RuntimeError("Distributed Image List requires enabled direct workers; serverless pools are unsupported")
|
||||
# Recheck worker capability and idle state immediately before the first item.
|
||||
worker_statuses = await preflight_image_list_workers(workers, require_idle=True)
|
||||
worker_statuses = await preflight_image_list_workers(workers, require_idle=not state.master_participates) if workers else {}
|
||||
workers = [worker for worker in workers if worker["id"] in worker_statuses]
|
||||
if not workers:
|
||||
if not workers and not state.master_participates:
|
||||
raise RuntimeError("No free, reachable, feature-compatible direct workers are available; master execution is disabled")
|
||||
session = await get_client_session()
|
||||
session = await get_client_session() if workers else None
|
||||
regions = validate_image_list_regions(state.original_prompt)
|
||||
region = next((entry for entry in regions if entry["start_id"] == state.start_id and entry["end_id"] == state.end_id), None)
|
||||
if region is None:
|
||||
@@ -770,11 +784,35 @@ async def _run_image_list_job(state: ImageListJob):
|
||||
use_websocket = bool((config.get("settings", {}) or {}).get("websocket_orchestration", False))
|
||||
prepared_worker_prompts: dict[str, dict[str, Any]] = {}
|
||||
owned_uploads: dict[str, dict[str, Any]] = {}
|
||||
# An opaque key cannot collide with any configured direct worker ID, even
|
||||
# a worker named "master" or a separate local worker process.
|
||||
master_slot = object()
|
||||
local_executor = None
|
||||
local_tokens: set[str] = set()
|
||||
if state.master_participates:
|
||||
from .local_executor import LocalExecutionSlot
|
||||
local_executor = LocalExecutionSlot(model_callback=state.local_model_callback)
|
||||
|
||||
async def dispatch(worker_id: str, index: int) -> dict[str, Any]:
|
||||
async def dispatch(worker_id: Any, index: int) -> dict[str, Any]:
|
||||
if worker_id is master_slot:
|
||||
assert local_executor is not None
|
||||
token = create_worker_task_token()
|
||||
register_assignment(context_id, index, "in-process-master", token)
|
||||
local_tokens.add(token)
|
||||
# Already normalized tensors stay in memory; no HTTP, serialization,
|
||||
# /prompt queue, or repeated execution of the original source.
|
||||
with _LOCK:
|
||||
_WORKER_ITEMS[token] = (time.monotonic() + _WORKER_ITEM_TTL_SECONDS, state.items[index])
|
||||
prompt = build_worker_prompt(state.original_prompt, region, token, context_id, index, "")
|
||||
prompt[state.start_id]["inputs"]["image_list_local_execution"] = True
|
||||
prompt[state.end_id]["inputs"]["image_list_local_execution"] = True
|
||||
future = local_executor.submit(prompt, f"{context_id}:{index}", state.end_id, state.workflow_meta)
|
||||
future.add_done_callback(lambda _: state.event_loop.call_soon_threadsafe(state.event.set))
|
||||
return {"index": index, "token": token, "future": future}
|
||||
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")
|
||||
assert session is not None
|
||||
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])}
|
||||
@@ -814,12 +852,26 @@ async def _run_image_list_job(state: ImageListJob):
|
||||
handle.update(prompt_id=prompt_id, last_seen=time.monotonic(), submission_pending=False)
|
||||
return handle
|
||||
|
||||
async def wait_next(inflight: dict[str, Any]):
|
||||
async def wait_next(inflight: dict[Any, Any]):
|
||||
while True:
|
||||
state.event.clear()
|
||||
if state.failed:
|
||||
raise RuntimeError(state.failed)
|
||||
local = inflight.get(master_slot)
|
||||
if local is not None and local["future"].done():
|
||||
# A posted result does not release this slot until execution has
|
||||
# actually finished, including restoration of parent state.
|
||||
local["future"].result()
|
||||
with _LOCK:
|
||||
result = state.results.pop(local["index"], None)
|
||||
if result is None:
|
||||
raise RuntimeError(f"Master completed item {local['index']} without posting a result")
|
||||
delete_worker_item(local["token"])
|
||||
local_tokens.discard(local["token"])
|
||||
return master_slot, local["index"], result
|
||||
# Treat worker-side execution errors as failures even when End never posts a result.
|
||||
async def history_status(worker_id, handle):
|
||||
assert session is not None
|
||||
url = build_worker_url(handle["worker"], f"/history/{handle['prompt_id']}")
|
||||
try:
|
||||
async with session.get(url, timeout=aiohttp.ClientTimeout(total=3)) as response:
|
||||
@@ -853,7 +905,8 @@ async def _run_image_list_job(state: ImageListJob):
|
||||
except (asyncio.TimeoutError, aiohttp.ClientError):
|
||||
return None
|
||||
return None
|
||||
errors = await asyncio.gather(*(history_status(wid, handle) for wid, handle in inflight.items()))
|
||||
errors = await asyncio.gather(*(history_status(wid, handle) for wid, handle in inflight.items()
|
||||
if wid is not master_slot))
|
||||
failure = next((error for error in errors if isinstance(error, str)), None)
|
||||
if failure:
|
||||
raise RuntimeError(failure)
|
||||
@@ -867,12 +920,22 @@ async def _run_image_list_job(state: ImageListJob):
|
||||
return completed
|
||||
now = time.monotonic()
|
||||
expired = [f"{wid}/item {handle['index']}" for wid, handle in inflight.items()
|
||||
if now - handle["last_seen"] > timeout_seconds]
|
||||
if wid is not master_slot and now - handle["last_seen"] > timeout_seconds]
|
||||
if expired:
|
||||
raise TimeoutError(f"Image-list owned-prompt liveness timeout after {timeout_seconds:.1f}s: {', '.join(expired)}")
|
||||
await asyncio.sleep(0.25)
|
||||
try:
|
||||
await asyncio.wait_for(state.event.wait(), timeout=0.25)
|
||||
except TimeoutError:
|
||||
pass
|
||||
|
||||
async def cleanup(worker_id: str, handle: dict[str, Any]) -> None:
|
||||
async def cleanup(worker_id: Any, handle: dict[str, Any]) -> None:
|
||||
if worker_id is master_slot:
|
||||
# The outer finally owns the non-abandonable stop/join. The generic
|
||||
# scheduler's bounded remote cleanup must never detach local work.
|
||||
assert local_executor is not None
|
||||
local_executor.request_stop()
|
||||
return
|
||||
assert session is not None
|
||||
worker = handle["worker"]
|
||||
prompt_id = handle.get("prompt_id")
|
||||
if prompt_id:
|
||||
@@ -906,10 +969,26 @@ async def _run_image_list_job(state: ImageListJob):
|
||||
return False
|
||||
|
||||
try:
|
||||
results = await dynamic_schedule(len(state.items), [worker["id"] for worker in workers],
|
||||
slots = ([master_slot] if local_executor is not None else []) + [worker["id"] for worker in workers]
|
||||
results = await dynamic_schedule(len(state.items), slots,
|
||||
dispatch, wait_next, cleanup, interrupted)
|
||||
return _collect_output_lists(results)
|
||||
finally:
|
||||
if local_executor is not None:
|
||||
# No timeout here: Python cannot safely kill a synchronous node.
|
||||
# Stop at a node/progress boundary and join BEFORE releasing End.
|
||||
# Shield against repeated cancellation as well as the job deadline.
|
||||
closing = asyncio.create_task(local_executor.close())
|
||||
try:
|
||||
while not closing.done():
|
||||
try:
|
||||
await asyncio.shield(closing)
|
||||
except asyncio.CancelledError:
|
||||
continue
|
||||
closing.result()
|
||||
finally:
|
||||
for token in local_tokens:
|
||||
delete_worker_item(token)
|
||||
# Upload ownership exists before dynamic_schedule has an acknowledged
|
||||
# handle. Cover rejection/lost upload responses as well as cancellation.
|
||||
pending = list(owned_uploads.values())
|
||||
|
||||
+33
-10
@@ -63,12 +63,17 @@ class DistributedImageListStart(io.ComfyNode):
|
||||
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):
|
||||
image_list_workflow_meta=None, image_list_master_participates=None,
|
||||
image_list_local_execution=False):
|
||||
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)
|
||||
if unwrap_single(image_list_local_execution):
|
||||
from ..image_list.runtime import consume_worker_item
|
||||
item = consume_worker_item(worker_item_token)
|
||||
else:
|
||||
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")
|
||||
@@ -99,11 +104,18 @@ class DistributedImageListStart(io.ComfyNode):
|
||||
raise RuntimeError(f"No safe paired Image List End was found for Start node {start_id}")
|
||||
from comfy_execution.utils import get_executing_context
|
||||
executing = get_executing_context()
|
||||
master_participates = unwrap_single(image_list_master_participates)
|
||||
if master_participates is None:
|
||||
from ..utils.config import load_config
|
||||
master_participates = not bool((load_config().get("settings", {}) or {}).get("master_delegate_only", False))
|
||||
if not isinstance(master_participates, bool):
|
||||
raise ValueError("image_list_master_participates must be a boolean")
|
||||
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,
|
||||
prompt_id=executing.prompt_id if executing is not None else None,
|
||||
master_participates=master_participates,
|
||||
)
|
||||
return io.NodeOutput(
|
||||
[item["image"] for item in items],
|
||||
@@ -148,7 +160,7 @@ class DistributedImageListEnd(io.ComfyNode):
|
||||
is_output_node=True,
|
||||
accept_all_inputs=True,
|
||||
not_idempotent=True,
|
||||
hidden=[io.Hidden.prompt, io.Hidden.unique_id],
|
||||
hidden=[io.Hidden.prompt, io.Hidden.unique_id, io.Hidden.execution_list],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@@ -177,7 +189,8 @@ class DistributedImageListEnd(io.ComfyNode):
|
||||
|
||||
@classmethod
|
||||
def execute(cls, images=None, masks=None, texts=None, context=None,
|
||||
worker_result_token="", worker_item_index=-1, master_callback_url=""):
|
||||
worker_result_token="", worker_item_index=-1, master_callback_url="",
|
||||
image_list_local_execution=False):
|
||||
worker_result_token = unwrap_single(worker_result_token)
|
||||
master_callback_url = unwrap_single(master_callback_url)
|
||||
if worker_result_token:
|
||||
@@ -186,11 +199,16 @@ class DistributedImageListEnd(io.ComfyNode):
|
||||
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,
|
||||
)
|
||||
if unwrap_single(image_list_local_execution):
|
||||
from ..image_list.runtime import set_master_result
|
||||
set_master_result(context_values[0], int(unwrap_single(worker_item_index)),
|
||||
str(worker_result_token), result)
|
||||
else:
|
||||
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])
|
||||
@@ -201,7 +219,12 @@ class DistributedImageListEnd(io.ComfyNode):
|
||||
from ..image_list.runtime import ImageListCancelledError
|
||||
try:
|
||||
# The runtime owns the configurable deadline and bounded remote cleanup.
|
||||
result_lists = run_async_in_server_loop(run_image_list_job(context_id))
|
||||
execution_list = getattr(cls.hidden, "execution_list", None)
|
||||
# Local model outputs belong to this still-running outer prompt.
|
||||
# Add them to its existing tracker; never start/end another tracker
|
||||
# that could unpin a model still needed by the parent.
|
||||
model_callback = getattr(execution_list, "output_link_callback", None)
|
||||
result_lists = run_async_in_server_loop(run_image_list_job(context_id, model_callback=model_callback))
|
||||
except ImageListCancelledError:
|
||||
import comfy.model_management
|
||||
raise comfy.model_management.InterruptProcessingException()
|
||||
|
||||
@@ -35,7 +35,8 @@ from aiohttp import web
|
||||
import aiohttp
|
||||
|
||||
torch.set_num_threads(1)
|
||||
SOURCE_CALLS, BRANCH_CALLS, CAPTURED = [], [], []
|
||||
SOURCE_CALLS, BRANCH_CALLS, CAPTURED, MODEL_CALLS = [], [], [], []
|
||||
ALLOW_MASTER = False
|
||||
COUNT = 7
|
||||
|
||||
|
||||
@@ -62,17 +63,29 @@ class TestImageSource:
|
||||
('image', 'reference_image', 'mask', 'text', 'strength', 'seed'))
|
||||
|
||||
|
||||
class TestCachedModel:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {}}
|
||||
RETURN_TYPES = ('MODEL',)
|
||||
FUNCTION = 'run'
|
||||
def run(self):
|
||||
MODEL_CALLS.append(True)
|
||||
return (object(),)
|
||||
|
||||
|
||||
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})}}
|
||||
'seed': ('INT', {'forceInput': True}), 'index': ('INT', {'forceInput': True}),
|
||||
'model': ('MODEL',)}}
|
||||
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'
|
||||
def run(self, image, reference_image, mask, text, strength, seed, index, model):
|
||||
assert WORKER or ALLOW_MASTER, 'per-image branch executed in orchestrator-only mode'
|
||||
BRANCH_CALLS.append(index)
|
||||
# Sleep in short chunks so the real Comfy interrupt signal can be tested.
|
||||
until = time.monotonic() + DELAY
|
||||
@@ -105,7 +118,8 @@ async def setup():
|
||||
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)
|
||||
TestImageOperation=TestImageOperation, TestCapture=TestCapture,
|
||||
TestCachedModel=TestCachedModel)
|
||||
# Match the CORS-enabled direct-worker launch contract for canvas probes.
|
||||
app = web.Application(middlewares=[server.create_cors_middleware('*')])
|
||||
# Register the actual feature routes, not test replacements.
|
||||
@@ -200,14 +214,16 @@ def graph(callback, worker_ids):
|
||||
'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)]}},
|
||||
('text', 3), ('strength', 4), ('seed', 5), ('index', 6)]} | {'model': ['6', 0]}},
|
||||
'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]}},
|
||||
'6': {'class_type': 'TestCachedModel', 'inputs': {}},
|
||||
}
|
||||
|
||||
|
||||
async def coordinator():
|
||||
global ALLOW_MASTER
|
||||
ps, assets, module, app = await setup()
|
||||
runtime = importlib.import_module(module.__name__ + '.image_list.runtime')
|
||||
config_mod = importlib.import_module(module.__name__ + '.utils.config')
|
||||
@@ -258,7 +274,7 @@ async def coordinator():
|
||||
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},
|
||||
config = {'workers': workers, 'settings': {'websocket_orchestration': False, 'master_delegate_only': True},
|
||||
'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)
|
||||
@@ -328,6 +344,54 @@ async def coordinator():
|
||||
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')
|
||||
# Same public queue path, now with master participation enabled. The
|
||||
# parent queue stays occupied while its local slot processes items.
|
||||
ALLOW_MASTER = True
|
||||
config['settings']['master_delegate_only'] = False
|
||||
participation_counts = {}
|
||||
for mode, selected in [('mixed', [w['id'] for w in workers]), ('master-only', [])]:
|
||||
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
|
||||
before_sources, before_local, before_models = len(SOURCE_CALLS), len(BRANCH_CALLS), len(MODEL_CALLS)
|
||||
async with aiohttp.ClientSession() as session:
|
||||
before_worker_calls = []
|
||||
for worker in workers:
|
||||
async with session.get(f"http://127.0.0.1:{worker['port']}/test-stats") as response:
|
||||
before_worker_calls.append(len((await response.json())['calls']))
|
||||
with patch.object(config_mod, 'load_config', return_value=config), \
|
||||
patch.object(orchestration, 'load_config', return_value=config), \
|
||||
patch.object(orchestration, 'queue_prompt_payload', enqueue), \
|
||||
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
|
||||
queued = await asyncio.wait_for(orchestration.orchestrate_distributed_execution(
|
||||
prompt, None, 'test', selected), 45)
|
||||
assert executor.success, (mode, executor.status_messages)
|
||||
assert queued[2] == len(selected)
|
||||
assert len(SOURCE_CALLS) == before_sources + 1, 'source was reevaluated per local item'
|
||||
local_calls = BRANCH_CALLS[before_local:]
|
||||
assert local_calls, f'{mode}: master did not process any items'
|
||||
assert len(MODEL_CALLS) == before_models + 1, 'local model dependency did not stay cached'
|
||||
remote_calls = []
|
||||
async with aiohttp.ClientSession() as session:
|
||||
for worker, offset in zip(workers, before_worker_calls):
|
||||
async with session.get(f"http://127.0.0.1:{worker['port']}/test-stats") as response:
|
||||
current = await response.json()
|
||||
remote_calls.extend(current['calls'][offset:])
|
||||
assert not current['running'] and not current['pending'], current
|
||||
if mode == 'mixed':
|
||||
assert remote_calls, 'mixed execution failed to use direct workers'
|
||||
else:
|
||||
assert not remote_calls and len(local_calls) == COUNT
|
||||
assert Counter(local_calls + remote_calls) == Counter(range(COUNT)), (local_calls, remote_calls)
|
||||
images, masks, texts = CAPTURED[-1]
|
||||
assert len(images) == len(masks) == len(texts) == COUNT
|
||||
for i, row in enumerate(source_rows()):
|
||||
assert torch.equal(images[i], row['image'] + row['strength']), (mode, i)
|
||||
assert torch.equal(masks[i], row['mask'] * 0.5), (mode, i)
|
||||
assert texts[i] == f'item-{i}:seed-{100 + i}', (mode, texts)
|
||||
assert not runtime._MASTER_JOBS and not runtime._WORKER_ITEMS
|
||||
participation_counts[mode] = {'master': len(local_calls), 'workers': len(remote_calls)}
|
||||
print('IMAGE_LIST_MASTER_OK ' + json.dumps(participation_counts))
|
||||
ALLOW_MASTER = False
|
||||
config['settings']['master_delegate_only'] = True
|
||||
# Real Comfy user-interrupt flag while a real worker executor is active.
|
||||
slow = workers[1]
|
||||
async with aiohttp.ClientSession() as session:
|
||||
|
||||
@@ -215,7 +215,7 @@ async def check_image_list_lazy_executor(mapping, module, prompt_server, asset_m
|
||||
original_contexts = set(runtime._MASTER_JOBS)
|
||||
image_list_nodes = sys.modules[module.__name__ + '.nodes.image_list']
|
||||
workflow_meta = {'nodes': [{'id': 2, 'type': 'DistributedImageListStart'}], 'extra': {'acceptance': True}}
|
||||
async def fake_worker_results(_context_id):
|
||||
async def fake_worker_results(_context_id, *, model_callback=None):
|
||||
state = runtime.get_master_job(_context_id)
|
||||
assert state.prompt_id == 'image-list-lazy', state.prompt_id
|
||||
assert state.workflow_meta == workflow_meta, state.workflow_meta
|
||||
|
||||
@@ -0,0 +1,433 @@
|
||||
"""Master slot regressions, plus an opt-in real CPU ComfyUI subprocess probe."""
|
||||
import asyncio
|
||||
import importlib
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import torch
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def graph():
|
||||
return {
|
||||
'source': {'class_type': 'Source', 'inputs': {}},
|
||||
'start': {'class_type': 'DistributedImageListStart', 'inputs': {'images': ['source', 0]}},
|
||||
'op': {'class_type': 'Operation', 'inputs': {'image': ['start', 0]}},
|
||||
'end': {'class_type': 'DistributedImageListEnd', 'inputs': {
|
||||
'images': ['op', 0], 'context': ['start', 7]}},
|
||||
}
|
||||
|
||||
|
||||
class Response:
|
||||
status = 200
|
||||
|
||||
def __init__(self, data=None):
|
||||
self.data = data
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
return None
|
||||
|
||||
async def json(self):
|
||||
return self.data
|
||||
|
||||
def raise_for_status(self):
|
||||
pass
|
||||
|
||||
|
||||
class MasterRuntimeTests(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self):
|
||||
name = '_master_slot_tests'
|
||||
self.modules = patch.dict(sys.modules)
|
||||
self.modules.start()
|
||||
self.addCleanup(self.modules.stop)
|
||||
for suffix in ('', '.image_list', '.api', '.api.orchestration', '.utils'):
|
||||
module = types.ModuleType(name + suffix)
|
||||
module.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
|
||||
sys.modules[name + suffix] = module
|
||||
self.rt = importlib.import_module(name + '.image_list.runtime')
|
||||
self.workers = []
|
||||
self.preflight = AsyncMock(side_effect=lambda workers, **kw: {w['id']: True for w in workers})
|
||||
self.dispatched = []
|
||||
self.local_calls = []
|
||||
self.tokens = []
|
||||
self.closed = 0
|
||||
self.local_error = None
|
||||
self.local_delay = 0.005
|
||||
self.interrupted = False
|
||||
self.local_started = asyncio.Event()
|
||||
self.local_finished = asyncio.Event()
|
||||
self.deadline = 0
|
||||
self.slot_instances = []
|
||||
self.config = {'workers': self.workers, 'settings': {}}
|
||||
fixture = self
|
||||
|
||||
class FakeSlot:
|
||||
def __init__(self, model_callback=None):
|
||||
fixture.slot_instances.append(self)
|
||||
self.tasks = []
|
||||
|
||||
def submit(self, prompt, prompt_id, end_id, workflow_meta=None):
|
||||
async def execute():
|
||||
inputs = prompt[end_id]['inputs']
|
||||
index = inputs['worker_item_index']
|
||||
token = inputs['worker_result_token']
|
||||
fixture.local_calls.append(index)
|
||||
fixture.tokens.append(token)
|
||||
fixture.local_started.set()
|
||||
try:
|
||||
await asyncio.sleep(fixture.local_delay)
|
||||
if fixture.local_error:
|
||||
raise fixture.local_error
|
||||
item = fixture.rt.consume_worker_item(token)
|
||||
fixture.rt.set_master_result(fixture.context, index, token, {
|
||||
'image': item['image'], 'mask': item.get('mask'), 'text': item.get('text')})
|
||||
finally:
|
||||
fixture.local_finished.set()
|
||||
task = asyncio.create_task(execute())
|
||||
self.tasks.append(task)
|
||||
return task
|
||||
|
||||
def request_stop(self):
|
||||
pass
|
||||
|
||||
async def close(self):
|
||||
# Model a synchronous node that must finish before its thread can join.
|
||||
await asyncio.gather(*self.tasks, return_exceptions=True)
|
||||
fixture.closed += 1
|
||||
|
||||
dependencies = {
|
||||
'.image_list.local_executor': types.SimpleNamespace(LocalExecutionSlot=FakeSlot),
|
||||
'.api.orchestration.dispatch': types.SimpleNamespace(
|
||||
preflight_image_list_workers=self.preflight,
|
||||
dispatch_worker_prompt=AsyncMock(side_effect=self.dispatch)),
|
||||
'.api.orchestration.prompt_transform': types.SimpleNamespace(
|
||||
_resolve_direct_workers_for_image_list=lambda config, ids: self.workers),
|
||||
'.api.orchestration.media_sync': types.SimpleNamespace(
|
||||
sync_worker_media=AsyncMock(), fetch_worker_path_separator=AsyncMock(return_value=None)),
|
||||
'.utils.config': types.SimpleNamespace(load_config=lambda: self.config,
|
||||
get_worker_timeout_seconds=lambda: 60,
|
||||
get_image_list_job_timeout_seconds=lambda: self.deadline),
|
||||
'.utils.network': types.SimpleNamespace(
|
||||
build_master_callback_url=lambda *a, **kw: 'http://master',
|
||||
build_worker_url=lambda worker, path: 'http://remote' + path,
|
||||
get_client_session=AsyncMock(return_value=self)),
|
||||
'.utils.logging': types.SimpleNamespace(log=lambda *args: None),
|
||||
}
|
||||
for suffix, module in dependencies.items():
|
||||
sys.modules[name + suffix] = module
|
||||
self.network = dependencies['.utils.network']
|
||||
sys.modules['aiohttp'] = types.SimpleNamespace(ClientTimeout=lambda **kw: kw, ClientError=ConnectionError)
|
||||
comfy = types.ModuleType('comfy')
|
||||
comfy.model_management = types.SimpleNamespace(processing_interrupted=lambda: self.interrupted)
|
||||
sys.modules['comfy'] = comfy
|
||||
sys.modules['comfy.model_management'] = comfy.model_management
|
||||
|
||||
def create(self, participates=True, count=5):
|
||||
self.items = [{'image': torch.full((1, i + 1, 2, 3), float(i)),
|
||||
'mask': torch.full((1, i + 1, 2), float(i)), 'text': str(i), 'index': i}
|
||||
for i in range(count)]
|
||||
self.context = self.rt.create_master_job('start', 'end', graph(), self.items,
|
||||
[w['id'] for w in self.workers], '', master_participates=participates)
|
||||
self.state = self.rt.get_master_job(self.context)
|
||||
return self.context
|
||||
|
||||
def post(self, url, *, data=None, **kwargs):
|
||||
if url.endswith('/distributed/image-list/item'):
|
||||
payload = json.loads(data)
|
||||
self.rt.store_worker_item(payload['token'], payload['item'])
|
||||
return Response({})
|
||||
|
||||
def delete(self, url, **kwargs):
|
||||
self.rt.delete_worker_item(url.rsplit('/', 1)[-1])
|
||||
return Response({})
|
||||
|
||||
def get(self, url, **kwargs):
|
||||
pid = url.rsplit('/', 1)[-1]
|
||||
return Response({pid: {'status': {'completed': True, 'status_str': 'success'}}})
|
||||
|
||||
async def dispatch(self, worker, prompt, meta, **kwargs):
|
||||
inputs = prompt['end']['inputs']
|
||||
index, token = inputs['worker_item_index'], inputs['worker_result_token']
|
||||
self.dispatched.append(index)
|
||||
# A remote callback is ready immediately; the local slot is slower.
|
||||
item = self.rt.consume_worker_item(token)
|
||||
self.rt.set_master_result(self.context, index, token, {
|
||||
'image': item['image'], 'mask': item['mask'], 'text': item['text']})
|
||||
return f'remote-{index}'
|
||||
|
||||
async def test_local_only_preserves_order_and_uses_no_http(self):
|
||||
result = await self.rt.run_image_list_job(self.create())
|
||||
self.assertEqual(result[2], [str(i) for i in range(5)])
|
||||
self.assertEqual(self.local_calls, list(range(5)))
|
||||
self.assertEqual(len(self.slot_instances), 1)
|
||||
self.assertEqual(len(set(self.tokens)), 5)
|
||||
self.network.get_client_session.assert_not_awaited()
|
||||
self.assertEqual(self.closed, 1)
|
||||
self.assertFalse(self.rt._WORKER_ITEMS)
|
||||
self.assertFalse(self.rt._MASTER_JOBS)
|
||||
for i, (image, mask) in enumerate(zip(*result[:2])):
|
||||
self.assertTrue(torch.equal(image, self.items[i]['image']))
|
||||
self.assertTrue(torch.equal(mask, self.items[i]['mask']))
|
||||
|
||||
async def test_mixed_pool_claims_from_one_queue_and_relaxes_remote_idle(self):
|
||||
self.workers.append({'id': 'master', 'type': 'local'}) # Not the in-process master slot.
|
||||
self.local_delay = 0.04
|
||||
result = await self.rt.run_image_list_job(self.create(count=9))
|
||||
self.assertTrue(self.local_calls)
|
||||
self.assertTrue(self.dispatched)
|
||||
self.assertEqual(sorted(self.local_calls + self.dispatched), list(range(9)))
|
||||
self.assertEqual(result[2], [str(i) for i in range(9)])
|
||||
self.preflight.assert_awaited_once_with(self.workers, require_idle=False)
|
||||
|
||||
async def test_orchestrator_mode_never_instantiates_local_slot(self):
|
||||
self.workers.append({'id': 'worker', 'type': 'remote'})
|
||||
await self.rt.run_image_list_job(self.create(participates=False))
|
||||
self.assertFalse(self.slot_instances)
|
||||
self.preflight.assert_awaited_once_with(self.workers, require_idle=True)
|
||||
with self.assertRaisesRegex(RuntimeError, 'worker|fallback'):
|
||||
self.rt.create_master_job('start', 'end', graph(), self.items, [], '')
|
||||
|
||||
async def test_local_failure_surfaces_and_reclaims_staged_token(self):
|
||||
self.local_error = RuntimeError('local branch failed')
|
||||
with self.assertRaisesRegex(RuntimeError, 'local branch failed'):
|
||||
await self.rt.run_image_list_job(self.create())
|
||||
self.assertEqual(self.closed, 1)
|
||||
self.assertFalse(self.rt._WORKER_ITEMS)
|
||||
self.assertFalse(self.rt._MASTER_JOBS)
|
||||
|
||||
async def test_deadline_joins_local_execution_before_removing_context(self):
|
||||
self.deadline, self.local_delay = 0.03, 0.1
|
||||
with self.assertRaisesRegex(TimeoutError, 'total deadline'):
|
||||
await self.rt.run_image_list_job(self.create())
|
||||
self.assertTrue(self.local_finished.is_set())
|
||||
self.assertEqual(self.closed, 1)
|
||||
self.assertFalse(self.rt._WORKER_ITEMS)
|
||||
self.assertFalse(self.rt._MASTER_JOBS)
|
||||
|
||||
|
||||
async def test_cancelled_scheduler_joins_owned_execution_even_if_cancelled_again(self):
|
||||
self.local_delay = 0.12
|
||||
task = asyncio.create_task(self.rt.run_image_list_job(self.create()))
|
||||
await self.local_started.wait()
|
||||
task.cancel()
|
||||
await asyncio.sleep(0.02)
|
||||
task.cancel()
|
||||
with self.assertRaises(self.rt.ImageListCancelledError):
|
||||
await task
|
||||
self.assertTrue(self.local_finished.is_set())
|
||||
self.assertFalse(self.rt._WORKER_ITEMS)
|
||||
self.assertFalse(self.rt._MASTER_JOBS)
|
||||
|
||||
|
||||
class RealFrameworkTests(unittest.TestCase):
|
||||
@unittest.skipUnless(os.environ.get('COMFYUI_SOURCE_ROOT'), 'set COMFYUI_SOURCE_ROOT for the real CPU framework probe')
|
||||
def test_real_framework_local_slot(self):
|
||||
result = subprocess.run([sys.executable, str(Path(__file__).resolve()), '--real-probe',
|
||||
os.environ['COMFYUI_SOURCE_ROOT']], text=True, capture_output=True, timeout=60)
|
||||
self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
|
||||
self.assertIn('MASTER_SLOT_REAL_OK', result.stdout)
|
||||
|
||||
|
||||
def real_framework_probe(comfy_root):
|
||||
import threading
|
||||
import time
|
||||
sys.path.insert(0, str(Path(comfy_root).resolve()))
|
||||
import comfy.cli_args
|
||||
comfy.cli_args.args.cpu = True
|
||||
comfy.cli_args.args.disable_assets = True
|
||||
import execution
|
||||
import nodes
|
||||
import comfy.memory_management as memory
|
||||
import comfy.model_management as models
|
||||
import comfy.utils as utils
|
||||
from comfy_execution import progress
|
||||
from comfy_execution.graph_utils import GraphBuilder
|
||||
from comfy_api.v0_0_2 import io
|
||||
|
||||
name = '_real_master_slot'
|
||||
for suffix in ('', '.image_list', '.utils', '.nodes'):
|
||||
module = types.ModuleType(name + suffix)
|
||||
module.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
|
||||
sys.modules[name + suffix] = module
|
||||
# No HTTP/service/extension bootstrap is needed by the local slot. Fail if
|
||||
# either boundary tries to use a server-loop bridge instead of direct memory.
|
||||
def forbidden(*args, **kwargs):
|
||||
raise AssertionError('local boundary used the server loop or /prompt')
|
||||
sys.modules[name + '.utils.async_helpers'] = types.SimpleNamespace(run_async_in_server_loop=forbidden)
|
||||
config = {'settings': {}}
|
||||
sys.modules[name + '.utils.config'] = types.SimpleNamespace(load_config=lambda: config)
|
||||
rt = importlib.import_module(name + '.image_list.runtime')
|
||||
local = importlib.import_module(name + '.image_list.local_executor')
|
||||
boundary = importlib.import_module(name + '.nodes.image_list')
|
||||
nodes.NODE_CLASS_MAPPINGS.update(DistributedImageListStart=boundary.DistributedImageListStart,
|
||||
DistributedImageListEnd=boundary.DistributedImageListEnd)
|
||||
loads, calls = [], []
|
||||
entered, stopped = threading.Event(), threading.Event()
|
||||
mode = ['normal']
|
||||
model = object()
|
||||
|
||||
class Loader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {}}
|
||||
RETURN_TYPES = ('MODEL',)
|
||||
FUNCTION = 'run'
|
||||
def run(self):
|
||||
loads.append(True)
|
||||
return (model,)
|
||||
|
||||
class Operation:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'image': ('IMAGE',), 'model': ('MODEL',),
|
||||
'mask': ('MASK',), 'text': ('STRING', {'forceInput': True})}}
|
||||
RETURN_TYPES = ('IMAGE', 'MASK', 'STRING')
|
||||
FUNCTION = 'run'
|
||||
async def run(self, image, model, mask, text):
|
||||
calls.append(text)
|
||||
entered.set()
|
||||
if mode[0] == 'error':
|
||||
raise ValueError('probe branch failure')
|
||||
if mode[0] in ('cancel', 'interrupt'):
|
||||
try:
|
||||
while True:
|
||||
utils.ProgressBar(1).update(1)
|
||||
await asyncio.sleep(0.01)
|
||||
finally:
|
||||
await asyncio.sleep(0.02)
|
||||
stopped.set()
|
||||
if mode[0] == 'sync':
|
||||
# No progress/interrupt checks: cleanup still must join, not
|
||||
# release the parent just because an asyncio Future is cancelled.
|
||||
time.sleep(0.08)
|
||||
stopped.set()
|
||||
await asyncio.sleep(0)
|
||||
return image + 1, mask + 1, text + ':done'
|
||||
|
||||
class Expanded:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return Operation.INPUT_TYPES()
|
||||
RETURN_TYPES = Operation.RETURN_TYPES
|
||||
FUNCTION = 'run'
|
||||
def run(self, image, model, mask, text):
|
||||
builder = GraphBuilder()
|
||||
node = builder.node('Operation', image=image, model=model, mask=mask, text=text)
|
||||
return {'result': (node.out(0), node.out(1), node.out(2)), 'expand': builder.finalize()}
|
||||
|
||||
nodes.NODE_CLASS_MAPPINGS.update(Loader=Loader, Operation=Operation, Expanded=Expanded)
|
||||
prompt = graph()
|
||||
prompt['loader'] = {'class_type': 'Loader', 'inputs': {}}
|
||||
prompt['op']['inputs'].update(model=['loader', 0], mask=['start', 2], text=['start', 3])
|
||||
prompt['end']['inputs'].update(masks=['op', 1], texts=['op', 2])
|
||||
items = [{'image': torch.full((1, i + 1, 2, 3), float(i)),
|
||||
'mask': torch.full((1, i + 1, 2), float(i)), 'text': str(i), 'index': i} for i in range(3)]
|
||||
region = rt.validate_image_list_regions(prompt)[0]
|
||||
parent_progress = progress.ProgressRegistry('outer', execution.DynamicPrompt({}))
|
||||
parent_progress.nodes['outer'] = {'state': progress.NodeState.Running, 'value': 3, 'max': 7}
|
||||
progress.global_progress_registry = parent_progress
|
||||
parent_hook = lambda *a, **kw: None
|
||||
utils.PROGRESS_BAR_HOOK = parent_hook
|
||||
ram_callback = lambda *a, **kw: None
|
||||
memory.set_ram_cache_release_state(ram_callback, 123)
|
||||
GraphBuilder.set_default_prefix('outer-prefix', 7, 9)
|
||||
|
||||
def restored():
|
||||
assert progress.global_progress_registry is parent_progress
|
||||
assert parent_progress.nodes == {'outer': {'state': progress.NodeState.Running, 'value': 3, 'max': 7}}
|
||||
assert utils.PROGRESS_BAR_HOOK is parent_hook
|
||||
assert memory.extra_ram_release_callback is ram_callback and memory.RAM_CACHE_HEADROOM == 123
|
||||
assert (GraphBuilder._default_prefix_root, GraphBuilder._default_prefix_call_index,
|
||||
GraphBuilder._default_prefix_graph_index) == ('outer-prefix', 7, 9)
|
||||
assert not any(t.name.startswith('image-list-master') for t in threading.enumerate())
|
||||
|
||||
async def probe():
|
||||
# Verify actual V3 input parsing and list mapping, not only direct execute.
|
||||
for supplied, delegated, expected in [(None, False, True), (None, True, False),
|
||||
(False, False, False), (True, True, True)]:
|
||||
config['settings']['master_delegate_only'] = delegated
|
||||
inputs = {'image_list_enabled_worker_ids': '["remote"]'}
|
||||
if supplied is not None:
|
||||
inputs['image_list_master_participates'] = supplied
|
||||
cls = boundary.DistributedImageListStart
|
||||
data, _, v3 = execution.get_input_data(inputs, cls, 'start', dynprompt=execution.DynamicPrompt(prompt))
|
||||
data['images'] = [torch.zeros((3, 2, 2, 3))]
|
||||
result = await execution._async_map_node_over_list('outer', 'start', cls, data, cls.FUNCTION, v3_data=v3)
|
||||
result = await execution.resolve_map_node_over_list_results(result)
|
||||
context = result[0].result[7][0]
|
||||
assert rt.get_master_job(context).master_participates is expected
|
||||
rt.remove_master_job(context)
|
||||
|
||||
for behavior in ('normal', 'expand', 'error', 'cancel', 'interrupt', 'sync', 'already-interrupted'):
|
||||
mode[0] = behavior
|
||||
prompt['op']['class_type'] = 'Expanded' if behavior == 'expand' else 'Operation'
|
||||
loads.clear()
|
||||
calls.clear()
|
||||
entered.clear()
|
||||
stopped.clear()
|
||||
tracked = []
|
||||
slot = local.LocalExecutionSlot(model_callback=tracked.append)
|
||||
context = rt.create_master_job('start', 'end', prompt, items, [], '', master_participates=True)
|
||||
try:
|
||||
for index in range(3 if behavior in ('normal', 'expand') else 1):
|
||||
token = rt.create_worker_task_token()
|
||||
rt.register_assignment(context, index, 'local', token)
|
||||
rt._WORKER_ITEMS[token] = (time.monotonic() + 60, items[index])
|
||||
worker_prompt = rt.build_worker_prompt(prompt, region, token, context, index, '')
|
||||
for key in ('start', 'end'):
|
||||
worker_prompt[key]['inputs']['image_list_local_execution'] = True
|
||||
if behavior == 'already-interrupted':
|
||||
models.interrupt_current_processing(True)
|
||||
future = slot.submit(worker_prompt, f'local-{index}', 'end')
|
||||
if behavior in ('cancel', 'interrupt', 'sync'):
|
||||
assert await asyncio.to_thread(entered.wait, 5)
|
||||
if behavior == 'interrupt':
|
||||
models.interrupt_current_processing(True)
|
||||
else:
|
||||
slot.request_stop()
|
||||
try:
|
||||
await asyncio.wait_for(asyncio.shield(asyncio.wrap_future(future)), 5)
|
||||
except rt.ImageListCancelledError:
|
||||
assert behavior in ('cancel', 'interrupt', 'sync', 'already-interrupted'), behavior
|
||||
except RuntimeError as exc:
|
||||
assert behavior == 'error' and 'probe branch failure' in str(exc), str(exc)
|
||||
else:
|
||||
assert behavior in ('normal', 'expand'), behavior
|
||||
result = rt.get_master_job(context).results[index]
|
||||
assert torch.equal(result['image'], items[index]['image'] + 1)
|
||||
assert torch.equal(result['mask'], items[index]['mask'] + 1)
|
||||
assert result['text'] == f'{index}:done'
|
||||
if behavior in ('normal', 'expand'):
|
||||
assert len(loads) == 1 and calls == ['0', '1', '2'], (loads, calls)
|
||||
assert any(model in outputs for outputs in tracked), tracked
|
||||
if behavior in ('cancel', 'interrupt', 'sync'):
|
||||
assert stopped.is_set(), behavior
|
||||
if behavior == 'already-interrupted':
|
||||
assert models.processing_interrupted() and not calls
|
||||
finally:
|
||||
await slot.close()
|
||||
for token in rt.get_master_job(context).expected_tokens.values():
|
||||
rt.delete_worker_item(token)
|
||||
rt.remove_master_job(context)
|
||||
models.interrupt_current_processing(False) # Only test setup/teardown.
|
||||
restored()
|
||||
assert not rt._MASTER_JOBS and not rt._WORKER_ITEMS
|
||||
asyncio.run(probe())
|
||||
print('MASTER_SLOT_REAL_OK V3 metadata; cache reuse; async/expanded graph; errors; cancellation; state restoration')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
if len(sys.argv) > 1 and sys.argv[1] == '--real-probe':
|
||||
real_framework_probe(sys.argv[2])
|
||||
else:
|
||||
unittest.main()
|
||||
@@ -0,0 +1,114 @@
|
||||
"""Image-list participation uses the existing master orchestrator setting."""
|
||||
import json
|
||||
import sys
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from test_serverless_integration import PKG, load_modules
|
||||
|
||||
|
||||
class MasterSelectionTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
self.q, modules, self.ps = load_modules()
|
||||
self.modules_patch = patch.dict(sys.modules, modules)
|
||||
self.modules_patch.start()
|
||||
self.addCleanup(self.modules_patch.stop)
|
||||
self.worker = {"id": "worker", "name": "Worker", "type": "local", "enabled": True}
|
||||
self.config = {"workers": [self.worker], "settings": {}}
|
||||
self.q.load_config = lambda: self.config
|
||||
# Model the legacy helper's fallback: image-list must not inherit it.
|
||||
self.q.select_active_workers = AsyncMock(side_effect=lambda workers, ws, delegate, **kw:
|
||||
(workers, delegate if workers else False))
|
||||
self.preflight = AsyncMock(return_value={"worker": True})
|
||||
sys.modules[PKG + ".api.orchestration.dispatch"].preflight_image_list_workers = self.preflight
|
||||
self.q.queue_prompt_payload = AsyncMock(return_value={"prompt_id": "master-prompt", "number": 1})
|
||||
self.prompt = {
|
||||
"1": {"class_type": "EmptyImage", "inputs": {}},
|
||||
"2": {"class_type": "DistributedImageListStart", "inputs": {"images": ["1", 0]}},
|
||||
"3": {"class_type": "DistributedImageListEnd", "inputs": {"images": ["2", 0], "context": ["2", 7]}},
|
||||
}
|
||||
|
||||
async def queue(self, **kwargs):
|
||||
return await self.q.orchestrate_distributed_execution(self.prompt, None, "client", **kwargs)
|
||||
|
||||
def prepared_inputs(self):
|
||||
return self.q.queue_prompt_payload.call_args.args[0]["2"]["inputs"]
|
||||
|
||||
async def test_default_includes_master_and_direct_workers(self):
|
||||
result = await self.queue()
|
||||
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
|
||||
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker"])
|
||||
self.assertEqual(result[2], 1) # Public worker_count still counts remote participants only.
|
||||
self.preflight.assert_awaited_once_with(
|
||||
[{"id": "worker", "name": "Worker", "type": "local", "host": None, "port": 8188}],
|
||||
require_idle=False)
|
||||
self.q.dispatch_worker_prompt.assert_not_awaited()
|
||||
self.assertEqual(self.ps.distributed_pending_jobs, {})
|
||||
|
||||
async def test_orchestrator_setting_excludes_master_processing(self):
|
||||
self.config["settings"]["master_delegate_only"] = True
|
||||
await self.queue()
|
||||
self.assertIs(self.prepared_inputs()["image_list_master_participates"], False)
|
||||
self.assertTrue(self.preflight.call_args.kwargs["require_idle"])
|
||||
# Upstream source stays on master even in orchestrator-only mode.
|
||||
self.assertIn("1", self.q.queue_prompt_payload.call_args.args[0])
|
||||
|
||||
async def test_explicit_override_wins_over_config_in_both_directions(self):
|
||||
for configured, requested in [(True, False), (False, True)]:
|
||||
with self.subTest(configured=configured, requested=requested):
|
||||
self.config["settings"]["master_delegate_only"] = configured
|
||||
await self.queue(delegate_master=requested)
|
||||
self.assertIs(self.prepared_inputs()["image_list_master_participates"], not requested)
|
||||
|
||||
async def test_master_only_with_explicit_empty_worker_selection(self):
|
||||
result = await self.queue(enabled_worker_ids=[])
|
||||
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
|
||||
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), [])
|
||||
self.assertEqual(result[2], 0)
|
||||
self.preflight.assert_not_awaited()
|
||||
|
||||
async def test_master_only_when_no_workers_are_configured(self):
|
||||
self.config["workers"] = []
|
||||
await self.queue()
|
||||
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
|
||||
|
||||
async def test_unavailable_workers_do_not_prevent_enabled_master(self):
|
||||
self.q.select_active_workers.return_value = ([], False)
|
||||
self.q.select_active_workers.side_effect = None
|
||||
self.preflight.return_value = {}
|
||||
await self.queue()
|
||||
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), [])
|
||||
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
|
||||
|
||||
async def test_busy_workers_do_not_prevent_enabled_master(self):
|
||||
self.preflight.return_value = {}
|
||||
await self.queue()
|
||||
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), [])
|
||||
|
||||
async def test_orchestrator_without_workers_fails_before_queue(self):
|
||||
self.config["settings"]["master_delegate_only"] = True
|
||||
with self.assertRaisesRegex(RuntimeError, "orchestrator|Orchestrator"):
|
||||
await self.queue(enabled_worker_ids=[])
|
||||
self.q.queue_prompt_payload.assert_not_awaited()
|
||||
|
||||
async def test_orchestrator_never_inherits_legacy_offline_master_fallback(self):
|
||||
self.config["settings"]["master_delegate_only"] = True
|
||||
self.q.select_active_workers.side_effect = None
|
||||
self.q.select_active_workers.return_value = ([], False)
|
||||
self.preflight.return_value = {}
|
||||
with self.assertRaisesRegex(RuntimeError, "orchestrator|Orchestrator"):
|
||||
await self.queue()
|
||||
self.q.queue_prompt_payload.assert_not_awaited()
|
||||
|
||||
async def test_serverless_only_uses_master_without_provider_dispatch(self):
|
||||
self.config["workers"] = [{"id": "pool", "type": "serverless", "enabled": True}]
|
||||
result = await self.queue()
|
||||
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
|
||||
self.assertEqual(result[2], 0)
|
||||
self.q.dispatch_worker_prompt.assert_not_awaited()
|
||||
|
||||
async def test_incompatible_worker_is_not_silently_hidden_by_master(self):
|
||||
self.preflight.side_effect = RuntimeError("upgrade worker")
|
||||
with self.assertRaisesRegex(RuntimeError, "upgrade worker"):
|
||||
await self.queue()
|
||||
self.q.queue_prompt_payload.assert_not_awaited()
|
||||
@@ -34,3 +34,4 @@ def test_image_list_network_runtime():
|
||||
assert result.returncode == 0, result.stdout + '\n' + result.stderr
|
||||
assert 'IMAGE_LIST_NETWORK_OK' in result.stdout
|
||||
assert 'IMAGE_LIST_CANCEL_OK' in result.stdout
|
||||
assert 'IMAGE_LIST_MASTER_OK' in result.stdout
|
||||
|
||||
@@ -37,12 +37,12 @@ export async function executeParallelDistributed(extension, promptWrapper) {
|
||||
// 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.");
|
||||
if (hasImageList && activeWorkers.length === 0 && extension.config?.settings?.master_delegate_only) {
|
||||
throw new Error("Distributed Image List in orchestrator-only mode requires an active direct worker.");
|
||||
}
|
||||
|
||||
// Case: Enabled workers but all offline
|
||||
if (activeWorkers.length === 0 && enabledWorkers.length > 0) {
|
||||
if (!hasImageList && activeWorkers.length === 0 && enabledWorkers.length > 0) {
|
||||
extension.log("No active workers found. All enabled workers are offline.");
|
||||
if (extension.ui?.showToast) {
|
||||
extension.ui.showToast(extension.app, "error", "All Workers Offline",
|
||||
|
||||
@@ -31,9 +31,10 @@ describe('image-list queue interception', () => {
|
||||
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 () => {
|
||||
it('does not fall back when orchestrator-only workers are offline', async () => {
|
||||
vi.useFakeTimers();
|
||||
const ext = extension([{ id: 'one', type: 'local' }], false);
|
||||
ext.config.settings.master_delegate_only = true;
|
||||
setupInterceptor(ext);
|
||||
expect(await api.queuePrompt(0, prompt)).toBeNull();
|
||||
expect(ext.originalQueuePrompt).not.toHaveBeenCalled();
|
||||
@@ -43,10 +44,36 @@ describe('image-list queue interception', () => {
|
||||
vi.useFakeTimers();
|
||||
const ext = extension([{ id: 'pool', type: 'serverless' }]);
|
||||
setupInterceptor(ext);
|
||||
expect(await api.queuePrompt(0, prompt)).toBeNull();
|
||||
expect(await api.queuePrompt(0, prompt)).toEqual({ prompt_id: 'distributed' });
|
||||
expect(ext.api.queueDistributed).toHaveBeenCalledWith(expect.objectContaining({ enabled_worker_ids: [], delegate_master: false }));
|
||||
expect(ext.api.getServerlessStatus).not.toHaveBeenCalled();
|
||||
expect(ext.originalQueuePrompt).not.toHaveBeenCalled();
|
||||
});
|
||||
it.each([{ workers: [] }, { workers: [{ id: 'one', type: 'local' }] }])('queues master-only through orchestration when no workers are online: $workers', async ({ workers }) => {
|
||||
vi.useFakeTimers();
|
||||
const ext = extension(workers, false);
|
||||
setupInterceptor(ext);
|
||||
expect(await api.queuePrompt(0, prompt)).toEqual({ prompt_id: 'distributed' });
|
||||
expect(ext.api.queueDistributed).toHaveBeenCalledWith(expect.objectContaining({ enabled_worker_ids: [], delegate_master: false }));
|
||||
expect(ext.originalQueuePrompt).not.toHaveBeenCalled();
|
||||
});
|
||||
it('passes orchestrator-only mode to backend when workers are online', async () => {
|
||||
vi.useFakeTimers();
|
||||
const ext = extension([{ id: 'one', type: 'local' }]);
|
||||
ext.config.settings.master_delegate_only = true;
|
||||
setupInterceptor(ext);
|
||||
await api.queuePrompt(0, prompt);
|
||||
expect(ext.api.queueDistributed).toHaveBeenCalledWith(expect.objectContaining({ delegate_master: true }));
|
||||
});
|
||||
it('does not bypass backend errors with local execution', async () => {
|
||||
vi.useFakeTimers();
|
||||
const ext = extension([]);
|
||||
ext.api.queueDistributed.mockRejectedValue(new Error('invalid region'));
|
||||
setupInterceptor(ext);
|
||||
expect(await api.queuePrompt(0, prompt)).toBeNull();
|
||||
expect(ext.api.queueDistributed).toHaveBeenCalled();
|
||||
expect(ext.originalQueuePrompt).not.toHaveBeenCalled();
|
||||
});
|
||||
it('preserves ordinary prompt handling', async () => {
|
||||
const ext = extension([]);
|
||||
setupInterceptor(ext);
|
||||
|
||||
Reference in New Issue
Block a user