fix: Bound embedding input length and attention memory per batch

tests/test-embed-token-budget.py

@@ -0,0 +1,146 @@
+"""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)]

tests/test-model-daemon-admission-297.py

@@ -201,7 +201,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(),
-                         "7057b1e5cdbd23d4de625e6e0c999d2a71455df5fec421842dd9841dabe1b69b")
+                         "adbe1e27ca5dff795f98b2ea3937b9174205ddd8ec327e8316dd9a56b5e3159a")
 
 
 if __name__ == "__main__":

truememory/_embed_budget.py

@@ -0,0 +1,61 @@
+"""Embedding input cap and token-budgeted batches that bound attention memory."""
+from __future__ import annotations
+
+from truememory._rerank_budget import 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
+    # Batches that fit even the smallest budget skip the memory read.
+    if lengths is None or budget_batch_size(max(lengths), batch_size, EMBED_TOKEN_BUDGET // 4) >= batch_size:
+        return model.encode(texts, **kwargs)
+    slices = budgeted_slices(lengths, batch_size, 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)

truememory/_rerank_budget.py

@@ -17,19 +17,24 @@ def cap_reranker_length(model: object) -> None:
         model.max_seq_length = 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
+
+
+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:

truememory/embedding_target.py

@@ -113,13 +113,18 @@ def admit() -> None:
     if target.model_id == "qwen3_256":
         model_kwargs = {"attn_implementation": "eager"} if sys.platform == "darwin" else None
         admit()
-        return SentenceTransformer(
+        model = SentenceTransformer(
             "Qwen/Qwen3-Embedding-0.6B", truncate_dim=256,
             config_kwargs={"use_cache": False},
             model_kwargs=model_kwargs, device=device,
         )
-    admit()
-    return 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,
+        )
+    if type(getattr(model, "max_seq_length", None)) is int:
+        from truememory._embed_budget import cap_embed_length
+        cap_embed_length(model)
+    return model

truememory/model_server.py

@@ -722,13 +722,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":
@@ -748,11 +752,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
@@ -1336,6 +1344,7 @@ def _handle_request_inner(self, request: dict) -> dict:
                     vectors = None
                     offset = 0
                     order = None
+                    steps = None
                     while offset < len(texts) or vectors is None:
                         retry = None
                         with deadline.locked(self._lock):
@@ -1352,18 +1361,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 budgeted_slices
+                                    lengths = text_token_lengths(model, texts)
+                                    slices = (None if lengths is None else
+                                              budgeted_slices(lengths, limit, 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):
@@ -1387,7 +1414,7 @@ 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)
+                            values = model.encode(retry, batch_size=len(retry) or limit, show_progress_bar=False)
                             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(

truememory/mps_utils.py

@@ -311,6 +311,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)
 
 
@@ -323,8 +326,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):
@@ -335,4 +342,7 @@ def encode_with_mps_fallback(model: object, texts: object, **kwargs: object) ->
         if hasattr(model, "to"):
             model.to("cpu")
         flush_mps_cache()
+        if bounded:
+            from truememory._embed_budget import budgeted_encode
+            return budgeted_encode(model, texts, **kwargs)
         return model.encode(texts, **kwargs)