diff --git a/dependencies.yaml b/dependencies.yaml index 970b173328..07529accc8 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -618,6 +618,7 @@ dependencies: packages: - click - cuvs==26.10.*,>=0.0.0a0 + - h5py>=3.8.0 - pandas - pyyaml - requests diff --git a/fern/pages/cuvs_bench/install.md b/fern/pages/cuvs_bench/install.md index 9c525b4973..bec57d8d29 100644 --- a/fern/pages/cuvs_bench/install.md +++ b/fern/pages/cuvs_bench/install.md @@ -49,6 +49,62 @@ Exact tags are listed on Docker Hub: **Note:** GPU containers use the CUDA toolkit inside the container. The host only needs a compatible driver, so CUDA 12 containers can run on systems with CUDA 13.x-capable drivers. GPU access also requires the NVIDIA Docker runtime from the [NVIDIA Container Toolkit](https://github.com/NVIDIA/nvidia-docker). +## PyLucene backend prerequisites + +The optional `pylucene` backend requires components that cuVS Bench does not install automatically: + +- An NVIDIA GPU supported by cuVS, plus matching CUDA and cuVS native libraries. +- JDK 22, `pytest`, and a [source-built PyLucene installation](https://lucene.apache.org/pylucene/install.html) compatible with the Lucene APIs used by the selected cuVS-Lucene checkout. The verified environment used PyLucene 10.0.0; the selected cuVS-Lucene POM compiles against Lucene 10.2.0, so validate the exact pair below. +- [Maven 3.9.6 or newer](https://maven.apache.org/download.cgi) to build cuVS-Lucene. +- The base `cuvs-java` JAR, standard `cuvs-lucene` JAR, and native libraries built from the pinned compatible revisions below. + +The required PyLucene service-provider and codec support is currently proposed in [NVIDIA/cuvs-lucene#174](https://github.com/NVIDIA/cuvs-lucene/pull/174). Until that work is merged and released, build cuVS-Lucene from that PR's branch; a release or `main` checkout without those changes is insufficient. + +Source compatibility is revision-specific while that PR is under review: the verified PR head `7d70d2f` does not compile against cuVS `main` at `f72199e3`. The commands below pin a source-compatible pair rather than a moving PR head. + +Build the dependency in a separate checkout so this does not change your cuVS Bench working tree: + +```bash +git clone https://github.com/NVIDIA/cuvs.git cuvs-pylucene-deps +cd cuvs-pylucene-deps +git switch --detach a59e2445 +./build.sh libcuvs java +cd .. +``` + +If matching native cuVS libraries are already built and installed, `./build.sh java` is sufficient. The Java build installs the base and native-classifier JARs into the local Maven repository; see the [cuVS Java build guide](https://github.com/NVIDIA/cuvs/blob/main/java/README.md). + +```bash +git clone https://github.com/NVIDIA/cuvs-lucene.git +cd cuvs-lucene +git fetch origin pull/174/head +git switch --detach 7d70d2f +mvn clean package -DskipTests +``` + +After the build, the conventional JAR paths are: + +```text +~/.m2/repository/com/nvidia/cuvs/cuvs-java//cuvs-java-.jar +/target/cuvs-lucene-.jar +``` + +Use the base `cuvs-java` JAR, not a native-classifier JAR. Use the standard cuVS-Lucene JAR, not its `-jar-with-dependencies`, sources, or Javadoc variants. Native-library paths must resolve `libcuvs.so`, `libcuvs_c.so`, their dependencies, and the CUDA runtime libraries from the pinned cuVS build. Then validate the artifacts from the cuVS-Lucene checkout: + +Use a clean environment without another cuVS native installation on its library path; otherwise, the JVM can load the other `libcuvs_c.so` first and reject the Java/native version mismatch. + +```bash +python -m pip install pytest + +CUVS_NATIVE_BUILD="$(cd ../cuvs-pylucene-deps/cpp/build && pwd)" +export JAVA_LIBRARY_PATH="$CUVS_NATIVE_BUILD:$CUVS_NATIVE_BUILD/c:/usr/local/cuda/lib64" +export LD_LIBRARY_PATH="$JAVA_LIBRARY_PATH${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" + +./test_pylucene.sh --full-e2e +``` + +See [Running the PyLucene backend](/user-guide/benchmarking-guide/cu-vs-bench-tool/usage#running-the-pylucene-backend) for a complete smoke workflow. + ## Build from Source Build cuVS Bench from source when you need local benchmark executables that match a development checkout, include custom algorithm targets, or use dependencies that are not available in the pre-built packages. diff --git a/fern/pages/cuvs_bench/pluggable_backend.md b/fern/pages/cuvs_bench/pluggable_backend.md index ce5eb2ec51..a72620350b 100644 --- a/fern/pages/cuvs_bench/pluggable_backend.md +++ b/fern/pages/cuvs_bench/pluggable_backend.md @@ -222,6 +222,10 @@ get_registry().register("elasticsearch", ElasticsearchBackend) | `BenchmarkBackend` | Abstract class whose `build(...)` method returns `BuildResult` and whose `search(...)` method returns `SearchResult`. Register with `BackendRegistry.register(name, backend_class)`. | | `BackendRegistry` | Singleton registry returned by `get_registry()`. It maps backend type names to backend classes. | +## PyLucene Backend + +The built-in `pylucene` loader expands algorithm YAML groups into one Lucene index per selected codec. The backend initializes PyLucene's process-global JVM and resolves cuVS-Lucene codecs through Lucene's service-provider interface. See [Installation](/user-guide/benchmarking-guide/cu-vs-bench-tool/installation#pylucene-backend-prerequisites) and [Usage](/user-guide/benchmarking-guide/cu-vs-bench-tool/usage#running-the-pylucene-backend) for the provisional dependency, configuration, and runtime limits. + ## C++ Backend The built-in `CppGoogleBenchmarkBackend` uses `backend_type="cpp_gbench"`. Its config loader reads YAML under `config/datasets` and `config/algos`, expands parameter combinations, and validates constraints. Its backend runs the C++ benchmark executables and merges their results. diff --git a/fern/pages/cuvs_bench/running.md b/fern/pages/cuvs_bench/running.md index 94606bb214..2dec4c8485 100644 --- a/fern/pages/cuvs_bench/running.md +++ b/fern/pages/cuvs_bench/running.md @@ -64,6 +64,7 @@ Create a custom YAML file with a `base` group to override the default benchmark | HNSWLIB | `hnswlib` | | DiskANN | `diskann_memory`, `diskann_ssd` | | NVIDIA cuVS | `cuvs_brute_force`, `cuvs_cagra`, `cuvs_ivf_flat`, `cuvs_ivf_pq`, `cuvs_cagra_hnswlib`, `cuvs_vamana` | +| PyLucene/cuVS | `pylucene_cuvs_hnsw`, `pylucene_cuvs_cagra` | ### Multi-GPU algorithms @@ -75,6 +76,70 @@ cuVS Bench includes single-node multi-GPU versions of IVF-Flat, IVF-PQ, and CAGR | IVF-PQ | `cuvs_mg_ivf_pq` | | CAGRA | `cuvs_mg_cagra` | +## Running the PyLucene backend + +The PyLucene backend type is `pylucene`. It runs through Python's embedded JVM, not the `*_ANN_BENCH` executables or `--executable-dir`. It accepts only FLOAT32 datasets with Euclidean distance and requires the prerequisites described in [Installation](/user-guide/benchmarking-guide/cu-vs-bench-tool/installation#pylucene-backend-prerequisites). + +Install the development checkout into the activated PyLucene environment: + +```bash +cd +python -m pip install -e ./python/cuvs_bench +``` + +Create `pylucene-backend.yaml` with absolute paths to the base `cuvs-java` JAR, the standard `cuvs-lucene` JAR, and a path-separated list of directories containing the matching cuVS and CUDA native libraries: + +```yaml +backend: pylucene +cuvs_java_jar: /home/user/.m2/repository/com/nvidia/cuvs/cuvs-java/VERSION/cuvs-java-VERSION.jar +cuvs_lucene_jar: /absolute/path/to/cuvs-lucene/target/cuvs-lucene-VERSION.jar +java_library_path: /work/cuvs-pylucene-deps/cpp/build:/work/cuvs-pylucene-deps/cpp/build/c:/usr/local/cuda/lib64 +``` + +Configuration sources are: + +| Backend configuration key | Environment variable | +| --- | --- | +| `cuvs_java_jar` | `CUVS_LUCENE_CUVS_JAVA_JAR` | +| `cuvs_lucene_jar` | `CUVS_LUCENE_JAR` | +| `java_library_path` | `JAVA_LIBRARY_PATH` or `LD_LIBRARY_PATH` | +| `jvm_args` | No environment alias; YAML list of additional JVM arguments. | + +Use one explicit dataset root for both preparation and execution. This small synthetic run is a functional smoke test: + +```bash +export DATASET_ROOT=/absolute/path/to/cuvs-bench-data + +python -m cuvs_bench.get_dataset \ + --dataset test-data \ + --dataset-path "$DATASET_ROOT" \ + --test-data-n-train 512 \ + --test-data-n-test 4 \ + --test-data-k 5 + +python -m cuvs_bench.run \ + --backend-config pylucene-backend.yaml \ + --dataset test-data \ + --dataset-path "$DATASET_ROOT" \ + --algorithms pylucene_cuvs_hnsw \ + --groups test \ + --batch-size 2 -k 5 \ + -m latency \ + --build --search --force +``` + +Do not use this tiny synthetic smoke run as performance or quality evidence; it exists only to verify the workflow. For a representative recall run, prepare `deep-image-96-angular` with `--normalize` into the same `DATASET_ROOT`, then run the backend against `deep-image-96-inner` with the same `--dataset-path`. + +The HNSW `base` group sweeps `Lucene101AcceleratedHNSWCodec`, `Lucene101AcceleratedHNSWBaseLayerCodec`, and `Lucene101AcceleratedHNSWMultiLayerCodec`; its `test` group uses `Lucene101AcceleratedHNSWCodec`. These codecs build the HNSW graph on the GPU with cuVS and use Lucene's CPU HNSW search. Use `--algorithms pylucene_cuvs_cagra` to run `CuVS2510GPUSearchCodec`, which builds and searches on the GPU. + +The supplied configurations use Lucene's public `KnnFloatVectorQuery` and do not expose search-time tuning parameters. The backend accepts FLOAT32 Euclidean datasets with at most 4096 dimensions. All supplied codecs require at least two indexed vectors because cuVS-Lucene bypasses cuVS for a single-vector index. CAGRA requires a GPU for build and search. Although cuVS-Lucene uses an effective `lucene_k` of `min(k, document_count)`, the backend conservatively requires `k <= 1024` to avoid cuVS-Lucene paths that can use brute-force search above that limit. HNSW requires a GPU for graph construction, and new builds reject cuVS-Lucene's CPU fallback. + +New HNSW and CAGRA builds atomically write commit-bound provenance manifests named `.cuvs-bench-pylucene-hnsw.json` and `.cuvs-bench-pylucene-cagra.json`, respectively. Reuse and search fail if the applicable manifest is missing, malformed, stale, or names a different codec. Indexes created outside this backend without the applicable manifest must be rebuilt with `--force`. CAGRA indexes also undergo structural and checksum-integrity verification. + +The backend currently supports latency mode with one search thread. `--batch-size` groups queries for measurement, and the reported latency percentiles are per batch. Throughput mode and multiple search threads are not implemented. + +PyLucene's JVM is process-global and can be initialized only once. Set the JAR, native-library locations, and `jvm_args` before the first PyLucene benchmark, and start a new Python process to change any of them. + ## Smaller-scale benchmarks (<1M to 10M vectors) Use `cuvs_bench.get_dataset` to prepare a built-in dataset. By default, datasets are stored under `RAPIDS_DATASET_ROOT_DIR` when that environment variable is set, or under a local `datasets` directory otherwise. @@ -192,7 +257,9 @@ Containers can also run in detached mode. ## Evaluating results -Build benchmarks report: +The tables below describe fields emitted by the default C++ Google Benchmark backend. Other backends report the fields that apply to their execution model, so not every field is present for every backend. In particular, PyLucene reports `index_size` as the total on-disk index size in bytes. + +C++ build benchmarks report: | Name | Description | | --- | --- | @@ -203,7 +270,7 @@ Build benchmarks report: | GPU | GPU time spent building. | | index_size | Number of vectors used to train the index. | -Search benchmarks report: +C++ search benchmarks report: | Name | Description | | --- | --- | diff --git a/python/cuvs_bench/cuvs_bench/backends/__init__.py b/python/cuvs_bench/cuvs_bench/backends/__init__.py index a09e6fd1fe..e2e7845953 100644 --- a/python/cuvs_bench/cuvs_bench/backends/__init__.py +++ b/python/cuvs_bench/cuvs_bench/backends/__init__.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ @@ -26,11 +26,13 @@ from .cpp_gbench import CppGoogleBenchmarkBackend from .opensearch import OpenSearchBackend +from .pylucene import PyLuceneBackend # Auto-register built-in backends _registry = get_registry() _registry.register("cpp_gbench", CppGoogleBenchmarkBackend) _registry.register("opensearch", OpenSearchBackend) +_registry.register("pylucene", PyLuceneBackend) __all__ = [ # Base classes and data structures @@ -46,4 +48,5 @@ # Built-in backends "CppGoogleBenchmarkBackend", "OpenSearchBackend", + "PyLuceneBackend", ] diff --git a/python/cuvs_bench/cuvs_bench/backends/base.py b/python/cuvs_bench/cuvs_bench/backends/base.py index bf802d2128..e6334c4588 100644 --- a/python/cuvs_bench/cuvs_bench/backends/base.py +++ b/python/cuvs_bench/cuvs_bench/backends/base.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ @@ -112,6 +112,13 @@ def training_vectors(self) -> np.ndarray: ) return self._training_vectors + @property + def loaded_training_vectors(self) -> Optional[np.ndarray]: + """Training vectors already in memory, without loading ``base_file``.""" + if self._training_vectors.size == 0: + return None + return self._training_vectors + @training_vectors.setter def training_vectors(self, value: Optional[np.ndarray]) -> None: """Set training vectors directly.""" @@ -243,15 +250,23 @@ def to_json(self) -> Dict[str, Any]: Dict[str, Any] Dictionary with benchmark results """ - return { - "name": f"{self.algorithm}/build", - "real_time": self.build_time_seconds, - "time_unit": "s", - "index_size": self.index_size_bytes, - "success": self.success, + result = { **self.build_params, **self.metadata, } + result.pop("error_message", None) + result.update( + { + "name": f"{self.algorithm}/build", + "real_time": self.build_time_seconds, + "time_unit": "s", + "index_size": self.index_size_bytes, + "success": self.success, + } + ) + if self.error_message is not None: + result["error_message"] = self.error_message + return result @dataclass @@ -288,6 +303,10 @@ class SearchResult: Whether the search succeeded error_message : Optional[str] Error message if search failed + latency_seconds : Optional[float] + Backend-reported mean latency in seconds. For batched backends, this + is mean batch latency. This field is last to preserve + compatibility with existing positional construction. """ neighbors: np.ndarray @@ -303,6 +322,7 @@ class SearchResult: metadata: Dict[str, Any] = field(default_factory=dict) success: bool = True error_message: Optional[str] = None + latency_seconds: Optional[float] = None def to_json(self) -> Dict[str, Any]: """ @@ -313,26 +333,39 @@ def to_json(self) -> Dict[str, Any]: Dict[str, Any] Dictionary with benchmark results """ - result = { - "name": f"{self.algorithm}/search", - "real_time": self.search_time_ms, - "time_unit": "ms", - "items_per_second": self.queries_per_second, - "Recall": self.recall, - "success": self.success, - "search_params": self.search_params, - **self.metadata, - } - + result = dict(self.metadata) + result.pop("error_message", None) if self.latency_percentiles: result.update(self.latency_percentiles) + real_time_ms = ( + self.latency_seconds * 1000.0 + if self.latency_seconds is not None + else self.search_time_ms + ) + result.update( + { + "name": f"{self.algorithm}/search", + "real_time": real_time_ms, + "time_unit": "ms", + "search_time_ms": self.search_time_ms, + "items_per_second": self.queries_per_second, + "Recall": self.recall, + "success": self.success, + "search_params": self.search_params, + } + ) + if self.latency_seconds is not None: + result["Latency"] = self.latency_seconds if self.gpu_time_seconds is not None: result["GPU"] = self.gpu_time_seconds if self.cpu_time_seconds is not None: result["cpu_time"] = self.cpu_time_seconds + if self.error_message is not None: + result["error_message"] = self.error_message + return result @@ -368,6 +401,11 @@ class BenchmarkBackend(ABC): - requires_network : bool - Whether backend requires network (default: False) """ + # Opt in when this backend's sweep results should be persisted by the + # orchestrator. Tune-mode and existing backends retain their own + # result-file behavior. + orchestrator_persists_results = False + def __init__(self, config: Dict[str, Any]): """Initialize backend with configuration.""" self.config = config diff --git a/python/cuvs_bench/cuvs_bench/backends/pylucene.py b/python/cuvs_bench/cuvs_bench/backends/pylucene.py new file mode 100644 index 0000000000..8a10fcdfbe --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/backends/pylucene.py @@ -0,0 +1,3029 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""PyLucene backend for cuVS-accelerated Lucene codecs.""" + +from __future__ import annotations + +import hashlib +import importlib +import json +import os +import shutil +import tempfile +import threading +import time +import zipfile +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +from types import MappingProxyType +from typing import ( + Any, + Callable, + Dict, + List, + Mapping, + NamedTuple, + Optional, + Tuple, + Union, +) + +import numpy as np + +from .._bin_format import read_bin_header +from ..orchestrator.config_loaders import ( + BenchmarkConfig, + ConfigLoader, + DatasetConfig, + IndexConfig, +) +from ._utils import dtype_from_filename +from .base import BenchmarkBackend, BuildResult, Dataset, SearchResult + +_ID_FIELD = "id" +_VECTOR_FIELD = "vector" +_MAX_DIMENSIONS = 4096 + +_HNSW_CODECS = frozenset( + { + "Lucene101AcceleratedHNSWCodec", + "Lucene101AcceleratedHNSWBaseLayerCodec", + "Lucene101AcceleratedHNSWMultiLayerCodec", + } +) +_CAGRA_CODEC = "CuVS2510GPUSearchCodec" +_SUPPORTED_CODECS = _HNSW_CODECS | {_CAGRA_CODEC} +_SUPPORTED_BUILD_KEYS = frozenset({"codec"}) +_EXPECTED_WRITER_PATH = { + **{codec: "gpu-hnsw" for codec in _HNSW_CODECS}, + _CAGRA_CODEC: "gpu-cagra", +} +_CAGRA_META_EXTENSION = ".vemc" +_CAGRA_META_CODEC_NAME = "Lucene102CuVSVectorsFormatMeta" +_CAGRA_INDEX_EXTENSION = ".vcag" +_CAGRA_INDEX_CODEC_NAME = "Lucene102CuVSVectorsFormatIndex" +_CAGRA_META_VERSION = 0 +_CAGRA_INDEX_VERSION = 0 +# Avoid caching checksums while coarse filesystem timestamps can still make a +# same-size rewrite look unchanged. +_CAGRA_CACHE_MIN_FILE_AGE_NS = 2_000_000_000 +_FLOAT32_ENCODING_ORDINAL = 1 +_EUCLIDEAN_SIMILARITY_ORDINAL = 0 + +_HNSW_PROVENANCE_FILE = ".cuvs-bench-pylucene-hnsw.json" +_CAGRA_PROVENANCE_FILE = ".cuvs-bench-pylucene-cagra.json" +_PROVENANCE_SCHEMA_VERSION = 1 +_PROVENANCE_KEYS = frozenset( + { + "schema_version", + "codec", + "writer_path", + "vector_count", + "dimensions", + "commit_fingerprints", + } +) +_SHA256_HEX_DIGITS = frozenset("0123456789abcdef") +_LUCENE_CORE_CLASS = "org/apache/lucene/index/IndexWriter.class" + +_JVM_INIT_LOCK = threading.Lock() +_INITIALIZED_CLASSPATH: Optional[str] = None +_INITIALIZED_VMARGS: Optional[Tuple[str, ...]] = None + + +def _attempt_cleanup( + cleanup: Callable[[], None], + description: str, + primary_error: Optional[BaseException], +) -> Optional[BaseException]: + try: + cleanup() + except BaseException as cleanup_error: + if primary_error is None: + return cleanup_error + if isinstance(primary_error, Exception) and not isinstance( + cleanup_error, Exception + ): + cleanup_error.add_note( + f"Raised while attempting to {description}; prior failure: " + f"{type(primary_error).__name__}: {primary_error}" + ) + return cleanup_error + primary_error.add_note( + f"Failed to {description}: " + f"{type(cleanup_error).__name__}: {cleanup_error}" + ) + return primary_error + + +class _CleanupStack: + """Close registered resources without losing the operation's failure.""" + + def __init__(self) -> None: + self._cleanups: List[Tuple[str, Callable[[], None]]] = [] + + def __enter__(self) -> _CleanupStack: + return self + + def add(self, description: str, cleanup: Callable[[], None]) -> None: + self._cleanups.append((description, cleanup)) + + def __exit__( + self, + _error_type: Any, + primary_error: Optional[BaseException], + _traceback: Any, + ) -> bool: + final_error = primary_error + for description, cleanup in reversed(self._cleanups): + final_error = _attempt_cleanup(cleanup, description, final_error) + if final_error is not None and final_error is not primary_error: + raise final_error + return False + + +def _exception_summary(error: Exception) -> str: + details = [f"{type(error).__name__}: {error}"] + if error.__cause__ is not None: + cause = error.__cause__ + details.append(f"caused by {type(cause).__name__}: {cause}") + details.extend(getattr(error, "__notes__", ())) + return "; ".join(details) + + +def _configured_jar( + config: Dict[str, Any], config_key: str, environment_key: str +) -> Path: + value = config.get(config_key) or os.environ.get(environment_key) + if not value: + raise RuntimeError( + f"PyLucene backend requires '{config_key}' or " + f"the {environment_key} environment variable" + ) + + path = Path(os.fspath(value)).expanduser().resolve() + if not path.is_file(): + raise FileNotFoundError(f"{config_key} does not exist: {path}") + return path + + +def _reject_bundled_lucene_classes(cuvs_lucene_jar: Path) -> None: + """Reject fat cuVS-Lucene jars before they can poison the process JVM.""" + if not zipfile.is_zipfile(cuvs_lucene_jar): + return + + with zipfile.ZipFile(cuvs_lucene_jar) as archive: + try: + archive.getinfo(_LUCENE_CORE_CLASS) + except KeyError: + return + + raise RuntimeError( + "cuvs_lucene_jar bundles Lucene classes and is incompatible with " + "PyLucene's process-wide JVM. Use the standard thin cuvs-lucene JAR, " + "not a '-jar-with-dependencies' artifact." + ) + + +def _load_pylucene() -> Any: + try: + return importlib.import_module("lucene") + except ImportError as exc: + raise ImportError( + "PyLucene's `lucene` module is required for the pylucene " + "cuvs-bench backend. PyLucene must be built and installed " + "separately; see the cuVS Bench installation documentation." + ) from exc + + +def _pylucene_classpath(config: Dict[str, Any], lucene: Any) -> str: + cuvs_java_jar = _configured_jar( + config, "cuvs_java_jar", "CUVS_LUCENE_CUVS_JAVA_JAR" + ) + cuvs_lucene_jar = _configured_jar( + config, "cuvs_lucene_jar", "CUVS_LUCENE_JAR" + ) + _reject_bundled_lucene_classes(cuvs_lucene_jar) + return os.pathsep.join( + [str(cuvs_java_jar), str(cuvs_lucene_jar), lucene.CLASSPATH] + ) + + +def _pylucene_vmargs(config: Dict[str, Any]) -> List[str]: + java_library_path = ( + config.get("java_library_path") + or os.environ.get("JAVA_LIBRARY_PATH") + or os.environ.get("LD_LIBRARY_PATH") + ) + vmargs = [ + "--enable-native-access=ALL-UNNAMED", + "--add-modules=jdk.incubator.vector", + ] + if java_library_path: + vmargs.append(f"-Djava.library.path={java_library_path}") + + extra_vmargs = config.get("jvm_args", []) + if isinstance(extra_vmargs, (str, bytes)) or not isinstance( + extra_vmargs, (list, tuple) + ): + raise TypeError("jvm_args must be a list or tuple of strings") + if not all(isinstance(arg, str) for arg in extra_vmargs): + raise TypeError("every jvm_args entry must be a string") + vmargs.extend(extra_vmargs) + return vmargs + + +def _attach_pylucene_jvm( + lucene: Any, classpath: str, vmargs: List[str] +) -> None: + with _JVM_INIT_LOCK: + vm_environment = lucene.getVMEnv() + if vm_environment is None: + vm_environment = _start_pylucene_jvm(lucene, classpath, vmargs) + else: + _validate_pylucene_jvm_configuration(classpath, vmargs) + + vm_environment.attachCurrentThread() + + +def _start_pylucene_jvm(lucene: Any, classpath: str, vmargs: List[str]) -> Any: + global _INITIALIZED_CLASSPATH, _INITIALIZED_VMARGS + vm_environment = lucene.initVM(classpath=classpath, vmargs=vmargs) + if vm_environment is None: + vm_environment = lucene.getVMEnv() + if vm_environment is None: + raise RuntimeError("PyLucene did not return a JVM environment") + + _INITIALIZED_CLASSPATH = classpath + _INITIALIZED_VMARGS = tuple(vmargs) + return vm_environment + + +def _validate_pylucene_jvm_configuration( + classpath: str, vmargs: List[str] +) -> None: + if _INITIALIZED_CLASSPATH is None or _INITIALIZED_VMARGS is None: + raise RuntimeError( + "PyLucene's process-wide JVM was initialized before the " + "pylucene backend, so the required cuVS classpath and JVM " + "arguments cannot be verified. Start a new Python process and " + "let cuVS Bench initialize PyLucene." + ) + if _INITIALIZED_CLASSPATH != classpath: + raise RuntimeError( + "PyLucene's process-wide JVM is already initialized with " + "different cuVS Java or cuVS-Lucene jars" + ) + if _INITIALIZED_VMARGS != tuple(vmargs): + raise RuntimeError( + "PyLucene's process-wide JVM is already initialized with " + "different JVM arguments or native-library paths" + ) + + +def _initialize_pylucene(config: Dict[str, Any]) -> Any: + """Initialize PyLucene once and attach the current Python thread.""" + lucene = _load_pylucene() + classpath = _pylucene_classpath(config, lucene) + vmargs = _pylucene_vmargs(config) + _attach_pylucene_jvm(lucene, classpath, vmargs) + return lucene + + +def _parse_writer_telemetry(description: str) -> Dict[str, str]: + """Parse cuVS-Lucene vector-format diagnostics.""" + _, separator, payload = description.partition("(") + if not separator or not payload.endswith(")"): + raise RuntimeError( + f"Malformed cuVS-Lucene writer telemetry: {description!r}" + ) + + telemetry: Dict[str, str] = {} + for item in payload[:-1].split(";"): + key, item_separator, value = item.partition("=") + if not item_separator or not key: + raise RuntimeError( + "Malformed cuVS-Lucene writer telemetry item " + f"{item!r} in {description!r}" + ) + telemetry[key] = value + return telemetry + + +def _score_to_squared_euclidean(score: float) -> float: + """Convert Lucene's Euclidean score to squared Euclidean distance.""" + if score <= 0.0: + return float("inf") + return max(0.0, (1.0 / score) - 1.0) + + +def _validate_float32_matrix(vectors: np.ndarray, name: str) -> np.ndarray: + array = np.asarray(vectors) + if array.ndim != 2: + raise ValueError(f"{name} must be a two-dimensional array") + if array.shape[0] == 0 or array.shape[1] == 0: + raise ValueError(f"{name} must contain at least one vector") + if array.dtype != np.float32: + raise TypeError(f"{name} must use float32 values, got {array.dtype}") + if array.shape[1] > _MAX_DIMENSIONS: + raise ValueError( + f"{name} dimensions must not exceed {_MAX_DIMENSIONS}, " + f"got {array.shape[1]}" + ) + if not np.isfinite(array).all(): + raise ValueError(f"{name} must contain only finite values") + return np.ascontiguousarray(array) + + +def _validate_metric(dataset: Dataset) -> None: + if dataset.distance_metric.lower() not in {"euclidean", "l2"}: + raise ValueError( + "PyLucene cuVS codecs currently support only Euclidean/L2 " + f"datasets, got {dataset.distance_metric!r}" + ) + + +def _validate_codec(codec_name: Any) -> str: + if not isinstance(codec_name, str) or codec_name not in _SUPPORTED_CODECS: + available = ", ".join(sorted(_SUPPORTED_CODECS)) + raise ValueError( + f"Unsupported PyLucene codec {codec_name!r}. " + f"Supported codecs: {available}" + ) + return codec_name + + +def _configured_codec( + build_params: Dict[str, Any], backend_config: Dict[str, Any] +) -> str: + if "codec" in build_params: + codec_name = build_params["codec"] + else: + codec_name = backend_config.get("codec") + return _validate_codec(codec_name) + + +def _validate_build_params(build_params: Any) -> Dict[str, Any]: + if not isinstance(build_params, dict): + raise TypeError("PyLucene build parameters must be a mapping") + unsupported = set(build_params) - _SUPPORTED_BUILD_KEYS + if unsupported: + names = ", ".join(sorted(str(name) for name in unsupported)) + raise ValueError(f"Unsupported PyLucene build parameter(s): {names}") + return build_params + + +def _validate_search_params(search_params: Any) -> List[Dict[str, Any]]: + if ( + not isinstance(search_params, list) + or len(search_params) != 1 + or search_params[0] != {} + ): + raise ValueError( + "PyLucene's public Lucene query API does not expose " + "cuVS-specific search parameters" + ) + return search_params + + +def _safe_remove_index(index_path: Path) -> None: + resolved = index_path.resolve() + forbidden = { + Path(resolved.anchor), + Path.cwd().resolve(), + Path.home().resolve(), + } + if resolved in forbidden: + raise ValueError(f"Refusing to remove unsafe index path: {resolved}") + if index_path.is_symlink() or not index_path.is_dir(): + raise ValueError( + f"PyLucene index path must be a directory: {index_path}" + ) + shutil.rmtree(index_path) + + +def _index_size(index_path: Path) -> int: + return sum( + path.stat().st_size for path in index_path.rglob("*") if path.is_file() + ) + + +@dataclass(frozen=True) +class _FileSignature: + resolved_path: str + device: int + inode: int + size: int + modified_at_ns: int + changed_at_ns: int + + +def _file_signature(path: Path) -> _FileSignature: + file_stat = path.stat() + return _FileSignature( + resolved_path=str(path.resolve()), + device=file_stat.st_dev, + inode=file_stat.st_ino, + size=file_stat.st_size, + modified_at_ns=file_stat.st_mtime_ns, + changed_at_ns=file_stat.st_ctime_ns, + ) + + +def _has_lucene_segments(index_path: Path) -> bool: + return any( + path.is_file() and path.name.startswith("segments_") + for path in index_path.iterdir() + ) + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as file: + for block in iter(lambda: file.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def _commit_fingerprints(index_path: Path) -> List[Dict[str, str]]: + commit_files = sorted( + path + for path in index_path.iterdir() + if path.name.startswith("segments_") + ) + if not commit_files: + raise RuntimeError( + "Lucene index has no segments_* commit file to fingerprint" + ) + + fingerprints = [] + for path in commit_files: + if path.is_symlink() or not path.is_file(): + raise RuntimeError( + "Lucene commit fingerprint target must be a regular file: " + f"{path}" + ) + fingerprints.append({"name": path.name, "sha256": _sha256_file(path)}) + return fingerprints + + +@dataclass(frozen=True) +class _IndexProvenanceVerification: + codec: str + writer_path: str + vector_count: int + dimensions: int + commit_file_count: int + + def to_metadata(self) -> Dict[str, Union[str, int]]: + return { + "status": f"{self.writer_path}-provenance", + "schema_version": _PROVENANCE_SCHEMA_VERSION, + "codec": self.codec, + "writer_path": self.writer_path, + "vector_count": self.vector_count, + "dimensions": self.dimensions, + "commit_file_count": self.commit_file_count, + } + + +@dataclass(frozen=True) +class _ProvenanceExpectation: + codec: str + manifest_name: str + label: str + vector_count: Optional[int] + dimensions: Optional[int] + + +class _IndexProvenanceError(RuntimeError): + """Raised when an index cannot be proven to be backend-built.""" + + +def _write_index_provenance( + index_path: Path, + codec: str, + vector_count: int, + dimensions: int, + manifest_name: str, +) -> None: + payload = { + "schema_version": _PROVENANCE_SCHEMA_VERSION, + "codec": codec, + "writer_path": _EXPECTED_WRITER_PATH[codec], + "vector_count": int(vector_count), + "dimensions": int(dimensions), + "commit_fingerprints": _commit_fingerprints(index_path), + } + manifest_path = index_path / manifest_name + with _CleanupStack() as cleanups: + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=index_path, + prefix=f"{manifest_name}.", + suffix=".tmp", + delete=False, + ) as file: + temporary_path = Path(file.name) + cleanups.add( + "remove temporary provenance file", + lambda: temporary_path.unlink(missing_ok=True), + ) + json.dump( + payload, + file, + allow_nan=False, + indent=2, + sort_keys=True, + ) + file.write("\n") + file.flush() + os.fsync(file.fileno()) + temporary_path.chmod(0o644) + temporary_path.replace(manifest_path) + + +def _write_hnsw_provenance( + index_path: Path, + codec: str, + vector_count: int, + dimensions: int, +) -> None: + _write_index_provenance( + index_path, + codec, + vector_count, + dimensions, + _HNSW_PROVENANCE_FILE, + ) + + +def _write_cagra_provenance( + index_path: Path, + vector_count: int, + dimensions: int, +) -> None: + _write_index_provenance( + index_path, + _CAGRA_CODEC, + vector_count, + dimensions, + _CAGRA_PROVENANCE_FILE, + ) + + +def _require_positive_int( + value: Any, provenance_label: str, field_name: str +) -> int: + if type(value) is not int or value < 1: + raise _IndexProvenanceError( + f"{provenance_label} field {field_name!r} must be a " + "positive integer" + ) + return value + + +def _read_index_provenance( + manifest_path: Path, provenance_label: str +) -> Dict[str, Any]: + try: + payload = json.loads(manifest_path.read_text(encoding="utf-8")) + except FileNotFoundError as exc: + raise _IndexProvenanceError( + f"{provenance_label} manifest is missing: {manifest_path}" + ) from exc + except (OSError, UnicodeError, json.JSONDecodeError) as exc: + raise _IndexProvenanceError( + f"{provenance_label} manifest cannot be read: " + f"{manifest_path}: {exc}" + ) from exc + + if not isinstance(payload, dict) or set(payload) != _PROVENANCE_KEYS: + raise _IndexProvenanceError( + f"{provenance_label} manifest has an unsupported schema" + ) + return payload + + +def _validate_provenance_identity( + payload: Dict[str, Any], + expected_codec: str, + provenance_label: str, +) -> str: + if ( + type(payload["schema_version"]) is not int + or payload["schema_version"] != _PROVENANCE_SCHEMA_VERSION + ): + raise _IndexProvenanceError( + f"{provenance_label} manifest has an unsupported schema version" + ) + if payload["codec"] != expected_codec: + raise _IndexProvenanceError( + f"{provenance_label} codec does not match the requested codec: " + f"{payload['codec']!r} != {expected_codec!r}" + ) + + writer_path = payload["writer_path"] + if writer_path != _EXPECTED_WRITER_PATH[expected_codec]: + raise _IndexProvenanceError( + f"{provenance_label} does not record the required GPU writer path" + ) + return writer_path + + +def _validate_provenance_shape( + payload: Dict[str, Any], + provenance_label: str, + expected_vector_count: Optional[int] = None, + expected_dimensions: Optional[int] = None, +) -> Tuple[int, int]: + vector_count = _require_positive_int( + payload["vector_count"], provenance_label, "vector_count" + ) + if vector_count < 2: + raise _IndexProvenanceError( + f"{provenance_label} vector count must be at least two" + ) + dimensions = _require_positive_int( + payload["dimensions"], provenance_label, "dimensions" + ) + if dimensions > _MAX_DIMENSIONS: + raise _IndexProvenanceError( + f"{provenance_label} dimensions exceed the supported maximum: " + f"{dimensions} > {_MAX_DIMENSIONS}" + ) + if ( + expected_vector_count is not None + and vector_count != expected_vector_count + ): + raise _IndexProvenanceError( + f"{provenance_label} vector count does not match the dataset: " + f"{vector_count} != {expected_vector_count}" + ) + if expected_dimensions is not None and dimensions != expected_dimensions: + raise _IndexProvenanceError( + f"{provenance_label} dimensions do not match the dataset: " + f"{dimensions} != {expected_dimensions}" + ) + return vector_count, dimensions + + +def _validate_commit_fingerprints( + stored_fingerprints: Any, provenance_label: str +) -> List[Dict[str, str]]: + if not isinstance(stored_fingerprints, list) or not stored_fingerprints: + raise _IndexProvenanceError( + f"{provenance_label} has no Lucene commit fingerprints" + ) + names = [] + for fingerprint in stored_fingerprints: + name, _ = _validate_commit_fingerprint(fingerprint, provenance_label) + names.append(name) + if names != sorted(set(names)): + raise _IndexProvenanceError( + f"{provenance_label} commit fingerprints must be unique and sorted" + ) + return stored_fingerprints + + +def _validate_commit_fingerprint( + fingerprint: Any, provenance_label: str +) -> Tuple[str, str]: + if not isinstance(fingerprint, dict) or set(fingerprint) != { + "name", + "sha256", + }: + raise _IndexProvenanceError( + f"{provenance_label} has a malformed Lucene commit fingerprint" + ) + + name = _validate_commit_filename(fingerprint["name"], provenance_label) + digest = _validate_sha256_digest(fingerprint["sha256"], provenance_label) + return name, digest + + +def _validate_commit_filename(value: Any, provenance_label: str) -> str: + if not isinstance(value, str): + raise _IndexProvenanceError( + f"{provenance_label} commit filename must be a string" + ) + if not value.startswith("segments_"): + raise _IndexProvenanceError( + f"{provenance_label} commit filename must start with 'segments_'" + ) + if Path(value).name != value: + raise _IndexProvenanceError( + f"{provenance_label} commit filename must not contain a path" + ) + return value + + +def _validate_sha256_digest(value: Any, provenance_label: str) -> str: + if not isinstance(value, str): + raise _IndexProvenanceError( + f"{provenance_label} commit SHA-256 must be a string" + ) + if len(value) != 64: + raise _IndexProvenanceError( + f"{provenance_label} commit SHA-256 must contain 64 characters" + ) + if not set(value).issubset(_SHA256_HEX_DIGITS): + raise _IndexProvenanceError( + f"{provenance_label} commit SHA-256 must be lowercase hexadecimal" + ) + return value + + +def _verify_index_provenance( + index_path: Path, + expectation: _ProvenanceExpectation, +) -> _IndexProvenanceVerification: + payload = _read_index_provenance( + index_path / expectation.manifest_name, expectation.label + ) + writer_path = _validate_provenance_identity( + payload, expectation.codec, expectation.label + ) + vector_count, dimensions = _validate_provenance_shape( + payload, + expectation.label, + expectation.vector_count, + expectation.dimensions, + ) + stored_fingerprints = _validate_commit_fingerprints( + payload["commit_fingerprints"], expectation.label + ) + try: + current_fingerprints = _commit_fingerprints(index_path) + except (OSError, RuntimeError) as exc: + raise _IndexProvenanceError( + f"{expectation.label} cannot fingerprint the Lucene commit: {exc}" + ) from exc + if stored_fingerprints != current_fingerprints: + raise _IndexProvenanceError( + f"{expectation.label} does not match the current Lucene commit" + ) + + return _IndexProvenanceVerification( + codec=expectation.codec, + writer_path=writer_path, + vector_count=vector_count, + dimensions=dimensions, + commit_file_count=len(current_fingerprints), + ) + + +def _verify_hnsw_provenance( + index_path: Path, + expected_codec: str, + *, + expected_vector_count: Optional[int] = None, + expected_dimensions: Optional[int] = None, +) -> _IndexProvenanceVerification: + return _verify_index_provenance( + index_path, + _ProvenanceExpectation( + codec=expected_codec, + manifest_name=_HNSW_PROVENANCE_FILE, + label="HNSW provenance", + vector_count=expected_vector_count, + dimensions=expected_dimensions, + ), + ) + + +def _verify_cagra_provenance( + index_path: Path, + *, + expected_vector_count: Optional[int] = None, + expected_dimensions: Optional[int] = None, +) -> _IndexProvenanceVerification: + return _verify_index_provenance( + index_path, + _ProvenanceExpectation( + codec=_CAGRA_CODEC, + manifest_name=_CAGRA_PROVENANCE_FILE, + label="CAGRA provenance", + vector_count=expected_vector_count, + dimensions=expected_dimensions, + ), + ) + + +def _validate_subset_size(value: Any) -> Optional[int]: + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise ValueError( + f"subset_size must be a positive integer, got {value!r}" + ) + return value + + +def _expected_training_shape(dataset: Dataset) -> Optional[Tuple[int, int]]: + vectors = dataset.loaded_training_vectors + if vectors is not None: + vectors = np.asarray(vectors) + if vectors.ndim != 2: + raise ValueError( + "training_vectors must be a two-dimensional array" + ) + if vectors.dtype != np.float32: + raise TypeError( + "training_vectors must use float32 values, " + f"got {vectors.dtype}" + ) + rows, dimensions = vectors.shape + elif dataset.base_file: + dtype = np.dtype(dtype_from_filename(dataset.base_file)) + if dtype != np.float32: + raise TypeError( + f"training_vectors must use float32 values, got {dtype}" + ) + rows, dimensions, _ = read_bin_header( + dataset.base_file, itemsize=dtype.itemsize + ) + subset_size = _validate_subset_size( + dataset.metadata.get("subset_size") + ) + if subset_size is not None: + rows = min(rows, subset_size) + else: + return None + + if rows < 1 or dimensions < 1: + raise ValueError("training_vectors must contain at least one vector") + if dimensions > _MAX_DIMENSIONS: + raise ValueError( + f"training_vectors dimensions must not exceed {_MAX_DIMENSIONS}, " + f"got {dimensions}" + ) + return int(rows), int(dimensions) + + +@dataclass(frozen=True) +class _SearchHit: + document_id: int + score: float + + +@dataclass(frozen=True) +class _RuntimeSearchResult: + hits: List[List[_SearchHit]] + batch_latencies_ms: List[float] + index_dimensions: int + document_count: int + + +@dataclass(frozen=True) +class _ProcessedSearchResult: + neighbors: np.ndarray + distances: np.ndarray + search_time_ms: float + queries_per_second: float + latency_seconds: float + latency_percentiles: Dict[str, float] + num_batches: int + + +@dataclass(frozen=True) +class _SearchInputs: + query_vectors: np.ndarray + expected_training_shape: Optional[Tuple[int, int]] + + +@dataclass(frozen=True) +class _SearchPlan: + index_path: Path + codec_name: str + search_params: List[Dict[str, Any]] + k: int + batch_size: int + mode: str + + +@dataclass(frozen=True) +class _ResolvedCodec: + codec_name: str + java_codec: Any + telemetry: Mapping[str, str] + writer_path: str + + def __post_init__(self) -> None: + object.__setattr__( + self, "telemetry", MappingProxyType(dict(self.telemetry)) + ) + + +class _ExistingIndexAction(Enum): + BUILD = "build" + REUSE = "reuse" + REJECT = "reject" + + +@dataclass(frozen=True) +class _ExistingIndexDecision: + action: _ExistingIndexAction + result: Optional[BuildResult] = None + + @classmethod + def build(cls) -> _ExistingIndexDecision: + return cls(action=_ExistingIndexAction.BUILD) + + @classmethod + def reuse(cls, result: BuildResult) -> _ExistingIndexDecision: + return cls(action=_ExistingIndexAction.REUSE, result=result) + + @classmethod + def reject(cls, result: BuildResult) -> _ExistingIndexDecision: + return cls(action=_ExistingIndexAction.REJECT, result=result) + + def completed_result(self) -> BuildResult: + if self.result is None: + raise RuntimeError( + "Existing-index decision is missing its build result" + ) + return self.result + + +def _validate_runtime_search_result( + runtime_result: _RuntimeSearchResult, + provenance: _IndexProvenanceVerification, + *, + query_count: int, +) -> None: + if runtime_result.document_count != provenance.vector_count: + raise RuntimeError( + "Lucene document count does not match index provenance: " + f"{runtime_result.document_count} != {provenance.vector_count}" + ) + if runtime_result.index_dimensions != provenance.dimensions: + raise RuntimeError( + "Lucene index dimensions do not match index provenance: " + f"{runtime_result.index_dimensions} != {provenance.dimensions}" + ) + if len(runtime_result.hits) != query_count: + raise RuntimeError( + "PyLucene returned an unexpected number of query results: " + f"{len(runtime_result.hits)} != {query_count}" + ) + + +def _validate_search_hit( + hit: _SearchHit, + *, + query_id: int, + document_count: int, + seen_document_ids: set[int], +) -> None: + if not 0 <= hit.document_id < document_count: + raise RuntimeError( + "PyLucene returned an out-of-range stored ID for " + f"query {query_id}: {hit.document_id}" + ) + if hit.document_id in seen_document_ids: + raise RuntimeError( + "PyLucene returned a duplicate stored ID for " + f"query {query_id}: {hit.document_id}" + ) + if not np.isfinite(hit.score): + raise RuntimeError( + "PyLucene returned a non-finite score for " + f"query {query_id}: {hit.score}" + ) + if not 0.0 <= hit.score <= 1.0: + raise RuntimeError( + "PyLucene returned a score outside the Euclidean range " + f"[0, 1] for query {query_id}: {hit.score}" + ) + seen_document_ids.add(hit.document_id) + + +def _convert_search_hits( + runtime_result: _RuntimeSearchResult, *, query_count: int, k: int +) -> Tuple[np.ndarray, np.ndarray]: + neighbors = np.full((query_count, k), -1, dtype=np.int64) + distances = np.full((query_count, k), np.inf, dtype=np.float32) + for query_id, hits in enumerate(runtime_result.hits): + query_neighbors, query_distances = _convert_query_hits( + hits, + query_id=query_id, + document_count=runtime_result.document_count, + k=k, + ) + hit_count = len(query_neighbors) + neighbors[query_id, :hit_count] = query_neighbors + distances[query_id, :hit_count] = query_distances + return neighbors, distances + + +def _convert_query_hits( + hits: List[_SearchHit], + *, + query_id: int, + document_count: int, + k: int, +) -> Tuple[List[int], List[float]]: + if len(hits) > min(k, document_count): + raise RuntimeError( + f"PyLucene returned too many hits for query {query_id}: " + f"{len(hits)}" + ) + + neighbors = [] + distances = [] + seen_document_ids = set() + for hit in hits: + _validate_search_hit( + hit, + query_id=query_id, + document_count=document_count, + seen_document_ids=seen_document_ids, + ) + neighbors.append(hit.document_id) + distances.append(_score_to_squared_euclidean(hit.score)) + return neighbors, distances + + +def _validate_batch_latencies( + runtime_result: _RuntimeSearchResult, + *, + query_count: int, + batch_size: int, +) -> Tuple[np.ndarray, int]: + latencies = np.asarray(runtime_result.batch_latencies_ms, dtype=np.float64) + num_batches = (query_count + batch_size - 1) // batch_size + if latencies.shape != (num_batches,): + raise RuntimeError( + "PyLucene returned an unexpected number of batch latencies: " + f"{latencies.size} != {num_batches}" + ) + if not np.isfinite(latencies).all() or np.any(latencies < 0.0): + raise RuntimeError( + "PyLucene returned invalid batch latency measurements" + ) + return latencies, num_batches + + +def _process_search_result( + runtime_result: _RuntimeSearchResult, + provenance: _IndexProvenanceVerification, + *, + query_count: int, + k: int, + batch_size: int, +) -> _ProcessedSearchResult: + """Validate the Java boundary result and convert it to benchmark arrays.""" + _validate_runtime_search_result( + runtime_result, provenance, query_count=query_count + ) + neighbors, distances = _convert_search_hits( + runtime_result, query_count=query_count, k=k + ) + latencies, num_batches = _validate_batch_latencies( + runtime_result, query_count=query_count, batch_size=batch_size + ) + + search_time_ms = float(latencies.sum()) + return _ProcessedSearchResult( + neighbors=neighbors, + distances=distances, + search_time_ms=search_time_ms, + queries_per_second=( + query_count / (search_time_ms / 1000.0) + if search_time_ms > 0.0 + else 0.0 + ), + latency_seconds=float(latencies.mean()) / 1000.0, + latency_percentiles={ + "p50": float(np.percentile(latencies, 50)), + "p95": float(np.percentile(latencies, 95)), + "p99": float(np.percentile(latencies, 99)), + }, + num_batches=num_batches, + ) + + +@dataclass(frozen=True) +class _CagraIndexVerification: + segment_count: int + field_count: int + vector_count: int + dimensions: int + + def to_metadata(self) -> Dict[str, Union[str, int]]: + return { + "status": "cagra-only", + "segment_count": self.segment_count, + "field_count": self.field_count, + "vector_count": self.vector_count, + "dimensions": self.dimensions, + } + + +class _CagraVerificationError(RuntimeError): + """Raised when persisted metadata cannot prove a CAGRA-only index.""" + + +@dataclass(frozen=True) +class _CagraFieldMetadata: + field_number: int + dimensions: int + vector_count: int + cagra_offset: int + cagra_length: int + + +@dataclass(frozen=True) +class _RawCagraFieldMetadata: + encoding: int + similarity: int + dimensions: int + vector_count: int + cagra_offset: int + cagra_length: int + brute_force_offset: int + brute_force_length: int + + +@dataclass(frozen=True) +class _CagraFieldSource: + metadata: _CagraFieldMetadata + metadata_file: str + + +@dataclass(frozen=True) +class _CagraDataFileContext: + index_path: Path + directory: Any + segment_info: Any + suffix: str + metadata_file: str + + @property + def data_file(self) -> str: + stem = self.metadata_file[: -len(_CAGRA_META_EXTENSION)] + return stem + _CAGRA_INDEX_EXTENSION + + @property + def data_path(self) -> Path: + return self.index_path / self.data_file + + +@dataclass(frozen=True) +class _CagraSegmentVerification: + field_count: int + vector_count: int + dimensions: frozenset[int] + + +class _CagraIndexVerifier: + """Own persisted CAGRA inspection and its verified-file cache.""" + + def __init__( + self, + *, + attach_current_thread: Callable[[], None], + paths: Any, + codec_util: Any, + field_info: Any, + segment_commit_info: Any, + segment_infos: Any, + vector_encoding: Any, + vector_similarity_function: Any, + fs_directory: Any, + io_context: Any, + ) -> None: + self._attach_current_thread = attach_current_thread + self.Paths = paths + self.CodecUtil = codec_util + self.FieldInfo = field_info + self.SegmentCommitInfo = segment_commit_info + self.SegmentInfos = segment_infos + self.VectorEncoding = vector_encoding + self.VectorSimilarityFunction = vector_similarity_function + self.FSDirectory = fs_directory + self.IOContext = io_context + self._verified_data_files: set[_FileSignature] = set() + + @staticmethod + def _segment_suffix(segment_name: str, metadata_file: str) -> str: + stem = metadata_file[: -len(_CAGRA_META_EXTENSION)] + if stem == segment_name: + return "" + + prefix = f"{segment_name}_" + if not stem.startswith(prefix) or len(stem) == len(prefix): + raise _CagraVerificationError( + "CAGRA-only verification found a metadata filename that " + f"does not match segment {segment_name!r}: {metadata_file!r}" + ) + return stem[len(prefix) :] + + @classmethod + def _read_cagra_field( + cls, + metadata_input: Any, + metadata_file: str, + field_number: int, + ) -> Optional[_CagraFieldMetadata]: + raw_field = cls._decode_cagra_field(metadata_input) + return cls._validate_cagra_field_metadata( + raw_field, metadata_file, field_number + ) + + @staticmethod + def _decode_cagra_field(metadata_input: Any) -> _RawCagraFieldMetadata: + return _RawCagraFieldMetadata( + encoding=int(metadata_input.readInt()), + similarity=int(metadata_input.readInt()), + dimensions=int(metadata_input.readInt()), + vector_count=int(metadata_input.readInt()), + cagra_offset=int(metadata_input.readVLong()), + cagra_length=int(metadata_input.readVLong()), + brute_force_offset=int(metadata_input.readVLong()), + brute_force_length=int(metadata_input.readVLong()), + ) + + @classmethod + def _validate_cagra_field_metadata( + cls, + field: _RawCagraFieldMetadata, + metadata_file: str, + field_number: int, + ) -> Optional[_CagraFieldMetadata]: + if field.encoding != _FLOAT32_ENCODING_ORDINAL: + raise _CagraVerificationError( + "CAGRA-only verification found unsupported vector " + f"encoding ordinal {field.encoding} in {metadata_file!r}" + ) + if field.similarity != _EUCLIDEAN_SIMILARITY_ORDINAL: + raise _CagraVerificationError( + "CAGRA-only verification found unsupported similarity " + f"ordinal {field.similarity} in {metadata_file!r}" + ) + if not 1 <= field.dimensions <= _MAX_DIMENSIONS: + raise _CagraVerificationError( + "CAGRA-only verification found invalid field metadata " + f"in {metadata_file!r}" + ) + if field.vector_count < 0: + raise _CagraVerificationError( + "CAGRA-only verification found invalid field metadata " + f"in {metadata_file!r}" + ) + if field.vector_count == 0: + if cls._empty_cagra_field_has_data(field): + raise _CagraVerificationError( + "CAGRA-only verification found index data for an " + f"empty field in {metadata_file!r}" + ) + return None + if field.brute_force_length != 0: + raise _CagraVerificationError( + "CAGRA-only verification found a persisted brute-force " + f"index for field {field_number} in {metadata_file!r}" + ) + if field.cagra_length <= 0: + raise _CagraVerificationError( + "CAGRA-only verification found no persisted CAGRA index " + f"for field {field_number} in {metadata_file!r}" + ) + return _CagraFieldMetadata( + field_number=field_number, + dimensions=field.dimensions, + vector_count=field.vector_count, + cagra_offset=field.cagra_offset, + cagra_length=field.cagra_length, + ) + + @staticmethod + def _empty_cagra_field_has_data(field: _RawCagraFieldMetadata) -> bool: + return any( + ( + field.cagra_offset, + field.cagra_length, + field.brute_force_offset, + field.brute_force_length, + ) + ) + + @classmethod + def _read_cagra_fields( + cls, metadata_input: Any, metadata_file: str + ) -> List[_CagraFieldMetadata]: + field_numbers = set() + fields = [] + while True: + field_number = int(metadata_input.readInt()) + if field_number == -1: + return fields + if field_number < 0 or field_number in field_numbers: + raise _CagraVerificationError( + "CAGRA-only verification found an invalid field number " + f"{field_number} in {metadata_file!r}" + ) + + field_numbers.add(field_number) + field = cls._read_cagra_field( + metadata_input, metadata_file, field_number + ) + if field is not None: + fields.append(field) + + def _read_segment_field_infos( + self, directory: Any, segment_info: Any + ) -> Any: + if segment_info.hasFieldUpdates(): + raise _CagraVerificationError( + "CAGRA-only verification does not accept segments with " + "field-info updates" + ) + try: + return ( + segment_info.info.getCodec() + .fieldInfosFormat() + .read( + directory, + segment_info.info, + "", + self.IOContext.READONCE, + ) + ) + except Exception as exc: + raise _CagraVerificationError( + "CAGRA-only verification cannot read Lucene field " + f"metadata for segment {segment_info.info.name!r}: {exc}" + ) from exc + + @classmethod + def _verify_cagra_payload_coverage( + cls, + fields: List[_CagraFieldMetadata], + payload_start: int, + payload_end: int, + data_file: str, + ) -> None: + intervals = [] + for field in fields: + intervals.append( + ( + field.cagra_offset, + field.cagra_offset + field.cagra_length, + ) + ) + intervals.sort() + + if not intervals: + if payload_start != payload_end: + raise _CagraVerificationError( + "CAGRA data file contains unreferenced payload: " + f"{data_file!r}" + ) + return + + coverage_error = ( + "CAGRA metadata ranges do not exactly cover data " + f"file {data_file!r}" + ) + expected_offset = payload_start + for interval_start, interval_end in intervals: + if interval_start != expected_offset: + raise _CagraVerificationError(coverage_error) + expected_offset = interval_end + if expected_offset != payload_end: + raise _CagraVerificationError(coverage_error) + + def _verify_checksum_with_signature_cache( + self, + data_path: Path, + data_input: Any, + signature_before: Optional[_FileSignature], + data_file: str, + ) -> None: + cache_hit = ( + signature_before is not None + and signature_before in self._verified_data_files + ) + if cache_hit: + self.CodecUtil.retrieveChecksum(data_input) + else: + self.CodecUtil.checksumEntireFile(data_input) + + if signature_before is None: + return + signature_after = _file_signature(data_path) + if signature_after != signature_before: + raise _CagraVerificationError( + f"CAGRA data file changed during verification: {data_file!r}" + ) + file_age = time.time_ns() - signature_before.changed_at_ns + if cache_hit or file_age < _CAGRA_CACHE_MIN_FILE_AGE_NS: + return + + self._verified_data_files = { + signature + for signature in self._verified_data_files + if signature.resolved_path != signature_before.resolved_path + } + self._verified_data_files.add(signature_before) + + def _verify_cagra_data_file( + self, + context: _CagraDataFileContext, + fields: List[_CagraFieldMetadata], + ) -> None: + try: + signature_before = _file_signature(context.data_path) + except OSError: + signature_before = None + + try: + with _CleanupStack() as cleanups: + data_input = context.directory.openInput( + context.data_file, self.IOContext.READONCE + ) + cleanups.add("close CAGRA data input", data_input.close) + self.CodecUtil.checkIndexHeader( + data_input, + _CAGRA_INDEX_CODEC_NAME, + _CAGRA_INDEX_VERSION, + _CAGRA_INDEX_VERSION, + context.segment_info.info.getId(), + context.suffix, + ) + payload_start = int(data_input.getFilePointer()) + payload_end = int(data_input.length()) - int( + self.CodecUtil.footerLength() + ) + if payload_end < payload_start: + raise _CagraVerificationError( + f"CAGRA data file is truncated: {context.data_file!r}" + ) + self._verify_cagra_payload_coverage( + fields, + payload_start, + payload_end, + context.data_file, + ) + self._verify_checksum_with_signature_cache( + context.data_path, + data_input, + signature_before, + context.data_file, + ) + except _CagraVerificationError: + raise + except Exception as exc: + raise _CagraVerificationError( + "CAGRA-only verification cannot read " + f"{context.data_file!r}: {exc}" + ) from exc + + def _verify_field_against_lucene_metadata( + self, + field_infos: Any, + field: _CagraFieldMetadata, + metadata_file: str, + ) -> None: + field_info = field_infos.fieldInfo(field.field_number) + if field_info is None: + raise _CagraVerificationError( + "CAGRA-only verification found an unknown field number " + f"{field.field_number} in {metadata_file!r}" + ) + field_name = str(field_info.getName()) + if field_name != _VECTOR_FIELD: + raise _CagraVerificationError( + "CAGRA-only verification found data for unexpected field " + f"{field_name!r} in {metadata_file!r}" + ) + if int(field_info.getVectorDimension()) != field.dimensions: + raise _CagraVerificationError( + "CAGRA-only verification found dimensions inconsistent " + f"with Lucene field metadata in {metadata_file!r}" + ) + if field_info.getVectorEncoding() != self.VectorEncoding.FLOAT32: + raise _CagraVerificationError( + "CAGRA-only verification found non-FLOAT32 Lucene field " + f"metadata in {metadata_file!r}" + ) + if ( + field_info.getVectorSimilarityFunction() + != self.VectorSimilarityFunction.EUCLIDEAN + ): + raise _CagraVerificationError( + "CAGRA-only verification found non-Euclidean Lucene field " + f"metadata in {metadata_file!r}" + ) + + @staticmethod + def _validate_segment_deletions( + segment_info: Any, segment_name: str + ) -> None: + deletion_count = int(segment_info.getDelCount()) + soft_deletion_count = int(segment_info.getSoftDelCount()) + if ( + segment_info.hasDeletions() + or deletion_count + or soft_deletion_count + ): + raise _CagraVerificationError( + "CAGRA-only verification does not accept committed " + f"deletions in segment {segment_name!r}: " + f"deleted={deletion_count}, " + f"soft_deleted={soft_deletion_count}" + ) + + @staticmethod + def _cagra_metadata_files( + segment_info: Any, segment_name: str + ) -> List[str]: + metadata_files = [] + for file_name in segment_info.files(): + file_name = str(file_name) + if file_name.endswith(_CAGRA_META_EXTENSION): + metadata_files.append(file_name) + metadata_files.sort() + if not metadata_files: + raise _CagraVerificationError( + "CAGRA-only verification found no " + f"{_CAGRA_META_EXTENSION} metadata for segment " + f"{segment_name!r}" + ) + return metadata_files + + def _read_and_verify_cagra_metadata_file( + self, + context: _CagraDataFileContext, + ) -> List[_CagraFieldMetadata]: + metadata_input = context.directory.openChecksumInput( + context.metadata_file + ) + try: + with _CleanupStack() as cleanups: + cleanups.add( + "close CAGRA metadata input", metadata_input.close + ) + self.CodecUtil.checkIndexHeader( + metadata_input, + _CAGRA_META_CODEC_NAME, + _CAGRA_META_VERSION, + _CAGRA_META_VERSION, + context.segment_info.info.getId(), + context.suffix, + ) + fields = self._read_cagra_fields( + metadata_input, context.metadata_file + ) + self.CodecUtil.checkFooter(metadata_input) + return fields + except _CagraVerificationError: + raise + except Exception as exc: + raise _CagraVerificationError( + "CAGRA-only verification cannot read " + f"{context.metadata_file!r} as cuVS-Lucene metadata " + f"format v{_CAGRA_META_VERSION}: {exc}" + ) from exc + + def _lucene_vector_field_numbers(self, field_infos: Any) -> set[int]: + field_numbers = set() + for raw_field_info in field_infos: + field_info = self.FieldInfo.cast_(raw_field_info) + if int(field_info.getVectorDimension()) > 0: + field_numbers.add(int(field_info.number)) + return field_numbers + + def _verify_cagra_segment( + self, + index_path: Path, + directory: Any, + segment_info: Any, + ) -> _CagraSegmentVerification: + segment_name = str(segment_info.info.name) + self._validate_segment_deletions(segment_info, segment_name) + metadata_files = self._cagra_metadata_files(segment_info, segment_name) + field_infos = self._read_segment_field_infos(directory, segment_info) + + # Verify every metadata/data pair before interpreting fields across + # the segment. This preserves fail-closed file-validation ordering. + verified_field_sources = [] + for metadata_file in metadata_files: + suffix = self._segment_suffix(segment_name, metadata_file) + context = _CagraDataFileContext( + index_path=index_path, + directory=directory, + segment_info=segment_info, + suffix=suffix, + metadata_file=metadata_file, + ) + fields = self._read_and_verify_cagra_metadata_file(context) + self._verify_cagra_data_file(context, fields) + for field in fields: + verified_field_sources.append( + _CagraFieldSource( + metadata=field, + metadata_file=metadata_file, + ) + ) + + # Compare the verified fields with Lucene's segment-wide view and + # aggregate the values needed for index-wide validation. + field_numbers = set() + vector_count = 0 + dimensions = set() + for field_source in verified_field_sources: + field = field_source.metadata + if field.field_number in field_numbers: + raise _CagraVerificationError( + "CAGRA-only verification found duplicate " + f"field {field.field_number} across metadata " + f"files for segment {segment_name!r}" + ) + field_numbers.add(field.field_number) + self._verify_field_against_lucene_metadata( + field_infos, field, field_source.metadata_file + ) + vector_count += field.vector_count + dimensions.add(field.dimensions) + + lucene_field_numbers = self._lucene_vector_field_numbers(field_infos) + if field_numbers != lucene_field_numbers: + raise _CagraVerificationError( + "CAGRA-only verification found vector fields without " + "matching CAGRA metadata in segment " + f"{segment_name!r}: metadata={sorted(field_numbers)}, " + f"Lucene={sorted(lucene_field_numbers)}" + ) + + max_documents = int(segment_info.info.maxDoc()) + if vector_count != max_documents: + raise _CagraVerificationError( + "CAGRA-only verification found " + f"{vector_count} vectors for " + f"{max_documents} documents in segment " + f"{segment_name!r}" + ) + return _CagraSegmentVerification( + field_count=len(field_numbers), + vector_count=vector_count, + dimensions=frozenset(dimensions), + ) + + @staticmethod + def _summarize_cagra_segments( + segments: List[_CagraSegmentVerification], + expected_vector_count: Optional[int], + expected_dimensions: Optional[int], + ) -> _CagraIndexVerification: + field_count = 0 + vector_count = 0 + dimensions = set() + for segment in segments: + field_count += segment.field_count + vector_count += segment.vector_count + dimensions.update(segment.dimensions) + if not segments or field_count == 0: + raise _CagraVerificationError( + "CAGRA-only verification found no nonempty vector fields" + ) + if len(dimensions) != 1: + raise _CagraVerificationError( + "CAGRA-only verification found inconsistent vector " + f"dimensions: {sorted(dimensions)}" + ) + + index_dimensions = next(iter(dimensions)) + if ( + expected_vector_count is not None + and vector_count != expected_vector_count + ): + raise _CagraVerificationError( + "CAGRA-only verification found " + f"{vector_count} vectors; expected {expected_vector_count}" + ) + if ( + expected_dimensions is not None + and index_dimensions != expected_dimensions + ): + raise _CagraVerificationError( + "CAGRA-only verification found " + f"{index_dimensions} dimensions; expected " + f"{expected_dimensions}" + ) + return _CagraIndexVerification( + segment_count=len(segments), + field_count=field_count, + vector_count=vector_count, + dimensions=index_dimensions, + ) + + def verify_index( + self, + index_path: Path, + *, + expected_vector_count: Optional[int] = None, + expected_dimensions: Optional[int] = None, + ) -> _CagraIndexVerification: + """Verify every committed vector field is persisted as CAGRA only.""" + self._attach_current_thread() + directory = self.FSDirectory.open(self.Paths.get(str(index_path))) + verified_segments = [] + with _CleanupStack() as cleanups: + cleanups.add("close Lucene directory", directory.close) + segment_infos = self.SegmentInfos.readLatestCommit(directory) + for raw_segment_info in segment_infos: + segment_info = self.SegmentCommitInfo.cast_(raw_segment_info) + verified_segments.append( + self._verify_cagra_segment( + index_path, directory, segment_info + ) + ) + return self._summarize_cagra_segments( + verified_segments, + expected_vector_count, + expected_dimensions, + ) + + +class _PyLuceneRuntime: + """Own generated PyLucene/Lucene bindings and index operations.""" + + def __init__(self, lucene: Any): + from java.lang import Class + from java.nio.file import Paths + from org.apache.lucene.codecs import Codec, CodecUtil + from org.apache.lucene.document import ( + Document, + KnnFloatVectorField, + StoredField, + ) + from org.apache.lucene.index import ( + DirectoryReader, + FieldInfo, + IndexWriter, + IndexWriterConfig, + SegmentCommitInfo, + SegmentInfos, + SerialMergeScheduler, + VectorEncoding, + VectorSimilarityFunction, + ) + from org.apache.lucene.search import ( + IndexSearcher, + KnnFloatVectorQuery, + ) + from org.apache.lucene.store import FSDirectory, IOContext + + self.lucene = lucene + self.Class = Class + self.Paths = Paths + self.Codec = Codec + self.Document = Document + self.KnnFloatVectorField = KnnFloatVectorField + self.StoredField = StoredField + self.DirectoryReader = DirectoryReader + self.IndexWriter = IndexWriter + self.IndexWriterConfig = IndexWriterConfig + self.SerialMergeScheduler = SerialMergeScheduler + self.VectorSimilarityFunction = VectorSimilarityFunction + self.IndexSearcher = IndexSearcher + self.KnnFloatVectorQuery = KnnFloatVectorQuery + self.FSDirectory = FSDirectory + self.IOContext = IOContext + self._codec_cache: Dict[str, Any] = {} + + # Resolve the base cuvs-java provider before codec construction so + # classpath failures have a direct error location. + self.Class.forName("com.nvidia.cuvs.spi.JDKProvider") + self._cagra_verifier = _CagraIndexVerifier( + attach_current_thread=self.attach_current_thread, + paths=Paths, + codec_util=CodecUtil, + field_info=FieldInfo, + segment_commit_info=SegmentCommitInfo, + segment_infos=SegmentInfos, + vector_encoding=VectorEncoding, + vector_similarity_function=VectorSimilarityFunction, + fs_directory=FSDirectory, + io_context=IOContext, + ) + + @classmethod + def create(cls, config: Dict[str, Any]) -> "_PyLuceneRuntime": + return cls(_initialize_pylucene(config)) + + @property + def pylucene_version(self) -> str: + return str(getattr(self.lucene, "VERSION", "unknown")) + + def attach_current_thread(self) -> None: + vm_environment = self.lucene.getVMEnv() + if vm_environment is None: + raise RuntimeError("PyLucene JVM is not initialized") + vm_environment.attachCurrentThread() + + def resolve_codec(self, codec_name: str) -> Any: + self.attach_current_thread() + cached = self._codec_cache.get(codec_name) + if cached is not None: + return cached + + available_codecs = self.Codec.availableCodecs() + if not available_codecs.contains(codec_name): + available = ", ".join(str(name) for name in available_codecs) + raise RuntimeError( + f"{codec_name} was not advertised by Lucene SPI. " + f"Available codecs: {available}" + ) + codec = self.Codec.forName(codec_name) + if str(codec.getName()) != codec_name: + raise RuntimeError( + f"Requested codec {codec_name}, got {codec.getName()}" + ) + self._codec_cache[codec_name] = codec + return codec + + def codec_telemetry(self, codec_name: str, codec: Any) -> Dict[str, str]: + """Inspect the selected writer path before indexing starts.""" + vectors_format = codec.knnVectorsFormat() + if vectors_format is None: + raise RuntimeError( + f"{codec_name} did not initialize a Lucene vector format" + ) + return _parse_writer_telemetry(str(vectors_format)) + + def _java_float_array(self, vector: np.ndarray) -> Any: + return self.lucene.JArray("float")( + tuple(float(value) for value in vector) + ) + + def verify_cagra_index( + self, + index_path: Path, + expected_vector_count: Optional[int] = None, + expected_dimensions: Optional[int] = None, + ) -> _CagraIndexVerification: + return self._cagra_verifier.verify_index( + index_path, + expected_vector_count=expected_vector_count, + expected_dimensions=expected_dimensions, + ) + + def build_index( + self, + index_path: Path, + vectors: np.ndarray, + resolved_codec: _ResolvedCodec, + ) -> Dict[str, str]: + self.attach_current_thread() + directory = self.FSDirectory.open(self.Paths.get(str(index_path))) + with _CleanupStack() as cleanups: + cleanups.add("close Lucene directory", directory.close) + writer_config = self.IndexWriterConfig() + writer_config.setOpenMode(self.IndexWriterConfig.OpenMode.CREATE) + writer_config.setCodec(resolved_codec.java_codec) + writer_config.setUseCompoundFile(False) + writer_config.setMergeScheduler(self.SerialMergeScheduler()) + writer = self.IndexWriter(directory, writer_config) + self._write_and_close_index(writer, vectors) + return dict(resolved_codec.telemetry) + + def _write_and_close_index(self, writer: Any, vectors: np.ndarray) -> None: + try: + for document_id, vector in enumerate(vectors): + writer.addDocument(self._vector_document(document_id, vector)) + writer.commit() + except BaseException as write_error: + self._rollback_failed_write(writer, write_error) + raise + writer.close() + + @staticmethod + def _rollback_failed_write( + writer: Any, write_error: BaseException + ) -> None: + try: + writer.rollback() + except BaseException as rollback_error: + if isinstance(write_error, Exception): + raise rollback_error from write_error + write_error.add_note( + "IndexWriter rollback also failed: " + f"{type(rollback_error).__name__}: {rollback_error}" + ) + + def _vector_document(self, document_id: int, vector: np.ndarray) -> Any: + document = self.Document() + document.add(self.StoredField(_ID_FIELD, str(document_id))) + document.add( + self.KnnFloatVectorField( + _VECTOR_FIELD, + self._java_float_array(vector), + self.VectorSimilarityFunction.EUCLIDEAN, + ) + ) + return document + + @staticmethod + def _index_vector_dimensions(reader: Any) -> int: + dimensions = set() + for leaf_context in reader.leaves(): + values = leaf_context.reader().getFloatVectorValues(_VECTOR_FIELD) + if values is not None and values.size() > 0: + dimensions.add(int(values.dimension())) + if not dimensions: + raise RuntimeError( + f"Lucene index contains no {_VECTOR_FIELD!r} vectors" + ) + if len(dimensions) != 1: + raise RuntimeError( + "Lucene index has inconsistent vector dimensions: " + f"{dimensions}" + ) + return dimensions.pop() + + def _search_vector( + self, + searcher: Any, + stored_fields: Any, + vector: np.ndarray, + k: int, + ) -> List[_SearchHit]: + query = self.KnnFloatVectorQuery( + _VECTOR_FIELD, + self._java_float_array(vector), + k, + ) + score_docs = searcher.search(query, k).scoreDocs + hits = [] + for score_doc in score_docs: + stored_id = stored_fields.document(score_doc.doc).get(_ID_FIELD) + if stored_id is None: + raise RuntimeError( + f"Lucene document {score_doc.doc} has no stored ID" + ) + hits.append( + _SearchHit( + document_id=int(stored_id), + score=float(score_doc.score), + ) + ) + return hits + + def search_index( + self, + index_path: Path, + query_vectors: np.ndarray, + k: int, + batch_size: int, + ) -> _RuntimeSearchResult: + self.attach_current_thread() + directory = self.FSDirectory.open(self.Paths.get(str(index_path))) + with _CleanupStack() as cleanups: + cleanups.add("close Lucene directory", directory.close) + reader = self.DirectoryReader.open(directory) + cleanups.add("close Lucene index reader", reader.close) + index_dimensions = self._index_vector_dimensions(reader) + if query_vectors.shape[1] != index_dimensions: + raise ValueError( + "Query vector dimensions do not match the Lucene index: " + f"{query_vectors.shape[1]} != {index_dimensions}" + ) + + document_count = int(reader.numDocs()) + if document_count < 1: + raise RuntimeError("Lucene index contains no documents") + lucene_k = min(k, document_count) + searcher = self.IndexSearcher(reader) + stored_fields = searcher.storedFields() + + all_hits: List[List[_SearchHit]] = [] + batch_latencies_ms: List[float] = [] + for batch_start in range(0, query_vectors.shape[0], batch_size): + start = time.perf_counter() + batch_end = min( + batch_start + batch_size, query_vectors.shape[0] + ) + all_hits.extend( + self._search_vectors( + searcher, + stored_fields, + query_vectors[batch_start:batch_end], + lucene_k, + ) + ) + batch_latencies_ms.append( + (time.perf_counter() - start) * 1000.0 + ) + + return _RuntimeSearchResult( + hits=all_hits, + batch_latencies_ms=batch_latencies_ms, + index_dimensions=index_dimensions, + document_count=document_count, + ) + + def _search_vectors( + self, + searcher: Any, + stored_fields: Any, + query_vectors: np.ndarray, + k: int, + ) -> List[List[_SearchHit]]: + hits = [] + for vector in query_vectors: + hits.append( + self._search_vector(searcher, stored_fields, vector, k) + ) + return hits + + +class _SelectedAlgorithmGroup(NamedTuple): + algorithm: str + group: str + configuration: dict + metadata: dict + + +@dataclass(frozen=True) +class _GlobalSelectionScope: + algorithms: frozenset[str] + groups: frozenset[str] + + +@dataclass(frozen=True) +class _AlgorithmSelection: + algorithms: Optional[frozenset[str]] + groups: Optional[frozenset[str]] + explicit_groups: frozenset[Tuple[str, str]] + + def _resolve_global_scope( + self, algorithm_configs: Dict[str, Dict[str, Any]] + ) -> Optional[_GlobalSelectionScope]: + available_algorithms = frozenset(algorithm_configs) + if self.algorithms is not None: + unknown_algorithms = self.algorithms - available_algorithms + if unknown_algorithms: + names = ", ".join(sorted(unknown_algorithms)) + raise ValueError( + f"Unknown PyLucene algorithm selector(s): {names}" + ) + + candidate_algorithms = ( + self.algorithms + if self.algorithms is not None + else available_algorithms + ) + if self.groups is not None: + available_groups = set() + for algorithm in candidate_algorithms: + available_groups.update(algorithm_configs[algorithm]) + unknown_groups = self.groups - available_groups + if unknown_groups: + names = ", ".join(sorted(unknown_groups)) + raise ValueError( + f"Unknown PyLucene group selector(s): {names}" + ) + + has_global_selector = ( + self.algorithms is not None or self.groups is not None + ) + if not has_global_selector and self.explicit_groups: + return None + return _GlobalSelectionScope( + algorithms=candidate_algorithms, + groups=( + self.groups if self.groups is not None else frozenset({"base"}) + ), + ) + + def _validate_explicit_groups( + self, algorithm_configs: Dict[str, Dict[str, Any]] + ) -> None: + for algorithm, group in sorted(self.explicit_groups): + if algorithm not in algorithm_configs: + raise ValueError( + f"Unknown PyLucene algorithm in --algo-groups: {algorithm}" + ) + if group not in algorithm_configs[algorithm]: + raise ValueError( + f"Unknown PyLucene group for {algorithm}: {group}" + ) + + def _selected_pairs( + self, + algorithm_configs: Dict[str, Dict[str, Any]], + global_scope: Optional[_GlobalSelectionScope], + ) -> set[Tuple[str, str]]: + selected_pairs = set(self.explicit_groups) + if global_scope is None: + return selected_pairs + + for algorithm, groups in algorithm_configs.items(): + if algorithm not in global_scope.algorithms: + continue + for group in groups: + if group in global_scope.groups: + selected_pairs.add((algorithm, group)) + return selected_pairs + + def resolve( + self, algorithm_configs: Dict[str, Dict[str, Any]] + ) -> List[_SelectedAlgorithmGroup]: + global_scope = self._resolve_global_scope(algorithm_configs) + self._validate_explicit_groups(algorithm_configs) + selected_pairs = self._selected_pairs(algorithm_configs, global_scope) + + selected_groups = [] + for algorithm, groups in algorithm_configs.items(): + for group, group_config in groups.items(): + if (algorithm, group) not in selected_pairs: + continue + selected_groups.append( + _SelectedAlgorithmGroup( + algorithm=algorithm, + group=group, + configuration=group_config, + metadata={}, + ) + ) + return selected_groups + + +@dataclass(frozen=True) +class _BenchmarkConfigContext: + dataset: str + dataset_path: str + subset_scope: Optional[str] + runtime_config: Dict[str, Any] + neighbor_count: int + batch_size: int + + +class PyLuceneConfigLoader(ConfigLoader): + """Load PyLucene algorithm configurations.""" + + def __init__(self, config_path: Optional[Union[str, os.PathLike]] = None): + self.config_path = ( + os.fspath(config_path) + if config_path is not None + else os.path.join( + os.path.dirname(os.path.realpath(__file__)), "../config" + ) + ) + + @property + def backend_type(self) -> str: + return "pylucene" + + @staticmethod + def _parse_name_filter( + value: Optional[str], + ) -> Optional[frozenset[str]]: + if not value: + return None + + names = set() + for raw_name in value.split(","): + name = raw_name.strip() + if name: + names.add(name) + return frozenset(names) or None + + @staticmethod + def _parse_algorithm_group_filter( + value: Optional[str], + ) -> frozenset[Tuple[str, str]]: + selections = set() + if not value: + return frozenset() + + for selection in value.split(","): + algorithm, separator, group = selection.strip().partition(".") + if not separator or not algorithm or not group: + raise ValueError( + "algo_groups entries must use ., " + f"got {selection!r}" + ) + selections.add((algorithm, group)) + return frozenset(selections) + + @classmethod + def _selection_from_options( + cls, options: Dict[str, Any] + ) -> _AlgorithmSelection: + return _AlgorithmSelection( + algorithms=cls._parse_name_filter(options.get("algorithms")), + groups=cls._parse_name_filter(options.get("groups")), + explicit_groups=cls._parse_algorithm_group_filter( + options.get("algo_groups") + ), + ) + + def _match_backend_algorithm_config( + self, config: Any + ) -> Optional[Tuple[str, Dict[str, Any]]]: + if not isinstance(config, dict): + return None + + algorithm = config.get("name") + if not isinstance(algorithm, str) or not algorithm: + return None + + declared_backend = config.get("backend") + if declared_backend is None: + belongs_to_backend = algorithm.startswith("pylucene_") + else: + belongs_to_backend = declared_backend == self.backend_type + if not belongs_to_backend: + return None + return algorithm, config.get("groups", {}) + + def _load_algorithm_configs( + self, algorithm_files: List[str] + ) -> Dict[str, Dict[str, Any]]: + algorithm_configs = {} + for algorithm_file in algorithm_files: + matched_config = self._match_backend_algorithm_config( + self.load_yaml_file(algorithm_file) + ) + if matched_config is None: + continue + algorithm, groups = matched_config + # Later files are explicit overrides and determine output order. + algorithm_configs.pop(algorithm, None) + algorithm_configs[algorithm] = groups + return algorithm_configs + + def _discover_algo_groups( + self, + dataset_conf: dict, + dataset: str, + dataset_path: str, + **kwargs, + ) -> List[Tuple[str, str, dict, dict]]: + algorithm_files = self.gather_algorithm_configs( + self.config_path, kwargs.get("algorithm_configuration") + ) + algorithm_configs = self._load_algorithm_configs(algorithm_files) + selection = self._selection_from_options(kwargs) + return selection.resolve(algorithm_configs) + + @staticmethod + def _runtime_config(options: Dict[str, Any]) -> Dict[str, Any]: + runtime_keys = ( + "cuvs_java_jar", + "cuvs_lucene_jar", + "java_library_path", + "jvm_args", + ) + runtime_config = {} + for key in runtime_keys: + value = options.get(key) + if value is not None: + runtime_config[key] = value + return runtime_config + + @staticmethod + def _subset_scope(subset_size: Any) -> Optional[str]: + subset_size = _validate_subset_size(subset_size) + if subset_size is None: + return None + return f"subset{subset_size}" + + @staticmethod + def _effective_parameters( + build_combinations: List[Dict[str, Any]], + search_combinations: List[Dict[str, Any]], + *, + tune_mode: bool, + tune_build_params: Optional[Dict[str, Any]], + tune_search_params: Optional[Dict[str, Any]], + ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + use_tuned_parameters = tune_mode and tune_build_params is not None + if not use_tuned_parameters: + return build_combinations, search_combinations + search_params = [tune_search_params] if tune_search_params else [{}] + return [tune_build_params], search_params + + @staticmethod + def _index_label( + algorithm: str, + group: str, + subset_scope: Optional[str], + build_params: Dict[str, Any], + ) -> str: + prefix = algorithm if group == "base" else f"{algorithm}_{group}" + label_parts = [prefix] + if subset_scope is not None: + label_parts.append(subset_scope) + for key, value in build_params.items(): + label_parts.append(f"{key}{value}") + return ".".join(label_parts) + + @classmethod + def _benchmark_config( + cls, + context: _BenchmarkConfigContext, + algorithm: str, + group: str, + build_params: Dict[str, Any], + search_params: List[Dict[str, Any]], + ) -> BenchmarkConfig: + index_label = cls._index_label( + algorithm, group, context.subset_scope, build_params + ) + index_path = os.path.join( + context.dataset_path, + context.dataset, + "index", + index_label, + ) + result_stem = f"{algorithm},{group}" + if context.subset_scope is not None: + result_stem = f"{result_stem},{context.subset_scope}" + + index_config = IndexConfig( + name=index_label, + algo=algorithm, + build_param=build_params, + search_params=search_params, + file=index_path, + ) + backend_config = { + "name": index_label, + "algo": algorithm, + "codec": build_params.get("codec"), + "requires_gpu": True, + "output_filename": ( + result_stem, + f"{result_stem},k{context.neighbor_count}," + f"bs{context.batch_size}", + ), + **context.runtime_config, + } + return BenchmarkConfig( + indexes=[index_config], + backend_config=backend_config, + ) + + @classmethod + def _group_benchmark_configs( + cls, + context: _BenchmarkConfigContext, + algorithm: str, + group: str, + build_params: List[Dict[str, Any]], + search_params: List[Dict[str, Any]], + ) -> List[BenchmarkConfig]: + _validate_search_params(search_params) + benchmark_configs = [] + for params in build_params: + _validate_build_params(params) + benchmark_configs.append( + cls._benchmark_config( + context, + algorithm, + group, + params, + search_params, + ) + ) + return benchmark_configs + + def _build_benchmark_configs( + self, + dataset_config: DatasetConfig, + dataset_conf: dict, + dataset: str, + dataset_path: str, + expanded_groups: List[Tuple[str, str, dict, List, List, dict]], + **kwargs, + ) -> List[BenchmarkConfig]: + context = _BenchmarkConfigContext( + dataset=dataset, + dataset_path=dataset_path, + subset_scope=self._subset_scope(dataset_config.subset_size), + runtime_config=self._runtime_config(kwargs), + neighbor_count=kwargs.get("count", 10), + batch_size=kwargs.get("batch_size", 10000), + ) + tune_mode = kwargs.get("_tune_mode", False) + tune_build_params = kwargs.get("_tune_build_params") + tune_search_params = kwargs.get("_tune_search_params") + + benchmark_configs = [] + for ( + algorithm, + group, + _group_config, + build_combinations, + search_combinations, + _group_metadata, + ) in expanded_groups: + build_params, search_params = self._effective_parameters( + build_combinations, + search_combinations, + tune_mode=tune_mode, + tune_build_params=tune_build_params, + tune_search_params=tune_search_params, + ) + benchmark_configs.extend( + self._group_benchmark_configs( + context, + algorithm, + group, + build_params, + search_params, + ) + ) + + return benchmark_configs + + +class PyLuceneBackend(BenchmarkBackend): + """Build and search cuVS-Lucene indexes through PyLucene.""" + + orchestrator_persists_results = True + + def __init__(self, config: Dict[str, Any]): + super().__init__(config) + self._runtime: Optional[_PyLuceneRuntime] = None + + @property + def algo(self) -> str: + return self.config.get("algo", "pylucene") + + def cleanup(self) -> None: + """Release Python references; the process-wide JVM remains active.""" + self._runtime = None + + def _get_runtime(self) -> _PyLuceneRuntime: + if self._runtime is None: + self._runtime = _PyLuceneRuntime.create(self.config) + return self._runtime + + def _failed_build_result( + self, + error_message: str, + index_path: str = "", + build_params: Optional[Dict[str, Any]] = None, + ) -> BuildResult: + return BuildResult( + index_path=index_path, + build_time_seconds=0.0, + index_size_bytes=0, + algorithm=self.algo, + build_params=build_params or {}, + success=False, + error_message=error_message, + ) + + def _failed_search_result( + self, + k: int, + error_message: str, + search_params: Optional[List[Dict[str, Any]]] = None, + ) -> SearchResult: + result_k = max(0, k) + return SearchResult( + neighbors=np.empty((0, result_k), dtype=np.int64), + distances=np.empty((0, result_k), dtype=np.float32), + search_time_ms=0.0, + queries_per_second=0.0, + recall=0.0, + algorithm=self.algo, + search_params=search_params or [], + success=False, + error_message=error_message, + ) + + def _verify_existing_index( + self, + index_path: Path, + codec_name: str, + *, + expected_vector_count: Optional[int] = None, + expected_dimensions: Optional[int] = None, + ) -> Tuple[_IndexProvenanceVerification, Dict[str, Any]]: + """Verify backend ownership and codec-specific persisted data.""" + if codec_name == _CAGRA_CODEC: + provenance = _verify_cagra_provenance( + index_path, + expected_vector_count=expected_vector_count, + expected_dimensions=expected_dimensions, + ) + cagra_verification = self._get_runtime().verify_cagra_index( + index_path, + expected_vector_count=provenance.vector_count, + expected_dimensions=provenance.dimensions, + ) + metadata = { + "cagra_provenance": provenance.to_metadata(), + "cagra_verification": cagra_verification.to_metadata(), + } + else: + provenance = _verify_hnsw_provenance( + index_path, + codec_name, + expected_vector_count=expected_vector_count, + expected_dimensions=expected_dimensions, + ) + metadata = {"hnsw_verification": provenance.to_metadata()} + return provenance, metadata + + def _dry_run_build_result( + self, + index_path: Path, + codec_name: str, + build_params: Dict[str, Any], + ) -> BuildResult: + print( + f"[dry_run] Would build PyLucene index '{index_path}' " + f"with codec={codec_name}" + ) + return BuildResult( + index_path=str(index_path), + build_time_seconds=0.0, + index_size_bytes=0, + algorithm=self.algo, + build_params=build_params, + metadata={"codec": codec_name}, + success=True, + ) + + def _decide_existing_index( + self, + dataset: Dataset, + index_path: Path, + codec_name: str, + build_params: Dict[str, Any], + force: bool, + ) -> _ExistingIndexDecision: + if not index_path.exists(): + return _ExistingIndexDecision.build() + if not index_path.is_dir(): + return _ExistingIndexDecision.reject( + self._failed_build_result( + f"PyLucene index path is not a directory: {index_path}", + str(index_path), + build_params, + ) + ) + if not _has_lucene_segments(index_path): + return _ExistingIndexDecision.reject( + self._failed_build_result( + "Existing PyLucene index directory does not contain " + f"a Lucene segments file: {index_path}", + str(index_path), + build_params, + ) + ) + if force: + return _ExistingIndexDecision.build() + + metadata: Dict[str, Any] = { + "codec": codec_name, + "skipped": True, + } + try: + _validate_metric(dataset) + expected_shape = _expected_training_shape(dataset) + _, verification_metadata = self._verify_existing_index( + index_path, + codec_name, + expected_vector_count=( + expected_shape[0] if expected_shape else None + ), + expected_dimensions=( + expected_shape[1] if expected_shape else None + ), + ) + metadata.update(verification_metadata) + result = BuildResult( + index_path=str(index_path), + build_time_seconds=0.0, + index_size_bytes=_index_size(index_path), + algorithm=self.algo, + build_params=build_params, + metadata=metadata, + success=True, + ) + except Exception as exc: + return _ExistingIndexDecision.reject( + self._failed_build_result( + _exception_summary(exc), + str(index_path), + build_params, + ) + ) + + return _ExistingIndexDecision.reuse(result) + + @staticmethod + def _training_vectors_for_build( + dataset: Dataset, codec_name: str + ) -> np.ndarray: + _validate_metric(dataset) + vectors = _validate_float32_matrix( + dataset.training_vectors, "training_vectors" + ) + if vectors.shape[0] < 2: + raise ValueError( + f"{codec_name} requires at least two training vectors; " + "cuVS-Lucene does not invoke cuVS for a single-vector index" + ) + return vectors + + @staticmethod + def _require_gpu_writer( + runtime: _PyLuceneRuntime, codec_name: str + ) -> _ResolvedCodec: + java_codec = runtime.resolve_codec(codec_name) + telemetry = runtime.codec_telemetry(codec_name, java_codec) + writer_path = telemetry.get("writerPath") + expected_writer_path = _EXPECTED_WRITER_PATH[codec_name] + if writer_path != expected_writer_path: + raise RuntimeError( + f"{codec_name} uses writerPath={writer_path!r}; " + f"expected {expected_writer_path!r}. Refusing to run " + "a silent CPU or alternate-index fallback as a cuVS build." + ) + return _ResolvedCodec( + codec_name=codec_name, + java_codec=java_codec, + telemetry=telemetry, + writer_path=writer_path, + ) + + @staticmethod + def _verify_and_persist_built_index_provenance( + runtime: _PyLuceneRuntime, + index_path: Path, + codec_name: str, + vectors: np.ndarray, + ) -> Dict[str, Any]: + vector_count = int(vectors.shape[0]) + dimensions = int(vectors.shape[1]) + if codec_name == _CAGRA_CODEC: + cagra_verification = runtime.verify_cagra_index( + index_path, + expected_vector_count=vector_count, + expected_dimensions=dimensions, + ) + _write_cagra_provenance( + index_path, + vector_count=vector_count, + dimensions=dimensions, + ) + provenance = _verify_cagra_provenance( + index_path, + expected_vector_count=vector_count, + expected_dimensions=dimensions, + ) + return { + "cagra_provenance": provenance.to_metadata(), + "cagra_verification": cagra_verification.to_metadata(), + } + + _write_hnsw_provenance( + index_path, + codec_name, + vector_count=vector_count, + dimensions=dimensions, + ) + verification = _verify_hnsw_provenance( + index_path, + codec_name, + expected_vector_count=vector_count, + expected_dimensions=dimensions, + ) + return {"hnsw_verification": verification.to_metadata()} + + @staticmethod + def _cleanup_partial_index( + index_path: Path, created_for_build: bool + ) -> Optional[Exception]: + if not created_for_build or not index_path.exists(): + return None + try: + _safe_remove_index(index_path) + except Exception as exc: + return exc + return None + + def build( + self, + dataset: Dataset, + indexes: List[IndexConfig], + force: bool = False, + dry_run: bool = False, + ) -> BuildResult: + """Build one local Lucene index with the selected cuVS codec.""" + if len(indexes) != 1: + return self._failed_build_result( + "PyLucene backend requires exactly one index configuration" + ) + + index_config = indexes[0] + build_params = index_config.build_param + index_path = Path(index_config.file) + try: + _validate_build_params(build_params) + codec_name = _configured_codec(build_params, self.config) + except Exception as exc: + return self._failed_build_result( + str(exc), str(index_path), build_params + ) + + if dry_run: + return self._dry_run_build_result( + index_path, codec_name, build_params + ) + + existing_index_decision = self._decide_existing_index( + dataset, index_path, codec_name, build_params, force + ) + if existing_index_decision.action is _ExistingIndexAction.BUILD: + return self._build_new_index( + dataset, index_path, codec_name, build_params + ) + + return existing_index_decision.completed_result() + + def _build_new_index( + self, + dataset: Dataset, + index_path: Path, + codec_name: str, + build_params: Dict[str, Any], + ) -> BuildResult: + """Build a validated index that is not eligible for reuse.""" + + created_for_build = False + try: + vectors = self._training_vectors_for_build(dataset, codec_name) + runtime = self._get_runtime() + # Preflight must succeed before an existing index is replaced. + resolved_codec = self._require_gpu_writer(runtime, codec_name) + + if index_path.exists(): + _safe_remove_index(index_path) + index_path.mkdir(parents=True) + created_for_build = True + + start = time.perf_counter() + telemetry = runtime.build_index( + index_path, vectors, resolved_codec + ) + build_time = time.perf_counter() - start + + metadata = { + "codec": codec_name, + "pylucene_version": runtime.pylucene_version, + "writer_path": resolved_codec.writer_path, + "writer_telemetry": telemetry, + } + metadata.update( + self._verify_and_persist_built_index_provenance( + runtime, index_path, codec_name, vectors + ) + ) + + return BuildResult( + index_path=str(index_path), + build_time_seconds=build_time, + index_size_bytes=_index_size(index_path), + algorithm=self.algo, + build_params=build_params, + metadata=metadata, + success=True, + ) + except Exception as exc: + cleanup_error = self._cleanup_partial_index( + index_path, created_for_build + ) + error_message = _exception_summary(exc) + if cleanup_error is not None: + error_message += ( + "; failed to remove partial index: " + f"{type(cleanup_error).__name__}: {cleanup_error}" + ) + return self._failed_build_result( + error_message, + str(index_path), + build_params, + ) + except BaseException as exc: + cleanup_error = self._cleanup_partial_index( + index_path, created_for_build + ) + if cleanup_error is not None: + exc.add_note( + "Failed to remove partial index: " + f"{type(cleanup_error).__name__}: {cleanup_error}" + ) + raise + + def _resolve_search_plan( + self, + index_config: IndexConfig, + *, + k: int, + batch_size: int, + mode: str, + search_threads: Optional[Union[int, str]], + ) -> _SearchPlan: + search_params = index_config.search_params or [{}] + _validate_build_params(index_config.build_param) + codec_name = _configured_codec(index_config.build_param, self.config) + if k < 1: + raise ValueError("k must be positive") + if batch_size < 1: + raise ValueError("batch_size must be positive") + if mode != "latency": + raise ValueError( + "PyLucene backend currently supports only latency mode" + ) + if search_threads not in (None, 1, "1"): + raise ValueError( + "PyLucene backend currently supports only one search thread" + ) + _validate_search_params(search_params) + if codec_name == _CAGRA_CODEC and k > 1024: + raise ValueError( + "CuVS2510GPUSearchCodec benchmarks support k <= 1024 " + "to avoid cuVS-Lucene search paths that can use GPU " + "brute-force search above that limit" + ) + return _SearchPlan( + index_path=Path(index_config.file), + codec_name=codec_name, + search_params=search_params, + k=k, + batch_size=batch_size, + mode=mode, + ) + + def _dry_run_search_result( + self, + plan: _SearchPlan, + ) -> SearchResult: + print( + f"[dry_run] Would search PyLucene index '{plan.index_path}' " + f"with codec={plan.codec_name}, k={plan.k}, " + f"batch_size={plan.batch_size}" + ) + return SearchResult( + neighbors=np.empty((0, plan.k), dtype=np.int64), + distances=np.empty((0, plan.k), dtype=np.float32), + search_time_ms=0.0, + queries_per_second=0.0, + recall=0.0, + algorithm=self.algo, + search_params=plan.search_params, + metadata={"codec": plan.codec_name}, + success=True, + ) + + @staticmethod + def _load_search_inputs(dataset: Dataset) -> _SearchInputs: + _validate_metric(dataset) + query_vectors = _validate_float32_matrix( + dataset.query_vectors, "query_vectors" + ) + expected_shape = _expected_training_shape(dataset) + if ( + expected_shape is not None + and expected_shape[1] != query_vectors.shape[1] + ): + raise ValueError( + "Query vector dimensions do not match the dataset: " + f"{query_vectors.shape[1]} != {expected_shape[1]}" + ) + return _SearchInputs( + query_vectors=query_vectors, + expected_training_shape=expected_shape, + ) + + def _execute_search( + self, dataset: Dataset, plan: _SearchPlan + ) -> SearchResult: + end_to_end_start = time.perf_counter() + inputs = self._load_search_inputs(dataset) + query_vectors = inputs.query_vectors + expected_shape = inputs.expected_training_shape + provenance, verification_metadata = self._verify_existing_index( + plan.index_path, + plan.codec_name, + expected_vector_count=( + expected_shape[0] if expected_shape else None + ), + expected_dimensions=int(query_vectors.shape[1]), + ) + runtime = self._get_runtime() + + runtime_result = runtime.search_index( + plan.index_path, + query_vectors, + plan.k, + plan.batch_size, + ) + end_to_end_time_ms = (time.perf_counter() - end_to_end_start) * 1000.0 + processed = _process_search_result( + runtime_result, + provenance, + query_count=int(query_vectors.shape[0]), + k=plan.k, + batch_size=plan.batch_size, + ) + metadata = { + "codec": plan.codec_name, + "pylucene_version": runtime.pylucene_version, + "index_dimensions": runtime_result.index_dimensions, + "document_count": runtime_result.document_count, + "batch_size": plan.batch_size, + "num_batches": processed.num_batches, + "mode": plan.mode, + "end_to_end_time_ms": end_to_end_time_ms, + "non_query_overhead_time_ms": max( + 0.0, end_to_end_time_ms - processed.search_time_ms + ), + } + metadata.update(verification_metadata) + + return SearchResult( + neighbors=processed.neighbors, + distances=processed.distances, + search_time_ms=processed.search_time_ms, + queries_per_second=processed.queries_per_second, + recall=0.0, + algorithm=self.algo, + search_params=plan.search_params, + latency_seconds=processed.latency_seconds, + latency_percentiles=processed.latency_percentiles, + metadata=metadata, + success=True, + ) + + def search( + self, + dataset: Dataset, + indexes: List[IndexConfig], + k: int, + batch_size: int = 10000, + mode: str = "latency", + force: bool = False, + search_threads: Optional[int] = None, + dry_run: bool = False, + ) -> SearchResult: + """Search one local Lucene index and report per-batch latency.""" + if len(indexes) != 1: + return self._failed_search_result( + k, "PyLucene backend requires exactly one index configuration" + ) + + index_config = indexes[0] + try: + plan = self._resolve_search_plan( + index_config, + k=k, + batch_size=batch_size, + mode=mode, + search_threads=search_threads, + ) + except Exception as exc: + return self._failed_search_result( + k, + str(exc), + index_config.search_params or [{}], + ) + + if dry_run: + return self._dry_run_search_result(plan) + + if not plan.index_path.is_dir(): + return self._failed_search_result( + k, + f"PyLucene index directory does not exist: {plan.index_path}", + plan.search_params, + ) + + try: + return self._execute_search(dataset, plan) + except Exception as exc: + return self._failed_search_result( + k, + _exception_summary(exc), + plan.search_params, + ) + + +__all__ = ["PyLuceneBackend", "PyLuceneConfigLoader"] diff --git a/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_cagra.yaml b/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_cagra.yaml new file mode 100644 index 0000000000..2e272895b7 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_cagra.yaml @@ -0,0 +1,10 @@ +name: pylucene_cuvs_cagra +groups: + base: + build: + codec: ["CuVS2510GPUSearchCodec"] + search: {} + test: + build: + codec: ["CuVS2510GPUSearchCodec"] + search: {} diff --git a/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_hnsw.yaml b/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_hnsw.yaml new file mode 100644 index 0000000000..a82700aeeb --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_hnsw.yaml @@ -0,0 +1,13 @@ +name: pylucene_cuvs_hnsw +groups: + base: + build: + codec: + - "Lucene101AcceleratedHNSWCodec" + - "Lucene101AcceleratedHNSWBaseLayerCodec" + - "Lucene101AcceleratedHNSWMultiLayerCodec" + search: {} + test: + build: + codec: ["Lucene101AcceleratedHNSWCodec"] + search: {} diff --git a/python/cuvs_bench/cuvs_bench/orchestrator/__init__.py b/python/cuvs_bench/cuvs_bench/orchestrator/__init__.py index 7600101439..36bc5539a7 100644 --- a/python/cuvs_bench/cuvs_bench/orchestrator/__init__.py +++ b/python/cuvs_bench/cuvs_bench/orchestrator/__init__.py @@ -1,10 +1,15 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # from .orchestrator import BenchmarkOrchestrator -from .config_loaders import ConfigLoader, BenchmarkConfig, DatasetConfig, CppGBenchConfigLoader +from .config_loaders import ( + ConfigLoader, + BenchmarkConfig, + DatasetConfig, + CppGBenchConfigLoader, +) from ..backends.registry import ( get_backend_class, list_backends, @@ -12,6 +17,7 @@ get_config_loader, ) from ..backends.opensearch import OpenSearchConfigLoader +from ..backends.pylucene import PyLuceneConfigLoader __all__ = [ # Main orchestrator @@ -22,6 +28,7 @@ "DatasetConfig", "CppGBenchConfigLoader", "OpenSearchConfigLoader", + "PyLuceneConfigLoader", # Registry functions "get_backend_class", "list_backends", @@ -34,10 +41,12 @@ # Register built-in config loaders # ============================================================================ + def _register_builtin_loaders(): """Register built-in config loaders.""" register_config_loader("cpp_gbench", CppGBenchConfigLoader) register_config_loader("opensearch", OpenSearchConfigLoader) + register_config_loader("pylucene", PyLuceneConfigLoader) # Auto-register when module is imported diff --git a/python/cuvs_bench/cuvs_bench/orchestrator/config_loaders.py b/python/cuvs_bench/cuvs_bench/orchestrator/config_loaders.py index e3af2fe58e..37be49f269 100644 --- a/python/cuvs_bench/cuvs_bench/orchestrator/config_loaders.py +++ b/python/cuvs_bench/cuvs_bench/orchestrator/config_loaders.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -91,6 +91,11 @@ def search_params_list(self) -> List[Dict[str, Any]]: def index_path(self) -> Path: return Path(self.indexes[0].file) if self.indexes else Path("") + @property + def output_filename(self) -> Tuple[str, str]: + """Return the build and search result filename stems.""" + return self.backend_config.get("output_filename", ("", "")) + @dataclass class DatasetConfig: @@ -362,7 +367,7 @@ def gather_algorithm_configs( list A list of paths to the algorithm configuration files. """ - algos_conf_fs = os.listdir(os.path.join(config_path, "algos")) + algos_conf_fs = sorted(os.listdir(os.path.join(config_path, "algos"))) algos_conf_fs = [ os.path.join(config_path, "algos", f) for f in algos_conf_fs @@ -373,7 +378,7 @@ def gather_algorithm_configs( if os.path.isdir(algorithm_configuration): algos_conf_fs += [ os.path.join(algorithm_configuration, f) - for f in os.listdir(algorithm_configuration) + for f in sorted(os.listdir(algorithm_configuration)) if f.endswith((".yaml", ".yml")) ] elif os.path.isfile(algorithm_configuration): diff --git a/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py b/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py index 580291c2f0..dce0cd1c6f 100644 --- a/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py +++ b/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -10,7 +10,7 @@ benchmark runs across different backends using the registry pattern. """ -from typing import List, Optional, Union +from typing import List, Optional, Set, Tuple, Union import numpy as np @@ -22,6 +22,11 @@ ) from ..backends._utils import compute_recall from .config_loaders import DatasetConfig +from .result_files import ( + ResultRecord, + validate_sweep_output_filenames, + write_result_files, +) class BenchmarkOrchestrator: @@ -217,6 +222,8 @@ def _run_sweep( # Collect results results: List[Union[BuildResult, SearchResult]] = [] + result_records: List[ResultRecord] = [] + stale_search_filenames: Set[str] = set() # Run benchmarks # Each config contains ALL indexes for one executable (matches runners.py) @@ -224,6 +231,17 @@ def _run_sweep( # Create backend instance backend = self.backend_class(config.backend_config) try: + output_filenames: Optional[Tuple[str, str]] = None + if backend.orchestrator_persists_results: + if len(config.indexes) != 1: + raise ValueError( + "Opt-in sweep result persistence requires exactly " + "one index per benchmark configuration" + ) + output_filenames = validate_sweep_output_filenames( + config.output_filename + ) + backend.initialize() if build: @@ -235,8 +253,18 @@ def _run_sweep( dry_run=dry_run, ) results.append(build_result) + if output_filenames is not None: + result_records.append( + ResultRecord( + result=build_result, + index_name=config.index_name, + output_filename=output_filenames[0], + ) + ) if not build_result.success: + if search and output_filenames is not None: + stale_search_filenames.add(output_filenames[1]) print( f"Build failed for {config.index_name}: {build_result.error_message}" ) @@ -274,6 +302,14 @@ def _run_sweep( ) results.append(search_result) + if output_filenames is not None: + result_records.append( + ResultRecord( + result=search_result, + index_name=config.index_name, + output_filename=output_filenames[1], + ) + ) if not search_result.success: print( @@ -282,6 +318,14 @@ def _run_sweep( finally: backend.cleanup() + if (result_records or stale_search_filenames) and not dry_run: + write_result_files( + result_records, + dataset=dataset_config.name, + dataset_path=loader_kwargs["dataset_path"], + stale_search_filenames=stale_search_filenames, + ) + return results def _run_tune( @@ -503,10 +547,17 @@ def objective(trial) -> float: print(f"Best params: {study.best_params}") print(f"Best {optimize_metric}: {study.best_value:.4f}") except ValueError: - print( - f"\n⚠️ All {n_trials} trials were pruned (constraints not met)" - ) - print(f" Consider relaxing constraints: {hard_constraints}") + if any(result.success for result in all_results): + print( + f"\n⚠️ All {n_trials} trials were pruned " + "(constraints not met)" + ) + print(f" Consider relaxing constraints: {hard_constraints}") + else: + print( + f"\n⚠️ All {n_trials} benchmark trials failed before " + "producing valid metrics" + ) return all_results diff --git a/python/cuvs_bench/cuvs_bench/orchestrator/result_files.py b/python/cuvs_bench/cuvs_bench/orchestrator/result_files.py new file mode 100644 index 0000000000..8985aa0051 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/orchestrator/result_files.py @@ -0,0 +1,232 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Persistence for benchmark backends that return in-process results.""" + +from __future__ import annotations + +import json +import re +import stat +import tempfile +from collections import defaultdict +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict, Iterable, List, Tuple, Union + +from ..backends.base import BuildResult, SearchResult + +BenchmarkResult = Union[BuildResult, SearchResult] +_SEARCH_SUFFIX_PATTERN = re.compile(r"k[1-9]\d*,bs[1-9]\d*") + + +@dataclass(frozen=True) +class ResultRecord: + """Associate a backend result with its benchmark artifact identity.""" + + result: BenchmarkResult + index_name: str + output_filename: str + + @property + def phase(self) -> str: + return "build" if isinstance(self.result, BuildResult) else "search" + + def to_json(self) -> Dict[str, Any]: + row = self.result.to_json() + row["name"] = f"{self.index_name}/{self.phase}" + return row + + +def _validate_path_component(value: object, description: str) -> str: + if not isinstance(value, str): + raise ValueError(f"Invalid {description}: {value!r}") + path = Path(value) + if not value or path.name != value or value in {".", ".."}: + raise ValueError(f"Invalid {description}: {value!r}") + return value + + +def validate_dataset_name(dataset: object) -> str: + """Validate a dataset name before using it as a path component.""" + return _validate_path_component(dataset, "benchmark dataset name") + + +def _validate_output_filename(output_filename: str) -> None: + _validate_path_component(output_filename, "benchmark result filename") + + +def validate_sweep_output_filenames( + output_filenames: object, +) -> Tuple[str, str]: + """Validate and normalize opt-in sweep result filename stems.""" + if ( + not isinstance(output_filenames, (list, tuple)) + or len(output_filenames) != 2 + ): + raise ValueError( + "Opt-in sweep persistence requires exactly two result filename " + "stems: (build_stem, search_stem)" + ) + + build_stem, search_stem = output_filenames + _validate_output_filename(build_stem) + _validate_output_filename(search_stem) + + build_parts = build_stem.split(",") + if len(build_parts) < 2 or any(not part for part in build_parts): + raise ValueError( + "Build result filename stem must contain at least " + "','" + ) + + expected_prefix = f"{build_stem}," + search_suffix = ( + search_stem[len(expected_prefix) :] + if search_stem.startswith(expected_prefix) + else "" + ) + if _SEARCH_SUFFIX_PATTERN.fullmatch(search_suffix) is None: + raise ValueError( + "Search result filename stem must be " + "',k,bs'" + ) + + return build_stem, search_stem + + +def _read_benchmarks(path: Path) -> List[Dict[str, Any]]: + if not path.exists(): + return [] + with path.open(encoding="utf-8") as file: + payload = json.load(file) + benchmarks = ( + payload.get("benchmarks") if isinstance(payload, dict) else None + ) + if not isinstance(benchmarks, list): + raise ValueError( + f"Benchmark result file must contain a 'benchmarks' list: {path}" + ) + seen_names = set() + for position, row in enumerate(benchmarks): + if not isinstance(row, dict) or not isinstance(row.get("name"), str): + raise ValueError( + f"Benchmark row {position} must be an object with a string " + f"'name': {path}" + ) + if row["name"] in seen_names: + raise ValueError( + f"Benchmark result file contains duplicate name " + f"{row['name']!r}: {path}" + ) + seen_names.add(row["name"]) + return benchmarks + + +def _merge_build_rows( + path: Path, records: List[ResultRecord] +) -> List[Dict[str, Any]]: + rows_by_name = {row["name"]: row for row in _read_benchmarks(path)} + for record in records: + row = record.to_json() + previous = rows_by_name.get(row["name"]) + if ( + record.result.success + and record.result.metadata.get("skipped") + and previous is not None + and previous.get("success", True) is True + ): + continue + rows_by_name[row["name"]] = row + return list(rows_by_name.values()) + + +def _write_payload(path: Path, rows: List[Dict[str, Any]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + file_mode = stat.S_IMODE(path.stat().st_mode) if path.exists() else 0o644 + temp_path = None + try: + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + delete=False, + ) as file: + temp_path = Path(file.name) + json.dump({"benchmarks": rows}, file, indent=2, allow_nan=False) + file.write("\n") + temp_path.chmod(file_mode) + temp_path.replace(path) + finally: + if temp_path is not None: + temp_path.unlink(missing_ok=True) + + +def _invalidate_derived_csvs(path: Path, phase: str) -> None: + suffixes = ( + (".csv",) + if phase == "build" + else (",raw.csv", ",throughput.csv", ",latency.csv") + ) + stem = path.with_suffix("") + for suffix in suffixes: + stem.with_name(f"{stem.name}{suffix}").unlink(missing_ok=True) + + +def _remove_result_artifacts(path: Path, phase: str) -> None: + path.unlink(missing_ok=True) + _invalidate_derived_csvs(path, phase) + + +def write_result_files( + records: List[ResultRecord], + dataset: str, + dataset_path: str, + *, + stale_search_filenames: Iterable[str] = (), +) -> None: + """Write grouped Google Benchmark-compatible JSON result files.""" + validate_dataset_name(dataset) + grouped: Dict[Tuple[str, str], List[ResultRecord]] = defaultdict(list) + for record in records: + _validate_output_filename(record.output_filename) + grouped[(record.phase, record.output_filename)].append(record) + + stale_search_filenames = set(stale_search_filenames) + for output_filename in stale_search_filenames: + _validate_output_filename(output_filename) + + result_root = Path(dataset_path) / dataset / "result" + current_search_filenames = { + output_filename + for phase, output_filename in grouped + if phase == "search" + } + for output_filename in stale_search_filenames - current_search_filenames: + _remove_result_artifacts( + result_root / "search" / f"{output_filename}.json", + "search", + ) + + for (phase, output_filename), file_records in grouped.items(): + output_path = result_root / phase / f"{output_filename}.json" + rows = ( + _merge_build_rows(output_path, file_records) + if phase == "build" + else [record.to_json() for record in file_records] + ) + if rows: + _write_payload(output_path, rows) + _invalidate_derived_csvs(output_path, phase) + + +__all__ = [ + "ResultRecord", + "validate_dataset_name", + "validate_sweep_output_filenames", + "write_result_files", +] diff --git a/python/cuvs_bench/cuvs_bench/run/__main__.py b/python/cuvs_bench/cuvs_bench/run/__main__.py index 6950ff7202..eb3726ffc9 100644 --- a/python/cuvs_bench/cuvs_bench/run/__main__.py +++ b/python/cuvs_bench/cuvs_bench/run/__main__.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -11,8 +11,12 @@ import click import yaml -from .data_export import convert_json_to_csv_build, convert_json_to_csv_search +from .data_export import ( + convert_json_to_csv_build, + convert_json_to_csv_search, +) from ..orchestrator import BenchmarkOrchestrator +from ..orchestrator.result_files import validate_dataset_name @click.command() @@ -257,6 +261,11 @@ def main( and any backend-specific connection parameters (host, port, etc.). """ + try: + validate_dataset_name(dataset) + except ValueError as exc: + raise click.BadParameter(str(exc), param_hint="--dataset") from exc + if not data_export: # Determine backend type and extra kwargs from --backend-config backend_type = "cpp_gbench" @@ -277,7 +286,7 @@ def main( backend_kwargs = cfg orchestrator = BenchmarkOrchestrator(backend_type=backend_type) - orchestrator.run_benchmark( + results = orchestrator.run_benchmark( mode=mode, constraints=json.loads(constraints) if constraints else None, n_trials=n_trials, @@ -300,6 +309,24 @@ def main( executable_dir=executable_dir, **backend_kwargs, ) + failures = [result for result in results if not result.success] + if mode == "sweep": + run_failed = not results or bool(failures) + else: + run_failed = not any(result.success for result in results) + if run_failed: + details = ( + "; ".join( + f"{result.algorithm}: " + f"{result.error_message or 'benchmark failed'}" + for result in failures + ) + if failures + else f"{mode} mode produced no benchmark results" + ) + raise click.ClickException(details) + if dry_run: + return convert_json_to_csv_build(dataset, dataset_path) convert_json_to_csv_search(dataset, dataset_path) diff --git a/python/cuvs_bench/cuvs_bench/run/data_export.py b/python/cuvs_bench/cuvs_bench/run/data_export.py index 707677a083..40078b2a9c 100644 --- a/python/cuvs_bench/cuvs_bench/run/data_export.py +++ b/python/cuvs_bench/cuvs_bench/run/data_export.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -9,6 +9,8 @@ import pandas as pd +from ..orchestrator.result_files import validate_dataset_name + skip_build_cols = set( [ "algo_name", @@ -25,11 +27,24 @@ "real_time", "time_unit", "index_size", + "success", + "error_message", + "skipped", ] ) skip_search_cols = ( - set(["recall", "qps", "latency", "items_per_second", "Recall", "Latency"]) + set( + [ + "recall", + "qps", + "latency", + "items_per_second", + "Recall", + "Latency", + "search_params", + ] + ) | skip_build_cols ) @@ -69,7 +84,11 @@ def read_json_files(dataset, dataset_path, method): A tuple containing the file path, algorithm name, and the DataFrame of JSON content. """ + validate_dataset_name(dataset) dir_path = os.path.join(dataset_path, dataset, "result", method) + if not os.path.isdir(dir_path): + return + for file in os.listdir(dir_path): if file.endswith(".json"): file_path = os.path.join(dir_path, file) @@ -77,6 +96,10 @@ def read_json_files(dataset, dataset_path, method): with open(file_path, "r", encoding="ISO-8859-1") as f: data = json.load(f) df = pd.DataFrame(data["benchmarks"]) + if "success" in df.columns: + df = df[df["success"].ne(False)].copy() + if method == "build" and "skipped" in df.columns: + df = df[df["skipped"].ne(True)].copy() algo_name = tuple(file.split(",")[:2]) yield file_path, algo_name, df except Exception as e: @@ -103,6 +126,15 @@ def clean_algo_name(algo_name): return name.removesuffix(".json") +def _remove_derived_csvs(file, suffixes): + stem = os.path.splitext(file)[0] + for suffix in suffixes: + try: + os.remove(f"{stem}{suffix}") + except FileNotFoundError: + pass + + def write_csv(file, algo_name, df, extra_columns=None, skip_cols=None): """ Write a DataFrame to CSV with specified columns skipped. @@ -153,6 +185,9 @@ def convert_json_to_csv_build(dataset, dataset_path): """ for file, algo_name, df in read_json_files(dataset, dataset_path, "build"): try: + if df.empty: + _remove_derived_csvs(file, (".csv",)) + continue algo_name = clean_algo_name(algo_name) write_csv(file, algo_name, df, skip_cols=skip_build_cols) except Exception as e: @@ -175,12 +210,26 @@ def convert_json_to_csv_search(dataset, dataset_path): dataset, dataset_path, "search" ): try: + if df.empty: + _remove_derived_csvs( + file, (",raw.csv", ",throughput.csv", ",latency.csv") + ) + continue + search_stem = os.path.splitext(os.path.basename(file))[0] + # Search stems end in k and batch-size components. Preserve every + # preceding component, including an optional dataset subset. + search_stem_parts = search_stem.rsplit(",", maxsplit=2) + build_stem = ( + search_stem_parts[0] + if len(search_stem_parts) == 3 + else ",".join(algo_name) + ) build_file = os.path.join( dataset_path, dataset, "result", "build", - f"{','.join(algo_name)}.csv", + f"{build_stem}.csv", ) algo_name = clean_algo_name(algo_name) df["name"] = df["name"].str.split("/").str[0] @@ -202,29 +251,46 @@ def convert_json_to_csv_search(dataset, dataset_path): ] if os.path.exists(build_file): build_df = pd.read_csv(build_file) - write_ncols = len(write.columns) write["build time"] = None write["build threads"] = None write["build cpu_time"] = None - start_idx = 5 + build_columns = { + "time": "build time", + "threads": "build threads", + "cpu_time": "build cpu_time", + } if "GPU" in build_df.columns: - start_idx = 6 write["build GPU"] = None - for col_idx in range(start_idx, len(build_df.columns)): - col_name = build_df.columns[col_idx] - write[col_name] = None - if col_name == "num_threads": - write["build_num_threads"] = None + build_columns["GPU"] = "build GPU" + + ignored_columns = {"algo_name", "index_name"} + for col_name in build_df.columns: + if ( + col_name in ignored_columns + or col_name in build_columns + ): + continue + target_name = ( + "build_num_threads" + if col_name == "num_threads" + else col_name + ) + if target_name not in write.columns: + write[target_name] = None + build_columns[col_name] = target_name + for s_index, search_row in write.iterrows(): - for b_index, build_row in build_df.iterrows(): + for _, build_row in build_df.iterrows(): if search_row["index_name"] == build_row["index_name"]: - write.iloc[s_index, write_ncols] = build_df.iloc[ - b_index, 2 - ] - write.iloc[s_index, write_ncols + 1 :] = ( - build_df.iloc[b_index, 3:] - ) + for ( + source_name, + target_name, + ) in build_columns.items(): + if source_name in build_df.columns: + write.at[s_index, target_name] = build_row[ + source_name + ] break # Write search data and compute frontiers write.to_csv(file.replace(".json", ",raw.csv"), index=False) diff --git a/python/cuvs_bench/cuvs_bench/tests/_pylucene_test_utils.py b/python/cuvs_bench/cuvs_bench/tests/_pylucene_test_utils.py new file mode 100644 index 0000000000..16b40affc3 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/_pylucene_test_utils.py @@ -0,0 +1,224 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Shared fakes and data builders for PyLucene unit tests.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np + +import cuvs_bench.backends.pylucene as pylucene_backend +from cuvs_bench._bin_format import write_bin_header +from cuvs_bench.backends.base import Dataset +from cuvs_bench.backends.pylucene import ( + PyLuceneBackend, + _CagraIndexVerification, + _RuntimeSearchResult, + _SearchHit, +) +from cuvs_bench.orchestrator.config_loaders import IndexConfig + +_HNSW_CODEC = "Lucene101AcceleratedHNSWCodec" +_HNSW_BASE_LAYER_CODEC = "Lucene101AcceleratedHNSWBaseLayerCodec" +_CAGRA_CODEC = "CuVS2510GPUSearchCodec" + + +class _FakeRuntime: + pylucene_version = "test" + + def __init__( + self, + *, + writer_path: str = "gpu-hnsw", + index_dimensions: int = 4, + document_count: int = 10, + hits: list[list[_SearchHit]] | None = None, + batch_latencies_ms: list[float] | None = None, + validate_query_dimensions: bool = True, + ): + self.writer_path = writer_path + self.index_dimensions = index_dimensions + self.document_count = document_count + self.hits = hits + self.batch_latencies_ms = batch_latencies_ms + self.validate_query_dimensions = validate_query_dimensions + self.resolve_calls = [] + self.telemetry_calls = [] + self.build_calls = [] + self.search_calls = [] + self.verification_calls = [] + self.resolve_error = None + self.build_error = None + self.search_error = None + self.verification_error = None + + def resolve_codec(self, codec_name): + self.resolve_calls.append(codec_name) + if self.resolve_error is not None: + raise self.resolve_error + return codec_name + + def codec_telemetry(self, codec_name, java_codec): + self.telemetry_calls.append(codec_name) + assert java_codec == codec_name + return {"writerPath": self.writer_path, "codec": codec_name} + + def build_index(self, index_path, vectors, resolved_codec): + self.build_calls.append((index_path, vectors.copy(), resolved_codec)) + if self.build_error is not None: + (index_path / "partial").write_bytes(b"partial") + raise self.build_error + self.document_count = int(vectors.shape[0]) + (index_path / "segments_1").write_bytes(b"index") + return dict(resolved_codec.telemetry) + + def verify_cagra_index( + self, + index_path, + expected_vector_count=None, + expected_dimensions=None, + ): + self.verification_calls.append( + (index_path, expected_vector_count, expected_dimensions) + ) + if self.verification_error is not None: + raise self.verification_error + return _CagraIndexVerification( + segment_count=1, + field_count=1, + vector_count=( + expected_vector_count + if expected_vector_count is not None + else 10 + ), + dimensions=( + expected_dimensions if expected_dimensions is not None else 4 + ), + ) + + def search_index(self, index_path, query_vectors, k, batch_size): + self.search_calls.append( + (index_path, query_vectors.copy(), k, batch_size) + ) + if self.search_error is not None: + raise self.search_error + if ( + self.validate_query_dimensions + and query_vectors.shape[1] != self.index_dimensions + ): + raise ValueError( + "Query vector dimensions do not match the Lucene index" + ) + hits = self.hits + if hits is None: + hits = [ + [_SearchHit(document_id=0, score=1.0)] + for _ in range(query_vectors.shape[0]) + ] + batch_latencies_ms = self.batch_latencies_ms + if batch_latencies_ms is None: + batch_latencies_ms = [ + 1.0 for _ in range(0, query_vectors.shape[0], batch_size) + ] + return _RuntimeSearchResult( + hits=hits, + batch_latencies_ms=batch_latencies_ms, + index_dimensions=self.index_dimensions, + document_count=self.document_count, + ) + + +def _dataset( + *, + n_base: int = 10, + n_queries: int = 2, + dimensions: int = 4, + dtype=np.float32, + distance_metric: str = "euclidean", +) -> Dataset: + rng = np.random.default_rng(7) + base = rng.random((n_base, dimensions)).astype(dtype) + queries = rng.random((n_queries, dimensions)).astype(dtype) + return Dataset( + name="test", + training_vectors=base, + query_vectors=queries, + distance_metric=distance_metric, + ) + + +def _index( + index_path: Path, + *, + codec: str = _HNSW_CODEC, + search_params: list[dict] | None = None, +) -> IndexConfig: + return IndexConfig( + name="pylucene-test", + algo="pylucene_cuvs_hnsw", + build_param={"codec": codec}, + search_params=[{}] if search_params is None else search_params, + file=str(index_path), + ) + + +def _backend( + runtime: _FakeRuntime | None = None, + *, + codec: str = _HNSW_CODEC, +) -> PyLuceneBackend: + backend = PyLuceneBackend( + { + "name": "pylucene-test", + "algo": "pylucene_cuvs_hnsw", + "codec": codec, + } + ) + if runtime is not None: + backend._runtime = runtime + return backend + + +def _prepare_hnsw_index( + index_path: Path, + *, + codec: str = _HNSW_CODEC, + vector_count: int = 10, + dimensions: int = 4, +) -> Path: + index_path.mkdir(parents=True, exist_ok=True) + (index_path / "segments_1").write_bytes(b"index") + pylucene_backend._write_hnsw_provenance( + index_path, + codec, + vector_count=vector_count, + dimensions=dimensions, + ) + return index_path / pylucene_backend._HNSW_PROVENANCE_FILE + + +def _prepare_cagra_index( + index_path: Path, + *, + vector_count: int = 10, + dimensions: int = 4, +) -> Path: + index_path.mkdir(parents=True, exist_ok=True) + (index_path / "segments_1").write_bytes(b"index") + pylucene_backend._write_cagra_provenance( + index_path, + vector_count=vector_count, + dimensions=dimensions, + ) + return index_path / pylucene_backend._CAGRA_PROVENANCE_FILE + + +def _write_test_bin(path: Path, data: np.ndarray) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("wb") as file: + write_bin_header(file, data.shape[0], data.shape[1]) + np.ascontiguousarray(data).tofile(file) diff --git a/python/cuvs_bench/cuvs_bench/tests/test_cli.py b/python/cuvs_bench/cuvs_bench/tests/test_cli.py index fa2a63d041..a1525d234d 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_cli.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_cli.py @@ -1,9 +1,10 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # from pathlib import Path +from types import SimpleNamespace import pandas as pd import pytest @@ -519,11 +520,182 @@ def test_plot_command_creates_png_files(temp_datasets_dir: Path): ) -# FIXME: Tests below use --dry-run to verify CLI flag parsing and orchestrator -# routing without requiring actual benchmark execution. Tune mode (--mode tune) -# requires Optuna and actual search results, so only flag acceptance is tested -# here. End-to-end tests for tune mode and non-C++ backends should be added -# when those features are exercised in integration testing. +# The mocked-result tests below isolate CLI outcome and export semantics. +# Later tests use --dry-run to verify flag parsing and orchestrator routing +# without requiring native benchmark execution. + + +def _invoke_run_with_results(monkeypatch, tmp_path, mode, results): + from cuvs_bench.run import __main__ as run_module + + orchestrator = SimpleNamespace( + run_benchmark=lambda **_kwargs: results, + ) + monkeypatch.setattr( + run_module, + "BenchmarkOrchestrator", + lambda backend_type: orchestrator, + ) + exported = [] + monkeypatch.setattr( + run_module, + "convert_json_to_csv_build", + lambda dataset, dataset_path: exported.append( + ("build", dataset, dataset_path) + ), + ) + monkeypatch.setattr( + run_module, + "convert_json_to_csv_search", + lambda dataset, dataset_path: exported.append( + ("search", dataset, dataset_path) + ), + ) + + result = CliRunner().invoke( + run_module.main, + [ + "--dataset", + "test-data", + "--dataset-path", + str(tmp_path), + "--algorithms", + "fake", + "--groups", + "base", + "--batch-size", + "1", + "-k", + "1", + "-m", + "latency", + "--mode", + mode, + ], + ) + return result, exported + + +def _benchmark_result(success, error_message=None): + return SimpleNamespace( + success=success, + algorithm="fake", + error_message=error_message, + ) + + +def test_tune_mixed_trial_results_exit_zero_and_export(monkeypatch, tmp_path): + result, exported = _invoke_run_with_results( + monkeypatch, + tmp_path, + "tune", + [ + _benchmark_result(False, "trial failed"), + _benchmark_result(True), + ], + ) + + assert result.exit_code == 0, result.output + assert exported == [ + ("build", "test-data", str(tmp_path)), + ("search", "test-data", str(tmp_path)), + ] + + +@pytest.mark.parametrize( + ("results", "expected_message"), + [ + ([_benchmark_result(False, "trial failed")], "trial failed"), + ([], "tune mode produced no benchmark results"), + ], +) +def test_tune_without_successful_results_exits_nonzero( + monkeypatch, tmp_path, results, expected_message +): + result, exported = _invoke_run_with_results( + monkeypatch, tmp_path, "tune", results + ) + + assert result.exit_code != 0 + assert expected_message in result.output + assert exported == [] + + +def test_tune_all_constraint_pruned_measurements_exit_zero_and_export( + monkeypatch, tmp_path +): + # _run_tune retains successful measurements even when Optuna prunes every + # trial for violating a hard constraint. + result, exported = _invoke_run_with_results( + monkeypatch, + tmp_path, + "tune", + [_benchmark_result(True), _benchmark_result(True)], + ) + + assert result.exit_code == 0, result.output + assert exported == [ + ("build", "test-data", str(tmp_path)), + ("search", "test-data", str(tmp_path)), + ] + + +def test_sweep_mixed_results_exit_nonzero(monkeypatch, tmp_path): + result, exported = _invoke_run_with_results( + monkeypatch, + tmp_path, + "sweep", + [ + _benchmark_result(False, "benchmark failed"), + _benchmark_result(True), + ], + ) + + assert result.exit_code != 0 + assert "benchmark failed" in result.output + assert exported == [] + + +def test_sweep_without_results_exits_nonzero(monkeypatch, tmp_path): + result, exported = _invoke_run_with_results( + monkeypatch, tmp_path, "sweep", [] + ) + + assert result.exit_code != 0 + assert "sweep mode produced no benchmark results" in result.output + assert exported == [] + + +@pytest.mark.parametrize( + "dataset", ["..", "../escape", "nested/dataset", "/absolute/dataset"] +) +def test_data_export_rejects_invalid_dataset_name(tmp_path, dataset): + from cuvs_bench.run.__main__ import main as run_main + + runner = CliRunner() + result = runner.invoke( + run_main, + [ + "--data-export", + "--dataset", + dataset, + "--dataset-path", + str(tmp_path), + "--count", + "10", + "--batch-size", + "100", + "--algorithms", + "cuvs_cagra", + "--groups", + "base", + "--search-mode", + "latency", + ], + ) + + assert result.exit_code == 2 + assert "Invalid benchmark dataset name" in result.output def test_run_with_mode_sweep(temp_datasets_dir): diff --git a/python/cuvs_bench/cuvs_bench/tests/test_data_export.py b/python/cuvs_bench/cuvs_bench/tests/test_data_export.py new file mode 100644 index 0000000000..3898c7f624 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_data_export.py @@ -0,0 +1,683 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Regression tests for in-process benchmark result persistence and export.""" + +import json +import stat + +import numpy as np +import pandas as pd +import pytest + +from cuvs_bench.backends.base import BuildResult, SearchResult +from cuvs_bench.orchestrator.result_files import ( + ResultRecord, + write_result_files, +) +from cuvs_bench.run.data_export import ( + convert_json_to_csv_build, + convert_json_to_csv_search, + read_json_files, +) + + +_DATASET = "test-data" +_ALGORITHM = "pylucene_cuvs_hnsw" +_BUILD_STEM = f"{_ALGORITHM},base" +_SEARCH_STEM = f"{_BUILD_STEM},k5,bs2" +_CODECS = ( + "Lucene101AcceleratedHNSWCodec", + "Lucene101AcceleratedHNSWBaseLayerCodec", + "Lucene101AcceleratedHNSWMultiLayerCodec", +) + + +def _index_name(codec): + return f"{_ALGORITHM}.codec{codec}" + + +def _build_result( + codec, + *, + build_time=1.25, + metadata=None, + success=True, + error_message=None, +): + return BuildResult( + index_path=f"/indexes/{_index_name(codec)}", + build_time_seconds=build_time, + index_size_bytes=4096, + algorithm=_ALGORITHM, + build_params={"codec": codec}, + metadata=metadata or {}, + success=success, + error_message=error_message, + ) + + +def _search_result( + codec, + *, + search_time_ms=6.0, + latency_seconds=0.003, + recall=0.9, + queries_per_second=500.0, + metadata=None, + success=True, + error_message=None, +): + result_metadata = {"codec": codec, "num_batches": 2} + if metadata: + result_metadata.update(metadata) + return SearchResult( + neighbors=np.empty((0, 5), dtype=np.int64), + distances=np.empty((0, 5), dtype=np.float32), + search_time_ms=search_time_ms, + queries_per_second=queries_per_second, + recall=recall, + algorithm=_ALGORITHM, + search_params=[{}], + latency_seconds=latency_seconds, + latency_percentiles={"p50": 2.5, "p95": 3.5, "p99": 3.9}, + metadata=result_metadata, + success=success, + error_message=error_message, + ) + + +def _record(result, codec, output_filename): + return ResultRecord( + result=result, + index_name=_index_name(codec), + output_filename=output_filename, + ) + + +def _load_payload(path): + with path.open(encoding="utf-8") as file: + return json.load(file) + + +def test_write_result_files_groups_three_codecs_with_expected_schema(tmp_path): + records = [] + for position, codec in enumerate(_CODECS, start=1): + records.extend( + [ + _record( + _build_result(codec, build_time=float(position)), + codec, + _BUILD_STEM, + ), + _record( + _search_result( + codec, + search_time_ms=float(position * 4), + latency_seconds=float(position) / 500.0, + recall=0.8 + position / 100.0, + queries_per_second=1000.0 / position, + ), + codec, + _SEARCH_STEM, + ), + ] + ) + + write_result_files(records, dataset=_DATASET, dataset_path=str(tmp_path)) + + result_root = tmp_path / _DATASET / "result" + build_path = result_root / "build" / f"{_BUILD_STEM}.json" + search_path = result_root / "search" / f"{_SEARCH_STEM}.json" + assert sorted( + str(path.relative_to(result_root)) + for path in result_root.rglob("*.json") + ) == [ + f"build/{_BUILD_STEM}.json", + f"search/{_SEARCH_STEM}.json", + ] + + build_rows = _load_payload(build_path)["benchmarks"] + search_rows = _load_payload(search_path)["benchmarks"] + assert len(build_rows) == len(_CODECS) + assert len(search_rows) == len(_CODECS) + assert {row["name"] for row in build_rows} == { + f"{_index_name(codec)}/build" for codec in _CODECS + } + assert {row["name"] for row in search_rows} == { + f"{_index_name(codec)}/search" for codec in _CODECS + } + + first_build = build_rows[0] + assert first_build == { + "codec": _CODECS[0], + "name": f"{_index_name(_CODECS[0])}/build", + "real_time": 1.0, + "time_unit": "s", + "index_size": 4096, + "success": True, + } + + first_search = search_rows[0] + assert first_search["real_time"] == pytest.approx(2.0) + assert first_search["time_unit"] == "ms" + assert first_search["search_time_ms"] == pytest.approx(4.0) + assert first_search["Latency"] == pytest.approx(0.002) + assert first_search["items_per_second"] == pytest.approx(1000.0) + assert first_search["Recall"] == pytest.approx(0.81) + assert first_search["p50"] == pytest.approx(2.5) + assert first_search["p95"] == pytest.approx(3.5) + assert first_search["p99"] == pytest.approx(3.9) + assert first_search["codec"] == _CODECS[0] + + +def test_skipped_build_preserves_existing_positive_result(tmp_path): + codec = _CODECS[0] + output_path = ( + tmp_path / _DATASET / "result" / "build" / f"{_BUILD_STEM}.json" + ) + write_result_files( + [_record(_build_result(codec, build_time=4.5), codec, _BUILD_STEM)], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + write_result_files( + [ + _record( + _build_result( + codec, + build_time=0.0, + metadata={"skipped": True}, + ), + codec, + _BUILD_STEM, + ) + ], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + rows = _load_payload(output_path)["benchmarks"] + assert len(rows) == 1 + assert rows[0]["real_time"] == pytest.approx(4.5) + assert "skipped" not in rows[0] + + +def test_skipped_build_replaces_existing_failed_result(tmp_path): + codec = _CODECS[0] + output_path = ( + tmp_path / _DATASET / "result" / "build" / f"{_BUILD_STEM}.json" + ) + write_result_files( + [ + _record( + _build_result( + codec, + build_time=0.0, + success=False, + error_message="old failure", + ), + codec, + _BUILD_STEM, + ) + ], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + write_result_files( + [ + _record( + _build_result( + codec, + build_time=0.0, + metadata={"skipped": True}, + ), + codec, + _BUILD_STEM, + ) + ], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + rows = _load_payload(output_path)["benchmarks"] + assert len(rows) == 1 + assert rows[0]["success"] is True + assert rows[0]["skipped"] is True + assert "error_message" not in rows[0] + + csv_path = output_path.with_suffix(".csv") + csv_path.write_text("stale\n", encoding="utf-8") + convert_json_to_csv_build(_DATASET, str(tmp_path)) + assert not csv_path.exists() + + +def test_build_upsert_replaces_by_name_without_duplicates(tmp_path): + initial_records = [ + _record(_build_result(codec, build_time=1.0), codec, _BUILD_STEM) + for codec in _CODECS + ] + write_result_files( + initial_records, dataset=_DATASET, dataset_path=str(tmp_path) + ) + + updated_codec = _CODECS[1] + write_result_files( + [ + _record( + _build_result( + updated_codec, + build_time=7.5, + metadata={"writer_path": "gpu-hnsw"}, + ), + updated_codec, + _BUILD_STEM, + ) + ], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + path = tmp_path / _DATASET / "result" / "build" / f"{_BUILD_STEM}.json" + rows = _load_payload(path)["benchmarks"] + assert len(rows) == len(_CODECS) + assert len({row["name"] for row in rows}) == len(_CODECS) + updated = next( + row + for row in rows + if row["name"] == f"{_index_name(updated_codec)}/build" + ) + assert updated["real_time"] == pytest.approx(7.5) + assert updated["writer_path"] == "gpu-hnsw" + + +def test_search_write_replaces_previous_rows(tmp_path): + write_result_files( + [ + _record(_search_result(codec), codec, _SEARCH_STEM) + for codec in _CODECS + ], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + replacement_codec = _CODECS[-1] + write_result_files( + [ + _record( + _search_result( + replacement_codec, + recall=0.99, + queries_per_second=750.0, + ), + replacement_codec, + _SEARCH_STEM, + ) + ], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + path = tmp_path / _DATASET / "result" / "search" / f"{_SEARCH_STEM}.json" + rows = _load_payload(path)["benchmarks"] + assert len(rows) == 1 + assert rows[0]["name"] == f"{_index_name(replacement_codec)}/search" + assert rows[0]["Recall"] == pytest.approx(0.99) + assert rows[0]["items_per_second"] == pytest.approx(750.0) + + +@pytest.mark.parametrize( + "output_filename", + ["", ".", "..", "../escape", "nested/result", "/absolute/result", None], +) +def test_write_result_files_rejects_invalid_output_filename( + tmp_path, output_filename +): + codec = _CODECS[0] + record = _record(_build_result(codec), codec, output_filename) + + with pytest.raises(ValueError, match="Invalid benchmark result filename"): + write_result_files( + [record], dataset=_DATASET, dataset_path=str(tmp_path) + ) + + assert not (tmp_path / _DATASET / "result").exists() + + +@pytest.mark.parametrize( + "dataset", + ["", ".", "..", "../escape", "nested/dataset", "/absolute/dataset", None], +) +def test_write_result_files_rejects_invalid_dataset_name(tmp_path, dataset): + codec = _CODECS[0] + record = _record(_build_result(codec), codec, _BUILD_STEM) + + with pytest.raises(ValueError, match="Invalid benchmark dataset name"): + write_result_files( + [record], dataset=dataset, dataset_path=str(tmp_path) + ) + + +@pytest.mark.parametrize( + "dataset", + ["", ".", "..", "../escape", "nested/dataset", "/absolute/dataset", None], +) +@pytest.mark.parametrize( + "converter", [convert_json_to_csv_build, convert_json_to_csv_search] +) +def test_data_export_rejects_invalid_dataset_name( + tmp_path, dataset, converter +): + with pytest.raises(ValueError, match="Invalid benchmark dataset name"): + converter(dataset, str(tmp_path)) + + +def test_build_merge_rejects_invalid_existing_payload(tmp_path): + build_dir = tmp_path / _DATASET / "result" / "build" + build_dir.mkdir(parents=True) + output_path = build_dir / f"{_BUILD_STEM}.json" + output_path.write_text('{"benchmarks": {}}\n', encoding="utf-8") + codec = _CODECS[0] + + with pytest.raises(ValueError, match="'benchmarks' list"): + write_result_files( + [_record(_build_result(codec), codec, _BUILD_STEM)], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + assert _load_payload(output_path) == {"benchmarks": {}} + + +@pytest.mark.parametrize( + "benchmarks", + [ + [None], + [{"real_time": 1.0}], + [{"name": 42}], + [{"name": "duplicate"}, {"name": "duplicate"}], + ], +) +def test_build_merge_rejects_invalid_existing_rows(tmp_path, benchmarks): + build_dir = tmp_path / _DATASET / "result" / "build" + build_dir.mkdir(parents=True) + output_path = build_dir / f"{_BUILD_STEM}.json" + original_payload = {"benchmarks": benchmarks} + output_path.write_text(json.dumps(original_payload), encoding="utf-8") + codec = _CODECS[0] + + with pytest.raises(ValueError, match="Benchmark row|duplicate name"): + write_result_files( + [_record(_build_result(codec), codec, _BUILD_STEM)], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + assert _load_payload(output_path) == original_payload + + +def test_atomic_result_write_uses_readable_and_preserved_file_modes(tmp_path): + codec = _CODECS[0] + output_path = ( + tmp_path / _DATASET / "result" / "build" / f"{_BUILD_STEM}.json" + ) + record = _record(_build_result(codec), codec, _BUILD_STEM) + + write_result_files([record], dataset=_DATASET, dataset_path=str(tmp_path)) + assert stat.S_IMODE(output_path.stat().st_mode) == 0o644 + + output_path.chmod(0o640) + write_result_files([record], dataset=_DATASET, dataset_path=str(tmp_path)) + assert stat.S_IMODE(output_path.stat().st_mode) == 0o640 + + +def test_result_write_invalidates_only_derived_csvs(tmp_path): + result_root = tmp_path / _DATASET / "result" + build_dir = result_root / "build" + search_dir = result_root / "search" + build_dir.mkdir(parents=True) + search_dir.mkdir(parents=True) + derived_paths = [ + build_dir / f"{_BUILD_STEM}.csv", + search_dir / f"{_SEARCH_STEM},raw.csv", + search_dir / f"{_SEARCH_STEM},throughput.csv", + search_dir / f"{_SEARCH_STEM},latency.csv", + ] + unrelated_path = search_dir / "unrelated.csv" + for path in [*derived_paths, unrelated_path]: + path.write_text("stale\n", encoding="utf-8") + + codec = _CODECS[0] + write_result_files( + [ + _record(_build_result(codec), codec, _BUILD_STEM), + _record(_search_result(codec), codec, _SEARCH_STEM), + ], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + assert not any(path.exists() for path in derived_paths) + assert unrelated_path.read_text(encoding="utf-8") == "stale\n" + + +def test_missing_phase_directories_are_noop(tmp_path): + assert list(read_json_files(_DATASET, str(tmp_path), "build")) == [] + assert list(read_json_files(_DATASET, str(tmp_path), "search")) == [] + + convert_json_to_csv_build(_DATASET, str(tmp_path)) + convert_json_to_csv_search(_DATASET, str(tmp_path)) + + assert not (tmp_path / _DATASET / "result").exists() + + +def test_failed_rows_remain_diagnostic_json_but_are_filtered_from_csv( + tmp_path, +): + good_codec, failed_codec = _CODECS[:2] + write_result_files( + [ + _record(_build_result(good_codec), good_codec, _BUILD_STEM), + _record( + _build_result( + failed_codec, + build_time=0.0, + success=False, + error_message="build failed", + ), + failed_codec, + _BUILD_STEM, + ), + _record(_search_result(good_codec), good_codec, _SEARCH_STEM), + _record( + _search_result( + failed_codec, + search_time_ms=0.0, + latency_seconds=None, + recall=0.0, + queries_per_second=0.0, + success=False, + error_message="search failed", + ), + failed_codec, + _SEARCH_STEM, + ), + ], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + result_root = tmp_path / _DATASET / "result" + build_json = result_root / "build" / f"{_BUILD_STEM}.json" + search_json = result_root / "search" / f"{_SEARCH_STEM}.json" + build_rows = _load_payload(build_json)["benchmarks"] + search_rows = _load_payload(search_json)["benchmarks"] + assert ( + next(row for row in build_rows if not row["success"])["error_message"] + == "build failed" + ) + assert ( + next(row for row in search_rows if not row["success"])["error_message"] + == "search failed" + ) + + convert_json_to_csv_build(_DATASET, str(tmp_path)) + convert_json_to_csv_search(_DATASET, str(tmp_path)) + + build_csv = pd.read_csv(build_json.with_suffix(".csv")) + search_csv = pd.read_csv(search_json.with_name(f"{_SEARCH_STEM},raw.csv")) + assert build_csv["index_name"].tolist() == [_index_name(good_codec)] + assert search_csv["index_name"].tolist() == [_index_name(good_codec)] + assert "success" not in build_csv.columns + assert "error_message" not in build_csv.columns + assert "success" not in search_csv.columns + assert "error_message" not in search_csv.columns + + +def test_all_failed_json_removes_only_matching_derived_csvs(tmp_path): + codec = _CODECS[0] + write_result_files( + [ + _record( + _build_result( + codec, + build_time=0.0, + success=False, + error_message="build failed", + ), + codec, + _BUILD_STEM, + ), + _record( + _search_result( + codec, + search_time_ms=0.0, + latency_seconds=None, + recall=0.0, + queries_per_second=0.0, + success=False, + error_message="search failed", + ), + codec, + _SEARCH_STEM, + ), + ], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + result_root = tmp_path / _DATASET / "result" + derived_paths = [ + result_root / "build" / f"{_BUILD_STEM}.csv", + result_root / "search" / f"{_SEARCH_STEM},raw.csv", + result_root / "search" / f"{_SEARCH_STEM},throughput.csv", + result_root / "search" / f"{_SEARCH_STEM},latency.csv", + ] + unrelated_path = result_root / "search" / "unrelated.csv" + for path in [*derived_paths, unrelated_path]: + path.write_text("stale\n", encoding="utf-8") + + convert_json_to_csv_build(_DATASET, str(tmp_path)) + convert_json_to_csv_search(_DATASET, str(tmp_path)) + + assert not any(path.exists() for path in derived_paths) + assert unrelated_path.read_text(encoding="utf-8") == "stale\n" + + +def test_named_build_join_preserves_search_codec_and_build_metadata(tmp_path): + codec = _CODECS[0] + index_name = _index_name(codec) + build_codec = "build-codec-sentinel" + search_codec = "search-codec-sentinel" + build_result = BuildResult( + index_path=f"/indexes/{index_name}", + build_time_seconds=2.75, + index_size_bytes=8192, + algorithm=_ALGORITHM, + build_params={"codec": build_codec}, + metadata={"writer_path": "gpu-hnsw"}, + ) + search_result = _search_result( + search_codec, + metadata={"mode": "latency"}, + ) + write_result_files( + [ + ResultRecord(build_result, index_name, _BUILD_STEM), + ResultRecord(search_result, index_name, _SEARCH_STEM), + ], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + convert_json_to_csv_build(_DATASET, str(tmp_path)) + convert_json_to_csv_search(_DATASET, str(tmp_path)) + + raw_path = ( + tmp_path / _DATASET / "result" / "search" / f"{_SEARCH_STEM},raw.csv" + ) + row = pd.read_csv(raw_path).iloc[0] + assert row["index_name"] == index_name + assert row["codec"] == search_codec + assert row["build time"] == pytest.approx(2.75) + assert row["writer_path"] == "gpu-hnsw" + assert row["mode"] == "latency" + assert pd.isna(row["build threads"]) + assert pd.isna(row["build cpu_time"]) + + +def test_scoped_search_joins_matching_build_and_preserves_legacy_stem( + tmp_path, +): + codec = _CODECS[0] + full_index_name = _index_name(codec) + subset_scope = "subset128" + subset_index_name = f"{_ALGORITHM}.{subset_scope}.codec{codec}" + subset_build_stem = f"{_BUILD_STEM},{subset_scope}" + subset_search_stem = f"{subset_build_stem},k5,bs2" + + write_result_files( + [ + ResultRecord( + _build_result(codec, build_time=1.25), + full_index_name, + _BUILD_STEM, + ), + ResultRecord( + _search_result(codec), + full_index_name, + _SEARCH_STEM, + ), + ResultRecord( + _build_result(codec, build_time=4.5), + subset_index_name, + subset_build_stem, + ), + ResultRecord( + _search_result(codec), + subset_index_name, + subset_search_stem, + ), + ], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + convert_json_to_csv_build(_DATASET, str(tmp_path)) + convert_json_to_csv_search(_DATASET, str(tmp_path)) + + result_root = tmp_path / _DATASET / "result" + full_row = pd.read_csv( + result_root / "search" / f"{_SEARCH_STEM},raw.csv" + ).iloc[0] + subset_row = pd.read_csv( + result_root / "search" / f"{subset_search_stem},raw.csv" + ).iloc[0] + assert full_row["index_name"] == full_index_name + assert full_row["build time"] == pytest.approx(1.25) + assert subset_row["index_name"] == subset_index_name + assert subset_row["build time"] == pytest.approx(4.5) diff --git a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_backend.py b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_backend.py new file mode 100644 index 0000000000..a5b75654c5 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_backend.py @@ -0,0 +1,1038 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Build and search behavior tests for the PyLucene backend.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import pytest + +import cuvs_bench.backends.pylucene as pylucene_backend +from cuvs_bench.backends.pylucene import _SearchHit +from cuvs_bench.tests._pylucene_test_utils import ( + _CAGRA_CODEC, + _HNSW_BASE_LAYER_CODEC, + _HNSW_CODEC, + _FakeRuntime, + _backend, + _dataset, + _index, + _prepare_cagra_index, + _prepare_hnsw_index, +) + + +@pytest.mark.parametrize( + ("score", "distance"), + [(1.0, 0.0), (0.5, 1.0), (0.2, 4.0), (0.0, np.inf)], +) +def test_score_to_squared_euclidean(score, distance): + assert pylucene_backend._score_to_squared_euclidean( + score + ) == pytest.approx(distance) + + +def test_build_dry_run_does_not_initialize_pylucene(tmp_path): + backend = _backend() + result = backend.build( + _dataset(), [_index(tmp_path / "index")], dry_run=True + ) + + assert result.success + assert backend._runtime is None + assert result.metadata["codec"] == _HNSW_CODEC + + +def test_build_creates_index_and_reports_gpu_writer(tmp_path): + runtime = _FakeRuntime() + backend = _backend(runtime) + index_path = tmp_path / "index" + + result = backend.build(_dataset(), [_index(index_path)], force=True) + + assert result.success + assert result.index_size_bytes > len(b"index") + assert result.metadata["writer_path"] == "gpu-hnsw" + assert result.metadata["writer_telemetry"]["writerPath"] == "gpu-hnsw" + assert result.metadata["hnsw_verification"] == { + "status": "gpu-hnsw-provenance", + "schema_version": 1, + "codec": _HNSW_CODEC, + "writer_path": "gpu-hnsw", + "vector_count": 10, + "dimensions": 4, + "commit_file_count": 1, + } + assert (index_path / pylucene_backend._HNSW_PROVENANCE_FILE).is_file() + assert runtime.resolve_calls == [_HNSW_CODEC] + assert runtime.telemetry_calls == [_HNSW_CODEC] + assert len(runtime.build_calls) == 1 + assert runtime.verification_calls == [] + + +def test_build_verifies_persisted_cagra_index(tmp_path): + runtime = _FakeRuntime(writer_path="gpu-cagra") + index_path = tmp_path / "index" + + result = _backend(runtime, codec=_CAGRA_CODEC).build( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + force=True, + ) + + assert result.success + assert runtime.verification_calls == [(index_path, 10, 4)] + assert result.metadata["cagra_provenance"] == { + "status": "gpu-cagra-provenance", + "schema_version": 1, + "codec": _CAGRA_CODEC, + "writer_path": "gpu-cagra", + "vector_count": 10, + "dimensions": 4, + "commit_file_count": 1, + } + assert result.metadata["cagra_verification"] == { + "status": "cagra-only", + "segment_count": 1, + "field_count": 1, + "vector_count": 10, + "dimensions": 4, + } + assert (index_path / pylucene_backend._CAGRA_PROVENANCE_FILE).is_file() + + +def test_build_rejects_persisted_cagra_fallback_and_removes_index(tmp_path): + runtime = _FakeRuntime(writer_path="gpu-cagra") + runtime.verification_error = RuntimeError("persisted brute-force index") + index_path = tmp_path / "index" + + result = _backend(runtime, codec=_CAGRA_CODEC).build( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + force=True, + ) + + assert not result.success + assert "persisted brute-force index" in result.error_message + assert runtime.verification_calls == [(index_path, 10, 4)] + assert len(runtime.build_calls) == 1 + assert not index_path.exists() + + +def test_build_reuses_existing_index_without_runtime(tmp_path): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + runtime = _FakeRuntime() + backend = _backend(runtime) + + result = backend.build(_dataset(), [_index(index_path)]) + + assert result.success + assert result.metadata["skipped"] is True + assert result.metadata["hnsw_verification"]["status"] == ( + "gpu-hnsw-provenance" + ) + assert runtime.resolve_calls == [] + assert runtime.telemetry_calls == [] + assert runtime.build_calls == [] + assert runtime.verification_calls == [] + + +def test_build_reports_reused_index_sizing_failure(tmp_path, monkeypatch): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + runtime = _FakeRuntime() + + def fail_index_size(_index_path): + raise PermissionError("cannot read index size") + + monkeypatch.setattr(pylucene_backend, "_index_size", fail_index_size) + + result = _backend(runtime).build(_dataset(), [_index(index_path)]) + + assert not result.success + assert "cannot read index size" in result.error_message + assert runtime.telemetry_calls == [] + assert runtime.build_calls == [] + + +def test_build_rejects_reused_hnsw_index_without_provenance(tmp_path): + index_path = tmp_path / "index" + index_path.mkdir() + (index_path / "segments_1").write_bytes(b"existing") + runtime = _FakeRuntime() + + result = _backend(runtime).build(_dataset(), [_index(index_path)]) + + assert not result.success + assert "provenance manifest is missing" in result.error_message + assert runtime.telemetry_calls == [] + assert runtime.build_calls == [] + + +def test_build_verifies_reused_cagra_index(tmp_path): + index_path = tmp_path / "index" + _prepare_cagra_index(index_path) + runtime = _FakeRuntime(writer_path="gpu-cagra") + + result = _backend(runtime, codec=_CAGRA_CODEC).build( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + ) + + assert result.success + assert result.metadata["skipped"] is True + assert runtime.verification_calls == [(index_path, 10, 4)] + assert result.metadata["cagra_provenance"]["status"] == ( + "gpu-cagra-provenance" + ) + assert result.metadata["cagra_verification"] == { + "status": "cagra-only", + "segment_count": 1, + "field_count": 1, + "vector_count": 10, + "dimensions": 4, + } + assert runtime.telemetry_calls == [] + assert runtime.build_calls == [] + + +def test_build_rejects_reused_cagra_index_that_cannot_be_verified(tmp_path): + index_path = tmp_path / "index" + _prepare_cagra_index(index_path) + segments_file = index_path / "segments_1" + runtime = _FakeRuntime(writer_path="gpu-cagra") + runtime.verification_error = RuntimeError("persisted brute-force index") + + result = _backend(runtime, codec=_CAGRA_CODEC).build( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + ) + + assert not result.success + assert "persisted brute-force index" in result.error_message + assert runtime.verification_calls == [(index_path, 10, 4)] + assert segments_file.read_bytes() == b"index" + assert runtime.telemetry_calls == [] + assert runtime.build_calls == [] + + +def test_build_rejects_reused_cagra_index_without_provenance_before_runtime( + tmp_path, +): + index_path = tmp_path / "index" + index_path.mkdir() + (index_path / "segments_1").write_bytes(b"existing") + runtime = _FakeRuntime(writer_path="gpu-cagra") + + result = _backend(runtime, codec=_CAGRA_CODEC).build( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + ) + + assert not result.success + assert "CAGRA provenance manifest is missing" in result.error_message + assert runtime.verification_calls == [] + assert runtime.search_calls == [] + + +@pytest.mark.parametrize("force", [False, True]) +def test_build_rejects_unrelated_existing_directory(tmp_path, force): + index_path = tmp_path / "index" + index_path.mkdir() + unrelated_file = index_path / "unrelated" + unrelated_file.write_bytes(b"not a Lucene index") + runtime = _FakeRuntime() + + result = _backend(runtime).build( + _dataset(), [_index(index_path)], force=force + ) + + assert not result.success + assert "does not contain a Lucene segments file" in result.error_message + assert unrelated_file.exists() + assert runtime.telemetry_calls == [] + + +def test_build_rejects_existing_index_path_that_is_a_file(tmp_path): + index_path = tmp_path / "index" + index_path.write_bytes(b"not a directory") + + result = _backend(_FakeRuntime()).build( + _dataset(), [_index(index_path)], force=True + ) + + assert not result.success + assert "not a directory" in result.error_message + assert index_path.read_bytes() == b"not a directory" + + +def test_build_force_replaces_existing_index(tmp_path): + index_path = tmp_path / "index" + index_path.mkdir() + (index_path / "segments_1").write_bytes(b"existing") + old_file = index_path / "old" + old_file.write_bytes(b"old") + runtime = _FakeRuntime() + + result = _backend(runtime).build( + _dataset(), [_index(index_path)], force=True + ) + + assert result.success + assert not old_file.exists() + assert (index_path / "segments_1").exists() + + +def test_build_preserves_existing_index_if_codec_preflight_fails(tmp_path): + index_path = tmp_path / "index" + index_path.mkdir() + (index_path / "segments_1").write_bytes(b"existing") + old_file = index_path / "old" + old_file.write_bytes(b"old") + runtime = _FakeRuntime() + runtime.resolve_error = RuntimeError("codec unavailable") + + result = _backend(runtime).build( + _dataset(), [_index(index_path)], force=True + ) + + assert not result.success + assert "codec unavailable" in result.error_message + assert old_file.exists() + + +def test_build_rejects_silent_cpu_fallback_and_removes_index(tmp_path): + runtime = _FakeRuntime(writer_path="cpu-hnsw-fallback") + index_path = tmp_path / "index" + + result = _backend(runtime).build( + _dataset(), [_index(index_path)], force=True + ) + + assert not result.success + assert "silent CPU" in result.error_message + assert not index_path.exists() + assert runtime.build_calls == [] + + +def test_build_preserves_existing_index_on_cpu_fallback(tmp_path): + runtime = _FakeRuntime(writer_path="cpu-hnsw-fallback") + index_path = tmp_path / "index" + index_path.mkdir() + (index_path / "segments_1").write_bytes(b"existing") + old_file = index_path / "old" + old_file.write_bytes(b"old") + + result = _backend(runtime).build( + _dataset(), [_index(index_path)], force=True + ) + + assert not result.success + assert "silent CPU" in result.error_message + assert old_file.read_bytes() == b"old" + assert runtime.build_calls == [] + + +@pytest.mark.parametrize( + ("dataset", "error"), + [ + (_dataset(dtype=np.float64), "float32"), + (_dataset(distance_metric="cosine"), "Euclidean"), + ], +) +def test_build_rejects_unsupported_dataset(dataset, error, tmp_path): + result = _backend(_FakeRuntime()).build( + dataset, [_index(tmp_path / "index")], force=True + ) + + assert not result.success + assert error in result.error_message + + +def test_build_rejects_nonfinite_vectors(tmp_path): + dataset = _dataset() + dataset.training_vectors[0, 0] = np.nan + + result = _backend(_FakeRuntime()).build( + dataset, [_index(tmp_path / "index")], force=True + ) + + assert not result.success + assert "finite" in result.error_message + + +@pytest.mark.parametrize( + ("codec", "writer_path"), + [ + (_HNSW_CODEC, "gpu-hnsw"), + (_CAGRA_CODEC, "gpu-cagra"), + ], + ids=["hnsw", "cagra"], +) +def test_build_rejects_single_vector_cuvs_bypass(tmp_path, codec, writer_path): + runtime = _FakeRuntime(writer_path=writer_path) + + result = _backend(runtime, codec=codec).build( + _dataset(n_base=1), + [_index(tmp_path / "index", codec=codec)], + force=True, + ) + + assert not result.success + assert "at least two training vectors" in result.error_message + assert "does not invoke cuVS" in result.error_message + assert runtime.telemetry_calls == [] + assert not (tmp_path / "index").exists() + + +def test_build_rejects_unknown_codec(tmp_path): + result = _backend(_FakeRuntime()).build( + _dataset(), + [_index(tmp_path / "index", codec="UnknownCodec")], + force=True, + ) + + assert not result.success + assert "Unsupported PyLucene codec" in result.error_message + + +@pytest.mark.parametrize( + "codec", + [None, "", 0, False], + ids=["none", "empty", "zero", "false"], +) +def test_explicit_invalid_codec_does_not_fall_back_to_backend_config( + tmp_path, codec +): + index = _index(tmp_path / "index") + index.build_param["codec"] = codec + backend = _backend(codec=_HNSW_CODEC) + + build_result = backend.build(_dataset(), [index], dry_run=True) + search_result = backend.search(_dataset(), [index], k=3, dry_run=True) + + for result in (build_result, search_result): + assert not result.success + assert f"Unsupported PyLucene codec {codec!r}" in result.error_message + + +def test_build_rejects_unsupported_parameter(tmp_path): + index = _index(tmp_path / "index") + index.build_param["ignored"] = 1 + + result = _backend(_FakeRuntime()).build(_dataset(), [index]) + + assert not result.success + assert "Unsupported PyLucene build parameter" in result.error_message + + +@pytest.mark.parametrize("index_count", [0, 2]) +def test_build_requires_exactly_one_index(index_count, tmp_path): + indexes = [ + _index(tmp_path / f"index-{index_id}") + for index_id in range(index_count) + ] + + result = _backend(_FakeRuntime()).build(_dataset(), indexes, force=True) + + assert not result.success + assert "exactly one" in result.error_message + + +@pytest.mark.parametrize( + ("vectors", "error"), + [ + (np.empty((0, 4), dtype=np.float32), "at least one vector"), + (np.ones(4, dtype=np.float32), "two-dimensional"), + (np.ones((1, 4097), dtype=np.float32), "4096"), + ], +) +def test_build_rejects_invalid_vector_shape(vectors, error, tmp_path): + dataset = _dataset() + dataset.training_vectors = vectors + + result = _backend(_FakeRuntime()).build( + dataset, [_index(tmp_path / "index")], force=True + ) + + assert not result.success + assert error in result.error_message + + +def test_build_removes_partial_index_after_runtime_failure(tmp_path): + runtime = _FakeRuntime() + runtime.build_error = RuntimeError("Java build failed") + index_path = tmp_path / "index" + + result = _backend(runtime).build( + _dataset(), [_index(index_path)], force=True + ) + + assert not result.success + assert "Java build failed" in result.error_message + assert not index_path.exists() + + +def test_build_reports_chained_runtime_failure(tmp_path): + runtime = _FakeRuntime() + runtime.build_error = RuntimeError("rollback failed") + runtime.build_error.__cause__ = RuntimeError("add failed") + index_path = tmp_path / "index" + + result = _backend(runtime).build( + _dataset(), [_index(index_path)], force=True + ) + + assert not result.success + assert "RuntimeError: rollback failed" in result.error_message + assert "caused by RuntimeError: add failed" in result.error_message + assert not index_path.exists() + + +def test_build_removes_partial_index_and_reraises_interrupt(tmp_path): + runtime = _FakeRuntime() + runtime.build_error = KeyboardInterrupt() + index_path = tmp_path / "index" + + with pytest.raises(KeyboardInterrupt): + _backend(runtime).build(_dataset(), [_index(index_path)], force=True) + + assert not index_path.exists() + + +def test_build_preserves_interrupt_when_partial_index_cleanup_fails( + tmp_path, monkeypatch +): + runtime = _FakeRuntime() + runtime.build_error = KeyboardInterrupt() + index_path = tmp_path / "index" + + def fail_cleanup(_index_path): + raise PermissionError("cleanup denied") + + monkeypatch.setattr(pylucene_backend, "_safe_remove_index", fail_cleanup) + with pytest.raises(KeyboardInterrupt) as exc_info: + _backend(runtime).build(_dataset(), [_index(index_path)], force=True) + + assert exc_info.value.__notes__ == [ + "Failed to remove partial index: PermissionError: cleanup denied" + ] + assert index_path.exists() + + +def test_build_reports_partial_index_cleanup_failure(tmp_path, monkeypatch): + runtime = _FakeRuntime() + runtime.build_error = RuntimeError("Java build failed") + index_path = tmp_path / "index" + + def fail_cleanup(_index_path): + raise PermissionError("cleanup denied") + + monkeypatch.setattr(pylucene_backend, "_safe_remove_index", fail_cleanup) + result = _backend(runtime).build( + _dataset(), [_index(index_path)], force=True + ) + + assert not result.success + assert "Java build failed" in result.error_message + assert "failed to remove partial index" in result.error_message + assert "cleanup denied" in result.error_message + assert index_path.exists() + + +def test_search_dry_run_does_not_initialize_pylucene(tmp_path): + backend = _backend() + + result = backend.search( + _dataset(), [_index(tmp_path / "index")], k=3, dry_run=True + ) + + assert result.success + assert backend._runtime is None + assert result.neighbors.shape == (0, 3) + + +def test_search_converts_hits_scores_and_padding(tmp_path): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + runtime = _FakeRuntime( + hits=[ + [ + _SearchHit(document_id=3, score=1.0), + _SearchHit(document_id=7, score=0.5), + ], + [_SearchHit(document_id=2, score=0.2)], + ] + ) + + result = _backend(runtime).search( + _dataset(), [_index(index_path)], k=3, batch_size=1 + ) + + assert result.success + np.testing.assert_array_equal( + result.neighbors, + np.array([[3, 7, -1], [2, -1, -1]], dtype=np.int64), + ) + np.testing.assert_allclose(result.distances[0, :2], [0.0, 1.0]) + assert result.distances[1, 0] == pytest.approx(4.0) + assert np.isinf(result.distances[:, -1]).all() + assert result.search_time_ms == pytest.approx(2.0) + assert result.latency_seconds == pytest.approx(0.001) + assert result.queries_per_second == pytest.approx(1000.0) + assert result.latency_percentiles == { + "p50": 1.0, + "p95": 1.0, + "p99": 1.0, + } + assert result.metadata["num_batches"] == 2 + assert result.metadata["hnsw_verification"]["status"] == ( + "gpu-hnsw-provenance" + ) + assert "per_search_param_results" not in result.metadata + assert runtime.search_calls[0][3] == 1 + assert runtime.verification_calls == [] + + +def test_search_end_to_end_timing_starts_before_input_verification( + tmp_path, monkeypatch +): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + backend = _backend(_FakeRuntime()) + events = [] + times = iter((10.0, 10.012)) + load_search_inputs = backend._load_search_inputs + + def record_time(): + events.append("time") + return next(times) + + def record_input_loading(dataset): + events.append("inputs") + return load_search_inputs(dataset) + + monkeypatch.setattr(pylucene_backend.time, "perf_counter", record_time) + monkeypatch.setattr(backend, "_load_search_inputs", record_input_loading) + + result = backend.search( + _dataset(), [_index(index_path)], k=3, batch_size=2 + ) + + assert result.success + assert events == ["time", "inputs", "time"] + assert result.metadata["end_to_end_time_ms"] == pytest.approx(12.0) + assert result.metadata["non_query_overhead_time_ms"] == pytest.approx(11.0) + + +@pytest.mark.parametrize( + ("runtime", "error"), + [ + ( + _FakeRuntime(document_count=9), + "document count does not match index provenance", + ), + ( + _FakeRuntime( + hits=[ + [ + _SearchHit(document_id=3, score=1.0), + _SearchHit(document_id=3, score=0.5), + ], + [_SearchHit(document_id=2, score=1.0)], + ] + ), + "duplicate stored ID", + ), + ( + _FakeRuntime( + hits=[ + [_SearchHit(document_id=10, score=1.0)], + [_SearchHit(document_id=2, score=1.0)], + ] + ), + "out-of-range stored ID", + ), + ( + _FakeRuntime(hits=[[_SearchHit(document_id=0, score=1.0)]]), + "unexpected number of query results", + ), + ( + _FakeRuntime( + hits=[ + [_SearchHit(document_id=0, score=float("nan"))], + [_SearchHit(document_id=1, score=1.0)], + ] + ), + "non-finite score", + ), + ( + _FakeRuntime( + hits=[ + [_SearchHit(document_id=0, score=1.1)], + [_SearchHit(document_id=1, score=1.0)], + ] + ), + "outside the Euclidean range", + ), + ( + _FakeRuntime( + hits=[ + [ + _SearchHit(document_id=0, score=1.0), + _SearchHit(document_id=1, score=0.9), + _SearchHit(document_id=2, score=0.8), + _SearchHit(document_id=3, score=0.7), + ], + [_SearchHit(document_id=1, score=1.0)], + ] + ), + "too many hits", + ), + ( + _FakeRuntime(batch_latencies_ms=[1.0, 2.0]), + "unexpected number of batch latencies", + ), + ( + _FakeRuntime(batch_latencies_ms=[float("nan")]), + "invalid batch latency", + ), + ], + ids=[ + "document-count", + "duplicate-id", + "out-of-range-id", + "query-count", + "score", + "score-range", + "hit-count", + "latency-count", + "batch-latency", + ], +) +def test_search_rejects_invalid_runtime_results(tmp_path, runtime, error): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + + result = _backend(runtime).search( + _dataset(), [_index(index_path)], k=3, batch_size=2 + ) + + assert not result.success + assert error in result.error_message + + +def test_search_verifies_cagra_index_before_querying(tmp_path): + index_path = tmp_path / "index" + _prepare_cagra_index(index_path) + runtime = _FakeRuntime(writer_path="gpu-cagra") + + result = _backend(runtime, codec=_CAGRA_CODEC).search( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + k=3, + ) + + assert result.success + assert runtime.verification_calls == [(index_path, 10, 4)] + assert len(runtime.search_calls) == 1 + assert result.metadata["cagra_provenance"]["status"] == ( + "gpu-cagra-provenance" + ) + assert result.metadata["cagra_verification"] == { + "status": "cagra-only", + "segment_count": 1, + "field_count": 1, + "vector_count": 10, + "dimensions": 4, + } + + +def test_search_rejects_unverified_cagra_index_before_querying(tmp_path): + index_path = tmp_path / "index" + _prepare_cagra_index(index_path) + runtime = _FakeRuntime(writer_path="gpu-cagra") + runtime.verification_error = RuntimeError("persisted brute-force index") + + result = _backend(runtime, codec=_CAGRA_CODEC).search( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + k=3, + ) + + assert not result.success + assert "persisted brute-force index" in result.error_message + assert runtime.verification_calls == [(index_path, 10, 4)] + assert runtime.search_calls == [] + + +@pytest.mark.parametrize("manifest_state", ["missing", "stale"]) +def test_search_rejects_invalid_cagra_provenance_before_runtime( + tmp_path, manifest_state +): + index_path = tmp_path / "index" + manifest_path = _prepare_cagra_index(index_path) + if manifest_state == "missing": + manifest_path.unlink() + else: + (index_path / "segments_1").write_bytes(b"changed commit") + runtime = _FakeRuntime(writer_path="gpu-cagra") + + result = _backend(runtime, codec=_CAGRA_CODEC).search( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + k=3, + ) + + assert not result.success + assert "CAGRA provenance" in result.error_message + assert runtime.verification_calls == [] + assert runtime.search_calls == [] + + +@pytest.mark.parametrize( + ("index", "k", "batch_size", "search_threads", "mode", "error"), + [ + ( + _index(Path("/tmp/index")), + 0, + 10000, + None, + "latency", + "k must be positive", + ), + ( + _index(Path("/tmp/index")), + 3, + 10000, + 2, + "latency", + "one search thread", + ), + ( + _index(Path("/tmp/index")), + 3, + 10000, + None, + "throughput", + "only latency mode", + ), + ( + _index(Path("/tmp/index")), + 3, + 0, + None, + "latency", + "batch_size must be positive", + ), + ( + _index(Path("/tmp/index"), search_params=[{"search_width": 16}]), + 3, + 10000, + None, + "latency", + "does not expose", + ), + ( + _index(Path("/tmp/index"), codec=_CAGRA_CODEC), + 1025, + 10000, + None, + "latency", + "GPU brute-force", + ), + ], +) +def test_search_rejects_unsupported_options( + index, k, batch_size, search_threads, mode, error +): + result = _backend(_FakeRuntime()).search( + _dataset(), + [index], + k=k, + batch_size=batch_size, + mode=mode, + search_threads=search_threads, + dry_run=True, + ) + + assert not result.success + assert error in result.error_message + + +def test_resolve_search_plan_normalizes_the_complete_request(tmp_path): + index = _index(tmp_path / "index", search_params=[]) + + plan = _backend()._resolve_search_plan( + index, + k=3, + batch_size=7, + mode="latency", + search_threads="1", + ) + + assert plan == pylucene_backend._SearchPlan( + index_path=tmp_path / "index", + codec_name=_HNSW_CODEC, + search_params=[{}], + k=3, + batch_size=7, + mode="latency", + ) + + +def test_search_validation_preserves_first_error_precedence(tmp_path): + result = _backend().search( + _dataset(), + [_index(tmp_path / "index", search_params=[{"unsupported": True}])], + k=0, + batch_size=0, + mode="throughput", + search_threads=2, + dry_run=True, + ) + + assert not result.success + assert result.error_message == "k must be positive" + + +@pytest.mark.parametrize("index_count", [0, 2]) +def test_search_requires_exactly_one_index(index_count, tmp_path): + indexes = [ + _index(tmp_path / f"index-{index_id}") + for index_id in range(index_count) + ] + + result = _backend(_FakeRuntime()).search(_dataset(), indexes, k=3) + + assert not result.success + assert "exactly one" in result.error_message + + +def test_search_rejects_missing_index(tmp_path): + result = _backend(_FakeRuntime()).search( + _dataset(), [_index(tmp_path / "missing")], k=3 + ) + + assert not result.success + assert "does not exist" in result.error_message + + +def test_search_rejects_empty_query_vectors(tmp_path): + index_path = tmp_path / "index" + index_path.mkdir() + dataset = _dataset() + dataset.query_vectors = np.empty((0, 4), dtype=np.float32) + + result = _backend(_FakeRuntime()).search( + dataset, [_index(index_path)], k=3 + ) + + assert not result.success + assert "at least one vector" in result.error_message + + +def test_search_reports_runtime_query_dimension_mismatch(tmp_path): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + runtime = _FakeRuntime(index_dimensions=8) + + result = _backend(runtime).search( + _dataset(dimensions=4), [_index(index_path)], k=3 + ) + + assert not result.success + assert "dimensions do not match the Lucene index" in result.error_message + + +def test_search_rejects_runtime_dimensions_that_disagree_with_provenance( + tmp_path, +): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path, dimensions=4) + runtime = _FakeRuntime( + index_dimensions=8, + validate_query_dimensions=False, + ) + + result = _backend(runtime).search( + _dataset(dimensions=4), [_index(index_path)], k=3 + ) + + assert not result.success + assert ( + "Lucene index dimensions do not match index provenance" + in result.error_message + ) + assert len(runtime.search_calls) == 1 + + +def test_search_rejects_query_dimensions_that_disagree_with_dataset(tmp_path): + index_path = tmp_path / "index" + index_path.mkdir() + dataset = _dataset(dimensions=4) + dataset.query_vectors = np.zeros((2, 8), dtype=np.float32) + runtime = _FakeRuntime() + + result = _backend(runtime).search(dataset, [_index(index_path)], k=3) + + assert not result.success + assert "Query vector dimensions do not match the dataset" in ( + result.error_message + ) + assert runtime.search_calls == [] + + +def test_search_reports_runtime_failure(tmp_path): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + runtime = _FakeRuntime() + runtime.search_error = RuntimeError("Java search failed") + + result = _backend(runtime).search(_dataset(), [_index(index_path)], k=3) + + assert not result.success + assert "Java search failed" in result.error_message + + +@pytest.mark.parametrize( + "manifest_state", + ["missing", "corrupt", "stale", "wrong-codec"], +) +def test_search_rejects_invalid_hnsw_provenance_before_runtime( + tmp_path, manifest_state +): + index_path = tmp_path / "index" + manifest_path = _prepare_hnsw_index(index_path) + if manifest_state == "missing": + manifest_path.unlink() + elif manifest_state == "corrupt": + manifest_path.write_text("{") + elif manifest_state == "stale": + (index_path / "segments_1").write_bytes(b"changed commit") + else: + payload = json.loads(manifest_path.read_text()) + payload["codec"] = _HNSW_BASE_LAYER_CODEC + manifest_path.write_text(json.dumps(payload)) + + runtime = _FakeRuntime() + result = _backend(runtime).search(_dataset(), [_index(index_path)], k=3) + + assert not result.success + assert "provenance" in result.error_message.lower() + assert runtime.search_calls == [] + + +def test_cleanup_releases_runtime_reference(): + backend = _backend(_FakeRuntime()) + + backend.cleanup() + + assert backend._runtime is None diff --git a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_cagra_verifier.py b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_cagra_verifier.py new file mode 100644 index 0000000000..fa320fdde4 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_cagra_verifier.py @@ -0,0 +1,815 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Persisted CAGRA verifier unit tests.""" + +from __future__ import annotations + +import os +from types import SimpleNamespace + +import pytest + +import cuvs_bench.backends.pylucene as pylucene_backend + + +class _FakeCagraMetadataInput: + def __init__(self, integers, variable_longs): + self._integers = iter(integers) + self._variable_longs = iter(variable_longs) + + def readInt(self): + return next(self._integers) + + def readVLong(self): + return next(self._variable_longs) + + +class _FakeChecksumInput(_FakeCagraMetadataInput): + def __init__(self, integers, variable_longs): + super().__init__(integers, variable_longs) + self.closed = False + + def close(self): + self.closed = True + + +class _FakeIndexInput: + def __init__(self, payload_start=57, payload_length=1000): + self.payload_start = payload_start + self.file_length = payload_start + payload_length + 16 + self.closed = False + + def getFilePointer(self): + return self.payload_start + + def length(self): + return self.file_length + + def close(self): + self.closed = True + + +class _FakeCodecUtil: + def __init__(self, error_at=None): + self.error_at = error_at + self.header_calls = [] + self.footer_calls = [] + self.checksum_calls = [] + self.retrieved_checksum_calls = [] + + def checkIndexHeader(self, *args): + self.header_calls.append(args) + if self.error_at == "header": + raise RuntimeError("unsupported metadata version") + return 0 + + def checkFooter(self, metadata_input): + self.footer_calls.append(metadata_input) + if self.error_at == "footer": + raise RuntimeError("invalid checksum footer") + + def footerLength(self): + return 16 + + def checksumEntireFile(self, data_input): + self.checksum_calls.append(data_input) + if self.error_at == "data-checksum": + raise RuntimeError("invalid data checksum") + return 0 + + def retrieveChecksum(self, data_input): + self.retrieved_checksum_calls.append(data_input) + return 0 + + +class _FakeFieldInfos: + def __init__(self, field_infos): + self._field_infos = { + field_info.number: field_info for field_info in field_infos + } + + def fieldInfo(self, field_number): + return self._field_infos.get(field_number) + + def __iter__(self): + return iter(self._field_infos.values()) + + +def _fake_cagra_index_verifier( + metadata_input, + *, + metadata_files=("_0.vemc",), + codec_util=None, + data_payload_length=1000, + field_name="vector", + field_dimensions=32, + field_updates=False, + max_documents=4, + deletion_count=0, + soft_deletion_count=0, + extra_vector_field=False, + data_error=None, +): + data_input = _FakeIndexInput(payload_length=data_payload_length) + vector_encoding = object() + vector_similarity = object() + field_info = SimpleNamespace( + number=1, + getName=lambda: field_name, + getVectorDimension=lambda: field_dimensions, + getVectorEncoding=lambda: vector_encoding, + getVectorSimilarityFunction=lambda: vector_similarity, + ) + all_field_infos = [field_info] + if extra_vector_field: + all_field_infos.append( + SimpleNamespace( + number=2, + getName=lambda: "extra-vector", + getVectorDimension=lambda: field_dimensions, + getVectorEncoding=lambda: vector_encoding, + getVectorSimilarityFunction=lambda: vector_similarity, + ) + ) + field_infos = _FakeFieldInfos(all_field_infos) + field_infos_format = SimpleNamespace( + read=lambda _directory, _info, _suffix, _context: field_infos + ) + + def open_data(_file_name, _context): + if data_error is not None: + raise data_error + return data_input + + directory = SimpleNamespace( + openChecksumInput=lambda _file_name: metadata_input, + openInput=open_data, + close=lambda: None, + ) + segment_info = SimpleNamespace( + info=SimpleNamespace( + name="_0", + getId=lambda: b"segment-id", + getCodec=lambda: SimpleNamespace( + fieldInfosFormat=lambda: field_infos_format + ), + maxDoc=lambda: max_documents, + ), + files=lambda: metadata_files, + hasFieldUpdates=lambda: field_updates, + hasDeletions=lambda: deletion_count > 0, + getDelCount=lambda: deletion_count, + getSoftDelCount=lambda: soft_deletion_count, + ) + verifier = pylucene_backend._CagraIndexVerifier( + attach_current_thread=lambda: None, + paths=SimpleNamespace(get=lambda path: path), + codec_util=codec_util or _FakeCodecUtil(), + field_info=SimpleNamespace( + cast_=lambda raw_field_info: raw_field_info + ), + segment_commit_info=SimpleNamespace( + cast_=lambda raw_segment_info: raw_segment_info + ), + segment_infos=SimpleNamespace( + readLatestCommit=lambda _directory: [segment_info] + ), + vector_encoding=SimpleNamespace(FLOAT32=vector_encoding), + vector_similarity_function=SimpleNamespace( + EUCLIDEAN=vector_similarity + ), + fs_directory=SimpleNamespace(open=lambda _path: directory), + io_context=SimpleNamespace(READONCE=object()), + ) + verifier._test_data_input = data_input + verifier._test_directory = directory + verifier._test_segment_info = segment_info + return verifier + + +def _cagra_data_context(verifier, index_path): + return pylucene_backend._CagraDataFileContext( + index_path=index_path, + directory=verifier._test_directory, + segment_info=verifier._test_segment_info, + suffix="", + metadata_file="_0.vemc", + ) + + +@pytest.mark.parametrize( + ("metadata_file", "expected_suffix"), + [("_0.vemc", ""), ("_0_CuVS_0.vemc", "CuVS_0")], +) +def test_cagra_segment_suffix(metadata_file, expected_suffix): + assert ( + pylucene_backend._CagraIndexVerifier._segment_suffix( + "_0", metadata_file + ) + == expected_suffix + ) + + +def test_cagra_segment_suffix_rejects_unrelated_metadata_file(): + with pytest.raises(RuntimeError, match="does not match segment"): + pylucene_backend._CagraIndexVerifier._segment_suffix("_0", "_1.vemc") + + +def test_read_cagra_fields_accepts_only_persisted_cagra_data(): + metadata_input = _FakeCagraMetadataInput( + integers=[ + 1, + 1, + 0, + 32, + 4, + 2, + 1, + 0, + 32, + 0, + -1, + ], + variable_longs=[57, 1000, 1057, 0, 0, 0, 0, 0], + ) + + assert pylucene_backend._CagraIndexVerifier._read_cagra_fields( + metadata_input, "_0.vemc" + ) == [ + pylucene_backend._CagraFieldMetadata( + field_number=1, + dimensions=32, + vector_count=4, + cagra_offset=57, + cagra_length=1000, + ) + ] + + +@pytest.mark.parametrize( + ("cagra_length", "brute_force_length", "error"), + [ + (0, 512, "persisted brute-force"), + (1000, 512, "persisted brute-force"), + ], +) +def test_read_cagra_fields_rejects_non_cagra_only_data( + cagra_length, brute_force_length, error +): + metadata_input = _FakeCagraMetadataInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[ + 57, + cagra_length, + 57 + cagra_length, + brute_force_length, + ], + ) + + with pytest.raises(RuntimeError, match=error): + pylucene_backend._CagraIndexVerifier._read_cagra_fields( + metadata_input, "_0.vemc" + ) + + +def test_read_cagra_fields_rejects_data_for_empty_field(): + metadata_input = _FakeCagraMetadataInput( + integers=[1, 1, 0, 32, 0, -1], + variable_longs=[57, 1, 58, 0], + ) + + with pytest.raises(RuntimeError, match="empty field"): + pylucene_backend._CagraIndexVerifier._read_cagra_fields( + metadata_input, "_0.vemc" + ) + + +@pytest.mark.parametrize( + ("encoding", "similarity", "dimensions", "error"), + [ + (0, 0, 32, "encoding ordinal 0"), + (1, 1, 32, "similarity ordinal 1"), + (1, 0, 0, "invalid field metadata"), + (1, 0, 4097, "invalid field metadata"), + ], +) +def test_read_cagra_fields_rejects_unsupported_vector_semantics( + encoding, similarity, dimensions, error +): + metadata_input = _FakeCagraMetadataInput( + integers=[1, encoding, similarity, dimensions, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ) + + with pytest.raises(RuntimeError, match=error): + pylucene_backend._CagraIndexVerifier._read_cagra_fields( + metadata_input, "_0.vemc" + ) + + +def test_verify_cagra_index_validates_header_footer_and_expected_count( + tmp_path, +): + metadata_input = _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ) + codec_util = _FakeCodecUtil() + verifier = _fake_cagra_index_verifier( + metadata_input, codec_util=codec_util + ) + + verification = verifier.verify_index( + tmp_path, expected_vector_count=4, expected_dimensions=32 + ) + + assert verification.to_metadata() == { + "status": "cagra-only", + "segment_count": 1, + "field_count": 1, + "vector_count": 4, + "dimensions": 32, + } + assert codec_util.header_calls == [ + ( + metadata_input, + "Lucene102CuVSVectorsFormatMeta", + 0, + 0, + b"segment-id", + "", + ), + ( + verifier._test_data_input, + "Lucene102CuVSVectorsFormatIndex", + 0, + 0, + b"segment-id", + "", + ), + ] + assert codec_util.footer_calls == [metadata_input] + assert codec_util.checksum_calls == [verifier._test_data_input] + assert metadata_input.closed is True + assert verifier._test_data_input.closed is True + + +def test_verify_cagra_index_traverses_segments_and_closes_each_input( + tmp_path, +): + events = [] + metadata_inputs = { + name: _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ) + for name in ("_0.vemc", "_1.vemc") + } + data_inputs = {name: _FakeIndexInput() for name in ("_0.vcag", "_1.vcag")} + verifier = _fake_cagra_index_verifier(metadata_inputs["_0.vemc"]) + + def segment(name): + field_info = SimpleNamespace( + number=1, + getName=lambda: "vector", + getVectorDimension=lambda: 32, + getVectorEncoding=lambda: verifier.VectorEncoding.FLOAT32, + getVectorSimilarityFunction=( + lambda: verifier.VectorSimilarityFunction.EUCLIDEAN + ), + ) + field_infos_format = SimpleNamespace( + read=lambda _directory, _info, _suffix, _context: ( + _FakeFieldInfos([field_info]) + ) + ) + return SimpleNamespace( + info=SimpleNamespace( + name=name, + getId=lambda: name.encode(), + getCodec=lambda: SimpleNamespace( + fieldInfosFormat=lambda: field_infos_format + ), + maxDoc=lambda: 4, + ), + files=lambda: (f"{name}.vemc",), + hasFieldUpdates=lambda: False, + hasDeletions=lambda: False, + getDelCount=lambda: 0, + getSoftDelCount=lambda: 0, + ) + + def open_metadata(file_name): + events.append(f"open {file_name}") + metadata_input = metadata_inputs[file_name] + + def close_metadata(): + metadata_input.closed = True + events.append(f"close {file_name}") + + metadata_input.close = close_metadata + return metadata_input + + def open_data(file_name, _context): + events.append(f"open {file_name}") + data_input = data_inputs[file_name] + + def close_data(): + data_input.closed = True + events.append(f"close {file_name}") + + data_input.close = close_data + return data_input + + verifier.SegmentInfos.readLatestCommit = lambda _directory: [ + segment("_0"), + segment("_1"), + ] + verifier._test_directory.openChecksumInput = open_metadata + verifier._test_directory.openInput = open_data + verifier._test_directory.close = lambda: events.append("close directory") + + verification = verifier.verify_index(tmp_path) + + assert verification == pylucene_backend._CagraIndexVerification( + segment_count=2, + field_count=2, + vector_count=8, + dimensions=32, + ) + assert events == [ + "open _0.vemc", + "close _0.vemc", + "open _0.vcag", + "close _0.vcag", + "open _1.vemc", + "close _1.vemc", + "open _1.vcag", + "close _1.vcag", + "close directory", + ] + assert all(item.closed for item in metadata_inputs.values()) + assert all(item.closed for item in data_inputs.values()) + + +@pytest.mark.parametrize("error_at", ["header", "footer"]) +def test_verify_cagra_index_fails_closed_on_invalid_format(error_at, tmp_path): + metadata_input = _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ) + verifier = _fake_cagra_index_verifier( + metadata_input, codec_util=_FakeCodecUtil(error_at=error_at) + ) + + with pytest.raises(RuntimeError, match="metadata format v0"): + verifier.verify_index(tmp_path) + + assert metadata_input.closed is True + + +def test_verify_cagra_index_rejects_missing_segment_metadata(tmp_path): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput([], []), metadata_files=() + ) + + with pytest.raises(RuntimeError, match="no .vemc metadata"): + verifier.verify_index(tmp_path) + + +def test_verify_cagra_index_rejects_only_empty_fields(tmp_path): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 0, -1], + variable_longs=[0, 0, 0, 0], + ), + data_payload_length=0, + ) + + with pytest.raises(RuntimeError, match="without matching CAGRA metadata"): + verifier.verify_index(tmp_path) + + +def test_verify_cagra_index_rejects_vector_count_mismatch(tmp_path): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ) + ) + + with pytest.raises(RuntimeError, match="4 vectors; expected 5"): + verifier.verify_index(tmp_path, expected_vector_count=5) + + +def test_verify_cagra_index_rejects_dimension_mismatch(tmp_path): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ) + ) + + with pytest.raises(RuntimeError, match="32 dimensions; expected 16"): + verifier.verify_index(tmp_path, expected_dimensions=16) + + +@pytest.mark.parametrize( + ("runtime_kwargs", "expected_counts"), + [ + ({"deletion_count": 1}, "deleted=1, soft_deleted=0"), + ({"soft_deletion_count": 1}, "deleted=0, soft_deleted=1"), + ], + ids=["hard-deletion", "soft-deletion"], +) +def test_verify_cagra_index_rejects_committed_deletions( + tmp_path, runtime_kwargs, expected_counts +): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ), + **runtime_kwargs, + ) + + with pytest.raises(RuntimeError, match=expected_counts): + verifier.verify_index(tmp_path) + + +def test_verify_cagra_index_rejects_segment_document_count_mismatch(tmp_path): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ), + max_documents=5, + ) + + with pytest.raises(RuntimeError, match="4 vectors for 5 documents"): + verifier.verify_index(tmp_path) + + +def test_verify_cagra_index_rejects_unaccounted_vector_field(tmp_path): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ), + extra_vector_field=True, + ) + + with pytest.raises(RuntimeError, match=r"metadata=\[1\], Lucene=\[1, 2\]"): + verifier.verify_index(tmp_path) + + +def test_verify_cagra_index_rejects_duplicate_field_across_metadata_files( + tmp_path, +): + metadata_inputs = { + name: _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ) + for name in ("_0.vemc", "_0_CuVS_0.vemc") + } + verifier = _fake_cagra_index_verifier( + metadata_inputs["_0.vemc"], + metadata_files=tuple(metadata_inputs), + ) + verifier._test_directory.openChecksumInput = metadata_inputs.__getitem__ + verifier._test_directory.openInput = ( + lambda _file_name, _context: _FakeIndexInput() + ) + + with pytest.raises(RuntimeError, match="duplicate field 1"): + verifier.verify_index(tmp_path) + + +def test_verify_cagra_index_rejects_inconsistent_dimensions_across_segments( + tmp_path, +): + metadata_inputs = { + "_0.vemc": _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ), + "_1.vemc": _FakeChecksumInput( + integers=[1, 1, 0, 16, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ), + } + verifier = _fake_cagra_index_verifier(metadata_inputs["_0.vemc"]) + + def segment(name, dimensions): + field_info = SimpleNamespace( + number=1, + getName=lambda: "vector", + getVectorDimension=lambda: dimensions, + getVectorEncoding=lambda: verifier.VectorEncoding.FLOAT32, + getVectorSimilarityFunction=( + lambda: verifier.VectorSimilarityFunction.EUCLIDEAN + ), + ) + field_infos = _FakeFieldInfos([field_info]) + field_infos_format = SimpleNamespace( + read=lambda _directory, _info, _suffix, _context: field_infos + ) + return SimpleNamespace( + info=SimpleNamespace( + name=name, + getId=lambda: name.encode(), + getCodec=lambda: SimpleNamespace( + fieldInfosFormat=lambda: field_infos_format + ), + maxDoc=lambda: 4, + ), + files=lambda: (f"{name}.vemc",), + hasFieldUpdates=lambda: False, + hasDeletions=lambda: False, + getDelCount=lambda: 0, + getSoftDelCount=lambda: 0, + ) + + segments = [segment("_0", 32), segment("_1", 16)] + verifier.SegmentInfos.readLatestCommit = lambda _directory: segments + verifier._test_directory.openChecksumInput = metadata_inputs.__getitem__ + verifier._test_directory.openInput = ( + lambda _file_name, _context: _FakeIndexInput() + ) + + with pytest.raises(RuntimeError, match=r"dimensions: \[16, 32\]"): + verifier.verify_index(tmp_path) + + +@pytest.mark.parametrize( + ("runtime_kwargs", "error"), + [ + ({"field_name": "other"}, "unexpected field"), + ({"field_dimensions": 16}, "field metadata"), + ({"field_updates": True}, "field-info updates"), + ({"data_payload_length": 999}, "do not exactly cover"), + ], +) +def test_verify_cagra_index_rejects_foreign_or_inconsistent_index( + tmp_path, runtime_kwargs, error +): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ), + **runtime_kwargs, + ) + + with pytest.raises(RuntimeError, match=error): + verifier.verify_index(tmp_path) + + +def test_verify_cagra_index_rejects_corrupt_data_file(tmp_path): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ), + codec_util=_FakeCodecUtil(error_at="data-checksum"), + ) + + with pytest.raises(RuntimeError, match="invalid data checksum"): + verifier.verify_index(tmp_path) + + assert verifier._test_data_input.closed is True + + +def test_verify_cagra_index_rejects_missing_data_file(tmp_path): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ), + data_error=FileNotFoundError("_0.vcag"), + ) + + with pytest.raises(RuntimeError, match="cannot read '_0.vcag'"): + verifier.verify_index(tmp_path) + + +def test_file_signature_exposes_named_stat_fields(tmp_path): + data_path = tmp_path / "data.vcag" + data_path.write_bytes(b"data") + file_stat = data_path.stat() + + signature = pylucene_backend._file_signature(data_path) + + assert signature == pylucene_backend._FileSignature( + resolved_path=str(data_path.resolve()), + device=file_stat.st_dev, + inode=file_stat.st_ino, + size=file_stat.st_size, + modified_at_ns=file_stat.st_mtime_ns, + changed_at_ns=file_stat.st_ctime_ns, + ) + + +def test_cagra_data_checksum_is_cached_for_unchanged_file( + tmp_path, monkeypatch +): + codec_util = _FakeCodecUtil() + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput([], []), codec_util=codec_util + ) + (tmp_path / "_0.vcag").write_bytes(b"data") + data_ctime_ns = (tmp_path / "_0.vcag").stat().st_ctime_ns + monkeypatch.setattr( + pylucene_backend.time, + "time_ns", + lambda: ( + data_ctime_ns + pylucene_backend._CAGRA_CACHE_MIN_FILE_AGE_NS + 1 + ), + ) + fields = [ + pylucene_backend._CagraFieldMetadata( + field_number=1, + dimensions=32, + vector_count=4, + cagra_offset=57, + cagra_length=1000, + ) + ] + context = _cagra_data_context(verifier, tmp_path) + + for _ in range(2): + verifier._verify_cagra_data_file(context, fields) + + assert codec_util.checksum_calls == [verifier._test_data_input] + assert codec_util.retrieved_checksum_calls == [verifier._test_data_input] + + +def test_cagra_data_checksum_cache_invalidates_on_size_change(tmp_path): + codec_util = _FakeCodecUtil() + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput([], []), codec_util=codec_util + ) + data_path = tmp_path / "_0.vcag" + data_path.write_bytes(b"data") + fields = [ + pylucene_backend._CagraFieldMetadata( + field_number=1, + dimensions=32, + vector_count=4, + cagra_offset=57, + cagra_length=1000, + ) + ] + context = _cagra_data_context(verifier, tmp_path) + + verifier._verify_cagra_data_file(context, fields) + data_path.write_bytes(b"changed-size") + verifier._verify_cagra_data_file(context, fields) + + assert codec_util.checksum_calls == [ + verifier._test_data_input, + verifier._test_data_input, + ] + assert codec_util.retrieved_checksum_calls == [] + + +def test_cagra_data_checksum_cache_invalidates_same_size_restored_mtime( + tmp_path, +): + codec_util = _FakeCodecUtil() + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput([], []), codec_util=codec_util + ) + data_path = tmp_path / "_0.vcag" + data_path.write_bytes(b"data") + fields = [ + pylucene_backend._CagraFieldMetadata( + field_number=1, + dimensions=32, + vector_count=4, + cagra_offset=57, + cagra_length=1000, + ) + ] + context = _cagra_data_context(verifier, tmp_path) + + verifier._verify_cagra_data_file(context, fields) + original_stat = data_path.stat() + data_path.write_bytes(b"evil") + os.utime( + data_path, + ns=(original_stat.st_atime_ns, original_stat.st_mtime_ns), + ) + assert data_path.stat().st_mtime_ns == original_stat.st_mtime_ns + + verifier._verify_cagra_data_file(context, fields) + + assert codec_util.checksum_calls == [ + verifier._test_data_input, + verifier._test_data_input, + ] + assert codec_util.retrieved_checksum_calls == [] diff --git a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_cli_config.py b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_cli_config.py new file mode 100644 index 0000000000..e3c2ddedc0 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_cli_config.py @@ -0,0 +1,899 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""CLI and configuration tests for the PyLucene backend.""" + +from __future__ import annotations + +import csv +import json +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pytest +from click.testing import CliRunner + +import cuvs_bench.backends.pylucene as pylucene_backend +from cuvs_bench.backends import CppGoogleBenchmarkBackend, get_registry +from cuvs_bench.backends.pylucene import ( + PyLuceneBackend, + PyLuceneConfigLoader, + _SearchHit, +) +from cuvs_bench.backends.registry import list_config_loaders +from cuvs_bench.tests._pylucene_test_utils import ( + _CAGRA_CODEC, + _HNSW_CODEC, + _FakeRuntime, + _write_test_bin, +) + + +def _prepare_cli_dataset(dataset_path: Path) -> None: + training_vectors = np.asarray( + [ + [0.0, 0.0, 0.0, 0.0], + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + ], + dtype=np.float32, + ) + query_vectors = training_vectors[:2].copy() + groundtruth_neighbors = np.asarray([[0], [1]], dtype=np.int32) + dataset_dir = dataset_path / "test-dataset" + _write_test_bin(dataset_dir / "base.fbin", training_vectors) + _write_test_bin(dataset_dir / "query.fbin", query_vectors) + _write_test_bin( + dataset_dir / "groundtruth.neighbors.ibin", groundtruth_neighbors + ) + + +def _pylucene_cli_args( + backend_config: Path, + config_dir: Path, + dataset_path: Path, + *, + dry_run: bool = False, +) -> list[str]: + args = [ + "--backend-config", + str(backend_config), + "--dataset-configuration", + str(config_dir / "datasets" / "datasets.yaml"), + "--configuration", + str(config_dir / "algos" / "pylucene_test.yaml"), + "--dataset", + "test-dataset", + "--dataset-path", + str(dataset_path), + "--algorithms", + "pylucene_test", + "--groups", + "test", + "--batch-size", + "2", + "-k", + "1", + "-m", + "latency", + "--build", + "--search", + "--force", + ] + if dry_run: + args.append("--dry-run") + return args + + +@pytest.fixture +def config_dir(tmp_path): + (tmp_path / "datasets").mkdir() + (tmp_path / "datasets" / "datasets.yaml").write_text( + """\ +- name: test-dataset + base_file: test-dataset/base.fbin + query_file: test-dataset/query.fbin + groundtruth_neighbors_file: test-dataset/groundtruth.neighbors.ibin + distance: euclidean + dims: 4 +""" + ) + (tmp_path / "algos").mkdir() + (tmp_path / "algos" / "pylucene_test.yaml").write_text( + f"""\ +name: pylucene_test +groups: + base: + build: + codec: [{_HNSW_CODEC}, {_CAGRA_CODEC}] + search: {{}} + test: + build: + codec: [{_HNSW_CODEC}] + search: {{}} +""" + ) + (tmp_path / "algos" / "unrelated.yaml").write_text( + """\ +name: unrelated +groups: + base: + build: {} + search: {} +""" + ) + return tmp_path + + +def test_backend_and_loader_are_registered(): + assert get_registry().is_registered("pylucene") + assert list_config_loaders()["pylucene"] is PyLuceneConfigLoader + + +def test_result_file_ownership_is_explicit(): + assert PyLuceneBackend.orchestrator_persists_results is True + assert CppGoogleBenchmarkBackend.orchestrator_persists_results is False + + +def test_import_does_not_load_pylucene(): + subprocess.run( + [ + sys.executable, + "-c", + ( + "import sys; import cuvs_bench.backends; " + "assert 'lucene' not in sys.modules" + ), + ], + check=True, + ) + + +def test_cli_backend_config_supports_pylucene_dry_run(config_dir, tmp_path): + from cuvs_bench.run.__main__ import main as run_main + + backend_config = tmp_path / "backend.yaml" + backend_config.write_text("backend: pylucene\n") + dataset_path = tmp_path / "runtime-datasets" + result_path = dataset_path / "test-dataset" / "result" + assert not result_path.exists() + + result = CliRunner().invoke( + run_main, + _pylucene_cli_args( + backend_config, + config_dir, + dataset_path, + dry_run=True, + ), + ) + + assert result.exit_code == 0, result.output + assert "Would build PyLucene index" in result.output + assert "Would search PyLucene index" in result.output + assert not result_path.exists() + + +def test_cli_build_search_persists_metrics_and_build_join( + config_dir, tmp_path, monkeypatch +): + from cuvs_bench.run.__main__ import main as run_main + + backend_config = tmp_path / "backend.yaml" + backend_config.write_text("backend: pylucene\n") + dataset_path = tmp_path / "runtime-datasets" + _prepare_cli_dataset(dataset_path) + result_path = dataset_path / "test-dataset" / "result" + assert not result_path.exists() + + runtime = _FakeRuntime( + hits=[ + [_SearchHit(document_id=0, score=1.0)], + [_SearchHit(document_id=1, score=1.0)], + ] + ) + monkeypatch.setattr( + pylucene_backend._PyLuceneRuntime, + "create", + staticmethod(lambda _config: runtime), + ) + + result = CliRunner().invoke( + run_main, + _pylucene_cli_args(backend_config, config_dir, dataset_path), + ) + + assert result.exit_code == 0, result.output + index_name = f"pylucene_test_test.codec{_HNSW_CODEC}" + index_path = dataset_path / "test-dataset" / "index" / index_name + assert (index_path / "segments_1").is_file() + assert (index_path / pylucene_backend._HNSW_PROVENANCE_FILE).is_file() + expected_index_size = sum( + path.stat().st_size for path in index_path.rglob("*") if path.is_file() + ) + + build_json = result_path / "build" / "pylucene_test,test.json" + search_json = result_path / "search" / "pylucene_test,test,k1,bs2.json" + build_rows = json.loads(build_json.read_text())["benchmarks"] + search_rows = json.loads(search_json.read_text())["benchmarks"] + assert len(build_rows) == 1 + assert len(search_rows) == 1 + + build_row = build_rows[0] + assert build_row["name"] == f"{index_name}/build" + assert build_row["real_time"] > 0.0 + assert build_row["index_size"] == expected_index_size + assert build_row["codec"] == _HNSW_CODEC + assert build_row["writer_path"] == "gpu-hnsw" + + search_row = search_rows[0] + assert search_row["name"] == f"{index_name}/search" + assert search_row["Recall"] == pytest.approx(1.0) + assert search_row["items_per_second"] == pytest.approx(2000.0) + assert search_row["search_time_ms"] == pytest.approx(1.0) + assert search_row["real_time"] == pytest.approx(1.0) + assert search_row["Latency"] == pytest.approx(0.001) + assert search_row["p50"] == pytest.approx(1.0) + + raw_csv = result_path / "search" / "pylucene_test,test,k1,bs2,raw.csv" + with raw_csv.open(newline="") as file: + csv_rows = list(csv.DictReader(file)) + assert len(csv_rows) == 1 + csv_row = csv_rows[0] + assert csv_row["index_name"] == index_name + assert float(csv_row["recall"]) == pytest.approx(1.0) + assert float(csv_row["throughput"]) == pytest.approx(2000.0) + assert float(csv_row["latency"]) == pytest.approx(0.001) + assert float(csv_row["build time"]) > 0.0 + assert (raw_csv.parent / "pylucene_test,test,k1,bs2,latency.csv").is_file() + assert ( + raw_csv.parent / "pylucene_test,test,k1,bs2,throughput.csv" + ).is_file() + + +def test_cli_returns_nonzero_and_persists_build_failure( + config_dir, tmp_path, monkeypatch +): + from cuvs_bench.run.__main__ import main as run_main + + backend_config = tmp_path / "backend.yaml" + backend_config.write_text("backend: pylucene\n") + dataset_path = tmp_path / "runtime-datasets" + _prepare_cli_dataset(dataset_path) + runtime = _FakeRuntime() + runtime.build_error = RuntimeError("intentional build failure") + monkeypatch.setattr( + pylucene_backend._PyLuceneRuntime, + "create", + staticmethod(lambda _config: runtime), + ) + + result = CliRunner().invoke( + run_main, + _pylucene_cli_args(backend_config, config_dir, dataset_path), + ) + + assert result.exit_code != 0 + assert "intentional build failure" in result.output + build_json = ( + dataset_path + / "test-dataset" + / "result" + / "build" + / "pylucene_test,test.json" + ) + build_rows = json.loads(build_json.read_text())["benchmarks"] + assert len(build_rows) == 1 + assert build_rows[0]["success"] is False + assert "intentional build failure" in build_rows[0]["error_message"] + assert not (dataset_path / "test-dataset" / "result" / "search").exists() + + +def test_config_loader_expands_codecs_and_forwards_runtime_config(config_dir): + loader = PyLuceneConfigLoader(config_path=config_dir) + dataset_config, configs = loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_test", + groups="base", + cuvs_java_jar="/artifacts/cuvs-java.jar", + cuvs_lucene_jar="/artifacts/cuvs-lucene.jar", + java_library_path="/native", + jvm_args=["-Xms1g"], + ) + + assert dataset_config.distance == "euclidean" + assert len(configs) == 2 + assert {config.indexes[0].build_param["codec"] for config in configs} == { + _HNSW_CODEC, + _CAGRA_CODEC, + } + for config in configs: + assert config.indexes[0].search_params == [{}] + assert config.backend_config["cuvs_java_jar"].endswith("cuvs-java.jar") + assert config.backend_config["cuvs_lucene_jar"].endswith( + "cuvs-lucene.jar" + ) + assert config.backend_config["java_library_path"] == "/native" + assert config.backend_config["jvm_args"] == ["-Xms1g"] + assert config.backend_config["requires_gpu"] is True + assert config.output_filename == ( + "pylucene_test,base", + "pylucene_test,base,k10,bs10000", + ) + + +@pytest.mark.parametrize( + ("build_yaml", "search_yaml", "error"), + [ + ( + f"codec: [{_HNSW_CODEC}]\n ignored: [1]", + "{}", + "Unsupported PyLucene build parameter.*ignored", + ), + ( + f"codec: [{_HNSW_CODEC}]", + "ignored: [1]", + "does not expose cuVS-specific search parameters", + ), + ], +) +def test_config_loader_rejects_unsupported_parameters( + config_dir, tmp_path, build_yaml, search_yaml, error +): + custom_config = tmp_path / "unsupported-parameters.yaml" + custom_config.write_text( + f"""\ +backend: pylucene +name: pylucene_unsupported +groups: + test: + build: + {build_yaml} + search: + {search_yaml} +""" + ) + + with pytest.raises(ValueError, match=error): + PyLuceneConfigLoader(config_path=config_dir).load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_unsupported", + groups="test", + algorithm_configuration=str(custom_config), + ) + + +def test_config_loader_scopes_artifact_identity_by_dataset_subset(config_dir): + loader = PyLuceneConfigLoader(config_path=config_dir) + identities = [] + + for subset_size, subset_scope in ( + (None, None), + (2, "subset2"), + (3, "subset3"), + ): + dataset_config, configs = loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_test", + groups="test", + subset_size=subset_size, + count=1, + batch_size=2, + ) + + assert dataset_config.subset_size == subset_size + assert len(configs) == 1 + config = configs[0] + index_prefix = "pylucene_test_test" + result_stem = "pylucene_test,test" + if subset_scope is not None: + index_prefix = f"{index_prefix}.{subset_scope}" + result_stem = f"{result_stem},{subset_scope}" + index_name = f"{index_prefix}.codec{_HNSW_CODEC}" + assert config.index_name == index_name + assert config.index_path == ( + Path("/datasets/test-dataset/index") / index_name + ) + assert config.output_filename == ( + result_stem, + f"{result_stem},k1,bs2", + ) + assert all(Path(stem).name == stem for stem in config.output_filename) + identities.append( + (config.index_name, config.index_path, config.output_filename) + ) + + assert len(set(identities)) == 3 + + +@pytest.mark.parametrize("subset_size", [0, -1, True, "../other"]) +def test_config_loader_rejects_unsafe_subset_identity(config_dir, subset_size): + loader = PyLuceneConfigLoader(config_path=config_dir) + + with pytest.raises(ValueError, match="positive integer"): + loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_test", + groups="test", + subset_size=subset_size, + ) + + +@pytest.mark.parametrize( + ("backend_declaration", "algorithm_name"), + [ + ("", "pylucene_custom"), + ("backend: pylucene\n", "custom_lucene"), + ], + ids=["parsed-name", "declared-backend"], +) +def test_config_loader_discovers_custom_filename_by_parsed_identity( + config_dir, + tmp_path, + backend_declaration, + algorithm_name, +): + custom_config = tmp_path / "arbitrary-custom-name.yaml" + custom_config.write_text( + f"""\ +{backend_declaration}name: {algorithm_name} +groups: + test: + build: + codec: [{_HNSW_CODEC}] + search: {{}} +""" + ) + loader = PyLuceneConfigLoader(config_path=config_dir) + + _, configs = loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms=algorithm_name, + groups="test", + algorithm_configuration=str(custom_config), + ) + + assert len(configs) == 1 + assert configs[0].indexes[0].algo == algorithm_name + assert configs[0].index_name.startswith(f"{algorithm_name}_test.") + + +def test_config_loader_excludes_explicit_other_backend(config_dir, tmp_path): + custom_config = tmp_path / "misleading-pylucene-name.yaml" + custom_config.write_text( + f"""\ +backend: cpp_gbench +name: pylucene_other_backend +groups: + test: + build: + codec: [{_HNSW_CODEC}] + search: {{}} +""" + ) + loader = PyLuceneConfigLoader(config_path=config_dir) + + with pytest.raises( + ValueError, + match="Unknown PyLucene algorithm selector.*pylucene_other_backend", + ): + loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_other_backend", + groups="test", + algorithm_configuration=str(custom_config), + ) + + +def test_config_loader_honors_algorithm_and_group_filters(config_dir): + loader = PyLuceneConfigLoader(config_path=config_dir) + + _, configs = loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_test", + groups="test", + ) + assert len(configs) == 1 + assert configs[0].index_name.startswith("pylucene_test_test") + + +def test_algorithm_config_discovery_is_deterministic(tmp_path): + config_path = tmp_path / "config" + bundled_algorithms = config_path / "algos" + bundled_algorithms.mkdir(parents=True) + bundled_zeta = bundled_algorithms / "zeta.yaml" + bundled_alpha = bundled_algorithms / "alpha.yml" + bundled_zeta.touch() + bundled_alpha.touch() + (bundled_algorithms / "ignored.txt").touch() + + custom_algorithms = tmp_path / "custom" + custom_algorithms.mkdir() + custom_zeta = custom_algorithms / "zeta.yml" + custom_alpha = custom_algorithms / "alpha.yaml" + custom_zeta.touch() + custom_alpha.touch() + + files = PyLuceneConfigLoader( + config_path=config_path + ).gather_algorithm_configs(config_path, str(custom_algorithms)) + + assert files == [ + str(bundled_alpha), + str(bundled_zeta), + str(custom_alpha), + str(custom_zeta), + ] + + +def test_duplicate_algorithm_config_uses_last_definition_and_position( + config_dir, monkeypatch +): + loader = PyLuceneConfigLoader(config_path=config_dir) + config_files = ["alpha-original", "beta", "alpha-override"] + configs_by_file = { + "alpha-original": { + "name": "pylucene_alpha", + "groups": {"original": {}}, + }, + "beta": { + "name": "pylucene_beta", + "groups": {"base": {}}, + }, + "alpha-override": { + "name": "pylucene_alpha", + "groups": {"override": {}}, + }, + } + monkeypatch.setattr(loader, "load_yaml_file", configs_by_file.__getitem__) + + algorithm_configs = loader._load_algorithm_configs(config_files) + + assert list(algorithm_configs) == ["pylucene_beta", "pylucene_alpha"] + assert algorithm_configs["pylucene_alpha"] == {"override": {}} + + +def test_config_loader_unions_global_and_algorithm_specific_groups(config_dir): + loader = PyLuceneConfigLoader(config_path=config_dir) + + _, configs = loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_test", + groups="base", + algo_groups="pylucene_test.test", + ) + + assert len(configs) == 3 + assert ( + sum( + config.index_name.startswith("pylucene_test_test.") + for config in configs + ) + == 1 + ) + + +def test_config_loader_selects_only_explicit_algorithm_group(config_dir): + loader = PyLuceneConfigLoader(config_path=config_dir) + + _, configs = loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algo_groups="pylucene_test.test", + ) + + assert len(configs) == 1 + assert configs[0].index_name.startswith("pylucene_test_test.") + + +@pytest.mark.parametrize( + ("selectors", "expected_groups"), + [ + ( + {}, + ["pylucene_alpha,base", "pylucene_beta,base"], + ), + ( + {"algorithms": "pylucene_beta,pylucene_alpha"}, + ["pylucene_alpha,base", "pylucene_beta,base"], + ), + ( + {"groups": "shared"}, + ["pylucene_alpha,shared", "pylucene_beta,shared"], + ), + ( + { + "algorithms": "pylucene_alpha,pylucene_beta", + "groups": "alpha_only,shared", + }, + [ + "pylucene_alpha,shared", + "pylucene_alpha,alpha_only", + "pylucene_beta,shared", + ], + ), + ( + {"algo_groups": "pylucene_beta.beta_only"}, + ["pylucene_beta,beta_only"], + ), + ( + { + "algorithms": "pylucene_alpha", + "algo_groups": "pylucene_beta.beta_only", + }, + ["pylucene_alpha,base", "pylucene_beta,beta_only"], + ), + ( + { + "groups": "alpha_only", + "algo_groups": "pylucene_beta.beta_only", + }, + ["pylucene_alpha,alpha_only", "pylucene_beta,beta_only"], + ), + ( + { + "algorithms": "pylucene_beta", + "groups": "shared", + "algo_groups": "pylucene_alpha.alpha_only", + }, + ["pylucene_alpha,alpha_only", "pylucene_beta,shared"], + ), + ( + { + "algorithms": "pylucene_alpha", + "groups": "shared", + "algo_groups": "pylucene_alpha.shared", + }, + ["pylucene_alpha,shared"], + ), + ( + { + "algo_groups": ( + "pylucene_beta.beta_only,pylucene_alpha.alpha_only" + ) + }, + ["pylucene_alpha,alpha_only", "pylucene_beta,beta_only"], + ), + ( + { + "algorithms": " pylucene_beta, pylucene_beta ", + "groups": " base,base ", + }, + ["pylucene_beta,base"], + ), + ], + ids=[ + "defaults", + "algorithm", + "group", + "algorithm-and-group", + "explicit-pair", + "algorithm-plus-explicit-pair", + "group-plus-explicit-pair", + "all-selectors", + "overlapping-selectors", + "explicit-pair-order-follows-config", + "duplicate-selectors", + ], +) +def test_config_loader_combines_global_and_explicit_selectors( + config_dir, monkeypatch, selectors, expected_groups +): + alpha_config = config_dir / "algos" / "pylucene_alpha.yaml" + alpha_config.write_text( + f"""\ +name: pylucene_alpha +groups: + base: + build: + codec: [{_HNSW_CODEC}] + search: {{}} + shared: + build: + codec: [{_HNSW_CODEC}] + search: {{}} + alpha_only: + build: + codec: [{_HNSW_CODEC}] + search: {{}} +""" + ) + beta_config = config_dir / "algos" / "pylucene_beta.yaml" + beta_config.write_text( + f"""\ +name: pylucene_beta +groups: + base: + build: + codec: [{_HNSW_CODEC}] + search: {{}} + shared: + build: + codec: [{_HNSW_CODEC}] + search: {{}} + beta_only: + build: + codec: [{_HNSW_CODEC}] + search: {{}} +""" + ) + loader = PyLuceneConfigLoader(config_path=config_dir) + monkeypatch.setattr( + loader, + "gather_algorithm_configs", + lambda *_args: [str(alpha_config), str(beta_config)], + ) + + _, configs = loader.load( + dataset="test-dataset", + dataset_path="/datasets", + **selectors, + ) + + selected_groups = [config.output_filename[0] for config in configs] + assert selected_groups == expected_groups + + +def test_algorithm_selection_resolution_preserves_inputs_and_source_order(): + alpha_base = {"build": {"codec": [_HNSW_CODEC]}} + beta_base = {"build": {"codec": [_HNSW_CODEC]}} + algorithm_configs = { + "pylucene_beta": {"base": beta_base}, + "pylucene_alpha": {"base": alpha_base}, + } + original_items = [ + (algorithm, list(groups.items())) + for algorithm, groups in algorithm_configs.items() + ] + selection = pylucene_backend._AlgorithmSelection( + algorithms=frozenset({"pylucene_alpha", "pylucene_beta"}), + groups=None, + explicit_groups=frozenset(), + ) + + selected = selection.resolve(algorithm_configs) + + assert [(group.algorithm, group.group) for group in selected] == [ + ("pylucene_beta", "base"), + ("pylucene_alpha", "base"), + ] + assert selected[0].configuration is beta_base + assert selected[1].configuration is alpha_base + assert [ + (algorithm, list(groups.items())) + for algorithm, groups in algorithm_configs.items() + ] == original_items + + +def test_algorithm_selection_rejects_group_outside_selected_algorithms(): + selection = pylucene_backend._AlgorithmSelection( + algorithms=frozenset({"pylucene_alpha"}), + groups=frozenset({"beta_only"}), + explicit_groups=frozenset(), + ) + algorithm_configs = { + "pylucene_alpha": {"base": {}}, + "pylucene_beta": {"beta_only": {}}, + } + + with pytest.raises( + ValueError, + match="Unknown PyLucene group selector\\(s\\): beta_only", + ): + selection.resolve(algorithm_configs) + + +def test_algorithm_selection_reports_explicit_errors_deterministically(): + selection = pylucene_backend._AlgorithmSelection( + algorithms=None, + groups=None, + explicit_groups=frozenset( + { + ("pylucene_zeta", "base"), + ("pylucene_beta", "base"), + } + ), + ) + + with pytest.raises(ValueError) as error: + selection.resolve({"pylucene_alpha": {"base": {}}}) + + assert str(error.value) == ( + "Unknown PyLucene algorithm in --algo-groups: pylucene_beta" + ) + + +def test_config_loader_rejects_malformed_algorithm_group(config_dir): + loader = PyLuceneConfigLoader(config_path=config_dir) + + with pytest.raises(ValueError, match="."): + loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algo_groups="pylucene_test", + ) + + +@pytest.mark.parametrize( + ("selectors", "error"), + [ + ( + {"algorithms": "not-pylucene", "groups": "base"}, + "Unknown PyLucene algorithm selector.*not-pylucene", + ), + ( + {"algorithms": "pylucene_test", "groups": "missing"}, + "Unknown PyLucene group selector.*missing", + ), + ( + { + "algorithms": "pylucene_test", + "groups": "base", + "algo_groups": "not-pylucene.test", + }, + "Unknown PyLucene algorithm in --algo-groups.*not-pylucene", + ), + ( + { + "algorithms": "pylucene_test", + "groups": "base", + "algo_groups": "pylucene_test.missing", + }, + "Unknown PyLucene group for pylucene_test.*missing", + ), + ], + ids=[ + "algorithm", + "global-group", + "algo-group-algorithm", + "algo-group-group", + ], +) +def test_config_loader_rejects_unknown_selectors(config_dir, selectors, error): + loader = PyLuceneConfigLoader(config_path=config_dir) + + with pytest.raises(ValueError, match=error): + loader.load( + dataset="test-dataset", + dataset_path="/datasets", + **selectors, + ) + + +@pytest.mark.parametrize( + ("algorithm", "expected_codecs"), + [ + ( + "pylucene_cuvs_hnsw", + { + "Lucene101AcceleratedHNSWCodec", + "Lucene101AcceleratedHNSWBaseLayerCodec", + "Lucene101AcceleratedHNSWMultiLayerCodec", + }, + ), + ("pylucene_cuvs_cagra", {_CAGRA_CODEC}), + ], +) +def test_shipped_algorithm_configs_load(algorithm, expected_codecs): + _, configs = PyLuceneConfigLoader().load( + dataset="test-data", + dataset_path="/datasets", + algorithms=algorithm, + groups="base", + ) + + assert { + config.indexes[0].build_param["codec"] for config in configs + } == expected_codecs diff --git a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_integration.py b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_integration.py new file mode 100644 index 0000000000..af886641f7 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_integration.py @@ -0,0 +1,704 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Opt-in integration coverage for the PyLucene benchmark backend.""" + +import csv +import importlib +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pytest + +from cuvs_bench._bin_format import write_bin_header +from cuvs_bench.backends._utils import compute_recall +from cuvs_bench.backends.base import Dataset +from cuvs_bench.backends.pylucene import ( + _CAGRA_PROVENANCE_FILE, + PyLuceneBackend, + _HNSW_PROVENANCE_FILE, + _PyLuceneRuntime, + _ResolvedCodec, +) +from cuvs_bench.orchestrator.config_loaders import IndexConfig + +pytestmark = [ + pytest.mark.pylucene, + pytest.mark.filterwarnings( + "ignore:builtin type .* has no __module__ attribute:DeprecationWarning" + ), +] + +_OPT_IN_ENV = "CUVS_BENCH_PYLUCENE_INTEGRATION" +_CUVS_JAVA_JAR_ENV = "CUVS_LUCENE_CUVS_JAVA_JAR" +_CUVS_LUCENE_JAR_ENV = "CUVS_LUCENE_JAR" +_CAGRA_CODEC = "CuVS2510GPUSearchCodec" +_HNSW_CODEC = "Lucene101AcceleratedHNSWCodec" + + +def _write_test_bin(path: Path, data: np.ndarray) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("wb") as file: + write_bin_header(file, data.shape[0], data.shape[1]) + np.ascontiguousarray(data).tofile(file) + + +def _required_jar(env_name: str) -> Path: + configured = os.environ.get(env_name) + if not configured: + pytest.fail(f"{env_name} must point to the required runtime jar") + + jar = Path(configured).expanduser() + if not jar.is_file(): + pytest.fail(f"{env_name} does not point to an existing file: {jar}") + return jar.resolve() + + +def _assert_cagra_verification( + metadata, expected_vector_count, expected_dimensions +): + verification = metadata["cagra_verification"] + assert verification["status"] == "cagra-only" + assert verification["segment_count"] >= 1 + assert verification["field_count"] >= 1 + assert verification["vector_count"] == expected_vector_count + assert verification["dimensions"] == expected_dimensions + + +def _assert_cagra_provenance( + metadata, expected_vector_count, expected_dimensions +): + assert metadata["cagra_provenance"] == { + "status": "gpu-cagra-provenance", + "schema_version": 1, + "codec": _CAGRA_CODEC, + "writer_path": "gpu-cagra", + "vector_count": expected_vector_count, + "dimensions": expected_dimensions, + "commit_file_count": 1, + } + + +def _commit_one_deletion(backend, index_path): + runtime = backend._get_runtime() + runtime.attach_current_thread() + directory = runtime.FSDirectory.open(runtime.Paths.get(str(index_path))) + writer = None + reader = None + try: + writer_config = runtime.IndexWriterConfig() + writer_config.setOpenMode(runtime.IndexWriterConfig.OpenMode.APPEND) + writer = runtime.IndexWriter(directory, writer_config) + reader = runtime.DirectoryReader.open(writer) + assert int(writer.tryDeleteDocument(reader, 0)) >= 0 + writer.commit() + finally: + try: + if reader is not None: + reader.close() + finally: + try: + if writer is not None: + writer.close() + finally: + directory.close() + + +def _assert_hnsw_verification( + metadata, codec, expected_vector_count, expected_dimensions +): + verification = metadata["hnsw_verification"] + assert verification == { + "status": "gpu-hnsw-provenance", + "schema_version": 1, + "codec": codec, + "writer_path": "gpu-hnsw", + "vector_count": expected_vector_count, + "dimensions": expected_dimensions, + "commit_file_count": 1, + } + + +@pytest.fixture(scope="module") +def pylucene_runtime_config(): + """Resolve the explicitly configured PyLucene/cuVS runtime.""" + if os.environ.get(_OPT_IN_ENV) != "1": + pytest.skip(f"set {_OPT_IN_ENV}=1 to run PyLucene integration tests") + + cuvs_java_jar = _required_jar(_CUVS_JAVA_JAR_ENV) + cuvs_lucene_jar = _required_jar(_CUVS_LUCENE_JAR_ENV) + + java_library_path = os.environ.get("JAVA_LIBRARY_PATH") or os.environ.get( + "LD_LIBRARY_PATH" + ) + if not java_library_path: + pytest.fail( + "JAVA_LIBRARY_PATH or LD_LIBRARY_PATH must provide the native " + "cuVS runtime libraries" + ) + + try: + importlib.import_module("lucene") + except ImportError as exc: + pytest.fail(f"PyLucene is not importable: {exc}") + + return { + "cuvs_java_jar": str(cuvs_java_jar), + "cuvs_lucene_jar": str(cuvs_lucene_jar), + "java_library_path": java_library_path, + } + + +@pytest.mark.parametrize( + ( + "algo", + "codec", + "expected_writer_telemetry", + "expected_index_suffix", + ), + [ + pytest.param( + "pylucene_cuvs_hnsw", + "Lucene101AcceleratedHNSWCodec", + { + "writerPath": "gpu-hnsw", + "hnswLayers": "1", + "cagraGraphBuildAlgo": "NN_DESCENT", + "cagraGraphDegree": "64", + "cagraIntermediateGraphDegree": "128", + }, + ".vex", + id="accelerated-hnsw", + ), + pytest.param( + "pylucene_cuvs_hnsw", + "Lucene101AcceleratedHNSWBaseLayerCodec", + { + "writerPath": "gpu-hnsw", + "hnswLayers": "1", + "cagraGraphBuildAlgo": "NN_DESCENT", + "cagraGraphDegree": "32", + "cagraIntermediateGraphDegree": "64", + }, + ".vex", + id="accelerated-hnsw-base-layer", + ), + pytest.param( + "pylucene_cuvs_hnsw", + "Lucene101AcceleratedHNSWMultiLayerCodec", + { + "writerPath": "gpu-hnsw", + "hnswLayers": "3", + "cagraGraphBuildAlgo": "NN_DESCENT", + "cagraGraphDegree": "32", + "cagraIntermediateGraphDegree": "64", + }, + ".vex", + id="accelerated-hnsw-multi-layer", + ), + pytest.param( + "pylucene_cuvs_cagra", + "CuVS2510GPUSearchCodec", + { + "writerPath": "gpu-cagra", + "cagraStrategy": "HEURISTIC", + "cagraGraphBuildAlgo": "NN_DESCENT", + "cagraGraphDegree": "64", + "cagraIntermediateGraphDegree": "128", + }, + ".vcag", + id="cagra", + ), + ], +) +def test_build_and_search_with_real_pylucene_runtime( + tmp_path, + pylucene_runtime_config, + algo, + codec, + expected_writer_telemetry, + expected_index_suffix, +): + """Build and search through Lucene's public codec and query SPI.""" + rng = np.random.default_rng(1907) + training_vectors = rng.standard_normal((512, 32)).astype(np.float32) + query_ids = np.asarray([0, 137, 259, 511], dtype=np.int64) + query_vectors = training_vectors[query_ids].copy() + k = 5 + + squared_distances = np.sum( + (query_vectors[:, np.newaxis, :] - training_vectors[np.newaxis, :, :]) + ** 2, + axis=2, + ) + groundtruth_neighbors = np.argsort(squared_distances, axis=1)[:, :k] + groundtruth_distances = np.take_along_axis( + squared_distances, groundtruth_neighbors, axis=1 + ) + dataset = Dataset( + name="pylucene-integration", + training_vectors=training_vectors, + query_vectors=query_vectors, + groundtruth_neighbors=groundtruth_neighbors.astype(np.int32), + groundtruth_distances=groundtruth_distances.astype(np.float32), + distance_metric="euclidean", + ) + + index_path = tmp_path / f"{algo}-index" + index = IndexConfig( + name=f"{algo}-integration", + algo=algo, + build_param={"codec": codec}, + search_params=[{}], + file=str(index_path), + ) + backend = PyLuceneBackend( + { + "name": index.name, + "algo": algo, + "codec": codec, + **pylucene_runtime_config, + } + ) + + try: + build_result = backend.build(dataset, [index], force=True) + assert build_result.success, build_result.error_message + assert build_result.build_time_seconds > 0 + assert build_result.index_size_bytes > 0 + assert build_result.metadata["codec"] == codec + assert build_result.metadata["pylucene_version"] != "unknown" + assert ( + build_result.metadata["writer_path"] + == expected_writer_telemetry["writerPath"] + ) + assert ( + build_result.metadata["writer_telemetry"] + == expected_writer_telemetry + ) + assert any( + path.suffix == expected_index_suffix + for path in index_path.iterdir() + ) + if codec == _CAGRA_CODEC: + _assert_cagra_provenance( + build_result.metadata, + training_vectors.shape[0], + training_vectors.shape[1], + ) + _assert_cagra_verification( + build_result.metadata, + training_vectors.shape[0], + training_vectors.shape[1], + ) + assert (index_path / _CAGRA_PROVENANCE_FILE).is_file() + else: + _assert_hnsw_verification( + build_result.metadata, + codec, + training_vectors.shape[0], + training_vectors.shape[1], + ) + assert (index_path / _HNSW_PROVENANCE_FILE).is_file() + + reused_result = backend.build(dataset, [index], force=False) + assert reused_result.success, reused_result.error_message + assert reused_result.metadata["skipped"] is True + if codec == _CAGRA_CODEC: + _assert_cagra_provenance( + reused_result.metadata, + training_vectors.shape[0], + training_vectors.shape[1], + ) + _assert_cagra_verification( + reused_result.metadata, + training_vectors.shape[0], + training_vectors.shape[1], + ) + else: + _assert_hnsw_verification( + reused_result.metadata, + codec, + training_vectors.shape[0], + training_vectors.shape[1], + ) + + search_result = backend.search(dataset, [index], k=k, batch_size=2) + assert search_result.success, search_result.error_message + assert search_result.neighbors.shape == (query_vectors.shape[0], k) + assert search_result.distances.shape == (query_vectors.shape[0], k) + np.testing.assert_array_equal(search_result.neighbors[:, 0], query_ids) + assert np.all(np.isfinite(search_result.distances)) + recall = compute_recall( + search_result.neighbors, groundtruth_neighbors, k + ) + assert recall >= 0.75 + returned_distances = np.take_along_axis( + squared_distances, search_result.neighbors, axis=1 + ) + np.testing.assert_allclose( + search_result.distances, + returned_distances, + rtol=1e-5, + atol=1e-5, + ) + assert np.all(np.diff(search_result.distances, axis=1) >= -1e-6) + assert search_result.search_time_ms > 0 + assert search_result.latency_seconds is not None + assert search_result.latency_seconds > 0 + assert search_result.queries_per_second > 0 + assert search_result.metadata["codec"] == codec + assert search_result.metadata["pylucene_version"] != "unknown" + assert search_result.metadata["num_batches"] == 2 + assert search_result.metadata["mode"] == "latency" + assert set(search_result.latency_percentiles) == {"p50", "p95", "p99"} + if codec == _CAGRA_CODEC: + _assert_cagra_provenance( + search_result.metadata, + training_vectors.shape[0], + training_vectors.shape[1], + ) + _assert_cagra_verification( + search_result.metadata, + training_vectors.shape[0], + training_vectors.shape[1], + ) + + wrong_count_dataset = Dataset( + name="pylucene-integration-wrong-count", + training_vectors=training_vectors[:-1], + query_vectors=query_vectors, + distance_metric="euclidean", + ) + rejected = backend.build(wrong_count_dataset, [index], force=False) + assert not rejected.success + assert "vector count does not match the dataset: 512 != 511" in ( + rejected.error_message + ) + rejected = backend.search(wrong_count_dataset, [index], k=k) + assert not rejected.success + assert "vector count does not match the dataset: 512 != 511" in ( + rejected.error_message + ) + + data_path = next(index_path.glob("*.vcag")) + original_data = data_path.read_bytes() + + data_path.unlink() + rejected = backend.search(dataset, [index], k=k) + assert not rejected.success + assert "cannot read" in rejected.error_message + data_path.write_bytes(original_data) + + data_path.write_bytes(original_data[:-1]) + rejected = backend.search(dataset, [index], k=k) + assert not rejected.success + assert "do not exactly cover" in rejected.error_message + data_path.write_bytes(original_data) + + corrupted_data = bytearray(original_data) + corrupted_data[len(corrupted_data) // 2] ^= 0xFF + original_stat = data_path.stat() + data_path.write_bytes(corrupted_data) + os.utime( + data_path, + ns=(original_stat.st_atime_ns, original_stat.st_mtime_ns), + ) + rejected = backend.search(dataset, [index], k=k) + assert not rejected.success + assert "checksum" in rejected.error_message.lower() + data_path.write_bytes(original_data) + + deleted_index_path = tmp_path / f"{algo}-deleted-index" + shutil.copytree(index_path, deleted_index_path) + deleted_index = IndexConfig( + name=f"{algo}-deleted-integration", + algo=algo, + build_param={"codec": codec}, + search_params=[{}], + file=str(deleted_index_path), + ) + _commit_one_deletion(backend, deleted_index_path) + with pytest.raises(RuntimeError, match="committed deletions"): + backend._get_runtime().verify_cagra_index(deleted_index_path) + rejected = backend.search(dataset, [deleted_index], k=k) + assert not rejected.success + assert "does not match the current Lucene commit" in ( + rejected.error_message + ) + + metadata_path = next(index_path.glob("*.vemc")) + contents = bytearray(metadata_path.read_bytes()) + contents[-1] ^= 0xFF + metadata_path.write_bytes(contents) + with pytest.raises(RuntimeError, match="metadata format v0"): + backend._get_runtime().verify_cagra_index(index_path) + else: + assert "cagra_verification" not in search_result.metadata + _assert_hnsw_verification( + search_result.metadata, + codec, + training_vectors.shape[0], + training_vectors.shape[1], + ) + if codec == _HNSW_CODEC: + manifest_path = index_path / _HNSW_PROVENANCE_FILE + original_manifest = manifest_path.read_bytes() + segment_path = next(index_path.glob("segments_*")) + original_segment = segment_path.read_bytes() + + manifest_path.unlink() + rejected = backend.search(dataset, [index], k=k) + assert not rejected.success + assert "provenance manifest is missing" in ( + rejected.error_message + ) + manifest_path.write_bytes(original_manifest) + + manifest_path.write_text("{") + rejected = backend.search(dataset, [index], k=k) + assert not rejected.success + assert "provenance manifest cannot be read" in ( + rejected.error_message + ) + manifest_path.write_bytes(original_manifest) + + segment_path.write_bytes(original_segment + b"stale") + rejected = backend.search(dataset, [index], k=k) + assert not rejected.success + assert "does not match the current Lucene commit" in ( + rejected.error_message + ) + segment_path.write_bytes(original_segment) + + payload = json.loads(original_manifest) + payload["codec"] = "Lucene101AcceleratedHNSWBaseLayerCodec" + manifest_path.write_text(json.dumps(payload)) + rejected = backend.search(dataset, [index], k=k) + assert not rejected.success + assert "does not match the requested codec" in ( + rejected.error_message + ) + manifest_path.write_bytes(original_manifest) + finally: + backend.cleanup() + + +def test_real_verifier_rejects_cagra_brute_force_fallback( + tmp_path, pylucene_runtime_config +): + """Reject cuVS-Lucene's real one-vector fallback before search.""" + runtime = _PyLuceneRuntime.create(pylucene_runtime_config) + vectors = np.ones((1, 32), dtype=np.float32) + index_path = tmp_path / "cagra-fallback-index" + index_path.mkdir() + java_codec = runtime.resolve_codec(_CAGRA_CODEC) + telemetry = runtime.codec_telemetry(_CAGRA_CODEC, java_codec) + assert telemetry["writerPath"] == "gpu-cagra" + runtime.build_index( + index_path, + vectors, + _ResolvedCodec( + codec_name=_CAGRA_CODEC, + java_codec=java_codec, + telemetry=telemetry, + writer_path=telemetry["writerPath"], + ), + ) + + with pytest.raises(RuntimeError, match="persisted brute-force"): + runtime.verify_cagra_index(index_path, expected_vector_count=1) + + dataset = Dataset( + name="pylucene-cagra-fallback", + training_vectors=vectors, + query_vectors=vectors.copy(), + distance_metric="euclidean", + ) + index = IndexConfig( + name="pylucene-cagra-fallback", + algo="pylucene_cuvs_cagra", + build_param={"codec": _CAGRA_CODEC}, + search_params=[{}], + file=str(index_path), + ) + backend = PyLuceneBackend( + { + "name": index.name, + "algo": index.algo, + "codec": _CAGRA_CODEC, + **pylucene_runtime_config, + } + ) + backend._runtime = runtime + try: + result = backend.search(dataset, [index], k=1) + assert not result.success + assert "CAGRA provenance manifest is missing" in result.error_message + finally: + backend.cleanup() + + +def test_cli_build_and_search_with_real_pylucene_runtime( + tmp_path, pylucene_runtime_config +): + """Exercise the documented module CLI in a fresh PyLucene process.""" + rng = np.random.default_rng(174) + training_vectors = rng.standard_normal((512, 32)).astype(np.float32) + query_ids = np.asarray([0, 137, 259, 511], dtype=np.int64) + query_vectors = training_vectors[query_ids].copy() + k = 5 + squared_distances = np.sum( + (query_vectors[:, np.newaxis, :] - training_vectors[np.newaxis, :, :]) + ** 2, + axis=2, + ) + groundtruth_neighbors = np.argsort(squared_distances, axis=1)[:, :k] + groundtruth_distances = np.take_along_axis( + squared_distances, groundtruth_neighbors, axis=1 + ) + + dataset_name = "pylucene-cli-integration" + dataset_path = tmp_path / "datasets" + dataset_dir = dataset_path / dataset_name + _write_test_bin(dataset_dir / "base.fbin", training_vectors) + _write_test_bin(dataset_dir / "query.fbin", query_vectors) + _write_test_bin( + dataset_dir / "groundtruth.neighbors.ibin", + groundtruth_neighbors.astype(np.int32), + ) + _write_test_bin( + dataset_dir / "groundtruth.distances.fbin", + groundtruth_distances.astype(np.float32), + ) + + dataset_config = tmp_path / "datasets.yaml" + dataset_config.write_text( + json.dumps( + [ + { + "name": dataset_name, + "base_file": f"{dataset_name}/base.fbin", + "query_file": f"{dataset_name}/query.fbin", + "groundtruth_neighbors_file": ( + f"{dataset_name}/groundtruth.neighbors.ibin" + ), + "groundtruth_distances_file": ( + f"{dataset_name}/groundtruth.distances.fbin" + ), + "distance": "euclidean", + "dims": training_vectors.shape[1], + } + ] + ) + ) + backend_config = tmp_path / "pylucene-backend.yaml" + backend_config.write_text( + json.dumps({"backend": "pylucene", **pylucene_runtime_config}) + ) + + environment = os.environ.copy() + source_root = str(Path(__file__).resolve().parents[2]) + environment["PYTHONPATH"] = os.pathsep.join( + value + for value in (source_root, environment.get("PYTHONPATH")) + if value + ) + completed = subprocess.run( + [ + sys.executable, + "-m", + "cuvs_bench.run", + "--backend-config", + str(backend_config), + "--dataset-configuration", + str(dataset_config), + "--dataset", + dataset_name, + "--dataset-path", + str(dataset_path), + "--algorithms", + "pylucene_cuvs_hnsw", + "--groups", + "test", + "--batch-size", + "2", + "-k", + str(k), + "-m", + "latency", + "--build", + "--search", + "--force", + ], + cwd=tmp_path, + env=environment, + capture_output=True, + text=True, + timeout=600, + check=False, + ) + assert completed.returncode == 0, ( + f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}" + ) + + codec = "Lucene101AcceleratedHNSWCodec" + index_name = f"pylucene_cuvs_hnsw_test.codec{codec}" + index_path = dataset_dir / "index" / index_name + assert any(path.suffix == ".vex" for path in index_path.iterdir()) + assert (index_path / _HNSW_PROVENANCE_FILE).is_file() + + result_path = dataset_dir / "result" + build_json = result_path / "build" / "pylucene_cuvs_hnsw,test.json" + search_stem = f"pylucene_cuvs_hnsw,test,k{k},bs2" + search_json = result_path / "search" / f"{search_stem}.json" + build_rows = json.loads(build_json.read_text())["benchmarks"] + search_rows = json.loads(search_json.read_text())["benchmarks"] + assert len(build_rows) == 1 + assert len(search_rows) == 1 + + build_row = build_rows[0] + assert build_row["name"] == f"{index_name}/build" + assert build_row["real_time"] > 0 + assert build_row["index_size"] > 0 + assert build_row["codec"] == codec + assert build_row["writer_path"] == "gpu-hnsw" + _assert_hnsw_verification( + build_row, + codec, + training_vectors.shape[0], + training_vectors.shape[1], + ) + + search_row = search_rows[0] + assert search_row["name"] == f"{index_name}/search" + assert search_row["Recall"] >= 0.75 + assert search_row["items_per_second"] > 0 + assert search_row["Latency"] > 0 + assert search_row["search_time_ms"] > 0 + _assert_hnsw_verification( + search_row, + codec, + training_vectors.shape[0], + training_vectors.shape[1], + ) + + raw_csv = result_path / "search" / f"{search_stem},raw.csv" + with raw_csv.open(newline="") as file: + csv_rows = list(csv.DictReader(file)) + assert len(csv_rows) == 1 + csv_row = csv_rows[0] + assert csv_row["index_name"] == index_name + assert float(csv_row["recall"]) >= 0.75 + assert float(csv_row["throughput"]) > 0 + assert float(csv_row["latency"]) > 0 + assert float(csv_row["build time"]) > 0 + assert (result_path / "search" / f"{search_stem},latency.csv").is_file() + assert (result_path / "search" / f"{search_stem},throughput.csv").is_file() diff --git a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_provenance.py b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_provenance.py new file mode 100644 index 0000000000..475aaba9c1 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_provenance.py @@ -0,0 +1,270 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Index provenance and training-shape tests for PyLucene.""" + +from __future__ import annotations + +import json +import stat + +import numpy as np +import pytest + +import cuvs_bench.backends.pylucene as pylucene_backend +from cuvs_bench.backends.base import Dataset +from cuvs_bench.tests._pylucene_test_utils import ( + _CAGRA_CODEC, + _HNSW_CODEC, + _FakeRuntime, + _backend, + _dataset, + _index, + _prepare_cagra_index, + _prepare_hnsw_index, + _write_test_bin, +) + + +def test_hnsw_provenance_round_trip(tmp_path): + index_path = tmp_path / "index" + manifest_path = _prepare_hnsw_index(index_path) + + verification = pylucene_backend._verify_hnsw_provenance( + index_path, + _HNSW_CODEC, + expected_vector_count=10, + expected_dimensions=4, + ) + + assert verification.to_metadata() == { + "status": "gpu-hnsw-provenance", + "schema_version": 1, + "codec": _HNSW_CODEC, + "writer_path": "gpu-hnsw", + "vector_count": 10, + "dimensions": 4, + "commit_file_count": 1, + } + payload = json.loads(manifest_path.read_text()) + assert payload["commit_fingerprints"] == [ + { + "name": "segments_1", + "sha256": ( + "1bc04b5291c26a46d918139138b992d2de976d6851d0893b" + "0476b85bfbdfc6e6" + ), + } + ] + assert stat.S_IMODE(manifest_path.stat().st_mode) == 0o644 + assert list(index_path.glob(f"{manifest_path.name}.*.tmp")) == [] + + +def test_cagra_provenance_round_trip(tmp_path): + index_path = tmp_path / "index" + manifest_path = _prepare_cagra_index(index_path) + + verification = pylucene_backend._verify_cagra_provenance( + index_path, + expected_vector_count=10, + expected_dimensions=4, + ) + + assert verification.to_metadata() == { + "status": "gpu-cagra-provenance", + "schema_version": 1, + "codec": _CAGRA_CODEC, + "writer_path": "gpu-cagra", + "vector_count": 10, + "dimensions": 4, + "commit_file_count": 1, + } + assert stat.S_IMODE(manifest_path.stat().st_mode) == 0o644 + + +@pytest.mark.parametrize( + ("expected_vector_count", "expected_dimensions", "error"), + [ + (11, 4, "vector count"), + (10, 8, "dimensions"), + ], +) +def test_hnsw_provenance_rejects_dataset_shape_mismatch( + tmp_path, expected_vector_count, expected_dimensions, error +): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + + with pytest.raises(RuntimeError, match=error): + pylucene_backend._verify_hnsw_provenance( + index_path, + _HNSW_CODEC, + expected_vector_count=expected_vector_count, + expected_dimensions=expected_dimensions, + ) + + +@pytest.mark.parametrize( + ("manifest_state", "error"), + [ + ("wrong-writer", "required GPU writer path"), + ("boolean-count", "positive integer"), + ("malformed-fingerprint", "malformed Lucene commit fingerprint"), + ("duplicate-fingerprint", "must be unique and sorted"), + ], +) +def test_hnsw_provenance_rejects_malformed_manifest_fields( + tmp_path, manifest_state, error +): + index_path = tmp_path / "index" + manifest_path = _prepare_hnsw_index(index_path) + payload = json.loads(manifest_path.read_text()) + + if manifest_state == "wrong-writer": + payload["writer_path"] = "cpu-hnsw" + elif manifest_state == "boolean-count": + payload["vector_count"] = True + elif manifest_state == "malformed-fingerprint": + payload["commit_fingerprints"] = [{"name": "segments_1"}] + else: + payload["commit_fingerprints"] *= 2 + manifest_path.write_text(json.dumps(payload)) + + with pytest.raises(RuntimeError, match=error): + pylucene_backend._verify_hnsw_provenance( + index_path, + _HNSW_CODEC, + ) + + +@pytest.mark.parametrize( + ("field_name", "field_value", "error"), + [ + ("name", 1, "commit filename must be a string"), + ("name", "commit_1", "must start with 'segments_'"), + ("name", "segments_1/child", "must not contain a path"), + ("sha256", 1, "commit SHA-256 must be a string"), + ("sha256", "0" * 63, "must contain 64 characters"), + ("sha256", "A" * 64, "must be lowercase hexadecimal"), + ], +) +def test_hnsw_provenance_reports_invalid_commit_fingerprint_field( + tmp_path, field_name, field_value, error +): + index_path = tmp_path / "index" + manifest_path = _prepare_hnsw_index(index_path) + payload = json.loads(manifest_path.read_text()) + payload["commit_fingerprints"][0][field_name] = field_value + manifest_path.write_text(json.dumps(payload)) + + with pytest.raises(RuntimeError, match=error): + pylucene_backend._verify_hnsw_provenance(index_path, _HNSW_CODEC) + + +def test_expected_training_shape_prefers_explicit_vectors_over_base_file( + tmp_path, +): + base_file = tmp_path / "different.fbin" + _write_test_bin(base_file, np.zeros((3, 8), dtype=np.float32)) + dataset = _dataset(n_base=10, dimensions=4) + dataset.base_file = str(base_file) + + assert pylucene_backend._expected_training_shape(dataset) == (10, 4) + + +def test_expected_training_shape_does_not_validate_unused_file_metadata( + tmp_path, +): + dataset = _dataset(n_base=10, dimensions=4) + dataset.base_file = str(tmp_path / "missing.ibin") + dataset.metadata["subset_size"] = True + + assert pylucene_backend._expected_training_shape(dataset) == (10, 4) + + +def test_expected_training_shape_returns_none_without_training_source(): + dataset = Dataset( + name="query-only", + query_vectors=np.zeros((2, 4), dtype=np.float32), + distance_metric="euclidean", + ) + + assert pylucene_backend._expected_training_shape(dataset) is None + + +def test_expected_training_shape_clamps_file_to_valid_subset(tmp_path): + base_file = tmp_path / "base.fbin" + _write_test_bin(base_file, np.zeros((10, 4), dtype=np.float32)) + dataset = Dataset( + name="file-backed", + query_vectors=np.zeros((2, 4), dtype=np.float32), + base_file=str(base_file), + distance_metric="euclidean", + metadata={"subset_size": 3}, + ) + + assert pylucene_backend._expected_training_shape(dataset) == (3, 4) + assert dataset.loaded_training_vectors is None + + +def test_build_reuses_file_backed_subset_without_loading_vectors(tmp_path): + base_file = tmp_path / "base.fbin" + _write_test_bin(base_file, np.zeros((10, 4), dtype=np.float32)) + dataset = Dataset( + name="file-backed", + query_vectors=np.zeros((2, 4), dtype=np.float32), + base_file=str(base_file), + distance_metric="euclidean", + metadata={"subset_size": 3}, + ) + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path, vector_count=3, dimensions=4) + backend = _backend() + + result = backend.build(dataset, [_index(index_path)], force=False) + + assert result.success + assert result.metadata["skipped"] is True + assert result.metadata["hnsw_verification"]["vector_count"] == 3 + assert dataset.loaded_training_vectors is None + assert backend._runtime is None + + +@pytest.mark.parametrize("subset_size", [0, -1, True, "3"]) +def test_expected_training_shape_rejects_invalid_file_subset( + tmp_path, subset_size +): + base_file = tmp_path / "base.fbin" + _write_test_bin(base_file, np.zeros((10, 4), dtype=np.float32)) + dataset = Dataset( + name="file-backed", + query_vectors=np.zeros((2, 4), dtype=np.float32), + base_file=str(base_file), + distance_metric="euclidean", + metadata={"subset_size": subset_size}, + ) + + with pytest.raises(ValueError, match="positive integer"): + pylucene_backend._expected_training_shape(dataset) + + +def test_reused_index_rejects_non_float32_base_file(tmp_path): + base_file = tmp_path / "base.ibin" + _write_test_bin(base_file, np.zeros((10, 4), dtype=np.int32)) + dataset = Dataset( + name="integer-data", + query_vectors=np.zeros((2, 4), dtype=np.float32), + base_file=str(base_file), + distance_metric="euclidean", + ) + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + + result = _backend(_FakeRuntime()).build( + dataset, [_index(index_path)], force=False + ) + + assert not result.success + assert "must use float32 values, got int32" in result.error_message diff --git a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_runtime.py b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_runtime.py new file mode 100644 index 0000000000..da62e13421 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_runtime.py @@ -0,0 +1,602 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""PyLucene JVM and Java-boundary unit tests.""" + +from __future__ import annotations + +import zipfile +from types import SimpleNamespace + +import numpy as np +import pytest + +import cuvs_bench.backends.pylucene as pylucene_backend +from cuvs_bench.backends.pylucene import _ResolvedCodec +from cuvs_bench.tests._pylucene_test_utils import _HNSW_CODEC + + +def _resolved_codec(java_codec=None, writer_path="gpu-hnsw") -> _ResolvedCodec: + return _ResolvedCodec( + codec_name=_HNSW_CODEC, + java_codec=java_codec if java_codec is not None else object(), + telemetry={"writerPath": writer_path}, + writer_path=writer_path, + ) + + +class _FakeVMEnvironment: + def __init__(self): + self.attach_count = 0 + + def attachCurrentThread(self): + self.attach_count += 1 + + +class _FakeLuceneModule: + CLASSPATH = "/pylucene/lucene-core.jar" + VERSION = "test" + + def __init__(self): + self.environment = None + self.init_calls = [] + + def getVMEnv(self): + return self.environment + + def initVM(self, **kwargs): + self.init_calls.append(kwargs) + self.environment = _FakeVMEnvironment() + return self.environment + + +class _FakeIndexWriter: + def __init__(self, error_at=None): + self.error_at = error_at + self.documents = [] + self.committed = False + self.rollback_called = False + self.close_called = False + + def addDocument(self, document): + if self.error_at in {"interrupt", "interrupt-rollback"}: + raise KeyboardInterrupt + if self.error_at in {"add", "rollback"}: + raise RuntimeError("add failed") + self.documents.append(document) + + def commit(self): + if self.error_at == "commit": + raise RuntimeError("commit failed") + self.committed = True + + def rollback(self): + self.rollback_called = True + if self.error_at in {"rollback", "interrupt-rollback"}: + raise RuntimeError("rollback failed") + + def close(self): + self.close_called = True + if self.error_at == "close": + raise RuntimeError("close failed") + + +class _FakeIndexWriterConfig: + OpenMode = SimpleNamespace(CREATE=object()) + + def setOpenMode(self, _mode): + pass + + def setCodec(self, _codec): + pass + + def setUseCompoundFile(self, _enabled): + pass + + def setMergeScheduler(self, _scheduler): + pass + + +def _fake_index_writer_runtime(error_at=None, directory_close_error=None): + writer = _FakeIndexWriter(error_at=error_at) + directory = SimpleNamespace(closed=False) + + def close_directory(): + directory.closed = True + if directory_close_error is not None: + raise directory_close_error + + directory.close = close_directory + runtime = pylucene_backend._PyLuceneRuntime.__new__( + pylucene_backend._PyLuceneRuntime + ) + runtime.attach_current_thread = lambda: None + runtime.resolve_codec = lambda _codec_name: object() + runtime.Paths = SimpleNamespace(get=lambda path: path) + runtime.FSDirectory = SimpleNamespace(open=lambda _path: directory) + runtime.IndexWriterConfig = _FakeIndexWriterConfig + runtime.SerialMergeScheduler = lambda: object() + runtime.IndexWriter = lambda _directory, _config: writer + runtime._vector_document = lambda document_id, vector: ( + document_id, + vector.copy(), + ) + runtime._test_writer = writer + runtime._test_directory = directory + runtime._test_codec = object() + return runtime + + +@pytest.fixture(autouse=True) +def _reset_jvm_tracking(monkeypatch): + monkeypatch.setattr(pylucene_backend, "_INITIALIZED_CLASSPATH", None) + monkeypatch.setattr(pylucene_backend, "_INITIALIZED_VMARGS", None) + + +def test_initialize_pylucene_uses_verified_classpath_and_vmargs( + tmp_path, monkeypatch +): + cuvs_java = tmp_path / "cuvs-java.jar" + cuvs_lucene = tmp_path / "cuvs-lucene.jar" + cuvs_java.touch() + cuvs_lucene.touch() + fake_lucene = _FakeLuceneModule() + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda name: fake_lucene, + ) + + returned = pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": cuvs_java, + "cuvs_lucene_jar": cuvs_lucene, + "java_library_path": "/native", + "jvm_args": ["-Xms1g"], + } + ) + + assert returned is fake_lucene + assert len(fake_lucene.init_calls) == 1 + init_call = fake_lucene.init_calls[0] + assert init_call["classpath"].split(":") == [ + str(cuvs_java), + str(cuvs_lucene), + fake_lucene.CLASSPATH, + ] + assert init_call["vmargs"] == [ + "--enable-native-access=ALL-UNNAMED", + "--add-modules=jdk.incubator.vector", + "-Djava.library.path=/native", + "-Xms1g", + ] + assert fake_lucene.environment.attach_count == 1 + + pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": cuvs_java, + "cuvs_lucene_jar": cuvs_lucene, + "java_library_path": "/native", + "jvm_args": ["-Xms1g"], + } + ) + assert len(fake_lucene.init_calls) == 1 + assert fake_lucene.environment.attach_count == 2 + + +@pytest.mark.parametrize( + ("library_paths", "expected_library_path"), + [ + ( + { + "JAVA_LIBRARY_PATH": "/java-native", + "LD_LIBRARY_PATH": "/ld-native", + }, + "/java-native", + ), + ({"LD_LIBRARY_PATH": "/ld-native"}, "/ld-native"), + ], + ids=["java-library-path-precedence", "ld-library-path-fallback"], +) +def test_initialize_pylucene_uses_environment_runtime_config( + library_paths, + expected_library_path, + tmp_path, + monkeypatch, +): + cuvs_java = tmp_path / "cuvs-java.jar" + cuvs_lucene = tmp_path / "cuvs-lucene.jar" + cuvs_java.touch() + cuvs_lucene.touch() + fake_lucene = _FakeLuceneModule() + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda _name: fake_lucene, + ) + monkeypatch.setenv("CUVS_LUCENE_CUVS_JAVA_JAR", str(cuvs_java)) + monkeypatch.setenv("CUVS_LUCENE_JAR", str(cuvs_lucene)) + monkeypatch.delenv("JAVA_LIBRARY_PATH", raising=False) + monkeypatch.delenv("LD_LIBRARY_PATH", raising=False) + for name, value in library_paths.items(): + monkeypatch.setenv(name, value) + + pylucene_backend._initialize_pylucene({}) + + init_call = fake_lucene.init_calls[0] + assert init_call["classpath"].split(":") == [ + str(cuvs_java), + str(cuvs_lucene), + fake_lucene.CLASSPATH, + ] + assert init_call["vmargs"] == [ + "--enable-native-access=ALL-UNNAMED", + "--add-modules=jdk.incubator.vector", + f"-Djava.library.path={expected_library_path}", + ] + + +def test_initialize_pylucene_rejects_fat_cuvs_lucene_jar_before_init( + tmp_path, monkeypatch +): + cuvs_java = tmp_path / "cuvs-java.jar" + cuvs_lucene = tmp_path / "cuvs-lucene-jar-with-dependencies.jar" + cuvs_java.touch() + with zipfile.ZipFile(cuvs_lucene, "w") as archive: + archive.writestr( + pylucene_backend._LUCENE_CORE_CLASS, + b"bundled Lucene bytecode", + ) + fake_lucene = _FakeLuceneModule() + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda _name: fake_lucene, + ) + + with pytest.raises(RuntimeError, match="standard thin cuvs-lucene JAR"): + pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": cuvs_java, + "cuvs_lucene_jar": cuvs_lucene, + } + ) + + assert fake_lucene.init_calls == [] + + +def test_initialize_pylucene_rejects_externally_started_jvm( + tmp_path, monkeypatch +): + cuvs_java = tmp_path / "cuvs-java.jar" + cuvs_lucene = tmp_path / "cuvs-lucene.jar" + cuvs_java.touch() + cuvs_lucene.touch() + fake_lucene = _FakeLuceneModule() + fake_lucene.environment = _FakeVMEnvironment() + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda name: fake_lucene, + ) + + with pytest.raises(RuntimeError, match="initialized before"): + pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": cuvs_java, + "cuvs_lucene_jar": cuvs_lucene, + } + ) + + +def test_initialize_pylucene_rejects_different_jars_after_start( + tmp_path, monkeypatch +): + first_java = tmp_path / "first-java.jar" + first_lucene = tmp_path / "first-lucene.jar" + second_java = tmp_path / "second-java.jar" + for path in (first_java, first_lucene, second_java): + path.touch() + fake_lucene = _FakeLuceneModule() + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda name: fake_lucene, + ) + + pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": first_java, + "cuvs_lucene_jar": first_lucene, + } + ) + + with pytest.raises(RuntimeError, match="different"): + pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": second_java, + "cuvs_lucene_jar": first_lucene, + } + ) + + +def test_initialize_pylucene_rejects_different_vmargs_after_start( + tmp_path, monkeypatch +): + cuvs_java = tmp_path / "cuvs-java.jar" + cuvs_lucene = tmp_path / "cuvs-lucene.jar" + cuvs_java.touch() + cuvs_lucene.touch() + fake_lucene = _FakeLuceneModule() + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda name: fake_lucene, + ) + + pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": cuvs_java, + "cuvs_lucene_jar": cuvs_lucene, + "java_library_path": "/first", + } + ) + + with pytest.raises(RuntimeError, match="different JVM arguments"): + pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": cuvs_java, + "cuvs_lucene_jar": cuvs_lucene, + "java_library_path": "/second", + } + ) + + +def test_initialize_pylucene_reports_missing_binding(monkeypatch): + def missing_import(_name): + raise ImportError("missing") + + monkeypatch.setattr( + pylucene_backend.importlib, "import_module", missing_import + ) + + with pytest.raises(ImportError, match="must be built"): + pylucene_backend._initialize_pylucene({}) + + +def test_initialize_pylucene_reports_missing_jar(monkeypatch): + monkeypatch.delenv("CUVS_LUCENE_CUVS_JAVA_JAR", raising=False) + monkeypatch.delenv("CUVS_LUCENE_JAR", raising=False) + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda _name: _FakeLuceneModule(), + ) + + with pytest.raises(RuntimeError, match="cuvs_java_jar"): + pylucene_backend._initialize_pylucene({}) + + +@pytest.mark.parametrize("jvm_args", ["-Xmx1g", ["-Xmx1g", 1]]) +def test_initialize_pylucene_rejects_invalid_jvm_args( + jvm_args, tmp_path, monkeypatch +): + cuvs_java = tmp_path / "cuvs-java.jar" + cuvs_lucene = tmp_path / "cuvs-lucene.jar" + cuvs_java.touch() + cuvs_lucene.touch() + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda _name: _FakeLuceneModule(), + ) + + with pytest.raises(TypeError, match="jvm_args"): + pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": cuvs_java, + "cuvs_lucene_jar": cuvs_lucene, + "jvm_args": jvm_args, + } + ) + + +@pytest.mark.parametrize( + ("description", "expected"), + [ + ( + "Format(writerPath=gpu-hnsw;hnswLayers=1)", + {"writerPath": "gpu-hnsw", "hnswLayers": "1"}, + ), + ("Format(writerPath=gpu-cagra)", {"writerPath": "gpu-cagra"}), + ], +) +def test_parse_writer_telemetry(description, expected): + assert pylucene_backend._parse_writer_telemetry(description) == expected + + +@pytest.mark.parametrize( + "description", + ["Format", "Format(writerPath)", "Format(=gpu-hnsw)"], +) +def test_parse_writer_telemetry_rejects_malformed_description(description): + with pytest.raises(RuntimeError, match="Malformed"): + pylucene_backend._parse_writer_telemetry(description) + + +def test_resolved_codec_copies_and_freezes_writer_telemetry(): + telemetry = {"writerPath": "gpu-hnsw"} + + resolved_codec = _ResolvedCodec( + codec_name=_HNSW_CODEC, + java_codec=object(), + telemetry=telemetry, + writer_path="gpu-hnsw", + ) + telemetry["writerPath"] = "changed" + + assert dict(resolved_codec.telemetry) == {"writerPath": "gpu-hnsw"} + with pytest.raises(TypeError): + resolved_codec.telemetry["writerPath"] = "changed" + + +def test_runtime_build_index_commits_and_closes_writer_and_directory(tmp_path): + runtime = _fake_index_writer_runtime() + vectors = np.zeros((2, 4), dtype=np.float32) + + telemetry = runtime.build_index( + tmp_path, + vectors, + _resolved_codec(runtime._test_codec), + ) + + assert telemetry == {"writerPath": "gpu-hnsw"} + assert len(runtime._test_writer.documents) == 2 + assert runtime._test_writer.committed is True + assert runtime._test_writer.rollback_called is False + assert runtime._test_writer.close_called is True + assert runtime._test_directory.closed is True + + +@pytest.mark.parametrize( + ("error_at", "expected_error", "expect_rollback", "expect_close"), + [ + ("add", "add failed", True, False), + ("commit", "commit failed", True, False), + ("rollback", "rollback failed", True, False), + ("close", "close failed", False, True), + ], +) +def test_runtime_build_index_preserves_transaction_cleanup( + tmp_path, error_at, expected_error, expect_rollback, expect_close +): + runtime = _fake_index_writer_runtime(error_at=error_at) + vectors = np.zeros((2, 4), dtype=np.float32) + + with pytest.raises(RuntimeError, match=expected_error): + runtime.build_index( + tmp_path, + vectors, + _resolved_codec(runtime._test_codec), + ) + + assert runtime._test_writer.rollback_called is expect_rollback + assert runtime._test_writer.close_called is expect_close + assert runtime._test_directory.closed is True + + +def test_runtime_build_index_rolls_back_on_interrupt(tmp_path): + runtime = _fake_index_writer_runtime(error_at="interrupt") + vectors = np.zeros((2, 4), dtype=np.float32) + + with pytest.raises(KeyboardInterrupt): + runtime.build_index( + tmp_path, + vectors, + _resolved_codec(runtime._test_codec), + ) + + assert runtime._test_writer.rollback_called is True + assert runtime._test_writer.close_called is False + assert runtime._test_directory.closed is True + + +def test_runtime_build_index_preserves_interrupt_when_rollback_fails(tmp_path): + runtime = _fake_index_writer_runtime(error_at="interrupt-rollback") + vectors = np.zeros((2, 4), dtype=np.float32) + + with pytest.raises(KeyboardInterrupt) as exc_info: + runtime.build_index( + tmp_path, + vectors, + _resolved_codec(runtime._test_codec), + ) + + assert exc_info.value.__notes__ == [ + "IndexWriter rollback also failed: RuntimeError: rollback failed" + ] + assert runtime._test_writer.rollback_called is True + assert runtime._test_directory.closed is True + + +def test_runtime_build_index_preserves_interrupt_when_directory_close_fails( + tmp_path, +): + runtime = _fake_index_writer_runtime( + error_at="interrupt", + directory_close_error=RuntimeError("directory close failed"), + ) + vectors = np.zeros((2, 4), dtype=np.float32) + + with pytest.raises(KeyboardInterrupt) as exc_info: + runtime.build_index( + tmp_path, + vectors, + _resolved_codec(runtime._test_codec), + ) + + assert exc_info.value.__notes__ == [ + "Failed to close Lucene directory: " + "RuntimeError: directory close failed" + ] + assert runtime._test_writer.rollback_called is True + assert runtime._test_directory.closed is True + + +def test_search_cleanup_attempts_directory_after_reader_close_fails(): + close_calls = [] + + def fail_reader_close(): + close_calls.append("reader") + raise RuntimeError("reader close failed") + + def fail_directory_close(): + close_calls.append("directory") + raise RuntimeError("directory close failed") + + reader = SimpleNamespace(close=fail_reader_close) + directory = SimpleNamespace(close=fail_directory_close) + + with pytest.raises(RuntimeError, match="reader close failed") as exc_info: + with pylucene_backend._CleanupStack() as cleanups: + cleanups.add("close Lucene directory", directory.close) + cleanups.add("close Lucene index reader", reader.close) + + assert close_calls == ["reader", "directory"] + assert exc_info.value.__notes__ == [ + "Failed to close Lucene directory: " + "RuntimeError: directory close failed" + ] + + +def test_cleanup_stack_ignores_unrelated_handled_exception(): + def fail_close(): + raise RuntimeError("close failed") + + unrelated_error = None + try: + raise ValueError("unrelated") + except ValueError as exc: + unrelated_error = exc + with pytest.raises(RuntimeError, match="close failed"): + with pylucene_backend._CleanupStack() as cleanups: + cleanups.add("close test resource", fail_close) + + assert getattr(unrelated_error, "__notes__", []) == [] + + +def test_cleanup_stack_prioritizes_cleanup_interrupt_over_operation_error(): + def interrupt_close(): + raise KeyboardInterrupt("stop") + + with pytest.raises(KeyboardInterrupt, match="stop") as exc_info: + with pylucene_backend._CleanupStack() as cleanups: + cleanups.add("close test resource", interrupt_close) + raise RuntimeError("operation failed") + + assert exc_info.value.__notes__ == [ + "Raised while attempting to close test resource; prior failure: " + "RuntimeError: operation failed" + ] diff --git a/python/cuvs_bench/cuvs_bench/tests/test_registry.py b/python/cuvs_bench/cuvs_bench/tests/test_registry.py index 960ea9cb25..d32b89048e 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_registry.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_registry.py @@ -1,11 +1,13 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ Unit tests for the backend registry system. """ +import json + import pytest import numpy as np @@ -19,6 +21,12 @@ register_backend, get_backend, ) +from cuvs_bench.orchestrator import ( + BenchmarkConfig, + BenchmarkOrchestrator, + DatasetConfig, +) +from cuvs_bench.orchestrator.config_loaders import IndexConfig class DummyBackend(BenchmarkBackend): @@ -197,6 +205,28 @@ def test_build_result_to_json(self): assert json_result["nlist"] == 1024 assert json_result["gpu_time"] == 4.2 + def test_build_result_core_fields_cannot_be_overridden(self): + result = BuildResult( + index_path="/path/to/index", + build_time_seconds=5.5, + index_size_bytes=1024000, + algorithm="test_algo", + build_params={"name": "wrong", "real_time": -1}, + metadata={ + "time_unit": "ms", + "success": False, + "error_message": "wrong", + }, + ) + + json_result = result.to_json() + + assert json_result["name"] == "test_algo/build" + assert json_result["real_time"] == 5.5 + assert json_result["time_unit"] == "s" + assert json_result["success"] is True + assert "error_message" not in json_result + class TestSearchResult: """Tests for SearchResult dataclass.""" @@ -220,6 +250,36 @@ def test_search_result_creation(self): assert result.queries_per_second == 20.0 assert result.success is True + def test_search_result_preserves_legacy_positional_arguments(self): + neighbors = np.array([[1]]) + distances = np.array([[0.1]]) + latency_percentiles = {"p50": 1.0} + metadata = {"source": "legacy"} + + result = SearchResult( + neighbors, + distances, + 2.0, + 500.0, + 0.9, + "test_algo", + [{}], + latency_percentiles, + 0.5, + 0.75, + metadata, + False, + "failed", + ) + + assert result.latency_percentiles is latency_percentiles + assert result.gpu_time_seconds == 0.5 + assert result.cpu_time_seconds == 0.75 + assert result.metadata is metadata + assert result.success is False + assert result.error_message == "failed" + assert result.latency_seconds is None + def test_search_result_to_json(self): """Test conversion to JSON format.""" neighbors = np.array([[1, 2, 3]]) @@ -246,6 +306,284 @@ def test_search_result_to_json(self): assert json_result["p50"] == 50.0 assert json_result["p95"] == 95.0 + def test_search_result_core_fields_cannot_be_overridden(self): + result = SearchResult( + neighbors=np.array([[1]]), + distances=np.array([[0.1]]), + search_time_ms=6.0, + queries_per_second=500.0, + recall=0.95, + algorithm="test_algo", + search_params=[{}], + latency_seconds=0.003, + gpu_time_seconds=0.004, + cpu_time_seconds=0.005, + metadata={ + "name": "wrong", + "real_time": -1, + "Latency": -1, + "GPU": -1, + "cpu_time": -1, + "Recall": -1, + "success": False, + "error_message": "wrong", + }, + ) + + json_result = result.to_json() + + assert json_result["name"] == "test_algo/search" + assert json_result["real_time"] == 3.0 + assert json_result["Latency"] == 0.003 + assert json_result["GPU"] == 0.004 + assert json_result["cpu_time"] == 0.005 + assert json_result["Recall"] == 0.95 + assert json_result["success"] is True + assert "error_message" not in json_result + + def test_search_result_optional_metadata_fields_are_fallbacks(self): + result = SearchResult( + neighbors=np.array([[1]]), + distances=np.array([[0.1]]), + search_time_ms=6.0, + queries_per_second=500.0, + recall=0.95, + algorithm="test_algo", + search_params=[{}], + metadata={ + "Latency": 0.006, + "GPU": 0.004, + "cpu_time": 0.005, + }, + ) + + json_result = result.to_json() + + assert json_result["Latency"] == 0.006 + assert json_result["GPU"] == 0.004 + assert json_result["cpu_time"] == 0.005 + + +def test_backend_without_result_stems_runs_without_generic_persistence( + tmp_path, monkeypatch +): + class DummyConfigLoader: + def load(self, **_kwargs): + return DatasetConfig(name="dummy"), [ + BenchmarkConfig( + indexes=[ + IndexConfig( + name="dummy-index", + algo="dummy", + build_param={}, + search_params=[{}], + file=str(tmp_path / "dummy-index"), + ) + ], + backend_config={"name": "dummy-index"}, + ) + ] + + monkeypatch.setattr( + "cuvs_bench.orchestrator.orchestrator.get_backend_class", + lambda _backend_type: DummyBackend, + ) + monkeypatch.setattr( + "cuvs_bench.orchestrator.orchestrator.get_config_loader", + lambda _backend_type: DummyConfigLoader, + ) + orchestrator = BenchmarkOrchestrator(backend_type="dummy") + + results = orchestrator.run_benchmark( + build=True, + search=False, + dataset="dummy", + dataset_path=str(tmp_path), + ) + + assert len(results) == 1 + assert results[0].success is True + assert not (tmp_path / "dummy" / "result").exists() + + +def test_failed_opt_in_build_clears_only_matching_stale_search_artifacts( + tmp_path, monkeypatch +): + build_stem = "dummy,base" + search_stem = f"{build_stem},k1,bs1" + + class DummyConfigLoader: + def load(self, **_kwargs): + return DatasetConfig(name="dummy"), [ + BenchmarkConfig( + indexes=[ + IndexConfig( + name="dummy-index", + algo="dummy", + build_param={}, + search_params=[{}], + file=str(tmp_path / "dummy-index"), + ) + ], + backend_config={ + "name": "dummy-index", + "output_filename": (build_stem, search_stem), + }, + ) + ] + + class FailedBuildBackend(DummyBackend): + orchestrator_persists_results = True + + def build(self, dataset, indexes, force=False, dry_run=False): + return BuildResult( + index_path=indexes[0].file, + build_time_seconds=0.0, + index_size_bytes=0, + algorithm=self.algo, + build_params={}, + success=False, + error_message="expected build failure", + ) + + def search(self, *args, **kwargs): + raise AssertionError("search must not run after a failed build") + + result_root = tmp_path / "dummy" / "result" + search_dir = result_root / "search" + search_dir.mkdir(parents=True) + stale_paths = [ + search_dir / f"{search_stem}.json", + search_dir / f"{search_stem},raw.csv", + search_dir / f"{search_stem},throughput.csv", + search_dir / f"{search_stem},latency.csv", + ] + unrelated_paths = [ + search_dir / "other,base,k1,bs1.json", + search_dir / "other,base,k1,bs1,raw.csv", + ] + for path in [*stale_paths, *unrelated_paths]: + path.write_text("stale\n", encoding="utf-8") + + monkeypatch.setattr( + "cuvs_bench.orchestrator.orchestrator.get_backend_class", + lambda _backend_type: FailedBuildBackend, + ) + monkeypatch.setattr( + "cuvs_bench.orchestrator.orchestrator.get_config_loader", + lambda _backend_type: DummyConfigLoader, + ) + + results = BenchmarkOrchestrator(backend_type="dummy").run_benchmark( + build=True, + search=True, + count=1, + batch_size=1, + dataset="dummy", + dataset_path=str(tmp_path), + ) + + assert len(results) == 1 + assert results[0].success is False + assert not any(path.exists() for path in stale_paths) + assert all( + path.read_text(encoding="utf-8") == "stale\n" + for path in unrelated_paths + ) + + build_path = result_root / "build" / f"{build_stem}.json" + build_rows = json.loads(build_path.read_text(encoding="utf-8"))[ + "benchmarks" + ] + assert len(build_rows) == 1 + assert build_rows[0]["success"] is False + assert build_rows[0]["error_message"] == "expected build failure" + + +@pytest.mark.parametrize( + ("index_count", "output_filenames", "expected_message"), + [ + ( + 2, + ("dummy,base", "dummy,base,k1,bs1"), + "exactly one index", + ), + ( + 1, + ("dummy", "dummy,k1,bs1"), + "Build result filename stem", + ), + ( + 1, + ("dummy,base", "other,base,k1,bs1"), + "Search result filename stem", + ), + ( + 1, + ("dummy,base",), + "exactly two result filename stems", + ), + ( + 1, + ("../dummy,base", "../dummy,base,k1,bs1"), + "Invalid benchmark result filename", + ), + ], +) +def test_opt_in_sweep_persistence_validates_artifact_identity( + tmp_path, + monkeypatch, + index_count, + output_filenames, + expected_message, +): + indexes = [ + IndexConfig( + name=f"dummy-index-{position}", + algo="dummy", + build_param={}, + search_params=[{}], + file=str(tmp_path / f"dummy-index-{position}"), + ) + for position in range(index_count) + ] + + class DummyConfigLoader: + def load(self, **_kwargs): + return DatasetConfig(name="dummy"), [ + BenchmarkConfig( + indexes=indexes, + backend_config={ + "name": "dummy-index", + "output_filename": output_filenames, + }, + ) + ] + + class PersistedDummyBackend(DummyBackend): + orchestrator_persists_results = True + + monkeypatch.setattr( + "cuvs_bench.orchestrator.orchestrator.get_backend_class", + lambda _backend_type: PersistedDummyBackend, + ) + monkeypatch.setattr( + "cuvs_bench.orchestrator.orchestrator.get_config_loader", + lambda _backend_type: DummyConfigLoader, + ) + + with pytest.raises(ValueError, match=expected_message): + BenchmarkOrchestrator(backend_type="dummy").run_benchmark( + build=True, + search=False, + count=1, + batch_size=1, + dataset="dummy", + dataset_path=str(tmp_path), + ) + + assert not (tmp_path / "dummy" / "result").exists() + class TestBackendRegistry: """Tests for BackendRegistry.""" diff --git a/python/cuvs_bench/cuvs_bench/tests/test_utils.py b/python/cuvs_bench/cuvs_bench/tests/test_utils.py index 01dd803a07..8fddb356b0 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_utils.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_utils.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ @@ -528,7 +528,9 @@ def test_lazy_load_training_vectors(self, tmp_path): _write_test_bin(path, data) dataset = Dataset(name="test", base_file=path) + assert dataset.loaded_training_vectors is None np.testing.assert_array_equal(dataset.training_vectors, data) + np.testing.assert_array_equal(dataset.loaded_training_vectors, data) def test_lazy_load_query_vectors(self, tmp_path): """Test that query vectors are loaded from file on first access.""" @@ -614,7 +616,7 @@ def test_file_path_access_does_not_trigger_loading(self, tmp_path): _ = dataset.name _ = dataset.distance_metric - assert dataset._training_vectors.size == 0 + assert dataset.loaded_training_vectors is None def test_dims_and_counts(self, tmp_path): """Test dims, n_base, and n_queries properties.""" diff --git a/python/cuvs_bench/pyproject.toml b/python/cuvs_bench/pyproject.toml index 2db75baf38..9f582bb313 100644 --- a/python/cuvs_bench/pyproject.toml +++ b/python/cuvs_bench/pyproject.toml @@ -21,6 +21,7 @@ requires-python = ">=3.11" dependencies = [ "click", "cuvs==26.10.*,>=0.0.0a0", + "h5py>=3.8.0", "matplotlib>=3.9", "pandas", "pyyaml", @@ -57,6 +58,7 @@ Homepage = "https://github.com/rapidsai/cuvs" [tool.pytest.ini_options] markers = [ "opensearch: tests that require a live OpenSearch node (run with '-m opensearch')", + "pylucene: real PyLucene/JVM/cuVS tests (set CUVS_BENCH_PYLUCENE_INTEGRATION=1 and run with '-m pylucene')", ] [tool.isort]