diff --git a/dependencies.yaml b/dependencies.yaml index a58e00cb58..d463647d01 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -631,6 +631,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..919d8b3031 100644 --- a/fern/pages/cuvs_bench/install.md +++ b/fern/pages/cuvs_bench/install.md @@ -49,6 +49,66 @@ 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: + +- For GPU indexing or search, an NVIDIA GPU supported by cuVS plus matching CUDA and cuVS native libraries. The accelerated HNSW codec intentionally falls back to Lucene's CPU writer when cuVS is unavailable; the CAGRA codec requires cuVS GPU support. +- JDK 22, including `javac`, and a custom PyLucene wrapper generated against Lucene 10.2.0. cuVS Bench compiles its small PyLucene codec adapter before starting the JVM. Apache does not publish PyLucene 10.2.0, and the official PyLucene 10.0.0 distribution is incompatible with the Lucene 10.2 APIs used here. The [PyLucene source-build instructions](https://lucene.apache.org/pylucene/install.html) describe the general build mechanics, but neither Apache nor cuVS currently provides a ready-to-use 10.2.0 wrapper. +- [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, for GPU execution, native libraries from the same cuVS version. The cuVS-Lucene PR declares cuVS Java 26.10.0. + +cuVS-Lucene now lives in this repository under `java/cuvs-lucene`. [NVIDIA/cuvs-lucene#174](https://github.com/NVIDIA/cuvs-lucene/pull/174), which predates that move, contains the remaining PyLucene 10.2 compatibility changes and its Python end-to-end coverage. Its standalone branch does not include the newer HNSW heuristic delegation from [NVIDIA/cuvs-lucene#177](https://github.com/NVIDIA/cuvs-lucene/pull/177), which is already present in cuVS and is required for the `m` and `ef_construction` benchmark parameters. Until the PyLucene changes are ported and merged, use the validated source combination below; neither source by itself is sufficient. + +Build the dependencies in separate checkouts so this does not change your cuVS Bench working tree. While the cuVS-Lucene PR is under review, record the exact cuVS and PR revisions used for a reproducible environment and keep their cuVS Java versions aligned. + +```bash +git clone https://github.com/NVIDIA/cuvs.git cuvs-pylucene-deps +cd cuvs-pylucene-deps +git switch --detach be8ab314d044aee0f80fbe2c2277a893561288de +./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 cuvs-lucene-pylucene +cd cuvs-lucene-pylucene +git fetch origin pull/174/head +git switch --detach 6fe2c2824408a4ff2ac8f201df05308fd2404b76 +git -C ../cuvs-pylucene-deps show \ + --relative=java/cuvs-lucene --format=email --binary \ + 65b4ae5f1ac5b5916d49709f6c16231df8d26188 \ + -- java/cuvs-lucene | git apply +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 matching cuVS build. + +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. + +The backend checks that `lucene.VERSION` is exactly `10.2.0` before starting the process-wide JVM. It also compiles its configured-codec adapter against the selected JAR, which fails early when the HNSW heuristic API is missing. Then validate the combined artifacts from the temporary cuVS-Lucene checkout. The upstream PyLucene suite is a pytest module; its Java test adapter is compiled into `target/test-classes` by the Maven build and is not part of the production JAR. + +```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}" + +python -m pytest -q -s src/test/python/test_pylucene_end_to_end.py +``` + +See [Running the PyLucene backend](/user-guide/benchmarking-guide/cu-vs-bench-tool/usage#running-the-pylucene-backend) for a smoke workflow after these prerequisites are prepared. + ## 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..0e096f653f 100644 --- a/fern/pages/cuvs_bench/pluggable_backend.md +++ b/fern/pages/cuvs_bench/pluggable_backend.md @@ -194,7 +194,7 @@ class ElasticsearchBackend(BenchmarkBackend): dry_run=False, ): n_queries = dataset.n_queries - return SearchResult( + return [SearchResult( neighbors=np.zeros((n_queries, k), dtype=np.int64), distances=np.zeros((n_queries, k), dtype=np.float32), search_time_ms=0.0, @@ -203,7 +203,7 @@ class ElasticsearchBackend(BenchmarkBackend): algorithm=self.algo, search_params=indexes[0].search_params if indexes else [], success=True, - ) + )] ``` ```python @@ -219,9 +219,13 @@ get_registry().register("elasticsearch", ElasticsearchBackend) | Component | Description | | --- | --- | | `ConfigLoader` | Abstract class whose `load(**kwargs)` method returns `(DatasetConfig, List[BenchmarkConfig])`. Register with `register_config_loader(backend_type, loader_class)`. | -| `BenchmarkBackend` | Abstract class whose `build(...)` method returns `BuildResult` and whose `search(...)` method returns `SearchResult`. Register with `BackendRegistry.register(name, backend_class)`. | +| `BenchmarkBackend` | Abstract class whose `build(...)` method returns `BuildResult` and whose `search(...)` method returns `List[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 the production `Lucene101AcceleratedHNSWCodec` and `CuVS2510GPUSearchCodec` through Lucene's service-provider interface. The HNSW codec intentionally permits Lucene's CPU writer fallback when cuVS is unavailable; CAGRA builds and searches require GPU support, and cuVS Bench verifies that committed CAGRA indexes contain only the expected CAGRA vector data. 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 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..e215a005e0 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,121 @@ 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-pylucene/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` and `test` groups use `Lucene101AcceleratedHNSWCodec`. It uses cuVS for HNSW graph construction when GPU support is available and intentionally falls back, with a warning, to Lucene's CPU HNSW writer otherwise. Search uses Lucene's CPU HNSW implementation in either case. Use `--algorithms pylucene_cuvs_cagra` to run `CuVS2510GPUSearchCodec`, which builds and searches on the GPU. + +The HNSW configuration maps benchmark parameters to cuVS-Lucene as follows: + +| Parameter | Scope | Default | Accepted values | Effect | +| --- | --- | --- | --- | --- | +| `codec` | Build | `Lucene101AcceleratedHNSWCodec` | That codec name | Selects the cuVS-accelerated HNSW codec. | +| `m` | Build | `32` | Integer from 1 through 512 | Sets `AcceleratedHNSWParams.maxConn`. | +| `ef_construction` | Build | `32` | Integer from 1 through 512 | Sets `AcceleratedHNSWParams.beamWidth`. | +| `direct_single_segment` | Build | `false` | Boolean | Requests one direct segment, as described below. | +| `num_candidates` | Search | `top_k` | Integer greater than or equal to `top_k` | Requests the candidate count passed to `KnnFloatVectorQuery`, capped at the index size; the backend returns `top_k` neighbors. | + +`m` and `ef_construction` are the HNSW-equivalent inputs used by cuVS-Lucene's `SAME_GRAPH_FOOTPRINT` heuristic; cuVS-Lucene derives its CAGRA build parameters from them. This requires a cuVS-Lucene build containing both the PyLucene 10.2 changes from NVIDIA/cuvs-lucene#174 and the HNSW heuristic delegation from NVIDIA/cuvs-lucene#177. The backend rejects an older JAR during its pre-JVM adapter compilation instead of silently building with unrelated defaults. `num_candidates` is Lucene's candidate budget, not a direct cuVS `ef_search` setting. + +Automatic tune mode samples `m` and `ef_construction` from 1 through 512 and `num_candidates` from `top_k` through 500. Explicit YAML sweeps may use larger candidate counts. + +Elasticsearch shard, replica, and field-name settings do not apply to this local Lucene index. Its HNSW index type maps to `codec`, and the backend validates Euclidean similarity from the dataset rather than accepting Elasticsearch's type, quantization, or similarity options. + +When `direct_single_segment` is true, the backend disables ordinary RAM-triggered flushes and merging, buffers the requested vectors for one flush, and fails unless the committed index has exactly one segment. Lucene's per-indexing-thread hard RAM limit remains in force (1945 MiB by default and less than 2048 MiB by contract); if it forces an earlier flush, the build fails instead of merging the segments. It does not call Lucene `forceMerge`; the cuVS Bench `--force` option only requests a rebuild. A larger JVM heap does not disable that per-thread limit. + +For example, save the following as `pylucene-deep1m.yaml` to sweep three index builds and three searches per index: + +```yaml +name: pylucene_cuvs_hnsw +groups: + deep1m: + build: + codec: ["Lucene101AcceleratedHNSWCodec"] + m: [16, 24, 32] + ef_construction: [32] + direct_single_segment: [true] + search: + num_candidates: [150, 200, 300] +``` + +Run it with `top_k=150` as follows, substituting the dataset name and paths from its dataset configuration: + +```bash +python -m cuvs_bench.run \ + --backend-config pylucene-backend.yaml \ + --configuration pylucene-deep1m.yaml \ + --dataset deep1b-1M \ + --dataset-configuration /absolute/path/to/deep1b-1M.yaml \ + --dataset-path /absolute/path/to/dataset-root \ + --algorithms pylucene_cuvs_hnsw \ + --groups deep1m \ + --batch-size 100 -k 150 \ + -m latency \ + --build --search --force +``` + +The backend accepts FLOAT32 Euclidean datasets with at most 4096 dimensions and at least two indexed vectors. 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` for CAGRA to avoid paths that can use brute-force search above that limit. + +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 different build parameters, writer policy, or compound-file policy. Indexes created outside this backend without the applicable manifest must be rebuilt with `--force`. For HNSW, the policy permits cuVS-Lucene's production CPU fallback. For CAGRA, the backend fails closed: it verifies the committed vector segments and checksums and rejects an index that is not CAGRA-only. + +For CAGRA, the backend disables compound files for both flushed and merged segments so the verifier can inspect the codec's `.vemc` and `.vcag` files directly. HNSW retains Lucene's default compound-file policy and its default index-writer scheduling unless `direct_single_segment` is enabled. + +The backend currently supports latency mode with one search thread. `--batch-size` groups queries for measurement, and the reported latency percentiles are milliseconds per batch. Throughput mode and multiple search threads are not implemented. + +PyLucene's JVM is process-global and can be initialized only once. The backend requires a wrapper generated against Lucene 10.2.0 and checks its version before JVM initialization. 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 +308,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. + +C++ build benchmarks report: | Name | Description | | --- | --- | @@ -203,7 +321,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/_validation.py b/python/cuvs_bench/cuvs_bench/_validation.py new file mode 100644 index 0000000000..0e71cd25f3 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/_validation.py @@ -0,0 +1,61 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Validation helpers for names used in benchmark artifacts.""" + +from __future__ import annotations + +import os +import re +from typing import Any + + +_RESULT_COMPONENT = re.compile(r"[A-Za-z0-9_.-]+") + + +def validate_path_component(value: Any, description: str) -> str: + """Return a safe, non-empty filesystem path component.""" + if ( + not isinstance(value, str) + or not value + or value in {".", ".."} + or os.path.basename(value) != value + or any( + ord(character) < 32 or ord(character) == 127 for character in value + ) + ): + raise ValueError(f"Invalid {description}: {value!r}") + return value + + +def validate_result_component(value: Any, description: str) -> str: + """Return a path component safe for comma-delimited result identities.""" + component = validate_path_component(value, description) + if _RESULT_COMPONENT.fullmatch(component) is None: + raise ValueError(f"Invalid {description}: {value!r}") + return component + + +def format_artifact_identity( + algorithm: Any, + group: Any, + scope: Any = None, + *, + description: str = "benchmark", +) -> str: + """Format an injective identity from validated artifact components.""" + algorithm = validate_result_component( + algorithm, f"{description} algorithm name" + ) + group = validate_result_component(group, f"{description} group name") + if scope is not None: + scope = validate_result_component(scope, f"{description} result scope") + + identity = algorithm + if group != "base": + identity += f"[group={group}]" + if scope is not None: + identity += f"[scope={scope}]" + return identity 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/_java/PyLuceneConfiguredHnswCodec.java b/python/cuvs_bench/cuvs_bench/backends/_java/PyLuceneConfiguredHnswCodec.java new file mode 100644 index 0000000000..1799bb1def --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/backends/_java/PyLuceneConfiguredHnswCodec.java @@ -0,0 +1,71 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.bench; + +import com.nvidia.cuvs.CagraIndexParams.HnswHeuristicType; +import com.nvidia.cuvs.lucene.AcceleratedHNSWParams; +import com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWCodec; + +/** + * PyLucene-compatible adapter for configuring the accelerated HNSW codec. + * + *

PyLucene instantiates codecs through a no-argument constructor. This adapter reads the two + * Lucene-equivalent HNSW build parameters from system properties and delegates all codec behavior + * to the production cuVS-Lucene implementation. + */ +public final class PyLuceneConfiguredHnswCodec extends Lucene101AcceleratedHNSWCodec { + + public static final String M_PROPERTY = "com.nvidia.cuvs.bench.pylucene.hnsw.m"; + public static final String EF_CONSTRUCTION_PROPERTY = + "com.nvidia.cuvs.bench.pylucene.hnsw.efConstruction"; + + private final int m; + private final int efConstruction; + + /** Constructs the production codec with parameters supplied through system properties. */ + public PyLuceneConfiguredHnswCodec() throws Exception { + this(configuredValues()); + } + + private PyLuceneConfiguredHnswCodec(ConfiguredValues values) throws Exception { + super(values.parameters()); + this.m = values.m(); + this.efConstruction = values.efConstruction(); + } + + private static ConfiguredValues configuredValues() { + int m = requiredIntegerProperty(M_PROPERTY); + int efConstruction = requiredIntegerProperty(EF_CONSTRUCTION_PROPERTY); + AcceleratedHNSWParams parameters = + new AcceleratedHNSWParams.Builder() + .withStrategy(AcceleratedHNSWParams.Strategy.HEURISTIC) + .withHnswHeuristicType(HnswHeuristicType.SAME_GRAPH_FOOTPRINT) + .withMaxConn(m) + .withBeamWidth(efConstruction) + .build(); + return new ConfiguredValues(m, efConstruction, parameters); + } + + private static int requiredIntegerProperty(String name) { + String value = System.getProperty(name); + if (value == null) { + throw new IllegalStateException("Required system property is not set: " + name); + } + + try { + return Integer.parseInt(value); + } catch (NumberFormatException error) { + throw new IllegalArgumentException( + "System property " + name + " must be an integer, found: " + value, error); + } + } + + @Override + public String toString() { + return getClass().getSimpleName() + "(m=" + m + ", efConstruction=" + efConstruction + ")"; + } + + private record ConfiguredValues(int m, int efConstruction, AcceleratedHNSWParams parameters) {} +} diff --git a/python/cuvs_bench/cuvs_bench/backends/_java/__init__.py b/python/cuvs_bench/cuvs_bench/backends/_java/__init__.py new file mode 100644 index 0000000000..5d56d9d52c --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/backends/_java/__init__.py @@ -0,0 +1,6 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Packaged Java sources used by the PyLucene backend.""" diff --git a/python/cuvs_bench/cuvs_bench/backends/_pylucene_java.py b/python/cuvs_bench/cuvs_bench/backends/_pylucene_java.py new file mode 100644 index 0000000000..07bd3437ee --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/backends/_pylucene_java.py @@ -0,0 +1,154 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Compile the no-argument codec adapter required by stock PyLucene.""" + +from __future__ import annotations + +import os +import shutil +import subprocess +import sys +import tempfile +import threading +from pathlib import Path +from typing import Optional, Union + +CONFIGURED_HNSW_CODEC_CLASS = ( + "com.nvidia.cuvs.bench.PyLuceneConfiguredHnswCodec" +) +M_PROPERTY = "com.nvidia.cuvs.bench.pylucene.hnsw.m" +EF_CONSTRUCTION_PROPERTY = "com.nvidia.cuvs.bench.pylucene.hnsw.efConstruction" + +_SOURCE_FILE = ( + Path(__file__).with_name("_java") / "PyLuceneConfiguredHnswCodec.java" +) +_COMPILE_LOCK = threading.Lock() +_COMPILED_CLASSPATH: Optional[str] = None +_CLASSES_DIRECTORY: Optional[Path] = None +_TEMPORARY_DIRECTORY: Optional[tempfile.TemporaryDirectory] = None + + +def configured_codec_classes_path( + cuvs_java_jar: Path, + cuvs_lucene_jar: Path, + pylucene_classpath: str, +) -> Path: + """Compile the codec adapter once and return its stable classes path. + + This function must be called before initializing PyLucene's process-wide + JVM so that the returned directory can be included in the JVM classpath. + """ + + compile_classpath = _compile_classpath( + cuvs_java_jar, cuvs_lucene_jar, pylucene_classpath + ) + with _COMPILE_LOCK: + if _CLASSES_DIRECTORY is not None: + _require_same_classpath(compile_classpath) + return _CLASSES_DIRECTORY + return _compile(compile_classpath) + + +def _compile_classpath( + cuvs_java_jar: Union[str, os.PathLike[str]], + cuvs_lucene_jar: Union[str, os.PathLike[str]], + pylucene_classpath: str, +) -> str: + if not isinstance(pylucene_classpath, str) or not pylucene_classpath: + raise ValueError("pylucene_classpath must be a non-empty string") + + jar_paths = ( + Path(cuvs_java_jar).resolve(), + Path(cuvs_lucene_jar).resolve(), + ) + for jar_path in jar_paths: + if not jar_path.is_file(): + raise FileNotFoundError( + f"Cannot compile the PyLucene codec adapter; JAR not found: {jar_path}" + ) + return os.pathsep.join((*map(str, jar_paths), pylucene_classpath)) + + +def _require_same_classpath(compile_classpath: str) -> None: + if compile_classpath != _COMPILED_CLASSPATH: + raise RuntimeError( + "The PyLucene codec adapter was already compiled with a different " + "dependency classpath. PyLucene JVM dependencies are process-wide." + ) + + +def _compile(compile_classpath: str) -> Path: + global _CLASSES_DIRECTORY, _COMPILED_CLASSPATH, _TEMPORARY_DIRECTORY + + if not _SOURCE_FILE.is_file(): + raise RuntimeError( + f"The packaged PyLucene codec adapter source is missing: {_SOURCE_FILE}" + ) + + javac = _find_javac() + temporary_directory = tempfile.TemporaryDirectory( + prefix="cuvs-bench-pylucene-codec-" + ) + classes_directory = Path(temporary_directory.name) / "classes" + classes_directory.mkdir() + try: + completed = subprocess.run( + [ + javac, + "--release", + "22", + "-classpath", + compile_classpath, + "-d", + str(classes_directory), + str(_SOURCE_FILE), + ], + capture_output=True, + check=False, + text=True, + ) + except OSError as error: + temporary_directory.cleanup() + raise RuntimeError( + f"Could not run JDK 22 javac at {javac}: {error}" + ) from error + + if completed.returncode != 0: + temporary_directory.cleanup() + details = completed.stderr.strip() or completed.stdout.strip() + raise RuntimeError( + "Could not compile the PyLucene codec adapter with JDK 22 javac. " + "Verify that cuvs-java, PyLucene 10.2, and the thin cuvs-lucene " + "JAR are compatible. The cuvs-lucene JAR must include both the " + "PyLucene 10.2 support from NVIDIA/cuvs-lucene#174 and the " + "HNSW heuristic delegation now present in cuVS. " + f"javac output:\n{details}" + ) + + _TEMPORARY_DIRECTORY = temporary_directory + _CLASSES_DIRECTORY = classes_directory + _COMPILED_CLASSPATH = compile_classpath + return classes_directory + + +def _find_javac() -> str: + candidates = [] + java_home = os.environ.get("JAVA_HOME") + if java_home: + candidates.append(Path(java_home) / "bin" / "javac") + candidates.append(Path(sys.prefix) / "lib" / "jvm" / "bin" / "javac") + + for candidate in candidates: + if candidate.is_file(): + return str(candidate) + + javac = shutil.which("javac") + if javac is not None: + return javac + raise RuntimeError( + "Configurable PyLucene HNSW builds require JDK 22 javac. Set " + "JAVA_HOME to a JDK 22 installation or put its javac on PATH." + ) diff --git a/python/cuvs_bench/cuvs_bench/backends/base.py b/python/cuvs_bench/cuvs_bench/backends/base.py index 15209decb6..1c9ebfeb55 100644 --- a/python/cuvs_bench/cuvs_bench/backends/base.py +++ b/python/cuvs_bench/cuvs_bench/backends/base.py @@ -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.""" 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..e62a478c20 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/backends/pylucene.py @@ -0,0 +1,3430 @@ +# +# 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 typing import ( + Any, + Callable, + Dict, + List, + NamedTuple, + Optional, + Tuple, + Union, +) + +import numpy as np + +from .._bin_format import read_bin_header +from .._validation import format_artifact_identity, validate_path_component +from ..orchestrator.config_loaders import ( + BenchmarkConfig, + ConfigLoader, + DatasetConfig, + IndexConfig, +) +from ._utils import dtype_from_filename +from ._pylucene_java import ( + CONFIGURED_HNSW_CODEC_CLASS, + EF_CONSTRUCTION_PROPERTY, + M_PROPERTY, + configured_codec_classes_path, +) +from .base import BenchmarkBackend, BuildResult, Dataset, SearchResult + +_ID_FIELD = "id" +_VECTOR_FIELD = "vector" +_MAX_DIMENSIONS = 4096 +_REQUIRED_PYLUCENE_VERSION = "10.2.0" + +_HNSW_CODEC = "Lucene101AcceleratedHNSWCodec" +_CAGRA_CODEC = "CuVS2510GPUSearchCodec" +_SUPPORTED_CODECS = frozenset({_HNSW_CODEC, _CAGRA_CODEC}) +_HNSW_BUILD_KEYS = frozenset( + {"codec", "m", "ef_construction", "direct_single_segment"} +) +_SUPPORTED_BUILD_KEYS = _HNSW_BUILD_KEYS +_SUPPORTED_SEARCH_KEYS = frozenset({"num_candidates"}) +_DEFAULT_M = 32 +_DEFAULT_EF_CONSTRUCTION = 32 +_MIN_HNSW_BUILD_PARAMETER = 1 +_MAX_HNSW_BUILD_PARAMETER = 512 +_EXPECTED_WRITER_POLICY = { + _HNSW_CODEC: "gpu-with-cpu-fallback", + _CAGRA_CODEC: "gpu-cagra", +} +_COMPOUND_FILE_POLICY = { + _HNSW_CODEC: "lucene-default", + _CAGRA_CODEC: "disabled", +} +_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 = 4 +_PROVENANCE_KEYS = frozenset( + { + "schema_version", + "codec", + "build_parameters", + "writer_policy", + "compound_file_policy", + "vector_count", + "dimensions", + "segment_count", + "commit_fingerprints", + } +) +_SHA256_HEX_DIGITS = frozenset("0123456789abcdef") +_LUCENE_CORE_CLASS = "org/apache/lucene/index/IndexWriter.class" + +_JVM_INIT_LOCK = threading.Lock() +_CONFIGURED_CODEC_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}: {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 _restore_java_property(system: Any, name: str, previous: Any) -> None: + if previous is None: + system.clearProperty(name) + else: + system.setProperty(name, str(previous)) + + +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 _validate_pylucene_version(lucene: Any) -> None: + actual_version = str(getattr(lucene, "VERSION", "")) + if actual_version != _REQUIRED_PYLUCENE_VERSION: + raise RuntimeError( + "PyLucene must be generated against the same Lucene version as " + "cuVS-Lucene: expected " + f"{_REQUIRED_PYLUCENE_VERSION}, found {actual_version}. " + "Activate a matching PyLucene build before initializing the JVM." + ) + + +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) + adapter_classes = configured_codec_classes_path( + cuvs_java_jar, + cuvs_lucene_jar, + str(lucene.CLASSPATH), + ) + return os.pathsep.join( + [ + str(adapter_classes), + str(cuvs_java_jar), + str(cuvs_lucene_jar), + str(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() + _validate_pylucene_version(lucene) + classpath = _pylucene_classpath(config, lucene) + vmargs = _pylucene_vmargs(config) + _attach_pylucene_jvm(lucene, classpath, vmargs) + return lucene + + +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}, 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}. 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 _bounded_hnsw_build_parameter(value: Any, name: str) -> int: + if type(value) is not int or not ( + _MIN_HNSW_BUILD_PARAMETER <= value <= _MAX_HNSW_BUILD_PARAMETER + ): + raise ValueError( + f"PyLucene {name} must be an integer in " + f"[{_MIN_HNSW_BUILD_PARAMETER}, {_MAX_HNSW_BUILD_PARAMETER}], " + f"got {value!r}" + ) + return value + + +def _direct_single_segment(value: Any) -> bool: + if type(value) is not bool: + raise TypeError( + f"PyLucene direct_single_segment must be a boolean, got {value!r}" + ) + return value + + +def _normalize_build_params( + build_params: Any, backend_config: Optional[Dict[str, Any]] = None +) -> 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}") + codec_name = _configured_codec(build_params, backend_config or {}) + hnsw_parameters = set(build_params) - {"codec"} + if codec_name != _HNSW_CODEC and hnsw_parameters: + names = ", ".join(sorted(hnsw_parameters)) + raise ValueError( + f"PyLucene build parameter(s) {names} apply only to {_HNSW_CODEC}" + ) + if codec_name == _CAGRA_CODEC: + return {"codec": codec_name} + + return { + "codec": codec_name, + "m": _bounded_hnsw_build_parameter( + build_params.get("m", _DEFAULT_M), "m" + ), + "ef_construction": _bounded_hnsw_build_parameter( + build_params.get("ef_construction", _DEFAULT_EF_CONSTRUCTION), + "ef_construction", + ), + "direct_single_segment": _direct_single_segment( + build_params.get("direct_single_segment", False) + ), + } + + +def _normalize_search_params( + search_params: Any, + *, + codec_name: str, + k: Optional[int] = None, +) -> List[Dict[str, Any]]: + if not isinstance(search_params, list) or not search_params: + raise ValueError( + "PyLucene search parameters must be a non-empty list of mappings" + ) + + normalized = [] + for parameters in search_params: + if not isinstance(parameters, dict): + raise TypeError( + "Every PyLucene search parameter must be a mapping" + ) + unsupported = set(parameters) - _SUPPORTED_SEARCH_KEYS + if unsupported: + names = ", ".join(sorted(str(name) for name in unsupported)) + raise ValueError( + f"Unsupported PyLucene search parameter(s): {names}" + ) + if codec_name != _HNSW_CODEC and parameters: + raise ValueError( + f"PyLucene num_candidates applies only to {_HNSW_CODEC}" + ) + if not parameters: + if codec_name == _HNSW_CODEC and k is not None: + normalized.append({"num_candidates": k}) + else: + normalized.append({}) + continue + + num_candidates = parameters["num_candidates"] + if type(num_candidates) is not int or num_candidates < 1: + raise ValueError( + "PyLucene num_candidates must be a positive integer, " + f"got {num_candidates!r}" + ) + if k is not None and num_candidates < k: + raise ValueError( + "PyLucene num_candidates must be greater than or equal to " + f"k ({k}), got {num_candidates}" + ) + normalized.append({"num_candidates": num_candidates}) + return normalized + + +def _safe_remove_index(index_path: Path, trusted_index_root: 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 resolved.parent != trusted_index_root.resolve(): + raise ValueError( + f"Refusing to remove PyLucene index outside its configured root: {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( + f"Lucene commit fingerprint target must be a regular file: {path}" + ) + fingerprints.append({"name": path.name, "sha256": _sha256_file(path)}) + return fingerprints + + +@dataclass(frozen=True) +class _IndexProvenanceVerification: + codec: str + build_parameters: Dict[str, Any] + writer_policy: str + compound_file_policy: str + vector_count: int + dimensions: int + segment_count: int + commit_file_count: int + + def to_metadata(self) -> Dict[str, Any]: + return { + "status": f"{self.writer_policy}-provenance", + "schema_version": _PROVENANCE_SCHEMA_VERSION, + "codec": self.codec, + "build_parameters": self.build_parameters, + "writer_policy": self.writer_policy, + "compound_file_policy": self.compound_file_policy, + "vector_count": self.vector_count, + "dimensions": self.dimensions, + "segment_count": self.segment_count, + "commit_file_count": self.commit_file_count, + } + + +@dataclass(frozen=True) +class _ProvenanceExpectation: + codec: str + build_parameters: Dict[str, Any] + 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, + build_parameters: Dict[str, Any], + vector_count: int, + dimensions: int, + segment_count: int, + manifest_name: str, +) -> None: + payload = { + "schema_version": _PROVENANCE_SCHEMA_VERSION, + "codec": codec, + "build_parameters": build_parameters, + "writer_policy": _EXPECTED_WRITER_POLICY[codec], + "compound_file_policy": _COMPOUND_FILE_POLICY[codec], + "vector_count": int(vector_count), + "dimensions": int(dimensions), + "segment_count": int(segment_count), + "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, + build_parameters: Optional[Dict[str, Any]] = None, + segment_count: int = 1, +) -> None: + effective_parameters = _normalize_build_params( + build_parameters or {"codec": codec} + ) + _write_index_provenance( + index_path, + codec, + effective_parameters, + vector_count, + dimensions, + segment_count, + _HNSW_PROVENANCE_FILE, + ) + + +def _write_cagra_provenance( + index_path: Path, + vector_count: int, + dimensions: int, + build_parameters: Optional[Dict[str, Any]] = None, + segment_count: int = 1, +) -> None: + effective_parameters = _normalize_build_params( + build_parameters or {"codec": _CAGRA_CODEC} + ) + _write_index_provenance( + index_path, + _CAGRA_CODEC, + effective_parameters, + vector_count, + dimensions, + segment_count, + _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: {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_build_parameters: Dict[str, Any], + provenance_label: str, +) -> Tuple[Dict[str, Any], 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" + ) + expected_codec = expected_build_parameters["codec"] + if payload["codec"] != expected_codec: + raise _IndexProvenanceError( + f"{provenance_label} codec does not match the requested codec: " + f"{payload['codec']!r} != {expected_codec!r}" + ) + try: + stored_build_parameters = _normalize_build_params( + payload["build_parameters"] + ) + except (TypeError, ValueError) as exc: + raise _IndexProvenanceError( + f"{provenance_label} has invalid build parameters: {exc}" + ) from exc + if stored_build_parameters != payload["build_parameters"]: + raise _IndexProvenanceError( + f"{provenance_label} build parameters are not canonical" + ) + if stored_build_parameters != expected_build_parameters: + raise _IndexProvenanceError( + f"{provenance_label} build parameters do not match the " + "requested index configuration" + ) + + writer_policy = payload["writer_policy"] + if writer_policy != _EXPECTED_WRITER_POLICY[expected_codec]: + raise _IndexProvenanceError( + f"{provenance_label} does not record the expected writer policy" + ) + compound_file_policy = payload["compound_file_policy"] + if compound_file_policy != _COMPOUND_FILE_POLICY[expected_codec]: + raise _IndexProvenanceError( + f"{provenance_label} does not record the expected compound-file policy" + ) + return stored_build_parameters, writer_policy, compound_file_policy + + +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_provenance_segment_count( + payload: Dict[str, Any], provenance_label: str +) -> int: + return _require_positive_int( + payload["segment_count"], provenance_label, "segment_count" + ) + + +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 + ) + ( + build_parameters, + writer_policy, + compound_file_policy, + ) = _validate_provenance_identity( + payload, expectation.build_parameters, expectation.label + ) + vector_count, dimensions = _validate_provenance_shape( + payload, + expectation.label, + expectation.vector_count, + expectation.dimensions, + ) + segment_count = _validate_provenance_segment_count( + payload, expectation.label + ) + if build_parameters.get("direct_single_segment") and segment_count != 1: + raise _IndexProvenanceError( + f"{expectation.label} direct_single_segment provenance must " + "record exactly one segment" + ) + 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, + build_parameters=build_parameters, + writer_policy=writer_policy, + compound_file_policy=compound_file_policy, + vector_count=vector_count, + dimensions=dimensions, + segment_count=segment_count, + commit_file_count=len(current_fingerprints), + ) + + +def _verify_hnsw_provenance( + index_path: Path, + expected_codec: str, + *, + expected_build_parameters: Optional[Dict[str, Any]] = None, + expected_vector_count: Optional[int] = None, + expected_dimensions: Optional[int] = None, +) -> _IndexProvenanceVerification: + build_parameters = _normalize_build_params( + expected_build_parameters or {"codec": expected_codec} + ) + return _verify_index_provenance( + index_path, + _ProvenanceExpectation( + codec=expected_codec, + build_parameters=build_parameters, + manifest_name=_HNSW_PROVENANCE_FILE, + label="HNSW provenance", + vector_count=expected_vector_count, + dimensions=expected_dimensions, + ), + ) + + +def _verify_cagra_provenance( + index_path: Path, + *, + expected_build_parameters: Optional[Dict[str, Any]] = None, + expected_vector_count: Optional[int] = None, + expected_dimensions: Optional[int] = None, +) -> _IndexProvenanceVerification: + build_parameters = _normalize_build_params( + expected_build_parameters or {"codec": _CAGRA_CODEC} + ) + return _verify_index_provenance( + index_path, + _ProvenanceExpectation( + codec=_CAGRA_CODEC, + build_parameters=build_parameters, + 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( + f"training_vectors must use float32 values, 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 + build_parameters: Dict[str, Any] + search_parameters: Dict[str, Any] + num_candidates: int + k: int + batch_size: int + mode: str + + +@dataclass(frozen=True) +class _BuildCodec: + codec_name: str + java_codec: Any + writer_policy: str + build_parameters: Dict[str, Any] + + @property + def direct_single_segment(self) -> bool: + return bool(self.build_parameters.get("direct_single_segment", False)) + + +@dataclass(frozen=True) +class _IndexTopology: + segment_document_counts: Tuple[int, ...] + segment_vector_counts: Tuple[int, ...] + + @property + def segment_count(self) -> int: + return len(self.segment_document_counts) + + def validate(self, expected_vector_count: int) -> None: + if not self.segment_document_counts: + raise RuntimeError("Lucene index contains no committed segments") + if sum(self.segment_document_counts) != expected_vector_count: + raise RuntimeError( + "Lucene segment document counts do not match the indexed " + f"vectors: {sum(self.segment_document_counts)} != " + f"{expected_vector_count}" + ) + if self.segment_document_counts != self.segment_vector_counts: + raise RuntimeError( + "Lucene segment document and vector counts do not match: " + f"{self.segment_document_counts} != " + f"{self.segment_vector_counts}" + ) + + def require_direct_single_segment( + self, expected_vector_count: int + ) -> None: + self.validate(expected_vector_count) + if self.segment_count != 1: + raise RuntimeError( + "direct_single_segment requested one committed Lucene " + f"segment, found {self.segment_count}" + ) + + def to_metadata(self) -> Dict[str, Any]: + return { + "segment_count": self.segment_count, + "document_count": sum(self.segment_document_counts), + "vector_count": sum(self.segment_vector_counts), + "segment_document_counts": list(self.segment_document_counts), + "segment_vector_counts": list(self.segment_vector_counts), + } + + +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( + f"PyLucene returned a non-finite score for 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}: {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( + f"CAGRA data file contains unreferenced payload: {data_file!r}" + ) + return + + coverage_error = f"CAGRA metadata ranges do not exactly cover data 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( + f"CAGRA-only verification cannot read {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, System + 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, + NoMergePolicy, + SegmentCommitInfo, + SegmentInfos, + 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.System = System + 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.NoMergePolicy = NoMergePolicy + 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) + self._validate_codec(codec, codec_name) + self._codec_cache[codec_name] = codec + return codec + + def resolve_configured_hnsw_codec( + self, m: int, ef_construction: int + ) -> Any: + self.attach_current_thread() + configured_values = { + M_PROPERTY: str(m), + EF_CONSTRUCTION_PROPERTY: str(ef_construction), + } + with _CONFIGURED_CODEC_LOCK: + with _CleanupStack() as cleanups: + for name, value in configured_values.items(): + previous = self.System.getProperty(name) + cleanups.add( + f"restore Java system property {name}", + lambda name=name, previous=previous: ( + _restore_java_property(self.System, name, previous) + ), + ) + self.System.setProperty(name, value) + reflected_codec = self.Class.forName( + CONFIGURED_HNSW_CODEC_CLASS + ).newInstance() + codec = self.Codec.cast_(reflected_codec) + self._validate_codec(codec, _HNSW_CODEC) + expected_diagnostics = f"PyLuceneConfiguredHnswCodec(m={m}, efConstruction={ef_construction})" + if str(codec) != expected_diagnostics: + raise RuntimeError( + "Configured PyLucene codec did not retain the requested " + f"parameters: expected {expected_diagnostics!r}, " + f"got {str(codec)!r}" + ) + return codec + + @staticmethod + def _validate_codec(codec: Any, codec_name: str) -> None: + if str(codec.getName()) != codec_name: + raise RuntimeError( + f"Requested codec {codec_name}, got {codec.getName()}" + ) + if codec.knnVectorsFormat() is None: + raise RuntimeError( + f"{codec_name} did not initialize a Lucene vector 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, + build_codec: _BuildCodec, + ) -> _IndexTopology: + 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._new_index_writer_config( + build_codec, vector_count=int(vectors.shape[0]) + ) + writer = self.IndexWriter(directory, writer_config) + self._write_and_close_index(writer, vectors) + topology = self._index_topology(directory) + if build_codec.direct_single_segment: + topology.require_direct_single_segment(int(vectors.shape[0])) + else: + topology.validate(int(vectors.shape[0])) + return topology + + def _new_index_writer_config( + self, build_codec: _BuildCodec, *, vector_count: int + ) -> Any: + writer_config = self.IndexWriterConfig() + writer_config.setOpenMode(self.IndexWriterConfig.OpenMode.CREATE) + writer_config.setCodec(build_codec.java_codec) + if build_codec.codec_name == _CAGRA_CODEC: + # The CAGRA verifier reads the codec's .vemc and .vcag files. + # Lucene controls compound files separately for flushes and merges. + writer_config.setUseCompoundFile(False) + writer_config.getMergePolicy().setNoCFSRatio(0.0) + if build_codec.direct_single_segment: + writer_config.setMaxBufferedDocs(vector_count + 1) + writer_config.setRAMBufferSizeMB( + float(self.IndexWriterConfig.DISABLE_AUTO_FLUSH) + ) + writer_config.setMergePolicy(self.NoMergePolicy.INSTANCE) + return writer_config + + def _index_topology(self, directory: Any) -> _IndexTopology: + reader = self.DirectoryReader.open(directory) + with _CleanupStack() as cleanups: + cleanups.add("close Lucene topology reader", reader.close) + document_counts = [] + vector_counts = [] + for leaf_context in reader.leaves(): + leaf_reader = leaf_context.reader() + values = leaf_reader.getFloatVectorValues(_VECTOR_FIELD) + if values is None: + raise RuntimeError( + "Lucene segment contains no vector values for " + f"{_VECTOR_FIELD!r}" + ) + document_counts.append(int(leaf_reader.numDocs())) + vector_counts.append(int(values.size())) + return _IndexTopology( + segment_document_counts=tuple(document_counts), + segment_vector_counts=tuple(vector_counts), + ) + + 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( + f"Lucene index has inconsistent vector dimensions: {dimensions}" + ) + return dimensions.pop() + + def _search_vector( + self, + searcher: Any, + stored_fields: Any, + vector: np.ndarray, + k: int, + num_candidates: int, + ) -> List[_SearchHit]: + query = self.KnnFloatVectorQuery( + _VECTOR_FIELD, + self._java_float_array(vector), + num_candidates, + ) + 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, + num_candidates: Optional[int] = None, + ) -> _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) + lucene_num_candidates = min( + num_candidates if num_candidates is not None else 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, + lucene_num_candidates, + ) + ) + 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, + num_candidates: int, + ) -> List[List[_SearchHit]]: + hits = [] + for vector in query_vectors: + hits.append( + self._search_vector( + searcher, + stored_fields, + vector, + k, + num_candidates, + ) + ) + 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] + + +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 + if not build_combinations: + raise ValueError("PyLucene tune mode requires build parameters") + build_params = { + **build_combinations[0], + **tune_build_params, + } + search_params = [tune_search_params] if tune_search_params else [{}] + return [build_params], search_params + + @staticmethod + def _index_label( + algorithm: str, + group: str, + subset_scope: Optional[str], + build_params: Dict[str, Any], + ) -> str: + identity = format_artifact_identity( + algorithm, group, subset_scope, description="PyLucene" + ) + parameters = _normalize_build_params(build_params) + label = f"{identity}[codec={parameters['codec']}]" + if parameters["codec"] == _HNSW_CODEC: + direct_single_segment = str( + parameters["direct_single_segment"] + ).lower() + label += ( + f"[m={parameters['m']}]" + f"[ef_construction={parameters['ef_construction']}]" + f"[direct_single_segment={direct_single_segment}]" + ) + return label + + @classmethod + def _benchmark_config( + cls, + context: _BenchmarkConfigContext, + algorithm: str, + group: str, + build_params: Dict[str, Any], + search_params: List[Dict[str, Any]], + ) -> BenchmarkConfig: + build_params = _normalize_build_params(build_params) + search_params = _normalize_search_params( + search_params, codec_name=build_params["codec"] + ) + dataset = validate_path_component( + context.dataset, "PyLucene dataset name" + ) + index_label = cls._index_label( + algorithm, group, context.subset_scope, build_params + ) + index_root = Path(context.dataset_path, dataset, "index") + index_path = index_root / index_label + index_config = IndexConfig( + name=index_label, + algo=algorithm, + build_param=build_params, + search_params=search_params, + file=str(index_path), + ) + backend_config = { + "name": index_label, + "algo": algorithm, + "group": group, + "index_name": index_label, + "index_root": str(index_root), + "result_scope": context.subset_scope, + "codec": build_params["codec"], + "requires_gpu": build_params["codec"] == _CAGRA_CODEC, + **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]: + benchmark_configs = [] + for params in build_params: + effective_params = _normalize_build_params(params) + benchmark_configs.append( + cls._benchmark_config( + context, + algorithm, + group, + effective_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), + ) + 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.""" + + 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 _result_identity(self) -> Dict[str, Any]: + identity = { + "group": self.config.get("group", "base"), + "index_name": self.config.get( + "index_name", self.config.get("name", self.algo) + ), + } + result_scope = self.config.get("result_scope") + if result_scope is not None: + identity["result_scope"] = result_scope + return identity + + def _index_metadata( + self, build_parameters: Dict[str, Any] + ) -> Dict[str, Any]: + codec_name = build_parameters["codec"] + return { + **self._result_identity(), + "compound_file_policy": _COMPOUND_FILE_POLICY[codec_name], + **build_parameters, + } + + 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 _trusted_index_root(self, index_path: Path) -> Path: + configured_root = self.config.get("index_root") + if configured_root is None: + # Direct backend users explicitly supply IndexConfig.file. The + # built-in loader supplies an independent root for CLI runs. + return index_path.parent + return Path(os.fspath(configured_root)) + + def _validate_index_location(self, index_path: Path) -> None: + trusted_root = self._trusted_index_root(index_path).resolve() + if index_path.resolve().parent != trusted_root: + raise ValueError( + "PyLucene index path must be an immediate child of its " + f"configured root {trusted_root}: {index_path.resolve()}" + ) + + def _remove_index(self, index_path: Path) -> None: + _safe_remove_index( + index_path, trusted_index_root=self._trusted_index_root(index_path) + ) + + 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 {}, + metadata=self._result_identity(), + 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 [], + metadata=self._result_identity(), + success=False, + error_message=error_message, + ) + + def _verify_existing_index( + self, + index_path: Path, + build_parameters: Dict[str, Any], + *, + expected_vector_count: Optional[int] = None, + expected_dimensions: Optional[int] = None, + ) -> Tuple[_IndexProvenanceVerification, Dict[str, Any]]: + """Verify backend ownership and codec-specific persisted data.""" + codec_name = build_parameters["codec"] + if codec_name == _CAGRA_CODEC: + provenance = _verify_cagra_provenance( + index_path, + expected_build_parameters=build_parameters, + 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_build_parameters=build_parameters, + 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, + build_parameters: Dict[str, Any], + ) -> BuildResult: + print( + f"[dry_run] Would build PyLucene index '{index_path}' " + f"with {build_parameters}" + ) + return BuildResult( + index_path=str(index_path), + build_time_seconds=0.0, + index_size_bytes=0, + algorithm=self.algo, + build_params=build_parameters, + metadata=self._index_metadata(build_parameters), + success=True, + ) + + def _decide_existing_index( + self, + dataset: Dataset, + index_path: Path, + build_parameters: 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_parameters, + ) + ) + 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_parameters, + ) + ) + if force: + return _ExistingIndexDecision.build() + + metadata: Dict[str, Any] = { + **self._index_metadata(build_parameters), + "skipped": True, + } + try: + _validate_metric(dataset) + expected_shape = _expected_training_shape(dataset) + provenance, verification_metadata = self._verify_existing_index( + index_path, + build_parameters, + expected_vector_count=( + expected_shape[0] if expected_shape else None + ), + expected_dimensions=( + expected_shape[1] if expected_shape else None + ), + ) + metadata["segment_count"] = provenance.segment_count + 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_parameters, + metadata=metadata, + success=True, + ) + except Exception as exc: + return _ExistingIndexDecision.reject( + self._failed_build_result( + _exception_summary(exc), + str(index_path), + build_parameters, + ) + ) + + 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 _resolve_build_codec( + runtime: _PyLuceneRuntime, build_parameters: Dict[str, Any] + ) -> _BuildCodec: + codec_name = build_parameters["codec"] + if codec_name == _HNSW_CODEC: + java_codec = runtime.resolve_configured_hnsw_codec( + build_parameters["m"], + build_parameters["ef_construction"], + ) + else: + java_codec = runtime.resolve_codec(codec_name) + return _BuildCodec( + codec_name=codec_name, + java_codec=java_codec, + writer_policy=_EXPECTED_WRITER_POLICY[codec_name], + build_parameters=build_parameters, + ) + + @staticmethod + def _verify_and_persist_built_index_provenance( + runtime: _PyLuceneRuntime, + index_path: Path, + build_parameters: Dict[str, Any], + vectors: np.ndarray, + topology: _IndexTopology, + ) -> Dict[str, Any]: + codec_name = build_parameters["codec"] + 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, + build_parameters=build_parameters, + segment_count=topology.segment_count, + ) + provenance = _verify_cagra_provenance( + index_path, + expected_build_parameters=build_parameters, + 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, + build_parameters=build_parameters, + segment_count=topology.segment_count, + ) + verification = _verify_hnsw_provenance( + index_path, + codec_name, + expected_build_parameters=build_parameters, + expected_vector_count=vector_count, + expected_dimensions=dimensions, + ) + return {"hnsw_verification": verification.to_metadata()} + + def _cleanup_partial_index( + self, index_path: Path, created_for_build: bool + ) -> Optional[Exception]: + if not created_for_build or not index_path.exists(): + return None + try: + self._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] + requested_build_params = index_config.build_param + index_path = Path(index_config.file) + try: + build_parameters = _normalize_build_params( + requested_build_params, self.config + ) + self._validate_index_location(index_path) + except Exception as exc: + return self._failed_build_result( + str(exc), str(index_path), requested_build_params + ) + + if dry_run: + return self._dry_run_build_result(index_path, build_parameters) + + existing_index_decision = self._decide_existing_index( + dataset, index_path, build_parameters, force + ) + if existing_index_decision.action is _ExistingIndexAction.BUILD: + return self._build_new_index(dataset, index_path, build_parameters) + + return existing_index_decision.completed_result() + + def _build_new_index( + self, + dataset: Dataset, + index_path: Path, + build_parameters: Dict[str, Any], + ) -> BuildResult: + """Build a validated index that is not eligible for reuse.""" + + created_for_build = False + try: + codec_name = build_parameters["codec"] + vectors = self._training_vectors_for_build(dataset, codec_name) + runtime = self._get_runtime() + # Preflight must succeed before an existing index is replaced. + build_codec = self._resolve_build_codec(runtime, build_parameters) + + if index_path.exists(): + self._remove_index(index_path) + index_path.mkdir(parents=True) + created_for_build = True + + start = time.perf_counter() + topology = runtime.build_index(index_path, vectors, build_codec) + build_time = time.perf_counter() - start + + metadata = { + **self._index_metadata(build_parameters), + "pylucene_version": runtime.pylucene_version, + "writer_policy": build_codec.writer_policy, + **topology.to_metadata(), + } + metadata.update( + self._verify_and_persist_built_index_provenance( + runtime, + index_path, + build_parameters, + vectors, + topology, + ) + ) + + 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_parameters, + 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_parameters, + ) + 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_plans( + self, + index_config: IndexConfig, + *, + k: int, + batch_size: int, + mode: str, + search_threads: Optional[Union[int, str]], + ) -> List[_SearchPlan]: + build_parameters = _normalize_build_params( + index_config.build_param, self.config + ) + codec_name = build_parameters["codec"] + 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" + ) + search_parameters = _normalize_search_params( + index_config.search_params or [{}], codec_name=codec_name, k=k + ) + 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, + build_parameters=build_parameters, + search_parameters=parameters, + num_candidates=parameters.get("num_candidates", k), + k=k, + batch_size=batch_size, + mode=mode, + ) + for parameters in search_parameters + ] + + 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"num_candidates={plan.num_candidates}, " + 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_parameters], + metadata={ + **self._index_metadata(plan.build_parameters), + "top_k": plan.k, + "num_candidates": plan.num_candidates, + "batch_size": plan.batch_size, + "mode": plan.mode, + }, + 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.build_parameters, + 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, + plan.num_candidates, + ) + 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 = { + **self._index_metadata(plan.build_parameters), + "pylucene_version": runtime.pylucene_version, + "index_dimensions": runtime_result.index_dimensions, + "document_count": runtime_result.document_count, + "segment_count": provenance.segment_count, + "top_k": plan.k, + "num_candidates": plan.num_candidates, + "batch_size": plan.batch_size, + "num_batches": processed.num_batches, + "mode": plan.mode, + "latency_seconds": processed.latency_seconds, + "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_parameters], + 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, + ) -> List[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: + plans = self._resolve_search_plans( + 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) for plan in plans] + + index_path = plans[0].index_path + if not index_path.is_dir(): + return [ + self._failed_search_result( + k, + f"PyLucene index directory does not exist: {index_path}", + [plan.search_parameters for plan in plans], + ) + ] + + results = [] + for plan in plans: + try: + results.append(self._execute_search(dataset, plan)) + except Exception as exc: + results.append( + self._failed_search_result( + k, + _exception_summary(exc), + [plan.search_parameters], + ) + ) + return results + + +__all__ = ["PyLuceneBackend", "PyLuceneConfigLoader"] diff --git a/python/cuvs_bench/cuvs_bench/backends/search_spaces.py b/python/cuvs_bench/cuvs_bench/backends/search_spaces.py index edb488f77e..e601abbb39 100644 --- a/python/cuvs_bench/cuvs_bench/backends/search_spaces.py +++ b/python/cuvs_bench/cuvs_bench/backends/search_spaces.py @@ -159,6 +159,22 @@ }, }, # ========================================================================= + # PyLucene/cuVS HNSW + # ========================================================================= + "pylucene_cuvs_hnsw": { + "build": { + "m": {"type": "int", "min": 1, "max": 512}, + "ef_construction": {"type": "int", "min": 1, "max": 512}, + }, + "search": { + "num_candidates": { + "type": "int", + "min": "top_k", + "max": 500, + }, + }, + }, + # ========================================================================= # Elasticsearch GPU HNSW (hnsw, int8_hnsw, int4_hnsw, bbq_hnsw) # Per ES-GPU-API-REFERENCE.md: index_options (m, ef_construction), # knn (num_candidates) 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..16c5bab34b --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_hnsw.yaml @@ -0,0 +1,20 @@ +name: pylucene_cuvs_hnsw +groups: + base: + build: + codec: + - "Lucene101AcceleratedHNSWCodec" + m: + - 32 + ef_construction: + - 32 + direct_single_segment: + - false + search: {} + test: + build: + codec: ["Lucene101AcceleratedHNSWCodec"] + m: [32] + ef_construction: [32] + direct_single_segment: [false] + 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..7d6375ce96 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 # @@ -362,7 +362,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 +373,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 abffad0217..826742bef3 100644 --- a/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py +++ b/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py @@ -10,7 +10,7 @@ benchmark runs across different backends using the registry pattern. """ -from typing import List, Optional, Union +from typing import Any, List, Mapping, Optional, Union import numpy as np @@ -29,6 +29,17 @@ def _should_compute_recall(result: SearchResult) -> bool: return result.success and result.neighbors.size > 0 +def _resolve_tune_bound( + bound: Any, bound_values: Optional[Mapping[str, Any]] +) -> Any: + """Resolve a symbolic tune bound from the current trial context.""" + if not isinstance(bound, str): + return bound + if bound_values is None or bound not in bound_values: + raise ValueError(f"Unable to resolve symbolic tune bound {bound!r}") + return bound_values[bound] + + class BenchmarkOrchestrator: """ Orchestrator for running benchmarks using the pluggable backend system. @@ -389,18 +400,23 @@ def _run_tune( all_results: List[Union[BuildResult, SearchResult]] = [] def suggest_params( - trial, param_space: dict, build_params: dict = None + trial, + param_space: dict, + bound_values: Optional[Mapping[str, Any]] = None, ) -> dict: """Suggest parameters from search space using Optuna trial.""" params = {} for param, spec in param_space.items(): if spec["type"] == "int": - max_val = spec["max"] - # Handle dynamic constraints (e.g., nprobe <= nlist) - if isinstance(max_val, str) and build_params: - max_val = build_params.get(max_val, 1000) + min_val = _resolve_tune_bound(spec["min"], bound_values) + max_val = _resolve_tune_bound(spec["max"], bound_values) + if min_val > max_val: + raise ValueError( + f"Invalid tune range for {param!r}: minimum " + f"{min_val} exceeds maximum {max_val}" + ) params[param] = trial.suggest_int( - param, spec["min"], max_val, log=spec.get("log", False) + param, min_val, max_val, log=spec.get("log", False) ) elif spec["type"] == "float": params[param] = trial.suggest_float( @@ -421,8 +437,9 @@ def objective(trial) -> float: build_params = suggest_params(trial, search_space.get("build", {})) # Suggest search parameters (may depend on build params) + search_bound_values = {**build_params, "top_k": count} search_params_dict = suggest_params( - trial, search_space.get("search", {}), build_params + trial, search_space.get("search", {}), search_bound_values ) # Run single trial with these specific parameters diff --git a/python/cuvs_bench/cuvs_bench/run/__main__.py b/python/cuvs_bench/cuvs_bench/run/__main__.py index 225a3d3f9c..4344550e76 100644 --- a/python/cuvs_bench/cuvs_bench/run/__main__.py +++ b/python/cuvs_bench/cuvs_bench/run/__main__.py @@ -14,11 +14,41 @@ from .data_export import ( convert_json_to_csv_build, convert_json_to_csv_search, + validate_dataset_name, write_results_to_csv, ) +from ..backends.base import SearchResult from ..orchestrator import BenchmarkOrchestrator +def _run_failed(results, mode, *, allow_empty=False): + if not results: + return not allow_empty + if mode == "sweep": + return any(not result.success for result in results) + return not any( + isinstance(result, SearchResult) and result.success + for result in results + ) + + +def _raise_for_failed_run(results, mode, *, allow_empty=False): + if not _run_failed(results, mode, allow_empty=allow_empty): + return + + failures = [result for result in results if not result.success] + if failures: + details = "; ".join( + f"{result.algorithm}: {result.error_message or 'benchmark failed'}" + for result in failures + ) + elif results: + details = f"{mode} mode produced no successful search result" + else: + details = f"{mode} mode produced no benchmark results" + raise click.ClickException(details) + + @click.command() @click.option( "--subset-size", @@ -257,6 +287,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 data_export: click.echo( "Warning: --data-export is deprecated because benchmark runs now " @@ -312,6 +347,10 @@ def main( ) if dry_run: + # Native dry runs print planned commands without result objects. + _raise_for_failed_run( + results, mode, allow_empty=backend_type == "cpp_gbench" + ) return if backend_type == "cpp_gbench": @@ -320,7 +359,16 @@ def main( if search: convert_json_to_csv_search(dataset, dataset_path) else: - write_results_to_csv(results, dataset, dataset_path, count, batch_size) + write_results_to_csv( + results, + dataset, + dataset_path, + count, + batch_size, + search_requested=search, + ) + + _raise_for_failed_run(results, mode) if __name__ == "__main__": diff --git a/python/cuvs_bench/cuvs_bench/run/data_export.py b/python/cuvs_bench/cuvs_bench/run/data_export.py index 50afa7a5ea..393103ed5a 100644 --- a/python/cuvs_bench/cuvs_bench/run/data_export.py +++ b/python/cuvs_bench/cuvs_bench/run/data_export.py @@ -10,6 +10,11 @@ import pandas as pd +from .._validation import ( + format_artifact_identity, + validate_path_component, + validate_result_component, +) from ..backends.base import BuildResult, SearchResult skip_build_cols = set( @@ -53,8 +58,37 @@ } -def write_results_to_csv(results, dataset, dataset_path, count, batch_size): +def validate_dataset_name(dataset): + """Validate a dataset name before using it as a path component.""" + return validate_path_component(dataset, "benchmark dataset name") + + +def _result_stem(algorithm, group, scope=None): + parts = [ + validate_result_component(algorithm, "benchmark algorithm name"), + validate_result_component(group, "benchmark group name"), + ] + if scope is not None: + parts.append( + validate_result_component(scope, "benchmark result scope") + ) + return ",".join(parts) + + +def _series_name(algorithm, group, scope): + return format_artifact_identity(algorithm, group, scope) + + +def write_results_to_csv( + results, + dataset, + dataset_path, + count, + batch_size, + search_requested=False, +): """Write Python-backend results using the existing plotting CSV schema.""" + validate_dataset_name(dataset) grouped = defaultdict(list) for result in results: group = result.metadata.get("group") @@ -64,29 +98,43 @@ def write_results_to_csv(results, dataset, dataset_path, count, batch_size): ): continue method = "build" if isinstance(result, BuildResult) else "search" - grouped[(method, result.algorithm, group)].append(result) + scope = result.metadata.get("result_scope") + grouped[(method, result.algorithm, group, scope)].append(result) - for (method, algorithm, group), group_results in grouped.items(): + for (method, algorithm, group, scope), group_results in grouped.items(): if method == "build": _write_build_results( - group_results, algorithm, group, dataset, dataset_path + group_results, + algorithm, + group, + scope, + dataset, + dataset_path, ) else: _write_search_results( group_results, algorithm, group, + scope, dataset, dataset_path, count, batch_size, ) + if search_requested: + _remove_search_artifacts_after_failed_builds( + grouped, dataset, dataset_path, count, batch_size + ) + -def _write_build_results(results, algorithm, group, dataset, dataset_path): +def _write_build_results( + results, algorithm, group, scope, dataset, dataset_path +): output_dir = os.path.join(dataset_path, dataset, "result", "build") os.makedirs(output_dir, exist_ok=True) - algo_name = algorithm if group == "base" else f"{algorithm}_{group}" + algo_name = _series_name(algorithm, group, scope) rows = [] for result in results: @@ -98,14 +146,18 @@ def _write_build_results(results, algorithm, group, dataset, dataset_path): **result.build_params, **metadata, "algo_name": algo_name, - "index_name": result.index_path, + "index_name": result.metadata.get( + "index_name", result.index_path + ), "time": result.build_time_seconds, } ) columns = ["algo_name", "index_name", "time"] dataframe = pd.DataFrame(rows) - build_file = os.path.join(output_dir, f"{algorithm},{group}.csv") + build_file = os.path.join( + output_dir, f"{_result_stem(algorithm, group, scope)}.csv" + ) complete_run = all( result.success and not result.metadata.get("skipped") @@ -131,11 +183,18 @@ def _write_build_results(results, algorithm, group, dataset, dataset_path): def _write_search_results( - results, algorithm, group, dataset, dataset_path, count, batch_size + results, + algorithm, + group, + scope, + dataset, + dataset_path, + count, + batch_size, ): output_dir = os.path.join(dataset_path, dataset, "result", "search") os.makedirs(output_dir, exist_ok=True) - algo_name = algorithm if group == "base" else f"{algorithm}_{group}" + algo_name = _series_name(algorithm, group, scope) rows = [] for result in results: @@ -148,6 +207,7 @@ def _write_search_results( rows.append( { **search_params, + **(result.latency_percentiles or {}), **metadata, "algo_name": algo_name, "index_name": result.metadata["index_name"], @@ -179,7 +239,7 @@ def _write_search_results( dataset, "result", "build", - f"{algorithm},{group}.csv", + f"{_result_stem(algorithm, group, scope)}.csv", ) if os.path.exists(build_file): build = pd.read_csv(build_file).drop_duplicates( @@ -194,7 +254,7 @@ def _write_search_results( how="left", ) - stem = f"{algorithm},{group},k{count},bs{batch_size}" + stem = f"{_result_stem(algorithm, group, scope)},k{count},bs{batch_size}" raw_file = os.path.join(output_dir, f"{stem},raw.csv") dataframe.to_csv(raw_file, index=False) frontier_file = os.path.join(output_dir, f"{stem}.json") @@ -203,7 +263,7 @@ def _write_search_results( def _scalar_metadata(metadata): - reserved = {"group", "index_name", "latency_seconds"} + reserved = {"group", "index_name", "latency_seconds", "result_scope"} return { key: value for key, value in metadata.items() @@ -212,6 +272,26 @@ def _scalar_metadata(metadata): } +def _remove_search_artifacts_after_failed_builds( + grouped, dataset, dataset_path, count, batch_size +): + search_dir = os.path.join(dataset_path, dataset, "result", "search") + for (method, algorithm, group, scope), results in grouped.items(): + if method != "build" or any(result.success for result in results): + continue + search_key = ("search", algorithm, group, scope) + if search_key in grouped: + continue + stem = ( + f"{_result_stem(algorithm, group, scope)},k{count},bs{batch_size}" + ) + for suffix in (",raw.csv", ",throughput.csv", ",latency.csv"): + try: + os.remove(os.path.join(search_dir, f"{stem}{suffix}")) + except FileNotFoundError: + pass + + def read_json_files(dataset, dataset_path, method): """ Yield file paths, algo names, and loaded JSON data as pandas DataFrames. @@ -231,6 +311,7 @@ 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 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..0b9cf6dc90 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/_pylucene_test_utils.py @@ -0,0 +1,244 @@ +# +# 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, + _IndexTopology, + _RuntimeSearchResult, + _SearchHit, +) +from cuvs_bench.orchestrator.config_loaders import IndexConfig + +_HNSW_CODEC = "Lucene101AcceleratedHNSWCodec" +_CAGRA_CODEC = "CuVS2510GPUSearchCodec" + + +class _FakeRuntime: + pylucene_version = "test" + + def __init__( + self, + *, + 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.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.configured_resolve_calls = [] + self.build_calls = [] + self.search_calls = [] + self.verification_calls = [] + self.resolve_error = None + self.build_error = None + self.topology = 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 resolve_configured_hnsw_codec(self, m, ef_construction): + self.configured_resolve_calls.append((m, ef_construction)) + if self.resolve_error is not None: + raise self.resolve_error + return (_HNSW_CODEC, m, ef_construction) + + def build_index(self, index_path, vectors, build_codec): + self.build_calls.append((index_path, vectors.copy(), build_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") + topology = self.topology or _IndexTopology( + segment_document_counts=(self.document_count,), + segment_vector_counts=(self.document_count,), + ) + if build_codec.direct_single_segment: + topology.require_direct_single_segment(self.document_count) + else: + topology.validate(self.document_count) + return topology + + 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, num_candidates=None + ): + self.search_calls.append( + ( + index_path, + query_vectors.copy(), + k, + batch_size, + num_candidates, + ) + ) + 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, + build_params: dict | None = None, + search_params: list[dict] | None = None, +) -> IndexConfig: + parameters = {"codec": codec} + if build_params: + parameters.update(build_params) + return IndexConfig( + name="pylucene-test", + algo="pylucene_cuvs_hnsw", + build_param=parameters, + 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/pylucene_cpu_fallback_probe.py b/python/cuvs_bench/cuvs_bench/tests/pylucene_cpu_fallback_probe.py new file mode 100644 index 0000000000..647e4797fa --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/pylucene_cpu_fallback_probe.py @@ -0,0 +1,111 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Fresh-process probe for the cuVS-Lucene HNSW CPU fallback.""" + +import os +import re +import tempfile +from pathlib import Path + +import lucene +import numpy as np + +from cuvs_bench.backends._pylucene_java import ( + EF_CONSTRUCTION_PROPERTY, + M_PROPERTY, +) +from cuvs_bench.backends.pylucene import ( + _BuildCodec, + _PyLuceneRuntime, + _validate_pylucene_version, +) + +_CPU_HNSW_WRITER = ( + "org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsWriter" +) +_HNSW_CODEC = "Lucene101AcceleratedHNSWCodec" +_WRITER_SELECTION_CODEC = "com.nvidia.cuvs.bench.PyLuceneWriterSelectionCodec" +_DEFAULT_HNSW_BUILD_PARAMETERS = { + "codec": _HNSW_CODEC, + "m": 32, + "ef_construction": 32, + "direct_single_segment": False, +} + + +def _writer_diagnostics(java_codec): + diagnostics = str(java_codec.knnVectorsFormat()) + match = re.search( + r"writerClass=([^,)]+), fieldsWriterCalls=(\d+)", diagnostics + ) + if match is None: + raise AssertionError(f"Unexpected writer diagnostics: {diagnostics}") + return match.group(1), int(match.group(2)) + + +def main(): + _validate_pylucene_version(lucene) + lucene.CLASSPATH = os.pathsep.join( + (os.environ["PYLUCENE_WRITER_SELECTION_CLASSES"], lucene.CLASSPATH) + ) + runtime = _PyLuceneRuntime.create( + { + "cuvs_java_jar": os.environ["CUVS_LUCENE_CUVS_JAVA_JAR"], + "cuvs_lucene_jar": os.environ["CUVS_LUCENE_JAR"], + "java_library_path": os.environ["JAVA_LIBRARY_PATH"], + } + ) + runtime.System.setProperty(M_PROPERTY, "16") + runtime.System.setProperty(EF_CONSTRUCTION_PROPERTY, "48") + reflected_codec = runtime.Class.forName( + _WRITER_SELECTION_CODEC + ).newInstance() + java_codec = runtime.Codec.cast_(reflected_codec) + vectors = ( + np.random.default_rng(174).standard_normal((4, 32)).astype(np.float32) + ) + + with tempfile.TemporaryDirectory(prefix="pylucene-cpu-fallback-") as temp: + index_path = Path(temp) + runtime.build_index( + index_path, + vectors, + _BuildCodec( + codec_name=_HNSW_CODEC, + java_codec=java_codec, + writer_policy="gpu-with-cpu-fallback", + build_parameters={ + **_DEFAULT_HNSW_BUILD_PARAMETERS, + "m": 16, + "ef_construction": 48, + }, + ), + ) + writer_class, writer_calls = _writer_diagnostics(java_codec) + if writer_class != _CPU_HNSW_WRITER: + raise AssertionError( + f"Expected {_CPU_HNSW_WRITER}, found {writer_class}" + ) + + search = runtime.search_index( + index_path, + vectors[[2]], + k=1, + batch_size=1, + num_candidates=1, + ) + first_hit = search.hits[0][0].document_id + + if first_hit != 2: + raise AssertionError(f"Expected document 2, found {first_hit}") + print( + f"writerClass={writer_class} fieldsWriterCalls={writer_calls} " + f"firstHit={first_hit}" + ) + + +if __name__ == "__main__": + main() diff --git a/python/cuvs_bench/cuvs_bench/tests/test_cli.py b/python/cuvs_bench/cuvs_bench/tests/test_cli.py index fa2a63d041..a9dd27e76d 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_cli.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_cli.py @@ -1,13 +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 pathlib import Path +from types import SimpleNamespace import pandas as pd import pytest from click.testing import CliRunner +from cuvs_bench.backends.base import BuildResult, SearchResult from cuvs_bench.get_dataset.__main__ import main @@ -519,11 +521,247 @@ 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, + *, + dry_run=False, + backend_type="fake", +): + 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, + "write_results_to_csv", + lambda *args, **kwargs: exported.append((args, kwargs)), + ) + backend_config = tmp_path / "backend.yaml" + backend_config.write_text(f"backend: {backend_type}\n") + + args = [ + "--dataset", + "test-data", + "--dataset-path", + str(tmp_path), + "--algorithms", + "fake", + "--groups", + "base", + "--batch-size", + "1", + "-k", + "1", + "-m", + "latency", + "--mode", + mode, + "--backend-config", + str(backend_config), + ] + if dry_run: + args.append("--dry-run") + + result = CliRunner().invoke( + run_module.main, + args, + ) + return result, exported + + +def _build_result(success=True, error_message=None): + return BuildResult( + index_path="test-index", + build_time_seconds=1.0, + index_size_bytes=1, + algorithm="fake", + build_params={}, + metadata={"group": "base", "index_name": "test-index"}, + success=success, + error_message=error_message, + ) + + +def _search_result(success=True, error_message=None): + return SearchResult( + neighbors=None, + distances=None, + search_time_ms=1.0, + queries_per_second=1.0, + recall=1.0, + success=success, + algorithm="fake", + search_params=[{}], + metadata={"group": "base", "index_name": "test-index"}, + 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", + [ + _search_result(False, "trial failed"), + _search_result(), + ], + ) + + assert result.exit_code == 0, result.output + assert len(exported) == 1 + assert exported[0][1]["search_requested"] is True + + +@pytest.mark.parametrize( + ("results", "expected_message"), + [ + ([_search_result(False, "trial failed")], "trial failed"), + ([_build_result()], "no successful search result"), + ([], "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 len(exported) == 1 + + +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", + [_search_result(), _search_result()], + ) + + assert result.exit_code == 0, result.output + assert len(exported) == 1 + + +def test_sweep_mixed_results_exit_nonzero(monkeypatch, tmp_path): + result, exported = _invoke_run_with_results( + monkeypatch, + tmp_path, + "sweep", + [ + _build_result(False, "benchmark failed"), + _search_result(), + ], + ) + + assert result.exit_code != 0 + assert "benchmark failed" in result.output + assert len(exported) == 1 + + +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 len(exported) == 1 + + +def test_dry_run_allows_native_backend_without_result_objects( + monkeypatch, tmp_path +): + result, exported = _invoke_run_with_results( + monkeypatch, + tmp_path, + "sweep", + [], + dry_run=True, + backend_type="cpp_gbench", + ) + + assert result.exit_code == 0, result.output + assert exported == [] + + +def test_dry_run_rejects_python_backend_without_result_objects( + monkeypatch, tmp_path +): + result, exported = _invoke_run_with_results( + monkeypatch, tmp_path, "sweep", [], dry_run=True + ) + + assert result.exit_code != 0 + assert "sweep mode produced no benchmark results" in result.output + assert exported == [] + + +def test_dry_run_reports_explicit_backend_failures(monkeypatch, tmp_path): + result, exported = _invoke_run_with_results( + monkeypatch, + tmp_path, + "sweep", + [_build_result(False, "invalid configuration")], + dry_run=True, + ) + + assert result.exit_code != 0 + assert "invalid configuration" 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 index 721d03c5ac..df43613131 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_data_export.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_data_export.py @@ -5,6 +5,7 @@ import numpy as np import pandas as pd +import pytest from cuvs_bench.backends.base import BuildResult, SearchResult from cuvs_bench.orchestrator.config_loaders import ( @@ -14,7 +15,10 @@ ) from cuvs_bench.orchestrator.orchestrator import BenchmarkOrchestrator from cuvs_bench.plot.__main__ import load_all_results -from cuvs_bench.run.data_export import write_results_to_csv +from cuvs_bench.run.data_export import ( + validate_dataset_name, + write_results_to_csv, +) def test_python_backend_csv_is_plot_compatible(tmp_path): @@ -23,12 +27,12 @@ def test_python_backend_csv_is_plot_compatible(tmp_path): index_name = "test-index" results = [ BuildResult( - index_path=index_name, + index_path="/tmp/physical-index-path", build_time_seconds=1.5, index_size_bytes=1024, algorithm=algorithm, build_params={"m": 16}, - metadata={"group": "base"}, + metadata={"group": "base", "index_name": index_name}, ), SearchResult( neighbors=np.empty((0, 2), dtype=np.int64), @@ -38,6 +42,7 @@ def test_python_backend_csv_is_plot_compatible(tmp_path): recall=0.5, algorithm=algorithm, search_params=[{"ef_search": 50}], + latency_percentiles={"p50": 8.0, "p95": 12.0, "p99": 15.0}, metadata={ "group": "base", "index_name": index_name, @@ -76,6 +81,7 @@ def test_python_backend_csv_is_plot_compatible(tmp_path): ] assert raw["ef_search"].tolist() == [50, 100] assert raw["build time"].tolist() == [1.5, 1.5] + assert raw.loc[0, ["p50", "p95", "p99"]].tolist() == [8.0, 12.0, 15.0] assert not list(result_path.rglob("*.json")) write_results_to_csv( @@ -197,3 +203,196 @@ def search(self, dataset, indexes, k, **kwargs): assert [type(result) for result in results] == [BuildResult, SearchResult] assert results[1].recall == 1.0 + + +def test_subset_scopes_use_independent_result_files(tmp_path): + dataset = "test-dataset" + algorithm = "pylucene_cuvs_hnsw" + + for scope in ("subset100", "subset200"): + index_name = f"test-index-{scope}" + write_results_to_csv( + [ + BuildResult( + index_path=f"/tmp/{index_name}", + build_time_seconds=1.0, + index_size_bytes=100, + algorithm=algorithm, + build_params={}, + metadata={ + "group": "base", + "index_name": index_name, + "result_scope": scope, + }, + ), + SearchResult( + neighbors=np.empty((0, 1), dtype=np.int64), + distances=np.empty((0, 1), dtype=np.float32), + search_time_ms=1.0, + queries_per_second=10.0, + recall=1.0, + algorithm=algorithm, + search_params=[{}], + metadata={ + "group": "base", + "index_name": index_name, + "result_scope": scope, + }, + ), + ], + dataset, + str(tmp_path), + count=1, + batch_size=1, + ) + + result_path = tmp_path / dataset / "result" + for scope in ("subset100", "subset200"): + assert ( + result_path / "build" / f"{algorithm},base,{scope}.csv" + ).is_file() + assert ( + result_path / "search" / f"{algorithm},base,{scope},k1,bs1,raw.csv" + ).is_file() + + plotted = load_all_results( + str(tmp_path / dataset), + algorithms=[algorithm], + groups=["base"], + algo_groups=[], + k=1, + batch_size=1, + method="search", + index_key="algo", + raw=True, + mode="latency", + time_unit="s", + ) + assert set(plotted) == { + f"{algorithm}[scope=subset100]", + f"{algorithm}[scope=subset200]", + } + + +def test_algorithm_and_group_names_remain_distinct_in_plot_series(tmp_path): + dataset = "test-dataset" + identities = ( + ("a_b", "c", "first-index"), + ("a", "b_c", "second-index"), + ) + + for algorithm, group, index_name in identities: + write_results_to_csv( + [ + SearchResult( + neighbors=np.empty((0, 1), dtype=np.int64), + distances=np.empty((0, 1), dtype=np.float32), + search_time_ms=1.0, + queries_per_second=10.0, + recall=1.0, + algorithm=algorithm, + search_params=[{}], + metadata={"group": group, "index_name": index_name}, + ) + ], + dataset, + str(tmp_path), + count=1, + batch_size=1, + ) + + plotted = load_all_results( + str(tmp_path / dataset), + algorithms=["a_b", "a"], + groups=["c", "b_c"], + algo_groups=[], + k=1, + batch_size=1, + method="search", + index_key="algo", + raw=True, + mode="latency", + time_unit="s", + ) + + assert set(plotted) == {"a_b[group=c]", "a[group=b_c]"} + + +def test_failed_build_removes_only_matching_stale_search_files(tmp_path): + dataset = "test-dataset" + algorithm = "pylucene_cuvs_hnsw" + search_dir = tmp_path / dataset / "result" / "search" + search_dir.mkdir(parents=True) + + matching_stem = f"{algorithm},base,subset100,k1,bs1" + other_stem = f"{algorithm},base,subset200,k1,bs1" + suffixes = (",raw.csv", ",throughput.csv", ",latency.csv") + for stem in (matching_stem, other_stem): + for suffix in suffixes: + (search_dir / f"{stem}{suffix}").write_text("stale\n") + + write_results_to_csv( + [ + BuildResult( + index_path="/tmp/failed-index", + build_time_seconds=0.0, + index_size_bytes=0, + algorithm=algorithm, + build_params={}, + metadata={ + "group": "base", + "index_name": "failed-index", + "result_scope": "subset100", + }, + success=False, + error_message="build failed", + ) + ], + dataset, + str(tmp_path), + count=1, + batch_size=1, + search_requested=True, + ) + + assert not list(search_dir.glob(f"{matching_stem}*")) + assert len(list(search_dir.glob(f"{other_stem}*"))) == 3 + + +@pytest.mark.parametrize("dataset", ["", ".", "..", "nested/dataset"]) +def test_dataset_name_must_be_a_safe_path_component(dataset): + with pytest.raises(ValueError, match="Invalid benchmark dataset name"): + validate_dataset_name(dataset) + + +@pytest.mark.parametrize( + ("algorithm", "metadata", "match"), + [ + ("../pylucene", {"group": "base"}, "benchmark algorithm name"), + ("pylucene,other", {"group": "base"}, "benchmark algorithm name"), + ("pylucene[other]", {"group": "base"}, "benchmark algorithm name"), + ("pylucene", {"group": "../base"}, "benchmark group name"), + ("pylucene", {"group": "base\nother"}, "benchmark group name"), + ( + "pylucene", + {"group": "base", "result_scope": "../subset100"}, + "benchmark result scope", + ), + ], +) +def test_result_identity_must_use_safe_path_components( + tmp_path, algorithm, metadata, match +): + result = BuildResult( + index_path="index", + build_time_seconds=1.0, + index_size_bytes=1, + algorithm=algorithm, + build_params={}, + metadata=metadata, + ) + + with pytest.raises(ValueError, match=match): + write_results_to_csv( + [result], "test-dataset", str(tmp_path), count=1, batch_size=1 + ) 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..5f1184ccea --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_backend.py @@ -0,0 +1,1390 @@ +# +# 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_CODEC, + _FakeRuntime, + _backend, + _dataset, + _index, + _prepare_cagra_index, + _prepare_hnsw_index, +) + + +def _only_search_result(results): + assert len(results) == 1 + return results[0] + + +@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 + assert result.metadata["compound_file_policy"] == "lucene-default" + + +def test_build_creates_index_and_records_hnsw_writer_policy(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_policy"] == "gpu-with-cpu-fallback" + assert result.metadata["compound_file_policy"] == "lucene-default" + assert result.metadata["hnsw_verification"] == { + "status": "gpu-with-cpu-fallback-provenance", + "schema_version": 4, + "codec": _HNSW_CODEC, + "build_parameters": { + "codec": _HNSW_CODEC, + "m": 32, + "ef_construction": 32, + "direct_single_segment": False, + }, + "writer_policy": "gpu-with-cpu-fallback", + "compound_file_policy": "lucene-default", + "vector_count": 10, + "dimensions": 4, + "segment_count": 1, + "commit_file_count": 1, + } + assert (index_path / pylucene_backend._HNSW_PROVENANCE_FILE).is_file() + assert runtime.configured_resolve_calls == [(32, 32)] + assert runtime.resolve_calls == [] + assert len(runtime.build_calls) == 1 + assert runtime.build_calls[0][2].writer_policy == ("gpu-with-cpu-fallback") + assert runtime.build_calls[0][2].build_parameters == { + "codec": _HNSW_CODEC, + "m": 32, + "ef_construction": 32, + "direct_single_segment": False, + } + assert result.metadata["segment_count"] == 1 + assert result.metadata["segment_document_counts"] == [10] + assert result.metadata["segment_vector_counts"] == [10] + assert runtime.verification_calls == [] + + +def test_build_verifies_persisted_cagra_index(tmp_path): + runtime = _FakeRuntime() + 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 result.metadata["compound_file_policy"] == "disabled" + assert runtime.verification_calls == [(index_path, 10, 4)] + assert result.metadata["cagra_provenance"] == { + "status": "gpu-cagra-provenance", + "schema_version": 4, + "codec": _CAGRA_CODEC, + "build_parameters": {"codec": _CAGRA_CODEC}, + "writer_policy": "gpu-cagra", + "compound_file_policy": "disabled", + "vector_count": 10, + "dimensions": 4, + "segment_count": 1, + "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() + 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-with-cpu-fallback-provenance" + ) + assert runtime.resolve_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.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.build_calls == [] + + +def test_build_verifies_reused_cagra_index(tmp_path): + index_path = tmp_path / "index" + _prepare_cagra_index(index_path) + runtime = _FakeRuntime() + + 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.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() + 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.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() + + 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() + + +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_refuses_index_outside_configured_root(tmp_path): + allowed_root = tmp_path / "dataset" / "index" + external_index = tmp_path / "external-index" + external_index.mkdir() + (external_index / "segments_1").write_bytes(b"existing") + old_file = external_index / "old" + old_file.write_bytes(b"must remain") + runtime = _FakeRuntime() + backend = _backend(runtime) + backend.config["index_root"] = str(allowed_root) + + result = backend.build(_dataset(), [_index(external_index)], force=True) + + assert not result.success + assert "immediate child" in result.error_message + assert old_file.read_bytes() == b"must remain" + assert runtime.resolve_calls == [] + assert runtime.build_calls == [] + + +def test_safe_remove_rechecks_configured_root(tmp_path): + allowed_root = tmp_path / "allowed" + external_index = tmp_path / "external-index" + external_index.mkdir() + old_file = external_index / "old" + old_file.write_bytes(b"must remain") + + with pytest.raises(ValueError, match="outside its configured root"): + pylucene_backend._safe_remove_index(external_index, allowed_root) + + assert old_file.read_bytes() == b"must remain" + + +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_accepts_production_hnsw_fallback_policy(tmp_path): + runtime = _FakeRuntime() + index_path = tmp_path / "index" + + result = _backend(runtime).build( + _dataset(), [_index(index_path)], force=True + ) + + assert result.success + assert result.metadata["writer_policy"] == "gpu-with-cpu-fallback" + assert runtime.build_calls[0][2].writer_policy == ("gpu-with-cpu-fallback") + + +@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", + [ + _HNSW_CODEC, + _CAGRA_CODEC, + ], + ids=["hnsw", "cagra"], +) +def test_build_rejects_single_vector_cuvs_bypass(tmp_path, codec): + runtime = _FakeRuntime() + + 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 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 = _only_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 + + +def test_hnsw_build_parameters_have_canonical_defaults(): + assert pylucene_backend._normalize_build_params( + {"codec": _HNSW_CODEC} + ) == { + "codec": _HNSW_CODEC, + "m": 32, + "ef_construction": 32, + "direct_single_segment": False, + } + + +@pytest.mark.parametrize("name", ["m", "ef_construction"]) +@pytest.mark.parametrize("value", [1, 512]) +def test_hnsw_build_parameter_boundaries_are_inclusive(name, value): + parameters = pylucene_backend._normalize_build_params( + {"codec": _HNSW_CODEC, name: value} + ) + + assert parameters[name] == value + + +@pytest.mark.parametrize("name", ["m", "ef_construction"]) +@pytest.mark.parametrize( + "value", + [0, 513, True, 1.5, "32", None], + ids=["below-min", "above-max", "bool", "float", "string", "none"], +) +def test_hnsw_build_parameters_reject_invalid_values(name, value): + with pytest.raises(ValueError, match=rf"{name} must be an integer"): + pylucene_backend._normalize_build_params( + {"codec": _HNSW_CODEC, name: value} + ) + + +@pytest.mark.parametrize("value", [0, 1, None, "true"]) +def test_direct_single_segment_rejects_non_boolean_values(value): + with pytest.raises(TypeError, match="must be a boolean"): + pylucene_backend._normalize_build_params( + { + "codec": _HNSW_CODEC, + "direct_single_segment": value, + } + ) + + +@pytest.mark.parametrize( + ("name", "value"), + [ + ("m", 32), + ("ef_construction", 64), + ("direct_single_segment", False), + ], +) +def test_cagra_rejects_hnsw_only_build_parameters(name, value): + with pytest.raises(ValueError, match="apply only to"): + pylucene_backend._normalize_build_params( + {"codec": _CAGRA_CODEC, name: value} + ) + + +@pytest.mark.parametrize("parameters", [None, [], "codec"]) +def test_build_parameters_must_be_a_mapping(parameters): + with pytest.raises(TypeError, match="must be a mapping"): + pylucene_backend._normalize_build_params(parameters) + + +def test_build_propagates_configured_hnsw_parameters_to_runtime(tmp_path): + runtime = _FakeRuntime() + index_path = tmp_path / "index" + index = _index( + index_path, + build_params={"m": 24, "ef_construction": 96}, + ) + + result = _backend(runtime).build(_dataset(), [index], force=True) + + assert result.success + assert runtime.configured_resolve_calls == [(24, 96)] + assert runtime.resolve_calls == [] + build_codec = runtime.build_calls[0][2] + assert build_codec.java_codec == (_HNSW_CODEC, 24, 96) + assert build_codec.build_parameters == { + "codec": _HNSW_CODEC, + "m": 24, + "ef_construction": 96, + "direct_single_segment": False, + } + assert result.build_params == build_codec.build_parameters + + +@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_removes_index_when_direct_segment_topology_is_not_one( + tmp_path, +): + runtime = _FakeRuntime() + runtime.topology = pylucene_backend._IndexTopology( + segment_document_counts=(5, 5), + segment_vector_counts=(5, 5), + ) + index_path = tmp_path / "index" + + result = _backend(runtime).build( + _dataset(), + [ + _index( + index_path, + build_params={"direct_single_segment": True}, + ) + ], + force=True, + ) + + assert not result.success + assert "requested one committed Lucene segment" 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, trusted_index_root=None): + 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, trusted_index_root=None): + 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 = _only_search_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) + assert result.search_params == [{"num_candidates": 3}] + + +def test_search_dry_run_returns_one_result_per_candidate_setting(tmp_path): + backend = _backend() + candidates = [150, 200, 300] + index = _index( + tmp_path / "index", + search_params=[ + {"num_candidates": num_candidates} for num_candidates in candidates + ], + ) + + results = backend.search(_dataset(), [index], k=150, dry_run=True) + + assert len(results) == 3 + assert all(result.success for result in results) + assert backend._runtime is None + assert [result.search_params for result in results] == [ + [{"num_candidates": num_candidates}] for num_candidates in candidates + ] + assert [result.metadata["num_candidates"] for result in results] == ( + candidates + ) + + +def test_search_associates_each_candidate_setting_with_its_result(tmp_path): + class _SelectiveFailureRuntime(_FakeRuntime): + def search_index( + self, + index_path, + query_vectors, + k, + batch_size, + num_candidates=None, + ): + result = super().search_index( + index_path, + query_vectors, + k, + batch_size, + num_candidates, + ) + if num_candidates == 4: + raise RuntimeError("candidate-specific failure") + return result + + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + runtime = _SelectiveFailureRuntime() + candidates = [3, 4, 5] + index = _index( + index_path, + search_params=[ + {"num_candidates": num_candidates} for num_candidates in candidates + ], + ) + + results = _backend(runtime).search(_dataset(), [index], k=3) + + assert len(results) == 3 + assert [result.success for result in results] == [True, False, True] + assert [result.search_params for result in results] == [ + [{"num_candidates": num_candidates}] for num_candidates in candidates + ] + assert "candidate-specific failure" in results[1].error_message + assert [call[4] for call in runtime.search_calls] == candidates + + +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 = _only_search_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.metadata["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["top_k"] == 3 + assert result.metadata["num_candidates"] == 3 + assert result.metadata["hnsw_verification"]["status"] == ( + "gpu-with-cpu-fallback-provenance" + ) + assert "per_search_param_results" not in result.metadata + assert runtime.search_calls[0][2:] == (3, 1, 3) + 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 = _only_search_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 = _only_search_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() + + result = _only_search_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() + runtime.verification_error = RuntimeError("persisted brute-force index") + + result = _only_search_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() + + result = _only_search_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", + "Unsupported PyLucene search parameter", + ), + ( + _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 = _only_search_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 + + +@pytest.mark.parametrize( + "num_candidates", + [0, -1, True, 1.5, "3", None], + ids=["zero", "negative", "bool", "float", "string", "none"], +) +def test_search_rejects_invalid_num_candidates(num_candidates, tmp_path): + result = _only_search_result( + _backend().search( + _dataset(), + [ + _index( + tmp_path / "index", + search_params=[{"num_candidates": num_candidates}], + ) + ], + k=3, + dry_run=True, + ) + ) + + assert not result.success + assert "num_candidates must be a positive integer" in result.error_message + + +def test_search_rejects_num_candidates_below_top_k(tmp_path): + result = _only_search_result( + _backend().search( + _dataset(), + [ + _index( + tmp_path / "index", + search_params=[{"num_candidates": 2}], + ) + ], + k=3, + dry_run=True, + ) + ) + + assert not result.success + assert "greater than or equal to k (3)" in result.error_message + + +def test_search_rejects_num_candidates_for_cagra(tmp_path): + result = _only_search_result( + _backend(codec=_CAGRA_CODEC).search( + _dataset(), + [ + _index( + tmp_path / "index", + codec=_CAGRA_CODEC, + search_params=[{"num_candidates": 3}], + ) + ], + k=3, + dry_run=True, + ) + ) + + assert not result.success + assert "num_candidates applies only to" in result.error_message + + +@pytest.mark.parametrize("search_parameters", [1, None, "candidates"]) +def test_search_parameter_entries_must_be_mappings( + search_parameters, tmp_path +): + result = _only_search_result( + _backend().search( + _dataset(), + [ + _index( + tmp_path / "index", + search_params=[search_parameters], + ) + ], + k=3, + dry_run=True, + ) + ) + + assert not result.success + assert "Every PyLucene search parameter must be a mapping" in ( + result.error_message + ) + + +def test_resolve_search_plans_normalizes_the_complete_request(tmp_path): + index = _index(tmp_path / "index", search_params=[]) + + plans = _backend()._resolve_search_plans( + index, + k=3, + batch_size=7, + mode="latency", + search_threads="1", + ) + + assert plans == [ + pylucene_backend._SearchPlan( + index_path=tmp_path / "index", + codec_name=_HNSW_CODEC, + build_parameters={ + "codec": _HNSW_CODEC, + "m": 32, + "ef_construction": 32, + "direct_single_segment": False, + }, + search_parameters={"num_candidates": 3}, + num_candidates=3, + k=3, + batch_size=7, + mode="latency", + ) + ] + + +def test_search_validation_preserves_first_error_precedence(tmp_path): + result = _only_search_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 = _only_search_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 = _only_search_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 = _only_search_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 = _only_search_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 = _only_search_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 = _only_search_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 = _only_search_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"] = "DifferentCodec" + manifest_path.write_text(json.dumps(payload)) + + runtime = _FakeRuntime() + result = _only_search_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..129f4758a1 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_cli_config.py @@ -0,0 +1,1204 @@ +# +# 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 subprocess +import sys +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +from click.testing import CliRunner + +import cuvs_bench.backends.pylucene as pylucene_backend +from cuvs_bench.backends import get_registry +from cuvs_bench.backends.base import SearchResult +from cuvs_bench.backends.pylucene import ( + PyLuceneConfigLoader, + _SearchHit, +) +from cuvs_bench.backends.registry import list_config_loaders +from cuvs_bench.backends.search_spaces import get_search_space +from cuvs_bench.orchestrator.orchestrator import BenchmarkOrchestrator +from cuvs_bench.tests._pylucene_test_utils import ( + _CAGRA_CODEC, + _HNSW_CODEC, + _FakeRuntime, + _write_test_bin, +) + + +_DEFAULT_HNSW_BUILD_PARAMETERS = { + "codec": _HNSW_CODEC, + "m": 32, + "ef_construction": 32, + "direct_single_segment": False, +} + + +def _install_single_trial_optuna(monkeypatch): + suggested_ranges = {} + selected_parameters = {} + + def suggest_int(name, minimum, maximum, *, log=False): + suggested_ranges[name] = (minimum, maximum, log) + selected_parameters[name] = minimum + return minimum + + trial = SimpleNamespace( + number=0, + suggest_int=suggest_int, + suggest_float=lambda name, minimum, maximum, **kwargs: minimum, + suggest_categorical=lambda name, choices: choices[0], + ) + study = SimpleNamespace( + best_trial=trial, + best_params=selected_parameters, + best_value=0.0, + ) + + def optimize(objective, **kwargs): + study.best_value = objective(trial) + + study.optimize = optimize + fake_optuna = SimpleNamespace( + TrialPruned=RuntimeError, + create_study=lambda **kwargs: study, + logging=SimpleNamespace(WARNING=0, set_verbosity=lambda level: None), + ) + monkeypatch.setitem(sys.modules, "optuna", fake_optuna) + return suggested_ranges + + +def _hnsw_index_name( + algorithm: str, + group: str, + *, + subset_scope: str | None = None, + m: int = 32, + ef_construction: int = 32, + direct_single_segment: bool = False, +) -> str: + identity = f"{algorithm}[group={group}]" + if subset_scope is not None: + identity = f"{identity}[scope={subset_scope}]" + return ( + f"{identity}[codec={_HNSW_CODEC}]" + f"[m={m}]" + f"[ef_construction={ef_construction}]" + f"[direct_single_segment={str(direct_single_segment).lower()}]" + ) + + +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_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 = _hnsw_index_name("pylucene_test", "test") + 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() + build_csv = result_path / "build" / "pylucene_test,test.csv" + with build_csv.open(newline="") as file: + build_rows = list(csv.DictReader(file)) + assert len(build_rows) == 1 + build_row = build_rows[0] + assert build_row["index_name"] == index_name + assert float(build_row["time"]) > 0.0 + assert build_row["codec"] == _HNSW_CODEC + assert build_row["writer_policy"] == "gpu-with-cpu-fallback" + + 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 float(csv_row["p50"]) == pytest.approx(1.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_without_persisting_failed_measurement( + 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_csv = ( + dataset_path + / "test-dataset" + / "result" + / "build" + / "pylucene_test,test.csv" + ) + assert not build_csv.exists() + 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"] + codec = config.indexes[0].build_param["codec"] + assert config.backend_config["requires_gpu"] is (codec == _CAGRA_CODEC) + assert config.backend_config["group"] == "base" + assert config.backend_config["index_name"] == config.index_name + assert config.backend_config["index_root"] == str( + Path("/datasets/test-dataset/index") + ) + assert config.backend_config["result_scope"] is None + + +def test_hnsw_build_defaults_are_canonical_and_share_one_identity(config_dir): + _, configs = PyLuceneConfigLoader(config_path=config_dir).load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_test", + groups="test", + ) + + assert len(configs) == 1 + config = configs[0] + assert config.indexes[0].build_param == _DEFAULT_HNSW_BUILD_PARAMETERS + assert config.index_name == _hnsw_index_name("pylucene_test", "test") + assert PyLuceneConfigLoader._index_label( + "pylucene_test", + "test", + None, + {"codec": _HNSW_CODEC}, + ) == PyLuceneConfigLoader._index_label( + "pylucene_test", + "test", + None, + _DEFAULT_HNSW_BUILD_PARAMETERS, + ) + + +def test_config_loader_expands_hnsw_build_and_search_sweeps( + config_dir, tmp_path +): + algorithm_config = tmp_path / "pylucene_hnsw_sweep.yaml" + algorithm_config.write_text( + f"""\ +backend: pylucene +name: pylucene_hnsw_sweep +groups: + requested: + build: + codec: [{_HNSW_CODEC}] + m: [16, 24, 32] + search: + num_candidates: [150, 200, 300, 600] + build_grid: + build: + codec: [{_HNSW_CODEC}] + m: [16, 32] + ef_construction: [48, 64] + search: {{}} +""" + ) + loader = PyLuceneConfigLoader(config_path=config_dir) + + _, requested_configs = loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_hnsw_sweep", + groups="requested", + algorithm_configuration=str(algorithm_config), + ) + + assert len(requested_configs) == 3 + assert len({config.index_path for config in requested_configs}) == 3 + assert ( + sum( + len(config.indexes[0].search_params) + for config in requested_configs + ) + == 12 + ) + for config, m in zip(requested_configs, (16, 24, 32), strict=True): + assert config.indexes[0].build_param == { + **_DEFAULT_HNSW_BUILD_PARAMETERS, + "m": m, + } + assert config.index_name == _hnsw_index_name( + "pylucene_hnsw_sweep", "requested", m=m + ) + assert config.indexes[0].search_params == [ + {"num_candidates": 150}, + {"num_candidates": 200}, + {"num_candidates": 300}, + {"num_candidates": 600}, + ] + + _, build_grid_configs = loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_hnsw_sweep", + groups="build_grid", + algorithm_configuration=str(algorithm_config), + ) + + assert len(build_grid_configs) == 4 + assert { + ( + config.indexes[0].build_param["m"], + config.indexes[0].build_param["ef_construction"], + ) + for config in build_grid_configs + } == {(16, 48), (16, 64), (32, 48), (32, 64)} + + +def test_pylucene_tune_space_uses_runtime_top_k_for_candidates(): + assert get_search_space("pylucene_cuvs_hnsw")["search"] == { + "num_candidates": { + "type": "int", + "min": "top_k", + "max": 500, + } + } + + +def test_pylucene_tune_resolves_candidate_minimum_from_top_k(monkeypatch): + suggested_ranges = _install_single_trial_optuna(monkeypatch) + orchestrator = BenchmarkOrchestrator.__new__(BenchmarkOrchestrator) + trial_arguments = {} + + def run_trial(**kwargs): + trial_arguments.update(kwargs) + return [ + SearchResult( + neighbors=np.empty((0, 0), dtype=np.int64), + distances=np.empty((0, 0), dtype=np.float32), + search_time_ms=1.0, + queries_per_second=1.0, + recall=1.0, + algorithm="pylucene_cuvs_hnsw", + search_params=[kwargs["search_params"]], + ) + ] + + orchestrator._run_trial = run_trial + results = orchestrator._run_tune( + constraints={"recall": "maximize"}, + n_trials=1, + build=True, + search=True, + force=False, + dry_run=False, + count=150, + batch_size=1, + search_mode="latency", + search_threads=1, + algorithms="pylucene_cuvs_hnsw", + ) + + assert suggested_ranges["num_candidates"] == (150, 500, False) + assert trial_arguments["search_params"] == {"num_candidates": 150} + assert results[0].search_params == [{"num_candidates": 150}] + + +def test_pylucene_tune_rejects_top_k_above_candidate_ceiling(monkeypatch): + _install_single_trial_optuna(monkeypatch) + orchestrator = BenchmarkOrchestrator.__new__(BenchmarkOrchestrator) + + with pytest.raises(ValueError, match="minimum 501 exceeds maximum 500"): + orchestrator._run_tune( + constraints={"recall": "maximize"}, + n_trials=1, + build=True, + search=True, + force=False, + dry_run=False, + count=501, + batch_size=1, + search_mode="latency", + search_threads=1, + algorithms="pylucene_cuvs_hnsw", + ) + + +@pytest.mark.parametrize( + ("field_name", "field_value"), + [ + ("m", True), + ("m", 0), + ("m", 513), + ("ef_construction", False), + ("ef_construction", 0), + ("ef_construction", 513), + ], +) +def test_hnsw_build_parameters_reject_booleans_and_values_outside_range( + field_name, field_value +): + with pytest.raises(ValueError, match=field_name): + pylucene_backend._normalize_build_params( + {"codec": _HNSW_CODEC, field_name: field_value} + ) + + +def test_hnsw_build_parameter_boundaries_are_supported(): + assert pylucene_backend._normalize_build_params( + { + "codec": _HNSW_CODEC, + "m": 1, + "ef_construction": 512, + } + ) == { + **_DEFAULT_HNSW_BUILD_PARAMETERS, + "m": 1, + "ef_construction": 512, + } + + +@pytest.mark.parametrize("field_name", ["m", "ef_construction"]) +def test_hnsw_build_parameters_are_rejected_for_cagra(field_name): + with pytest.raises(ValueError, match=field_name): + pylucene_backend._normalize_build_params( + {"codec": _CAGRA_CODEC, field_name: 32} + ) + + +@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]", + "Unsupported PyLucene search parameter.*ignored", + ), + ], +) +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), + ) + + +@pytest.mark.parametrize( + ("algorithm_name", "group_name", "match"), + [ + ("../../../escaped", "test", "PyLucene algorithm name"), + ("pylucene_safe", "../escaped", "PyLucene group name"), + ("pylucene,ambiguous", "test", "PyLucene algorithm name"), + ], +) +def test_config_loader_rejects_unsafe_artifact_identity( + config_dir, tmp_path, algorithm_name, group_name, match +): + custom_config = tmp_path / "unsafe-identity.yaml" + custom_config.write_text( + f"""\ +backend: pylucene +name: {algorithm_name} +groups: + {group_name}: + build: + codec: [{_HNSW_CODEC}] + search: {{}} +""" + ) + + with pytest.raises(ValueError, match=match): + PyLuceneConfigLoader(config_path=config_dir).load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms=(None if "," in algorithm_name else algorithm_name), + groups=group_name, + 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_name = _hnsw_index_name( + "pylucene_test", "test", subset_scope=subset_scope + ) + assert config.index_name == index_name + assert config.index_path == ( + Path("/datasets/test-dataset/index") / index_name + ) + assert config.backend_config["group"] == "test" + assert config.backend_config["index_name"] == index_name + assert config.backend_config["result_scope"] == subset_scope + identities.append( + ( + config.index_name, + config.index_path, + config.backend_config["result_scope"], + ) + ) + + assert len(set(identities)) == 3 + + +@pytest.mark.parametrize( + ("first", "second"), + [ + (("foo.subset2", "base", None), ("foo", "base", "subset2")), + (("a_b", "c", None), ("a", "b_c", None)), + ], +) +def test_index_labels_are_injective(first, second): + build_params = _DEFAULT_HNSW_BUILD_PARAMETERS + + first_label = PyLuceneConfigLoader._index_label( + first[0], first[1], first[2], build_params + ) + second_label = PyLuceneConfigLoader._index_label( + second[0], second[1], second[2], build_params + ) + + assert first_label != second_label + + +@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 == _hnsw_index_name(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 == _hnsw_index_name("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[group=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[group=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 = [ + f"{config.algo},{config.backend_config['group']}" 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", + {_HNSW_CODEC}, + ), + ("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..fdfc4b04fb --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_integration.py @@ -0,0 +1,1307 @@ +# +# 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 re +import shutil +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import pytest + +from cuvs_bench._bin_format import write_bin_header +from cuvs_bench.backends._pylucene_java import ( + EF_CONSTRUCTION_PROPERTY, + M_PROPERTY, +) +from cuvs_bench.backends._utils import compute_recall +from cuvs_bench.backends.base import BuildResult, Dataset +from cuvs_bench.backends.pylucene import ( + _BuildCodec, + _CAGRA_PROVENANCE_FILE, + PyLuceneBackend, + _HNSW_PROVENANCE_FILE, + _PyLuceneRuntime, + _restore_java_property, + _validate_pylucene_version, +) +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" +_HNSW_WRITER_POLICY = "gpu-with-cpu-fallback" +_COMPOUND_FILE_POLICY = { + _HNSW_CODEC: "lucene-default", + _CAGRA_CODEC: "disabled", +} +_GPU_HNSW_WRITER = ( + "com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsWriter" +) +_CPU_HNSW_WRITER = ( + "org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsWriter" +) +_WRITER_SELECTION_CODEC = "com.nvidia.cuvs.bench.PyLuceneWriterSelectionCodec" +_DEFAULT_HNSW_BUILD_PARAMETERS = { + "codec": _HNSW_CODEC, + "m": 32, + "ef_construction": 32, + "direct_single_segment": False, +} +_CAGRA_BUILD_PARAMETERS = {"codec": _CAGRA_CODEC} +_WRITER_SELECTION_SOURCE = ( + Path(__file__).resolve().parents[2] + / "tests" + / "java" + / "com" + / "nvidia" + / "cuvs" + / "bench" + / "PyLuceneWriterSelectionCodec.java" +) +_CPU_FALLBACK_PROBE = Path(__file__).with_name( + "pylucene_cpu_fallback_probe.py" +) + + +@dataclass(frozen=True) +class _DatasetCase: + dataset: Dataset + query_ids: np.ndarray + squared_distances: np.ndarray + k: int + + +@dataclass(frozen=True) +class _RuntimeFixture: + backend_config: dict[str, str] + writer_selection_classes: Path + + +@dataclass(frozen=True) +class _BaselineIndex: + algo: str + codec: str + writer_policy: str + build_parameters: dict[str, object] + index_path: Path + index: IndexConfig + build_result: BuildResult + dataset_case: _DatasetCase + runtime_config: dict[str, str] + + +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 _single_search_result(results): + assert len(results) == 1 + return results[0] + + +def _writer_diagnostics(java_codec): + diagnostics = str(java_codec.knnVectorsFormat()) + match = re.search( + r"writerClass=([^,)]+), fieldsWriterCalls=(\d+)", diagnostics + ) + assert match is not None, diagnostics + return match.group(1), int(match.group(2)) + + +def _configured_writer_selection_codec(runtime, m, ef_construction): + properties = { + M_PROPERTY: str(m), + EF_CONSTRUCTION_PROPERTY: str(ef_construction), + } + previous = {name: runtime.System.getProperty(name) for name in properties} + try: + for name, value in properties.items(): + runtime.System.setProperty(name, value) + reflected_codec = runtime.Class.forName( + _WRITER_SELECTION_CODEC + ).newInstance() + return runtime.Codec.cast_(reflected_codec) + finally: + for name, value in previous.items(): + _restore_java_property(runtime.System, name, value) + + +def _backend(algo, codec, runtime_config): + return PyLuceneBackend( + { + "name": f"{algo}-integration", + "algo": algo, + "codec": codec, + **runtime_config, + } + ) + + +def _index_at(baseline, index_path): + return IndexConfig( + name=baseline.index.name, + algo=baseline.algo, + build_param=baseline.build_parameters, + search_params=[{}], + file=str(index_path), + ) + + +def _copy_baseline(baseline, tmp_path): + index_path = tmp_path / f"{baseline.algo}-index" + shutil.copytree(baseline.index_path, index_path) + return index_path, _index_at(baseline, index_path) + + +def _search_copied_index(baseline, index, dataset=None): + backend = _backend(baseline.algo, baseline.codec, baseline.runtime_config) + search_dataset = ( + baseline.dataset_case.dataset if dataset is None else dataset + ) + try: + return _single_search_result( + backend.search( + search_dataset, + [index], + k=baseline.dataset_case.k, + ) + ) + finally: + backend.cleanup() + + +def _compile_writer_selection_codec( + output_dir: Path, + cuvs_java_jar: Path, + cuvs_lucene_jar: Path, + pylucene_classpath: str, +) -> None: + environment_javac = Path(sys.prefix) / "lib" / "jvm" / "bin" / "javac" + javac = shutil.which("javac") + if javac is None and environment_javac.is_file(): + javac = str(environment_javac) + if javac is None: + pytest.fail("javac is required for the PyLucene integration tests") + if not _WRITER_SELECTION_SOURCE.is_file(): + pytest.fail( + "PyLucene writer-selection test source is missing: " + f"{_WRITER_SELECTION_SOURCE}" + ) + + output_dir.mkdir(parents=True, exist_ok=True) + compile_classpath = os.pathsep.join( + (str(cuvs_java_jar), str(cuvs_lucene_jar), pylucene_classpath) + ) + completed = subprocess.run( + [ + javac, + "--release", + "22", + "-classpath", + compile_classpath, + "-d", + str(output_dir), + str(_WRITER_SELECTION_SOURCE), + ], + capture_output=True, + text=True, + check=False, + ) + if completed.returncode != 0: + pytest.fail( + "Could not compile the PyLucene writer-selection test codec:\n" + f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}" + ) + + +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": 4, + "codec": _CAGRA_CODEC, + "build_parameters": _CAGRA_BUILD_PARAMETERS, + "writer_policy": "gpu-cagra", + "compound_file_policy": "disabled", + "vector_count": expected_vector_count, + "dimensions": expected_dimensions, + "segment_count": metadata["segment_count"], + "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, + expected_build_parameters, +): + verification = metadata["hnsw_verification"] + assert verification == { + "status": "gpu-with-cpu-fallback-provenance", + "schema_version": 4, + "codec": codec, + "build_parameters": expected_build_parameters, + "writer_policy": _HNSW_WRITER_POLICY, + "compound_file_policy": "lucene-default", + "vector_count": expected_vector_count, + "dimensions": expected_dimensions, + "segment_count": metadata["segment_count"], + "commit_file_count": 1, + } + + +@pytest.fixture(scope="module") +def pylucene_runtime_config(tmp_path_factory): + """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: + lucene = importlib.import_module("lucene") + except ImportError as exc: + pytest.fail(f"PyLucene is not importable: {exc}") + _validate_pylucene_version(lucene) + + test_classes = tmp_path_factory.mktemp("pylucene-java-test-classes") + _compile_writer_selection_codec( + test_classes, + cuvs_java_jar, + cuvs_lucene_jar, + lucene.CLASSPATH, + ) + lucene.CLASSPATH = os.pathsep.join((str(test_classes), lucene.CLASSPATH)) + + return _RuntimeFixture( + backend_config={ + "cuvs_java_jar": str(cuvs_java_jar), + "cuvs_lucene_jar": str(cuvs_lucene_jar), + "java_library_path": java_library_path, + }, + writer_selection_classes=test_classes, + ) + + +@pytest.fixture(scope="module") +def integration_dataset_case(): + 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 + ) + return _DatasetCase( + 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", + ), + query_ids=query_ids, + squared_distances=squared_distances, + k=k, + ) + + +def _build_baseline( + tmp_path_factory, + runtime_fixture, + dataset_case, + *, + algo, + codec, + writer_policy, +): + build_parameters = ( + _DEFAULT_HNSW_BUILD_PARAMETERS + if codec == _HNSW_CODEC + else _CAGRA_BUILD_PARAMETERS + ) + index_path = tmp_path_factory.mktemp(f"{algo}-baseline") / "index" + index = IndexConfig( + name=f"{algo}-integration", + algo=algo, + build_param=build_parameters, + search_params=[{}], + file=str(index_path), + ) + backend = _backend(algo, codec, runtime_fixture.backend_config) + try: + build_result = backend.build(dataset_case.dataset, [index], force=True) + finally: + backend.cleanup() + assert build_result.success, build_result.error_message + return _BaselineIndex( + algo=algo, + codec=codec, + writer_policy=writer_policy, + build_parameters=build_parameters, + index_path=index_path, + index=index, + build_result=build_result, + dataset_case=dataset_case, + runtime_config=runtime_fixture.backend_config, + ) + + +@pytest.fixture(scope="module") +def hnsw_baseline( + tmp_path_factory, pylucene_runtime_config, integration_dataset_case +): + return _build_baseline( + tmp_path_factory, + pylucene_runtime_config, + integration_dataset_case, + algo="pylucene_cuvs_hnsw", + codec=_HNSW_CODEC, + writer_policy=_HNSW_WRITER_POLICY, + ) + + +@pytest.fixture(scope="module") +def cagra_baseline( + tmp_path_factory, pylucene_runtime_config, integration_dataset_case +): + return _build_baseline( + tmp_path_factory, + pylucene_runtime_config, + integration_dataset_case, + algo="pylucene_cuvs_cagra", + codec=_CAGRA_CODEC, + writer_policy="gpu-cagra", + ) + + +def _assert_build_contract(baseline): + result = baseline.build_result + case = baseline.dataset_case + vectors = case.dataset.training_vectors + assert result.build_time_seconds > 0 + assert result.index_size_bytes > 0 + assert result.build_params == baseline.build_parameters + assert result.metadata["codec"] == baseline.codec + assert result.metadata["pylucene_version"] != "unknown" + assert result.metadata["writer_policy"] == baseline.writer_policy + assert result.metadata["segment_count"] >= 1 + assert ( + result.metadata["compound_file_policy"] + == _COMPOUND_FILE_POLICY[baseline.codec] + ) + if baseline.codec == _CAGRA_CODEC: + assert any( + path.suffix == ".vemc" for path in baseline.index_path.iterdir() + ) + assert any( + path.suffix == ".vcag" for path in baseline.index_path.iterdir() + ) + _assert_cagra_provenance( + result.metadata, vectors.shape[0], vectors.shape[1] + ) + _assert_cagra_verification( + result.metadata, vectors.shape[0], vectors.shape[1] + ) + assert (baseline.index_path / _CAGRA_PROVENANCE_FILE).is_file() + else: + _assert_hnsw_verification( + result.metadata, + baseline.codec, + vectors.shape[0], + vectors.shape[1], + baseline.build_parameters, + ) + assert (baseline.index_path / _HNSW_PROVENANCE_FILE).is_file() + + +def _assert_search_contract(baseline, result): + case = baseline.dataset_case + dataset = case.dataset + query_vectors = dataset.query_vectors + training_vectors = dataset.training_vectors + assert result.success, result.error_message + assert result.neighbors.shape == (query_vectors.shape[0], case.k) + assert result.distances.shape == (query_vectors.shape[0], case.k) + np.testing.assert_array_equal(result.neighbors[:, 0], case.query_ids) + assert np.all(np.isfinite(result.distances)) + recall = compute_recall( + result.neighbors, dataset.groundtruth_neighbors, case.k + ) + assert recall >= 0.75 + returned_distances = np.take_along_axis( + case.squared_distances, result.neighbors, axis=1 + ) + np.testing.assert_allclose( + result.distances, returned_distances, rtol=1e-5, atol=1e-5 + ) + assert np.all(np.diff(result.distances, axis=1) >= -1e-6) + assert result.search_time_ms > 0 + assert result.metadata["latency_seconds"] > 0 + assert result.queries_per_second > 0 + assert result.metadata["codec"] == baseline.codec + assert result.metadata["segment_count"] >= 1 + assert result.metadata["pylucene_version"] != "unknown" + assert ( + result.metadata["compound_file_policy"] + == _COMPOUND_FILE_POLICY[baseline.codec] + ) + assert result.metadata["num_batches"] == 2 + assert result.metadata["mode"] == "latency" + assert set(result.latency_percentiles) == {"p50", "p95", "p99"} + if baseline.codec == _CAGRA_CODEC: + assert result.search_params == [{}] + _assert_cagra_provenance( + result.metadata, + training_vectors.shape[0], + training_vectors.shape[1], + ) + _assert_cagra_verification( + result.metadata, + training_vectors.shape[0], + training_vectors.shape[1], + ) + else: + assert result.search_params == [{"num_candidates": case.k}] + assert result.metadata["num_candidates"] == case.k + assert "cagra_verification" not in result.metadata + _assert_hnsw_verification( + result.metadata, + baseline.codec, + training_vectors.shape[0], + training_vectors.shape[1], + baseline.build_parameters, + ) + + +def test_build_hnsw_with_real_pylucene_runtime(hnsw_baseline): + _assert_build_contract(hnsw_baseline) + + +def test_build_cagra_with_real_pylucene_runtime(cagra_baseline): + _assert_build_contract(cagra_baseline) + + +def _assert_reuse_contract(baseline_index): + backend = _backend( + baseline_index.algo, + baseline_index.codec, + baseline_index.runtime_config, + ) + try: + result = backend.build( + baseline_index.dataset_case.dataset, + [baseline_index.index], + force=False, + ) + finally: + backend.cleanup() + assert result.success, result.error_message + assert result.build_params == baseline_index.build_parameters + assert result.metadata["skipped"] is True + assert result.metadata["segment_count"] >= 1 + assert ( + result.metadata["compound_file_policy"] + == _COMPOUND_FILE_POLICY[baseline_index.codec] + ) + vectors = baseline_index.dataset_case.dataset.training_vectors + if baseline_index.codec == _CAGRA_CODEC: + _assert_cagra_provenance( + result.metadata, vectors.shape[0], vectors.shape[1] + ) + _assert_cagra_verification( + result.metadata, vectors.shape[0], vectors.shape[1] + ) + else: + _assert_hnsw_verification( + result.metadata, + baseline_index.codec, + vectors.shape[0], + vectors.shape[1], + baseline_index.build_parameters, + ) + + +def test_reuse_real_hnsw_index(hnsw_baseline): + _assert_reuse_contract(hnsw_baseline) + + +def test_reuse_real_cagra_index(cagra_baseline): + _assert_reuse_contract(cagra_baseline) + + +def _assert_baseline_search(baseline_index): + backend = _backend( + baseline_index.algo, + baseline_index.codec, + baseline_index.runtime_config, + ) + try: + result = _single_search_result( + backend.search( + baseline_index.dataset_case.dataset, + [baseline_index.index], + k=baseline_index.dataset_case.k, + batch_size=2, + ) + ) + finally: + backend.cleanup() + _assert_search_contract(baseline_index, result) + + +def test_search_real_hnsw_index(hnsw_baseline): + _assert_baseline_search(hnsw_baseline) + + +def test_search_real_cagra_index(cagra_baseline): + _assert_baseline_search(cagra_baseline) + + +def test_cagra_rejects_dataset_vector_count_mismatch(tmp_path, cagra_baseline): + _, index = _copy_baseline(cagra_baseline, tmp_path) + case = cagra_baseline.dataset_case + dataset = Dataset( + name="pylucene-integration", + training_vectors=case.dataset.training_vectors[:-1], + query_vectors=case.dataset.query_vectors, + distance_metric="euclidean", + ) + backend = _backend( + cagra_baseline.algo, + cagra_baseline.codec, + cagra_baseline.runtime_config, + ) + try: + build_result = backend.build(dataset, [index], force=False) + search_result = _single_search_result( + backend.search(dataset, [index], k=case.k) + ) + finally: + backend.cleanup() + assert not build_result.success + assert "vector count does not match the dataset: 512 != 511" in ( + build_result.error_message + ) + assert not search_result.success + assert "vector count does not match the dataset: 512 != 511" in ( + search_result.error_message + ) + + +def test_cagra_rejects_missing_vector_data(tmp_path, cagra_baseline): + index_path, index = _copy_baseline(cagra_baseline, tmp_path) + next(index_path.glob("*.vcag")).unlink() + result = _search_copied_index(cagra_baseline, index) + assert not result.success + assert "cannot read" in result.error_message + + +def test_cagra_rejects_truncated_vector_data(tmp_path, cagra_baseline): + index_path, index = _copy_baseline(cagra_baseline, tmp_path) + data_path = next(index_path.glob("*.vcag")) + data = data_path.read_bytes() + data_path.write_bytes(data[:-1]) + result = _search_copied_index(cagra_baseline, index) + assert not result.success + assert "do not exactly cover" in result.error_message + + +def test_cagra_rejects_corrupted_vector_data(tmp_path, cagra_baseline): + index_path, index = _copy_baseline(cagra_baseline, tmp_path) + data_path = next(index_path.glob("*.vcag")) + data = bytearray(data_path.read_bytes()) + data[len(data) // 2] ^= 0xFF + original_stat = data_path.stat() + data_path.write_bytes(data) + os.utime( + data_path, + ns=(original_stat.st_atime_ns, original_stat.st_mtime_ns), + ) + result = _search_copied_index(cagra_baseline, index) + assert not result.success + assert "checksum" in result.error_message.lower() + + +def test_cagra_rejects_committed_deletion(tmp_path, cagra_baseline): + index_path, index = _copy_baseline(cagra_baseline, tmp_path) + backend = _backend( + cagra_baseline.algo, + cagra_baseline.codec, + cagra_baseline.runtime_config, + ) + try: + _commit_one_deletion(backend, index_path) + with pytest.raises(RuntimeError, match="committed deletions"): + backend._get_runtime().verify_cagra_index(index_path) + result = _single_search_result( + backend.search( + cagra_baseline.dataset_case.dataset, + [index], + k=cagra_baseline.dataset_case.k, + ) + ) + finally: + backend.cleanup() + assert not result.success + assert "does not match the current Lucene commit" in result.error_message + + +def test_cagra_rejects_corrupted_metadata(tmp_path, cagra_baseline): + index_path, _ = _copy_baseline(cagra_baseline, tmp_path) + metadata_path = next(index_path.glob("*.vemc")) + contents = bytearray(metadata_path.read_bytes()) + contents[-1] ^= 0xFF + metadata_path.write_bytes(contents) + runtime = _PyLuceneRuntime.create(cagra_baseline.runtime_config) + with pytest.raises(RuntimeError, match="metadata format v0"): + runtime.verify_cagra_index(index_path) + + +def test_hnsw_rejects_missing_provenance(tmp_path, hnsw_baseline): + index_path, index = _copy_baseline(hnsw_baseline, tmp_path) + (index_path / _HNSW_PROVENANCE_FILE).unlink() + result = _search_copied_index(hnsw_baseline, index) + assert not result.success + assert "provenance manifest is missing" in result.error_message + + +def test_hnsw_rejects_unreadable_provenance(tmp_path, hnsw_baseline): + index_path, index = _copy_baseline(hnsw_baseline, tmp_path) + (index_path / _HNSW_PROVENANCE_FILE).write_text("{") + result = _search_copied_index(hnsw_baseline, index) + assert not result.success + assert "provenance manifest cannot be read" in result.error_message + + +def test_hnsw_rejects_changed_lucene_commit(tmp_path, hnsw_baseline): + index_path, index = _copy_baseline(hnsw_baseline, tmp_path) + segment_path = next(index_path.glob("segments_*")) + segment_path.write_bytes(segment_path.read_bytes() + b"stale") + result = _search_copied_index(hnsw_baseline, index) + assert not result.success + assert "does not match the current Lucene commit" in result.error_message + + +def test_hnsw_rejects_wrong_codec_provenance(tmp_path, hnsw_baseline): + index_path, index = _copy_baseline(hnsw_baseline, tmp_path) + manifest_path = index_path / _HNSW_PROVENANCE_FILE + payload = json.loads(manifest_path.read_bytes()) + payload["codec"] = _CAGRA_CODEC + manifest_path.write_text(json.dumps(payload)) + result = _search_copied_index(hnsw_baseline, index) + assert not result.success + assert "does not match the requested codec" in result.error_message + + +def _build_and_force_merge(runtime, index_path, vectors, writer_config): + directory = runtime.FSDirectory.open(runtime.Paths.get(str(index_path))) + writer = None + reader = None + try: + writer = runtime.IndexWriter(directory, writer_config) + for document_id, vector in enumerate(vectors): + writer.addDocument(runtime._vector_document(document_id, vector)) + writer.commit() + + reader = runtime.DirectoryReader.open(writer) + initial_segment_count = int(reader.leaves().size()) + reader.close() + reader = None + + writer.forceMerge(1) + writer.commit() + reader = runtime.DirectoryReader.open(writer) + final_segment_count = int(reader.leaves().size()) + finally: + try: + if reader is not None: + reader.close() + finally: + try: + if writer is not None: + writer.close() + finally: + directory.close() + return initial_segment_count, final_segment_count + + +def _assert_default_merge_scheduler(writer_config): + assert ( + str(writer_config.getMergeScheduler().getClass().getName()) + == "org.apache.lucene.index.ConcurrentMergeScheduler" + ) + + +def _committed_segment_compound_flags(runtime, index_path): + verifier = runtime._cagra_verifier + runtime.attach_current_thread() + directory = verifier.FSDirectory.open(verifier.Paths.get(str(index_path))) + try: + segment_infos = verifier.SegmentInfos.readLatestCommit(directory) + return [ + bool( + verifier.SegmentCommitInfo.cast_( + raw_segment + ).info.getUseCompoundFile() + ) + for raw_segment in segment_infos + ] + finally: + directory.close() + + +def test_configured_hnsw_codec_selects_gpu_writer( + tmp_path, pylucene_runtime_config +): + """Exercise configured GPU writer selection during default-scheduler merges.""" + runtime = _PyLuceneRuntime.create(pylucene_runtime_config.backend_config) + runtime.attach_current_thread() + m = 16 + ef_construction = 48 + java_codec = _configured_writer_selection_codec( + runtime, m, ef_construction + ) + assert str(java_codec) == ( + "PyLuceneWriterSelectionCodec(" + "PyLuceneConfiguredHnswCodec(m=16, efConstruction=48))" + ) + index_path = tmp_path / "hnsw-writer-selection-index" + index_path.mkdir() + vectors = ( + np.random.default_rng(1907).standard_normal((6, 32)).astype(np.float32) + ) + default_config = runtime.IndexWriterConfig() + writer_config = runtime._new_index_writer_config( + _BuildCodec( + codec_name=_HNSW_CODEC, + java_codec=java_codec, + writer_policy=_HNSW_WRITER_POLICY, + build_parameters={ + **_DEFAULT_HNSW_BUILD_PARAMETERS, + "m": m, + "ef_construction": ef_construction, + }, + ), + vector_count=vectors.shape[0], + ) + assert bool(writer_config.getUseCompoundFile()) == bool( + default_config.getUseCompoundFile() + ) + assert float( + writer_config.getMergePolicy().getNoCFSRatio() + ) == pytest.approx(float(default_config.getMergePolicy().getNoCFSRatio())) + _assert_default_merge_scheduler(writer_config) + writer_config.setMaxBufferedDocs(2) + + initial_segments, final_segments = _build_and_force_merge( + runtime, index_path, vectors, writer_config + ) + assert initial_segments >= 2 + assert final_segments == 1 + + writer_class, writer_calls = _writer_diagnostics(java_codec) + assert writer_class == _GPU_HNSW_WRITER + assert writer_calls >= 4 + search = runtime.search_index( + index_path, + vectors[[3]], + k=2, + batch_size=1, + num_candidates=2, + ) + assert search.hits[0][0].document_id == 3 + + +def test_configured_hnsw_codecs_retain_sequential_parameters( + pylucene_runtime_config, +): + """Resolve independent configured codecs in the process-wide JVM.""" + runtime = _PyLuceneRuntime.create(pylucene_runtime_config.backend_config) + + first = runtime.resolve_configured_hnsw_codec(16, 48) + assert str(first) == "PyLuceneConfiguredHnswCodec(m=16, efConstruction=48)" + + second = runtime.resolve_configured_hnsw_codec(24, 96) + assert ( + str(second) == "PyLuceneConfiguredHnswCodec(m=24, efConstruction=96)" + ) + assert str(first) == "PyLuceneConfiguredHnswCodec(m=16, efConstruction=48)" + assert str(first.getName()) == _HNSW_CODEC + assert str(second.getName()) == _HNSW_CODEC + + +def test_configured_hnsw_direct_single_segment_and_num_candidates( + tmp_path, pylucene_runtime_config, integration_dataset_case +): + """Build one committed leaf and over-fetch without changing top-k.""" + runtime = _PyLuceneRuntime.create(pylucene_runtime_config.backend_config) + build_parameters = { + **_DEFAULT_HNSW_BUILD_PARAMETERS, + "m": 16, + "ef_construction": 48, + "direct_single_segment": True, + } + java_codec = runtime.resolve_configured_hnsw_codec(16, 48) + index_path = tmp_path / "configured-hnsw-single-segment" + index_path.mkdir() + vectors = integration_dataset_case.dataset.training_vectors + + topology = runtime.build_index( + index_path, + vectors, + _BuildCodec( + codec_name=_HNSW_CODEC, + java_codec=java_codec, + writer_policy=_HNSW_WRITER_POLICY, + build_parameters=build_parameters, + ), + ) + + expected_vector_count = vectors.shape[0] + assert topology.segment_count == 1 + assert topology.segment_document_counts == (expected_vector_count,) + assert topology.segment_vector_counts == (expected_vector_count,) + + top_k = 5 + num_candidates = 11 + search = runtime.search_index( + index_path, + vectors[[137]], + k=top_k, + batch_size=1, + num_candidates=num_candidates, + ) + hits = search.hits[0] + document_ids = [hit.document_id for hit in hits] + assert len(hits) == top_k + assert document_ids[0] == 137 + assert len(set(document_ids)) == top_k + assert all( + 0 <= document_id < expected_vector_count + for document_id in document_ids + ) + + +def test_real_cagra_codec_merges_without_compound_files( + tmp_path, pylucene_runtime_config, integration_dataset_case +): + """Exercise CAGRA flush and merge layout through production configuration.""" + runtime = _PyLuceneRuntime.create(pylucene_runtime_config.backend_config) + java_codec = runtime.resolve_codec(_CAGRA_CODEC) + index_path = tmp_path / "cagra-merge-index" + index_path.mkdir() + vectors = integration_dataset_case.dataset.training_vectors + writer_config = runtime._new_index_writer_config( + _BuildCodec( + codec_name=_CAGRA_CODEC, + java_codec=java_codec, + writer_policy="gpu-cagra", + build_parameters=_CAGRA_BUILD_PARAMETERS, + ), + vector_count=vectors.shape[0], + ) + assert not bool(writer_config.getUseCompoundFile()) + assert float(writer_config.getMergePolicy().getNoCFSRatio()) == 0.0 + _assert_default_merge_scheduler(writer_config) + writer_config.setMaxBufferedDocs(128) + + initial_segments, final_segments = _build_and_force_merge( + runtime, index_path, vectors, writer_config + ) + assert initial_segments >= 2 + assert final_segments == 1 + assert _committed_segment_compound_flags(runtime, index_path) == [False] + + suffixes = {path.suffix for path in index_path.iterdir()} + assert ".vemc" in suffixes + assert ".vcag" in suffixes + assert suffixes.isdisjoint({".cfs", ".cfe"}) + verification = runtime.verify_cagra_index( + index_path, + expected_vector_count=vectors.shape[0], + expected_dimensions=vectors.shape[1], + ) + assert verification.segment_count == 1 + search = runtime.search_index( + index_path, + vectors[[137]], + k=1, + batch_size=1, + num_candidates=1, + ) + assert search.hits[0][0].document_id == 137 + + +def test_real_hnsw_codec_falls_back_to_cpu_in_fresh_process( + pylucene_runtime_config, +): + """Prove production fallback when CUDA devices are hidden at startup.""" + environment = os.environ.copy() + source_root = str(Path(__file__).resolve().parents[2]) + environment.update( + { + "CUDA_VISIBLE_DEVICES": "", + "CUVS_LUCENE_CUVS_JAVA_JAR": ( + pylucene_runtime_config.backend_config["cuvs_java_jar"] + ), + "CUVS_LUCENE_JAR": pylucene_runtime_config.backend_config[ + "cuvs_lucene_jar" + ], + "JAVA_LIBRARY_PATH": pylucene_runtime_config.backend_config[ + "java_library_path" + ], + "PYLUCENE_WRITER_SELECTION_CLASSES": str( + pylucene_runtime_config.writer_selection_classes + ), + "PYTHONPATH": os.pathsep.join( + value + for value in (source_root, environment.get("PYTHONPATH")) + if value + ), + } + ) + completed = subprocess.run( + [sys.executable, str(_CPU_FALLBACK_PROBE)], + env=environment, + capture_output=True, + text=True, + timeout=600, + check=False, + ) + assert completed.returncode == 0, ( + f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}" + ) + assert f"writerClass={_CPU_HNSW_WRITER}" in completed.stdout + assert "firstHit=2" in completed.stdout + + +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.backend_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) + runtime.build_index( + index_path, + vectors, + _BuildCodec( + codec_name=_CAGRA_CODEC, + java_codec=java_codec, + writer_policy="gpu-cagra", + build_parameters=_CAGRA_BUILD_PARAMETERS, + ), + ) + + 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_config, + } + ) + backend._runtime = runtime + try: + result = _single_search_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 a build/search parameter sweep 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.backend_config, + } + ) + ) + algorithm_config = tmp_path / "pylucene-hnsw-sweep.yaml" + algorithm_config.write_text( + json.dumps( + { + "name": "pylucene_cuvs_hnsw", + "groups": { + "test": { + "build": { + "codec": [_HNSW_CODEC], + "m": [16, 24], + "ef_construction": [32], + "direct_single_segment": [False], + }, + "search": {"num_candidates": [k, 11]}, + } + }, + } + ) + ) + + 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), + "--configuration", + str(algorithm_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 = _HNSW_CODEC + index_names = { + ( + f"pylucene_cuvs_hnsw[group=test][codec={codec}]" + f"[m={m}][ef_construction=32][direct_single_segment=false]" + ) + for m in (16, 24) + } + for index_name in index_names: + index_path = dataset_dir / "index" / index_name + assert any(index_path.glob("segments_*")) + assert (index_path / _HNSW_PROVENANCE_FILE).is_file() + + result_path = dataset_dir / "result" + build_csv = result_path / "build" / "pylucene_cuvs_hnsw,test.csv" + search_stem = f"pylucene_cuvs_hnsw,test,k{k},bs2" + with build_csv.open(newline="") as file: + build_rows = list(csv.DictReader(file)) + assert len(build_rows) == 2 + assert {row["index_name"] for row in build_rows} == index_names + assert {int(row["m"]) for row in build_rows} == {16, 24} + for build_row in build_rows: + assert float(build_row["time"]) > 0 + assert build_row["codec"] == codec + assert build_row["writer_policy"] == _HNSW_WRITER_POLICY + assert build_row["compound_file_policy"] == "lucene-default" + + 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) == 4 + assert {row["index_name"] for row in csv_rows} == index_names + assert {int(row["num_candidates"]) for row in csv_rows} == {k, 11} + assert { + (int(row["m"]), int(row["num_candidates"])) for row in csv_rows + } == { + (16, k), + (16, 11), + (24, k), + (24, 11), + } + for csv_row in csv_rows: + 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 float(csv_row["p50"]) > 0 + assert float(csv_row["p95"]) > 0 + assert float(csv_row["p99"]) > 0 + assert csv_row["codec"] == codec + assert csv_row["compound_file_policy"] == "lucene-default" + 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..8d29a54d1c --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_provenance.py @@ -0,0 +1,352 @@ +# +# 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, +) + + +_DEFAULT_HNSW_BUILD_PARAMETERS = { + "codec": _HNSW_CODEC, + "m": 32, + "ef_construction": 32, + "direct_single_segment": False, +} + + +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-with-cpu-fallback-provenance", + "schema_version": 4, + "codec": _HNSW_CODEC, + "build_parameters": _DEFAULT_HNSW_BUILD_PARAMETERS, + "writer_policy": "gpu-with-cpu-fallback", + "compound_file_policy": "lucene-default", + "vector_count": 10, + "dimensions": 4, + "segment_count": 1, + "commit_file_count": 1, + } + payload = json.loads(manifest_path.read_text()) + assert payload["build_parameters"] == _DEFAULT_HNSW_BUILD_PARAMETERS + assert payload["commit_fingerprints"] == [ + { + "name": "segments_1", + "sha256": ( + "1bc04b5291c26a46d918139138b992d2de976d6851d0893b0476b85bfbdfc6e6" + ), + } + ] + 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": 4, + "codec": _CAGRA_CODEC, + "build_parameters": {"codec": _CAGRA_CODEC}, + "writer_policy": "gpu-cagra", + "compound_file_policy": "disabled", + "vector_count": 10, + "dimensions": 4, + "segment_count": 1, + "commit_file_count": 1, + } + assert stat.S_IMODE(manifest_path.stat().st_mode) == 0o644 + + +def test_hnsw_provenance_rejects_build_parameter_mismatch(tmp_path): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + + with pytest.raises(RuntimeError, match="build parameter"): + pylucene_backend._verify_hnsw_provenance( + index_path, + _HNSW_CODEC, + expected_build_parameters={ + **_DEFAULT_HNSW_BUILD_PARAMETERS, + "m": 24, + }, + ) + + +def test_hnsw_provenance_rejects_inconsistent_direct_segment_count(tmp_path): + index_path = tmp_path / "index" + manifest_path = _prepare_hnsw_index(index_path) + build_parameters = { + **_DEFAULT_HNSW_BUILD_PARAMETERS, + "direct_single_segment": True, + } + payload = json.loads(manifest_path.read_text()) + payload["build_parameters"] = build_parameters + payload["segment_count"] = 2 + manifest_path.write_text(json.dumps(payload)) + + with pytest.raises(RuntimeError, match="must record exactly one segment"): + pylucene_backend._verify_hnsw_provenance( + index_path, + _HNSW_CODEC, + expected_build_parameters=build_parameters, + ) + + +@pytest.mark.parametrize( + "malformed_build_parameters", + [ + { + **_DEFAULT_HNSW_BUILD_PARAMETERS, + "m": True, + }, + { + **_DEFAULT_HNSW_BUILD_PARAMETERS, + "ef_construction": 513, + }, + { + **_DEFAULT_HNSW_BUILD_PARAMETERS, + "direct_single_segment": 0, + }, + ], +) +def test_hnsw_provenance_rejects_malformed_build_parameters( + tmp_path, malformed_build_parameters +): + index_path = tmp_path / "index" + manifest_path = _prepare_hnsw_index(index_path) + payload = json.loads(manifest_path.read_text()) + payload["build_parameters"] = malformed_build_parameters + manifest_path.write_text(json.dumps(payload)) + + with pytest.raises(RuntimeError, match="build parameter"): + pylucene_backend._verify_hnsw_provenance(index_path, _HNSW_CODEC) + + +@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-policy", "expected writer policy"), + ("wrong-compound-policy", "expected compound-file policy"), + ("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-policy": + payload["writer_policy"] = "gpu-cagra" + elif manifest_state == "wrong-compound-policy": + payload["compound_file_policy"] = "disabled" + 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..ddd9daec7b --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_runtime.py @@ -0,0 +1,1061 @@ +# +# 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_java as pylucene_java +import cuvs_bench.backends.pylucene as pylucene_backend +from cuvs_bench.backends.pylucene import _BuildCodec, _IndexTopology +from cuvs_bench.tests._pylucene_test_utils import _CAGRA_CODEC, _HNSW_CODEC + + +def _build_codec( + java_codec=None, + codec_name=_HNSW_CODEC, + *, + build_parameters=None, +) -> _BuildCodec: + if build_parameters is None: + build_parameters = ( + {"codec": _CAGRA_CODEC} + if codec_name == _CAGRA_CODEC + else { + "codec": _HNSW_CODEC, + "m": 32, + "ef_construction": 32, + "direct_single_segment": False, + } + ) + return _BuildCodec( + codec_name=codec_name, + java_codec=java_codec if java_codec is not None else object(), + writer_policy=( + "gpu-cagra" + if codec_name == _CAGRA_CODEC + else "gpu-with-cpu-fallback" + ), + build_parameters=build_parameters, + ) + + +class _FakeVMEnvironment: + def __init__(self): + self.attach_count = 0 + + def attachCurrentThread(self): + self.attach_count += 1 + + +class _FakeLuceneModule: + CLASSPATH = "/pylucene/lucene-core.jar" + VERSION = pylucene_backend._REQUIRED_PYLUCENE_VERSION + + 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 + self.force_merge_calls = [] + + 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") + + def forceMerge(self, segment_count): + self.force_merge_calls.append(segment_count) + + +class _FakeMergePolicy: + def __init__(self): + self.no_cfs_ratio = None + + def setNoCFSRatio(self, ratio): + self.no_cfs_ratio = ratio + + +class _FakeIndexWriterConfig: + OpenMode = SimpleNamespace(CREATE=object()) + DISABLE_AUTO_FLUSH = -1 + + def __init__(self): + self.use_compound_file = None + self.merge_policy = _FakeMergePolicy() + self.max_buffered_docs = None + self.ram_buffer_size_mb = None + + def setOpenMode(self, _mode): + pass + + def setCodec(self, _codec): + pass + + def setUseCompoundFile(self, enabled): + self.use_compound_file = enabled + + def getMergePolicy(self): + return self.merge_policy + + def setMaxBufferedDocs(self, max_buffered_docs): + self.max_buffered_docs = max_buffered_docs + + def setRAMBufferSizeMB(self, ram_buffer_size_mb): + self.ram_buffer_size_mb = ram_buffer_size_mb + + def setMergePolicy(self, merge_policy): + self.merge_policy = merge_policy + + def setMergeScheduler(self, _scheduler): + raise AssertionError("PyLucene must use Lucene's default scheduler") + + +class _FakeAvailableCodecs: + def __init__(self, names): + self.names = names + + def contains(self, name): + return name in self.names + + def __iter__(self): + return iter(self.names) + + +def _fake_codec_runtime(codec, available_names=(_HNSW_CODEC,)): + registry = SimpleNamespace( + calls=[], + availableCodecs=lambda: _FakeAvailableCodecs(available_names), + ) + + def for_name(codec_name): + registry.calls.append(codec_name) + return codec + + registry.forName = for_name + runtime = pylucene_backend._PyLuceneRuntime.__new__( + pylucene_backend._PyLuceneRuntime + ) + runtime.attach_current_thread = lambda: None + runtime.Codec = registry + runtime._codec_cache = {} + return runtime, registry + + +def _fake_configured_codec_runtime(diagnostics, initial_properties=None): + properties = dict(initial_properties or {}) + property_snapshots = [] + class_names = [] + + class _System: + @staticmethod + def getProperty(name): + return properties.get(name) + + @staticmethod + def setProperty(name, value): + properties[name] = value + + @staticmethod + def clearProperty(name): + properties.pop(name, None) + + class _Codec: + @staticmethod + def getName(): + return _HNSW_CODEC + + @staticmethod + def knnVectorsFormat(): + return object() + + def __str__(self): + return diagnostics + + codec = _Codec() + + def new_instance(): + property_snapshots.append(dict(properties)) + return codec + + def for_name(class_name): + class_names.append(class_name) + return SimpleNamespace(newInstance=new_instance) + + runtime = pylucene_backend._PyLuceneRuntime.__new__( + pylucene_backend._PyLuceneRuntime + ) + runtime.attach_current_thread = lambda: None + runtime.System = _System + runtime.Class = SimpleNamespace(forName=for_name) + runtime.Codec = SimpleNamespace(cast_=lambda reflected: reflected) + return runtime, properties, property_snapshots, class_names, codec + + +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.NoMergePolicy = SimpleNamespace(INSTANCE=object()) + + def create_writer(_directory, config): + runtime._test_writer_config = config + return writer + + runtime.IndexWriter = create_writer + runtime._vector_document = lambda document_id, vector: ( + document_id, + vector.copy(), + ) + runtime._test_writer = writer + runtime._test_directory = directory + runtime._test_codec = object() + runtime._index_topology = lambda _directory: _IndexTopology( + segment_document_counts=(len(writer.documents),), + segment_vector_counts=(len(writer.documents),), + ) + 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) + monkeypatch.setattr( + pylucene_backend, + "configured_codec_classes_path", + lambda *_args: "/configured-codec-classes", + ) + + +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(":") == [ + "/configured-codec-classes", + 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 + + +def test_configured_codec_compile_error_identifies_required_cuvs_lucene_api( + monkeypatch, +): + completed = SimpleNamespace( + returncode=1, + stderr="cannot find symbol: withHnswHeuristicType", + stdout="", + ) + monkeypatch.setattr(pylucene_java, "_find_javac", lambda: "/jdk/bin/javac") + monkeypatch.setattr( + pylucene_java.subprocess, + "run", + lambda *_args, **_kwargs: completed, + ) + + with pytest.raises( + RuntimeError, + match="PyLucene 10.2 support.*HNSW heuristic delegation", + ): + pylucene_java._compile("/dependencies") + + +@pytest.mark.parametrize( + "lucene", + [SimpleNamespace(VERSION="10.0.0"), SimpleNamespace()], + ids=["mismatched", "missing"], +) +def test_validate_pylucene_version_rejects_incompatible_binding(lucene): + with pytest.raises( + RuntimeError, + match="expected 10[.]2[.]0, found .*Activate a matching PyLucene build", + ): + pylucene_backend._validate_pylucene_version(lucene) + + +def test_initialize_pylucene_checks_version_before_starting_jvm(monkeypatch): + fake_lucene = _FakeLuceneModule() + fake_lucene.VERSION = "10.0.0" + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda _name: fake_lucene, + ) + + with pytest.raises(RuntimeError, match="expected 10[.]2[.]0"): + pylucene_backend._initialize_pylucene({}) + + assert fake_lucene.init_calls == [] + + +@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(":") == [ + "/configured-codec-classes", + 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({}) + + +def test_resolve_codec_validates_and_caches_initialized_vector_format(): + vectors_format = object() + codec = SimpleNamespace( + getName=lambda: _HNSW_CODEC, + knnVectorsFormat=lambda: vectors_format, + ) + runtime, registry = _fake_codec_runtime(codec) + + assert runtime.resolve_codec(_HNSW_CODEC) is codec + assert runtime.resolve_codec(_HNSW_CODEC) is codec + assert registry.calls == [_HNSW_CODEC] + + +@pytest.mark.parametrize( + ("available_names", "returned_name", "vectors_format", "error"), + [ + ((), _HNSW_CODEC, object(), "was not advertised by Lucene SPI"), + ((_HNSW_CODEC,), "DifferentCodec", object(), "Requested codec"), + ( + (_HNSW_CODEC,), + _HNSW_CODEC, + None, + "did not initialize a Lucene vector format", + ), + ], + ids=["missing-spi", "wrong-name", "missing-vector-format"], +) +def test_resolve_codec_rejects_unusable_codec( + available_names, returned_name, vectors_format, error +): + codec = SimpleNamespace( + getName=lambda: returned_name, + knnVectorsFormat=lambda: vectors_format, + ) + runtime, _ = _fake_codec_runtime(codec, available_names) + + with pytest.raises(RuntimeError, match=error): + runtime.resolve_codec(_HNSW_CODEC) + + +def test_resolve_configured_hnsw_codec_passes_parameters_and_restores_state(): + diagnostics = "PyLuceneConfiguredHnswCodec(m=24, efConstruction=96)" + runtime, properties, snapshots, class_names, codec = ( + _fake_configured_codec_runtime( + diagnostics, + {pylucene_backend.M_PROPERTY: "previous-m"}, + ) + ) + + assert runtime.resolve_configured_hnsw_codec(24, 96) is codec + assert snapshots == [ + { + pylucene_backend.M_PROPERTY: "24", + pylucene_backend.EF_CONSTRUCTION_PROPERTY: "96", + } + ] + assert properties == {pylucene_backend.M_PROPERTY: "previous-m"} + assert class_names == [pylucene_backend.CONFIGURED_HNSW_CODEC_CLASS] + + +def test_resolve_configured_hnsw_codec_rejects_diagnostic_mismatch_and_restores_state(): + runtime, properties, snapshots, _, _ = _fake_configured_codec_runtime( + "PyLuceneConfiguredHnswCodec(m=16, efConstruction=48)", + {pylucene_backend.EF_CONSTRUCTION_PROPERTY: "previous-ef"}, + ) + + with pytest.raises(RuntimeError, match="did not retain the requested"): + runtime.resolve_configured_hnsw_codec(24, 96) + + assert snapshots == [ + { + pylucene_backend.M_PROPERTY: "24", + pylucene_backend.EF_CONSTRUCTION_PROPERTY: "96", + } + ] + assert properties == { + pylucene_backend.EF_CONSTRUCTION_PROPERTY: "previous-ef" + } + + +@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, + } + ) + + +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) + + result = runtime.build_index( + tmp_path, + vectors, + _build_codec(runtime._test_codec), + ) + + assert result == _IndexTopology( + segment_document_counts=(2,), + segment_vector_counts=(2,), + ) + 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 + assert runtime._test_writer_config.use_compound_file is None + assert runtime._test_writer_config.merge_policy.no_cfs_ratio is None + assert runtime._test_writer.force_merge_calls == [] + + +@pytest.mark.parametrize( + ("codec_name", "use_compound_file", "no_cfs_ratio"), + [ + (_HNSW_CODEC, None, None), + (_CAGRA_CODEC, False, 0.0), + ], +) +def test_runtime_writer_config_applies_codec_compound_file_policy( + codec_name, use_compound_file, no_cfs_ratio +): + runtime = _fake_index_writer_runtime() + + config = runtime._new_index_writer_config( + _build_codec(runtime._test_codec, codec_name), + vector_count=10, + ) + + assert config.use_compound_file is use_compound_file + assert config.merge_policy.no_cfs_ratio == no_cfs_ratio + + +def test_runtime_writer_config_builds_one_segment_without_force_merge(): + runtime = _fake_index_writer_runtime() + build_parameters = { + "codec": _HNSW_CODEC, + "m": 24, + "ef_construction": 96, + "direct_single_segment": True, + } + + config = runtime._new_index_writer_config( + _build_codec( + runtime._test_codec, + build_parameters=build_parameters, + ), + vector_count=1_000_000, + ) + + assert config.max_buffered_docs == 1_000_001 + assert config.ram_buffer_size_mb == config.DISABLE_AUTO_FLUSH + assert type(config.ram_buffer_size_mb) is float + assert config.merge_policy is runtime.NoMergePolicy.INSTANCE + + +def test_runtime_direct_single_segment_build_does_not_force_merge(tmp_path): + runtime = _fake_index_writer_runtime() + vectors = np.zeros((2, 4), dtype=np.float32) + build_parameters = { + "codec": _HNSW_CODEC, + "m": 32, + "ef_construction": 32, + "direct_single_segment": True, + } + + topology = runtime.build_index( + tmp_path, + vectors, + _build_codec( + runtime._test_codec, + build_parameters=build_parameters, + ), + ) + + assert topology.segment_count == 1 + assert runtime._test_writer.force_merge_calls == [] + assert runtime._test_writer.close_called is True + assert runtime._test_directory.closed is True + + +def test_runtime_reads_committed_segment_topology_and_closes_reader(): + reader = SimpleNamespace(closed=False) + + def close_reader(): + reader.closed = True + + def leaf(document_count, vector_count): + leaf_reader = SimpleNamespace( + numDocs=lambda: document_count, + getFloatVectorValues=lambda _field: SimpleNamespace( + size=lambda: vector_count + ), + ) + return SimpleNamespace(reader=lambda: leaf_reader) + + reader.close = close_reader + reader.leaves = lambda: [leaf(4, 4), leaf(6, 6)] + runtime = pylucene_backend._PyLuceneRuntime.__new__( + pylucene_backend._PyLuceneRuntime + ) + runtime.DirectoryReader = SimpleNamespace(open=lambda _directory: reader) + + topology = runtime._index_topology(object()) + + assert topology == _IndexTopology( + segment_document_counts=(4, 6), + segment_vector_counts=(4, 6), + ) + assert reader.closed is True + + +def test_runtime_topology_rejects_segment_without_vectors_and_closes_reader(): + leaf_reader = SimpleNamespace( + numDocs=lambda: 2, + getFloatVectorValues=lambda _field: None, + ) + reader = SimpleNamespace( + leaves=lambda: [SimpleNamespace(reader=lambda: leaf_reader)], + closed=False, + ) + + def close_reader(): + reader.closed = True + + reader.close = close_reader + runtime = pylucene_backend._PyLuceneRuntime.__new__( + pylucene_backend._PyLuceneRuntime + ) + runtime.DirectoryReader = SimpleNamespace(open=lambda _directory: reader) + + with pytest.raises(RuntimeError, match="contains no vector values"): + runtime._index_topology(object()) + + assert reader.closed is True + + +@pytest.mark.parametrize( + ("topology", "direct_single_segment", "error"), + [ + (_IndexTopology((), ()), False, "no committed segments"), + ( + _IndexTopology((1,), (1,)), + False, + "document counts do not match", + ), + ( + _IndexTopology((2,), (1,)), + False, + "document and vector counts do not match", + ), + ( + _IndexTopology((1, 1), (1, 1)), + True, + "requested one committed Lucene segment", + ), + ], + ids=[ + "no-segments", + "wrong-document-count", + "wrong-vector-count", + "multiple-direct-segments", + ], +) +def test_runtime_build_rejects_invalid_topology_and_closes_resources( + tmp_path, topology, direct_single_segment, error +): + runtime = _fake_index_writer_runtime() + runtime._index_topology = lambda _directory: topology + vectors = np.zeros((2, 4), dtype=np.float32) + build_parameters = { + "codec": _HNSW_CODEC, + "m": 32, + "ef_construction": 32, + "direct_single_segment": direct_single_segment, + } + + with pytest.raises(RuntimeError, match=error): + runtime.build_index( + tmp_path, + vectors, + _build_codec( + runtime._test_codec, + build_parameters=build_parameters, + ), + ) + + 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, + _build_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, + _build_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, + _build_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, + _build_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_runtime_search_uses_candidates_for_query_and_top_k_for_results(): + query_calls = [] + search_calls = [] + query = object() + + def new_query(field, vector, num_candidates): + query_calls.append((field, vector, num_candidates)) + return query + + def search(received_query, top_k): + search_calls.append((received_query, top_k)) + return SimpleNamespace(scoreDocs=[SimpleNamespace(doc=4, score=0.5)]) + + runtime = pylucene_backend._PyLuceneRuntime.__new__( + pylucene_backend._PyLuceneRuntime + ) + runtime.KnnFloatVectorQuery = new_query + runtime._java_float_array = lambda vector: tuple(vector.tolist()) + searcher = SimpleNamespace(search=search) + stored_fields = SimpleNamespace( + document=lambda _document_id: SimpleNamespace(get=lambda _field: "17") + ) + + hits = runtime._search_vector( + searcher, + stored_fields, + np.array([1.0, 2.0], dtype=np.float32), + k=150, + num_candidates=300, + ) + + assert query_calls == [(pylucene_backend._VECTOR_FIELD, (1.0, 2.0), 300)] + assert search_calls == [(query, 150)] + assert hits == [pylucene_backend._SearchHit(document_id=17, score=0.5)] + + +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_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 24e54ff28f..6dbe62d041 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", @@ -68,6 +69,7 @@ elastic = "cuvs_bench.backends.elasticsearch:register" [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] diff --git a/python/cuvs_bench/tests/java/com/nvidia/cuvs/bench/PyLuceneWriterSelectionCodec.java b/python/cuvs_bench/tests/java/com/nvidia/cuvs/bench/PyLuceneWriterSelectionCodec.java new file mode 100644 index 0000000000..934238b73d --- /dev/null +++ b/python/cuvs_bench/tests/java/com/nvidia/cuvs/bench/PyLuceneWriterSelectionCodec.java @@ -0,0 +1,99 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.bench; + +import java.io.IOException; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.FilterCodec; +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; + +/** Test-only codec that reports the writer selected by the configured production HNSW codec. */ +public final class PyLuceneWriterSelectionCodec extends FilterCodec { + + private static final String CONFIGURED_CODEC_CLASS = + "com.nvidia.cuvs.bench.PyLuceneConfiguredHnswCodec"; + + private final KnnVectorsFormat knnVectorsFormat; + private final String configuredCodecDiagnostics; + + public PyLuceneWriterSelectionCodec() throws Exception { + this(configuredCodec()); + } + + private PyLuceneWriterSelectionCodec(Codec delegate) { + super(delegate.getName(), delegate); + knnVectorsFormat = new WriterSelectionFormat(delegate.knnVectorsFormat()); + configuredCodecDiagnostics = delegate.toString(); + } + + private static Codec configuredCodec() throws Exception { + return (Codec) + Class.forName(CONFIGURED_CODEC_CLASS).getConstructor().newInstance(); + } + + @Override + public KnnVectorsFormat knnVectorsFormat() { + return knnVectorsFormat; + } + + @Override + public String toString() { + return getClass().getSimpleName() + "(" + configuredCodecDiagnostics + ")"; + } + + private static final class WriterSelectionFormat extends KnnVectorsFormat { + + private static final String NOT_SELECTED = "not-selected"; + + private final KnnVectorsFormat delegate; + private String writerClass = NOT_SELECTED; + private int fieldsWriterCalls; + + private WriterSelectionFormat(KnnVectorsFormat delegate) { + super(delegate.getName()); + this.delegate = delegate; + } + + @Override + public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + KnnVectorsWriter writer = delegate.fieldsWriter(state); + recordWriter(writer.getClass().getName()); + return writer; + } + + private synchronized void recordWriter(String selectedClass) { + if (!NOT_SELECTED.equals(writerClass) && !writerClass.equals(selectedClass)) { + throw new AssertionError( + "Vector writer selection changed from " + writerClass + " to " + selectedClass); + } + writerClass = selectedClass; + fieldsWriterCalls++; + } + + @Override + public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { + return delegate.fieldsReader(state); + } + + @Override + public int getMaxDimensions(String fieldName) { + return delegate.getMaxDimensions(fieldName); + } + + @Override + public synchronized String toString() { + return getName() + + "(writerClass=" + + writerClass + + ", fieldsWriterCalls=" + + fieldsWriterCalls + + ")"; + } + } +}