191 lines
7.4 KiB
Python
191 lines
7.4 KiB
Python
"""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 _without_runtime_fingerprints(prompt: dict) -> dict:
|
|
"""Compare submitted structure without ComfyUI's per-node cache annotation.
|
|
|
|
PromptQueue snapshots the graph before IsChangedCache adds ``is_changed``
|
|
to executing nodes. Ignore only that top-level field, never inputs or nested
|
|
metadata, and leave both the queued and executing graphs untouched.
|
|
"""
|
|
return {
|
|
node_id: {key: value for key, value in node.items() if key != "is_changed"}
|
|
if isinstance(node, dict) else node
|
|
for node_id, node in prompt.items()
|
|
}
|
|
|
|
|
|
def has_safe_execution_targets(prompt_server: Any, prompt_id: str | None, prompt: dict) -> bool:
|
|
"""Prove outer roots from the running queue, not a forgeable prompt marker.
|
|
|
|
Plain /prompt remains strict when ComfyUI selected all outputs, including a
|
|
flagged per-item dependency. No saved workflow field grants this exemption.
|
|
"""
|
|
if not prompt_id:
|
|
return False
|
|
from .execution_plan import prepare_image_list_execution
|
|
execution_prompt = _without_runtime_fingerprints(prompt)
|
|
prepared, targets = prepare_image_list_execution(execution_prompt)
|
|
if prepared != execution_prompt or not targets:
|
|
return False
|
|
queue = getattr(prompt_server, "prompt_queue", None)
|
|
for name in ("get_current_queue_volatile", "get_current_queue"):
|
|
reader = getattr(queue, name, None)
|
|
if not callable(reader):
|
|
continue
|
|
try:
|
|
running, _pending = reader()
|
|
for entry in running:
|
|
if (isinstance(entry, (tuple, list)) and len(entry) > 4
|
|
and entry[1] == prompt_id and isinstance(entry[2], dict)
|
|
and _without_runtime_fingerprints(entry[2]) == execution_prompt):
|
|
roots = entry[4]
|
|
return (isinstance(roots, (list, tuple))
|
|
and all(isinstance(root, str) for root in roots)
|
|
and len(roots) == len(targets) and set(roots) == set(targets))
|
|
return False
|
|
except Exception:
|
|
continue
|
|
return False
|
|
|
|
|
|
def _queue_prompt_ids(prompt_server: Any) -> set[str]:
|
|
"""Read running/pending prompt IDs, preferring ComfyUI's volatile view."""
|
|
queue = getattr(prompt_server, "prompt_queue", None)
|
|
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)
|