fix(image-list): own job lifecycle and preserve worker preparation

This commit is contained in:
Robert Wojciechowski
2026-10-11 21:20:22 +00:00
parent 8e0ac3befe
commit a34969e6f3
12 changed files with 847 additions and 64 deletions
+1 -2
View File
@@ -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):
+2
View File
@@ -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
+1
View File
@@ -413,6 +413,7 @@ async def orchestrate_distributed_execution(
master_url,
delegate_master,
prompt_index,
workflow_meta=workflow_meta,
)
if delegate_master:
+20 -4
View File
@@ -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.
+142
View File
@@ -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)
+129 -53
View File
@@ -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}")
+8 -4
View File
@@ -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:
+3
View File
@@ -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)
+3 -1
View File
@@ -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()
+21
View File
@@ -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):
+177
View File
@@ -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()
+340
View File
@@ -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)