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.
533 lines
20 KiB
Python
533 lines
20 KiB
Python
"""Count a cached PromptManager node like the execution it reuses (#347).
|
|
|
|
The queue hook (``usage_tracking``) counts text it can read from the queued
|
|
graph. Text only known at run time (a PreviewAny or LLM link, a list node) is
|
|
counted by the node itself when it executes. When ComfyUI later serves the node
|
|
from its output cache the node does not execute, so without this module that
|
|
job is not counted.
|
|
|
|
When a PromptManager-family node executes it records a *receipt item*: the
|
|
role-independent facts its own count is made from (the stored text hashes)
|
|
and digests of the STRING outputs it returned. When ComfyUI reports the node
|
|
as cached for a later job (``execution_cached``), the receipt whose digests
|
|
equal the text in the cache entry is replayed once for that job. The node's
|
|
role (positive, negative, which Pair outputs are used) and what the hook
|
|
already saw are recomputed from that job's own graph, and both paths count
|
|
through the same ``own_counts`` so they cannot drift.
|
|
|
|
Fail closed: when the cached text cannot be matched to a receipt recorded in
|
|
this process, nothing is counted (the behavior before #347). In memory only and
|
|
bounded; starts no threads.
|
|
"""
|
|
|
|
import hashlib
|
|
import threading
|
|
from collections import OrderedDict
|
|
from typing import Any, Callable, Dict, FrozenSet, List, NamedTuple, Optional, Tuple
|
|
|
|
try:
|
|
from .hashing import generate_prompt_hash
|
|
from .logging_config import get_logger
|
|
from .prompt_graph import (
|
|
PROMPT_MANAGER_TYPES,
|
|
intrinsic_pair_nodes,
|
|
resolve_text,
|
|
run_prompt_components,
|
|
run_prompt_nodes,
|
|
)
|
|
from .prompt_pair_resolver import resolve_text_input
|
|
from .queue_observers import subscribe_dequeue, subscribe_task_done
|
|
from .run_keys import prompt_node_key, upstream_nodes
|
|
from .usage_tracking import find_pair_id_by_hashes, survive_switch
|
|
except ImportError:
|
|
from utils.hashing import generate_prompt_hash
|
|
from utils.logging_config import get_logger
|
|
from utils.prompt_graph import (
|
|
PROMPT_MANAGER_TYPES,
|
|
intrinsic_pair_nodes,
|
|
resolve_text,
|
|
run_prompt_components,
|
|
run_prompt_nodes,
|
|
)
|
|
from utils.prompt_pair_resolver import resolve_text_input
|
|
from utils.queue_observers import subscribe_dequeue, subscribe_task_done
|
|
from utils.run_keys import prompt_node_key, upstream_nodes
|
|
from utils.usage_tracking import find_pair_id_by_hashes, survive_switch
|
|
|
|
logger = get_logger("prompt_manager.cached_runs")
|
|
|
|
PAIR = "PromptManagerPair"
|
|
MAX_KEYS = 256 # run keys remembered (the same bound as MAX_KEYED_RUNS)
|
|
MAX_RECEIPTS = 4 # receipts per run key: a recurring upstream value (A, B, A)
|
|
MAX_ITEMS = 64 # list items per receipt (the cache reader's bound)
|
|
MAX_CACHED_NODES = 500 # the execution_cached guard the history capture uses
|
|
_EMPTY_HASH = generate_prompt_hash("")
|
|
|
|
Count = Tuple[Any, ...] # ("prompt", hash) or ("pair", positive, negative or None)
|
|
|
|
|
|
class Facts(NamedTuple):
|
|
"""What one execution stored: hashes of the stored text, None when nothing.
|
|
|
|
``positive_hash`` is None for an empty editor (or saving disabled);
|
|
``negative_hash`` (Pair only) is None for an empty negative, which the pair
|
|
stores as ``negative_prompt_id IS NULL``.
|
|
"""
|
|
|
|
positive_hash: Optional[str]
|
|
negative_hash: Optional[str] = None
|
|
|
|
|
|
class Role(NamedTuple):
|
|
"""How a job's graph uses the node: which editors reach a sampler's positive."""
|
|
|
|
used_fields: FrozenSet[str]
|
|
pair_used: bool = False
|
|
|
|
|
|
class Seen(NamedTuple):
|
|
"""Which of the node's texts the queue hook resolved, and so already counted."""
|
|
|
|
text: bool = False
|
|
negative_text: bool = False
|
|
pair: bool = False
|
|
|
|
|
|
class _Item(NamedTuple):
|
|
outputs: Tuple[str, ...]
|
|
facts: Facts
|
|
|
|
|
|
class _Receipt:
|
|
__slots__ = ("item_id", "items")
|
|
|
|
def __init__(self, item_id: int) -> None:
|
|
self.item_id = item_id
|
|
self.items: Dict[int, _Item] = {}
|
|
|
|
def outputs(self) -> Dict[int, Tuple[str, ...]]:
|
|
return {index: item.outputs for index, item in self.items.items()}
|
|
|
|
|
|
class _Attempt(NamedTuple):
|
|
item_id: int
|
|
prompt_id: str
|
|
graph: Dict[str, Any]
|
|
nodes: Dict[str, Tuple[str, str]] # node id -> (class, run key)
|
|
replayed: set
|
|
|
|
|
|
def text_digest(value: str) -> str:
|
|
"""Digest of one returned STRING value, exactly as returned."""
|
|
return hashlib.sha256(value.encode("utf-8", "surrogatepass")).hexdigest()
|
|
|
|
|
|
def stored_facts(text: Any, negative_text: Any = None) -> Facts:
|
|
"""The facts of an execution that stores ``text`` (and a Pair's negative)."""
|
|
positive = text if isinstance(text, str) and text.strip() else None
|
|
negative = (
|
|
negative_text
|
|
if isinstance(negative_text, str) and negative_text.strip()
|
|
else None
|
|
)
|
|
return Facts(
|
|
generate_prompt_hash(positive) if positive is not None else None,
|
|
generate_prompt_hash(negative) if negative is not None else None,
|
|
)
|
|
|
|
|
|
def pair_components(facts: Facts, role: Role) -> List[Tuple[str, str]]:
|
|
"""``(field, hash)`` of each used, stored Pair component; each hash once."""
|
|
found, counted = [], set()
|
|
for field, prompt_hash in (
|
|
("text", facts.positive_hash),
|
|
("negative_text", facts.negative_hash),
|
|
):
|
|
if field not in role.used_fields or prompt_hash is None:
|
|
continue
|
|
if prompt_hash in counted:
|
|
continue
|
|
counted.add(prompt_hash)
|
|
found.append((field, prompt_hash))
|
|
return found
|
|
|
|
|
|
def own_counts(node_type: str, facts: Facts, role: Role, seen: Seen) -> List[Count]:
|
|
"""What one execution of a node counts itself: text the hook could not see.
|
|
|
|
Shared by the executing node and by replay. Excludes the pending first-use
|
|
catch-up, which only an execution that saves the prompt can consume.
|
|
"""
|
|
if facts.positive_hash is None:
|
|
return []
|
|
if node_type != PAIR:
|
|
if "text" in role.used_fields and not seen.text:
|
|
return [("prompt", facts.positive_hash)]
|
|
return []
|
|
counts: List[Count] = [
|
|
("prompt", prompt_hash)
|
|
for field, prompt_hash in pair_components(facts, role)
|
|
if not getattr(seen, field)
|
|
]
|
|
if role.pair_used and not seen.pair:
|
|
counts.append(("pair", facts.positive_hash, facts.negative_hash))
|
|
return counts
|
|
|
|
|
|
def graph_role(graph: Any, node_id: str, node_type: str) -> Role:
|
|
"""The node's role in a queued graph, by the functions the nodes use."""
|
|
if node_type != PAIR:
|
|
used = {"text"} if node_id in run_prompt_nodes(graph) else set()
|
|
return Role(frozenset(used))
|
|
used = {
|
|
field for source, field in run_prompt_components(graph) if source == node_id
|
|
}
|
|
return Role(frozenset(used), node_id in intrinsic_pair_nodes(graph))
|
|
|
|
|
|
def _same(resolved: Optional[str], prompt_hash: Optional[str]) -> bool:
|
|
return resolved is not None and generate_prompt_hash(resolved) == prompt_hash
|
|
|
|
|
|
def graph_seen(graph: Any, node_id: str, node_type: str, facts: Facts) -> Seen:
|
|
"""Which texts the queue hook resolved from ``graph`` (as the node decides)."""
|
|
if node_type != PAIR:
|
|
return Seen(text=_same(resolve_text(graph, node_id), facts.positive_hash))
|
|
positive = _same(resolve_text_input(graph, node_id, "text"), facts.positive_hash)
|
|
negative = _same(
|
|
resolve_text_input(graph, node_id, "negative_text"),
|
|
facts.negative_hash or _EMPTY_HASH,
|
|
)
|
|
return Seen(positive, negative, positive and negative)
|
|
|
|
|
|
def _host_context() -> Tuple[Optional[str], Optional[int]]:
|
|
"""(prompt id, list index) of the node executing on this thread, if known."""
|
|
try:
|
|
from comfy_execution.utils import get_executing_context
|
|
except Exception:
|
|
return None, None
|
|
try:
|
|
context = get_executing_context()
|
|
except Exception:
|
|
return None, None
|
|
prompt_id = getattr(context, "prompt_id", None)
|
|
index = getattr(context, "list_index", None)
|
|
return (
|
|
prompt_id if isinstance(prompt_id, str) and prompt_id else None,
|
|
index if type(index) is int else None,
|
|
)
|
|
|
|
|
|
def write_count(db: Any, count: Count) -> bool:
|
|
"""Record one counted run; False when the prompt or pair no longer exists."""
|
|
|
|
def write() -> bool:
|
|
if count[0] == "prompt":
|
|
row = db.get_prompt_by_hash(count[1])
|
|
if row:
|
|
db.record_prompt_use(row["id"])
|
|
return bool(row)
|
|
pair_id = find_pair_id_by_hashes(db, count[1], count[2])
|
|
if pair_id:
|
|
db.record_prompt_pair_use(pair_id)
|
|
return bool(pair_id)
|
|
|
|
return bool(survive_switch(db, write))
|
|
|
|
|
|
class CachedRuns:
|
|
"""Receipts of executed nodes and the one job currently executing."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
max_keys: int = MAX_KEYS,
|
|
max_receipts: int = MAX_RECEIPTS,
|
|
max_items: int = MAX_ITEMS,
|
|
context: Callable[[], Tuple[Optional[str], Optional[int]]] = _host_context,
|
|
) -> None:
|
|
self._max_keys = max_keys
|
|
self._max_receipts = max_receipts
|
|
self._max_items = max_items
|
|
self._context = context
|
|
self._receipts: "OrderedDict[str, List[_Receipt]]" = OrderedDict()
|
|
self._attempt: Optional[_Attempt] = None
|
|
self._lock = threading.Lock()
|
|
|
|
# -- the job ------------------------------------------------------------
|
|
|
|
def begin_attempt(self, result: Any) -> None:
|
|
"""A job was dequeued: ``result`` is ``PromptQueue.get``'s return value."""
|
|
attempt = None
|
|
try:
|
|
item, item_id = result[0], result[1]
|
|
prompt_id, graph, outputs = item[1], item[2], item[4]
|
|
if type(item_id) is int and isinstance(prompt_id, str):
|
|
if isinstance(graph, dict):
|
|
attempt = _Attempt(
|
|
item_id, prompt_id, graph, _job_nodes(graph, outputs), set()
|
|
)
|
|
except Exception as error:
|
|
logger.debug(f"Dequeued job not tracked for cached runs: {error}")
|
|
with self._lock:
|
|
self._attempt = attempt
|
|
|
|
def end_attempt(self, item_id: Any) -> None:
|
|
"""The job finished: its receipts are complete and replace equal older ones.
|
|
|
|
Deduplicating only complete receipts keeps a list execution that is
|
|
still recording ([A] so far of [A, B]) from dropping an older [A].
|
|
"""
|
|
with self._lock:
|
|
if self._attempt is not None and self._attempt.item_id == item_id:
|
|
self._attempt = None
|
|
for run_key, receipts in self._receipts.items():
|
|
if receipts and receipts[0].item_id == item_id:
|
|
self._receipts[run_key] = _deduplicated(receipts)
|
|
|
|
def _current(self) -> Tuple[Optional[_Attempt], Optional[int]]:
|
|
"""The attempt the executing node belongs to, and its list index."""
|
|
prompt_id, index = self._context()
|
|
with self._lock:
|
|
attempt = self._attempt
|
|
if attempt is None or prompt_id is None or prompt_id != attempt.prompt_id:
|
|
return None, None
|
|
return attempt, index
|
|
|
|
# -- node execution -----------------------------------------------------
|
|
|
|
def replayed(self, node_id: Any) -> bool:
|
|
"""Whether this job already counted the executing node by replay."""
|
|
attempt, _ = self._current()
|
|
if attempt is None:
|
|
return False
|
|
with self._lock:
|
|
return str(node_id) in attempt.replayed
|
|
|
|
def record_item(self, node_id: Any, outputs: Tuple[Any, ...], facts: Facts) -> None:
|
|
"""Remember what this execution returned and was entitled to count."""
|
|
try:
|
|
attempt, index = self._current()
|
|
node = attempt.nodes.get(str(node_id)) if attempt else None
|
|
if node is None:
|
|
return
|
|
index = 0 if index is None else index
|
|
if not 0 <= index < self._max_items:
|
|
return
|
|
if any(not isinstance(value, str) for value in outputs):
|
|
return
|
|
item = _Item(tuple(text_digest(value) for value in outputs), facts)
|
|
self._store(node[1], attempt.item_id, index, item)
|
|
except Exception as error:
|
|
logger.debug(f"Could not record a cached-run receipt: {error}")
|
|
|
|
def _store(self, run_key: str, item_id: int, index: int, item: _Item) -> None:
|
|
with self._lock:
|
|
receipts = self._receipts.pop(run_key, [])
|
|
if receipts and receipts[0].item_id == item_id:
|
|
receipt = receipts[0]
|
|
else:
|
|
# The previous job's receipt is complete (its task_done may
|
|
# never have fired): it replaces older receipts it equals.
|
|
receipts = _deduplicated(receipts)
|
|
receipt = _Receipt(item_id)
|
|
receipts.insert(0, receipt)
|
|
receipt.items[index] = item
|
|
self._receipts[run_key] = receipts[: self._max_receipts]
|
|
while len(self._receipts) > self._max_keys:
|
|
self._receipts.popitem(last=False)
|
|
|
|
# -- cached nodes -------------------------------------------------------
|
|
|
|
def replay(
|
|
self,
|
|
executor: Any,
|
|
prompt_id: Any,
|
|
nodes: Any,
|
|
db_factory: Callable[[], Any],
|
|
reader: Optional[Callable[[Any, str, str], Any]] = None,
|
|
) -> None:
|
|
"""Count the job's cached nodes whose cached text a receipt proves."""
|
|
if reader is None:
|
|
reader = _host().cached_string_outputs
|
|
with self._lock:
|
|
attempt = self._attempt
|
|
if attempt is None or prompt_id != attempt.prompt_id:
|
|
return
|
|
if type(nodes) not in (list, tuple) or len(nodes) > MAX_CACHED_NODES:
|
|
return
|
|
db = None
|
|
for node_id in nodes:
|
|
node = attempt.nodes.get(node_id) if type(node_id) is str else None
|
|
if node is None:
|
|
continue
|
|
node_type, run_key = node
|
|
items = self._match(run_key, reader(executor, node_id, node_type))
|
|
if items is None:
|
|
logger.debug(
|
|
f"Cached node {node_id} not counted: no matching execution"
|
|
)
|
|
continue
|
|
with self._lock:
|
|
if node_id in attempt.replayed:
|
|
continue # a repeated message or id never counts twice
|
|
attempt.replayed.add(node_id)
|
|
counts: List[Count] = []
|
|
role = graph_role(attempt.graph, node_id, node_type)
|
|
for item in items:
|
|
seen = graph_seen(attempt.graph, node_id, node_type, item.facts)
|
|
counts.extend(own_counts(node_type, item.facts, role, seen))
|
|
if not counts:
|
|
continue
|
|
try:
|
|
db = db if db is not None else db_factory()
|
|
for count in counts:
|
|
if not write_count(db, count):
|
|
logger.info(
|
|
f"Cached node {node_id}: its {count[0]} no longer exists, "
|
|
"not counted"
|
|
)
|
|
except Exception as error:
|
|
logger.warning(f"Could not count cached node {node_id}: {error}")
|
|
|
|
def _match(self, run_key: str, values: Any) -> Optional[List[_Item]]:
|
|
"""The items of the one receipt whose outputs equal ``values``, in order."""
|
|
if not values:
|
|
return None
|
|
length = len(values[0])
|
|
if length == 0 or length > self._max_items:
|
|
return None
|
|
if any(len(slot) != length for slot in values):
|
|
return None
|
|
wanted = {
|
|
index: tuple(text_digest(slot[index]) for slot in values)
|
|
for index in range(length)
|
|
}
|
|
with self._lock:
|
|
receipts = self._receipts.get(run_key)
|
|
if not receipts:
|
|
return None
|
|
matches = [receipt for receipt in receipts if receipt.outputs() == wanted]
|
|
if not matches:
|
|
return None
|
|
self._receipts.move_to_end(run_key)
|
|
facts = {
|
|
tuple(receipt.items[index].facts for index in range(length))
|
|
for receipt in matches
|
|
}
|
|
if len(facts) != 1:
|
|
return None
|
|
return [matches[0].items[index] for index in range(length)]
|
|
|
|
|
|
def _deduplicated(receipts: List[_Receipt]) -> List[_Receipt]:
|
|
"""The newest receipt replaces older ones with the same outputs."""
|
|
if not receipts:
|
|
return receipts
|
|
outputs = receipts[0].outputs()
|
|
return [receipts[0]] + [r for r in receipts[1:] if r.outputs() != outputs]
|
|
|
|
|
|
def _job_nodes(graph: Dict[str, Any], outputs: Any) -> Dict[str, Tuple[str, str]]:
|
|
"""PromptManager-family nodes the job's outputs need, with their run keys."""
|
|
starts = (
|
|
[str(node) for node in outputs] if isinstance(outputs, (list, tuple)) else []
|
|
)
|
|
needed = upstream_nodes(graph, starts)
|
|
nodes = {}
|
|
for node_id in needed:
|
|
node = graph.get(node_id)
|
|
node_type = node.get("class_type") if isinstance(node, dict) else None
|
|
if node_type not in PROMPT_MANAGER_TYPES:
|
|
continue
|
|
run_key = prompt_node_key(graph, node_id)
|
|
if run_key is not None:
|
|
nodes[node_id] = (node_type, run_key)
|
|
return nodes
|
|
|
|
|
|
def _host() -> Any:
|
|
"""The host adapter module (imported late: it needs ComfyUI's modules)."""
|
|
try:
|
|
from . import execution_context_host
|
|
except ImportError:
|
|
from utils import execution_context_host
|
|
return execution_context_host
|
|
|
|
|
|
# Shared between the nodes and the host observers
|
|
CACHED_RUNS = CachedRuns()
|
|
|
|
_registered = threading.Event()
|
|
_register_lock = threading.Lock()
|
|
|
|
|
|
def register_cached_runs(
|
|
server: Any, db_factory: Callable[[], Any], runs: Optional[CachedRuns] = None
|
|
) -> bool:
|
|
"""Observe dequeue, ``execution_cached`` and ``task_done``; once per process.
|
|
|
|
Registers nothing when the host lacks any of them: counting is then as
|
|
before #347.
|
|
"""
|
|
runs = CACHED_RUNS if runs is None else runs
|
|
with _register_lock:
|
|
if _registered.is_set():
|
|
return False
|
|
queue = getattr(server, "prompt_queue", None)
|
|
if not callable(getattr(queue, "get", None)) or not callable(
|
|
getattr(queue, "task_done", None)
|
|
):
|
|
logger.info("Cached runs not counted: the prompt queue is not observable")
|
|
return False
|
|
try:
|
|
from execution import PromptExecutor
|
|
except Exception:
|
|
logger.info("Cached runs not counted: no ComfyUI executor")
|
|
return False
|
|
if not callable(getattr(PromptExecutor, "add_message", None)):
|
|
logger.info("Cached runs not counted: executor messages unavailable")
|
|
return False
|
|
subscribe_executor_messages = _host().subscribe_executor_messages
|
|
|
|
def on_message(executor: Any, args: tuple, kwargs: dict) -> None:
|
|
try:
|
|
event = args[0] if args else kwargs.get("event")
|
|
if event != "execution_cached":
|
|
return
|
|
server_of = getattr(executor, "server", None)
|
|
if getattr(server_of, "prompt_queue", None) is not queue:
|
|
return
|
|
data = args[1] if len(args) > 1 else kwargs.get("data")
|
|
if type(data) is not dict:
|
|
return
|
|
runs.replay(
|
|
executor, data.get("prompt_id"), data.get("nodes"), db_factory
|
|
)
|
|
except Exception as error:
|
|
logger.warning(f"Could not count cached nodes: {error}")
|
|
|
|
unsubscribers = []
|
|
try:
|
|
unsubscribers.append(
|
|
subscribe_dequeue(queue, "cached_runs", runs.begin_attempt)
|
|
)
|
|
unsubscribers.append(
|
|
subscribe_task_done(
|
|
queue,
|
|
"cached_runs",
|
|
lambda args, kwargs: runs.end_attempt(
|
|
args[0] if args else kwargs.get("item_id")
|
|
),
|
|
)
|
|
)
|
|
unsubscribers.append(subscribe_executor_messages("cached_runs", on_message))
|
|
except Exception:
|
|
for unsubscribe in unsubscribers:
|
|
unsubscribe()
|
|
raise
|
|
_registered.set()
|
|
logger.info("Counting cached runs of runtime-linked prompt text")
|
|
return True
|