Vito Sansevero
94d50d15f2
History writes deliberately give up after a bounded 250 ms wait so they never stall ComfyUI; the existing bounded-timeout test still proves that. The concurrent replay tests assert that every peer succeeds, which is a serialization and first-capture property, so they now raise the budget to the shared 30 s busy timeout while the first writer is held open. The budget is one named constant, EXECUTION_BUSY_TIMEOUT_MS, shared by the transaction helper and the background writer. Before: 20/20 'database is locked' with the held writer. After: 0/20.
440 lines
16 KiB
Python
440 lines
16 KiB
Python
"""Frozen output occurrences belong to an exact stored execution attempt."""
|
|
|
|
import sqlite3
|
|
import unittest
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from contextlib import contextmanager
|
|
from dataclasses import replace
|
|
from threading import Barrier
|
|
from unittest.mock import patch
|
|
|
|
from database.execution_history import HistoryEpochChanged
|
|
from utils.execution_outputs import OutputBundle, OutputReference
|
|
|
|
try:
|
|
from tests import test_execution_history_database as history_tests
|
|
except ImportError:
|
|
import test_execution_history_database as history_tests
|
|
|
|
|
|
def bundle(
|
|
filename="first.png", *, root="a" * 64, node="9", ordinal=0, path_style="posix"
|
|
):
|
|
return OutputBundle(
|
|
(OutputReference(node, ordinal, "output", "", filename, path_style, root),),
|
|
"complete" if root else "partial",
|
|
"reported" if root else "root_unavailable",
|
|
)
|
|
|
|
|
|
class OutputDatabaseTests(unittest.TestCase):
|
|
setUp = history_tests.HistoryDatabaseTests.setUp
|
|
close_databases = history_tests.HistoryDatabaseTests.close_databases
|
|
database = history_tests.HistoryDatabaseTests.database
|
|
start = history_tests.HistoryDatabaseTests.start
|
|
finish = history_tests.HistoryDatabaseTests.finish
|
|
|
|
def capture(self, db, value=None, item=0):
|
|
start = self.start(db, item)
|
|
db.write_execution_start(start)
|
|
self.assertTrue(
|
|
db.write_execution_outputs_for_finish(self.finish(start), value or bundle())
|
|
)
|
|
return start
|
|
|
|
def rows(self, db):
|
|
return [
|
|
tuple(row)
|
|
for row in db.model.get_connection().execute(
|
|
"SELECT * FROM execution_run_outputs ORDER BY run_id,"
|
|
"output_node_id,image_ordinal"
|
|
)
|
|
]
|
|
|
|
def test_exact_attempt_parent_not_recomputed_uuid(self):
|
|
db = self.database()
|
|
stored_id = "55555555-5555-4555-8555-555555555555"
|
|
start = replace(self.start(db), run_id=stored_id)
|
|
db.write_execution_start(start)
|
|
self.assertTrue(
|
|
db.write_execution_outputs_for_finish(self.finish(start), bundle())
|
|
)
|
|
self.assertEqual(self.rows(db)[0][0], stored_id)
|
|
self.assertEqual(db.get_execution_outputs(stored_id)["state"], "complete")
|
|
other = self.capture(db, item=1)
|
|
self.assertEqual(len(self.rows(db)), 2)
|
|
self.assertEqual(
|
|
db.get_execution_outputs(other.run_id)["references"][0]["filename"],
|
|
"first.png",
|
|
)
|
|
self.assertEqual(db.get_execution_run(stored_id)["outcome"], "observed")
|
|
|
|
def test_exact_replay_noop_and_conflict_sticky(self):
|
|
db = self.database()
|
|
first = bundle()
|
|
refs = (replace(first.references[0], output_node_id="2"),) + (
|
|
replace(first.references[0], image_ordinal=0),
|
|
replace(first.references[0], image_ordinal=1),
|
|
)
|
|
first = replace(first, references=refs)
|
|
start = self.capture(db, first)
|
|
finish = self.finish(start)
|
|
original = self.rows(db)
|
|
# Construct identical DTOs from dictionaries with different insertion orders.
|
|
reordered = tuple(
|
|
OutputReference(**dict(reversed(list(vars(ref).items())))) for ref in refs
|
|
)
|
|
self.assertTrue(
|
|
db.write_execution_outputs_for_finish(
|
|
finish, replace(first, references=reordered)
|
|
)
|
|
)
|
|
self.assertEqual(db.get_execution_outputs(start.run_id)["state"], "complete")
|
|
self.assertEqual(self.rows(db), original)
|
|
self.assertTrue(
|
|
db.write_execution_outputs_for_finish(finish, bundle("other.png"))
|
|
)
|
|
capture = db.get_execution_outputs(start.run_id)
|
|
self.assertEqual(
|
|
(capture["state"], capture["reason"]), ("conflicting", "conflicting")
|
|
)
|
|
self.assertEqual(
|
|
[r["filename"] for r in capture["references"]], ["first.png"] * 3
|
|
)
|
|
self.assertEqual(capture["references"][0]["context_membership"], "unknown")
|
|
for replay in (first, bundle("third.png")):
|
|
db.write_execution_outputs_for_finish(finish, replay)
|
|
self.assertEqual(db.get_execution_outputs(start.run_id), capture)
|
|
self.assertEqual(self.rows(db), original)
|
|
|
|
def test_header_and_every_reference_scalar_participate_in_replay(self):
|
|
variants = (
|
|
replace(bundle(), state="partial", reason="unsupported"),
|
|
bundle(root="b" * 64),
|
|
bundle(node="8"),
|
|
bundle(ordinal=1),
|
|
replace(
|
|
bundle(),
|
|
references=(replace(bundle().references[0], folder_type="temp"),),
|
|
),
|
|
replace(
|
|
bundle(),
|
|
references=(replace(bundle().references[0], subfolder="nested"),),
|
|
),
|
|
replace(
|
|
bundle(),
|
|
references=(replace(bundle().references[0], path_style="windows"),),
|
|
),
|
|
OutputBundle((), "complete", "none_reported"),
|
|
)
|
|
db = self.database()
|
|
for item, variant in enumerate(variants):
|
|
with self.subTest(item=item):
|
|
start = self.capture(db, item=item)
|
|
db.write_execution_outputs_for_finish(self.finish(start), variant)
|
|
capture = db.get_execution_outputs(start.run_id)
|
|
self.assertEqual(capture["state"], "conflicting")
|
|
self.assertEqual(capture["references"][0]["root_key"], "a" * 64)
|
|
self.assertEqual(capture["references"][0]["output_node_id"], "9")
|
|
|
|
def test_concurrent_replay_keeps_first_rows(self):
|
|
db = self.database()
|
|
start = self.start(db)
|
|
db.write_execution_start(start)
|
|
finish = self.finish(start)
|
|
peers = [self.database() for _ in range(4)]
|
|
barrier = Barrier(len(peers))
|
|
hold = history_tests.FirstWriterHold()
|
|
|
|
@contextmanager
|
|
def held_open():
|
|
hold()
|
|
yield True
|
|
|
|
def write(index):
|
|
barrier.wait(timeout=60)
|
|
return peers[index].write_execution_outputs_for_finish(
|
|
finish,
|
|
bundle("first.png" if index % 2 else "second.png"),
|
|
generation_guard=held_open,
|
|
)
|
|
|
|
with (
|
|
history_tests.queued_writer_budget(),
|
|
ThreadPoolExecutor(max_workers=4) as executor,
|
|
):
|
|
self.assertEqual(list(executor.map(write, range(4))), [True] * 4)
|
|
capture = db.get_execution_outputs(start.run_id)
|
|
self.assertEqual(
|
|
(capture["state"], capture["reason"]), ("conflicting", "conflicting")
|
|
)
|
|
self.assertEqual(len(capture["references"]), 1)
|
|
frozen = self.rows(db)
|
|
for value in (bundle(), bundle("second.png")):
|
|
db.write_execution_outputs_for_finish(finish, value)
|
|
self.assertEqual(self.rows(db), frozen)
|
|
|
|
def test_epoch_and_generation_guard_cover_commit(self):
|
|
db = self.database()
|
|
start = self.start(db)
|
|
db.write_execution_start(start)
|
|
observations = []
|
|
conn = db.model.get_connection()
|
|
active = False
|
|
|
|
@contextmanager
|
|
def guard():
|
|
nonlocal active
|
|
self.assertTrue(conn.in_transaction)
|
|
active = True
|
|
try:
|
|
yield True
|
|
finally:
|
|
observations.append(("exit", conn.in_transaction))
|
|
active = False
|
|
|
|
def trace(sql):
|
|
if sql == "COMMIT":
|
|
observations.append(("commit", active))
|
|
|
|
conn.set_trace_callback(trace)
|
|
self.addCleanup(conn.set_trace_callback, None)
|
|
self.assertTrue(
|
|
db.write_execution_outputs_for_finish(
|
|
self.finish(start), bundle(), generation_guard=guard
|
|
)
|
|
)
|
|
self.assertEqual(observations, [("commit", True), ("exit", False)])
|
|
self.assertFalse(active)
|
|
|
|
@contextmanager
|
|
def stale():
|
|
yield False
|
|
|
|
other = self.start(db, item=1)
|
|
db.write_execution_start(other)
|
|
self.assertFalse(
|
|
db.write_execution_outputs_for_finish(
|
|
self.finish(other), bundle(), generation_guard=stale
|
|
)
|
|
)
|
|
self.assertEqual(
|
|
db.get_execution_outputs(other.run_id)["state"], "not_captured"
|
|
)
|
|
old_epoch = start.storage_epoch
|
|
with conn:
|
|
conn.execute(
|
|
"UPDATE execution_history_metadata SET storage_epoch=?",
|
|
("66666666-6666-4666-8666-666666666666",),
|
|
)
|
|
with self.assertRaises(HistoryEpochChanged):
|
|
db.write_execution_outputs_for_finish(
|
|
self.finish(start), bundle("late.png")
|
|
)
|
|
self.assertEqual(db.get_execution_run(start.run_id)["storage_epoch"], old_epoch)
|
|
self.assertEqual(
|
|
db.get_execution_outputs(start.run_id)["references"][0]["filename"],
|
|
"first.png",
|
|
)
|
|
|
|
def test_failure_rolls_back_header_and_keeps_guard_until_rollback(self):
|
|
db = self.database()
|
|
start = self.start(db)
|
|
db.write_execution_start(start)
|
|
conn = db.model.get_connection()
|
|
with conn:
|
|
conn.execute(
|
|
"CREATE TRIGGER reject_output BEFORE INSERT ON execution_run_outputs "
|
|
"BEGIN SELECT RAISE(ABORT,'injected'); END"
|
|
)
|
|
states = []
|
|
|
|
@contextmanager
|
|
def guard():
|
|
states.append(conn.in_transaction)
|
|
try:
|
|
yield True
|
|
finally:
|
|
states.append(conn.in_transaction)
|
|
|
|
with self.assertRaises(sqlite3.IntegrityError):
|
|
db.write_execution_outputs_for_finish(
|
|
self.finish(start), bundle(), generation_guard=guard
|
|
)
|
|
self.assertEqual(states, [True, False])
|
|
self.assertEqual(
|
|
db.get_execution_outputs(start.run_id)["state"], "not_captured"
|
|
)
|
|
self.assertEqual(self.rows(db), [])
|
|
|
|
def test_missing_header_or_parent_rejected(self):
|
|
db = self.database()
|
|
start = self.start(db)
|
|
self.assertFalse(
|
|
db.write_execution_outputs_for_finish(self.finish(start), bundle())
|
|
)
|
|
self.assertEqual(
|
|
db.get_execution_outputs(start.run_id),
|
|
{
|
|
"state": "not_captured",
|
|
"reason": None,
|
|
"format_version": None,
|
|
"source": None,
|
|
"references": [],
|
|
},
|
|
)
|
|
db.write_execution_start(start)
|
|
conn = db.model.get_connection()
|
|
conn.execute("PRAGMA foreign_keys=OFF")
|
|
self.assertEqual(conn.execute("PRAGMA foreign_keys").fetchone()[0], 0)
|
|
with conn:
|
|
conn.execute(
|
|
"INSERT INTO execution_run_outputs "
|
|
"VALUES(?, '9',0,'output','','orphan.png','posix',?)",
|
|
(start.run_id, "a" * 64),
|
|
)
|
|
with self.assertRaises(sqlite3.DatabaseError):
|
|
db.get_execution_outputs(start.run_id)
|
|
with self.assertRaises(sqlite3.DatabaseError):
|
|
db.write_execution_outputs_for_finish(self.finish(start), bundle())
|
|
self.assertEqual(
|
|
conn.execute("SELECT count(*) FROM execution_output_capture").fetchone()[0],
|
|
0,
|
|
)
|
|
with conn:
|
|
conn.execute("DELETE FROM execution_run_outputs")
|
|
db.write_execution_outputs_for_finish(self.finish(start), bundle())
|
|
with conn:
|
|
conn.execute("DELETE FROM execution_runs")
|
|
with self.assertRaises(sqlite3.DatabaseError):
|
|
db.get_execution_outputs(start.run_id)
|
|
|
|
def test_reads_never_access_files_or_rebind(self):
|
|
db = self.database()
|
|
value = bundle(root=None)
|
|
value = replace(
|
|
value,
|
|
references=(
|
|
replace(
|
|
value.references[0],
|
|
folder_type="temp",
|
|
subfolder=" Exact Ω ",
|
|
filename="Keep Case.png",
|
|
),
|
|
),
|
|
)
|
|
start = self.capture(db, value)
|
|
conn = db.model.get_connection()
|
|
before = list(conn.iterdump())
|
|
with (
|
|
patch("builtins.open", side_effect=AssertionError("file read")),
|
|
patch("os.stat", side_effect=AssertionError("stat")),
|
|
patch("os.path.realpath", side_effect=AssertionError("resolve")),
|
|
):
|
|
capture = db.get_execution_outputs(start.run_id)
|
|
self.assertEqual(
|
|
capture,
|
|
{
|
|
"state": "partial",
|
|
"reason": "root_unavailable",
|
|
"format_version": 1,
|
|
"source": "host_final_history",
|
|
"references": [
|
|
{
|
|
"output_node_id": "9",
|
|
"image_ordinal": 0,
|
|
"folder_type": "temp",
|
|
"subfolder": " Exact Ω ",
|
|
"filename": "Keep Case.png",
|
|
"path_style": "posix",
|
|
"root_key": None,
|
|
"file_identity": "unverified",
|
|
"newly_generated": "unknown",
|
|
"context_membership": "unknown",
|
|
"transient": True,
|
|
}
|
|
],
|
|
},
|
|
)
|
|
self.assertEqual(list(conn.iterdump()), before)
|
|
for table in (
|
|
"prompts",
|
|
"prompt_pairs",
|
|
"generated_images",
|
|
"execution_run_contexts",
|
|
):
|
|
self.assertEqual(
|
|
conn.execute(f"SELECT count(*) FROM {table}").fetchone()[0], 0
|
|
)
|
|
|
|
def test_read_uses_consistent_snapshot_and_preserves_callers_transaction(self):
|
|
db = self.database()
|
|
start = self.capture(db)
|
|
peer = self.database()
|
|
conn = db.model.get_connection()
|
|
changed = False
|
|
|
|
def trace(sql):
|
|
nonlocal changed
|
|
if "FROM execution_run_outputs" in sql and not changed:
|
|
changed = True
|
|
with peer.model.get_connection() as writer:
|
|
writer.execute(
|
|
"UPDATE execution_run_outputs SET filename='later.png'"
|
|
)
|
|
|
|
conn.set_trace_callback(trace)
|
|
self.addCleanup(conn.set_trace_callback, None)
|
|
self.assertEqual(
|
|
db.get_execution_outputs(start.run_id)["references"][0]["filename"],
|
|
"first.png",
|
|
)
|
|
conn.set_trace_callback(None)
|
|
self.assertTrue(changed)
|
|
self.assertFalse(conn.in_transaction)
|
|
self.assertEqual(
|
|
db.get_execution_outputs(start.run_id)["references"][0]["filename"],
|
|
"later.png",
|
|
)
|
|
conn.execute("BEGIN")
|
|
db.get_execution_outputs(start.run_id)
|
|
self.assertTrue(conn.in_transaction)
|
|
conn.rollback()
|
|
|
|
def test_empty_unavailable_and_partial_headers_remain_distinct(self):
|
|
db = self.database()
|
|
for item, (state, reason) in enumerate(
|
|
(
|
|
("complete", "none_reported"),
|
|
("unavailable", "unavailable"),
|
|
("partial", "unsupported"),
|
|
)
|
|
):
|
|
start = self.capture(db, OutputBundle((), state, reason), item)
|
|
capture = db.get_execution_outputs(start.run_id)
|
|
self.assertEqual(
|
|
(capture["state"], capture["reason"], capture["references"]),
|
|
(state, reason, []),
|
|
)
|
|
|
|
def test_invalid_scalars_rejected_before_transaction(self):
|
|
db = self.database()
|
|
start = self.start(db)
|
|
db.write_execution_start(start)
|
|
trace = []
|
|
conn = db.model.get_connection()
|
|
conn.set_trace_callback(trace.append)
|
|
self.addCleanup(conn.set_trace_callback, None)
|
|
for invalid in (
|
|
replace(
|
|
bundle(),
|
|
references=(replace(bundle().references[0], filename="../unsafe.png"),),
|
|
),
|
|
replace(bundle(), source="filesystem_scan"),
|
|
):
|
|
with self.assertRaises(ValueError):
|
|
db.write_execution_outputs_for_finish(self.finish(start), invalid)
|
|
with self.assertRaises(ValueError):
|
|
db.write_execution_outputs_for_finish(start, bundle())
|
|
self.assertEqual(trace, [])
|