diff --git a/tests/test-embed-token-budget.py b/tests/test-embed-token-budget.py new file mode 100644 index 00000000..2e714fa6 --- /dev/null +++ b/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)] diff --git a/tests/test-model-daemon-admission-297.py b/tests/test-model-daemon-admission-297.py index abd47447..3fdbf6e4 100644 --- a/tests/test-model-daemon-admission-297.py +++ b/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__": diff --git a/truememory/_embed_budget.py b/truememory/_embed_budget.py new file mode 100644 index 00000000..dba6dcd1 --- /dev/null +++ b/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) diff --git a/truememory/_rerank_budget.py b/truememory/_rerank_budget.py index c2688921..b21b90fb 100644 --- a/truememory/_rerank_budget.py +++ b/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: diff --git a/truememory/embedding_target.py b/truememory/embedding_target.py index 5604881b..511068b1 100644 --- a/truememory/embedding_target.py +++ b/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 diff --git a/truememory/model_server.py b/truememory/model_server.py index bed0fbdb..e51f1d12 100644 --- a/truememory/model_server.py +++ b/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( diff --git a/truememory/mps_utils.py b/truememory/mps_utils.py index e8a9e30e..1ae19cb4 100644 --- a/truememory/mps_utils.py +++ b/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) diff --git a/truememory/vector_search.py b/truememory/vector_search.py index 9510b5f9..5590715e 100644 --- a/truememory/vector_search.py +++ b/truememory/vector_search.py @@ -374,6 +374,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: @@ -615,10 +618,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)