Vito
55fa2c4870
* fix(usage): count cached re-runs of runtime-linked prompt text (#347) A PromptManager-family node whose text comes from a runtime link (PreviewAny, an LLM, a list node) counted only the job that executed it; later jobs served from ComfyUI's output cache counted nothing. Each execution now records a receipt (digests of its STRING outputs and hashes of the text it stores) per run key and queue item. On execution_cached the receipt whose digests equal the cache entry's text is replayed once for the job, with the node's role and what the queue hook already saw recomputed from the job's own graph. The executing node and the replay count through one shared own_counts. Unmatched cache entries count nothing, as before; a node that executes after its replay in the same job skips its own count. No schema change, no backfill. * fix(usage): never replay a node twice in one job; dedupe receipts once complete (#347) A repeated execution_cached id or message for the same job now counts the node once. Receipts replace equal older ones only when their job finishes (or the next job starts), so a list still recording [A] of [A, B] no longer drops an older [A] receipt. Node-level tests check the receipt each execution records. * test(comfyui): count #347's cached-run cases in the execution suite pin (#347) The pin listed 19 execution-text-link cases per project; #347 adds five node kinds and eight one-renderer cases, so the suite lists 32 per project and 118 in total with #332's image-attribution case from main.
218 lines
7.8 KiB
Python
218 lines
7.8 KiB
Python
"""Pinned-host v1 adapter for local cached STRING evidence only."""
|
|
|
|
import functools
|
|
import threading
|
|
from typing import Callable
|
|
|
|
from .execution_context_capture import active_capture
|
|
from .execution_contexts import PromptSnapshot, TextEvidence
|
|
|
|
_STATE_ATTRIBUTE = "_prompt_manager_executor_messages_v1"
|
|
_registration_lock = threading.Lock()
|
|
|
|
|
|
class _Messages:
|
|
def __init__(self, executor_class: type) -> None:
|
|
self.executor_class = executor_class
|
|
self.lock = threading.Lock()
|
|
self.callbacks = {}
|
|
self.wrapper = None
|
|
self.original = None
|
|
self.installed = False
|
|
|
|
def subscribe(self, key: str, callback: Callable) -> Callable[[], None]:
|
|
with self.lock:
|
|
if not self.installed:
|
|
original = self.executor_class.add_message
|
|
if not callable(original):
|
|
raise ValueError("Executor message capability unavailable")
|
|
|
|
@functools.wraps(original)
|
|
def observe(executor, *args, **kwargs):
|
|
result = original(executor, *args, **kwargs)
|
|
with self.lock:
|
|
callbacks = tuple(entry[1] for entry in self.callbacks.values())
|
|
for notify in callbacks:
|
|
try:
|
|
notify(executor, args, kwargs)
|
|
except Exception:
|
|
# Best-effort observers cannot change host results.
|
|
pass
|
|
return result
|
|
|
|
self.executor_class.add_message = observe
|
|
self.original, self.wrapper, self.installed = original, observe, True
|
|
existing = self.callbacks.get(key)
|
|
token = existing[0] if existing is not None else object()
|
|
if existing is None:
|
|
self.callbacks[key] = (token, callback)
|
|
|
|
def unsubscribe() -> None:
|
|
with self.lock:
|
|
current = self.callbacks.get(key)
|
|
if current is None or current[0] is not token:
|
|
return
|
|
del self.callbacks[key]
|
|
if (
|
|
not self.callbacks
|
|
and self.executor_class.add_message is self.wrapper
|
|
):
|
|
self.executor_class.add_message = self.original
|
|
self.installed = False
|
|
|
|
return unsubscribe
|
|
|
|
|
|
def subscribe_executor_messages(key: str, callback: Callable) -> Callable[[], None]:
|
|
"""Observe successful original calls; preserve earlier and later wrappers."""
|
|
from execution import PromptExecutor
|
|
|
|
with _registration_lock:
|
|
state = vars(PromptExecutor).get(_STATE_ATTRIBUTE)
|
|
if state is None:
|
|
state = _Messages(PromptExecutor)
|
|
setattr(PromptExecutor, _STATE_ATTRIBUTE, state)
|
|
return state.subscribe(key, callback)
|
|
|
|
|
|
# STRING output slots of each PromptManager-family node, in output order
|
|
STRING_SLOTS = {
|
|
"PromptManager": (1,),
|
|
"PromptManagerFocused": (1,),
|
|
"PromptManagerSplit": (1,),
|
|
"PromptManagerText": (0,),
|
|
"PromptManagerPair": (2, 3),
|
|
}
|
|
MAX_CACHED_VALUES = 64
|
|
|
|
|
|
def _local_cache_reader(executor: object) -> Callable | None:
|
|
"""``node_id -> outputs`` of the executor's local output cache, or None.
|
|
|
|
Only the supported cache classes, only ``get_local`` (never a provider or
|
|
an async lookup), only a real ``CacheEntry`` with at most four outputs.
|
|
"""
|
|
from comfy_execution import caching
|
|
from execution import CacheEntry
|
|
|
|
cache = executor.caches.outputs
|
|
supported = tuple(
|
|
cls
|
|
for name in ("HierarchicalCache", "LRUCache", "RAMPressureCache")
|
|
if isinstance(cls := getattr(caching, name, None), type)
|
|
)
|
|
if type(cache) not in supported:
|
|
return None
|
|
local = getattr(cache, "get_local", None)
|
|
if not callable(local):
|
|
return None
|
|
|
|
def read(node_id: str) -> list | tuple | None:
|
|
entry = local(node_id)
|
|
if type(entry) is not CacheEntry:
|
|
return None
|
|
outputs = entry.outputs
|
|
if type(outputs) not in (list, tuple) or not 0 < len(outputs) <= 4:
|
|
return None
|
|
return outputs
|
|
|
|
return read
|
|
|
|
|
|
def cached_string_outputs(
|
|
executor: object, node_id: str, node_type: str
|
|
) -> tuple[tuple[str, ...], ...] | None:
|
|
"""One tuple per STRING slot of a cached node's local entry, or None (#347).
|
|
|
|
None when the cache or entry is not supported, a slot is missing or holds
|
|
more than 64 values, any value is not a ``str``, or the slots differ in
|
|
length. Never runs a provider or an async lookup.
|
|
"""
|
|
try:
|
|
slots = STRING_SLOTS.get(node_type)
|
|
if slots is None or type(node_id) is not str:
|
|
return None
|
|
local = _local_cache_reader(executor)
|
|
outputs = local(node_id) if local is not None else None
|
|
if outputs is None or not max(slots) < len(outputs):
|
|
return None
|
|
found = []
|
|
for slot in slots:
|
|
values = outputs[slot]
|
|
if type(values) not in (list, tuple) or len(values) > MAX_CACHED_VALUES:
|
|
return None
|
|
if any(type(value) is not str for value in values):
|
|
return None
|
|
found.append(tuple(values))
|
|
if len({len(values) for values in found}) != 1:
|
|
return None
|
|
return tuple(found)
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def observe_cached_context(
|
|
executor: object, message_type: object, message_data: object
|
|
) -> None:
|
|
"""Copy bounded local strings without cache/provider execution or retention."""
|
|
try:
|
|
binding = active_capture()
|
|
if binding is None or type(message_type) is not str:
|
|
return
|
|
manager, cap = binding
|
|
if (
|
|
type(message_data) is not dict
|
|
or type(message_data.get("prompt_id")) is not str
|
|
or message_data["prompt_id"] != cap.job_id
|
|
or id(executor.server.prompt_queue) != manager.queue_identity
|
|
):
|
|
return
|
|
if message_type == "execution_start":
|
|
manager.bind_executor(cap, id(executor))
|
|
return
|
|
if message_type != "execution_cached" or not manager.matches_executor(
|
|
cap, id(executor)
|
|
):
|
|
return
|
|
nodes = message_data.get("nodes")
|
|
if type(nodes) not in (list, tuple) or len(nodes) > 500:
|
|
return
|
|
local = _local_cache_reader(executor)
|
|
if local is None:
|
|
return
|
|
# Only source IDs retained by the bounded submitted graph are queried.
|
|
for source in cap.sources:
|
|
if not any(type(node) is str and node == source.node_id for node in nodes):
|
|
continue
|
|
slot = (
|
|
(3 if source.output_slot in (1, 3) else 2)
|
|
if source.node_type == "PromptManagerPair"
|
|
else 0 if source.node_type == "PromptManagerText" else 1
|
|
)
|
|
outputs = local(source.node_id)
|
|
if outputs is None or not slot < len(outputs):
|
|
continue
|
|
values = outputs[slot]
|
|
if type(values) not in (list, tuple) or len(values) > 64:
|
|
continue
|
|
inferred = manager.source_snapshot(cap, source)
|
|
if inferred is None:
|
|
continue
|
|
for index, value in enumerate(values):
|
|
if type(value) is not str:
|
|
continue
|
|
manager.observe(
|
|
cap,
|
|
source,
|
|
index,
|
|
PromptSnapshot(
|
|
inferred.raw,
|
|
TextEvidence(value, "recorded", "host_cache_output", None),
|
|
inferred.prepend,
|
|
inferred.append,
|
|
),
|
|
)
|
|
except Exception:
|
|
# Unknown host versions and malformed entries leave evidence unavailable.
|
|
return
|