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>
115 lines
4.3 KiB
Python
115 lines
4.3 KiB
Python
"""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
|