fix: bound MPS embedding memory on Apple Silicon

[fml09]

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.

[badmonster0]

@fml09 thanks for the fix! @georgeh0 can help take a look!

[georgeh0]

thanks for the improvement!

Actually the underlying cocoindex engine has a subprocess mode for GPU workload. Currently we can enable it by setting COCOINDEX_RUN_GPU_IN_SUBPROCESS=1 (doc). I'm wondering will it make the code change here simpler by using that (e.g. we may set the env at initialization time, and happy to explore easier ways to toggle this)

[fml09]

Thanks for catching this earlier. I reworked the implementation around the upstream primitives and rerun the Apple Silicon benchmarks.

The experiments confirmed that process isolation and active memory reclamation are separate concerns:

• The CocoIndex GPU subprocess runner maintains a single persistent worker. While COCOINDEX_RUN_GPU_IN_SUBPROCESS=1 shields the main daemon process from Metal allocations, it does not automatically reclaim memory from a healthy worker.
• Under a variable-shape workload, the baseline runner retained ~7.32 GB versus ~3.30 GB on our custom worker.
• Most of that delta turned out to be reclaimable allocator cache. Running gc.collect() followed by torch.mps.empty_cache() directly inside the GPU child process reduced driver_allocated_memory() from ~6.19 GB to 1.22 GB, and process RSS dropped from 6.79 GB to 1.74 GB without changing the worker PID. This is critical because driver_allocated_memory() captures cached allocator and MPSGraph allocations that process isolation alone leaves untouched.

Updated Implementation

Based on these findings, the revised implementation:

1. Adopts standard SentenceTransformerEmbedder with COCOINDEX_RUN_GPU_IN_SUBPROCESS=1.
2. Configures conservative MPS watermarks prior to the first GPU invocation (explicit env vars still take precedence).
3. Triggers best-effort allocator cleanup inside the GPU child process after each indexing run.
4. Bumps the CocoIndex minimum version to 1.0.17 to consume native RetryWithSmallerBatch support.

This allowed us to strip out our custom worker, IPC protocol, manual batch-halving, recycling thresholds, and the daemon-wide indexing gate entirely.

Benchmark Results & Validation

• End-to-End Test: Completed across 3 projects (62 files, 811 chunks each) on a single, stable subprocess PID.
• Memory Footprint: Peak process footprint reached ~4.12 GB, and post-index cleanup brought RSS back down to ~1.79 GB (matching warm-model baseline levels). High-pressure workloads also completed within configured hard watermarks.
• Test Suite: All 287 tests are passing along with Ruff, mypy, and lockfile validation.

Conclusion & Follow-up

Given these measurements, active worker recycling is not required for our memory safety goals. Combining the standard runner, conservative watermarks, native smaller-batch retries, and post-index cleanup provides the leanest implementation while preserving warm-model latency across requests.

Follow-up Scope:
The current CocoIndex runner does not provide bounded per-call timeouts or recovery from repeated hard crashes. Rather than re-introducing private process-pool supervision in this PR, I propose addressing timeout and crash-recovery guarantees separately, ideally via an upstream runner API proposal.

[georgeh0]

Thanks for the improvement!