Skip to content

fix: bound MPS embedding memory on Apple Silicon - #244

Open
fml09 wants to merge 1 commit into
cocoindex-io:mainfrom
fml09:fix/bound-mps-embedding-memory
Open

fix: bound MPS embedding memory on Apple Silicon#244
fml09 wants to merge 1 commit into
cocoindex-io:mainfrom
fml09:fix/bound-mps-embedding-memory

Conversation

@fml09

@fml09 fml09 commented Jul 30, 2026

Copy link
Copy Markdown

Problem

On a 48 GiB Apple Silicon machine, one long-lived daemon was observed at roughly 54 GiB of process footprint. About 53.6 GiB was attributed to IOAccelerator/Metal while the Python heap remained small. Stack samples consistently pointed at SentenceTransformer encode, an MPS-to-CPU copy, and Metal command-buffer completion.

The existing failure path only called torch.mps.empty_cache() after an OOM. That is too late here: the default MPS allocator watermarks scale from recommended_max_memory, so system pressure can become severe before PyTorch raises. Cache cleanup also cannot guarantee that a long-lived process returns all driver-owned Metal allocations.

Multiple projects could compound the problem by indexing concurrently through the same daemon.

Approach

  • Run MPS SentenceTransformer inference in a supervised spawn worker. The child owns the model and Metal allocations, so recycling it gives the daemon a deterministic reclamation boundary.
  • Start with an inner batch size of 8, retry MPS OOMs with batch-size halving, and replace a crashed or timed-out worker once.
  • Clear caches at the configured memory threshold and recycle the worker when driver-owned memory remains high.
  • Install conservative allocator defaults before Torch initializes: 40% soft watermark and 50% hard watermark. Explicit environment variables continue to win.
  • Serialize project indexing through a daemon-wide gate and coalesce concurrent initial-index requests, preventing several projects from driving the local embedder at once.
  • Keep non-MPS SentenceTransformer and LiteLLM execution paths unchanged.

Behavior and compatibility

No configuration migration is required. The guard applies to SentenceTransformer models on explicit MPS, or automatic device selection on macOS. All controls are optional and available under embedding:

  • batch_size defaults to 8 on the worker path
  • mps_memory_limit_ratio defaults to 0.35
  • mps_low_watermark_ratio defaults to 0.40
  • mps_high_watermark_ratio defaults to 0.50
  • worker_timeout_seconds defaults to 300

Cross-project indexing is intentionally serialized. This trades aggregate indexing throughput for bounded unified-memory pressure. Search remains available; a project waits only when it has not completed its initial index.

Validation

  • uv run prek run --all-files --verbose --show-diff-on-failure
    • Ruff, Ruff format, lockfile validation, strict mypy, secret checks, and the full test hook passed
    • 289 tests passed; 8 Docker-only tests were deselected by the existing suite configuration
  • Apple Silicon smoke test with codefuse-ai/F2LLM-v2-0.6B
    • a forced low recycle threshold advanced worker generations across four requests
    • every request returned 64 embeddings with dimension 1024
    • driver allocation returned to roughly 1.2 GiB after recycling
  • Three-project live reindex on the same machine
    • one embedding worker was active
    • Metal allocation held near 5.3 GiB with stricter configured thresholds, versus roughly 53.6 GiB in the original failure

Scope

This PR is intentionally independent of #243. That PR prevents a multi-client daemon-startup stampede; this PR bounds memory inside the daemon after startup. They address separate failure modes and can land in either order.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant