diff --git a/tests/test-cpu-fallback-token-budget.py b/tests/test-cpu-fallback-token-budget.py new file mode 100644 index 00000000..8a265c56 --- /dev/null +++ b/tests/test-cpu-fallback-token-budget.py @@ -0,0 +1,122 @@ +"""After an MPS OOM, CPU pieces are sized against the RSS budget, not free system memory.""" +from __future__ import annotations + +import sys +import types + +import numpy as np +import pytest + +from truememory import memory_budget +from truememory import model_server as ms +from truememory import mps_utils +from truememory._embed_budget import EMBED_TOKEN_BUDGET +from truememory._rerank_budget import accelerator_token_budget, budget_batch_size + +GIB = 1024 ** 3 +RESIDENT = 3_560_000_000 # the soak's resident set beside a 4 GiB budget (16 GiB Mac) + + +class FakeRSS: + def __init__(self, rss: int) -> None: + self.rss = rss + + def Process(self) -> "FakeRSS": + return self + + def virtual_memory(self) -> types.SimpleNamespace: + # Plenty of free system memory: the old sizing saw no pressure at all. + return types.SimpleNamespace(available=16 * GIB, total=16 * GIB) + + def memory_info(self) -> types.SimpleNamespace: + return types.SimpleNamespace(rss=self.rss) + + +class Embedder: + """Qwen-shaped: 16 heads, word count stands in for tokens, MPS can OOM.""" + + config = types.SimpleNamespace(num_attention_heads=16) + + def __init__(self, device: str) -> None: + self.device = device + self.max_seq_length = 2048 + self.calls: list[tuple[str, int, int]] = [] + self.oom_on_mps = True + + def tokenizer(self, texts, truncation, max_length): + return {"input_ids": [[0] * min(len(text.split()), max_length) for text in texts]} + + def encode(self, texts, batch_size=32, show_progress_bar=False, **kwargs): + longest = max(len(text.split()) for text in texts) + self.calls.append((self.device, len(texts), longest)) + if self.device == "mps" and self.oom_on_mps: + raise RuntimeError("MPS backend out of memory (MPS allocated: 3.9 GiB, max allowed: 4.0 GiB)") + return np.array([[len(text.split()), index] for index, text in enumerate(texts)], dtype=np.float32) + + def to(self, device: str) -> "Embedder": + self.device = device + return self + + +@pytest.fixture +def rss_budget(monkeypatch): + sensor = FakeRSS(RESIDENT) + monkeypatch.setitem(sys.modules, "psutil", sensor) + monkeypatch.setattr(memory_budget, "process_rss_limit", lambda: 4 * GIB, raising=False) + return sensor + + +def _cpu_fit() -> int: + return int((4 * GIB - RESIDENT) * 0.8) // (16 * 8) + + +def test_cpu_budget_is_the_rss_room_not_free_system_memory(rss_budget): + model = Embedder("cpu") + fit = accelerator_token_budget(model, EMBED_TOKEN_BUDGET) + assert fit == _cpu_fit() == 4_593_545 + # The soak's first CPU piece was 11 texts of 1,718 tokens: 4.16 GB of attention. + assert budget_batch_size(1718, 32, EMBED_TOKEN_BUDGET) == 11 + assert budget_batch_size(1718, 32, fit) == 1 + assert 1 * 1718 ** 2 * 16 * 8 < (4 * GIB - RESIDENT) + + +def test_cpu_budget_floors_at_one_text_and_keeps_room_when_plentiful(rss_budget, monkeypatch): + model = Embedder("cpu") + rss_budget.rss = 5 * GIB + assert accelerator_token_budget(model, EMBED_TOKEN_BUDGET) == 1 + rss_budget.rss = 1 * GIB + monkeypatch.setattr(memory_budget, "process_rss_limit", lambda: 8 * GIB) + assert accelerator_token_budget(model, EMBED_TOKEN_BUDGET) == EMBED_TOKEN_BUDGET + + +def test_models_without_known_heads_keep_their_budget(rss_budget): + model = types.SimpleNamespace(device="cpu") + assert accelerator_token_budget(model, EMBED_TOKEN_BUDGET) == EMBED_TOKEN_BUDGET + + +def test_server_retries_the_oom_piece_and_the_rest_in_rss_sized_pieces(rss_budget, monkeypatch): + monkeypatch.delenv("TRUEMEMORY_DEVICE", raising=False) + monkeypatch.setattr(mps_utils, "auto_detect_device", lambda: "mps") + monkeypatch.setattr(mps_utils, "ensure_mps_memory_budget", lambda device: None) + # No configured MPS cap: the MPS pieces keep the plain 32 x 1,024^2 budget. + monkeypatch.setattr(mps_utils, "get_mps_memory_budget", lambda: None) + server = ms.ModelServer() + server._write_status_file = lambda: None + server._embed_device = lambda: "mps" + model = Embedder("mps") + monkeypatch.setattr(server, "_build_embed_model", lambda model_id, device: model) + texts = [" ".join(["memory"] * 1718) for _ in range(22)] + [" ".join(["short"] * 20) for _ in range(10)] + try: + response = server.handle_request({"op": "embed", "texts": texts, "tier": "qwen3_256"}) + finally: + server._workers.shutdown(wait=False) + assert response["ok"] is True + # Row order is preserved: each vector carries its own text's length. + assert [int(row[0]) for row in response["vectors"]] == [1718] * 22 + [20] * 10 + mps_calls = [call for call in model.calls if call[0] == "mps"] + cpu_calls = [call for call in model.calls if call[0] == "cpu"] + assert mps_calls == [("mps", 11, 1718)], "The MPS piece is sized by the MPS budget" + fit = _cpu_fit() + assert cpu_calls and all(size * longest ** 2 <= fit or size == 1 for _, size, longest in cpu_calls) + assert all(size == 1 for _, size, longest in cpu_calls if longest == 1718) + assert sum(size for _, size, _ in cpu_calls) == 32 diff --git a/tests/test-embed-token-budget.py b/tests/test-embed-token-budget.py new file mode 100644 index 00000000..974c74bb --- /dev/null +++ b/tests/test-embed-token-budget.py @@ -0,0 +1,220 @@ +"""Embedding input cap and token-budgeted batches bound attention memory.""" +from __future__ import annotations + +import numpy as np +import pytest + +from truememory import _embed_budget as embed_budget +from truememory._embed_budget import ( + EMBED_MAX_TOKENS, EMBED_TOKEN_BUDGET, budgeted_encode, cap_embed_length, embed_token_budget, + text_token_lengths, +) + + +class FakeEmbedder: + """Word count stands in for tokens; each vector depends only on its text.""" + + def __init__(self, max_seq_length: int | None = 32768) -> None: + self.max_seq_length = max_seq_length + self.calls: list[tuple[int, int]] = [] + + def tokenizer(self, texts, truncation, max_length): + assert truncation is True + self.tokenized = getattr(self, "tokenized", 0) + 1 + return {"input_ids": [[0] * min(len(text.split()) + 2, max_length) for text in texts]} + + def encode(self, texts, batch_size=32, show_progress_bar=False, **kwargs): + self.calls.append((len(texts), batch_size)) + return np.array([[len(text.split()), len(text) % 7] for text in texts], dtype=np.float32) + + +def _texts(*word_counts: int) -> list[str]: + return [" ".join(["word"] * count) for count in word_counts] + + +def test_cap_lowers_long_limits_only(): + for current, expected in ((32768, EMBED_MAX_TOKENS), (512, 512), (None, None)): + model = FakeEmbedder(current) + cap_embed_length(model) + assert model.max_seq_length == expected + cap_embed_length(object()) + + +def test_budget_shrinks_under_memory_pressure(): + assert embed_token_budget(16 * 1024 ** 3) == EMBED_TOKEN_BUDGET + assert embed_token_budget(4 * 1024 ** 3) == EMBED_TOKEN_BUDGET // 2 + assert embed_token_budget(1024 ** 3) == EMBED_TOKEN_BUDGET // 4 + + +def test_texts_up_to_1024_tokens_keep_one_unchanged_call(monkeypatch): + monkeypatch.setattr(embed_budget, "embed_token_budget", lambda: EMBED_TOKEN_BUDGET) + model = FakeEmbedder(EMBED_MAX_TOKENS) + texts = _texts(*([1000] * 40 + [5] * 40)) + values = budgeted_encode(model, texts, batch_size=32, show_progress_bar=False) + assert model.calls == [(80, 32)] + assert values.shape == (80, 2) + + +def test_long_texts_split_longest_first_and_keep_order(monkeypatch): + monkeypatch.setattr(embed_budget, "embed_token_budget", lambda: EMBED_TOKEN_BUDGET) + model = FakeEmbedder(EMBED_MAX_TOKENS) + texts = _texts(*([20] * 60 + [5000] * 12)) + values = budgeted_encode(model, texts, batch_size=32, show_progress_bar=False) + assert np.array_equal(values, FakeEmbedder().encode(texts)) + assert sum(size for size, _ in model.calls) == 72 + assert model.calls[0] == (EMBED_TOKEN_BUDGET // EMBED_MAX_TOKENS ** 2,) * 2 + assert max(text_token_lengths(model, texts)) == EMBED_MAX_TOKENS + assert all(batch == size <= 32 for size, batch in model.calls) + + +def test_tensor_output_and_static_models_keep_single_call(): + model = FakeEmbedder(EMBED_MAX_TOKENS) + budgeted_encode(model, _texts(5000, 5000), batch_size=32, convert_to_tensor=True) + assert model.calls == [(2, 32)] + + class Static: + calls = [] + + def encode(self, texts, **kwargs): + self.calls.append(len(texts)) + return np.zeros((len(texts), 2), dtype=np.float32) + + static = Static() + from truememory.mps_utils import encode_with_model_ownership + encode_with_model_ownership(static, _texts(5000, 5000), batch_size=32) + assert static.calls == [2] + + +def test_local_encode_wrappers_use_the_budget(monkeypatch): + monkeypatch.setattr(embed_budget, "embed_token_budget", lambda: EMBED_TOKEN_BUDGET) + from truememory.mps_utils import encode_with_model_ownership, encode_with_mps_fallback + for wrapper in (encode_with_model_ownership, encode_with_mps_fallback): + model = FakeEmbedder(EMBED_MAX_TOKENS) + texts = _texts(*([10] * 30 + [4000] * 9)) + values = wrapper(model, texts, batch_size=32, show_progress_bar=False) + assert np.array_equal(values, FakeEmbedder().encode(texts)) + assert model.calls[0][0] == EMBED_TOKEN_BUDGET // EMBED_MAX_TOKENS ** 2 + + +def test_model_server_embed_uses_budgeted_slices(monkeypatch): + from truememory.model_server import ModelServer, _EmbedState + + monkeypatch.setattr(embed_budget, "embed_token_budget", lambda: EMBED_TOKEN_BUDGET) + server = ModelServer() + model = FakeEmbedder(EMBED_MAX_TOKENS) + server._embed_state = _EmbedState(model=model, tier="", model_id="model2vec") + texts = _texts(*([30] * 50 + [3000] * 10)) + response = server.handle_request({"op": "embed", "texts": texts, "tier": ""}) + assert response["ok"] is True + assert np.array_equal(np.asarray(response["vectors"]), FakeEmbedder().encode(texts)) + assert model.calls[0] == (EMBED_TOKEN_BUDGET // EMBED_MAX_TOKENS ** 2,) * 2 + assert all(size <= 32 for size, _ in model.calls) + + +def test_model_server_short_embed_unchanged(): + from truememory.model_server import ModelServer, _EmbedState + + server = ModelServer() + model = FakeEmbedder(EMBED_MAX_TOKENS) + server._embed_state = _EmbedState(model=model, tier="", model_id="model2vec") + response = server.handle_request({"op": "embed", "texts": _texts(*range(5, 45)), "tier": ""}) + assert response["ok"] is True + assert model.calls == [(32, 32), (8, 32)] + + +@pytest.mark.parametrize("length", [1, 1024, 1025, 2048]) +def test_batch_size_bounds(length): + from truememory._rerank_budget import budget_batch_size + assert 1 <= budget_batch_size(length, 32, EMBED_TOKEN_BUDGET) <= 32 + + +def test_single_text_and_list_outputs_skip_measurement(): + model = FakeEmbedder(EMBED_MAX_TOKENS) + budgeted_encode(model, _texts(5000), batch_size=32) + budgeted_encode(model, _texts(5000, 5000), batch_size=32, convert_to_numpy=False) + assert getattr(model, "tokenized", 0) == 0 + assert model.calls == [(1, 32), (2, 32)] + + +def test_short_batches_skip_the_memory_read(monkeypatch): + def unexpected(): + raise AssertionError("memory read for a batch that always fits") + + monkeypatch.setattr(embed_budget, "embed_token_budget", unexpected) + model = FakeEmbedder(EMBED_MAX_TOKENS) + budgeted_encode(model, _texts(*([500] * 32)), batch_size=32) + assert model.calls == [(32, 32)] + + +def _on_mps(monkeypatch, model, allocated_mib, cap_mib=1536, heads=16): + import sys + import types + + from truememory import mps_utils + model.device = types.SimpleNamespace(type="mps") + model.config = types.SimpleNamespace(num_attention_heads=heads) + monkeypatch.setattr(mps_utils, "get_mps_memory_budget", + lambda: types.SimpleNamespace(effective_bytes=cap_mib * 1024 ** 2)) + monkeypatch.setitem(sys.modules, "torch", types.SimpleNamespace( + mps=types.SimpleNamespace(current_allocated_memory=lambda: allocated_mib * 1024 ** 2))) + + +def test_mps_batches_fit_the_free_allocation(monkeypatch): + monkeypatch.setattr(embed_budget, "embed_token_budget", lambda: EMBED_TOKEN_BUDGET) + model = FakeEmbedder(EMBED_MAX_TOKENS) + _on_mps(monkeypatch, model, allocated_mib=1190) + texts = _texts(*([45] * 31 + [1000])) + values = budgeted_encode(model, texts, batch_size=32) + # 80% of the free 346 MiB over 16 heads x 48 bytes per score: 377,924 tokens^2. + fit = int(346 * 1024 ** 2 * 0.8) // (16 * 48) + # Below one long text's 1,004,004 tokens^2, so it is embedded alone. + assert fit < 1002 ** 2 + assert model.calls[0] == (1, 1) + assert sum(size for size, _ in model.calls) == 32 + assert np.array_equal(values, FakeEmbedder().encode(texts)) + + +def test_mps_short_batches_keep_one_call(monkeypatch): + model = FakeEmbedder(EMBED_MAX_TOKENS) + _on_mps(monkeypatch, model, allocated_mib=1190) + budgeted_encode(model, _texts(*([45] * 32)), batch_size=32) + assert model.calls == [(32, 32)] + + +def test_cpu_models_ignore_the_mps_allocation(monkeypatch): + from truememory._rerank_budget import accelerator_token_budget + model = FakeEmbedder(EMBED_MAX_TOKENS) + _on_mps(monkeypatch, model, allocated_mib=1500) + model.device = "cpu" + # CPU pieces are sized against the process RSS budget; give it ample room + # so the result does not depend on how much the test process holds. + from truememory import memory_budget + monkeypatch.setattr(memory_budget, "process_rss_limit", lambda: 1 << 40) + assert accelerator_token_budget(model, EMBED_TOKEN_BUDGET) == EMBED_TOKEN_BUDGET + + +def test_server_embed_loop_fits_the_free_mps_allocation(monkeypatch): + from truememory.model_server import ModelServer, _EmbedState + + monkeypatch.setattr(embed_budget, "embed_token_budget", lambda: EMBED_TOKEN_BUDGET) + model = FakeEmbedder(EMBED_MAX_TOKENS) + _on_mps(monkeypatch, model, allocated_mib=1190, cap_mib=4096) + server = ModelServer() + server._embed_state = _EmbedState(model=model, tier="", model_id="model2vec") + texts = _texts(*([30] * 40 + [1000] * 8)) + response = server.handle_request({"op": "embed", "texts": texts, "tier": ""}) + assert response["ok"] is True + assert np.array_equal(np.asarray(response["vectors"]), FakeEmbedder().encode(texts)) + # 80% of the free 2,906 MiB over 16 heads x 48 bytes: 3,174,345 tokens^2, so three 1,002-token texts. + fit = int(2906 * 1024 ** 2 * 0.8) // (16 * 48) + assert model.calls[0][0] == fit // 1002 ** 2 == 3 + + +def test_attention_heads_are_read_through_sentence_transformer_modules(): + import types + + from truememory._rerank_budget import _attention_heads + module = types.SimpleNamespace(auto_model=types.SimpleNamespace(config=types.SimpleNamespace(num_attention_heads=16))) + model = types.SimpleNamespace(_first_module=lambda: module) + assert _attention_heads(model) == 16 + assert _attention_heads(types.SimpleNamespace()) is None diff --git a/tests/test-model-daemon-admission-297.py b/tests/test-model-daemon-admission-297.py index 89e8ba71..5c041f03 100644 --- a/tests/test-model-daemon-admission-297.py +++ b/tests/test-model-daemon-admission-297.py @@ -209,7 +209,7 @@ def test_original_module_and_console_launch_contract_is_preserved(self) -> None: self.assertIn('[executable, "-m", "truememory.model_server"]', client) self.assertIn('truememory-model-server = "truememory.model_server:main"', (ROOT / "pyproject.toml").read_text()) self.assertEqual(hashlib.sha256((ROOT / "truememory/mps_utils.py").read_bytes()).hexdigest(), - "0a0e2e8e1a24c307d3a612f87e7ea05c67462fd1b2612ed753b0b668392592e8") + "1433f88ae29d56c2206fc22c2cd5ffeb58b30016a8ba4addbde14691751b5b23") if __name__ == "__main__": diff --git a/tests/test-model-server-memory-pressure.py b/tests/test-model-server-memory-pressure.py index aa4336eb..a2d83deb 100644 --- a/tests/test-model-server-memory-pressure.py +++ b/tests/test-model-server-memory-pressure.py @@ -349,6 +349,30 @@ def test_session_release_returns_one_stage(budgeted): +def test_short_deadline_bounds_the_bookkeeping_wait(plain, monkeypatch): + plain._publish_embed_state(ms._EmbedState(model=Model("cpu"), tier="", model_id="model2vec")) + monkeypatch.setattr(plain, "_build_embed_model", lambda model_id, device: Model(device)) + held, done = threading.Event(), threading.Event() + + def bookkeeping(): + with plain._lock: + held.set() + done.wait(5) + + worker = threading.Thread(target=bookkeeping) + worker.start() + held.wait() + try: + started = time.monotonic() + plain._handle_fast_embed(["synthetic"], "", ms._RequestDeadline(time.monotonic() + 0.05)) + assert time.monotonic() - started < 0.25 + except ms._RequestDeadlineExceeded: + assert time.monotonic() - started < 0.25 + finally: + done.set() + worker.join() + + def test_loading_model_does_not_reserve_its_size_against_concurrent_work(budgeted, monkeypatch): server, rss = budgeted monkeypatch.setattr(mps_utils, "auto_detect_device", lambda: "cpu") diff --git a/tests/test-reranker-token-budget.py b/tests/test-reranker-token-budget.py index 68db22f3..54669a90 100644 --- a/tests/test-reranker-token-budget.py +++ b/tests/test-reranker-token-budget.py @@ -198,3 +198,28 @@ def test_model_server_caps_and_slices_legacy_crossencoder(monkeypatch): assert response["ok"] is True assert server._reranker.max_length == RERANK_MAX_TOKENS assert server._reranker.calls[0][0] == RERANK_TOKEN_BUDGET // RERANK_MAX_TOKENS ** 2 + + +def test_server_rerank_loop_fits_the_free_mps_allocation(monkeypatch): + import sys + import types + + from truememory import mps_utils + from truememory.model_server import ModelServer + + fake = FakeCrossEncoder() + fake.device = types.SimpleNamespace(type="mps") + fake.config = types.SimpleNamespace(num_attention_heads=12) + monkeypatch.setattr(mps_utils, "get_mps_memory_budget", + lambda: types.SimpleNamespace(effective_bytes=4096 * 1024 ** 2)) + monkeypatch.setitem(sys.modules, "torch", types.SimpleNamespace( + mps=types.SimpleNamespace(current_allocated_memory=lambda: 1790 * 1024 ** 2))) + monkeypatch.setattr(ModelServer, "_get_reranker", lambda self, model_name=None, deadline=None: fake) + monkeypatch.setattr(rerank_budget, "rerank_token_budget", lambda: RERANK_TOKEN_BUDGET) + pairs = _pairs(*([30] * 90 + [2000] * 10)) + response = ModelServer().handle_request({"op": "rerank", "pairs": pairs, "model_name": "m"}) + assert response["ok"] is True + assert np.allclose(response["scores"], FakeCrossEncoder().predict(pairs, batch_size=len(pairs)), rtol=0, atol=0) + # 80% of the free 2,306 MiB over 12 heads x 48 bytes: 3,358,665 tokens^2, three 1,024-token pairs. + fit = int(2306 * 1024 ** 2 * 0.8) // (12 * 48) + assert fake.calls[0][0] == fit // RERANK_MAX_TOKENS ** 2 == 3 diff --git a/truememory/_embed_budget.py b/truememory/_embed_budget.py new file mode 100644 index 00000000..00e91030 --- /dev/null +++ b/truememory/_embed_budget.py @@ -0,0 +1,66 @@ +"""Embedding input cap and token-budgeted batches that bound attention memory.""" +from __future__ import annotations + +from truememory._rerank_budget import ( + accelerator_token_budget, budget_batch_size, budgeted_slices, memory_pressure_divisor, +) + +# Above the longest stored memory, so existing vectors are unchanged; only +# oversized texts such as pasted documents are truncated. +EMBED_MAX_TOKENS = 2048 +# Keeps the default 32-text batch for every text up to 1,024 tokens. +EMBED_TOKEN_BUDGET = 32 * 1024 * 1024 + + +def cap_embed_length(model: object) -> None: + """Lower an embedder's integer input limit to the cap; static models are untouched.""" + current = getattr(model, "max_seq_length", None) + if type(current) is int and current > EMBED_MAX_TOKENS: + model.max_seq_length = EMBED_MAX_TOKENS + + +def embed_token_budget(available: int | None = None) -> int: + return EMBED_TOKEN_BUDGET // memory_pressure_divisor(available) + + +def measurable(model: object) -> bool: + return getattr(model, "tokenizer", None) is not None and type(getattr(model, "max_seq_length", None)) is int + + +def text_token_lengths(model: object, texts: object) -> list[int] | None: + """Truncated token length of each text, or None when it cannot be measured.""" + if not measurable(model) or not isinstance(texts, list) or not all(isinstance(text, str) for text in texts): + return None + if not texts: + return [] + encoded = model.tokenizer(texts, truncation=True, max_length=model.max_seq_length) + return [len(ids) for ids in encoded["input_ids"]] + + +def budgeted_encode(model: object, texts: object, **kwargs: object): + """SentenceTransformer.encode with long texts embedded in smaller batches. + + Inputs whose default batches already fit keep the original single call. + """ + batch_size = kwargs.get("batch_size", 32) + # One text is always one batch, so single queries skip measurement entirely. + plain = (isinstance(texts, list) and len(texts) > 1 and type(batch_size) is int and batch_size > 0 + and not kwargs.get("convert_to_tensor") and kwargs.get("convert_to_numpy", True) is not False + and kwargs.get("output_value", "sentence_embedding") == "sentence_embedding") + lengths = text_token_lengths(model, texts) if plain else None + if lengths is None: + return model.encode(texts, **kwargs) + # Batches that fit even the smallest budget skip the memory read. + if budget_batch_size(max(lengths), batch_size, + accelerator_token_budget(model, EMBED_TOKEN_BUDGET // 4)) >= batch_size: + return model.encode(texts, **kwargs) + slices = budgeted_slices(lengths, batch_size, accelerator_token_budget(model, embed_token_budget())) + if slices is None: + return model.encode(texts, **kwargs) + import numpy as np + results: list = [None] * len(texts) + for indices in slices: + values = model.encode([texts[index] for index in indices], **{**kwargs, "batch_size": len(indices)}) + for position, index in enumerate(indices): + results[index] = values[position] + return np.stack(results) diff --git a/truememory/_rerank_budget.py b/truememory/_rerank_budget.py index 925196e9..21b1a5de 100644 --- a/truememory/_rerank_budget.py +++ b/truememory/_rerank_budget.py @@ -37,19 +37,98 @@ def cap_reranker_length(model: object) -> None: setattr(model, name, RERANK_MAX_TOKENS) -def rerank_token_budget(available: int | None = None) -> int: - """Per-batch attention budget, reduced while the host is short of memory.""" +def memory_pressure_divisor(available: int | None = None) -> int: + """1 normally, 2 below 6 GiB available and 4 below 2 GiB.""" if available is None: try: import psutil available = psutil.virtual_memory().available except (ImportError, OSError): - return RERANK_TOKEN_BUDGET + return 1 if available < 2 * 1024 ** 3: - return RERANK_TOKEN_BUDGET // 4 + return 4 if available < 6 * 1024 ** 3: - return RERANK_TOKEN_BUDGET // 2 - return RERANK_TOKEN_BUDGET + return 2 + return 1 + + +# Peak MPS footprint per batch x tokens^2 x attention head, measured on Apple +# silicon for Qwen3-Embedding and the ModernBERT reranker with both resident +# (scores, masks, softmax and the allocator's heaps), with a fifth in reserve. +_MPS_BYTES_PER_SCORE = 48 +_MPS_HEADROOM_SHARE = 0.8 + + +def _attention_heads(model: object) -> int | None: + configs = [getattr(model, "config", None), getattr(getattr(model, "model", None), "config", None)] + first = getattr(model, "_first_module", None) + if callable(first): + try: + configs.append(getattr(getattr(first(), "auto_model", None), "config", None)) + except (AttributeError, IndexError, KeyError, StopIteration, TypeError): + pass + for config in configs: + heads = getattr(config, "num_attention_heads", None) + if type(heads) is int and heads > 0: + return heads + return None + + +# Peak CPU bytes per batch x tokens^2 x attention head: eager attention keeps +# two float32 score tensors alive at its softmax (11 x 1,718 tokens x 16 heads +# is 4.16 GB per layer). +_CPU_BYTES_PER_SCORE = 8 + + +def _cpu_token_budget(model: object, budget: int) -> int: + """Lower *budget* so one CPU batch's attention fits the process RSS budget. + + CPU attention lands in resident memory. Pieces are sized against the RSS + budget minus what is already resident, not against free system memory. + """ + heads = _attention_heads(model) + if heads is None: + return budget + try: + import psutil + from truememory.memory_budget import process_rss_limit + room = process_rss_limit() - psutil.Process().memory_info().rss + except Exception: + return budget + fit = int(max(0, room) * _MPS_HEADROOM_SHARE) // (heads * _CPU_BYTES_PER_SCORE) + return max(1, min(budget, fit)) + + +def accelerator_token_budget(model: object, budget: int) -> int: + """Lower *budget* so one batch's attention fits where the model runs. + + MPS allocations are capped per process: a batch above the remaining cap + fails with an OOM that degrades the model to CPU. On CPU the attention is + resident memory, so a batch must fit the process RSS budget. + """ + device = getattr(model, "device", None) + kind = getattr(device, "type", device) + if kind == "cpu": + return _cpu_token_budget(model, budget) + if kind != "mps": + return budget + from truememory.mps_utils import get_mps_memory_budget + policy = get_mps_memory_budget() + heads = _attention_heads(model) + if policy is None or not policy.effective_bytes or heads is None: + return budget + try: + import torch + free = policy.effective_bytes - torch.mps.current_allocated_memory() + except (ImportError, AttributeError, RuntimeError): + return budget + fit = int(max(0, free) * _MPS_HEADROOM_SHARE) // (heads * _MPS_BYTES_PER_SCORE) + return max(1, min(budget, fit)) + + +def rerank_token_budget(available: int | None = None) -> int: + """Per-batch attention budget, reduced while the host is short of memory.""" + return RERANK_TOKEN_BUDGET // memory_pressure_divisor(available) def pair_token_lengths(model: object, pairs: list) -> list[int] | None: @@ -92,7 +171,8 @@ def budgeted_predict(model: object, pairs: list, batch_size: int): """CrossEncoder.predict with long pairs scored in smaller batches.""" lengths = (pair_token_lengths(model, pairs) if type(batch_size) is int and batch_size > 0 else None) - slices = None if lengths is None else budgeted_slices(lengths, batch_size, rerank_token_budget()) + slices = (None if lengths is None else + budgeted_slices(lengths, batch_size, accelerator_token_budget(model, rerank_token_budget()))) if slices is None: return model.predict(pairs, batch_size=batch_size, show_progress_bar=False) import numpy as np diff --git a/truememory/embedding_target.py b/truememory/embedding_target.py index 0fac7673..f37294b4 100644 --- a/truememory/embedding_target.py +++ b/truememory/embedding_target.py @@ -122,13 +122,19 @@ def admit() -> None: if target.model_id == "qwen3_256": model_kwargs = {"attn_implementation": "eager"} if sys.platform == "darwin" else None admit() - return prefer_float32_on_cpu(SentenceTransformer( + model = SentenceTransformer( model_repository(target.model_id), truncate_dim=256, config_kwargs={"use_cache": False}, model_kwargs=model_kwargs, device=device, - )) - admit() - return prefer_float32_on_cpu(SentenceTransformer( - target.model_id, truncate_dim=target.dimension, - trust_remote_code=False, device=device, - )) + ) + else: + admit() + model = SentenceTransformer( + target.model_id, truncate_dim=target.dimension, + trust_remote_code=False, device=device, + ) + prefer_float32_on_cpu(model) + if type(getattr(model, "max_seq_length", None)) is int: + from truememory._embed_budget import cap_embed_length + cap_embed_length(model) + return model diff --git a/truememory/memory_budget.py b/truememory/memory_budget.py index 168967d5..29748498 100644 --- a/truememory/memory_budget.py +++ b/truememory/memory_budget.py @@ -414,6 +414,11 @@ def _record(self, token: int) -> _OwnerRecord: raise StaleBudgetOwner("Owner has retired") return record + def ceiling(self) -> int: + """Tightest enrolled owner limit in bytes, or 0 when none is enrolled.""" + with self._lock: + return self._ceiling() + def _retired(self, owner: BudgetOwner) -> bool: """Whether an owner of this registry is closing or gone (not inherited or foreign).""" with self._lock: @@ -714,6 +719,12 @@ def resolve_allocation_plan(key: ProfileKey) -> AllocationPlan: return (registry or get_process_registry()).resolve_profile(key) +def process_rss_limit() -> int: + """The process RSS ceiling: the tightest enrolled owner, else the automatic budget.""" + ceiling = get_process_registry().ceiling() + return ceiling if ceiling else automatic_budget_bytes() + + def automatic_budget_bytes() -> int: try: import psutil diff --git a/truememory/model_server.py b/truememory/model_server.py index d3fbcd6b..115b5f49 100644 --- a/truememory/model_server.py +++ b/truememory/model_server.py @@ -125,6 +125,40 @@ def __init__(self, message: str, *, oom: bool) -> None: self.mps_oom = oom +def _run_in_pieces(run: Callable[[list, int], object], items: list, lengths: list[int] | None, + budget: int) -> object: + """Run *items* in longest-first pieces whose batch x tokens^2 fits *budget*. + + Results come back in item order. Without lengths, one call runs them all. + """ + from truememory._rerank_budget import budgeted_slices + slices = None if lengths is None else budgeted_slices(lengths, max(1, len(items)), budget) + if slices is None: + return run(items, max(1, len(items))) + results: list = [None] * len(items) + for group in slices: + values = run([items[index] for index in group], len(group)) + for position, index in enumerate(group): + results[index] = values[position] + return np.stack(results) + + +def _reslice_remaining(order: np.ndarray | None, lengths: list[int], start: int, total: int, limit: int, + budget: int) -> tuple[np.ndarray, dict[int, int]]: + """Order and piece sizes for positions from *start* on, under a new budget.""" + from truememory._rerank_budget import budgeted_slices + done = [int(index) for index in order[:start]] if order is not None else list(range(start)) + remaining = [int(index) for index in order[start:]] if order is not None else list(range(start, total)) + slices = budgeted_slices([lengths[index] for index in remaining], limit, budget) + groups = ([[remaining[position] for position in group] for group in slices] if slices is not None + else [remaining[offset:offset + limit] for offset in range(0, len(remaining), limit)]) + steps, position = {}, start + for group in groups: + steps[position] = len(group) + position += len(group) + return np.array(done + [index for group in groups for index in group], dtype=np.intp), steps + + def _array_metadata(shape: tuple, encoded: str = "") -> dict: return {"__ndarray__": encoded, "dtype": "float32", "shape": list(shape)} @@ -815,13 +849,17 @@ def _build_embed_model(model_id: str, device: str | None): mkwargs = {} if sys.platform == "darwin": mkwargs["attn_implementation"] = "eager" - return SentenceTransformer( + model = SentenceTransformer( "Qwen/Qwen3-Embedding-0.6B", truncate_dim=256, config_kwargs={"use_cache": False}, model_kwargs=mkwargs or None, device=device, ) + if type(getattr(model, "max_seq_length", None)) is int: + from truememory._embed_budget import cap_embed_length + cap_embed_length(model) + return model if model_id not in ("model2vec", "minilm", "bge-small", "qwen3_256"): # Custom model: require explicit opt-in for arbitrary downloads if os.environ.get("TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD", "").strip() != "1": @@ -841,11 +879,15 @@ def _build_embed_model(model_id: str, device: str | None): from truememory.mps_utils import ensure_mps_memory_budget ensure_mps_memory_budget(device) from sentence_transformers import SentenceTransformer - return SentenceTransformer( + model = SentenceTransformer( model_id, truncate_dim=custom_dim, trust_remote_code=False, device=device, ) + if type(getattr(model, "max_seq_length", None)) is int: + from truememory._embed_budget import cap_embed_length + cap_embed_length(model) + return model from model2vec import StaticModel return StaticModel.from_pretrained( "minishlab/potion-base-8M", force_download=False @@ -2083,6 +2125,8 @@ def _handle_request_inner(self, request: dict) -> dict: vectors = None offset = 0 order = None + steps = None + lengths = None while offset < len(texts) or vectors is None: retry = None with deadline.locked(self._lock): @@ -2101,18 +2145,36 @@ def _handle_request_inner(self, request: dict) -> dict: else self._get_prepared_embed_model(target, deadline)) self._check_serving_embed_model(model, expected_target) order = self._embed_global_order(model, texts, limit, deadline) + if (getattr(model, "tokenizer", None) is not None + and type(getattr(model, "max_seq_length", None)) is int): + from truememory._embed_budget import embed_token_budget, text_token_lengths + from truememory._rerank_budget import accelerator_token_budget, budgeted_slices + lengths = text_token_lengths(model, texts) + slices = (None if lengths is None else budgeted_slices( + lengths, limit, accelerator_token_budget(model, embed_token_budget()))) + if slices is not None: + # Long texts: smaller longest-first batches bound attention memory. + order = np.array([index for group in slices for index in group], dtype=np.intp) + steps, start = {}, 0 + for group in slices: + steps[start] = len(group) + start += len(group) while offset < len(texts) or vectors is None: deadline.check() - indices = None if order is None else self._embed_slice_indices( - model, texts, order, offset, limit, deadline, - ) + step = limit if steps is None else steps[offset] + if steps is not None: + indices = order[offset:offset + step] + else: + indices = None if order is None else self._embed_slice_indices( + model, texts, order, offset, limit, deadline, + ) 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: - values = model.encode(batch, batch_size=limit, show_progress_bar=False) + values = model.encode(batch, batch_size=step, show_progress_bar=False) except RuntimeError as exc: from truememory.mps_utils import is_mps_oom if not is_mps_oom(exc): @@ -2136,7 +2198,20 @@ def _handle_request_inner(self, request: dict) -> dict: 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) + if lengths is None: + values = model.encode(retry, batch_size=len(retry) or limit, show_progress_bar=False) + else: + # MPS-sized pieces do not fit CPU memory: size the retry and + # every remaining piece against the RSS budget instead. + from truememory._embed_budget import embed_token_budget + from truememory._rerank_budget import accelerator_token_budget + budget = accelerator_token_budget(model, embed_token_budget()) + retry_lengths = ([lengths[int(index)] for index in indices] if indices is not None + else lengths[offset:offset + len(retry)]) + values = _run_in_pieces( + lambda items, size: model.encode(items, batch_size=size, show_progress_bar=False), + retry, retry_lengths, budget, + ) if target is not None or expected_target is not None: check_target_vectors(target or expected_target, values, len(retry)) vectors = _store_batch_result( @@ -2144,6 +2219,8 @@ def _handle_request_inner(self, request: dict) -> dict: **result_options, ) offset += len(retry) + if lengths is not None and offset < len(texts): + order, steps = _reslice_remaining(order, lengths, offset, len(texts), limit, budget) self._after_request_batches(throttler, len(texts), encode_start, deadline) with deadline.locked(self._lock): @@ -2181,6 +2258,7 @@ def _handle_request_inner(self, request: dict) -> dict: offset = 0 order = None steps = None + lengths = None recovery_name = model_name while offset < len(pairs) or scores is None: retry = None @@ -2200,11 +2278,12 @@ def _handle_request_inner(self, request: dict) -> dict: type(getattr(reranker, "max_seq_length", None)) is int or hasattr(reranker, "max_length")): from truememory._rerank_budget import ( - budgeted_slices, pair_token_lengths, rerank_token_budget, + accelerator_token_budget, budgeted_slices, pair_token_lengths, + rerank_token_budget, ) lengths = pair_token_lengths(reranker, pairs) - slices = (None if lengths is None else - budgeted_slices(lengths, limit, rerank_token_budget())) + slices = (None if lengths is None else budgeted_slices( + lengths, limit, accelerator_token_budget(reranker, rerank_token_budget()))) if slices is not None: # Long pairs: smaller longest-first batches bound attention scratch. order = np.array([index for group in slices for index in group], dtype=np.intp) @@ -2267,7 +2346,18 @@ def _handle_request_inner(self, request: dict) -> dict: if retry is not None: self._check_process_memory("reranker CPU retry", deadline) deadline.check() - values = reranker.predict(retry, batch_size=len(retry) or limit, show_progress_bar=False) + if lengths is None: + values = reranker.predict(retry, batch_size=len(retry) or limit, show_progress_bar=False) + else: + # Size the retry and remaining pieces for the CPU replacement. + from truememory._rerank_budget import accelerator_token_budget, rerank_token_budget + budget = accelerator_token_budget(reranker, rerank_token_budget()) + retry_lengths = ([lengths[int(index)] for index in indices] if indices is not None + else lengths[offset:offset + len(retry)]) + values = _run_in_pieces( + lambda items, size: reranker.predict(items, batch_size=size, show_progress_bar=False), + retry, retry_lengths, budget, + ) if indices is None: scores = _store_batch_result(scores, values, offset, len(retry), len(pairs), "scores", **result_options) @@ -2277,6 +2367,8 @@ def _handle_request_inner(self, request: dict) -> dict: **result_options, ) offset += len(retry) + if lengths is not None and offset < len(pairs): + order, steps = _reslice_remaining(order, lengths, offset, len(pairs), limit, budget) self._after_request_batches(throttler, len(pairs), predict_start, deadline) return ({"ok": True, "scores": scores} if expected_name is None else {"ok": True, "scores": scores, "protocol": PROTOCOL_VERSION, diff --git a/truememory/mps_utils.py b/truememory/mps_utils.py index 5442ad39..49a6a5ca 100644 --- a/truememory/mps_utils.py +++ b/truememory/mps_utils.py @@ -397,6 +397,9 @@ def encode_with_model_ownership(model: object, texts: object, **kwargs: object) if isinstance(model, EmbeddingProxy): return model.encode(texts, **kwargs) with _own_model(model): + if getattr(model, "tokenizer", None) is not None and type(getattr(model, "max_seq_length", None)) is int: + from truememory._embed_budget import budgeted_encode + return budgeted_encode(model, texts, **kwargs) return model.encode(texts, **kwargs) @@ -409,8 +412,12 @@ def encode_with_mps_fallback(model: object, texts: object, **kwargs: object) -> # prevent recall from using the daemon's independent fast lane. return model.encode(texts, **kwargs) + bounded = getattr(model, "tokenizer", None) is not None and type(getattr(model, "max_seq_length", None)) is int with _own_model(model): try: + if bounded: + from truememory._embed_budget import budgeted_encode + return budgeted_encode(model, texts, **kwargs) return model.encode(texts, **kwargs) except RuntimeError as exc: if not is_mps_oom(exc): @@ -422,4 +429,7 @@ def encode_with_mps_fallback(model: object, texts: object, **kwargs: object) -> model.to("cpu") prefer_float32_on_cpu(model) flush_mps_cache() + if bounded: + from truememory._embed_budget import budgeted_encode + return budgeted_encode(model, texts, **kwargs) return model.encode(texts, **kwargs) diff --git a/truememory/vector_search.py b/truememory/vector_search.py index 962c4c24..3bc8cfe3 100644 --- a/truememory/vector_search.py +++ b/truememory/vector_search.py @@ -381,6 +381,9 @@ def get_model(): from model2vec import StaticModel _model = StaticModel.from_pretrained("minishlab/potion-base-8M", force_download=False) _embedding_dim = 256 + if type(getattr(_model, "max_seq_length", None)) is int: + from truememory._embed_budget import cap_embed_length + cap_embed_length(_model) if budget_owner is not None: bind_model_owner(_model, budget_owner) except BaseException: @@ -622,10 +625,16 @@ def encode( native_tensor_shape=(count, self.target.dimension)), known_bytes=payload_bytes, retained_bytes=payload_bytes, deadline=lambda: _target_remaining(expires), group="result assembly") - values = encoder.encode( - [TARGET_PROBE_TEXT] if probe else texts, - batch_size=batch_size, show_progress_bar=False, - ) + if (getattr(encoder, "tokenizer", None) is not None + and type(getattr(encoder, "max_seq_length", None)) is int): + from truememory._embed_budget import budgeted_encode + values = budgeted_encode(encoder, [TARGET_PROBE_TEXT] if probe else texts, + batch_size=batch_size, show_progress_bar=False) + else: + 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, count)