feat(history): observe actual queue attempts without altering execution (#209)
This commit is contained in:
Vito Sansevero
+28
@@ -158,6 +158,34 @@ except Exception as e:
|
||||
f"Failed to register usage hook: {e}"
|
||||
)
|
||||
|
||||
|
||||
def _history_db_factory() -> "PromptDatabase":
|
||||
"""Initialize history on its worker even if monitoring initialization failed."""
|
||||
if _global_db is not None:
|
||||
return _global_db
|
||||
from .database.operations import PromptDatabase
|
||||
|
||||
return PromptDatabase()
|
||||
|
||||
|
||||
_history_registration = None
|
||||
try:
|
||||
from .py.config import server_instance
|
||||
from .utils.execution_history import register_history_hook
|
||||
|
||||
_history_registration = register_history_hook(server_instance, _history_db_factory)
|
||||
except Exception:
|
||||
try:
|
||||
from .utils.logging_config import get_logger
|
||||
|
||||
get_logger("prompt_manager.init").error("Failed to register execution history")
|
||||
except Exception:
|
||||
import logging
|
||||
|
||||
logging.getLogger("prompt_manager.init").error(
|
||||
"Failed to register execution history"
|
||||
)
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
|
||||
|
||||
# Print startup message with loaded tools
|
||||
|
||||
@@ -0,0 +1,451 @@
|
||||
"""Actual dequeue and completion capture through the owned SQLite worker."""
|
||||
|
||||
from collections import namedtuple
|
||||
import logging
|
||||
import gc
|
||||
import contextlib
|
||||
import io
|
||||
import os
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import sys
|
||||
import types
|
||||
import weakref
|
||||
from types import SimpleNamespace
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
from database import PromptDatabase
|
||||
from utils import usage_tracking
|
||||
from utils.execution_history_writer import HistoryWriter
|
||||
|
||||
try:
|
||||
from utils.execution_history import register_history_hook
|
||||
except ImportError:
|
||||
register_history_hook = None
|
||||
|
||||
Status = namedtuple("ExecutionStatus", "status_str completed messages")
|
||||
JOB = "a0c26da3-eb26-405c-b143-c662a84c3112"
|
||||
|
||||
|
||||
class HostQueue:
|
||||
def __init__(self):
|
||||
self.results = []
|
||||
self.done_result = object()
|
||||
self.done_calls = []
|
||||
|
||||
def get(self, *args, **kwargs):
|
||||
return self.results.pop(0) if self.results else None
|
||||
|
||||
def task_done(self, *args, **kwargs):
|
||||
self.done_calls.append((args, kwargs))
|
||||
return self.done_result
|
||||
|
||||
|
||||
def item(item_id, job=JOB, graph=None, extra=None, sensitive=None):
|
||||
return ((0, job, graph or {}, extra or {}, [], sensitive or {}), item_id)
|
||||
|
||||
|
||||
def status(kind="execution_success", job=JOB):
|
||||
return Status("success", True, [(kind, {"prompt_id": job, "timestamp": 20})])
|
||||
|
||||
|
||||
class HistoryHooksTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.directory = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.directory.cleanup)
|
||||
self.db = PromptDatabase(os.path.join(self.directory.name, "history.db"))
|
||||
self.addCleanup(self.db.model.reset_schema_cache)
|
||||
self.addCleanup(self.db.model.close_all)
|
||||
self.queue = HostQueue()
|
||||
self.server = SimpleNamespace(prompt_queue=self.queue)
|
||||
self.registrations, self.releases = [], []
|
||||
self.addCleanup(self.cleanup_registrations)
|
||||
for name in ("execution_history_writer", "execution_history"):
|
||||
logger = logging.getLogger("prompt_manager." + name)
|
||||
previous = logger.handlers, logger.propagate, logger.level
|
||||
logger.handlers, logger.propagate, logger.level = (
|
||||
[logging.NullHandler()],
|
||||
False,
|
||||
logging.WARNING,
|
||||
)
|
||||
self.addCleanup(self.restore_logger, logger, previous)
|
||||
|
||||
def restore_logger(self, logger, previous):
|
||||
logger.handlers, logger.propagate, logger.level = previous
|
||||
|
||||
def cleanup_registrations(self):
|
||||
for release in self.releases:
|
||||
release.set()
|
||||
for registration in self.registrations:
|
||||
self.assertTrue(registration.dispose(3), "owned worker leaked")
|
||||
|
||||
def wait_for(self, predicate):
|
||||
deadline = time.monotonic() + 3
|
||||
while time.monotonic() < deadline:
|
||||
if predicate():
|
||||
return
|
||||
threading.Event().wait(0.005)
|
||||
self.fail("finite deadline expired waiting for recorded history")
|
||||
|
||||
def register(self, factory=None, ready=True):
|
||||
self.assertTrue(
|
||||
callable(register_history_hook), "history registration is absent"
|
||||
)
|
||||
registration = register_history_hook(self.server, factory or (lambda: self.db))
|
||||
self.assertIsNotNone(registration)
|
||||
self.registrations.append(registration)
|
||||
if ready:
|
||||
self.wait_for(lambda: registration.writer.status().accepting)
|
||||
return registration
|
||||
|
||||
def dequeue(self, item_id=0, **kwargs):
|
||||
result = item(item_id, **kwargs)
|
||||
self.queue.results.append(result)
|
||||
self.assertIs(self.queue.get(), result)
|
||||
return result
|
||||
|
||||
def rows(self):
|
||||
return [
|
||||
dict(row)
|
||||
for row in self.db.model.get_connection()
|
||||
.execute("SELECT * FROM execution_runs ORDER BY queue_item_id")
|
||||
.fetchall()
|
||||
]
|
||||
|
||||
def drain(self, registration):
|
||||
self.wait_for(lambda: registration.writer.status().pending == 0)
|
||||
|
||||
def test_package_history_initializes_when_monitor_and_usage_setup_fail(self):
|
||||
try:
|
||||
from tests.test_init_package import (
|
||||
close_every_database,
|
||||
fake_server_module,
|
||||
load_package,
|
||||
)
|
||||
except ImportError:
|
||||
from test_init_package import (
|
||||
close_every_database,
|
||||
fake_server_module,
|
||||
load_package,
|
||||
)
|
||||
name = "prompt_manager_pkg_history_independent"
|
||||
server = fake_server_module()
|
||||
server.PromptServer.instance.prompt_queue = self.queue
|
||||
server.PromptServer.instance.app = SimpleNamespace(middlewares=[])
|
||||
|
||||
def reject_usage(_):
|
||||
raise RuntimeError("usage registration unavailable")
|
||||
|
||||
server.PromptServer.instance.add_on_prompt_handler = reject_usage
|
||||
output = os.path.join(self.directory.name, "output")
|
||||
os.makedirs(output)
|
||||
monitor_module = types.ModuleType(name + ".utils.image_monitor")
|
||||
# Missing monitor import fails before shared database initialization.
|
||||
with (
|
||||
patch.dict(sys.modules, {monitor_module.__name__: monitor_module}),
|
||||
contextlib.redirect_stdout(io.StringIO()),
|
||||
self.assertLogs("prompt_manager.init", level="ERROR"),
|
||||
):
|
||||
package, _ = load_package(name, self.db.model.db_path, output, server)
|
||||
self.addCleanup(sys.modules.pop, name, None)
|
||||
self.addCleanup(close_every_database)
|
||||
registration = getattr(package, "_history_registration", None)
|
||||
self.assertIsNotNone(registration, "history setup depended on another hook")
|
||||
self.registrations.append(registration)
|
||||
self.addCleanup(registration.dispose)
|
||||
self.assertIsNone(package._global_db)
|
||||
self.wait_for(lambda: registration.writer.status().accepting)
|
||||
self.dequeue(0)
|
||||
self.queue.task_done(0, {}, status())
|
||||
self.drain(registration)
|
||||
self.assertEqual(self.rows()[0]["outcome"], "succeeded")
|
||||
|
||||
def test_distinct_actual_attempts_for_reused_job_and_cached_rerun(self):
|
||||
registration = self.register()
|
||||
self.dequeue(0)
|
||||
self.queue.task_done(0, {"outputs": {}}, status("execution_error"))
|
||||
self.dequeue(1)
|
||||
self.queue.task_done(item_id=1, history_result={}, status=status())
|
||||
self.drain(registration)
|
||||
rows = self.rows()
|
||||
self.assertEqual(len(rows), 2)
|
||||
self.assertEqual([r["job_id"] for r in rows], [JOB, JOB])
|
||||
self.assertEqual([r["outcome"] for r in rows], ["failed", "succeeded"])
|
||||
self.assertNotEqual(rows[0]["run_id"], rows[1]["run_id"])
|
||||
|
||||
def test_later_conflicting_completion_resolves_completed_stored_identity(self):
|
||||
registration = self.register()
|
||||
self.dequeue(0)
|
||||
self.queue.task_done(0, {}, status())
|
||||
self.drain(registration)
|
||||
self.queue.task_done(0, {}, status("execution_error"))
|
||||
self.drain(registration)
|
||||
self.assertEqual(self.rows()[0]["outcome"], "unknown")
|
||||
self.assertEqual(self.rows()[0]["reason"], "conflicting_terminal")
|
||||
|
||||
def test_repeated_registration_cleanup_and_reregistration_keep_session_and_wrappers(
|
||||
self,
|
||||
):
|
||||
before = self.queue.get
|
||||
calls = []
|
||||
|
||||
def unrelated_before(*a, **k):
|
||||
calls.append("before")
|
||||
return before(*a, **k)
|
||||
|
||||
self.queue.get = unrelated_before
|
||||
first = self.register()
|
||||
self.assertIs(register_history_hook(self.server, lambda: self.db), first)
|
||||
observed = self.queue.get
|
||||
|
||||
def unrelated_after(*a, **k):
|
||||
calls.append("after")
|
||||
return observed(*a, **k)
|
||||
|
||||
self.queue.get = unrelated_after
|
||||
self.dequeue(0)
|
||||
self.queue.task_done(0, {}, status())
|
||||
self.assertTrue(first.dispose())
|
||||
self.assertTrue(first.dispose())
|
||||
self.assertIs(self.queue.get, unrelated_after)
|
||||
second = self.register()
|
||||
self.dequeue(1)
|
||||
self.queue.task_done(1, {}, status())
|
||||
self.drain(second)
|
||||
self.assertEqual(len(self.rows()), 2)
|
||||
self.assertEqual(self.rows()[0]["session_id"], self.rows()[1]["session_id"])
|
||||
self.assertEqual(calls, ["after", "before"] * 2)
|
||||
|
||||
def test_concurrent_handle_capacity_never_reassigns_active_attempts(self):
|
||||
registration = self.register()
|
||||
self.queue.results.extend(item(index) for index in range(129))
|
||||
ready = threading.Barrier(130)
|
||||
errors = []
|
||||
|
||||
def dequeue():
|
||||
try:
|
||||
ready.wait(3)
|
||||
self.queue.get()
|
||||
except Exception as error:
|
||||
errors.append(error)
|
||||
|
||||
threads = [threading.Thread(target=dequeue) for _ in range(129)]
|
||||
for thread in threads:
|
||||
thread.start()
|
||||
ready.wait(3)
|
||||
for thread in threads:
|
||||
thread.join(3)
|
||||
self.assertFalse(thread.is_alive())
|
||||
self.assertEqual(errors, [])
|
||||
self.drain(registration)
|
||||
rows = self.rows()
|
||||
admitted = {row["queue_item_id"] for row in rows}
|
||||
self.assertEqual(len(rows), 128)
|
||||
missing = (set(range(129)) - admitted).pop()
|
||||
oldest = min(admitted)
|
||||
self.queue.task_done(oldest, {}, status())
|
||||
self.queue.task_done(missing, {}, status())
|
||||
self.drain(registration)
|
||||
by_item = {row["queue_item_id"]: row for row in self.rows()}
|
||||
self.assertEqual(by_item[oldest]["outcome"], "succeeded")
|
||||
self.assertNotIn(missing, by_item)
|
||||
self.assertGreater(registration.writer.status().recording_gaps, 0)
|
||||
self.dequeue(129)
|
||||
self.queue.task_done(129, {}, status())
|
||||
self.drain(registration)
|
||||
self.assertEqual(self.rows()[-1]["queue_item_id"], 129)
|
||||
|
||||
def test_active_handles_do_not_retain_graph_extra_or_sensitive_envelopes(self):
|
||||
class Envelope(dict):
|
||||
pass
|
||||
|
||||
registration = self.register()
|
||||
graph, extra, sensitive = (
|
||||
Envelope({"g": object()}),
|
||||
Envelope({"e": object()}),
|
||||
Envelope({"s": object()}),
|
||||
)
|
||||
references = [weakref.ref(value) for value in (graph, extra, sensitive)]
|
||||
self.dequeue(0, graph=graph, extra=extra, sensitive=sensitive)
|
||||
self.drain(registration)
|
||||
del graph, extra, sensitive
|
||||
gc.collect()
|
||||
self.assertEqual([reference() for reference in references], [None, None, None])
|
||||
self.queue.task_done(0, {}, status())
|
||||
self.drain(registration)
|
||||
self.assertEqual(self.rows()[0]["outcome"], "succeeded")
|
||||
|
||||
def test_capture_uses_no_database_methods_on_the_host_thread(self):
|
||||
registration = self.register()
|
||||
host_thread = threading.get_ident()
|
||||
original = self.db.model.get_connection
|
||||
|
||||
def connection():
|
||||
if threading.get_ident() == host_thread:
|
||||
self.fail("history queried SQLite on the host thread")
|
||||
return original()
|
||||
|
||||
with patch.object(self.db.model, "get_connection", connection):
|
||||
self.dequeue(0)
|
||||
self.queue.task_done(0, {}, status())
|
||||
self.drain(registration)
|
||||
self.assertEqual(self.rows()[0]["outcome"], "succeeded")
|
||||
|
||||
def test_capture_never_copies_or_serializes_sensitive_objects(self):
|
||||
class Sensitive:
|
||||
def __deepcopy__(self, memo):
|
||||
raise AssertionError("sensitive copy")
|
||||
|
||||
def __repr__(self):
|
||||
raise AssertionError("sensitive serialization")
|
||||
|
||||
registration = self.register()
|
||||
secret = Sensitive()
|
||||
graph, extra, sensitive = {"graph": secret}, {"extra": secret}, {"auth": secret}
|
||||
result = self.dequeue(0, graph=graph, extra=extra, sensitive=sensitive)
|
||||
payload = {"outputs": secret, "meta": secret}
|
||||
completion = status()
|
||||
completion.messages[0][1].update(
|
||||
{"traceback": secret, "current_inputs": secret}
|
||||
)
|
||||
self.assertIs(
|
||||
self.queue.task_done(0, payload, completion), self.queue.done_result
|
||||
)
|
||||
self.drain(registration)
|
||||
self.assertIs(result[0][2], graph)
|
||||
self.assertIs(result[0][3], extra)
|
||||
self.assertIs(result[0][5], sensitive)
|
||||
self.assertEqual(len(self.rows()), 1)
|
||||
self.assertEqual(self.rows()[0]["outcome"], "succeeded")
|
||||
self.assertEqual(
|
||||
set(completion.messages[0][1]),
|
||||
{"prompt_id", "timestamp", "traceback", "current_inputs"},
|
||||
)
|
||||
|
||||
def test_callbacks_return_while_writer_blocked_and_disposal_is_bounded(self):
|
||||
entered, release = threading.Event(), threading.Event()
|
||||
self.releases.append(release)
|
||||
original = self.db.write_execution_start
|
||||
|
||||
def blocked(command):
|
||||
entered.set()
|
||||
release.wait(5)
|
||||
return original(command)
|
||||
|
||||
with patch.object(self.db, "write_execution_start", blocked):
|
||||
registration = self.register()
|
||||
self.dequeue(0)
|
||||
self.assertTrue(entered.wait(1))
|
||||
began = time.monotonic()
|
||||
self.queue.task_done(0, {}, status())
|
||||
self.dequeue(1)
|
||||
self.assertLess(time.monotonic() - began, 0.2)
|
||||
began = time.monotonic()
|
||||
self.assertFalse(registration.dispose())
|
||||
self.assertGreaterEqual(time.monotonic() - began, 1.8)
|
||||
self.assertLess(time.monotonic() - began, 2.3)
|
||||
release.set()
|
||||
self.assertTrue(registration.dispose(2))
|
||||
|
||||
def test_restore_disables_capture_and_clears_handles_before_fresh_epoch(self):
|
||||
try:
|
||||
from tests.test_execution_history_writer import ManualClock
|
||||
except ImportError:
|
||||
from test_execution_history_writer import ManualClock
|
||||
clock = ManualClock()
|
||||
with patch(
|
||||
"utils.execution_history.HistoryWriter",
|
||||
lambda factory: HistoryWriter(factory, clock=clock),
|
||||
):
|
||||
registration = self.register()
|
||||
for index in range(128):
|
||||
self.dequeue(index)
|
||||
self.drain(registration)
|
||||
old_epoch = registration.writer.status().epoch
|
||||
backup = os.path.join(self.directory.name, "backup.db")
|
||||
self.assertTrue(self.db.model.backup_database(backup))
|
||||
self.db.model.restore_from_file(backup)
|
||||
self.assertFalse(registration.writer.status().accepting)
|
||||
self.assertIsNone(registration.writer.status().epoch)
|
||||
self.queue.task_done(0, {}, status())
|
||||
self.dequeue(128)
|
||||
self.assertGreater(registration.writer.status().recording_gaps, 0)
|
||||
clock.advance(5)
|
||||
self.wait_for(lambda: registration.writer.status().accepting)
|
||||
self.assertNotEqual(registration.writer.status().epoch, old_epoch)
|
||||
# A late final descriptor cannot complete rows from the restored epoch.
|
||||
self.queue.task_done(0, {}, status())
|
||||
self.dequeue(129)
|
||||
self.queue.task_done(129, {}, status())
|
||||
self.drain(registration)
|
||||
rows = self.rows()
|
||||
self.assertEqual(len(rows), 129)
|
||||
self.assertEqual(rows[0]["outcome"], "observed")
|
||||
self.assertEqual(rows[-1]["queue_item_id"], 129)
|
||||
self.assertEqual(rows[-1]["outcome"], "succeeded")
|
||||
|
||||
def test_unavailable_database_and_invalid_identity_never_escape_host(self):
|
||||
entered = threading.Event()
|
||||
|
||||
def broken():
|
||||
entered.set()
|
||||
raise RuntimeError("secret unavailable database")
|
||||
|
||||
registration = self.register(broken, ready=False)
|
||||
self.assertTrue(entered.wait(1))
|
||||
result = self.dequeue(0)
|
||||
self.assertEqual(result[1], 0)
|
||||
self.queue.task_done(0, {}, status())
|
||||
self.assertEqual(self.rows(), [])
|
||||
self.assertGreater(registration.writer.status().recording_gaps, 0)
|
||||
self.assertTrue(registration.dispose())
|
||||
registration = self.register()
|
||||
self.dequeue(0, job="invalid id")
|
||||
self.queue.task_done(True, {}, object())
|
||||
self.drain(registration)
|
||||
self.assertEqual(self.rows(), [])
|
||||
self.assertGreater(registration.writer.status().recording_gaps, 0)
|
||||
|
||||
def test_invalid_evidence_is_recorded_unknown_and_raising_status_is_isolated(self):
|
||||
registration = self.register()
|
||||
self.dequeue(0)
|
||||
self.queue.task_done(0, {}, object())
|
||||
self.drain(registration)
|
||||
self.assertEqual(self.rows()[0]["outcome"], "unknown")
|
||||
self.assertEqual(self.rows()[0]["reason"], "invalid_evidence")
|
||||
|
||||
class BrokenStatus:
|
||||
@property
|
||||
def messages(self):
|
||||
raise RuntimeError("secret failure")
|
||||
|
||||
self.dequeue(1)
|
||||
self.queue.task_done(1, {}, BrokenStatus())
|
||||
self.drain(registration)
|
||||
self.assertEqual(self.rows()[1]["outcome"], "observed")
|
||||
self.assertGreater(registration.writer.status().recording_gaps, 0)
|
||||
|
||||
def test_unsupported_queue_does_not_create_writer_or_submit_fallback(self):
|
||||
self.assertTrue(
|
||||
callable(register_history_hook), "history registration is absent"
|
||||
)
|
||||
for host in (
|
||||
object(),
|
||||
SimpleNamespace(prompt_queue=SimpleNamespace(get=lambda: None)),
|
||||
):
|
||||
self.assertIsNone(
|
||||
register_history_hook(host, lambda: self.fail("DB touched"))
|
||||
)
|
||||
|
||||
def test_existing_usage_callback_counts_once_alongside_history(self):
|
||||
registration = self.register()
|
||||
seen = []
|
||||
self.assertTrue(usage_tracking._wrap_queue_get(self.queue, seen.append))
|
||||
self.assertFalse(usage_tracking._wrap_queue_get(self.queue, seen.append))
|
||||
self.dequeue(0)
|
||||
self.queue.task_done(0, {}, status())
|
||||
self.drain(registration)
|
||||
self.assertEqual(seen, [JOB])
|
||||
self.assertEqual(len(self.rows()), 1)
|
||||
@@ -467,6 +467,20 @@ class HistoryWriterTests(unittest.TestCase):
|
||||
self.assertFalse(writer.status().alive)
|
||||
self.assertFalse(writer.submit(self.start_command()))
|
||||
|
||||
def test_external_capture_gaps_are_nonblocking_and_saturate_all_rejections(self):
|
||||
writer = HistoryWriter(lambda: self.fail("producer accessed database"))
|
||||
self.assertTrue(callable(getattr(writer, "record_gap", None)))
|
||||
writer.record_gap()
|
||||
self.assertEqual(writer.status().recording_gaps, 1)
|
||||
with writer._lock:
|
||||
writer._gaps = 2**63 - 2
|
||||
writer.record_gap()
|
||||
writer.record_gap()
|
||||
writer.submit(self.start_command())
|
||||
writer.submit(object())
|
||||
self.assertEqual(writer.status().recording_gaps, 2**63 - 1)
|
||||
self.assertTrue(writer.close())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,132 @@
|
||||
"""Cooperative queue subscriptions preserve the host and other wrappers."""
|
||||
|
||||
import queue
|
||||
import unittest
|
||||
|
||||
try:
|
||||
from utils.queue_observers import subscribe_dequeue, subscribe_task_done
|
||||
except ImportError:
|
||||
subscribe_dequeue = subscribe_task_done = None
|
||||
|
||||
|
||||
class FakeQueue:
|
||||
def __init__(self):
|
||||
self.result = ((0, "job", {}, {}, [], {}), 0)
|
||||
self.calls = []
|
||||
self.failure = None
|
||||
self.done_result = object()
|
||||
|
||||
def get(self, *args, **kwargs):
|
||||
self.calls.append(("get", args, kwargs))
|
||||
if self.failure:
|
||||
raise self.failure
|
||||
return self.result
|
||||
|
||||
def task_done(self, *args, **kwargs):
|
||||
self.calls.append(("done", args, kwargs))
|
||||
if self.failure:
|
||||
raise self.failure
|
||||
return self.done_result
|
||||
|
||||
|
||||
class QueueObserversTests(unittest.TestCase):
|
||||
def require_api(self):
|
||||
self.assertTrue(callable(subscribe_dequeue), "queue subscriptions are absent")
|
||||
self.assertTrue(callable(subscribe_task_done), "task subscriptions are absent")
|
||||
|
||||
def test_dequeue_preserves_identity_arguments_empty_and_exception(self):
|
||||
self.require_api()
|
||||
host, seen, sentinel = FakeQueue(), [], object()
|
||||
subscribe_dequeue(host, "capture", seen.append)
|
||||
original = host.result
|
||||
self.assertIs(host.get(sentinel, timeout=sentinel), original)
|
||||
self.assertEqual(host.calls, [("get", (sentinel,), {"timeout": sentinel})])
|
||||
self.assertEqual(seen, [original])
|
||||
host.result = None
|
||||
self.assertIsNone(host.get())
|
||||
self.assertEqual(seen, [original])
|
||||
for error in (queue.Empty(), RuntimeError("host failure")):
|
||||
host.failure = error
|
||||
with self.assertRaises(type(error)) as raised:
|
||||
host.get(timeout=sentinel)
|
||||
self.assertIs(raised.exception, error)
|
||||
self.assertEqual(seen, [original])
|
||||
|
||||
def test_task_done_observes_success_only_and_preserves_host_values(self):
|
||||
self.require_api()
|
||||
host, seen, sentinel = FakeQueue(), [], object()
|
||||
subscribe_task_done(host, "capture", lambda a, k: seen.append((a, k)))
|
||||
self.assertIs(host.task_done(0, sentinel, status=sentinel), host.done_result)
|
||||
self.assertEqual(seen, [((0, sentinel), {"status": sentinel})])
|
||||
error = ValueError("host failure")
|
||||
host.failure = error
|
||||
with self.assertRaises(ValueError) as raised:
|
||||
host.task_done(1)
|
||||
self.assertIs(raised.exception, error)
|
||||
self.assertEqual(len(seen), 1)
|
||||
|
||||
def test_keys_are_idempotent_and_independently_removable(self):
|
||||
self.require_api()
|
||||
host, first, second = FakeQueue(), [], []
|
||||
unsubscribe = subscribe_dequeue(host, "a", first.append)
|
||||
subscribe_dequeue(host, "a", lambda value: self.fail("duplicate callback"))
|
||||
remove_second = subscribe_dequeue(host, "b", second.append)
|
||||
host.get()
|
||||
unsubscribe()
|
||||
unsubscribe()
|
||||
host.get()
|
||||
self.assertEqual(len(first), 1)
|
||||
self.assertEqual(len(second), 2)
|
||||
remove_second()
|
||||
subscribe_dequeue(host, "a", first.append)
|
||||
unsubscribe() # old cleanup must not remove a new subscriber with this key
|
||||
host.get()
|
||||
self.assertEqual(len(first), 2)
|
||||
|
||||
def test_other_wrappers_before_after_cleanup_and_resubscription_survive(self):
|
||||
self.require_api()
|
||||
host, calls, seen = FakeQueue(), [], []
|
||||
original = host.get
|
||||
|
||||
def before(*a, **k):
|
||||
calls.append("before")
|
||||
return original(*a, **k)
|
||||
|
||||
host.get = before
|
||||
unsubscribe = subscribe_dequeue(host, "capture", seen.append)
|
||||
observed = host.get
|
||||
|
||||
def after(*a, **k):
|
||||
calls.append("after")
|
||||
return observed(*a, **k)
|
||||
|
||||
host.get = after
|
||||
host.get()
|
||||
unsubscribe()
|
||||
self.assertIs(host.get, after)
|
||||
host.get()
|
||||
subscribe_dequeue(host, "capture", seen.append)
|
||||
host.get()
|
||||
self.assertIs(host.get, after)
|
||||
self.assertEqual(calls, ["after", "before"] * 3)
|
||||
self.assertEqual(len(seen), 2)
|
||||
|
||||
def test_failing_subscriber_cannot_disable_other_subscribers_or_host(self):
|
||||
self.require_api()
|
||||
host, seen = FakeQueue(), []
|
||||
|
||||
def fail(_):
|
||||
raise ValueError("secret error")
|
||||
|
||||
subscribe_dequeue(host, "bad", fail)
|
||||
subscribe_dequeue(host, "good", seen.append)
|
||||
with self.assertLogs("prompt_manager.queue_observers", level="WARNING") as logs:
|
||||
self.assertIs(host.get(), host.result)
|
||||
self.assertEqual(seen, [host.result])
|
||||
self.assertNotIn("secret error", "\n".join(logs.output))
|
||||
|
||||
def test_missing_methods_are_rejected_without_installing_a_fallback(self):
|
||||
self.require_api()
|
||||
for subscribe in (subscribe_dequeue, subscribe_task_done):
|
||||
with self.assertRaises(ValueError):
|
||||
subscribe(object(), "key", lambda *a: None)
|
||||
@@ -373,6 +373,22 @@ class TestCountAtDequeue(DbTestCase):
|
||||
self.assertEqual(self._runs(positive), 1)
|
||||
self.assertEqual(calls, [])
|
||||
|
||||
def test_usage_subscription_remains_once_after_an_unrelated_wrapper(self):
|
||||
positive = self._save("a cat")
|
||||
prompt_queue = self.server.prompt_queue
|
||||
original = prompt_queue.get
|
||||
|
||||
def unrelated(*args, **kwargs):
|
||||
return original(*args, **kwargs)
|
||||
|
||||
prompt_queue.get = unrelated
|
||||
self.assertFalse(
|
||||
_wrap_queue_get(prompt_queue, lambda _: self.fail("duplicate"))
|
||||
)
|
||||
self.assertIs(prompt_queue.get, unrelated)
|
||||
self._dequeue(self._submit(workflow()))
|
||||
self.assertEqual(self._runs(positive), 1)
|
||||
|
||||
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))
|
||||
|
||||
@@ -0,0 +1,145 @@
|
||||
"""Own bounded queue-attempt capture without retaining host workflow payloads."""
|
||||
|
||||
import atexit
|
||||
import threading
|
||||
import time
|
||||
from typing import Callable
|
||||
from uuid import uuid4
|
||||
|
||||
try:
|
||||
from ..database import PromptDatabase
|
||||
except ImportError:
|
||||
from database import PromptDatabase
|
||||
|
||||
from .execution_history_events import make_finish, make_start
|
||||
from .execution_history_writer import HistoryWriter
|
||||
from .logging_config import get_logger
|
||||
from .queue_observers import subscribe_dequeue, subscribe_task_done
|
||||
|
||||
_logger = get_logger("prompt_manager.execution_history")
|
||||
_STATE_ATTRIBUTE = "_prompt_manager_execution_history"
|
||||
_registration_lock = threading.Lock()
|
||||
_MAX_ACTIVE = 128
|
||||
|
||||
|
||||
class HistoryRegistration:
|
||||
"""Own subscriptions, one asynchronous writer and bounded attempt identities."""
|
||||
|
||||
def __init__(
|
||||
self, queue: object, session_id: str, db_factory: Callable[[], PromptDatabase]
|
||||
) -> None:
|
||||
self.writer = HistoryWriter(db_factory)
|
||||
self._session_id = session_id
|
||||
self._lock = threading.Lock()
|
||||
self._active: dict[int, str] = {}
|
||||
self._disposed = False
|
||||
self._unsubscribers: list[Callable[[], None]] = []
|
||||
try:
|
||||
self._unsubscribers.append(
|
||||
subscribe_dequeue(queue, "execution_history", self._dequeue)
|
||||
)
|
||||
self._unsubscribers.append(
|
||||
subscribe_task_done(queue, "execution_history", self._task_done)
|
||||
)
|
||||
self.writer.start()
|
||||
atexit.register(self.dispose)
|
||||
except Exception:
|
||||
self.dispose()
|
||||
raise
|
||||
|
||||
def _ready_epoch(self) -> str | None:
|
||||
status = self.writer.status()
|
||||
epoch = status.epoch if status.accepting else None
|
||||
if any(value != epoch for value in self._active.values()):
|
||||
self._active.clear()
|
||||
return epoch
|
||||
|
||||
def _dequeue(self, result: object) -> None:
|
||||
try:
|
||||
with self._lock:
|
||||
if self._disposed:
|
||||
return
|
||||
epoch = self._ready_epoch()
|
||||
if epoch is None:
|
||||
self.writer.record_gap()
|
||||
return
|
||||
command = make_start(
|
||||
self._session_id, result, epoch, time.time_ns() // 1_000_000
|
||||
)
|
||||
if command is None:
|
||||
self.writer.record_gap()
|
||||
return
|
||||
item_id = command.queue_item_id
|
||||
if item_id in self._active:
|
||||
return
|
||||
if len(self._active) >= _MAX_ACTIVE:
|
||||
self.writer.record_gap()
|
||||
return
|
||||
if self.writer.submit(command):
|
||||
self._active[item_id] = epoch
|
||||
except Exception:
|
||||
self.writer.record_gap()
|
||||
|
||||
def _task_done(self, args: tuple, kwargs: dict) -> None:
|
||||
try:
|
||||
with self._lock:
|
||||
if self._disposed:
|
||||
return
|
||||
epoch = self._ready_epoch()
|
||||
if epoch is None:
|
||||
self.writer.record_gap()
|
||||
return
|
||||
item_id = args[0] if args else kwargs.get("item_id")
|
||||
status = args[2] if len(args) > 2 else kwargs.get("status")
|
||||
command = make_finish(
|
||||
self._session_id,
|
||||
item_id,
|
||||
epoch,
|
||||
status,
|
||||
time.time_ns() // 1_000_000,
|
||||
)
|
||||
if command is None:
|
||||
self.writer.record_gap()
|
||||
return
|
||||
self._active.pop(command.queue_item_id, None)
|
||||
self.writer.submit(command)
|
||||
except Exception:
|
||||
self.writer.record_gap()
|
||||
|
||||
def dispose(self, timeout: float = 2.0) -> bool:
|
||||
"""Stop capture, release subscriptions and drain the owned writer boundedly."""
|
||||
with self._lock:
|
||||
if not self._disposed:
|
||||
self._disposed = True
|
||||
self._active.clear()
|
||||
for unsubscribe in self._unsubscribers:
|
||||
unsubscribe()
|
||||
self._unsubscribers.clear()
|
||||
atexit.unregister(self.dispose)
|
||||
return self.writer.close(timeout)
|
||||
|
||||
|
||||
def register_history_hook(
|
||||
server: object, db_factory: Callable[[], PromptDatabase]
|
||||
) -> HistoryRegistration | None:
|
||||
"""Register only on queues supporting actual dequeue and final callbacks."""
|
||||
try:
|
||||
queue = getattr(server, "prompt_queue", None)
|
||||
if not all(
|
||||
callable(getattr(queue, method, None)) for method in ("get", "task_done")
|
||||
):
|
||||
return None
|
||||
with _registration_lock:
|
||||
state = getattr(queue, _STATE_ATTRIBUTE, None)
|
||||
if state is None:
|
||||
state = {"session_id": str(uuid4()), "registration": None}
|
||||
setattr(queue, _STATE_ATTRIBUTE, state)
|
||||
previous = state["registration"]
|
||||
if previous is not None and not previous._disposed:
|
||||
return previous
|
||||
registration = HistoryRegistration(queue, state["session_id"], db_factory)
|
||||
state["registration"] = registration
|
||||
return registration
|
||||
except Exception:
|
||||
_logger.warning("Execution history queue capability unavailable")
|
||||
return None
|
||||
@@ -26,6 +26,7 @@ from .logging_config import get_logger
|
||||
_logger = get_logger("prompt_manager.execution_history_writer")
|
||||
_MAX_PENDING = 1024
|
||||
_PROBE_INTERVAL = 5.0
|
||||
_MAX_GAPS = 2**63 - 1
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -75,7 +76,7 @@ class HistoryWriter:
|
||||
encode_command(command)
|
||||
except ValueError:
|
||||
with self._lock:
|
||||
self._gaps += 1
|
||||
self._record_gap_locked()
|
||||
self._wake.set()
|
||||
return False
|
||||
with self._lock:
|
||||
@@ -84,7 +85,7 @@ class HistoryWriter:
|
||||
or command.storage_epoch != self._epoch
|
||||
or self._pending >= _MAX_PENDING
|
||||
):
|
||||
self._gaps += 1
|
||||
self._record_gap_locked()
|
||||
accepted = False
|
||||
else:
|
||||
self._queue.put_nowait(command)
|
||||
@@ -93,6 +94,15 @@ class HistoryWriter:
|
||||
self._wake.set()
|
||||
return accepted
|
||||
|
||||
def _record_gap_locked(self) -> None:
|
||||
self._gaps = min(_MAX_GAPS, self._gaps + 1)
|
||||
|
||||
def record_gap(self) -> None:
|
||||
"""Account for unavailable capture without constructing a fake command."""
|
||||
with self._lock:
|
||||
self._record_gap_locked()
|
||||
self._wake.set()
|
||||
|
||||
def invalidate(self) -> None:
|
||||
with self._lock:
|
||||
self._accepting = False
|
||||
@@ -134,7 +144,7 @@ class HistoryWriter:
|
||||
except queue.Empty:
|
||||
return
|
||||
self._pending -= 1
|
||||
self._gaps += 1
|
||||
self._record_gap_locked()
|
||||
|
||||
def _run(self) -> None:
|
||||
db: PromptDatabase | None = None
|
||||
@@ -187,7 +197,7 @@ class HistoryWriter:
|
||||
if db is not None:
|
||||
db.model.close()
|
||||
with self._lock:
|
||||
self._gaps += 1
|
||||
self._record_gap_locked()
|
||||
if command is not None:
|
||||
failed = False
|
||||
try:
|
||||
@@ -206,7 +216,7 @@ class HistoryWriter:
|
||||
with self._lock:
|
||||
self._pending -= 1
|
||||
if failed:
|
||||
self._gaps += 1
|
||||
self._record_gap_locked()
|
||||
self._accepting = False
|
||||
self._epoch = None
|
||||
self._generation += 1
|
||||
|
||||
@@ -0,0 +1,110 @@
|
||||
"""Keyed, independently removable observers for successful queue operations."""
|
||||
|
||||
import functools
|
||||
import threading
|
||||
import time
|
||||
from typing import Callable
|
||||
|
||||
from .logging_config import get_logger
|
||||
|
||||
_logger = get_logger("prompt_manager.queue_observers")
|
||||
_STATE_ATTRIBUTE = "_prompt_manager_queue_observers"
|
||||
_registration_lock = threading.Lock()
|
||||
|
||||
|
||||
class _Observation:
|
||||
def __init__(self, queue: object, method: str) -> None:
|
||||
self.queue = queue
|
||||
self.method = method
|
||||
self.lock = threading.Lock()
|
||||
self.callbacks: dict[str, tuple[object, Callable]] = {}
|
||||
self.original = None
|
||||
self.wrapper = None
|
||||
self.installed = False
|
||||
self.next_warning = 0.0
|
||||
|
||||
def install(self) -> None:
|
||||
original = getattr(self.queue, self.method, None)
|
||||
if not callable(original):
|
||||
raise ValueError("Queue capability unavailable")
|
||||
|
||||
@functools.wraps(original)
|
||||
def observe(*args: object, **kwargs: object) -> object:
|
||||
result = original(*args, **kwargs)
|
||||
if self.method == "task_done" or result is not None:
|
||||
with self.lock:
|
||||
callbacks = tuple(entry[1] for entry in self.callbacks.values())
|
||||
for callback in callbacks:
|
||||
try:
|
||||
if self.method == "get":
|
||||
callback(result)
|
||||
else:
|
||||
callback(args, kwargs)
|
||||
except Exception:
|
||||
with self.lock:
|
||||
now = time.monotonic()
|
||||
warn = now >= self.next_warning
|
||||
if warn:
|
||||
self.next_warning = now + 5.0
|
||||
if warn:
|
||||
_logger.warning("Queue observation callback failed")
|
||||
return result
|
||||
|
||||
setattr(self.queue, self.method, observe)
|
||||
self.original, self.wrapper, self.installed = original, observe, True
|
||||
|
||||
def subscribe(self, key: str, callback: Callable) -> Callable[[], None]:
|
||||
with self.lock:
|
||||
if not self.installed:
|
||||
self.install()
|
||||
existing = self.callbacks.get(key)
|
||||
token = existing[0] if existing is not None else object()
|
||||
if existing is None:
|
||||
self.callbacks[key] = (token, callback)
|
||||
|
||||
def unsubscribe() -> None:
|
||||
with self.lock:
|
||||
current = self.callbacks.get(key)
|
||||
if current is None or current[0] is not token:
|
||||
return
|
||||
del self.callbacks[key]
|
||||
if (
|
||||
not self.callbacks
|
||||
and getattr(self.queue, self.method) is self.wrapper
|
||||
):
|
||||
setattr(self.queue, self.method, self.original)
|
||||
self.installed = False
|
||||
|
||||
return unsubscribe
|
||||
|
||||
|
||||
def _subscribe(
|
||||
queue: object, method: str, key: str, callback: Callable
|
||||
) -> Callable[[], None]:
|
||||
with _registration_lock:
|
||||
states = getattr(queue, _STATE_ATTRIBUTE, None)
|
||||
if states is None:
|
||||
states = {}
|
||||
try:
|
||||
setattr(queue, _STATE_ATTRIBUTE, states)
|
||||
except (AttributeError, TypeError) as error:
|
||||
raise ValueError("Queue does not support observation") from error
|
||||
state = states.get(method)
|
||||
if state is None:
|
||||
state = _Observation(queue, method)
|
||||
states[method] = state
|
||||
return state.subscribe(key, callback)
|
||||
|
||||
|
||||
def subscribe_dequeue(
|
||||
queue: object, key: str, callback: Callable[[object], None]
|
||||
) -> Callable[[], None]:
|
||||
"""Observe each actual get result once per key, after the host returns."""
|
||||
return _subscribe(queue, "get", key, callback)
|
||||
|
||||
|
||||
def subscribe_task_done(
|
||||
queue: object, key: str, callback: Callable[[tuple, dict], None]
|
||||
) -> Callable[[], None]:
|
||||
"""Observe successful task_done calls with their original argument objects."""
|
||||
return _subscribe(queue, "task_done", key, callback)
|
||||
+13
-13
@@ -19,7 +19,6 @@ 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
|
||||
@@ -32,6 +31,7 @@ try:
|
||||
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
|
||||
@@ -42,6 +42,7 @@ except ImportError:
|
||||
)
|
||||
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")
|
||||
|
||||
@@ -294,21 +295,20 @@ def _wrap_queue_get(prompt_queue: Any, on_dequeued: Callable[[str], None]) -> bo
|
||||
original = getattr(prompt_queue, "get", None)
|
||||
if not callable(original):
|
||||
return False
|
||||
if getattr(original, _WRAPPED_MARKER, False):
|
||||
if getattr(prompt_queue, _WRAPPED_MARKER, False) or 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
|
||||
def report(result: object) -> None:
|
||||
try:
|
||||
on_dequeued(str(result[0][1]))
|
||||
except Exception as e:
|
||||
logger.warning(f"Dequeue hook failed: {e}")
|
||||
|
||||
setattr(get, _WRAPPED_MARKER, True)
|
||||
prompt_queue.get = get
|
||||
subscribe_dequeue(prompt_queue, "usage_tracking", report)
|
||||
setattr(prompt_queue.get, _WRAPPED_MARKER, True)
|
||||
setattr(prompt_queue, _WRAPPED_MARKER, True)
|
||||
return True
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user