Vito Sansevero
d605a6a758
The prompt hash folds case, so a run of "A Lighthouse" counts toward the saved "a lighthouse" but never joined its history. History now compares the hash's own key (prompt_key = strip + str.lower()); SQL calls Python for it through pm_prompt_key/pm_prompt_hash registered on every managed connection, since SQLite's lower() folds ASCII only. Hints match by text first (#438) and fall back to the dedup hash only when no text matches, so no existing hint becomes ambiguous. A history position from before this change still resumes.
959 lines
38 KiB
Python
959 lines
38 KiB
Python
"""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 tests.history_waits import drain_writer
|
|
from database import PromptDatabase
|
|
from utils import usage_tracking
|
|
from utils.hashing import generate_prompt_hash
|
|
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"
|
|
# A cross-thread wait or a duration check fails only past this (#370).
|
|
CEILING_SECONDS = 30.0
|
|
|
|
|
|
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 HooksCase(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 registration in self.registrations:
|
|
registration.writer.invalidate()
|
|
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):
|
|
drain_writer(registration.writer, clock=time.monotonic)
|
|
|
|
|
|
class HistoryHooksTests(HooksCase):
|
|
def test_package_history_initializes_when_monitor_and_usage_setup_fail(self):
|
|
try:
|
|
from tests.test_init_package import (
|
|
fake_server_module,
|
|
load_package,
|
|
)
|
|
except ImportError:
|
|
from test_init_package import (
|
|
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"),
|
|
):
|
|
with load_package(name, self.db.model.db_path, output, server) as (
|
|
package,
|
|
_,
|
|
):
|
|
registration = getattr(package, "_history_registration", None)
|
|
self.assertIsNotNone(
|
|
registration, "history setup depended on another hook"
|
|
)
|
|
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())
|
|
# The writer drains first: this test checks wrappers and the session,
|
|
# not the 2 s bound, which a slow runner can spend on the write (#370).
|
|
self.drain(first)
|
|
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)
|
|
# Ceilings, not tight wall clocks (#370): the return values show
|
|
# the bound; a ceiling still fails a call that waits unbounded.
|
|
self.assertTrue(entered.wait(CEILING_SECONDS))
|
|
began = time.monotonic()
|
|
self.queue.task_done(0, {}, status())
|
|
self.dequeue(1)
|
|
self.assertLess(time.monotonic() - began, CEILING_SECONDS)
|
|
began = time.monotonic()
|
|
self.assertFalse(registration.dispose())
|
|
self.assertGreaterEqual(time.monotonic() - began, 1.8)
|
|
self.assertLess(time.monotonic() - began, CEILING_SECONDS)
|
|
release.set()
|
|
# Drain first: the 2 s bound is not under test here (#370).
|
|
self.drain(registration)
|
|
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, **kwargs: HistoryWriter(factory, clock=clock, **kwargs),
|
|
):
|
|
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)
|
|
|
|
def test_fixture_cleanup_discards_saturation_before_release_and_database_close(
|
|
self,
|
|
):
|
|
from utils.execution_history_events import make_start
|
|
|
|
entered, release, rescue = (threading.Event() for _ in range(3))
|
|
self.releases.append(release)
|
|
pending_at_release, writes = [], []
|
|
original = self.db.write_execution_start
|
|
|
|
def blocked(command):
|
|
entered.set()
|
|
release.wait(5)
|
|
pending = registration.writer.status().pending
|
|
pending_at_release.append(pending)
|
|
# Model a slow drain without executing 1,024 SQLite transactions.
|
|
if pending > 1:
|
|
rescue.wait(5)
|
|
writes.append(command.run_id)
|
|
return original(command)
|
|
|
|
with patch.object(self.db, "write_execution_start", blocked):
|
|
registration = self.register()
|
|
try:
|
|
self.dequeue(1)
|
|
self.assertTrue(entered.wait(1))
|
|
epoch = registration.writer.status().epoch
|
|
for i in range(1023):
|
|
command = make_start(
|
|
registration._session_id, item(1000 + i), epoch, 0
|
|
)
|
|
self.assertTrue(registration.writer.submit(command))
|
|
self.assertEqual(registration.writer.status().pending, 1024)
|
|
before = time.monotonic()
|
|
self.cleanup_registrations()
|
|
self.assertLess(time.monotonic() - before, 2)
|
|
self.assertEqual(pending_at_release, [1])
|
|
self.assertEqual(len(writes), 1)
|
|
self.assertFalse(registration.writer.status().alive)
|
|
self.assertEqual(registration.writer.status().pending, 0)
|
|
self.db.model.close_all()
|
|
finally:
|
|
registration.writer.invalidate()
|
|
rescue.set()
|
|
release.set()
|
|
self.assertTrue(
|
|
registration.dispose(3), "regression rescue leaked worker"
|
|
)
|
|
|
|
def test_rejected_new_start_revokes_previous_worker_capability(self):
|
|
from utils.execution_context_capture import active_capture
|
|
from utils.execution_history_events import make_start
|
|
|
|
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(1)
|
|
self.assertTrue(entered.wait(1))
|
|
previous = active_capture()
|
|
self.assertIsNotNone(previous)
|
|
epoch = registration.writer.status().epoch
|
|
for i in range(1023):
|
|
command = make_start(registration._session_id, item(1000 + i), epoch, 0)
|
|
self.assertTrue(registration.writer.submit(command))
|
|
self.assertEqual(registration.writer.status().pending, 1024)
|
|
self.dequeue(2)
|
|
self.assertIsNone(active_capture())
|
|
self.assertFalse(previous[0].valid(previous[1]))
|
|
registration.writer.invalidate()
|
|
release.set()
|
|
self.assertTrue(registration.dispose(3), "owned worker leaked")
|
|
|
|
def test_late_finish_does_not_clear_new_attempt_binding(self):
|
|
from utils.execution_context_capture import active_capture
|
|
|
|
registration = self.register()
|
|
self.dequeue(1)
|
|
self.dequeue(2)
|
|
current = active_capture()
|
|
self.assertIsNotNone(current)
|
|
self.queue.task_done(1, {}, status())
|
|
self.assertEqual(active_capture(), current)
|
|
# An identical callback replay keeps the already accepted binding.
|
|
registration._dequeue(item(2))
|
|
self.assertEqual(active_capture(), current)
|
|
self.queue.task_done(2, {}, status())
|
|
self.assertIsNone(active_capture())
|
|
|
|
def test_invalidation_releases_capture_immediately_outside_writer_locks(self):
|
|
from utils.execution_context_capture import active_capture
|
|
|
|
registration = self.register()
|
|
self.dequeue(1)
|
|
current = active_capture()
|
|
pool = registration.writer.context_budget
|
|
self.assertGreater(pool.charged_bytes, 0)
|
|
registration.writer.invalidate()
|
|
self.assertEqual(pool.charged_bytes, 0)
|
|
self.assertFalse(current[0].valid(current[1]))
|
|
|
|
def test_replayed_older_accepted_dequeue_cannot_revoke_current_binding(self):
|
|
from utils.execution_context_capture import active_capture
|
|
|
|
registration = self.register()
|
|
old = self.dequeue(1)
|
|
self.dequeue(2)
|
|
current = active_capture()
|
|
registration._dequeue(old)
|
|
self.assertEqual(active_capture(), current)
|
|
|
|
def test_broken_finish_releases_active_capture(self):
|
|
from utils.execution_context_capture import active_capture
|
|
|
|
registration = self.register()
|
|
self.dequeue(1)
|
|
|
|
class Broken:
|
|
@property
|
|
def messages(self):
|
|
raise RuntimeError("broken status")
|
|
|
|
self.queue.task_done(1, {}, Broken())
|
|
self.assertEqual(registration.writer.context_budget.charged_bytes, 0)
|
|
self.assertIsNone(active_capture())
|
|
|
|
def test_finished_accepted_replay_preserves_new_same_uuid_attempt(self):
|
|
from utils.execution_context_capture import active_capture
|
|
|
|
registration = self.register()
|
|
old = self.dequeue(1)
|
|
self.queue.task_done(1, {}, status())
|
|
self.dequeue(2)
|
|
current = active_capture()
|
|
registration._dequeue(old)
|
|
self.assertEqual(active_capture(), current)
|
|
self.queue.task_done(2, {}, status())
|
|
self.drain(registration)
|
|
self.assertEqual(len(self.rows()), 2)
|
|
|
|
def test_observed_upstream_survives_downstream_failure_in_database(self):
|
|
from utils.execution_context_capture import record_runtime_context
|
|
|
|
g = {
|
|
"p": {
|
|
"class_type": "PromptManagerPair",
|
|
"inputs": {"text": "raw", "negative_text": ""},
|
|
},
|
|
"s": {
|
|
"class_type": "KSampler",
|
|
"inputs": {"positive": ["p", 0], "negative": ["p", 1]},
|
|
},
|
|
}
|
|
registration = self.register()
|
|
self.queue.results.append(((0, JOB, g, {}, ["s"]), 1))
|
|
self.queue.get()
|
|
host = types.ModuleType("comfy_execution.utils")
|
|
host.get_executing_context = lambda: SimpleNamespace(
|
|
prompt_id=JOB, node_id="p", list_index=0
|
|
)
|
|
with patch.dict(sys.modules, {"comfy_execution.utils": host}):
|
|
record_runtime_context(
|
|
"PromptManagerPair",
|
|
" raw ",
|
|
" pre ",
|
|
"",
|
|
"pre raw TRIGGER",
|
|
role="positive",
|
|
)
|
|
record_runtime_context("PromptManagerPair", "", "", "", "", role="negative")
|
|
self.queue.task_done(1, {}, status("execution_error"))
|
|
self.drain(registration)
|
|
row = self.rows()[0]
|
|
self.assertEqual(row["outcome"], "failed")
|
|
result = self.db.get_execution_contexts(row["run_id"])
|
|
self.assertEqual(result["state"], "complete")
|
|
self.assertEqual(
|
|
result["contexts"][0]["positive"],
|
|
{
|
|
"raw": " raw ",
|
|
"effective": "pre raw TRIGGER",
|
|
"prepend": " pre ",
|
|
"append": "",
|
|
},
|
|
)
|
|
self.assertEqual(result["contexts"][0]["negative"]["effective"], "")
|
|
|
|
def test_close_releases_capture_while_writer_is_blocked(self):
|
|
entered, release = threading.Event(), threading.Event()
|
|
self.releases.append(release)
|
|
original = self.db.write_execution_start
|
|
|
|
def blocked(command):
|
|
entered.set()
|
|
release.wait(3)
|
|
return original(command)
|
|
|
|
with patch.object(self.db, "write_execution_start", blocked):
|
|
registration = self.register()
|
|
self.dequeue(1)
|
|
self.assertTrue(entered.wait(1))
|
|
self.assertGreater(registration.writer.context_budget.charged_bytes, 0)
|
|
self.assertFalse(registration.writer.close(0))
|
|
self.assertEqual(registration.writer.context_budget.charged_bytes, 0)
|
|
release.set()
|
|
|
|
def test_reregistration_cannot_rebind_already_observed_queue_item(self):
|
|
from utils.execution_context_capture import active_capture
|
|
|
|
first = self.register()
|
|
self.dequeue(10)
|
|
self.queue.task_done(10, {}, status())
|
|
self.drain(first)
|
|
self.assertTrue(first.dispose())
|
|
second = self.register()
|
|
self.dequeue(10)
|
|
self.assertIsNone(active_capture())
|
|
self.queue.task_done(10, {}, status())
|
|
self.drain(second)
|
|
self.assertEqual(len(self.rows()), 1)
|
|
|
|
def test_cached_executor_subscription_persists_strings_and_disposes_cooperatively(
|
|
self,
|
|
):
|
|
from utils.execution_context_host import subscribe_executor_messages
|
|
|
|
calls = []
|
|
execution = types.ModuleType("execution")
|
|
caching = types.ModuleType("comfy_execution.caching")
|
|
package = types.ModuleType("comfy_execution")
|
|
entry_type = namedtuple("CacheEntry", "ui outputs")
|
|
|
|
class PromptExecutor:
|
|
def add_message(self, event, data, broadcast):
|
|
calls.append(event)
|
|
return "host-result"
|
|
|
|
class HierarchicalCache:
|
|
def get_local(self, node_id):
|
|
return entry_type(None, [object(), object(), ["old trigger"], [""]])
|
|
|
|
execution.PromptExecutor, execution.CacheEntry = PromptExecutor, entry_type
|
|
caching.HierarchicalCache = HierarchicalCache
|
|
package.caching = caching
|
|
with patch.dict(
|
|
sys.modules,
|
|
{
|
|
"execution": execution,
|
|
"comfy_execution": package,
|
|
"comfy_execution.caching": caching,
|
|
},
|
|
):
|
|
observer_calls = []
|
|
independent = subscribe_executor_messages(
|
|
"independent", lambda *args: observer_calls.append(args)
|
|
)
|
|
try:
|
|
registration = self.register()
|
|
g = {
|
|
"p": {
|
|
"class_type": "PromptManagerPair",
|
|
"inputs": {"text": "raw", "negative_text": ""},
|
|
},
|
|
"s": {
|
|
"class_type": "KSampler",
|
|
"inputs": {"positive": ["p", 0], "negative": ["p", 1]},
|
|
},
|
|
}
|
|
self.queue.results.append(((0, JOB, g, {}, ["s"]), 1))
|
|
self.queue.get()
|
|
executor = PromptExecutor()
|
|
executor.server = self.server
|
|
executor.caches = SimpleNamespace(outputs=HierarchicalCache())
|
|
self.assertEqual(
|
|
executor.add_message("execution_start", {"prompt_id": JOB}, False),
|
|
"host-result",
|
|
)
|
|
executor.add_message(
|
|
event="execution_cached",
|
|
data={"prompt_id": JOB, "nodes": ["p"]},
|
|
broadcast=False,
|
|
)
|
|
self.queue.task_done(1, {}, status())
|
|
self.drain(registration)
|
|
result = self.db.get_execution_contexts(self.rows()[0]["run_id"])
|
|
self.assertEqual(
|
|
result["contexts"][0]["positive"]["effective"], "old trigger"
|
|
)
|
|
self.assertEqual(result["contexts"][0]["negative"]["effective"], "")
|
|
self.assertTrue(registration.dispose())
|
|
executor.add_message(
|
|
"execution_cached", {"prompt_id": JOB, "nodes": ["p"]}, False
|
|
)
|
|
self.assertEqual(len(observer_calls), 3)
|
|
finally:
|
|
independent()
|
|
self.assertEqual(
|
|
calls, ["execution_start", "execution_cached", "execution_cached"]
|
|
)
|
|
|
|
|
|
class CachedRuntimeHistoryTests(HooksCase):
|
|
"""A cached runtime-linked prompt is a tracked run, like an executed one (#403).
|
|
|
|
The real history hook and writer, a real ``CachedRuns`` and the real host
|
|
adapter over fakes of ComfyUI's executor, cache and executing context.
|
|
"""
|
|
|
|
TEXT = "a lighthouse the LLM wrote"
|
|
GRAPH = {
|
|
"str": {"class_type": "PrimitiveString", "inputs": {"value": "seed"}},
|
|
"llm": {"class_type": "PreviewAny", "inputs": {"source": ["str", 0]}},
|
|
"p": {"class_type": "PromptManager", "inputs": {"text": ["llm", 0]}},
|
|
"s": {"class_type": "KSampler", "inputs": {"positive": ["p", 0]}},
|
|
}
|
|
|
|
def setUp(self):
|
|
super().setUp()
|
|
from utils import cached_runs
|
|
from utils.cached_runs import CachedRuns
|
|
|
|
self.context = None
|
|
self.runs = CachedRuns(context=lambda: self.context or (None, None))
|
|
self.cache_entries = {}
|
|
entries = self.cache_entries
|
|
execution = types.ModuleType("execution")
|
|
caching = types.ModuleType("comfy_execution.caching")
|
|
package = types.ModuleType("comfy_execution")
|
|
utils_module = types.ModuleType("comfy_execution.utils")
|
|
entry_type = namedtuple("CacheEntry", "ui outputs")
|
|
|
|
class PromptExecutor:
|
|
def add_message(self, event, data, broadcast):
|
|
return None
|
|
|
|
class HierarchicalCache:
|
|
def get_local(self, node_id):
|
|
return entries.get(node_id)
|
|
|
|
execution.PromptExecutor, execution.CacheEntry = PromptExecutor, entry_type
|
|
caching.HierarchicalCache = HierarchicalCache
|
|
package.caching, package.utils = caching, utils_module
|
|
utils_module.get_executing_context = lambda: (
|
|
SimpleNamespace(prompt_id=self.context[0], node_id="p", list_index=0)
|
|
if self.context
|
|
else None
|
|
)
|
|
self.entry_type = entry_type
|
|
modules = patch.dict(
|
|
sys.modules,
|
|
{
|
|
"execution": execution,
|
|
"comfy_execution": package,
|
|
"comfy_execution.caching": caching,
|
|
"comfy_execution.utils": utils_module,
|
|
},
|
|
)
|
|
modules.start()
|
|
self.addCleanup(modules.stop)
|
|
swap = patch.object(cached_runs, "CACHED_RUNS", self.runs)
|
|
swap.start()
|
|
self.addCleanup(swap.stop)
|
|
self.executor = PromptExecutor()
|
|
self.executor.server = self.server
|
|
self.executor.caches = SimpleNamespace(outputs=HierarchicalCache())
|
|
self.prompt_id = self.db.save_prompt(text=self.TEXT)
|
|
self.item_id = 0
|
|
|
|
def job(self, cached):
|
|
"""One job; the node executes unless ``cached`` (ComfyUI's cache serves it)."""
|
|
from uuid import uuid4
|
|
|
|
from utils.cached_runs import stored_facts
|
|
from utils.execution_context_capture import record_runtime_context
|
|
|
|
self.item_id += 1
|
|
job = str(uuid4())
|
|
result = ((0, job, self.GRAPH, {}, ["s"]), self.item_id)
|
|
self.queue.results.append(result)
|
|
self.queue.get()
|
|
self.runs.begin_attempt(result)
|
|
self.context = (job, 0)
|
|
self.executor.add_message("execution_start", {"prompt_id": job}, False)
|
|
if cached:
|
|
self.executor.add_message(
|
|
"execution_cached", {"prompt_id": job, "nodes": ["p"]}, False
|
|
)
|
|
else:
|
|
# What PromptManager.encode_prompt records for this text
|
|
record_runtime_context("PromptManager", self.TEXT, "", "", self.TEXT)
|
|
self.runs.record_item(
|
|
"p", (self.TEXT,), stored_facts(self.TEXT), ((self.TEXT, "", ""),)
|
|
)
|
|
self.cache_entries["p"] = self.entry_type(None, [object(), [self.TEXT]])
|
|
self.context = None
|
|
self.queue.task_done(self.item_id, {}, status(job=job))
|
|
self.runs.end_attempt(self.item_id)
|
|
|
|
def tracked_runs(self):
|
|
with self.db.recorded_history_snapshot():
|
|
return self.db.list_history_window(self.prompt_id)["recorded_run_count"]
|
|
|
|
def test_three_jobs_of_runtime_text_are_three_tracked_runs(self):
|
|
registration = self.register()
|
|
for cached in (False, True, True):
|
|
self.job(cached)
|
|
self.drain(registration)
|
|
self.assertEqual(self.tracked_runs(), 3)
|
|
rows = self.rows()
|
|
for row in rows:
|
|
contexts = self.db.get_execution_contexts(row["run_id"])["contexts"]
|
|
self.assertEqual(contexts[0]["positive"]["raw"], self.TEXT)
|
|
self.assertEqual(contexts[0]["positive_prompt_id"], self.prompt_id)
|
|
cached = self.db.get_execution_contexts(rows[1]["run_id"])["contexts"][0]
|
|
self.assertEqual(
|
|
cached["evidence"]["positive"]["raw"],
|
|
{"provenance": "inferred", "source": "graph_semantics", "reason": None},
|
|
)
|
|
self.assertEqual(
|
|
cached["evidence"]["positive"]["effective"]["source"], "host_cache_output"
|
|
)
|
|
# The stored history passes the whole-database check older builds run too
|
|
from database.execution_history_schema import history_schema_error
|
|
|
|
self.assertIsNone(history_schema_error(self.db.model.get_connection()))
|
|
# An image of a cached job can now be attributed through its run (#342)
|
|
from database.run_attribution import named_prompt
|
|
|
|
self.assertEqual(
|
|
[named_prompt(self.db, row["run_id"]) for row in rows],
|
|
[self.prompt_id] * 3,
|
|
)
|
|
|
|
def test_without_a_receipt_a_cached_job_stays_untracked(self):
|
|
registration = self.register()
|
|
self.job(False)
|
|
self.runs._receipts.clear() # as after a restart with a provider entry
|
|
self.job(True)
|
|
self.drain(registration)
|
|
self.assertEqual(self.tracked_runs(), 1)
|
|
cached = self.db.get_execution_contexts(self.rows()[1]["run_id"])["contexts"]
|
|
self.assertIsNone(cached[0]["positive"]["raw"])
|
|
self.assertIsNone(cached[0]["positive_prompt_id"])
|
|
from database.run_attribution import named_prompt
|
|
|
|
self.assertIsNone(named_prompt(self.db, self.rows()[1]["run_id"]))
|
|
|
|
|
|
class CachedRuntimeCaseHistoryTests(CachedRuntimeHistoryTests):
|
|
"""Raw text that differs only in case joins the saved prompt's history (#460).
|
|
|
|
The hash folds case, so the run counts toward the lower-case prompt saved
|
|
first; the run keeps its exact raw spelling.
|
|
"""
|
|
|
|
TEXT = "\u00c9 LIGHTHOUSE the LLM wrote"
|
|
|
|
def setUp(self):
|
|
super().setUp()
|
|
self.db.delete_prompt(self.prompt_id)
|
|
lower = self.TEXT.lower()
|
|
self.prompt_id = self.db.save_prompt(
|
|
text=lower, prompt_hash=generate_prompt_hash(lower)
|
|
)
|
|
|
|
def test_the_library_counts_the_case_variant_as_this_prompt(self):
|
|
self.assertEqual(
|
|
self.db.get_prompt_by_hash(generate_prompt_hash(self.TEXT))["id"],
|
|
self.prompt_id,
|
|
)
|
|
|
|
|
|
class CachedRuntimeWhitespaceHistoryTests(CachedRuntimeHistoryTests):
|
|
"""Raw text with outer whitespace joins the history of the prompt it saves (#438).
|
|
|
|
The Library stores ``"a lighthouse the LLM wrote"``; the run keeps the exact
|
|
frozen raw text, newline included, yet is a tracked run of that prompt.
|
|
"""
|
|
|
|
TEXT = "\u3000 a lighthouse the LLM wrote\n"
|
|
|
|
def test_the_library_stores_the_stripped_text(self):
|
|
self.assertEqual(
|
|
self.db.get_prompt_by_id(self.prompt_id)["text"], self.TEXT.strip()
|
|
)
|