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.
This commit is contained in:
Vito Sansevero
@@ -0,0 +1,515 @@
|
||||
"""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:
|
||||
with self._lock:
|
||||
if self._attempt is not None and self._attempt.item_id == item_id:
|
||||
self._attempt = None
|
||||
|
||||
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:
|
||||
receipt = _Receipt(item_id)
|
||||
receipts.insert(0, receipt)
|
||||
receipt.items[index] = item
|
||||
outputs = receipt.outputs()
|
||||
receipts = [receipt] + [
|
||||
older for older in receipts[1:] if older.outputs() != outputs
|
||||
]
|
||||
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:
|
||||
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 _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
|
||||
@@ -75,6 +75,82 @@ def subscribe_executor_messages(key: str, callback: Callable) -> Callable[[], No
|
||||
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:
|
||||
@@ -101,19 +177,8 @@ def observe_cached_context(
|
||||
nodes = message_data.get("nodes")
|
||||
if type(nodes) not in (list, tuple) or len(nodes) > 500:
|
||||
return
|
||||
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
|
||||
local = getattr(cache, "get_local", None)
|
||||
if not callable(local):
|
||||
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:
|
||||
@@ -124,11 +189,8 @@ def observe_cached_context(
|
||||
if source.node_type == "PromptManagerPair"
|
||||
else 0 if source.node_type == "PromptManagerText" else 1
|
||||
)
|
||||
entry = local(source.node_id)
|
||||
if type(entry) is not CacheEntry:
|
||||
continue
|
||||
outputs = entry.outputs
|
||||
if type(outputs) not in (list, tuple) or not slot < len(outputs) <= 4:
|
||||
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:
|
||||
|
||||
+50
-14
@@ -51,6 +51,53 @@ def _links(inputs: Any) -> List[str]:
|
||||
return [value[0] for value in inputs.values() if is_link(value)]
|
||||
|
||||
|
||||
def upstream_closure(
|
||||
graph: Dict[str, Any], starts: List[str], limit: Optional[int] = None
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""``{node_id: [class_type, inputs]}`` of ``starts`` and everything upstream.
|
||||
|
||||
None when a node or an upstream link's node is missing, or the closure
|
||||
exceeds ``limit`` nodes.
|
||||
"""
|
||||
closure: Dict[str, Any] = {}
|
||||
stack = list(starts)
|
||||
while stack:
|
||||
current = stack.pop()
|
||||
if current in closure:
|
||||
continue
|
||||
node = graph.get(current)
|
||||
if not isinstance(node, dict):
|
||||
return None
|
||||
if limit is not None and len(closure) >= limit:
|
||||
return None
|
||||
inputs = node.get("inputs")
|
||||
closure[current] = [node.get("class_type"), inputs]
|
||||
stack.extend(_links(inputs))
|
||||
return closure
|
||||
|
||||
|
||||
def upstream_nodes(graph: Any, starts: List[str]) -> List[str]:
|
||||
"""Ids of ``starts`` and every node upstream of them that exists in ``graph``.
|
||||
|
||||
Used to tell which nodes a job's outputs to execute need ("run selected
|
||||
output" jobs leave the rest of the graph out). Missing nodes are skipped.
|
||||
"""
|
||||
if not isinstance(graph, dict):
|
||||
return []
|
||||
found: Dict[str, None] = {}
|
||||
stack = [str(node) for node in starts]
|
||||
while stack:
|
||||
current = stack.pop()
|
||||
if current in found:
|
||||
continue
|
||||
node = graph.get(current)
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
found[current] = None
|
||||
stack.extend(_links(node.get("inputs")))
|
||||
return list(found)
|
||||
|
||||
|
||||
def prompt_node_key(graph: Any, node_id: Any) -> Optional[str]:
|
||||
"""Digest of ``node_id`` and its whole upstream closure, or None.
|
||||
|
||||
@@ -62,20 +109,9 @@ def prompt_node_key(graph: Any, node_id: Any) -> Optional[str]:
|
||||
if not isinstance(graph, dict) or node_id is None:
|
||||
return None
|
||||
start = str(node_id)
|
||||
closure: Dict[str, Any] = {}
|
||||
stack = [start]
|
||||
while stack:
|
||||
current = stack.pop()
|
||||
if current in closure:
|
||||
continue
|
||||
node = graph.get(current)
|
||||
if not isinstance(node, dict):
|
||||
return None
|
||||
if len(closure) >= MAX_VISITS:
|
||||
return None
|
||||
inputs = node.get("inputs")
|
||||
closure[current] = [node.get("class_type"), inputs]
|
||||
stack.extend(_links(inputs))
|
||||
closure = upstream_closure(graph, [start], MAX_VISITS)
|
||||
if closure is None:
|
||||
return None
|
||||
try:
|
||||
encoded = json.dumps(
|
||||
[RUN_KEY_VERSION, start, closure],
|
||||
|
||||
+18
-7
@@ -16,7 +16,9 @@ it never counts. Servers without a ``prompt_queue`` fall back to counting at
|
||||
submit time.
|
||||
|
||||
Text from known pure string nodes is resolved at queue time too. Anything else
|
||||
(PromptSearchList batches, unknown nodes) is counted by the node as it executes.
|
||||
(PromptSearchList batches, unknown nodes) is counted by the node as it executes,
|
||||
and when ComfyUI serves such a node from its cache, the job replays what that
|
||||
execution counted (``cached_runs``, #347).
|
||||
"""
|
||||
|
||||
import threading
|
||||
@@ -167,13 +169,22 @@ def queued_pair_runs(graph: Any) -> List[Tuple[str, str, str]]:
|
||||
|
||||
def find_pair_id(db: Any, positive: str, negative: str) -> Optional[int]:
|
||||
"""Find a saved pair without creating or overwriting its metadata."""
|
||||
positive_row = db.get_prompt_by_hash(generate_prompt_hash(positive))
|
||||
negative_row = (
|
||||
db.get_prompt_by_hash(generate_prompt_hash(negative))
|
||||
if negative.strip()
|
||||
else None
|
||||
return find_pair_id_by_hashes(
|
||||
db,
|
||||
generate_prompt_hash(positive),
|
||||
generate_prompt_hash(negative) if negative.strip() else None,
|
||||
)
|
||||
if not positive_row or (negative.strip() and not negative_row):
|
||||
|
||||
|
||||
def find_pair_id_by_hashes(
|
||||
db: Any, positive_hash: str, negative_hash: Optional[str]
|
||||
) -> Optional[int]:
|
||||
"""The saved pair of two prompt hashes; ``None`` negative is the empty one."""
|
||||
positive_row = db.get_prompt_by_hash(positive_hash)
|
||||
negative_row = (
|
||||
db.get_prompt_by_hash(negative_hash) if negative_hash is not None else None
|
||||
)
|
||||
if not positive_row or (negative_hash is not None and not negative_row):
|
||||
return None
|
||||
row = (
|
||||
db.model.get_connection()
|
||||
|
||||
Reference in New Issue
Block a user