diff --git a/docs/resource-budgets.md b/docs/resource-budgets.md index 96803fd6..6db977a3 100644 --- a/docs/resource-budgets.md +++ b/docs/resource-budgets.md @@ -1860,3 +1860,100 @@ Synthetic tests use stub models and in-memory SQLite transaction controls. They do not measure native loading, MPS memory, allocator release or latency. The tier manager still retains its complete source list; bounded source enrollment, truthful activation and the original Mac memory incident remain unresolved. + +### Streamed inactive-pair worker primitive + +`RebuildWorker.run_source(target, plan)` consumes an already initialized source +plan and frozen embedding target. It prepares an encode-only lease outside SQL, +retains one pending source page, and publishes only an encoded prefix of that +page. Prefix text/vector buffers and the adapter's writer/retry validation reads +are also bounded by the requested page row limit. That limit defaults to 64; +it does not bound the bytes in an individual message. Completion text, separation +formatting and float serialization are unchanged. Inner encode batch size remains +32, with no added normalization, alternate model or reranker changes. + +The caller must own the clean connection, verified database identity, canonical +maintenance lease and exclusive inactive-pair writer. This entrypoint neither +initializes schemas nor clears data. Its durable progress is the paired manifest; +optional callbacks receive live counts and metrics without worker status SQL. +There are no registry, configuration, active-model or serving-pair writes. The +structured result reports commits during this call separately from the plan's +durable cursor and distinguishes captured-range coverage from current-source +coverage. `activated` is always false. Empty ranges still prepare and certify the +target and finish the source checkpoint. + +OOM recovery retains the actual failed prefix and the existing halving credits: +at a fixed committed cursor, initial failed size N allows at most +`2 * (floor(log2(N)) + 1)` failures. N=1 permits 2; N=8 permits 8. Failed arrays +and tracebacks are released before recovery hooks. Only a committed pair resets +credits. A clean rolled-back SQL OOM may require one bounded restart audit before +its retry; the generation, manifest and exact pending rows must still agree. +Revalidation does not reset credits. Ordinary errors, changed sources and +uncertain rollback propagate, without status writes or additional database work +after rollback uncertainty. Exceptions from SQLite, leases and hooks are not +translated into a result that could hide an ambiguous commit. +SQL-origin OOM requires checkpoint revalidation before any ordinary result, +including terminal exhaustion. If cancellation or deadline expiry prevents that +validation, publication uncertainty propagates without additional SQL. + +Cancellation and the 9000-second maximum deadline are checked at worker-owned +boundaries, including after preparation, throttle admission, encoding, recovery, +callbacks and publication. Remaining time is passed into each prepared lease +call. Cancellation cannot interrupt an active native or SQLite call, including +its internal ownership/writer wait. If a call commits before returning, the +returned plan preserves that committed prefix even when cancellation is then +observed. A stronger cancellable admission contract remains necessary before +claiming interruption after every internal wait. + +The existing manager still calls the original full-list worker path. This unused +primitive does not yet reduce production tier-switch source residency or provide +activation, background reopen ownership, inactive reset or active-force staging. +Native allocator behavior, latency and the original Mac memory incident remain +unproven by these synthetic tests. + +### Background tier job: immutable database admission + +`tier_switch/job.py` adds a frozen `TierJob` containing a generated job ID, the +explicit prepared embedding target, canonical existing database path, physical +file identity, tracker kind, source epoch and source-schema signature. It retains +no messages, connection, model or maintenance ownership token. Nothing calls this +API from the manager or public tier-switch path yet; this prerequisite does not +reduce their retained source list or change active configuration or retrieval. + +`capture_tier_job(conn, target)` requires a known-current, transaction-clean file +connection with source tracking already installed. It performs no tracking, +status, registry or nonce writes. TEMP objects, attached databases, missing files +and private/shared-memory databases are refused. Source identity reads use one +short owned read snapshot and end it before returning. Uncertain rollback raises; +the caller must abort and close that connection rather than publish status on it. + +`open_tier_job(job)` acquires the existing cooperative maintenance owner in its +calling worker thread, then opens an existing database through a quoted SQLite +file URI with `mode=rw`, never SQLite's create-on-open default. It checks the main +database pathname, source tracker/epoch/schema and read-only file observations +around reopening and before yielding a transaction-clean connection. Both the +owned connection and maintenance owner are released on `BaseException`; owner +release still runs if closing the connection fails. Guard errors contain fixed +categories rather than paths or source contents. `check_tier_job(conn, job)` +repeats admission before a future activation transaction. The descriptor's job ID +does not by itself attest any durable status row or currently selected job. + +Physical identity is `(st_dev, st_ino)` from read-only `os.open`/`os.fstat` +observations only. Mutable size, modification time and ctime are excluded; path +stat and descriptor stat ctime are never compared across Windows semantics. +Nonregular files and unavailable/zero inode identities fail closed. `O_NONBLOCK` +where available prevents waiting for a FIFO writer before its type is checked; +`O_BINARY` is used where available. This does not bound filesystem or SQLite I/O +duration, and unsupported filesystems must not be treated as certified. + +The caller must cooperate in keeping the pathname stable and supply a connection +known to refer to that current file. Python's SQLite API exposes no database file +descriptor, so these checks cannot prove that association for a stale connection, +a copied same-epoch database, or an adversarial replace-and-restore race. Future +public adoption must bind a durable job/status marker through the accepted +connection and validate it on reopen, or establish a stronger equivalent binding. +Final writer ownership, current-source freshness, job selection and atomic +activation remain separate requirements. Ordinary edits within the same epoch +remain the source adapter's responsibility. Local tests use mocked file handles +and in-memory SQLite; the synthetic file-backed worker test is opt-in for GPUBox. +No native-memory reduction or resolution of the original Mac incident is claimed. diff --git a/tests/test-tier-job-identity-795.py b/tests/test-tier-job-identity-795.py new file mode 100644 index 00000000..58a73e4c --- /dev/null +++ b/tests/test-tier-job-identity-795.py @@ -0,0 +1,528 @@ +"""Synthetic job admission; local tests use mocked files and in-memory SQLite.""" + +import builtins +import dataclasses +import os +import sqlite3 +import stat +import sys +import threading +import types +import unittest +from contextlib import contextmanager +from pathlib import Path +from unittest.mock import patch + +ROOT = Path(__file__).resolve().parents[1] +CONNECT = sqlite3.connect +SYNTHETIC_PATH = Path(Path.cwd().anchor) / "synthetic" / "tier database?#.db" + + +def load_modules() -> dict[str, types.ModuleType]: + modules = {} + + def safe_import(name, globals=None, locals=None, fromlist=(), level=0): + if name in {"numpy", "torch", "sentence_transformers", "model2vec", "sqlite_vec"}: + raise AssertionError("Native/model imports are excluded from job admission") + if name.startswith("truememory."): + short = name.removeprefix("truememory.") + if short in modules: + return modules[short] + raise AssertionError("Unexpected application import: " + name) + return builtins.__import__(name, globals, locals, fromlist, level) + + for name in ("storage", "_platform", "maintenance", "rebuild_source", "embedding_target", "tier_switch.job"): + module = types.ModuleType("synthetic_tier_job_" + name.replace(".", "_")) + sys.modules[module.__name__] = module + module.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) + path = ROOT / "truememory" / (name.replace(".", "/") + ".py") + exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), module.__dict__) + modules[name] = module + return modules + + +class MockFileConnection(sqlite3.Connection): + closed = False + fake_path = str(SYNTHETIC_PATH) + extra_database = False + rollback_failure = None + close_failure = None + + def execute(self, sql: str, parameters: tuple = ()): + if sql == "PRAGMA database_list": + rows = [(0, "main", self.fake_path)] + if self.extra_database: + rows.append((2, "extra", "/synthetic/attached.db")) + return rows + return super().execute(sql, parameters) + + def close(self) -> None: + self.closed = True + super().close() + if self.close_failure is not None: + raise self.close_failure + + def rollback(self) -> None: + if self.rollback_failure == "retained": + return + if self.rollback_failure is not None: + raise self.rollback_failure + super().rollback() + + +class Cancelled(BaseException): + pass + + +class TestTierJobIdentity(unittest.TestCase): + def setUp(self) -> None: + self.modules = load_modules() + self.api = self.modules["tier_switch.job"] + self.source = self.modules["rebuild_source"] + self.target = self.modules["embedding_target"].EmbeddingTarget("base", "qwen3_256", 256, "basepro") + self.conn = CONNECT(":memory:", factory=MockFileConnection) + self.addCleanup(self.conn.close) + self.conn.execute("CREATE TABLE messages(id INTEGER PRIMARY KEY, content TEXT)") + self.source.ensure_rebuild_tracking(self.conn) + self.identity = (7, 19) + self.observations = 0 + self.events = [] + + @contextmanager + def observe(path): + self.assertEqual(path, SYNTHETIC_PATH) + self.observations += 1 + self.events.append("file") + yield self.identity + + self.observe = observe + self.addCleanup(patch.stopall) + patch.object(self.api, "_observe_file", observe).start() + patch.object(self.api.Path, "resolve", lambda path, strict=False: path).start() + + def capture(self): + return self.api.capture_tier_job(self.conn, self.target) + + def reopen_hooks(self, *, before_open=None, opened=None): + @contextmanager + def owner(path): + self.assertEqual(path, SYNTHETIC_PATH) + self.events.append(("owner", threading.get_ident())) + try: + yield object() + finally: + self.events.append("release") + + reopened = [] + + def connect(database, **kwargs): + self.events.append("connect") + self.assertEqual(database, SYNTHETIC_PATH.as_uri() + "?mode=rw") + self.assertEqual(kwargs, {"uri": True}) + if before_open is not None: + before_open() + conn = CONNECT(":memory:", factory=MockFileConnection) + self.conn.backup(conn) + reopened.append(conn) + if opened is not None: + opened(conn) + return conn + + patch.object(self.api, "maintenance_owner", owner).start() + patch.object(self.api.sqlite3, "connect", connect).start() + return reopened + + def test_capture_is_frozen_read_only_and_retains_no_connection(self) -> None: + changes = self.conn.total_changes + job = self.capture() + self.assertEqual(self.conn.total_changes, changes) + self.assertFalse(self.conn.in_transaction) + self.assertEqual(job.target, self.target) + self.assertEqual(job.database_path, SYNTHETIC_PATH) + self.assertEqual(job.file_identity, (7, 19)) + self.assertEqual(job.tracker, "rebuild-bridge-v1") + self.assertNotEqual(job.job_id, self.capture().job_id) + self.assertNotIn(str(SYNTHETIC_PATH), repr(job)) + self.assertEqual({field.name for field in dataclasses.fields(job)}, { + "job_id", "target", "database_path", "file_identity", "tracker", "source_epoch", "source_schema_signature", + }) + with self.assertRaises(dataclasses.FrozenInstanceError): + job.job_id = "changed" + + def test_missing_tracking_never_installs_it(self) -> None: + conn = CONNECT(":memory:", factory=MockFileConnection) + self.addCleanup(conn.close) + conn.execute("CREATE TABLE messages(id INTEGER PRIMARY KEY, content TEXT)") + before = conn.total_changes + with self.assertRaises(self.api.TierJobIdentityError): + self.api.capture_tier_job(conn, self.target) + self.assertEqual(conn.total_changes, before) + self.assertFalse(conn.in_transaction) + self.assertEqual(conn.execute("SELECT name FROM sqlite_master WHERE type='table'").fetchall(), [("messages",)]) + + def test_borrowed_transaction_is_untouched(self) -> None: + job = self.capture() + self.conn.execute("BEGIN") + for call in (lambda: self.capture(), lambda: self.api.check_tier_job(self.conn, job)): + with self.assertRaises(self.api.TierJobIdentityError): + call() + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + + def test_memory_and_attachments_fail_before_file_observation(self) -> None: + for path, attached in (("", False), (str(SYNTHETIC_PATH), True)): + with self.subTest(path=path, attached=attached): + self.conn.fake_path = path + self.conn.extra_database = attached + with self.assertRaises(self.api.TierJobIdentityError): + self.capture() + self.assertEqual(self.observations, 0) + with CONNECT(":memory:") as real_memory: + with self.assertRaises(self.api.TierJobIdentityError): + self.api.capture_tier_job(real_memory, self.target) + real_memory.close() + + def test_temp_shadow_and_unrelated_temp_objects_fail_closed(self) -> None: + for ddl in ("CREATE TEMP TABLE MESSAGES(id)", "CREATE TEMP VIEW unrelated AS SELECT 1"): + with self.subTest(ddl=ddl): + conn = CONNECT(":memory:", factory=MockFileConnection) + try: + self.conn.backup(conn) + conn.execute(ddl) + with self.assertRaises(self.api.TierJobIdentityError): + self.api.capture_tier_job(conn, self.target) + finally: + conn.close() + self.assertEqual(self.observations, 0) + + def test_source_mutations_with_same_epoch_are_adapter_responsibility(self) -> None: + job = self.capture() + self.conn.execute("INSERT INTO messages VALUES (1,'synthetic source')") + self.conn.commit() + self.api.check_tier_job(self.conn, job) + self.conn.execute("UPDATE messages SET content='synthetic correction'") + self.conn.commit() + self.api.check_tier_job(self.conn, job) + self.assertFalse(self.conn.in_transaction) + + def test_epoch_schema_tracker_and_trigger_changes_refuse(self) -> None: + job = self.capture() + actions = ( + "UPDATE truememory_rebuild_source_v1 SET epoch='replacement'", + "ALTER TABLE messages ADD COLUMN synthetic TEXT", + "DROP TRIGGER truememory_rebuild_messages_v1_ai", + ) + for sql in actions: + with self.subTest(sql=sql): + conn = CONNECT(":memory:", factory=MockFileConnection) + try: + self.conn.backup(conn) + conn.execute(sql) + conn.commit() + with self.assertRaises(self.api.TierJobIdentityError): + self.api.check_tier_job(conn, job) + self.assertFalse(conn.in_transaction) + finally: + conn.close() + altered = dataclasses.replace(job, tracker="canonical-v1") + with self.assertRaises(self.api.TierJobIdentityError): + self.api.check_tier_job(self.conn, altered) + + def test_replaced_file_refuses_before_connect(self) -> None: + job = self.capture() + reopened = self.reopen_hooks() + self.identity = (7, 20) + with self.assertRaises(self.api.TierJobIdentityError): + with self.api.open_tier_job(job): + self.fail("Replacement must not yield") + self.assertEqual(reopened, []) + self.assertNotIn("connect", self.events) + self.assertEqual(self.events[-1], "release") + + def test_change_during_open_closes_before_yield(self) -> None: + job = self.capture() + reopened = self.reopen_hooks(before_open=lambda: setattr(self, "identity", (7, 20))) + with self.assertRaises(self.api.TierJobIdentityError): + with self.api.open_tier_job(job): + self.fail("Raced open must not yield") + self.assertEqual(len(reopened), 1) + self.assertTrue(reopened[0].closed) + self.assertEqual(self.events[-1], "release") + + def test_reopened_epoch_and_pragma_path_checked_before_yield(self) -> None: + job = self.capture() + for change in ("epoch", "path"): + def opened(conn): + if change == "epoch": + conn.execute("UPDATE truememory_rebuild_source_v1 SET epoch='replacement'") + conn.commit() + else: + conn.fake_path = "/synthetic/other.db" + reopened = self.reopen_hooks(opened=opened) + with self.assertRaises(self.api.TierJobIdentityError): + with self.api.open_tier_job(job): + self.fail("Unmatched reopened identity must not yield") + self.assertTrue(reopened[0].closed) + + def test_reopen_is_rw_only_owner_is_in_calling_worker_thread_and_cleanup_is_exact(self) -> None: + job = self.capture() + self.events.clear() + reopened = self.reopen_hooks() + sentinel = Cancelled("synthetic cancellation") + with self.assertRaises(Cancelled) as raised: + with self.api.open_tier_job(job) as conn: + self.assertFalse(conn.in_transaction) + self.api.check_tier_job(conn, job) + conn.execute("BEGIN") + conn.execute("INSERT INTO messages VALUES (1,'uncommitted synthetic row')") + raise sentinel + self.assertIs(raised.exception, sentinel) + self.assertEqual(self.events[0], ("owner", threading.get_ident())) + self.assertTrue(reopened[0].closed) + self.assertEqual(self.events[-1], "release") + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 0) + + def test_body_oserror_is_not_relabelled_as_identity_failure(self) -> None: + job = self.capture() + self.reopen_hooks() + sentinel = OSError("synthetic caller operation") + with self.assertRaises(OSError) as raised: + with self.api.open_tier_job(job): + raise sentinel + self.assertIs(raised.exception, sentinel) + + def test_open_and_source_failure_close_and_release(self) -> None: + job = self.capture() + self.reopen_hooks() + for error in (sqlite3.OperationalError("/private/path synthetic"), Cancelled("synthetic cancellation")): + with patch.object(self.api.sqlite3, "connect", side_effect=error): + expected = Cancelled if isinstance(error, Cancelled) else self.api.TierJobIdentityError + with self.assertRaises(expected) as caught: + with self.api.open_tier_job(job): + self.fail("Failed connect must not yield") + self.assertEqual(self.events[-1], "release") + if expected is self.api.TierJobIdentityError: + self.assertNotIn("/private", str(caught.exception)) + self.assertTrue(caught.exception.__suppress_context__) + reopened = self.reopen_hooks() + with patch.object(self.api, "read_rebuild_revision", side_effect=Cancelled("synthetic snapshot cancellation")): + with self.assertRaises(Cancelled): + with self.api.open_tier_job(job): + self.fail("Cancelled source read must not yield") + self.assertTrue(reopened[0].closed) + self.assertEqual(self.events[-1], "release") + + def test_snapshot_failure_rolls_back_and_does_not_log_source_error(self) -> None: + sentinel = sqlite3.OperationalError("synthetic identifying source detail") + with patch.object(self.api, "read_rebuild_revision", side_effect=sentinel): + with self.assertRaises(self.api.TierJobIdentityError) as caught: + self.capture() + self.assertFalse(self.conn.in_transaction) + self.assertNotIn("identifying", str(caught.exception)) + self.assertTrue(caught.exception.__suppress_context__) + + def test_uncertain_snapshot_rollback_propagates_and_owned_reopen_closes(self) -> None: + job = self.capture() + for failure in ("retained", sqlite3.OperationalError("synthetic rollback failure")): + self.conn.rollback_failure = failure + with self.assertRaises(self.api.TierJobIdentityError): + self.capture() + self.assertTrue(self.conn.in_transaction) + self.conn.rollback_failure = None + self.conn.rollback() + reopened = self.reopen_hooks(opened=lambda conn: setattr(conn, "rollback_failure", failure)) + with self.assertRaises(self.api.TierJobIdentityError): + with self.api.open_tier_job(job): + self.fail("Uncertain snapshot cleanup must not yield") + self.assertTrue(reopened[0].closed) + self.assertEqual(self.events[-1], "release") + + def test_connection_close_failure_still_releases_owner_and_is_sanitized(self) -> None: + job = self.capture() + reopened = self.reopen_hooks(opened=lambda conn: setattr( + conn, "close_failure", sqlite3.OperationalError("/private/synthetic.db"), + )) + with self.assertRaises(self.api.TierJobIdentityError) as caught: + with self.api.open_tier_job(job): + pass + self.assertTrue(reopened[0].closed) + self.assertEqual(self.events[-1], "release") + self.assertNotIn("/private", str(caught.exception)) + + def test_descriptor_handoff_acquires_owner_only_inside_worker(self) -> None: + job = self.capture() + dump = "\n".join(self.conn.iterdump()) + caller_thread = threading.get_ident() + owners = [] + results = [] + opened = [] + @contextmanager + def owner(path): + owners.append(threading.get_ident()) + try: + yield object() + finally: + owners.append("released") + def connect(database, **kwargs): + conn = CONNECT(":memory:", factory=MockFileConnection) + conn.executescript(dump) + opened.append(conn) + return conn + def worker(): + try: + with self.api.open_tier_job(job) as conn: + results.append((threading.get_ident(), conn.in_transaction)) + except BaseException as error: + results.append(type(error).__name__) + with patch.object(self.api, "maintenance_owner", owner), patch.object(self.api.sqlite3, "connect", connect): + thread = threading.Thread(target=worker) + thread.start() + thread.join(timeout=3) + self.assertFalse(thread.is_alive()) + self.assertEqual(len(results), 1) + self.assertEqual(results[0], (thread.ident, False)) + self.assertEqual(owners, [thread.ident, "released"]) + self.assertNotEqual(thread.ident, caller_thread) + self.assertTrue(opened[0].closed) + + def test_busy_owner_prevents_all_file_and_database_opening(self) -> None: + job = self.capture() + self.events.clear() + reopened = self.reopen_hooks() + with patch.object(self.api, "maintenance_owner", side_effect=self.modules["maintenance"].MaintenanceBusyError("synthetic busy")): + with self.assertRaises(self.modules["maintenance"].MaintenanceBusyError): + with self.api.open_tier_job(job): + self.fail("Busy owner must not yield") + self.assertEqual(reopened, []) + self.assertEqual(self.events, []) + + def test_missing_path_is_sanitized_and_no_reopen_occurs(self) -> None: + job = self.capture() + reopened = self.reopen_hooks() + with patch.object(self.api, "_observe_file", side_effect=FileNotFoundError("/private/synthetic.db")): + with self.assertRaises(self.api.TierJobIdentityError) as caught: + with self.api.open_tier_job(job): + self.fail("Missing database must not yield") + self.assertEqual(reopened, []) + self.assertNotIn("/private", str(caught.exception)) + + def test_capture_and_final_check_detect_second_observation_change(self) -> None: + original = self.api.read_rebuild_revision + def replace_after_read(conn): + result = original(conn) + self.identity = (7, 20) + return result + job = self.capture() + for operation in (lambda: self.capture(), lambda: self.api.check_tier_job(self.conn, job)): + self.identity = (7, 19) + with patch.object(self.api, "read_rebuild_revision", replace_after_read): + with self.assertRaises(self.api.TierJobIdentityError): + operation() + self.assertFalse(self.conn.in_transaction) + + def test_same_epoch_stale_connection_cannot_be_disproved_by_path_observations(self) -> None: + # Explicit limitation: capture assumes this connection still names the + # current file. A copied epoch plus a stale handle violates that contract. + self.identity = (7, 99) + job = self.capture() + reopened = self.reopen_hooks() + with self.api.open_tier_job(job) as conn: + self.api.check_tier_job(conn, job) + self.assertEqual(job.file_identity, (7, 99)) + self.assertTrue(reopened[0].closed) + + +class TestFileObservation(unittest.TestCase): + def test_same_handle_api_ignores_mutable_and_cross_api_ctime(self) -> None: + api = load_modules()["tier_switch.job"] + value = types.SimpleNamespace(st_mode=stat.S_IFREG, st_dev=5, st_ino=19, st_ctime_ns=10) + with patch.object(api.os, "open", return_value=81) as opened, patch.object(api.os, "fstat", return_value=value) as fstat, patch.object(api.os, "close") as closed: + with api._observe_file(SYNTHETIC_PATH) as first: + value.st_ctime_ns = 999 + with api._observe_file(SYNTHETIC_PATH) as second: + self.assertEqual(first, second) + self.assertEqual(first, (5, 19)) + self.assertEqual(fstat.call_count, 2) + self.assertEqual(opened.call_args.args, (SYNTHETIC_PATH, os.O_RDONLY | getattr(os, "O_NONBLOCK", 0) | getattr(os, "O_BINARY", 0))) + self.assertEqual(closed.call_count, 2) + self.assertEqual(closed.call_args.args, (81,)) + + def test_unsupported_identity_and_open_errors_fail_closed_without_path(self) -> None: + api = load_modules()["tier_switch.job"] + for mode, device, inode in ((stat.S_IFDIR, 5, 19), (stat.S_IFIFO, 5, 19), (stat.S_IFREG, 5, 0), (stat.S_IFREG, -1, 19), (stat.S_IFREG, True, 19)): + value = types.SimpleNamespace(st_mode=mode, st_dev=device, st_ino=inode) + with patch.object(api.os, "open", return_value=81), patch.object(api.os, "fstat", return_value=value), patch.object(api.os, "close") as closed: + with self.assertRaises(api.TierJobIdentityError): + with api._observe_file(SYNTHETIC_PATH): + self.fail("Unsupported identity must not yield") + self.assertEqual(closed.call_count, 1) + with patch.object(api.os, "open", side_effect=OSError("/private/synthetic.db")): + with self.assertRaises(api.TierJobIdentityError) as caught: + with api._observe_file(SYNTHETIC_PATH): + self.fail("Failed file observation must not yield") + self.assertNotIn("/private", str(caught.exception)) + + def test_observer_preserves_caller_baseexception_and_oserror(self) -> None: + api = load_modules()["tier_switch.job"] + value = types.SimpleNamespace(st_mode=stat.S_IFREG, st_dev=5, st_ino=19) + for sentinel in (Cancelled("synthetic"), OSError("synthetic"), ValueError("synthetic")): + with patch.object(api.os, "open", return_value=81), patch.object(api.os, "fstat", return_value=value), patch.object(api.os, "close") as closed: + with self.assertRaises(type(sentinel)) as raised: + with api._observe_file(SYNTHETIC_PATH): + raise sentinel + self.assertIs(raised.exception, sentinel) + self.assertEqual(closed.call_count, 1) + + def test_fstat_cancellation_and_failure_close_the_readonly_descriptor(self) -> None: + api = load_modules()["tier_switch.job"] + for sentinel in (Cancelled("synthetic"), OSError("synthetic")): + with patch.object(api.os, "open", return_value=81), patch.object(api.os, "fstat", side_effect=sentinel), patch.object(api.os, "close") as closed: + expected = Cancelled if isinstance(sentinel, Cancelled) else api.TierJobIdentityError + with self.assertRaises(expected): + with api._observe_file(SYNTHETIC_PATH): + self.fail("Failed observation must not yield") + self.assertEqual(closed.call_args.args, (81,)) + + +@unittest.skipUnless(os.environ.get("TRUEMEMORY_TIER_JOB_FILE_TEST") == "1", "opt-in synthetic file DB gate only") +class TestTierJobFileDatabase(unittest.TestCase): + def test_existing_database_reopen_replacement_missing_and_worker_ownership(self) -> None: + import tempfile + + modules = load_modules() + api = modules["tier_switch.job"] + target = modules["embedding_target"].EmbeddingTarget("base", "qwen3_256", 256, "basepro") + with tempfile.TemporaryDirectory(prefix="truememory-synthetic-job-") as directory: + path = Path(directory) / "synthetic.db" + conn = CONNECT(path) + try: + conn.execute("CREATE TABLE messages(id INTEGER PRIMARY KEY, content TEXT)") + modules["rebuild_source"].ensure_rebuild_tracking(conn) + job = api.capture_tier_job(conn, target) + finally: + conn.close() + outcomes = [] + def worker(): + with api.open_tier_job(job) as opened: + api.check_tier_job(opened, job) + outcomes.append(not opened.in_transaction) + thread = threading.Thread(target=worker) + thread.start() + thread.join(timeout=5) + self.assertFalse(thread.is_alive()) + self.assertEqual(outcomes, [True]) + archived = path.with_suffix(".old") + path.rename(archived) + with self.assertRaises(api.TierJobIdentityError): + with api.open_tier_job(job): + self.fail("Missing file must not yield") + self.assertFalse(path.exists()) + replacement = CONNECT(path) + replacement.close() + with self.assertRaises(api.TierJobIdentityError): + with api.open_tier_job(job): + self.fail("Replaced file must not yield") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-streamed-worker-756.py b/tests/test-tier-streamed-worker-756.py new file mode 100644 index 00000000..88a32e9c --- /dev/null +++ b/tests/test-tier-streamed-worker-756.py @@ -0,0 +1,646 @@ +"""Actual source checkpoints and worker flow with in-memory SQL and stub leases.""" + +import ast +import builtins +import gc +import math +import runpy +import sqlite3 +import struct +import sys +import types +import unittest +import weakref +from contextlib import contextmanager, nullcontext +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +SOURCE_TEST = runpy.run_path(str(Path(__file__).with_name("test-tier-rebuild-source-756.py"))) +TABLES = ("vec_messages_custom", "vec_messages_sep_custom") + + +class Workspace: + pass + + +class Vectors(list): + pass + + +class TestStreamedWorker(unittest.TestCase): + def setUp(self): + self.modules = SOURCE_TEST["load_modules"]() + self.source = self.modules["tier_switch.source"] + self.tracking = self.modules["rebuild_source"] + self.now = 0.0 + self.sizes = [64] + self.admissions = [] + self.attempts = [] + self.forwards = [] + self.deadlines = [] + self.refs = [] + self.cleanup = [] + self.sql = [] + self.committed = [] + self.live = [] + self.prepared = self.closed = 0 + self.fail_stage = "" + self.failed_attempts = {1} + self.error_factory = lambda: RuntimeError("synthetic out of memory") + self.rollback_mode = "" + self.commit_after_success_timeout = False + self.commit_after_success_oom = "" + self.armed = False + self.stage = "" + self.before_hook = lambda: None + self.after_hook = lambda: None + self.prepare_hook = lambda: None + self.encode_hook = lambda: None + self.cleanup_hook = lambda kind: None + self.write_hook = lambda: None + self.finish_hook = lambda: None + owner = self + + class Connection(sqlite3.Connection): + tier_write = False + + def execute(self, sql, parameters=()): + owner.sql.append(sql) + result = super().execute(sql, parameters) + if owner.armed: + if sql.startswith(f'INSERT INTO main."{TABLES[0]}"'): + self.tier_write = True + owner.write_hook() + owner.maybe_fail("insert_completion") + elif sql.startswith(f'INSERT INTO main."{TABLES[1]}"'): + owner.maybe_fail("insert_separation") + elif sql.startswith("INSERT OR REPLACE INTO main.metadata"): + owner.maybe_fail("checkpoint") + return result + + def commit(self): + owner.sql.append("COMMIT") + wrote = self.tier_write + if owner.armed and self.tier_write: + owner.maybe_fail("commit") + super().commit() + self.tier_write = False + if owner.armed and wrote and owner.commit_after_success_timeout: + owner.now = 9001.0 + raise TimeoutError("synthetic timeout after committed pair") + if owner.armed and wrote and owner.commit_after_success_oom: + if owner.commit_after_success_oom == "cancel": + owner.worker.cancel() + elif owner.commit_after_success_oom == "timeout": + owner.now = 9001.0 + raise RuntimeError("synthetic out of memory after committed pair") + + def rollback(self): + owner.sql.append("ROLLBACK") + if owner.rollback_mode == "dirty_raise": + raise MemoryError("synthetic out of memory during rollback") + if owner.rollback_mode == "dirty_return": + return + super().rollback() + self.tier_write = False + if owner.rollback_mode == "clean_raise": + raise MemoryError("synthetic out of memory after rollback") + if owner.rollback_mode == "clean_timeout": + owner.now = 9001.0 + raise TimeoutError("synthetic rollback timeout after cleanup") + + self.conn = sqlite3.connect(":memory:", factory=Connection) + self.addCleanup(self.conn.close) + self.conn.executescript(""" + CREATE TABLE messages(id INTEGER PRIMARY KEY, content TEXT, sender TEXT, recipient TEXT, timestamp TEXT); + CREATE TABLE metadata(key TEXT PRIMARY KEY, value TEXT NOT NULL); + CREATE TABLE vec_messages_custom(embedding BLOB); + CREATE TABLE vec_messages_sep_custom(embedding BLOB); + CREATE TABLE serving_pair(embedding BLOB); + INSERT INTO serving_pair(rowid,embedding) VALUES (42,x'0102'); + CREATE TABLE rebuild_status(id INTEGER PRIMARY KEY, status TEXT); + INSERT INTO rebuild_status VALUES (1,'untouched'); + """) + self.tracking.ensure_rebuild_tracking(self.conn) + + class Throttler: + batch_size = 64 + + def before_batch(self): + if len(owner.admissions) > 2000: + raise AssertionError("Unbounded synthetic loop") + if owner.sizes: + self.batch_size = owner.sizes.pop(0) + owner.admissions.append(self.batch_size) + owner.before_hook() + return self.batch_size, {"synthetic": True} + + def after_batch(self, count, duration): + owner.committed.append(count) + owner.after_hook() + + def on_oom(self): + owner.observe_cleanup("backoff") + + def should_flush_cache(self): + return False + + @staticmethod + def flush_gpu_cache(): + owner.observe_cleanup("flush") + + registry = types.ModuleType("synthetic_registry") + registry.VectorCacheRegistry = type("ForbiddenRegistry", (), {}) + throttler_module = types.ModuleType("synthetic_throttler") + throttler_module.DynamicThrottler = Throttler + self.modules["tier_switch.cache"] = registry + self.modules["tier_switch.throttler"] = throttler_module + self.torch = types.SimpleNamespace(no_grad=nullcontext, cuda=types.SimpleNamespace(OutOfMemoryError=type("CudaOOM", (RuntimeError,), {}))) + + def safe_import(name, globals=None, locals=None, fromlist=(), level=0): + if name == "torch": + return self.torch + if name in ("numpy", "sqlite_vec", "sentence_transformers"): + raise AssertionError("No native imports in streamed worker tests") + if name.startswith("truememory."): + short = name.removeprefix("truememory.") + if short in self.modules: + return self.modules[short] + raise AssertionError("Unexpected application import: " + name) + return builtins.__import__(name, globals, locals, fromlist, level) + + def load(name): + module = types.ModuleType("synthetic_streamed_" + name.replace(".", "_")) + sys.modules[module.__name__] = module + module.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) + path = ROOT / "truememory" / (name.replace(".", "/") + ".py") + exec(compile(path.read_text(), str(path), "exec"), module.__dict__) + self.modules[name] = module + return module + + target_module = load("embedding_target") + self.target = target_module.EmbeddingTarget("custom", "synthetic/worker", 2, "custom") + self.vector = types.ModuleType("synthetic_streamed_vector") + self.vector.__dict__.update(math=math, struct=struct, np=types.SimpleNamespace(ndarray=type("NativeArray", (), {}))) + tree = ast.parse((ROOT / "truememory/vector_search.py").read_text()) + tree.body = [node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name in {"serialize_f32", "_build_sep_text"}] + exec(compile(tree, "actual-serialization-and-separation", "exec"), self.vector.__dict__) + serialize = self.vector.serialize_f32 + + def serialized(value): + self.maybe_fail("serialize_" + self.stage) + return serialize(value) + + self.vector.serialize_f32 = serialized + self.active_model = object() + self.vector._model = self.active_model + self.vector.EMBEDDING_MODEL = "serving-model" + self.vector._embedding_dim = 99 + self.vector._model_generation = 17 + + class Lease: + def encode(self, texts, *, timeout, batch_size): + owner.assertFalse(owner.conn.in_transaction) + owner.assertGreater(timeout, 0) + owner.assertEqual(batch_size, 32) + owner.deadlines.append(timeout) + owner.stage = "completion" if texts[0].startswith("content-") else "separation" + if owner.stage == "completion": + owner.attempts.append(list(texts)) + owner.forwards.append((owner.stage, list(texts))) + owner.encode_hook() + owner.maybe_fail(owner.stage) + values = Vectors([[3., 4.]] * len(texts)) + owner.refs.append(weakref.ref(values)) + return values + + @contextmanager + def prepare(target, *, timeout): + self.assertEqual(target, self.target) + self.assertFalse(self.conn.in_transaction) + self.assertGreater(timeout, 0) + self.prepared += 1 + self.prepare_hook() + try: + yield Lease() + finally: + self.closed += 1 + + self.vector.prepare_embedding_target = prepare + self.modules["vector_search"] = self.vector + self.worker_module = load("tier_switch.worker") + self.worker_module.time = types.SimpleNamespace(monotonic=lambda: self.now, time=lambda: self.now) + self.worker = self.worker_module.RebuildWorker(self.conn, "custom", "custom", Throttler(), status_id=1, + status_callback=lambda count, total, metrics: self.live.append((count, total, metrics))) + self.worker._update_status = lambda *args, **kwargs: self.fail("Streamed worker wrote status") + self.worker._process_batch = lambda *args, **kwargs: self.fail("Streamed worker invoked full-list batch path") + self.worker._count_vectors = lambda *args, **kwargs: self.fail("Streamed worker recounted target vectors") + + def seed(self, ids): + self.conn.executemany("INSERT INTO messages VALUES (?,?,?,?,?)", [ + (mid, f"content-{mid}", "sender", "recipient", "2026-01-02T03:04:05") for mid in ids + ]) + self.conn.commit() + + def start(self): + return self.source.initialize_tier_source(self.conn, self.source.plan_tier_source( + self.conn, model_id=self.target.model_id, dimension=self.target.dimension, targets=TABLES, + )) + + def run_worker(self, count=8, *, page_size=64, plan=None): + if plan is None: + self.seed(range(1, count + 1)) + plan = self.start() + self.armed = True + return self.worker.run_source(self.target, plan, page_size=page_size) + + def maybe_fail(self, stage): + if stage == self.fail_stage and (self.failed_attempts is None or len(self.attempts) in self.failed_attempts): + workspace = Workspace() + self.refs.append(weakref.ref(workspace)) + raise self.error_factory() + + def observe_cleanup(self, kind): + gc.collect() + self.cleanup.append((kind, sys.exc_info()[0] is None, all(ref() is None for ref in self.refs))) + self.cleanup_hook(kind) + + def ids(self): + return tuple([row[0] for row in self.conn.execute(f'SELECT rowid FROM "{table}" ORDER BY rowid')] for table in TABLES) + + def test_order_metadata_exact_vectors_no_normalization_or_activation(self): + self.seed((-8, 0, 3, 9)) + result = self.run_worker(plan=self.start(), page_size=3) + self.assertEqual((result.status, result.processed, result.captured_range_complete, result.current_source_complete, result.activated), + ("complete", 4, True, True, False)) + self.assertEqual(self.ids(), ([-8, 0, 3, 9], [-8, 0, 3, 9])) + self.assertEqual(self.forwards, [ + ("completion", ["content--8", "content-0", "content-3"]), + ("separation", [f"sender to recipient on 2026-01-02: content-{mid}" for mid in (-8, 0, 3)]), + ("completion", ["content-9"]), ("separation", ["sender to recipient on 2026-01-02: content-9"]), + ]) + for table in TABLES: + self.assertEqual({row[0] for row in self.conn.execute(f'SELECT embedding FROM "{table}"')}, {struct.pack("2f", 3., 4.)}) + self.assertIs(self.vector._model, self.active_model) + self.assertEqual((self.vector.EMBEDDING_MODEL, self.vector._embedding_dim, self.vector._model_generation), ("serving-model", 99, 17)) + self.assertEqual(self.conn.execute("SELECT rowid,embedding FROM serving_pair").fetchone(), (42, b"\x01\x02")) + self.assertEqual(self.conn.execute("SELECT status FROM rebuild_status").fetchone()[0], "untouched") + self.assertTrue(all(not metrics["activated"] for _, _, metrics in self.live)) + self.assertEqual((self.prepared, self.closed), (1, 1)) + + def test_long_source_retains_only_bounded_pending_pages_without_full_fetch(self): + self.seed(range(1, 513)) + plan = self.start() + pages = [] + weak_pages = [] + read = self.source.read_tier_page + audit = self.source._audit_pair + audits = [] + + def bounded(conn, candidate, *, size=64): + self.assertLessEqual(size, 7) + gc.collect() + self.assertLessEqual(sum(ref() is not None for ref in weak_pages), 1) + page = read(conn, candidate, size=size) + pages.append(len(page.rows)) + weak_pages.append(weakref.ref(page)) + return page + + self.source.read_tier_page = bounded + self.source._audit_pair = lambda *args: (audits.append(True), audit(*args))[1] + result = self.run_worker(plan=plan, page_size=7) + self.assertEqual(result.processed, 512) + self.assertLessEqual(max(pages), 7) + self.assertLessEqual(max(map(len, self.attempts)), 7) + self.assertEqual(audits, []) + source_queries = [sql for sql in self.sql if sql.startswith("SELECT id, content")] + self.assertTrue(source_queries) + self.assertTrue(all("LIMIT ?" in sql for sql in source_queries)) + self.assertEqual(self.ids(), (list(range(1, 513)), list(range(1, 513)))) + + def test_empty_input_still_prepares_certifies_and_finishes(self): + result = self.run_worker(0) + self.assertEqual((self.prepared, self.closed, self.attempts), (1, 1, [])) + self.assertEqual((result.status, result.processed, result.current_source_complete), ("complete", 0, True)) + + def test_suffix_pair_failure_replans_exact_checkpoint_and_recovers(self): + self.sizes = [2, 4] + self.fail_stage = "insert_separation" + self.failed_attempts = {2} + plans = [] + original = self.source.plan_tier_source + + def replan(*args, **kwargs): + plan = original(*args, **kwargs) + plans.append((plan.manifest.generation, plan.manifest.cursor, plan.manifest.consumed)) + return plan + + self.seed(range(1, 7)) + plan = self.start() + self.source.plan_tier_source = replan + result = self.run_worker(plan=plan, page_size=6) + self.assertEqual(result.processed, 6) + self.assertEqual([batch[0] for batch in self.attempts], ["content-1", "content-3", "content-3"]) + self.assertEqual([len(batch) for batch in self.attempts], [2, 4, 4]) + self.assertEqual(plans, [(plan.manifest.generation, 2, 2)]) + self.assertEqual(self.ids(), (list(range(1, 7)), list(range(1, 7)))) + self.assertTrue(all(clear and released for _, clear, released in self.cleanup)) + + def test_tracebacks_and_arrays_release_before_all_recovery_cleanup(self): + for stage in ("completion", "separation", "serialize_completion", "serialize_separation", "insert_completion", "insert_separation", "checkpoint", "commit"): + with self.subTest(stage=stage): + case = TestStreamedWorker() + case.setUp() + try: + case.fail_stage = stage + result = case.run_worker(1) + self.assertEqual((result.status, result.processed), ("complete", 1)) + self.assertEqual(case.cleanup, [("backoff", True, True), ("flush", True, True)]) + finally: + case.doCleanups() + + def test_finite_halving_uses_actual_failed_prefix_and_cannot_renew_credits(self): + for admissions, expected in ( + ([32] * 20, [8, 8, 4, 4, 2, 2, 1, 1]), + (list(range(16, 0, -1)), None), + ([16, 32, 15, 32, 7, 32, 3, 32, 1, 32], None), + ([16, 1, 32], [8, 1, 1]), + ): + with self.subTest(admissions=admissions): + case = TestStreamedWorker() + case.setUp() + try: + case.sizes = admissions + case.fail_stage = "completion" + case.failed_attempts = None + result = case.run_worker(8) + counts = list(map(len, case.attempts)) + self.assertEqual((result.status, result.processed), ("oom_exhausted", 0)) + self.assertLessEqual(len(counts), 2 * (8).bit_length()) + self.assertEqual(counts[-2:], [1, 1]) + self.assertTrue(all(right <= left for left, right in zip(counts, counts[1:]))) + if expected is not None: + self.assertEqual(counts, expected) + finally: + case.doCleanups() + + def test_sql_oom_revalidation_does_not_reset_halving_credits(self): + self.fail_stage = "insert_separation" + self.failed_attempts = None + result = self.run_worker(8) + self.assertEqual((result.status, result.processed), ("oom_exhausted", 0)) + self.assertEqual(list(map(len, self.attempts)), [8, 8, 4, 4, 2, 2, 1, 1]) + self.assertEqual(self.ids(), ([], [])) + + def test_oom_budget_resets_only_after_committed_pair(self): + self.sizes = [8] * 20 + self.fail_stage = "completion" + self.failed_attempts = {1, 2, 4, 5, 6, 7, 8, 9} + result = self.run_worker(8) + self.assertEqual((result.status, result.processed), ("oom_exhausted", 4)) + self.assertEqual(list(map(len, self.attempts)), [8, 8, 4, 4, 4, 2, 2, 1, 1]) + self.assertEqual(self.ids(), ([1, 2, 3, 4], [1, 2, 3, 4])) + + def test_rollback_uncertainty_never_classifies_or_runs_later_sql(self): + for mode in ("dirty_raise", "clean_raise", "dirty_return"): + with self.subTest(mode=mode): + case = TestStreamedWorker() + case.setUp() + try: + case.seed((1,)) + plan = case.start() + case.fail_stage = "insert_separation" + case.rollback_mode = mode + case.worker._is_oom_error = lambda error: self.fail("Classified an uncertain rollback") + with self.assertRaises((MemoryError, RuntimeError)): + case.run_worker(plan=plan) + self.assertEqual(case.sql[-1], "ROLLBACK") + self.assertEqual(case.cleanup, []) + self.assertEqual((case.prepared, case.closed), (1, 1)) + finally: + case.doCleanups() + + def test_bare_memory_error_is_not_reclassified_as_retryable(self): + self.fail_stage = "separation" + self.error_factory = MemoryError + with self.assertRaises(MemoryError): + self.run_worker(1) + self.assertEqual(self.cleanup, []) + self.assertEqual(self.ids(), ([], [])) + + def test_expired_deadline_does_not_translate_uncertain_rollback_timeout(self): + self.fail_stage = "insert_separation" + self.rollback_mode = "clean_timeout" + with self.assertRaisesRegex(TimeoutError, "rollback timeout after cleanup"): + self.run_worker(1) + self.assertEqual(self.sql[-1], "ROLLBACK") + self.assertFalse(self.conn.in_transaction) + self.assertEqual(self.cleanup, []) + self.assertEqual((self.prepared, self.closed), (1, 1)) + + def test_commit_then_timeout_propagates_instead_of_returning_old_cursor(self): + self.seed((1,)) + plan = self.start() + self.commit_after_success_timeout = True + with self.assertRaisesRegex(TimeoutError, "timeout after committed pair"): + self.run_worker(plan=plan) + self.assertEqual(self.sql[-1], "ROLLBACK") + self.assertEqual(self.ids(), ([1], [1])) + durable = self.tracking.load_manifest(self.conn, self.tracking.manifest_key(TABLES)) + self.assertEqual((durable.cursor, durable.consumed), (1, 1)) + self.assertEqual(plan.manifest.consumed, 0) + self.assertEqual(self.cleanup, []) + self.assertEqual((self.prepared, self.closed), (1, 1)) + + def test_terminal_singleton_oom_cannot_hide_successful_pair_commit(self): + self.fail_stage = "completion" + self.failed_attempts = {1} + self.commit_after_success_oom = "exhausted" + with self.assertRaisesRegex(self.source.TierSourceUntrusted, "same source checkpoint"): + self.run_worker(1) + self.assertEqual(list(map(len, self.attempts)), [1, 1]) + self.assertEqual(self.ids(), ([1], [1])) + self.assertEqual(self.tracking.load_manifest(self.conn, self.tracking.manifest_key(TABLES)).consumed, 1) + + def test_cancel_or_deadline_after_ambiguous_commit_propagates_without_more_sql(self): + for stop in ("cancel", "timeout"): + with self.subTest(stop=stop): + case = TestStreamedWorker() + case.setUp() + try: + case.commit_after_success_oom = stop + with self.assertRaisesRegex(case.source.TierSourceUntrusted, "durable checkpoint revalidation"): + case.run_worker(1) + self.assertEqual(case.sql[-1], "ROLLBACK") + self.assertEqual(case.cleanup, []) + self.assertEqual(case.ids(), ([1], [1])) + self.assertEqual(len(case.attempts), 1) + finally: + case.doCleanups() + + def test_preparation_and_each_external_boundary_observe_cancellation(self): + for point in ("before_prepare", "prepare", "admission", "completion", "separation", "after_batch", "callback"): + with self.subTest(point=point): + case = TestStreamedWorker() + case.setUp() + try: + if point == "before_prepare": + case.worker.cancel() + elif point == "prepare": + case.prepare_hook = case.worker.cancel + elif point == "admission": + case.before_hook = case.worker.cancel + elif point in ("completion", "separation"): + case.encode_hook = lambda: case.worker.cancel() if case.stage == point else None + elif point == "after_batch": + case.after_hook = case.worker.cancel + else: + case.worker.status_callback = lambda *args: case.worker.cancel() + result = case.run_worker(2) + self.assertEqual(result.status, "cancelled") + self.assertFalse(result.activated) + self.assertEqual(result.processed, 2 if point in ("after_batch", "callback") else 0) + expected_calls = {"before_prepare": 0, "prepare": 0, "admission": 0, "completion": 1, "separation": 2, "after_batch": 2, "callback": 2} + self.assertEqual(len(case.forwards), expected_calls[point]) + self.assertEqual(case.closed, case.prepared) + finally: + case.doCleanups() + + def test_cancel_inside_sql_call_reports_committed_cursor_without_next_native(self): + self.write_hook = self.worker.cancel + self.sizes = [1] + result = self.run_worker(2) + self.assertEqual((result.status, result.processed, result.plan.manifest.cursor), ("cancelled", 1, 1)) + self.assertEqual(self.ids(), ([1], [1])) + self.assertEqual(len(self.forwards), 2) + + def test_cancellation_and_timeout_after_cleanup_prevent_retry(self): + for stop in ("cancelled", "timeout"): + for point in ("backoff", "flush", "admission"): + with self.subTest(stop=stop, point=point): + case = TestStreamedWorker() + case.setUp() + try: + case.fail_stage = "completion" + case.failed_attempts = None + + def stopped(): + if stop == "cancelled": + case.worker.cancel() + else: + case.now = 9001.0 + + case.cleanup_hook = lambda kind: stopped() if kind == point else None + case.before_hook = lambda: stopped() if point == "admission" and len(case.admissions) == 2 else None + result = case.run_worker(1) + self.assertEqual(result.status, stop) + self.assertEqual(len(case.attempts), 1) + self.assertEqual(result.processed, 0) + finally: + case.doCleanups() + + def test_timeout_after_prepare_and_during_encode_blocks_followup(self): + self.prepare_hook = lambda: setattr(self, "now", 9001.0) + result = self.run_worker(1) + self.assertEqual((result.status, len(self.forwards), self.closed), ("timeout", 0, 1)) + + def test_timeout_during_completion_stops_before_separation_and_sql(self): + self.encode_hook = lambda: setattr(self, "now", 9001.0) + result = self.run_worker(1) + self.assertEqual((result.status, len(self.forwards), result.processed), ("timeout", 1, 0)) + self.assertEqual(self.ids(), ([], [])) + self.assertEqual(self.closed, 1) + + def test_remaining_deadline_reaches_both_lease_calls(self): + self.encode_hook = lambda: setattr(self, "now", self.now + 10.0) + result = self.run_worker(1) + self.assertEqual(self.deadlines, [9000.0, 8990.0]) + self.assertEqual(result.status, "complete") + + def test_finish_cancellation_preserves_completed_manifest_truth(self): + finish = self.source.finish_tier_source + + def finish_then_cancel(*args): + result = finish(*args) + self.worker.cancel() + return result + + self.source.finish_tier_source = finish_then_cancel + result = self.run_worker(1) + self.assertEqual((result.status, result.processed, result.plan.manifest.complete), ("cancelled", 1, True)) + self.assertTrue(result.captured_range_complete) + self.assertTrue(result.current_source_complete) + self.assertFalse(result.activated) + + def test_revalidation_cannot_resume_another_generation(self): + self.seed((1,)) + plan = self.start() + self.fail_stage = "insert_separation" + original = self.source.plan_tier_source + + def switched(*args, **kwargs): + from dataclasses import replace + resumed = original(*args, **kwargs) + return replace(resumed, manifest=replace(resumed.manifest, generation="0" * 32)) + + self.source.plan_tier_source = switched + with self.assertRaisesRegex(self.source.TierSourceUntrusted, "same source checkpoint"): + self.run_worker(plan=plan) + self.assertEqual(len(self.attempts), 1) + self.assertEqual(self.ids(), ([], [])) + + def test_lease_call_cancellation_witness_does_not_claim_in_call_interruption(self): + admitted_native = [] + + def simulated_internal_wait_then_native(): + self.worker.cancel() + admitted_native.append("already-admitted-forward") + + self.encode_hook = simulated_internal_wait_then_native + result = self.run_worker(1) + self.assertEqual(admitted_native, ["already-admitted-forward"]) + self.assertEqual(len(self.forwards), 1) + self.assertEqual((result.status, result.processed), ("cancelled", 0)) + self.assertEqual(self.ids(), ([], [])) + + def test_borrowed_connection_and_identity_mismatch_fail_before_model_or_sql(self): + plan = self.start() + self.conn.execute("INSERT INTO metadata VALUES ('borrowed','pending')") + before = len(self.sql) + with self.assertRaisesRegex(RuntimeError, "clean owned"): + self.worker.run_source(self.target, plan) + self.assertEqual(len(self.sql), before) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.prepared, 0) + self.conn.rollback() + self.worker.target_group = "different" + with self.assertRaisesRegex(ValueError, "identities"): + self.worker.run_source(self.target, plan) + self.assertEqual(self.prepared, 0) + + def test_source_change_after_native_failure_refuses_before_retry(self): + self.fail_stage = "completion" + + def correct(kind): + if kind == "backoff": + self.conn.execute("UPDATE messages SET content='corrected' WHERE id=1") + self.conn.commit() + + self.cleanup_hook = correct + with self.assertRaises(self.tracking.RebuildSourceChanged): + self.run_worker(1) + self.assertEqual(len(self.attempts), 1) + self.assertEqual(self.ids(), ([], [])) + + def test_captured_completion_does_not_claim_new_tail_or_activation(self): + self.seed((1, 2)) + self.start() + self.seed((3,)) + plan = self.source.plan_tier_source(self.conn, model_id=self.target.model_id, dimension=2, targets=TABLES) + result = self.run_worker(plan=plan) + self.assertEqual((result.captured_range_complete, result.current_source_complete, result.activated), (True, False, False)) + self.assertEqual(self.ids(), ([1, 2], [1, 2])) + + +if __name__ == "__main__": + unittest.main() diff --git a/truememory/tier_switch/job.py b/truememory/tier_switch/job.py new file mode 100644 index 00000000..66e888d1 --- /dev/null +++ b/truememory/tier_switch/job.py @@ -0,0 +1,161 @@ +"""Immutable background job identity for a cooperatively stable database path. + +This is an admission guard, not activation or a lock against arbitrary pathname +replacement. Capture requires a known-current caller connection. SQLite's Python +API exposes no database file descriptor to prove that historical association. +""" + +from __future__ import annotations + +import os +import re +import sqlite3 +import stat +import uuid +from collections.abc import Iterator +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass, field +from pathlib import Path + +from truememory.embedding_target import EmbeddingTarget +from truememory.maintenance import MaintenanceUnavailableError, maintenance_owner +from truememory.rebuild_source import read_rebuild_revision + + +class TierJobIdentityError(RuntimeError): + """The background database identity cannot be established safely.""" + + +@dataclass(frozen=True) +class TierJob: + job_id: str + target: EmbeddingTarget + database_path: Path = field(repr=False) + file_identity: tuple[int, int] + tracker: str + source_epoch: str + source_schema_signature: str + + def __post_init__(self) -> None: + if (not isinstance(self.job_id, str) or re.fullmatch(r"[0-9a-f]{32}", self.job_id) is None + or not isinstance(self.target, EmbeddingTarget) + or not isinstance(self.database_path, Path) or not self.database_path.is_absolute() + or type(self.file_identity) is not tuple or len(self.file_identity) != 2 + or any(type(value) is not int for value in self.file_identity) + or self.file_identity[0] < 0 or self.file_identity[1] <= 0 + or self.tracker not in {"canonical-v1", "rebuild-bridge-v1", "rebuild-bridge-conservative-v1"} + or not isinstance(self.source_epoch, str) or not self.source_epoch + or not isinstance(self.source_schema_signature, str) + or re.fullmatch(r"[0-9a-f]{64}", self.source_schema_signature) is None): + raise TierJobIdentityError("Invalid background tier job descriptor") + + +@contextmanager +def _observe_file(path: Path) -> Iterator[tuple[int, int]]: + # Use fstat on every observation, including a fresh pathname reopen. + # ctime/mtime/size can change during legitimate writes and are not identity. + with ExitStack() as stack: + try: + flags = os.O_RDONLY | getattr(os, "O_NONBLOCK", 0) | getattr(os, "O_BINARY", 0) + descriptor = os.open(path, flags) + stack.callback(os.close, descriptor) + value = os.fstat(descriptor) + if (not stat.S_ISREG(value.st_mode) or type(value.st_dev) is not int + or type(value.st_ino) is not int or value.st_dev < 0 or value.st_ino <= 0): + raise TierJobIdentityError("Stable database file identity is unavailable") + except (OSError, ValueError): + raise TierJobIdentityError("Database file identity is unavailable") from None + yield value.st_dev, value.st_ino + + +def _connection_path(conn: sqlite3.Connection) -> Path: + if conn.in_transaction: + raise TierJobIdentityError("Background identity requires a clean connection") + databases = tuple(conn.execute("PRAGMA database_list")) + mains = [row[2] for row in databases if row[1] == "main"] + if len(mains) != 1 or not mains[0] or any(row[1] not in {"main", "temp"} for row in databases): + raise TierJobIdentityError("Background identity requires one existing file database") + # The tracker uses unqualified source tables. Refuse every TEMP object so + # temporary tables, views and triggers cannot redirect identity reads. + if conn.execute("SELECT 1 FROM sqlite_temp_master LIMIT 1").fetchone() is not None: + raise TierJobIdentityError("Temporary database objects make background identity ambiguous") + path = Path(mains[0]).expanduser().resolve(strict=True) + return path + + +def _source_identity(conn: sqlite3.Connection) -> tuple[str, str, str]: + conn.execute("BEGIN") + try: + tracker, revision, signature = read_rebuild_revision(conn) + return tracker, revision.epoch, signature + finally: + conn.rollback() + if conn.in_transaction: + raise TierJobIdentityError("Background identity could not release its read transaction") + + +def capture_tier_job(conn: sqlite3.Connection, target: EmbeddingTarget) -> TierJob: + """Capture a clean, known-current connection with tracking already installed. + + No connection or ownership token is retained. Private/shared-memory databases + cannot use this background reopen protocol. Caller owns stable-path admission. + """ + try: + path = _connection_path(conn) + with _observe_file(path) as identity: + tracker, epoch, signature = _source_identity(conn) + with _observe_file(path) as after: + if identity != after or _connection_path(conn) != path: + raise TierJobIdentityError("Database identity changed during job capture") + return TierJob(uuid.uuid4().hex, target, path, identity, tracker, epoch, signature) + except (OSError, ValueError, sqlite3.Error, MaintenanceUnavailableError): + raise TierJobIdentityError("Background database identity could not be captured") from None + + +def check_tier_job(conn: sqlite3.Connection, job: TierJob) -> None: + """Recheck before model admission or a future activation transaction. + + This short read snapshot ends here. Source-content freshness, current job + selection and final activation serialization remain separate caller duties. + """ + try: + if _connection_path(conn) != job.database_path: + raise TierJobIdentityError("Background database path changed") + with _observe_file(job.database_path) as before: + if before != job.file_identity: + raise TierJobIdentityError("Background database file changed") + if _source_identity(conn) != (job.tracker, job.source_epoch, job.source_schema_signature): + raise TierJobIdentityError("Background database source identity changed") + with _observe_file(job.database_path) as after: + if before != after or _connection_path(conn) != job.database_path: + raise TierJobIdentityError("Background database identity changed during validation") + except (OSError, ValueError, sqlite3.Error, MaintenanceUnavailableError): + raise TierJobIdentityError("Background database identity could not be validated") from None + + +def _close_connection(conn: sqlite3.Connection) -> None: + try: + conn.close() + except (OSError, sqlite3.Error): + raise TierJobIdentityError("Background database connection could not be closed") from None + + +@contextmanager +def open_tier_job(job: TierJob) -> Iterator[sqlite3.Connection]: + """Own and reopen an existing database in this worker thread, never create it. + + The connection and cooperative maintenance owner both end on BaseException. + The descriptor conveys no caller-thread maintenance ownership. + """ + with ExitStack() as stack: + try: + stack.enter_context(maintenance_owner(job.database_path)) + identity = stack.enter_context(_observe_file(job.database_path)) + if identity != job.file_identity: + raise TierJobIdentityError("Background database file changed before reopen") + conn = sqlite3.connect(job.database_path.as_uri() + "?mode=rw", uri=True) + stack.callback(_close_connection, conn) + check_tier_job(conn, job) + except (OSError, ValueError, sqlite3.Error, MaintenanceUnavailableError): + raise TierJobIdentityError("Background database could not be reopened") from None + yield conn diff --git a/truememory/tier_switch/worker.py b/truememory/tier_switch/worker.py index 7b21e3e7..b219346c 100644 --- a/truememory/tier_switch/worker.py +++ b/truememory/tier_switch/worker.py @@ -7,15 +7,21 @@ """ import logging +import math import sqlite3 import time from collections.abc import Iterator from contextlib import contextmanager -from typing import Callable +from typing import TYPE_CHECKING, Callable, NamedTuple from truememory.tier_switch.cache import VectorCacheRegistry from truememory.tier_switch.throttler import DynamicThrottler +if TYPE_CHECKING: + from truememory.embedding_target import EmbeddingTarget + from truememory.tier_switch.source import TierSourcePlan + from truememory.vector_search import _PreparedEmbeddingLease + log = logging.getLogger(__name__) _HARD_TIMEOUT = 9000 # 2.5 hours @@ -23,6 +29,43 @@ StatusCallback = Callable[[int, int, dict], None] +class SourceBuildResult(NamedTuple): + status: str + plan: "TierSourcePlan" + processed: int + captured_range_complete: bool = False + current_source_complete: bool = False + activated: bool = False + + +class _SourceStopped(Exception): + pass + + +def _encode_source_prefix(lease: "_PreparedEmbeddingLease", rows: tuple, count: int, + admit: Callable[[], float]) -> tuple[list[bytes], list[bytes]]: + from truememory.vector_search import _build_sep_text, serialize_f32 + + # Keep arrays inside this frame so a failed attempt releases its traceback + # before the outer loop performs allocator cleanup or revalidates SQL. + texts = [row[1] for row in rows[:count]] + embeddings = lease.encode(texts, timeout=admit(), batch_size=32) + admit() + if len(embeddings) != count: + raise ValueError("Completion embedding count does not match the rebuild prefix") + completion = [serialize_f32(vector) for vector in embeddings] + del embeddings, texts + sep_texts = [_build_sep_text(row[2], row[3], row[4], row[1]) for row in rows[:count]] + embeddings = lease.encode(sep_texts, timeout=admit(), batch_size=32) + admit() + if len(embeddings) != count: + raise ValueError("Separation embedding count does not match the rebuild prefix") + separation = [serialize_f32(vector) for vector in embeddings] + del embeddings, sep_texts + admit() + return completion, separation + + @contextmanager def _rebuild_transaction(conn: sqlite3.Connection) -> Iterator[None]: owned = not conn.in_transaction @@ -64,6 +107,190 @@ def cancel(self): """Signal the worker to stop after the current batch.""" self._cancelled = True + def run_source( + self, target: "EmbeddingTarget", plan: "TierSourcePlan", *, + page_size: int = 64, timeout: float = _HARD_TIMEOUT, + ) -> SourceBuildResult: + """Build an initialized inactive pair without activating or clearing it. + + The caller owns maintenance, the clean connection and exclusive pair + writer, and has validated target schema and database identity. At most + one source page plus encoded prefix buffers is retained. Cancellation is + cooperative between calls: it cannot interrupt native/SQLite calls or + their internal waits. Deadlines also flow into the prepared lease. + """ + from truememory.embedding_target import EmbeddingTarget + from truememory.tier_switch.source import ( + TierSourcePlan, TierSourceUntrusted, finish_tier_source, plan_tier_source, + publish_tier_prefix, read_tier_page, + ) + from truememory.vector_search import prepare_embedding_target + + if not isinstance(target, EmbeddingTarget) or not isinstance(plan, TierSourcePlan): + raise TypeError("An explicit frozen target and initialized source plan are required") + if (plan.connection is not self.conn or plan.action not in ("resume", "complete") + or plan.manifest.model != target.model_id or plan.manifest.dimension != target.dimension + or plan.manifest.targets != target.tables or self.target_tier != target.tier + or self.target_group != target.tier_group): + raise ValueError("Worker, initialized source and prepared target identities must agree") + if type(page_size) is not int or page_size < 1: + raise ValueError("Source page size must be a positive integer") + if isinstance(timeout, bool) or not isinstance(timeout, (int, float)) or not math.isfinite(timeout) or timeout <= 0: + raise ValueError("Source timeout must be finite and positive") + if self.conn.in_transaction: + raise RuntimeError("RebuildWorker.run_source requires a clean owned connection") + expires = time.monotonic() + min(timeout, _HARD_TIMEOUT) + initial_consumed = plan.manifest.consumed + phase_ceiling = effective_ceiling = None + phase_failures = 0 + captured_complete = current_complete = False + publication_uncertain = False + + def admit() -> float: + if self.conn.in_transaction: + raise RuntimeError("Source worker connection did not return to a clean boundary") + if self._cancelled: + raise _SourceStopped("cancelled") + remaining = expires - time.monotonic() + if remaining <= 0: + raise _SourceStopped("timeout") + return remaining + + def result(status: str) -> SourceBuildResult: + if publication_uncertain: + raise TierSourceUntrusted("Failed publication requires durable checkpoint revalidation") + return SourceBuildResult(status, plan, plan.manifest.consumed - initial_consumed, + captured_complete, current_complete) + + try: + admit() + # A stale plan must fail before preparation, even for an empty range. + pending = read_tier_page(self.conn, plan, size=page_size) + admit() + try: + import torch + no_grad = torch.no_grad() + except ImportError: + from contextlib import nullcontext + no_grad = nullcontext() + with no_grad, prepare_embedding_target(target, timeout=admit()) as lease: + admit() + while pending.rows: + batch_size, metrics = self.throttler.before_batch() + admit() + if type(batch_size) is not int or batch_size < 1: + raise ValueError("Throttle admission must grant a positive integer prefix") + count = min(batch_size, len(pending.rows)) + if effective_ceiling is not None: + count = min(count, effective_ceiling) + started = time.monotonic() + oom = False + publishing = False + completion = separation = None + try: + completion, separation = _encode_source_prefix(lease, pending.rows, count, admit) + admit() + publishing = True + plan, pending = publish_tier_prefix( + self.conn, plan, pending, completion=completion, separation=separation, + ) + except _SourceStopped: + raise + except Exception as error: + # No DB queries, status updates or retries after uncertain + # rollback. The caller must discard this connection. + if (self.conn.in_transaction or (publishing and + (error.__context__ is not None or error.__cause__ is not None))): + raise + if not self._is_oom_error(error): + raise + publication_uncertain = publishing + oom = True + finally: + completion = separation = None + + if oom: + if phase_ceiling is None: + phase_ceiling = effective_ceiling = count + phase_failures = 1 + else: + effective_ceiling = min(effective_ceiling, count) + if count <= phase_ceiling // 2: + phase_ceiling = count + phase_failures = 1 + else: + phase_failures += 1 + exhausted = phase_ceiling == 1 and phase_failures == 2 + if phase_failures == 2 and not exhausted: + phase_ceiling //= 2 + effective_ceiling = min(effective_ceiling, phase_ceiling) + phase_failures = 0 + admit() + self.throttler.on_oom() + admit() + DynamicThrottler.flush_gpu_cache() + admit() + if exhausted and not publication_uncertain: + return result("oom_exhausted") + if self.conn.total_changes != plan.changes: + # Rolled-back SQL contributes to total_changes. One + # audit per failed attempt, bounded by OOM credits. + resumed = plan_tier_source( + self.conn, model_id=target.model_id, dimension=target.dimension, + targets=target.tables, + ) + admit() + if resumed.action != "resume" or resumed.manifest != plan.manifest: + raise TierSourceUntrusted("Failed publication no longer has the same source checkpoint") + plan = resumed + checked = read_tier_page(self.conn, plan, size=len(pending.rows)) + admit() + # The adapter checks exact SQLite types again under the + # writer. Retain the original immutable pending suffix. + from truememory.tier_switch.source import _same_rows + if checked.generation != pending.generation or checked.cursor != pending.cursor or not _same_rows(checked.rows, pending.rows): + raise TierSourceUntrusted("Failed publication no longer has the same pending source rows") + del checked + publication_uncertain = False + if exhausted: + return result("oom_exhausted") + continue + + # Capture the committed checkpoint before observing a + # cancellation that arrived during SQLite's writer wait. + phase_ceiling = effective_ceiling = None + phase_failures = 0 + admit() + self.throttler.after_batch(count, time.monotonic() - started) + admit() + flush = self.throttler.should_flush_cache() + admit() + if flush: + DynamicThrottler.flush_gpu_cache() + admit() + if self.status_callback is not None: + live = dict(metrics, batch_size=count, activated=False) + try: + self.status_callback(plan.manifest.consumed, plan.manifest.total, live) + except Exception: + pass + admit() + if not pending.rows: + pending = read_tier_page(self.conn, plan, size=page_size) + admit() + admit() + finished = finish_tier_source(self.conn, plan) + plan = finished.plan + captured_complete = finished.captured_range_complete + current_complete = finished.current_source_complete + admit() + return SourceBuildResult( + "complete", plan, plan.manifest.consumed - initial_consumed, + finished.captured_range_complete, finished.current_source_complete, + ) + except _SourceStopped as stopped: + return result(str(stopped)) + def run( self, messages: list[dict],