Vito
763e6b37b1
* feat(history): normalize bounded execution evidence (#209) * fix(history): bound lifecycle kind checks before hashing (#209) * feat(history): persist execution attempts with restore fencing (#209) * test(api): bootstrap standalone backup restore fixtures (#209) * feat(history): isolate bounded asynchronous ledger writes (#209) * feat(history): observe actual queue attempts without altering execution (#209) * fix(database): honor package-scoped configuration (#211) * test(history): exercise real cached and headless executions (#209) * fix(history): reject unsupported uniqueness and isolate package fixtures (#209) * test(history): support Python 3.10 package fixture cleanup (#209)
356 lines
13 KiB
Python
356 lines
13 KiB
Python
"""Count prompt runs when ComfyUI dequeues a workflow, not when nodes execute.
|
|
|
|
ComfyUI skips nodes whose inputs are unchanged, so a PromptManager node does not
|
|
execute when the same prompt is re-run. Instead:
|
|
|
|
1. The on-prompt hook sees every submitted workflow before validation. It resolves
|
|
the text of each PromptManager node feeding a sampler's positive input and
|
|
stashes the hashes under the request's prompt_id (minting one when the client
|
|
sent none; ComfyUI keeps a client-supplied id).
|
|
2. ``PromptQueue.get`` is wrapped. When ComfyUI dequeues a prompt for execution,
|
|
its stashed hashes are counted: one run per prompt already in the database,
|
|
otherwise a pending first use that the node consumes when it first saves it.
|
|
|
|
A prompt that fails validation or is deleted from the queue is never dequeued, so
|
|
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.
|
|
"""
|
|
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from collections import OrderedDict
|
|
from typing import Any, Callable, List, Optional, Tuple
|
|
|
|
try:
|
|
from .hashing import generate_prompt_hash
|
|
from .logging_config import get_logger
|
|
from .prompt_graph import intrinsic_pair_nodes, resolve_text, run_prompt_components
|
|
from .prompt_pair_resolver import resolve_text_input
|
|
from .pair_execution import clear_pair_execution
|
|
from .queue_observers import subscribe_dequeue
|
|
except ImportError:
|
|
from utils.hashing import generate_prompt_hash
|
|
from utils.logging_config import get_logger
|
|
from utils.prompt_graph import (
|
|
intrinsic_pair_nodes,
|
|
resolve_text,
|
|
run_prompt_components,
|
|
)
|
|
from utils.prompt_pair_resolver import resolve_text_input
|
|
from utils.pair_execution import clear_pair_execution
|
|
from utils.queue_observers import subscribe_dequeue
|
|
|
|
logger = get_logger("prompt_manager.usage_tracking")
|
|
|
|
PENDING_TTL_SECONDS = 3600 # a queued prompt that never executes is forgotten
|
|
PENDING_MAX_ENTRIES = 1000
|
|
QUEUED_MAX_ENTRIES = 1000 # prompts submitted but not yet dequeued
|
|
|
|
Run = Tuple[str, str] # (prompt hash, resolved text)
|
|
|
|
|
|
class PendingFirstUse:
|
|
"""Dequeued runs of positive prompts that were not in the database yet.
|
|
|
|
Counts per hash: the same new prompt can be run several times before its
|
|
first job saves it, and later jobs may reuse ComfyUI's cached node.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
ttl_seconds: float = PENDING_TTL_SECONDS,
|
|
max_entries: int = PENDING_MAX_ENTRIES,
|
|
):
|
|
self._ttl = ttl_seconds
|
|
self._max = max_entries
|
|
self._entries: OrderedDict[str, Tuple[int, float]] = OrderedDict()
|
|
self._lock = threading.Lock()
|
|
|
|
def add(self, prompt_hash: str) -> None:
|
|
with self._lock:
|
|
count, _ = self._entries.pop(prompt_hash, (0, 0.0))
|
|
self._entries[prompt_hash] = (count + 1, time.monotonic())
|
|
while len(self._entries) > self._max:
|
|
self._entries.popitem(last=False)
|
|
|
|
def consume(self, prompt_hash: str) -> int:
|
|
"""Number of recent queued runs not yet counted (0 if none); clears them."""
|
|
with self._lock:
|
|
count, added = self._entries.pop(prompt_hash, (0, 0.0))
|
|
return count if count and time.monotonic() - added <= self._ttl else 0
|
|
|
|
|
|
class QueuedRuns:
|
|
"""Runs resolved at submit time, keyed by prompt_id until ComfyUI dequeues them.
|
|
|
|
Bounded: a prompt rejected by validation or deleted from the queue is never
|
|
dequeued, so the oldest entries are evicted once the bound is reached.
|
|
"""
|
|
|
|
def __init__(self, max_entries: int = QUEUED_MAX_ENTRIES):
|
|
self._max = max_entries
|
|
self._entries: OrderedDict[str, List[Run]] = OrderedDict()
|
|
self._lock = threading.Lock()
|
|
|
|
def put(self, prompt_id: str, runs: List[Run]) -> None:
|
|
with self._lock:
|
|
self._entries.pop(prompt_id, None)
|
|
if not runs:
|
|
return
|
|
self._entries[prompt_id] = list(runs)
|
|
while len(self._entries) > self._max:
|
|
self._entries.popitem(last=False)
|
|
|
|
def pop(self, prompt_id: str) -> List[Run]:
|
|
with self._lock:
|
|
return self._entries.pop(prompt_id, [])
|
|
|
|
def __len__(self) -> int:
|
|
with self._lock:
|
|
return len(self._entries)
|
|
|
|
|
|
# Shared between the queue hook and the nodes
|
|
PENDING_FIRST_USE = PendingFirstUse()
|
|
QUEUED_RUNS = QueuedRuns()
|
|
PENDING_PAIR_FIRST_USE = PendingFirstUse()
|
|
QUEUED_PAIR_RUNS = QueuedRuns()
|
|
|
|
|
|
def queued_prompt_runs(graph: Any) -> List[Run]:
|
|
"""(hash, text) of each positive PromptManager node whose text is knowable.
|
|
|
|
Each hash once: the same text in two positive nodes is one run of the prompt.
|
|
"""
|
|
runs: List[Run] = []
|
|
seen = set()
|
|
for node_id, field in run_prompt_components(graph):
|
|
node = graph.get(node_id, {}) if isinstance(graph, dict) else {}
|
|
text = (
|
|
resolve_text_input(graph, node_id, field)
|
|
if node.get("class_type") == "PromptManagerPair"
|
|
else resolve_text(graph, node_id)
|
|
)
|
|
if not text:
|
|
continue
|
|
prompt_hash = generate_prompt_hash(text)
|
|
if prompt_hash in seen:
|
|
continue
|
|
seen.add(prompt_hash)
|
|
runs.append((prompt_hash, text))
|
|
return runs
|
|
|
|
|
|
def pair_run_key(positive: str, negative: str) -> str:
|
|
"""Stable ordered identity; an empty negative remains an explicit component."""
|
|
return generate_prompt_hash(positive) + ":" + generate_prompt_hash(negative)
|
|
|
|
|
|
def queued_pair_runs(graph: Any) -> List[Tuple[str, str, str]]:
|
|
"""Resolve raw paired recipes only when both outputs form the used pair."""
|
|
found, seen = [], set()
|
|
for node_id in intrinsic_pair_nodes(graph):
|
|
positive = resolve_text_input(graph, node_id, "text")
|
|
negative = resolve_text_input(graph, node_id, "negative_text")
|
|
if positive is None or not positive.strip() or negative is None:
|
|
continue
|
|
key = pair_run_key(positive, negative)
|
|
if key not in seen:
|
|
seen.add(key)
|
|
found.append((key, positive, negative))
|
|
return found
|
|
|
|
|
|
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
|
|
)
|
|
if not positive_row or (negative.strip() and not negative_row):
|
|
return None
|
|
row = (
|
|
db.model.get_connection()
|
|
.execute(
|
|
"SELECT id FROM prompt_pairs "
|
|
"WHERE positive_prompt_id=? AND negative_prompt_id IS ?",
|
|
(positive_row["id"], negative_row["id"] if negative_row else None),
|
|
)
|
|
.fetchone()
|
|
)
|
|
return row["id"] if row else None
|
|
|
|
|
|
def count_pair_runs(
|
|
runs: List[Tuple[str, str, str]],
|
|
db: Any,
|
|
pending: PendingFirstUse,
|
|
) -> None:
|
|
"""Count existing recipes at dequeue and defer new recipes until first save."""
|
|
for key, positive, negative in runs:
|
|
pair_id = find_pair_id(db, positive, negative)
|
|
if pair_id:
|
|
db.record_prompt_pair_use(pair_id)
|
|
else:
|
|
pending.add(key)
|
|
|
|
|
|
def count_runs(runs: List[Run], db: Any, pending: PendingFirstUse) -> None:
|
|
"""Record one use per prompt in the database; remember the rest for first save."""
|
|
for prompt_hash, _text in runs:
|
|
existing = db.get_prompt_by_hash(prompt_hash)
|
|
if existing:
|
|
db.record_prompt_use(existing["id"])
|
|
else:
|
|
pending.add(prompt_hash)
|
|
|
|
|
|
def handle_queued_prompt(
|
|
json_data: Any,
|
|
queued: QueuedRuns = QUEUED_RUNS,
|
|
*,
|
|
pair_queued: QueuedRuns = QUEUED_PAIR_RUNS,
|
|
) -> Any:
|
|
"""ComfyUI on-prompt handler body: stash the runs, always return the request.
|
|
|
|
Adds a canonical UUID4 ``prompt_id`` when the client sent none, so the stash
|
|
can be matched when ComfyUI dequeues the prompt. The database is not touched:
|
|
the request has not been validated yet.
|
|
"""
|
|
try:
|
|
clear_pair_execution(json_data)
|
|
graph = json_data.get("prompt") if isinstance(json_data, dict) else None
|
|
runs = queued_prompt_runs(graph)
|
|
pair_runs = queued_pair_runs(graph)
|
|
if not runs and not pair_runs:
|
|
if isinstance(json_data, dict) and json_data.get("prompt_id") is not None:
|
|
queued.put(str(json_data["prompt_id"]), [])
|
|
pair_queued.put(str(json_data["prompt_id"]), [])
|
|
return json_data
|
|
prompt_id = json_data.get("prompt_id")
|
|
if prompt_id is None:
|
|
prompt_id = str(uuid.uuid4())
|
|
json_data["prompt_id"] = prompt_id
|
|
queued.put(str(prompt_id), runs)
|
|
pair_queued.put(str(prompt_id), pair_runs)
|
|
except Exception as e:
|
|
logger.warning(f"Could not resolve queued prompt runs: {e}")
|
|
return json_data
|
|
|
|
|
|
def on_dequeue(
|
|
prompt_id: str,
|
|
queued: QueuedRuns,
|
|
db_factory: Callable[[], Any],
|
|
pending: PendingFirstUse,
|
|
*,
|
|
pair_queued: QueuedRuns = QUEUED_PAIR_RUNS,
|
|
pair_pending: PendingFirstUse = PENDING_PAIR_FIRST_USE,
|
|
) -> None:
|
|
"""Count the runs stashed for a prompt ComfyUI is about to execute."""
|
|
try:
|
|
runs = queued.pop(str(prompt_id))
|
|
pair_runs = pair_queued.pop(str(prompt_id))
|
|
if runs or pair_runs:
|
|
db = db_factory()
|
|
count_runs(runs, db, pending)
|
|
count_pair_runs(pair_runs, db, pair_pending)
|
|
except Exception as e:
|
|
logger.warning(f"Could not count runs of dequeued prompt {prompt_id}: {e}")
|
|
|
|
|
|
def count_queued_prompt(
|
|
json_data: Any, db_factory: Callable[[], Any], pending: PendingFirstUse
|
|
) -> Any:
|
|
"""Legacy on-prompt handler body: count at submit time (no prompt_queue)."""
|
|
try:
|
|
clear_pair_execution(json_data)
|
|
graph = json_data.get("prompt") if isinstance(json_data, dict) else None
|
|
runs = queued_prompt_runs(graph)
|
|
pair_runs = queued_pair_runs(graph)
|
|
if runs or pair_runs:
|
|
db = db_factory()
|
|
count_runs(runs, db, pending)
|
|
count_pair_runs(pair_runs, db, PENDING_PAIR_FIRST_USE)
|
|
except Exception as e:
|
|
logger.warning(f"Could not count queued prompt runs: {e}")
|
|
return json_data
|
|
|
|
|
|
_WRAPPED_MARKER = "_prompt_manager_dequeue_hook"
|
|
|
|
|
|
def _wrap_queue_get(prompt_queue: Any, on_dequeued: Callable[[str], None]) -> bool:
|
|
"""Wrap ``prompt_queue.get`` so each dequeued prompt_id is reported once.
|
|
|
|
The wrapper forwards every argument and result unchanged. Idempotent: a queue
|
|
whose ``get`` is already wrapped is left alone.
|
|
"""
|
|
original = getattr(prompt_queue, "get", None)
|
|
if not callable(original):
|
|
return False
|
|
if getattr(prompt_queue, _WRAPPED_MARKER, False) or getattr(
|
|
original, _WRAPPED_MARKER, False
|
|
):
|
|
return False
|
|
|
|
def report(result: object) -> None:
|
|
try:
|
|
on_dequeued(str(result[0][1]))
|
|
except Exception as e:
|
|
logger.warning(f"Dequeue hook failed: {e}")
|
|
|
|
subscribe_dequeue(prompt_queue, "usage_tracking", report)
|
|
setattr(prompt_queue.get, _WRAPPED_MARKER, True)
|
|
setattr(prompt_queue, _WRAPPED_MARKER, True)
|
|
return True
|
|
|
|
|
|
_hook_registered = threading.Event()
|
|
_register_lock = threading.Lock()
|
|
|
|
|
|
def register_queue_hook(
|
|
server: Any,
|
|
db_factory: Callable[[], Any],
|
|
pending: PendingFirstUse = PENDING_FIRST_USE,
|
|
queued: Optional[QueuedRuns] = None,
|
|
) -> bool:
|
|
"""Register the on-prompt handler once and wrap the prompt queue's ``get``.
|
|
|
|
``db_factory`` is called lazily, inside the guarded handlers, so a database
|
|
that is unavailable never raises into ComfyUI. Without a usable
|
|
``server.prompt_queue`` runs are counted at submit time instead.
|
|
"""
|
|
if queued is None:
|
|
queued = QUEUED_RUNS
|
|
with _register_lock:
|
|
if _hook_registered.is_set():
|
|
return False
|
|
prompt_queue = getattr(server, "prompt_queue", None)
|
|
if callable(getattr(prompt_queue, "get", None)):
|
|
server.add_on_prompt_handler(
|
|
lambda json_data: handle_queued_prompt(json_data, queued)
|
|
)
|
|
_wrap_queue_get(
|
|
prompt_queue,
|
|
lambda prompt_id: on_dequeue(prompt_id, queued, db_factory, pending),
|
|
)
|
|
logger.info("Registered prompt usage hook (counting at dequeue)")
|
|
else:
|
|
logger.warning(
|
|
"PromptServer has no prompt_queue; counting prompt runs at queue "
|
|
"time, so rejected prompts may be counted"
|
|
)
|
|
server.add_on_prompt_handler(
|
|
lambda json_data: count_queued_prompt(json_data, db_factory, pending)
|
|
)
|
|
_hook_registered.set()
|
|
return True
|