fix: Bound reranker input length and attention memory per batch

tests/test-reranker-token-budget.py

@@ -0,0 +1,145 @@
+"""Reranker input cap and token-budgeted batches bound attention scratch."""
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from truememory import _rerank_budget as rerank_budget
+from truememory._rerank_budget import (
+    RERANK_MAX_TOKENS, RERANK_TOKEN_BUDGET, budget_batch_size, budgeted_predict,
+    budgeted_slices, cap_reranker_length, pair_token_lengths, rerank_token_budget,
+)
+
+
+class FakeCrossEncoder:
+    """Word count stands in for tokens; scores depend only on each pair."""
+
+    def __init__(self, max_seq_length: int | None = 1024) -> None:
+        self.max_seq_length = max_seq_length
+        self.calls: list[tuple[int, int]] = []
+
+    def tokenizer(self, first, second, truncation, max_length):
+        assert truncation == "longest_first"
+        return {"input_ids": [[0] * min(len(a.split()) + len(b.split()) + 3, max_length)
+                              for a, b in zip(first, second)]}
+
+    @staticmethod
+    def _input_length(pair) -> int:
+        return len(pair[0]) + len(pair[1])
+
+    def predict(self, pairs, batch_size, show_progress_bar=False):
+        self.calls.append((len(pairs), batch_size))
+        return np.array([len(b.split()) / 1000 + len(a) / 1e6 for a, b in pairs], dtype=np.float32)
+
+
+def _pairs(*word_counts: int) -> list[list[str]]:
+    return [["synthetic query", " ".join(["word"] * count)] for count in word_counts]
+
+
+def test_cap_lowers_long_limits_only():
+    for current, expected in ((8192, RERANK_MAX_TOKENS), (None, None), (512, 512)):
+        model = FakeCrossEncoder(current)
+        cap_reranker_length(model)
+        assert model.max_seq_length == expected
+    cap_reranker_length(object())
+
+
+def test_budget_shrinks_under_memory_pressure():
+    assert rerank_token_budget(16 * 1024 ** 3) == RERANK_TOKEN_BUDGET
+    assert rerank_token_budget(4 * 1024 ** 3) == RERANK_TOKEN_BUDGET // 2
+    assert rerank_token_budget(1024 ** 3) == RERANK_TOKEN_BUDGET // 4
+
+
+def test_default_batches_kept_for_pairs_up_to_512_tokens():
+    assert budget_batch_size(512, 64, RERANK_TOKEN_BUDGET) == 64
+    assert budgeted_slices([512, 40, 300] * 30, 64, RERANK_TOKEN_BUDGET) is None
+
+
+def test_long_pairs_go_longest_first_in_budgeted_slices():
+    lengths = [50] * 80 + [1024] * 20
+    slices = budgeted_slices(lengths, 64, RERANK_TOKEN_BUDGET)
+    assert sorted(index for group in slices for index in group) == list(range(100))
+    for group in slices:
+        assert len(group) * max(lengths[index] for index in group) ** 2 <= RERANK_TOKEN_BUDGET or len(group) == 1
+        assert len(group) <= 64
+    assert [len(group) for group in slices] == [16, 16, 64, 4]
+
+
+def test_short_search_is_one_unchanged_predict_call():
+    model = FakeCrossEncoder()
+    pairs = _pairs(*range(10, 110))
+    scores = budgeted_predict(model, pairs, 64)
+    assert model.calls == [(100, 64)]
+    assert scores.shape == (100,)
+
+
+def test_long_search_scores_match_and_respect_budget(monkeypatch):
+    monkeypatch.setattr(rerank_budget, "rerank_token_budget", lambda: RERANK_TOKEN_BUDGET)
+    model = FakeCrossEncoder()
+    pairs = _pairs(*([20] * 70 + [2000] * 30))
+    scores = budgeted_predict(model, pairs, 64)
+    expected = FakeCrossEncoder().predict(pairs, batch_size=len(pairs))
+    assert np.array_equal(scores, expected)
+    assert sum(size for size, _ in model.calls) == 100
+    lengths = pair_token_lengths(model, pairs)
+    assert max(lengths) == RERANK_MAX_TOKENS
+    assert all(batch == size and size <= 64 for size, batch in model.calls)
+    assert model.calls[0][0] == RERANK_TOKEN_BUDGET // RERANK_MAX_TOKENS ** 2
+
+
+def test_unmeasurable_model_keeps_single_call():
+    class Proxy:
+        calls = []
+
+        def predict(self, pairs, batch_size, show_progress_bar=False):
+            self.calls.append(len(pairs))
+            return np.zeros(len(pairs), dtype=np.float32)
+
+    proxy = Proxy()
+    budgeted_predict(proxy, _pairs(5000, 5000), 64)
+    assert proxy.calls == [2]
+
+
+def test_model_server_rerank_uses_budgeted_slices(monkeypatch):
+    from truememory.model_server import ModelServer
+
+    fake = FakeCrossEncoder()
+    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] * 50 + [1500] * 10))
+    response = ModelServer().handle_request({"op": "rerank", "pairs": pairs, "model_name": "m"})
+    assert response["ok"] is True
+    expected = FakeCrossEncoder().predict(pairs, batch_size=len(pairs))
+    assert np.allclose(response["scores"], expected, rtol=0, atol=0)
+    assert max(size for size, _ in fake.calls) <= 64
+    assert fake.calls[0] == (RERANK_TOKEN_BUDGET // RERANK_MAX_TOKENS ** 2, RERANK_TOKEN_BUDGET // RERANK_MAX_TOKENS ** 2)
+
+
+def test_model_server_short_rerank_unchanged(monkeypatch):
+    from truememory.model_server import ModelServer
+
+    fake = FakeCrossEncoder()
+    monkeypatch.setattr(ModelServer, "_get_reranker", lambda self, model_name=None, deadline=None: fake)
+    response = ModelServer().handle_request({"op": "rerank", "pairs": _pairs(*range(5, 45)), "model_name": "m"})
+    assert response["ok"] is True
+    assert fake.calls == [(40, 64)]
+
+
+@pytest.mark.parametrize("length", [1, 512, 513, 1024])
+def test_batch_size_never_zero_or_above_limit(length):
+    assert 1 <= budget_batch_size(length, 64, RERANK_TOKEN_BUDGET) <= 64
+
+
+def test_model_server_budget_path_needs_no_input_length_helper(monkeypatch):
+    from truememory.model_server import ModelServer
+
+    class WithoutInputLength(FakeCrossEncoder):
+        _input_length = None
+
+    fake = WithoutInputLength()
+    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] * 20 + [1500] * 3))
+    response = ModelServer().handle_request({"op": "rerank", "pairs": pairs, "model_name": "m"})
+    assert response["ok"] is True
+    assert np.array_equal(np.asarray(response["scores"]), FakeCrossEncoder().predict(pairs, batch_size=len(pairs)))

truememory/_rerank_budget.py

@@ -0,0 +1,85 @@
+"""Reranker input cap and token-budgeted batches that bound attention scratch."""
+from __future__ import annotations
+
+
+# Longer pairs are truncated longest side first. Stored memories are far
+# shorter; the cap bounds oversized queries.
+RERANK_MAX_TOKENS = 1024
+# Attention scratch grows with batch x length^2. This budget keeps the default
+# 64-pair batch for every pair up to 512 tokens.
+RERANK_TOKEN_BUDGET = 64 * 512 * 512
+
+
+def cap_reranker_length(model: object) -> None:
+    """Bound a cross-encoder's input length without raising a smaller limit."""
+    current = getattr(model, "max_seq_length", None)
+    if type(current) is int and current > RERANK_MAX_TOKENS:
+        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."""
+    if available is None:
+        try:
+            import psutil
+            available = psutil.virtual_memory().available
+        except (ImportError, OSError):
+            return RERANK_TOKEN_BUDGET
+    if available < 2 * 1024 ** 3:
+        return RERANK_TOKEN_BUDGET // 4
+    if available < 6 * 1024 ** 3:
+        return RERANK_TOKEN_BUDGET // 2
+    return RERANK_TOKEN_BUDGET
+
+
+def pair_token_lengths(model: object, pairs: list) -> list[int] | None:
+    """Truncated token length of each pair, or None when it cannot be measured."""
+    tokenizer = getattr(model, "tokenizer", None)
+    limit = getattr(model, "max_seq_length", None)
+    if (tokenizer is None or type(limit) is not int or limit <= 0
+            or not all(isinstance(pair, (list, tuple)) and len(pair) == 2
+                       and all(isinstance(text, str) for text in pair) for pair in pairs)):
+        return None
+    if not pairs:
+        return []
+    encoded = tokenizer([pair[0] for pair in pairs], [pair[1] for pair in pairs],
+                        truncation="longest_first", max_length=limit)
+    return [len(ids) for ids in encoded["input_ids"]]
+
+
+def budget_batch_size(length: int, limit: int, budget: int) -> int:
+    return max(1, min(limit, budget // max(1, length) ** 2))
+
+
+def budgeted_slices(lengths: list[int], limit: int, budget: int) -> list[list[int]] | None:
+    """Longest-first index groups sized to the budget.
+
+    None means every default batch already fits, so callers keep their
+    original single call and its exact results.
+    """
+    if budget_batch_size(max(lengths, default=0), limit, budget) >= limit:
+        return None
+    order = sorted(range(len(lengths)), key=lambda index: -lengths[index])
+    slices, start = [], 0
+    while start < len(order):
+        size = budget_batch_size(lengths[order[start]], limit, budget)
+        slices.append(order[start:start + size])
+        start += size
+    return slices
+
+
+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())
+    if slices is None:
+        return model.predict(pairs, batch_size=batch_size, show_progress_bar=False)
+    import numpy as np
+    results: list = [None] * len(pairs)
+    for indices in slices:
+        values = model.predict([pairs[index] for index in indices], batch_size=len(indices),
+                               show_progress_bar=False)
+        for position, index in enumerate(indices):
+            results[index] = values[position]
+    return np.stack(results)

truememory/model_server.py

@@ -1078,6 +1078,9 @@ def _get_reranker(
                 self._memory_stage("reranker construction", deadline,
                                    model_identity=(name, 0, "reranker"), device=device)
             model = CrossEncoder(name, device=device)
+            if type(getattr(model, "max_seq_length", None)) is int:
+                from truememory._rerank_budget import cap_reranker_length
+                cap_reranker_length(model)
             self._check_process_memory("reranker construction completion", deadline)
             self._reranker = model
             self._reranker_name = name
@@ -1428,6 +1431,7 @@ def _handle_request_inner(self, request: dict) -> dict:
                     scores = None
                     offset = 0
                     order = None
+                    steps = None
                     recovery_name = model_name
                     while offset < len(pairs) or scores is None:
                         retry = None
@@ -1441,18 +1445,39 @@ def _handle_request_inner(self, request: dict) -> dict:
                                 order = self._rerank_global_order(reranker, pairs, limit, deadline)
                                 if order is not None:
                                     recovery_name = self._reranker_name
+                                slices = None
+                                if (getattr(reranker, "tokenizer", None) is not None
+                                        and type(getattr(reranker, "max_seq_length", None)) is int):
+                                    from truememory._rerank_budget import (
+                                        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()))
+                                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)
+                                    steps, start = {}, 0
+                                    for group in slices:
+                                        steps[start] = len(group)
+                                        start += len(group)
+                                    recovery_name = self._reranker_name
                             while offset < len(pairs) or scores is None:
                                 deadline.check()
-                                indices = None if order is None else self._rerank_slice_indices(
-                                    reranker, pairs, order, offset, limit, deadline,
-                                )
-                                batch = (pairs[offset:offset + limit] if indices is None else
+                                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._rerank_slice_indices(
+                                        reranker, pairs, order, offset, step, deadline,
+                                    )
+                                batch = (pairs[offset:offset + step] if indices is None else
                                          [pairs[index] for index in indices])
                                 self._check_process_memory("reranker slice", deadline)
                                 deadline.check()
                                 recover = False
                                 try:
-                                    values = reranker.predict(batch, batch_size=limit, show_progress_bar=False)
+                                    values = reranker.predict(batch, batch_size=step, show_progress_bar=False)
                                 except RuntimeError as exc:
                                     from truememory.mps_utils import is_mps_oom, flush_mps_cache
                                     if not is_mps_oom(exc):
@@ -1485,7 +1510,7 @@ 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=limit, show_progress_bar=False)
+                            values = reranker.predict(retry, batch_size=len(retry) or limit, show_progress_bar=False)
                             if indices is None:
                                 scores = _store_batch_result(scores, values, offset, len(retry), len(pairs),
                                                              "scores", **result_options)

truememory/reranker.py

@@ -339,6 +339,9 @@ def get_reranker(
                     from sentence_transformers import CrossEncoder
                     _check(control)
                     _model = CrossEncoder(name, device=device)
+                    if type(getattr(_model, "max_seq_length", None)) is int:
+                        from truememory._rerank_budget import cap_reranker_length
+                        cap_reranker_length(_model)
                     if budget_owner is not None:
                         bind_model_owner(_model, budget_owner)
             except BaseException:
@@ -447,7 +450,11 @@ def rerank(
             scope = nullcontext()
         with scope as memory:
             with _own_model(model):
-                scores = model.predict(pairs, batch_size=batch_size, show_progress_bar=False)
+                if getattr(model, "tokenizer", None) is not None and type(getattr(model, "max_seq_length", None)) is int:
+                    from truememory._rerank_budget import budgeted_predict
+                    scores = budgeted_predict(model, pairs, batch_size)
+                else:
+                    scores = model.predict(pairs, batch_size=batch_size, show_progress_bar=False)
                 if memory is not None:
                     memory.retain_buffer(scores)