Vito Sansevero
6559de8255
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.
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
|