2 commits

Author	SHA1	Message	Date
Krishna Chaitanya	b84afbfb83	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>	2026-07-01 17:12:02 -05:00
Logan Kang	c71592d421	feat(memory): add opt-in Apple-GPU (MPS) embedding runtime (#766) ... ## Description On Apple-Silicon Macs — especially fanless models like the MacBook Air (M5) — running the proxy with memory context injection can pin the CPU while embedding. The embedding work runs an uncapped session on the CPU, saturating multiple cores, which starves the proxy's asyncio loop and leads to request timeouts. This PR adds an **opt-in** runtime that offloads the memory embedder to the Apple GPU (MPS). Setting `HEADROOM_EMBEDDER_RUNTIME=pytorch_mps` routes embedding through the torch `sentence-transformers` backend on MPS instead of the default ONNX CPU embedder, moving the work off the CPU and keeping the proxy responsive. The default behavior is unchanged — the feature is strictly opt-in, env-var only, and falls through to the existing default embedder selection (with a warning) whenever MPS or the torch dependencies are unavailable. Fixes: N/A — no tracking issue (surfaced while running codex auto-review through the proxy on a fanless MacBook Air (M5)). ## Type of Change - [ ] Bug fix (non-breaking change that fixes an issue) - [x] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [x] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - **Runtime selection** (`headroom/proxy/memory_handler.py`): read `HEADROOM_EMBEDDER_RUNTIME`; when set to `pytorch_mps` **and** MPS is actually available, route the memory embedder to the torch `sentence-transformers` backend (Apple GPU). - If MPS is unavailable or torch/sentence-transformers is not installed, log a warning and fall through to the existing default embedder selection (ONNX when available, else the pre-existing local sentence-transformers fallback) — no crash. The default (env var unset) is unchanged. Env-var only - **MPS serialization** (`headroom/memory/adapters/embedders.py`): `LocalEmbedder` now funnels every `encode()` through a dedicated single-worker `ThreadPoolExecutor` when the resolved device is MPS. torch-MPS is not thread-safe, and the existing `run_in_executor(None, ...)` dispatch would otherwise let concurrent proxy requests call MPS from multiple threads. CPU/CUDA keep the shared default executor (behavior unchanged). `close()` also drops the cached model so re-use after close re-initializes cleanly. - **Packaging** (`pyproject.toml`): new `pytorch-mps` extra (`torch` + `sentence-transformers`), **platform-gated to macOS** (`; sys_platform == 'darwin'`) since MPS is Apple-Silicon-only. Deliberately left out of `[all]` (its deps already arrive via `[ml]`/`[memory]`). - **Tests** (`tests/test_memory/test_embedder_mps_serialization.py`): regression coverage for the serialized executor, concurrency safety (no SIGABRT), CPU-path default behavior, and close/re-use re-initialization. - **Docs**: `wiki/{configuration,memory,macos-deployment}.md`, `docs/content/docs/{configuration,installation,memory}.mdx`, `README.md`, `CHANGELOG.md`. ## 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 (CPU-offload + concurrency profiling on Apple Silicon) ## Test Output ``` $ pytest -v tests/test_memory/test_embedder_mps_serialization.py collected 4 items tests/test_memory/test_embedder_mps_serialization.py::test_cpu_uses_shared_executor PASSED [ 25%] tests/test_memory/test_embedder_mps_serialization.py::test_mps_creates_single_worker_executor PASSED [ 50%] tests/test_memory/test_embedder_mps_serialization.py::test_mps_concurrent_embeds_do_not_crash PASSED [ 75%] tests/test_memory/test_embedder_mps_serialization.py::test_mps_reembed_after_close_recreates_executor PASSED [100%] ============================== 4 passed in 6.92s =============================== $ ruff check headroom/memory/adapters/embedders.py headroom/proxy/memory_handler.py tests/test_memory/test_embedder_mps_serialization.py All checks passed! $ mypy headroom/memory/adapters/embedders.py headroom/proxy/memory_handler.py Success: no issues found in 2 source files $ pytest -q tests/test_memory/ tests/test_memory_handler_concurrent_init.py tests/test_memory_handler_native_ops.py 553 passed, 1 skipped in 13.51s ``` ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md if applicable ## Additional Notes **Why MPS (and not CoreML or a thread cap):** measured on an Apple-Silicon Mac, the default uncapped CPU embedding session saturates the cores; the same model on MPS runs at a fraction of the CPU (≈8x lower sustained CPU utilization in profiling) while producing **byte-identical embeddings** (cosine distance ≈ 0 across runtimes), so relevance/ranking is unchanged. A CoreML execution-provider path was evaluated and rejected: the default optimized ONNX model uses fused ops that fall back to CPU under CoreML (no offload), and a full-precision re-export was impractical (very low throughput + multi-GB memory). MPS via `sentence-transformers` was the only practical GPU offload. **Why serialization is mandatory:** torch-MPS is not thread-safe — concurrent encode calls from a multi-worker executor abort with `-[IOGPUMetalCommandBuffer validate]: failed assertion 'commit an already committed command buffer'` (reproduced deterministically; a single-worker executor resolves it). Under concurrent load the serialized single-GPU-stream throughput meets or exceeds the parallel CPU path while using a fraction of the cores. **Scope / boundary:** this targets the Python **memory** embedder, which is live on the proxy request path (memory context injection). The Rust-backed SmartCrusher compression path is unaffected and remains non-configurable from Python by design. **Safety:** default behavior is unchanged (ONNX, no torch). The feature is opt-in, env-var only, macOS-gated at the packaging layer, and degrades gracefully (warn + the existing default embedder selection) when MPS or the dependencies are unavailable.	2026-06-11 12:59:20 -05:00