From 4bc096095d834095db962a20ee1547b1426a7cd0 Mon Sep 17 00:00:00 2001 From: buildingjoshbetter <236675935+buildingjoshbetter@users.noreply.github.com> Date: Thu, 8 Oct 2026 20:30:28 -0500 Subject: [PATCH] fix: Gate model stages on current process memory --- docs/resource-budgets.md | 50 ++ tests/test-model-admission-737.py | 4 +- tests/test-model-allocation-preflight-297.py | 14 +- tests/test-model-microbatches-735.py | 22 +- tests/test-model-process-admission-297.py | 564 +++++++++++++++++++ tests/test-model-request-deadlines-736.py | 12 +- tests/test-mps-budget-743.py | 9 +- tests/test_issue_577_model_server_oom.py | 6 +- tests/test_issue_646_model_server.py | 6 +- tests/test_platform_compat.py | 17 +- truememory/model_server.py | 95 +++- 11 files changed, 754 insertions(+), 45 deletions(-) create mode 100644 tests/test-model-process-admission-297.py diff --git a/docs/resource-budgets.md b/docs/resource-budgets.md index d023e152..ad6c12e6 100644 --- a/docs/resource-budgets.md +++ b/docs/resource-budgets.md @@ -468,3 +468,53 @@ unchanged. Stdlib tests exercise the actual loaders and handlers with weakrefs, constructor and inference barriers, failed replacement, deadlines and identity generation changes. Native memory and whole-process admission remain separate validation work. + +### Cooperative process RSS admission checkpoint + +The shared model daemon accepts `TRUEMEMORY_MODEL_SERVER_MAX_RSS_MB` as an +integer MiB setting, from 0 through 2,147,483,647. Its default is 0, which +disables this admission policy and performs no RSS measurements. This separate +setting does not change MCP's historical `ru_maxrss` reporting or the existing +`TRUEMEMORY_MAX_RSS_MB` setting. A finite deployment default remains uncalibrated. + +For a nonzero setting, the byte budget is B = configured MiB * 1,048,576. +Each admission reads current process RSS U from `psutil.Process().memory_info()`. +U < B admits the stage; U >= B refuses it. For example, a synthetic 1 MiB budget +admits 1,048,575 bytes and refuses both 1,048,576 and 1,048,577 bytes. An unavailable +sensor, sensor exception, missing value, boolean, noninteger or nonpositive RSS +also refuses when enabled. A later current reading below B can admit work even +after a higher historical peak. RSS is not added to GPU allocator or physical +footprint measurements, which can overlap it. + +Checks occur under the existing main inference/state or fast encoder owner: +after obsolete cache references are released and before model construction, +after construction before cache publication, before each encode/predict slice, +before embedding CPU transfer and before CPU retries. Recovery records sticky +CPU selection before possible refusal. A refused embedding transfer invalidates +the failed accelerator cache, so the next request cannot bypass CPU selection +through a cache hit. Reranker recovery retains its release-before-replacement +behavior. Refused requests return `server_busy` with the existing 250 ms retry +hint, without partial vectors or scores. Fast-lane refusal propagates directly +instead of falling through to main inference. Ownership release, deadlines, +fast generation fencing, models, dtypes, dimensions, result ordering and +existing request/batch/result limits are preserved. + +RSS sensor work consumes the request's existing monotonic deadline. The daemon +checks that deadline before and after sampling, and again after device setup +before main constructors. If sensing both expires the request and returns an +unavailable or overbudget value, deadline expiry takes precedence. No constructor, +transfer or inference retry starts for that expired request. Expiry while +admitting an embedding CPU transfer still invalidates its failed MPS cache and +retains sticky CPU selection. Disabled admission performs no sensor reads and +retains the constructor deadline checks. + +This is opt-in sampled admission, not a hard memory ceiling or a reservation. +Concurrent main and fast lanes can each pass a sample, and one constructor, +transfer, native call or result allocation can grow beyond B before the next +sample. Native allocator caches and unified-memory accounting remain outside +this policy. Synthetic tests cover thresholds, invalid sensors, constructor and +slice crossings, both CPU recovery routes, refusal without fallback, and zero +sensor reads when disabled. Finite-default calibration, workspace reservations, +system headroom, bounded recovery and actual Mac/MPS memory and thermal gates +remain pending. This checkpoint does not establish that the original 26.24 GB +Mac symptom is fixed. diff --git a/tests/test-model-admission-737.py b/tests/test-model-admission-737.py index c01de0e5..81fd0a5e 100644 --- a/tests/test-model-admission-737.py +++ b/tests/test-model-admission-737.py @@ -330,7 +330,7 @@ def predict(pairs: list, **kwargs: object) -> list: self.assertTrue(release.wait(2)) return [0.5] * len(pairs) - self.server._get_reranker = lambda name: types.SimpleNamespace(predict=predict) + self.server._get_reranker = lambda name, deadline=None: types.SimpleNamespace(predict=predict) self.server_module._store_batch_result = lambda result, values, *args: values self.server._queue_timeout = 0.05 payload = b'{"op":"rerank","pairs":[["synthetic","document"]]}' @@ -480,7 +480,7 @@ def close(self) -> None: def test_admitted_fast_query_still_runs_while_main_ownership_is_busy(self) -> None: self.server._inference_lock.acquire() - self.server._get_fast_encoder = lambda tier: types.SimpleNamespace(encode=lambda texts, **kwargs: [[0.25]]) + self.server._get_fast_encoder = lambda tier, deadline=None: types.SimpleNamespace(encode=lambda texts, **kwargs: [[0.25]]) try: client, future = self.send(b'{"op":"embed","texts":["synthetic"]}') self.assertEqual(self.response(client)["vectors"], [[0.25]]) diff --git a/tests/test-model-allocation-preflight-297.py b/tests/test-model-allocation-preflight-297.py index a2f264fa..a5a7ff38 100644 --- a/tests/test-model-allocation-preflight-297.py +++ b/tests/test-model-allocation-preflight-297.py @@ -246,7 +246,7 @@ def forward(items: list, **kwargs: object) -> Array: model = types.SimpleNamespace(**{method: forward}) setattr(self.server, "_get_embed_model" if op == "embed" else "_get_reranker", - lambda _name: model) + lambda _name, deadline=None: model) response = self.server.handle_request(self.request(op)) self.assertTrue(response["ok"]) self.assertEqual(calls, [[0, 1], [2, 3], [4]]) @@ -265,8 +265,8 @@ def test_native_count_and_later_dimension_changes_remain_errors(self) -> None: def test_empty_native_results_keep_their_existing_rank(self) -> None: model = types.SimpleNamespace(encode=lambda *args, **kwargs: Array((0,), []), predict=lambda *args, **kwargs: Array((0,), [])) - self.server._get_embed_model = lambda tier: model - self.server._get_reranker = lambda name: model + self.server._get_embed_model = lambda tier, deadline=None: model + self.server._get_reranker = lambda name, deadline=None: model for op, field in (("embed", "vectors"), ("rerank", "scores")): response = self.server.handle_request(self.request(op, count=0)) self.assertEqual(response[field].shape, (0,)) @@ -330,7 +330,7 @@ def to(self, device: str) -> None: moves.append(device) recovered.set() - self.server._get_embed_model = lambda tier: Model() + self.server._get_embed_model = lambda tier, deadline=None: Model() self.mps.flush_mps_cache.side_effect = lambda: self.assertIsNone(refs[-1]()) response = self.server.handle_request(self.request("embed", count=count)) self.assertTrue(response["ok"]) @@ -357,7 +357,7 @@ def predict(self, pairs: list, **kwargs: object) -> list: raise RuntimeError("MPS backend out of memory") return [float(value) for value in ids] - def get_model(name: str) -> Model: + def get_model(name: str, deadline: object = None) -> Model: if self.server._reranker is None: replacing = "rerank" in self.server._sticky_cpu if replacing: @@ -421,7 +421,7 @@ def run(request: dict, started: threading.Event | None = None) -> None: except Exception as error: outcomes.append(error) - def load_embed(tier: str) -> object: + def load_embed(tier: str, deadline: object = None) -> object: loaded.set() return types.SimpleNamespace(encode=lambda texts, **kwargs: [[0.5] for _ in texts]) @@ -460,7 +460,7 @@ class Model: predict = staticmethod(fail) to = Mock(side_effect=AssertionError("Expired CPU transfer")) - def load(name: str) -> Model: + def load(name: str, deadline: object = None) -> Model: model = Model() models.append(weakref.ref(model)) if op == "embed": diff --git a/tests/test-model-microbatches-735.py b/tests/test-model-microbatches-735.py index f88484aa..bd74ff9d 100644 --- a/tests/test-model-microbatches-735.py +++ b/tests/test-model-microbatches-735.py @@ -83,9 +83,9 @@ def runtime(monkeypatch: pytest.MonkeyPatch) -> tuple[ms.ModelServer, RecordingM server = ms.ModelServer() model = RecordingModel() - monkeypatch.setattr(server, "_get_embed_model", lambda _tier: model) - monkeypatch.setattr(server, "_get_fast_encoder", lambda _tier: model) - monkeypatch.setattr(server, "_get_reranker", lambda _name: model) + monkeypatch.setattr(server, "_get_embed_model", lambda _tier, deadline=None: model) + monkeypatch.setattr(server, "_get_fast_encoder", lambda _tier, deadline=None: model) + monkeypatch.setattr(server, "_get_reranker", lambda _name, deadline=None: model) monkeypatch.setattr(server, "_write_status_file", lambda: None) monkeypatch.setattr(server, "_flush_mps_cache", lambda: None) monkeypatch.setattr(mps_utils, "flush_mps_cache", lambda: None) @@ -331,7 +331,7 @@ def release(self) -> None: server._lock = ObservingLock() - def load(identity: str) -> RecordingModel: + def load(identity: str, deadline: object = None) -> RecordingModel: loads.append(identity) return model if identity == "primary" else alternate @@ -598,12 +598,14 @@ def load_order_runtime() -> dict: path = Path(__file__).resolve().parents[1] / "truememory/model_server.py" functions = {"_array_metadata", "_array_base64_size", "_result_wire_size", "_check_result_size", "_batch_limit", "_checked_batch", "_store_batch_result"} - classes = {"_EmbedState", "_RequestDeadline", "_RequestDeadlineExceeded", "_ResultTooLarge"} + classes = {"_EmbedState", "_RequestDeadline", "_RequestDeadlineExceeded", "_ResultTooLarge", + "_ProcessMemoryRefused"} methods = {"handle_request", "_handle_request_inner", "_handle_fast_embed", "_embed_global_order", "_embed_slice_indices", "_preflight_embed_result", "_resolve_embed_cache", "_request_batch_limit", "_after_request_batches", "_recover_embed_oom_locked", "_check_embed_recovery_deadline_locked", "_rerank_global_order", "_rerank_slice_indices", - "_preflight_rerank_result", "_publish_embed_state", "_retire_stale_fast_encoder"} + "_preflight_rerank_result", "_publish_embed_state", "_retire_stale_fast_encoder", + "_check_process_memory"} body = [ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0)] for node in ast.parse(path.read_text(encoding="utf-8")).body: if isinstance(node, ast.FunctionDef) and node.name in functions: @@ -641,6 +643,7 @@ class TestQwenGlobalOrder(unittest.TestCase): def runtime(self) -> tuple[dict, object, OrderModel]: module = load_order_runtime() server = object.__new__(module["ModelServer"]) + server._max_rss_bytes = 0 model = OrderModel(module["np"]) server._embed_state = module["_EmbedState"](model, "base", "qwen3_256") server._lock = threading.Lock() @@ -659,7 +662,7 @@ def runtime(self) -> tuple[dict, object, OrderModel]: server._sticky_cpu = set() server.loads = 0 - def load(_tier: str) -> OrderModel: + def load(_tier: str, deadline: object = None) -> OrderModel: server.loads += 1 return model @@ -943,7 +946,7 @@ def expire() -> None: if phase == "entry": expires = clock.now elif phase == "model-load": - server._get_embed_model = lambda _tier: (expire(), model)[1] + server._get_embed_model = lambda _tier, deadline=None: (expire(), model)[1] elif phase == "flatten-probe": model._can_flatten_inputs = lambda: (expire(), False)[1] elif phase == "lengths": @@ -1089,6 +1092,7 @@ def runtime(self) -> tuple[dict, object, PairOrderModel]: module = load_order_runtime() module["np"] = PairOrderNP() server = object.__new__(module["ModelServer"]) + server._max_rss_bytes = 0 model = PairOrderModel(module["np"]) server._lock = threading.Lock() server._inference_lock = threading.Lock() @@ -1103,7 +1107,7 @@ def runtime(self) -> tuple[dict, object, PairOrderModel]: server.loads, server.models = [], [] server.after_load = lambda _model: None - def load(name: str | None) -> PairOrderModel: + def load(name: str | None, deadline: object = None) -> PairOrderModel: server.loads.append(name) loaded = model if not server.models else PairOrderModel(module["np"]) loaded.default_prompt = model.default_prompt diff --git a/tests/test-model-process-admission-297.py b/tests/test-model-process-admission-297.py new file mode 100644 index 00000000..2c0dbb3c --- /dev/null +++ b/tests/test-model-process-admission-297.py @@ -0,0 +1,564 @@ +"""Current-RSS stage admission with synthetic models and stdlib-only sensors.""" +from __future__ import annotations + +import ast +import importlib.util +import os +import sys +import types +import unittest +import weakref +from collections.abc import Callable +from pathlib import Path +from unittest.mock import Mock, patch + + +SOURCE = Path(__file__).resolve().parents[1] / "truememory" / "model_server.py" +SOURCE_TEXT: str | None = None +SUPPORT_PATH = Path(__file__).with_name("test-model-allocation-preflight-297.py") +SPEC = importlib.util.spec_from_file_location("process_admission_array_support", SUPPORT_PATH) +SUPPORT = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(SUPPORT) +BUDGET = 1048576 + + +def load_server() -> types.ModuleType: + tree = ast.parse(SOURCE_TEXT if SOURCE_TEXT is not None else SOURCE.read_text(encoding="utf-8")) + tree.body = [node for node in tree.body if not ( + isinstance(node, ast.Import) and any(alias.name == "numpy" for alias in node.names) + or isinstance(node, ast.ImportFrom) and (node.module or "").startswith("truememory.") + or isinstance(node, ast.Try) and any( + isinstance(child, ast.Import) and any(alias.name == "psutil" for alias in child.names) + for child in node.body + ) + or isinstance(node, ast.Expr) and isinstance(node.value, ast.Call) + and isinstance(node.value.func, ast.Name) and node.value.func.id == "_set_mps_memory_cap" + )] + module = types.ModuleType("synthetic_process_admission_server") + module.__dict__.update({ + "np": SUPPORT.Numpy(), "psutil": None, "_USE_UNIX": True, "_LOOPBACK_HOST": "127.0.0.1", + "_env_int": lambda name, default, **kwargs: default, "pid_is_alive": lambda pid: False, + }) + with patch.dict(sys.modules, {module.__name__: module}): + exec(compile(ast.fix_missing_locations(tree), str(SOURCE), "exec"), module.__dict__) + module.log.disabled = True + return module + + +class Model: + def __init__(self, owner: TestProcessAdmission, identity: str, device: str) -> None: + self.owner, self.identity, self.device = owner, identity, device + + def encode(self, texts: list, **kwargs: object) -> list: + self.owner.calls.append(("embed", self.device, list(texts), kwargs)) + self.owner.inference_hook(self) + return [[float(text), 0.5] for text in texts] + + def predict(self, pairs: list, **kwargs: object) -> list: + self.owner.calls.append(("rerank", self.device, list(pairs), kwargs)) + self.owner.inference_hook(self) + return [float(pair[1]) for pair in pairs] + + def to(self, device: str) -> None: + self.owner.assertTrue(self.owner.server._inference_lock.locked()) + self.owner.assertTrue(self.owner.server._lock.locked()) + self.owner.moves.append(device) + self.device = device + self.owner.transfer_hook() + + +class TestProcessAdmission(unittest.TestCase): + def setUp(self) -> None: + package = types.ModuleType("truememory") + package.__path__ = [] + vector = types.ModuleType("truememory.vector_search") + vector.EMBEDDING_MODEL = "model2vec" + vector._TIER_ALIASES = {"edge": "model2vec", "base": "qwen3_256"} + reranker = types.ModuleType("truememory.reranker") + reranker.get_current_reranker_name = lambda: "rank-a" + self.mps = types.ModuleType("truememory.mps_utils") + self.mps.auto_detect_device = lambda: "mps" + self.mps.resolve_device = lambda device: device or "mps" + self.mps.ensure_mps_memory_budget = lambda device: None + self.mps.is_mps_oom = lambda error: "MPS backend out of memory" in str(error) + self.mps.flush_mps_cache = Mock() + transformers = types.ModuleType("sentence_transformers") + transformers.CrossEncoder = lambda identity, device: self.build(identity, device) + self.modules = patch.dict(sys.modules, { + "truememory": package, "truememory.vector_search": vector, + "truememory.reranker": reranker, "truememory.mps_utils": self.mps, + "sentence_transformers": transformers, + }) + self.modules.start() + self.addCleanup(self.modules.stop) + self.environment = patch.dict(os.environ, {"TRUEMEMORY_MODEL_SERVER_MAX_RSS_MB": "1"}) + self.environment.start() + self.addCleanup(self.environment.stop) + self.ms = load_server() + self.server = self.ms.ModelServer() + self.addCleanup(self.server._workers.shutdown, wait=True, cancel_futures=True) + self.server._write_status_file = lambda: None + self.server._build_embed_model = self.build + self.rss: object = BUDGET - 1 + self.readings: list[object] = [] + self.measurement_owners: list[tuple[bool, bool]] = [] + self.sensor = Mock(side_effect=self.measure) + self.ms.psutil = types.SimpleNamespace(Process=lambda: types.SimpleNamespace(memory_info=self.sensor)) + self.builds: list[tuple[str, str]] = [] + self.calls: list[tuple] = [] + self.moves: list[str] = [] + self.refs: list[weakref.ReferenceType] = [] + self.build_hook: Callable[[Model], None] = lambda model: None + self.inference_hook: Callable[[Model], None] = lambda model: None + self.transfer_hook: Callable[[], None] = lambda: None + + def tearDown(self) -> None: + self.assertFalse(self.server._lock.locked()) + self.assertFalse(self.server._inference_lock.locked()) + self.assertFalse(self.server._fast_lock.locked()) + self.assertEqual(self.server._inflight, 0) + self.assertTrue(all(owned and not metadata for owned, metadata in self.measurement_owners)) + + def measure(self) -> types.SimpleNamespace: + self.measurement_owners.append(( + self.server._inference_lock.locked() or self.server._fast_lock.locked(), + self.server._residency_lock.locked(), + )) + self.readings.append(self.rss) + return types.SimpleNamespace(rss=self.rss) + + def build(self, identity: str, device: str) -> Model: + self.assertTrue(self.server._inference_lock.locked() or self.server._fast_lock.locked()) + self.builds.append((identity, device)) + model = Model(self, identity, device) + self.refs.append(weakref.ref(model)) + self.build_hook(model) + return model + + def request(self, op: str = "embed", count: int = 3, **options: object) -> dict: + request = { + "op": op, "tier": "edge", "model_name": "rank-a", "batch_size": 2, + "texts": [str(index) for index in range(count)], + "pairs": [("query", str(index)) for index in range(count)], + } + request.update(options) + return self.server.handle_request(request) + + def assert_busy(self, response: dict) -> None: + self.assertFalse(response["ok"], response) + self.assertEqual(response["error_code"], "server_busy") + self.assertEqual(response["retry_after_ms"], 250) + self.assertNotIn("vectors", response) + self.assertNotIn("scores", response) + + def clear_models(self) -> None: + self.server._embed_state = None + self.server._reranker = None + self.server._reranker_name = None + self.server._fast_encoder = None + self.server._fast_model_id = None + self.server._fast_generation = None + self.builds.clear() + self.calls.clear() + self.server._embed_timestamps.clear() + + def test_setting_uses_mib_and_rejects_malformed_values(self) -> None: + self.assertEqual(self.server._max_rss_bytes, 1 * 1048576) + for raw in ("-1", "1.0", "nan", "inf", "true", "", " 1", "١", str(2**31)): + with self.subTest(raw=raw), patch.dict(os.environ, {"TRUEMEMORY_MODEL_SERVER_MAX_RSS_MB": raw}): + with self.assertRaises(ValueError): + self.ms.ModelServer() + + def test_cold_constructor_boundary_all_main_paths(self) -> None: + for op, count in (("embed", 1), ("embed_batched", 3), ("rerank", 3), ("rerank_batched", 3)): + for usage in (BUDGET - 1, BUDGET, BUDGET + 1): + with self.subTest(op=op, count=count, usage=usage): + self.clear_models() + self.rss = usage + response = self.request(op, count) + if usage < BUDGET: + self.assertTrue(response["ok"]) + self.assertEqual(len(self.builds), 1) + else: + self.assert_busy(response) + self.assertEqual(self.builds, []) + self.assertEqual(self.calls, []) + + def test_cached_slice_boundary_all_main_paths(self) -> None: + for op, count in (("embed", 1), ("embed", 3), ("rerank", 3)): + with self.subTest(op=op, count=count): + self.clear_models() + self.rss = BUDGET - 1 + self.assertTrue(self.request(op, count)["ok"]) + built = len(self.builds) + self.calls.clear() + for usage in (BUDGET, BUDGET + 1): + self.rss = usage + self.assert_busy(self.request(op, count)) + self.assertEqual(self.calls, []) + self.assertEqual(len(self.builds), built) + + def test_invalid_and_missing_measurements_refuse_before_construction(self) -> None: + for invalid in (None, True, False, 0, -1, "1048575", 1.0, float("nan"), float("inf")): + with self.subTest(invalid=invalid): + # Each sensor case is independent of sustained-workload policy. + self.server._embed_timestamps.clear() + self.rss = invalid + self.assert_busy(self.request()) + self.assertEqual(self.builds, []) + for error in (OSError("synthetic sensor failure"), RuntimeError("synthetic sensor failure")): + with self.subTest(error=type(error).__name__): + self.server._embed_timestamps.clear() + self.sensor.side_effect = error + self.assert_busy(self.request()) + self.assertEqual(self.builds, []) + self.ms.psutil = None + self.assert_busy(self.request()) + self.assertEqual(self.builds, []) + + def test_constructor_crossing_refuses_before_publication_or_inference(self) -> None: + for op, count in (("embed", 1), ("embed", 3), ("rerank", 3)): + with self.subTest(op=op, count=count): + self.clear_models() + self.rss = BUDGET - 1 + self.build_hook = lambda model: setattr(self, "rss", BUDGET) + self.assert_busy(self.request(op, count)) + self.assertEqual(len(self.builds), 1) + self.assertEqual(self.calls, []) + self.assertIsNone(self.server._embed_state) + self.assertIsNone(self.server._reranker) + self.assertIsNone(self.refs[-1]()) + + def test_slice_crossing_never_returns_partial_output(self) -> None: + for op in ("embed", "rerank"): + with self.subTest(op=op): + self.clear_models() + self.rss = BUDGET - 1 + self.inference_hook = lambda model: setattr(self, "rss", BUDGET) + self.assert_busy(self.request(op, count=5)) + self.assertEqual(len(self.calls), 1) + self.assertEqual(len(self.calls[0][2]), 2) + + def test_fast_constructor_and_cache_refusal_never_fall_through_to_main(self) -> None: + for cached in (False, True): + with self.subTest(cached=cached): + self.clear_models() + self.rss = BUDGET - 1 + self.assertTrue(self.request()["ok"]) + with self.server._inference_lock: + if cached: + self.assertTrue(self.request(count=1)["ok"]) + self.rss = BUDGET + before_builds, before_calls = len(self.builds), len(self.calls) + with patch.object(self.server, "_request_batch_limit", side_effect=AssertionError("Main fallback")): + self.assert_busy(self.request(count=1)) + self.assertEqual(len(self.builds), before_builds) + self.assertEqual(len(self.calls), before_calls) + + def test_fast_constructor_crossing_drops_unpublished_clone(self) -> None: + self.assertTrue(self.request()["ok"]) + self.build_hook = lambda model: setattr(self, "rss", BUDGET) + with self.server._inference_lock: + with patch.object(self.server, "_request_batch_limit", side_effect=AssertionError("Main fallback")): + self.assert_busy(self.request(count=1)) + self.assertIsNone(self.server._fast_encoder) + self.assertIsNone(self.refs[-1]()) + self.assertEqual(len(self.calls), 2) + + def fail_mps(self, model: Model) -> None: + if model.device == "mps": + raise RuntimeError("MPS backend out of memory") + + def test_embed_refused_transfer_is_sticky_and_invalidates_failed_cache(self) -> None: + for count in (1, 3): + with self.subTest(count=count): + self.clear_models() + self.server._sticky_cpu.clear() + self.rss = BUDGET - 1 + + def fail(model: Model) -> None: + if model.device == "mps": + self.rss = BUDGET + self.fail_mps(model) + + self.inference_hook = fail + self.assert_busy(self.request(count=count)) + self.assertIn("embed", self.server._sticky_cpu) + self.assertIsNone(self.server._embed_state) + self.assertIsNone(self.refs[-1]()) + self.assertEqual(self.moves, []) + self.rss = BUDGET - 1 + self.assertTrue(self.request(count=count)["ok"]) + self.assertEqual(self.builds[-1], ("model2vec", "cpu")) + + def test_embed_transfer_crossing_refuses_cpu_retry(self) -> None: + for count in (1, 3): + with self.subTest(count=count): + self.clear_models() + self.server._sticky_cpu.clear() + self.moves.clear() + self.rss = BUDGET - 1 + self.inference_hook = self.fail_mps + self.transfer_hook = lambda: setattr(self, "rss", BUDGET) + self.assert_busy(self.request(count=count)) + self.assertIn("embed", self.server._sticky_cpu) + self.assertEqual(self.moves, ["cpu"]) + self.assertEqual(len(self.calls), 1) + self.assertEqual(self.server._embed_state.model.device, "cpu") + + def test_rerank_refused_replacement_is_sticky_and_releases_failed_model(self) -> None: + def fail(model: Model) -> None: + if model.device == "mps": + self.rss = BUDGET + self.fail_mps(model) + + self.inference_hook = fail + self.assert_busy(self.request("rerank")) + self.assertIn("rerank", self.server._sticky_cpu) + self.assertEqual(len(self.builds), 1) + self.assertIsNone(self.server._reranker) + self.assertIsNone(self.server._reranker_name) + self.assertIsNone(self.refs[-1]()) + self.rss = BUDGET - 1 + self.assertTrue(self.request("rerank")["ok"]) + self.assertEqual(self.builds[-1], ("rank-a", "cpu")) + + def test_rerank_cpu_constructor_crossing_refuses_retry(self) -> None: + self.inference_hook = self.fail_mps + + def constructed(model: Model) -> None: + if model.device == "cpu": + self.rss = BUDGET + + self.build_hook = constructed + self.assert_busy(self.request("rerank")) + self.assertEqual(self.builds, [("rank-a", "mps"), ("rank-a", "cpu")]) + self.assertEqual(len(self.calls), 1) + self.assertIsNone(self.server._reranker) + self.assertTrue(all(ref() is None for ref in self.refs)) + + def test_rerank_retry_samples_after_constructor_admission(self) -> None: + self.inference_hook = self.fail_mps + original = self.server._get_reranker + + def loaded(name: str, deadline: object = None) -> Model: + model = original(name, deadline=deadline) + if model.device == "cpu": + self.rss = BUDGET + return model + + self.server._get_reranker = loaded + self.assert_busy(self.request("rerank")) + self.assertEqual(len(self.builds), 2) + self.assertEqual(len(self.calls), 1) + + def test_admitted_recovery_retains_complete_order_and_batch_controls(self) -> None: + for op in ("embed", "rerank"): + with self.subTest(op=op): + self.clear_models() + self.server._sticky_cpu.clear() + self.inference_hook = self.fail_mps + response = self.request(op, count=5) + self.assertTrue(response["ok"]) + self.assertEqual([len(call[2]) for call in self.calls], [2, 2, 2, 1]) + self.assertTrue(all(call[3] == {"batch_size": 2, "show_progress_bar": False} + for call in self.calls)) + actual = response["vectors" if op == "embed" else "scores"] + self.assertEqual(actual.dtype, "float32") + self.assertEqual(actual.data, [0.0, 0.5, 1.0, 0.5, 2.0, 0.5, 3.0, 0.5, 4.0, 0.5] + if op == "embed" else [0.0, 1.0, 2.0, 3.0, 4.0]) + + def test_current_usage_can_fall_after_historical_peak(self) -> None: + resource = types.ModuleType("resource") + resource.getrusage = Mock(return_value=types.SimpleNamespace(ru_maxrss=10 * BUDGET)) + with patch.dict(sys.modules, {"resource": resource}): + self.rss = BUDGET + 1 + self.assert_busy(self.request()) + self.rss = BUDGET - 1 + self.assertTrue(self.request()["ok"]) + resource.getrusage.assert_not_called() + self.assertIn(BUDGET + 1, self.readings) + self.assertIn(BUDGET - 1, self.readings) + + def test_disabled_budget_performs_zero_sensor_reads_including_recovery(self) -> None: + self.server._max_rss_bytes = 0 + self.ms.psutil = types.SimpleNamespace(Process=Mock(side_effect=AssertionError("Disabled sensor"))) + self.inference_hook = self.fail_mps + self.assertTrue(self.request()["ok"]) + self.assertTrue(self.request("rerank")["ok"]) + with self.server._inference_lock: + self.assertTrue(self.request(count=1)["ok"]) + self.assertEqual(self.readings, []) + self.ms.psutil.Process.assert_not_called() + + def test_default_budget_is_disabled(self) -> None: + with patch.dict(os.environ): + os.environ.pop("TRUEMEMORY_MODEL_SERVER_MAX_RSS_MB", None) + server = self.ms.ModelServer() + try: + self.assertEqual(server._max_rss_bytes, 0) + finally: + server._workers.shutdown(wait=True, cancel_futures=True) + + def test_replacement_releases_obsolete_model_before_measuring(self) -> None: + self.assertTrue(self.request()["ok"]) + old = self.refs[-1] + + def measurement() -> types.SimpleNamespace: + self.assertIsNone(old(), "Old residency retained at admission") + return types.SimpleNamespace(rss=BUDGET - 1) + + self.sensor.side_effect = measurement + self.assertTrue(self.request(tier="base")["ok"]) + + def test_expired_request_never_samples_or_loads(self) -> None: + response = self.request(deadline=1) + self.assertFalse(response["ok"]) + self.assertIn("deadline", response["error"]) + self.assertEqual(self.readings, []) + self.assertEqual(self.builds, []) + + def test_sensor_time_expiry_prevents_all_five_constructor_routes(self) -> None: + routes = (("embed", 1, "cold"), ("embed", 3, "cold"), ("rerank", 3, "cold"), + ("embed", 1, "fast"), ("rerank", 3, "recovery")) + for op, count, route in routes: + for outcome in ("under", "equal", "over", "invalid", "error"): + with self.subTest(op=op, count=count, route=route, outcome=outcome): + self.clear_models() + self.server._sticky_cpu.clear() + self.rss = BUDGET - 1 + self.sensor.side_effect = self.measure + self.inference_hook = self.fail_mps if route == "recovery" else lambda model: None + if route == "fast": + self.assertTrue(self.request()["ok"]) + self.calls.clear() + self.builds.clear() + tick, expired_samples = [1000.0], [] + clock = types.SimpleNamespace(time=lambda: tick[0], monotonic=lambda: tick[0]) + + def measure() -> types.SimpleNamespace: + value = self.measure() + if route != "recovery" or "rerank" in self.server._sticky_cpu: + expired_samples.append(value) + tick[0] = 1002.0 + if outcome == "error": + raise OSError("synthetic expired RSS read") + value.rss = {"under": BUDGET - 1, "equal": BUDGET, + "over": BUDGET + 1, "invalid": None}[outcome] + return value + + self.sensor.side_effect = measure + with patch.object(self.ms, "time", clock): + if route == "fast": + with self.server._inference_lock, patch.object( + self.server, "_request_batch_limit", side_effect=AssertionError("Main fallback"), + ): + response = self.request(op, count, deadline=1001.0) + else: + response = self.request(op, count, deadline=1001.0) + self.assertFalse(response["ok"]) + self.assertIn("deadline", response["error"]) + self.assertNotIn("error_code", response) + self.assertEqual(len(expired_samples), 1) + self.assertEqual(self.builds, [("rank-a", "mps")] if route == "recovery" else []) + self.assertEqual(len(self.calls), 1 if route == "recovery" else 0) + if route == "recovery": + self.assertIn("rerank", self.server._sticky_cpu) + self.assertIsNone(self.server._reranker) + self.assertIsNone(self.server._reranker_name) + self.assertIsNone(self.refs[-1]()) + if route == "fast": + self.assertIsNone(self.server._fast_encoder) + self.assertFalse(self.server._lock.locked()) + self.assertFalse(self.server._inference_lock.locked()) + self.assertFalse(self.server._fast_lock.locked()) + self.assertEqual(self.server._inflight, 0) + + def test_sensor_expiry_after_construction_cannot_publish_or_infer(self) -> None: + for op, count, fast in (("embed", 1, False), ("embed", 3, False), + ("rerank", 3, False), ("embed", 1, True)): + with self.subTest(op=op, count=count, fast=fast): + self.clear_models() + self.sensor.side_effect = self.measure + if fast: + self.assertTrue(self.request()["ok"]) + self.calls.clear() + self.builds.clear() + tick = [1000.0] + clock = types.SimpleNamespace(time=lambda: tick[0], monotonic=lambda: tick[0]) + + def measure() -> types.SimpleNamespace: + value = self.measure() + if self.builds: + tick[0] = 1002.0 + return value + + self.sensor.side_effect = measure + with patch.object(self.ms, "time", clock): + if fast: + with self.server._inference_lock: + response = self.request(op, count, deadline=1001.0) + else: + response = self.request(op, count, deadline=1001.0) + self.assertIn("deadline", response["error"]) + self.assertEqual(len(self.builds), 1) + self.assertEqual(self.calls, []) + self.assertIsNone(self.refs[-1]()) + self.assertIsNone(self.server._fast_encoder) + if not fast: + self.assertIsNone(self.server._embed_state) + self.assertIsNone(self.server._reranker) + + def test_embed_transfer_sensor_expiry_invalidates_failed_accelerator(self) -> None: + for count in (1, 3): + with self.subTest(count=count): + self.clear_models() + self.server._sticky_cpu.clear() + tick = [1000.0] + clock = types.SimpleNamespace(time=lambda: tick[0], monotonic=lambda: tick[0]) + + def measure() -> types.SimpleNamespace: + value = self.measure() + if "embed" in self.server._sticky_cpu: + tick[0] = 1002.0 + return value + + self.sensor.side_effect = measure + self.inference_hook = self.fail_mps + with patch.object(self.ms, "time", clock): + response = self.request(count=count, deadline=1001.0) + self.assertIn("deadline", response["error"]) + self.assertIn("embed", self.server._sticky_cpu) + self.assertEqual(len(self.calls), 1) + self.assertEqual(self.moves, []) + self.assertIsNone(self.server._embed_state) + self.assertIsNone(self.refs[-1]()) + + def test_disabled_admission_retains_deadline_checks_without_sensor_reads(self) -> None: + self.server._max_rss_bytes = 0 + self.ms.psutil = types.SimpleNamespace(Process=Mock(side_effect=AssertionError("Disabled sensor"))) + for op, count in (("embed", 1), ("embed", 3), ("rerank", 3)): + with self.subTest(op=op, count=count): + self.clear_models() + tick = [1000.0] + clock = types.SimpleNamespace(time=lambda: tick[0], monotonic=lambda: tick[0]) + + def expire_device(device: str | None) -> str: + tick[0] = 1002.0 + return "mps" + + with patch.object(self.ms, "time", clock), patch.object(self.mps, "resolve_device", expire_device): + response = self.request(op, count, deadline=1001.0) + self.assertIn("deadline", response["error"]) + self.assertEqual(self.builds, []) + self.assertEqual(self.calls, []) + self.ms.psutil.Process.assert_not_called() + self.assertEqual(self.readings, []) + + def test_invalid_sensor_cases_do_not_activate_unrelated_throttler(self) -> None: + with patch.object(self.server, "_activate_throttler", side_effect=AssertionError("Unrelated throttler")): + self.test_invalid_and_missing_measurements_refuse_before_construction() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-model-request-deadlines-736.py b/tests/test-model-request-deadlines-736.py index 92eb8df9..415ce62f 100644 --- a/tests/test-model-request-deadlines-736.py +++ b/tests/test-model-request-deadlines-736.py @@ -100,7 +100,7 @@ def setUp(self) -> None: self.loads: list[str] = [] self.recoveries: list[str] = [] - def load(tier: str | None) -> FakeModel: + def load(tier: str | None, deadline: object = None) -> FakeModel: self.loads.append(tier or "fixture") return self.model @@ -217,7 +217,7 @@ def test_model_loading_consumes_the_same_budget(self) -> None: with self.subTest(op=op, single=single, contended=contended): self.server._lock = FakeLock(self.clock, contended=contended) - def slow_load(_identity: str | None) -> FakeModel: + def slow_load(_identity: str | None, deadline: object = None) -> FakeModel: self.clock.advance(2.0) return self.model @@ -339,7 +339,7 @@ def test_rerank_reload_consumes_budget_before_retry(self) -> None: self.model = FakeModel(self.clock, oom=True) cpu_model = FakeModel(self.clock) - def load(_name: str | None) -> FakeModel: + def load(_name: str | None, deadline: object = None) -> FakeModel: if self.recoveries: self.clock.advance(2.0) return cpu_model @@ -361,7 +361,7 @@ def test_rerank_cache_flush_expiry_prevents_cpu_model_load(self) -> None: self.assertIsNone(self.server._reranker_name) def test_wall_clock_rollback_cannot_extend_accepted_budget(self) -> None: - def slow_load(_tier: str) -> FakeModel: + def slow_load(_tier: str, deadline: object = None) -> FakeModel: self.clock.tick += 2.0 self.clock.wall -= 100.0 return self.model @@ -371,7 +371,7 @@ def slow_load(_tier: str) -> FakeModel: self.assertEqual(self.model.calls, 0) def test_wall_clock_jump_cannot_expire_unspent_monotonic_budget(self) -> None: - def load(_tier: str) -> FakeModel: + def load(_tier: str, deadline: object = None) -> FakeModel: self.clock.wall += 100.0 return self.model @@ -404,7 +404,7 @@ def test_distant_valid_deadline_cannot_overflow_native_lock_timeout(self) -> Non def test_fast_model_failure_before_expiry_still_uses_main_fallback(self) -> None: self.server._lock = FakeLock(self.clock, contended=True) - def fail(_tier: str) -> FakeModel: + def fail(_tier: str, deadline: object = None) -> FakeModel: raise ValueError("synthetic model failure") self.server._get_fast_encoder = fail diff --git a/tests/test-mps-budget-743.py b/tests/test-mps-budget-743.py index c595e47c..49200a8f 100644 --- a/tests/test-mps-budget-743.py +++ b/tests/test-mps-budget-743.py @@ -274,13 +274,18 @@ def local_embed(self, model_id: str) -> types.ModuleType: def server(self): tree = ast.parse((SOURCE / "model_server.py").read_text(encoding="utf-8")) server = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "ModelServer") - names = {"_build_embed_model", "_get_reranker", "_mark_sticky_cpu", "_flush_mps_cache"} + names = {"_build_embed_model", "_get_reranker", "_mark_sticky_cpu", "_flush_mps_cache", + "_check_process_memory"} server.body = [node for node in server.body if isinstance(node, ast.FunctionDef) and node.name in names] + refusal = next(node for node in tree.body if isinstance(node, ast.ClassDef) + and node.name == "_ProcessMemoryRefused") namespace = {"sys": sys, "os": types.SimpleNamespace(environ=self.env), "log": logging.getLogger("synthetic743"), "gc": types.SimpleNamespace(collect=Mock())} - exec(compile(ast.fix_missing_locations(ast.Module(body=[server], type_ignores=[])), + future = ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0) + exec(compile(ast.fix_missing_locations(ast.Module(body=[future, refusal, server], type_ignores=[])), "model_server.py", "exec"), namespace) instance = namespace["ModelServer"]() + instance._max_rss_bytes = 0 instance._reranker = None instance._reranker_name = None instance._sticky_cpu = set() diff --git a/tests/test_issue_577_model_server_oom.py b/tests/test_issue_577_model_server_oom.py index 28c270eb..a441dc65 100644 --- a/tests/test_issue_577_model_server_oom.py +++ b/tests/test_issue_577_model_server_oom.py @@ -332,7 +332,7 @@ def predict(self, pairs, **kwargs): raise _make_mps_oom_error() return np.full(len(pairs), 0.5, dtype=np.float32) - def fake_get_reranker(self, model_name=None): + def fake_get_reranker(self, model_name=None, deadline=None): # Before sticky degradation the loaded model is on MPS and OOMs; # after, the reload is CPU-resident and succeeds. return FakeReranker(fail="rerank" not in self._sticky_cpu) @@ -371,7 +371,7 @@ def test_issue_577_rerank_non_oom_error_propagates(self, no_flush, monkeypatch): broken = MagicMock() broken.predict = MagicMock(side_effect=RuntimeError("tokenizer exploded")) monkeypatch.setattr( - ModelServer, "_get_reranker", lambda self, model_name=None: broken + ModelServer, "_get_reranker", lambda self, model_name=None, deadline=None: broken ) server = ModelServer() with pytest.raises(RuntimeError, match="tokenizer"): @@ -408,7 +408,7 @@ def encode(self, texts, **kwargs): return np.zeros((len(texts), 4), dtype=np.float32) monkeypatch.setattr( - ModelServer, "_get_fast_encoder", lambda self, tier: FakeFastEncoder() + ModelServer, "_get_fast_encoder", lambda self, tier, deadline=None: FakeFastEncoder() ) # Simulate an in-flight ingestion batch occupying the heavy path. diff --git a/tests/test_issue_646_model_server.py b/tests/test_issue_646_model_server.py index 88e32250..18d73a60 100644 --- a/tests/test_issue_646_model_server.py +++ b/tests/test_issue_646_model_server.py @@ -172,7 +172,7 @@ def test_embed_with_active_throttler_does_not_crash(self, monkeypatch): fake_model = MagicMock() fake_model.encode.return_value = np.zeros((2, 4), dtype=np.float32) - monkeypatch.setattr(server, "_get_embed_model", lambda tier: fake_model) + monkeypatch.setattr(server, "_get_embed_model", lambda tier, deadline=None: fake_model) class FakeThrottler: def before_batch(self): @@ -206,7 +206,7 @@ def test_expired_deadline_rejected_before_encode(self, monkeypatch): encode_called = {"n": 0} - def boom(tier): + def boom(tier, deadline=None): encode_called["n"] += 1 raise AssertionError("encode must not run for expired request") @@ -229,7 +229,7 @@ def test_future_deadline_runs_normally(self, monkeypatch): server = ModelServer() fake_model = MagicMock() fake_model.encode.return_value = np.zeros((1, 4), dtype=np.float32) - monkeypatch.setattr(server, "_get_embed_model", lambda tier: fake_model) + monkeypatch.setattr(server, "_get_embed_model", lambda tier, deadline=None: fake_model) resp = server.handle_request({ "op": "embed", diff --git a/tests/test_platform_compat.py b/tests/test_platform_compat.py index 6cdfcd24..de48b509 100644 --- a/tests/test_platform_compat.py +++ b/tests/test_platform_compat.py @@ -85,6 +85,7 @@ def _tmp_config_dir(tmp_path, monkeypatch, ms): (home / ".truememory").mkdir() monkeypatch.setattr(ms, "_TRUEMEMORY_DIR", home / ".truememory") monkeypatch.setattr(ms, "_CONFIG_PATH", home / ".truememory" / "config.json") + monkeypatch.setattr(ms, "_CONFIG_LOCK_PATH", home / ".truememory" / "config.json.lock") return home @@ -158,24 +159,36 @@ def test_ingest_cli_save_warns_on_windows_when_api_key_present( tmp_path, monkeypatch, capsys ): from truememory.ingest import cli as ic - cfg_path = tmp_path / ".truememory" / "config.json" + import truememory.mcp_server as ms + home = _tmp_config_dir(tmp_path, monkeypatch, ms) + cfg_path = home / ".truememory" / "config.json" monkeypatch.setattr(ic, "_TRUEMEMORY_CONFIG_PATH", cfg_path) monkeypatch.setattr(ic.sys, "platform", "win32") + assert ic._TRUEMEMORY_CONFIG_PATH == ms._CONFIG_PATH == cfg_path + assert ms._TRUEMEMORY_DIR == cfg_path.parent + assert ms._CONFIG_LOCK_PATH == cfg_path.parent / "config.json.lock" ic._save_truememory_config({"tier": "pro", "openrouter_api_key": "sk-or-fake"}) captured = capsys.readouterr() assert "windows" in captured.err.lower() assert "environment variable" in captured.err.lower() + assert json.loads(cfg_path.read_text(encoding="utf-8")) == {"tier": "pro", "openrouter_api_key": "sk-or-fake"} def test_ingest_cli_save_silent_on_posix(tmp_path, monkeypatch, capsys): from truememory.ingest import cli as ic - cfg_path = tmp_path / ".truememory" / "config.json" + import truememory.mcp_server as ms + home = _tmp_config_dir(tmp_path, monkeypatch, ms) + cfg_path = home / ".truememory" / "config.json" monkeypatch.setattr(ic, "_TRUEMEMORY_CONFIG_PATH", cfg_path) monkeypatch.setattr(ic.sys, "platform", "linux") + assert ic._TRUEMEMORY_CONFIG_PATH == ms._CONFIG_PATH == cfg_path + assert ms._TRUEMEMORY_DIR == cfg_path.parent + assert ms._CONFIG_LOCK_PATH == cfg_path.parent / "config.json.lock" ic._save_truememory_config({"tier": "pro", "openrouter_api_key": "sk-or-fake"}) captured = capsys.readouterr() assert captured.err == "" + assert json.loads(cfg_path.read_text(encoding="utf-8")) == {"tier": "pro", "openrouter_api_key": "sk-or-fake"} diff --git a/truememory/model_server.py b/truememory/model_server.py index 55c232d4..f8727c24 100644 --- a/truememory/model_server.py +++ b/truememory/model_server.py @@ -49,6 +49,11 @@ def _set_mps_memory_cap(): import numpy as np # noqa: E402 +try: + import psutil # noqa: E402 +except ImportError: + psutil = None + from truememory._platform import ( # noqa: E402 _LOOPBACK_HOST, _MODEL_SERVER_GENERATION_ENV, _USE_UNIX, _env_int, pid_is_alive, read_start_claim, try_file_lock, @@ -84,6 +89,10 @@ class _ResultTooLarge(ValueError): """The complete float32 result cannot fit the existing response frame.""" +class _ProcessMemoryRefused(RuntimeError): + """Current process memory cannot admit another model stage.""" + + def _array_metadata(shape: tuple, encoded: str = "") -> dict: return {"__ndarray__": encoded, "dtype": "float32", "shape": list(shape)} @@ -316,6 +325,9 @@ class ModelServer: _FAST_LANE_MAX_TEXTS = 1 def __init__(self): + self._max_rss_bytes = _admission_setting( + "TRUEMEMORY_MODEL_SERVER_MAX_RSS_MB", 0, 0, 2**31 - 1, + ) * 1048576 self._max_handlers = _admission_setting("TRUEMEMORY_MODEL_SERVER_MAX_HANDLERS", 16, 1, 128) self._max_request_bytes = _admission_setting( "TRUEMEMORY_MODEL_SERVER_MAX_REQUEST_BYTES", 32 * 1024**2, _MAX_MESSAGE_SIZE, 1024**3, @@ -382,6 +394,32 @@ def __init__(self): # kills a request mid-encode. Guarded by _activity_lock. self._inflight = 0 + def _check_process_memory( + self, stage: str, deadline: _RequestDeadline | None = None, + ) -> None: + """Sample current RSS under the caller's model owner, never peak RSS.""" + if deadline is not None: + deadline.check() + budget = self._max_rss_bytes + if budget == 0: + return + try: + usage = psutil.Process().memory_info().rss if psutil is not None else None + except Exception: + usage = None + # Sensor work consumes the same request budget. Expiry takes precedence + # over a memory refusal and cannot admit another native stage. + if deadline is not None: + deadline.check() + if isinstance(usage, bool) or not isinstance(usage, int) or usage <= 0: + raise _ProcessMemoryRefused( + f"Model server busy: current process RSS unavailable before {stage}" + ) + if usage >= budget: + raise _ProcessMemoryRefused( + f"Model server busy: process RSS budget exhausted before {stage}" + ) + def _mark_sticky_cpu(self, kind: str) -> bool: """Permanently degrade *kind* ("embed"/"rerank") to CPU after an MPS OOM (issue #577). Re-promoting to MPS after recovery guaranteed @@ -438,6 +476,17 @@ def _recover_embed_oom_locked( from truememory.mps_utils import flush_mps_cache self._mark_sticky_cpu("embed") flush_mps_cache() + if deadline is not None: + self._check_embed_recovery_deadline_locked(model, deadline) + try: + self._check_process_memory("embed CPU transfer", deadline) + except (_ProcessMemoryRefused, _RequestDeadlineExceeded): + # A refused transfer leaves the failed accelerator instance behind. + # Invalidate it so a cache hit cannot bypass sticky CPU selection. + state = self._embed_state + if state is not None and state.model is model: + self._publish_embed_state(None) + raise if deadline is not None: self._check_embed_recovery_deadline_locked(model, deadline) if hasattr(model, "to"): @@ -704,7 +753,7 @@ def _publish_embed_state(self, state: _EmbedState | None) -> None: del retired self._retire_stale_fast_encoder() - def _get_embed_model(self, tier: str): + def _get_embed_model(self, tier: str, deadline: _RequestDeadline | None = None): model_id, state = self._resolve_embed_cache(tier) if state is not None: if state.tier != tier: @@ -715,14 +764,19 @@ def _get_embed_model(self, tier: str): # Memory-first replacement: failed construction leaves the cache empty. self._publish_embed_state(None) - model = self._build_embed_model(model_id, self._embed_device()) + self._check_process_memory("embed construction", deadline) + device = self._embed_device() + if deadline is not None: + deadline.check() + model = self._build_embed_model(model_id, device) + self._check_process_memory("embed construction completion", deadline) # ONE atomic reference assignment of an immutable snapshot — the # fast lane can never observe a torn (model, tier, model_id) triple. self._publish_embed_state(_EmbedState(model=model, tier=tier, model_id=model_id)) log.info("Loaded embedding model for tier=%s (model=%s)", tier, model_id) return model - def _get_fast_encoder(self, tier: str): + def _get_fast_encoder(self, tier: str, deadline: _RequestDeadline | None = None): """CPU-resident encoder for the single-text fast lane (issue #577). Loaded lazily on the first single-text request that finds the global @@ -756,7 +810,9 @@ def _get_fast_encoder(self, tier: str): self._fast_encoder = None self._fast_model_id = None self._fast_generation = None + self._check_process_memory("fast embed construction", deadline) model = self._build_embed_model(model_id, "cpu") + self._check_process_memory("fast embed construction completion", deadline) # A changed generation, including A -> B -> A, must not publish this # construction. Its already-captured request may still finish on it. if self._residency_lock.acquire(blocking=False): @@ -775,7 +831,9 @@ def _get_fast_encoder(self, tier: str): ) return model - def _get_reranker(self, model_name: str | None = None): + def _get_reranker( + self, model_name: str | None = None, deadline: _RequestDeadline | None = None, + ): from truememory.reranker import get_current_reranker_name name = model_name or get_current_reranker_name() @@ -784,6 +842,7 @@ def _get_reranker(self, model_name: str | None = None): self._reranker = None self._reranker_name = None + self._check_process_memory("reranker construction", deadline) from truememory.mps_utils import auto_detect_device, ensure_mps_memory_budget, resolve_device if "rerank" in self._sticky_cpu: device = "cpu" @@ -792,7 +851,11 @@ def _get_reranker(self, model_name: str | None = None): ensure_mps_memory_budget(device) from sentence_transformers import CrossEncoder - self._reranker = CrossEncoder(name, device=device) + if deadline is not None: + deadline.check() + model = CrossEncoder(name, device=device) + self._check_process_memory("reranker construction completion", deadline) + self._reranker = model self._reranker_name = name log.info("Loaded reranker model=%s device=%s", name, device) return self._reranker @@ -823,7 +886,8 @@ def _handle_fast_embed( deadline.start_inference() self._preflight_embed_result(tier, len(texts)) deadline.check() - model = self._get_embed_model(tier) + model = self._get_embed_model(tier, deadline=deadline) + self._check_process_memory("embed slice", deadline) deadline.check() recover = False try: @@ -838,6 +902,7 @@ def _handle_fast_embed( # failed forward's traceback no longer owns tensors. self._check_embed_recovery_deadline_locked(model, deadline) self._recover_embed_oom_locked(model, deadline) + self._check_process_memory("embed CPU retry", deadline) deadline.check() vectors = model.encode(texts, batch_size=1, show_progress_bar=False) return {"ok": True, "vectors": _checked_batch(vectors, len(texts), len(texts), "vectors")} @@ -856,13 +921,15 @@ def _handle_fast_embed( model = None try: deadline.start_inference() - model = self._get_fast_encoder(tier) + model = self._get_fast_encoder(tier, deadline=deadline) deadline.check() if model is None: if deadline.transport is not None: deadline.transport.started = was_started # Custom identities decline rather than drift. return None + self._check_process_memory("fast embed slice", deadline) + deadline.check() vectors = model.encode(texts, batch_size=1, show_progress_bar=False) finally: model = None @@ -870,7 +937,7 @@ def _handle_fast_embed( if fast_owned: self._retire_stale_fast_encoder() return {"ok": True, "vectors": _checked_batch(vectors, len(texts), len(texts), "vectors")} - except (_RequestDeadlineExceeded, _ResultTooLarge): + except (_RequestDeadlineExceeded, _ResultTooLarge, _ProcessMemoryRefused): raise except Exception: if deadline.transport is not None: @@ -896,6 +963,8 @@ def handle_request(self, request: dict) -> dict: return {"ok": False, "error": str(exc)} except _ResultTooLarge as exc: return {"ok": False, "error": str(exc), "error_code": "result_too_large"} + except _ProcessMemoryRefused as exc: + return {"ok": False, "error": str(exc), "error_code": "server_busy", "retry_after_ms": 250} finally: with self._activity_lock: self._inflight -= 1 @@ -960,7 +1029,7 @@ def _handle_request_inner(self, request: dict) -> dict: deadline.start_inference() self._preflight_embed_result(tier, len(texts)) deadline.check() - model = self._get_embed_model(tier) + model = self._get_embed_model(tier, deadline=deadline) order = self._embed_global_order(model, texts, limit, deadline) while offset < len(texts) or vectors is None: deadline.check() @@ -969,6 +1038,7 @@ def _handle_request_inner(self, request: dict) -> dict: ) batch = (texts[offset:offset + limit] if indices is None else [texts[index] for index in indices]) + self._check_process_memory("embed slice", deadline) deadline.check() recover = False try: @@ -990,6 +1060,7 @@ def _handle_request_inner(self, request: dict) -> dict: if retry is not None: # Release only the state lock; inference ownership # still covers the retry and remaining slices. + self._check_process_memory("embed CPU retry", deadline) deadline.check() log.warning("MPS OOM during encoding; retrying microbatch on CPU") values = model.encode(retry, batch_size=limit, show_progress_bar=False) @@ -1032,7 +1103,7 @@ def _handle_request_inner(self, request: dict) -> dict: deadline.start_inference() self._preflight_rerank_result(model_name, len(pairs)) deadline.check() - reranker = self._get_reranker(model_name) + reranker = self._get_reranker(model_name, deadline=deadline) order = self._rerank_global_order(reranker, pairs, limit, deadline) if order is not None: recovery_name = self._reranker_name @@ -1043,6 +1114,7 @@ def _handle_request_inner(self, request: dict) -> dict: ) batch = (pairs[offset:offset + limit] if indices is None else [pairs[index] for index in indices]) + self._check_process_memory("reranker slice", deadline) deadline.check() recover = False try: @@ -1063,7 +1135,7 @@ def _handle_request_inner(self, request: dict) -> dict: deadline.check() flush_mps_cache() deadline.check() - reranker = self._get_reranker(recovery_name) + reranker = self._get_reranker(recovery_name, deadline=deadline) retry = batch break if indices is None: @@ -1074,6 +1146,7 @@ def _handle_request_inner(self, request: dict) -> dict: ) offset += len(batch) if retry is not None: + self._check_process_memory("reranker CPU retry", deadline) deadline.check() values = reranker.predict(retry, batch_size=limit, show_progress_bar=False) if indices is None: