From 0a41eec163c0bacc2564d3f42fe0071b7c1f0112 Mon Sep 17 00:00:00 2001 From: buildingjoshbetter Date: Fri, 9 Oct 2026 08:25:54 -0500 Subject: [PATCH] fix: Prepare certified inactive embedding targets --- docs/resource-budgets.md | 67 ++- tests/test-prepared-target-795.py | 653 ++++++++++++++++++++++++++++++ truememory/embedding_target.py | 124 ++++++ truememory/model_client.py | 47 +++ truememory/model_server.py | 72 +++- truememory/mps_utils.py | 29 +- truememory/vector_search.py | 235 ++++++++++- 7 files changed, 1214 insertions(+), 13 deletions(-) create mode 100644 tests/test-prepared-target-795.py create mode 100644 truememory/embedding_target.py diff --git a/docs/resource-budgets.md b/docs/resource-budgets.md index aa3f2f53..8cf5e638 100644 --- a/docs/resource-budgets.md +++ b/docs/resource-budgets.md @@ -322,8 +322,9 @@ default alone does not invalidate an already published empty-tier snapshot. `base`, `pro` and `qwen3_256` reuse the same loaded Qwen3 instance; `edge` and `model2vec` reuse the same loaded Model2Vec instance. These are the two deterministic identities already allowed by the fast lane. Alias updates -publish one immutable cache snapshot. A failed distinct load preserves the -previous snapshot. +publish one immutable cache snapshot. The later resident-replacement checkpoint +supersedes the original failure behavior: distinct loads release the old cache +before construction, and a failed replacement leaves that cache empty. The removed `qwen3` name still raises the existing migration error, including case and whitespace variants. This validation does not mutate the default or @@ -533,3 +534,65 @@ Synthetic tests cover ordinary errors, memory errors and cancellation, retained identity, ownership release and weak references after exception tracebacks are released. This bounds failed Python cache retention, not native allocator reclamation, CPU-transfer peak memory or the original Mac memory incident. + +### Certified inactive embedding targets (#795) + +`EmbeddingTarget.capture(tier)` freezes the requested public tier, effective +model ID, dimension and table group. Its exact completion/separation table pair +is derived from that validated group. `prepare_embedding_target(target)` yields +an encode-only context lease without changing the active `vector_search` model, +dimension or generation, config, environment, metadata or registry. Daemon cache +residency generations still change when resident models are replaced. Call preparation outside SQLite +transactions; this API has no connection argument and cannot discover an +unrelated transaction held by its caller. No tier-switch caller adopts it yet. + +Shared preparation uses additive `prepare_embed_target_v1` and +`embed_target_v1` operations. Older daemons must reject these operations; the +client raises a protocol error and never falls back to ordinary embedding. +Every successful response carries the exact descriptor. Preparation encodes one fixed +synthetic validation string under the existing inference/state owners, batch +policy, cooperative RSS admission, deadline and CPU-recovery path. Every actual +output slice must have the captured width, including CPU retries. Model dimension +getters and configured truncation bounds alone cannot certify a target. Empty +encode requests also validate the descriptor and a probe, returning shape +`(0, dimension)` only after successful certification. + +The daemon uses its existing single main cache and release-before-load policy; +prepared requests do not add a target cache or fast CPU clone. Deterministic +built-ins can reuse the same resident weights. Legacy custom cache labels are +insufficient evidence because their constructor may have fallen back. Strict +custom construction uses frozen arguments and rechecks configuration/download +permission after ownership waits and before construction. It never relabels a +fallback as the requested model. Legacy `minilm`, `bge-small` and removed `qwen3` +identities are rejected by this new API; their ordinary loader behavior is +unchanged. A receipt certifies the embedding identity at that operation, not +persistent residency, a source generation or database readiness. Later requests +may reload after eviction or reject changed configuration. + +Explicit local mode permits one outstanding preparation lease per process. It +borrows a matching active built-in model, or owns at most one additional target. +Construction and inference do not hold the active-model state lock. Local +encoding shares the existing exact-instance ownership registry; optional +deadlines cover state, lease and model-owner admission and are checked again +before native work. The ownership helper's ordinary callers retain unbounded +waiting. Cleanup unregisters users even after timeout. Local target encoding +does not introduce another OOM-retry policy. Closing rejects new calls before +waiting for active encoding, clears its model reference, then releases the +preparation slot. An interrupted closer can retry, and active encoding drains +the pending close itself. Same-target leases do not get independent native +owners. This is a reference-count policy, not a byte bound or a guarantee that +external references and exception tracebacks have disappeared. + +`init_prepared_target_tables(conn, target)` creates or validates only the exact +named pair in `main`. It rejects temporary objects shadowing either name and +requires the canonical `vec0` float dimension and cosine metric. Existing table +contents are neither cleared nor certified. A failed pair rolls back together; +an existing caller transaction remains open under a savepoint. No active-table +migration or active metadata write occurs. Future publishers must retain main +schema qualification or recheck shadowing and must separately attest source, +target generation and content identity. + +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. diff --git a/tests/test-prepared-target-795.py b/tests/test-prepared-target-795.py new file mode 100644 index 00000000..c4930d5c --- /dev/null +++ b/tests/test-prepared-target-795.py @@ -0,0 +1,653 @@ +"""Prepared targets through real control flow, stdlib stubs and in-memory SQL.""" +from __future__ import annotations + +import ast +from contextlib import contextmanager +from dataclasses import FrozenInstanceError +import gc +import math +import os +from pathlib import Path +import re +import runpy +import sqlite3 +import sys +import threading +import time +import types +import unittest +from unittest.mock import Mock, patch +import weakref + + +ROOT = Path(__file__).resolve().parents[1] +SOURCE_TEXTS: dict[str, str] = {} +BASE = runpy.run_path(str(Path(__file__).with_name("test-tier-target-model-795.py"))) + + +def load(name: str, *, definitions: bool = False, namespace: dict | None = None) -> types.ModuleType: + path = ROOT / "truememory" / f"{name}.py" + tree = ast.parse(SOURCE_TEXTS.get(name, path.read_text())) + if definitions: + tree.body = [ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), + *[node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.ClassDef))]] + tree = ast.fix_missing_locations(tree) + module = types.ModuleType(f"truememory.{name}") + if namespace: + module.__dict__.update(namespace) + sys.modules[module.__name__] = module + exec(compile(tree, str(path), "exec"), module.__dict__) + return module + + +class Array(BASE["Array"]): + def __getitem__(self, key: slice) -> Array: + return Array(self.rows[key], (len(self.rows[key]), self.shape[1])) + + +class Numpy(BASE["Numpy"]): + @staticmethod + def empty(shape: tuple[int, int], **_kwargs: object) -> Array: + return Array([[] for _ in range(shape[0])], shape) + + +class VecConnection: + """Actual SQLite transactions, with only native vec0 DDL replaced by a table.""" + + def __init__(self) -> None: + self.db = sqlite3.connect(":memory:") + self.schemas: dict[str, str] = {} + self.fail_table: str | None = None + self.extensions: list[bool] = [] + + @property + def in_transaction(self) -> bool: + return self.db.in_transaction + + def enable_load_extension(self, enabled: bool) -> None: + self.extensions.append(enabled) + + def execute(self, sql: str, parameters: tuple = ()): + match = re.fullmatch(r'CREATE VIRTUAL TABLE main."(\w+)" USING (vec0\(.+\))', sql) + if match: + name = match[1] + if name == self.fail_table: + raise sqlite3.OperationalError("synthetic second-table denial") + result = self.db.execute(f'CREATE TABLE main."{name}"(embedding BLOB)') + self.schemas[name] = f'CREATE VIRTUAL TABLE "{name}" USING {match[2]}' + return result + result = self.db.execute(sql, parameters) + if sql == "SELECT type, sql FROM main.sqlite_master WHERE name = ?": + row = result.fetchone() + if row and parameters[0] in self.schemas: + row = (row[0], self.schemas[parameters[0]]) + return types.SimpleNamespace(fetchone=lambda: row) + return result + + def commit(self) -> None: + self.db.commit() + + def rollback(self) -> None: + self.db.rollback() + + def close(self) -> None: + self.db.close() + + +class TestPreparedTargets(unittest.TestCase): + def setUp(self) -> None: + self.fixture = BASE["TestTierTargetIdentity"]() + self.fixture.setUp() + self.addCleanup(self.fixture.doCleanups) + self.target_module = load("embedding_target") + self.config = load("tier_config") + environment = patch.dict("os.environ", { + "TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD": "1", + "TRUEMEMORY_CUSTOM_EMBED_MODEL": "synthetic/encoder", + "TRUEMEMORY_CUSTOM_EMBED_DIM": "384", + }) + environment.start() + self.addCleanup(environment.stop) + self.Target = self.target_module.EmbeddingTarget + self.target = self.Target.capture("base") + self.calls: list[tuple[str, dict]] = [] + self.encoded: list[tuple[str, list[str], dict]] = [] + self.width: int | None = None + self.build_hook = lambda: None + self.encode_hook = lambda: None + self.refs: list[weakref.ReferenceType] = [] + owner = self + + class Model: + def __init__(self, identity: str, **kwargs: object): + owner.refs.append(weakref.ref(self)) + owner.build_hook() + owner.calls.append((identity, kwargs)) + self.identity = identity + self.dimension = kwargs.get("truncate_dim", 256) + self.device = kwargs.get("device", "cpu") + self.oom_once = False + + def encode(self, texts: list[str], **kwargs: object) -> Array: + owner.encode_hook() + owner.encoded.append((self.identity, list(texts), kwargs)) + if self.oom_once: + self.oom_once = False + raise RuntimeError("MPS backend out of memory") + width = self.dimension if owner.width is None else owner.width + return Array([[float(index)] * width for index in range(len(texts))], (len(texts), width)) + + def to(self, device: str) -> None: + self.device = device + + self.Model = Model + self.mps = load("mps_utils") + self.mps.ensure_mps_memory_budget = Mock() + self.mps.resolve_device = lambda _device: "cpu" + self.mps.flush_mps_cache = Mock() + modules = patch.dict(sys.modules, { + "model2vec": types.SimpleNamespace(StaticModel=types.SimpleNamespace(from_pretrained=Model)), + "sentence_transformers": types.SimpleNamespace(SentenceTransformer=Model), + "sqlite_vec": types.SimpleNamespace(load=Mock()), + }) + modules.start() + self.addCleanup(modules.stop) + self.ms = self.fixture.ms + self.ms.np = Numpy() + if "model_server" in SOURCE_TEXTS: + tree = ast.parse(SOURCE_TEXTS["model_server"]) + tree.body = [node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.ClassDef))] + with patch.dict(sys.modules, {self.ms.__name__: self.ms}): + exec(compile(tree, "synthetic-server-override", "exec"), self.ms.__dict__) + self.server = self.ms.ModelServer() + self.addCleanup(self.server._workers.shutdown, wait=True) + self.server._write_status_file = lambda: None + self.server._SUSTAINED_THRESHOLD = 1000 + self.client = load("model_client", definitions=True, namespace={"np": Numpy()}) + self.client._request_with_autostart = lambda request, timeout=None: self.server.handle_request(request) + self.client.use_model_server = lambda: False + self.vector = load("vector_search", definitions=True, namespace={ + "np": Numpy(), "threading": threading, "contextmanager": contextmanager, + "time": time, "math": math, "re": re, "sqlite3": sqlite3, + "_lock": threading.Lock(), "_prepared_target_slot": threading.Lock(), + "_model": None, "EMBEDDING_MODEL": "model2vec", "_embedding_dim": 256, + "_model_generation": 71, "_VEC_DISTANCE_METRIC": "cosine", + }) + self.vector._TIER_ALIASES = {"base": "qwen3_256", "pro": "qwen3_256", "edge": "model2vec"} + self.assertTrue(hasattr(self.vector, "prepare_embedding_target"), "Prepared lease API missing") + self.assertTrue(hasattr(self.client, "PreparedEmbeddingProxy"), "Certified client protocol missing") + self.assertTrue(hasattr(self.server, "_get_prepared_embed_model"), "Certified server ownership missing") + + def request(self, *, target=None, prepare: bool = True, texts=None, **fields: object) -> dict: + target = self.target if target is None else target + request = {"op": "prepare_embed_target_v1" if prepare else "embed_target_v1", + "target": target.to_wire(), "batch_size": 8, **fields} + if not prepare: + request["texts"] = ["synthetic"] if texts is None else texts + return self.server.handle_request(request) + + def globals(self) -> tuple: + return (self.vector.EMBEDDING_MODEL, self.vector._embedding_dim, + self.vector._model_generation, self.vector._model) + + def test_descriptor_is_frozen_strict_and_aliases_share_effective_identity(self) -> None: + self.assertEqual(self.target.identity, self.Target.capture("pro").identity) + self.assertEqual(self.target.tables, ("vec_messages_basepro", "vec_messages_sep_basepro")) + self.assertEqual(self.Target.from_wire(self.target.to_wire()), self.target) + with self.assertRaises(FrozenInstanceError): + self.target.dimension = 4 + for mutation in ({"dimension": True}, {"version": True}, {"tier_group": "x; DROP TABLE x"}, + {"model_id": "synthetic/other"}, {"extra": 1}): + with self.subTest(mutation=mutation), self.assertRaises(ValueError): + self.Target.from_wire({**self.target.to_wire(), **mutation}) + + def test_server_certifies_probe_and_every_batch_without_global_mutation(self) -> None: + before = self.globals() + response = self.request() + self.assertTrue(response["ok"]) + self.assertEqual(response["target"], self.target.to_wire()) + self.assertEqual(self.encoded[0][1], [self.target_module.TARGET_PROBE_TEXT]) + result = self.request(prepare=False, texts=["synthetic"] * 19) + self.assertEqual(result["vectors"].shape, (19, 256)) + self.assertEqual([len(entry[1]) for entry in self.encoded], [1, 8, 8, 3]) + self.assertEqual(len(self.calls), 1) + self.assertEqual(self.globals(), before) + self.assertFalse(self.server._inference_lock.locked()) + + def test_builtin_alias_reuses_resident_model_without_second_constructor(self) -> None: + self.request() + resident = self.server._embed_state.model + self.request(target=self.Target.capture("pro")) + self.assertIs(self.server._embed_state.model, resident) + self.assertEqual(len(self.calls), 1) + + def test_custom_legacy_fallback_is_replaced_and_frozen_args_are_used(self) -> None: + target = self.Target.capture("custom") + stale = self.Model("wrong-fallback") + self.server._embed_state = self.ms._EmbedState(stale, "custom", target.model_id) + stale_ref = weakref.ref(stale) + stale = None + self.build_hook = lambda: self.assertIsNone(stale_ref()) + response = self.request(target=target) + self.assertTrue(response["ok"]) + identity, kwargs = self.calls[-1] + self.assertEqual(identity, "synthetic/encoder") + self.assertEqual(kwargs["truncate_dim"], 384) + self.assertIs(kwargs["trust_remote_code"], False) + + def test_wrong_actual_width_cannot_certify_or_return_vectors(self) -> None: + target = self.Target.capture("custom") + self.width = 128 + for preparing in (True, False): + with self.subTest(preparing=preparing), self.assertRaises(self.target_module.EmbeddingTargetError): + self.request(target=target, prepare=preparing) + self.assertFalse(self.server._inference_lock.locked()) + self.assertFalse(self.server._lock.locked()) + + def test_custom_config_mismatch_and_revoked_permission_refuse_even_empty_input(self) -> None: + target = self.Target.capture("custom") + for setting in ({"TRUEMEMORY_CUSTOM_EMBED_DIM": "512"}, {"TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD": "0"}): + with self.subTest(setting=setting), patch.dict("os.environ", setting), self.assertRaises(ValueError): + self.request(target=target, prepare=False, texts=[]) + self.assertEqual(self.calls, []) + + def test_constructor_direct_call_rechecks_permission_and_deadline(self) -> None: + target = self.Target.capture("custom") + with patch.dict("os.environ", {"TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD": "0"}), self.assertRaises(ValueError): + self.target_module.build_target_model(target, "cpu") + with self.assertRaises(TimeoutError): + self.target_module.build_target_model(target, "cpu", check=Mock(side_effect=TimeoutError("expired"))) + self.assertEqual(self.calls, []) + + def test_empty_encode_still_certifies_and_preserves_requested_width(self) -> None: + proxy = self.client.PreparedEmbeddingProxy(self.target) + result = proxy.encode([]) + self.assertEqual(result.shape, (0, 256)) + self.assertEqual(len(self.encoded), 1) + self.width = 2 + with self.assertRaises(self.target_module.EmbeddingTargetError): + proxy.encode([]) + + def test_old_daemon_missing_or_mismatched_receipt_never_falls_back(self) -> None: + proxy = self.client.PreparedEmbeddingProxy(self.target) + for response in ({"ok": False, "error": "Unknown op: prepare_embed_target_v1"}, + {"ok": True}, {"ok": True, "target": self.Target.capture("edge").to_wire()}): + with self.subTest(response=response): + request = Mock(return_value=response) + self.client._request_with_autostart = request + with self.assertRaises(self.client.ProtocolMismatchError): + proxy.prepare() + self.assertEqual(request.call_count, 1) + self.assertEqual(request.call_args.args[0]["op"], "prepare_embed_target_v1") + + def test_expired_or_rss_refused_preparation_starts_no_constructor(self) -> None: + expired = self.request(deadline=time.time() - 1) + self.assertFalse(expired["ok"]) + self.server._max_rss_bytes = 1 + self.ms.psutil = types.SimpleNamespace(Process=lambda: types.SimpleNamespace(memory_info=lambda: types.SimpleNamespace(rss=1))) + refused = self.request() + self.assertEqual(refused["error_code"], "server_busy") + self.assertEqual(self.calls, []) + self.assertFalse(self.server._inference_lock.locked()) + + def test_expiry_during_rss_sensor_prevents_strict_constructor(self) -> None: + clock = [1000.0] + self.ms.time = types.SimpleNamespace(time=lambda: clock[0], monotonic=lambda: clock[0]) + self.server._max_rss_bytes = 2 + + def sample(): + clock[0] = 1002.0 + return types.SimpleNamespace(rss=1) + + self.ms.psutil = types.SimpleNamespace(Process=lambda: types.SimpleNamespace(memory_info=sample)) + result = self.request(deadline=1001.0) + self.assertFalse(result["ok"]) + self.assertIn("deadline", result["error"]) + self.assertEqual(self.calls, []) + self.assertFalse(self.server._lock.locked()) + self.assertFalse(self.server._inference_lock.locked()) + + def test_prepared_server_recovery_uses_sticky_cpu_and_one_resident_model(self) -> None: + self.request() + resident = self.server._embed_state.model + resident.device = "mps" + resident.oom_once = True + result = self.request(prepare=False) + self.assertTrue(result["ok"]) + self.assertIs(self.server._embed_state.model, resident) + self.assertEqual(resident.device, "cpu") + self.assertIn("embed", self.server._sticky_cpu) + self.assertEqual(len(self.calls), 1) + + def test_local_lease_is_exclusive_and_never_changes_active_globals(self) -> None: + before = self.globals() + self.build_hook = lambda: self.assertFalse(self.vector._lock.locked()) + self.encode_hook = lambda: self.assertFalse(self.vector._lock.locked()) + with self.vector.prepare_embedding_target(self.target) as lease: + self.assertEqual(lease.target, self.target) + self.assertEqual(lease.encode(["synthetic"]).shape, (1, 256)) + with self.assertRaises(RuntimeError): + with self.vector.prepare_embedding_target(self.target): + self.fail("Second lease admitted") + lease.close() + with self.assertRaises(RuntimeError): + lease.encode(["synthetic"]) + self.assertEqual(self.globals(), before) + self.assertFalse(self.vector._prepared_target_slot.locked()) + gc.collect() + self.assertTrue(all(ref() is None for ref in self.refs)) + + def test_shared_context_lease_uses_certified_protocol_without_local_slot(self) -> None: + self.client.use_model_server = lambda: True + before = self.globals() + with self.vector.prepare_embedding_target(self.target) as lease: + self.assertFalse(self.vector._prepared_target_slot.locked()) + self.assertEqual(lease.encode(["synthetic"] * 3, batch_size=2).shape, (3, 256)) + self.assertEqual([len(entry[1]) for entry in self.encoded], [1, 2, 1]) + self.assertIsNone(lease._encoder) + self.assertIsNotNone(self.server._embed_state) + self.assertEqual(self.globals(), before) + + def test_matching_local_active_builtin_is_borrowed_without_constructor(self) -> None: + active = self.Model("Qwen/Qwen3-Embedding-0.6B", truncate_dim=256) + self.vector._model = active + self.vector.EMBEDDING_MODEL = "qwen3_256" + before = self.globals() + with self.vector.prepare_embedding_target(self.target) as lease: + lease.encode(["synthetic"]) + self.assertEqual(len(self.calls), 1) + self.assertEqual(self.globals(), before) + + def test_local_constructor_and_probe_baseexception_release_slot(self) -> None: + for stage in ("construction", "probe"): + with self.subTest(stage=stage): + error = KeyboardInterrupt("synthetic cancellation") + failing = Mock(side_effect=error) + self.build_hook = failing if stage == "construction" else lambda: None + self.encode_hook = failing if stage == "probe" else lambda: None + with self.assertRaises(KeyboardInterrupt) as caught: + with self.vector.prepare_embedding_target(self.target): + self.fail("Cancelled preparation returned a lease") + self.assertIs(caught.exception, error) + self.assertFalse(self.vector._prepared_target_slot.locked()) + self.assertFalse(self.vector._lock.locked()) + error.__traceback__ = None + failing.side_effect = None + self.build_hook = self.encode_hook = lambda: None + gc.collect() + self.assertTrue(all(ref() is None for ref in self.refs)) + + def test_active_state_wait_uses_preparation_deadline(self) -> None: + lock = Mock() + lock.acquire.return_value = False + self.vector._lock = lock + with self.assertRaises(TimeoutError): + with self.vector.prepare_embedding_target(self.target, timeout=0.25): + self.fail("Busy active state admitted") + self.assertGreater(lock.acquire.call_args.kwargs["timeout"], 0) + self.assertLessEqual(lock.acquire.call_args.kwargs["timeout"], 0.25) + self.assertEqual(self.calls, []) + self.assertFalse(self.vector._prepared_target_slot.locked()) + + def test_local_native_owner_and_registry_waits_respect_same_deadline(self) -> None: + with self.vector.prepare_embedding_target(self.target) as lease: + model = lease._encoder + encoded = len(self.encoded) + with self.mps._own_model(model): + with self.assertRaises(TimeoutError): + lease.encode(["synthetic"], timeout=0.01) + self.assertEqual(self.mps._model_owners[id(model)].users, 1) + self.assertEqual(self.mps._model_owners, {}) + with self.mps._device_lock: + with self.assertRaises(TimeoutError): + lease.encode(["synthetic"], timeout=0.01) + self.assertEqual(self.mps._model_owners, {}) + self.assertEqual(len(self.encoded), encoded) + self.assertFalse(lease._encode_lock.locked()) + self.assertEqual(lease.encode(["synthetic"]).shape, (1, 256)) + + def test_legacy_model_owner_preserves_context_manager_lock_contract(self) -> None: + entered: list[str] = [] + + class ContextOnlyLock: + def __enter__(self) -> None: + entered.append("enter") + + def __exit__(self, *_exception: object) -> None: + entered.append("exit") + + model = object() + factory = self.mps._ModelOwnership + with patch.object(self.mps, "_ModelOwnership", lambda: factory(lock=ContextOnlyLock())): + with self.mps._own_model(model): + self.assertEqual(entered, ["enter"]) + self.assertEqual(self.mps._model_owners[id(model)].users, 1) + self.assertEqual(entered, ["enter", "exit"]) + self.assertEqual(self.mps._model_owners, {}) + + def test_permission_revoked_during_admission_or_budget_setup_starts_no_constructor(self) -> None: + target = self.Target.capture("custom") + def revoke(limit: int, _deadline: object) -> tuple[int, None]: + os.environ["TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD"] = "0" + return limit, None + + self.server._request_batch_limit = revoke + with self.assertRaises(ValueError): + self.request(target=target) + os.environ["TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD"] = "1" + self.mps.ensure_mps_memory_budget.side_effect = lambda _device: os.environ.update( + TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD="0", + ) + with self.assertRaises(ValueError): + self.target_module.build_target_model(target, "mps") + self.assertEqual(self.calls, []) + + def test_late_slice_width_mismatch_returns_no_partial_output(self) -> None: + def vary_width() -> None: + if self.encoded: + self.width = 128 + + self.encode_hook = vary_width + with self.assertRaises(self.target_module.EmbeddingTargetError): + self.request(prepare=False, texts=["synthetic"] * 9) + self.assertEqual([len(entry[1]) for entry in self.encoded], [8, 1]) + self.assertFalse(self.server._inference_lock.locked()) + + def test_preparation_deadline_consumed_by_constructor_prevents_probe(self) -> None: + clock = [100.0] + self.vector.time = types.SimpleNamespace(monotonic=lambda: clock[0]) + self.build_hook = lambda: clock.__setitem__(0, 102.0) + with self.assertRaises(TimeoutError): + with self.vector.prepare_embedding_target(self.target, timeout=1.0): + self.fail("Expired construction started a preparation probe") + self.assertEqual(self.encoded, []) + self.assertFalse(self.vector._prepared_target_slot.locked()) + + def test_close_marks_closed_before_wait_and_holds_slot_until_encode_drains(self) -> None: + entered, finish, closed = threading.Event(), threading.Event(), threading.Event() + errors: list[BaseException] = [] + with self.vector.prepare_embedding_target(self.target) as lease: + def encode_hook() -> None: + entered.set() + if not finish.wait(3): + raise AssertionError("Synthetic encode barrier timed out") + + self.encode_hook = encode_hook + + def encode() -> None: + try: + lease.encode(["synthetic"]) + except BaseException as error: + errors.append(error) + + thread = threading.Thread(target=encode) + thread.start() + self.assertTrue(entered.wait(3)) + original_close = lease._finish_close + lease._finish_close = lambda: (original_close(), closed.set()) + closer = threading.Thread(target=lease.close) + closer.start() + deadline = time.monotonic() + 3 + while not lease._closed and time.monotonic() < deadline: + closed.wait(0.001) + try: + self.assertTrue(lease._closed) + self.assertFalse(closed.is_set()) + with self.assertRaises(RuntimeError): + lease.encode(["queued"]) + with self.assertRaises(RuntimeError): + with self.vector.prepare_embedding_target(self.target): + self.fail("Closing target released residency early") + finally: + finish.set() + thread.join(3) + closer.join(3) + self.assertFalse(thread.is_alive()) + self.assertFalse(closer.is_alive()) + self.assertEqual(errors, []) + self.assertTrue(closed.is_set()) + self.assertIsNone(lease._encoder) + self.assertFalse(self.vector._prepared_target_slot.locked()) + + def test_interrupted_close_can_be_retried_and_does_not_reopen_lease(self) -> None: + with self.vector.prepare_embedding_target(self.target) as lease: + actual = lease._encode_lock + failing = Mock() + failing.__enter__ = Mock(side_effect=KeyboardInterrupt("synthetic close cancellation")) + failing.__exit__ = Mock(return_value=False) + failing.acquire.return_value = False + lease._encode_lock = failing + with self.assertRaises(KeyboardInterrupt): + lease.close() + self.assertTrue(self.vector._prepared_target_slot.locked()) + with self.assertRaises(RuntimeError): + lease.encode(["synthetic"]) + lease._encode_lock = actual + lease.close() + self.assertFalse(self.vector._prepared_target_slot.locked()) + self.assertIsNone(lease._encoder) + + def test_named_pair_is_idempotent_without_active_metadata_or_registry(self) -> None: + conn = VecConnection() + self.addCleanup(conn.close) + before = self.globals() + self.assertEqual(self.vector.init_prepared_target_tables(conn, self.target), self.target.tables) + self.vector.init_prepared_target_tables(conn, self.target) + names = {row[0] for row in conn.db.execute("SELECT name FROM sqlite_master")} + self.assertEqual(names, set(self.target.tables)) + self.assertFalse(conn.in_transaction) + self.assertEqual(self.globals(), before) + self.assertEqual(conn.extensions, [True, False, True, False]) + + def test_target_ddl_preserves_borrowed_transaction_and_rolls_back_failed_pair(self) -> None: + for borrowed in (False, True): + with self.subTest(borrowed=borrowed): + conn = VecConnection() + self.addCleanup(conn.close) + conn.db.execute("CREATE TABLE unrelated(value TEXT)") + if borrowed: + conn.db.execute("INSERT INTO unrelated VALUES('synthetic')") + conn.fail_table = self.target.tables[1] + with self.assertRaises(sqlite3.OperationalError): + self.vector.init_prepared_target_tables(conn, self.target) + self.assertEqual(conn.in_transaction, borrowed) + self.assertEqual(conn.db.execute("SELECT count(*) FROM unrelated").fetchone()[0], int(borrowed)) + self.assertIsNone(conn.db.execute("SELECT 1 FROM sqlite_master WHERE name=?", (self.target.tables[0],)).fetchone()) + conn.fail_table = None + self.vector.init_prepared_target_tables(conn, self.target) + self.assertEqual(conn.in_transaction, borrowed) + if borrowed: + conn.rollback() + self.assertIsNone(conn.db.execute("SELECT 1 FROM sqlite_master WHERE name=?", (self.target.tables[0],)).fetchone()) + + def test_target_schema_rejects_wrong_width_metric_ordinary_table_and_temp_shadow(self) -> None: + for kind in ("width", "metric", "ordinary", "temp", "temp_upper", "temp_mixed"): + with self.subTest(kind=kind): + conn = VecConnection() + self.addCleanup(conn.close) + name = self.target.tables[0] + if kind == "ordinary": + conn.db.execute(f"CREATE TABLE {name}(embedding BLOB)") + else: + self.vector.init_prepared_target_tables(conn, self.target) + if kind == "width": + conn.schemas[name] = conn.schemas[name].replace("float[256]", "float[384]") + elif kind == "metric": + conn.schemas[name] = conn.schemas[name].replace("cosine", "l2") + else: + shadow = name.upper() if kind == "temp_upper" else name + if kind == "temp_mixed": + shadow = self.target.tables[1].title() + conn.db.execute(f"CREATE TEMP TABLE {shadow}(embedding BLOB)") + with self.assertRaises(ValueError): + self.vector.init_prepared_target_tables(conn, self.target) + self.assertFalse(conn.in_transaction) + + +@unittest.skipUnless(os.environ.get("TRUEMEMORY_PREPARED_NATIVE_SQLITE_TEST") == "1", + "Actual sqlite-vec is an explicit isolated native gate") +class TestNativePreparedTables(unittest.TestCase): + def test_real_vec0_pair_owned_and_borrowed_transactions(self) -> None: + # Run only in the isolated native test container, never in local checks. + import sqlite_vec + + package = types.ModuleType("truememory") + package.__path__ = [] + with patch.dict(sys.modules, {"truememory": package, "sqlite_vec": sqlite_vec}): + target_module = load("embedding_target") + load("tier_config") + vector = load("vector_search", definitions=True, namespace={ + "np": Numpy(), "threading": threading, "contextmanager": contextmanager, + "time": time, "math": math, "re": re, "sqlite3": sqlite3, + "_VEC_DISTANCE_METRIC": "cosine", + }) + target = target_module.EmbeddingTarget.capture("base") + for borrowed in (False, True): + with self.subTest(borrowed=borrowed): + conn = sqlite3.connect(":memory:") + try: + conn.execute("CREATE TABLE unrelated(value TEXT)") + if borrowed: + conn.execute("INSERT INTO unrelated VALUES('synthetic')") + self.assertEqual(vector.init_prepared_target_tables(conn, target), target.tables) + self.assertEqual(vector.init_prepared_target_tables(conn, target), target.tables) + self.assertEqual(conn.in_transaction, borrowed) + for name in target.tables: + self.assertEqual(conn.execute(f'SELECT count(*) FROM main."{name}"').fetchone()[0], 0) + if borrowed: + conn.rollback() + self.assertIsNone(conn.execute("SELECT 1 FROM main.sqlite_master WHERE name=?", + (target.tables[0],)).fetchone()) + else: + conn.execute(f"CREATE TEMP TABLE {target.tables[0]}(value TEXT)") + with self.assertRaises(ValueError): + vector.init_prepared_target_tables(conn, target) + finally: + conn.close() + for mismatched in target.tables: + with self.subTest(mismatched=mismatched): + conn = sqlite3.connect(":memory:") + try: + conn.enable_load_extension(True) + sqlite_vec.load(conn) + conn.enable_load_extension(False) + conn.execute( + f'CREATE VIRTUAL TABLE main."{mismatched}" ' + 'USING vec0(embedding float[384] distance_metric=cosine)' + ) + conn.commit() + with self.assertRaises(ValueError): + vector.init_prepared_target_tables(conn, target) + self.assertFalse(conn.in_transaction) + for name in target.tables: + row = conn.execute("SELECT sql FROM main.sqlite_master WHERE name=?", (name,)).fetchone() + if name == mismatched: + self.assertIn("float[384]", row[0]) + else: + self.assertIsNone(row) + finally: + conn.close() + + +if __name__ == "__main__": + unittest.main() diff --git a/truememory/embedding_target.py b/truememory/embedding_target.py new file mode 100644 index 00000000..7ee2be9f --- /dev/null +++ b/truememory/embedding_target.py @@ -0,0 +1,124 @@ +"""Explicit embedding identity for inactive target preparation. + +This module imports no model framework. A descriptor identifies an embedding +space, not resident weights, database readiness, or an active configuration. +""" +from __future__ import annotations + +from dataclasses import dataclass +from collections.abc import Callable +import re +import sys + + +TARGET_PROBE_TEXT = "TrueMemory embedding target validation." + + +class EmbeddingTargetError(ValueError): + """The requested target cannot be certified without changing its identity.""" + + +@dataclass(frozen=True) +class EmbeddingTarget: + tier: str + model_id: str + dimension: int + tier_group: str + + def __post_init__(self) -> None: + builtins = { + "edge": ("model2vec", 256, "edge"), + "base": ("qwen3_256", 256, "basepro"), + "pro": ("qwen3_256", 256, "basepro"), + } + if not isinstance(self.tier, str) or self.tier not in (*builtins, "custom"): + raise EmbeddingTargetError("Prepared targets require an explicit public tier") + if (not isinstance(self.model_id, str) + or re.fullmatch(r"[\w][\w.\-]*(/[\w][\w.\-]*)?", self.model_id) is None): + raise EmbeddingTargetError("Invalid prepared embedding model identity") + if type(self.dimension) is not int or not 1 <= self.dimension <= 4096: + raise EmbeddingTargetError("Prepared embedding dimension must be an integer from 1 to 4096") + if self.tier in builtins: + if self.identity != builtins[self.tier]: + raise EmbeddingTargetError("Prepared target does not match its built-in tier") + elif self.tier_group != "custom": + raise EmbeddingTargetError("Custom prepared targets require the custom table group") + if self.model_id in {"minilm", "bge-small", "qwen3"}: + raise EmbeddingTargetError("This legacy model identity cannot be certified by prepared targets") + if self.model_id in {"model2vec", "qwen3_256"} and self.dimension != 256: + raise EmbeddingTargetError("This built-in embedding model requires 256 dimensions") + + @property + def identity(self) -> tuple[str, int, str]: + return self.model_id, self.dimension, self.tier_group + + @property + def tables(self) -> tuple[str, str]: + return f"vec_messages_{self.tier_group}", f"vec_messages_sep_{self.tier_group}" + + @classmethod + def capture(cls, tier: str) -> EmbeddingTarget: + from truememory.tier_config import get_tier_config + + if not isinstance(tier, str): + raise EmbeddingTargetError("Prepared targets require an explicit public tier") + tier = tier.strip().lower() + config = get_tier_config(tier) + return cls(tier, config["embed_model"], config["embed_dim"], config["tier_group"]) + + def check_configuration(self) -> None: + if self != self.capture(self.tier): + raise EmbeddingTargetError("Prepared target configuration changed; capture a new target") + + def to_wire(self) -> dict: + return {"version": 1, "tier": self.tier, "model_id": self.model_id, + "dimension": self.dimension, "tier_group": self.tier_group} + + @classmethod + def from_wire(cls, value: object) -> EmbeddingTarget: + if (not isinstance(value, dict) + or set(value) != {"version", "tier", "model_id", "dimension", "tier_group"} + or type(value["version"]) is not int or value["version"] != 1): + raise EmbeddingTargetError("Unsupported prepared target descriptor") + return cls(value["tier"], value["model_id"], value["dimension"], value["tier_group"]) + + +def check_target_vectors(target: EmbeddingTarget, vectors: object, count: int) -> None: + if getattr(vectors, "shape", None) != (count, target.dimension): + raise EmbeddingTargetError("Actual embedding output does not match the prepared target dimension") + + +def build_target_model( + target: EmbeddingTarget, device: str | None, *, check: Callable[[], None] | None = None, +) -> object: + """Construct from captured arguments, never the legacy custom fallback.""" + def admit() -> None: + # Revalidate permission after ownership waits and imports, but never + # replace captured constructor arguments with newly read settings. + target.check_configuration() + if check is not None: + check() + + admit() + if target.model_id == "model2vec": + from model2vec import StaticModel + + admit() + return StaticModel.from_pretrained("minishlab/potion-base-8M", force_download=False) + + from truememory.mps_utils import ensure_mps_memory_budget + ensure_mps_memory_budget(device) + from sentence_transformers import SentenceTransformer + + if target.model_id == "qwen3_256": + model_kwargs = {"attn_implementation": "eager"} if sys.platform == "darwin" else None + admit() + return SentenceTransformer( + "Qwen/Qwen3-Embedding-0.6B", truncate_dim=256, + model_kwargs=model_kwargs, device=device, + ) + admit() + return SentenceTransformer( + target.model_id, truncate_dim=target.dimension, + trust_remote_code=False, device=device, + ) diff --git a/truememory/model_client.py b/truememory/model_client.py index 77efe075..778c001b 100644 --- a/truememory/model_client.py +++ b/truememory/model_client.py @@ -24,9 +24,13 @@ import sys import time from pathlib import Path +from typing import TYPE_CHECKING import numpy as np +if TYPE_CHECKING: + from truememory.embedding_target import EmbeddingTarget + from truememory._platform import ( _LOOPBACK_HOST, _MODEL_SERVER_GENERATION_ENV, @@ -626,6 +630,49 @@ def encode(self, texts, timeout: float | None = None, **kwargs) -> np.ndarray: return resp["vectors"] +class PreparedEmbeddingProxy: + """Strict target requests; no fallback to the ordinary embedding protocol.""" + + def __init__(self, target: "EmbeddingTarget") -> None: + self.target = target + + def _request(self, operation: str, timeout: float | None, **fields: object) -> dict: + from truememory.embedding_target import EmbeddingTarget, EmbeddingTargetError + + request = {"op": operation, "target": self.target.to_wire(), **fields} + response = _request_with_autostart(request, timeout=timeout) + if response.get("error") == f"Unknown op: {operation}": + raise ProtocolMismatchError( + "Model server does not support certified embedding targets; restart it after upgrading" + ) + _check_model_response(response, request) + try: + effective = EmbeddingTarget.from_wire(response.get("target")) + except EmbeddingTargetError as exc: + raise ProtocolMismatchError("Model server omitted a valid embedding target receipt") from exc + if effective != self.target: + raise ProtocolMismatchError("Model server returned a different embedding target") + return response + + def prepare(self, timeout: float | None = None) -> None: + self._request("prepare_embed_target_v1", timeout) + + def encode( + self, texts: str | list[str], *, timeout: float | None = None, batch_size: int = 32, + ) -> np.ndarray: + from truememory.embedding_target import check_target_vectors + + if isinstance(texts, str): + texts = [texts] + texts = list(texts) + if isinstance(batch_size, bool) or not isinstance(batch_size, int) or batch_size <= 0: + raise ValueError("batch_size must be a positive integer") + response = self._request("embed_target_v1", timeout, texts=texts, batch_size=batch_size) + vectors = response["vectors"] + check_target_vectors(self.target, vectors, len(texts)) + return vectors + + class RerankerProxy: """Drop-in replacement for CrossEncoder with .predict() method.""" diff --git a/truememory/model_server.py b/truememory/model_server.py index af92bee2..30f41f42 100644 --- a/truememory/model_server.py +++ b/truememory/model_server.py @@ -46,9 +46,13 @@ def _set_mps_memory_cap(): from contextlib import contextmanager # noqa: E402 from dataclasses import dataclass # noqa: E402 from pathlib import Path # noqa: E402 +from typing import TYPE_CHECKING # noqa: E402 import numpy as np # noqa: E402 +if TYPE_CHECKING: + from truememory.embedding_target import EmbeddingTarget + try: import psutil # noqa: E402 except ImportError: @@ -234,6 +238,7 @@ class _EmbedState: tier: str model_id: str generation: object = None + prepared_identity: tuple[str, int, str] | None = None def __post_init__(self) -> None: if self.generation is None: @@ -767,6 +772,7 @@ def _get_embed_model(self, tier: str, deadline: _RequestDeadline | None = None): if state.tier != tier: self._publish_embed_state(_EmbedState( model=state.model, tier=tier, model_id=model_id, generation=state.generation, + prepared_identity=state.prepared_identity, )) return state.model @@ -784,6 +790,32 @@ def _get_embed_model(self, tier: str, deadline: _RequestDeadline | None = None): log.info("Loaded embedding model for tier=%s (model=%s)", tier, model_id) return model + def _get_prepared_embed_model(self, target: "EmbeddingTarget", deadline: _RequestDeadline) -> object: + from truememory.embedding_target import build_target_model + + target.check_configuration() + deadline.check() + state = self._embed_state + if state is not None and ( + state.prepared_identity == target.identity + or (state.model_id == target.model_id and target.model_id in self._FAST_LANE_SAFE_MODEL_IDS) + ): + return state.model + # A legacy custom label is not evidence that its constructor did not + # fall back. Only strict construction can establish that identity. + state = None + self._publish_embed_state(None) + self._check_process_memory("prepared embed construction", deadline) + device = self._embed_device() + deadline.check() + model = build_target_model(target, device, check=deadline.check) + self._check_process_memory("prepared embed construction completion", deadline) + self._publish_embed_state(_EmbedState( + model=model, tier=target.tier, model_id=target.model_id, + prepared_identity=target.identity, + )) + return model + def _get_fast_encoder(self, tier: str, deadline: _RequestDeadline | None = None): """CPU-resident encoder for the single-text fast lane (issue #577). @@ -990,9 +1022,25 @@ def _handle_request_inner(self, request: dict) -> dict: if op == "ping": return {"ok": True} - if op in ("embed", "embed_batched"): - texts = request["texts"] - tier = request.get("tier", "") + if op in ("embed", "embed_batched", "prepare_embed_target_v1", "embed_target_v1"): + target = None + preparing = op == "prepare_embed_target_v1" + if op in ("prepare_embed_target_v1", "embed_target_v1"): + from truememory.embedding_target import ( + EmbeddingTarget, TARGET_PROBE_TEXT, check_target_vectors, + ) + target = EmbeddingTarget.from_wire(request.get("target")) + target.check_configuration() + texts = [TARGET_PROBE_TEXT] if preparing else request["texts"] + if not isinstance(texts, list) or not all(isinstance(text, str) for text in texts): + raise ValueError("Prepared embedding inputs must be a list of strings") + empty_target_input = not texts + if empty_target_input: + texts = [TARGET_PROBE_TEXT] + tier = target.tier + else: + texts = request["texts"] + tier = request.get("tier", "") try: batch_limit = _batch_limit(request.get("batch_size", _EMBED_BATCH_LIMIT), _EMBED_BATCH_LIMIT) @@ -1001,7 +1049,7 @@ def _handle_request_inner(self, request: dict) -> dict: # Single-text fast lane (issue #577): hook recall queries must # never queue behind batch ingestion work or OOM recovery. - if len(texts) <= self._FAST_LANE_MAX_TEXTS: + if target is None and len(texts) <= self._FAST_LANE_MAX_TEXTS: fast = self._handle_fast_embed(texts, tier, deadline) if fast is not None: return fast @@ -1035,9 +1083,13 @@ def _handle_request_inner(self, request: dict) -> dict: with deadline.locked(self._lock): if model is None: deadline.start_inference() - self._preflight_embed_result(tier, len(texts)) + if target is None: + self._preflight_embed_result(tier, len(texts)) + else: + _check_result_size((len(texts), target.dimension), "vectors") deadline.check() - model = self._get_embed_model(tier, deadline=deadline) + model = (self._get_embed_model(tier, deadline=deadline) if target is None + else self._get_prepared_embed_model(target, deadline)) order = self._embed_global_order(model, texts, limit, deadline) while offset < len(texts) or vectors is None: deadline.check() @@ -1061,6 +1113,8 @@ def _handle_request_inner(self, request: dict) -> dict: self._recover_embed_oom_locked(model, deadline) retry = batch break + if target is not None: + check_target_vectors(target, values, len(batch)) vectors = _store_batch_result( vectors, values, offset, len(batch), len(texts), indices=indices, ) @@ -1072,6 +1126,8 @@ def _handle_request_inner(self, request: dict) -> dict: deadline.check() log.warning("MPS OOM during encoding; retrying microbatch on CPU") values = model.encode(retry, batch_size=limit, show_progress_bar=False) + if target is not None: + check_target_vectors(target, values, len(retry)) vectors = _store_batch_result( vectors, values, offset, len(retry), len(texts), indices=indices, ) @@ -1083,6 +1139,10 @@ def _handle_request_inner(self, request: dict) -> dict: if should_deactivate: self._deactivate_throttler() + if target is not None: + return ({"ok": True, "target": target.to_wire()} if preparing else + {"ok": True, "target": target.to_wire(), + "vectors": vectors[:0] if empty_target_input else vectors}) return {"ok": True, "vectors": vectors} finally: model = None diff --git a/truememory/mps_utils.py b/truememory/mps_utils.py index cf50717c..e8a9e30e 100644 --- a/truememory/mps_utils.py +++ b/truememory/mps_utils.py @@ -12,7 +12,7 @@ import os import threading from _thread import LockType -from collections.abc import Iterator +from collections.abc import Callable, Iterator from contextlib import contextmanager from dataclasses import dataclass, field @@ -33,18 +33,39 @@ class _ModelOwnership: @contextmanager -def _own_model(model: object) -> Iterator[None]: +def _own_model( + model: object, *, remaining: Callable[[], float | None] | None = None, +) -> Iterator[None]: + @contextmanager + def locked(lock: LockType) -> Iterator[None]: + if remaining is None: + with lock: + yield + return + budget = remaining() + if budget is None: + lock.acquire() + elif budget <= 0 or not lock.acquire(timeout=min(budget, threading.TIMEOUT_MAX)): + raise TimeoutError("Model ownership deadline exceeded") + try: + remaining() + yield + finally: + lock.release() + key = id(model) - with _device_lock: + with locked(_device_lock): owner = _model_owners.get(key) if owner is None: owner = _ModelOwnership() _model_owners[key] = owner owner.users += 1 try: - with owner.lock: + with locked(owner.lock): yield finally: + # Cleanup cannot abandon a registered user when its request expires. + # The registry lock covers bookkeeping only, never native work. with _device_lock: owner.users -= 1 if owner.users == 0: diff --git a/truememory/vector_search.py b/truememory/vector_search.py index 537623d1..49f80989 100644 --- a/truememory/vector_search.py +++ b/truememory/vector_search.py @@ -45,6 +45,9 @@ import struct import sqlite3 import threading +import time +from collections.abc import Callable, Iterator +from contextlib import contextmanager from typing import TYPE_CHECKING import numpy as np @@ -53,7 +56,7 @@ from truememory.storage import _deserialize_metadata, select_message_cols if TYPE_CHECKING: - pass + from truememory.embedding_target import EmbeddingTarget logger = logging.getLogger(__name__) @@ -317,6 +320,236 @@ def get_model(): return _model +_prepared_target_slot = threading.Lock() + + +class _PreparedEmbeddingLease: + """An encode-only handle; closing drains work before releasing residency.""" + + def __init__( + self, target: EmbeddingTarget, encoder: object, *, shared: bool, + release_slot: Callable[[], None] | None = None, + ) -> None: + self._target = target + self._encoder = encoder + self._shared = shared + self._release_slot = release_slot + self._state_lock = threading.Lock() + self._encode_lock = threading.Lock() + self._closed = False + self._released = False + + @property + def target(self) -> EmbeddingTarget: + return self._target + + def encode( + self, texts: str | list[str], *, timeout: float | None = None, batch_size: int = 32, + ) -> np.ndarray: + from truememory.embedding_target import TARGET_PROBE_TEXT, check_target_vectors + from truememory.mps_utils import _own_model + + expires = _target_deadline(timeout) + if type(batch_size) is not int or batch_size <= 0: + raise ValueError("batch_size must be a positive integer") + if isinstance(texts, str): + texts = [texts] + texts = list(texts) + with self._state_lock: + if self._closed: + raise RuntimeError("Prepared embedding target is closed") + remaining = _target_remaining(expires) + if remaining is None: + self._encode_lock.acquire() + elif not self._encode_lock.acquire(timeout=min(remaining, threading.TIMEOUT_MAX)): + raise TimeoutError("Prepared embedding deadline exceeded") + encoder = None + try: + with self._state_lock: + if self._closed: + raise RuntimeError("Prepared embedding target is closed") + encoder = self._encoder + remaining = _target_remaining(expires) + if self._shared: + return encoder.encode(texts, timeout=remaining, batch_size=batch_size) + # The local tier worker owns its retry policy. Preserve exact-model + # exclusion without adding an independent native recovery loop. + with _own_model(encoder, remaining=lambda: _target_remaining(expires)): + self.target.check_configuration() + _target_remaining(expires) + probe = not texts + values = encoder.encode( + [TARGET_PROBE_TEXT] if probe else texts, + batch_size=batch_size, show_progress_bar=False, + ) + _target_remaining(expires) + values = np.asarray(values, dtype=np.float32) + check_target_vectors(self.target, values, 1 if probe else len(texts)) + return values[:0] if probe else values + finally: + encoder = None + with self._state_lock: + closed = self._closed + if closed: + self._finish_close() + self._encode_lock.release() + + def _finish_close(self) -> None: + # Called only by the encode-lock owner, including an encode that drains + # after its closing thread was cancelled while waiting. + self._encoder = None + release_slot, self._release_slot = self._release_slot, None + if release_slot is not None: + release_slot() + with self._state_lock: + self._released = True + + def close(self) -> None: + with self._state_lock: + if self._released: + return + # Queued calls must see closure before the current call drains. + self._closed = True + try: + with self._encode_lock: + self._finish_close() + except BaseException: + if self._encode_lock.acquire(blocking=False): + try: + self._finish_close() + finally: + self._encode_lock.release() + raise + + +def _target_deadline(timeout: float | None) -> float | None: + if timeout is None: + return None + if isinstance(timeout, bool) or not isinstance(timeout, (int, float)) or not math.isfinite(timeout): + raise ValueError("Prepared embedding timeout must be finite") + if timeout <= 0: + raise TimeoutError("Prepared embedding deadline exceeded") + return time.monotonic() + timeout + + +def _target_remaining(expires: float | None) -> float | None: + if expires is None: + return None + remaining = expires - time.monotonic() + if remaining <= 0: + raise TimeoutError("Prepared embedding deadline exceeded") + return remaining + + +@contextmanager +def prepare_embedding_target( + target: EmbeddingTarget, *, timeout: float | None = None, +) -> Iterator[_PreparedEmbeddingLease]: + """Prepare outside SQLite transactions; never activate a tier or expose weights.""" + from truememory.embedding_target import EmbeddingTarget, TARGET_PROBE_TEXT, build_target_model + from truememory.model_client import PreparedEmbeddingProxy, use_model_server + from truememory.mps_utils import resolve_device + + if not isinstance(target, EmbeddingTarget): + raise TypeError("Expected an immutable EmbeddingTarget") + expires = _target_deadline(timeout) + target.check_configuration() + shared = use_model_server() + owns_slot = False + encoder = lease = None + try: + if shared: + encoder = PreparedEmbeddingProxy(target) + encoder.prepare(timeout=_target_remaining(expires)) + else: + if not _prepared_target_slot.acquire(blocking=False): + raise RuntimeError("Another prepared embedding target lease is active") + owns_slot = True + remaining = _target_remaining(expires) + if remaining is None: + _lock.acquire() + elif not _lock.acquire(timeout=min(remaining, threading.TIMEOUT_MAX)): + raise TimeoutError("Prepared embedding deadline exceeded") + try: + _target_remaining(expires) + # Legacy custom loaders may have fallen back despite their + # label. Only deterministic built-ins can borrow active weights. + if (target.model_id in {"model2vec", "qwen3_256"} + and EMBEDDING_MODEL == target.model_id and _embedding_dim == target.dimension): + encoder = _model + finally: + _lock.release() + if encoder is None: + device = resolve_device(None) + _target_remaining(expires) + encoder = build_target_model(target, device, check=lambda: _target_remaining(expires)) + _target_remaining(expires) + lease = _PreparedEmbeddingLease( + target, encoder, shared=shared, + release_slot=_prepared_target_slot.release if owns_slot else None, + ) + owns_slot = False + encoder = None + if not shared: + lease.encode([TARGET_PROBE_TEXT], timeout=_target_remaining(expires)) + yield lease + finally: + encoder = None + if lease is not None: + lease.close() + elif owns_slot: + _prepared_target_slot.release() + + +def init_prepared_target_tables(conn: sqlite3.Connection, target: EmbeddingTarget) -> tuple[str, str]: + """Initialize only the named main-schema pair; preserve a caller transaction.""" + from truememory.embedding_target import EmbeddingTarget + import sqlite_vec + + if not isinstance(target, EmbeddingTarget): + raise TypeError("Expected an immutable EmbeddingTarget") + conn.enable_load_extension(True) + try: + sqlite_vec.load(conn) + finally: + conn.enable_load_extension(False) + owned = not conn.in_transaction + conn.execute("BEGIN IMMEDIATE" if owned else "SAVEPOINT prepared_embedding_target") + try: + for name in target.tables: + if conn.execute("SELECT 1 FROM temp.sqlite_master WHERE name = ? COLLATE NOCASE", (name,)).fetchone(): + raise ValueError("A temporary object shadows the prepared target table") + existing = conn.execute( + "SELECT type, sql FROM main.sqlite_master WHERE name = ?", (name,), + ).fetchone() + if existing is not None: + declaration = ( + r"CREATE\s+VIRTUAL\s+TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?" + + r'(?:"' + name + r'"|`' + name + r'`|\[' + name + r'\]|' + name + r')' + + r"\s+USING\s+vec0\s*\(\s*embedding\s+float\[" + + str(target.dimension) + r"\]\s+distance_metric\s*=\s*cosine\s*\)\s*" + ) + if (existing[0] != "table" or not isinstance(existing[1], str) + or re.fullmatch(declaration, existing[1], re.IGNORECASE) is None): + raise ValueError("Existing table does not match the prepared target schema") + else: + conn.execute( + f'CREATE VIRTUAL TABLE main."{name}" USING vec0({_vec0_column_decl(target.dimension)})' + ) + if owned: + conn.commit() + else: + conn.execute("RELEASE SAVEPOINT prepared_embedding_target") + except BaseException: + if owned: + conn.rollback() + else: + conn.execute("ROLLBACK TO SAVEPOINT prepared_embedding_target") + conn.execute("RELEASE SAVEPOINT prepared_embedding_target") + raise + return target.tables + + # --------------------------------------------------------------------------- # Serialization helpers # ---------------------------------------------------------------------------