From ebd35e44441663a74ebc8d30c29202eab887cf57 Mon Sep 17 00:00:00 2001 From: buildingjoshbetter <236675935+buildingjoshbetter@users.noreply.github.com> Date: Thu, 8 Oct 2026 10:08:15 -0500 Subject: [PATCH 1/2] fix: Require MCP compatible with current Pydantic --- .github/workflows/ci.yml | 17 ++++++++++++++++- .github/workflows/publish.yml | 11 ++++++++++- pyproject.toml | 5 +++-- tests/mcp-wheel-smoke.py | 10 ++++++++++ 4 files changed, 39 insertions(+), 4 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b3b8e2ee..762659bc 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -125,10 +125,22 @@ jobs: os: [ubuntu-latest, macos-latest, windows-latest] python-version: ["3.12"] mcp-version: [minimum, latest] + pydantic-version: [latest] include: - os: ubuntu-latest python-version: "3.10" mcp-version: minimum + pydantic-version: latest + # Exercise the MCP/Pydantic floor pair and the release that removed + # eval_type_backport. Other dependencies still use normal resolution. + - os: ubuntu-latest + python-version: "3.10" + mcp-version: minimum + pydantic-version: "2.11.0" + - os: ubuntu-latest + python-version: "3.12" + mcp-version: minimum + pydantic-version: "2.14.0" steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Set up Python ${{ matrix.python-version }} @@ -149,6 +161,7 @@ jobs: shell: python env: MCP_VERSION: ${{ matrix.mcp-version }} + PYDANTIC_VERSION: ${{ matrix.pydantic-version }} run: | import os from pathlib import Path @@ -161,7 +174,9 @@ jobs: python = env_dir / ("Scripts/python.exe" if os.name == "nt" else "bin/python") requirements = [str(wheel.resolve())] if os.environ["MCP_VERSION"] == "minimum": - requirements.append("mcp[cli]==1.19.0") + requirements.append("mcp[cli]==1.21.1") + if os.environ["PYDANTIC_VERSION"] != "latest": + requirements.append("pydantic==" + os.environ["PYDANTIC_VERSION"]) subprocess.run([str(python), "-m", "pip", "install", *requirements], check=True) subprocess.run([str(python), "-m", "pip", "check"], check=True) subprocess.run([str(python), "-I", "tests/mcp-wheel-smoke.py"], check=True) diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 9fda4e0e..90bdc63c 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -37,6 +37,12 @@ jobs: fail-fast: false matrix: mcp-version: [minimum, latest] + pydantic-version: [latest] + include: + - mcp-version: minimum + pydantic-version: "2.11.0" + - mcp-version: minimum + pydantic-version: "2.14.0" steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 @@ -54,6 +60,7 @@ jobs: shell: python env: MCP_VERSION: ${{ matrix.mcp-version }} + PYDANTIC_VERSION: ${{ matrix.pydantic-version }} run: | import os from pathlib import Path @@ -66,7 +73,9 @@ jobs: python = env_dir / ("Scripts/python.exe" if os.name == "nt" else "bin/python") requirements = [str(wheel.resolve())] if os.environ["MCP_VERSION"] == "minimum": - requirements.append("mcp[cli]==1.19.0") + requirements.append("mcp[cli]==1.21.1") + if os.environ["PYDANTIC_VERSION"] != "latest": + requirements.append("pydantic==" + os.environ["PYDANTIC_VERSION"]) subprocess.run([str(python), "-m", "pip", "install", *requirements], check=True) subprocess.run([str(python), "-m", "pip", "check"], check=True) subprocess.run([str(python), "-I", "tests/mcp-wheel-smoke.py"], check=True) diff --git a/pyproject.toml b/pyproject.toml index 50d5c171..d512714b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,8 +29,9 @@ dependencies = [ "model2vec>=0.6.0,<1.0", "sqlite-vec>=0.1.0,<1.0", "numpy>=1.24.0,<3.0", - # Tool metadata requires 1.19; SDK 2 removes the FastMCP import path. - "mcp[cli]>=1.19.0,<2.0", + # 1.21.1 removes a private Pydantic import broken by Pydantic 2.14. + # SDK 2 removes the FastMCP import path. + "mcp[cli]>=1.21.1,<2.0", "httpx>=0.27.0,<1.0", "sentence-transformers>=3.0.0,<7.0", "torch>=2.0.0,<3.0", diff --git a/tests/mcp-wheel-smoke.py b/tests/mcp-wheel-smoke.py index 06388fe9..6a0f2522 100644 --- a/tests/mcp-wheel-smoke.py +++ b/tests/mcp-wheel-smoke.py @@ -182,6 +182,16 @@ async def _check_protocol(root: Path, env: dict[str, str], origin: Path) -> None def main() -> None: _require(bool(sys.flags.isolated), "Run this probe with python -I") distribution = importlib.metadata.distribution("truememory") + # Report public package versions before subprocess checks, whose output is + # intentionally redacted. Dependency import failures then remain traceable. + print(json.dumps({ + "stage": "dependency_versions", + "python_version": ".".join(str(part) for part in sys.version_info[:3]), + "truememory_version": distribution.version, + "mcp_version": importlib.metadata.version("mcp"), + "pydantic_version": importlib.metadata.version("pydantic"), + "pydantic_core_version": importlib.metadata.version("pydantic-core"), + }), flush=True) origin = Path(distribution.locate_file("truememory/__init__.py")).resolve() spec = importlib.util.find_spec("truememory") _require(origin.is_file() and spec is not None and spec.origin is not None, From 82110b55ff2162beeeef8e6299f154a4fec5a450 Mon Sep 17 00:00:00 2001 From: buildingjoshbetter <236675935+buildingjoshbetter@users.noreply.github.com> Date: Thu, 8 Oct 2026 11:18:15 -0500 Subject: [PATCH 2/2] fix: Resolve daemon tier requests without changing defaults --- docs/resource-budgets.md | 35 ++ tests/test-tier-target-model-795.py | 520 +++++++++++++++++++++++ tests/test_issue_577_model_server_oom.py | 5 +- truememory/model_server.py | 45 +- 4 files changed, 587 insertions(+), 18 deletions(-) create mode 100644 tests/test-tier-target-model-795.py diff --git a/docs/resource-budgets.md b/docs/resource-budgets.md index 18674934..dd70bfd9 100644 --- a/docs/resource-budgets.md +++ b/docs/resource-budgets.md @@ -310,3 +310,38 @@ It does not change the tier worker's adaptive/OOM policy, models, precision or retrieval depth. The initial comparison's numerical failures and latency regressions remain recorded until separate native validation resolves them. Linux/CUDA behavior does not establish MPS memory or thermal performance. + +## Server target identity prerequisite (#795) + +An explicit embedding request no longer changes the daemon's default model, +dimension or generation. The same raw tier keeps its loaded snapshot, including +an empty tier whose default changed while its initial load was in progress. +Main and fast requests retain that captured identity. A cache miss or a +different tier request resolves the current selection; changing the global +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. + +The removed `qwen3` name still raises the existing migration error, including +case and whitespace variants. This validation does not mutate the default or +adopt broader alias normalization from the active-model setter. + +Explicit custom and other configuration-dependent cache entries retain their +existing raw-key reuse policy, including a custom entry that loaded a built-in +identity or Model2Vec fallback before its configuration changed. They gain no +new identity or dimension attestation, and the legacy server fallback for `minilm` and `bge-small` +remains unchanged. Result-size preflight uses the same cache-selection rule +as loading. Fast-lane inference continues to follow its captured loaded-model +snapshot, without resolving current global settings. + +This correction preserves inference/state lock ownership, request deadlines, +native batch ceilings and sticky CPU recovery. It does not bound the lifetime +of separately retained old model references: the previous main instance can +still be alive while a distinct replacement is constructed, and the optional +fast CPU instance remains separately owned. A managed preparation lifetime and +an explicit protocol guarantee for older daemons are subsequent prerequisites. +No tier-switch caller, local loader, target table or active configuration is +changed by this checkpoint; issue #795 remains open. diff --git a/tests/test-tier-target-model-795.py b/tests/test-tier-target-model-795.py new file mode 100644 index 00000000..47e7164e --- /dev/null +++ b/tests/test-tier-target-model-795.py @@ -0,0 +1,520 @@ +"""Server target identity regressions using real control flow and stdlib fakes.""" +from __future__ import annotations + +import ast +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from dataclasses import dataclass +import json +import logging +import math +import os +from pathlib import Path +import sys +import threading +import time +import types +import unittest +from unittest.mock import Mock, patch +import weakref + + +SOURCE = Path(__file__).resolve().parents[1] / "truememory" / "model_server.py" + + +class Array: + def __init__(self, rows: list[list[float]], shape: tuple[int, int]) -> None: + self.rows = rows + self.shape = shape + self.ndim = 2 + + def __setitem__(self, key: slice, value: Array) -> None: + self.rows[key] = value.rows + + +class Numpy: + ndarray = Array + float32 = "float32" + + @staticmethod + def asarray(values: Array, **_kwargs: object) -> Array: + return values + + @staticmethod + def empty(shape: tuple[int, int], **_kwargs: object) -> Array: + return Array([[] for _ in range(shape[0])], shape) + + +class FakeModel: + def __init__(self, identity: str, device: str | None) -> None: + self.identity = identity + self.device = device + self.calls: list[int] = [] + self.before_encode = lambda: None + self.oom_once = False + + def encode(self, texts: list[str], **kwargs: object) -> Array: + self.before_encode() + self.calls.append(len(texts)) + if self.oom_once: + self.oom_once = False + raise RuntimeError("MPS backend out of memory") + if len(texts) > kwargs["batch_size"]: + raise AssertionError("Native microbatch ceiling exceeded") + return Array([[1.0] * 256 for _ in texts], (len(texts), 256)) + + def to(self, device: str) -> None: + self.device = device + + +def load_server() -> types.ModuleType: + tree = ast.parse(SOURCE.read_text(encoding="utf-8")) + definitions = [node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.ClassDef))] + tree = ast.fix_missing_locations(ast.Module(body=[ + ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), + *definitions, + ], type_ignores=[])) + module = types.ModuleType("synthetic_target_server_795") + module.__dict__.update( + dataclass=dataclass, contextmanager=contextmanager, ThreadPoolExecutor=ThreadPoolExecutor, + np=Numpy(), json=json, math=math, os=os, sys=sys, threading=threading, time=time, + log=logging.getLogger("synthetic-target-795"), _MAX_MESSAGE_SIZE=100 * 1024**2, + PROTOCOL_VERSION=1, _EMBED_BATCH_LIMIT=32, _RERANK_BATCH_LIMIT=64, + ) + with patch.dict(sys.modules, {module.__name__: module}): + exec(compile(tree, str(SOURCE), "exec"), module.__dict__) + return module + + +class TestTierTargetIdentity(unittest.TestCase): + def setUp(self) -> None: + package = types.ModuleType("truememory") + package.__path__ = [] + self.vector = types.ModuleType("truememory.vector_search") + self.vector._TIER_ALIASES = {"edge": "model2vec", "base": "qwen3_256", "pro": "qwen3_256"} + self.vector.EMBEDDING_MODEL = "model2vec" + self.vector._embedding_dim = 256 + self.vector._model_generation = 7 + self.vector._model = object() + + def set_model(value: str) -> None: + self.vector.EMBEDDING_MODEL = self.vector._resolve_model_name(value) + self.vector._embedding_dim = 256 + self.vector._model_generation += 1 + self.vector._model = None + + self.vector.set_embedding_model = Mock(side_effect=set_model) + self.custom = types.ModuleType("truememory.tier_config") + self.custom.get_embed_model = Mock(return_value="synthetic/encoder-one") + vector_tree = ast.parse(SOURCE.with_name("vector_search.py").read_text(encoding="utf-8")) + resolver_nodes = [node for node in vector_tree.body if ( + isinstance(node, ast.FunctionDef) and node.name == "_resolve_model_name" + or isinstance(node, ast.Assign) and any( + isinstance(target, ast.Name) and target.id == "_REMOVED_MODELS" + for target in node.targets + ) + )] + self.vector.__dict__.update( + logger=logging.getLogger("synthetic-vector-795"), + _MODEL_DIMS={"model2vec": 256, "qwen3_256": 256, "minilm": 384, "bge-small": 384}, + _cfg_get_embed_model=self.custom.get_embed_model, + ) + exec(compile(ast.Module(body=resolver_nodes, type_ignores=[]), "synthetic-real-vector-resolver", "exec"), + self.vector.__dict__) + self.mps = types.ModuleType("truememory.mps_utils") + self.mps.resolve_device = lambda _device: "cpu" + self.mps.is_mps_oom = lambda error: "MPS backend out of memory" in str(error) + self.mps.flush_mps_cache = Mock() + modules = patch.dict(sys.modules, { + "truememory": package, "truememory.vector_search": self.vector, + "truememory.tier_config": self.custom, "truememory.mps_utils": self.mps, + }) + modules.start() + self.addCleanup(modules.stop) + environment = patch.dict(os.environ, {"TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD": "1"}) + environment.start() + self.addCleanup(environment.stop) + self.ms = load_server() + 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.builds: list[tuple[str, str | None]] = [] + self.server._build_embed_model = self.build + + def build(self, identity: str, device: str | None) -> FakeModel: + self.assertTrue(self.server._lock.locked()) + self.assertTrue(self.server._inference_lock.locked()) + self.builds.append((identity, device)) + return FakeModel(identity, device) + + def request(self, tier: str, count: int = 2, **extra: object) -> dict: + return self.server.handle_request({ + "op": "embed_batched", "tier": tier, "texts": ["synthetic"] * count, + "batch_size": 32, **extra, + }) + + def globals(self) -> tuple: + return (self.vector.EMBEDDING_MODEL, self.vector._embedding_dim, + self.vector._model_generation, self.vector._model) + + def test_all_builtin_alias_pairs_reuse_one_model_without_global_mutation(self) -> None: + for aliases in (("edge", "model2vec"), ("base", "pro", "qwen3_256")): + for first in aliases: + for second in aliases: + with self.subTest(first=first, second=second): + self.server._embed_state = None + self.builds.clear() + before = self.globals() + self.assertTrue(self.request(first)["ok"]) + snapshot = self.server._embed_state + self.assertTrue(self.request(second)["ok"]) + self.assertIs(self.server._embed_state.model, snapshot.model) + self.assertEqual(self.server._embed_state.tier, second) + self.assertEqual(snapshot.tier, first) + self.assertEqual(len(self.builds), 1) + self.assertEqual(self.globals(), before) + self.vector.set_embedding_model.assert_not_called() + + def test_explicit_target_does_not_change_later_default_request(self) -> None: + before = self.globals() + for tier, identity in (("", "model2vec"), ("base", "qwen3_256"), ("", "model2vec")): + self.assertTrue(self.request(tier)["ok"]) + self.assertEqual(self.server._embed_state.model_id, identity) + self.assertEqual(self.globals(), before) + self.assertEqual([identity for identity, _ in self.builds], + ["model2vec", "qwen3_256", "model2vec"]) + + def test_main_default_snapshot_survives_global_change_until_invalidated(self) -> None: + self.request("") + old = self.server._embed_state.model + self.vector.set_embedding_model("pro") + self.assertTrue(self.request("")["ok"]) + self.assertIs(self.server._embed_state.model, old) + self.assertEqual(self.server._embed_state.model_id, "model2vec") + self.assertEqual(len(self.builds), 1) + self.server._embed_state = None + self.assertTrue(self.request("")["ok"]) + self.assertIsNot(self.server._embed_state.model, old) + self.assertEqual(self.server._embed_state.model_id, "qwen3_256") + self.vector.set_embedding_model.assert_called_once_with("pro") + + def test_failed_distinct_load_preserves_old_snapshot_and_global_state(self) -> None: + self.request("edge") + old = self.server._embed_state + before = self.globals() + self.server._build_embed_model = Mock(side_effect=RuntimeError("synthetic load failure")) + with self.assertRaisesRegex(RuntimeError, "synthetic load failure"): + self.request("pro") + self.assertIs(self.server._embed_state, old) + self.assertEqual(self.globals(), before) + self.assertEqual(self.server._inflight, 0) + self.assertFalse(self.server._lock.locked()) + self.assertFalse(self.server._inference_lock.locked()) + + def test_removed_model_rejection_matches_real_resolver_without_mutation(self) -> None: + self.request("base") + before = self.globals() + snapshot = self.server._embed_state + builds = list(self.builds) + for removed in self.vector._REMOVED_MODELS: + for spelling in (removed, removed.upper(), " " + removed + " ", "\t" + removed.upper() + "\n"): + with self.subTest(spelling=spelling): + with self.assertRaises(ValueError) as reference: + self.vector._resolve_model_name(spelling) + with self.assertRaises(ValueError) as direct: + self.server._resolve_embed_model_id(spelling) + self.assertEqual(str(direct.exception), str(reference.exception)) + for count in (1, 2): + with self.assertRaises(ValueError) as request: + self.request(spelling, count=count) + self.assertEqual(str(request.exception), str(reference.exception)) + self.assertIs(self.server._embed_state, snapshot) + self.assertEqual(self.globals(), before) + self.assertEqual(self.builds, builds) + self.assertFalse(self.server._lock.locked()) + self.assertFalse(self.server._inference_lock.locked()) + self.vector.set_embedding_model.assert_not_called() + + def test_other_resolver_spellings_keep_existing_peek_behavior(self) -> None: + before = self.globals() + for spelling in ("base", "pro", "qwen3_256", "custom", " PRO ", "invalid-tier", "Synthetic/CaseSensitive"): + with self.subTest(spelling=spelling): + self.assertEqual(self.server._resolve_embed_model_id(spelling), + self.server._peek_embed_model_id(spelling)) + self.assertEqual(self.globals(), before) + self.vector.set_embedding_model.assert_not_called() + + def test_explicit_custom_cache_keeps_original_config_and_opt_in_semantics(self) -> None: + self.request("custom") + original = self.server._embed_state + self.custom.get_embed_model.side_effect = ValueError("synthetic invalid custom settings") + with patch.dict(os.environ, {"TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD": "0"}): + self.assertTrue(self.request("custom")["ok"]) + self.assertIs(self.server._embed_state, original) + self.assertEqual(len(self.builds), 1) + self.custom.get_embed_model.side_effect = None + self.custom.get_embed_model.return_value = "synthetic/encoder-two" + self.assertTrue(self.request("custom")["ok"]) + self.assertIs(self.server._embed_state, original) + # A different explicit key retains the legacy custom reload boundary. + self.assertTrue(self.request("synthetic/encoder-one")["ok"]) + self.assertIsNot(self.server._embed_state.model, original.model) + self.assertEqual(len(self.builds), 2) + + def test_legacy_fallback_identities_are_not_canonicalized_as_builtin_aliases(self) -> None: + for tier in ("minilm", "bge-small", "edge"): + self.assertTrue(self.request(tier)["ok"]) + self.assertEqual([identity for identity, _ in self.builds], ["minilm", "bge-small", "model2vec"]) + static = types.ModuleType("model2vec") + static.StaticModel = types.SimpleNamespace(from_pretrained=Mock(return_value=object())) + with patch.dict(sys.modules, {"model2vec": static}): + for identity in ("minilm", "bge-small"): + self.ms.ModelServer._build_embed_model(identity, "cpu") + self.assertEqual(static.StaticModel.from_pretrained.call_count, 2) + for call in static.StaticModel.from_pretrained.call_args_list: + self.assertEqual(call.args, ("minishlab/potion-base-8M",)) + self.assertEqual(call.kwargs, {"force_download": False}) + + def test_raw_custom_builtin_or_fallback_snapshot_survives_config_changes(self) -> None: + for initial in ("qwen3_256", "model2vec", None): + for opt_out in (False, True): + with self.subTest(initial=initial, opt_out=opt_out): + self.server._embed_state = None + self.builds.clear() + self.custom.get_embed_model.return_value = initial + self.custom.get_embed_model.side_effect = ( + ValueError("synthetic missing custom settings") if initial is None else None + ) + self.assertTrue(self.request("custom")["ok"]) + original = self.server._embed_state + self.assertEqual(original.model_id, initial or "model2vec") + self.custom.get_embed_model.reset_mock() + self.custom.get_embed_model.return_value = "synthetic/encoder-two" + self.custom.get_embed_model.side_effect = ( + ValueError("synthetic custom opt-out") if opt_out else None + ) + with patch.dict(os.environ, {"TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD": "0" if opt_out else "1"}): + with patch.object(self.ms, "_check_result_size") as check: + self.server._preflight_embed_result("custom", 2) + check.assert_called_once_with((2, 256), "vectors") + self.assertTrue(self.request("custom")["ok"]) + self.assertIs(self.server._embed_state, original) + self.assertEqual(len(self.builds), 1) + self.custom.get_embed_model.assert_not_called() + + def test_preflight_uses_same_resolved_identity_and_custom_cache_boundary(self) -> None: + self.request("") + self.vector.EMBEDDING_MODEL = "synthetic/encoder-one" + with patch.object(self.ms, "_check_result_size") as check: + self.server._preflight_embed_result("", 2) + check.assert_called_once_with((2, 256), "vectors") + check.reset_mock() + self.server._embed_state = None + self.server._preflight_embed_result("", 2) + check.assert_not_called() + self.request("custom") + with patch.dict(os.environ, {"TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD": "0"}): + with patch.object(self.ms, "_check_result_size") as check: + self.server._preflight_embed_result("custom", 2) + check.assert_not_called() + self.server._preflight_embed_result("synthetic/other", 2) + check.assert_called_once_with((2, 256), "vectors") + self.request("base") + with patch.object(self.ms, "_check_result_size") as check: + self.server._preflight_embed_result("pro", 2) + check.assert_called_once_with((2, 256), "vectors") + + def test_fast_and_main_lanes_keep_captured_default_identity(self) -> None: + self.vector.EMBEDDING_MODEL = "base" + self.request("") + snapshot = self.server._embed_state + self.vector.EMBEDDING_MODEL = "edge" + + def fast_build(identity: str, device: str | None) -> FakeModel: + self.assertTrue(self.server._fast_lock.locked()) + self.assertEqual((identity, device), ("qwen3_256", "cpu")) + return FakeModel(identity, device) + + self.server._build_embed_model = fast_build + with self.server._lock: + response = self.server._handle_fast_embed(["synthetic"], "") + self.assertTrue(response["ok"]) + self.assertIs(self.server._embed_state, snapshot) + self.assertEqual(self.server._fast_model_id, "qwen3_256") + self.server._build_embed_model = self.build + self.request("") + self.assertIs(self.server._embed_state, snapshot) + # A different explicit key leaves the original raw-default cache. + self.request("edge") + self.request("") + self.assertEqual(self.server._embed_state.model_id, "model2vec") + + def test_waiting_default_query_reuses_identity_captured_during_load(self) -> None: + self.vector.EMBEDDING_MODEL = "base" + build_started = threading.Event() + release_build = threading.Event() + query_waiting = threading.Event() + responses: dict[str, dict] = {} + failures: list[BaseException] = [] + main_lock = self.server._lock + + class ObservedLock: + def acquire(inner_self, blocking: bool = True, timeout: float = -1) -> bool: + if threading.current_thread().name == "synthetic-waiting-query" and blocking: + query_waiting.set() + return main_lock.acquire(blocking=blocking, timeout=timeout) + + def release(inner_self) -> None: + main_lock.release() + + def locked(inner_self) -> bool: + return main_lock.locked() + + self.server._lock = ObservedLock() + + def build(identity: str, device: str | None) -> FakeModel: + model = self.build(identity, device) + if threading.current_thread().name == "synthetic-default-loader": + build_started.set() + if not release_build.wait(2): + raise AssertionError("Synthetic model load was not released") + return model + + def request(key: str, count: int) -> None: + try: + responses[key] = self.request("", count=count, deadline=time.time() + 3) + except BaseException as error: + failures.append(error) + + self.server._build_embed_model = build + loader = threading.Thread(name="synthetic-default-loader", target=request, args=("batch", 2)) + query = threading.Thread(name="synthetic-waiting-query", target=request, args=("query", 1)) + loader.start() + try: + self.assertTrue(build_started.wait(2)) + self.assertIsNone(self.server._embed_state) + self.vector.EMBEDDING_MODEL = "edge" + query.start() + self.assertTrue(query_waiting.wait(2)) + self.assertNotIn("query", responses) + finally: + release_build.set() + loader.join(5) + if query.ident is not None: + query.join(5) + self.assertFalse(loader.is_alive()) + self.assertFalse(query.is_alive()) + self.assertEqual(failures, []) + self.assertEqual(set(responses), {"batch", "query"}) + self.assertTrue(all(response["ok"] for response in responses.values())) + self.assertEqual(self.builds, [("qwen3_256", "cpu")]) + self.assertEqual(self.server._embed_state.model.calls, [2, 1]) + self.assertEqual(self.server._embed_state.model_id, "qwen3_256") + self.assertFalse(main_lock.locked()) + self.assertFalse(self.server._inference_lock.locked()) + + def test_alias_publication_preserves_snapshot_for_fast_lane(self) -> None: + self.request("base") + snapshot = self.server._embed_state + self.request("pro") + with self.server._lock: + self.server._fast_encoder = FakeModel("qwen3_256", "cpu") + self.server._fast_model_id = "qwen3_256" + self.assertTrue(self.server._handle_fast_embed(["synthetic"], "pro")["ok"]) + self.assertEqual(snapshot.tier, "base") + self.assertIs(snapshot.model, self.server._embed_state.model) + + def test_expired_request_cannot_load_or_relabel_cached_alias(self) -> None: + self.request("base") + snapshot = self.server._embed_state + response = self.request("pro", deadline=time.time() - 1) + self.assertFalse(response["ok"]) + self.assertIn("deadline", response["error"]) + self.assertIs(self.server._embed_state, snapshot) + self.assertEqual(len(self.builds), 1) + + def test_distinct_target_cannot_load_while_existing_main_inference_is_owned(self) -> None: + self.request("edge") + entered = threading.Event() + release = threading.Event() + results: list[dict] = [] + + def blocked_encode() -> None: + self.assertTrue(self.server._inference_lock.locked()) + entered.set() + if not release.wait(2): + raise AssertionError("Synthetic inference was not released") + + self.server._embed_state.model.before_encode = blocked_encode + worker = threading.Thread(target=lambda: results.append(self.request("edge"))) + worker.start() + try: + self.assertTrue(entered.wait(2)) + response = self.request("pro", deadline=time.time() + 0.05) + self.assertFalse(response["ok"]) + self.assertIn("deadline", response["error"]) + self.assertEqual(self.builds, [("model2vec", "cpu")]) + finally: + release.set() + worker.join(2) + self.assertFalse(worker.is_alive()) + self.assertEqual(len(results), 1) + self.assertTrue(results[0]["ok"]) + self.request("pro") + self.assertEqual(self.builds[-1], ("qwen3_256", "cpu")) + + def test_alias_reuse_retains_native_batch_ceiling_and_sticky_recovery(self) -> None: + self.request("base") + model = self.server._embed_state.model + model.oom_once = True + self.assertTrue(self.request("pro", count=35)["ok"]) + self.assertEqual(model.calls, [2, 32, 32, 3]) + self.assertEqual(self.server._sticky_cpu, {"embed"}) + self.assertEqual(model.device, "cpu") + self.assertEqual(len(self.builds), 1) + + def test_expired_oom_invalidation_reloads_current_default_on_sticky_cpu(self) -> None: + now = [1000.0] + clock = types.SimpleNamespace(time=lambda: now[0], monotonic=lambda: now[0]) + self.vector.EMBEDDING_MODEL = "base" + self.mps.resolve_device = lambda _device: "mps" + with patch.object(self.ms, "time", clock): + self.request("") + failed = self.server._embed_state.model + failed.oom_once = True + + def expired_native_work() -> None: + now[0] += 2 + self.vector.EMBEDDING_MODEL = "edge" + + failed.before_encode = expired_native_work + response = self.request("", deadline=now[0] + 1) + self.assertFalse(response["ok"]) + self.assertIn("deadline", response["error"]) + self.assertIsNone(self.server._embed_state) + self.assertEqual(self.server._sticky_cpu, {"embed"}) + self.mps.flush_mps_cache.assert_not_called() + self.assertEqual(failed.device, "mps") + self.assertTrue(self.request("")["ok"]) + self.assertEqual(self.builds, [("qwen3_256", "mps"), ("model2vec", "cpu")]) + self.assertEqual(self.server._embed_state.model_id, "model2vec") + self.assertEqual(self.server._embed_state.model.device, "cpu") + + def test_old_instance_can_still_be_retained_during_distinct_load(self) -> None: + self.request("edge") + old = weakref.ref(self.server._embed_state.model) + observed: list[bool] = [] + + def build(identity: str, device: str | None) -> FakeModel: + observed.append(old() is not None) + return self.build(identity, device) + + self.server._build_embed_model = build + self.request("base") + self.assertEqual(observed, [True]) + self.assertIsNone(old()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_issue_577_model_server_oom.py b/tests/test_issue_577_model_server_oom.py index c5e93046..28c270eb 100644 --- a/tests/test_issue_577_model_server_oom.py +++ b/tests/test_issue_577_model_server_oom.py @@ -567,11 +567,10 @@ def fake_build(model_id, device): monkeypatch.setattr(vs, "EMBEDDING_MODEL", "edge") server = ModelServer() - # Arm the snapshot via a main-path load (the historical mutating - # sync happens exactly here, under the lock). + # Arm the snapshot without changing the daemon's default selection. resp = server.handle_request({"op": "embed", "texts": ["a", "b"], "tier": "base"}) assert resp["ok"] is True - assert calls == ["base"] + assert calls == [] calls.clear() # Contended fast lane: serves from the snapshot, zero mutation. diff --git a/truememory/model_server.py b/truememory/model_server.py index a9801601..1ce6cadb 100644 --- a/truememory/model_server.py +++ b/truememory/model_server.py @@ -473,13 +473,12 @@ def _peek_embed_model_id(tier: str) -> str: @staticmethod def _resolve_embed_model_id(tier: str) -> str: - """Resolve a tier name to the internal embedding model ID, keeping - vector_search's globals in sync (main-path behavior, pre-#577).""" - from truememory.vector_search import EMBEDDING_MODEL, set_embedding_model - - if tier and tier != EMBEDDING_MODEL: - set_embedding_model(tier) - + """Resolve this request without changing the daemon's default model.""" + if tier and tier.strip().lower() == "qwen3": + # Retain the removed-model error formerly raised by the setter, + # without adopting its broader normalization or changing globals. + from truememory.vector_search import _resolve_model_name + _resolve_model_name(tier) return ModelServer._peek_embed_model_id(tier) @staticmethod @@ -541,13 +540,28 @@ def _build_embed_model(model_id: str, device: str | None): # declines all of those and falls through to the main path. _FAST_LANE_SAFE_MODEL_IDS = frozenset({"model2vec", "qwen3_256"}) + def _resolve_embed_cache(self, tier: str) -> tuple[str, _EmbedState | None]: + """Return the requested identity and its compatible cached snapshot.""" + state = self._embed_state + if state is not None and state.tier == tier: + # A queued empty-tier request must follow the load it waited for, + # even if the default changed mid-load. Custom keys likewise keep + # the identity actually built rather than re-reading their config. + return state.model_id, state + + model_id = self._resolve_embed_model_id(tier) + if state is not None and state.model_id == model_id and ( + model_id in self._FAST_LANE_SAFE_MODEL_IDS or state.tier == tier + ): + return model_id, state + return model_id, None + def _preflight_embed_result(self, tier: str, count: int) -> None: """Caller holds the state lock; mirror the actual server constructors.""" if not count: return - state = self._embed_state - cached = state is not None and state.tier == tier - model_id = state.model_id if cached else self._peek_embed_model_id(tier) + model_id, state = self._resolve_embed_cache(tier) + cached = state is not None known = model_id in ("model2vec", "qwen3_256", "minilm", "bge-small") fallback = not cached and os.environ.get("TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD", "").strip() != "1" if known or fallback: @@ -571,11 +585,12 @@ def _preflight_rerank_result(model_name: str | None, count: int) -> None: # output may contain several labels per pair. Inspect its real shape. def _get_embed_model(self, tier: str): - state = self._embed_state - if state is not None and state.tier == tier: + model_id, state = self._resolve_embed_cache(tier) + if state is not None: + if state.tier != tier: + self._embed_state = _EmbedState(model=state.model, tier=tier, model_id=model_id) return state.model - model_id = self._resolve_embed_model_id(tier) model = self._build_embed_model(model_id, self._embed_device()) # ONE atomic reference assignment of an immutable snapshot — the # fast lane can never observe a torn (model, tier, model_id) triple. @@ -601,8 +616,8 @@ def _get_fast_encoder(self, tier: str): state = self._embed_state # single atomic reference read if state is None or state.tier != tier: # No main model loaded for this tier yet — declining means the - # one-time load happens exactly once, on the main path, with - # its usual global sync. + # one-time load happens exactly once, on the main path, + # without changing the daemon's default selection. return None if state.model_id not in self._FAST_LANE_SAFE_MODEL_IDS: