fix: bound MPS embedding memory on Apple Silicon

EMBEDDINGS.md

@@ -21,7 +21,7 @@
 
 | Path | Best For... | Key Advantage | Trade-off |
 | :--- | :--- | :--- | :--- |
-| **Local Sentence-Transformers** | Most users, laptops, quick setup. | **Fastest** (in-process), private, offline. | Larger initial pip install (`[full]`). |
+| **Local Sentence-Transformers** | Most users, laptops, quick setup. | **Fastest**, private, offline. | Larger initial pip install (`[full]`). |
 | **Cloud LiteLLM Remote** | Large codebases, weak local hardware. | Top performance, zero local resource usage. | Per-token costs, data leaves machine. |
 | **Local LiteLLM** | Power users, shared GPU resources. | Flexibility, unified model management. | Requires managing a separate server. |
 
@@ -31,7 +31,7 @@
 
 ### Speed & Latency
 
-- **Local Sentence-Transformers**: Typically the **fastest** option for small-to-medium **encoder** models. Because it runs directly inside the `cocoindex-code` process, it avoids the network latency of Cloud APIs and the communication overhead of Local Servers (Ollama). Decoder-based models (e.g. `harrier-oss-v1`) are also in-process, but process tokens sequentially — expect 3–10× slower indexing per chunk on CPU unless you have a GPU.
+- **Local Sentence-Transformers**: Typically the **fastest** option for small-to-medium **encoder** models. It avoids network latency and the server overhead of Local Servers (Ollama). On Apple Silicon, MPS inference runs in a supervised child process so Metal allocations can be reclaimed independently; other devices use the normal in-process path. Decoder-based models (e.g. `harrier-oss-v1`) process tokens sequentially — expect 3–10× slower indexing per chunk on CPU unless you have a GPU.
 - **Local Servers (Ollama)**: Ideal for running **heavy models** (like `mxbai-embed-large`) on a GPU. While it has slight overhead compared to in-process execution, it is much faster than running large models on a CPU.
 - **Cloud APIs**: Slower per-request due to network latency, but highly parallel. Best for the initial indexing of massive repositories.
 
@@ -54,6 +54,19 @@ Indexing speed depends on **both parameter count and model architecture**:
 
 For encoder models at medium or large sizes, a GPU will still accelerate indexing. Add `device: cuda` (or `mps` on Mac) to `global_settings.yml`.
 
+On Apple Silicon, SentenceTransformer calls run through [CocoIndex's isolated
+GPU subprocess](https://github.com/cocoindex-io/cocoindex/blob/v1.0.18/python/cocoindex/_internal/runner.py).
+This keeps the model warm across batches while separating Metal allocations
+from the daemon. cocoindex-code configures PyTorch to begin adaptive allocation
+cleanup at 40% of its recommended maximum working set and to enforce a hard
+limit at 50%; tune these values with `mps_low_watermark_ratio` and
+`mps_high_watermark_ratio`. At the end of each index run, unused allocator cache
+is released inside the GPU subprocess. CocoIndex also [retries MPS OOM failures
+with smaller batches](https://github.com/cocoindex-io/cocoindex/blob/v1.0.18/python/cocoindex/ops/sentence_transformers.py).
+See the [PyTorch MPS environment variable reference](https://docs.pytorch.org/docs/stable/mps_environment_variables.html)
+for the allocator semantics. Explicit environment variables take precedence
+over the YAML defaults.
+
 ---
 
 ## The `ccc init` Wizard

README.md

@@ -505,6 +505,8 @@ embedding:
   model: Snowflake/snowflake-arctic-embed-xs
   device: mps                                        # optional: cpu, cuda, mps (auto-detected if omitted)
   min_interval_ms: 300                               # optional: pace LiteLLM embedding requests to reduce 429s; defaults to 5 for LiteLLM
+  mps_low_watermark_ratio: 0.4                       # optional: PyTorch allocator soft limit
+  mps_high_watermark_ratio: 0.5                      # optional: PyTorch allocator hard limit
 
   # Optional extra kwargs passed to the embedder, separately for indexing vs query.
   # `ccc init` auto-populates these for known models (e.g. Cohere, Voyage, Nvidia NIM,
@@ -523,6 +525,10 @@ daemon:
 
 > **Note:** The daemon inherits your shell environment. If an API key (e.g. `OPENAI_API_KEY`) is already set as an environment variable, you don't need to duplicate it in `envs`. The `envs` field is only for values that aren't in your environment.
 
+> **Apple Silicon memory safety:** MPS SentenceTransformer calls use [CocoIndex's isolated GPU subprocess](https://github.com/cocoindex-io/cocoindex/blob/v1.0.18/python/cocoindex/_internal/runner.py), keeping the model loaded while isolating Metal allocations from the daemon. The low and high watermarks are ratios of PyTorch's recommended maximum working set; they default here to `0.4` and `0.5`. CocoIndex retries MPS out-of-memory failures with progressively smaller batches, and cocoindex-code [releases unused allocator cache](https://docs.pytorch.org/docs/stable/generated/torch.mps.empty_cache.html) after each index run. Explicit `COCOINDEX_RUN_GPU_IN_SUBPROCESS`, `PYTORCH_MPS_LOW_WATERMARK_RATIO`, and `PYTORCH_MPS_HIGH_WATERMARK_RATIO` environment variables take precedence over these defaults.
+
+> **Indexing concurrency:** Multiple projects may prepare indexes concurrently, while CocoIndex serializes their GPU calls through its single MPS subprocess. A search waits only when its own project still needs the initial index.
+
 > **Idle timeout:** the background daemon holds the embedding model in RAM, so it exits after `daemon.idle_timeout_minutes` without client activity and is restarted automatically on your next `ccc` command or MCP search. A live MCP session sends periodic heartbeats, so the daemon never idles out while your coding agent is connected. Set `0` to keep the daemon running forever.
 
 > **Custom location:** set `COCOINDEX_CODE_DIR` to place `global_settings.yml` somewhere other than `~/.cocoindex_code/` — useful if you want the file to live alongside your projects (e.g. on a synced folder).

pyproject.toml

@@ -23,7 +23,7 @@ classifiers = [
 
 dependencies = [
     "mcp>=1.0.0",
-    "cocoindex[litellm]>=1.0.13,<1.1.0",
+    "cocoindex[litellm]>=1.0.17,<1.1.0",
     "sqlite-vec>=0.1.0",
     "pydantic>=2.0.0",
     "numpy>=1.24.0",
@@ -39,23 +39,23 @@ dependencies = [
 # `embeddings-local` is the primary feature extra: it pulls in
 # `sentence-transformers` (via cocoindex) so local embeddings work without
 # an API key.
-embeddings-local = ["cocoindex[sentence-transformers]>=1.0.13,<1.1.0"]
+embeddings-local = ["cocoindex[sentence-transformers]>=1.0.17,<1.1.0"]
 # `full` is the umbrella "batteries-included" alias. Today it's just
 # `embeddings-local`, but we expect to bundle more optional niceties under
 # it over time — users who want everything can keep using `[full]` and pick
 # up the additions automatically. The name also matches the Docker
 # `:full` image variant for consistency across install paths. Contents are
 # inlined rather than self-referencing `cocoindex-code[embeddings-local]`
 # to avoid resolver edge cases with older pip.
-full = ["cocoindex[sentence-transformers]>=1.0.13,<1.1.0"]
+full = ["cocoindex[sentence-transformers]>=1.0.17,<1.1.0"]
 dev = [
     "pytest>=7.0.0",
     "pytest-asyncio>=0.21.0",
     "pytest-cov>=4.0.0",
     "ruff>=0.1.0",
     "mypy>=1.0.0",
     "prek>=0.1.0",
-    "cocoindex[sentence-transformers]>=1.0.13,<1.1.0",
+    "cocoindex[sentence-transformers]>=1.0.17,<1.1.0",
 ]
 
 [project.scripts]
@@ -89,7 +89,7 @@ dev = [
     "mypy>=1.0.0",
     "prek>=0.1.0",
     "types-pyyaml>=6.0.12.20250915",
-    "cocoindex[sentence-transformers]>=1.0.13,<1.1.0",
+    "cocoindex[sentence-transformers]>=1.0.17,<1.1.0",
 ]
 
 [tool.ruff]

src/cocoindex_code/daemon.py

@@ -73,7 +73,7 @@
     target_sqlite_db_path,
     user_settings_path,
 )
-from .shared import Embedder, check_embedding, create_embedder
+from .shared import Embedder, check_embedding, configure_mps_environment, create_embedder
 
 logger = logging.getLogger(__name__)
 
@@ -135,6 +135,7 @@ class ProjectRegistry:
 
     _projects: dict[str, Project]
     _embedder: Embedder | None
+    _clear_mps_cache_after_index: bool
     indexing_params: dict[str, Any]
     query_params: dict[str, Any]
 
@@ -143,9 +144,11 @@ def __init__(
         embedder: Embedder | None,
         indexing_params: dict[str, Any] | None = None,
         query_params: dict[str, Any] | None = None,
+        clear_mps_cache_after_index: bool = False,
     ) -> None:
         self._projects = {}
         self._embedder = embedder
+        self._clear_mps_cache_after_index = clear_mps_cache_after_index
         self.indexing_params = dict(indexing_params) if indexing_params else {}
         self.query_params = dict(query_params) if query_params else {}
 
@@ -165,6 +168,7 @@ async def get_project(self, project_root: str) -> Project:
                 indexing_params=self.indexing_params,
                 query_params=self.query_params,
                 chunker_registry=chunker_registry,
+                clear_mps_cache_after_index=self._clear_mps_cache_after_index,
             )
             self._projects[project_root] = project
         return self._projects[project_root]
@@ -636,6 +640,7 @@ def run_daemon(
     embedder: Embedder | None
     indexing_params: dict[str, Any] = {}
     query_params: dict[str, Any] = {}
+    clear_mps_cache_after_index = False
     handshake_warnings: list[str] = []
     daemon_settings = DaemonSettings()
     if user_settings_path().is_file():
@@ -644,6 +649,7 @@ def run_daemon(
         settings_env_keys = list(user_settings.envs.keys())
         for key, value in user_settings.envs.items():
             os.environ[key] = value
+        clear_mps_cache_after_index = configure_mps_environment(user_settings.embedding)
         # Resolve params BEFORE constructing the embedder so invalid configs
         # fail fast without paying the model-load cost.
         try:
@@ -684,6 +690,7 @@ def run_daemon(
         embedder,
         indexing_params=indexing_params,
         query_params=query_params,
+        clear_mps_cache_after_index=clear_mps_cache_after_index,
     )
 
     sock_path = daemon_socket_path()

src/cocoindex_code/project.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import asyncio
+import logging
 import sqlite3
 from collections.abc import AsyncIterator, Callable
 from pathlib import Path
@@ -39,15 +40,21 @@
     QUERY_EMBED_PARAMS,
     SQLITE_DB,
     Embedder,
+    clear_mps_allocator_cache,
 )
 
+logger = logging.getLogger(__name__)
+
 
 class Project:
     _env: coco.Environment
     _app: coco.App[[], None]
     _project_root: Path
     _index_lock: asyncio.Lock
+    _clear_mps_cache_after_index: bool
     _initial_index_done: asyncio.Event
+    _initial_index_task: asyncio.Task[None] | None
+    _initial_index_started: asyncio.Event | None
     _indexing_stats: IndexingProgress | None = None
 
     def close(self) -> None:
@@ -67,13 +74,15 @@ async def run_index(
         on_progress: Callable[[IndexingProgress], None] | None = None,
         on_started: asyncio.Event | None = None,
     ) -> None:
-        """Acquire the index lock, run indexing, and release.
+        """Acquire the project lock, run indexing, and release it.
 
-        If *on_started* is provided, it is set once the lock is acquired
-        (i.e. indexing has truly begun).  On completion (success or failure)
-        ``_initial_index_done`` is set.
+        If *on_started* is provided, it is set once the project lock is
+        acquired. On completion (success or failure) ``_initial_index_done``
+        is set.
         """
         async with self._index_lock:
+            if on_started is not None:
+                on_started.set()
             self._indexing_stats = IndexingProgress(
                 num_execution_starts=0,
                 num_unchanged=0,
@@ -82,8 +91,6 @@ async def run_index(
                 num_reprocesses=0,
                 num_errors=0,
             )
-            if on_started is not None:
-                on_started.set()
             await self._run_index_inner(on_progress=on_progress)
 
     async def _run_index_inner(
@@ -109,20 +116,32 @@ async def _run_index_inner(
                         on_progress(progress)
                     await asyncio.sleep(0.1)
         finally:
-            self._initial_index_done.set()
-            self._indexing_stats = None
+            try:
+                if self._clear_mps_cache_after_index:
+                    try:
+                        await clear_mps_allocator_cache()
+                    except Exception:
+                        logger.exception("Unable to clear the MPS allocator cache after indexing")
+            finally:
+                self._initial_index_done.set()
+                self._indexing_stats = None
 
     async def ensure_indexing_started(self) -> None:
-        """Kick off background indexing and wait until it has actually started.
+        """Kick off background indexing and wait until it is active or queued.
 
-        Returns once the indexing task holds the lock.  Safe to call multiple
-        times — only the first call spawns a task; subsequent calls return
-        immediately.
+        Returns once the indexing task holds the project lock. Safe to call
+        multiple times — only the first call spawns a task; subsequent calls
+        return immediately.
         """
         if self._initial_index_done.is_set() or self._index_lock.locked():
             return
-        started = asyncio.Event()
-        asyncio.create_task(self.run_index(on_started=started))
+        if self._initial_index_task is None or self._initial_index_task.done():
+            self._initial_index_started = asyncio.Event()
+            self._initial_index_task = asyncio.create_task(
+                self.run_index(on_started=self._initial_index_started)
+            )
+        started = self._initial_index_started
+        assert started is not None
         await started.wait()
 
     async def stream_index(self) -> AsyncIterator[IndexStreamResponse]:
@@ -263,6 +282,7 @@ async def create(
         indexing_params: dict[str, Any],
         query_params: dict[str, Any],
         chunker_registry: dict[str, ChunkerFn] | None = None,
+        clear_mps_cache_after_index: bool = False,
     ) -> Project:
         """Create a project with explicit embedder and per-call params.
 
@@ -281,6 +301,8 @@ async def create(
             chunker_registry: Optional mapping of file suffix (e.g. ``".toml"``)
                 to a ``ChunkerFn``. When a suffix matches, the registered
                 chunker is called instead of the built-in splitter.
+            clear_mps_cache_after_index: Whether to release unused MPS allocator
+                memory in CocoIndex's GPU subprocess after each index run.
         """
         settings_dir = project_root / ".cocoindex_code"
         settings_dir.mkdir(parents=True, exist_ok=True)
@@ -315,5 +337,8 @@ async def create(
         result._app = app
         result._project_root = project_root
         result._index_lock = asyncio.Lock()
+        result._clear_mps_cache_after_index = clear_mps_cache_after_index
         result._initial_index_done = asyncio.Event()
+        result._initial_index_task = None
+        result._initial_index_started = None
         return result

src/cocoindex_code/settings.py

@@ -98,12 +98,30 @@ class EmbeddingSettings:
     provider: str = "litellm"
     device: str | None = None
     min_interval_ms: int | None = None
+    # PyTorch MPS allocator limits used by CocoIndex's isolated GPU runner.
+    mps_low_watermark_ratio: float = 0.4
+    mps_high_watermark_ratio: float = 0.5
     # Extra kwargs spread into ``embedder.embed()`` during indexing/query.
     # ``None`` means the user did not set the key; ``{}`` is an explicit empty
     # dict (used to opt out of the legacy-bridge warning).
     indexing_params: dict[str, Any] | None = None
     query_params: dict[str, Any] | None = None
 
+    def __post_init__(self) -> None:
+        ratios = (
+            ("mps_low_watermark_ratio", self.mps_low_watermark_ratio),
+            ("mps_high_watermark_ratio", self.mps_high_watermark_ratio),
+        )
+        for name, value in ratios:
+            if isinstance(value, bool) or not 0 < value <= 1:
+                raise ValueError(f"embedding.{name} must be greater than 0 and at most 1")
+
+        if self.mps_low_watermark_ratio > self.mps_high_watermark_ratio:
+            raise ValueError(
+                "embedding.mps_low_watermark_ratio must be at most "
+                "embedding.mps_high_watermark_ratio"
+            )
+
 
 @dataclass
 class DaemonSettings:
@@ -430,6 +448,11 @@ def _embedding_settings_to_dict(embedding: EmbeddingSettings) -> dict[str, Any]:
         d["device"] = embedding.device
     if embedding.min_interval_ms is not None:
         d["min_interval_ms"] = embedding.min_interval_ms
+    defaults = EmbeddingSettings(model=embedding.model)
+    if embedding.mps_low_watermark_ratio != defaults.mps_low_watermark_ratio:
+        d["mps_low_watermark_ratio"] = embedding.mps_low_watermark_ratio
+    if embedding.mps_high_watermark_ratio != defaults.mps_high_watermark_ratio:
+        d["mps_high_watermark_ratio"] = embedding.mps_high_watermark_ratio
     if embedding.indexing_params is not None:
         d["indexing_params"] = dict(embedding.indexing_params)
     if embedding.query_params is not None:
@@ -458,6 +481,10 @@ def _user_settings_from_dict(d: dict[str, Any]) -> UserSettings:
         emb_kwargs["device"] = emb_dict["device"]
     if "min_interval_ms" in emb_dict:
         emb_kwargs["min_interval_ms"] = emb_dict["min_interval_ms"]
+    if "mps_low_watermark_ratio" in emb_dict:
+        emb_kwargs["mps_low_watermark_ratio"] = float(emb_dict["mps_low_watermark_ratio"])
+    if "mps_high_watermark_ratio" in emb_dict:
+        emb_kwargs["mps_high_watermark_ratio"] = float(emb_dict["mps_high_watermark_ratio"])
     # indexing_params / query_params: missing → None (dataclass default);
     # present-but-null → {} (treat the same as an empty dict, since both mean
     # "user acknowledged the key and wants no extra kwargs").
@@ -573,6 +600,13 @@ def save_user_settings(settings: UserSettings) -> Path:
     ),
 }
 
+_MPS_SAFETY_COMMENT = (
+    "  #\n"
+    "  # Apple Silicon MPS allocator limits for CocoIndex's GPU subprocess.\n"
+    "  # mps_low_watermark_ratio: 0.4\n"
+    "  # mps_high_watermark_ratio: 0.5\n"
+)
+
 
 def save_initial_user_settings(
     embedding: EmbeddingSettings,
@@ -595,6 +629,8 @@ def save_initial_user_settings(
         sort_keys=False,
     )
     content = _INITIAL_HEADER + emb_block
+    if embedding.provider == "sentence-transformers":
+        content += _MPS_SAFETY_COMMENT
     if not defaults_applied:
         hint = _PARAMS_COMMENT_BY_PROVIDER.get(embedding.provider)
         if hint is not None: