From 955c11a4274847b0f79b1378ecd7dc07ac41eb2e Mon Sep 17 00:00:00 2001 From: sha2fiddy <103975074+sha2fiddy@users.noreply.github.com> Date: Wed, 29 Apr 2026 19:01:34 -0400 Subject: [PATCH 1/2] fix: cap ONNX intra_op threads via MEMPAL_MAX_THREADS MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Without this, ORT spawns ~physical-core-count workers in its intra_op pool and a background mine pegs 400-500% CPU. OMP_NUM_THREADS does not control the ORT pool — ORT has its own. Cap is applied in mempalace.embedding by overriding the model cached_property on the EF subclass so the InferenceSession is built with explicit SessionOptions (intra_op_num_threads=N, inter_op=1). Default cap is 2; "0"/"off"/"default"/"none" disables it. HNSW thread pinning is already handled by _pin_hnsw_threads on develop, so this PR only addresses the ONNX side. The auto_save and precompact hooks export MEMPAL_MAX_THREADS=2 so background mines stay throttled. Closes #1068. --- hooks/mempal_precompact_hook.sh | 6 ++ hooks/mempal_save_hook.sh | 6 ++ mempalace/embedding.py | 91 +++++++++++++++++++++--- tests/test_embedding.py | 118 +++++++++++++++++++++++++++++++- 4 files changed, 211 insertions(+), 10 deletions(-) diff --git a/hooks/mempal_precompact_hook.sh b/hooks/mempal_precompact_hook.sh index b9585876ca..fda58650e5 100755 --- a/hooks/mempal_precompact_hook.sh +++ b/hooks/mempal_precompact_hook.sh @@ -189,6 +189,12 @@ echo "[$(date '+%H:%M:%S')] PRE-COMPACT triggered for session $SESSION_ID" >> "$ # independent targets — both run if both are set: # 1. TRANSCRIPT_PATH (from Claude Code) → parent dir, --mode convos # 2. MEMPAL_DIR → --mode projects +# +# Cap ONNX intra_op threads so the synchronous mine doesn't pin every +# core on multi-core hosts. See ``_read_thread_cap()`` in +# ``mempalace/embedding.py``. ``TOKENIZERS_PARALLELISM=false`` silences +# the Hugging Face fork warning. +export MEMPAL_MAX_THREADS=2 TOKENIZERS_PARALLELISM=false if is_valid_transcript_path "$TRANSCRIPT_PATH" && [ -f "$TRANSCRIPT_PATH" ]; then mempalace mine "$(dirname "$TRANSCRIPT_PATH")" --mode convos \ >> "$STATE_DIR/hook.log" 2>&1 diff --git a/hooks/mempal_save_hook.sh b/hooks/mempal_save_hook.sh index 58c61e9980..62f27d384d 100755 --- a/hooks/mempal_save_hook.sh +++ b/hooks/mempal_save_hook.sh @@ -291,6 +291,12 @@ if [ "$SINCE_LAST" -ge "$SAVE_INTERVAL" ] && [ "$EXCHANGE_COUNT" -gt 0 ]; then # (code, notes, docs) # MEMPAL_DIR is *additive*, not an override: a user with MEMPAL_DIR # pointed at their project still gets the active conversation mined. + # + # Cap ONNX intra_op threads so the background mine doesn't pin every + # core on multi-core hosts. See ``_read_thread_cap()`` in + # ``mempalace/embedding.py``. ``TOKENIZERS_PARALLELISM=false`` silences + # the Hugging Face fork warning. + export MEMPAL_MAX_THREADS=2 TOKENIZERS_PARALLELISM=false if is_valid_transcript_path "$TRANSCRIPT_PATH" && [ -f "$TRANSCRIPT_PATH" ]; then mempalace mine "$(dirname "$TRANSCRIPT_PATH")" --mode convos \ >> "$STATE_DIR/hook.log" 2>&1 & diff --git a/mempalace/embedding.py b/mempalace/embedding.py index b22b1f52b4..b28efb3bf5 100644 --- a/mempalace/embedding.py +++ b/mempalace/embedding.py @@ -27,15 +27,50 @@ Requesting an unavailable accelerator emits a warning and falls back to CPU rather than hard-failing — mining must still work on a laptop without CUDA. + +ONNX Runtime's intra-op thread pool is capped via ``MEMPAL_MAX_THREADS`` +(default 2; ``0``/``off``/``default``/``none`` disables the cap). Without +this, ORT spawns ≈physical-core-count workers and a background mine can +peg 400-500%% CPU. ``OMP_NUM_THREADS`` does not control the ORT pool. """ from __future__ import annotations import logging +import os +from functools import cached_property from typing import Optional logger = logging.getLogger(__name__) +_DEFAULT_THREAD_CAP = 2 + + +def _read_thread_cap() -> int: + """Return the ONNX intra-op thread cap from ``MEMPAL_MAX_THREADS``. + + * Unset → ``_DEFAULT_THREAD_CAP`` (2). + * ``"0"`` / ``"off"`` / ``"default"`` / ``"none"`` / empty → 0 (no cap). + * Positive int → that int. + * Anything else → 0 with a warning (fail-open: never break mining on a typo). + """ + raw = os.environ.get("MEMPAL_MAX_THREADS") + if raw is None: + return _DEFAULT_THREAD_CAP + raw = raw.strip().lower() + if raw in ("", "0", "off", "default", "none"): + return 0 + try: + n = int(raw) + except ValueError: + logger.warning( + "MEMPAL_MAX_THREADS=%r is not an integer; leaving ORT defaults", + raw, + ) + return 0 + return n if n > 0 else 0 + + _PROVIDER_MAP = { "cpu": ["CPUExecutionProvider"], "cuda": ["CUDAExecutionProvider", "CPUExecutionProvider"], @@ -107,7 +142,7 @@ def _resolve_providers(device: str) -> tuple[list, str]: return (requested, device) -def _build_ef_class(): +def _build_ef_class(thread_cap: int = 0): """Subclass ``ONNXMiniLM_L6_V2`` with name ``"default"``. Why the rename: ChromaDB 1.5 persists the EF identity on the collection @@ -116,6 +151,12 @@ def _build_ef_class(): ``name()`` tag differs — so spoofing the name lets one EF class serve palaces created with ``DefaultEmbeddingFunction`` *and* palaces we create ourselves, with the same GPU-capable ``preferred_providers``. + + When ``thread_cap > 0`` the ``model`` cached property is overridden so + the ORT ``InferenceSession`` is built with explicit ``SessionOptions`` + capping ``intra_op_num_threads`` and ``inter_op_num_threads=1``. This + keeps the background mine from pinning every core on multi-core hosts; + ``OMP_NUM_THREADS`` does not control ORT's intra-op pool. """ from chromadb.utils.embedding_functions import ONNXMiniLM_L6_V2 @@ -124,7 +165,27 @@ class _MempalaceONNX(ONNXMiniLM_L6_V2): def name() -> str: return "default" - return _MempalaceONNX + if thread_cap <= 0: + return _MempalaceONNX + + class _CappedMempalaceONNX(_MempalaceONNX): + _mempal_thread_cap = thread_cap + + @cached_property + def model(self): # type: ignore[override] + so = self.ort.SessionOptions() + so.log_severity_level = 3 + so.intra_op_num_threads = self._mempal_thread_cap + so.inter_op_num_threads = 1 + if not self._preferred_providers: + self._preferred_providers = ["CPUExecutionProvider"] + return self.ort.InferenceSession( + os.path.join(self.DOWNLOAD_PATH, self.EXTRACTED_FOLDER_NAME, "model.onnx"), + providers=self._preferred_providers, + sess_options=so, + ) + + return _CappedMempalaceONNX # Embeddinggemma-300m ONNX (q8) — 100+ languages, MRL-truncated to 384 dims so @@ -158,10 +219,11 @@ def name() -> str: # when switching models. Keep it stable. return "embeddinggemma_300m" - def __init__(self, preferred_providers=None): + def __init__(self, preferred_providers=None, thread_cap: int = 0): self._providers = ( list(preferred_providers) if preferred_providers else ["CPUExecutionProvider"] ) + self._thread_cap = thread_cap if thread_cap and thread_cap > 0 else 0 self._session = None self._tokenizer = None self._np = None @@ -193,7 +255,15 @@ def _lazy_load(self) -> None: ) tok_path = hf_hub_download(_EMBEDDINGGEMMA_REPO, filename="tokenizer.json") - self._session = ort.InferenceSession(model_path, providers=self._providers) + sess_options = None + if self._thread_cap: + sess_options = ort.SessionOptions() + sess_options.log_severity_level = 3 + sess_options.intra_op_num_threads = self._thread_cap + sess_options.inter_op_num_threads = 1 + self._session = ort.InferenceSession( + model_path, providers=self._providers, sess_options=sess_options + ) out_names = [o.name for o in self._session.get_outputs()] # Model card: sentence_embedding is the pooled output (last_hidden_state # is the per-token output we don't want). @@ -230,7 +300,8 @@ def get_embedding_function(device: Optional[str] = None, model: Optional[str] = ``device=None`` reads :attr:`MempalaceConfig.embedding_device`; ``model=None`` reads :attr:`MempalaceConfig.embedding_model`. The returned function is shared across calls with the same resolved - provider list + model so we only pay model-load cost once per process. + provider list, model, and ``MEMPAL_MAX_THREADS`` cap so we only pay + model-load cost once per process. """ if device is None or model is None: from .config import MempalaceConfig @@ -242,24 +313,26 @@ def get_embedding_function(device: Optional[str] = None, model: Optional[str] = model = cfg.embedding_model providers, effective = _resolve_providers(device) - cache_key = (model, tuple(providers)) + thread_cap = _read_thread_cap() + cache_key = (model, tuple(providers), thread_cap) cached = _EF_CACHE.get(cache_key) if cached is not None: return cached if model == "embeddinggemma": - ef = EmbeddinggemmaONNX(preferred_providers=providers) + ef = EmbeddinggemmaONNX(preferred_providers=providers, thread_cap=thread_cap) else: # Default: minilm (or anything we don't recognize — back-compat win). - ef_cls = _build_ef_class() + ef_cls = _build_ef_class(thread_cap) ef = ef_cls(preferred_providers=providers) _EF_CACHE[cache_key] = ef logger.info( - "Embedding function initialized (model=%s device=%s providers=%s)", + "Embedding function initialized (model=%s device=%s providers=%s thread_cap=%d)", model, effective, providers, + thread_cap, ) return ef diff --git a/tests/test_embedding.py b/tests/test_embedding.py index d05075d69c..688aaad8ac 100644 --- a/tests/test_embedding.py +++ b/tests/test_embedding.py @@ -76,10 +76,11 @@ class DummyEF: def __init__(self, preferred_providers): self.preferred_providers = preferred_providers - monkeypatch.setattr(embedding, "_build_ef_class", lambda: DummyEF) + monkeypatch.setattr(embedding, "_build_ef_class", lambda thread_cap=0: DummyEF) monkeypatch.setattr( embedding, "_resolve_providers", lambda device: (["CPUExecutionProvider"], "cpu") ) + monkeypatch.delenv("MEMPAL_MAX_THREADS", raising=False) first = embedding.get_embedding_function("cpu") second = embedding.get_embedding_function("auto") @@ -96,3 +97,118 @@ def test_describe_device_uses_resolved_effective_device(monkeypatch): ) assert embedding.describe_device("auto") == "cuda" + + +# ── MEMPAL_MAX_THREADS: ONNX intra-op cap ──────────────────────────────── + + +def test_read_thread_cap_default(monkeypatch): + """Unset MEMPAL_MAX_THREADS → default cap of 2.""" + monkeypatch.delenv("MEMPAL_MAX_THREADS", raising=False) + assert embedding._read_thread_cap() == 2 + + +@pytest.mark.parametrize("value", ["0", "off", "default", "none", "", " "]) +def test_read_thread_cap_disabled_values(monkeypatch, value): + """Sentinel strings disable the cap (return 0 → ORT defaults).""" + monkeypatch.setenv("MEMPAL_MAX_THREADS", value) + assert embedding._read_thread_cap() == 0 + + +@pytest.mark.parametrize("value,expected", [("1", 1), ("2", 2), ("4", 4), ("16", 16)]) +def test_read_thread_cap_positive_int(monkeypatch, value, expected): + monkeypatch.setenv("MEMPAL_MAX_THREADS", value) + assert embedding._read_thread_cap() == expected + + +def test_read_thread_cap_bad_value_is_disabled(monkeypatch, caplog): + """Non-integer input falls back to 0 (no cap) and logs a warning.""" + monkeypatch.setenv("MEMPAL_MAX_THREADS", "banana") + with caplog.at_level("WARNING"): + assert embedding._read_thread_cap() == 0 + assert any("MEMPAL_MAX_THREADS" in r.message for r in caplog.records) + + +def test_read_thread_cap_negative_is_disabled(monkeypatch): + monkeypatch.setenv("MEMPAL_MAX_THREADS", "-1") + assert embedding._read_thread_cap() == 0 + + +def test_build_ef_class_no_cap_returns_uncapped_subclass(): + """thread_cap <= 0 → no model override; chromadb defaults apply.""" + cls = embedding._build_ef_class(0) + assert cls.name() == "default" + assert not hasattr(cls, "_mempal_thread_cap") + + +def test_build_ef_class_with_cap_overrides_model(): + """thread_cap > 0 → subclass with capped model and _mempal_thread_cap.""" + cls = embedding._build_ef_class(3) + assert cls.name() == "default" + assert cls._mempal_thread_cap == 3 + + +def test_build_ef_class_capped_model_uses_session_options(): + """The capped ``model`` cached_property builds an ORT session with our cap. + + Monkey-patches ``self.ort`` on the embedder instance so we don't download + the real 90 MB ONNX model; instead we capture the SessionOptions passed to + ``InferenceSession`` and assert on them. + """ + cls = embedding._build_ef_class(2) + instance = cls.__new__(cls) + instance._preferred_providers = None + + captured: dict = {} + + class _FakeSessionOptions: + def __init__(self): + self.intra_op_num_threads = None + self.inter_op_num_threads = None + self.log_severity_level = None + + class _FakeORT: + SessionOptions = _FakeSessionOptions + + @staticmethod + def InferenceSession(*args, **kwargs): + captured["args"] = args + captured["kwargs"] = kwargs + return "fake-session" + + instance.ort = _FakeORT + + session = instance.model + assert session == "fake-session" + so = captured["kwargs"]["sess_options"] + assert so.intra_op_num_threads == 2 + assert so.inter_op_num_threads == 1 + assert so.log_severity_level == 3 + assert captured["kwargs"]["providers"] == ["CPUExecutionProvider"] + + +def test_get_embedding_function_cache_keyed_by_thread_cap(monkeypatch): + """Different MEMPAL_MAX_THREADS values must not share a cached EF.""" + + class DummyEF: + def __init__(self, preferred_providers): + self.preferred_providers = preferred_providers + + built: list = [] + + def fake_build(thread_cap=0): + built.append(thread_cap) + return DummyEF + + monkeypatch.setattr(embedding, "_build_ef_class", fake_build) + monkeypatch.setattr( + embedding, "_resolve_providers", lambda device: (["CPUExecutionProvider"], "cpu") + ) + + monkeypatch.setenv("MEMPAL_MAX_THREADS", "2") + first = embedding.get_embedding_function("cpu") + monkeypatch.setenv("MEMPAL_MAX_THREADS", "4") + second = embedding.get_embedding_function("cpu") + + assert first is not second + assert built == [2, 4] From e78074903e0492daa9f76bbc281f8c0becdf0c23 Mon Sep 17 00:00:00 2001 From: sha2fiddy <103975074+sha2fiddy@users.noreply.github.com> Date: Wed, 27 May 2026 13:55:40 -0400 Subject: [PATCH 2/2] test(embeddinggemma): adapt cache-key test to thread_cap signature PR #1071 changes _build_ef_class to accept thread_cap=N. The test in tests/test_embeddinggemma.py (added on develop) stubs _build_ef_class with a zero-arg lambda, which now fails after the rebase. Update the lambda and clear MEMPAL_MAX_THREADS so the cap branch is exercised deterministically. --- tests/test_embeddinggemma.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_embeddinggemma.py b/tests/test_embeddinggemma.py index 3100fce7a2..7e66c5cd61 100644 --- a/tests/test_embeddinggemma.py +++ b/tests/test_embeddinggemma.py @@ -183,10 +183,11 @@ class DummyMiniLM: def __init__(self, preferred_providers=None): self.kind = "minilm" - monkeypatch.setattr(embedding, "_build_ef_class", lambda: DummyMiniLM) + monkeypatch.setattr(embedding, "_build_ef_class", lambda thread_cap=0: DummyMiniLM) monkeypatch.setattr( embedding, "_resolve_providers", lambda device: (["CPUExecutionProvider"], "cpu") ) + monkeypatch.delenv("MEMPAL_MAX_THREADS", raising=False) ml = embedding.get_embedding_function(device="cpu", model="minilm") eg = embedding.get_embedding_function(device="cpu", model="embeddinggemma")