fix: cap ONNX intra_op threads via MEMPAL_MAX_THREADS

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

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 &

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
 

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]

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