fix: bound MPS embedding memory on Apple Silicon - #244
Open
fml09 wants to merge 1 commit into
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 fromrecommended_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
spawnworker. The child owns the model and Metal allocations, so recycling it gives the daemon a deterministic reclamation boundary.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_sizedefaults to 8 on the worker pathmps_memory_limit_ratiodefaults to 0.35mps_low_watermark_ratiodefaults to 0.40mps_high_watermark_ratiodefaults to 0.50worker_timeout_secondsdefaults to 300Cross-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-failurecodefuse-ai/F2LLM-v2-0.6BScope
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.