mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
b84afbfb83
## 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>
72 lines
2.9 KiB
Python
72 lines
2.9 KiB
Python
"""Regression tests for the LocalEmbedder MPS serialization fix.
|
|
|
|
torch-MPS is not thread-safe: concurrent encode() calls from the default
|
|
multi-worker executor abort with "commit an already committed command buffer".
|
|
LocalEmbedder funnels every encode through a dedicated single-worker executor
|
|
when (and only when) the resolved device is MPS. CPU/CUDA keep the shared pool.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
|
|
import pytest
|
|
|
|
torch = pytest.importorskip("torch")
|
|
pytest.importorskip("sentence_transformers")
|
|
|
|
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_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 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")
|
|
async def test_mps_creates_single_worker_executor() -> None:
|
|
"""On MPS a dedicated max_workers=1 executor is created for serialization."""
|
|
emb = LocalEmbedder(device="mps")
|
|
await emb.embed("warmup")
|
|
assert emb._device == "mps"
|
|
assert emb._executor is not None
|
|
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")
|
|
async def test_mps_concurrent_embeds_do_not_crash() -> None:
|
|
"""Concurrent embeds on MPS must not SIGABRT — the serialization guarantees it."""
|
|
emb = LocalEmbedder(device="mps")
|
|
await emb.embed("warmup")
|
|
batches = [emb.embed_batch([f"text {i} " * 20] * 8) for i in range(16)]
|
|
results = await asyncio.gather(*batches)
|
|
assert len(results) == 16
|
|
assert all(len(r[0]) == emb.dimension for r in results)
|
|
await emb.close()
|
|
|
|
|
|
@pytest.mark.skipif(not _HAS_MPS, reason="requires Apple-Silicon MPS")
|
|
async def test_mps_reembed_after_close_recreates_executor() -> None:
|
|
"""close() drops the cached model so a later embed() re-initializes and
|
|
re-creates the serialized executor — never encodes on the torn-down pool."""
|
|
emb = LocalEmbedder(device="mps")
|
|
await emb.embed("warmup")
|
|
await emb.close()
|
|
assert emb._executor is None
|
|
assert emb._model is None
|
|
# Re-use after close must re-initialize cleanly and stay serialized.
|
|
await emb.embed("again")
|
|
assert emb._executor is not None
|
|
assert emb._executor._max_workers == 1 # type: ignore[attr-defined]
|
|
await emb.close()
|