diff --git a/README.md b/README.md index b062ce7f..545f378e 100644 --- a/README.md +++ b/README.md @@ -232,6 +232,8 @@ result.chunk.end_line # 150 result.chunk.content # "def save_pretrained(self, path: PathLike, ..." ``` +When many files need indexing, the CLI and MCP server chunk them in 4 worker processes (`SEMBLE_INDEX_WORKERS` sets the number, `0` turns them off). In the Python API they are off by default. To turn them on, set `SEMBLE_INDEX_WORKERS` and call Semble from inside an `if __name__ == "__main__":` block, since each worker re-imports your script. + ## MCP Server diff --git a/src/semble/cli.py b/src/semble/cli.py index faa094e8..94e7123c 100644 --- a/src/semble/cli.py +++ b/src/semble/cli.py @@ -95,6 +95,8 @@ def main() -> None: for stream in (sys.stdout, sys.stderr): if isinstance(stream, io.TextIOWrapper): stream.reconfigure(errors="replace") + # Worker processes re-import __main__, so only guarded entry points like this one turn them on by default. + os.environ.setdefault("SEMBLE_INDEX_WORKERS", "4") if len(sys.argv) > 1 and sys.argv[1] in _CLI_DISPATCH_ARGS: _cli_main() else: diff --git a/src/semble/index/create.py b/src/semble/index/create.py index 9026cf1b..4c354779 100644 --- a/src/semble/index/create.py +++ b/src/semble/index/create.py @@ -1,7 +1,15 @@ import contextlib +import itertools import logging +import multiprocessing +import multiprocessing.connection +import os +import signal import sys -from collections.abc import Sequence +import threading +from collections import deque +from collections.abc import Generator, Sequence +from concurrent.futures import ProcessPoolExecutor from pathlib import Path import numpy as np @@ -28,6 +36,14 @@ logger = logging.getLogger(__name__) +# Below this many files to chunk, SEMBLE_INDEX_WORKERS is ignored and files are chunked in the main process. +_MIN_FILES_FOR_PROCESSES = 200 +# Files handed to the workers but not yet consumed by the main process. +_MAX_FILES_IN_FLIGHT = 256 + +# A file's chunks and their BM25 tokens. +_ChunkedFile = tuple[list[Chunk], list[list[str]]] + def _warn_skipped_large(skipped_large: list[str]) -> None: """Warn about files skipped for exceeding the maximum indexable file size.""" @@ -42,18 +58,83 @@ def _warn_skipped_large(skipped_large: list[str]) -> None: ) +def _chunk_file(file_path: Path, indexed_path: str) -> _ChunkedFile | None: + """Chunk a file and tokenize its chunks for BM25, or return None if it can't be read.""" + try: + source = read_file_text(file_path) + except OSError: + return None + file_chunks = chunk_source(source, indexed_path, detect_language(file_path)) + return file_chunks, [tokenize(enrich_for_bm25(chunk)) for chunk in file_chunks] + + +def _init_worker() -> None: + """Leave Ctrl-C to the parent, and exit this worker once the parent is gone, even if it never shut the pool down.""" + signal.signal(signal.SIGINT, signal.SIG_IGN) + parent = multiprocessing.parent_process() + assert parent is not None + + def wait_and_exit() -> None: + multiprocessing.connection.wait([parent.sentinel]) + os._exit(0) + + threading.Thread(target=wait_and_exit, daemon=True).start() + + +def _chunk_files(files: list[tuple[Path, str]]) -> Generator[_ChunkedFile | None, None, None]: + """Chunk files in order, in SEMBLE_INDEX_WORKERS worker processes when there are enough files.""" + workers = int(os.environ.get("SEMBLE_INDEX_WORKERS", 0)) + if workers == 0 or len(files) < _MIN_FILES_FOR_PROCESSES: + yield from itertools.starmap(_chunk_file, files) + else: + executor = ProcessPoolExecutor( + workers, mp_context=multiprocessing.get_context("spawn"), initializer=_init_worker + ) + try: + submitted = (executor.submit(_chunk_file, *file) for file in files) + pending = deque(itertools.islice(submitted, _MAX_FILES_IN_FLIGHT)) + while pending: + yield pending.popleft().result() + pending.extend(itertools.islice(submitted, 1)) + finally: + executor.shutdown(cancel_futures=True) + + +def _stat_files( + path: Path, + extensions: Sequence[str], + display_root: Path | None, + previous_manifest: dict[str, FileManifestEntry], + skipped_large: list[str], +) -> list[tuple[Path, str, int, FileManifestEntry | None]]: + """Return each indexable file's path, indexed path, mtime and previous manifest entry.""" + files = [] + for file_path in walk_files(path, extensions): + with contextlib.suppress(OSError): + stat = file_path.stat() + file_status = get_file_status(file_path, stat) + if file_status is FileStatus.TOO_LARGE: + skipped_large.append(str(file_path)) + if file_status != FileStatus.VALID: + continue + + indexed_path = str(file_path.relative_to(display_root) if display_root else file_path) + files.append((file_path, indexed_path, stat.st_mtime_ns, previous_manifest.get(indexed_path))) + return files + + def _reindex_file( bm25_index: BM25, indexed_path: str, - file_chunks: list[Chunk], + file_tokens: list[list[str]], previous_entry: FileManifestEntry | None, ) -> None: """Replace a file's BM25 postings: remove its old slots (if any), then add its new ones.""" if previous_entry is not None: for slot in range(previous_entry.count): bm25_index.remove_document(make_chunk_id(indexed_path, slot)) - for slot, chunk in enumerate(file_chunks): - bm25_index.add_document(make_chunk_id(indexed_path, slot), tokenize(enrich_for_bm25(chunk))) + for slot, tokens in enumerate(file_tokens): + bm25_index.add_document(make_chunk_id(indexed_path, slot), tokens) def _has_same_vector_layout( @@ -102,37 +183,33 @@ def create_index_from_path( skipped_large: list[str] = [] - files = list(walk_files(path, resolved_extensions)) - for file_path in tqdm( - files, - desc="Indexing", - unit="file", - file=sys.stderr, - leave=False, - colour="green", - miniters=1, # tqdm's adaptive miniters stalls the bar when fast files are followed by slow ones - disable=not show_progress_bar, - ): - language = detect_language(file_path) - with contextlib.suppress(OSError): - stat = file_path.stat() - file_status = get_file_status(file_path, stat) - if file_status is FileStatus.TOO_LARGE: - skipped_large.append(str(file_path)) - if file_status != FileStatus.VALID: - continue - - indexed_path = str(file_path.relative_to(display_root) if display_root else file_path) - mtime_ns = stat.st_mtime_ns - previous_entry = previous_manifest.get(indexed_path) - + files = _stat_files(path, resolved_extensions, display_root, previous_manifest, skipped_large) + stale = [ + (file_path, indexed_path) + for file_path, indexed_path, mtime_ns, previous_entry in files + if previous_entry is None or previous_entry.mtime_ns != mtime_ns + ] + + with contextlib.closing(_chunk_files(stale)) as chunked: + for _, indexed_path, mtime_ns, previous_entry in tqdm( + files, + desc="Indexing", + unit="file", + file=sys.stderr, + leave=False, + colour="green", + miniters=1, # tqdm's adaptive miniters stalls the bar when fast files are followed by slow ones + disable=not show_progress_bar, + ): if previous is not None and previous_entry is not None and previous_entry.mtime_ns == mtime_ns: file_chunks = previous.chunks[previous_entry.start : previous_entry.end] vector_parts.append(previous.vectors[previous_entry.start : previous_entry.end]) else: - source = read_file_text(file_path) - file_chunks = chunk_source(source, indexed_path, language) - _reindex_file(bm25_index, indexed_path, file_chunks, previous_entry) + result = next(chunked) + if result is None: + continue + file_chunks, file_tokens = result + _reindex_file(bm25_index, indexed_path, file_tokens, previous_entry) embedding_parts.append((len(vector_parts), len(chunks), len(file_chunks))) vector_parts.append(embed_chunks(model, file_chunks)) diff --git a/tests/index/test_create.py b/tests/index/test_create.py index cb280a3e..15b0936f 100644 --- a/tests/index/test_create.py +++ b/tests/index/test_create.py @@ -1,4 +1,7 @@ +import multiprocessing +import threading from pathlib import Path +from types import SimpleNamespace from typing import Any from unittest.mock import patch @@ -7,6 +10,7 @@ import pytest from semble.cache import load_previous_for_incremental +from semble.index import create from semble.index.bm25 import BM25 from semble.index.create import create_index_from_path from semble.index.index import SembleIndex @@ -21,8 +25,14 @@ def _write_files(root: Path, files: dict[str, str]) -> None: path.write_text(content) -def test_incremental_reindex_reuses_updates_and_prunes(mock_model: Any, tmp_path: Path) -> None: +@pytest.mark.parametrize("in_processes", [False, True]) +def test_incremental_reindex_reuses_updates_and_prunes( + mock_model: Any, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, in_processes: bool +) -> None: """One incremental pass reuses unchanged vectors, re-embeds changes, and keeps BM25 slots current.""" + if in_processes: + monkeypatch.setattr(create, "_MIN_FILES_FOR_PROCESSES", 0) + monkeypatch.setenv("SEMBLE_INDEX_WORKERS", "2") _write_files( tmp_path, { @@ -89,6 +99,35 @@ def test_incremental_reindex_reuses_updates_and_prunes(mock_model: Any, tmp_path assert set(bm25_after.doc_order) == expected_ids +def test_unreadable_file_is_skipped(mock_model: Any, tmp_path: Path) -> None: + """A file that can't be read is left out of the index.""" + _write_files(tmp_path, {"a.py": "def a():\n return 1\n", "b.py": "def b():\n return 2\n"}) + read_file_text = create.read_file_text + + def read_or_fail(file_path: Path) -> str: + if file_path.name == "b.py": + raise PermissionError(file_path) + return read_file_text(file_path) + + with patch.object(create, "read_file_text", read_or_fail): + _, _, _, manifest = create_index_from_path(tmp_path, mock_model, display_root=tmp_path) + assert list(manifest) == ["a.py"] + + +def test_worker_exits_with_parent(monkeypatch: pytest.MonkeyPatch) -> None: + """A chunking worker exits once its parent process is gone.""" + parent_end, worker_end = multiprocessing.Pipe() + exited = threading.Event() + monkeypatch.setattr(multiprocessing, "parent_process", lambda: SimpleNamespace(sentinel=worker_end)) + monkeypatch.setattr(create.os, "_exit", lambda _: exited.set()) + monkeypatch.setattr(create.signal, "signal", lambda *_: None) + + create._init_worker() + assert not exited.wait(0.1) + parent_end.close() + assert exited.wait(5) + + def _build_valid_cache(index_path: Path, mock_model: Any) -> dict: """Build a real, well-formed on-disk index and return its metadata dict for mutation.""" src = index_path.parent / "src"