From 80d1036f55454b96897a1efc669f99974bb37947 Mon Sep 17 00:00:00 2001 From: fml09 Date: Thu, 30 Jul 2026 16:07:12 +0900 Subject: [PATCH 1/2] fix: bound and recycle MPS embedding memory --- EMBEDDINGS.md | 13 +- README.md | 9 + src/cocoindex_code/daemon.py | 7 + src/cocoindex_code/mps_embedder.py | 541 ++++++++++++++++++++++++++++ src/cocoindex_code/project.py | 54 ++- src/cocoindex_code/settings.py | 72 ++++ src/cocoindex_code/shared.py | 48 ++- tests/test_mps_embedder.py | 178 +++++++++ tests/test_project_indexing_gate.py | 146 ++++++++ tests/test_settings.py | 59 +++ tests/test_shared.py | 47 +++ 11 files changed, 1147 insertions(+), 27 deletions(-) create mode 100644 src/cocoindex_code/mps_embedder.py create mode 100644 tests/test_mps_embedder.py create mode 100644 tests/test_project_indexing_gate.py diff --git a/EMBEDDINGS.md b/EMBEDDINGS.md index ba87b0a..b5465b3 100644 --- a/EMBEDDINGS.md +++ b/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,15 @@ 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, the default MPS policy uses inner batches of 8, asks PyTorch +to begin allocator cleanup at 40% of its recommended memory, and enforces a +hard allocator limit at 50%. The embedding worker is recycled if driver-owned +memory remains above 35% after cache cleanup. These values can be tuned under +`embedding` with `batch_size`, `mps_memory_limit_ratio`, +`mps_low_watermark_ratio`, `mps_high_watermark_ratio`, and +`worker_timeout_seconds`. Keep the recycle ratio below the low watermark and +the low watermark at or below the high watermark. + --- ## The `ccc init` Wizard diff --git a/README.md b/README.md index 6594c82..2f49aa5 100644 --- a/README.md +++ b/README.md @@ -505,6 +505,11 @@ 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 + batch_size: 8 # optional: local MPS inner batch size (default 8) + mps_memory_limit_ratio: 0.35 # recycle the MPS worker above 35% of PyTorch's recommended memory + mps_low_watermark_ratio: 0.4 # PyTorch allocator soft limit + mps_high_watermark_ratio: 0.5 # PyTorch allocator hard limit + worker_timeout_seconds: 300 # replace an unresponsive embedding worker # 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 +528,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 models run in a child process so Metal driver allocations can be reclaimed without restarting the daemon. The worker clears cached allocations at the configured memory threshold and is replaced if memory remains high. MPS out-of-memory errors retry with progressively smaller batches. Explicit `PYTORCH_MPS_LOW_WATERMARK_RATIO` or `PYTORCH_MPS_HIGH_WATERMARK_RATIO` environment variables take precedence over the YAML defaults. + +> **Indexing concurrency:** A daemon indexes one project at a time. Requests for other projects queue behind the active project, which prevents multiple local embedding workloads from competing for unified memory. Search requests continue to use the shared daemon and wait only when their project still needs its 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). diff --git a/src/cocoindex_code/daemon.py b/src/cocoindex_code/daemon.py index d2b38ab..4919897 100644 --- a/src/cocoindex_code/daemon.py +++ b/src/cocoindex_code/daemon.py @@ -135,6 +135,7 @@ class ProjectRegistry: _projects: dict[str, Project] _embedder: Embedder | None + _indexing_gate: asyncio.Lock indexing_params: dict[str, Any] query_params: dict[str, Any] @@ -146,6 +147,7 @@ def __init__( ) -> None: self._projects = {} self._embedder = embedder + self._indexing_gate = asyncio.Lock() self.indexing_params = dict(indexing_params) if indexing_params else {} self.query_params = dict(query_params) if query_params else {} @@ -165,6 +167,7 @@ async def get_project(self, project_root: str) -> Project: indexing_params=self.indexing_params, query_params=self.query_params, chunker_registry=chunker_registry, + indexing_gate=self._indexing_gate, ) self._projects[project_root] = project return self._projects[project_root] @@ -188,6 +191,10 @@ def close_all(self) -> None: for project in self._projects.values(): project.close() self._projects.clear() + if self._embedder is not None: + close = getattr(self._embedder, "close", None) + if callable(close): + close() gc.collect() def list_projects(self) -> list[DaemonProjectInfo]: diff --git a/src/cocoindex_code/mps_embedder.py b/src/cocoindex_code/mps_embedder.py new file mode 100644 index 0000000..9782de8 --- /dev/null +++ b/src/cocoindex_code/mps_embedder.py @@ -0,0 +1,541 @@ +"""SentenceTransformer embedding isolated behind a recyclable MPS worker.""" + +from __future__ import annotations + +import gc +import logging +import multiprocessing +import os +import threading +import traceback +from collections.abc import Callable +from dataclasses import dataclass +from multiprocessing.connection import Connection +from multiprocessing.process import BaseProcess +from typing import Any, Protocol, cast + +import cocoindex as coco +import numpy as np +from cocoindex.resources import schema +from numpy.typing import NDArray + +logger = logging.getLogger(__name__) + +_MPS_DEFAULT_BATCH_SIZE = 8 +_WORKER_JOIN_TIMEOUT_SECONDS = 2.0 + + +@dataclass(frozen=True) +class MPSWorkerConfig: + """Serializable configuration owned by the embedding worker process.""" + + model_name: str + device: str | None + trust_remote_code: bool + batch_size: int + memory_limit_ratio: float + + +class _SentenceTransformerModel(Protocol): + @property + def device(self) -> Any: ... + + def encode( + self, + texts: list[str], + *, + batch_size: int, + prompt_name: str | None, + convert_to_numpy: bool, + normalize_embeddings: bool, + show_progress_bar: bool, + ) -> NDArray[np.float32]: ... + + def get_sentence_embedding_dimension(self) -> int | None: ... + + +WorkerTarget = Callable[[Connection, MPSWorkerConfig], None] + + +class EmbeddingWorkerError(RuntimeError): + """Base error for failures at the embedding process boundary.""" + + +class EmbeddingWorkerTimeoutError(EmbeddingWorkerError): + """The worker did not answer before its configured deadline.""" + + +class EmbeddingWorkerCrashedError(EmbeddingWorkerError): + """The worker exited or broke its IPC channel before answering.""" + + +class EmbeddingWorkerRemoteError(EmbeddingWorkerError): + """The worker answered with an embedding/model error.""" + + +def apply_mps_allocator_guards(*, low_watermark_ratio: float, high_watermark_ratio: float) -> None: + """Set conservative PyTorch MPS allocator defaults without overriding user env. + + PyTorch reads these values from the process environment. The daemon calls + this before importing ``torch`` or ``sentence_transformers``; spawned + workers then inherit the same bounds. + """ + + os.environ.setdefault("PYTORCH_MPS_LOW_WATERMARK_RATIO", str(low_watermark_ratio)) + os.environ.setdefault("PYTORCH_MPS_HIGH_WATERMARK_RATIO", str(high_watermark_ratio)) + + +def _is_oom_error(exc: RuntimeError) -> bool: + message = str(exc).lower() + return "out of memory" in message or "mps backend oom" in message + + +def _encode_with_oom_backoff( + model: _SentenceTransformerModel, + texts: list[str], + *, + batch_size: int, + prompt_name: str | None, + normalize_embeddings: bool, + empty_cache: Callable[[], None], +) -> NDArray[np.float32]: + """Encode, halving the inner batch after recoverable MPS OOM failures.""" + + current_batch_size = max(1, batch_size) + while True: + try: + return model.encode( + texts, + batch_size=current_batch_size, + prompt_name=prompt_name, + convert_to_numpy=True, + normalize_embeddings=normalize_embeddings, + show_progress_bar=False, + ) + except RuntimeError as exc: + if not _is_oom_error(exc) or current_batch_size == 1: + raise + empty_cache() + next_batch_size = max(1, current_batch_size // 2) + logger.warning( + "MPS embedding OOM at batch size %d; retrying with %d", + current_batch_size, + next_batch_size, + ) + current_batch_size = next_batch_size + + +def _should_recycle_mps_worker( + *, + driver_memory: int, + recommended_memory: int, + memory_limit_ratio: float, +) -> bool: + """Return whether driver-owned bytes crossed the worker recycle threshold.""" + + return recommended_memory > 0 and driver_memory >= recommended_memory * memory_limit_ratio + + +def _mps_memory_after_cache_cleanup( + model: _SentenceTransformerModel, + *, + memory_limit_ratio: float, +) -> tuple[bool, int | None, int | None]: + """Inspect MPS driver memory and clear cache only when near the recycle bound.""" + + if str(model.device).split(":", 1)[0] != "mps": + return False, None, None + + import torch + + torch.mps.synchronize() + driver_memory = int(torch.mps.driver_allocated_memory()) + recommended_memory = int(torch.mps.recommended_max_memory()) + if not _should_recycle_mps_worker( + driver_memory=driver_memory, + recommended_memory=recommended_memory, + memory_limit_ratio=memory_limit_ratio, + ): + return False, driver_memory, recommended_memory + + gc.collect() + torch.mps.empty_cache() + torch.mps.synchronize() + driver_memory = int(torch.mps.driver_allocated_memory()) + return ( + _should_recycle_mps_worker( + driver_memory=driver_memory, + recommended_memory=recommended_memory, + memory_limit_ratio=memory_limit_ratio, + ), + driver_memory, + recommended_memory, + ) + + +def _sentence_transformer_worker_main(conn: Connection, config: MPSWorkerConfig) -> None: + """Own the model and all Metal allocations until recycling or shutdown.""" + + try: + from sentence_transformers import SentenceTransformer + + model = cast( + _SentenceTransformerModel, + SentenceTransformer( + config.model_name, + device=config.device, + trust_remote_code=config.trust_remote_code, + ), + ) + + while True: + request = cast(dict[str, Any], conn.recv()) + operation = request.get("operation") + if operation == "close": + return + + payload = request.get("payload") + if operation == "dimension": + result: Any = model.get_sentence_embedding_dimension() + if result is None: + raise RuntimeError( + f"Embedding dimension is unknown for model {config.model_name}." + ) + result = int(result) + elif operation == "embed": + if not isinstance(payload, dict): + raise TypeError("embed worker payload must be a mapping") + result = _encode_with_oom_backoff( + model, + list(payload["texts"]), + batch_size=int(payload.get("batch_size", config.batch_size)), + prompt_name=payload.get("prompt_name"), + normalize_embeddings=bool(payload.get("normalize_embeddings", True)), + empty_cache=_empty_mps_cache, + ) + else: + raise ValueError(f"Unknown embedding worker operation: {operation!r}") + + try: + recycle, driver_memory, recommended_memory = _mps_memory_after_cache_cleanup( + model, + memory_limit_ratio=config.memory_limit_ratio, + ) + except Exception: + # The allocator watermarks remain the hard safety boundary when + # a PyTorch build cannot report one of the advisory metrics. + logger.exception("Unable to inspect MPS worker memory") + recycle, driver_memory, recommended_memory = False, None, None + + conn.send( + { + "ok": True, + "result": result, + "recycle": recycle, + "driver_memory": driver_memory, + "recommended_memory": recommended_memory, + } + ) + if recycle: + return + except EOFError: + return + except BaseException as exc: + try: + conn.send( + { + "ok": False, + "error_type": type(exc).__name__, + "error": str(exc), + "traceback": traceback.format_exc(), + "recycle": True, + } + ) + except (BrokenPipeError, EOFError, OSError): + pass + finally: + conn.close() + + +def _empty_mps_cache() -> None: + import torch + + gc.collect() + torch.mps.empty_cache() + + +class EmbeddingWorkerClient: + """Synchronous, thread-safe owner of one disposable worker process.""" + + def __init__( + self, + config: MPSWorkerConfig, + *, + timeout_seconds: float, + worker_target: WorkerTarget = _sentence_transformer_worker_main, + ) -> None: + self._config = config + self._timeout_seconds = timeout_seconds + self._worker_target = worker_target + self._context = multiprocessing.get_context("spawn") + self._lock = threading.Lock() + self._process: BaseProcess | None = None + self._conn: Connection | None = None + self._generation = 0 + + @property + def generation(self) -> int: + return self._generation + + @property + def worker_alive(self) -> bool: + with self._lock: + return self._process is not None and self._process.is_alive() + + def _start_locked(self) -> None: + if self._process is not None and self._process.is_alive(): + return + self._discard_locked(graceful=False) + parent_conn, child_conn = self._context.Pipe(duplex=True) + process = self._context.Process( + target=self._worker_target, + args=(child_conn, self._config), + name="cocoindex-mps-embedder", + daemon=True, + ) + process.start() + child_conn.close() + self._process = process + self._conn = parent_conn + self._generation += 1 + logger.info( + "Started MPS embedding worker PID %s (generation %d)", + process.pid, + self._generation, + ) + + def request(self, operation: str, payload: Any) -> Any: + with self._lock: + self._start_locked() + assert self._conn is not None + try: + self._conn.send({"operation": operation, "payload": payload}) + if not self._conn.poll(self._timeout_seconds): + self._discard_locked(graceful=False, force=True) + raise EmbeddingWorkerTimeoutError( + f"Embedding worker exceeded {self._timeout_seconds:g}s timeout" + ) + response = cast(dict[str, Any], self._conn.recv()) + except EmbeddingWorkerTimeoutError: + raise + except (BrokenPipeError, EOFError, OSError) as exc: + self._discard_locked(graceful=False, force=True) + raise EmbeddingWorkerCrashedError( + f"Embedding worker exited before answering: {exc}" + ) from exc + + if not response.get("ok"): + self._discard_locked(graceful=False) + error_type = response.get("error_type", "Exception") + message = response.get("error", "unknown worker error") + remote_traceback = response.get("traceback", "") + raise EmbeddingWorkerRemoteError( + f"{error_type}: {message}\nWorker traceback:\n{remote_traceback}" + ) + + if response.get("recycle"): + driver_memory = response.get("driver_memory") + recommended_memory = response.get("recommended_memory") + logger.warning( + "Recycling MPS embedding worker after driver memory reached %s/%s bytes", + driver_memory, + recommended_memory, + ) + self._discard_locked(graceful=False) + + return response.get("result") + + def _discard_locked(self, *, graceful: bool, force: bool = False) -> None: + process = self._process + conn = self._conn + self._process = None + self._conn = None + + if graceful and conn is not None and process is not None and process.is_alive(): + try: + conn.send({"operation": "close", "payload": None}) + except (BrokenPipeError, EOFError, OSError): + pass + if conn is not None: + conn.close() + + if process is None: + return + if force and process.is_alive(): + process.terminate() + process.join(timeout=_WORKER_JOIN_TIMEOUT_SECONDS) + if process.is_alive(): + process.terminate() + process.join(timeout=_WORKER_JOIN_TIMEOUT_SECONDS) + try: + process.close() + except ValueError: + pass + + def close(self) -> None: + if self._lock.acquire(timeout=0.1): + try: + self._discard_locked(graceful=True) + finally: + self._lock.release() + return + + # A request can hold the lock while blocked in a Metal kernel or IPC + # wait. Killing the child is safe without mutating the connection + # fields; the request thread observes EOF and performs normal cleanup. + process = self._process + if process is not None and process.is_alive(): + logger.warning("Force-terminating busy MPS embedding worker during shutdown") + process.terminate() + process.join(timeout=_WORKER_JOIN_TIMEOUT_SECONDS) + + +class MPSWorkerSentenceTransformerEmbedder(schema.VectorSchemaProvider): + """SentenceTransformer embedder whose Metal allocations live in a child process.""" + + def __init__( + self, + model_name_or_path: str, + *, + device: str | None = "mps", + trust_remote_code: bool = False, + batch_size: int = _MPS_DEFAULT_BATCH_SIZE, + memory_limit_ratio: float = 0.35, + worker_timeout_seconds: float = 300.0, + ) -> None: + self._model_name_or_path = model_name_or_path + self._device = device + self._trust_remote_code = trust_remote_code + self._batch_size = batch_size + self._memory_limit_ratio = memory_limit_ratio + self._worker_timeout_seconds = worker_timeout_seconds + self._worker: EmbeddingWorkerClient | None = None + self._worker_lock = threading.Lock() + + @property + def batch_size(self) -> int: + return self._batch_size + + @property + def memory_limit_ratio(self) -> float: + return self._memory_limit_ratio + + @property + def worker_timeout_seconds(self) -> float: + return self._worker_timeout_seconds + + def __getstate__(self) -> dict[str, Any]: + return { + "model_name_or_path": self._model_name_or_path, + "device": self._device, + "trust_remote_code": self._trust_remote_code, + "batch_size": self._batch_size, + "memory_limit_ratio": self._memory_limit_ratio, + "worker_timeout_seconds": self._worker_timeout_seconds, + } + + def __setstate__(self, state: dict[str, Any]) -> None: + self._model_name_or_path = state["model_name_or_path"] + self._device = state["device"] + self._trust_remote_code = state["trust_remote_code"] + self._batch_size = state["batch_size"] + self._memory_limit_ratio = state["memory_limit_ratio"] + self._worker_timeout_seconds = state["worker_timeout_seconds"] + self._worker = None + self._worker_lock = threading.Lock() + + def _get_worker(self) -> EmbeddingWorkerClient: + if self._worker is None: + with self._worker_lock: + if self._worker is None: + self._worker = EmbeddingWorkerClient( + MPSWorkerConfig( + model_name=self._model_name_or_path, + device=self._device, + trust_remote_code=self._trust_remote_code, + batch_size=self._batch_size, + memory_limit_ratio=self._memory_limit_ratio, + ), + timeout_seconds=self._worker_timeout_seconds, + ) + return self._worker + + def _request_with_crash_retry(self, operation: str, payload: Any) -> Any: + for attempt in range(2): + try: + return self._get_worker().request(operation, payload) + except (EmbeddingWorkerCrashedError, EmbeddingWorkerTimeoutError): + if attempt == 1: + raise + logger.warning("Restarting embedding worker after %s failure", operation) + if operation == "embed" and isinstance(payload, dict): + payload = dict(payload) + payload["batch_size"] = max(1, int(payload["batch_size"]) // 2) + raise AssertionError("unreachable") + + @coco.fn.as_async(batching=True, runner=coco.GPU, max_batch_size=64) + def _embed( + self, + texts: list[str], + prompt_name: str | None = None, + normalize_embeddings: bool = True, + ) -> list[NDArray[np.float32]]: + result = cast( + NDArray[np.float32], + self._request_with_crash_retry( + "embed", + { + "texts": texts, + "batch_size": self._batch_size, + "prompt_name": prompt_name, + "normalize_embeddings": normalize_embeddings, + }, + ), + ) + return list(result) + + @coco.fn(memo=True, version=2, logic_tracking="self") + async def embed( + self, + text: str, + prompt_name: str | None = None, + normalize_embeddings: bool = True, + ) -> NDArray[np.float32]: + result: NDArray[np.float32] = await self._embed( + text, + prompt_name, + normalize_embeddings, + ) + return result + + @coco.fn.as_async(runner=coco.GPU, memo=True) + def dimension(self) -> int: + return int(self._request_with_crash_retry("dimension", None)) + + async def __coco_vector_schema__(self) -> schema.VectorSchema: + return schema.VectorSchema(dtype=np.dtype(np.float32), size=await self.dimension()) + + def __coco_memo_key__(self) -> object: + return ( + self._model_name_or_path, + self._device, + self._trust_remote_code, + self._batch_size, + self._memory_limit_ratio, + ) + + def close(self) -> None: + worker = self._worker + if worker is not None: + worker.close() + self._worker = None diff --git a/src/cocoindex_code/project.py b/src/cocoindex_code/project.py index f661c21..c966b9a 100644 --- a/src/cocoindex_code/project.py +++ b/src/cocoindex_code/project.py @@ -47,7 +47,10 @@ class Project: _app: coco.App[[], None] _project_root: Path _index_lock: asyncio.Lock + _indexing_gate: asyncio.Lock _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,24 +70,26 @@ 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, queue on the daemon gate, run, and release. - If *on_started* is provided, it is set once the lock is acquired - (i.e. indexing has truly begun). On completion (success or failure) + If *on_started* is provided, it is set once the project lock is + acquired. At that point indexing is either active or queued behind + another project. On completion (success or failure) ``_initial_index_done`` is set. """ async with self._index_lock: - self._indexing_stats = IndexingProgress( - num_execution_starts=0, - num_unchanged=0, - num_adds=0, - num_deletes=0, - 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 with self._indexing_gate: + self._indexing_stats = IndexingProgress( + num_execution_starts=0, + num_unchanged=0, + num_adds=0, + num_deletes=0, + num_reprocesses=0, + num_errors=0, + ) + await self._run_index_inner(on_progress=on_progress) async def _run_index_inner( self, @@ -113,16 +118,21 @@ async def _run_index_inner( 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]: @@ -132,7 +142,7 @@ async def stream_index(self) -> AsyncIterator[IndexStreamResponse]: The actual indexing runs in a separate task so that client disconnects (``GeneratorExit``) do not abort the indexing. """ - if self._index_lock.locked(): + if self._index_lock.locked() or self._indexing_gate.locked(): yield IndexWaitingNotice() progress_queue: asyncio.Queue[IndexingProgress] = asyncio.Queue() @@ -263,6 +273,7 @@ async def create( indexing_params: dict[str, Any], query_params: dict[str, Any], chunker_registry: dict[str, ChunkerFn] | None = None, + indexing_gate: asyncio.Lock | None = None, ) -> Project: """Create a project with explicit embedder and per-call params. @@ -281,6 +292,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. + indexing_gate: Daemon-wide lock shared by every loaded project. + When omitted, the project gets a private lock for standalone use. """ settings_dir = project_root / ".cocoindex_code" settings_dir.mkdir(parents=True, exist_ok=True) @@ -315,5 +328,8 @@ async def create( result._app = app result._project_root = project_root result._index_lock = asyncio.Lock() + result._indexing_gate = indexing_gate if indexing_gate is not None else asyncio.Lock() result._initial_index_done = asyncio.Event() + result._initial_index_task = None + result._initial_index_started = None return result diff --git a/src/cocoindex_code/settings.py b/src/cocoindex_code/settings.py index e228839..93b97ef 100644 --- a/src/cocoindex_code/settings.py +++ b/src/cocoindex_code/settings.py @@ -98,12 +98,51 @@ class EmbeddingSettings: provider: str = "litellm" device: str | None = None min_interval_ms: int | None = None + # SentenceTransformer safety controls. ``batch_size=None`` resolves to 8 + # on MPS and keeps the upstream default on other devices. + batch_size: int | None = None + # Recycle the isolated MPS embedding worker before allocator pressure + # reaches PyTorch's soft/hard watermarks. + mps_memory_limit_ratio: float = 0.35 + mps_low_watermark_ratio: float = 0.4 + mps_high_watermark_ratio: float = 0.5 + worker_timeout_seconds: float = 300.0 # 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: + if self.batch_size is not None and ( + not isinstance(self.batch_size, int) + or isinstance(self.batch_size, bool) + or not 1 <= self.batch_size <= 64 + ): + raise ValueError("embedding.batch_size must be an integer between 1 and 64") + + ratios = ( + ("mps_memory_limit_ratio", self.mps_memory_limit_ratio), + ("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_memory_limit_ratio >= self.mps_low_watermark_ratio: + raise ValueError( + "embedding.mps_memory_limit_ratio must be lower than " + "embedding.mps_low_watermark_ratio" + ) + 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" + ) + if isinstance(self.worker_timeout_seconds, bool) or self.worker_timeout_seconds <= 0: + raise ValueError("embedding.worker_timeout_seconds must be greater than 0") + @dataclass class DaemonSettings: @@ -430,6 +469,17 @@ 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 + if embedding.batch_size is not None: + d["batch_size"] = embedding.batch_size + defaults = EmbeddingSettings(model=embedding.model) + if embedding.mps_memory_limit_ratio != defaults.mps_memory_limit_ratio: + d["mps_memory_limit_ratio"] = embedding.mps_memory_limit_ratio + 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.worker_timeout_seconds != defaults.worker_timeout_seconds: + d["worker_timeout_seconds"] = embedding.worker_timeout_seconds if embedding.indexing_params is not None: d["indexing_params"] = dict(embedding.indexing_params) if embedding.query_params is not None: @@ -458,6 +508,16 @@ 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 "batch_size" in emb_dict: + emb_kwargs["batch_size"] = emb_dict["batch_size"] + if "mps_memory_limit_ratio" in emb_dict: + emb_kwargs["mps_memory_limit_ratio"] = float(emb_dict["mps_memory_limit_ratio"]) + 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"]) + if "worker_timeout_seconds" in emb_dict: + emb_kwargs["worker_timeout_seconds"] = float(emb_dict["worker_timeout_seconds"]) # 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 +633,16 @@ def save_user_settings(settings: UserSettings) -> Path: ), } +_MPS_SAFETY_COMMENT = ( + " #\n" + " # Apple Silicon MPS safety defaults. The model runs in a recyclable worker.\n" + " # batch_size: 8\n" + " # mps_memory_limit_ratio: 0.35\n" + " # mps_low_watermark_ratio: 0.4\n" + " # mps_high_watermark_ratio: 0.5\n" + " # worker_timeout_seconds: 300\n" +) + def save_initial_user_settings( embedding: EmbeddingSettings, @@ -595,6 +665,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: diff --git a/src/cocoindex_code/shared.py b/src/cocoindex_code/shared.py index f607755..67acbef 100644 --- a/src/cocoindex_code/shared.py +++ b/src/cocoindex_code/shared.py @@ -5,6 +5,7 @@ import importlib.util import logging import pathlib +import sys import traceback as _tb from dataclasses import dataclass from typing import TYPE_CHECKING, Annotated, Any, NamedTuple, Union @@ -18,6 +19,8 @@ from cocoindex.ops.litellm import LiteLLMEmbedder from cocoindex.ops.sentence_transformers import SentenceTransformerEmbedder + from .mps_embedder import MPSWorkerSentenceTransformerEmbedder + from .settings import EmbeddingSettings logger = logging.getLogger(__name__) @@ -26,7 +29,11 @@ DEFAULT_LITELLM_MIN_INTERVAL_MS = 5 # Type alias -Embedder = Union["SentenceTransformerEmbedder", "LiteLLMEmbedder"] +Embedder = Union[ + "SentenceTransformerEmbedder", + "MPSWorkerSentenceTransformerEmbedder", + "LiteLLMEmbedder", +] # Context keys EMBEDDER = coco.ContextKey[Embedder]("embedder", detect_change=True) @@ -100,6 +107,7 @@ def create_embedder( only and the indexing default is supplied at the call site via :data:`INDEXING_EMBED_PARAMS`. """ + instance: Embedder if settings.provider == "sentence-transformers": from cocoindex.ops.sentence_transformers import SentenceTransformerEmbedder @@ -108,12 +116,40 @@ def create_embedder( if model_name.startswith(SBERT_PREFIX): model_name = model_name[len(SBERT_PREFIX) :] - instance: Embedder = SentenceTransformerEmbedder( - model_name, - device=settings.device, - trust_remote_code=True, + use_mps_worker = settings.device == "mps" or ( + settings.device is None and sys.platform == "darwin" ) - logger.info("Embedding model: %s | device: %s", settings.model, settings.device) + if use_mps_worker: + from .mps_embedder import ( + MPSWorkerSentenceTransformerEmbedder, + apply_mps_allocator_guards, + ) + + apply_mps_allocator_guards( + low_watermark_ratio=settings.mps_low_watermark_ratio, + high_watermark_ratio=settings.mps_high_watermark_ratio, + ) + instance = MPSWorkerSentenceTransformerEmbedder( + model_name, + device=settings.device, + trust_remote_code=True, + batch_size=settings.batch_size or 8, + memory_limit_ratio=settings.mps_memory_limit_ratio, + worker_timeout_seconds=settings.worker_timeout_seconds, + ) + logger.info( + "Embedding model: %s | device: %s | isolated MPS worker | batch size: %d", + settings.model, + settings.device, + settings.batch_size or 8, + ) + else: + instance = SentenceTransformerEmbedder( + model_name, + device=settings.device, + trust_remote_code=True, + ) + logger.info("Embedding model: %s | device: %s", settings.model, settings.device) else: from .litellm_embedder import PacedLiteLLMEmbedder diff --git a/tests/test_mps_embedder.py b/tests/test_mps_embedder.py new file mode 100644 index 0000000..96b188b --- /dev/null +++ b/tests/test_mps_embedder.py @@ -0,0 +1,178 @@ +"""MPS embedding memory and worker-lifecycle safety tests.""" + +from __future__ import annotations + +import threading +import time +from multiprocessing.connection import Connection +from typing import Any + +import numpy as np +import pytest + +from cocoindex_code.mps_embedder import ( + EmbeddingWorkerClient, + EmbeddingWorkerCrashedError, + EmbeddingWorkerTimeoutError, + MPSWorkerConfig, + _encode_with_oom_backoff, + _should_recycle_mps_worker, +) + + +class _OomUntilSmallBatchModel: + def __init__(self) -> None: + self.batch_sizes: list[int] = [] + + def encode(self, texts: list[str], *, batch_size: int, **kwargs: Any) -> np.ndarray: + self.batch_sizes.append(batch_size) + if batch_size > 2: + raise RuntimeError("MPS backend out of memory") + return np.zeros((len(texts), 3), dtype=np.float32) + + +def _recycling_worker(conn: Connection, config: MPSWorkerConfig) -> None: + del config + request = conn.recv() + conn.send( + { + "ok": True, + "result": request["payload"], + "recycle": True, + "driver_memory": 100, + "recommended_memory": 200, + } + ) + conn.close() + + +def _hanging_worker(conn: Connection, config: MPSWorkerConfig) -> None: + del config + conn.recv() + time.sleep(10) + + +def _worker_config() -> MPSWorkerConfig: + return MPSWorkerConfig( + model_name="unused-in-boundary-test", + device="mps", + trust_remote_code=False, + batch_size=8, + memory_limit_ratio=0.35, + ) + + +def test_mps_oom_retries_with_smaller_inner_batches() -> None: + model = _OomUntilSmallBatchModel() + empty_cache_calls = 0 + + def empty_cache() -> None: + nonlocal empty_cache_calls + empty_cache_calls += 1 + + result = _encode_with_oom_backoff( + model, + ["a", "b", "c"], + batch_size=8, + prompt_name=None, + normalize_embeddings=True, + empty_cache=empty_cache, + ) + + assert result.shape == (3, 3) + assert model.batch_sizes == [8, 4, 2] + assert empty_cache_calls == 2 + + +def test_non_oom_embedding_error_is_not_retried() -> None: + class BrokenModel: + def encode(self, texts: list[str], *, batch_size: int, **kwargs: Any) -> np.ndarray: + raise RuntimeError("invalid prompt") + + with pytest.raises(RuntimeError, match="invalid prompt"): + _encode_with_oom_backoff( + BrokenModel(), + ["a"], + batch_size=8, + prompt_name=None, + normalize_embeddings=True, + empty_cache=lambda: None, + ) + + +def test_worker_recycles_only_after_driver_memory_crosses_limit() -> None: + assert ( + _should_recycle_mps_worker( + driver_memory=349, + recommended_memory=1000, + memory_limit_ratio=0.35, + ) + is False + ) + assert ( + _should_recycle_mps_worker( + driver_memory=350, + recommended_memory=1000, + memory_limit_ratio=0.35, + ) + is True + ) + + +def test_worker_recycle_response_starts_a_fresh_process_for_next_request() -> None: + client = EmbeddingWorkerClient( + _worker_config(), + timeout_seconds=2, + worker_target=_recycling_worker, + ) + try: + assert client.request("echo", [1]) == [1] + assert client.generation == 1 + assert client.request("echo", [2]) == [2] + assert client.generation == 2 + finally: + client.close() + + +def test_worker_timeout_terminates_stuck_process() -> None: + client = EmbeddingWorkerClient( + _worker_config(), + timeout_seconds=0.1, + worker_target=_hanging_worker, + ) + try: + with pytest.raises(EmbeddingWorkerTimeoutError, match="0.1"): + client.request("embed", ["stuck"]) + assert client.worker_alive is False + finally: + client.close() + + +def test_close_terminates_worker_even_while_request_is_stuck() -> None: + client = EmbeddingWorkerClient( + _worker_config(), + timeout_seconds=2, + worker_target=_hanging_worker, + ) + request_finished = threading.Event() + + def request() -> None: + try: + client.request("embed", ["stuck"]) + except EmbeddingWorkerCrashedError: + pass + finally: + request_finished.set() + + thread = threading.Thread(target=request) + thread.start() + time.sleep(0.3) + + started = time.monotonic() + client.close() + elapsed = time.monotonic() - started + + thread.join(timeout=2) + assert elapsed < 1 + assert request_finished.is_set() is True + assert client.worker_alive is False diff --git a/tests/test_project_indexing_gate.py b/tests/test_project_indexing_gate.py new file mode 100644 index 0000000..295cb67 --- /dev/null +++ b/tests/test_project_indexing_gate.py @@ -0,0 +1,146 @@ +"""Cross-project indexing concurrency tests.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Callable + +from cocoindex_code.project import Project + + +class _ControlledProject(Project): # type: ignore[misc] + def __init__( + self, + *, + indexing_gate: asyncio.Lock, + on_enter: Callable[[], None], + on_exit: Callable[[], None], + release: asyncio.Event | None = None, + ) -> None: + self._index_lock = asyncio.Lock() + self._indexing_gate = indexing_gate + self._initial_index_done = asyncio.Event() + self._initial_index_task = None + self._initial_index_started = None + self._indexing_stats = None + self._on_enter = on_enter + self._on_exit = on_exit + self._release = release + + async def _run_index_inner( + self, + on_progress: Callable[..., None] | None = None, + ) -> None: + del on_progress + try: + self._on_enter() + if self._release is not None: + await self._release.wait() + finally: + self._on_exit() + self._initial_index_done.set() + self._indexing_stats = None + + +async def test_projects_share_one_indexing_gate() -> None: + indexing_gate = asyncio.Lock() + release_first = asyncio.Event() + second_entered = asyncio.Event() + active = 0 + max_active = 0 + + def enter_first() -> None: + nonlocal active, max_active + active += 1 + max_active = max(max_active, active) + + def enter_second() -> None: + nonlocal active, max_active + active += 1 + max_active = max(max_active, active) + second_entered.set() + + def exit() -> None: + nonlocal active + active -= 1 + + first = _ControlledProject( + indexing_gate=indexing_gate, + on_enter=enter_first, + on_exit=exit, + release=release_first, + ) + second = _ControlledProject( + indexing_gate=indexing_gate, + on_enter=enter_second, + on_exit=exit, + ) + + first_task = asyncio.create_task(first.run_index()) + while active == 0: + await asyncio.sleep(0) + + second_task = asyncio.create_task(second.run_index()) + await asyncio.sleep(0) + + assert second._index_lock.locked() is True + assert second_entered.is_set() is False + assert active == 1 + + release_first.set() + await first_task + await second_entered.wait() + await second_task + + assert max_active == 1 + assert active == 0 + + +async def test_initial_index_reports_started_when_queued_behind_another_project() -> None: + indexing_gate = asyncio.Lock() + await indexing_gate.acquire() + project = _ControlledProject( + indexing_gate=indexing_gate, + on_enter=lambda: None, + on_exit=lambda: None, + ) + + ensure_task = asyncio.create_task(project.ensure_indexing_started()) + + try: + await asyncio.wait_for(asyncio.shield(ensure_task), timeout=0.5) + assert project._index_lock.locked() is True + assert ensure_task.done() is True + assert project.should_wait_for_indexing is True + finally: + indexing_gate.release() + + await ensure_task + await project.wait_for_indexing_done() + + +async def test_concurrent_initial_index_requests_share_one_background_task() -> None: + indexing_gate = asyncio.Lock() + await indexing_gate.acquire() + execution_count = 0 + + def enter() -> None: + nonlocal execution_count + execution_count += 1 + + project = _ControlledProject( + indexing_gate=indexing_gate, + on_enter=enter, + on_exit=lambda: None, + ) + first = asyncio.create_task(project.ensure_indexing_started()) + second = asyncio.create_task(project.ensure_indexing_started()) + + try: + await asyncio.wait_for(asyncio.gather(first, second), timeout=0.5) + assert project._index_lock.locked() is True + finally: + indexing_gate.release() + + await project.wait_for_indexing_done() + assert execution_count == 1 diff --git a/tests/test_settings.py b/tests/test_settings.py index f9be64b..3307d81 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -55,6 +55,11 @@ def test_default_user_settings() -> None: assert s.embedding.model == "Snowflake/snowflake-arctic-embed-xs" assert s.embedding.device is None assert s.embedding.min_interval_ms is None + assert s.embedding.batch_size is None + assert s.embedding.mps_memory_limit_ratio == 0.35 + assert s.embedding.mps_low_watermark_ratio == 0.4 + assert s.embedding.mps_high_watermark_ratio == 0.5 + assert s.embedding.worker_timeout_seconds == 300.0 assert s.envs == {} @@ -77,6 +82,11 @@ def test_save_and_load_user_settings(tmp_path: Path) -> None: model="gemini/text-embedding-004", device="cpu", min_interval_ms=300, + batch_size=8, + mps_memory_limit_ratio=0.3, + mps_low_watermark_ratio=0.35, + mps_high_watermark_ratio=0.45, + worker_timeout_seconds=120.0, ), envs={"GEMINI_API_KEY": "test-key"}, ) @@ -86,9 +96,53 @@ def test_save_and_load_user_settings(tmp_path: Path) -> None: assert loaded.embedding.model == settings.embedding.model assert loaded.embedding.device == settings.embedding.device assert loaded.embedding.min_interval_ms == settings.embedding.min_interval_ms + assert loaded.embedding.batch_size == settings.embedding.batch_size + assert loaded.embedding.mps_memory_limit_ratio == settings.embedding.mps_memory_limit_ratio + assert loaded.embedding.mps_low_watermark_ratio == settings.embedding.mps_low_watermark_ratio + assert loaded.embedding.mps_high_watermark_ratio == settings.embedding.mps_high_watermark_ratio + assert loaded.embedding.worker_timeout_seconds == settings.embedding.worker_timeout_seconds assert loaded.envs == settings.envs +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("batch_size", 0, "batch_size"), + ("batch_size", 3.5, "batch_size"), + ("batch_size", 65, "batch_size"), + ("mps_memory_limit_ratio", 0, "mps_memory_limit_ratio"), + ("mps_low_watermark_ratio", 0, "mps_low_watermark_ratio"), + ("mps_high_watermark_ratio", 1.1, "mps_high_watermark_ratio"), + ("worker_timeout_seconds", 0, "worker_timeout_seconds"), + ], +) +def test_embedding_safety_settings_reject_invalid_values( + field: str, + value: int | float, + message: str, +) -> None: + with pytest.raises(ValueError, match=message): + EmbeddingSettings(model="model", **{field: value}) + + +def test_embedding_safety_settings_require_ordered_mps_limits() -> None: + with pytest.raises(ValueError, match="mps_memory_limit_ratio"): + EmbeddingSettings( + model="model", + mps_memory_limit_ratio=0.45, + mps_low_watermark_ratio=0.4, + mps_high_watermark_ratio=0.5, + ) + + with pytest.raises(ValueError, match="mps_low_watermark_ratio"): + EmbeddingSettings( + model="model", + mps_memory_limit_ratio=0.3, + mps_low_watermark_ratio=0.6, + mps_high_watermark_ratio=0.5, + ) + + def test_save_and_load_project_settings(tmp_path: Path) -> None: settings = ProjectSettings( include_patterns=["**/*.py", "**/*.rs"], @@ -361,6 +415,11 @@ def test_save_initial_user_settings_round_trip() -> None: # Hint comment and the four commented env-var examples. assert "ccc doctor" in content + assert "# batch_size: 8" in content + assert "# mps_memory_limit_ratio: 0.35" in content + assert "# mps_low_watermark_ratio: 0.4" in content + assert "# mps_high_watermark_ratio: 0.5" in content + assert "# worker_timeout_seconds: 300" in content assert "# envs:" in content for key in ("OPENAI_API_KEY", "GEMINI_API_KEY", "ANTHROPIC_API_KEY", "VOYAGE_API_KEY"): assert f"# {key}:" in content diff --git a/tests/test_shared.py b/tests/test_shared.py index 2af86f0..1248016 100644 --- a/tests/test_shared.py +++ b/tests/test_shared.py @@ -2,12 +2,14 @@ from __future__ import annotations +import os from typing import Any import numpy as np import pytest from cocoindex_code.litellm_embedder import PacedLiteLLMEmbedder +from cocoindex_code.mps_embedder import MPSWorkerSentenceTransformerEmbedder from cocoindex_code.settings import EmbeddingSettings from cocoindex_code.shared import ( check_embedding, @@ -66,6 +68,51 @@ def test_create_embedder_sentence_transformers_ignores_indexing_params() -> None assert not isinstance(embedder, PacedLiteLLMEmbedder) +def test_create_embedder_uses_isolated_worker_and_allocator_guards_for_mps( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("PYTORCH_MPS_HIGH_WATERMARK_RATIO", raising=False) + monkeypatch.delenv("PYTORCH_MPS_LOW_WATERMARK_RATIO", raising=False) + + embedder = create_embedder( + EmbeddingSettings( + provider="sentence-transformers", + model="sentence-transformers/all-MiniLM-L6-v2", + device="mps", + batch_size=4, + mps_memory_limit_ratio=0.3, + mps_low_watermark_ratio=0.4, + mps_high_watermark_ratio=0.5, + worker_timeout_seconds=90, + ) + ) + + assert isinstance(embedder, MPSWorkerSentenceTransformerEmbedder) + assert embedder.batch_size == 4 + assert embedder.memory_limit_ratio == 0.3 + assert embedder.worker_timeout_seconds == 90 + assert os.environ["PYTORCH_MPS_LOW_WATERMARK_RATIO"] == "0.4" + assert os.environ["PYTORCH_MPS_HIGH_WATERMARK_RATIO"] == "0.5" + + +def test_explicit_mps_allocator_env_takes_precedence( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("PYTORCH_MPS_LOW_WATERMARK_RATIO", "0.25") + monkeypatch.setenv("PYTORCH_MPS_HIGH_WATERMARK_RATIO", "0.3") + + create_embedder( + EmbeddingSettings( + provider="sentence-transformers", + model="sentence-transformers/all-MiniLM-L6-v2", + device="mps", + ) + ) + + assert os.environ["PYTORCH_MPS_LOW_WATERMARK_RATIO"] == "0.25" + assert os.environ["PYTORCH_MPS_HIGH_WATERMARK_RATIO"] == "0.3" + + def test_is_sentence_transformers_installed_true_in_dev() -> None: # Dev env pulls in sentence-transformers via the `dev` extras group. assert is_sentence_transformers_installed() is True From 861aa0809b41dd8abefa02e0ba5296dcbfba341d Mon Sep 17 00:00:00 2001 From: fml09 Date: Sat, 1 Aug 2026 21:29:41 +0900 Subject: [PATCH 2/2] refactor: use CocoIndex GPU runner for MPS memory safety --- EMBEDDINGS.md | 20 +- README.md | 11 +- pyproject.toml | 10 +- src/cocoindex_code/daemon.py | 16 +- src/cocoindex_code/mps_embedder.py | 541 ---------------------------- src/cocoindex_code/project.py | 53 +-- src/cocoindex_code/settings.py | 40 +- src/cocoindex_code/shared.py | 82 +++-- tests/test_mps_embedder.py | 178 --------- tests/test_project_indexing.py | 180 +++++++++ tests/test_project_indexing_gate.py | 146 -------- tests/test_settings.py | 47 +-- tests/test_shared.py | 69 +++- uv.lock | 82 ++--- 14 files changed, 405 insertions(+), 1070 deletions(-) delete mode 100644 src/cocoindex_code/mps_embedder.py delete mode 100644 tests/test_mps_embedder.py create mode 100644 tests/test_project_indexing.py delete mode 100644 tests/test_project_indexing_gate.py diff --git a/EMBEDDINGS.md b/EMBEDDINGS.md index b5465b3..71d7fb4 100644 --- a/EMBEDDINGS.md +++ b/EMBEDDINGS.md @@ -54,14 +54,18 @@ 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, the default MPS policy uses inner batches of 8, asks PyTorch -to begin allocator cleanup at 40% of its recommended memory, and enforces a -hard allocator limit at 50%. The embedding worker is recycled if driver-owned -memory remains above 35% after cache cleanup. These values can be tuned under -`embedding` with `batch_size`, `mps_memory_limit_ratio`, -`mps_low_watermark_ratio`, `mps_high_watermark_ratio`, and -`worker_timeout_seconds`. Keep the recycle ratio below the low watermark and -the low watermark at or below the high watermark. +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. --- diff --git a/README.md b/README.md index 2f49aa5..716b128 100644 --- a/README.md +++ b/README.md @@ -505,11 +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 - batch_size: 8 # optional: local MPS inner batch size (default 8) - mps_memory_limit_ratio: 0.35 # recycle the MPS worker above 35% of PyTorch's recommended memory - mps_low_watermark_ratio: 0.4 # PyTorch allocator soft limit - mps_high_watermark_ratio: 0.5 # PyTorch allocator hard limit - worker_timeout_seconds: 300 # replace an unresponsive embedding worker + 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, @@ -528,9 +525,9 @@ 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 models run in a child process so Metal driver allocations can be reclaimed without restarting the daemon. The worker clears cached allocations at the configured memory threshold and is replaced if memory remains high. MPS out-of-memory errors retry with progressively smaller batches. Explicit `PYTORCH_MPS_LOW_WATERMARK_RATIO` or `PYTORCH_MPS_HIGH_WATERMARK_RATIO` environment variables take precedence over the YAML defaults. +> **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:** A daemon indexes one project at a time. Requests for other projects queue behind the active project, which prevents multiple local embedding workloads from competing for unified memory. Search requests continue to use the shared daemon and wait only when their project still needs its initial index. +> **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. diff --git a/pyproject.toml b/pyproject.toml index 7e3bb6d..9c578b3 100644 --- a/pyproject.toml +++ b/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,7 +39,7 @@ 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 @@ -47,7 +47,7 @@ embeddings-local = ["cocoindex[sentence-transformers]>=1.0.13,<1.1.0"] # `: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", @@ -55,7 +55,7 @@ dev = [ "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] diff --git a/src/cocoindex_code/daemon.py b/src/cocoindex_code/daemon.py index 4919897..a291956 100644 --- a/src/cocoindex_code/daemon.py +++ b/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,7 +135,7 @@ class ProjectRegistry: _projects: dict[str, Project] _embedder: Embedder | None - _indexing_gate: asyncio.Lock + _clear_mps_cache_after_index: bool indexing_params: dict[str, Any] query_params: dict[str, Any] @@ -144,10 +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._indexing_gate = asyncio.Lock() + 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 {} @@ -167,7 +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, - indexing_gate=self._indexing_gate, + clear_mps_cache_after_index=self._clear_mps_cache_after_index, ) self._projects[project_root] = project return self._projects[project_root] @@ -191,10 +192,6 @@ def close_all(self) -> None: for project in self._projects.values(): project.close() self._projects.clear() - if self._embedder is not None: - close = getattr(self._embedder, "close", None) - if callable(close): - close() gc.collect() def list_projects(self) -> list[DaemonProjectInfo]: @@ -643,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(): @@ -651,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: @@ -691,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() diff --git a/src/cocoindex_code/mps_embedder.py b/src/cocoindex_code/mps_embedder.py deleted file mode 100644 index 9782de8..0000000 --- a/src/cocoindex_code/mps_embedder.py +++ /dev/null @@ -1,541 +0,0 @@ -"""SentenceTransformer embedding isolated behind a recyclable MPS worker.""" - -from __future__ import annotations - -import gc -import logging -import multiprocessing -import os -import threading -import traceback -from collections.abc import Callable -from dataclasses import dataclass -from multiprocessing.connection import Connection -from multiprocessing.process import BaseProcess -from typing import Any, Protocol, cast - -import cocoindex as coco -import numpy as np -from cocoindex.resources import schema -from numpy.typing import NDArray - -logger = logging.getLogger(__name__) - -_MPS_DEFAULT_BATCH_SIZE = 8 -_WORKER_JOIN_TIMEOUT_SECONDS = 2.0 - - -@dataclass(frozen=True) -class MPSWorkerConfig: - """Serializable configuration owned by the embedding worker process.""" - - model_name: str - device: str | None - trust_remote_code: bool - batch_size: int - memory_limit_ratio: float - - -class _SentenceTransformerModel(Protocol): - @property - def device(self) -> Any: ... - - def encode( - self, - texts: list[str], - *, - batch_size: int, - prompt_name: str | None, - convert_to_numpy: bool, - normalize_embeddings: bool, - show_progress_bar: bool, - ) -> NDArray[np.float32]: ... - - def get_sentence_embedding_dimension(self) -> int | None: ... - - -WorkerTarget = Callable[[Connection, MPSWorkerConfig], None] - - -class EmbeddingWorkerError(RuntimeError): - """Base error for failures at the embedding process boundary.""" - - -class EmbeddingWorkerTimeoutError(EmbeddingWorkerError): - """The worker did not answer before its configured deadline.""" - - -class EmbeddingWorkerCrashedError(EmbeddingWorkerError): - """The worker exited or broke its IPC channel before answering.""" - - -class EmbeddingWorkerRemoteError(EmbeddingWorkerError): - """The worker answered with an embedding/model error.""" - - -def apply_mps_allocator_guards(*, low_watermark_ratio: float, high_watermark_ratio: float) -> None: - """Set conservative PyTorch MPS allocator defaults without overriding user env. - - PyTorch reads these values from the process environment. The daemon calls - this before importing ``torch`` or ``sentence_transformers``; spawned - workers then inherit the same bounds. - """ - - os.environ.setdefault("PYTORCH_MPS_LOW_WATERMARK_RATIO", str(low_watermark_ratio)) - os.environ.setdefault("PYTORCH_MPS_HIGH_WATERMARK_RATIO", str(high_watermark_ratio)) - - -def _is_oom_error(exc: RuntimeError) -> bool: - message = str(exc).lower() - return "out of memory" in message or "mps backend oom" in message - - -def _encode_with_oom_backoff( - model: _SentenceTransformerModel, - texts: list[str], - *, - batch_size: int, - prompt_name: str | None, - normalize_embeddings: bool, - empty_cache: Callable[[], None], -) -> NDArray[np.float32]: - """Encode, halving the inner batch after recoverable MPS OOM failures.""" - - current_batch_size = max(1, batch_size) - while True: - try: - return model.encode( - texts, - batch_size=current_batch_size, - prompt_name=prompt_name, - convert_to_numpy=True, - normalize_embeddings=normalize_embeddings, - show_progress_bar=False, - ) - except RuntimeError as exc: - if not _is_oom_error(exc) or current_batch_size == 1: - raise - empty_cache() - next_batch_size = max(1, current_batch_size // 2) - logger.warning( - "MPS embedding OOM at batch size %d; retrying with %d", - current_batch_size, - next_batch_size, - ) - current_batch_size = next_batch_size - - -def _should_recycle_mps_worker( - *, - driver_memory: int, - recommended_memory: int, - memory_limit_ratio: float, -) -> bool: - """Return whether driver-owned bytes crossed the worker recycle threshold.""" - - return recommended_memory > 0 and driver_memory >= recommended_memory * memory_limit_ratio - - -def _mps_memory_after_cache_cleanup( - model: _SentenceTransformerModel, - *, - memory_limit_ratio: float, -) -> tuple[bool, int | None, int | None]: - """Inspect MPS driver memory and clear cache only when near the recycle bound.""" - - if str(model.device).split(":", 1)[0] != "mps": - return False, None, None - - import torch - - torch.mps.synchronize() - driver_memory = int(torch.mps.driver_allocated_memory()) - recommended_memory = int(torch.mps.recommended_max_memory()) - if not _should_recycle_mps_worker( - driver_memory=driver_memory, - recommended_memory=recommended_memory, - memory_limit_ratio=memory_limit_ratio, - ): - return False, driver_memory, recommended_memory - - gc.collect() - torch.mps.empty_cache() - torch.mps.synchronize() - driver_memory = int(torch.mps.driver_allocated_memory()) - return ( - _should_recycle_mps_worker( - driver_memory=driver_memory, - recommended_memory=recommended_memory, - memory_limit_ratio=memory_limit_ratio, - ), - driver_memory, - recommended_memory, - ) - - -def _sentence_transformer_worker_main(conn: Connection, config: MPSWorkerConfig) -> None: - """Own the model and all Metal allocations until recycling or shutdown.""" - - try: - from sentence_transformers import SentenceTransformer - - model = cast( - _SentenceTransformerModel, - SentenceTransformer( - config.model_name, - device=config.device, - trust_remote_code=config.trust_remote_code, - ), - ) - - while True: - request = cast(dict[str, Any], conn.recv()) - operation = request.get("operation") - if operation == "close": - return - - payload = request.get("payload") - if operation == "dimension": - result: Any = model.get_sentence_embedding_dimension() - if result is None: - raise RuntimeError( - f"Embedding dimension is unknown for model {config.model_name}." - ) - result = int(result) - elif operation == "embed": - if not isinstance(payload, dict): - raise TypeError("embed worker payload must be a mapping") - result = _encode_with_oom_backoff( - model, - list(payload["texts"]), - batch_size=int(payload.get("batch_size", config.batch_size)), - prompt_name=payload.get("prompt_name"), - normalize_embeddings=bool(payload.get("normalize_embeddings", True)), - empty_cache=_empty_mps_cache, - ) - else: - raise ValueError(f"Unknown embedding worker operation: {operation!r}") - - try: - recycle, driver_memory, recommended_memory = _mps_memory_after_cache_cleanup( - model, - memory_limit_ratio=config.memory_limit_ratio, - ) - except Exception: - # The allocator watermarks remain the hard safety boundary when - # a PyTorch build cannot report one of the advisory metrics. - logger.exception("Unable to inspect MPS worker memory") - recycle, driver_memory, recommended_memory = False, None, None - - conn.send( - { - "ok": True, - "result": result, - "recycle": recycle, - "driver_memory": driver_memory, - "recommended_memory": recommended_memory, - } - ) - if recycle: - return - except EOFError: - return - except BaseException as exc: - try: - conn.send( - { - "ok": False, - "error_type": type(exc).__name__, - "error": str(exc), - "traceback": traceback.format_exc(), - "recycle": True, - } - ) - except (BrokenPipeError, EOFError, OSError): - pass - finally: - conn.close() - - -def _empty_mps_cache() -> None: - import torch - - gc.collect() - torch.mps.empty_cache() - - -class EmbeddingWorkerClient: - """Synchronous, thread-safe owner of one disposable worker process.""" - - def __init__( - self, - config: MPSWorkerConfig, - *, - timeout_seconds: float, - worker_target: WorkerTarget = _sentence_transformer_worker_main, - ) -> None: - self._config = config - self._timeout_seconds = timeout_seconds - self._worker_target = worker_target - self._context = multiprocessing.get_context("spawn") - self._lock = threading.Lock() - self._process: BaseProcess | None = None - self._conn: Connection | None = None - self._generation = 0 - - @property - def generation(self) -> int: - return self._generation - - @property - def worker_alive(self) -> bool: - with self._lock: - return self._process is not None and self._process.is_alive() - - def _start_locked(self) -> None: - if self._process is not None and self._process.is_alive(): - return - self._discard_locked(graceful=False) - parent_conn, child_conn = self._context.Pipe(duplex=True) - process = self._context.Process( - target=self._worker_target, - args=(child_conn, self._config), - name="cocoindex-mps-embedder", - daemon=True, - ) - process.start() - child_conn.close() - self._process = process - self._conn = parent_conn - self._generation += 1 - logger.info( - "Started MPS embedding worker PID %s (generation %d)", - process.pid, - self._generation, - ) - - def request(self, operation: str, payload: Any) -> Any: - with self._lock: - self._start_locked() - assert self._conn is not None - try: - self._conn.send({"operation": operation, "payload": payload}) - if not self._conn.poll(self._timeout_seconds): - self._discard_locked(graceful=False, force=True) - raise EmbeddingWorkerTimeoutError( - f"Embedding worker exceeded {self._timeout_seconds:g}s timeout" - ) - response = cast(dict[str, Any], self._conn.recv()) - except EmbeddingWorkerTimeoutError: - raise - except (BrokenPipeError, EOFError, OSError) as exc: - self._discard_locked(graceful=False, force=True) - raise EmbeddingWorkerCrashedError( - f"Embedding worker exited before answering: {exc}" - ) from exc - - if not response.get("ok"): - self._discard_locked(graceful=False) - error_type = response.get("error_type", "Exception") - message = response.get("error", "unknown worker error") - remote_traceback = response.get("traceback", "") - raise EmbeddingWorkerRemoteError( - f"{error_type}: {message}\nWorker traceback:\n{remote_traceback}" - ) - - if response.get("recycle"): - driver_memory = response.get("driver_memory") - recommended_memory = response.get("recommended_memory") - logger.warning( - "Recycling MPS embedding worker after driver memory reached %s/%s bytes", - driver_memory, - recommended_memory, - ) - self._discard_locked(graceful=False) - - return response.get("result") - - def _discard_locked(self, *, graceful: bool, force: bool = False) -> None: - process = self._process - conn = self._conn - self._process = None - self._conn = None - - if graceful and conn is not None and process is not None and process.is_alive(): - try: - conn.send({"operation": "close", "payload": None}) - except (BrokenPipeError, EOFError, OSError): - pass - if conn is not None: - conn.close() - - if process is None: - return - if force and process.is_alive(): - process.terminate() - process.join(timeout=_WORKER_JOIN_TIMEOUT_SECONDS) - if process.is_alive(): - process.terminate() - process.join(timeout=_WORKER_JOIN_TIMEOUT_SECONDS) - try: - process.close() - except ValueError: - pass - - def close(self) -> None: - if self._lock.acquire(timeout=0.1): - try: - self._discard_locked(graceful=True) - finally: - self._lock.release() - return - - # A request can hold the lock while blocked in a Metal kernel or IPC - # wait. Killing the child is safe without mutating the connection - # fields; the request thread observes EOF and performs normal cleanup. - process = self._process - if process is not None and process.is_alive(): - logger.warning("Force-terminating busy MPS embedding worker during shutdown") - process.terminate() - process.join(timeout=_WORKER_JOIN_TIMEOUT_SECONDS) - - -class MPSWorkerSentenceTransformerEmbedder(schema.VectorSchemaProvider): - """SentenceTransformer embedder whose Metal allocations live in a child process.""" - - def __init__( - self, - model_name_or_path: str, - *, - device: str | None = "mps", - trust_remote_code: bool = False, - batch_size: int = _MPS_DEFAULT_BATCH_SIZE, - memory_limit_ratio: float = 0.35, - worker_timeout_seconds: float = 300.0, - ) -> None: - self._model_name_or_path = model_name_or_path - self._device = device - self._trust_remote_code = trust_remote_code - self._batch_size = batch_size - self._memory_limit_ratio = memory_limit_ratio - self._worker_timeout_seconds = worker_timeout_seconds - self._worker: EmbeddingWorkerClient | None = None - self._worker_lock = threading.Lock() - - @property - def batch_size(self) -> int: - return self._batch_size - - @property - def memory_limit_ratio(self) -> float: - return self._memory_limit_ratio - - @property - def worker_timeout_seconds(self) -> float: - return self._worker_timeout_seconds - - def __getstate__(self) -> dict[str, Any]: - return { - "model_name_or_path": self._model_name_or_path, - "device": self._device, - "trust_remote_code": self._trust_remote_code, - "batch_size": self._batch_size, - "memory_limit_ratio": self._memory_limit_ratio, - "worker_timeout_seconds": self._worker_timeout_seconds, - } - - def __setstate__(self, state: dict[str, Any]) -> None: - self._model_name_or_path = state["model_name_or_path"] - self._device = state["device"] - self._trust_remote_code = state["trust_remote_code"] - self._batch_size = state["batch_size"] - self._memory_limit_ratio = state["memory_limit_ratio"] - self._worker_timeout_seconds = state["worker_timeout_seconds"] - self._worker = None - self._worker_lock = threading.Lock() - - def _get_worker(self) -> EmbeddingWorkerClient: - if self._worker is None: - with self._worker_lock: - if self._worker is None: - self._worker = EmbeddingWorkerClient( - MPSWorkerConfig( - model_name=self._model_name_or_path, - device=self._device, - trust_remote_code=self._trust_remote_code, - batch_size=self._batch_size, - memory_limit_ratio=self._memory_limit_ratio, - ), - timeout_seconds=self._worker_timeout_seconds, - ) - return self._worker - - def _request_with_crash_retry(self, operation: str, payload: Any) -> Any: - for attempt in range(2): - try: - return self._get_worker().request(operation, payload) - except (EmbeddingWorkerCrashedError, EmbeddingWorkerTimeoutError): - if attempt == 1: - raise - logger.warning("Restarting embedding worker after %s failure", operation) - if operation == "embed" and isinstance(payload, dict): - payload = dict(payload) - payload["batch_size"] = max(1, int(payload["batch_size"]) // 2) - raise AssertionError("unreachable") - - @coco.fn.as_async(batching=True, runner=coco.GPU, max_batch_size=64) - def _embed( - self, - texts: list[str], - prompt_name: str | None = None, - normalize_embeddings: bool = True, - ) -> list[NDArray[np.float32]]: - result = cast( - NDArray[np.float32], - self._request_with_crash_retry( - "embed", - { - "texts": texts, - "batch_size": self._batch_size, - "prompt_name": prompt_name, - "normalize_embeddings": normalize_embeddings, - }, - ), - ) - return list(result) - - @coco.fn(memo=True, version=2, logic_tracking="self") - async def embed( - self, - text: str, - prompt_name: str | None = None, - normalize_embeddings: bool = True, - ) -> NDArray[np.float32]: - result: NDArray[np.float32] = await self._embed( - text, - prompt_name, - normalize_embeddings, - ) - return result - - @coco.fn.as_async(runner=coco.GPU, memo=True) - def dimension(self) -> int: - return int(self._request_with_crash_retry("dimension", None)) - - async def __coco_vector_schema__(self) -> schema.VectorSchema: - return schema.VectorSchema(dtype=np.dtype(np.float32), size=await self.dimension()) - - def __coco_memo_key__(self) -> object: - return ( - self._model_name_or_path, - self._device, - self._trust_remote_code, - self._batch_size, - self._memory_limit_ratio, - ) - - def close(self) -> None: - worker = self._worker - if worker is not None: - worker.close() - self._worker = None diff --git a/src/cocoindex_code/project.py b/src/cocoindex_code/project.py index c966b9a..9849f77 100644 --- a/src/cocoindex_code/project.py +++ b/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,18 @@ 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 - _indexing_gate: 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 @@ -70,26 +74,24 @@ async def run_index( on_progress: Callable[[IndexingProgress], None] | None = None, on_started: asyncio.Event | None = None, ) -> None: - """Acquire the project lock, queue on the daemon gate, run, and release. + """Acquire the project lock, run indexing, and release it. If *on_started* is provided, it is set once the project lock is - acquired. At that point indexing is either active or queued behind - another project. On completion (success or failure) - ``_initial_index_done`` is set. + 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() - async with self._indexing_gate: - self._indexing_stats = IndexingProgress( - num_execution_starts=0, - num_unchanged=0, - num_adds=0, - num_deletes=0, - num_reprocesses=0, - num_errors=0, - ) - await self._run_index_inner(on_progress=on_progress) + self._indexing_stats = IndexingProgress( + num_execution_starts=0, + num_unchanged=0, + num_adds=0, + num_deletes=0, + num_reprocesses=0, + num_errors=0, + ) + await self._run_index_inner(on_progress=on_progress) async def _run_index_inner( self, @@ -114,8 +116,15 @@ 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 is active or queued. @@ -142,7 +151,7 @@ async def stream_index(self) -> AsyncIterator[IndexStreamResponse]: The actual indexing runs in a separate task so that client disconnects (``GeneratorExit``) do not abort the indexing. """ - if self._index_lock.locked() or self._indexing_gate.locked(): + if self._index_lock.locked(): yield IndexWaitingNotice() progress_queue: asyncio.Queue[IndexingProgress] = asyncio.Queue() @@ -273,7 +282,7 @@ async def create( indexing_params: dict[str, Any], query_params: dict[str, Any], chunker_registry: dict[str, ChunkerFn] | None = None, - indexing_gate: asyncio.Lock | None = None, + clear_mps_cache_after_index: bool = False, ) -> Project: """Create a project with explicit embedder and per-call params. @@ -292,8 +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. - indexing_gate: Daemon-wide lock shared by every loaded project. - When omitted, the project gets a private lock for standalone use. + 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) @@ -328,7 +337,7 @@ async def create( result._app = app result._project_root = project_root result._index_lock = asyncio.Lock() - result._indexing_gate = indexing_gate if indexing_gate is not None else 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 diff --git a/src/cocoindex_code/settings.py b/src/cocoindex_code/settings.py index 93b97ef..2bda4a4 100644 --- a/src/cocoindex_code/settings.py +++ b/src/cocoindex_code/settings.py @@ -98,15 +98,9 @@ class EmbeddingSettings: provider: str = "litellm" device: str | None = None min_interval_ms: int | None = None - # SentenceTransformer safety controls. ``batch_size=None`` resolves to 8 - # on MPS and keeps the upstream default on other devices. - batch_size: int | None = None - # Recycle the isolated MPS embedding worker before allocator pressure - # reaches PyTorch's soft/hard watermarks. - mps_memory_limit_ratio: float = 0.35 + # 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 - worker_timeout_seconds: float = 300.0 # 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). @@ -114,15 +108,7 @@ class EmbeddingSettings: query_params: dict[str, Any] | None = None def __post_init__(self) -> None: - if self.batch_size is not None and ( - not isinstance(self.batch_size, int) - or isinstance(self.batch_size, bool) - or not 1 <= self.batch_size <= 64 - ): - raise ValueError("embedding.batch_size must be an integer between 1 and 64") - ratios = ( - ("mps_memory_limit_ratio", self.mps_memory_limit_ratio), ("mps_low_watermark_ratio", self.mps_low_watermark_ratio), ("mps_high_watermark_ratio", self.mps_high_watermark_ratio), ) @@ -130,18 +116,11 @@ def __post_init__(self) -> None: 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_memory_limit_ratio >= self.mps_low_watermark_ratio: - raise ValueError( - "embedding.mps_memory_limit_ratio must be lower than " - "embedding.mps_low_watermark_ratio" - ) 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" ) - if isinstance(self.worker_timeout_seconds, bool) or self.worker_timeout_seconds <= 0: - raise ValueError("embedding.worker_timeout_seconds must be greater than 0") @dataclass @@ -469,17 +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 - if embedding.batch_size is not None: - d["batch_size"] = embedding.batch_size defaults = EmbeddingSettings(model=embedding.model) - if embedding.mps_memory_limit_ratio != defaults.mps_memory_limit_ratio: - d["mps_memory_limit_ratio"] = embedding.mps_memory_limit_ratio 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.worker_timeout_seconds != defaults.worker_timeout_seconds: - d["worker_timeout_seconds"] = embedding.worker_timeout_seconds if embedding.indexing_params is not None: d["indexing_params"] = dict(embedding.indexing_params) if embedding.query_params is not None: @@ -508,16 +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 "batch_size" in emb_dict: - emb_kwargs["batch_size"] = emb_dict["batch_size"] - if "mps_memory_limit_ratio" in emb_dict: - emb_kwargs["mps_memory_limit_ratio"] = float(emb_dict["mps_memory_limit_ratio"]) 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"]) - if "worker_timeout_seconds" in emb_dict: - emb_kwargs["worker_timeout_seconds"] = float(emb_dict["worker_timeout_seconds"]) # 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"). @@ -635,12 +602,9 @@ def save_user_settings(settings: UserSettings) -> Path: _MPS_SAFETY_COMMENT = ( " #\n" - " # Apple Silicon MPS safety defaults. The model runs in a recyclable worker.\n" - " # batch_size: 8\n" - " # mps_memory_limit_ratio: 0.35\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" - " # worker_timeout_seconds: 300\n" ) diff --git a/src/cocoindex_code/shared.py b/src/cocoindex_code/shared.py index 67acbef..08a71ae 100644 --- a/src/cocoindex_code/shared.py +++ b/src/cocoindex_code/shared.py @@ -4,6 +4,7 @@ import importlib.util import logging +import os import pathlib import sys import traceback as _tb @@ -19,8 +20,6 @@ from cocoindex.ops.litellm import LiteLLMEmbedder from cocoindex.ops.sentence_transformers import SentenceTransformerEmbedder - from .mps_embedder import MPSWorkerSentenceTransformerEmbedder - from .settings import EmbeddingSettings logger = logging.getLogger(__name__) @@ -31,7 +30,6 @@ # Type alias Embedder = Union[ "SentenceTransformerEmbedder", - "MPSWorkerSentenceTransformerEmbedder", "LiteLLMEmbedder", ] @@ -52,6 +50,46 @@ def is_sentence_transformers_installed() -> bool: return importlib.util.find_spec("sentence_transformers") is not None +def configure_mps_environment(settings: EmbeddingSettings) -> bool: + """Install MPS safety defaults before CocoIndex's first GPU invocation. + + CocoIndex reads ``COCOINDEX_RUN_GPU_IN_SUBPROCESS`` lazily on the first + ``coco.GPU`` call, while PyTorch reads the allocator watermarks when the + child initializes MPS. Explicit user environment variables always win. + + Returns whether this embedding configuration uses the guarded MPS path. + """ + + use_mps = settings.provider == "sentence-transformers" and ( + settings.device == "mps" or (settings.device is None and sys.platform == "darwin") + ) + if not use_mps: + return False + + os.environ.setdefault("COCOINDEX_RUN_GPU_IN_SUBPROCESS", "1") + os.environ.setdefault("PYTORCH_MPS_LOW_WATERMARK_RATIO", str(settings.mps_low_watermark_ratio)) + os.environ.setdefault( + "PYTORCH_MPS_HIGH_WATERMARK_RATIO", str(settings.mps_high_watermark_ratio) + ) + return os.environ.get("COCOINDEX_RUN_GPU_IN_SUBPROCESS") == "1" + + +@coco.fn.as_async(runner=coco.GPU) +def clear_mps_allocator_cache() -> None: + """Release unused MPS allocator cache inside CocoIndex's GPU child.""" + + import gc + + import torch + + if not torch.backends.mps.is_available(): + return + torch.mps.synchronize() + gc.collect() + torch.mps.empty_cache() + torch.mps.synchronize() + + class EmbeddingCheckResult(NamedTuple): """Outcome of a single embed-test call. See `check_embedding`. @@ -116,40 +154,12 @@ def create_embedder( if model_name.startswith(SBERT_PREFIX): model_name = model_name[len(SBERT_PREFIX) :] - use_mps_worker = settings.device == "mps" or ( - settings.device is None and sys.platform == "darwin" + instance = SentenceTransformerEmbedder( + model_name, + device=settings.device, + trust_remote_code=True, ) - if use_mps_worker: - from .mps_embedder import ( - MPSWorkerSentenceTransformerEmbedder, - apply_mps_allocator_guards, - ) - - apply_mps_allocator_guards( - low_watermark_ratio=settings.mps_low_watermark_ratio, - high_watermark_ratio=settings.mps_high_watermark_ratio, - ) - instance = MPSWorkerSentenceTransformerEmbedder( - model_name, - device=settings.device, - trust_remote_code=True, - batch_size=settings.batch_size or 8, - memory_limit_ratio=settings.mps_memory_limit_ratio, - worker_timeout_seconds=settings.worker_timeout_seconds, - ) - logger.info( - "Embedding model: %s | device: %s | isolated MPS worker | batch size: %d", - settings.model, - settings.device, - settings.batch_size or 8, - ) - else: - instance = SentenceTransformerEmbedder( - model_name, - device=settings.device, - trust_remote_code=True, - ) - logger.info("Embedding model: %s | device: %s", settings.model, settings.device) + logger.info("Embedding model: %s | device: %s", settings.model, settings.device) else: from .litellm_embedder import PacedLiteLLMEmbedder diff --git a/tests/test_mps_embedder.py b/tests/test_mps_embedder.py deleted file mode 100644 index 96b188b..0000000 --- a/tests/test_mps_embedder.py +++ /dev/null @@ -1,178 +0,0 @@ -"""MPS embedding memory and worker-lifecycle safety tests.""" - -from __future__ import annotations - -import threading -import time -from multiprocessing.connection import Connection -from typing import Any - -import numpy as np -import pytest - -from cocoindex_code.mps_embedder import ( - EmbeddingWorkerClient, - EmbeddingWorkerCrashedError, - EmbeddingWorkerTimeoutError, - MPSWorkerConfig, - _encode_with_oom_backoff, - _should_recycle_mps_worker, -) - - -class _OomUntilSmallBatchModel: - def __init__(self) -> None: - self.batch_sizes: list[int] = [] - - def encode(self, texts: list[str], *, batch_size: int, **kwargs: Any) -> np.ndarray: - self.batch_sizes.append(batch_size) - if batch_size > 2: - raise RuntimeError("MPS backend out of memory") - return np.zeros((len(texts), 3), dtype=np.float32) - - -def _recycling_worker(conn: Connection, config: MPSWorkerConfig) -> None: - del config - request = conn.recv() - conn.send( - { - "ok": True, - "result": request["payload"], - "recycle": True, - "driver_memory": 100, - "recommended_memory": 200, - } - ) - conn.close() - - -def _hanging_worker(conn: Connection, config: MPSWorkerConfig) -> None: - del config - conn.recv() - time.sleep(10) - - -def _worker_config() -> MPSWorkerConfig: - return MPSWorkerConfig( - model_name="unused-in-boundary-test", - device="mps", - trust_remote_code=False, - batch_size=8, - memory_limit_ratio=0.35, - ) - - -def test_mps_oom_retries_with_smaller_inner_batches() -> None: - model = _OomUntilSmallBatchModel() - empty_cache_calls = 0 - - def empty_cache() -> None: - nonlocal empty_cache_calls - empty_cache_calls += 1 - - result = _encode_with_oom_backoff( - model, - ["a", "b", "c"], - batch_size=8, - prompt_name=None, - normalize_embeddings=True, - empty_cache=empty_cache, - ) - - assert result.shape == (3, 3) - assert model.batch_sizes == [8, 4, 2] - assert empty_cache_calls == 2 - - -def test_non_oom_embedding_error_is_not_retried() -> None: - class BrokenModel: - def encode(self, texts: list[str], *, batch_size: int, **kwargs: Any) -> np.ndarray: - raise RuntimeError("invalid prompt") - - with pytest.raises(RuntimeError, match="invalid prompt"): - _encode_with_oom_backoff( - BrokenModel(), - ["a"], - batch_size=8, - prompt_name=None, - normalize_embeddings=True, - empty_cache=lambda: None, - ) - - -def test_worker_recycles_only_after_driver_memory_crosses_limit() -> None: - assert ( - _should_recycle_mps_worker( - driver_memory=349, - recommended_memory=1000, - memory_limit_ratio=0.35, - ) - is False - ) - assert ( - _should_recycle_mps_worker( - driver_memory=350, - recommended_memory=1000, - memory_limit_ratio=0.35, - ) - is True - ) - - -def test_worker_recycle_response_starts_a_fresh_process_for_next_request() -> None: - client = EmbeddingWorkerClient( - _worker_config(), - timeout_seconds=2, - worker_target=_recycling_worker, - ) - try: - assert client.request("echo", [1]) == [1] - assert client.generation == 1 - assert client.request("echo", [2]) == [2] - assert client.generation == 2 - finally: - client.close() - - -def test_worker_timeout_terminates_stuck_process() -> None: - client = EmbeddingWorkerClient( - _worker_config(), - timeout_seconds=0.1, - worker_target=_hanging_worker, - ) - try: - with pytest.raises(EmbeddingWorkerTimeoutError, match="0.1"): - client.request("embed", ["stuck"]) - assert client.worker_alive is False - finally: - client.close() - - -def test_close_terminates_worker_even_while_request_is_stuck() -> None: - client = EmbeddingWorkerClient( - _worker_config(), - timeout_seconds=2, - worker_target=_hanging_worker, - ) - request_finished = threading.Event() - - def request() -> None: - try: - client.request("embed", ["stuck"]) - except EmbeddingWorkerCrashedError: - pass - finally: - request_finished.set() - - thread = threading.Thread(target=request) - thread.start() - time.sleep(0.3) - - started = time.monotonic() - client.close() - elapsed = time.monotonic() - started - - thread.join(timeout=2) - assert elapsed < 1 - assert request_finished.is_set() is True - assert client.worker_alive is False diff --git a/tests/test_project_indexing.py b/tests/test_project_indexing.py new file mode 100644 index 0000000..0db593a --- /dev/null +++ b/tests/test_project_indexing.py @@ -0,0 +1,180 @@ +"""Project indexing lifecycle tests.""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator, Callable +from typing import Any, cast +from unittest.mock import AsyncMock + +import pytest + +from cocoindex_code.project import Project + + +class _WatchHandle: + def __init__( + self, + *, + on_enter: Callable[[], None], + on_exit: Callable[[], None], + release: asyncio.Event | None, + ) -> None: + self._on_enter = on_enter + self._on_exit = on_exit + self._release = release + + async def watch(self) -> AsyncIterator[Any]: + self._on_enter() + try: + if self._release is not None: + await self._release.wait() + finally: + self._on_exit() + if False: # pragma: no cover - makes this an async generator + yield None + + +class _ControlledApp: + def __init__(self, handle: _WatchHandle) -> None: + self._handle = handle + + def update(self) -> _WatchHandle: + return self._handle + + +class _ControlledProject(Project): # type: ignore[misc] + def __init__( + self, + *, + on_enter: Callable[[], None], + on_exit: Callable[[], None], + release: asyncio.Event | None = None, + clear_mps_cache_after_index: bool = False, + ) -> None: + self._app = cast( + Any, + _ControlledApp(_WatchHandle(on_enter=on_enter, on_exit=on_exit, release=release)), + ) + self._index_lock = asyncio.Lock() + self._clear_mps_cache_after_index = clear_mps_cache_after_index + self._initial_index_done = asyncio.Event() + self._initial_index_task = None + self._initial_index_started = None + self._indexing_stats = None + + +async def test_projects_can_prepare_indexes_concurrently() -> None: + release = asyncio.Event() + both_entered = asyncio.Event() + active = 0 + max_active = 0 + + def enter() -> None: + nonlocal active, max_active + active += 1 + max_active = max(max_active, active) + if active == 2: + both_entered.set() + + def exit() -> None: + nonlocal active + active -= 1 + + first = _ControlledProject(on_enter=enter, on_exit=exit, release=release) + second = _ControlledProject(on_enter=enter, on_exit=exit, release=release) + + first_task = asyncio.create_task(first.run_index()) + second_task = asyncio.create_task(second.run_index()) + await asyncio.wait_for(both_entered.wait(), timeout=0.5) + + assert max_active == 2 + + release.set() + await asyncio.gather(first_task, second_task) + assert active == 0 + + +async def test_concurrent_initial_index_requests_share_one_background_task() -> None: + release = asyncio.Event() + entered = asyncio.Event() + execution_count = 0 + + def enter() -> None: + nonlocal execution_count + execution_count += 1 + entered.set() + + project = _ControlledProject(on_enter=enter, on_exit=lambda: None, release=release) + first = asyncio.create_task(project.ensure_indexing_started()) + second = asyncio.create_task(project.ensure_indexing_started()) + + await asyncio.wait_for(asyncio.gather(first, second), timeout=0.5) + await asyncio.wait_for(entered.wait(), timeout=0.5) + assert project._index_lock.locked() is True + assert execution_count == 1 + + release.set() + await project.wait_for_indexing_done() + assert execution_count == 1 + + +async def test_mps_allocator_cache_is_cleared_after_indexing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + clear_cache = AsyncMock() + monkeypatch.setattr("cocoindex_code.project.clear_mps_allocator_cache", clear_cache) + project = _ControlledProject( + on_enter=lambda: None, + on_exit=lambda: None, + clear_mps_cache_after_index=True, + ) + + await project.run_index() + + clear_cache.assert_awaited_once_with() + assert project._initial_index_done.is_set() is True + + +async def test_mps_cache_cleanup_failure_does_not_fail_completed_index( + monkeypatch: pytest.MonkeyPatch, +) -> None: + clear_cache = AsyncMock(side_effect=RuntimeError("cleanup failed")) + monkeypatch.setattr("cocoindex_code.project.clear_mps_allocator_cache", clear_cache) + project = _ControlledProject( + on_enter=lambda: None, + on_exit=lambda: None, + clear_mps_cache_after_index=True, + ) + + await project.run_index() + + assert project._initial_index_done.is_set() is True + assert project.indexing_stats is None + + +async def test_cancelling_mps_cache_cleanup_still_restores_indexing_state( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cleanup_started = asyncio.Event() + + async def blocked_cleanup() -> None: + cleanup_started.set() + await asyncio.Event().wait() + + monkeypatch.setattr("cocoindex_code.project.clear_mps_allocator_cache", blocked_cleanup) + project = _ControlledProject( + on_enter=lambda: None, + on_exit=lambda: None, + clear_mps_cache_after_index=True, + ) + task = asyncio.create_task(project.run_index()) + await asyncio.wait_for(cleanup_started.wait(), timeout=0.5) + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert project._initial_index_done.is_set() is True + assert project.indexing_stats is None + await asyncio.wait_for(project.wait_for_indexing_done(), timeout=0.5) diff --git a/tests/test_project_indexing_gate.py b/tests/test_project_indexing_gate.py deleted file mode 100644 index 295cb67..0000000 --- a/tests/test_project_indexing_gate.py +++ /dev/null @@ -1,146 +0,0 @@ -"""Cross-project indexing concurrency tests.""" - -from __future__ import annotations - -import asyncio -from collections.abc import Callable - -from cocoindex_code.project import Project - - -class _ControlledProject(Project): # type: ignore[misc] - def __init__( - self, - *, - indexing_gate: asyncio.Lock, - on_enter: Callable[[], None], - on_exit: Callable[[], None], - release: asyncio.Event | None = None, - ) -> None: - self._index_lock = asyncio.Lock() - self._indexing_gate = indexing_gate - self._initial_index_done = asyncio.Event() - self._initial_index_task = None - self._initial_index_started = None - self._indexing_stats = None - self._on_enter = on_enter - self._on_exit = on_exit - self._release = release - - async def _run_index_inner( - self, - on_progress: Callable[..., None] | None = None, - ) -> None: - del on_progress - try: - self._on_enter() - if self._release is not None: - await self._release.wait() - finally: - self._on_exit() - self._initial_index_done.set() - self._indexing_stats = None - - -async def test_projects_share_one_indexing_gate() -> None: - indexing_gate = asyncio.Lock() - release_first = asyncio.Event() - second_entered = asyncio.Event() - active = 0 - max_active = 0 - - def enter_first() -> None: - nonlocal active, max_active - active += 1 - max_active = max(max_active, active) - - def enter_second() -> None: - nonlocal active, max_active - active += 1 - max_active = max(max_active, active) - second_entered.set() - - def exit() -> None: - nonlocal active - active -= 1 - - first = _ControlledProject( - indexing_gate=indexing_gate, - on_enter=enter_first, - on_exit=exit, - release=release_first, - ) - second = _ControlledProject( - indexing_gate=indexing_gate, - on_enter=enter_second, - on_exit=exit, - ) - - first_task = asyncio.create_task(first.run_index()) - while active == 0: - await asyncio.sleep(0) - - second_task = asyncio.create_task(second.run_index()) - await asyncio.sleep(0) - - assert second._index_lock.locked() is True - assert second_entered.is_set() is False - assert active == 1 - - release_first.set() - await first_task - await second_entered.wait() - await second_task - - assert max_active == 1 - assert active == 0 - - -async def test_initial_index_reports_started_when_queued_behind_another_project() -> None: - indexing_gate = asyncio.Lock() - await indexing_gate.acquire() - project = _ControlledProject( - indexing_gate=indexing_gate, - on_enter=lambda: None, - on_exit=lambda: None, - ) - - ensure_task = asyncio.create_task(project.ensure_indexing_started()) - - try: - await asyncio.wait_for(asyncio.shield(ensure_task), timeout=0.5) - assert project._index_lock.locked() is True - assert ensure_task.done() is True - assert project.should_wait_for_indexing is True - finally: - indexing_gate.release() - - await ensure_task - await project.wait_for_indexing_done() - - -async def test_concurrent_initial_index_requests_share_one_background_task() -> None: - indexing_gate = asyncio.Lock() - await indexing_gate.acquire() - execution_count = 0 - - def enter() -> None: - nonlocal execution_count - execution_count += 1 - - project = _ControlledProject( - indexing_gate=indexing_gate, - on_enter=enter, - on_exit=lambda: None, - ) - first = asyncio.create_task(project.ensure_indexing_started()) - second = asyncio.create_task(project.ensure_indexing_started()) - - try: - await asyncio.wait_for(asyncio.gather(first, second), timeout=0.5) - assert project._index_lock.locked() is True - finally: - indexing_gate.release() - - await project.wait_for_indexing_done() - assert execution_count == 1 diff --git a/tests/test_settings.py b/tests/test_settings.py index 3307d81..267fff1 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -21,6 +21,7 @@ UserSettings, _reset_db_path_mapping_cache, _reset_host_path_mapping_cache, + _user_settings_from_dict, default_project_settings, default_user_settings, find_parent_with_marker, @@ -55,11 +56,8 @@ def test_default_user_settings() -> None: assert s.embedding.model == "Snowflake/snowflake-arctic-embed-xs" assert s.embedding.device is None assert s.embedding.min_interval_ms is None - assert s.embedding.batch_size is None - assert s.embedding.mps_memory_limit_ratio == 0.35 assert s.embedding.mps_low_watermark_ratio == 0.4 assert s.embedding.mps_high_watermark_ratio == 0.5 - assert s.embedding.worker_timeout_seconds == 300.0 assert s.envs == {} @@ -82,11 +80,8 @@ def test_save_and_load_user_settings(tmp_path: Path) -> None: model="gemini/text-embedding-004", device="cpu", min_interval_ms=300, - batch_size=8, - mps_memory_limit_ratio=0.3, mps_low_watermark_ratio=0.35, mps_high_watermark_ratio=0.45, - worker_timeout_seconds=120.0, ), envs={"GEMINI_API_KEY": "test-key"}, ) @@ -96,24 +91,16 @@ def test_save_and_load_user_settings(tmp_path: Path) -> None: assert loaded.embedding.model == settings.embedding.model assert loaded.embedding.device == settings.embedding.device assert loaded.embedding.min_interval_ms == settings.embedding.min_interval_ms - assert loaded.embedding.batch_size == settings.embedding.batch_size - assert loaded.embedding.mps_memory_limit_ratio == settings.embedding.mps_memory_limit_ratio assert loaded.embedding.mps_low_watermark_ratio == settings.embedding.mps_low_watermark_ratio assert loaded.embedding.mps_high_watermark_ratio == settings.embedding.mps_high_watermark_ratio - assert loaded.embedding.worker_timeout_seconds == settings.embedding.worker_timeout_seconds assert loaded.envs == settings.envs @pytest.mark.parametrize( ("field", "value", "message"), [ - ("batch_size", 0, "batch_size"), - ("batch_size", 3.5, "batch_size"), - ("batch_size", 65, "batch_size"), - ("mps_memory_limit_ratio", 0, "mps_memory_limit_ratio"), ("mps_low_watermark_ratio", 0, "mps_low_watermark_ratio"), ("mps_high_watermark_ratio", 1.1, "mps_high_watermark_ratio"), - ("worker_timeout_seconds", 0, "worker_timeout_seconds"), ], ) def test_embedding_safety_settings_reject_invalid_values( @@ -126,23 +113,31 @@ def test_embedding_safety_settings_reject_invalid_values( def test_embedding_safety_settings_require_ordered_mps_limits() -> None: - with pytest.raises(ValueError, match="mps_memory_limit_ratio"): - EmbeddingSettings( - model="model", - mps_memory_limit_ratio=0.45, - mps_low_watermark_ratio=0.4, - mps_high_watermark_ratio=0.5, - ) - with pytest.raises(ValueError, match="mps_low_watermark_ratio"): EmbeddingSettings( model="model", - mps_memory_limit_ratio=0.3, mps_low_watermark_ratio=0.6, mps_high_watermark_ratio=0.5, ) +def test_removed_custom_mps_worker_settings_are_ignored() -> None: + settings = _user_settings_from_dict( + { + "embedding": { + "provider": "sentence-transformers", + "model": "model", + "batch_size": 8, + "mps_memory_limit_ratio": 0.35, + "worker_timeout_seconds": 300, + } + } + ) + + assert settings.embedding.mps_low_watermark_ratio == 0.4 + assert settings.embedding.mps_high_watermark_ratio == 0.5 + + def test_save_and_load_project_settings(tmp_path: Path) -> None: settings = ProjectSettings( include_patterns=["**/*.py", "**/*.rs"], @@ -413,13 +408,11 @@ def test_save_initial_user_settings_round_trip() -> None: path = save_initial_user_settings(emb, defaults_applied=False) content = path.read_text() - # Hint comment and the four commented env-var examples. + # Hint comment, MPS allocator defaults, and env-var examples. assert "ccc doctor" in content - assert "# batch_size: 8" in content - assert "# mps_memory_limit_ratio: 0.35" in content assert "# mps_low_watermark_ratio: 0.4" in content assert "# mps_high_watermark_ratio: 0.5" in content - assert "# worker_timeout_seconds: 300" in content + assert "CocoIndex's GPU subprocess" in content assert "# envs:" in content for key in ("OPENAI_API_KEY", "GEMINI_API_KEY", "ANTHROPIC_API_KEY", "VOYAGE_API_KEY"): assert f"# {key}:" in content diff --git a/tests/test_shared.py b/tests/test_shared.py index 1248016..88288a2 100644 --- a/tests/test_shared.py +++ b/tests/test_shared.py @@ -7,12 +7,13 @@ import numpy as np import pytest +from cocoindex.ops.sentence_transformers import SentenceTransformerEmbedder from cocoindex_code.litellm_embedder import PacedLiteLLMEmbedder -from cocoindex_code.mps_embedder import MPSWorkerSentenceTransformerEmbedder from cocoindex_code.settings import EmbeddingSettings from cocoindex_code.shared import ( check_embedding, + configure_mps_environment, create_embedder, is_sentence_transformers_installed, ) @@ -68,29 +69,35 @@ def test_create_embedder_sentence_transformers_ignores_indexing_params() -> None assert not isinstance(embedder, PacedLiteLLMEmbedder) -def test_create_embedder_uses_isolated_worker_and_allocator_guards_for_mps( +def test_create_embedder_uses_cocoindex_sentence_transformer_for_mps() -> None: + embedder = create_embedder( + EmbeddingSettings( + provider="sentence-transformers", + model="sentence-transformers/all-MiniLM-L6-v2", + device="mps", + ) + ) + + assert isinstance(embedder, SentenceTransformerEmbedder) + + +def test_configure_mps_environment_enables_cocoindex_gpu_subprocess( monkeypatch: pytest.MonkeyPatch, ) -> None: + monkeypatch.delenv("COCOINDEX_RUN_GPU_IN_SUBPROCESS", raising=False) monkeypatch.delenv("PYTORCH_MPS_HIGH_WATERMARK_RATIO", raising=False) monkeypatch.delenv("PYTORCH_MPS_LOW_WATERMARK_RATIO", raising=False) - embedder = create_embedder( + enabled = configure_mps_environment( EmbeddingSettings( provider="sentence-transformers", model="sentence-transformers/all-MiniLM-L6-v2", device="mps", - batch_size=4, - mps_memory_limit_ratio=0.3, - mps_low_watermark_ratio=0.4, - mps_high_watermark_ratio=0.5, - worker_timeout_seconds=90, ) ) - assert isinstance(embedder, MPSWorkerSentenceTransformerEmbedder) - assert embedder.batch_size == 4 - assert embedder.memory_limit_ratio == 0.3 - assert embedder.worker_timeout_seconds == 90 + assert enabled is True + assert os.environ["COCOINDEX_RUN_GPU_IN_SUBPROCESS"] == "1" assert os.environ["PYTORCH_MPS_LOW_WATERMARK_RATIO"] == "0.4" assert os.environ["PYTORCH_MPS_HIGH_WATERMARK_RATIO"] == "0.5" @@ -98,10 +105,11 @@ def test_create_embedder_uses_isolated_worker_and_allocator_guards_for_mps( def test_explicit_mps_allocator_env_takes_precedence( monkeypatch: pytest.MonkeyPatch, ) -> None: + monkeypatch.setenv("COCOINDEX_RUN_GPU_IN_SUBPROCESS", "0") monkeypatch.setenv("PYTORCH_MPS_LOW_WATERMARK_RATIO", "0.25") monkeypatch.setenv("PYTORCH_MPS_HIGH_WATERMARK_RATIO", "0.3") - create_embedder( + enabled = configure_mps_environment( EmbeddingSettings( provider="sentence-transformers", model="sentence-transformers/all-MiniLM-L6-v2", @@ -109,10 +117,45 @@ def test_explicit_mps_allocator_env_takes_precedence( ) ) + assert enabled is False + assert os.environ["COCOINDEX_RUN_GPU_IN_SUBPROCESS"] == "0" assert os.environ["PYTORCH_MPS_LOW_WATERMARK_RATIO"] == "0.25" assert os.environ["PYTORCH_MPS_HIGH_WATERMARK_RATIO"] == "0.3" +def test_configure_mps_environment_does_not_change_non_mps_provider( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("COCOINDEX_RUN_GPU_IN_SUBPROCESS", raising=False) + monkeypatch.delenv("PYTORCH_MPS_HIGH_WATERMARK_RATIO", raising=False) + monkeypatch.delenv("PYTORCH_MPS_LOW_WATERMARK_RATIO", raising=False) + + enabled = configure_mps_environment( + EmbeddingSettings(provider="litellm", model="text-embedding-3-small") + ) + + assert enabled is False + assert "COCOINDEX_RUN_GPU_IN_SUBPROCESS" not in os.environ + assert "PYTORCH_MPS_LOW_WATERMARK_RATIO" not in os.environ + assert "PYTORCH_MPS_HIGH_WATERMARK_RATIO" not in os.environ + + +def test_configure_mps_environment_auto_detects_macos( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr("cocoindex_code.shared.sys.platform", "darwin") + monkeypatch.delenv("COCOINDEX_RUN_GPU_IN_SUBPROCESS", raising=False) + + enabled = configure_mps_environment( + EmbeddingSettings( + provider="sentence-transformers", + model="sentence-transformers/all-MiniLM-L6-v2", + ) + ) + + assert enabled is True + + def test_is_sentence_transformers_installed_true_in_dev() -> None: # Dev env pulls in sentence-transformers via the `dev` extras group. assert is_sentence_transformers_installed() is True diff --git a/uv.lock b/uv.lock index c92f9e3..c6e4253 100644 --- a/uv.lock +++ b/uv.lock @@ -355,7 +355,7 @@ wheels = [ [[package]] name = "cocoindex" -version = "1.0.13" +version = "1.0.18" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, @@ -367,17 +367,17 @@ dependencies = [ { name = "typing-extensions" }, { name = "watchdog" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/6e/b9/5fe83e2e0fd4adc54fb4d526c7197942861d184bd006c4a4591b229e290e/cocoindex-1.0.13.tar.gz", hash = "sha256:8e52558f02cbfe2cc5505bea214354f147a8a0385ed7b9a47a0cdca20b8673be", size = 647804, upload-time = "2026-06-22T16:46:17.563Z" } +sdist = { url = "https://files.pythonhosted.org/packages/c3/99/9e055aaa2ecf803d3bff3295ddc7613baba8073b23b4528776655758cf5a/cocoindex-1.0.18.tar.gz", hash = "sha256:f366a14438f3c581abf94a31b6cb96897e69f3d0776fd52da3ae20a9cc10a98f", size = 741567, upload-time = "2026-07-23T17:44:16.551Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/24/0c/0cd5e797b108b50080dd17fb7f5761a23d41432b694355554609b392e191/cocoindex-1.0.13-cp311-abi3-macosx_10_12_x86_64.whl", hash = "sha256:186aef33ea46fab459188c1ef7c697334a6a7299e726d38d46b01bd5543cf6f7", size = 9109005, upload-time = "2026-06-22T16:46:15.304Z" }, - { url = "https://files.pythonhosted.org/packages/ff/6b/7b1511d1a4036e14fd599f163cd84103b96f4140135295bdf3b75966eed4/cocoindex-1.0.13-cp311-abi3-macosx_11_0_arm64.whl", hash = "sha256:e48570953f4919f3061043d56835d629836052385648d6b874432918bdbf811c", size = 9183196, upload-time = "2026-06-22T16:46:09.447Z" }, - { url = "https://files.pythonhosted.org/packages/f2/a3/4be3a6b0257e27a437edba3942812537b191fe3c662f0c514207b92e831f/cocoindex-1.0.13-cp311-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:7970d7ca83d90a15f417e49c09a8f253b6d231c8f3571e0c325efc1857dc449b", size = 9059825, upload-time = "2026-06-22T16:45:59.309Z" }, - { url = "https://files.pythonhosted.org/packages/bb/75/1d7c019746d3dac105eec5f639328f5ccccc1e750d03f3ecb47161460dde/cocoindex-1.0.13-cp311-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:b4682c8d84a4a80c1bca10731bc6ab9d2c97f6451194526a39ddd4d0ad984458", size = 9282739, upload-time = "2026-06-22T16:46:04.456Z" }, - { url = "https://files.pythonhosted.org/packages/cd/08/f40e17233c70452b94a4a3791689d43793a11cae8f06f81d8a4186628cab/cocoindex-1.0.13-cp311-abi3-win_amd64.whl", hash = "sha256:47787cd38e7e0b8f74c2ea0b74030564c0405d4795a6af1a48ccc4055150cf39", size = 9497162, upload-time = "2026-06-22T16:46:19.827Z" }, - { url = "https://files.pythonhosted.org/packages/98/06/a0f0766f84f818fc4f79509ef4de0deeacbf7e6e13b74958498fb64ba316/cocoindex-1.0.13-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:dec6e4df7d62b56221427aae1c4f067a18ecf5c8eb50d0fc09cc9d033eed3048", size = 9184066, upload-time = "2026-06-22T16:46:12.68Z" }, - { url = "https://files.pythonhosted.org/packages/1a/71/e265116df84ac5338e0ae7d8e8b6e3f177dc6f28df0a5c7308662333ba37/cocoindex-1.0.13-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:111a5bc1d49e20e06659a5323019f8b7ca195c983e0e6700c66c104973efb689", size = 9062342, upload-time = "2026-06-22T16:46:02.011Z" }, - { url = "https://files.pythonhosted.org/packages/a2/61/3bf28b73b256be57626f60bbc0e360aa5707ddf0625a55c3dd9869f0403d/cocoindex-1.0.13-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:5dbaa672014ddf68bd0ca8ae10c4f6e36381cacc1f482a14b500aa4a458633b4", size = 9279283, upload-time = "2026-06-22T16:46:06.797Z" }, - { url = "https://files.pythonhosted.org/packages/64/30/a8bc89c9c99fe783bab6e79334bc4bc780a18b7bb4d1aa5f69d31bb751c8/cocoindex-1.0.13-cp314-cp314t-win_amd64.whl", hash = "sha256:d070e450e1ed8cb9b234a447dc6171611007ed66b9ef80bcaf328ccae6577ad5", size = 9489393, upload-time = "2026-06-22T16:46:22.353Z" }, + { url = "https://files.pythonhosted.org/packages/ca/c3/1c2f230578bbe3b207eed5b2c32b6baa86dec8be6f76165d4ff88d00dbb4/cocoindex-1.0.18-cp311-abi3-macosx_10_12_x86_64.whl", hash = "sha256:30e669eb179781a08c300a0d868ad0e624d8383fee8f009df1adb3325e6d79f5", size = 9393082, upload-time = "2026-07-23T17:44:13.976Z" }, + { url = "https://files.pythonhosted.org/packages/0d/fd/51e464e9705cd4ede622bd610efe4ec86fea6a1bec9e28ed697e86425732/cocoindex-1.0.18-cp311-abi3-macosx_11_0_arm64.whl", hash = "sha256:b9102afba6e6f66d7958bb4fc01bd04db7b4807ef25b1101f39be879efcae600", size = 9483580, upload-time = "2026-07-23T17:44:09.762Z" }, + { url = "https://files.pythonhosted.org/packages/52/7c/d44c4c62da545d3eaa97691100dcd164c793c2bc0486a5082bbd40a9aa0b/cocoindex-1.0.18-cp311-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:e6f20760676a39ea069e7659c8e8fc8d9c296165097839e5806a80c4e376ebc0", size = 9363332, upload-time = "2026-07-23T17:44:01.1Z" }, + { url = "https://files.pythonhosted.org/packages/de/d0/83ddd620431c31c3224d8efa01475d7e92fb982aa2eaf453ab64f8d697e9/cocoindex-1.0.18-cp311-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:3830c748675672b55cd6edaa9c222d28ccd1a782ebc29ab08c3a187ee4b82059", size = 9591860, upload-time = "2026-07-23T17:44:05.396Z" }, + { url = "https://files.pythonhosted.org/packages/94/52/bfc4fd2dd838067d3b78ed30c7a8b8c2b4b5c17470737ab14e1588496cf5/cocoindex-1.0.18-cp311-abi3-win_amd64.whl", hash = "sha256:9546e2a5cf050692eb14accbcccfa9d235367ecd7efb071f43a8de46f7b981db", size = 9821212, upload-time = "2026-07-23T17:44:18.006Z" }, + { url = "https://files.pythonhosted.org/packages/2e/95/96876ea3f3d832df6cb70196b6de6f96c3c1f1cf8ca1b0c36dd8c6be60d3/cocoindex-1.0.18-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:74e137d1496db81ceee2ef600cb91d79d22bec68b5a1ca19fdc410c03e4d8f57", size = 9465860, upload-time = "2026-07-23T17:44:12.127Z" }, + { url = "https://files.pythonhosted.org/packages/57/c4/f2e1d9ea6561e254187a62f9021119d2ffedc88c48311813981e4f14c805/cocoindex-1.0.18-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:49150befb5878fb7b99e325cbab0b35c46cacf3184c92ec6a44a7bf3725d4a58", size = 9362158, upload-time = "2026-07-23T17:44:03.366Z" }, + { url = "https://files.pythonhosted.org/packages/c1/1c/e2700f861d0ab2be62aa08e3c08cde193967afafec8bf8c1f7a938e5b92e/cocoindex-1.0.18-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:28be1d33bb70592e1b989afd60346ac587e62f50d7b1ed3846c6855d885c93b7", size = 9588061, upload-time = "2026-07-23T17:44:07.602Z" }, + { url = "https://files.pythonhosted.org/packages/a7/3e/5754546e3fe77dc7ced2a5dca7c3e6732a4fa495cc2eed7d3d949b88d0f3/cocoindex-1.0.18-cp314-cp314t-win_amd64.whl", hash = "sha256:70c0448cc6281ac77188f37995b49a1a94c91cbd1fd9e2b1b6ae8d4ec1e72096", size = 9812541, upload-time = "2026-07-23T17:44:19.928Z" }, ] [package.optional-dependencies] @@ -436,10 +436,10 @@ dev = [ [package.metadata] requires-dist = [ - { name = "cocoindex", extras = ["litellm"], specifier = ">=1.0.13,<1.1.0" }, - { name = "cocoindex", extras = ["sentence-transformers"], marker = "extra == 'dev'", specifier = ">=1.0.13,<1.1.0" }, - { name = "cocoindex", extras = ["sentence-transformers"], marker = "extra == 'embeddings-local'", specifier = ">=1.0.13,<1.1.0" }, - { name = "cocoindex", extras = ["sentence-transformers"], marker = "extra == 'full'", specifier = ">=1.0.13,<1.1.0" }, + { name = "cocoindex", extras = ["litellm"], specifier = ">=1.0.17,<1.1.0" }, + { name = "cocoindex", extras = ["sentence-transformers"], marker = "extra == 'dev'", specifier = ">=1.0.17,<1.1.0" }, + { name = "cocoindex", extras = ["sentence-transformers"], marker = "extra == 'embeddings-local'", specifier = ">=1.0.17,<1.1.0" }, + { name = "cocoindex", extras = ["sentence-transformers"], marker = "extra == 'full'", specifier = ">=1.0.17,<1.1.0" }, { name = "einops", specifier = ">=0.8.2" }, { name = "mcp", specifier = ">=1.0.0" }, { name = "msgspec", specifier = ">=0.19.0" }, @@ -461,7 +461,7 @@ provides-extras = ["dev", "embeddings-local", "full"] [package.metadata.requires-dev] dev = [ - { name = "cocoindex", extras = ["sentence-transformers"], specifier = ">=1.0.13,<1.1.0" }, + { name = "cocoindex", extras = ["sentence-transformers"], specifier = ">=1.0.17,<1.1.0" }, { name = "mypy", specifier = ">=1.0.0" }, { name = "prek", specifier = ">=0.1.0" }, { name = "pytest", specifier = ">=7.0.0" }, @@ -636,7 +636,7 @@ name = "cuda-bindings" version = "13.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cuda-pathfinder", marker = "python_full_version < '3.12' or sys_platform != 'win32'" }, + { name = "cuda-pathfinder" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/51/6b/457ca12dad3ee9bfcc9a545cfd6b64b359ba49de40f776f6e028e678f262/cuda_bindings-13.3.1-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c5879712accf6e14bb01aa5e67440eb84998b8d104b509cc7a6dc0b8f656a474", size = 6053539, upload-time = "2026-05-29T23:11:43.19Z" }, @@ -669,43 +669,43 @@ wheels = [ [package.optional-dependencies] cublas = [ - { name = "nvidia-cublas", marker = "(python_full_version < '3.12' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-cuda-nvrtc", marker = "(python_full_version < '3.12' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cublas", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cuda-nvrtc", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] cudart = [ - { name = "nvidia-cuda-runtime", marker = "(python_full_version < '3.12' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cuda-runtime", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] cufft = [ - { name = "nvidia-cufft", marker = "(python_full_version < '3.12' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.12' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cufft", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] cufile = [ - { name = "nvidia-cufile", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cufile", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] cupti = [ - { name = "nvidia-cuda-cupti", marker = "(python_full_version < '3.12' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cuda-cupti", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] curand = [ - { name = "nvidia-curand", marker = "(python_full_version < '3.12' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-curand", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] cusolver = [ - { name = "nvidia-cublas", marker = "(python_full_version < '3.12' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-cusolver", marker = "(python_full_version < '3.12' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-cusparse", marker = "(python_full_version < '3.12' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.12' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cublas", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cusolver", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] cusparse = [ - { name = "nvidia-cusparse", marker = "(python_full_version < '3.12' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.12' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] nvjitlink = [ - { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.12' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] nvrtc = [ - { name = "nvidia-cuda-nvrtc", marker = "(python_full_version < '3.12' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cuda-nvrtc", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] nvtx = [ - { name = "nvidia-nvtx", marker = "(python_full_version < '3.12' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-nvtx", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, ] [[package]] @@ -1675,7 +1675,7 @@ name = "nvidia-cublas" version = "13.1.1.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cuda-nvrtc", marker = "python_full_version < '3.12' or sys_platform != 'win32'" }, + { name = "nvidia-cuda-nvrtc" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/a7/a1/0bd24ee8c8d03adac032fd2909426a00c88f8c57961b1277ded97f91119f/nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b7a210458267ac818974c53038fbec2e969d5c99f305ab15c72522fa9f001dd5", size = 542848918, upload-time = "2026-04-08T18:46:22.985Z" }, @@ -1714,7 +1714,7 @@ name = "nvidia-cudnn-cu13" version = "9.20.0.48" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas", marker = "python_full_version < '3.12' or sys_platform != 'win32'" }, + { name = "nvidia-cublas" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/56/c5/83384d846b2fd17c44bd499b36c75a45ed4f095fbbb2252294e89cea5c5c/nvidia_cudnn_cu13-9.20.0.48-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:e31454ae00094b0c55319d9d15b6fa2fc50a9e1c0f5c8c80fb75258234e731e1", size = 444574296, upload-time = "2026-03-09T19:28:27.751Z" }, @@ -1726,7 +1726,7 @@ name = "nvidia-cufft" version = "12.0.0.61" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink", marker = "python_full_version < '3.12' or sys_platform != 'win32'" }, + { name = "nvidia-nvjitlink" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" }, @@ -1756,9 +1756,9 @@ name = "nvidia-cusolver" version = "12.0.4.66" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas", marker = "python_full_version < '3.12' or sys_platform != 'win32'" }, - { name = "nvidia-cusparse", marker = "python_full_version < '3.12' or sys_platform != 'win32'" }, - { name = "nvidia-nvjitlink", marker = "python_full_version < '3.12' or sys_platform != 'win32'" }, + { name = "nvidia-cublas" }, + { name = "nvidia-cusparse" }, + { name = "nvidia-nvjitlink" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" }, @@ -1770,7 +1770,7 @@ name = "nvidia-cusparse" version = "12.6.3.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink", marker = "python_full_version < '3.12' or sys_platform != 'win32'" }, + { name = "nvidia-nvjitlink" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" },