fix: Bound embedding input length and attention memory per batch

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

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

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__":

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")