mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
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.
This commit is contained in:
parent
b4b50253f1
commit
c71592d421
12 changed files with 234 additions and 15 deletions
|
|
@ -12,6 +12,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
|||
|
||||
* **proxy:** per-project savings breakdown on the dashboard for all wrapped agents — Claude Code, Codex, aider, Copilot, and Cursor ([#802](https://github.com/chopratejas/headroom/issues/802)). `headroom wrap claude`/`codex` tag requests with an `X-Headroom-Project` header (launch-directory name); `wrap aider`/`copilot`/`cursor` — whose clients cannot send custom headers — use a `/p/<name>` base-URL prefix the proxy strips. Savings are aggregated per project (persisted, schema v3 with transparent v2 migration), exposed as `savings.per_project` in `/stats` and `projects` in `/stats-history`, and shown in a Per-Project Savings dashboard table.
|
||||
|
||||
### Features
|
||||
|
||||
* **memory:** opt-in Apple-GPU (MPS) embedding offload via `HEADROOM_EMBEDDER_RUNTIME=pytorch_mps`. When set (and Apple MPS is available), the memory embedder runs on the torch sentence-transformers backend on the Apple GPU instead of the default ONNX CPU embedder, freeing the CPU under load. If MPS or the dependencies are unavailable, Headroom logs a warning and uses the existing default embedder selection path (ONNX when available, then the pre-existing local fallback). MPS encode calls are serialized internally (torch-MPS is not thread-safe). Adds the new `[pytorch-mps]` extra (`pip install 'headroom-ai[pytorch-mps]'`). Default behavior is unchanged.
|
||||
|
||||
### Bug Fixes
|
||||
|
||||
* **ccr:** make retrieval store TTL configurable with `HEADROOM_CCR_TTL_SECONDS`, expose the effective TTL in `/v1/retrieve/stats`, and distinguish expired retrievals from missing hashes.
|
||||
|
|
|
|||
|
|
@ -99,7 +99,7 @@ headroom proxy --port 8787 # drop-in proxy, zero code changes
|
|||
headroom perf
|
||||
```
|
||||
|
||||
Granular extras: `[proxy]`, `[mcp]`, `[ml]`, `[code]`, `[memory]`, `[relevance]`, `[image]`, `[agno]`, `[langchain]`, `[evals]`. Requires **Python 3.10+**.
|
||||
Granular extras: `[proxy]`, `[mcp]`, `[ml]`, `[code]`, `[memory]`, `[relevance]`, `[image]`, `[agno]`, `[langchain]`, `[evals]`, `[pytorch-mps]` (Apple-GPU memory-embedder offload — set `HEADROOM_EMBEDDER_RUNTIME=pytorch_mps`). Requires **Python 3.10+**.
|
||||
|
||||
## Proof
|
||||
|
||||
|
|
@ -228,7 +228,7 @@ npm install headroom-ai # TypeScript / Node
|
|||
docker pull ghcr.io/chopratejas/headroom:latest
|
||||
```
|
||||
|
||||
Granular extras: `[proxy]`, `[mcp]`, `[ml]` (Kompress-base), `[code]`, `[memory]`, `[relevance]`, `[image]`, `[agno]`, `[langchain]`, `[evals]`. Requires **Python 3.10+**.
|
||||
Granular extras: `[proxy]`, `[mcp]`, `[ml]` (Kompress-base), `[code]`, `[memory]`, `[relevance]`, `[image]`, `[agno]`, `[langchain]`, `[evals]`, `[pytorch-mps]` (Apple-GPU memory-embedder offload — set `HEADROOM_EMBEDDER_RUNTIME=pytorch_mps`). Requires **Python 3.10+**.
|
||||
|
||||
Using `pipx`? Choose a supported interpreter explicitly:
|
||||
|
||||
|
|
|
|||
|
|
@ -234,6 +234,7 @@ headroom proxy --learn --min-evidence 3
|
|||
| `HEADROOM_PROXY_PYTHON_FORWARDER_MODE` | Python forwarder serialization mode. `byte_faithful` (default) forwards original request bytes verbatim when no transform mutated the body and re-serializes canonically only when needed — keeps Anthropic prompt-cache hit-rate intact. `legacy_json_kwarg` is an explicit operator opt-in for emergency rollback to the historical `httpx ... json=body` behavior. NOT a fallback — only flip on explicit operator decision. | `byte_faithful` |
|
||||
| `HEADROOM_STRIP_INTERNAL_HEADERS` | Python proxy: whether to strip internal `x-headroom-*` request headers (e.g. `x-headroom-bypass`, `x-headroom-mode`, `x-headroom-user-id`, `x-headroom-stack`, `x-headroom-base-url`) before every upstream forwarder call (PR-A5, fixes P5-49). `enabled` (default) stops fingerprinting / leakage. `disabled` is an explicit operator opt-in for diagnostic shadow tracing — NOT a fallback. Inbound reads of these headers (bypass gating, memory user-id resolution) are unaffected because they read `request.headers` directly. | `enabled` |
|
||||
| `HEADROOM_PROXY_STRIP_INTERNAL_HEADERS` | Rust proxy: same policy as `HEADROOM_STRIP_INTERNAL_HEADERS` but for the Rust transparent proxy. Stripping happens inside `build_forward_request_headers` so both HTTP and WebSocket upstream calls are gated by one flag. `enabled` default; `disabled` operator opt-in for diagnostic shadow tracing. Response-side `X-Headroom-*` injection (e.g. `x-headroom-tokens-saved`) is unrelated and stays. | `enabled` |
|
||||
| `HEADROOM_EMBEDDER_RUNTIME` | Set to `pytorch_mps` to run the memory embedder via the torch sentence-transformers backend on the Apple GPU (MPS). Only engages when Apple MPS is actually available; otherwise it logs a warning and uses the existing default embedder selection path. `pytorch_mps` is the only accepted value. Requires the `[pytorch-mps]` extra. See [Memory](/docs/memory#embedding-runtime--gpu-offload-apple-silicon). | default embedder selection |
|
||||
|
||||
### Filesystem Contract
|
||||
|
||||
|
|
|
|||
|
|
@ -51,6 +51,7 @@ pip install "headroom-ai[all]"
|
|||
| `langchain` | LangChain `HeadroomChatModel` wrapper | `pip install "headroom-ai[langchain]"` |
|
||||
| `agno` | Agno `HeadroomAgnoModel` wrapper | `pip install "headroom-ai[agno]"` |
|
||||
| `evals` | Evaluation framework (GSM8K, SQuAD, BFCL benchmarks) | `pip install "headroom-ai[evals]"` |
|
||||
| `pytorch-mps` | Apple-GPU (MPS) memory-embedder offload — **macOS only**, not in `[all]` (torch + sentence-transformers); opt in with `HEADROOM_EMBEDDER_RUNTIME=pytorch_mps` | `pip install "headroom-ai[pytorch-mps]"` |
|
||||
| `all` | Everything above | `pip install "headroom-ai[all]"` |
|
||||
|
||||
You can combine extras:
|
||||
|
|
|
|||
|
|
@ -195,6 +195,36 @@ config = MemoryConfig(
|
|||
)
|
||||
```
|
||||
|
||||
### Embedding Runtime / GPU Offload (Apple Silicon)
|
||||
|
||||
By default the proxy's memory embedder runs on the **ONNX CPU** backend -- fast
|
||||
and dependency-light, but CPU-only. Under sustained load the embedding step can
|
||||
saturate the CPU and make the proxy less responsive.
|
||||
|
||||
On Apple Silicon you can opt in to running the embedder on the **Apple GPU
|
||||
(MPS)** instead, which offloads that work off the CPU and keeps the proxy
|
||||
responsive. Install the extra and set the env var:
|
||||
|
||||
```bash
|
||||
pip install "headroom-ai[pytorch-mps]" # also works as [pytorch_mps]
|
||||
export HEADROOM_EMBEDDER_RUNTIME=pytorch_mps
|
||||
```
|
||||
|
||||
When set, the embedder runs via the torch sentence-transformers backend on the
|
||||
Apple GPU instead of the default ONNX CPU embedder. Notes:
|
||||
|
||||
- **Strictly opt-in.** `pytorch_mps` is the only accepted value; anything else
|
||||
(or unset) keeps the default ONNX CPU embedder. Default behavior is unchanged,
|
||||
and there is no CLI flag -- it is env-var only.
|
||||
- **Auto-fallback.** It only activates when Apple MPS is actually available
|
||||
(Apple Silicon + torch). If MPS is unavailable or torch/sentence-transformers
|
||||
is not installed, it logs a warning and uses the existing default embedder
|
||||
selection path: ONNX when available, then the pre-existing local
|
||||
sentence-transformers fallback.
|
||||
- **MPS serialization.** torch-MPS is not thread-safe, so the embedder
|
||||
serializes MPS encode calls internally via a single-worker executor. This is
|
||||
automatic -- there is nothing to configure.
|
||||
|
||||
### Storage
|
||||
|
||||
Storage uses **SQLite** for CRUD and filtering, **HNSW** for vector similarity search, and **FTS5** for full-text keyword search. All embedded -- no external services required.
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ import asyncio
|
|||
import logging
|
||||
import os
|
||||
import warnings
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from functools import cached_property
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
|
||||
|
|
@ -123,6 +124,10 @@ class LocalEmbedder:
|
|||
self._device: str | None = None
|
||||
self._dimension: int | None = None
|
||||
self._lock = asyncio.Lock()
|
||||
# Dedicated single-worker executor, created only when the resolved device
|
||||
# is MPS (see _load_model). torch-MPS is not thread-safe, so every encode()
|
||||
# must run on one thread. Stays None for CPU/CUDA → default shared executor.
|
||||
self._executor: ThreadPoolExecutor | None = None
|
||||
|
||||
def _check_dependencies(self) -> None:
|
||||
"""Check that required dependencies are installed."""
|
||||
|
|
@ -166,6 +171,13 @@ class LocalEmbedder:
|
|||
else:
|
||||
self._device = self._detect_device()
|
||||
|
||||
# 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
|
||||
# MPS so calls serialize; other devices keep the shared default executor.
|
||||
if self._device == "mps" and self._executor is None:
|
||||
self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="mps-embed")
|
||||
|
||||
# Use centralized registry for shared model instances
|
||||
self._model = MLModelRegistry.get_sentence_transformer(self._model_name, self._device)
|
||||
|
||||
|
|
@ -199,7 +211,7 @@ class LocalEmbedder:
|
|||
model = self._model # Local reference for lambda closure
|
||||
loop = asyncio.get_event_loop()
|
||||
embedding = await loop.run_in_executor(
|
||||
None,
|
||||
self._executor,
|
||||
lambda: model.encode(text, convert_to_numpy=True, normalize_embeddings=False),
|
||||
)
|
||||
|
||||
|
|
@ -242,7 +254,7 @@ class LocalEmbedder:
|
|||
model = self._model # Local reference for lambda closure
|
||||
loop = asyncio.get_event_loop()
|
||||
embeddings = await loop.run_in_executor(
|
||||
None,
|
||||
self._executor,
|
||||
lambda: model.encode(
|
||||
non_empty_texts, convert_to_numpy=True, normalize_embeddings=False
|
||||
),
|
||||
|
|
@ -276,9 +288,14 @@ class LocalEmbedder:
|
|||
return self.DEFAULT_MAX_TOKENS
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close resources (no-op for local embedder)."""
|
||||
# LocalEmbedder doesn't hold persistent connections
|
||||
pass
|
||||
"""Close resources: shut down the MPS serialization executor and drop the
|
||||
cached model reference so a later embed() fully re-initializes (and
|
||||
re-creates the serialized executor) instead of encoding on a torn-down pool.
|
||||
"""
|
||||
if self._executor is not None:
|
||||
self._executor.shutdown(wait=False)
|
||||
self._executor = None
|
||||
self._model = None
|
||||
|
||||
|
||||
# =============================================================================
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ import enum
|
|||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
|
|
@ -327,15 +328,44 @@ class MemoryHandler:
|
|||
embedder_model = "all-MiniLM-L6-v2"
|
||||
vector_dimension = 384
|
||||
|
||||
# Opt-in GPU offload: HEADROOM_EMBEDDER_RUNTIME=pytorch_mps routes embedding
|
||||
# through the torch sentence-transformers backend on the Apple GPU (MPS).
|
||||
# LocalEmbedder serializes MPS encode calls (torch-MPS is not thread-safe).
|
||||
# We switch only when MPS is actually available; otherwise keep the
|
||||
# existing default embedder selection path (ONNX when available, then
|
||||
# the pre-existing local sentence-transformers fallback).
|
||||
if os.environ.get("HEADROOM_EMBEDDER_RUNTIME", "").strip().lower() == "pytorch_mps":
|
||||
try:
|
||||
import sentence_transformers # noqa: F401
|
||||
import torch
|
||||
|
||||
if torch.backends.mps.is_available():
|
||||
embedder_backend = "local"
|
||||
logger.info(
|
||||
"Memory: HEADROOM_EMBEDDER_RUNTIME=pytorch_mps → "
|
||||
"torch embedder on Apple GPU (MPS)"
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
"Memory: HEADROOM_EMBEDDER_RUNTIME=pytorch_mps requested but "
|
||||
"MPS is not available; using default embedder selection"
|
||||
)
|
||||
except ImportError:
|
||||
logger.warning(
|
||||
"Memory: HEADROOM_EMBEDDER_RUNTIME=pytorch_mps requested but "
|
||||
"torch/sentence-transformers not installed; using default embedder selection"
|
||||
)
|
||||
|
||||
# Check if ONNX runtime is available (should be — it's in proxy deps)
|
||||
try:
|
||||
import onnxruntime # noqa: F401
|
||||
except ImportError:
|
||||
# Fall back to sentence-transformers (requires torch)
|
||||
embedder_backend = "local"
|
||||
logger.info(
|
||||
"Memory: onnxruntime not available, falling back to sentence-transformers"
|
||||
)
|
||||
if embedder_backend == "onnx":
|
||||
try:
|
||||
import onnxruntime # noqa: F401
|
||||
except ImportError:
|
||||
# Fall back to sentence-transformers (requires torch)
|
||||
embedder_backend = "local"
|
||||
logger.info(
|
||||
"Memory: onnxruntime not available, falling back to sentence-transformers"
|
||||
)
|
||||
|
||||
backend_config = LocalBackendConfig(
|
||||
db_path=self.config.db_path,
|
||||
|
|
|
|||
|
|
@ -109,6 +109,12 @@ memory-stack = [
|
|||
"qdrant-client>=1.9.0,<2.0",
|
||||
"neo4j>=5.20.0,<7.0",
|
||||
]
|
||||
# Apple-Silicon GPU (MPS) offload for the memory embedder. Opt in at runtime with
|
||||
# HEADROOM_EMBEDDER_RUNTIME=pytorch_mps. macOS-only; intentionally excluded from [all].
|
||||
pytorch-mps = [
|
||||
"torch>=2.0.0; sys_platform == 'darwin'",
|
||||
"sentence-transformers>=2.2.0; sys_platform == 'darwin'",
|
||||
]
|
||||
# Semantic relevance scoring with embeddings.
|
||||
# Uses `fastembed` (BAAI/bge-small-en-v1.5 by default — 33M params,
|
||||
# 384 dims, ~30 MB int8-quantized ONNX). Same library + model used by
|
||||
|
|
|
|||
69
tests/test_memory/test_embedder_mps_serialization.py
Normal file
69
tests/test_memory/test_embedder_mps_serialization.py
Normal file
|
|
@ -0,0 +1,69 @@
|
|||
"""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_shared_executor() -> None:
|
||||
"""On CPU the dedicated executor stays None (unchanged default-pool behavior)."""
|
||||
emb = LocalEmbedder(device="cpu")
|
||||
await emb.embed("hello world")
|
||||
assert emb._device == "cpu"
|
||||
assert emb._executor is None # default shared executor, not serialized
|
||||
await emb.close()
|
||||
|
||||
|
||||
@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()
|
||||
|
|
@ -311,6 +311,7 @@ Some settings can be configured via environment variables:
|
|||
| `HEADROOM_SAVINGS_PATH` | Full path to the proxy savings JSON ledger. Always wins when set. | derived from `${HEADROOM_WORKSPACE_DIR}` |
|
||||
| `HEADROOM_TOIN_PATH` | Full path to the TOIN telemetry JSON file. Always wins when set. | derived from `${HEADROOM_WORKSPACE_DIR}` |
|
||||
| `HEADROOM_SUBSCRIPTION_STATE_PATH` | Full path to the subscription tracker state. Always wins when set. | derived from `${HEADROOM_WORKSPACE_DIR}` |
|
||||
| `HEADROOM_EMBEDDER_RUNTIME` | Set to `pytorch_mps` to run the memory embedder via the torch sentence-transformers backend on the Apple GPU (MPS). Only engages when Apple MPS is actually available; otherwise it logs a warning and uses the existing default embedder selection path. `pytorch_mps` is the only accepted value. Requires the `[pytorch-mps]` extra. See [Memory](memory.md#embedding-runtime--gpu-offload-apple-silicon). | default embedder selection |
|
||||
|
||||
## Filesystem Contract
|
||||
|
||||
|
|
|
|||
|
|
@ -648,6 +648,34 @@ Limit CPU and memory usage:
|
|||
</dict>
|
||||
```
|
||||
|
||||
## Apple GPU (MPS) Embedding Offload
|
||||
|
||||
On Apple Silicon, the proxy's memory embedder can run on the Apple GPU (MPS)
|
||||
instead of the default ONNX CPU backend. Offloading embedding to the GPU frees
|
||||
the CPU under load, keeping the proxy responsive — useful on fanless Macs (e.g.
|
||||
the M5 Air) that are prone to CPU-saturation timeouts.
|
||||
|
||||
Enable it by installing the extra and setting the env var:
|
||||
|
||||
```bash
|
||||
pip install 'headroom-ai[pytorch-mps]' # also works as [pytorch_mps]
|
||||
export HEADROOM_EMBEDDER_RUNTIME=pytorch_mps
|
||||
```
|
||||
|
||||
Under a LaunchAgent, set the env var in the plist `EnvironmentVariables`
|
||||
section:
|
||||
|
||||
```xml
|
||||
<key>HEADROOM_EMBEDDER_RUNTIME</key>
|
||||
<string>pytorch_mps</string>
|
||||
```
|
||||
|
||||
It only engages when Apple MPS is actually available (Apple Silicon + torch).
|
||||
If MPS is unavailable or the dependencies are missing, the proxy logs a warning
|
||||
and uses the existing default embedder selection path. This is strictly opt-in;
|
||||
default behavior is unchanged. See [Memory](memory.md#embedding-runtime--gpu-offload-apple-silicon)
|
||||
for details.
|
||||
|
||||
## FAQ
|
||||
|
||||
**Q: Why LaunchAgent instead of running `headroom proxy` manually?**
|
||||
|
|
|
|||
|
|
@ -446,6 +446,38 @@ config = MemoryConfig(
|
|||
)
|
||||
```
|
||||
|
||||
### Embedding Runtime / GPU Offload (Apple Silicon)
|
||||
|
||||
By default the proxy's memory embedder runs on the **ONNX CPU** backend. This
|
||||
is fast and dependency-light, but it is CPU-only — under sustained load the
|
||||
embedding step can saturate the CPU and make the proxy less responsive.
|
||||
|
||||
On Apple Silicon you can opt in to running the embedder on the **Apple GPU
|
||||
(MPS)** instead, which offloads that work off the CPU and keeps the proxy
|
||||
responsive. This is especially useful on fanless Macs (e.g. the M5 Air) that
|
||||
are prone to CPU-saturation timeouts.
|
||||
|
||||
Enable it by installing the extra and setting the env var:
|
||||
|
||||
```bash
|
||||
pip install 'headroom-ai[pytorch-mps]' # also works as [pytorch_mps]
|
||||
export HEADROOM_EMBEDDER_RUNTIME=pytorch_mps
|
||||
```
|
||||
|
||||
When set, the embedder runs via the torch sentence-transformers backend on the
|
||||
Apple GPU instead of the default ONNX CPU embedder. Notes:
|
||||
|
||||
- **Strictly opt-in.** `pytorch_mps` is the only accepted value; anything else
|
||||
(or unset) keeps the default ONNX CPU embedder. Default behavior is unchanged.
|
||||
- **Auto-fallback.** It only activates when Apple MPS is actually available
|
||||
(Apple Silicon + torch). If MPS is unavailable or torch/sentence-transformers
|
||||
is not installed, it logs a warning and uses the existing default embedder
|
||||
selection path: ONNX when available, then the pre-existing local
|
||||
sentence-transformers fallback.
|
||||
- **MPS serialization.** torch-MPS is not thread-safe, so the embedder
|
||||
serializes MPS encode calls internally via a single-worker executor. This is
|
||||
automatic — there is nothing to configure.
|
||||
|
||||
### Storage Configuration
|
||||
|
||||
```python
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue