cocoindex-code 0.2.38__tar.gz → 0.2.40__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- cocoindex_code-0.2.38/README.md → cocoindex_code-0.2.40/PKG-INFO +63 -2
- cocoindex_code-0.2.38/PKG-INFO → cocoindex_code-0.2.40/README.md +19 -46
- {cocoindex_code-0.2.38 → cocoindex_code-0.2.40}/pyproject.toml +7 -7
- {cocoindex_code-0.2.38 → cocoindex_code-0.2.40}/src/cocoindex_code/_daemon_paths.py +36 -10
- {cocoindex_code-0.2.38 → cocoindex_code-0.2.40}/src/cocoindex_code/_version.py +2 -2
- {cocoindex_code-0.2.38 → cocoindex_code-0.2.40}/src/cocoindex_code/cli.py +18 -2
- {cocoindex_code-0.2.38 → cocoindex_code-0.2.40}/src/cocoindex_code/client.py +40 -3
- {cocoindex_code-0.2.38 → cocoindex_code-0.2.40}/src/cocoindex_code/daemon.py +8 -1
- {cocoindex_code-0.2.38 → cocoindex_code-0.2.40}/src/cocoindex_code/file_walk.py +36 -2
- {cocoindex_code-0.2.38 → cocoindex_code-0.2.40}/src/cocoindex_code/grep.py +3 -1
- {cocoindex_code-0.2.38 → cocoindex_code-0.2.40}/src/cocoindex_code/indexer.py +3 -1
- {cocoindex_code-0.2.38 → cocoindex_code-0.2.40}/src/cocoindex_code/project.py +39 -14
- {cocoindex_code-0.2.38 → cocoindex_code-0.2.40}/src/cocoindex_code/protocol.py +9 -1
- {cocoindex_code-0.2.38 → cocoindex_code-0.2.40}/src/cocoindex_code/server.py +18 -14
- {cocoindex_code-0.2.38 → cocoindex_code-0.2.40}/src/cocoindex_code/settings.py +102 -2
- {cocoindex_code-0.2.38 → cocoindex_code-0.2.40}/src/cocoindex_code/shared.py +48 -2
- {cocoindex_code-0.2.38 → cocoindex_code-0.2.40}/.gitignore +0 -0
- {cocoindex_code-0.2.38 → cocoindex_code-0.2.40}/LICENSE +0 -0
- {cocoindex_code-0.2.38 → cocoindex_code-0.2.40}/src/cocoindex_code/__init__.py +0 -0
- {cocoindex_code-0.2.38 → cocoindex_code-0.2.40}/src/cocoindex_code/__main__.py +0 -0
- {cocoindex_code-0.2.38 → cocoindex_code-0.2.40}/src/cocoindex_code/chunking.py +0 -0
- {cocoindex_code-0.2.38 → cocoindex_code-0.2.40}/src/cocoindex_code/embedder_defaults.py +0 -0
- {cocoindex_code-0.2.38 → cocoindex_code-0.2.40}/src/cocoindex_code/embedder_params.py +0 -0
- {cocoindex_code-0.2.38 → cocoindex_code-0.2.40}/src/cocoindex_code/litellm_embedder.py +0 -0
- {cocoindex_code-0.2.38 → cocoindex_code-0.2.40}/src/cocoindex_code/query.py +0 -0
- {cocoindex_code-0.2.38 → cocoindex_code-0.2.40}/src/cocoindex_code/schema.py +0 -0
|
@@ -1,3 +1,47 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: cocoindex-code
|
|
3
|
+
Version: 0.2.40
|
|
4
|
+
Summary: MCP server for indexing and querying codebases using CocoIndex
|
|
5
|
+
Project-URL: Homepage, https://github.com/cocoindex-io/cocoindex-code
|
|
6
|
+
Project-URL: Repository, https://github.com/cocoindex-io/cocoindex-code
|
|
7
|
+
Project-URL: Issues, https://github.com/cocoindex-io/cocoindex-code/issues
|
|
8
|
+
License-Expression: Apache-2.0
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
Keywords: cocoindex,codebase,indexing,mcp,vector-search
|
|
11
|
+
Classifier: Development Status :: 3 - Alpha
|
|
12
|
+
Classifier: Intended Audience :: Developers
|
|
13
|
+
Classifier: License :: OSI Approved :: Apache Software License
|
|
14
|
+
Classifier: Programming Language :: Python :: 3
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
18
|
+
Classifier: Topic :: Software Development :: Libraries :: Python Modules
|
|
19
|
+
Requires-Python: >=3.11
|
|
20
|
+
Requires-Dist: cocoindex[litellm]<1.1.0,>=1.0.17
|
|
21
|
+
Requires-Dist: einops>=0.8.2
|
|
22
|
+
Requires-Dist: mcp<3,>=2.0.0
|
|
23
|
+
Requires-Dist: msgspec>=0.19.0
|
|
24
|
+
Requires-Dist: numpy>=1.24.0
|
|
25
|
+
Requires-Dist: pathspec<1.0,>=0.12.1
|
|
26
|
+
Requires-Dist: pydantic>=2.0.0
|
|
27
|
+
Requires-Dist: pyyaml>=6.0
|
|
28
|
+
Requires-Dist: questionary>=2.0.0
|
|
29
|
+
Requires-Dist: sqlite-vec>=0.1.0
|
|
30
|
+
Requires-Dist: typer>=0.9.0
|
|
31
|
+
Provides-Extra: dev
|
|
32
|
+
Requires-Dist: cocoindex[sentence-transformers]<1.1.0,>=1.0.17; extra == 'dev'
|
|
33
|
+
Requires-Dist: mypy>=1.0.0; extra == 'dev'
|
|
34
|
+
Requires-Dist: prek>=0.1.0; extra == 'dev'
|
|
35
|
+
Requires-Dist: pytest-asyncio>=0.21.0; extra == 'dev'
|
|
36
|
+
Requires-Dist: pytest-cov>=4.0.0; extra == 'dev'
|
|
37
|
+
Requires-Dist: pytest>=7.0.0; extra == 'dev'
|
|
38
|
+
Requires-Dist: ruff>=0.1.0; extra == 'dev'
|
|
39
|
+
Provides-Extra: embeddings-local
|
|
40
|
+
Requires-Dist: cocoindex[sentence-transformers]<1.1.0,>=1.0.17; extra == 'embeddings-local'
|
|
41
|
+
Provides-Extra: full
|
|
42
|
+
Requires-Dist: cocoindex[sentence-transformers]<1.1.0,>=1.0.17; extra == 'full'
|
|
43
|
+
Description-Content-Type: text/markdown
|
|
44
|
+
|
|
1
45
|
<p align="center">
|
|
2
46
|
<img width="2428" alt="cocoindex code" src="https://github.com/user-attachments/assets/d05961b4-0b7b-42ea-834a-59c3c01717ca" />
|
|
3
47
|
</p>
|
|
@@ -253,6 +297,7 @@ The background daemon starts automatically on first use.
|
|
|
253
297
|
| `ccc mcp` | Run as MCP server in stdio mode |
|
|
254
298
|
| `ccc doctor` | Run diagnostics — checks settings, daemon, model, file matching, and index health |
|
|
255
299
|
| `ccc reset` | Delete index databases. `--all` also removes settings. `-f` skips confirmation. |
|
|
300
|
+
| `ccc version` | Print the CLI version |
|
|
256
301
|
| `ccc daemon status` | Show daemon version, uptime, and loaded projects |
|
|
257
302
|
| `ccc daemon restart` | Restart the background daemon |
|
|
258
303
|
| `ccc daemon stop` | Stop the daemon |
|
|
@@ -504,6 +549,8 @@ embedding:
|
|
|
504
549
|
model: Snowflake/snowflake-arctic-embed-xs
|
|
505
550
|
device: mps # optional: cpu, cuda, mps (auto-detected if omitted)
|
|
506
551
|
min_interval_ms: 300 # optional: pace LiteLLM embedding requests to reduce 429s; defaults to 5 for LiteLLM
|
|
552
|
+
mps_low_watermark_ratio: 0.4 # optional: PyTorch allocator soft limit
|
|
553
|
+
mps_high_watermark_ratio: 0.5 # optional: PyTorch allocator hard limit
|
|
507
554
|
|
|
508
555
|
# Optional extra kwargs passed to the embedder, separately for indexing vs query.
|
|
509
556
|
# `ccc init` auto-populates these for known models (e.g. Cohere, Voyage, Nvidia NIM,
|
|
@@ -518,11 +565,16 @@ envs: # extra environment variabl
|
|
|
518
565
|
|
|
519
566
|
daemon:
|
|
520
567
|
idle_timeout_minutes: 180 # optional: exit the daemon after this long without client activity (default 180, 0 = never)
|
|
568
|
+
keep_alive_with_mcp: true # optional: keep the daemon warm while an MCP client is connected (default true)
|
|
521
569
|
```
|
|
522
570
|
|
|
523
571
|
> **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.
|
|
524
572
|
|
|
525
|
-
> **
|
|
573
|
+
> **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.
|
|
574
|
+
|
|
575
|
+
> **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.
|
|
576
|
+
|
|
577
|
+
> **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. By default, a live MCP session sends periodic heartbeats so the daemon remains warm while your coding agent is connected. Set `daemon.keep_alive_with_mcp: false` to let the daemon idle-exit during long-lived MCP sessions and release the model between real requests. Set `idle_timeout_minutes: 0` to keep the daemon running forever.
|
|
526
578
|
|
|
527
579
|
> **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).
|
|
528
580
|
|
|
@@ -570,6 +622,8 @@ exclude_patterns:
|
|
|
570
622
|
- "**/dist"
|
|
571
623
|
# ...
|
|
572
624
|
|
|
625
|
+
max_file_size: 500KB # skip files larger than this (default: no limit)
|
|
626
|
+
|
|
573
627
|
language_overrides:
|
|
574
628
|
- ext: inc # treat .inc files as PHP
|
|
575
629
|
lang: php
|
|
@@ -581,7 +635,14 @@ chunkers:
|
|
|
581
635
|
|
|
582
636
|
> `.cocoindex_code/` is automatically added to `.gitignore` during init.
|
|
583
637
|
|
|
584
|
-
|
|
638
|
+
`max_file_size` keeps bundled or generated files out of the index without having
|
|
639
|
+
to enumerate them in `exclude_patterns`. It accepts a plain byte count
|
|
640
|
+
(`1048576`) or a size with a binary unit suffix: `B`, `KB`, `MB`, `GB`
|
|
641
|
+
(case-insensitive, so `500KB` and `500 kb` are the same). The limit is
|
|
642
|
+
inclusive, and omitting the key indexes files of any size. It applies wherever
|
|
643
|
+
project file matching does, so `ccc grep` skips the same files.
|
|
644
|
+
|
|
645
|
+
After editing `include_patterns`, `exclude_patterns`, `max_file_size`, or `language_overrides`:
|
|
585
646
|
|
|
586
647
|
- Run `ccc doctor` to preview which files match.
|
|
587
648
|
- Run `ccc index` or `ccc search --refresh ...` to update the existing index.
|
|
@@ -1,47 +1,3 @@
|
|
|
1
|
-
Metadata-Version: 2.4
|
|
2
|
-
Name: cocoindex-code
|
|
3
|
-
Version: 0.2.38
|
|
4
|
-
Summary: MCP server for indexing and querying codebases using CocoIndex
|
|
5
|
-
Project-URL: Homepage, https://github.com/cocoindex-io/cocoindex-code
|
|
6
|
-
Project-URL: Repository, https://github.com/cocoindex-io/cocoindex-code
|
|
7
|
-
Project-URL: Issues, https://github.com/cocoindex-io/cocoindex-code/issues
|
|
8
|
-
License-Expression: Apache-2.0
|
|
9
|
-
License-File: LICENSE
|
|
10
|
-
Keywords: cocoindex,codebase,indexing,mcp,vector-search
|
|
11
|
-
Classifier: Development Status :: 3 - Alpha
|
|
12
|
-
Classifier: Intended Audience :: Developers
|
|
13
|
-
Classifier: License :: OSI Approved :: Apache Software License
|
|
14
|
-
Classifier: Programming Language :: Python :: 3
|
|
15
|
-
Classifier: Programming Language :: Python :: 3.11
|
|
16
|
-
Classifier: Programming Language :: Python :: 3.12
|
|
17
|
-
Classifier: Programming Language :: Python :: 3.13
|
|
18
|
-
Classifier: Topic :: Software Development :: Libraries :: Python Modules
|
|
19
|
-
Requires-Python: >=3.11
|
|
20
|
-
Requires-Dist: cocoindex[litellm]<1.1.0,>=1.0.13
|
|
21
|
-
Requires-Dist: einops>=0.8.2
|
|
22
|
-
Requires-Dist: mcp>=1.0.0
|
|
23
|
-
Requires-Dist: msgspec>=0.19.0
|
|
24
|
-
Requires-Dist: numpy>=1.24.0
|
|
25
|
-
Requires-Dist: pathspec>=0.12.1
|
|
26
|
-
Requires-Dist: pydantic>=2.0.0
|
|
27
|
-
Requires-Dist: pyyaml>=6.0
|
|
28
|
-
Requires-Dist: questionary>=2.0.0
|
|
29
|
-
Requires-Dist: sqlite-vec>=0.1.0
|
|
30
|
-
Requires-Dist: typer>=0.9.0
|
|
31
|
-
Provides-Extra: dev
|
|
32
|
-
Requires-Dist: cocoindex[sentence-transformers]<1.1.0,>=1.0.13; extra == 'dev'
|
|
33
|
-
Requires-Dist: mypy>=1.0.0; extra == 'dev'
|
|
34
|
-
Requires-Dist: prek>=0.1.0; extra == 'dev'
|
|
35
|
-
Requires-Dist: pytest-asyncio>=0.21.0; extra == 'dev'
|
|
36
|
-
Requires-Dist: pytest-cov>=4.0.0; extra == 'dev'
|
|
37
|
-
Requires-Dist: pytest>=7.0.0; extra == 'dev'
|
|
38
|
-
Requires-Dist: ruff>=0.1.0; extra == 'dev'
|
|
39
|
-
Provides-Extra: embeddings-local
|
|
40
|
-
Requires-Dist: cocoindex[sentence-transformers]<1.1.0,>=1.0.13; extra == 'embeddings-local'
|
|
41
|
-
Provides-Extra: full
|
|
42
|
-
Requires-Dist: cocoindex[sentence-transformers]<1.1.0,>=1.0.13; extra == 'full'
|
|
43
|
-
Description-Content-Type: text/markdown
|
|
44
|
-
|
|
45
1
|
<p align="center">
|
|
46
2
|
<img width="2428" alt="cocoindex code" src="https://github.com/user-attachments/assets/d05961b4-0b7b-42ea-834a-59c3c01717ca" />
|
|
47
3
|
</p>
|
|
@@ -297,6 +253,7 @@ The background daemon starts automatically on first use.
|
|
|
297
253
|
| `ccc mcp` | Run as MCP server in stdio mode |
|
|
298
254
|
| `ccc doctor` | Run diagnostics — checks settings, daemon, model, file matching, and index health |
|
|
299
255
|
| `ccc reset` | Delete index databases. `--all` also removes settings. `-f` skips confirmation. |
|
|
256
|
+
| `ccc version` | Print the CLI version |
|
|
300
257
|
| `ccc daemon status` | Show daemon version, uptime, and loaded projects |
|
|
301
258
|
| `ccc daemon restart` | Restart the background daemon |
|
|
302
259
|
| `ccc daemon stop` | Stop the daemon |
|
|
@@ -548,6 +505,8 @@ embedding:
|
|
|
548
505
|
model: Snowflake/snowflake-arctic-embed-xs
|
|
549
506
|
device: mps # optional: cpu, cuda, mps (auto-detected if omitted)
|
|
550
507
|
min_interval_ms: 300 # optional: pace LiteLLM embedding requests to reduce 429s; defaults to 5 for LiteLLM
|
|
508
|
+
mps_low_watermark_ratio: 0.4 # optional: PyTorch allocator soft limit
|
|
509
|
+
mps_high_watermark_ratio: 0.5 # optional: PyTorch allocator hard limit
|
|
551
510
|
|
|
552
511
|
# Optional extra kwargs passed to the embedder, separately for indexing vs query.
|
|
553
512
|
# `ccc init` auto-populates these for known models (e.g. Cohere, Voyage, Nvidia NIM,
|
|
@@ -562,11 +521,16 @@ envs: # extra environment variabl
|
|
|
562
521
|
|
|
563
522
|
daemon:
|
|
564
523
|
idle_timeout_minutes: 180 # optional: exit the daemon after this long without client activity (default 180, 0 = never)
|
|
524
|
+
keep_alive_with_mcp: true # optional: keep the daemon warm while an MCP client is connected (default true)
|
|
565
525
|
```
|
|
566
526
|
|
|
567
527
|
> **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.
|
|
568
528
|
|
|
569
|
-
> **
|
|
529
|
+
> **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.
|
|
530
|
+
|
|
531
|
+
> **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.
|
|
532
|
+
|
|
533
|
+
> **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. By default, a live MCP session sends periodic heartbeats so the daemon remains warm while your coding agent is connected. Set `daemon.keep_alive_with_mcp: false` to let the daemon idle-exit during long-lived MCP sessions and release the model between real requests. Set `idle_timeout_minutes: 0` to keep the daemon running forever.
|
|
570
534
|
|
|
571
535
|
> **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).
|
|
572
536
|
|
|
@@ -614,6 +578,8 @@ exclude_patterns:
|
|
|
614
578
|
- "**/dist"
|
|
615
579
|
# ...
|
|
616
580
|
|
|
581
|
+
max_file_size: 500KB # skip files larger than this (default: no limit)
|
|
582
|
+
|
|
617
583
|
language_overrides:
|
|
618
584
|
- ext: inc # treat .inc files as PHP
|
|
619
585
|
lang: php
|
|
@@ -625,7 +591,14 @@ chunkers:
|
|
|
625
591
|
|
|
626
592
|
> `.cocoindex_code/` is automatically added to `.gitignore` during init.
|
|
627
593
|
|
|
628
|
-
|
|
594
|
+
`max_file_size` keeps bundled or generated files out of the index without having
|
|
595
|
+
to enumerate them in `exclude_patterns`. It accepts a plain byte count
|
|
596
|
+
(`1048576`) or a size with a binary unit suffix: `B`, `KB`, `MB`, `GB`
|
|
597
|
+
(case-insensitive, so `500KB` and `500 kb` are the same). The limit is
|
|
598
|
+
inclusive, and omitting the key indexes files of any size. It applies wherever
|
|
599
|
+
project file matching does, so `ccc grep` skips the same files.
|
|
600
|
+
|
|
601
|
+
After editing `include_patterns`, `exclude_patterns`, `max_file_size`, or `language_overrides`:
|
|
629
602
|
|
|
630
603
|
- Run `ccc doctor` to preview which files match.
|
|
631
604
|
- Run `ccc index` or `ccc search --refresh ...` to update the existing index.
|
|
@@ -22,15 +22,15 @@ classifiers = [
|
|
|
22
22
|
]
|
|
23
23
|
|
|
24
24
|
dependencies = [
|
|
25
|
-
"mcp>=
|
|
26
|
-
"cocoindex[litellm]>=1.0.
|
|
25
|
+
"mcp>=2.0.0,<3",
|
|
26
|
+
"cocoindex[litellm]>=1.0.17,<1.1.0",
|
|
27
27
|
"sqlite-vec>=0.1.0",
|
|
28
28
|
"pydantic>=2.0.0",
|
|
29
29
|
"numpy>=1.24.0",
|
|
30
30
|
"einops>=0.8.2",
|
|
31
31
|
"typer>=0.9.0",
|
|
32
32
|
"msgspec>=0.19.0",
|
|
33
|
-
"pathspec>=0.12.1",
|
|
33
|
+
"pathspec>=0.12.1,<1.0",
|
|
34
34
|
"pyyaml>=6.0",
|
|
35
35
|
"questionary>=2.0.0",
|
|
36
36
|
]
|
|
@@ -39,7 +39,7 @@ dependencies = [
|
|
|
39
39
|
# `embeddings-local` is the primary feature extra: it pulls in
|
|
40
40
|
# `sentence-transformers` (via cocoindex) so local embeddings work without
|
|
41
41
|
# an API key.
|
|
42
|
-
embeddings-local = ["cocoindex[sentence-transformers]>=1.0.
|
|
42
|
+
embeddings-local = ["cocoindex[sentence-transformers]>=1.0.17,<1.1.0"]
|
|
43
43
|
# `full` is the umbrella "batteries-included" alias. Today it's just
|
|
44
44
|
# `embeddings-local`, but we expect to bundle more optional niceties under
|
|
45
45
|
# 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"]
|
|
|
47
47
|
# `:full` image variant for consistency across install paths. Contents are
|
|
48
48
|
# inlined rather than self-referencing `cocoindex-code[embeddings-local]`
|
|
49
49
|
# to avoid resolver edge cases with older pip.
|
|
50
|
-
full = ["cocoindex[sentence-transformers]>=1.0.
|
|
50
|
+
full = ["cocoindex[sentence-transformers]>=1.0.17,<1.1.0"]
|
|
51
51
|
dev = [
|
|
52
52
|
"pytest>=7.0.0",
|
|
53
53
|
"pytest-asyncio>=0.21.0",
|
|
@@ -55,7 +55,7 @@ dev = [
|
|
|
55
55
|
"ruff>=0.1.0",
|
|
56
56
|
"mypy>=1.0.0",
|
|
57
57
|
"prek>=0.1.0",
|
|
58
|
-
"cocoindex[sentence-transformers]>=1.0.
|
|
58
|
+
"cocoindex[sentence-transformers]>=1.0.17,<1.1.0",
|
|
59
59
|
]
|
|
60
60
|
|
|
61
61
|
[project.scripts]
|
|
@@ -89,7 +89,7 @@ dev = [
|
|
|
89
89
|
"mypy>=1.0.0",
|
|
90
90
|
"prek>=0.1.0",
|
|
91
91
|
"types-pyyaml>=6.0.12.20250915",
|
|
92
|
-
"cocoindex[sentence-transformers]>=1.0.
|
|
92
|
+
"cocoindex[sentence-transformers]>=1.0.17,<1.1.0",
|
|
93
93
|
]
|
|
94
94
|
|
|
95
95
|
[tool.ruff]
|
|
@@ -6,15 +6,22 @@ can import these without pulling in the full daemon stack.
|
|
|
6
6
|
|
|
7
7
|
from __future__ import annotations
|
|
8
8
|
|
|
9
|
+
import hashlib
|
|
9
10
|
import json
|
|
10
11
|
import os
|
|
11
12
|
import sys
|
|
13
|
+
import tempfile
|
|
12
14
|
import time
|
|
13
15
|
from dataclasses import dataclass
|
|
14
16
|
from pathlib import Path
|
|
15
17
|
|
|
16
18
|
from .settings import user_settings_dir
|
|
17
19
|
|
|
20
|
+
# Max bytes in sockaddr_un.sun_path, including the NUL terminator: 104 on the
|
|
21
|
+
# BSDs (incl. macOS), 108 on Linux. Exceeding it makes bind() fail with
|
|
22
|
+
# "AF_UNIX path too long" — unlike ordinary files, which get PATH_MAX (~1024).
|
|
23
|
+
_SUN_PATH_MAX = 104 if sys.platform == "darwin" else 108
|
|
24
|
+
|
|
18
25
|
|
|
19
26
|
def daemon_runtime_dir() -> Path:
|
|
20
27
|
"""Return the directory that holds daemon runtime artifacts.
|
|
@@ -39,18 +46,37 @@ def connection_family() -> str:
|
|
|
39
46
|
return "AF_PIPE" if sys.platform == "win32" else "AF_UNIX"
|
|
40
47
|
|
|
41
48
|
|
|
49
|
+
def _runtime_dir_hash() -> str:
|
|
50
|
+
"""Stable short id for the current runtime dir.
|
|
51
|
+
|
|
52
|
+
Distinguishes daemon instances that differ only by
|
|
53
|
+
``COCOINDEX_CODE_RUNTIME_DIR`` / ``COCOINDEX_CODE_DIR``, so tests, users,
|
|
54
|
+
and containers never collide on one address.
|
|
55
|
+
"""
|
|
56
|
+
return hashlib.md5(str(daemon_runtime_dir()).encode()).hexdigest()[:12]
|
|
57
|
+
|
|
58
|
+
|
|
42
59
|
def daemon_socket_path() -> str:
|
|
43
|
-
"""Return the daemon socket/pipe address.
|
|
60
|
+
"""Return the daemon socket/pipe address.
|
|
61
|
+
|
|
62
|
+
Normally ``<runtime dir>/daemon.sock``. When that exceeds the platform's
|
|
63
|
+
``sun_path`` limit — a deep ``$HOME``, a sandbox, a container path — it
|
|
64
|
+
would fail at bind() with a message that reads like a broken install, so
|
|
65
|
+
fall back to a short path under the temp dir keyed by the runtime dir.
|
|
66
|
+
Client and daemon both resolve through here, so they agree either way.
|
|
67
|
+
"""
|
|
44
68
|
if sys.platform == "win32":
|
|
45
|
-
|
|
46
|
-
|
|
47
|
-
|
|
48
|
-
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
|
|
69
|
+
return rf"\\.\pipe\cocoindex_code_{_runtime_dir_hash()}"
|
|
70
|
+
|
|
71
|
+
path = daemon_runtime_dir() / "daemon.sock"
|
|
72
|
+
if len(str(path).encode()) < _SUN_PATH_MAX:
|
|
73
|
+
return str(path)
|
|
74
|
+
|
|
75
|
+
# gettempdir() honors $TMPDIR, which is per-user and private on macOS; on
|
|
76
|
+
# Linux it is usually shared /tmp, so include the uid to avoid a foreign
|
|
77
|
+
# socket sitting at our address.
|
|
78
|
+
short = Path(tempfile.gettempdir()) / f"ccc-{os.getuid()}-{_runtime_dir_hash()}.sock"
|
|
79
|
+
return str(short)
|
|
54
80
|
|
|
55
81
|
|
|
56
82
|
def daemon_pid_path() -> Path:
|
|
@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
|
|
|
18
18
|
commit_id: str | None
|
|
19
19
|
__commit_id__: str | None
|
|
20
20
|
|
|
21
|
-
__version__ = version = '0.2.
|
|
22
|
-
__version_tuple__ = version_tuple = (0, 2,
|
|
21
|
+
__version__ = version = '0.2.40'
|
|
22
|
+
__version_tuple__ = version_tuple = (0, 2, 40)
|
|
23
23
|
|
|
24
24
|
__commit_id__ = commit_id = None
|
|
@@ -22,6 +22,7 @@ if TYPE_CHECKING:
|
|
|
22
22
|
SearchResponse,
|
|
23
23
|
)
|
|
24
24
|
|
|
25
|
+
from ._version import __version__
|
|
25
26
|
from .settings import (
|
|
26
27
|
DEFAULT_ST_MODEL,
|
|
27
28
|
EmbeddingSettings,
|
|
@@ -195,6 +196,15 @@ def print_index_stats(status: ProjectStatusResponse) -> None:
|
|
|
195
196
|
_typer.echo(f" {lang}: {count} chunks")
|
|
196
197
|
|
|
197
198
|
|
|
199
|
+
def _echo_search_text(text: str) -> None:
|
|
200
|
+
"""Echo result text, replacing characters unsupported by the console codec."""
|
|
201
|
+
try:
|
|
202
|
+
_typer.echo(text)
|
|
203
|
+
except UnicodeEncodeError as error:
|
|
204
|
+
safe_text = text.encode(error.encoding, errors="replace").decode(error.encoding)
|
|
205
|
+
_typer.echo(safe_text)
|
|
206
|
+
|
|
207
|
+
|
|
198
208
|
def print_search_results(response: SearchResponse) -> None:
|
|
199
209
|
"""Print formatted search results."""
|
|
200
210
|
if not response.success:
|
|
@@ -207,8 +217,8 @@ def print_search_results(response: SearchResponse) -> None:
|
|
|
207
217
|
|
|
208
218
|
for i, r in enumerate(response.results, 1):
|
|
209
219
|
_typer.echo(f"\n--- Result {i} (score: {r.score:.3f}) ---")
|
|
210
|
-
|
|
211
|
-
|
|
220
|
+
_echo_search_text(f"File: {r.file_path}:{r.start_line}-{r.end_line} [{r.language}]")
|
|
221
|
+
_echo_search_text(r.content)
|
|
212
222
|
|
|
213
223
|
|
|
214
224
|
def _run_index_with_progress(project_root: str) -> None:
|
|
@@ -1120,6 +1130,12 @@ def daemon_stop() -> None:
|
|
|
1120
1130
|
_typer.echo("Daemon stopped.")
|
|
1121
1131
|
|
|
1122
1132
|
|
|
1133
|
+
@app.command()
|
|
1134
|
+
def version() -> None:
|
|
1135
|
+
"""Print the CLI version."""
|
|
1136
|
+
_typer.echo(__version__)
|
|
1137
|
+
|
|
1138
|
+
|
|
1123
1139
|
@app.command("run-daemon", hidden=True)
|
|
1124
1140
|
def run_daemon_cmd() -> None:
|
|
1125
1141
|
"""Internal: run the daemon process."""
|
|
@@ -18,6 +18,8 @@ from multiprocessing.connection import Client, Connection
|
|
|
18
18
|
from pathlib import Path
|
|
19
19
|
from typing import NamedTuple
|
|
20
20
|
|
|
21
|
+
import msgspec
|
|
22
|
+
|
|
21
23
|
from ._daemon_paths import (
|
|
22
24
|
connection_family,
|
|
23
25
|
daemon_log_path,
|
|
@@ -158,6 +160,16 @@ def _connect_and_handshake() -> Connection:
|
|
|
158
160
|
if _daemon_ensured and not e.resp.ok:
|
|
159
161
|
raise
|
|
160
162
|
stop_daemon()
|
|
163
|
+
except DaemonProtocolError:
|
|
164
|
+
# The reply didn't even decode — a daemon whose wire protocol
|
|
165
|
+
# drifted beyond what the version-mismatch path can express (e.g.
|
|
166
|
+
# still running a pre-upgrade binary). Same treatment as a version
|
|
167
|
+
# mismatch: restart it below; once a matching daemon was already
|
|
168
|
+
# ensured this cannot legitimately happen, so fail fast instead of
|
|
169
|
+
# looping on restarts.
|
|
170
|
+
if _daemon_ensured:
|
|
171
|
+
raise
|
|
172
|
+
stop_daemon()
|
|
161
173
|
except (ConnectionRefusedError, OSError):
|
|
162
174
|
# No daemon answered. Normal on the first call of this process (start
|
|
163
175
|
# one below). If we had already ensured one, decide crash vs graceful
|
|
@@ -238,7 +250,15 @@ def _raw_connect_and_handshake() -> _HandshakeResult:
|
|
|
238
250
|
conn.close()
|
|
239
251
|
raise ConnectionRefusedError(f"Handshake failed: {e}") from e
|
|
240
252
|
|
|
241
|
-
|
|
253
|
+
try:
|
|
254
|
+
resp = decode_response(data)
|
|
255
|
+
except msgspec.DecodeError as e:
|
|
256
|
+
# The daemon speaks an incompatible wire protocol — e.g. a stale
|
|
257
|
+
# pre-upgrade daemon whose handshake reply no longer matches this
|
|
258
|
+
# client's schema. Must not escape as a raw decode error: the caller
|
|
259
|
+
# handles DaemonProtocolError by restarting the daemon (issue #237).
|
|
260
|
+
conn.close()
|
|
261
|
+
raise DaemonProtocolError(f"Undecodable handshake reply from daemon: {e}") from e
|
|
242
262
|
if isinstance(resp, ErrorResponse):
|
|
243
263
|
conn.close()
|
|
244
264
|
raise RuntimeError(f"Daemon error: {resp.message}")
|
|
@@ -274,6 +294,19 @@ class DaemonVersionError(RuntimeError):
|
|
|
274
294
|
super().__init__(message)
|
|
275
295
|
|
|
276
296
|
|
|
297
|
+
class DaemonProtocolError(RuntimeError):
|
|
298
|
+
"""Raised when the daemon's handshake reply cannot be decoded at all.
|
|
299
|
+
|
|
300
|
+
A version-mismatched daemon normally reports itself via ``resp.ok`` being
|
|
301
|
+
False (→ ``DaemonVersionError``), but that requires its reply to decode
|
|
302
|
+
against this client's schema. When even decoding fails, the daemon is
|
|
303
|
+
from a version whose wire protocol drifted too far — handled the same
|
|
304
|
+
way: restart on first contact, fail fast once a matching daemon was
|
|
305
|
+
already ensured. Subclasses RuntimeError so broad daemon-failure handlers
|
|
306
|
+
(e.g. in ``stop_daemon``) cover it.
|
|
307
|
+
"""
|
|
308
|
+
|
|
309
|
+
|
|
277
310
|
class DaemonStartError(RuntimeError):
|
|
278
311
|
"""Raised when the daemon process fails to start.
|
|
279
312
|
|
|
@@ -427,7 +460,7 @@ def send_heartbeat() -> bool:
|
|
|
427
460
|
"""
|
|
428
461
|
try:
|
|
429
462
|
conn, _resp = _raw_connect_and_handshake()
|
|
430
|
-
except (ConnectionRefusedError, DaemonVersionError, OSError):
|
|
463
|
+
except (ConnectionRefusedError, DaemonVersionError, DaemonProtocolError, OSError):
|
|
431
464
|
return False
|
|
432
465
|
try:
|
|
433
466
|
conn.send_bytes(encode_request(HeartbeatRequest()))
|
|
@@ -608,7 +641,11 @@ def stop_daemon() -> None:
|
|
|
608
641
|
conn.recv_bytes()
|
|
609
642
|
finally:
|
|
610
643
|
conn.close()
|
|
611
|
-
except (ConnectionRefusedError, OSError, RuntimeError, DaemonVersionError):
|
|
644
|
+
except (ConnectionRefusedError, OSError, RuntimeError, DaemonVersionError, DaemonProtocolError):
|
|
645
|
+
# A version-mismatched or protocol-incompatible daemon refuses the
|
|
646
|
+
# handshake, so the graceful StopRequest is impossible — fall through
|
|
647
|
+
# to the SIGTERM escalation below, which must never be skipped just
|
|
648
|
+
# because the daemon speaks a different protocol (issue #237).
|
|
612
649
|
pass
|
|
613
650
|
|
|
614
651
|
if _wait_for_daemon_exit(timeout=3.0):
|
|
@@ -73,7 +73,7 @@ from .settings import (
|
|
|
73
73
|
target_sqlite_db_path,
|
|
74
74
|
user_settings_path,
|
|
75
75
|
)
|
|
76
|
-
from .shared import Embedder, check_embedding, create_embedder
|
|
76
|
+
from .shared import Embedder, check_embedding, configure_mps_environment, create_embedder
|
|
77
77
|
|
|
78
78
|
logger = logging.getLogger(__name__)
|
|
79
79
|
|
|
@@ -135,6 +135,7 @@ class ProjectRegistry:
|
|
|
135
135
|
|
|
136
136
|
_projects: dict[str, Project]
|
|
137
137
|
_embedder: Embedder | None
|
|
138
|
+
_clear_mps_cache_after_index: bool
|
|
138
139
|
indexing_params: dict[str, Any]
|
|
139
140
|
query_params: dict[str, Any]
|
|
140
141
|
|
|
@@ -143,9 +144,11 @@ class ProjectRegistry:
|
|
|
143
144
|
embedder: Embedder | None,
|
|
144
145
|
indexing_params: dict[str, Any] | None = None,
|
|
145
146
|
query_params: dict[str, Any] | None = None,
|
|
147
|
+
clear_mps_cache_after_index: bool = False,
|
|
146
148
|
) -> None:
|
|
147
149
|
self._projects = {}
|
|
148
150
|
self._embedder = embedder
|
|
151
|
+
self._clear_mps_cache_after_index = clear_mps_cache_after_index
|
|
149
152
|
self.indexing_params = dict(indexing_params) if indexing_params else {}
|
|
150
153
|
self.query_params = dict(query_params) if query_params else {}
|
|
151
154
|
|
|
@@ -165,6 +168,7 @@ class ProjectRegistry:
|
|
|
165
168
|
indexing_params=self.indexing_params,
|
|
166
169
|
query_params=self.query_params,
|
|
167
170
|
chunker_registry=chunker_registry,
|
|
171
|
+
clear_mps_cache_after_index=self._clear_mps_cache_after_index,
|
|
168
172
|
)
|
|
169
173
|
self._projects[project_root] = project
|
|
170
174
|
return self._projects[project_root]
|
|
@@ -636,6 +640,7 @@ def run_daemon(
|
|
|
636
640
|
embedder: Embedder | None
|
|
637
641
|
indexing_params: dict[str, Any] = {}
|
|
638
642
|
query_params: dict[str, Any] = {}
|
|
643
|
+
clear_mps_cache_after_index = False
|
|
639
644
|
handshake_warnings: list[str] = []
|
|
640
645
|
daemon_settings = DaemonSettings()
|
|
641
646
|
if user_settings_path().is_file():
|
|
@@ -644,6 +649,7 @@ def run_daemon(
|
|
|
644
649
|
settings_env_keys = list(user_settings.envs.keys())
|
|
645
650
|
for key, value in user_settings.envs.items():
|
|
646
651
|
os.environ[key] = value
|
|
652
|
+
clear_mps_cache_after_index = configure_mps_environment(user_settings.embedding)
|
|
647
653
|
# Resolve params BEFORE constructing the embedder so invalid configs
|
|
648
654
|
# fail fast without paying the model-load cost.
|
|
649
655
|
try:
|
|
@@ -684,6 +690,7 @@ def run_daemon(
|
|
|
684
690
|
embedder,
|
|
685
691
|
indexing_params=indexing_params,
|
|
686
692
|
query_params=query_params,
|
|
693
|
+
clear_mps_cache_after_index=clear_mps_cache_after_index,
|
|
687
694
|
)
|
|
688
695
|
|
|
689
696
|
sock_path = daemon_socket_path()
|
|
@@ -120,6 +120,33 @@ class GitignoreAwareMatcher(FilePathMatcher):
|
|
|
120
120
|
return self._delegate.is_file_included(path)
|
|
121
121
|
|
|
122
122
|
|
|
123
|
+
class SizeLimitedMatcher(FilePathMatcher):
|
|
124
|
+
"""Wraps another matcher and drops files larger than ``max_bytes``.
|
|
125
|
+
|
|
126
|
+
The matcher protocol only sees paths relative to the project root, so this
|
|
127
|
+
resolves each candidate against ``project_root`` to stat it. Files that
|
|
128
|
+
cannot be stat'd (a broken symlink, a race with a delete) are left to the
|
|
129
|
+
delegate rather than silently dropped, since a size limit should not be the
|
|
130
|
+
thing that decides an unreadable file's fate.
|
|
131
|
+
"""
|
|
132
|
+
|
|
133
|
+
def __init__(self, delegate: FilePathMatcher, project_root: Path, max_bytes: int) -> None:
|
|
134
|
+
self._delegate = delegate
|
|
135
|
+
self._root = project_root
|
|
136
|
+
self._max_bytes = max_bytes
|
|
137
|
+
|
|
138
|
+
def is_dir_included(self, path: PurePath) -> bool:
|
|
139
|
+
return self._delegate.is_dir_included(path)
|
|
140
|
+
|
|
141
|
+
def is_file_included(self, path: PurePath) -> bool:
|
|
142
|
+
if not self._delegate.is_file_included(path):
|
|
143
|
+
return False
|
|
144
|
+
try:
|
|
145
|
+
return (self._root / path).stat().st_size <= self._max_bytes
|
|
146
|
+
except OSError:
|
|
147
|
+
return True
|
|
148
|
+
|
|
149
|
+
|
|
123
150
|
def find_git_root(start: Path) -> Path | None:
|
|
124
151
|
"""Walk up from ``start`` to the nearest directory holding a ``.git`` entry — a
|
|
125
152
|
directory for a normal repo, or a *file* for a submodule or linked worktree.
|
|
@@ -140,14 +167,21 @@ def build_matcher(
|
|
|
140
167
|
project_root: Path,
|
|
141
168
|
included_patterns: list[str],
|
|
142
169
|
excluded_patterns: list[str],
|
|
170
|
+
max_file_size: int | None = None,
|
|
143
171
|
) -> FilePathMatcher:
|
|
144
172
|
"""Build the project's file matcher: include/exclude globs plus nested
|
|
145
|
-
``.gitignore`` awareness anchored at ``project_root
|
|
173
|
+
``.gitignore`` awareness anchored at ``project_root``, and an optional
|
|
174
|
+
``max_file_size`` cap in bytes."""
|
|
146
175
|
base_matcher = PatternFilePathMatcher(
|
|
147
176
|
included_patterns=included_patterns,
|
|
148
177
|
excluded_patterns=excluded_patterns,
|
|
149
178
|
)
|
|
150
|
-
|
|
179
|
+
matcher: FilePathMatcher = GitignoreAwareMatcher(
|
|
180
|
+
base_matcher, load_gitignore_spec(project_root), project_root
|
|
181
|
+
)
|
|
182
|
+
if max_file_size is not None:
|
|
183
|
+
matcher = SizeLimitedMatcher(matcher, project_root, max_file_size)
|
|
184
|
+
return matcher
|
|
151
185
|
|
|
152
186
|
|
|
153
187
|
def iter_included_files(
|
|
@@ -214,17 +214,19 @@ def _iter_targets(req: GrepRequest, compiler: _PatternCompiler) -> Iterator[_Tar
|
|
|
214
214
|
if project_root is not None:
|
|
215
215
|
ps = load_project_settings(project_root)
|
|
216
216
|
included, excluded = ps.include_patterns, ps.exclude_patterns
|
|
217
|
+
max_file_size = ps.max_file_size
|
|
217
218
|
ext_overrides = {f".{lo.ext}": lo.lang for lo in ps.language_overrides}
|
|
218
219
|
base = project_root
|
|
219
220
|
else:
|
|
220
221
|
included = list(DEFAULT_INCLUDED_PATTERNS)
|
|
221
222
|
excluded = list(DEFAULT_EXCLUDED_PATTERNS)
|
|
222
223
|
ext_overrides = {}
|
|
224
|
+
max_file_size = None
|
|
223
225
|
# Anchor at the enclosing git repo so grepping a subdirectory still honors the
|
|
224
226
|
# repo-root (and intervening) .gitignore; fall back to the target dir itself.
|
|
225
227
|
base = find_git_root(root) or root
|
|
226
228
|
|
|
227
|
-
matcher = build_matcher(base, included, excluded)
|
|
229
|
+
matcher = build_matcher(base, included, excluded, max_file_size)
|
|
228
230
|
path_filter = (
|
|
229
231
|
PatternFilePathMatcher(included_patterns=[req.path_glob]) if req.path_glob else None
|
|
230
232
|
)
|
|
@@ -110,7 +110,9 @@ async def indexer_main() -> None:
|
|
|
110
110
|
),
|
|
111
111
|
)
|
|
112
112
|
|
|
113
|
-
matcher = build_matcher(
|
|
113
|
+
matcher = build_matcher(
|
|
114
|
+
project_root, ps.include_patterns, ps.exclude_patterns, ps.max_file_size
|
|
115
|
+
)
|
|
114
116
|
|
|
115
117
|
files = localfs.walk_dir(
|
|
116
118
|
CODEBASE_DIR,
|
|
@@ -3,6 +3,7 @@
|
|
|
3
3
|
from __future__ import annotations
|
|
4
4
|
|
|
5
5
|
import asyncio
|
|
6
|
+
import logging
|
|
6
7
|
import sqlite3
|
|
7
8
|
from collections.abc import AsyncIterator, Callable
|
|
8
9
|
from pathlib import Path
|
|
@@ -39,15 +40,21 @@ from .shared import (
|
|
|
39
40
|
QUERY_EMBED_PARAMS,
|
|
40
41
|
SQLITE_DB,
|
|
41
42
|
Embedder,
|
|
43
|
+
clear_mps_allocator_cache,
|
|
42
44
|
)
|
|
43
45
|
|
|
46
|
+
logger = logging.getLogger(__name__)
|
|
47
|
+
|
|
44
48
|
|
|
45
49
|
class Project:
|
|
46
50
|
_env: coco.Environment
|
|
47
51
|
_app: coco.App[[], None]
|
|
48
52
|
_project_root: Path
|
|
49
53
|
_index_lock: asyncio.Lock
|
|
54
|
+
_clear_mps_cache_after_index: bool
|
|
50
55
|
_initial_index_done: asyncio.Event
|
|
56
|
+
_initial_index_task: asyncio.Task[None] | None
|
|
57
|
+
_initial_index_started: asyncio.Event | None
|
|
51
58
|
_indexing_stats: IndexingProgress | None = None
|
|
52
59
|
|
|
53
60
|
def close(self) -> None:
|
|
@@ -67,13 +74,15 @@ class Project:
|
|
|
67
74
|
on_progress: Callable[[IndexingProgress], None] | None = None,
|
|
68
75
|
on_started: asyncio.Event | None = None,
|
|
69
76
|
) -> None:
|
|
70
|
-
"""Acquire the
|
|
77
|
+
"""Acquire the project lock, run indexing, and release it.
|
|
71
78
|
|
|
72
|
-
If *on_started* is provided, it is set once the lock is
|
|
73
|
-
|
|
74
|
-
|
|
79
|
+
If *on_started* is provided, it is set once the project lock is
|
|
80
|
+
acquired. On completion (success or failure) ``_initial_index_done``
|
|
81
|
+
is set.
|
|
75
82
|
"""
|
|
76
83
|
async with self._index_lock:
|
|
84
|
+
if on_started is not None:
|
|
85
|
+
on_started.set()
|
|
77
86
|
self._indexing_stats = IndexingProgress(
|
|
78
87
|
num_execution_starts=0,
|
|
79
88
|
num_unchanged=0,
|
|
@@ -82,8 +91,6 @@ class Project:
|
|
|
82
91
|
num_reprocesses=0,
|
|
83
92
|
num_errors=0,
|
|
84
93
|
)
|
|
85
|
-
if on_started is not None:
|
|
86
|
-
on_started.set()
|
|
87
94
|
await self._run_index_inner(on_progress=on_progress)
|
|
88
95
|
|
|
89
96
|
async def _run_index_inner(
|
|
@@ -109,20 +116,32 @@ class Project:
|
|
|
109
116
|
on_progress(progress)
|
|
110
117
|
await asyncio.sleep(0.1)
|
|
111
118
|
finally:
|
|
112
|
-
|
|
113
|
-
|
|
119
|
+
try:
|
|
120
|
+
if self._clear_mps_cache_after_index:
|
|
121
|
+
try:
|
|
122
|
+
await clear_mps_allocator_cache()
|
|
123
|
+
except Exception:
|
|
124
|
+
logger.exception("Unable to clear the MPS allocator cache after indexing")
|
|
125
|
+
finally:
|
|
126
|
+
self._initial_index_done.set()
|
|
127
|
+
self._indexing_stats = None
|
|
114
128
|
|
|
115
129
|
async def ensure_indexing_started(self) -> None:
|
|
116
|
-
"""Kick off background indexing and wait until it
|
|
130
|
+
"""Kick off background indexing and wait until it is active or queued.
|
|
117
131
|
|
|
118
|
-
Returns once the indexing task holds the lock.
|
|
119
|
-
times — only the first call spawns a task; subsequent calls
|
|
120
|
-
immediately.
|
|
132
|
+
Returns once the indexing task holds the project lock. Safe to call
|
|
133
|
+
multiple times — only the first call spawns a task; subsequent calls
|
|
134
|
+
return immediately.
|
|
121
135
|
"""
|
|
122
136
|
if self._initial_index_done.is_set() or self._index_lock.locked():
|
|
123
137
|
return
|
|
124
|
-
|
|
125
|
-
|
|
138
|
+
if self._initial_index_task is None or self._initial_index_task.done():
|
|
139
|
+
self._initial_index_started = asyncio.Event()
|
|
140
|
+
self._initial_index_task = asyncio.create_task(
|
|
141
|
+
self.run_index(on_started=self._initial_index_started)
|
|
142
|
+
)
|
|
143
|
+
started = self._initial_index_started
|
|
144
|
+
assert started is not None
|
|
126
145
|
await started.wait()
|
|
127
146
|
|
|
128
147
|
async def stream_index(self) -> AsyncIterator[IndexStreamResponse]:
|
|
@@ -263,6 +282,7 @@ class Project:
|
|
|
263
282
|
indexing_params: dict[str, Any],
|
|
264
283
|
query_params: dict[str, Any],
|
|
265
284
|
chunker_registry: dict[str, ChunkerFn] | None = None,
|
|
285
|
+
clear_mps_cache_after_index: bool = False,
|
|
266
286
|
) -> Project:
|
|
267
287
|
"""Create a project with explicit embedder and per-call params.
|
|
268
288
|
|
|
@@ -281,6 +301,8 @@ class Project:
|
|
|
281
301
|
chunker_registry: Optional mapping of file suffix (e.g. ``".toml"``)
|
|
282
302
|
to a ``ChunkerFn``. When a suffix matches, the registered
|
|
283
303
|
chunker is called instead of the built-in splitter.
|
|
304
|
+
clear_mps_cache_after_index: Whether to release unused MPS allocator
|
|
305
|
+
memory in CocoIndex's GPU subprocess after each index run.
|
|
284
306
|
"""
|
|
285
307
|
settings_dir = project_root / ".cocoindex_code"
|
|
286
308
|
settings_dir.mkdir(parents=True, exist_ok=True)
|
|
@@ -315,5 +337,8 @@ class Project:
|
|
|
315
337
|
result._app = app
|
|
316
338
|
result._project_root = project_root
|
|
317
339
|
result._index_lock = asyncio.Lock()
|
|
340
|
+
result._clear_mps_cache_after_index = clear_mps_cache_after_index
|
|
318
341
|
result._initial_index_done = asyncio.Event()
|
|
342
|
+
result._initial_index_task = None
|
|
343
|
+
result._initial_index_started = None
|
|
319
344
|
return result
|
|
@@ -79,13 +79,21 @@ Request = (
|
|
|
79
79
|
|
|
80
80
|
|
|
81
81
|
class HandshakeResponse(_msgspec.Struct, tag="handshake"):
|
|
82
|
+
# The handshake is the one message exchanged between *mismatched*
|
|
83
|
+
# client/daemon versions — it is how a mismatch gets detected in the
|
|
84
|
+
# first place. Every field added after a release MUST have a default:
|
|
85
|
+
# a required field makes an older daemon's reply undecodable, which
|
|
86
|
+
# breaks the restart-on-mismatch path and bricks every command
|
|
87
|
+
# (issue #237).
|
|
82
88
|
ok: bool
|
|
83
89
|
daemon_version: str
|
|
84
90
|
# The daemon's process id. The client remembers it so that, when the
|
|
85
91
|
# daemon later vanishes, the graceful-exit marker (written on shutdown
|
|
86
92
|
# with the same pid) can be matched race-free against the exact process
|
|
87
93
|
# the client was talking to — distinguishing a graceful exit from a crash.
|
|
88
|
-
|
|
94
|
+
# None only when decoding the reply of a pre-0.2.38 daemon, whose
|
|
95
|
+
# handshake never satisfies ``ok`` and version checks anyway.
|
|
96
|
+
pid: int | None = None
|
|
89
97
|
global_settings_mtime_us: int | None = None
|
|
90
98
|
# Non-fatal daemon-side warnings surfaced to the client on every handshake.
|
|
91
99
|
# The client dedupes and prints them to stderr (see client._print_handshake_warnings).
|
|
@@ -14,9 +14,10 @@ import json
|
|
|
14
14
|
import os
|
|
15
15
|
from pathlib import Path
|
|
16
16
|
|
|
17
|
-
from mcp.server.
|
|
17
|
+
from mcp.server.mcpserver import MCPServer
|
|
18
18
|
from pydantic import BaseModel, Field
|
|
19
19
|
|
|
20
|
+
from ._version import __version__
|
|
20
21
|
from .settings import DaemonSettings, load_user_settings
|
|
21
22
|
|
|
22
23
|
_MCP_INSTRUCTIONS = (
|
|
@@ -58,9 +59,9 @@ class SearchResultModel(BaseModel):
|
|
|
58
59
|
# === Daemon-backed MCP server factory ===
|
|
59
60
|
|
|
60
61
|
|
|
61
|
-
def create_mcp_server(project_root: str) ->
|
|
62
|
+
def create_mcp_server(project_root: str) -> MCPServer:
|
|
62
63
|
"""Create a lightweight MCP server that delegates to the daemon."""
|
|
63
|
-
mcp =
|
|
64
|
+
mcp = MCPServer("cocoindex-code", instructions=_MCP_INSTRUCTIONS, version=__version__)
|
|
64
65
|
|
|
65
66
|
@mcp.tool(
|
|
66
67
|
name="search",
|
|
@@ -174,24 +175,27 @@ def heartbeat_interval_s(idle_timeout_minutes: int) -> int:
|
|
|
174
175
|
|
|
175
176
|
async def run_heartbeat_loop() -> None:
|
|
176
177
|
"""Periodically heartbeat the daemon so it never idle-exits under this
|
|
177
|
-
live MCP session
|
|
178
|
-
|
|
178
|
+
live MCP session when ``daemon.keep_alive_with_mcp`` is enabled. When the
|
|
179
|
+
option is disabled, the daemon may idle-exit and the next real request
|
|
180
|
+
restarts it transparently. When this process dies (even SIGKILL), the
|
|
181
|
+
heartbeats stop and the daemon exits on its normal idle timeout.
|
|
179
182
|
|
|
180
183
|
Reads the timeout from the same ``global_settings.yml`` the daemon reads
|
|
181
184
|
(dataclass default when the file is missing or invalid). Returns
|
|
182
|
-
immediately when the timeout is 0 (daemon never idle-exits)
|
|
183
|
-
heartbeat itself never starts or restarts a
|
|
184
|
-
``client.send_heartbeat``.
|
|
185
|
+
immediately when the timeout is 0 (daemon never idle-exits) or MCP
|
|
186
|
+
keep-alive is disabled. The heartbeat itself never starts or restarts a
|
|
187
|
+
daemon — see ``client.send_heartbeat``.
|
|
185
188
|
"""
|
|
186
|
-
from .client import send_heartbeat
|
|
187
|
-
|
|
188
189
|
try:
|
|
189
|
-
|
|
190
|
+
daemon_settings = load_user_settings().daemon
|
|
190
191
|
except (FileNotFoundError, ValueError):
|
|
191
|
-
|
|
192
|
-
|
|
192
|
+
daemon_settings = DaemonSettings()
|
|
193
|
+
timeout_minutes = daemon_settings.idle_timeout_minutes
|
|
194
|
+
if timeout_minutes <= 0 or not daemon_settings.keep_alive_with_mcp:
|
|
193
195
|
return
|
|
194
196
|
|
|
197
|
+
from .client import send_heartbeat
|
|
198
|
+
|
|
195
199
|
interval = heartbeat_interval_s(timeout_minutes)
|
|
196
200
|
loop = asyncio.get_event_loop()
|
|
197
201
|
while True:
|
|
@@ -201,7 +205,7 @@ async def run_heartbeat_loop() -> None:
|
|
|
201
205
|
|
|
202
206
|
|
|
203
207
|
# Keep the old `mcp` global for backward compatibility in __init__.py
|
|
204
|
-
mcp:
|
|
208
|
+
mcp: MCPServer | None = None
|
|
205
209
|
|
|
206
210
|
|
|
207
211
|
# === Backward-compatible entry point ===
|
|
@@ -98,18 +98,38 @@ class EmbeddingSettings:
|
|
|
98
98
|
provider: str = "litellm"
|
|
99
99
|
device: str | None = None
|
|
100
100
|
min_interval_ms: int | None = None
|
|
101
|
+
# PyTorch MPS allocator limits used by CocoIndex's isolated GPU runner.
|
|
102
|
+
mps_low_watermark_ratio: float = 0.4
|
|
103
|
+
mps_high_watermark_ratio: float = 0.5
|
|
101
104
|
# Extra kwargs spread into ``embedder.embed()`` during indexing/query.
|
|
102
105
|
# ``None`` means the user did not set the key; ``{}`` is an explicit empty
|
|
103
106
|
# dict (used to opt out of the legacy-bridge warning).
|
|
104
107
|
indexing_params: dict[str, Any] | None = None
|
|
105
108
|
query_params: dict[str, Any] | None = None
|
|
106
109
|
|
|
110
|
+
def __post_init__(self) -> None:
|
|
111
|
+
ratios = (
|
|
112
|
+
("mps_low_watermark_ratio", self.mps_low_watermark_ratio),
|
|
113
|
+
("mps_high_watermark_ratio", self.mps_high_watermark_ratio),
|
|
114
|
+
)
|
|
115
|
+
for name, value in ratios:
|
|
116
|
+
if isinstance(value, bool) or not 0 < value <= 1:
|
|
117
|
+
raise ValueError(f"embedding.{name} must be greater than 0 and at most 1")
|
|
118
|
+
|
|
119
|
+
if self.mps_low_watermark_ratio > self.mps_high_watermark_ratio:
|
|
120
|
+
raise ValueError(
|
|
121
|
+
"embedding.mps_low_watermark_ratio must be at most "
|
|
122
|
+
"embedding.mps_high_watermark_ratio"
|
|
123
|
+
)
|
|
124
|
+
|
|
107
125
|
|
|
108
126
|
@dataclass
|
|
109
127
|
class DaemonSettings:
|
|
110
128
|
# Minutes without client activity before the daemon exits (0 = never exit).
|
|
111
129
|
# Clients auto-restart the daemon on the next request, so exiting is cheap.
|
|
112
130
|
idle_timeout_minutes: int = 180
|
|
131
|
+
# Keep the daemon warm while a long-lived MCP client is connected.
|
|
132
|
+
keep_alive_with_mcp: bool = True
|
|
113
133
|
|
|
114
134
|
|
|
115
135
|
@dataclass
|
|
@@ -137,6 +157,8 @@ class ProjectSettings:
|
|
|
137
157
|
exclude_patterns: list[str] = field(default_factory=lambda: list(DEFAULT_EXCLUDED_PATTERNS))
|
|
138
158
|
language_overrides: list[LanguageOverride] = field(default_factory=list)
|
|
139
159
|
chunkers: list[ChunkerMapping] = field(default_factory=list)
|
|
160
|
+
#: Files larger than this many bytes are excluded. ``None`` means no limit.
|
|
161
|
+
max_file_size: int | None = None
|
|
140
162
|
|
|
141
163
|
|
|
142
164
|
# ---------------------------------------------------------------------------
|
|
@@ -430,6 +452,11 @@ def _embedding_settings_to_dict(embedding: EmbeddingSettings) -> dict[str, Any]:
|
|
|
430
452
|
d["device"] = embedding.device
|
|
431
453
|
if embedding.min_interval_ms is not None:
|
|
432
454
|
d["min_interval_ms"] = embedding.min_interval_ms
|
|
455
|
+
defaults = EmbeddingSettings(model=embedding.model)
|
|
456
|
+
if embedding.mps_low_watermark_ratio != defaults.mps_low_watermark_ratio:
|
|
457
|
+
d["mps_low_watermark_ratio"] = embedding.mps_low_watermark_ratio
|
|
458
|
+
if embedding.mps_high_watermark_ratio != defaults.mps_high_watermark_ratio:
|
|
459
|
+
d["mps_high_watermark_ratio"] = embedding.mps_high_watermark_ratio
|
|
433
460
|
if embedding.indexing_params is not None:
|
|
434
461
|
d["indexing_params"] = dict(embedding.indexing_params)
|
|
435
462
|
if embedding.query_params is not None:
|
|
@@ -441,8 +468,14 @@ def _user_settings_to_dict(settings: UserSettings) -> dict[str, Any]:
|
|
|
441
468
|
d: dict[str, Any] = {"embedding": _embedding_settings_to_dict(settings.embedding)}
|
|
442
469
|
if settings.envs:
|
|
443
470
|
d["envs"] = dict(settings.envs)
|
|
444
|
-
|
|
445
|
-
|
|
471
|
+
daemon_defaults = DaemonSettings()
|
|
472
|
+
daemon_dict: dict[str, Any] = {}
|
|
473
|
+
if settings.daemon.idle_timeout_minutes != daemon_defaults.idle_timeout_minutes:
|
|
474
|
+
daemon_dict["idle_timeout_minutes"] = settings.daemon.idle_timeout_minutes
|
|
475
|
+
if settings.daemon.keep_alive_with_mcp != daemon_defaults.keep_alive_with_mcp:
|
|
476
|
+
daemon_dict["keep_alive_with_mcp"] = settings.daemon.keep_alive_with_mcp
|
|
477
|
+
if daemon_dict:
|
|
478
|
+
d["daemon"] = daemon_dict
|
|
446
479
|
return d
|
|
447
480
|
|
|
448
481
|
|
|
@@ -458,6 +491,10 @@ def _user_settings_from_dict(d: dict[str, Any]) -> UserSettings:
|
|
|
458
491
|
emb_kwargs["device"] = emb_dict["device"]
|
|
459
492
|
if "min_interval_ms" in emb_dict:
|
|
460
493
|
emb_kwargs["min_interval_ms"] = emb_dict["min_interval_ms"]
|
|
494
|
+
if "mps_low_watermark_ratio" in emb_dict:
|
|
495
|
+
emb_kwargs["mps_low_watermark_ratio"] = float(emb_dict["mps_low_watermark_ratio"])
|
|
496
|
+
if "mps_high_watermark_ratio" in emb_dict:
|
|
497
|
+
emb_kwargs["mps_high_watermark_ratio"] = float(emb_dict["mps_high_watermark_ratio"])
|
|
461
498
|
# indexing_params / query_params: missing → None (dataclass default);
|
|
462
499
|
# present-but-null → {} (treat the same as an empty dict, since both mean
|
|
463
500
|
# "user acknowledged the key and wants no extra kwargs").
|
|
@@ -472,15 +509,66 @@ def _user_settings_from_dict(d: dict[str, Any]) -> UserSettings:
|
|
|
472
509
|
daemon_kwargs: dict[str, Any] = {}
|
|
473
510
|
if "idle_timeout_minutes" in daemon_dict:
|
|
474
511
|
daemon_kwargs["idle_timeout_minutes"] = int(daemon_dict["idle_timeout_minutes"])
|
|
512
|
+
if "keep_alive_with_mcp" in daemon_dict:
|
|
513
|
+
keep_alive_with_mcp = daemon_dict["keep_alive_with_mcp"]
|
|
514
|
+
if not isinstance(keep_alive_with_mcp, bool):
|
|
515
|
+
raise ValueError("daemon.keep_alive_with_mcp must be a boolean")
|
|
516
|
+
daemon_kwargs["keep_alive_with_mcp"] = keep_alive_with_mcp
|
|
475
517
|
daemon = DaemonSettings(**daemon_kwargs)
|
|
476
518
|
return UserSettings(embedding=embedding, envs=envs, daemon=daemon)
|
|
477
519
|
|
|
478
520
|
|
|
521
|
+
_SIZE_UNITS: dict[str, int] = {
|
|
522
|
+
"B": 1,
|
|
523
|
+
"KB": 1024,
|
|
524
|
+
"MB": 1024**2,
|
|
525
|
+
"GB": 1024**3,
|
|
526
|
+
}
|
|
527
|
+
|
|
528
|
+
|
|
529
|
+
def parse_file_size(value: Any) -> int:
|
|
530
|
+
"""Parse a ``max_file_size`` value into bytes.
|
|
531
|
+
|
|
532
|
+
Accepts a plain number of bytes (``1048576``) or a string with an optional
|
|
533
|
+
unit suffix (``500KB``, ``1.5 MB``, ``2048``). Units are binary (1KB = 1024
|
|
534
|
+
bytes) and case-insensitive.
|
|
535
|
+
"""
|
|
536
|
+
if isinstance(value, bool): # bool is an int subclass; reject it explicitly.
|
|
537
|
+
raise ValueError(f"invalid max_file_size: {value!r}")
|
|
538
|
+
if isinstance(value, int):
|
|
539
|
+
size = value
|
|
540
|
+
elif isinstance(value, float):
|
|
541
|
+
size = int(value)
|
|
542
|
+
elif isinstance(value, str):
|
|
543
|
+
text = value.strip().upper()
|
|
544
|
+
if not text:
|
|
545
|
+
raise ValueError("invalid max_file_size: empty value")
|
|
546
|
+
for suffix in ("GB", "MB", "KB", "B"):
|
|
547
|
+
if text.endswith(suffix):
|
|
548
|
+
number = text[: -len(suffix)].strip()
|
|
549
|
+
multiplier = _SIZE_UNITS[suffix]
|
|
550
|
+
break
|
|
551
|
+
else:
|
|
552
|
+
number, multiplier = text, 1
|
|
553
|
+
try:
|
|
554
|
+
size = int(float(number) * multiplier)
|
|
555
|
+
except ValueError as exc:
|
|
556
|
+
raise ValueError(f"invalid max_file_size: {value!r}") from exc
|
|
557
|
+
else:
|
|
558
|
+
raise ValueError(f"invalid max_file_size: {value!r}")
|
|
559
|
+
|
|
560
|
+
if size <= 0:
|
|
561
|
+
raise ValueError(f"max_file_size must be positive, got {value!r}")
|
|
562
|
+
return size
|
|
563
|
+
|
|
564
|
+
|
|
479
565
|
def _project_settings_to_dict(settings: ProjectSettings) -> dict[str, Any]:
|
|
480
566
|
d: dict[str, Any] = {
|
|
481
567
|
"include_patterns": settings.include_patterns,
|
|
482
568
|
"exclude_patterns": settings.exclude_patterns,
|
|
483
569
|
}
|
|
570
|
+
if settings.max_file_size is not None:
|
|
571
|
+
d["max_file_size"] = settings.max_file_size
|
|
484
572
|
if settings.language_overrides:
|
|
485
573
|
d["language_overrides"] = [
|
|
486
574
|
{"ext": lo.ext, "lang": lo.lang} for lo in settings.language_overrides
|
|
@@ -495,11 +583,14 @@ def _project_settings_from_dict(d: dict[str, Any]) -> ProjectSettings:
|
|
|
495
583
|
LanguageOverride(ext=lo["ext"], lang=lo["lang"]) for lo in d.get("language_overrides", [])
|
|
496
584
|
]
|
|
497
585
|
chunkers = [ChunkerMapping(ext=cm["ext"], module=cm["module"]) for cm in d.get("chunkers", [])]
|
|
586
|
+
raw_max_size = d.get("max_file_size")
|
|
587
|
+
max_file_size = None if raw_max_size is None else parse_file_size(raw_max_size)
|
|
498
588
|
return ProjectSettings(
|
|
499
589
|
include_patterns=d.get("include_patterns", list(DEFAULT_INCLUDED_PATTERNS)),
|
|
500
590
|
exclude_patterns=d.get("exclude_patterns", list(DEFAULT_EXCLUDED_PATTERNS)),
|
|
501
591
|
language_overrides=overrides,
|
|
502
592
|
chunkers=chunkers,
|
|
593
|
+
max_file_size=max_file_size,
|
|
503
594
|
)
|
|
504
595
|
|
|
505
596
|
|
|
@@ -573,6 +664,13 @@ _PARAMS_COMMENT_BY_PROVIDER: dict[str, str] = {
|
|
|
573
664
|
),
|
|
574
665
|
}
|
|
575
666
|
|
|
667
|
+
_MPS_SAFETY_COMMENT = (
|
|
668
|
+
" #\n"
|
|
669
|
+
" # Apple Silicon MPS allocator limits for CocoIndex's GPU subprocess.\n"
|
|
670
|
+
" # mps_low_watermark_ratio: 0.4\n"
|
|
671
|
+
" # mps_high_watermark_ratio: 0.5\n"
|
|
672
|
+
)
|
|
673
|
+
|
|
576
674
|
|
|
577
675
|
def save_initial_user_settings(
|
|
578
676
|
embedding: EmbeddingSettings,
|
|
@@ -595,6 +693,8 @@ def save_initial_user_settings(
|
|
|
595
693
|
sort_keys=False,
|
|
596
694
|
)
|
|
597
695
|
content = _INITIAL_HEADER + emb_block
|
|
696
|
+
if embedding.provider == "sentence-transformers":
|
|
697
|
+
content += _MPS_SAFETY_COMMENT
|
|
598
698
|
if not defaults_applied:
|
|
599
699
|
hint = _PARAMS_COMMENT_BY_PROVIDER.get(embedding.provider)
|
|
600
700
|
if hint is not None:
|
|
@@ -4,7 +4,9 @@ from __future__ import annotations
|
|
|
4
4
|
|
|
5
5
|
import importlib.util
|
|
6
6
|
import logging
|
|
7
|
+
import os
|
|
7
8
|
import pathlib
|
|
9
|
+
import sys
|
|
8
10
|
import traceback as _tb
|
|
9
11
|
from dataclasses import dataclass
|
|
10
12
|
from typing import TYPE_CHECKING, Annotated, Any, NamedTuple, Union
|
|
@@ -26,7 +28,10 @@ SBERT_PREFIX = "sbert/"
|
|
|
26
28
|
DEFAULT_LITELLM_MIN_INTERVAL_MS = 5
|
|
27
29
|
|
|
28
30
|
# Type alias
|
|
29
|
-
Embedder = Union[
|
|
31
|
+
Embedder = Union[
|
|
32
|
+
"SentenceTransformerEmbedder",
|
|
33
|
+
"LiteLLMEmbedder",
|
|
34
|
+
]
|
|
30
35
|
|
|
31
36
|
# Context keys
|
|
32
37
|
EMBEDDER = coco.ContextKey[Embedder]("embedder", detect_change=True)
|
|
@@ -45,6 +50,46 @@ def is_sentence_transformers_installed() -> bool:
|
|
|
45
50
|
return importlib.util.find_spec("sentence_transformers") is not None
|
|
46
51
|
|
|
47
52
|
|
|
53
|
+
def configure_mps_environment(settings: EmbeddingSettings) -> bool:
|
|
54
|
+
"""Install MPS safety defaults before CocoIndex's first GPU invocation.
|
|
55
|
+
|
|
56
|
+
CocoIndex reads ``COCOINDEX_RUN_GPU_IN_SUBPROCESS`` lazily on the first
|
|
57
|
+
``coco.GPU`` call, while PyTorch reads the allocator watermarks when the
|
|
58
|
+
child initializes MPS. Explicit user environment variables always win.
|
|
59
|
+
|
|
60
|
+
Returns whether this embedding configuration uses the guarded MPS path.
|
|
61
|
+
"""
|
|
62
|
+
|
|
63
|
+
use_mps = settings.provider == "sentence-transformers" and (
|
|
64
|
+
settings.device == "mps" or (settings.device is None and sys.platform == "darwin")
|
|
65
|
+
)
|
|
66
|
+
if not use_mps:
|
|
67
|
+
return False
|
|
68
|
+
|
|
69
|
+
os.environ.setdefault("COCOINDEX_RUN_GPU_IN_SUBPROCESS", "1")
|
|
70
|
+
os.environ.setdefault("PYTORCH_MPS_LOW_WATERMARK_RATIO", str(settings.mps_low_watermark_ratio))
|
|
71
|
+
os.environ.setdefault(
|
|
72
|
+
"PYTORCH_MPS_HIGH_WATERMARK_RATIO", str(settings.mps_high_watermark_ratio)
|
|
73
|
+
)
|
|
74
|
+
return os.environ.get("COCOINDEX_RUN_GPU_IN_SUBPROCESS") == "1"
|
|
75
|
+
|
|
76
|
+
|
|
77
|
+
@coco.fn.as_async(runner=coco.GPU)
|
|
78
|
+
def clear_mps_allocator_cache() -> None:
|
|
79
|
+
"""Release unused MPS allocator cache inside CocoIndex's GPU child."""
|
|
80
|
+
|
|
81
|
+
import gc
|
|
82
|
+
|
|
83
|
+
import torch
|
|
84
|
+
|
|
85
|
+
if not torch.backends.mps.is_available():
|
|
86
|
+
return
|
|
87
|
+
torch.mps.synchronize()
|
|
88
|
+
gc.collect()
|
|
89
|
+
torch.mps.empty_cache()
|
|
90
|
+
torch.mps.synchronize()
|
|
91
|
+
|
|
92
|
+
|
|
48
93
|
class EmbeddingCheckResult(NamedTuple):
|
|
49
94
|
"""Outcome of a single embed-test call. See `check_embedding`.
|
|
50
95
|
|
|
@@ -100,6 +145,7 @@ def create_embedder(
|
|
|
100
145
|
only and the indexing default is supplied at the call site via
|
|
101
146
|
:data:`INDEXING_EMBED_PARAMS`.
|
|
102
147
|
"""
|
|
148
|
+
instance: Embedder
|
|
103
149
|
if settings.provider == "sentence-transformers":
|
|
104
150
|
from cocoindex.ops.sentence_transformers import SentenceTransformerEmbedder
|
|
105
151
|
|
|
@@ -108,7 +154,7 @@ def create_embedder(
|
|
|
108
154
|
if model_name.startswith(SBERT_PREFIX):
|
|
109
155
|
model_name = model_name[len(SBERT_PREFIX) :]
|
|
110
156
|
|
|
111
|
-
instance
|
|
157
|
+
instance = SentenceTransformerEmbedder(
|
|
112
158
|
model_name,
|
|
113
159
|
device=settings.device,
|
|
114
160
|
trust_remote_code=True,
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|