fix: count prompt runs when ComfyUI dequeues the prompt, not when it is submitted
The on-prompt hook runs before ComfyUI validates the request, so counting there bumped run_count for prompts that were rejected or deleted from the queue. The hook now only resolves the positive PromptManager texts, mints a UUID4 prompt_id when the client sent none, and stashes the hashes under that id (QueuedRuns, bounded to 1000 entries). PromptQueue.get is wrapped idempotently; when a prompt is dequeued for execution its stashed runs are counted (or remembered as pending first use). Servers without a prompt_queue fall back to counting at submit time with one warning. __init__ defines _global_db before the monitoring block and passes a factory that raises RuntimeError while the database is unavailable; the package is now imported under test with server and folder_paths stubbed.
This commit is contained in:
Vito Sansevero
+187
-31
@@ -1,19 +1,30 @@
|
||||
"""Count prompt runs when a workflow is queued, not when nodes execute.
|
||||
"""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. The on-prompt hook sees every queued
|
||||
workflow instead: it counts a run for each PromptManager node feeding a sampler's
|
||||
positive input. Prompts not yet in the database are remembered as pending and
|
||||
counted by the node when it first saves them.
|
||||
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 functools
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from collections import OrderedDict
|
||||
from typing import Any, Tuple
|
||||
from typing import Any, Callable, List, Optional, Tuple
|
||||
|
||||
try:
|
||||
from .hashing import generate_prompt_hash
|
||||
@@ -28,13 +39,16 @@ 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:
|
||||
"""Queued runs of positive prompts that were not in the database yet.
|
||||
"""Dequeued runs of positive prompts that were not in the database yet.
|
||||
|
||||
Counts per hash: the same new prompt can be queued several times before its
|
||||
first job runs, and later jobs may reuse ComfyUI's cached node.
|
||||
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__(
|
||||
@@ -61,47 +75,189 @@ class PendingFirstUse:
|
||||
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:
|
||||
if not runs:
|
||||
return
|
||||
with self._lock:
|
||||
self._entries.pop(prompt_id, None)
|
||||
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()
|
||||
|
||||
|
||||
def handle_queued_prompt(json_data: Any, db: Any, pending: PendingFirstUse) -> Any:
|
||||
"""ComfyUI on-prompt handler body: count runs, always return the request as-is."""
|
||||
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 in run_prompt_nodes(graph):
|
||||
text = 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 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) -> 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:
|
||||
graph = json_data.get("prompt") if isinstance(json_data, dict) else None
|
||||
counted = set()
|
||||
for node_id in run_prompt_nodes(graph):
|
||||
text = resolve_text(graph, node_id)
|
||||
if not text:
|
||||
continue
|
||||
prompt_hash = generate_prompt_hash(text)
|
||||
if prompt_hash in counted:
|
||||
continue
|
||||
counted.add(prompt_hash)
|
||||
existing = db.get_prompt_by_hash(prompt_hash)
|
||||
if existing:
|
||||
db.record_prompt_use(existing["id"])
|
||||
else:
|
||||
pending.add(prompt_hash)
|
||||
runs = queued_prompt_runs(graph)
|
||||
if not runs:
|
||||
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)
|
||||
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,
|
||||
) -> None:
|
||||
"""Count the runs stashed for a prompt ComfyUI is about to execute."""
|
||||
try:
|
||||
runs = queued.pop(str(prompt_id))
|
||||
if runs:
|
||||
count_runs(runs, db_factory(), 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:
|
||||
graph = json_data.get("prompt") if isinstance(json_data, dict) else None
|
||||
runs = queued_prompt_runs(graph)
|
||||
if runs:
|
||||
count_runs(runs, db_factory(), pending)
|
||||
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(original, _WRAPPED_MARKER, False):
|
||||
return False
|
||||
|
||||
@functools.wraps(original)
|
||||
def get(*args, **kwargs):
|
||||
result = original(*args, **kwargs)
|
||||
if result is not None:
|
||||
try:
|
||||
on_dequeued(str(result[0][1]))
|
||||
except Exception as e:
|
||||
logger.warning(f"Dequeue hook failed: {e}")
|
||||
return result
|
||||
|
||||
setattr(get, _WRAPPED_MARKER, True)
|
||||
prompt_queue.get = get
|
||||
return True
|
||||
|
||||
|
||||
_hook_registered = threading.Event()
|
||||
_register_lock = threading.Lock()
|
||||
|
||||
|
||||
def register_queue_hook(
|
||||
server: Any, db_factory: Any, pending: PendingFirstUse = PENDING_FIRST_USE
|
||||
server: Any,
|
||||
db_factory: Callable[[], Any],
|
||||
pending: PendingFirstUse = PENDING_FIRST_USE,
|
||||
queued: Optional[QueuedRuns] = None,
|
||||
) -> bool:
|
||||
"""Register the on-prompt handler once; db_factory is called lazily per request."""
|
||||
"""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
|
||||
server.add_on_prompt_handler(
|
||||
lambda json_data: handle_queued_prompt(json_data, db_factory(), pending)
|
||||
)
|
||||
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()
|
||||
logger.info("Registered prompt usage hook")
|
||||
return True
|
||||
|
||||
Reference in New Issue
Block a user