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
@@ -1,14 +1,18 @@
|
||||
"""Run counting and node roles: only positive PromptManager nodes count and link.
|
||||
|
||||
Runs are counted when a prompt is queued (ComfyUI's on-prompt hook), because
|
||||
ComfyUI skips unchanged nodes, so node execution can't be relied on to count.
|
||||
Runs are resolved when a prompt is queued (ComfyUI's on-prompt hook, which runs
|
||||
before validation) and counted when ComfyUI dequeues the prompt for execution,
|
||||
because ComfyUI skips unchanged nodes, so node execution can't be relied on to
|
||||
count, and a rejected or cancelled queue entry must not count either.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
import unittest.mock as mock
|
||||
import uuid
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
@@ -16,7 +20,16 @@ from database.operations import PromptDatabase
|
||||
from prompt_manager_base import PromptManagerBase
|
||||
from utils.hashing import generate_prompt_hash
|
||||
from utils.prompt_tracker import PromptTracker
|
||||
from utils.usage_tracking import PendingFirstUse, handle_queued_prompt
|
||||
from utils import usage_tracking
|
||||
from utils.usage_tracking import (
|
||||
PendingFirstUse,
|
||||
QueuedRuns,
|
||||
_wrap_queue_get,
|
||||
count_queued_prompt,
|
||||
handle_queued_prompt,
|
||||
on_dequeue,
|
||||
register_queue_hook,
|
||||
)
|
||||
|
||||
LOADER = {
|
||||
"class_type": "CheckpointLoaderSimple",
|
||||
@@ -41,14 +54,52 @@ def workflow(positive_text="a cat", negative_text="blurry"):
|
||||
}
|
||||
|
||||
|
||||
class FakePromptQueue:
|
||||
"""Stand-in for execution.PromptQueue: get() pops preset results."""
|
||||
|
||||
def __init__(self):
|
||||
self.results = []
|
||||
self.get_calls = []
|
||||
|
||||
def get(self, timeout=None):
|
||||
self.get_calls.append(timeout)
|
||||
return self.results.pop(0) if self.results else None
|
||||
|
||||
|
||||
class FakeServer:
|
||||
"""Stand-in for PromptServer: on-prompt handlers plus an optional queue."""
|
||||
|
||||
def __init__(self, with_queue=True):
|
||||
self.handlers = []
|
||||
if with_queue:
|
||||
self.prompt_queue = FakePromptQueue()
|
||||
|
||||
def add_on_prompt_handler(self, handler):
|
||||
self.handlers.append(handler)
|
||||
|
||||
def on_prompt(self, json_data):
|
||||
for handler in self.handlers:
|
||||
json_data = handler(json_data)
|
||||
return json_data
|
||||
|
||||
|
||||
def queue_item(prompt_id, graph):
|
||||
"""What PromptQueue.get returns for a prompt that passed validation."""
|
||||
return ((0, prompt_id, graph, {}, [], False), 1)
|
||||
|
||||
|
||||
class DbTestCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
fd, self.path = tempfile.mkstemp(suffix=".db")
|
||||
os.close(fd)
|
||||
self.db = PromptDatabase(self.path)
|
||||
self.pending = PendingFirstUse()
|
||||
self.queued = QueuedRuns()
|
||||
|
||||
def tearDown(self):
|
||||
close = getattr(self.db, "close", None)
|
||||
if callable(close):
|
||||
close()
|
||||
for suffix in ("", "-wal", "-shm"):
|
||||
if os.path.exists(self.path + suffix):
|
||||
os.unlink(self.path + suffix)
|
||||
@@ -60,9 +111,12 @@ class DbTestCase(unittest.TestCase):
|
||||
return self.db.get_prompt_by_id(prompt_id)["run_count"]
|
||||
|
||||
def _queue(self, graph):
|
||||
return handle_queued_prompt(
|
||||
{"prompt": graph, "client_id": "x"}, self.db, self.pending
|
||||
"""Queue a graph and let ComfyUI dequeue it for execution."""
|
||||
request = handle_queued_prompt({"prompt": graph, "client_id": "x"}, self.queued)
|
||||
on_dequeue(
|
||||
request.get("prompt_id", ""), self.queued, lambda: self.db, self.pending
|
||||
)
|
||||
return request
|
||||
|
||||
|
||||
class TestQueueHook(DbTestCase):
|
||||
@@ -121,14 +175,243 @@ class TestQueueHook(DbTestCase):
|
||||
self._queue(graph)
|
||||
self.assertEqual(self._runs(positive), 0)
|
||||
|
||||
def test_returns_the_request_unchanged_and_never_raises(self):
|
||||
def test_returns_the_request_and_never_raises(self):
|
||||
request = {"prompt": workflow(), "client_id": "x", "extra_data": {}}
|
||||
self.assertIs(handle_queued_prompt(request, self.db, self.pending), request)
|
||||
self.assertIs(handle_queued_prompt(request, self.queued), request)
|
||||
for bad in (None, {}, {"prompt": "x"}, {"prompt": None}, []):
|
||||
self.assertIs(handle_queued_prompt(bad, self.db, self.pending), bad)
|
||||
self.assertIs(handle_queued_prompt(bad, self.queued), bad)
|
||||
broken_db = mock.Mock()
|
||||
broken_db.get_prompt_by_hash.side_effect = RuntimeError("db locked")
|
||||
self.assertIs(handle_queued_prompt(request, broken_db, self.pending), request)
|
||||
with self.assertLogs("prompt_manager.usage_tracking", level="WARNING"):
|
||||
on_dequeue(
|
||||
request["prompt_id"], self.queued, lambda: broken_db, self.pending
|
||||
)
|
||||
|
||||
def test_graph_without_prompt_manager_nodes_is_not_stashed_or_tagged(self):
|
||||
request = {"prompt": {"20": LOADER}}
|
||||
self.assertIs(handle_queued_prompt(request, self.queued), request)
|
||||
self.assertNotIn("prompt_id", request)
|
||||
self.assertEqual(len(self.queued), 0)
|
||||
|
||||
|
||||
class TestCountAtDequeue(DbTestCase):
|
||||
"""End to end through the registered hook and the wrapped queue."""
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self.server = FakeServer()
|
||||
self._register(self.server)
|
||||
|
||||
def _register(self, server):
|
||||
with mock.patch.object(usage_tracking, "_hook_registered", threading.Event()):
|
||||
return register_queue_hook(
|
||||
server, lambda: self.db, pending=self.pending, queued=self.queued
|
||||
)
|
||||
|
||||
def _submit(self, graph, prompt_id=None):
|
||||
request = {"prompt": graph, "client_id": "x"}
|
||||
if prompt_id is not None:
|
||||
request["prompt_id"] = prompt_id
|
||||
return self.server.on_prompt(request)
|
||||
|
||||
def _dequeue(self, request):
|
||||
self.server.prompt_queue.results.append(
|
||||
queue_item(request["prompt_id"], request["prompt"])
|
||||
)
|
||||
return self.server.prompt_queue.get(timeout=1.0)
|
||||
|
||||
def test_hook_adds_a_lowercase_uuid4_prompt_id_when_absent(self):
|
||||
request = self._submit(workflow())
|
||||
prompt_id = request["prompt_id"]
|
||||
self.assertEqual(str(uuid.UUID(prompt_id, version=4)), prompt_id)
|
||||
self.assertEqual(prompt_id, prompt_id.lower())
|
||||
|
||||
def test_hook_keeps_an_existing_prompt_id(self):
|
||||
request = self._submit(workflow(), prompt_id="client-chosen")
|
||||
self.assertEqual(request["prompt_id"], "client-chosen")
|
||||
|
||||
def test_hook_alone_does_not_count(self):
|
||||
positive = self._save("a cat")
|
||||
self._submit(workflow())
|
||||
self.assertEqual(self._runs(positive), 0)
|
||||
|
||||
def test_dequeue_counts_exactly_once(self):
|
||||
positive = self._save("a cat")
|
||||
request = self._submit(workflow())
|
||||
result = self._dequeue(request)
|
||||
self.assertEqual(result[0][1], request["prompt_id"])
|
||||
self.assertEqual(self.server.prompt_queue.get_calls, [1.0])
|
||||
self.assertEqual(self._runs(positive), 1)
|
||||
self._dequeue(request) # a second dequeue of the same id has nothing left
|
||||
self.assertEqual(self._runs(positive), 1)
|
||||
|
||||
def test_two_queued_prompts_are_each_counted_once(self):
|
||||
positive = self._save("a cat")
|
||||
first = self._submit(workflow())
|
||||
second = self._submit(workflow())
|
||||
self.assertNotEqual(first["prompt_id"], second["prompt_id"])
|
||||
self._dequeue(first)
|
||||
self.assertEqual(self._runs(positive), 1)
|
||||
self._dequeue(second)
|
||||
self.assertEqual(self._runs(positive), 2)
|
||||
|
||||
def test_rejected_prompt_is_never_counted(self):
|
||||
positive = self._save("a cat")
|
||||
self._submit(workflow()) # validation fails: never dequeued
|
||||
other = self._submit(workflow())
|
||||
self._dequeue(other)
|
||||
self.assertEqual(self._runs(positive), 1)
|
||||
|
||||
def test_get_returning_none_is_a_noop(self):
|
||||
positive = self._save("a cat")
|
||||
self._submit(workflow())
|
||||
self.assertIsNone(self.server.prompt_queue.get(timeout=0.1))
|
||||
self.assertEqual(self._runs(positive), 0)
|
||||
|
||||
def test_unknown_text_is_pending_only_after_dequeue(self):
|
||||
prompt_hash = generate_prompt_hash("never saved")
|
||||
request = self._submit(workflow(positive_text="never saved"))
|
||||
self.assertEqual(self.pending.consume(prompt_hash), 0)
|
||||
self._dequeue(request)
|
||||
self.assertEqual(self.pending.consume(prompt_hash), 1)
|
||||
|
||||
def test_shorter_and_longer_queue_items_are_tolerated(self):
|
||||
positive = self._save("a cat")
|
||||
request = self._submit(workflow())
|
||||
self.server.prompt_queue.results.append(
|
||||
((0, request["prompt_id"], request["prompt"], {}, []), 1)
|
||||
)
|
||||
self.server.prompt_queue.get()
|
||||
self.assertEqual(self._runs(positive), 1)
|
||||
request = self._submit(workflow())
|
||||
self.server.prompt_queue.results.append(
|
||||
((0, request["prompt_id"], request["prompt"], {}, [], False, "x"), 2)
|
||||
)
|
||||
self.server.prompt_queue.get()
|
||||
self.assertEqual(self._runs(positive), 2)
|
||||
|
||||
def test_malformed_queue_item_is_logged_and_returned(self):
|
||||
item = (("not", "a", "tuple"),)
|
||||
self.server.prompt_queue.results.append(item)
|
||||
self.server.prompt_queue.results.append("garbage")
|
||||
with self.assertLogs("prompt_manager.usage_tracking", level="WARNING"):
|
||||
self.assertEqual(self.server.prompt_queue.get(), item)
|
||||
self.assertEqual(self.server.prompt_queue.get(), "garbage")
|
||||
|
||||
def test_db_factory_raising_at_dequeue_is_logged_not_raised(self):
|
||||
def broken():
|
||||
raise RuntimeError("PromptManager database unavailable")
|
||||
|
||||
queued = QueuedRuns()
|
||||
server = FakeServer()
|
||||
with mock.patch.object(usage_tracking, "_hook_registered", threading.Event()):
|
||||
register_queue_hook(server, broken, pending=self.pending, queued=queued)
|
||||
request = server.on_prompt({"prompt": workflow()})
|
||||
server.prompt_queue.results.append(queue_item(request["prompt_id"], {}))
|
||||
with self.assertLogs("prompt_manager.usage_tracking", level="WARNING") as logs:
|
||||
self.assertIsNotNone(server.prompt_queue.get())
|
||||
self.assertIn("database unavailable", "\n".join(logs.output))
|
||||
|
||||
def test_server_without_prompt_queue_counts_at_queue_time_and_warns_once(self):
|
||||
positive = self._save("a cat")
|
||||
server = FakeServer(with_queue=False)
|
||||
with self.assertLogs("prompt_manager.usage_tracking", level="WARNING") as logs:
|
||||
self.assertTrue(self._register(server))
|
||||
warnings = [line for line in logs.output if line.startswith("WARNING")]
|
||||
self.assertEqual(len(warnings), 1)
|
||||
server.on_prompt({"prompt": workflow()})
|
||||
self.assertEqual(self._runs(positive), 1)
|
||||
server.on_prompt({"prompt": workflow(positive_text="new text")})
|
||||
self.assertEqual(self.pending.consume(generate_prompt_hash("new text")), 1)
|
||||
|
||||
def test_server_whose_queue_has_no_get_uses_legacy_counting(self):
|
||||
positive = self._save("a cat")
|
||||
server = FakeServer(with_queue=False)
|
||||
server.prompt_queue = object()
|
||||
with self.assertLogs("prompt_manager.usage_tracking", level="WARNING"):
|
||||
self._register(server)
|
||||
server.on_prompt({"prompt": workflow()})
|
||||
self.assertEqual(self._runs(positive), 1)
|
||||
|
||||
def test_legacy_counting_calls_db_factory_inside_the_guard(self):
|
||||
def broken():
|
||||
raise RuntimeError("PromptManager database unavailable")
|
||||
|
||||
request = {"prompt": workflow()}
|
||||
with self.assertLogs("prompt_manager.usage_tracking", level="WARNING"):
|
||||
self.assertIs(count_queued_prompt(request, broken, self.pending), request)
|
||||
|
||||
def test_wrapping_twice_does_not_double_count(self):
|
||||
positive = self._save("a cat")
|
||||
prompt_queue = self.server.prompt_queue
|
||||
wrapped_get = prompt_queue.get
|
||||
calls = []
|
||||
self.assertFalse(_wrap_queue_get(prompt_queue, calls.append))
|
||||
self.assertIs(prompt_queue.get, wrapped_get)
|
||||
request = self._submit(workflow())
|
||||
self._dequeue(request)
|
||||
self.assertEqual(self._runs(positive), 1)
|
||||
self.assertEqual(calls, [])
|
||||
|
||||
def test_wrap_rejects_queues_without_a_callable_get(self):
|
||||
self.assertFalse(_wrap_queue_get(object(), lambda prompt_id: None))
|
||||
self.assertFalse(_wrap_queue_get(None, lambda prompt_id: None))
|
||||
|
||||
def test_wrapper_reports_the_dequeued_prompt_id(self):
|
||||
prompt_queue = FakePromptQueue()
|
||||
seen = []
|
||||
self.assertTrue(_wrap_queue_get(prompt_queue, seen.append))
|
||||
prompt_queue.results.append(queue_item("abc", {}))
|
||||
self.assertEqual(prompt_queue.get(timeout=2.5), queue_item("abc", {}))
|
||||
self.assertEqual(prompt_queue.get_calls, [2.5])
|
||||
self.assertEqual(seen, ["abc"])
|
||||
|
||||
def test_dequeue_callback_errors_are_logged_not_raised(self):
|
||||
prompt_queue = FakePromptQueue()
|
||||
|
||||
def explode(prompt_id):
|
||||
raise ValueError("boom")
|
||||
|
||||
_wrap_queue_get(prompt_queue, explode)
|
||||
prompt_queue.results.append(queue_item("abc", {}))
|
||||
with self.assertLogs("prompt_manager.usage_tracking", level="WARNING"):
|
||||
self.assertIsNotNone(prompt_queue.get())
|
||||
|
||||
|
||||
class TestQueuedRuns(unittest.TestCase):
|
||||
def test_pop_returns_runs_once(self):
|
||||
queued = QueuedRuns()
|
||||
queued.put("p1", [("h1", "text one")])
|
||||
self.assertEqual(queued.pop("p1"), [("h1", "text one")])
|
||||
self.assertEqual(queued.pop("p1"), [])
|
||||
self.assertEqual(queued.pop("missing"), [])
|
||||
|
||||
def test_bound_evicts_oldest(self):
|
||||
queued = QueuedRuns(max_entries=3)
|
||||
for i in range(4):
|
||||
queued.put(f"p{i}", [(f"h{i}", f"text {i}")])
|
||||
self.assertEqual(len(queued), 3)
|
||||
self.assertEqual(queued.pop("p0"), [])
|
||||
self.assertEqual(queued.pop("p3"), [("h3", "text 3")])
|
||||
|
||||
def test_default_bound_is_one_thousand(self):
|
||||
queued = QueuedRuns()
|
||||
for i in range(1001):
|
||||
queued.put(str(i), [("h", "t")])
|
||||
self.assertEqual(len(queued), 1000)
|
||||
self.assertEqual(queued.pop("0"), [])
|
||||
self.assertEqual(queued.pop("1000"), [("h", "t")])
|
||||
|
||||
def test_requeue_of_the_same_id_replaces_the_runs(self):
|
||||
queued = QueuedRuns()
|
||||
queued.put("p1", [("h1", "one")])
|
||||
queued.put("p1", [("h2", "two")])
|
||||
self.assertEqual(queued.pop("p1"), [("h2", "two")])
|
||||
|
||||
def test_empty_runs_are_not_stored(self):
|
||||
queued = QueuedRuns()
|
||||
queued.put("p1", [])
|
||||
self.assertEqual(len(queued), 0)
|
||||
|
||||
|
||||
class TestPendingFirstUse(unittest.TestCase):
|
||||
|
||||
Reference in New Issue
Block a user