diff --git a/EMBEDDINGS.md b/EMBEDDINGS.md index ba87b0a..71d7fb4 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,19 @@ Indexing speed depends on **both parameter count and model architecture**: For encoder models at medium or large sizes, a GPU will still accelerate indexing. Add `device: cuda` (or `mps` on Mac) to `global_settings.yml`. +On Apple Silicon, SentenceTransformer calls run through [CocoIndex's isolated +GPU subprocess](https://github.com/cocoindex-io/cocoindex/blob/v1.0.18/python/cocoindex/_internal/runner.py). +This keeps the model warm across batches while separating Metal allocations +from the daemon. cocoindex-code configures PyTorch to begin adaptive allocation +cleanup at 40% of its recommended maximum working set and to enforce a hard +limit at 50%; tune these values with `mps_low_watermark_ratio` and +`mps_high_watermark_ratio`. At the end of each index run, unused allocator cache +is released inside the GPU subprocess. CocoIndex also [retries MPS OOM failures +with smaller batches](https://github.com/cocoindex-io/cocoindex/blob/v1.0.18/python/cocoindex/ops/sentence_transformers.py). +See the [PyTorch MPS environment variable reference](https://docs.pytorch.org/docs/stable/mps_environment_variables.html) +for the allocator semantics. Explicit environment variables take precedence +over the YAML defaults. + --- ## The `ccc init` Wizard diff --git a/README.md b/README.md index 6594c82..716b128 100644 --- a/README.md +++ b/README.md @@ -505,6 +505,8 @@ embedding: model: Snowflake/snowflake-arctic-embed-xs device: mps # optional: cpu, cuda, mps (auto-detected if omitted) min_interval_ms: 300 # optional: pace LiteLLM embedding requests to reduce 429s; defaults to 5 for LiteLLM + mps_low_watermark_ratio: 0.4 # optional: PyTorch allocator soft limit + mps_high_watermark_ratio: 0.5 # optional: PyTorch allocator hard limit # Optional extra kwargs passed to the embedder, separately for indexing vs query. # `ccc init` auto-populates these for known models (e.g. Cohere, Voyage, Nvidia NIM, @@ -523,6 +525,10 @@ daemon: > **Note:** The daemon inherits your shell environment. If an API key (e.g. `OPENAI_API_KEY`) is already set as an environment variable, you don't need to duplicate it in `envs`. The `envs` field is only for values that aren't in your environment. +> **Apple Silicon memory safety:** MPS SentenceTransformer calls use [CocoIndex's isolated GPU subprocess](https://github.com/cocoindex-io/cocoindex/blob/v1.0.18/python/cocoindex/_internal/runner.py), keeping the model loaded while isolating Metal allocations from the daemon. The low and high watermarks are ratios of PyTorch's recommended maximum working set; they default here to `0.4` and `0.5`. CocoIndex retries MPS out-of-memory failures with progressively smaller batches, and cocoindex-code [releases unused allocator cache](https://docs.pytorch.org/docs/stable/generated/torch.mps.empty_cache.html) after each index run. Explicit `COCOINDEX_RUN_GPU_IN_SUBPROCESS`, `PYTORCH_MPS_LOW_WATERMARK_RATIO`, and `PYTORCH_MPS_HIGH_WATERMARK_RATIO` environment variables take precedence over these defaults. + +> **Indexing concurrency:** Multiple projects may prepare indexes concurrently, while CocoIndex serializes their GPU calls through its single MPS subprocess. A search waits only when its own project still needs the initial index. + > **Idle timeout:** the background daemon holds the embedding model in RAM, so it exits after `daemon.idle_timeout_minutes` without client activity and is restarted automatically on your next `ccc` command or MCP search. A live MCP session sends periodic heartbeats, so the daemon never idles out while your coding agent is connected. Set `0` to keep the daemon running forever. > **Custom location:** set `COCOINDEX_CODE_DIR` to place `global_settings.yml` somewhere other than `~/.cocoindex_code/` — useful if you want the file to live alongside your projects (e.g. on a synced folder). 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 d2b38ab..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,6 +135,7 @@ class ProjectRegistry: _projects: dict[str, Project] _embedder: Embedder | None + _clear_mps_cache_after_index: bool indexing_params: dict[str, Any] query_params: dict[str, Any] @@ -143,9 +144,11 @@ def __init__( embedder: Embedder | None, indexing_params: dict[str, Any] | None = None, query_params: dict[str, Any] | None = None, + clear_mps_cache_after_index: bool = False, ) -> None: self._projects = {} self._embedder = embedder + self._clear_mps_cache_after_index = clear_mps_cache_after_index self.indexing_params = dict(indexing_params) if indexing_params else {} self.query_params = dict(query_params) if query_params else {} @@ -165,6 +168,7 @@ async def get_project(self, project_root: str) -> Project: indexing_params=self.indexing_params, query_params=self.query_params, chunker_registry=chunker_registry, + clear_mps_cache_after_index=self._clear_mps_cache_after_index, ) self._projects[project_root] = project return self._projects[project_root] @@ -636,6 +640,7 @@ def run_daemon( embedder: Embedder | None indexing_params: dict[str, Any] = {} query_params: dict[str, Any] = {} + clear_mps_cache_after_index = False handshake_warnings: list[str] = [] daemon_settings = DaemonSettings() if user_settings_path().is_file(): @@ -644,6 +649,7 @@ def run_daemon( settings_env_keys = list(user_settings.envs.keys()) for key, value in user_settings.envs.items(): os.environ[key] = value + clear_mps_cache_after_index = configure_mps_environment(user_settings.embedding) # Resolve params BEFORE constructing the embedder so invalid configs # fail fast without paying the model-load cost. try: @@ -684,6 +690,7 @@ def run_daemon( embedder, indexing_params=indexing_params, query_params=query_params, + clear_mps_cache_after_index=clear_mps_cache_after_index, ) sock_path = daemon_socket_path() diff --git a/src/cocoindex_code/project.py b/src/cocoindex_code/project.py index f661c21..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,21 @@ QUERY_EMBED_PARAMS, SQLITE_DB, Embedder, + clear_mps_allocator_cache, ) +logger = logging.getLogger(__name__) + class Project: _env: coco.Environment _app: coco.App[[], None] _project_root: Path _index_lock: asyncio.Lock + _clear_mps_cache_after_index: bool _initial_index_done: asyncio.Event + _initial_index_task: asyncio.Task[None] | None + _initial_index_started: asyncio.Event | None _indexing_stats: IndexingProgress | None = None def close(self) -> None: @@ -67,13 +74,15 @@ async def run_index( on_progress: Callable[[IndexingProgress], None] | None = None, on_started: asyncio.Event | None = None, ) -> None: - """Acquire the index lock, run indexing, and release. + """Acquire the project lock, run indexing, and release it. - If *on_started* is provided, it is set once the lock is acquired - (i.e. indexing has truly begun). On completion (success or failure) - ``_initial_index_done`` is set. + If *on_started* is provided, it is set once the project lock is + acquired. On completion (success or failure) ``_initial_index_done`` + is set. """ async with self._index_lock: + if on_started is not None: + on_started.set() self._indexing_stats = IndexingProgress( num_execution_starts=0, num_unchanged=0, @@ -82,8 +91,6 @@ async def run_index( num_reprocesses=0, num_errors=0, ) - if on_started is not None: - on_started.set() await self._run_index_inner(on_progress=on_progress) async def _run_index_inner( @@ -109,20 +116,32 @@ async def _run_index_inner( on_progress(progress) await asyncio.sleep(0.1) finally: - self._initial_index_done.set() - self._indexing_stats = None + try: + if self._clear_mps_cache_after_index: + try: + await clear_mps_allocator_cache() + except Exception: + logger.exception("Unable to clear the MPS allocator cache after indexing") + finally: + self._initial_index_done.set() + self._indexing_stats = None async def ensure_indexing_started(self) -> None: - """Kick off background indexing and wait until it has actually started. + """Kick off background indexing and wait until it is active or queued. - Returns once the indexing task holds the lock. Safe to call multiple - times — only the first call spawns a task; subsequent calls return - immediately. + Returns once the indexing task holds the project lock. Safe to call + multiple times — only the first call spawns a task; subsequent calls + return immediately. """ if self._initial_index_done.is_set() or self._index_lock.locked(): return - started = asyncio.Event() - asyncio.create_task(self.run_index(on_started=started)) + if self._initial_index_task is None or self._initial_index_task.done(): + self._initial_index_started = asyncio.Event() + self._initial_index_task = asyncio.create_task( + self.run_index(on_started=self._initial_index_started) + ) + started = self._initial_index_started + assert started is not None await started.wait() async def stream_index(self) -> AsyncIterator[IndexStreamResponse]: @@ -263,6 +282,7 @@ async def create( indexing_params: dict[str, Any], query_params: dict[str, Any], chunker_registry: dict[str, ChunkerFn] | None = None, + clear_mps_cache_after_index: bool = False, ) -> Project: """Create a project with explicit embedder and per-call params. @@ -281,6 +301,8 @@ async def create( chunker_registry: Optional mapping of file suffix (e.g. ``".toml"``) to a ``ChunkerFn``. When a suffix matches, the registered chunker is called instead of the built-in splitter. + clear_mps_cache_after_index: Whether to release unused MPS allocator + memory in CocoIndex's GPU subprocess after each index run. """ settings_dir = project_root / ".cocoindex_code" settings_dir.mkdir(parents=True, exist_ok=True) @@ -315,5 +337,8 @@ async def create( result._app = app result._project_root = project_root result._index_lock = asyncio.Lock() + result._clear_mps_cache_after_index = clear_mps_cache_after_index result._initial_index_done = asyncio.Event() + result._initial_index_task = None + result._initial_index_started = None return result diff --git a/src/cocoindex_code/settings.py b/src/cocoindex_code/settings.py index e228839..2bda4a4 100644 --- a/src/cocoindex_code/settings.py +++ b/src/cocoindex_code/settings.py @@ -98,12 +98,30 @@ class EmbeddingSettings: provider: str = "litellm" device: str | None = None min_interval_ms: int | None = None + # PyTorch MPS allocator limits used by CocoIndex's isolated GPU runner. + mps_low_watermark_ratio: float = 0.4 + mps_high_watermark_ratio: float = 0.5 # Extra kwargs spread into ``embedder.embed()`` during indexing/query. # ``None`` means the user did not set the key; ``{}`` is an explicit empty # dict (used to opt out of the legacy-bridge warning). indexing_params: dict[str, Any] | None = None query_params: dict[str, Any] | None = None + def __post_init__(self) -> None: + ratios = ( + ("mps_low_watermark_ratio", self.mps_low_watermark_ratio), + ("mps_high_watermark_ratio", self.mps_high_watermark_ratio), + ) + for name, value in ratios: + if isinstance(value, bool) or not 0 < value <= 1: + raise ValueError(f"embedding.{name} must be greater than 0 and at most 1") + + if self.mps_low_watermark_ratio > self.mps_high_watermark_ratio: + raise ValueError( + "embedding.mps_low_watermark_ratio must be at most " + "embedding.mps_high_watermark_ratio" + ) + @dataclass class DaemonSettings: @@ -430,6 +448,11 @@ def _embedding_settings_to_dict(embedding: EmbeddingSettings) -> dict[str, Any]: d["device"] = embedding.device if embedding.min_interval_ms is not None: d["min_interval_ms"] = embedding.min_interval_ms + defaults = EmbeddingSettings(model=embedding.model) + if embedding.mps_low_watermark_ratio != defaults.mps_low_watermark_ratio: + d["mps_low_watermark_ratio"] = embedding.mps_low_watermark_ratio + if embedding.mps_high_watermark_ratio != defaults.mps_high_watermark_ratio: + d["mps_high_watermark_ratio"] = embedding.mps_high_watermark_ratio if embedding.indexing_params is not None: d["indexing_params"] = dict(embedding.indexing_params) if embedding.query_params is not None: @@ -458,6 +481,10 @@ def _user_settings_from_dict(d: dict[str, Any]) -> UserSettings: emb_kwargs["device"] = emb_dict["device"] if "min_interval_ms" in emb_dict: emb_kwargs["min_interval_ms"] = emb_dict["min_interval_ms"] + if "mps_low_watermark_ratio" in emb_dict: + emb_kwargs["mps_low_watermark_ratio"] = float(emb_dict["mps_low_watermark_ratio"]) + if "mps_high_watermark_ratio" in emb_dict: + emb_kwargs["mps_high_watermark_ratio"] = float(emb_dict["mps_high_watermark_ratio"]) # indexing_params / query_params: missing → None (dataclass default); # present-but-null → {} (treat the same as an empty dict, since both mean # "user acknowledged the key and wants no extra kwargs"). @@ -573,6 +600,13 @@ def save_user_settings(settings: UserSettings) -> Path: ), } +_MPS_SAFETY_COMMENT = ( + " #\n" + " # Apple Silicon MPS allocator limits for CocoIndex's GPU subprocess.\n" + " # mps_low_watermark_ratio: 0.4\n" + " # mps_high_watermark_ratio: 0.5\n" +) + def save_initial_user_settings( embedding: EmbeddingSettings, @@ -595,6 +629,8 @@ def save_initial_user_settings( sort_keys=False, ) content = _INITIAL_HEADER + emb_block + if embedding.provider == "sentence-transformers": + content += _MPS_SAFETY_COMMENT if not defaults_applied: hint = _PARAMS_COMMENT_BY_PROVIDER.get(embedding.provider) if hint is not None: diff --git a/src/cocoindex_code/shared.py b/src/cocoindex_code/shared.py index f607755..08a71ae 100644 --- a/src/cocoindex_code/shared.py +++ b/src/cocoindex_code/shared.py @@ -4,7 +4,9 @@ import importlib.util import logging +import os import pathlib +import sys import traceback as _tb from dataclasses import dataclass from typing import TYPE_CHECKING, Annotated, Any, NamedTuple, Union @@ -26,7 +28,10 @@ DEFAULT_LITELLM_MIN_INTERVAL_MS = 5 # Type alias -Embedder = Union["SentenceTransformerEmbedder", "LiteLLMEmbedder"] +Embedder = Union[ + "SentenceTransformerEmbedder", + "LiteLLMEmbedder", +] # Context keys EMBEDDER = coco.ContextKey[Embedder]("embedder", detect_change=True) @@ -45,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`. @@ -100,6 +145,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,7 +154,7 @@ def create_embedder( if model_name.startswith(SBERT_PREFIX): model_name = model_name[len(SBERT_PREFIX) :] - instance: Embedder = SentenceTransformerEmbedder( + instance = SentenceTransformerEmbedder( model_name, device=settings.device, trust_remote_code=True, 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_settings.py b/tests/test_settings.py index f9be64b..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,6 +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.mps_low_watermark_ratio == 0.4 + assert s.embedding.mps_high_watermark_ratio == 0.5 assert s.envs == {} @@ -77,6 +80,8 @@ def test_save_and_load_user_settings(tmp_path: Path) -> None: model="gemini/text-embedding-004", device="cpu", min_interval_ms=300, + mps_low_watermark_ratio=0.35, + mps_high_watermark_ratio=0.45, ), envs={"GEMINI_API_KEY": "test-key"}, ) @@ -86,9 +91,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.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.envs == settings.envs +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("mps_low_watermark_ratio", 0, "mps_low_watermark_ratio"), + ("mps_high_watermark_ratio", 1.1, "mps_high_watermark_ratio"), + ], +) +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_low_watermark_ratio"): + EmbeddingSettings( + model="model", + 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"], @@ -359,8 +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 "# mps_low_watermark_ratio: 0.4" in content + assert "# mps_high_watermark_ratio: 0.5" 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 2af86f0..88288a2 100644 --- a/tests/test_shared.py +++ b/tests/test_shared.py @@ -2,15 +2,18 @@ from __future__ import annotations +import os from typing import Any import numpy as np import pytest +from cocoindex.ops.sentence_transformers import SentenceTransformerEmbedder from cocoindex_code.litellm_embedder import PacedLiteLLMEmbedder from cocoindex_code.settings import EmbeddingSettings from cocoindex_code.shared import ( check_embedding, + configure_mps_environment, create_embedder, is_sentence_transformers_installed, ) @@ -66,6 +69,93 @@ def test_create_embedder_sentence_transformers_ignores_indexing_params() -> None assert not isinstance(embedder, PacedLiteLLMEmbedder) +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) + + enabled = configure_mps_environment( + EmbeddingSettings( + provider="sentence-transformers", + model="sentence-transformers/all-MiniLM-L6-v2", + device="mps", + ) + ) + + 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" + + +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") + + enabled = configure_mps_environment( + EmbeddingSettings( + provider="sentence-transformers", + model="sentence-transformers/all-MiniLM-L6-v2", + device="mps", + ) + ) + + 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" },