Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6db5692d78 | ||
|
|
591ffc7396 |
@@ -53,8 +53,9 @@ async def image_list_item_endpoint(request):
|
||||
if not isinstance(context_id, str) or not context_id:
|
||||
raise ValueError("context_id must be a non-empty string")
|
||||
item = data.get("item")
|
||||
store_worker_item(data.get("token"), item)
|
||||
return web.json_response({"status": "stored", "item_index": item["index"]})
|
||||
identity = store_worker_item(data.get("token"), item)
|
||||
return web.json_response({"status": "stored", "item_index": item["index"],
|
||||
"item_identity": identity})
|
||||
except (ValueError, TypeError) as exc:
|
||||
return web.json_response({"error": str(exc)}, status=400)
|
||||
except Exception as exc:
|
||||
|
||||
@@ -86,6 +86,17 @@ End's optional `texts` is a connection-only socket. For workflows saved before t
|
||||
|
||||
File-based dependencies keep ComfyUI's `[output]` / `[temp]` root annotations and are 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.
|
||||
|
||||
## Native cache reuse
|
||||
|
||||
Repeated items can reuse **ComfyUI's native processing cache** on the same participant. Start uses a deterministic identity of the exact image, optional mask, all supplied value slots and source index; tensor dtype, shape and bytes are preserved. Fresh job IDs, assignment tokens and callback addresses live at End, so they no longer invalidate otherwise unchanged upstream processing. End still validates the current assignment and delivers results in source order on every run.
|
||||
|
||||
- Processing settings, seeds, links and node fingerprints still participate in ComfyUI's normal signatures. Changed content invalidates reuse; nodes that declare themselves noncacheable still run.
|
||||
- The master's private executor retains one native LRU output cache between successful slots, capped at **128 entries when retained**. An active graph may temporarily exceed that target; oversized generations are not retained. Failed or cancelled local slots do not seed reuse, and `--cache-none` disables the local output cache. Direct workers use their own configured native caches.
|
||||
- Cache hits are best-effort, not guaranteed: eviction, restarts, changed ordering/index, changed settings and assignment to a different worker can cause misses. There is no worker affinity, disk cache or separate completed-result store.
|
||||
- Reuse retains tensor/model references and can increase RAM/VRAM use; the entry cap is **not a byte limit**. Exact identity calculation also adds CPU hashing and, for GPU-resident inputs, a CPU copy.
|
||||
- Remote items are still uploaded. Each upload returns a worker-computed identity; if a compatible older worker does not return a valid identity, that assignment uses the legacy graph without stable-input reuse. A worker downgrade between items also falls back safely. Assignment expiry and current-item validation apply even when processing hits the cache.
|
||||
- Keep side-effect nodes outside the region when they must run on every submission. A cached processing node follows ordinary ComfyUI semantics and does not repeat its save/preview side effects; End's delivery and downstream saving still run.
|
||||
|
||||
## Scope and failure behavior
|
||||
|
||||
- This version supports the master and direct workers, not Runpod serverless pools. Unsupported pools are excluded without contacting the provider for this job.
|
||||
|
||||
@@ -24,6 +24,34 @@ from .runtime import ImageListCancelledError
|
||||
|
||||
|
||||
_EXECUTION_LOCK = threading.Lock()
|
||||
_CACHE_LOCK = threading.Lock()
|
||||
# One process-local native output cache, not a second completed-result cache.
|
||||
# Core's LRU can temporarily exceed its target for the active graph; oversized
|
||||
# generations are discarded rather than retained after an item/run finishes.
|
||||
NATIVE_CACHE_MAX_ENTRIES = 128
|
||||
_RETAINED_OUTPUTS = None
|
||||
|
||||
|
||||
def _take_outputs(execution):
|
||||
global _RETAINED_OUTPUTS
|
||||
with _CACHE_LOCK:
|
||||
outputs, _RETAINED_OUTPUTS = _RETAINED_OUTPUTS, None
|
||||
if getattr(execution.args, "cache_none", False):
|
||||
return execution.NullCache()
|
||||
return outputs if outputs is not None else execution.LRUCache(
|
||||
execution.CacheKeySetInputSignature, max_size=NATIVE_CACHE_MAX_ENTRIES)
|
||||
|
||||
|
||||
def _retain_outputs(outputs):
|
||||
global _RETAINED_OUTPUTS
|
||||
# Only the bounded native LRU is eligible. NullCache and test doubles are not.
|
||||
if getattr(outputs, "max_size", None) != NATIVE_CACHE_MAX_ENTRIES:
|
||||
return
|
||||
outputs.clean_unused()
|
||||
if len(outputs.cache) > NATIVE_CACHE_MAX_ENTRIES:
|
||||
return
|
||||
with _CACHE_LOCK:
|
||||
_RETAINED_OUTPUTS = outputs
|
||||
|
||||
|
||||
class LocalExecutionSlot:
|
||||
@@ -32,6 +60,7 @@ class LocalExecutionSlot:
|
||||
self._stop = threading.Event()
|
||||
self._future = None
|
||||
self._caches = None
|
||||
self._failed = False
|
||||
self._model_callback = model_callback
|
||||
self._closed = False
|
||||
self._loop = None
|
||||
@@ -55,6 +84,10 @@ class LocalExecutionSlot:
|
||||
self._loop.call_soon_threadsafe(self._task.cancel)
|
||||
|
||||
async def close(self):
|
||||
if self._closed:
|
||||
return
|
||||
# A stop requested before close indicates cancellation, not normal drain.
|
||||
cancelled = self._stop.is_set()
|
||||
self.request_stop()
|
||||
try:
|
||||
if self._future is not None:
|
||||
@@ -72,6 +105,8 @@ class LocalExecutionSlot:
|
||||
waiter.exception()
|
||||
finally:
|
||||
self._pool.shutdown(wait=True, cancel_futures=True)
|
||||
if not cancelled and not self._failed and self._caches is not None:
|
||||
_retain_outputs(self._caches.outputs)
|
||||
self._caches = None
|
||||
self._closed = True
|
||||
|
||||
@@ -84,7 +119,11 @@ class LocalExecutionSlot:
|
||||
try:
|
||||
asyncio.run(self._execute(prompt, prompt_id, end_id, workflow_meta))
|
||||
except asyncio.CancelledError as exc:
|
||||
self._failed = True
|
||||
raise ImageListCancelledError("Local image-list execution cancelled") from exc
|
||||
except BaseException:
|
||||
self._failed = True
|
||||
raise
|
||||
finally:
|
||||
with self._state_lock:
|
||||
self._loop = self._task = None
|
||||
@@ -125,7 +164,7 @@ class LocalExecutionSlot:
|
||||
try:
|
||||
with torch.inference_mode():
|
||||
if self._caches is None:
|
||||
outputs = execution.HierarchicalCache(execution.CacheKeySetInputSignature)
|
||||
outputs = _take_outputs(execution)
|
||||
objects = execution.HierarchicalCache(execution.CacheKeySetID)
|
||||
self._caches = SimpleNamespace(outputs=outputs, objects=objects, all=[outputs, objects])
|
||||
caches = self._caches
|
||||
|
||||
+134
-15
@@ -7,9 +7,11 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import secrets
|
||||
import struct
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
@@ -242,6 +244,64 @@ def normalize_image_list_inputs(images: Any, mask: Any = None, values: Any = Non
|
||||
]
|
||||
|
||||
|
||||
def is_valid_item_identity(identity: Any) -> bool:
|
||||
return (isinstance(identity, str) and len(identity) == 71
|
||||
and identity.startswith("sha256:")
|
||||
and all(char in "0123456789abcdef" for char in identity[7:]))
|
||||
|
||||
|
||||
def item_content_identity(item: dict[str, Any]) -> str:
|
||||
"""Hash normalized content, not transport bytes or dispatch metadata.
|
||||
|
||||
Length-framed, typed fields distinguish scalars, absent masks, slot names,
|
||||
shapes and exact tensor bits (including bfloat16). CPU copying changes neither
|
||||
dtype nor values; local items deliberately have no HTTP transport-size limit.
|
||||
"""
|
||||
digest = hashlib.sha256(b"distributed-image-list-item-v1\0")
|
||||
|
||||
def field(tag, raw):
|
||||
digest.update(tag.encode("ascii") + b"\0")
|
||||
digest.update(len(raw).to_bytes(8, "big"))
|
||||
digest.update(raw)
|
||||
|
||||
def value(raw, name):
|
||||
if raw is None:
|
||||
field("NONE", b"")
|
||||
return
|
||||
kind = value_kind(raw, name)
|
||||
if kind == "TENSOR":
|
||||
cpu = raw.detach().to("cpu").contiguous()
|
||||
field("DTYPE", str(cpu.dtype).encode("ascii"))
|
||||
field("SHAPE", json.dumps(list(cpu.shape), separators=(",", ":")).encode("ascii"))
|
||||
field("TENSOR", memoryview(cpu.view(_torch().uint8).numpy()).cast("B"))
|
||||
elif kind == "FLOAT":
|
||||
field(kind, struct.pack("!d", raw))
|
||||
elif kind == "BOOLEAN":
|
||||
field(kind, b"1" if raw else b"0")
|
||||
elif kind == "INT":
|
||||
field(kind, str(raw).encode("ascii"))
|
||||
else:
|
||||
field(kind, raw.encode("utf-8", errors="surrogatepass"))
|
||||
|
||||
index = item["index"]
|
||||
if isinstance(index, bool) or not isinstance(index, int) or index < 0:
|
||||
raise ValueError("item.index must be a non-negative integer")
|
||||
field("INDEX", str(index).encode("ascii"))
|
||||
field("IMAGE", b"")
|
||||
value(item["image"], "image")
|
||||
field("MASK", b"")
|
||||
value(item.get("mask"), "mask")
|
||||
values = item.get("values") or {}
|
||||
if not isinstance(values, dict) or any(name not in VALUE_NAMES for name in values):
|
||||
raise ValueError("item.values must be a mapping of value_N slots")
|
||||
# Canonical socket order, not the producer's mapping insertion order.
|
||||
for name in VALUE_NAMES:
|
||||
if name in values:
|
||||
field("SLOT", name.encode("ascii"))
|
||||
value(values[name], name)
|
||||
return "sha256:" + digest.hexdigest()
|
||||
|
||||
|
||||
def _tensor_dtype_name(tensor: Any) -> str:
|
||||
name = str(tensor.dtype).removeprefix("torch.")
|
||||
if name not in _ALLOWED_FLOAT_DTYPES:
|
||||
@@ -391,26 +451,64 @@ def _sweep_worker_items() -> None:
|
||||
del _WORKER_ITEMS[token]
|
||||
|
||||
|
||||
def store_worker_item(token: str, item: dict[str, Any]) -> None:
|
||||
def store_worker_item(token: str, item: dict[str, Any]) -> str:
|
||||
if not isinstance(token, str) or len(token) < 32 or len(token) > 128:
|
||||
raise ValueError("token must be a 32-128 character capability")
|
||||
clean = decode_item(item)
|
||||
# Never trust a master-supplied digest; hash the decoded worker-owned item.
|
||||
identity = item_content_identity(clean)
|
||||
_sweep_worker_items()
|
||||
with _LOCK:
|
||||
if token in _WORKER_ITEMS:
|
||||
raise ValueError("duplicate item token")
|
||||
_WORKER_ITEMS[token] = (time.monotonic() + _WORKER_ITEM_TTL_SECONDS, clean)
|
||||
return identity
|
||||
|
||||
|
||||
def consume_worker_item(token: str) -> dict[str, Any]:
|
||||
def _pop_worker_item(token: str) -> tuple[float, dict[str, Any]]:
|
||||
if not isinstance(token, str) or not token:
|
||||
raise ValueError("worker item token is required")
|
||||
_sweep_worker_items()
|
||||
with _LOCK:
|
||||
entry = _WORKER_ITEMS.pop(token, None)
|
||||
if entry is None:
|
||||
if entry is None or entry[0] <= time.monotonic():
|
||||
raise RuntimeError("Image-list item is missing, expired, or already consumed")
|
||||
return entry[1]
|
||||
return entry
|
||||
|
||||
|
||||
def _check_worker_item_identity(item: dict[str, Any], identity: str, index: int) -> None:
|
||||
if not is_valid_item_identity(identity):
|
||||
raise ValueError("Invalid worker item identity")
|
||||
actual = item_content_identity(item) if "image" in item else item.get("item_identity")
|
||||
if actual != identity:
|
||||
raise ValueError("Worker item identity does not match the staged item")
|
||||
if isinstance(index, bool) or not isinstance(index, int) or item["index"] != index:
|
||||
raise ValueError("Worker item index does not match the staged item")
|
||||
|
||||
|
||||
def consume_worker_item(token: str, *, item_identity: str | None = None,
|
||||
item_index: int | None = None) -> dict[str, Any]:
|
||||
with _LOCK:
|
||||
expires, item = _pop_worker_item(token)
|
||||
if "image" not in item:
|
||||
raise RuntimeError("Image-list item is already consumed")
|
||||
if item_identity is not None:
|
||||
_check_worker_item_identity(item, item_identity, item_index)
|
||||
if expires <= time.monotonic():
|
||||
raise RuntimeError("Image-list item is expired")
|
||||
# Start misses release staged tensors, but retain a tiny per-token
|
||||
# receipt until End. This is NOT a completed-result/content cache.
|
||||
# Keeping the original deadline prevents cache hits bypassing TTL.
|
||||
_WORKER_ITEMS[token] = (expires, {"item_identity": item_identity, "index": item["index"]})
|
||||
return item
|
||||
|
||||
|
||||
def finish_worker_item(token: str, item_identity: str, item_index: int) -> None:
|
||||
"""Validate and release this dispatch, whether Start ran or hit native cache."""
|
||||
expires, item = _pop_worker_item(token)
|
||||
_check_worker_item_identity(item, item_identity, item_index)
|
||||
if expires <= time.monotonic():
|
||||
raise RuntimeError("Image-list item is expired")
|
||||
|
||||
|
||||
def delete_worker_item(token: str) -> bool:
|
||||
@@ -638,18 +736,28 @@ def validate_image_list_regions(prompt: dict[str, Any], *, explicit_targets: boo
|
||||
|
||||
|
||||
def build_worker_prompt(prompt: dict[str, Any], region: dict[str, Any], item_token: str,
|
||||
context_id: str, item_index: int, callback_url: str) -> dict[str, Any]:
|
||||
"""Copy only the per-image branch and its worker-side dependencies."""
|
||||
context_id: str, item_index: int, callback_url: str, *,
|
||||
item_identity: str | None = None) -> dict[str, Any]:
|
||||
"""Copy the branch; opt into stable ancestry only after a worker handshake."""
|
||||
start_id, end_id = region["start_id"], region["end_id"]
|
||||
_direct_link(prompt[end_id].get("inputs", {}).get("context"), prompt, end_id, "context")
|
||||
kept = {node_id: json.loads(json.dumps(prompt[node_id])) for node_id in region["node_ids"]}
|
||||
start_inputs = kept[start_id].setdefault("inputs", {})
|
||||
# Per-item data reaches workers only through the staged item token.
|
||||
for name in list(start_inputs):
|
||||
if name in ("images", "mask") or name.startswith(VALUES_INPUT + "."):
|
||||
start_inputs.pop(name)
|
||||
start_inputs.update({"worker_item_token": item_token, "worker_context_id": context_id})
|
||||
end_inputs = kept[end_id].setdefault("inputs", {})
|
||||
if item_identity is not None:
|
||||
if not is_valid_item_identity(item_identity):
|
||||
raise ValueError("Invalid worker item identity")
|
||||
# Native ancestry signatures include EVERY literal, not just a custom
|
||||
# fingerprint. Remove inherited execution IDs, callbacks and all source
|
||||
# inputs wholesale; transient delivery state belongs downstream at End.
|
||||
kept[start_id]["inputs"] = {"worker_item_identity": item_identity}
|
||||
end_inputs.update(worker_context_id=context_id, worker_item_identity=item_identity)
|
||||
else:
|
||||
# Preserve old raw /prompt and older-worker behavior byte-for-byte.
|
||||
start_inputs = kept[start_id].setdefault("inputs", {})
|
||||
for name in list(start_inputs):
|
||||
if name in ("images", "mask") or name.startswith(VALUES_INPUT + "."):
|
||||
start_inputs.pop(name)
|
||||
start_inputs.update({"worker_item_token": item_token, "worker_context_id": context_id})
|
||||
end_inputs.update({"worker_result_token": item_token,
|
||||
"worker_item_index": int(item_index),
|
||||
"master_callback_url": callback_url})
|
||||
@@ -957,7 +1065,8 @@ async def _run_image_list_job(state: ImageListJob, *, include_values: bool = Fal
|
||||
# /prompt queue, or repeated execution of the original source.
|
||||
with _LOCK:
|
||||
_WORKER_ITEMS[token] = (time.monotonic() + _WORKER_ITEM_TTL_SECONDS, state.items[index])
|
||||
prompt = build_worker_prompt(state.original_prompt, region, token, context_id, index, "")
|
||||
prompt = build_worker_prompt(state.original_prompt, region, token, context_id, index, "",
|
||||
item_identity=item_content_identity(state.items[index]))
|
||||
prompt[state.start_id]["inputs"]["image_list_local_execution"] = True
|
||||
prompt[state.end_id]["inputs"]["image_list_local_execution"] = True
|
||||
future = local_executor.submit(prompt, f"{context_id}:{index}", state.end_id, state.workflow_meta)
|
||||
@@ -987,6 +1096,14 @@ async def _run_image_list_job(state: ImageListJob, *, include_values: bool = Fal
|
||||
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}")
|
||||
try:
|
||||
stored = await response.json()
|
||||
except (ValueError, aiohttp.ClientError):
|
||||
stored = None
|
||||
# A per-upload acknowledgement gates the new graph. Older workers may
|
||||
# return no identity; do not change global protocol/preflight contracts.
|
||||
identity = stored.get("item_identity") if isinstance(stored, dict) else None
|
||||
identity_options = {"item_identity": identity} if is_valid_item_identity(identity) else {}
|
||||
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,
|
||||
@@ -999,9 +1116,11 @@ async def _run_image_list_job(state: ImageListJob, *, include_values: bool = Fal
|
||||
from ..api.orchestration.media_sync import sync_worker_media
|
||||
await asyncio.wait_for(sync_worker_media(worker, worker_prompt), timeout=60)
|
||||
prepared_worker_prompts[worker_id] = worker_prompt
|
||||
# Reuse media rewrites, never assignment identity or a mutable submitted graph.
|
||||
# Reuse legacy media rewrites, never assignment identity or a mutable
|
||||
# submitted graph. Keeping this template legacy also permits a worker to
|
||||
# downgrade between items without inheriting a previous stable identity.
|
||||
worker_prompt = build_worker_prompt(prepared_worker_prompts[worker_id], region, token,
|
||||
context_id, index, callback_url)
|
||||
context_id, index, callback_url, **identity_options)
|
||||
handle["submission_pending"] = True
|
||||
prompt_id = await dispatch_worker_prompt(
|
||||
worker, worker_prompt, state.workflow_meta,
|
||||
|
||||
+81
-7
@@ -18,9 +18,40 @@ from ..image_list.runtime import (
|
||||
encode_result,
|
||||
encode_request_payload,
|
||||
unwrap_single,
|
||||
is_valid_item_identity,
|
||||
)
|
||||
|
||||
|
||||
def _stable_worker_binding(prompt, start_id, identity):
|
||||
"""Resolve transient state only from the exact Start -> End context wire."""
|
||||
if not is_valid_item_identity(identity):
|
||||
raise ValueError("Invalid worker item identity")
|
||||
start = prompt.get(start_id) if isinstance(prompt, dict) else None
|
||||
if (not isinstance(start, dict) or start.get("class_type") != "DistributedImageListStart"
|
||||
or start.get("inputs", {}).get("worker_item_identity") != identity):
|
||||
raise RuntimeError("Worker item identity is missing its paired Start")
|
||||
candidates = []
|
||||
for end_id, node in prompt.items():
|
||||
if not isinstance(node, dict) or node.get("class_type") != "DistributedImageListEnd":
|
||||
continue
|
||||
inputs = node.get("inputs", {})
|
||||
link = inputs.get("context")
|
||||
if isinstance(link, (list, tuple)) and list(link) == [start_id, 3]:
|
||||
candidates.append((end_id, inputs))
|
||||
if len(candidates) != 1:
|
||||
raise RuntimeError("Worker Start must have exactly one paired End context wire")
|
||||
end_id, inputs = candidates[0]
|
||||
if inputs.get("worker_item_identity") != identity:
|
||||
raise ValueError("Paired End worker item identity does not match Start")
|
||||
for name in ("worker_result_token", "worker_context_id"):
|
||||
if not isinstance(inputs.get(name), str) or not inputs[name]:
|
||||
raise RuntimeError(f"Worker paired End is missing {name}")
|
||||
index = inputs.get("worker_item_index")
|
||||
if isinstance(index, bool) or not isinstance(index, int) or index < 0:
|
||||
raise ValueError("Worker paired End item index is invalid")
|
||||
return end_id, inputs
|
||||
|
||||
|
||||
# Types the transport can carry exactly. Anything else (models, latents, ...)
|
||||
# must be loaded inside the region or shared as a model dependency instead.
|
||||
_VALUE_TYPES = [io.Image, io.Mask, io.String, io.Int, io.Float, io.Boolean]
|
||||
@@ -70,10 +101,26 @@ class DistributedImageListStart(io.ComfyNode):
|
||||
worker_item_token="", worker_context_id="", image_list_enabled_worker_ids="[]",
|
||||
image_list_master_callback="", image_list_execution_id="",
|
||||
image_list_workflow_meta=None, image_list_master_participates=None,
|
||||
image_list_local_execution=False):
|
||||
image_list_local_execution=False, worker_item_identity=None):
|
||||
worker_item_token = unwrap_single(worker_item_token)
|
||||
worker_context_id = unwrap_single(worker_context_id)
|
||||
worker_item_identity = unwrap_single(worker_item_identity)
|
||||
image_list_workflow_meta = unwrap_single(image_list_workflow_meta)
|
||||
if worker_item_identity is not None:
|
||||
prompt = getattr(cls.hidden, "prompt", None)
|
||||
start_id = str(getattr(cls.hidden, "unique_id", ""))
|
||||
_, delivery = _stable_worker_binding(prompt, start_id, worker_item_identity)
|
||||
# Only cache misses execute Start. Never put this token/context in
|
||||
# Start's literal inputs or returned context: both would poison reuse.
|
||||
from ..image_list.runtime import consume_worker_item
|
||||
item = consume_worker_item(delivery["worker_result_token"],
|
||||
item_identity=worker_item_identity,
|
||||
item_index=delivery["worker_item_index"])
|
||||
item_values = item.get("values") or {}
|
||||
return io.NodeOutput(
|
||||
[item["image"]], [item.get("mask")], [item["index"]], [worker_item_identity],
|
||||
*[[item_values.get(name)] for name in VALUE_NAMES],
|
||||
)
|
||||
if worker_item_token:
|
||||
if unwrap_single(image_list_local_execution):
|
||||
from ..image_list.runtime import consume_worker_item
|
||||
@@ -203,18 +250,45 @@ class DistributedImageListEnd(io.ComfyNode):
|
||||
@classmethod
|
||||
def execute(cls, images=None, masks=None, texts=None, context=None,
|
||||
worker_result_token="", worker_item_index=-1, master_callback_url="",
|
||||
image_list_local_execution=False, **values):
|
||||
image_list_local_execution=False, worker_context_id="",
|
||||
worker_item_identity=None, **values):
|
||||
worker_result_token = unwrap_single(worker_result_token)
|
||||
master_callback_url = unwrap_single(master_callback_url)
|
||||
worker_context_id = unwrap_single(worker_context_id)
|
||||
worker_item_identity = unwrap_single(worker_item_identity)
|
||||
if worker_item_identity is not None and not worker_result_token:
|
||||
raise RuntimeError("Worker item identity is missing its delivery token")
|
||||
if worker_result_token:
|
||||
context_values = [item for item in (context if isinstance(context, (list, tuple)) else [context])
|
||||
if isinstance(item, str) and item]
|
||||
if not context_values or len(set(context_values)) != 1:
|
||||
raise RuntimeError("Worker Image List End is missing its paired context")
|
||||
prompt = getattr(cls.hidden, "prompt", {}) or {}
|
||||
end_inputs = next((node.get("inputs", {}) for node in prompt.values()
|
||||
if isinstance(node, dict) and node.get("class_type") == "DistributedImageListEnd"
|
||||
and node.get("inputs", {}).get("worker_result_token") == worker_result_token), {})
|
||||
context_id = context_values[0]
|
||||
if worker_item_identity is not None:
|
||||
end_id = str(getattr(cls.hidden, "unique_id", ""))
|
||||
node = prompt.get(end_id, {})
|
||||
link = node.get("inputs", {}).get("context")
|
||||
if (node.get("class_type") != "DistributedImageListEnd"
|
||||
or not isinstance(link, (list, tuple)) or len(link) != 2):
|
||||
raise RuntimeError("Worker End is missing its paired Start context wire")
|
||||
paired_id, end_inputs = _stable_worker_binding(prompt, link[0], worker_item_identity)
|
||||
if (paired_id != end_id or context_values != [worker_item_identity]
|
||||
or end_inputs["worker_context_id"] != worker_context_id
|
||||
or end_inputs["worker_result_token"] != worker_result_token
|
||||
or end_inputs["worker_item_index"] != unwrap_single(worker_item_index)
|
||||
or end_inputs.get("master_callback_url", "") != master_callback_url):
|
||||
raise RuntimeError("Worker End identity/context does not match its current delivery")
|
||||
from ..image_list.runtime import finish_worker_item
|
||||
# On a miss this consumes Start's receipt; on a hit it validates
|
||||
# and drops the current staged tensors. Both paths enforce TTL.
|
||||
finish_worker_item(worker_result_token, worker_item_identity,
|
||||
unwrap_single(worker_item_index))
|
||||
context_id = worker_context_id
|
||||
else:
|
||||
end_inputs = next((node.get("inputs", {}) for node in prompt.values()
|
||||
if isinstance(node, dict) and node.get("class_type") == "DistributedImageListEnd"
|
||||
and node.get("inputs", {}).get("worker_result_token") == worker_result_token), {})
|
||||
expected = end_override_names(end_inputs)
|
||||
# Missing connected lazy values must fail instead of becoming an
|
||||
# implicit fallback to the original Start value.
|
||||
@@ -223,13 +297,13 @@ class DistributedImageListEnd(io.ComfyNode):
|
||||
result = normalize_worker_result(images, masks, texts, supplied)
|
||||
if unwrap_single(image_list_local_execution):
|
||||
from ..image_list.runtime import set_master_result
|
||||
set_master_result(context_values[0], int(unwrap_single(worker_item_index)),
|
||||
set_master_result(context_id, int(unwrap_single(worker_item_index)),
|
||||
str(worker_result_token), result)
|
||||
else:
|
||||
payload = encode_result(result)
|
||||
run_async_in_server_loop(
|
||||
submit_worker_result(str(master_callback_url), str(worker_result_token),
|
||||
context_values[0], int(unwrap_single(worker_item_index)), payload), timeout=60,
|
||||
context_id, int(unwrap_single(worker_item_index)), payload), timeout=60,
|
||||
)
|
||||
return io.NodeOutput([result["image"]], [result.get("mask")], [result.get("text")],
|
||||
[result.get("values", {})])
|
||||
|
||||
@@ -0,0 +1,172 @@
|
||||
"""Offline CPU proof against ComfyUI's real signatures and PromptExecutor.
|
||||
|
||||
Exercises one native worker executor and separate private master-slot lifetimes,
|
||||
including their shared stable Start / fresh End contract.
|
||||
"""
|
||||
import asyncio
|
||||
import copy
|
||||
import importlib
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import types
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[2]
|
||||
|
||||
|
||||
def main(core):
|
||||
sys.path.insert(0, str(Path(core).resolve()))
|
||||
import comfy.cli_args
|
||||
comfy.cli_args.args.cpu = True # Must precede all model-management imports.
|
||||
comfy.cli_args.args.disable_assets = True
|
||||
import execution
|
||||
import nodes
|
||||
import torch
|
||||
from comfy_execution.caching import CacheKeySetInputSignature
|
||||
|
||||
name = '_native_cache_identity_probe'
|
||||
for suffix in ('', '.image_list', '.nodes', '.utils'):
|
||||
package = types.ModuleType(name + suffix)
|
||||
package.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
|
||||
sys.modules[name + suffix] = package
|
||||
|
||||
def forbidden(*args, **kwargs):
|
||||
raise AssertionError('offline native proof attempted a server-loop bridge')
|
||||
|
||||
sys.modules[name + '.utils.async_helpers'] = types.SimpleNamespace(run_async_in_server_loop=forbidden)
|
||||
rt = importlib.import_module(name + '.image_list.runtime')
|
||||
boundary = importlib.import_module(name + '.nodes.image_list')
|
||||
calls = []
|
||||
|
||||
class Operation:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'image': ('IMAGE',), 'mask': ('MASK',),
|
||||
'text': ('STRING', {'forceInput': True})}}
|
||||
RETURN_TYPES = ('IMAGE', 'MASK', 'STRING')
|
||||
FUNCTION = 'run'
|
||||
|
||||
def run(self, image, mask, text):
|
||||
calls.append(text)
|
||||
return image + 1, mask + 1, text + ':done'
|
||||
|
||||
nodes.NODE_CLASS_MAPPINGS.update(DistributedImageListStart=boundary.DistributedImageListStart,
|
||||
DistributedImageListEnd=boundary.DistributedImageListEnd,
|
||||
Operation=Operation)
|
||||
prompt = {
|
||||
'start': {'class_type': rt.START_CLASS, 'inputs': {'image_list_execution_id': 'outer-A'}},
|
||||
'op': {'class_type': 'Operation', 'inputs': {
|
||||
'image': ['start', 0], 'mask': ['start', 1], 'text': ['start', 4]}},
|
||||
'end': {'class_type': rt.END_CLASS, 'inputs': {
|
||||
'images': ['op', 0], 'masks': ['op', 1], 'texts': ['op', 2],
|
||||
'context': ['start', 3], 'value_3': ['start', 2]}},
|
||||
}
|
||||
data = rt.normalize_image_list_inputs(torch.zeros(1, 2, 3, 3, dtype=torch.float64),
|
||||
torch.ones(1, 2, 3, dtype=torch.bfloat16),
|
||||
{'value_1': 'original', 'value_2': 3})[0]
|
||||
region = rt.validate_image_list_regions(prompt)[0]
|
||||
|
||||
def prepare(data):
|
||||
token = rt.create_worker_task_token()
|
||||
context = rt.create_master_job('start', 'end', prompt, [data], ['worker'], '')
|
||||
rt.register_assignment(context, data['index'], 'worker', token)
|
||||
identity = rt.store_worker_item(token, rt.encode_item(data))
|
||||
result = rt.build_worker_prompt(prompt, region, token, context, data['index'], '', item_identity=identity)
|
||||
result['start']['inputs']['image_list_local_execution'] = True
|
||||
result['end']['inputs']['image_list_local_execution'] = True
|
||||
return token, context, result
|
||||
|
||||
async def signature(graph, executor):
|
||||
dynamic = execution.DynamicPrompt(graph)
|
||||
changed = execution.IsChangedCache('signature-only', dynamic, executor.caches.outputs)
|
||||
keys = CacheKeySetInputSignature(dynamic, graph.keys(), changed)
|
||||
await keys.add_keys(graph.keys())
|
||||
return {node: keys.get_data_key(node) for node in graph}
|
||||
|
||||
async def probe():
|
||||
server = types.SimpleNamespace(client_id=None, last_node_id=None, send_sync=lambda *a, **kw: None)
|
||||
executor = execution.PromptExecutor(server, cache_type=execution.CacheType.LRU,
|
||||
cache_args={'lru': 64, 'ram': 0, 'ram_inactive': 0})
|
||||
ta, ca, a = prepare(data)
|
||||
prompt['start']['inputs']['image_list_execution_id'] = 'outer-B'
|
||||
tb, cb, b = prepare(data)
|
||||
assert ta != tb and ca != cb
|
||||
ka, kb = await signature(a, executor), await signature(b, executor)
|
||||
assert ka['start'] == kb['start'] and ka['op'] == kb['op']
|
||||
assert ka['end'] != kb['end']
|
||||
assert a['start']['inputs'] == b['start']['inputs']
|
||||
# Exercise actual V3 list mapping, hidden.prompt/unique_id, lazy inputs,
|
||||
# native node signatures and cache hits, not a reimplementation of them.
|
||||
for number, (token, context, graph) in enumerate(((ta, ca, a), (tb, cb, b))):
|
||||
valid = await execution.validate_prompt(f'native-{number}', graph, ['end'])
|
||||
assert valid[0], valid
|
||||
await executor.execute_async(graph, f'native-{number}', {}, ['end'])
|
||||
assert executor.success, executor.status_messages
|
||||
assert token not in rt._WORKER_ITEMS
|
||||
result = rt.get_master_job(context).results[0]
|
||||
assert torch.equal(result['image'], data['image'] + 1)
|
||||
assert torch.equal(result['mask'], data['mask'] + 1)
|
||||
assert result['text'] == 'original:done'
|
||||
assert result['values'] == {'value_3': 0}
|
||||
cached = next(entry['nodes'] for event, entry in executor.status_messages if event == 'execution_cached')
|
||||
if number:
|
||||
assert 'start' in cached and 'op' in cached and 'end' not in cached, cached
|
||||
assert calls == ['original'], calls
|
||||
for failure in ('expired', 'mismatched'):
|
||||
token, context, graph = prepare(data)
|
||||
expires, staged = rt._WORKER_ITEMS[token]
|
||||
if failure == 'expired':
|
||||
rt._WORKER_ITEMS[token] = (0, staged)
|
||||
else:
|
||||
staged['image'] = staged['image'] + 9
|
||||
await executor.execute_async(graph, f'rejected-{failure}', {}, ['end'])
|
||||
assert not executor.success, failure
|
||||
assert not rt.get_master_job(context).received_indexes
|
||||
assert token not in rt._WORKER_ITEMS
|
||||
cached = next(entry['nodes'] for event, entry in executor.status_messages if event == 'execution_cached')
|
||||
assert 'start' in cached and 'op' in cached and 'end' not in cached, cached
|
||||
assert calls == ['original'], calls
|
||||
for part in ('image', 'mask', 'value', 'index'):
|
||||
changed = copy.deepcopy(data)
|
||||
if part in ('image', 'mask'):
|
||||
changed[part].reshape(-1)[0] += 1
|
||||
elif part == 'value':
|
||||
changed['values']['value_2'] = 4 # Even an unconnected slot is content.
|
||||
else:
|
||||
changed['index'] = 1
|
||||
token, context, graph = prepare(changed)
|
||||
keys = await signature(graph, executor)
|
||||
assert keys['start'] != ka['start'] and keys['op'] != ka['op'], part
|
||||
await executor.execute_async(graph, f'changed-{part}', {}, ['end'])
|
||||
assert executor.success, executor.status_messages
|
||||
assert token not in rt._WORKER_ITEMS
|
||||
assert rt.get_master_job(context).received_indexes == {changed['index']}
|
||||
assert len(calls) == 5, calls
|
||||
# Combine the real Start/End identity path with cache transfer between
|
||||
# private master slots, not just independent signature/retention probes.
|
||||
local = importlib.import_module(name + '.image_list.local_executor')
|
||||
before = len(calls)
|
||||
for number in range(2):
|
||||
token, context, graph = prepare(data)
|
||||
slot = local.LocalExecutionSlot()
|
||||
try:
|
||||
await asyncio.wrap_future(slot.submit(graph, f'local-{number}', 'end'))
|
||||
result = rt.get_master_job(context).results[0]
|
||||
assert torch.equal(result['image'], data['image'] + 1)
|
||||
assert result['values'] == {'value_3': 0}
|
||||
assert token not in rt._WORKER_ITEMS
|
||||
finally:
|
||||
await slot.close()
|
||||
rt.remove_master_job(context)
|
||||
assert local._RETAINED_OUTPUTS is not None
|
||||
assert len(calls) == before + 1, calls
|
||||
assert not rt._WORKER_ITEMS
|
||||
assert not torch.cuda.is_initialized()
|
||||
print('NATIVE_CACHE_IDENTITY_OK stable Start/op keys; fresh End key; '
|
||||
'1 operation for 2 same-data jobs; 4 changed-content misses; '
|
||||
'expired/mismatched cached inputs rejected; current delivery and cleanup')
|
||||
|
||||
asyncio.run(probe())
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main(sys.argv[1])
|
||||
@@ -92,7 +92,9 @@ def assert_item_execution(prompt, index):
|
||||
from comfy_execution.utils import get_executing_context
|
||||
context = get_executing_context()
|
||||
assert context is not None
|
||||
assert prompt['2']['inputs'].get('worker_item_token'), 'per-item node ran in outer graph'
|
||||
assert prompt['4']['inputs'].get('worker_result_token'), 'per-item node ran in outer graph'
|
||||
assert (prompt['2']['inputs'].get('worker_item_identity')
|
||||
or prompt['2']['inputs'].get('worker_item_token')), 'per-item source binding missing'
|
||||
running, _ = server.PromptServer.instance.prompt_queue.get_current_queue_volatile()
|
||||
assert len(running) == 1, running
|
||||
entry = running[0]
|
||||
@@ -513,6 +515,8 @@ async def coordinator():
|
||||
config['settings']['master_delegate_only'] = False
|
||||
participation_counts = {}
|
||||
for mode, selected in [('mixed', [w['id'] for w in workers]), ('master-only', [])]:
|
||||
# Keep cold-run exact-once assertions independent of other modes.
|
||||
importlib.import_module(module.__name__ + '.image_list.local_executor')._RETAINED_OUTPUTS = None
|
||||
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
|
||||
before_sources, before_local, before_models = len(SOURCE_CALLS), len(BRANCH_CALLS), len(MODEL_CALLS)
|
||||
before_utilities, before_events = len(UTILITY_CALLS), len(ITEM_EXECUTIONS)
|
||||
@@ -567,9 +571,36 @@ async def coordinator():
|
||||
assert texts[i] == f'item-{i}:seed-{100 + i}', (mode, texts)
|
||||
assert not runtime._MASTER_JOBS and not runtime._WORKER_ITEMS
|
||||
participation_counts[mode] = {'master': len(local_calls), 'workers': len(remote_calls)}
|
||||
# Repeat the completed master-only queue with a fresh outer executor.
|
||||
# Every item must be delivered again without recomputing its branch.
|
||||
before_sources, before_local = len(SOURCE_CALLS), len(BRANCH_CALLS)
|
||||
before_models, before_utilities = len(MODEL_CALLS), len(UTILITY_CALLS)
|
||||
before_captures = len(CAPTURED)
|
||||
previous_outputs = CAPTURED[-1]
|
||||
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
|
||||
with patch.object(config_mod, 'load_config', return_value=config), \
|
||||
patch.object(orchestration, 'load_config', return_value=config), \
|
||||
patch.object(orchestration, 'queue_prompt_payload', enqueue), \
|
||||
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
|
||||
await asyncio.wait_for(orchestration.orchestrate_distributed_execution(
|
||||
prompt, None, 'test', []), 45)
|
||||
assert executor.success, executor.status_messages
|
||||
assert len(SOURCE_CALLS) == before_sources + 1
|
||||
assert len(BRANCH_CALLS) == before_local
|
||||
assert len(MODEL_CALLS) == before_models
|
||||
assert len(UTILITY_CALLS) == before_utilities
|
||||
assert len(CAPTURED) == before_captures + 1
|
||||
for before, after in zip(previous_outputs[:2], CAPTURED[-1][:2]):
|
||||
assert len(before) == len(after) == COUNT
|
||||
assert all(torch.equal(a, b) for a, b in zip(before, after))
|
||||
assert previous_outputs[2] == CAPTURED[-1][2]
|
||||
check_carried_values()
|
||||
assert not runtime._MASTER_JOBS and not runtime._WORKER_ITEMS
|
||||
assert prompt == saved_prompt
|
||||
print('IMAGE_LIST_CACHE_REQUEUE_OK unchanged items reused; fresh ordered delivery and cleanup')
|
||||
print('IMAGE_LIST_MASTER_OK ' + json.dumps(participation_counts))
|
||||
print('IMAGE_LIST_VALUES_OK passthrough all six types; sparse override/add; Get; sequential Get-to-Start; direct/mixed/master-only')
|
||||
assert len(outer_executions) == 4 and all(row['captures'] == 1 for row in outer_executions)
|
||||
assert len(outer_executions) == 5 and all(row['captures'] == 1 for row in outer_executions)
|
||||
print('IMAGE_LIST_TARGETS_OK ' + json.dumps({
|
||||
'modes': ['delegate-only', 'mixed', 'master-only'],
|
||||
'outer_targets': outer_executions[0]['targets'], 'worker_targets': ['4'],
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
"""CPU-only native Comfy execution across private-slot lifetimes."""
|
||||
import asyncio
|
||||
import importlib
|
||||
import os
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import types
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[2]
|
||||
sys.path.insert(0, sys.argv[1])
|
||||
import comfy.cli_args
|
||||
comfy.cli_args.args.cpu = True
|
||||
comfy.cli_args.args.disable_assets = True
|
||||
import execution
|
||||
import nodes
|
||||
|
||||
pkg = '_local_retention_probe'
|
||||
for name, path in [(pkg, ROOT), (pkg + '.image_list', ROOT / 'image_list')]:
|
||||
module = types.ModuleType(name)
|
||||
module.__path__ = [str(path)]
|
||||
sys.modules[name] = module
|
||||
local = importlib.import_module(pkg + '.image_list.local_executor')
|
||||
|
||||
CALLS = []
|
||||
OUTPUTS = []
|
||||
|
||||
class Counted:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'value': ('INT',)}}
|
||||
RETURN_TYPES = ('INT',)
|
||||
FUNCTION = 'run'
|
||||
def run(self, value):
|
||||
if value < 0:
|
||||
raise ValueError('intentional retention failure')
|
||||
CALLS.append(value)
|
||||
return (value,)
|
||||
|
||||
class Uncached(Counted):
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kwargs):
|
||||
return float('nan')
|
||||
|
||||
class Capture:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'value': ('INT',), 'delivery': ('STRING',)}}
|
||||
RETURN_TYPES = ()
|
||||
OUTPUT_NODE = True
|
||||
FUNCTION = 'run'
|
||||
def run(self, value, delivery):
|
||||
OUTPUTS.append((value, delivery))
|
||||
return ()
|
||||
|
||||
nodes.NODE_CLASS_MAPPINGS.update(RetentionCounted=Counted, RetentionUncached=Uncached, RetentionCapture=Capture)
|
||||
|
||||
async def run(slot, value, delivery, uncached=False):
|
||||
prompt = {
|
||||
'source': {'class_type': 'RetentionUncached' if uncached else 'RetentionCounted', 'inputs': {'value': value}},
|
||||
'end': {'class_type': 'RetentionCapture', 'inputs': {'value': ['source', 0], 'delivery': delivery}},
|
||||
}
|
||||
await asyncio.wrap_future(slot.submit(prompt, delivery, 'end'))
|
||||
|
||||
async def main():
|
||||
first = local.LocalExecutionSlot()
|
||||
await run(first, 1, 'first-1')
|
||||
await run(first, 2, 'first-2')
|
||||
await first.close()
|
||||
assert local._RETAINED_OUTPUTS is not None
|
||||
second = local.LocalExecutionSlot()
|
||||
await run(second, 1, 'second-1')
|
||||
await run(second, 2, 'second-2')
|
||||
await run(second, 3, 'changed')
|
||||
await second.close()
|
||||
assert CALLS == [1, 2, 3], CALLS
|
||||
assert OUTPUTS == [(1, 'first-1'), (2, 'first-2'), (1, 'second-1'), (2, 'second-2'), (3, 'changed')], OUTPUTS
|
||||
third = local.LocalExecutionSlot()
|
||||
await run(third, 4, 'uncached-1', True)
|
||||
await run(third, 4, 'uncached-2', True)
|
||||
await third.close()
|
||||
assert CALLS[-2:] == [4, 4]
|
||||
comfy.cli_args.args.cache_none = True
|
||||
disabled = local.LocalExecutionSlot()
|
||||
await run(disabled, 1, 'disabled-1')
|
||||
await run(disabled, 1, 'disabled-2')
|
||||
await disabled.close()
|
||||
assert CALLS[-2:] == [1, 1]
|
||||
assert local._RETAINED_OUTPUTS is None
|
||||
comfy.cli_args.args.cache_none = False
|
||||
failed = local.LocalExecutionSlot()
|
||||
await run(failed, 7, 'before-failure')
|
||||
try:
|
||||
await run(failed, -1, 'failure')
|
||||
except RuntimeError:
|
||||
pass
|
||||
else:
|
||||
raise AssertionError('failure fixture unexpectedly succeeded')
|
||||
await failed.close()
|
||||
assert local._RETAINED_OUTPUTS is None
|
||||
cancelled = local.LocalExecutionSlot()
|
||||
await run(cancelled, 8, 'cancelled')
|
||||
cancelled.request_stop()
|
||||
await cancelled.close()
|
||||
assert local._RETAINED_OUTPUTS is None
|
||||
oversized = execution.LRUCache(execution.CacheKeySetInputSignature, max_size=local.NATIVE_CACHE_MAX_ENTRIES)
|
||||
empty = execution.DynamicPrompt({})
|
||||
await oversized.set_prompt(empty, [], execution.IsChangedCache('bound', empty, oversized))
|
||||
oversized.min_generation = oversized.generation
|
||||
oversized.cache = {n: n for n in range(local.NATIVE_CACHE_MAX_ENTRIES + 1)}
|
||||
local._retain_outputs(oversized)
|
||||
assert local._RETAINED_OUTPUTS is None
|
||||
print('PASS: cross-run native cache reuse, fresh delivery, input change, noncacheable nodes, cache_none, cancellation, retention bound')
|
||||
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,20 @@
|
||||
"""Run the real framework's CPU-only private-cache retention regression."""
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def test_native_local_cache_retention():
|
||||
core = os.environ.get('COMFYUI_SOURCE_ROOT')
|
||||
if not core:
|
||||
pytest.skip('Set COMFYUI_SOURCE_ROOT for CPU-only real-framework validation')
|
||||
helper = Path(__file__).parent / 'helpers' / 'local_cache_retention_check.py'
|
||||
env = dict(os.environ, CUDA_VISIBLE_DEVICES='', HIP_VISIBLE_DEVICES='')
|
||||
env.pop('PYTHONPATH', None)
|
||||
result = subprocess.run([sys.executable, str(helper), core], env=env,
|
||||
text=True, capture_output=True, timeout=90)
|
||||
assert result.returncode == 0, result.stdout + result.stderr
|
||||
assert 'PASS: cross-run native cache reuse' in result.stdout
|
||||
@@ -440,6 +440,8 @@ def real_framework_probe(comfy_root):
|
||||
rt.remove_master_job(context)
|
||||
|
||||
for behavior in ('normal', 'expand', 'error', 'cancel', 'interrupt', 'sync', 'already-interrupted'):
|
||||
# Independent lifecycle cases start cold; retention is tested separately.
|
||||
local._RETAINED_OUTPUTS = None
|
||||
mode[0] = behavior
|
||||
prompt['op']['class_type'] = 'Expanded' if behavior == 'expand' else 'Operation'
|
||||
loads.clear()
|
||||
|
||||
@@ -0,0 +1,527 @@
|
||||
"""Stable content signatures with fresh, expiring delivery capabilities."""
|
||||
import asyncio
|
||||
import copy
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
import types
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from image_list import runtime as rt
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def item():
|
||||
return rt.normalize_image_list_inputs(
|
||||
torch.arange(18, dtype=torch.float64).reshape(1, 2, 3, 3),
|
||||
torch.zeros(1, 2, 3, dtype=torch.bfloat16),
|
||||
{'value_1': 'original', 'value_2': True, 'value_3': 7, 'value_4': 0.25,
|
||||
'value_10': torch.ones(2, 3, dtype=torch.float16)})[0]
|
||||
|
||||
|
||||
def graph():
|
||||
return {
|
||||
'source': {'class_type': 'Source', 'inputs': {}},
|
||||
'start': {'class_type': rt.START_CLASS, 'inputs': {
|
||||
'images': ['source', 0], 'image_list_execution_id': 'transient',
|
||||
'image_list_workflow_meta': {'job': 'metadata'},
|
||||
'image_list_enabled_worker_ids': '["worker"]',
|
||||
'image_list_master_callback': 'http://old-master',
|
||||
'values.value_1': 'inherited', 'worker_context_id': 'old-context'}},
|
||||
'op': {'class_type': 'Operation', 'inputs': {'image': ['start', 0]}},
|
||||
'end': {'class_type': rt.END_CLASS, 'inputs': {
|
||||
'images': ['op', 0], 'context': ['start', 3], 'value_3': ['start', 2]}},
|
||||
}
|
||||
|
||||
|
||||
def worker_prompt(token, context, data, *, local=True):
|
||||
prompt = graph()
|
||||
result = rt.build_worker_prompt(prompt, rt.validate_image_list_regions(prompt)[0],
|
||||
token, context, data['index'], 'http://master',
|
||||
item_identity=rt.item_content_identity(data))
|
||||
if local:
|
||||
result['start']['inputs']['image_list_local_execution'] = True
|
||||
result['end']['inputs']['image_list_local_execution'] = True
|
||||
return result
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def clean_state():
|
||||
yield
|
||||
rt._WORKER_ITEMS.clear()
|
||||
rt._MASTER_JOBS.clear()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def boundaries(monkeypatch):
|
||||
"""Load only the boundary module; restore every replaced module on exit."""
|
||||
name = '_native_cache_identity_tests'
|
||||
for suffix in ('', '.nodes', '.utils', '.image_list', '.api'):
|
||||
module = types.ModuleType(name + suffix)
|
||||
module.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
|
||||
monkeypatch.setitem(sys.modules, name + suffix, module)
|
||||
monkeypatch.setitem(sys.modules, name + '.image_list.runtime', rt)
|
||||
io = types.SimpleNamespace(ComfyNode=type('ComfyNode', (), {}),
|
||||
NodeOutput=lambda *args: args)
|
||||
for cls in ('Image', 'Mask', 'String', 'Int', 'Float', 'Boolean'):
|
||||
setattr(io, cls, object())
|
||||
monkeypatch.setitem(sys.modules, 'comfy_api', types.ModuleType('comfy_api'))
|
||||
monkeypatch.setitem(sys.modules, 'comfy_api.v0_0_2', types.SimpleNamespace(io=io))
|
||||
monkeypatch.setitem(sys.modules, name + '.utils.async_helpers',
|
||||
types.SimpleNamespace(run_async_in_server_loop=lambda coro, **kw: asyncio.run(coro)))
|
||||
spec = importlib.util.spec_from_file_location(name + '.nodes.image_list', ROOT / 'nodes/image_list.py')
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
monkeypatch.setitem(sys.modules, spec.name, module)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def start(boundaries, prompt):
|
||||
boundaries.DistributedImageListStart.hidden = types.SimpleNamespace(prompt=prompt, unique_id='start')
|
||||
return boundaries.DistributedImageListStart.execute(**prompt['start']['inputs'])
|
||||
|
||||
|
||||
def end(boundaries, prompt, cached):
|
||||
boundaries.DistributedImageListEnd.hidden = types.SimpleNamespace(prompt=prompt, unique_id='end')
|
||||
inputs = {k: v for k, v in prompt['end']['inputs'].items()
|
||||
if k not in ('images', 'context', 'value_3')}
|
||||
return boundaries.DistributedImageListEnd.execute(
|
||||
images=cached[0], context=cached[3], value_3=cached[2], **inputs)
|
||||
|
||||
|
||||
def stage(data):
|
||||
token = rt.create_worker_task_token()
|
||||
rt.store_worker_item(token, rt.encode_item(data))
|
||||
return token
|
||||
|
||||
|
||||
def master(data, token):
|
||||
context = rt.create_master_job('start', 'end', graph(), [data], ['worker'], '')
|
||||
rt.register_assignment(context, data['index'], 'worker', token)
|
||||
return context
|
||||
|
||||
|
||||
def test_identity_is_exact_typed_order_independent_and_not_transport_bounded(monkeypatch):
|
||||
original = item()
|
||||
identity = rt.item_content_identity(original)
|
||||
assert rt.is_valid_item_identity(identity)
|
||||
reordered = copy.deepcopy(original)
|
||||
reordered['values'] = dict(reversed(list(reordered['values'].items())))
|
||||
assert rt.item_content_identity(reordered) == identity
|
||||
strided = copy.deepcopy(original)
|
||||
strided['image'] = original['image'].transpose(1, 2).contiguous().transpose(1, 2)
|
||||
assert not strided['image'].is_contiguous()
|
||||
assert rt.item_content_identity(strided) == identity
|
||||
monkeypatch.setattr(rt, 'MAX_PAYLOAD_BYTES', 1)
|
||||
with pytest.raises(ValueError, match='transport limit'):
|
||||
rt.encode_item(original)
|
||||
assert rt.item_content_identity(original) == identity
|
||||
scalar_identities = set()
|
||||
for value in (True, 1, 1.0, '1', 0.0, -0.0):
|
||||
changed = copy.deepcopy(original)
|
||||
changed['values']['value_2'] = value
|
||||
scalar_identities.add(rt.item_content_identity(changed))
|
||||
assert len(scalar_identities) == 6
|
||||
|
||||
|
||||
@pytest.mark.parametrize('change', ['image', 'mask', 'value', 'index', 'dtype', 'shape', 'slot', 'missing-mask'])
|
||||
def test_every_item_component_invalidates_identity(change):
|
||||
original, changed = item(), item()
|
||||
if change in ('image', 'mask'):
|
||||
changed[change].reshape(-1)[0] = 1
|
||||
elif change == 'value':
|
||||
changed['values']['value_10'][0, 0] = 2
|
||||
elif change == 'index':
|
||||
changed['index'] = 1
|
||||
elif change == 'dtype':
|
||||
changed['image'] = changed['image'].float()
|
||||
elif change == 'shape':
|
||||
changed['values']['value_10'] = changed['values']['value_10'].reshape(3, 2)
|
||||
elif change == 'slot':
|
||||
changed['values']['value_9'] = changed['values'].pop('value_10')
|
||||
else:
|
||||
changed['mask'] = None
|
||||
assert rt.item_content_identity(original) != rt.item_content_identity(changed)
|
||||
|
||||
|
||||
def test_float64_low_bits_and_bfloat16_bytes_are_not_quantized():
|
||||
data, changed = item(), item()
|
||||
changed['image'].reshape(-1)[0] = torch.nextafter(torch.tensor(0., dtype=torch.float64),
|
||||
torch.tensor(1., dtype=torch.float64))
|
||||
assert rt.item_content_identity(data) != rt.item_content_identity(changed)
|
||||
changed = item()
|
||||
changed['mask'].view(torch.int16).reshape(-1)[0] = 1
|
||||
assert rt.item_content_identity(data) != rt.item_content_identity(changed)
|
||||
|
||||
|
||||
def test_stable_start_replaced_wholesale_and_legacy_graph_preserved():
|
||||
prompt = graph()
|
||||
frozen = copy.deepcopy(prompt)
|
||||
region = rt.validate_image_list_regions(prompt)[0]
|
||||
legacy = rt.build_worker_prompt(prompt, region, 'a', 'ctx-a', 0, 'http://a')
|
||||
assert legacy['start']['inputs']['worker_item_token'] == 'a'
|
||||
assert legacy['start']['inputs']['image_list_execution_id'] == 'transient'
|
||||
identity = rt.item_content_identity(item())
|
||||
a = rt.build_worker_prompt(prompt, region, 'a', 'ctx-a', 0, 'http://a', item_identity=identity)
|
||||
b = rt.build_worker_prompt(prompt, region, 'b', 'ctx-b', 0, 'http://b', item_identity=identity)
|
||||
assert a['start']['inputs'] == b['start']['inputs'] == {'worker_item_identity': identity}
|
||||
assert a['end']['inputs']['worker_context_id'] == 'ctx-a'
|
||||
assert b['end']['inputs']['worker_result_token'] == 'b'
|
||||
assert b['end']['inputs']['worker_item_identity'] == identity
|
||||
assert b['end']['inputs']['value_3'] == ['start', 2]
|
||||
assert prompt == frozen
|
||||
|
||||
|
||||
@pytest.mark.parametrize('bad', ['', 'sha256:' + 'z' * 64, 'a' * 64, False, ['bad']])
|
||||
def test_malformed_identity_rejected(bad):
|
||||
assert not rt.is_valid_item_identity(bad)
|
||||
prompt = graph()
|
||||
with pytest.raises((ValueError, TypeError), match='identity'):
|
||||
rt.build_worker_prompt(prompt, rt.validate_image_list_regions(prompt)[0],
|
||||
'token', 'context', 0, '', item_identity=bad)
|
||||
|
||||
|
||||
def test_cache_miss_then_hit_delivers_current_context_token_and_values(boundaries):
|
||||
data = item()
|
||||
a = stage(data)
|
||||
context_a = master(data, a)
|
||||
prompt_a = worker_prompt(a, context_a, data)
|
||||
cached = start(boundaries, prompt_a)
|
||||
assert cached[3] == [rt.item_content_identity(data)]
|
||||
assert 'image' not in rt._WORKER_ITEMS[a][1] # only a small, expiring receipt remains
|
||||
end(boundaries, prompt_a, cached)
|
||||
assert a not in rt._WORKER_ITEMS
|
||||
assert rt.get_master_job(context_a).results[0]['values'] == {'value_3': 0}
|
||||
b = stage(data)
|
||||
context_b = master(data, b)
|
||||
prompt_b = worker_prompt(b, context_b, data)
|
||||
assert prompt_a['start'] == prompt_b['start']
|
||||
# A native hit does NOT execute Start; End must consume this current upload.
|
||||
end(boundaries, prompt_b, cached)
|
||||
assert b not in rt._WORKER_ITEMS
|
||||
state = rt.get_master_job(context_b)
|
||||
assert state.expected_tokens[0] == b
|
||||
assert torch.equal(state.results[0]['image'], data['image'])
|
||||
assert rt._collect_output_lists([state.results[0]], state=state, include_values=True)[3][0]['value_3'] == 0
|
||||
with pytest.raises(RuntimeError, match='missing|expired|consumed'):
|
||||
end(boundaries, prompt_b, cached)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('miss', [False, True])
|
||||
def test_expired_item_cannot_deliver_even_with_cached_start(boundaries, miss):
|
||||
data = item()
|
||||
token = stage(data)
|
||||
context = master(data, token)
|
||||
prompt = worker_prompt(token, context, data)
|
||||
cached = start(boundaries, prompt) if miss else ([data['image']], [data['mask']], [0], [rt.item_content_identity(data)])
|
||||
_, staged = rt._WORKER_ITEMS[token]
|
||||
rt._WORKER_ITEMS[token] = (0, staged)
|
||||
with pytest.raises(RuntimeError, match='missing|expired|consumed'):
|
||||
end(boundaries, prompt, cached)
|
||||
assert not rt.get_master_job(context).received_indexes
|
||||
assert token not in rt._WORKER_ITEMS
|
||||
|
||||
|
||||
def test_identity_mismatch_on_start_and_on_cache_hit_fails_closed(boundaries):
|
||||
data, different = item(), item()
|
||||
different['image'][0, 0, 0, 0] = 3
|
||||
for hit in (False, True):
|
||||
token = stage(different)
|
||||
context = master(different, token)
|
||||
prompt = worker_prompt(token, context, data)
|
||||
cached = ([data['image']], [data['mask']], [0], [rt.item_content_identity(data)])
|
||||
with pytest.raises((ValueError, RuntimeError), match='identity'):
|
||||
end(boundaries, prompt, cached) if hit else start(boundaries, prompt)
|
||||
assert not rt.get_master_job(context).received_indexes
|
||||
rt.delete_worker_item(token)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('change', ['ambiguous', 'wrong-link', 'wrong-start-id', 'end-identity', 'end-index'])
|
||||
def test_start_requires_exact_paired_end_binding(boundaries, change):
|
||||
data = item()
|
||||
token = stage(data)
|
||||
prompt = worker_prompt(token, 'context', data)
|
||||
if change == 'ambiguous':
|
||||
prompt['second-end'] = copy.deepcopy(prompt['end'])
|
||||
elif change == 'wrong-link':
|
||||
prompt['end']['inputs']['context'] = ['start', 2]
|
||||
elif change == 'wrong-start-id':
|
||||
prompt['renamed'] = prompt.pop('start')
|
||||
elif change == 'end-identity':
|
||||
prompt['end']['inputs']['worker_item_identity'] = 'sha256:' + '0' * 64
|
||||
else:
|
||||
prompt['end']['inputs']['worker_item_index'] = 1
|
||||
boundaries.DistributedImageListStart.hidden = types.SimpleNamespace(prompt=prompt, unique_id='start')
|
||||
with pytest.raises((RuntimeError, ValueError), match='paired|identity|index'):
|
||||
boundaries.DistributedImageListStart.execute(worker_item_identity=rt.item_content_identity(data),
|
||||
image_list_local_execution=True)
|
||||
|
||||
|
||||
def test_end_checks_stable_context_before_delivering(boundaries):
|
||||
data = item()
|
||||
token = stage(data)
|
||||
context = master(data, token)
|
||||
prompt = worker_prompt(token, context, data)
|
||||
cached = list(start(boundaries, prompt))
|
||||
cached[3] = ['sha256:' + '0' * 64]
|
||||
with pytest.raises((RuntimeError, ValueError), match='identity|context'):
|
||||
end(boundaries, prompt, cached)
|
||||
assert not rt.get_master_job(context).received_indexes
|
||||
|
||||
|
||||
def test_legacy_worker_path_still_consumes_and_delivers_job_context(boundaries):
|
||||
data = item()
|
||||
token = stage(data)
|
||||
context = master(data, token)
|
||||
prompt = graph()
|
||||
prompt = rt.build_worker_prompt(prompt, rt.validate_image_list_regions(prompt)[0], token, context, 0, '')
|
||||
prompt['start']['inputs']['image_list_local_execution'] = True
|
||||
prompt['end']['inputs']['image_list_local_execution'] = True
|
||||
cached = start(boundaries, prompt)
|
||||
assert cached[3] == [context]
|
||||
assert token not in rt._WORKER_ITEMS
|
||||
end(boundaries, prompt, cached)
|
||||
assert rt.get_master_job(context).received_indexes == {0}
|
||||
|
||||
|
||||
def test_remote_cache_hit_submits_fresh_callback_metadata(boundaries, monkeypatch):
|
||||
data = item()
|
||||
token = stage(data)
|
||||
prompt = worker_prompt(token, 'fresh-context', data, local=False)
|
||||
submit = AsyncMock()
|
||||
monkeypatch.setattr(boundaries, 'submit_worker_result', submit)
|
||||
cached = ([data['image']], [data['mask']], [0], [rt.item_content_identity(data)])
|
||||
end(boundaries, prompt, cached)
|
||||
assert submit.await_args.args[:4] == ('http://master', token, 'fresh-context', 0)
|
||||
assert submit.await_args.args[4]['values']['value_3'] == {'type': 'INT', 'value': 0}
|
||||
assert token not in rt._WORKER_ITEMS
|
||||
|
||||
|
||||
def test_item_route_returns_server_computed_identity(boundaries, monkeypatch):
|
||||
from aiohttp import web
|
||||
name = '_native_cache_identity_tests'
|
||||
server = types.SimpleNamespace(PromptServer=types.SimpleNamespace(instance=types.SimpleNamespace(routes=web.RouteTableDef())))
|
||||
monkeypatch.setitem(sys.modules, 'server', server)
|
||||
spec = importlib.util.spec_from_file_location(name + '.api.image_list_routes', ROOT / 'api/image_list_routes.py')
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
monkeypatch.setitem(sys.modules, spec.name, module)
|
||||
spec.loader.exec_module(module)
|
||||
data = item()
|
||||
token = rt.create_worker_task_token()
|
||||
monkeypatch.setattr(module, '_json_body', AsyncMock(return_value={
|
||||
'token': token, 'context_id': 'context', 'item': rt.encode_item(data),
|
||||
'item_identity': 'sha256:' + '0' * 64}))
|
||||
response = asyncio.run(module.image_list_item_endpoint(object()))
|
||||
assert response.status == 200
|
||||
assert json.loads(response.body)['item_identity'] == rt.item_content_identity(data)
|
||||
assert len(rt._WORKER_ITEMS[token]) == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize('responses', [
|
||||
('valid', 'valid'), ('valid', 'absent'), ('invalid', 'valid'),
|
||||
('bad-json', 'valid'), ('non-object', 'valid'),
|
||||
])
|
||||
def test_runtime_opts_in_per_upload_and_can_downgrade(monkeypatch, responses):
|
||||
name = '_native_cache_identity_runtime'
|
||||
for suffix in ('', '.image_list', '.api', '.api.orchestration', '.utils'):
|
||||
module = types.ModuleType(name + suffix)
|
||||
module.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
|
||||
monkeypatch.setitem(sys.modules, name + suffix, module)
|
||||
spec = importlib.util.spec_from_file_location(name + '.image_list.runtime', ROOT / 'image_list/runtime.py')
|
||||
runtime = importlib.util.module_from_spec(spec)
|
||||
monkeypatch.setitem(sys.modules, spec.name, runtime)
|
||||
spec.loader.exec_module(runtime)
|
||||
workers = [{'id': 'worker', 'type': 'local'}]
|
||||
dispatched, uploads = [], []
|
||||
preflight = AsyncMock(return_value={'worker': True})
|
||||
sync_media = AsyncMock()
|
||||
|
||||
class Response:
|
||||
status = 200
|
||||
|
||||
def __init__(self, data):
|
||||
self.data = data
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
pass
|
||||
|
||||
async def json(self):
|
||||
if self.data == 'bad-json':
|
||||
raise ValueError('old response is not JSON')
|
||||
return self.data
|
||||
|
||||
def raise_for_status(self):
|
||||
pass
|
||||
|
||||
class Session:
|
||||
def post(self, url, *, data, **kw):
|
||||
assert url.endswith('/distributed/image-list/item')
|
||||
body = json.loads(data)
|
||||
identity = runtime.store_worker_item(body['token'], body['item'])
|
||||
mode = responses[len(uploads)]
|
||||
uploads.append(body)
|
||||
response = {'valid': {'item_identity': identity}, 'absent': {'status': 'stored'},
|
||||
'invalid': {'item_identity': 'bad'}, 'bad-json': 'bad-json', 'non-object': []}[mode]
|
||||
return Response(response)
|
||||
|
||||
def get(self, url, **kw):
|
||||
pid = url.rsplit('/', 1)[-1]
|
||||
return Response({pid: {'status': {'completed': True, 'status_str': 'success'}}})
|
||||
|
||||
def delete(self, url, **kw):
|
||||
runtime.delete_worker_item(url.rsplit('/', 1)[-1])
|
||||
return Response({})
|
||||
|
||||
async def dispatch(worker, prompt, meta, **kwargs):
|
||||
dispatched.append(copy.deepcopy(prompt))
|
||||
delivery = prompt['end']['inputs']
|
||||
token, index = delivery['worker_result_token'], delivery['worker_item_index']
|
||||
identity = prompt['start']['inputs'].get('worker_item_identity')
|
||||
options = {'item_identity': identity, 'item_index': index} if identity else {}
|
||||
data = runtime.consume_worker_item(token, **options)
|
||||
if identity:
|
||||
runtime.finish_worker_item(token, identity, index)
|
||||
runtime.set_master_result(context, index, token, {
|
||||
'image': data['image'], 'text': str(index), 'values': {'value_3': index}})
|
||||
assert kwargs['partial_execution_targets'] == ['end']
|
||||
return f'worker-{index}'
|
||||
|
||||
dependencies = {
|
||||
'.api.orchestration.dispatch': types.SimpleNamespace(
|
||||
preflight_image_list_workers=preflight, dispatch_worker_prompt=dispatch),
|
||||
'.api.orchestration.prompt_transform': types.SimpleNamespace(
|
||||
_resolve_direct_workers_for_image_list=lambda *args: workers),
|
||||
'.api.orchestration.media_sync': types.SimpleNamespace(sync_worker_media=sync_media),
|
||||
'.utils.config': types.SimpleNamespace(load_config=lambda: {'settings': {}},
|
||||
get_worker_timeout_seconds=lambda: 60),
|
||||
'.utils.network': types.SimpleNamespace(build_master_callback_url=lambda *a, **kw: 'http://master',
|
||||
build_worker_url=lambda worker, path: 'http://worker' + path,
|
||||
get_client_session=AsyncMock(return_value=Session())),
|
||||
'.utils.logging': types.SimpleNamespace(log=lambda *args: None),
|
||||
}
|
||||
for suffix, module in dependencies.items():
|
||||
monkeypatch.setitem(sys.modules, name + suffix, module)
|
||||
monkeypatch.setitem(sys.modules, 'aiohttp', types.SimpleNamespace(
|
||||
ClientTimeout=lambda **kwargs: kwargs, ClientError=ConnectionError))
|
||||
comfy = types.ModuleType('comfy')
|
||||
comfy.model_management = types.SimpleNamespace(processing_interrupted=lambda: False)
|
||||
monkeypatch.setitem(sys.modules, 'comfy', comfy)
|
||||
monkeypatch.setitem(sys.modules, 'comfy.model_management', comfy.model_management)
|
||||
items = [item(), item()]
|
||||
items[1]['index'] = 1
|
||||
context = runtime.create_master_job('start', 'end', graph(), items, ['worker'], '')
|
||||
state = runtime.get_master_job(context)
|
||||
try:
|
||||
result = asyncio.run(runtime._run_image_list_job(state, include_values=True))
|
||||
assert result[2] == ['0', '1']
|
||||
assert [row['value_3'] for row in result[3]] == [0, 1]
|
||||
assert len(uploads) == len(dispatched) == 2
|
||||
assert sync_media.await_count == 1
|
||||
for index, prompt in enumerate(dispatched):
|
||||
start_inputs = prompt['start']['inputs']
|
||||
if responses[index] == 'valid':
|
||||
assert start_inputs == {'worker_item_identity': runtime.item_content_identity(items[index])}
|
||||
assert prompt['end']['inputs']['worker_context_id'] == context
|
||||
else:
|
||||
assert 'worker_item_identity' not in start_inputs
|
||||
assert start_inputs['worker_item_token'] == uploads[index]['token']
|
||||
assert start_inputs['worker_context_id'] == context
|
||||
assert 'worker_item_identity' not in prompt['end']['inputs']
|
||||
assert prompt['end']['inputs']['worker_result_token'] == uploads[index]['token']
|
||||
assert not runtime._WORKER_ITEMS
|
||||
finally:
|
||||
runtime.remove_master_job(context)
|
||||
runtime._WORKER_ITEMS.clear()
|
||||
|
||||
|
||||
def test_runtime_local_path_hashes_without_transport(monkeypatch):
|
||||
name = '_native_cache_identity_local_runtime'
|
||||
for suffix in ('', '.image_list', '.api', '.api.orchestration', '.utils'):
|
||||
module = types.ModuleType(name + suffix)
|
||||
module.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
|
||||
monkeypatch.setitem(sys.modules, name + suffix, module)
|
||||
spec = importlib.util.spec_from_file_location(name + '.image_list.runtime', ROOT / 'image_list/runtime.py')
|
||||
runtime = importlib.util.module_from_spec(spec)
|
||||
monkeypatch.setitem(sys.modules, spec.name, runtime)
|
||||
spec.loader.exec_module(runtime)
|
||||
dispatched = []
|
||||
|
||||
class Slot:
|
||||
def __init__(self, **kw):
|
||||
pass
|
||||
|
||||
def submit(self, prompt, prompt_id, end_id, workflow_meta):
|
||||
dispatched.append(prompt)
|
||||
delivery = prompt[end_id]['inputs']
|
||||
identity = prompt['start']['inputs']['worker_item_identity']
|
||||
data = runtime.consume_worker_item(delivery['worker_result_token'], item_identity=identity,
|
||||
item_index=delivery['worker_item_index'])
|
||||
runtime.finish_worker_item(delivery['worker_result_token'], identity, data['index'])
|
||||
runtime.set_master_result(delivery['worker_context_id'], data['index'], delivery['worker_result_token'],
|
||||
{'image': data['image'], 'values': {'value_3': data['index']}})
|
||||
future = asyncio.get_running_loop().create_future()
|
||||
future.set_result(None)
|
||||
return future
|
||||
|
||||
def request_stop(self):
|
||||
pass
|
||||
|
||||
async def close(self):
|
||||
pass
|
||||
|
||||
def forbidden(*args, **kwargs):
|
||||
raise AssertionError('local path attempted HTTP/transport')
|
||||
|
||||
dependencies = {
|
||||
'.image_list.local_executor': types.SimpleNamespace(LocalExecutionSlot=Slot),
|
||||
'.api.orchestration.dispatch': types.SimpleNamespace(
|
||||
preflight_image_list_workers=forbidden, dispatch_worker_prompt=forbidden),
|
||||
'.api.orchestration.prompt_transform': types.SimpleNamespace(
|
||||
_resolve_direct_workers_for_image_list=lambda *args: []),
|
||||
'.utils.config': types.SimpleNamespace(load_config=lambda: {'settings': {}},
|
||||
get_worker_timeout_seconds=lambda: 60),
|
||||
'.utils.network': types.SimpleNamespace(build_master_callback_url=forbidden,
|
||||
build_worker_url=forbidden, get_client_session=forbidden),
|
||||
'.utils.logging': types.SimpleNamespace(log=lambda *args: None),
|
||||
}
|
||||
for suffix, module in dependencies.items():
|
||||
monkeypatch.setitem(sys.modules, name + suffix, module)
|
||||
monkeypatch.setitem(sys.modules, 'aiohttp', types.SimpleNamespace(
|
||||
ClientTimeout=lambda **kwargs: kwargs, ClientError=ConnectionError))
|
||||
comfy = types.ModuleType('comfy')
|
||||
comfy.model_management = types.SimpleNamespace(processing_interrupted=lambda: False)
|
||||
monkeypatch.setitem(sys.modules, 'comfy', comfy)
|
||||
monkeypatch.setitem(sys.modules, 'comfy.model_management', comfy.model_management)
|
||||
monkeypatch.setattr(runtime, 'MAX_PAYLOAD_BYTES', 1)
|
||||
monkeypatch.setattr(runtime, 'encode_item', forbidden)
|
||||
data = item()
|
||||
context = runtime.create_master_job('start', 'end', graph(), [data], [], '', master_participates=True)
|
||||
try:
|
||||
result = asyncio.run(runtime._run_image_list_job(runtime.get_master_job(context), include_values=True))
|
||||
assert dispatched[0]['start']['inputs'] == {'worker_item_identity': runtime.item_content_identity(data),
|
||||
'image_list_local_execution': True}
|
||||
assert result[3][0]['value_3'] == 0
|
||||
assert not runtime._WORKER_ITEMS
|
||||
finally:
|
||||
runtime.remove_master_job(context)
|
||||
runtime._WORKER_ITEMS.clear()
|
||||
|
||||
|
||||
def test_real_native_signatures_and_cached_delivery():
|
||||
core = Path(os.environ.get('COMFYUI_SOURCE_ROOT', str(Path(sys.executable).parents[2])))
|
||||
if not (core / 'comfy_execution/caching.py').exists():
|
||||
pytest.skip('real ComfyUI source not available beside this interpreter')
|
||||
result = subprocess.run([sys.executable, str(ROOT / 'tests/helpers/image_list_native_cache_identity.py'), str(core)],
|
||||
text=True, capture_output=True, timeout=90)
|
||||
assert result.returncode == 0, result.stdout + result.stderr
|
||||
assert 'NATIVE_CACHE_IDENTITY_OK' in result.stdout
|
||||
Reference in New Issue
Block a user