fix(memory): cap local embedder CPU thread oversubscription (#198) (#1559)

## Description

The torch/sentence-transformers `LocalEmbedder` ran encodes on the
shared default executor with **no BLAS/OpenMP thread cap**. Under
concurrent load each `encode()` fans out to ~`os.cpu_count()`
BLAS/OpenMP threads, so N in-flight encodes spawn ~`N × cpu_count` OS
threads — oversubscribing the CPU, slowing the `memory_context` stage
and (on smaller boxes) starving the asyncio event loop. The ONNX
embedder already bounds its threads
(`create_cpu_session_options(intra_op_num_threads=1,
inter_op_num_threads=1)`); this brings the torch path to parity.

Supersedes #691 by @oxura — closed only for the open-PR cap, with an
explicit invitation to resubmit; no technical objection was raised, and
its CI was fully green. Credit to @oxura for the original diagnosis and
fix. That PR capped threads by setting BLAS/OpenMP env vars at import
time plus `torch.set_num_threads`; this PR instead runs CPU encodes on a
dedicated, size-limited executor whose workers each pin their thread
pool — which additionally bounds in-flight encode concurrency (the
issue's Fix B/C) and keeps the cap contained to the embedder rather than
mutating process-global env at import.

Closes #198

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)
- [x] Performance improvement

## Changes Made

- CPU encodes now run on a **dedicated, size-limited executor** whose
worker `initializer` pins each worker's torch intra-op pool (and sets
BLAS/OpenMP env defaults). torch's OpenMP thread count is per-thread, so
a one-shot cap misses pooled executor workers — the per-worker
initializer caps every worker deterministically.
- Total embedding threads are bounded by `HEADROOM_EMBED_CONCURRENCY`
(default `min(4, os.cpu_count())`) × `HEADROOM_EMBED_NUM_THREADS`
(default `1`); invalid/non-positive values fall back safely (≥1).
- Mirrors the existing MPS dedicated-single-worker-executor pattern;
CUDA keeps the shared default executor (GPU compute is off-CPU).
`setdefault` never overrides an operator's explicit `OMP_NUM_THREADS`.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [x] Manual testing performed

### Test Output

```text
$ uv run pytest tests/test_memory/test_embedder_thread_cap.py tests/test_memory/test_embedder_mps_serialization.py -q
13 passed

$ uv run pytest tests/test_memory/ tests/test_cli_proxy_embedding_server.py -q
533 passed        # no regressions from the executor change

$ uv run ruff check .  &&  uv run ruff format --check .
All checks passed!   /   1016 files already formatted

$ uv run mypy headroom --ignore-missing-imports
Success: no issues found in 404 source files
```

New `tests/test_memory/test_embedder_thread_cap.py`: env resolution for
both knobs (default / positive / invalid / clamped), worker-init env
application + operator-override safety, and a behavioral test that loads
the real CPU embedder and asserts every executor worker is pinned to the
configured intra-op thread count. Updated
`test_embedder_mps_serialization.py` to the new CPU contract.

## Real Behavior Proof

- Environment: built this branch into a CPU-only Linux container,
removed `onnxruntime` so the proxy falls back to the torch
`LocalEmbedder`; a container has no MPS/CUDA, so it resolves to
`device=cpu` — the deployment where #198 occurs. Python 3.12, torch
2.12.1, `all-MiniLM-L6-v2`, container capped to 4 CPUs, 32 concurrent
clients.
- Exact command / steps: `headroom proxy --host 0.0.0.0 --memory`
in-container; a concurrent `/v1/messages` driver from the host (invalid
key — `memory_context` runs before the upstream call); measured the
`memory_context` stage from `/metrics` before vs after the cap.
- Observed result: the embedder stage this PR targets improved —
`memory_context` avg 73.5 ms → 58.7 ms and max 279 ms → 242 ms (uncapped
12×8 = 96 threads vs fix 4×1): ~20% faster and steadier inside the real
proxy. Isolated component benchmarks (heavy concurrent `embed_batch`;
`LocalBackend.search_memories`) show a larger effect — tail event-loop
stall ~16–24 ms → ~3 ms, and search throughput +57%. Unit/regression: 13
new tests + 533 memory-suite tests pass; `ruff` + `mypy` clean.
- Not tested: the issue's absolute multi-second `/livez` spike. On my
hardware/synthetic load, `/livez` stalls were dominated by the
upstream-connection path (invalid-key DNS/TLS), not the ~250 ms
`memory_context` stage, so I can't attribute the multi-second figure to
the embedder here — the original report was on an 8-core box with real
Claude Code transcripts that drove `memory_context` itself to several
seconds. Linux/CUDA hardware not exercised; no live LLM provider used;
ONNX path unchanged. This PR removes the documented thread
oversubscription and brings the torch path to ONNX parity; it does not
claim to single-handedly resolve the 4 s figure.

Measured `memory_context` stage timing (real containerized proxy, torch
CPU embedder, 4 CPUs, 32 concurrent clients):

| `memory_context` | avg | max |
|---|---|---|
| Before (uncapped, 12×8 = 96 threads) | 73.5 ms | 279 ms |
| After (fix, 4×1) | 58.7 ms | 242 ms |

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Additional Notes

Default-behavior change: CPU encodes use a dedicated bounded pool
instead of the shared default executor (`close()` tears it down). Both
knobs are opt-in overrides with safe defaults. No new dependencies.

Signed-off-by: Krishnachaitanyakc <krishnabkc15@gmail.com>

CHANGELOG.md

@@ -20,6 +20,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   `<headroom_proactive_expansion>` XML tags, giving downstream consumers
   (LLMs, loggers, attribution parsers) a machine-readable provenance
   boundary and preventing misattribution in multi-agent threads.
+- **memory/embedder:** cap CPU thread oversubscription in the local
+  torch/sentence-transformers embedder. Concurrent encodes previously each
+  fanned out to ~`os.cpu_count()` BLAS/OpenMP threads, so under load the memory
+  path starved the asyncio event loop and spiked `/livez` latency to several
+  seconds. CPU encodes now run on a dedicated, size-limited executor whose
+  workers each pin their thread pool, bounding total embedding threads to
+  `HEADROOM_EMBED_CONCURRENCY` × `HEADROOM_EMBED_NUM_THREADS` (defaults
+  `min(4, cpu)` × 1). The ONNX embedder already capped its threads; this brings
+  the torch path to parity
+  ([#198](https://github.com/headroomlabs-ai/headroom/issues/198)).
 
 ### Changed
 

headroom/memory/adapters/embedders.py

@@ -42,6 +42,102 @@ logging.getLogger("httpx").setLevel(logging.WARNING)
 logger = logging.getLogger(__name__)
 
 
+# =============================================================================
+# Local (torch / sentence-transformers) CPU thread cap — issue #198
+# =============================================================================
+# A long-lived proxy serves many requests concurrently. Each torch ``encode()``
+# fans out to BLAS (MKL / OpenBLAS / Accelerate) + OpenMP worker threads, which
+# default to roughly ``os.cpu_count()``. Under concurrency this oversubscribes
+# the CPU — N in-flight encodes x ~cpu_count threads each thrash the scheduler
+# and starve the asyncio event loop, so liveness probes (``/livez``) spike to
+# multiple seconds even though the loop itself is idle (issue #198).
+#
+# Capping intra-op parallelism makes a single encode modestly slower but lets
+# concurrent encodes scale linearly without thread-pool thrash — the standard
+# trade-off for serving torch models inside an async server. The ONNX embedder
+# already caps its threads (see ``onnx_runtime.create_cpu_session_options``);
+# this brings the torch path to parity.
+#
+# torch's OpenMP thread count is per-thread, and encodes run on executor worker
+# threads, so a one-shot cap would miss most workers. Instead, CPU encodes run
+# on a dedicated, size-limited executor whose ``initializer`` pins each worker's
+# thread pool once. Total embedding threads are then bounded by
+# ``workers (HEADROOM_EMBED_CONCURRENCY) x threads-per-encode
+# (HEADROOM_EMBED_NUM_THREADS)``. Applies to the CPU device only (GPU/MPS do
+# their compute off-CPU).
+_EMBED_THREADS_ENV = "HEADROOM_EMBED_NUM_THREADS"
+_DEFAULT_EMBED_THREADS = 1
+_EMBED_CONCURRENCY_ENV = "HEADROOM_EMBED_CONCURRENCY"
+_DEFAULT_EMBED_CONCURRENCY = 4
+_BLAS_THREAD_ENV_VARS = (
+    "OMP_NUM_THREADS",
+    "OPENBLAS_NUM_THREADS",
+    "MKL_NUM_THREADS",
+    "NUMEXPR_NUM_THREADS",
+    "VECLIB_MAXIMUM_THREADS",
+)
+
+
+def _resolve_positive_int_env(env_var: str, default: int) -> int:
+    """Read a positive integer from ``env_var``, falling back to ``default``.
+
+    A non-positive or unparseable value logs a warning and returns a safe value
+    (>= 1) rather than disabling the limit.
+    """
+    raw = os.environ.get(env_var)
+    if raw is None:
+        return default
+    try:
+        value = int(raw)
+    except (TypeError, ValueError):
+        logger.warning("Invalid %s=%r; falling back to %d.", env_var, raw, default)
+        return default
+    if value < 1:
+        logger.warning("%s=%d is below 1; using 1.", env_var, value)
+        return 1
+    return value
+
+
+def _resolve_embed_thread_cap() -> int:
+    """Resolve the per-encode CPU thread cap (``HEADROOM_EMBED_NUM_THREADS``)."""
+    return _resolve_positive_int_env(_EMBED_THREADS_ENV, _DEFAULT_EMBED_THREADS)
+
+
+def _resolve_embed_concurrency() -> int:
+    """Resolve the max concurrent CPU encodes (``HEADROOM_EMBED_CONCURRENCY``).
+
+    Defaults to ``min(4, os.cpu_count())`` so embedding cannot occupy every core
+    and starve the event loop, while still allowing useful parallelism.
+    """
+    cpu = os.cpu_count() or 1
+    raw = os.environ.get(_EMBED_CONCURRENCY_ENV)
+    if raw is None:
+        return max(1, min(_DEFAULT_EMBED_CONCURRENCY, cpu))
+    return _resolve_positive_int_env(
+        _EMBED_CONCURRENCY_ENV, max(1, min(_DEFAULT_EMBED_CONCURRENCY, cpu))
+    )
+
+
+def _init_cpu_embed_worker() -> None:
+    """Pin a CPU embed worker's thread pool (runs once per worker; issue #198).
+
+    Sets BLAS/OpenMP env defaults (``setdefault`` never overrides an operator's
+    explicit setting) and bounds torch's intra-op pool for this worker thread.
+    Best-effort: failures never block embedding.
+    """
+    n = _resolve_embed_thread_cap()
+    for var in _BLAS_THREAD_ENV_VARS:
+        os.environ.setdefault(var, str(n))
+    try:
+        import torch
+
+        torch.set_num_threads(n)
+    except ImportError:
+        pass
+    except Exception as exc:  # pragma: no cover - defensive, never block embedding
+        logger.debug("Could not cap torch intra-op thread pool: %s", exc)
+
+
 def _normalize_embedding(embedding: np.ndarray) -> np.ndarray:
     """Normalize embedding to unit vector for cosine similarity.
 
@@ -172,6 +268,19 @@ class LocalEmbedder:
         else:
             self._device = self._detect_device()
 
+        # CPU: run encodes on a dedicated, size-limited executor whose workers
+        # each pin their torch/BLAS/OpenMP thread pool (issue #198). Without this,
+        # N concurrent encodes on the shared default executor each fan out to
+        # ~os.cpu_count() BLAS threads and starve the asyncio event loop, spiking
+        # /livez latency. Total embed threads are bounded by workers x per-encode
+        # threads.
+        if self._device == "cpu" and self._executor is None:
+            self._executor = ThreadPoolExecutor(
+                max_workers=_resolve_embed_concurrency(),
+                thread_name_prefix="cpu-embed",
+                initializer=_init_cpu_embed_worker,
+            )
+
         # torch-MPS is not thread-safe: concurrent encode() calls from the default
         # multi-worker executor abort with "commit an already committed command
         # buffer" (verified). Funnel every encode through one worker thread when on

tests/test_memory/test_embedder_mps_serialization.py

@@ -20,13 +20,16 @@ from headroom.memory.adapters.embedders import LocalEmbedder  # noqa: E402
 _HAS_MPS = bool(getattr(torch.backends, "mps", None)) and torch.backends.mps.is_available()
 
 
-async def test_cpu_uses_shared_executor() -> None:
-    """On CPU the dedicated executor stays None (unchanged default-pool behavior)."""
+async def test_cpu_uses_dedicated_thread_capped_executor() -> None:
+    """On CPU a dedicated, size-limited executor is used so encodes run with a
+    bounded thread pool instead of oversubscribing BLAS/OMP threads (issue #198)."""
     emb = LocalEmbedder(device="cpu")
     await emb.embed("hello world")
     assert emb._device == "cpu"
-    assert emb._executor is None  # default shared executor, not serialized
+    assert emb._executor is not None  # dedicated capped pool, not the shared default
+    assert emb._executor._max_workers >= 1  # type: ignore[attr-defined]
     await emb.close()
+    assert emb._executor is None  # close() tears it down
 
 
 @pytest.mark.skipif(not _HAS_MPS, reason="requires Apple-Silicon MPS")

tests/test_memory/test_embedder_thread_cap.py

@@ -0,0 +1,115 @@
+"""Regression tests for the LocalEmbedder CPU thread cap (issue #198).
+
+Under concurrent load the torch/sentence-transformers embedder oversubscribes
+BLAS/OpenMP threads (≈ ``os.cpu_count()`` per ``encode()``), starving the
+asyncio event loop and spiking ``/livez`` latency. ``LocalEmbedder`` now runs
+CPU encodes on a dedicated, size-limited executor whose workers each pin their
+torch/BLAS/OpenMP thread pool, bounding total embedding threads to
+``HEADROOM_EMBED_CONCURRENCY x HEADROOM_EMBED_NUM_THREADS``. The ONNX embedder
+already caps its threads; this brings the torch path to parity.
+"""
+
+from __future__ import annotations
+
+import os
+
+import pytest
+
+from headroom.memory.adapters import embedders
+from headroom.memory.adapters.embedders import (
+    _BLAS_THREAD_ENV_VARS,
+    _init_cpu_embed_worker,
+    _resolve_embed_concurrency,
+    _resolve_embed_thread_cap,
+)
+
+# ---------------------------------------------------------------------------
+# Env resolution (no torch required)
+# ---------------------------------------------------------------------------
+
+
+def test_thread_cap_default_when_unset(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.delenv("HEADROOM_EMBED_NUM_THREADS", raising=False)
+    assert _resolve_embed_thread_cap() == 1
+
+
+def test_thread_cap_reads_positive_int(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("HEADROOM_EMBED_NUM_THREADS", "3")
+    assert _resolve_embed_thread_cap() == 3
+
+
+def test_thread_cap_invalid_falls_back(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("HEADROOM_EMBED_NUM_THREADS", "not-a-number")
+    assert _resolve_embed_thread_cap() == 1
+
+
+def test_thread_cap_non_positive_is_clamped(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("HEADROOM_EMBED_NUM_THREADS", "0")
+    assert _resolve_embed_thread_cap() == 1
+
+
+def test_concurrency_default_is_bounded(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.delenv("HEADROOM_EMBED_CONCURRENCY", raising=False)
+    value = _resolve_embed_concurrency()
+    assert 1 <= value <= 4
+    assert value <= (os.cpu_count() or 1)
+
+
+def test_concurrency_reads_positive_int(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("HEADROOM_EMBED_CONCURRENCY", "7")
+    assert _resolve_embed_concurrency() == 7
+
+
+# ---------------------------------------------------------------------------
+# Worker initializer env application (no torch required)
+# ---------------------------------------------------------------------------
+
+
+def test_worker_init_sets_blas_env_defaults(monkeypatch: pytest.MonkeyPatch) -> None:
+    for var in _BLAS_THREAD_ENV_VARS:
+        monkeypatch.delenv(var, raising=False)
+    monkeypatch.setenv("HEADROOM_EMBED_NUM_THREADS", "2")
+
+    _init_cpu_embed_worker()
+
+    for var in _BLAS_THREAD_ENV_VARS:
+        assert os.environ[var] == "2", var
+
+
+def test_worker_init_does_not_override_operator_env(monkeypatch: pytest.MonkeyPatch) -> None:
+    """An explicit operator setting must win over our default (setdefault)."""
+    monkeypatch.setenv("OMP_NUM_THREADS", "8")
+    monkeypatch.setenv("HEADROOM_EMBED_NUM_THREADS", "1")
+
+    _init_cpu_embed_worker()
+
+    assert os.environ["OMP_NUM_THREADS"] == "8"
+
+
+# ---------------------------------------------------------------------------
+# Behavioral: real CPU load path bounds every encode worker's thread pool
+# ---------------------------------------------------------------------------
+
+
+async def test_cpu_embed_workers_are_thread_capped(monkeypatch: pytest.MonkeyPatch) -> None:
+    """CPU encodes run on a dedicated, size-limited executor and every worker
+    pins its torch intra-op thread pool to the configured cap."""
+    torch = pytest.importorskip("torch")
+    pytest.importorskip("sentence_transformers")
+
+    monkeypatch.setenv("HEADROOM_EMBED_NUM_THREADS", "1")
+    monkeypatch.setenv("HEADROOM_EMBED_CONCURRENCY", "2")
+
+    emb = embedders.LocalEmbedder(device="cpu")
+    await emb.embed("hello world")
+
+    assert emb._device == "cpu"
+    assert emb._executor is not None
+    assert emb._executor._max_workers == 2  # type: ignore[attr-defined]
+
+    # Probe the actual encode workers: each was pinned to 1 intra-op thread.
+    futures = [emb._executor.submit(torch.get_num_threads) for _ in range(4)]
+    assert [f.result() for f in futures] == [1, 1, 1, 1]
+
+    await emb.close()
+    assert emb._executor is None  # close() tears the executor down