Compare commits

..
11 changed files with 1137 additions and 27 deletions
+3 -2
View File
@@ -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:
+11
View File
@@ -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.
+40 -1
View File
@@ -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
View File
@@ -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
View File
@@ -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])
+33 -2
View File
@@ -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
+2
View File
@@ -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