From a34969e6f36362b5acdb7c0b958d5ac78aff35eb Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sun, 11 Oct 2026 21:20:22 +0000 Subject: [PATCH] fix(image-list): own job lifecycle and preserve worker preparation --- api/image_list_routes.py | 3 +- api/orchestration/prompt_transform.py | 2 + api/queue_orchestration.py | 1 + docs/image-list-queue.md | 24 +- image_list/lifecycle.py | 142 ++++++++ image_list/runtime.py | 182 ++++++++--- nodes/image_list.py | 12 +- runtime/bootstrap.py | 3 + tests/helpers/image_list_network_check.py | 4 +- tests/helpers/v3_runtime_check.py | 21 ++ tests/test_image_list_lifecycle.py | 177 ++++++++++ tests/test_image_list_runtime_boundaries.py | 340 ++++++++++++++++++++ 12 files changed, 847 insertions(+), 64 deletions(-) create mode 100644 image_list/lifecycle.py create mode 100644 tests/test_image_list_lifecycle.py create mode 100644 tests/test_image_list_runtime_boundaries.py diff --git a/api/image_list_routes.py b/api/image_list_routes.py index 2e97b7e..3a43db2 100644 --- a/api/image_list_routes.py +++ b/api/image_list_routes.py @@ -5,7 +5,7 @@ from aiohttp import web import server from ..image_list.runtime import ( - MAX_PAYLOAD_BYTES, + MAX_REQUEST_BYTES, decode_result, delete_worker_item, set_master_result, @@ -13,7 +13,6 @@ from ..image_list.runtime import ( ) prompt_server = server.PromptServer.instance -MAX_REQUEST_BYTES = MAX_PAYLOAD_BYTES * 2 + 1024 * 1024 async def _json_body(request): diff --git a/api/orchestration/prompt_transform.py b/api/orchestration/prompt_transform.py index c03a039..ded449a 100644 --- a/api/orchestration/prompt_transform.py +++ b/api/orchestration/prompt_transform.py @@ -607,6 +607,7 @@ def apply_participant_overrides( master_url, delegate_master, prompt_index, + workflow_meta=None, ): """Return a prompt copy with hidden inputs configured for master/worker.""" is_master = participant_id == "master" @@ -641,6 +642,7 @@ def apply_participant_overrides( enabled_worker_ids, job_id_map, master_url, + workflow_meta=workflow_meta, ) return prompt_copy diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index 28b8082..6f517d2 100644 --- a/api/queue_orchestration.py +++ b/api/queue_orchestration.py @@ -413,6 +413,7 @@ async def orchestrate_distributed_execution( master_url, delegate_master, prompt_index, + workflow_meta=workflow_meta, ) if delegate_master: diff --git a/docs/image-list-queue.md b/docs/image-list-queue.md index 75dbc24..493586b 100644 --- a/docs/image-list-queue.md +++ b/docs/image-list-queue.md @@ -44,19 +44,35 @@ 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. Per-image data must pass through Start; do not bypass it with a second image list connected directly into the processing region. +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. + +Per-image data must pass through Start; do not bypass it with upstream images, latents or generated batches connected directly into the processing region. Shared dependencies across Start are restricted to registered model outputs or scalar constants, not arbitrary image-producing leaf nodes. + +End's optional `texts` is a connection-only socket. For workflows saved before this fix, recreate End and reconnect its wires if an old text widget or literal `texts: ""` remains in the compiled prompt. Arbitrary text literals are not accepted at End. + +File-based dependencies keep ComfyUI's `[output]` / `[temp]` root annotations while being read on the master. Their worker copies use separate content-addressed paths under `input/distributed-media/`; each worker's rewritten prompt is reused across items without reusing item tokens. ## 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. +- 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. - 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. +- 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. +- 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 + +`settings.image_list_job_timeout_seconds` defaults to `0` (no arbitrary total-run cap). A positive finite number opts into an overall deadline with owned-work cleanup; negative, non-finite or malformed settings fail explicitly. There is no fixed one-hour region deadline. The setting is separate from the existing worker liveness timeout. + +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. 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. diff --git a/image_list/lifecycle.py b/image_list/lifecycle.py new file mode 100644 index 0000000..fe4c0c1 --- /dev/null +++ b/image_list/lifecycle.py @@ -0,0 +1,142 @@ +"""Server-owned background cleanup for image-list runtime state.""" +from __future__ import annotations + +import asyncio +import logging +import math +import time +from collections.abc import Iterable +from typing import Any, cast + +LOGGER = logging.getLogger(__name__) +DEFAULT_SWEEP_INTERVAL_SECONDS = 30.0 +DEFAULT_ORPHAN_GRACE_SECONDS = 300.0 +_LIFECYCLE_MARKER = object() + + +def _queue_prompt_ids(prompt_server: Any) -> set[str]: + """Read running/pending prompt IDs, preferring ComfyUI's volatile view.""" + queue = getattr(prompt_server, "prompt_queue", None) + readers = ( + getattr(queue, "get_current_queue_volatile", None), + getattr(queue, "get_current_queue", None), + ) + errors: list[Exception] = [] + for reader in readers: + if not callable(reader): + continue + try: + snapshot = reader() + if not isinstance(snapshot, (tuple, list)) or len(snapshot) < 2: + raise TypeError("prompt queue snapshot must contain running and pending entries") + running, pending = snapshot[0], snapshot[1] + prompt_ids: set[str] = set() + for entries in (running, pending): + if not isinstance(entries, Iterable): + raise TypeError("prompt queue entries must be iterable") + for entry in entries: + if isinstance(entry, (tuple, list)) and len(entry) > 1: + prompt_id = entry[1] + if isinstance(prompt_id, str): + prompt_ids.add(prompt_id) + return prompt_ids + except Exception as exc: + errors.append(exc) + + if errors: + raise RuntimeError("unable to read the prompt queue") from errors[-1] + raise RuntimeError("prompt queue has no supported queue reader") + + +def _sweep_orphan_contexts(prompt_server: Any, runtime: Any, orphan_grace_seconds: float) -> None: + queue_read_ok = True + try: + queued_prompt_ids = _queue_prompt_ids(prompt_server) + except Exception: + # A failed queue read must not make a submitted context look orphaned. + LOGGER.warning("Could not read prompt queue during image-list lifecycle sweep", exc_info=True) + queued_prompt_ids = set() + queue_read_ok = False + + now = time.monotonic() + with runtime._LOCK: + for context_id, state in list(runtime._MASTER_JOBS.items()): + if getattr(state, "active", False): + continue + prompt_id = getattr(state, "prompt_id", None) + if prompt_id is not None: + if not queue_read_ok or prompt_id in queued_prompt_ids: + continue + created_at = getattr(state, "created_at", None) + expired = ( + isinstance(created_at, (int, float)) + and not isinstance(created_at, bool) + and math.isfinite(created_at) + and now - created_at >= orphan_grace_seconds + ) + if expired and runtime._MASTER_JOBS.get(context_id) is state: + runtime.remove_master_job(context_id) + + +def _sweep_once(prompt_server: Any, runtime: Any, orphan_grace_seconds: float) -> None: + try: + runtime._sweep_worker_items() + except Exception: + LOGGER.exception("Could not sweep expired image-list worker staging") + _sweep_orphan_contexts(prompt_server, runtime, orphan_grace_seconds) + + +async def _sweep_loop( + prompt_server: Any, + runtime: Any, + interval_seconds: float, + orphan_grace_seconds: float, +) -> None: + while True: + try: + _sweep_once(prompt_server, runtime, orphan_grace_seconds) + except Exception: + # Keep the server-owned task alive if one pass encounters unexpected state. + LOGGER.exception("Image-list lifecycle sweep failed") + await asyncio.sleep(interval_seconds) + + +def install_lifecycle( + prompt_server: Any, + runtime: Any, + *, + interval_seconds: float = DEFAULT_SWEEP_INTERVAL_SECONDS, + orphan_grace_seconds: float = DEFAULT_ORPHAN_GRACE_SECONDS, +) -> None: + """Register one cancellable periodic sweep in the prompt server's app context.""" + interval_seconds = float(interval_seconds) + orphan_grace_seconds = float(orphan_grace_seconds) + if not math.isfinite(interval_seconds) or interval_seconds <= 0: + raise ValueError("interval_seconds must be a positive finite value") + if not math.isfinite(orphan_grace_seconds) or orphan_grace_seconds < 0: + raise ValueError("orphan_grace_seconds must be a non-negative finite value") + + app = getattr(prompt_server, "app", None) + cleanup_ctx = getattr(app, "cleanup_ctx", None) + if cleanup_ctx is None: + raise TypeError("prompt_server must expose an app with cleanup_ctx") + if any(getattr(factory, "_image_list_lifecycle_marker", None) is _LIFECYCLE_MARKER + for factory in cleanup_ctx): + return + + async def lifecycle_context(_app): + task = asyncio.create_task( + _sweep_loop(prompt_server, runtime, interval_seconds, orphan_grace_seconds), + name="image-list-lifecycle-sweeper", + ) + try: + yield + finally: + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + cast(Any, lifecycle_context)._image_list_lifecycle_marker = _LIFECYCLE_MARKER + cleanup_ctx.append(lifecycle_context) diff --git a/image_list/runtime.py b/image_list/runtime.py index 448651d..c1cf781 100644 --- a/image_list/runtime.py +++ b/image_list/runtime.py @@ -23,7 +23,11 @@ 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_REQUEST_BYTES = MAX_PAYLOAD_BYTES * 2 + 1024 * 1024 MAX_ITEMS = 4096 +SUBMISSION_TIMEOUT_SECONDS = 190.0 +SUBMISSION_CANCEL_GRACE_SECONDS = 1.0 +CLEANUP_TIMEOUT_SECONDS = 16.0 _ALLOWED_FLOAT_DTYPES = {"float16", "bfloat16", "float32", "float64"} _LOCK = threading.RLock() @@ -53,6 +57,8 @@ class ImageListJob: event: asyncio.Event | None = None created_at: float = field(default_factory=time.monotonic) failed: str | None = None + prompt_id: str | None = None + active: bool = False def _torch(): @@ -241,6 +247,17 @@ def decode_value(data: Any, expected_kind: str) -> Any: return encode_value(value, expected_kind) and value +def encode_request_payload(payload: dict[str, Any]) -> bytes: + """Budget the final JSON envelope, including base64, before HTTP upload.""" + body = json.dumps(payload, separators=(",", ":"), allow_nan=False).encode("utf-8") + if len(body) > MAX_REQUEST_BYTES: + raise ValueError( + f"Image-list complete JSON/base64 request is {len(body)} bytes; " + f"limit is {MAX_REQUEST_BYTES} bytes (image/reference/mask combined)" + ) + return body + + 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"), @@ -311,7 +328,8 @@ 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) -> str: + workflow_meta: dict[str, Any] | None = None, + prompt_id: str | 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: @@ -319,7 +337,7 @@ def create_master_job(start_id: str, end_id: str, prompt: dict[str, Any], items: 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) + workflow_meta, prompt_id=prompt_id) with _LOCK: _MASTER_JOBS[context_id] = state return context_id @@ -556,16 +574,30 @@ async def dynamic_schedule(item_count: int, workers: list[str], dispatch, wait_n 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)) + # Bound the shield itself, and inspect the real Comfy interrupt flag while + # awaiting an acknowledgement. Briefly retain a late acknowledged handle. + task = asyncio.create_task(asyncio.wait_for( + dispatch(worker, index), SUBMISSION_TIMEOUT_SECONDS)) try: - return await asyncio.shield(task) - except asyncio.CancelledError: + while True: + check_interrupted() + done, _ = await asyncio.wait({task}, timeout=0.1) + if done: + handle = task.result() + inflight[worker] = handle + check_interrupted() + return handle + except (asyncio.CancelledError, ImageListCancelledError): try: - inflight[worker] = await task - except Exception: + inflight[worker] = await asyncio.wait_for( + asyncio.shield(task), SUBMISSION_CANCEL_GRACE_SECONDS) + except (Exception, asyncio.CancelledError): pass raise + finally: + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) try: for worker in workers: @@ -594,7 +626,8 @@ async def dynamic_schedule(item_count: int, workers: list[str], dispatch, wait_n 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) + await asyncio.gather(*(asyncio.wait_for(cleanup(worker, handle), CLEANUP_TIMEOUT_SECONDS) + for worker, handle in inflight.items()), return_exceptions=True) raise @@ -690,8 +723,19 @@ async def run_image_list_job(context_id: str) -> tuple[list[Any], list[Any], lis 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 try: - return await _run_image_list_job(state) + from ..utils.config import get_image_list_job_timeout_seconds + deadline = get_image_list_job_timeout_seconds() + timer = asyncio.timeout(deadline if deadline > 0 else None) + try: + async with timer: + return await _run_image_list_job(state) + except TimeoutError: + if timer.expired(): + raise TimeoutError(f"Image-list total deadline exceeded ({deadline:g}s); cleanup attempted") from None + raise finally: remove_master_job(context_id) @@ -704,7 +748,7 @@ async def _run_image_list_job(state: ImageListJob): ) 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.network import build_master_callback_url, build_worker_url, get_client_session from ..utils.logging import log import aiohttp @@ -718,15 +762,14 @@ async def _run_image_list_job(state: ImageListJob): 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] = {} + prepared_worker_prompts: dict[str, dict[str, Any]] = {} + owned_uploads: dict[str, dict[str, Any]] = {} async def dispatch(worker_id: str, index: int) -> dict[str, Any]: worker = next((entry for entry in workers if entry["id"] == worker_id), None) @@ -735,26 +778,32 @@ async def _run_image_list_job(state: ImageListJob): 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])} + body = encode_request_payload({"token": token, **item_payload}) + handle = {"prompt_id": None, "token": token, "index": index, "worker": worker, + "submission_pending": False} + # Own even a lost upload response; DELETE is idempotent. + owned_uploads[token] = handle input_url = build_worker_url(worker, "/distributed/image-list/item") - async with session.post(input_url, json={"token": token, **item_payload}, + async with session.post(input_url, data=body, headers={"Content-Type": "application/json"}, 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 + callback_url = build_master_callback_url(worker, config=config) + if worker_id not in prepared_worker_prompts: + worker_prompt = build_worker_prompt(state.original_prompt, region, token, context_id, + index, callback_url) 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 fetch_worker_path_separator, convert_paths_for_platform + separator = await fetch_worker_path_separator(worker) + if separator: + worker_prompt = convert_paths_for_platform(worker_prompt, separator) 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) + prepared_worker_prompts[worker_id] = worker_prompt + # Reuse media rewrites, never assignment identity or a mutable submitted graph. + worker_prompt = build_worker_prompt(prepared_worker_prompts[worker_id], region, token, + context_id, index, callback_url) + handle["submission_pending"] = True prompt_id = await dispatch_worker_prompt( worker, worker_prompt, state.workflow_meta, client_id=f"image-list-{context_id[:20]}", use_websocket=use_websocket, @@ -762,18 +811,13 @@ async def _run_image_list_job(state: ImageListJob): ) 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} + handle.update(prompt_id=prompt_id, last_seen=time.monotonic(), submission_pending=False) + return handle 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']}") @@ -784,6 +828,16 @@ async def _run_image_list_job(state: ImageListJob): data = await response.json() entry = data.get(handle["prompt_id"]) if not isinstance(entry, dict): + async with session.get(build_worker_url(handle["worker"], "/queue"), + timeout=aiohttp.ClientTimeout(total=3)) as response: + response.raise_for_status() + queue = await response.json() + if isinstance(queue, dict) and any( + isinstance(row, (list, tuple)) and len(row) > 1 and row[1] == handle["prompt_id"] + for key in ("queue_running", "queue_pending") + for row in (queue.get(key) if isinstance(queue.get(key), list) else []) + ): + handle["last_seen"] = time.monotonic() return None status = entry.get("status", {}) or {} if status.get("status_str") == "error": @@ -809,29 +863,40 @@ async def _run_image_list_job(state: ImageListJob): # successful history before releasing the slot / next region. with _LOCK: state.results.pop(completed[1], None) + owned_uploads.pop(inflight[completed[0]]["token"], None) 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 expired: + raise TimeoutError(f"Image-list owned-prompt liveness timeout after {timeout_seconds:.1f}s: {', '.join(expired)}") 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}") + prompt_id = handle.get("prompt_id") + if prompt_id: + # Try both independently: failure of queue deletion must not suppress + # the targeted interrupt. No global interrupt or queue clearing. + for route, body in (("/queue", {"delete": [prompt_id]}), + ("/interrupt", {"prompt_id": prompt_id})): + try: + async with session.post(build_worker_url(worker, route), json=body, + timeout=aiohttp.ClientTimeout(total=5)) as response: + response.raise_for_status() + except Exception as exc: + log(f"[Distributed Image List] Owned prompt cleanup unconfirmed for worker {worker_id} ({route}): {exc}") + elif handle.get("submission_pending"): + log(f"[Distributed Image List] Worker {worker_id} submission status is ambiguous; " + "no prompt ID to target. Check worker manually; not retrying.") 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 + timeout=aiohttp.ClientTimeout(total=5)) as response: + response.raise_for_status() + except Exception as exc: + log(f"[Distributed Image List] Staged input cleanup unconfirmed for worker {worker_id}: {exc}") + finally: + owned_uploads.pop(handle["token"], None) def interrupted() -> bool: try: @@ -840,6 +905,17 @@ async def _run_image_list_job(state: ImageListJob): 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) + try: + results = await dynamic_schedule(len(state.items), [worker["id"] for worker in workers], + dispatch, wait_next, cleanup, interrupted) + return _collect_output_lists(results) + finally: + # Upload ownership exists before dynamic_schedule has an acknowledged + # handle. Cover rejection/lost upload responses as well as cancellation. + pending = list(owned_uploads.values()) + outcomes = await asyncio.gather(*( + asyncio.wait_for(cleanup(handle["worker"]["id"], handle), CLEANUP_TIMEOUT_SECONDS) + for handle in pending), return_exceptions=True) + for handle, outcome in zip(pending, outcomes): + if isinstance(outcome, BaseException): + log(f"[Distributed Image List] Cleanup unconfirmed for worker {handle['worker']['id']}: {outcome}") diff --git a/nodes/image_list.py b/nodes/image_list.py index b5be4b2..2421dda 100644 --- a/nodes/image_list.py +++ b/nodes/image_list.py @@ -12,6 +12,7 @@ from ..image_list.runtime import ( normalize_worker_result, run_image_list_job, encode_result, + encode_request_payload, unwrap_single, ) @@ -96,10 +97,13 @@ class DistributedImageListStart(io.ComfyNode): region = next((row for row in regions if row["start_id"] == start_id), None) if region is None: raise RuntimeError(f"No safe paired Image List End was found for Start node {start_id}") + from comfy_execution.utils import get_executing_context + executing = get_executing_context() 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, ) return io.NodeOutput( [item["image"] for item in items], @@ -196,9 +200,8 @@ class DistributedImageListEnd(io.ComfyNode): 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)) + # The runtime owns the configurable deadline and bounded remote cleanup. + result_lists = run_async_in_server_loop(run_image_list_job(context_id)) except ImageListCancelledError: import comfy.model_management raise comfy.model_management.InterruptProcessingException() @@ -218,7 +221,8 @@ async def submit_worker_result(callback_url, item_token, context_id, item_index, "item_index": item_index, "result": result} async with session.post( f"{callback_url.rstrip('/')}/distributed/image-list/result", - json=payload, + data=encode_request_payload(payload), + headers={"Content-Type": "application/json"}, timeout=aiohttp.ClientTimeout(total=60), ) as response: if response.status >= 400: diff --git a/runtime/bootstrap.py b/runtime/bootstrap.py index 630b021..995a94c 100644 --- a/runtime/bootstrap.py +++ b/runtime/bootstrap.py @@ -24,6 +24,9 @@ def initialize(): ensure_distributed_state(server.PromptServer.instance) ensure_tile_jobs_initialized() + from ..image_list import runtime as image_list_runtime + from ..image_list.lifecycle import install_lifecycle + install_lifecycle(server.PromptServer.instance, image_list_runtime) if not os.environ.get('COMFYUI_IS_WORKER'): atexit.register(sync_cleanup) diff --git a/tests/helpers/image_list_network_check.py b/tests/helpers/image_list_network_check.py index b0fbc32..f231aa0 100644 --- a/tests/helpers/image_list_network_check.py +++ b/tests/helpers/image_list_network_check.py @@ -106,7 +106,8 @@ async def setup(): module = sys.modules[str(ROOT).replace('.', '_x_')] nodes.NODE_CLASS_MAPPINGS.update(TestImageSource=TestImageSource, TestImageOperation=TestImageOperation, TestCapture=TestCapture) - app = web.Application() + # 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. app.add_routes([route for route in ps.routes if route.path.startswith('/distributed/image-list/')]) return ps, assets, module, app @@ -212,6 +213,7 @@ async def coordinator(): config_mod = importlib.import_module(module.__name__ + '.utils.config') orchestration = importlib.import_module(module.__name__ + '.api.queue_orchestration') runner, port = await listen(app) + ps.port = port # Match the listener used by local-worker callback resolution. processes, stderr_files = [], [] # Old worker capability refusal is read-only and precedes input evaluation. legacy_app = web.Application() diff --git a/tests/helpers/v3_runtime_check.py b/tests/helpers/v3_runtime_check.py index d42eade..a4d7160 100644 --- a/tests/helpers/v3_runtime_check.py +++ b/tests/helpers/v3_runtime_check.py @@ -214,7 +214,12 @@ async def check_image_list_lazy_executor(mapping, module, prompt_server, asset_m runtime = sys.modules[module.__name__ + '.image_list.runtime'] 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): + 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 + assert len(prompt_server.app.cleanup_ctx) >= 1, 'lifecycle cleanup was not registered' 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], @@ -227,6 +232,7 @@ async def check_image_list_lazy_executor(mapping, module, prompt_server, asset_m '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', + 'image_list_workflow_meta': workflow_meta, }}, '3': {'class_type': 'ImageListTestProbe', 'inputs': {'image': ['2', 0]}}, '4': {'class_type': 'DistributedImageListEnd', 'inputs': { @@ -258,6 +264,21 @@ async def check_image_list_lazy_executor(mapping, module, prompt_server, asset_m import folder_paths scratch = Path(os.environ.get('TMPDIR', Path.home() / '.hermes/cache/scratch')) try: + # Resolve genuine Comfy output/temp annotations against the pinned core, + # not only the controlled folder_paths seam in unit tests. + media = sys.modules[module.__name__ + '.api.orchestration.media_sync'] + with tempfile.TemporaryDirectory(prefix='media-roots-', dir=scratch) as media_root: + for kind in ('input', 'output', 'temp'): + source = Path(media_root) / kind / 'shared/ref.png' + source.parent.mkdir(parents=True) + source.write_bytes(kind.encode()) + with patch.object(folder_paths, 'input_directory', str(Path(media_root) / 'input')), \ + patch.object(folder_paths, 'output_directory', str(Path(media_root) / 'output')), \ + patch.object(folder_paths, 'temp_directory', str(Path(media_root) / 'temp')): + for kind in ('output', 'temp'): + refs = media._find_media_references({'1': {'inputs': {'image': f'shared/ref.png [{kind}]'}}}) + assert len(refs) == 1 + assert media._load_media_file_sync(refs[0])[0] == kind.encode() 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): diff --git a/tests/test_image_list_lifecycle.py b/tests/test_image_list_lifecycle.py new file mode 100644 index 0000000..998a617 --- /dev/null +++ b/tests/test_image_list_lifecycle.py @@ -0,0 +1,177 @@ +"""Tests for periodic image-list server lifecycle cleanup.""" +import asyncio +import time +import unittest +from types import SimpleNamespace +from typing import Any, cast +from unittest.mock import patch + +from aiohttp import web + +from image_list import lifecycle, runtime + + +class ImageListLifecycleTests(unittest.IsolatedAsyncioTestCase): + def setUp(self): + with runtime._LOCK: + runtime._MASTER_JOBS.clear() + runtime._WORKER_ITEMS.clear() + + def tearDown(self): + with runtime._LOCK: + runtime._MASTER_JOBS.clear() + runtime._WORKER_ITEMS.clear() + + @staticmethod + def _state(prompt_id=None, *, active=False, age=10_000.0): + return SimpleNamespace( + prompt_id=prompt_id, + active=active, + created_at=time.monotonic() - age, + ) + + @staticmethod + def _server(queue): + return SimpleNamespace(app=web.Application(), prompt_queue=queue) + + async def _start(self, server, *, interval=0.01, grace=0.01): + lifecycle.install_lifecycle( + server, + runtime, + interval_seconds=interval, + orphan_grace_seconds=grace, + ) + runner = web.AppRunner(server.app) + await runner.setup() + return runner + + async def test_unattended_sweep_removes_orphan_and_expired_staging(self): + context_id = "orphan" + with runtime._LOCK: + cast(Any, runtime._MASTER_JOBS)[context_id] = self._state() + runtime._WORKER_ITEMS["expired-token"] = (time.monotonic() - 1, {"image": "stale"}) + + queue = SimpleNamespace( + get_current_queue_volatile=lambda: ([], []), + ) + server = self._server(queue) + swept = asyncio.Event() + sweep_worker_items = runtime._sweep_worker_items + + def track_worker_sweep(): + swept.set() + sweep_worker_items() + + with patch.object(runtime, "_sweep_worker_items", side_effect=track_worker_sweep): + runner = await self._start(server) + try: + await asyncio.wait_for(swept.wait(), timeout=0.5) + self.assertNotIn(context_id, runtime._MASTER_JOBS) + self.assertNotIn("expired-token", runtime._WORKER_ITEMS) + with self.assertRaisesRegex(ValueError, "Unknown or expired"): + runtime.set_master_result(context_id, 0, "late-token", {}) + finally: + await runner.cleanup() + + async def test_old_active_and_queued_contexts_survive_but_orphans_expire(self): + states = { + "active": self._state("active-prompt", active=True), + "queued": self._state("queued-prompt"), + "orphan": self._state("missing-prompt"), + "unbound": self._state(None), + } + with runtime._LOCK: + cast(Any, runtime._MASTER_JOBS).update(states) + queue = SimpleNamespace( + get_current_queue_volatile=lambda: ( + [[1, "active-prompt", {}]], + [(2, "queued-prompt", {}, {}, None, None)], + ), + ) + runner = await self._start(self._server(queue)) + try: + await asyncio.sleep(0.04) + self.assertIs(runtime._MASTER_JOBS.get("active"), states["active"]) + self.assertIs(runtime._MASTER_JOBS.get("queued"), states["queued"]) + self.assertNotIn("orphan", runtime._MASTER_JOBS) + self.assertNotIn("unbound", runtime._MASTER_JOBS) + finally: + await runner.cleanup() + + async def test_falls_back_to_nonvolatile_queue_reader(self): + state = self._state("queued-prompt") + with runtime._LOCK: + cast(Any, runtime._MASTER_JOBS)["queued"] = state + fallback_reads = 0 + + def volatile_failure(): + raise RuntimeError("volatile view unavailable") + + def current_queue(): + nonlocal fallback_reads + fallback_reads += 1 + return ([], [[1, "queued-prompt", {}]]) + + queue = SimpleNamespace( + get_current_queue_volatile=volatile_failure, + get_current_queue=current_queue, + ) + runner = await self._start(self._server(queue)) + try: + await asyncio.sleep(0.04) + self.assertGreaterEqual(fallback_reads, 1) + self.assertIs(runtime._MASTER_JOBS.get("queued"), state) + finally: + await runner.cleanup() + + async def test_queue_read_error_protects_bound_contexts(self): + with runtime._LOCK: + cast(Any, runtime._MASTER_JOBS)["bound"] = self._state("bound-prompt") + cast(Any, runtime._MASTER_JOBS)["unbound"] = self._state(None) + + def fail_read(): + raise RuntimeError("queue unavailable") + + queue = SimpleNamespace( + get_current_queue_volatile=fail_read, + get_current_queue=fail_read, + ) + runner = await self._start(self._server(queue)) + try: + await asyncio.sleep(0.04) + self.assertIn("bound", runtime._MASTER_JOBS) + self.assertNotIn("unbound", runtime._MASTER_JOBS) + finally: + await runner.cleanup() + + async def test_install_and_removal_are_idempotent_and_task_stops_on_cleanup(self): + queue_reads = 0 + + def read_queue(): + nonlocal queue_reads + queue_reads += 1 + return ([], []) + + server = self._server(SimpleNamespace(get_current_queue_volatile=read_queue)) + lifecycle.install_lifecycle(server, runtime, interval_seconds=0.01, orphan_grace_seconds=0.01) + lifecycle.install_lifecycle(server, runtime, interval_seconds=0.01, orphan_grace_seconds=0.01) + self.assertEqual(len(server.app.cleanup_ctx), 1) + + with runtime._LOCK: + cast(Any, runtime._MASTER_JOBS)["removed"] = self._state() + runtime.remove_master_job("removed") + runtime.remove_master_job("removed") + self.assertNotIn("removed", runtime._MASTER_JOBS) + + runner = web.AppRunner(server.app) + await runner.setup() + await asyncio.sleep(0.035) + self.assertGreaterEqual(queue_reads, 1) + await runner.cleanup() + reads_after_cleanup = queue_reads + await asyncio.sleep(0.035) + self.assertEqual(queue_reads, reads_after_cleanup) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_image_list_runtime_boundaries.py b/tests/test_image_list_runtime_boundaries.py new file mode 100644 index 0000000..70d87be --- /dev/null +++ b/tests/test_image_list_runtime_boundaries.py @@ -0,0 +1,340 @@ +"""Exercise production per-item orchestration through controlled provider seams.""" +import asyncio +import importlib.util +import json +from pathlib import Path +import sys +import types +import unittest +from unittest.mock import AsyncMock, Mock, patch + +import torch + +from test_image_list import graph_one_region + +ROOT = Path(__file__).resolve().parents[1] + + +class Response: + def __init__(self, data=None, status=200): + self.data, self.status = data, status + + 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): + if self.status >= 400: + raise RuntimeError(f'HTTP {self.status}') + + +class RuntimeBoundaryTests(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + name = '_image_list_boundaries' + self.modules = {} + for suffix in ('', '.image_list', '.api', '.api.orchestration', '.utils'): + module = types.ModuleType(name + suffix) + module.__path__ = [] + self.modules[name + suffix] = module + self.patch_modules = patch.dict(sys.modules, self.modules) + self.patch_modules.start() + self.addCleanup(self.patch_modules.stop) + spec = importlib.util.spec_from_file_location(name + '.image_list.runtime', ROOT / 'image_list/runtime.py') + self.rt = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = self.rt + spec.loader.exec_module(self.rt) + self.worker = {'id': 'worker', 'host': 'worker', 'type': 'remote', 'port': 8188} + self.config = {'workers': [self.worker], 'settings': {}} + self.prompts, self.uploads, self.deletions, self.cancel_requests = [], [], [], [] + self.reject = False + self.interrupted = False + self.completed = True + self.queue = {'queue_running': [], 'queue_pending': []} + self.callback = Mock(side_effect=lambda worker, **_kwargs: + 'http://127.0.0.1:8188' if worker['type'] == 'local' else 'http://public-master') + self.sync = AsyncMock(side_effect=self.sync_media) + self.dispatch = AsyncMock(side_effect=self.dispatch_prompt) + self.logs = [] + dependencies = { + '.api.orchestration.dispatch': types.SimpleNamespace( + preflight_image_list_workers=AsyncMock(return_value={'worker': True}), + dispatch_worker_prompt=self.dispatch), + '.api.orchestration.prompt_transform': types.SimpleNamespace( + _resolve_direct_workers_for_image_list=lambda config, ids: config['workers']), + '.api.orchestration.media_sync': types.SimpleNamespace( + sync_worker_media=self.sync, fetch_worker_path_separator=AsyncMock(return_value='/'), + convert_paths_for_platform=lambda prompt, separator: prompt), + '.utils.config': types.SimpleNamespace(load_config=lambda: self.config, + get_worker_timeout_seconds=lambda: 60, + get_image_list_job_timeout_seconds=lambda config=None: 0), + '.utils.network': types.SimpleNamespace(build_master_url=lambda **kwargs: 'http://public-master', + build_master_callback_url=self.callback, + build_worker_url=lambda worker, path='': f"http://{worker['id']}{path}", + get_client_session=AsyncMock(return_value=self)), + '.utils.logging': types.SimpleNamespace(log=self.logs.append), + } + for suffix, module in dependencies.items(): + sys.modules[name + suffix] = module + self.config_module = dependencies['.utils.config'] + # Other legacy test modules replace aiohttp during collection. This + # controlled provider fixture needs only its explicit timeout/error seam. + sys.modules['aiohttp'] = types.SimpleNamespace( + ClientTimeout=lambda **kwargs: kwargs, ClientError=ConnectionError) + 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 + + async def sync_media(self, worker, prompt): + prompt['7']['inputs']['image'] = 'shared/ref.png' + + async def dispatch_prompt(self, worker, prompt, meta, **kwargs): + if self.reject: + raise RuntimeError('prompt rejected') + self.prompts.append((json.loads(json.dumps(prompt)), meta)) + inputs = prompt['5']['inputs'] + token = inputs['worker_result_token'] + item = self.rt.consume_worker_item(token) + self.rt.set_master_result(self.context, inputs['worker_item_index'], token, {'image': item['image']}) + return f'prompt-{len(self.prompts)}' + + def post(self, url, *, json=None, data=None, **kwargs): + if data is not None: + json = __import__('json').loads(data) + if url.endswith('/distributed/image-list/item'): + self.uploads.append(json['token']) + self.rt.store_worker_item(json['token'], json['item']) + else: + self.cancel_requests.append((url, json)) + return Response({}) + + def delete(self, url, **kwargs): + token = url.rsplit('/', 1)[-1] + self.deletions.append(token) + self.rt.delete_worker_item(token) + return Response({}) + + def get(self, url, **kwargs): + if url.endswith('/queue'): + return Response(self.queue) + pid = url.rsplit('/', 1)[-1] + return Response({pid: {'status': {'completed': True, 'status_str': 'success'}}} + if self.completed else {}) + + def create(self, count=2): + prompt = graph_one_region() + prompt['7'] = {'class_type': 'LoadImage', 'inputs': {'image': 'shared/ref.png [output]'}} + prompt['4']['inputs']['reference'] = ['7', 0] + self.metadata = {'nodes': [{'id': 42}], 'extra': {'label': 'generic fixture'}} + self.context = self.rt.create_master_job('3', '5', prompt, + [{'image': torch.zeros((1, 2, 2, 3)), 'index': i} for i in range(count)], + ['worker'], 'http://public-master', self.metadata) + return self.context + + async def test_each_assignment_reuses_media_paths_with_fresh_tokens_and_metadata(self): + context = self.create() + await self.rt.run_image_list_job(context) + self.assertEqual(self.sync.await_count, 1) + self.assertEqual([p['7']['inputs']['image'] for p, _ in self.prompts], + ['shared/ref.png', 'shared/ref.png']) + tokens = [p['5']['inputs']['worker_result_token'] for p, _ in self.prompts] + self.assertEqual(len(set(tokens)), 2) + self.assertEqual([p['5']['inputs']['worker_item_index'] for p, _ in self.prompts], [0, 1]) + self.assertEqual([meta for _, meta in self.prompts], [self.metadata, self.metadata]) + + async def test_annotated_media_roots_survive_windows_prep_and_repeated_items(self): + import tempfile + + context = self.create(2) + state = self.rt.get_master_job(context) + state.original_prompt['8'] = {'class_type': 'LoadImage', 'inputs': {'image': 'shared\\ref.png [temp]'}} + state.original_prompt['4']['inputs']['second_reference'] = ['8', 0] + name = '_image_list_boundaries.api.orchestration.media_sync' + spec = importlib.util.spec_from_file_location(name, ROOT / 'api/orchestration/media_sync.py') + assert spec is not None and spec.loader is not None + media = importlib.util.module_from_spec(spec) + sys.modules['_image_list_boundaries.utils.logging'].debug_log = self.logs.append + sys.modules['_image_list_boundaries.utils.trace_logger'] = types.SimpleNamespace( + trace_debug=lambda *args: None, trace_info=lambda *args: None) + spec.loader.exec_module(media) + media.fetch_worker_path_separator = AsyncMock(return_value='\\') + converted = Mock(wraps=media.convert_paths_for_platform) + media.convert_paths_for_platform = converted + uploads = [] + + async def upload(worker, filename, content, digest, mime, **kwargs): + uploads.append((filename, content)) + return True, f'worker-input/{len(uploads)}.png' + + media._upload_media_to_worker = AsyncMock(side_effect=upload) + media.sync_worker_media = AsyncMock(wraps=media.sync_worker_media) + with tempfile.TemporaryDirectory(prefix='image-list-media-') as directory: + root = Path(directory) + for kind in ('input', 'output', 'temp'): + path = root / kind / 'shared/ref.png' + path.parent.mkdir(parents=True) + path.write_bytes(kind.encode()) + + def resolve_annotated(filename): + for kind in ('output', 'temp', 'input'): + suffix = f' [{kind}]' + if filename.endswith(suffix): + return str(root / kind / filename[:-len(suffix)]) + return str(root / 'input' / filename) + + with patch.dict(sys.modules, {name: media, 'folder_paths': types.SimpleNamespace( + get_annotated_filepath=resolve_annotated)}): + await self.rt.run_image_list_job(context) + self.assertEqual(uploads, [('shared/ref.png [output]', b'output'), + ('shared/ref.png [temp]', b'temp')]) + media.sync_worker_media.assert_awaited_once() + media.fetch_worker_path_separator.assert_awaited_once() + converted.assert_called_once() + self.assertEqual(converted.call_args.args[1], '\\') + for prompt, _ in self.prompts: + self.assertEqual(prompt['7']['inputs']['image'], 'worker-input/1.png') + self.assertEqual(prompt['8']['inputs']['image'], 'worker-input/2.png') + self.assertEqual(state.original_prompt['7']['inputs']['image'], 'shared/ref.png [output]') + + async def test_local_worker_callback_does_not_use_public_master_address(self): + self.worker['type'] = 'local' + await self.rt.run_image_list_job(self.create(1)) + self.assertEqual(self.prompts[0][0]['5']['inputs']['master_callback_url'], 'http://127.0.0.1:8188') + self.callback.assert_called() + + async def test_remote_worker_uses_topology_callback_helper(self): + await self.rt.run_image_list_job(self.create(1)) + self.assertEqual(self.prompts[0][0]['5']['inputs']['master_callback_url'], 'http://public-master') + self.callback.assert_called() + + def test_metadata_is_forwarded_into_master_start_prompt(self): + from test_prompt_transform import pt + prompt = graph_one_region() + metadata = {'nodes': [{'id': 1}]} + index = pt.PromptIndex(prompt) + transformed = pt.apply_participant_overrides( + prompt, 'master', ['worker'], {'3': 'run'}, 'http://public-master', False, + index, workflow_meta=metadata) + self.assertEqual(transformed['3']['inputs']['image_list_workflow_meta'], metadata) + + async def test_aggregate_input_limit_fails_before_upload(self): + self.rt.MAX_REQUEST_BYTES = 64 + with self.assertRaisesRegex(ValueError, 'complete.*request'): + await self.rt.run_image_list_job(self.create(1)) + self.assertFalse(self.uploads) + self.assertFalse(self.rt._WORKER_ITEMS) + + def test_tensor_and_combined_request_limits_preserve_dtype(self): + item = {'image': torch.zeros((1, 4, 4, 3)), 'index': 0} + encoded = self.rt.encode_item(item) + payload = {'context_id': 'example', 'token': 't' * 48, 'item': encoded} + single = self.rt.encode_request_payload(payload) + self.assertEqual(json.loads(single)['item']['image']['dtype'], 'float32') + self.rt.MAX_REQUEST_BYTES = len(single) + self.assertEqual(self.rt.encode_request_payload(payload), single) + item['reference_image'] = item['image'] + item['mask'] = torch.zeros((1, 4, 4)) + payload['item'] = self.rt.encode_item(item) + with self.assertRaisesRegex(ValueError, 'complete.*request'): + self.rt.encode_request_payload(payload) + # Size is checked before allocating/copying a 4K tensor onto the CPU. + with self.assertRaisesRegex(ValueError, 'transport limit'): + self.rt.encode_value(torch.empty((1, 2160, 3840, 3), device='meta'), 'IMAGE') + + async def test_rejected_submission_reclaims_uploaded_token_without_prompt_id(self): + self.reject = True + with self.assertRaisesRegex(RuntimeError, 'prompt rejected'): + await self.rt.run_image_list_job(self.create(1)) + self.assertEqual(self.deletions, self.uploads) + self.assertFalse(self.rt._WORKER_ITEMS) + self.assertFalse(self.rt._MASTER_JOBS) + self.assertFalse(self.cancel_requests) + + async def test_user_interrupt_during_pending_submission_is_bounded(self): + entered = asyncio.Event() + async def never_ack(*args, **kwargs): + entered.set() + await asyncio.Event().wait() + self.dispatch.side_effect = never_ack + self.rt.SUBMISSION_CANCEL_GRACE_SECONDS = 0.02 + task = asyncio.create_task(self.rt.run_image_list_job(self.create(2))) + await asyncio.wait_for(entered.wait(), 1) + self.interrupted = True + with self.assertRaises(self.rt.ImageListCancelledError): + await asyncio.wait_for(task, 0.8) + self.assertEqual(len(self.uploads), 1) + self.assertEqual(self.deletions, self.uploads) + self.assertFalse(self.rt._MASTER_JOBS) + self.assertTrue(any('ambiguous' in line for line in self.logs), self.logs) + + async def test_owned_running_prompt_survives_cold_load_and_old_hour_cap(self): + self.completed = False + context = self.create(2) + clock = [0.0] + calls = [0] + original_get = self.get + def get(url, **kwargs): + if '/history/' in url: + calls[0] += 1 + clock[0] += 3700 + self.completed = calls[0] >= 3 + if url.endswith('/queue'): + return Response({'queue_running': [[0, 'prompt-1']] if calls[0] == 2 else [], + 'queue_pending': [[0, 'prompt-1']] if calls[0] == 1 else []}) + return original_get(url, **kwargs) + self.rt.time = types.SimpleNamespace(monotonic=lambda: clock[0]) + with patch.object(self, 'get', side_effect=get): + result = await self.rt.run_image_list_job(context) + self.assertEqual(len(result[0]), 2) + self.assertGreaterEqual(clock[0], 7200) + self.assertFalse(self.cancel_requests) + + async def test_unrelated_busy_prompt_does_not_refresh_owned_liveness(self): + self.completed = False + clock = [0.0] + original_get = self.get + def get(url, **kwargs): + if '/history/' in url: + clock[0] += 70 + if url.endswith('/queue'): + return Response({'queue_running': [[0, 'not-our-prompt']], 'queue_pending': []}) + return original_get(url, **kwargs) + self.rt.time = types.SimpleNamespace(monotonic=lambda: clock[0]) + with patch.object(self, 'get', side_effect=get): + with self.assertRaisesRegex(TimeoutError, 'liveness'): + await self.rt.run_image_list_job(self.create(1)) + self.assertEqual(self.cancel_requests, [ + ('http://worker/queue', {'delete': ['prompt-1']}), + ('http://worker/interrupt', {'prompt_id': 'prompt-1'})]) + + async def test_failed_queue_cleanup_still_interrupts_and_reclaims_staging(self): + self.completed = False + self.config_module.get_image_list_job_timeout_seconds = lambda config=None: 0.03 + original_post = self.post + def post(url, **kwargs): + result = original_post(url, **kwargs) + return Response(status=500) if url.endswith('/queue') else result + with patch.object(self, 'post', side_effect=post): + with self.assertRaisesRegex(TimeoutError, 'total deadline'): + await self.rt.run_image_list_job(self.create(2)) + self.assertEqual(self.cancel_requests, [ + ('http://worker/queue', {'delete': ['prompt-1']}), + ('http://worker/interrupt', {'prompt_id': 'prompt-1'})]) + self.assertEqual(self.deletions, self.uploads) + self.assertTrue(any('cleanup unconfirmed' in line for line in self.logs), self.logs) + self.assertEqual(self.dispatch.await_count, 1) + + async def test_configured_total_deadline_cleans_owned_work(self): + self.completed = False + self.config_module.get_image_list_job_timeout_seconds = lambda config=None: 0.03 + self.queue = {'queue_running': [[0, 'prompt-1']], 'queue_pending': []} + with self.assertRaisesRegex(TimeoutError, 'total deadline'): + await asyncio.wait_for(self.rt.run_image_list_job(self.create(2)), 1) + self.assertEqual(len(self.uploads), 1) + self.assertEqual(self.deletions, self.uploads) + self.assertFalse(self.rt._MASTER_JOBS)