From af71bf1f0bc1469bdb80e6c9aa9b2f0c1b0ed9a9 Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Fri, 31 Jul 2026 13:27:39 +0000 Subject: [PATCH 01/12] Declare h5py as a cuvs-bench dependency Dataset preparation imports h5py at runtime. Declare it in both the dependency manifest and project metadata so supported environments install it consistently. --- dependencies.yaml | 1 + python/cuvs_bench/pyproject.toml | 1 + 2 files changed, 2 insertions(+) diff --git a/dependencies.yaml b/dependencies.yaml index 970b173328..07529accc8 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -618,6 +618,7 @@ dependencies: packages: - click - cuvs==26.10.*,>=0.0.0a0 + - h5py>=3.8.0 - pandas - pyyaml - requests diff --git a/python/cuvs_bench/pyproject.toml b/python/cuvs_bench/pyproject.toml index 2db75baf38..162564fcf4 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", From 9b4312583f90538febaf7fccb1b9349870196a27 Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Fri, 31 Jul 2026 13:29:06 +0000 Subject: [PATCH 02/12] Persist results from in-process benchmark backends Add opt-in atomic JSON persistence for Python-native backends, preserve canonical result fields, and keep derived CSV artifacts synchronized. Propagate sweep failures through the CLI and cover result identity, export, and cleanup behavior. --- python/cuvs_bench/cuvs_bench/backends/base.py | 67 +- .../cuvs_bench/orchestrator/config_loaders.py | 7 +- .../cuvs_bench/orchestrator/orchestrator.py | 63 +- .../cuvs_bench/orchestrator/result_files.py | 232 ++++++ python/cuvs_bench/cuvs_bench/run/__main__.py | 33 +- .../cuvs_bench/cuvs_bench/run/data_export.py | 102 ++- .../cuvs_bench/cuvs_bench/tests/test_cli.py | 184 ++++- .../cuvs_bench/tests/test_data_export.py | 683 ++++++++++++++++++ .../cuvs_bench/tests/test_registry.py | 340 ++++++++- 9 files changed, 1658 insertions(+), 53 deletions(-) create mode 100644 python/cuvs_bench/cuvs_bench/orchestrator/result_files.py create mode 100644 python/cuvs_bench/cuvs_bench/tests/test_data_export.py diff --git a/python/cuvs_bench/cuvs_bench/backends/base.py b/python/cuvs_bench/cuvs_bench/backends/base.py index bf802d2128..a897fa2f0c 100644 --- a/python/cuvs_bench/cuvs_bench/backends/base.py +++ b/python/cuvs_bench/cuvs_bench/backends/base.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ @@ -243,15 +243,23 @@ def to_json(self) -> Dict[str, Any]: Dict[str, Any] Dictionary with benchmark results """ - return { - "name": f"{self.algorithm}/build", - "real_time": self.build_time_seconds, - "time_unit": "s", - "index_size": self.index_size_bytes, - "success": self.success, + result = { **self.build_params, **self.metadata, } + result.pop("error_message", None) + result.update( + { + "name": f"{self.algorithm}/build", + "real_time": self.build_time_seconds, + "time_unit": "s", + "index_size": self.index_size_bytes, + "success": self.success, + } + ) + if self.error_message is not None: + result["error_message"] = self.error_message + return result @dataclass @@ -288,6 +296,10 @@ class SearchResult: Whether the search succeeded error_message : Optional[str] Error message if search failed + latency_seconds : Optional[float] + Backend-reported mean latency in seconds. For batched backends, this + is mean batch latency. This field is last to preserve + compatibility with existing positional construction. """ neighbors: np.ndarray @@ -303,6 +315,7 @@ class SearchResult: metadata: Dict[str, Any] = field(default_factory=dict) success: bool = True error_message: Optional[str] = None + latency_seconds: Optional[float] = None def to_json(self) -> Dict[str, Any]: """ @@ -313,26 +326,39 @@ def to_json(self) -> Dict[str, Any]: Dict[str, Any] Dictionary with benchmark results """ - result = { - "name": f"{self.algorithm}/search", - "real_time": self.search_time_ms, - "time_unit": "ms", - "items_per_second": self.queries_per_second, - "Recall": self.recall, - "success": self.success, - "search_params": self.search_params, - **self.metadata, - } - + result = dict(self.metadata) + result.pop("error_message", None) if self.latency_percentiles: result.update(self.latency_percentiles) + real_time_ms = ( + self.latency_seconds * 1000.0 + if self.latency_seconds is not None + else self.search_time_ms + ) + result.update( + { + "name": f"{self.algorithm}/search", + "real_time": real_time_ms, + "time_unit": "ms", + "search_time_ms": self.search_time_ms, + "items_per_second": self.queries_per_second, + "Recall": self.recall, + "success": self.success, + "search_params": self.search_params, + } + ) + if self.latency_seconds is not None: + result["Latency"] = self.latency_seconds if self.gpu_time_seconds is not None: result["GPU"] = self.gpu_time_seconds if self.cpu_time_seconds is not None: result["cpu_time"] = self.cpu_time_seconds + if self.error_message is not None: + result["error_message"] = self.error_message + return result @@ -368,6 +394,11 @@ class BenchmarkBackend(ABC): - requires_network : bool - Whether backend requires network (default: False) """ + # Opt in when this backend's sweep results should be persisted by the + # orchestrator. Tune-mode and existing backends retain their own + # result-file behavior. + orchestrator_persists_results = False + def __init__(self, config: Dict[str, Any]): """Initialize backend with configuration.""" self.config = config diff --git a/python/cuvs_bench/cuvs_bench/orchestrator/config_loaders.py b/python/cuvs_bench/cuvs_bench/orchestrator/config_loaders.py index e3af2fe58e..5e3a51b7f7 100644 --- a/python/cuvs_bench/cuvs_bench/orchestrator/config_loaders.py +++ b/python/cuvs_bench/cuvs_bench/orchestrator/config_loaders.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -91,6 +91,11 @@ def search_params_list(self) -> List[Dict[str, Any]]: def index_path(self) -> Path: return Path(self.indexes[0].file) if self.indexes else Path("") + @property + def output_filename(self) -> Tuple[str, str]: + """Return the build and search result filename stems.""" + return self.backend_config.get("output_filename", ("", "")) + @dataclass class DatasetConfig: diff --git a/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py b/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py index 580291c2f0..dce0cd1c6f 100644 --- a/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py +++ b/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -10,7 +10,7 @@ benchmark runs across different backends using the registry pattern. """ -from typing import List, Optional, Union +from typing import List, Optional, Set, Tuple, Union import numpy as np @@ -22,6 +22,11 @@ ) from ..backends._utils import compute_recall from .config_loaders import DatasetConfig +from .result_files import ( + ResultRecord, + validate_sweep_output_filenames, + write_result_files, +) class BenchmarkOrchestrator: @@ -217,6 +222,8 @@ def _run_sweep( # Collect results results: List[Union[BuildResult, SearchResult]] = [] + result_records: List[ResultRecord] = [] + stale_search_filenames: Set[str] = set() # Run benchmarks # Each config contains ALL indexes for one executable (matches runners.py) @@ -224,6 +231,17 @@ def _run_sweep( # Create backend instance backend = self.backend_class(config.backend_config) try: + output_filenames: Optional[Tuple[str, str]] = None + if backend.orchestrator_persists_results: + if len(config.indexes) != 1: + raise ValueError( + "Opt-in sweep result persistence requires exactly " + "one index per benchmark configuration" + ) + output_filenames = validate_sweep_output_filenames( + config.output_filename + ) + backend.initialize() if build: @@ -235,8 +253,18 @@ def _run_sweep( dry_run=dry_run, ) results.append(build_result) + if output_filenames is not None: + result_records.append( + ResultRecord( + result=build_result, + index_name=config.index_name, + output_filename=output_filenames[0], + ) + ) if not build_result.success: + if search and output_filenames is not None: + stale_search_filenames.add(output_filenames[1]) print( f"Build failed for {config.index_name}: {build_result.error_message}" ) @@ -274,6 +302,14 @@ def _run_sweep( ) results.append(search_result) + if output_filenames is not None: + result_records.append( + ResultRecord( + result=search_result, + index_name=config.index_name, + output_filename=output_filenames[1], + ) + ) if not search_result.success: print( @@ -282,6 +318,14 @@ def _run_sweep( finally: backend.cleanup() + if (result_records or stale_search_filenames) and not dry_run: + write_result_files( + result_records, + dataset=dataset_config.name, + dataset_path=loader_kwargs["dataset_path"], + stale_search_filenames=stale_search_filenames, + ) + return results def _run_tune( @@ -503,10 +547,17 @@ def objective(trial) -> float: print(f"Best params: {study.best_params}") print(f"Best {optimize_metric}: {study.best_value:.4f}") except ValueError: - print( - f"\n⚠️ All {n_trials} trials were pruned (constraints not met)" - ) - print(f" Consider relaxing constraints: {hard_constraints}") + if any(result.success for result in all_results): + print( + f"\n⚠️ All {n_trials} trials were pruned " + "(constraints not met)" + ) + print(f" Consider relaxing constraints: {hard_constraints}") + else: + print( + f"\n⚠️ All {n_trials} benchmark trials failed before " + "producing valid metrics" + ) return all_results diff --git a/python/cuvs_bench/cuvs_bench/orchestrator/result_files.py b/python/cuvs_bench/cuvs_bench/orchestrator/result_files.py new file mode 100644 index 0000000000..8985aa0051 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/orchestrator/result_files.py @@ -0,0 +1,232 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Persistence for benchmark backends that return in-process results.""" + +from __future__ import annotations + +import json +import re +import stat +import tempfile +from collections import defaultdict +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict, Iterable, List, Tuple, Union + +from ..backends.base import BuildResult, SearchResult + +BenchmarkResult = Union[BuildResult, SearchResult] +_SEARCH_SUFFIX_PATTERN = re.compile(r"k[1-9]\d*,bs[1-9]\d*") + + +@dataclass(frozen=True) +class ResultRecord: + """Associate a backend result with its benchmark artifact identity.""" + + result: BenchmarkResult + index_name: str + output_filename: str + + @property + def phase(self) -> str: + return "build" if isinstance(self.result, BuildResult) else "search" + + def to_json(self) -> Dict[str, Any]: + row = self.result.to_json() + row["name"] = f"{self.index_name}/{self.phase}" + return row + + +def _validate_path_component(value: object, description: str) -> str: + if not isinstance(value, str): + raise ValueError(f"Invalid {description}: {value!r}") + path = Path(value) + if not value or path.name != value or value in {".", ".."}: + raise ValueError(f"Invalid {description}: {value!r}") + return value + + +def validate_dataset_name(dataset: object) -> str: + """Validate a dataset name before using it as a path component.""" + return _validate_path_component(dataset, "benchmark dataset name") + + +def _validate_output_filename(output_filename: str) -> None: + _validate_path_component(output_filename, "benchmark result filename") + + +def validate_sweep_output_filenames( + output_filenames: object, +) -> Tuple[str, str]: + """Validate and normalize opt-in sweep result filename stems.""" + if ( + not isinstance(output_filenames, (list, tuple)) + or len(output_filenames) != 2 + ): + raise ValueError( + "Opt-in sweep persistence requires exactly two result filename " + "stems: (build_stem, search_stem)" + ) + + build_stem, search_stem = output_filenames + _validate_output_filename(build_stem) + _validate_output_filename(search_stem) + + build_parts = build_stem.split(",") + if len(build_parts) < 2 or any(not part for part in build_parts): + raise ValueError( + "Build result filename stem must contain at least " + "','" + ) + + expected_prefix = f"{build_stem}," + search_suffix = ( + search_stem[len(expected_prefix) :] + if search_stem.startswith(expected_prefix) + else "" + ) + if _SEARCH_SUFFIX_PATTERN.fullmatch(search_suffix) is None: + raise ValueError( + "Search result filename stem must be " + "',k,bs'" + ) + + return build_stem, search_stem + + +def _read_benchmarks(path: Path) -> List[Dict[str, Any]]: + if not path.exists(): + return [] + with path.open(encoding="utf-8") as file: + payload = json.load(file) + benchmarks = ( + payload.get("benchmarks") if isinstance(payload, dict) else None + ) + if not isinstance(benchmarks, list): + raise ValueError( + f"Benchmark result file must contain a 'benchmarks' list: {path}" + ) + seen_names = set() + for position, row in enumerate(benchmarks): + if not isinstance(row, dict) or not isinstance(row.get("name"), str): + raise ValueError( + f"Benchmark row {position} must be an object with a string " + f"'name': {path}" + ) + if row["name"] in seen_names: + raise ValueError( + f"Benchmark result file contains duplicate name " + f"{row['name']!r}: {path}" + ) + seen_names.add(row["name"]) + return benchmarks + + +def _merge_build_rows( + path: Path, records: List[ResultRecord] +) -> List[Dict[str, Any]]: + rows_by_name = {row["name"]: row for row in _read_benchmarks(path)} + for record in records: + row = record.to_json() + previous = rows_by_name.get(row["name"]) + if ( + record.result.success + and record.result.metadata.get("skipped") + and previous is not None + and previous.get("success", True) is True + ): + continue + rows_by_name[row["name"]] = row + return list(rows_by_name.values()) + + +def _write_payload(path: Path, rows: List[Dict[str, Any]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + file_mode = stat.S_IMODE(path.stat().st_mode) if path.exists() else 0o644 + temp_path = None + try: + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + delete=False, + ) as file: + temp_path = Path(file.name) + json.dump({"benchmarks": rows}, file, indent=2, allow_nan=False) + file.write("\n") + temp_path.chmod(file_mode) + temp_path.replace(path) + finally: + if temp_path is not None: + temp_path.unlink(missing_ok=True) + + +def _invalidate_derived_csvs(path: Path, phase: str) -> None: + suffixes = ( + (".csv",) + if phase == "build" + else (",raw.csv", ",throughput.csv", ",latency.csv") + ) + stem = path.with_suffix("") + for suffix in suffixes: + stem.with_name(f"{stem.name}{suffix}").unlink(missing_ok=True) + + +def _remove_result_artifacts(path: Path, phase: str) -> None: + path.unlink(missing_ok=True) + _invalidate_derived_csvs(path, phase) + + +def write_result_files( + records: List[ResultRecord], + dataset: str, + dataset_path: str, + *, + stale_search_filenames: Iterable[str] = (), +) -> None: + """Write grouped Google Benchmark-compatible JSON result files.""" + validate_dataset_name(dataset) + grouped: Dict[Tuple[str, str], List[ResultRecord]] = defaultdict(list) + for record in records: + _validate_output_filename(record.output_filename) + grouped[(record.phase, record.output_filename)].append(record) + + stale_search_filenames = set(stale_search_filenames) + for output_filename in stale_search_filenames: + _validate_output_filename(output_filename) + + result_root = Path(dataset_path) / dataset / "result" + current_search_filenames = { + output_filename + for phase, output_filename in grouped + if phase == "search" + } + for output_filename in stale_search_filenames - current_search_filenames: + _remove_result_artifacts( + result_root / "search" / f"{output_filename}.json", + "search", + ) + + for (phase, output_filename), file_records in grouped.items(): + output_path = result_root / phase / f"{output_filename}.json" + rows = ( + _merge_build_rows(output_path, file_records) + if phase == "build" + else [record.to_json() for record in file_records] + ) + if rows: + _write_payload(output_path, rows) + _invalidate_derived_csvs(output_path, phase) + + +__all__ = [ + "ResultRecord", + "validate_dataset_name", + "validate_sweep_output_filenames", + "write_result_files", +] diff --git a/python/cuvs_bench/cuvs_bench/run/__main__.py b/python/cuvs_bench/cuvs_bench/run/__main__.py index 6950ff7202..eb3726ffc9 100644 --- a/python/cuvs_bench/cuvs_bench/run/__main__.py +++ b/python/cuvs_bench/cuvs_bench/run/__main__.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -11,8 +11,12 @@ import click import yaml -from .data_export import convert_json_to_csv_build, convert_json_to_csv_search +from .data_export import ( + convert_json_to_csv_build, + convert_json_to_csv_search, +) from ..orchestrator import BenchmarkOrchestrator +from ..orchestrator.result_files import validate_dataset_name @click.command() @@ -257,6 +261,11 @@ def main( and any backend-specific connection parameters (host, port, etc.). """ + try: + validate_dataset_name(dataset) + except ValueError as exc: + raise click.BadParameter(str(exc), param_hint="--dataset") from exc + if not data_export: # Determine backend type and extra kwargs from --backend-config backend_type = "cpp_gbench" @@ -277,7 +286,7 @@ def main( backend_kwargs = cfg orchestrator = BenchmarkOrchestrator(backend_type=backend_type) - orchestrator.run_benchmark( + results = orchestrator.run_benchmark( mode=mode, constraints=json.loads(constraints) if constraints else None, n_trials=n_trials, @@ -300,6 +309,24 @@ def main( executable_dir=executable_dir, **backend_kwargs, ) + failures = [result for result in results if not result.success] + if mode == "sweep": + run_failed = not results or bool(failures) + else: + run_failed = not any(result.success for result in results) + if run_failed: + details = ( + "; ".join( + f"{result.algorithm}: " + f"{result.error_message or 'benchmark failed'}" + for result in failures + ) + if failures + else f"{mode} mode produced no benchmark results" + ) + raise click.ClickException(details) + if dry_run: + return convert_json_to_csv_build(dataset, dataset_path) convert_json_to_csv_search(dataset, dataset_path) diff --git a/python/cuvs_bench/cuvs_bench/run/data_export.py b/python/cuvs_bench/cuvs_bench/run/data_export.py index 707677a083..40078b2a9c 100644 --- a/python/cuvs_bench/cuvs_bench/run/data_export.py +++ b/python/cuvs_bench/cuvs_bench/run/data_export.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -9,6 +9,8 @@ import pandas as pd +from ..orchestrator.result_files import validate_dataset_name + skip_build_cols = set( [ "algo_name", @@ -25,11 +27,24 @@ "real_time", "time_unit", "index_size", + "success", + "error_message", + "skipped", ] ) skip_search_cols = ( - set(["recall", "qps", "latency", "items_per_second", "Recall", "Latency"]) + set( + [ + "recall", + "qps", + "latency", + "items_per_second", + "Recall", + "Latency", + "search_params", + ] + ) | skip_build_cols ) @@ -69,7 +84,11 @@ def read_json_files(dataset, dataset_path, method): A tuple containing the file path, algorithm name, and the DataFrame of JSON content. """ + validate_dataset_name(dataset) dir_path = os.path.join(dataset_path, dataset, "result", method) + if not os.path.isdir(dir_path): + return + for file in os.listdir(dir_path): if file.endswith(".json"): file_path = os.path.join(dir_path, file) @@ -77,6 +96,10 @@ def read_json_files(dataset, dataset_path, method): with open(file_path, "r", encoding="ISO-8859-1") as f: data = json.load(f) df = pd.DataFrame(data["benchmarks"]) + if "success" in df.columns: + df = df[df["success"].ne(False)].copy() + if method == "build" and "skipped" in df.columns: + df = df[df["skipped"].ne(True)].copy() algo_name = tuple(file.split(",")[:2]) yield file_path, algo_name, df except Exception as e: @@ -103,6 +126,15 @@ def clean_algo_name(algo_name): return name.removesuffix(".json") +def _remove_derived_csvs(file, suffixes): + stem = os.path.splitext(file)[0] + for suffix in suffixes: + try: + os.remove(f"{stem}{suffix}") + except FileNotFoundError: + pass + + def write_csv(file, algo_name, df, extra_columns=None, skip_cols=None): """ Write a DataFrame to CSV with specified columns skipped. @@ -153,6 +185,9 @@ def convert_json_to_csv_build(dataset, dataset_path): """ for file, algo_name, df in read_json_files(dataset, dataset_path, "build"): try: + if df.empty: + _remove_derived_csvs(file, (".csv",)) + continue algo_name = clean_algo_name(algo_name) write_csv(file, algo_name, df, skip_cols=skip_build_cols) except Exception as e: @@ -175,12 +210,26 @@ def convert_json_to_csv_search(dataset, dataset_path): dataset, dataset_path, "search" ): try: + if df.empty: + _remove_derived_csvs( + file, (",raw.csv", ",throughput.csv", ",latency.csv") + ) + continue + search_stem = os.path.splitext(os.path.basename(file))[0] + # Search stems end in k and batch-size components. Preserve every + # preceding component, including an optional dataset subset. + search_stem_parts = search_stem.rsplit(",", maxsplit=2) + build_stem = ( + search_stem_parts[0] + if len(search_stem_parts) == 3 + else ",".join(algo_name) + ) build_file = os.path.join( dataset_path, dataset, "result", "build", - f"{','.join(algo_name)}.csv", + f"{build_stem}.csv", ) algo_name = clean_algo_name(algo_name) df["name"] = df["name"].str.split("/").str[0] @@ -202,29 +251,46 @@ def convert_json_to_csv_search(dataset, dataset_path): ] if os.path.exists(build_file): build_df = pd.read_csv(build_file) - write_ncols = len(write.columns) write["build time"] = None write["build threads"] = None write["build cpu_time"] = None - start_idx = 5 + build_columns = { + "time": "build time", + "threads": "build threads", + "cpu_time": "build cpu_time", + } if "GPU" in build_df.columns: - start_idx = 6 write["build GPU"] = None - for col_idx in range(start_idx, len(build_df.columns)): - col_name = build_df.columns[col_idx] - write[col_name] = None - if col_name == "num_threads": - write["build_num_threads"] = None + build_columns["GPU"] = "build GPU" + + ignored_columns = {"algo_name", "index_name"} + for col_name in build_df.columns: + if ( + col_name in ignored_columns + or col_name in build_columns + ): + continue + target_name = ( + "build_num_threads" + if col_name == "num_threads" + else col_name + ) + if target_name not in write.columns: + write[target_name] = None + build_columns[col_name] = target_name + for s_index, search_row in write.iterrows(): - for b_index, build_row in build_df.iterrows(): + for _, build_row in build_df.iterrows(): if search_row["index_name"] == build_row["index_name"]: - write.iloc[s_index, write_ncols] = build_df.iloc[ - b_index, 2 - ] - write.iloc[s_index, write_ncols + 1 :] = ( - build_df.iloc[b_index, 3:] - ) + for ( + source_name, + target_name, + ) in build_columns.items(): + if source_name in build_df.columns: + write.at[s_index, target_name] = build_row[ + source_name + ] break # Write search data and compute frontiers write.to_csv(file.replace(".json", ",raw.csv"), index=False) diff --git a/python/cuvs_bench/cuvs_bench/tests/test_cli.py b/python/cuvs_bench/cuvs_bench/tests/test_cli.py index fa2a63d041..a1525d234d 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_cli.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_cli.py @@ -1,9 +1,10 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # from pathlib import Path +from types import SimpleNamespace import pandas as pd import pytest @@ -519,11 +520,182 @@ def test_plot_command_creates_png_files(temp_datasets_dir: Path): ) -# FIXME: Tests below use --dry-run to verify CLI flag parsing and orchestrator -# routing without requiring actual benchmark execution. Tune mode (--mode tune) -# requires Optuna and actual search results, so only flag acceptance is tested -# here. End-to-end tests for tune mode and non-C++ backends should be added -# when those features are exercised in integration testing. +# The mocked-result tests below isolate CLI outcome and export semantics. +# Later tests use --dry-run to verify flag parsing and orchestrator routing +# without requiring native benchmark execution. + + +def _invoke_run_with_results(monkeypatch, tmp_path, mode, results): + from cuvs_bench.run import __main__ as run_module + + orchestrator = SimpleNamespace( + run_benchmark=lambda **_kwargs: results, + ) + monkeypatch.setattr( + run_module, + "BenchmarkOrchestrator", + lambda backend_type: orchestrator, + ) + exported = [] + monkeypatch.setattr( + run_module, + "convert_json_to_csv_build", + lambda dataset, dataset_path: exported.append( + ("build", dataset, dataset_path) + ), + ) + monkeypatch.setattr( + run_module, + "convert_json_to_csv_search", + lambda dataset, dataset_path: exported.append( + ("search", dataset, dataset_path) + ), + ) + + result = CliRunner().invoke( + run_module.main, + [ + "--dataset", + "test-data", + "--dataset-path", + str(tmp_path), + "--algorithms", + "fake", + "--groups", + "base", + "--batch-size", + "1", + "-k", + "1", + "-m", + "latency", + "--mode", + mode, + ], + ) + return result, exported + + +def _benchmark_result(success, error_message=None): + return SimpleNamespace( + success=success, + algorithm="fake", + error_message=error_message, + ) + + +def test_tune_mixed_trial_results_exit_zero_and_export(monkeypatch, tmp_path): + result, exported = _invoke_run_with_results( + monkeypatch, + tmp_path, + "tune", + [ + _benchmark_result(False, "trial failed"), + _benchmark_result(True), + ], + ) + + assert result.exit_code == 0, result.output + assert exported == [ + ("build", "test-data", str(tmp_path)), + ("search", "test-data", str(tmp_path)), + ] + + +@pytest.mark.parametrize( + ("results", "expected_message"), + [ + ([_benchmark_result(False, "trial failed")], "trial failed"), + ([], "tune mode produced no benchmark results"), + ], +) +def test_tune_without_successful_results_exits_nonzero( + monkeypatch, tmp_path, results, expected_message +): + result, exported = _invoke_run_with_results( + monkeypatch, tmp_path, "tune", results + ) + + assert result.exit_code != 0 + assert expected_message in result.output + assert exported == [] + + +def test_tune_all_constraint_pruned_measurements_exit_zero_and_export( + monkeypatch, tmp_path +): + # _run_tune retains successful measurements even when Optuna prunes every + # trial for violating a hard constraint. + result, exported = _invoke_run_with_results( + monkeypatch, + tmp_path, + "tune", + [_benchmark_result(True), _benchmark_result(True)], + ) + + assert result.exit_code == 0, result.output + assert exported == [ + ("build", "test-data", str(tmp_path)), + ("search", "test-data", str(tmp_path)), + ] + + +def test_sweep_mixed_results_exit_nonzero(monkeypatch, tmp_path): + result, exported = _invoke_run_with_results( + monkeypatch, + tmp_path, + "sweep", + [ + _benchmark_result(False, "benchmark failed"), + _benchmark_result(True), + ], + ) + + assert result.exit_code != 0 + assert "benchmark failed" in result.output + assert exported == [] + + +def test_sweep_without_results_exits_nonzero(monkeypatch, tmp_path): + result, exported = _invoke_run_with_results( + monkeypatch, tmp_path, "sweep", [] + ) + + assert result.exit_code != 0 + assert "sweep mode produced no benchmark results" in result.output + assert exported == [] + + +@pytest.mark.parametrize( + "dataset", ["..", "../escape", "nested/dataset", "/absolute/dataset"] +) +def test_data_export_rejects_invalid_dataset_name(tmp_path, dataset): + from cuvs_bench.run.__main__ import main as run_main + + runner = CliRunner() + result = runner.invoke( + run_main, + [ + "--data-export", + "--dataset", + dataset, + "--dataset-path", + str(tmp_path), + "--count", + "10", + "--batch-size", + "100", + "--algorithms", + "cuvs_cagra", + "--groups", + "base", + "--search-mode", + "latency", + ], + ) + + assert result.exit_code == 2 + assert "Invalid benchmark dataset name" in result.output def test_run_with_mode_sweep(temp_datasets_dir): diff --git a/python/cuvs_bench/cuvs_bench/tests/test_data_export.py b/python/cuvs_bench/cuvs_bench/tests/test_data_export.py new file mode 100644 index 0000000000..3898c7f624 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_data_export.py @@ -0,0 +1,683 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Regression tests for in-process benchmark result persistence and export.""" + +import json +import stat + +import numpy as np +import pandas as pd +import pytest + +from cuvs_bench.backends.base import BuildResult, SearchResult +from cuvs_bench.orchestrator.result_files import ( + ResultRecord, + write_result_files, +) +from cuvs_bench.run.data_export import ( + convert_json_to_csv_build, + convert_json_to_csv_search, + read_json_files, +) + + +_DATASET = "test-data" +_ALGORITHM = "pylucene_cuvs_hnsw" +_BUILD_STEM = f"{_ALGORITHM},base" +_SEARCH_STEM = f"{_BUILD_STEM},k5,bs2" +_CODECS = ( + "Lucene101AcceleratedHNSWCodec", + "Lucene101AcceleratedHNSWBaseLayerCodec", + "Lucene101AcceleratedHNSWMultiLayerCodec", +) + + +def _index_name(codec): + return f"{_ALGORITHM}.codec{codec}" + + +def _build_result( + codec, + *, + build_time=1.25, + metadata=None, + success=True, + error_message=None, +): + return BuildResult( + index_path=f"/indexes/{_index_name(codec)}", + build_time_seconds=build_time, + index_size_bytes=4096, + algorithm=_ALGORITHM, + build_params={"codec": codec}, + metadata=metadata or {}, + success=success, + error_message=error_message, + ) + + +def _search_result( + codec, + *, + search_time_ms=6.0, + latency_seconds=0.003, + recall=0.9, + queries_per_second=500.0, + metadata=None, + success=True, + error_message=None, +): + result_metadata = {"codec": codec, "num_batches": 2} + if metadata: + result_metadata.update(metadata) + return SearchResult( + neighbors=np.empty((0, 5), dtype=np.int64), + distances=np.empty((0, 5), dtype=np.float32), + search_time_ms=search_time_ms, + queries_per_second=queries_per_second, + recall=recall, + algorithm=_ALGORITHM, + search_params=[{}], + latency_seconds=latency_seconds, + latency_percentiles={"p50": 2.5, "p95": 3.5, "p99": 3.9}, + metadata=result_metadata, + success=success, + error_message=error_message, + ) + + +def _record(result, codec, output_filename): + return ResultRecord( + result=result, + index_name=_index_name(codec), + output_filename=output_filename, + ) + + +def _load_payload(path): + with path.open(encoding="utf-8") as file: + return json.load(file) + + +def test_write_result_files_groups_three_codecs_with_expected_schema(tmp_path): + records = [] + for position, codec in enumerate(_CODECS, start=1): + records.extend( + [ + _record( + _build_result(codec, build_time=float(position)), + codec, + _BUILD_STEM, + ), + _record( + _search_result( + codec, + search_time_ms=float(position * 4), + latency_seconds=float(position) / 500.0, + recall=0.8 + position / 100.0, + queries_per_second=1000.0 / position, + ), + codec, + _SEARCH_STEM, + ), + ] + ) + + write_result_files(records, dataset=_DATASET, dataset_path=str(tmp_path)) + + result_root = tmp_path / _DATASET / "result" + build_path = result_root / "build" / f"{_BUILD_STEM}.json" + search_path = result_root / "search" / f"{_SEARCH_STEM}.json" + assert sorted( + str(path.relative_to(result_root)) + for path in result_root.rglob("*.json") + ) == [ + f"build/{_BUILD_STEM}.json", + f"search/{_SEARCH_STEM}.json", + ] + + build_rows = _load_payload(build_path)["benchmarks"] + search_rows = _load_payload(search_path)["benchmarks"] + assert len(build_rows) == len(_CODECS) + assert len(search_rows) == len(_CODECS) + assert {row["name"] for row in build_rows} == { + f"{_index_name(codec)}/build" for codec in _CODECS + } + assert {row["name"] for row in search_rows} == { + f"{_index_name(codec)}/search" for codec in _CODECS + } + + first_build = build_rows[0] + assert first_build == { + "codec": _CODECS[0], + "name": f"{_index_name(_CODECS[0])}/build", + "real_time": 1.0, + "time_unit": "s", + "index_size": 4096, + "success": True, + } + + first_search = search_rows[0] + assert first_search["real_time"] == pytest.approx(2.0) + assert first_search["time_unit"] == "ms" + assert first_search["search_time_ms"] == pytest.approx(4.0) + assert first_search["Latency"] == pytest.approx(0.002) + assert first_search["items_per_second"] == pytest.approx(1000.0) + assert first_search["Recall"] == pytest.approx(0.81) + assert first_search["p50"] == pytest.approx(2.5) + assert first_search["p95"] == pytest.approx(3.5) + assert first_search["p99"] == pytest.approx(3.9) + assert first_search["codec"] == _CODECS[0] + + +def test_skipped_build_preserves_existing_positive_result(tmp_path): + codec = _CODECS[0] + output_path = ( + tmp_path / _DATASET / "result" / "build" / f"{_BUILD_STEM}.json" + ) + write_result_files( + [_record(_build_result(codec, build_time=4.5), codec, _BUILD_STEM)], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + write_result_files( + [ + _record( + _build_result( + codec, + build_time=0.0, + metadata={"skipped": True}, + ), + codec, + _BUILD_STEM, + ) + ], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + rows = _load_payload(output_path)["benchmarks"] + assert len(rows) == 1 + assert rows[0]["real_time"] == pytest.approx(4.5) + assert "skipped" not in rows[0] + + +def test_skipped_build_replaces_existing_failed_result(tmp_path): + codec = _CODECS[0] + output_path = ( + tmp_path / _DATASET / "result" / "build" / f"{_BUILD_STEM}.json" + ) + write_result_files( + [ + _record( + _build_result( + codec, + build_time=0.0, + success=False, + error_message="old failure", + ), + codec, + _BUILD_STEM, + ) + ], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + write_result_files( + [ + _record( + _build_result( + codec, + build_time=0.0, + metadata={"skipped": True}, + ), + codec, + _BUILD_STEM, + ) + ], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + rows = _load_payload(output_path)["benchmarks"] + assert len(rows) == 1 + assert rows[0]["success"] is True + assert rows[0]["skipped"] is True + assert "error_message" not in rows[0] + + csv_path = output_path.with_suffix(".csv") + csv_path.write_text("stale\n", encoding="utf-8") + convert_json_to_csv_build(_DATASET, str(tmp_path)) + assert not csv_path.exists() + + +def test_build_upsert_replaces_by_name_without_duplicates(tmp_path): + initial_records = [ + _record(_build_result(codec, build_time=1.0), codec, _BUILD_STEM) + for codec in _CODECS + ] + write_result_files( + initial_records, dataset=_DATASET, dataset_path=str(tmp_path) + ) + + updated_codec = _CODECS[1] + write_result_files( + [ + _record( + _build_result( + updated_codec, + build_time=7.5, + metadata={"writer_path": "gpu-hnsw"}, + ), + updated_codec, + _BUILD_STEM, + ) + ], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + path = tmp_path / _DATASET / "result" / "build" / f"{_BUILD_STEM}.json" + rows = _load_payload(path)["benchmarks"] + assert len(rows) == len(_CODECS) + assert len({row["name"] for row in rows}) == len(_CODECS) + updated = next( + row + for row in rows + if row["name"] == f"{_index_name(updated_codec)}/build" + ) + assert updated["real_time"] == pytest.approx(7.5) + assert updated["writer_path"] == "gpu-hnsw" + + +def test_search_write_replaces_previous_rows(tmp_path): + write_result_files( + [ + _record(_search_result(codec), codec, _SEARCH_STEM) + for codec in _CODECS + ], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + replacement_codec = _CODECS[-1] + write_result_files( + [ + _record( + _search_result( + replacement_codec, + recall=0.99, + queries_per_second=750.0, + ), + replacement_codec, + _SEARCH_STEM, + ) + ], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + path = tmp_path / _DATASET / "result" / "search" / f"{_SEARCH_STEM}.json" + rows = _load_payload(path)["benchmarks"] + assert len(rows) == 1 + assert rows[0]["name"] == f"{_index_name(replacement_codec)}/search" + assert rows[0]["Recall"] == pytest.approx(0.99) + assert rows[0]["items_per_second"] == pytest.approx(750.0) + + +@pytest.mark.parametrize( + "output_filename", + ["", ".", "..", "../escape", "nested/result", "/absolute/result", None], +) +def test_write_result_files_rejects_invalid_output_filename( + tmp_path, output_filename +): + codec = _CODECS[0] + record = _record(_build_result(codec), codec, output_filename) + + with pytest.raises(ValueError, match="Invalid benchmark result filename"): + write_result_files( + [record], dataset=_DATASET, dataset_path=str(tmp_path) + ) + + assert not (tmp_path / _DATASET / "result").exists() + + +@pytest.mark.parametrize( + "dataset", + ["", ".", "..", "../escape", "nested/dataset", "/absolute/dataset", None], +) +def test_write_result_files_rejects_invalid_dataset_name(tmp_path, dataset): + codec = _CODECS[0] + record = _record(_build_result(codec), codec, _BUILD_STEM) + + with pytest.raises(ValueError, match="Invalid benchmark dataset name"): + write_result_files( + [record], dataset=dataset, dataset_path=str(tmp_path) + ) + + +@pytest.mark.parametrize( + "dataset", + ["", ".", "..", "../escape", "nested/dataset", "/absolute/dataset", None], +) +@pytest.mark.parametrize( + "converter", [convert_json_to_csv_build, convert_json_to_csv_search] +) +def test_data_export_rejects_invalid_dataset_name( + tmp_path, dataset, converter +): + with pytest.raises(ValueError, match="Invalid benchmark dataset name"): + converter(dataset, str(tmp_path)) + + +def test_build_merge_rejects_invalid_existing_payload(tmp_path): + build_dir = tmp_path / _DATASET / "result" / "build" + build_dir.mkdir(parents=True) + output_path = build_dir / f"{_BUILD_STEM}.json" + output_path.write_text('{"benchmarks": {}}\n', encoding="utf-8") + codec = _CODECS[0] + + with pytest.raises(ValueError, match="'benchmarks' list"): + write_result_files( + [_record(_build_result(codec), codec, _BUILD_STEM)], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + assert _load_payload(output_path) == {"benchmarks": {}} + + +@pytest.mark.parametrize( + "benchmarks", + [ + [None], + [{"real_time": 1.0}], + [{"name": 42}], + [{"name": "duplicate"}, {"name": "duplicate"}], + ], +) +def test_build_merge_rejects_invalid_existing_rows(tmp_path, benchmarks): + build_dir = tmp_path / _DATASET / "result" / "build" + build_dir.mkdir(parents=True) + output_path = build_dir / f"{_BUILD_STEM}.json" + original_payload = {"benchmarks": benchmarks} + output_path.write_text(json.dumps(original_payload), encoding="utf-8") + codec = _CODECS[0] + + with pytest.raises(ValueError, match="Benchmark row|duplicate name"): + write_result_files( + [_record(_build_result(codec), codec, _BUILD_STEM)], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + assert _load_payload(output_path) == original_payload + + +def test_atomic_result_write_uses_readable_and_preserved_file_modes(tmp_path): + codec = _CODECS[0] + output_path = ( + tmp_path / _DATASET / "result" / "build" / f"{_BUILD_STEM}.json" + ) + record = _record(_build_result(codec), codec, _BUILD_STEM) + + write_result_files([record], dataset=_DATASET, dataset_path=str(tmp_path)) + assert stat.S_IMODE(output_path.stat().st_mode) == 0o644 + + output_path.chmod(0o640) + write_result_files([record], dataset=_DATASET, dataset_path=str(tmp_path)) + assert stat.S_IMODE(output_path.stat().st_mode) == 0o640 + + +def test_result_write_invalidates_only_derived_csvs(tmp_path): + result_root = tmp_path / _DATASET / "result" + build_dir = result_root / "build" + search_dir = result_root / "search" + build_dir.mkdir(parents=True) + search_dir.mkdir(parents=True) + derived_paths = [ + build_dir / f"{_BUILD_STEM}.csv", + search_dir / f"{_SEARCH_STEM},raw.csv", + search_dir / f"{_SEARCH_STEM},throughput.csv", + search_dir / f"{_SEARCH_STEM},latency.csv", + ] + unrelated_path = search_dir / "unrelated.csv" + for path in [*derived_paths, unrelated_path]: + path.write_text("stale\n", encoding="utf-8") + + codec = _CODECS[0] + write_result_files( + [ + _record(_build_result(codec), codec, _BUILD_STEM), + _record(_search_result(codec), codec, _SEARCH_STEM), + ], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + assert not any(path.exists() for path in derived_paths) + assert unrelated_path.read_text(encoding="utf-8") == "stale\n" + + +def test_missing_phase_directories_are_noop(tmp_path): + assert list(read_json_files(_DATASET, str(tmp_path), "build")) == [] + assert list(read_json_files(_DATASET, str(tmp_path), "search")) == [] + + convert_json_to_csv_build(_DATASET, str(tmp_path)) + convert_json_to_csv_search(_DATASET, str(tmp_path)) + + assert not (tmp_path / _DATASET / "result").exists() + + +def test_failed_rows_remain_diagnostic_json_but_are_filtered_from_csv( + tmp_path, +): + good_codec, failed_codec = _CODECS[:2] + write_result_files( + [ + _record(_build_result(good_codec), good_codec, _BUILD_STEM), + _record( + _build_result( + failed_codec, + build_time=0.0, + success=False, + error_message="build failed", + ), + failed_codec, + _BUILD_STEM, + ), + _record(_search_result(good_codec), good_codec, _SEARCH_STEM), + _record( + _search_result( + failed_codec, + search_time_ms=0.0, + latency_seconds=None, + recall=0.0, + queries_per_second=0.0, + success=False, + error_message="search failed", + ), + failed_codec, + _SEARCH_STEM, + ), + ], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + result_root = tmp_path / _DATASET / "result" + build_json = result_root / "build" / f"{_BUILD_STEM}.json" + search_json = result_root / "search" / f"{_SEARCH_STEM}.json" + build_rows = _load_payload(build_json)["benchmarks"] + search_rows = _load_payload(search_json)["benchmarks"] + assert ( + next(row for row in build_rows if not row["success"])["error_message"] + == "build failed" + ) + assert ( + next(row for row in search_rows if not row["success"])["error_message"] + == "search failed" + ) + + convert_json_to_csv_build(_DATASET, str(tmp_path)) + convert_json_to_csv_search(_DATASET, str(tmp_path)) + + build_csv = pd.read_csv(build_json.with_suffix(".csv")) + search_csv = pd.read_csv(search_json.with_name(f"{_SEARCH_STEM},raw.csv")) + assert build_csv["index_name"].tolist() == [_index_name(good_codec)] + assert search_csv["index_name"].tolist() == [_index_name(good_codec)] + assert "success" not in build_csv.columns + assert "error_message" not in build_csv.columns + assert "success" not in search_csv.columns + assert "error_message" not in search_csv.columns + + +def test_all_failed_json_removes_only_matching_derived_csvs(tmp_path): + codec = _CODECS[0] + write_result_files( + [ + _record( + _build_result( + codec, + build_time=0.0, + success=False, + error_message="build failed", + ), + codec, + _BUILD_STEM, + ), + _record( + _search_result( + codec, + search_time_ms=0.0, + latency_seconds=None, + recall=0.0, + queries_per_second=0.0, + success=False, + error_message="search failed", + ), + codec, + _SEARCH_STEM, + ), + ], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + result_root = tmp_path / _DATASET / "result" + derived_paths = [ + result_root / "build" / f"{_BUILD_STEM}.csv", + result_root / "search" / f"{_SEARCH_STEM},raw.csv", + result_root / "search" / f"{_SEARCH_STEM},throughput.csv", + result_root / "search" / f"{_SEARCH_STEM},latency.csv", + ] + unrelated_path = result_root / "search" / "unrelated.csv" + for path in [*derived_paths, unrelated_path]: + path.write_text("stale\n", encoding="utf-8") + + convert_json_to_csv_build(_DATASET, str(tmp_path)) + convert_json_to_csv_search(_DATASET, str(tmp_path)) + + assert not any(path.exists() for path in derived_paths) + assert unrelated_path.read_text(encoding="utf-8") == "stale\n" + + +def test_named_build_join_preserves_search_codec_and_build_metadata(tmp_path): + codec = _CODECS[0] + index_name = _index_name(codec) + build_codec = "build-codec-sentinel" + search_codec = "search-codec-sentinel" + build_result = BuildResult( + index_path=f"/indexes/{index_name}", + build_time_seconds=2.75, + index_size_bytes=8192, + algorithm=_ALGORITHM, + build_params={"codec": build_codec}, + metadata={"writer_path": "gpu-hnsw"}, + ) + search_result = _search_result( + search_codec, + metadata={"mode": "latency"}, + ) + write_result_files( + [ + ResultRecord(build_result, index_name, _BUILD_STEM), + ResultRecord(search_result, index_name, _SEARCH_STEM), + ], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + convert_json_to_csv_build(_DATASET, str(tmp_path)) + convert_json_to_csv_search(_DATASET, str(tmp_path)) + + raw_path = ( + tmp_path / _DATASET / "result" / "search" / f"{_SEARCH_STEM},raw.csv" + ) + row = pd.read_csv(raw_path).iloc[0] + assert row["index_name"] == index_name + assert row["codec"] == search_codec + assert row["build time"] == pytest.approx(2.75) + assert row["writer_path"] == "gpu-hnsw" + assert row["mode"] == "latency" + assert pd.isna(row["build threads"]) + assert pd.isna(row["build cpu_time"]) + + +def test_scoped_search_joins_matching_build_and_preserves_legacy_stem( + tmp_path, +): + codec = _CODECS[0] + full_index_name = _index_name(codec) + subset_scope = "subset128" + subset_index_name = f"{_ALGORITHM}.{subset_scope}.codec{codec}" + subset_build_stem = f"{_BUILD_STEM},{subset_scope}" + subset_search_stem = f"{subset_build_stem},k5,bs2" + + write_result_files( + [ + ResultRecord( + _build_result(codec, build_time=1.25), + full_index_name, + _BUILD_STEM, + ), + ResultRecord( + _search_result(codec), + full_index_name, + _SEARCH_STEM, + ), + ResultRecord( + _build_result(codec, build_time=4.5), + subset_index_name, + subset_build_stem, + ), + ResultRecord( + _search_result(codec), + subset_index_name, + subset_search_stem, + ), + ], + dataset=_DATASET, + dataset_path=str(tmp_path), + ) + + convert_json_to_csv_build(_DATASET, str(tmp_path)) + convert_json_to_csv_search(_DATASET, str(tmp_path)) + + result_root = tmp_path / _DATASET / "result" + full_row = pd.read_csv( + result_root / "search" / f"{_SEARCH_STEM},raw.csv" + ).iloc[0] + subset_row = pd.read_csv( + result_root / "search" / f"{subset_search_stem},raw.csv" + ).iloc[0] + assert full_row["index_name"] == full_index_name + assert full_row["build time"] == pytest.approx(1.25) + assert subset_row["index_name"] == subset_index_name + assert subset_row["build time"] == pytest.approx(4.5) diff --git a/python/cuvs_bench/cuvs_bench/tests/test_registry.py b/python/cuvs_bench/cuvs_bench/tests/test_registry.py index 960ea9cb25..d32b89048e 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_registry.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_registry.py @@ -1,11 +1,13 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ Unit tests for the backend registry system. """ +import json + import pytest import numpy as np @@ -19,6 +21,12 @@ register_backend, get_backend, ) +from cuvs_bench.orchestrator import ( + BenchmarkConfig, + BenchmarkOrchestrator, + DatasetConfig, +) +from cuvs_bench.orchestrator.config_loaders import IndexConfig class DummyBackend(BenchmarkBackend): @@ -197,6 +205,28 @@ def test_build_result_to_json(self): assert json_result["nlist"] == 1024 assert json_result["gpu_time"] == 4.2 + def test_build_result_core_fields_cannot_be_overridden(self): + result = BuildResult( + index_path="/path/to/index", + build_time_seconds=5.5, + index_size_bytes=1024000, + algorithm="test_algo", + build_params={"name": "wrong", "real_time": -1}, + metadata={ + "time_unit": "ms", + "success": False, + "error_message": "wrong", + }, + ) + + json_result = result.to_json() + + assert json_result["name"] == "test_algo/build" + assert json_result["real_time"] == 5.5 + assert json_result["time_unit"] == "s" + assert json_result["success"] is True + assert "error_message" not in json_result + class TestSearchResult: """Tests for SearchResult dataclass.""" @@ -220,6 +250,36 @@ def test_search_result_creation(self): assert result.queries_per_second == 20.0 assert result.success is True + def test_search_result_preserves_legacy_positional_arguments(self): + neighbors = np.array([[1]]) + distances = np.array([[0.1]]) + latency_percentiles = {"p50": 1.0} + metadata = {"source": "legacy"} + + result = SearchResult( + neighbors, + distances, + 2.0, + 500.0, + 0.9, + "test_algo", + [{}], + latency_percentiles, + 0.5, + 0.75, + metadata, + False, + "failed", + ) + + assert result.latency_percentiles is latency_percentiles + assert result.gpu_time_seconds == 0.5 + assert result.cpu_time_seconds == 0.75 + assert result.metadata is metadata + assert result.success is False + assert result.error_message == "failed" + assert result.latency_seconds is None + def test_search_result_to_json(self): """Test conversion to JSON format.""" neighbors = np.array([[1, 2, 3]]) @@ -246,6 +306,284 @@ def test_search_result_to_json(self): assert json_result["p50"] == 50.0 assert json_result["p95"] == 95.0 + def test_search_result_core_fields_cannot_be_overridden(self): + result = SearchResult( + neighbors=np.array([[1]]), + distances=np.array([[0.1]]), + search_time_ms=6.0, + queries_per_second=500.0, + recall=0.95, + algorithm="test_algo", + search_params=[{}], + latency_seconds=0.003, + gpu_time_seconds=0.004, + cpu_time_seconds=0.005, + metadata={ + "name": "wrong", + "real_time": -1, + "Latency": -1, + "GPU": -1, + "cpu_time": -1, + "Recall": -1, + "success": False, + "error_message": "wrong", + }, + ) + + json_result = result.to_json() + + assert json_result["name"] == "test_algo/search" + assert json_result["real_time"] == 3.0 + assert json_result["Latency"] == 0.003 + assert json_result["GPU"] == 0.004 + assert json_result["cpu_time"] == 0.005 + assert json_result["Recall"] == 0.95 + assert json_result["success"] is True + assert "error_message" not in json_result + + def test_search_result_optional_metadata_fields_are_fallbacks(self): + result = SearchResult( + neighbors=np.array([[1]]), + distances=np.array([[0.1]]), + search_time_ms=6.0, + queries_per_second=500.0, + recall=0.95, + algorithm="test_algo", + search_params=[{}], + metadata={ + "Latency": 0.006, + "GPU": 0.004, + "cpu_time": 0.005, + }, + ) + + json_result = result.to_json() + + assert json_result["Latency"] == 0.006 + assert json_result["GPU"] == 0.004 + assert json_result["cpu_time"] == 0.005 + + +def test_backend_without_result_stems_runs_without_generic_persistence( + tmp_path, monkeypatch +): + class DummyConfigLoader: + def load(self, **_kwargs): + return DatasetConfig(name="dummy"), [ + BenchmarkConfig( + indexes=[ + IndexConfig( + name="dummy-index", + algo="dummy", + build_param={}, + search_params=[{}], + file=str(tmp_path / "dummy-index"), + ) + ], + backend_config={"name": "dummy-index"}, + ) + ] + + monkeypatch.setattr( + "cuvs_bench.orchestrator.orchestrator.get_backend_class", + lambda _backend_type: DummyBackend, + ) + monkeypatch.setattr( + "cuvs_bench.orchestrator.orchestrator.get_config_loader", + lambda _backend_type: DummyConfigLoader, + ) + orchestrator = BenchmarkOrchestrator(backend_type="dummy") + + results = orchestrator.run_benchmark( + build=True, + search=False, + dataset="dummy", + dataset_path=str(tmp_path), + ) + + assert len(results) == 1 + assert results[0].success is True + assert not (tmp_path / "dummy" / "result").exists() + + +def test_failed_opt_in_build_clears_only_matching_stale_search_artifacts( + tmp_path, monkeypatch +): + build_stem = "dummy,base" + search_stem = f"{build_stem},k1,bs1" + + class DummyConfigLoader: + def load(self, **_kwargs): + return DatasetConfig(name="dummy"), [ + BenchmarkConfig( + indexes=[ + IndexConfig( + name="dummy-index", + algo="dummy", + build_param={}, + search_params=[{}], + file=str(tmp_path / "dummy-index"), + ) + ], + backend_config={ + "name": "dummy-index", + "output_filename": (build_stem, search_stem), + }, + ) + ] + + class FailedBuildBackend(DummyBackend): + orchestrator_persists_results = True + + def build(self, dataset, indexes, force=False, dry_run=False): + return BuildResult( + index_path=indexes[0].file, + build_time_seconds=0.0, + index_size_bytes=0, + algorithm=self.algo, + build_params={}, + success=False, + error_message="expected build failure", + ) + + def search(self, *args, **kwargs): + raise AssertionError("search must not run after a failed build") + + result_root = tmp_path / "dummy" / "result" + search_dir = result_root / "search" + search_dir.mkdir(parents=True) + stale_paths = [ + search_dir / f"{search_stem}.json", + search_dir / f"{search_stem},raw.csv", + search_dir / f"{search_stem},throughput.csv", + search_dir / f"{search_stem},latency.csv", + ] + unrelated_paths = [ + search_dir / "other,base,k1,bs1.json", + search_dir / "other,base,k1,bs1,raw.csv", + ] + for path in [*stale_paths, *unrelated_paths]: + path.write_text("stale\n", encoding="utf-8") + + monkeypatch.setattr( + "cuvs_bench.orchestrator.orchestrator.get_backend_class", + lambda _backend_type: FailedBuildBackend, + ) + monkeypatch.setattr( + "cuvs_bench.orchestrator.orchestrator.get_config_loader", + lambda _backend_type: DummyConfigLoader, + ) + + results = BenchmarkOrchestrator(backend_type="dummy").run_benchmark( + build=True, + search=True, + count=1, + batch_size=1, + dataset="dummy", + dataset_path=str(tmp_path), + ) + + assert len(results) == 1 + assert results[0].success is False + assert not any(path.exists() for path in stale_paths) + assert all( + path.read_text(encoding="utf-8") == "stale\n" + for path in unrelated_paths + ) + + build_path = result_root / "build" / f"{build_stem}.json" + build_rows = json.loads(build_path.read_text(encoding="utf-8"))[ + "benchmarks" + ] + assert len(build_rows) == 1 + assert build_rows[0]["success"] is False + assert build_rows[0]["error_message"] == "expected build failure" + + +@pytest.mark.parametrize( + ("index_count", "output_filenames", "expected_message"), + [ + ( + 2, + ("dummy,base", "dummy,base,k1,bs1"), + "exactly one index", + ), + ( + 1, + ("dummy", "dummy,k1,bs1"), + "Build result filename stem", + ), + ( + 1, + ("dummy,base", "other,base,k1,bs1"), + "Search result filename stem", + ), + ( + 1, + ("dummy,base",), + "exactly two result filename stems", + ), + ( + 1, + ("../dummy,base", "../dummy,base,k1,bs1"), + "Invalid benchmark result filename", + ), + ], +) +def test_opt_in_sweep_persistence_validates_artifact_identity( + tmp_path, + monkeypatch, + index_count, + output_filenames, + expected_message, +): + indexes = [ + IndexConfig( + name=f"dummy-index-{position}", + algo="dummy", + build_param={}, + search_params=[{}], + file=str(tmp_path / f"dummy-index-{position}"), + ) + for position in range(index_count) + ] + + class DummyConfigLoader: + def load(self, **_kwargs): + return DatasetConfig(name="dummy"), [ + BenchmarkConfig( + indexes=indexes, + backend_config={ + "name": "dummy-index", + "output_filename": output_filenames, + }, + ) + ] + + class PersistedDummyBackend(DummyBackend): + orchestrator_persists_results = True + + monkeypatch.setattr( + "cuvs_bench.orchestrator.orchestrator.get_backend_class", + lambda _backend_type: PersistedDummyBackend, + ) + monkeypatch.setattr( + "cuvs_bench.orchestrator.orchestrator.get_config_loader", + lambda _backend_type: DummyConfigLoader, + ) + + with pytest.raises(ValueError, match=expected_message): + BenchmarkOrchestrator(backend_type="dummy").run_benchmark( + build=True, + search=False, + count=1, + batch_size=1, + dataset="dummy", + dataset_path=str(tmp_path), + ) + + assert not (tmp_path / "dummy" / "result").exists() + class TestBackendRegistry: """Tests for BackendRegistry.""" From 6aa35703a0375fdcee71d145baec00f61807c7a6 Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Fri, 31 Jul 2026 13:29:37 +0000 Subject: [PATCH 03/12] Add PyLucene backend to cuvs-bench Register a local PyLucene backend for cuVS-Lucene HNSW and CAGRA codecs. Add deterministic config selection, lazy JVM and codec resolution, GPU writer validation, safe index lifecycle handling, commit-bound provenance, CAGRA integrity checks, and focused unit coverage. --- .../cuvs_bench/backends/__init__.py | 5 +- python/cuvs_bench/cuvs_bench/backends/base.py | 7 + .../cuvs_bench/backends/pylucene.py | 3029 +++++++++++++++++ .../config/algos/pylucene_cuvs_cagra.yaml | 10 + .../config/algos/pylucene_cuvs_hnsw.yaml | 13 + .../cuvs_bench/orchestrator/__init__.py | 13 +- .../cuvs_bench/orchestrator/config_loaders.py | 4 +- .../cuvs_bench/tests/_pylucene_test_utils.py | 224 ++ .../cuvs_bench/tests/test_pylucene_backend.py | 1038 ++++++ .../tests/test_pylucene_cagra_verifier.py | 815 +++++ .../tests/test_pylucene_cli_config.py | 899 +++++ .../tests/test_pylucene_provenance.py | 270 ++ .../cuvs_bench/tests/test_pylucene_runtime.py | 602 ++++ .../cuvs_bench/cuvs_bench/tests/test_utils.py | 6 +- 14 files changed, 6928 insertions(+), 7 deletions(-) create mode 100644 python/cuvs_bench/cuvs_bench/backends/pylucene.py create mode 100644 python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_cagra.yaml create mode 100644 python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_hnsw.yaml create mode 100644 python/cuvs_bench/cuvs_bench/tests/_pylucene_test_utils.py create mode 100644 python/cuvs_bench/cuvs_bench/tests/test_pylucene_backend.py create mode 100644 python/cuvs_bench/cuvs_bench/tests/test_pylucene_cagra_verifier.py create mode 100644 python/cuvs_bench/cuvs_bench/tests/test_pylucene_cli_config.py create mode 100644 python/cuvs_bench/cuvs_bench/tests/test_pylucene_provenance.py create mode 100644 python/cuvs_bench/cuvs_bench/tests/test_pylucene_runtime.py diff --git a/python/cuvs_bench/cuvs_bench/backends/__init__.py b/python/cuvs_bench/cuvs_bench/backends/__init__.py index a09e6fd1fe..e2e7845953 100644 --- a/python/cuvs_bench/cuvs_bench/backends/__init__.py +++ b/python/cuvs_bench/cuvs_bench/backends/__init__.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ @@ -26,11 +26,13 @@ from .cpp_gbench import CppGoogleBenchmarkBackend from .opensearch import OpenSearchBackend +from .pylucene import PyLuceneBackend # Auto-register built-in backends _registry = get_registry() _registry.register("cpp_gbench", CppGoogleBenchmarkBackend) _registry.register("opensearch", OpenSearchBackend) +_registry.register("pylucene", PyLuceneBackend) __all__ = [ # Base classes and data structures @@ -46,4 +48,5 @@ # Built-in backends "CppGoogleBenchmarkBackend", "OpenSearchBackend", + "PyLuceneBackend", ] diff --git a/python/cuvs_bench/cuvs_bench/backends/base.py b/python/cuvs_bench/cuvs_bench/backends/base.py index a897fa2f0c..e6334c4588 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..8a10fcdfbe --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/backends/pylucene.py @@ -0,0 +1,3029 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""PyLucene backend for cuVS-accelerated Lucene codecs.""" + +from __future__ import annotations + +import hashlib +import importlib +import json +import os +import shutil +import tempfile +import threading +import time +import zipfile +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +from types import MappingProxyType +from typing import ( + Any, + Callable, + Dict, + List, + Mapping, + NamedTuple, + Optional, + Tuple, + Union, +) + +import numpy as np + +from .._bin_format import read_bin_header +from ..orchestrator.config_loaders import ( + BenchmarkConfig, + ConfigLoader, + DatasetConfig, + IndexConfig, +) +from ._utils import dtype_from_filename +from .base import BenchmarkBackend, BuildResult, Dataset, SearchResult + +_ID_FIELD = "id" +_VECTOR_FIELD = "vector" +_MAX_DIMENSIONS = 4096 + +_HNSW_CODECS = frozenset( + { + "Lucene101AcceleratedHNSWCodec", + "Lucene101AcceleratedHNSWBaseLayerCodec", + "Lucene101AcceleratedHNSWMultiLayerCodec", + } +) +_CAGRA_CODEC = "CuVS2510GPUSearchCodec" +_SUPPORTED_CODECS = _HNSW_CODECS | {_CAGRA_CODEC} +_SUPPORTED_BUILD_KEYS = frozenset({"codec"}) +_EXPECTED_WRITER_PATH = { + **{codec: "gpu-hnsw" for codec in _HNSW_CODECS}, + _CAGRA_CODEC: "gpu-cagra", +} +_CAGRA_META_EXTENSION = ".vemc" +_CAGRA_META_CODEC_NAME = "Lucene102CuVSVectorsFormatMeta" +_CAGRA_INDEX_EXTENSION = ".vcag" +_CAGRA_INDEX_CODEC_NAME = "Lucene102CuVSVectorsFormatIndex" +_CAGRA_META_VERSION = 0 +_CAGRA_INDEX_VERSION = 0 +# Avoid caching checksums while coarse filesystem timestamps can still make a +# same-size rewrite look unchanged. +_CAGRA_CACHE_MIN_FILE_AGE_NS = 2_000_000_000 +_FLOAT32_ENCODING_ORDINAL = 1 +_EUCLIDEAN_SIMILARITY_ORDINAL = 0 + +_HNSW_PROVENANCE_FILE = ".cuvs-bench-pylucene-hnsw.json" +_CAGRA_PROVENANCE_FILE = ".cuvs-bench-pylucene-cagra.json" +_PROVENANCE_SCHEMA_VERSION = 1 +_PROVENANCE_KEYS = frozenset( + { + "schema_version", + "codec", + "writer_path", + "vector_count", + "dimensions", + "commit_fingerprints", + } +) +_SHA256_HEX_DIGITS = frozenset("0123456789abcdef") +_LUCENE_CORE_CLASS = "org/apache/lucene/index/IndexWriter.class" + +_JVM_INIT_LOCK = threading.Lock() +_INITIALIZED_CLASSPATH: Optional[str] = None +_INITIALIZED_VMARGS: Optional[Tuple[str, ...]] = None + + +def _attempt_cleanup( + cleanup: Callable[[], None], + description: str, + primary_error: Optional[BaseException], +) -> Optional[BaseException]: + try: + cleanup() + except BaseException as cleanup_error: + if primary_error is None: + return cleanup_error + if isinstance(primary_error, Exception) and not isinstance( + cleanup_error, Exception + ): + cleanup_error.add_note( + f"Raised while attempting to {description}; prior failure: " + f"{type(primary_error).__name__}: {primary_error}" + ) + return cleanup_error + primary_error.add_note( + f"Failed to {description}: " + f"{type(cleanup_error).__name__}: {cleanup_error}" + ) + return primary_error + + +class _CleanupStack: + """Close registered resources without losing the operation's failure.""" + + def __init__(self) -> None: + self._cleanups: List[Tuple[str, Callable[[], None]]] = [] + + def __enter__(self) -> _CleanupStack: + return self + + def add(self, description: str, cleanup: Callable[[], None]) -> None: + self._cleanups.append((description, cleanup)) + + def __exit__( + self, + _error_type: Any, + primary_error: Optional[BaseException], + _traceback: Any, + ) -> bool: + final_error = primary_error + for description, cleanup in reversed(self._cleanups): + final_error = _attempt_cleanup(cleanup, description, final_error) + if final_error is not None and final_error is not primary_error: + raise final_error + return False + + +def _exception_summary(error: Exception) -> str: + details = [f"{type(error).__name__}: {error}"] + if error.__cause__ is not None: + cause = error.__cause__ + details.append(f"caused by {type(cause).__name__}: {cause}") + details.extend(getattr(error, "__notes__", ())) + return "; ".join(details) + + +def _configured_jar( + config: Dict[str, Any], config_key: str, environment_key: str +) -> Path: + value = config.get(config_key) or os.environ.get(environment_key) + if not value: + raise RuntimeError( + f"PyLucene backend requires '{config_key}' or " + f"the {environment_key} environment variable" + ) + + path = Path(os.fspath(value)).expanduser().resolve() + if not path.is_file(): + raise FileNotFoundError(f"{config_key} does not exist: {path}") + return path + + +def _reject_bundled_lucene_classes(cuvs_lucene_jar: Path) -> None: + """Reject fat cuVS-Lucene jars before they can poison the process JVM.""" + if not zipfile.is_zipfile(cuvs_lucene_jar): + return + + with zipfile.ZipFile(cuvs_lucene_jar) as archive: + try: + archive.getinfo(_LUCENE_CORE_CLASS) + except KeyError: + return + + raise RuntimeError( + "cuvs_lucene_jar bundles Lucene classes and is incompatible with " + "PyLucene's process-wide JVM. Use the standard thin cuvs-lucene JAR, " + "not a '-jar-with-dependencies' artifact." + ) + + +def _load_pylucene() -> Any: + try: + return importlib.import_module("lucene") + except ImportError as exc: + raise ImportError( + "PyLucene's `lucene` module is required for the pylucene " + "cuvs-bench backend. PyLucene must be built and installed " + "separately; see the cuVS Bench installation documentation." + ) from exc + + +def _pylucene_classpath(config: Dict[str, Any], lucene: Any) -> str: + cuvs_java_jar = _configured_jar( + config, "cuvs_java_jar", "CUVS_LUCENE_CUVS_JAVA_JAR" + ) + cuvs_lucene_jar = _configured_jar( + config, "cuvs_lucene_jar", "CUVS_LUCENE_JAR" + ) + _reject_bundled_lucene_classes(cuvs_lucene_jar) + return os.pathsep.join( + [str(cuvs_java_jar), str(cuvs_lucene_jar), lucene.CLASSPATH] + ) + + +def _pylucene_vmargs(config: Dict[str, Any]) -> List[str]: + java_library_path = ( + config.get("java_library_path") + or os.environ.get("JAVA_LIBRARY_PATH") + or os.environ.get("LD_LIBRARY_PATH") + ) + vmargs = [ + "--enable-native-access=ALL-UNNAMED", + "--add-modules=jdk.incubator.vector", + ] + if java_library_path: + vmargs.append(f"-Djava.library.path={java_library_path}") + + extra_vmargs = config.get("jvm_args", []) + if isinstance(extra_vmargs, (str, bytes)) or not isinstance( + extra_vmargs, (list, tuple) + ): + raise TypeError("jvm_args must be a list or tuple of strings") + if not all(isinstance(arg, str) for arg in extra_vmargs): + raise TypeError("every jvm_args entry must be a string") + vmargs.extend(extra_vmargs) + return vmargs + + +def _attach_pylucene_jvm( + lucene: Any, classpath: str, vmargs: List[str] +) -> None: + with _JVM_INIT_LOCK: + vm_environment = lucene.getVMEnv() + if vm_environment is None: + vm_environment = _start_pylucene_jvm(lucene, classpath, vmargs) + else: + _validate_pylucene_jvm_configuration(classpath, vmargs) + + vm_environment.attachCurrentThread() + + +def _start_pylucene_jvm(lucene: Any, classpath: str, vmargs: List[str]) -> Any: + global _INITIALIZED_CLASSPATH, _INITIALIZED_VMARGS + vm_environment = lucene.initVM(classpath=classpath, vmargs=vmargs) + if vm_environment is None: + vm_environment = lucene.getVMEnv() + if vm_environment is None: + raise RuntimeError("PyLucene did not return a JVM environment") + + _INITIALIZED_CLASSPATH = classpath + _INITIALIZED_VMARGS = tuple(vmargs) + return vm_environment + + +def _validate_pylucene_jvm_configuration( + classpath: str, vmargs: List[str] +) -> None: + if _INITIALIZED_CLASSPATH is None or _INITIALIZED_VMARGS is None: + raise RuntimeError( + "PyLucene's process-wide JVM was initialized before the " + "pylucene backend, so the required cuVS classpath and JVM " + "arguments cannot be verified. Start a new Python process and " + "let cuVS Bench initialize PyLucene." + ) + if _INITIALIZED_CLASSPATH != classpath: + raise RuntimeError( + "PyLucene's process-wide JVM is already initialized with " + "different cuVS Java or cuVS-Lucene jars" + ) + if _INITIALIZED_VMARGS != tuple(vmargs): + raise RuntimeError( + "PyLucene's process-wide JVM is already initialized with " + "different JVM arguments or native-library paths" + ) + + +def _initialize_pylucene(config: Dict[str, Any]) -> Any: + """Initialize PyLucene once and attach the current Python thread.""" + lucene = _load_pylucene() + classpath = _pylucene_classpath(config, lucene) + vmargs = _pylucene_vmargs(config) + _attach_pylucene_jvm(lucene, classpath, vmargs) + return lucene + + +def _parse_writer_telemetry(description: str) -> Dict[str, str]: + """Parse cuVS-Lucene vector-format diagnostics.""" + _, separator, payload = description.partition("(") + if not separator or not payload.endswith(")"): + raise RuntimeError( + f"Malformed cuVS-Lucene writer telemetry: {description!r}" + ) + + telemetry: Dict[str, str] = {} + for item in payload[:-1].split(";"): + key, item_separator, value = item.partition("=") + if not item_separator or not key: + raise RuntimeError( + "Malformed cuVS-Lucene writer telemetry item " + f"{item!r} in {description!r}" + ) + telemetry[key] = value + return telemetry + + +def _score_to_squared_euclidean(score: float) -> float: + """Convert Lucene's Euclidean score to squared Euclidean distance.""" + if score <= 0.0: + return float("inf") + return max(0.0, (1.0 / score) - 1.0) + + +def _validate_float32_matrix(vectors: np.ndarray, name: str) -> np.ndarray: + array = np.asarray(vectors) + if array.ndim != 2: + raise ValueError(f"{name} must be a two-dimensional array") + if array.shape[0] == 0 or array.shape[1] == 0: + raise ValueError(f"{name} must contain at least one vector") + if array.dtype != np.float32: + raise TypeError(f"{name} must use float32 values, got {array.dtype}") + if array.shape[1] > _MAX_DIMENSIONS: + raise ValueError( + f"{name} dimensions must not exceed {_MAX_DIMENSIONS}, " + f"got {array.shape[1]}" + ) + if not np.isfinite(array).all(): + raise ValueError(f"{name} must contain only finite values") + return np.ascontiguousarray(array) + + +def _validate_metric(dataset: Dataset) -> None: + if dataset.distance_metric.lower() not in {"euclidean", "l2"}: + raise ValueError( + "PyLucene cuVS codecs currently support only Euclidean/L2 " + f"datasets, got {dataset.distance_metric!r}" + ) + + +def _validate_codec(codec_name: Any) -> str: + if not isinstance(codec_name, str) or codec_name not in _SUPPORTED_CODECS: + available = ", ".join(sorted(_SUPPORTED_CODECS)) + raise ValueError( + f"Unsupported PyLucene codec {codec_name!r}. " + f"Supported codecs: {available}" + ) + return codec_name + + +def _configured_codec( + build_params: Dict[str, Any], backend_config: Dict[str, Any] +) -> str: + if "codec" in build_params: + codec_name = build_params["codec"] + else: + codec_name = backend_config.get("codec") + return _validate_codec(codec_name) + + +def _validate_build_params(build_params: Any) -> Dict[str, Any]: + if not isinstance(build_params, dict): + raise TypeError("PyLucene build parameters must be a mapping") + unsupported = set(build_params) - _SUPPORTED_BUILD_KEYS + if unsupported: + names = ", ".join(sorted(str(name) for name in unsupported)) + raise ValueError(f"Unsupported PyLucene build parameter(s): {names}") + return build_params + + +def _validate_search_params(search_params: Any) -> List[Dict[str, Any]]: + if ( + not isinstance(search_params, list) + or len(search_params) != 1 + or search_params[0] != {} + ): + raise ValueError( + "PyLucene's public Lucene query API does not expose " + "cuVS-specific search parameters" + ) + return search_params + + +def _safe_remove_index(index_path: Path) -> None: + resolved = index_path.resolve() + forbidden = { + Path(resolved.anchor), + Path.cwd().resolve(), + Path.home().resolve(), + } + if resolved in forbidden: + raise ValueError(f"Refusing to remove unsafe index path: {resolved}") + if index_path.is_symlink() or not index_path.is_dir(): + raise ValueError( + f"PyLucene index path must be a directory: {index_path}" + ) + shutil.rmtree(index_path) + + +def _index_size(index_path: Path) -> int: + return sum( + path.stat().st_size for path in index_path.rglob("*") if path.is_file() + ) + + +@dataclass(frozen=True) +class _FileSignature: + resolved_path: str + device: int + inode: int + size: int + modified_at_ns: int + changed_at_ns: int + + +def _file_signature(path: Path) -> _FileSignature: + file_stat = path.stat() + return _FileSignature( + resolved_path=str(path.resolve()), + device=file_stat.st_dev, + inode=file_stat.st_ino, + size=file_stat.st_size, + modified_at_ns=file_stat.st_mtime_ns, + changed_at_ns=file_stat.st_ctime_ns, + ) + + +def _has_lucene_segments(index_path: Path) -> bool: + return any( + path.is_file() and path.name.startswith("segments_") + for path in index_path.iterdir() + ) + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as file: + for block in iter(lambda: file.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def _commit_fingerprints(index_path: Path) -> List[Dict[str, str]]: + commit_files = sorted( + path + for path in index_path.iterdir() + if path.name.startswith("segments_") + ) + if not commit_files: + raise RuntimeError( + "Lucene index has no segments_* commit file to fingerprint" + ) + + fingerprints = [] + for path in commit_files: + if path.is_symlink() or not path.is_file(): + raise RuntimeError( + "Lucene commit fingerprint target must be a regular file: " + f"{path}" + ) + fingerprints.append({"name": path.name, "sha256": _sha256_file(path)}) + return fingerprints + + +@dataclass(frozen=True) +class _IndexProvenanceVerification: + codec: str + writer_path: str + vector_count: int + dimensions: int + commit_file_count: int + + def to_metadata(self) -> Dict[str, Union[str, int]]: + return { + "status": f"{self.writer_path}-provenance", + "schema_version": _PROVENANCE_SCHEMA_VERSION, + "codec": self.codec, + "writer_path": self.writer_path, + "vector_count": self.vector_count, + "dimensions": self.dimensions, + "commit_file_count": self.commit_file_count, + } + + +@dataclass(frozen=True) +class _ProvenanceExpectation: + codec: str + manifest_name: str + label: str + vector_count: Optional[int] + dimensions: Optional[int] + + +class _IndexProvenanceError(RuntimeError): + """Raised when an index cannot be proven to be backend-built.""" + + +def _write_index_provenance( + index_path: Path, + codec: str, + vector_count: int, + dimensions: int, + manifest_name: str, +) -> None: + payload = { + "schema_version": _PROVENANCE_SCHEMA_VERSION, + "codec": codec, + "writer_path": _EXPECTED_WRITER_PATH[codec], + "vector_count": int(vector_count), + "dimensions": int(dimensions), + "commit_fingerprints": _commit_fingerprints(index_path), + } + manifest_path = index_path / manifest_name + with _CleanupStack() as cleanups: + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=index_path, + prefix=f"{manifest_name}.", + suffix=".tmp", + delete=False, + ) as file: + temporary_path = Path(file.name) + cleanups.add( + "remove temporary provenance file", + lambda: temporary_path.unlink(missing_ok=True), + ) + json.dump( + payload, + file, + allow_nan=False, + indent=2, + sort_keys=True, + ) + file.write("\n") + file.flush() + os.fsync(file.fileno()) + temporary_path.chmod(0o644) + temporary_path.replace(manifest_path) + + +def _write_hnsw_provenance( + index_path: Path, + codec: str, + vector_count: int, + dimensions: int, +) -> None: + _write_index_provenance( + index_path, + codec, + vector_count, + dimensions, + _HNSW_PROVENANCE_FILE, + ) + + +def _write_cagra_provenance( + index_path: Path, + vector_count: int, + dimensions: int, +) -> None: + _write_index_provenance( + index_path, + _CAGRA_CODEC, + vector_count, + dimensions, + _CAGRA_PROVENANCE_FILE, + ) + + +def _require_positive_int( + value: Any, provenance_label: str, field_name: str +) -> int: + if type(value) is not int or value < 1: + raise _IndexProvenanceError( + f"{provenance_label} field {field_name!r} must be a " + "positive integer" + ) + return value + + +def _read_index_provenance( + manifest_path: Path, provenance_label: str +) -> Dict[str, Any]: + try: + payload = json.loads(manifest_path.read_text(encoding="utf-8")) + except FileNotFoundError as exc: + raise _IndexProvenanceError( + f"{provenance_label} manifest is missing: {manifest_path}" + ) from exc + except (OSError, UnicodeError, json.JSONDecodeError) as exc: + raise _IndexProvenanceError( + f"{provenance_label} manifest cannot be read: " + f"{manifest_path}: {exc}" + ) from exc + + if not isinstance(payload, dict) or set(payload) != _PROVENANCE_KEYS: + raise _IndexProvenanceError( + f"{provenance_label} manifest has an unsupported schema" + ) + return payload + + +def _validate_provenance_identity( + payload: Dict[str, Any], + expected_codec: str, + provenance_label: str, +) -> str: + if ( + type(payload["schema_version"]) is not int + or payload["schema_version"] != _PROVENANCE_SCHEMA_VERSION + ): + raise _IndexProvenanceError( + f"{provenance_label} manifest has an unsupported schema version" + ) + if payload["codec"] != expected_codec: + raise _IndexProvenanceError( + f"{provenance_label} codec does not match the requested codec: " + f"{payload['codec']!r} != {expected_codec!r}" + ) + + writer_path = payload["writer_path"] + if writer_path != _EXPECTED_WRITER_PATH[expected_codec]: + raise _IndexProvenanceError( + f"{provenance_label} does not record the required GPU writer path" + ) + return writer_path + + +def _validate_provenance_shape( + payload: Dict[str, Any], + provenance_label: str, + expected_vector_count: Optional[int] = None, + expected_dimensions: Optional[int] = None, +) -> Tuple[int, int]: + vector_count = _require_positive_int( + payload["vector_count"], provenance_label, "vector_count" + ) + if vector_count < 2: + raise _IndexProvenanceError( + f"{provenance_label} vector count must be at least two" + ) + dimensions = _require_positive_int( + payload["dimensions"], provenance_label, "dimensions" + ) + if dimensions > _MAX_DIMENSIONS: + raise _IndexProvenanceError( + f"{provenance_label} dimensions exceed the supported maximum: " + f"{dimensions} > {_MAX_DIMENSIONS}" + ) + if ( + expected_vector_count is not None + and vector_count != expected_vector_count + ): + raise _IndexProvenanceError( + f"{provenance_label} vector count does not match the dataset: " + f"{vector_count} != {expected_vector_count}" + ) + if expected_dimensions is not None and dimensions != expected_dimensions: + raise _IndexProvenanceError( + f"{provenance_label} dimensions do not match the dataset: " + f"{dimensions} != {expected_dimensions}" + ) + return vector_count, dimensions + + +def _validate_commit_fingerprints( + stored_fingerprints: Any, provenance_label: str +) -> List[Dict[str, str]]: + if not isinstance(stored_fingerprints, list) or not stored_fingerprints: + raise _IndexProvenanceError( + f"{provenance_label} has no Lucene commit fingerprints" + ) + names = [] + for fingerprint in stored_fingerprints: + name, _ = _validate_commit_fingerprint(fingerprint, provenance_label) + names.append(name) + if names != sorted(set(names)): + raise _IndexProvenanceError( + f"{provenance_label} commit fingerprints must be unique and sorted" + ) + return stored_fingerprints + + +def _validate_commit_fingerprint( + fingerprint: Any, provenance_label: str +) -> Tuple[str, str]: + if not isinstance(fingerprint, dict) or set(fingerprint) != { + "name", + "sha256", + }: + raise _IndexProvenanceError( + f"{provenance_label} has a malformed Lucene commit fingerprint" + ) + + name = _validate_commit_filename(fingerprint["name"], provenance_label) + digest = _validate_sha256_digest(fingerprint["sha256"], provenance_label) + return name, digest + + +def _validate_commit_filename(value: Any, provenance_label: str) -> str: + if not isinstance(value, str): + raise _IndexProvenanceError( + f"{provenance_label} commit filename must be a string" + ) + if not value.startswith("segments_"): + raise _IndexProvenanceError( + f"{provenance_label} commit filename must start with 'segments_'" + ) + if Path(value).name != value: + raise _IndexProvenanceError( + f"{provenance_label} commit filename must not contain a path" + ) + return value + + +def _validate_sha256_digest(value: Any, provenance_label: str) -> str: + if not isinstance(value, str): + raise _IndexProvenanceError( + f"{provenance_label} commit SHA-256 must be a string" + ) + if len(value) != 64: + raise _IndexProvenanceError( + f"{provenance_label} commit SHA-256 must contain 64 characters" + ) + if not set(value).issubset(_SHA256_HEX_DIGITS): + raise _IndexProvenanceError( + f"{provenance_label} commit SHA-256 must be lowercase hexadecimal" + ) + return value + + +def _verify_index_provenance( + index_path: Path, + expectation: _ProvenanceExpectation, +) -> _IndexProvenanceVerification: + payload = _read_index_provenance( + index_path / expectation.manifest_name, expectation.label + ) + writer_path = _validate_provenance_identity( + payload, expectation.codec, expectation.label + ) + vector_count, dimensions = _validate_provenance_shape( + payload, + expectation.label, + expectation.vector_count, + expectation.dimensions, + ) + stored_fingerprints = _validate_commit_fingerprints( + payload["commit_fingerprints"], expectation.label + ) + try: + current_fingerprints = _commit_fingerprints(index_path) + except (OSError, RuntimeError) as exc: + raise _IndexProvenanceError( + f"{expectation.label} cannot fingerprint the Lucene commit: {exc}" + ) from exc + if stored_fingerprints != current_fingerprints: + raise _IndexProvenanceError( + f"{expectation.label} does not match the current Lucene commit" + ) + + return _IndexProvenanceVerification( + codec=expectation.codec, + writer_path=writer_path, + vector_count=vector_count, + dimensions=dimensions, + commit_file_count=len(current_fingerprints), + ) + + +def _verify_hnsw_provenance( + index_path: Path, + expected_codec: str, + *, + expected_vector_count: Optional[int] = None, + expected_dimensions: Optional[int] = None, +) -> _IndexProvenanceVerification: + return _verify_index_provenance( + index_path, + _ProvenanceExpectation( + codec=expected_codec, + manifest_name=_HNSW_PROVENANCE_FILE, + label="HNSW provenance", + vector_count=expected_vector_count, + dimensions=expected_dimensions, + ), + ) + + +def _verify_cagra_provenance( + index_path: Path, + *, + expected_vector_count: Optional[int] = None, + expected_dimensions: Optional[int] = None, +) -> _IndexProvenanceVerification: + return _verify_index_provenance( + index_path, + _ProvenanceExpectation( + codec=_CAGRA_CODEC, + manifest_name=_CAGRA_PROVENANCE_FILE, + label="CAGRA provenance", + vector_count=expected_vector_count, + dimensions=expected_dimensions, + ), + ) + + +def _validate_subset_size(value: Any) -> Optional[int]: + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise ValueError( + f"subset_size must be a positive integer, got {value!r}" + ) + return value + + +def _expected_training_shape(dataset: Dataset) -> Optional[Tuple[int, int]]: + vectors = dataset.loaded_training_vectors + if vectors is not None: + vectors = np.asarray(vectors) + if vectors.ndim != 2: + raise ValueError( + "training_vectors must be a two-dimensional array" + ) + if vectors.dtype != np.float32: + raise TypeError( + "training_vectors must use float32 values, " + f"got {vectors.dtype}" + ) + rows, dimensions = vectors.shape + elif dataset.base_file: + dtype = np.dtype(dtype_from_filename(dataset.base_file)) + if dtype != np.float32: + raise TypeError( + f"training_vectors must use float32 values, got {dtype}" + ) + rows, dimensions, _ = read_bin_header( + dataset.base_file, itemsize=dtype.itemsize + ) + subset_size = _validate_subset_size( + dataset.metadata.get("subset_size") + ) + if subset_size is not None: + rows = min(rows, subset_size) + else: + return None + + if rows < 1 or dimensions < 1: + raise ValueError("training_vectors must contain at least one vector") + if dimensions > _MAX_DIMENSIONS: + raise ValueError( + f"training_vectors dimensions must not exceed {_MAX_DIMENSIONS}, " + f"got {dimensions}" + ) + return int(rows), int(dimensions) + + +@dataclass(frozen=True) +class _SearchHit: + document_id: int + score: float + + +@dataclass(frozen=True) +class _RuntimeSearchResult: + hits: List[List[_SearchHit]] + batch_latencies_ms: List[float] + index_dimensions: int + document_count: int + + +@dataclass(frozen=True) +class _ProcessedSearchResult: + neighbors: np.ndarray + distances: np.ndarray + search_time_ms: float + queries_per_second: float + latency_seconds: float + latency_percentiles: Dict[str, float] + num_batches: int + + +@dataclass(frozen=True) +class _SearchInputs: + query_vectors: np.ndarray + expected_training_shape: Optional[Tuple[int, int]] + + +@dataclass(frozen=True) +class _SearchPlan: + index_path: Path + codec_name: str + search_params: List[Dict[str, Any]] + k: int + batch_size: int + mode: str + + +@dataclass(frozen=True) +class _ResolvedCodec: + codec_name: str + java_codec: Any + telemetry: Mapping[str, str] + writer_path: str + + def __post_init__(self) -> None: + object.__setattr__( + self, "telemetry", MappingProxyType(dict(self.telemetry)) + ) + + +class _ExistingIndexAction(Enum): + BUILD = "build" + REUSE = "reuse" + REJECT = "reject" + + +@dataclass(frozen=True) +class _ExistingIndexDecision: + action: _ExistingIndexAction + result: Optional[BuildResult] = None + + @classmethod + def build(cls) -> _ExistingIndexDecision: + return cls(action=_ExistingIndexAction.BUILD) + + @classmethod + def reuse(cls, result: BuildResult) -> _ExistingIndexDecision: + return cls(action=_ExistingIndexAction.REUSE, result=result) + + @classmethod + def reject(cls, result: BuildResult) -> _ExistingIndexDecision: + return cls(action=_ExistingIndexAction.REJECT, result=result) + + def completed_result(self) -> BuildResult: + if self.result is None: + raise RuntimeError( + "Existing-index decision is missing its build result" + ) + return self.result + + +def _validate_runtime_search_result( + runtime_result: _RuntimeSearchResult, + provenance: _IndexProvenanceVerification, + *, + query_count: int, +) -> None: + if runtime_result.document_count != provenance.vector_count: + raise RuntimeError( + "Lucene document count does not match index provenance: " + f"{runtime_result.document_count} != {provenance.vector_count}" + ) + if runtime_result.index_dimensions != provenance.dimensions: + raise RuntimeError( + "Lucene index dimensions do not match index provenance: " + f"{runtime_result.index_dimensions} != {provenance.dimensions}" + ) + if len(runtime_result.hits) != query_count: + raise RuntimeError( + "PyLucene returned an unexpected number of query results: " + f"{len(runtime_result.hits)} != {query_count}" + ) + + +def _validate_search_hit( + hit: _SearchHit, + *, + query_id: int, + document_count: int, + seen_document_ids: set[int], +) -> None: + if not 0 <= hit.document_id < document_count: + raise RuntimeError( + "PyLucene returned an out-of-range stored ID for " + f"query {query_id}: {hit.document_id}" + ) + if hit.document_id in seen_document_ids: + raise RuntimeError( + "PyLucene returned a duplicate stored ID for " + f"query {query_id}: {hit.document_id}" + ) + if not np.isfinite(hit.score): + raise RuntimeError( + "PyLucene returned a non-finite score for " + f"query {query_id}: {hit.score}" + ) + if not 0.0 <= hit.score <= 1.0: + raise RuntimeError( + "PyLucene returned a score outside the Euclidean range " + f"[0, 1] for query {query_id}: {hit.score}" + ) + seen_document_ids.add(hit.document_id) + + +def _convert_search_hits( + runtime_result: _RuntimeSearchResult, *, query_count: int, k: int +) -> Tuple[np.ndarray, np.ndarray]: + neighbors = np.full((query_count, k), -1, dtype=np.int64) + distances = np.full((query_count, k), np.inf, dtype=np.float32) + for query_id, hits in enumerate(runtime_result.hits): + query_neighbors, query_distances = _convert_query_hits( + hits, + query_id=query_id, + document_count=runtime_result.document_count, + k=k, + ) + hit_count = len(query_neighbors) + neighbors[query_id, :hit_count] = query_neighbors + distances[query_id, :hit_count] = query_distances + return neighbors, distances + + +def _convert_query_hits( + hits: List[_SearchHit], + *, + query_id: int, + document_count: int, + k: int, +) -> Tuple[List[int], List[float]]: + if len(hits) > min(k, document_count): + raise RuntimeError( + f"PyLucene returned too many hits for query {query_id}: " + f"{len(hits)}" + ) + + neighbors = [] + distances = [] + seen_document_ids = set() + for hit in hits: + _validate_search_hit( + hit, + query_id=query_id, + document_count=document_count, + seen_document_ids=seen_document_ids, + ) + neighbors.append(hit.document_id) + distances.append(_score_to_squared_euclidean(hit.score)) + return neighbors, distances + + +def _validate_batch_latencies( + runtime_result: _RuntimeSearchResult, + *, + query_count: int, + batch_size: int, +) -> Tuple[np.ndarray, int]: + latencies = np.asarray(runtime_result.batch_latencies_ms, dtype=np.float64) + num_batches = (query_count + batch_size - 1) // batch_size + if latencies.shape != (num_batches,): + raise RuntimeError( + "PyLucene returned an unexpected number of batch latencies: " + f"{latencies.size} != {num_batches}" + ) + if not np.isfinite(latencies).all() or np.any(latencies < 0.0): + raise RuntimeError( + "PyLucene returned invalid batch latency measurements" + ) + return latencies, num_batches + + +def _process_search_result( + runtime_result: _RuntimeSearchResult, + provenance: _IndexProvenanceVerification, + *, + query_count: int, + k: int, + batch_size: int, +) -> _ProcessedSearchResult: + """Validate the Java boundary result and convert it to benchmark arrays.""" + _validate_runtime_search_result( + runtime_result, provenance, query_count=query_count + ) + neighbors, distances = _convert_search_hits( + runtime_result, query_count=query_count, k=k + ) + latencies, num_batches = _validate_batch_latencies( + runtime_result, query_count=query_count, batch_size=batch_size + ) + + search_time_ms = float(latencies.sum()) + return _ProcessedSearchResult( + neighbors=neighbors, + distances=distances, + search_time_ms=search_time_ms, + queries_per_second=( + query_count / (search_time_ms / 1000.0) + if search_time_ms > 0.0 + else 0.0 + ), + latency_seconds=float(latencies.mean()) / 1000.0, + latency_percentiles={ + "p50": float(np.percentile(latencies, 50)), + "p95": float(np.percentile(latencies, 95)), + "p99": float(np.percentile(latencies, 99)), + }, + num_batches=num_batches, + ) + + +@dataclass(frozen=True) +class _CagraIndexVerification: + segment_count: int + field_count: int + vector_count: int + dimensions: int + + def to_metadata(self) -> Dict[str, Union[str, int]]: + return { + "status": "cagra-only", + "segment_count": self.segment_count, + "field_count": self.field_count, + "vector_count": self.vector_count, + "dimensions": self.dimensions, + } + + +class _CagraVerificationError(RuntimeError): + """Raised when persisted metadata cannot prove a CAGRA-only index.""" + + +@dataclass(frozen=True) +class _CagraFieldMetadata: + field_number: int + dimensions: int + vector_count: int + cagra_offset: int + cagra_length: int + + +@dataclass(frozen=True) +class _RawCagraFieldMetadata: + encoding: int + similarity: int + dimensions: int + vector_count: int + cagra_offset: int + cagra_length: int + brute_force_offset: int + brute_force_length: int + + +@dataclass(frozen=True) +class _CagraFieldSource: + metadata: _CagraFieldMetadata + metadata_file: str + + +@dataclass(frozen=True) +class _CagraDataFileContext: + index_path: Path + directory: Any + segment_info: Any + suffix: str + metadata_file: str + + @property + def data_file(self) -> str: + stem = self.metadata_file[: -len(_CAGRA_META_EXTENSION)] + return stem + _CAGRA_INDEX_EXTENSION + + @property + def data_path(self) -> Path: + return self.index_path / self.data_file + + +@dataclass(frozen=True) +class _CagraSegmentVerification: + field_count: int + vector_count: int + dimensions: frozenset[int] + + +class _CagraIndexVerifier: + """Own persisted CAGRA inspection and its verified-file cache.""" + + def __init__( + self, + *, + attach_current_thread: Callable[[], None], + paths: Any, + codec_util: Any, + field_info: Any, + segment_commit_info: Any, + segment_infos: Any, + vector_encoding: Any, + vector_similarity_function: Any, + fs_directory: Any, + io_context: Any, + ) -> None: + self._attach_current_thread = attach_current_thread + self.Paths = paths + self.CodecUtil = codec_util + self.FieldInfo = field_info + self.SegmentCommitInfo = segment_commit_info + self.SegmentInfos = segment_infos + self.VectorEncoding = vector_encoding + self.VectorSimilarityFunction = vector_similarity_function + self.FSDirectory = fs_directory + self.IOContext = io_context + self._verified_data_files: set[_FileSignature] = set() + + @staticmethod + def _segment_suffix(segment_name: str, metadata_file: str) -> str: + stem = metadata_file[: -len(_CAGRA_META_EXTENSION)] + if stem == segment_name: + return "" + + prefix = f"{segment_name}_" + if not stem.startswith(prefix) or len(stem) == len(prefix): + raise _CagraVerificationError( + "CAGRA-only verification found a metadata filename that " + f"does not match segment {segment_name!r}: {metadata_file!r}" + ) + return stem[len(prefix) :] + + @classmethod + def _read_cagra_field( + cls, + metadata_input: Any, + metadata_file: str, + field_number: int, + ) -> Optional[_CagraFieldMetadata]: + raw_field = cls._decode_cagra_field(metadata_input) + return cls._validate_cagra_field_metadata( + raw_field, metadata_file, field_number + ) + + @staticmethod + def _decode_cagra_field(metadata_input: Any) -> _RawCagraFieldMetadata: + return _RawCagraFieldMetadata( + encoding=int(metadata_input.readInt()), + similarity=int(metadata_input.readInt()), + dimensions=int(metadata_input.readInt()), + vector_count=int(metadata_input.readInt()), + cagra_offset=int(metadata_input.readVLong()), + cagra_length=int(metadata_input.readVLong()), + brute_force_offset=int(metadata_input.readVLong()), + brute_force_length=int(metadata_input.readVLong()), + ) + + @classmethod + def _validate_cagra_field_metadata( + cls, + field: _RawCagraFieldMetadata, + metadata_file: str, + field_number: int, + ) -> Optional[_CagraFieldMetadata]: + if field.encoding != _FLOAT32_ENCODING_ORDINAL: + raise _CagraVerificationError( + "CAGRA-only verification found unsupported vector " + f"encoding ordinal {field.encoding} in {metadata_file!r}" + ) + if field.similarity != _EUCLIDEAN_SIMILARITY_ORDINAL: + raise _CagraVerificationError( + "CAGRA-only verification found unsupported similarity " + f"ordinal {field.similarity} in {metadata_file!r}" + ) + if not 1 <= field.dimensions <= _MAX_DIMENSIONS: + raise _CagraVerificationError( + "CAGRA-only verification found invalid field metadata " + f"in {metadata_file!r}" + ) + if field.vector_count < 0: + raise _CagraVerificationError( + "CAGRA-only verification found invalid field metadata " + f"in {metadata_file!r}" + ) + if field.vector_count == 0: + if cls._empty_cagra_field_has_data(field): + raise _CagraVerificationError( + "CAGRA-only verification found index data for an " + f"empty field in {metadata_file!r}" + ) + return None + if field.brute_force_length != 0: + raise _CagraVerificationError( + "CAGRA-only verification found a persisted brute-force " + f"index for field {field_number} in {metadata_file!r}" + ) + if field.cagra_length <= 0: + raise _CagraVerificationError( + "CAGRA-only verification found no persisted CAGRA index " + f"for field {field_number} in {metadata_file!r}" + ) + return _CagraFieldMetadata( + field_number=field_number, + dimensions=field.dimensions, + vector_count=field.vector_count, + cagra_offset=field.cagra_offset, + cagra_length=field.cagra_length, + ) + + @staticmethod + def _empty_cagra_field_has_data(field: _RawCagraFieldMetadata) -> bool: + return any( + ( + field.cagra_offset, + field.cagra_length, + field.brute_force_offset, + field.brute_force_length, + ) + ) + + @classmethod + def _read_cagra_fields( + cls, metadata_input: Any, metadata_file: str + ) -> List[_CagraFieldMetadata]: + field_numbers = set() + fields = [] + while True: + field_number = int(metadata_input.readInt()) + if field_number == -1: + return fields + if field_number < 0 or field_number in field_numbers: + raise _CagraVerificationError( + "CAGRA-only verification found an invalid field number " + f"{field_number} in {metadata_file!r}" + ) + + field_numbers.add(field_number) + field = cls._read_cagra_field( + metadata_input, metadata_file, field_number + ) + if field is not None: + fields.append(field) + + def _read_segment_field_infos( + self, directory: Any, segment_info: Any + ) -> Any: + if segment_info.hasFieldUpdates(): + raise _CagraVerificationError( + "CAGRA-only verification does not accept segments with " + "field-info updates" + ) + try: + return ( + segment_info.info.getCodec() + .fieldInfosFormat() + .read( + directory, + segment_info.info, + "", + self.IOContext.READONCE, + ) + ) + except Exception as exc: + raise _CagraVerificationError( + "CAGRA-only verification cannot read Lucene field " + f"metadata for segment {segment_info.info.name!r}: {exc}" + ) from exc + + @classmethod + def _verify_cagra_payload_coverage( + cls, + fields: List[_CagraFieldMetadata], + payload_start: int, + payload_end: int, + data_file: str, + ) -> None: + intervals = [] + for field in fields: + intervals.append( + ( + field.cagra_offset, + field.cagra_offset + field.cagra_length, + ) + ) + intervals.sort() + + if not intervals: + if payload_start != payload_end: + raise _CagraVerificationError( + "CAGRA data file contains unreferenced payload: " + f"{data_file!r}" + ) + return + + coverage_error = ( + "CAGRA metadata ranges do not exactly cover data " + f"file {data_file!r}" + ) + expected_offset = payload_start + for interval_start, interval_end in intervals: + if interval_start != expected_offset: + raise _CagraVerificationError(coverage_error) + expected_offset = interval_end + if expected_offset != payload_end: + raise _CagraVerificationError(coverage_error) + + def _verify_checksum_with_signature_cache( + self, + data_path: Path, + data_input: Any, + signature_before: Optional[_FileSignature], + data_file: str, + ) -> None: + cache_hit = ( + signature_before is not None + and signature_before in self._verified_data_files + ) + if cache_hit: + self.CodecUtil.retrieveChecksum(data_input) + else: + self.CodecUtil.checksumEntireFile(data_input) + + if signature_before is None: + return + signature_after = _file_signature(data_path) + if signature_after != signature_before: + raise _CagraVerificationError( + f"CAGRA data file changed during verification: {data_file!r}" + ) + file_age = time.time_ns() - signature_before.changed_at_ns + if cache_hit or file_age < _CAGRA_CACHE_MIN_FILE_AGE_NS: + return + + self._verified_data_files = { + signature + for signature in self._verified_data_files + if signature.resolved_path != signature_before.resolved_path + } + self._verified_data_files.add(signature_before) + + def _verify_cagra_data_file( + self, + context: _CagraDataFileContext, + fields: List[_CagraFieldMetadata], + ) -> None: + try: + signature_before = _file_signature(context.data_path) + except OSError: + signature_before = None + + try: + with _CleanupStack() as cleanups: + data_input = context.directory.openInput( + context.data_file, self.IOContext.READONCE + ) + cleanups.add("close CAGRA data input", data_input.close) + self.CodecUtil.checkIndexHeader( + data_input, + _CAGRA_INDEX_CODEC_NAME, + _CAGRA_INDEX_VERSION, + _CAGRA_INDEX_VERSION, + context.segment_info.info.getId(), + context.suffix, + ) + payload_start = int(data_input.getFilePointer()) + payload_end = int(data_input.length()) - int( + self.CodecUtil.footerLength() + ) + if payload_end < payload_start: + raise _CagraVerificationError( + f"CAGRA data file is truncated: {context.data_file!r}" + ) + self._verify_cagra_payload_coverage( + fields, + payload_start, + payload_end, + context.data_file, + ) + self._verify_checksum_with_signature_cache( + context.data_path, + data_input, + signature_before, + context.data_file, + ) + except _CagraVerificationError: + raise + except Exception as exc: + raise _CagraVerificationError( + "CAGRA-only verification cannot read " + f"{context.data_file!r}: {exc}" + ) from exc + + def _verify_field_against_lucene_metadata( + self, + field_infos: Any, + field: _CagraFieldMetadata, + metadata_file: str, + ) -> None: + field_info = field_infos.fieldInfo(field.field_number) + if field_info is None: + raise _CagraVerificationError( + "CAGRA-only verification found an unknown field number " + f"{field.field_number} in {metadata_file!r}" + ) + field_name = str(field_info.getName()) + if field_name != _VECTOR_FIELD: + raise _CagraVerificationError( + "CAGRA-only verification found data for unexpected field " + f"{field_name!r} in {metadata_file!r}" + ) + if int(field_info.getVectorDimension()) != field.dimensions: + raise _CagraVerificationError( + "CAGRA-only verification found dimensions inconsistent " + f"with Lucene field metadata in {metadata_file!r}" + ) + if field_info.getVectorEncoding() != self.VectorEncoding.FLOAT32: + raise _CagraVerificationError( + "CAGRA-only verification found non-FLOAT32 Lucene field " + f"metadata in {metadata_file!r}" + ) + if ( + field_info.getVectorSimilarityFunction() + != self.VectorSimilarityFunction.EUCLIDEAN + ): + raise _CagraVerificationError( + "CAGRA-only verification found non-Euclidean Lucene field " + f"metadata in {metadata_file!r}" + ) + + @staticmethod + def _validate_segment_deletions( + segment_info: Any, segment_name: str + ) -> None: + deletion_count = int(segment_info.getDelCount()) + soft_deletion_count = int(segment_info.getSoftDelCount()) + if ( + segment_info.hasDeletions() + or deletion_count + or soft_deletion_count + ): + raise _CagraVerificationError( + "CAGRA-only verification does not accept committed " + f"deletions in segment {segment_name!r}: " + f"deleted={deletion_count}, " + f"soft_deleted={soft_deletion_count}" + ) + + @staticmethod + def _cagra_metadata_files( + segment_info: Any, segment_name: str + ) -> List[str]: + metadata_files = [] + for file_name in segment_info.files(): + file_name = str(file_name) + if file_name.endswith(_CAGRA_META_EXTENSION): + metadata_files.append(file_name) + metadata_files.sort() + if not metadata_files: + raise _CagraVerificationError( + "CAGRA-only verification found no " + f"{_CAGRA_META_EXTENSION} metadata for segment " + f"{segment_name!r}" + ) + return metadata_files + + def _read_and_verify_cagra_metadata_file( + self, + context: _CagraDataFileContext, + ) -> List[_CagraFieldMetadata]: + metadata_input = context.directory.openChecksumInput( + context.metadata_file + ) + try: + with _CleanupStack() as cleanups: + cleanups.add( + "close CAGRA metadata input", metadata_input.close + ) + self.CodecUtil.checkIndexHeader( + metadata_input, + _CAGRA_META_CODEC_NAME, + _CAGRA_META_VERSION, + _CAGRA_META_VERSION, + context.segment_info.info.getId(), + context.suffix, + ) + fields = self._read_cagra_fields( + metadata_input, context.metadata_file + ) + self.CodecUtil.checkFooter(metadata_input) + return fields + except _CagraVerificationError: + raise + except Exception as exc: + raise _CagraVerificationError( + "CAGRA-only verification cannot read " + f"{context.metadata_file!r} as cuVS-Lucene metadata " + f"format v{_CAGRA_META_VERSION}: {exc}" + ) from exc + + def _lucene_vector_field_numbers(self, field_infos: Any) -> set[int]: + field_numbers = set() + for raw_field_info in field_infos: + field_info = self.FieldInfo.cast_(raw_field_info) + if int(field_info.getVectorDimension()) > 0: + field_numbers.add(int(field_info.number)) + return field_numbers + + def _verify_cagra_segment( + self, + index_path: Path, + directory: Any, + segment_info: Any, + ) -> _CagraSegmentVerification: + segment_name = str(segment_info.info.name) + self._validate_segment_deletions(segment_info, segment_name) + metadata_files = self._cagra_metadata_files(segment_info, segment_name) + field_infos = self._read_segment_field_infos(directory, segment_info) + + # Verify every metadata/data pair before interpreting fields across + # the segment. This preserves fail-closed file-validation ordering. + verified_field_sources = [] + for metadata_file in metadata_files: + suffix = self._segment_suffix(segment_name, metadata_file) + context = _CagraDataFileContext( + index_path=index_path, + directory=directory, + segment_info=segment_info, + suffix=suffix, + metadata_file=metadata_file, + ) + fields = self._read_and_verify_cagra_metadata_file(context) + self._verify_cagra_data_file(context, fields) + for field in fields: + verified_field_sources.append( + _CagraFieldSource( + metadata=field, + metadata_file=metadata_file, + ) + ) + + # Compare the verified fields with Lucene's segment-wide view and + # aggregate the values needed for index-wide validation. + field_numbers = set() + vector_count = 0 + dimensions = set() + for field_source in verified_field_sources: + field = field_source.metadata + if field.field_number in field_numbers: + raise _CagraVerificationError( + "CAGRA-only verification found duplicate " + f"field {field.field_number} across metadata " + f"files for segment {segment_name!r}" + ) + field_numbers.add(field.field_number) + self._verify_field_against_lucene_metadata( + field_infos, field, field_source.metadata_file + ) + vector_count += field.vector_count + dimensions.add(field.dimensions) + + lucene_field_numbers = self._lucene_vector_field_numbers(field_infos) + if field_numbers != lucene_field_numbers: + raise _CagraVerificationError( + "CAGRA-only verification found vector fields without " + "matching CAGRA metadata in segment " + f"{segment_name!r}: metadata={sorted(field_numbers)}, " + f"Lucene={sorted(lucene_field_numbers)}" + ) + + max_documents = int(segment_info.info.maxDoc()) + if vector_count != max_documents: + raise _CagraVerificationError( + "CAGRA-only verification found " + f"{vector_count} vectors for " + f"{max_documents} documents in segment " + f"{segment_name!r}" + ) + return _CagraSegmentVerification( + field_count=len(field_numbers), + vector_count=vector_count, + dimensions=frozenset(dimensions), + ) + + @staticmethod + def _summarize_cagra_segments( + segments: List[_CagraSegmentVerification], + expected_vector_count: Optional[int], + expected_dimensions: Optional[int], + ) -> _CagraIndexVerification: + field_count = 0 + vector_count = 0 + dimensions = set() + for segment in segments: + field_count += segment.field_count + vector_count += segment.vector_count + dimensions.update(segment.dimensions) + if not segments or field_count == 0: + raise _CagraVerificationError( + "CAGRA-only verification found no nonempty vector fields" + ) + if len(dimensions) != 1: + raise _CagraVerificationError( + "CAGRA-only verification found inconsistent vector " + f"dimensions: {sorted(dimensions)}" + ) + + index_dimensions = next(iter(dimensions)) + if ( + expected_vector_count is not None + and vector_count != expected_vector_count + ): + raise _CagraVerificationError( + "CAGRA-only verification found " + f"{vector_count} vectors; expected {expected_vector_count}" + ) + if ( + expected_dimensions is not None + and index_dimensions != expected_dimensions + ): + raise _CagraVerificationError( + "CAGRA-only verification found " + f"{index_dimensions} dimensions; expected " + f"{expected_dimensions}" + ) + return _CagraIndexVerification( + segment_count=len(segments), + field_count=field_count, + vector_count=vector_count, + dimensions=index_dimensions, + ) + + def verify_index( + self, + index_path: Path, + *, + expected_vector_count: Optional[int] = None, + expected_dimensions: Optional[int] = None, + ) -> _CagraIndexVerification: + """Verify every committed vector field is persisted as CAGRA only.""" + self._attach_current_thread() + directory = self.FSDirectory.open(self.Paths.get(str(index_path))) + verified_segments = [] + with _CleanupStack() as cleanups: + cleanups.add("close Lucene directory", directory.close) + segment_infos = self.SegmentInfos.readLatestCommit(directory) + for raw_segment_info in segment_infos: + segment_info = self.SegmentCommitInfo.cast_(raw_segment_info) + verified_segments.append( + self._verify_cagra_segment( + index_path, directory, segment_info + ) + ) + return self._summarize_cagra_segments( + verified_segments, + expected_vector_count, + expected_dimensions, + ) + + +class _PyLuceneRuntime: + """Own generated PyLucene/Lucene bindings and index operations.""" + + def __init__(self, lucene: Any): + from java.lang import Class + from java.nio.file import Paths + from org.apache.lucene.codecs import Codec, CodecUtil + from org.apache.lucene.document import ( + Document, + KnnFloatVectorField, + StoredField, + ) + from org.apache.lucene.index import ( + DirectoryReader, + FieldInfo, + IndexWriter, + IndexWriterConfig, + SegmentCommitInfo, + SegmentInfos, + SerialMergeScheduler, + VectorEncoding, + VectorSimilarityFunction, + ) + from org.apache.lucene.search import ( + IndexSearcher, + KnnFloatVectorQuery, + ) + from org.apache.lucene.store import FSDirectory, IOContext + + self.lucene = lucene + self.Class = Class + self.Paths = Paths + self.Codec = Codec + self.Document = Document + self.KnnFloatVectorField = KnnFloatVectorField + self.StoredField = StoredField + self.DirectoryReader = DirectoryReader + self.IndexWriter = IndexWriter + self.IndexWriterConfig = IndexWriterConfig + self.SerialMergeScheduler = SerialMergeScheduler + self.VectorSimilarityFunction = VectorSimilarityFunction + self.IndexSearcher = IndexSearcher + self.KnnFloatVectorQuery = KnnFloatVectorQuery + self.FSDirectory = FSDirectory + self.IOContext = IOContext + self._codec_cache: Dict[str, Any] = {} + + # Resolve the base cuvs-java provider before codec construction so + # classpath failures have a direct error location. + self.Class.forName("com.nvidia.cuvs.spi.JDKProvider") + self._cagra_verifier = _CagraIndexVerifier( + attach_current_thread=self.attach_current_thread, + paths=Paths, + codec_util=CodecUtil, + field_info=FieldInfo, + segment_commit_info=SegmentCommitInfo, + segment_infos=SegmentInfos, + vector_encoding=VectorEncoding, + vector_similarity_function=VectorSimilarityFunction, + fs_directory=FSDirectory, + io_context=IOContext, + ) + + @classmethod + def create(cls, config: Dict[str, Any]) -> "_PyLuceneRuntime": + return cls(_initialize_pylucene(config)) + + @property + def pylucene_version(self) -> str: + return str(getattr(self.lucene, "VERSION", "unknown")) + + def attach_current_thread(self) -> None: + vm_environment = self.lucene.getVMEnv() + if vm_environment is None: + raise RuntimeError("PyLucene JVM is not initialized") + vm_environment.attachCurrentThread() + + def resolve_codec(self, codec_name: str) -> Any: + self.attach_current_thread() + cached = self._codec_cache.get(codec_name) + if cached is not None: + return cached + + available_codecs = self.Codec.availableCodecs() + if not available_codecs.contains(codec_name): + available = ", ".join(str(name) for name in available_codecs) + raise RuntimeError( + f"{codec_name} was not advertised by Lucene SPI. " + f"Available codecs: {available}" + ) + codec = self.Codec.forName(codec_name) + if str(codec.getName()) != codec_name: + raise RuntimeError( + f"Requested codec {codec_name}, got {codec.getName()}" + ) + self._codec_cache[codec_name] = codec + return codec + + def codec_telemetry(self, codec_name: str, codec: Any) -> Dict[str, str]: + """Inspect the selected writer path before indexing starts.""" + vectors_format = codec.knnVectorsFormat() + if vectors_format is None: + raise RuntimeError( + f"{codec_name} did not initialize a Lucene vector format" + ) + return _parse_writer_telemetry(str(vectors_format)) + + def _java_float_array(self, vector: np.ndarray) -> Any: + return self.lucene.JArray("float")( + tuple(float(value) for value in vector) + ) + + def verify_cagra_index( + self, + index_path: Path, + expected_vector_count: Optional[int] = None, + expected_dimensions: Optional[int] = None, + ) -> _CagraIndexVerification: + return self._cagra_verifier.verify_index( + index_path, + expected_vector_count=expected_vector_count, + expected_dimensions=expected_dimensions, + ) + + def build_index( + self, + index_path: Path, + vectors: np.ndarray, + resolved_codec: _ResolvedCodec, + ) -> Dict[str, str]: + self.attach_current_thread() + directory = self.FSDirectory.open(self.Paths.get(str(index_path))) + with _CleanupStack() as cleanups: + cleanups.add("close Lucene directory", directory.close) + writer_config = self.IndexWriterConfig() + writer_config.setOpenMode(self.IndexWriterConfig.OpenMode.CREATE) + writer_config.setCodec(resolved_codec.java_codec) + writer_config.setUseCompoundFile(False) + writer_config.setMergeScheduler(self.SerialMergeScheduler()) + writer = self.IndexWriter(directory, writer_config) + self._write_and_close_index(writer, vectors) + return dict(resolved_codec.telemetry) + + def _write_and_close_index(self, writer: Any, vectors: np.ndarray) -> None: + try: + for document_id, vector in enumerate(vectors): + writer.addDocument(self._vector_document(document_id, vector)) + writer.commit() + except BaseException as write_error: + self._rollback_failed_write(writer, write_error) + raise + writer.close() + + @staticmethod + def _rollback_failed_write( + writer: Any, write_error: BaseException + ) -> None: + try: + writer.rollback() + except BaseException as rollback_error: + if isinstance(write_error, Exception): + raise rollback_error from write_error + write_error.add_note( + "IndexWriter rollback also failed: " + f"{type(rollback_error).__name__}: {rollback_error}" + ) + + def _vector_document(self, document_id: int, vector: np.ndarray) -> Any: + document = self.Document() + document.add(self.StoredField(_ID_FIELD, str(document_id))) + document.add( + self.KnnFloatVectorField( + _VECTOR_FIELD, + self._java_float_array(vector), + self.VectorSimilarityFunction.EUCLIDEAN, + ) + ) + return document + + @staticmethod + def _index_vector_dimensions(reader: Any) -> int: + dimensions = set() + for leaf_context in reader.leaves(): + values = leaf_context.reader().getFloatVectorValues(_VECTOR_FIELD) + if values is not None and values.size() > 0: + dimensions.add(int(values.dimension())) + if not dimensions: + raise RuntimeError( + f"Lucene index contains no {_VECTOR_FIELD!r} vectors" + ) + if len(dimensions) != 1: + raise RuntimeError( + "Lucene index has inconsistent vector dimensions: " + f"{dimensions}" + ) + return dimensions.pop() + + def _search_vector( + self, + searcher: Any, + stored_fields: Any, + vector: np.ndarray, + k: int, + ) -> List[_SearchHit]: + query = self.KnnFloatVectorQuery( + _VECTOR_FIELD, + self._java_float_array(vector), + k, + ) + score_docs = searcher.search(query, k).scoreDocs + hits = [] + for score_doc in score_docs: + stored_id = stored_fields.document(score_doc.doc).get(_ID_FIELD) + if stored_id is None: + raise RuntimeError( + f"Lucene document {score_doc.doc} has no stored ID" + ) + hits.append( + _SearchHit( + document_id=int(stored_id), + score=float(score_doc.score), + ) + ) + return hits + + def search_index( + self, + index_path: Path, + query_vectors: np.ndarray, + k: int, + batch_size: int, + ) -> _RuntimeSearchResult: + self.attach_current_thread() + directory = self.FSDirectory.open(self.Paths.get(str(index_path))) + with _CleanupStack() as cleanups: + cleanups.add("close Lucene directory", directory.close) + reader = self.DirectoryReader.open(directory) + cleanups.add("close Lucene index reader", reader.close) + index_dimensions = self._index_vector_dimensions(reader) + if query_vectors.shape[1] != index_dimensions: + raise ValueError( + "Query vector dimensions do not match the Lucene index: " + f"{query_vectors.shape[1]} != {index_dimensions}" + ) + + document_count = int(reader.numDocs()) + if document_count < 1: + raise RuntimeError("Lucene index contains no documents") + lucene_k = min(k, document_count) + searcher = self.IndexSearcher(reader) + stored_fields = searcher.storedFields() + + all_hits: List[List[_SearchHit]] = [] + batch_latencies_ms: List[float] = [] + for batch_start in range(0, query_vectors.shape[0], batch_size): + start = time.perf_counter() + batch_end = min( + batch_start + batch_size, query_vectors.shape[0] + ) + all_hits.extend( + self._search_vectors( + searcher, + stored_fields, + query_vectors[batch_start:batch_end], + lucene_k, + ) + ) + batch_latencies_ms.append( + (time.perf_counter() - start) * 1000.0 + ) + + return _RuntimeSearchResult( + hits=all_hits, + batch_latencies_ms=batch_latencies_ms, + index_dimensions=index_dimensions, + document_count=document_count, + ) + + def _search_vectors( + self, + searcher: Any, + stored_fields: Any, + query_vectors: np.ndarray, + k: int, + ) -> List[List[_SearchHit]]: + hits = [] + for vector in query_vectors: + hits.append( + self._search_vector(searcher, stored_fields, vector, k) + ) + return hits + + +class _SelectedAlgorithmGroup(NamedTuple): + algorithm: str + group: str + configuration: dict + metadata: dict + + +@dataclass(frozen=True) +class _GlobalSelectionScope: + algorithms: frozenset[str] + groups: frozenset[str] + + +@dataclass(frozen=True) +class _AlgorithmSelection: + algorithms: Optional[frozenset[str]] + groups: Optional[frozenset[str]] + explicit_groups: frozenset[Tuple[str, str]] + + def _resolve_global_scope( + self, algorithm_configs: Dict[str, Dict[str, Any]] + ) -> Optional[_GlobalSelectionScope]: + available_algorithms = frozenset(algorithm_configs) + if self.algorithms is not None: + unknown_algorithms = self.algorithms - available_algorithms + if unknown_algorithms: + names = ", ".join(sorted(unknown_algorithms)) + raise ValueError( + f"Unknown PyLucene algorithm selector(s): {names}" + ) + + candidate_algorithms = ( + self.algorithms + if self.algorithms is not None + else available_algorithms + ) + if self.groups is not None: + available_groups = set() + for algorithm in candidate_algorithms: + available_groups.update(algorithm_configs[algorithm]) + unknown_groups = self.groups - available_groups + if unknown_groups: + names = ", ".join(sorted(unknown_groups)) + raise ValueError( + f"Unknown PyLucene group selector(s): {names}" + ) + + has_global_selector = ( + self.algorithms is not None or self.groups is not None + ) + if not has_global_selector and self.explicit_groups: + return None + return _GlobalSelectionScope( + algorithms=candidate_algorithms, + groups=( + self.groups if self.groups is not None else frozenset({"base"}) + ), + ) + + def _validate_explicit_groups( + self, algorithm_configs: Dict[str, Dict[str, Any]] + ) -> None: + for algorithm, group in sorted(self.explicit_groups): + if algorithm not in algorithm_configs: + raise ValueError( + f"Unknown PyLucene algorithm in --algo-groups: {algorithm}" + ) + if group not in algorithm_configs[algorithm]: + raise ValueError( + f"Unknown PyLucene group for {algorithm}: {group}" + ) + + def _selected_pairs( + self, + algorithm_configs: Dict[str, Dict[str, Any]], + global_scope: Optional[_GlobalSelectionScope], + ) -> set[Tuple[str, str]]: + selected_pairs = set(self.explicit_groups) + if global_scope is None: + return selected_pairs + + for algorithm, groups in algorithm_configs.items(): + if algorithm not in global_scope.algorithms: + continue + for group in groups: + if group in global_scope.groups: + selected_pairs.add((algorithm, group)) + return selected_pairs + + def resolve( + self, algorithm_configs: Dict[str, Dict[str, Any]] + ) -> List[_SelectedAlgorithmGroup]: + global_scope = self._resolve_global_scope(algorithm_configs) + self._validate_explicit_groups(algorithm_configs) + selected_pairs = self._selected_pairs(algorithm_configs, global_scope) + + selected_groups = [] + for algorithm, groups in algorithm_configs.items(): + for group, group_config in groups.items(): + if (algorithm, group) not in selected_pairs: + continue + selected_groups.append( + _SelectedAlgorithmGroup( + algorithm=algorithm, + group=group, + configuration=group_config, + metadata={}, + ) + ) + return selected_groups + + +@dataclass(frozen=True) +class _BenchmarkConfigContext: + dataset: str + dataset_path: str + subset_scope: Optional[str] + runtime_config: Dict[str, Any] + neighbor_count: int + batch_size: int + + +class PyLuceneConfigLoader(ConfigLoader): + """Load PyLucene algorithm configurations.""" + + def __init__(self, config_path: Optional[Union[str, os.PathLike]] = None): + self.config_path = ( + os.fspath(config_path) + if config_path is not None + else os.path.join( + os.path.dirname(os.path.realpath(__file__)), "../config" + ) + ) + + @property + def backend_type(self) -> str: + return "pylucene" + + @staticmethod + def _parse_name_filter( + value: Optional[str], + ) -> Optional[frozenset[str]]: + if not value: + return None + + names = set() + for raw_name in value.split(","): + name = raw_name.strip() + if name: + names.add(name) + return frozenset(names) or None + + @staticmethod + def _parse_algorithm_group_filter( + value: Optional[str], + ) -> frozenset[Tuple[str, str]]: + selections = set() + if not value: + return frozenset() + + for selection in value.split(","): + algorithm, separator, group = selection.strip().partition(".") + if not separator or not algorithm or not group: + raise ValueError( + "algo_groups entries must use ., " + f"got {selection!r}" + ) + selections.add((algorithm, group)) + return frozenset(selections) + + @classmethod + def _selection_from_options( + cls, options: Dict[str, Any] + ) -> _AlgorithmSelection: + return _AlgorithmSelection( + algorithms=cls._parse_name_filter(options.get("algorithms")), + groups=cls._parse_name_filter(options.get("groups")), + explicit_groups=cls._parse_algorithm_group_filter( + options.get("algo_groups") + ), + ) + + def _match_backend_algorithm_config( + self, config: Any + ) -> Optional[Tuple[str, Dict[str, Any]]]: + if not isinstance(config, dict): + return None + + algorithm = config.get("name") + if not isinstance(algorithm, str) or not algorithm: + return None + + declared_backend = config.get("backend") + if declared_backend is None: + belongs_to_backend = algorithm.startswith("pylucene_") + else: + belongs_to_backend = declared_backend == self.backend_type + if not belongs_to_backend: + return None + return algorithm, config.get("groups", {}) + + def _load_algorithm_configs( + self, algorithm_files: List[str] + ) -> Dict[str, Dict[str, Any]]: + algorithm_configs = {} + for algorithm_file in algorithm_files: + matched_config = self._match_backend_algorithm_config( + self.load_yaml_file(algorithm_file) + ) + if matched_config is None: + continue + algorithm, groups = matched_config + # Later files are explicit overrides and determine output order. + algorithm_configs.pop(algorithm, None) + algorithm_configs[algorithm] = groups + return algorithm_configs + + def _discover_algo_groups( + self, + dataset_conf: dict, + dataset: str, + dataset_path: str, + **kwargs, + ) -> List[Tuple[str, str, dict, dict]]: + algorithm_files = self.gather_algorithm_configs( + self.config_path, kwargs.get("algorithm_configuration") + ) + algorithm_configs = self._load_algorithm_configs(algorithm_files) + selection = self._selection_from_options(kwargs) + return selection.resolve(algorithm_configs) + + @staticmethod + def _runtime_config(options: Dict[str, Any]) -> Dict[str, Any]: + runtime_keys = ( + "cuvs_java_jar", + "cuvs_lucene_jar", + "java_library_path", + "jvm_args", + ) + runtime_config = {} + for key in runtime_keys: + value = options.get(key) + if value is not None: + runtime_config[key] = value + return runtime_config + + @staticmethod + def _subset_scope(subset_size: Any) -> Optional[str]: + subset_size = _validate_subset_size(subset_size) + if subset_size is None: + return None + return f"subset{subset_size}" + + @staticmethod + def _effective_parameters( + build_combinations: List[Dict[str, Any]], + search_combinations: List[Dict[str, Any]], + *, + tune_mode: bool, + tune_build_params: Optional[Dict[str, Any]], + tune_search_params: Optional[Dict[str, Any]], + ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + use_tuned_parameters = tune_mode and tune_build_params is not None + if not use_tuned_parameters: + return build_combinations, search_combinations + search_params = [tune_search_params] if tune_search_params else [{}] + return [tune_build_params], search_params + + @staticmethod + def _index_label( + algorithm: str, + group: str, + subset_scope: Optional[str], + build_params: Dict[str, Any], + ) -> str: + prefix = algorithm if group == "base" else f"{algorithm}_{group}" + label_parts = [prefix] + if subset_scope is not None: + label_parts.append(subset_scope) + for key, value in build_params.items(): + label_parts.append(f"{key}{value}") + return ".".join(label_parts) + + @classmethod + def _benchmark_config( + cls, + context: _BenchmarkConfigContext, + algorithm: str, + group: str, + build_params: Dict[str, Any], + search_params: List[Dict[str, Any]], + ) -> BenchmarkConfig: + index_label = cls._index_label( + algorithm, group, context.subset_scope, build_params + ) + index_path = os.path.join( + context.dataset_path, + context.dataset, + "index", + index_label, + ) + result_stem = f"{algorithm},{group}" + if context.subset_scope is not None: + result_stem = f"{result_stem},{context.subset_scope}" + + index_config = IndexConfig( + name=index_label, + algo=algorithm, + build_param=build_params, + search_params=search_params, + file=index_path, + ) + backend_config = { + "name": index_label, + "algo": algorithm, + "codec": build_params.get("codec"), + "requires_gpu": True, + "output_filename": ( + result_stem, + f"{result_stem},k{context.neighbor_count}," + f"bs{context.batch_size}", + ), + **context.runtime_config, + } + return BenchmarkConfig( + indexes=[index_config], + backend_config=backend_config, + ) + + @classmethod + def _group_benchmark_configs( + cls, + context: _BenchmarkConfigContext, + algorithm: str, + group: str, + build_params: List[Dict[str, Any]], + search_params: List[Dict[str, Any]], + ) -> List[BenchmarkConfig]: + _validate_search_params(search_params) + benchmark_configs = [] + for params in build_params: + _validate_build_params(params) + benchmark_configs.append( + cls._benchmark_config( + context, + algorithm, + group, + params, + search_params, + ) + ) + return benchmark_configs + + def _build_benchmark_configs( + self, + dataset_config: DatasetConfig, + dataset_conf: dict, + dataset: str, + dataset_path: str, + expanded_groups: List[Tuple[str, str, dict, List, List, dict]], + **kwargs, + ) -> List[BenchmarkConfig]: + context = _BenchmarkConfigContext( + dataset=dataset, + dataset_path=dataset_path, + subset_scope=self._subset_scope(dataset_config.subset_size), + runtime_config=self._runtime_config(kwargs), + neighbor_count=kwargs.get("count", 10), + batch_size=kwargs.get("batch_size", 10000), + ) + tune_mode = kwargs.get("_tune_mode", False) + tune_build_params = kwargs.get("_tune_build_params") + tune_search_params = kwargs.get("_tune_search_params") + + benchmark_configs = [] + for ( + algorithm, + group, + _group_config, + build_combinations, + search_combinations, + _group_metadata, + ) in expanded_groups: + build_params, search_params = self._effective_parameters( + build_combinations, + search_combinations, + tune_mode=tune_mode, + tune_build_params=tune_build_params, + tune_search_params=tune_search_params, + ) + benchmark_configs.extend( + self._group_benchmark_configs( + context, + algorithm, + group, + build_params, + search_params, + ) + ) + + return benchmark_configs + + +class PyLuceneBackend(BenchmarkBackend): + """Build and search cuVS-Lucene indexes through PyLucene.""" + + orchestrator_persists_results = True + + def __init__(self, config: Dict[str, Any]): + super().__init__(config) + self._runtime: Optional[_PyLuceneRuntime] = None + + @property + def algo(self) -> str: + return self.config.get("algo", "pylucene") + + def cleanup(self) -> None: + """Release Python references; the process-wide JVM remains active.""" + self._runtime = None + + def _get_runtime(self) -> _PyLuceneRuntime: + if self._runtime is None: + self._runtime = _PyLuceneRuntime.create(self.config) + return self._runtime + + def _failed_build_result( + self, + error_message: str, + index_path: str = "", + build_params: Optional[Dict[str, Any]] = None, + ) -> BuildResult: + return BuildResult( + index_path=index_path, + build_time_seconds=0.0, + index_size_bytes=0, + algorithm=self.algo, + build_params=build_params or {}, + success=False, + error_message=error_message, + ) + + def _failed_search_result( + self, + k: int, + error_message: str, + search_params: Optional[List[Dict[str, Any]]] = None, + ) -> SearchResult: + result_k = max(0, k) + return SearchResult( + neighbors=np.empty((0, result_k), dtype=np.int64), + distances=np.empty((0, result_k), dtype=np.float32), + search_time_ms=0.0, + queries_per_second=0.0, + recall=0.0, + algorithm=self.algo, + search_params=search_params or [], + success=False, + error_message=error_message, + ) + + def _verify_existing_index( + self, + index_path: Path, + codec_name: str, + *, + expected_vector_count: Optional[int] = None, + expected_dimensions: Optional[int] = None, + ) -> Tuple[_IndexProvenanceVerification, Dict[str, Any]]: + """Verify backend ownership and codec-specific persisted data.""" + if codec_name == _CAGRA_CODEC: + provenance = _verify_cagra_provenance( + index_path, + expected_vector_count=expected_vector_count, + expected_dimensions=expected_dimensions, + ) + cagra_verification = self._get_runtime().verify_cagra_index( + index_path, + expected_vector_count=provenance.vector_count, + expected_dimensions=provenance.dimensions, + ) + metadata = { + "cagra_provenance": provenance.to_metadata(), + "cagra_verification": cagra_verification.to_metadata(), + } + else: + provenance = _verify_hnsw_provenance( + index_path, + codec_name, + expected_vector_count=expected_vector_count, + expected_dimensions=expected_dimensions, + ) + metadata = {"hnsw_verification": provenance.to_metadata()} + return provenance, metadata + + def _dry_run_build_result( + self, + index_path: Path, + codec_name: str, + build_params: Dict[str, Any], + ) -> BuildResult: + print( + f"[dry_run] Would build PyLucene index '{index_path}' " + f"with codec={codec_name}" + ) + return BuildResult( + index_path=str(index_path), + build_time_seconds=0.0, + index_size_bytes=0, + algorithm=self.algo, + build_params=build_params, + metadata={"codec": codec_name}, + success=True, + ) + + def _decide_existing_index( + self, + dataset: Dataset, + index_path: Path, + codec_name: str, + build_params: Dict[str, Any], + force: bool, + ) -> _ExistingIndexDecision: + if not index_path.exists(): + return _ExistingIndexDecision.build() + if not index_path.is_dir(): + return _ExistingIndexDecision.reject( + self._failed_build_result( + f"PyLucene index path is not a directory: {index_path}", + str(index_path), + build_params, + ) + ) + if not _has_lucene_segments(index_path): + return _ExistingIndexDecision.reject( + self._failed_build_result( + "Existing PyLucene index directory does not contain " + f"a Lucene segments file: {index_path}", + str(index_path), + build_params, + ) + ) + if force: + return _ExistingIndexDecision.build() + + metadata: Dict[str, Any] = { + "codec": codec_name, + "skipped": True, + } + try: + _validate_metric(dataset) + expected_shape = _expected_training_shape(dataset) + _, verification_metadata = self._verify_existing_index( + index_path, + codec_name, + expected_vector_count=( + expected_shape[0] if expected_shape else None + ), + expected_dimensions=( + expected_shape[1] if expected_shape else None + ), + ) + metadata.update(verification_metadata) + result = BuildResult( + index_path=str(index_path), + build_time_seconds=0.0, + index_size_bytes=_index_size(index_path), + algorithm=self.algo, + build_params=build_params, + metadata=metadata, + success=True, + ) + except Exception as exc: + return _ExistingIndexDecision.reject( + self._failed_build_result( + _exception_summary(exc), + str(index_path), + build_params, + ) + ) + + return _ExistingIndexDecision.reuse(result) + + @staticmethod + def _training_vectors_for_build( + dataset: Dataset, codec_name: str + ) -> np.ndarray: + _validate_metric(dataset) + vectors = _validate_float32_matrix( + dataset.training_vectors, "training_vectors" + ) + if vectors.shape[0] < 2: + raise ValueError( + f"{codec_name} requires at least two training vectors; " + "cuVS-Lucene does not invoke cuVS for a single-vector index" + ) + return vectors + + @staticmethod + def _require_gpu_writer( + runtime: _PyLuceneRuntime, codec_name: str + ) -> _ResolvedCodec: + java_codec = runtime.resolve_codec(codec_name) + telemetry = runtime.codec_telemetry(codec_name, java_codec) + writer_path = telemetry.get("writerPath") + expected_writer_path = _EXPECTED_WRITER_PATH[codec_name] + if writer_path != expected_writer_path: + raise RuntimeError( + f"{codec_name} uses writerPath={writer_path!r}; " + f"expected {expected_writer_path!r}. Refusing to run " + "a silent CPU or alternate-index fallback as a cuVS build." + ) + return _ResolvedCodec( + codec_name=codec_name, + java_codec=java_codec, + telemetry=telemetry, + writer_path=writer_path, + ) + + @staticmethod + def _verify_and_persist_built_index_provenance( + runtime: _PyLuceneRuntime, + index_path: Path, + codec_name: str, + vectors: np.ndarray, + ) -> Dict[str, Any]: + vector_count = int(vectors.shape[0]) + dimensions = int(vectors.shape[1]) + if codec_name == _CAGRA_CODEC: + cagra_verification = runtime.verify_cagra_index( + index_path, + expected_vector_count=vector_count, + expected_dimensions=dimensions, + ) + _write_cagra_provenance( + index_path, + vector_count=vector_count, + dimensions=dimensions, + ) + provenance = _verify_cagra_provenance( + index_path, + expected_vector_count=vector_count, + expected_dimensions=dimensions, + ) + return { + "cagra_provenance": provenance.to_metadata(), + "cagra_verification": cagra_verification.to_metadata(), + } + + _write_hnsw_provenance( + index_path, + codec_name, + vector_count=vector_count, + dimensions=dimensions, + ) + verification = _verify_hnsw_provenance( + index_path, + codec_name, + expected_vector_count=vector_count, + expected_dimensions=dimensions, + ) + return {"hnsw_verification": verification.to_metadata()} + + @staticmethod + def _cleanup_partial_index( + index_path: Path, created_for_build: bool + ) -> Optional[Exception]: + if not created_for_build or not index_path.exists(): + return None + try: + _safe_remove_index(index_path) + except Exception as exc: + return exc + return None + + def build( + self, + dataset: Dataset, + indexes: List[IndexConfig], + force: bool = False, + dry_run: bool = False, + ) -> BuildResult: + """Build one local Lucene index with the selected cuVS codec.""" + if len(indexes) != 1: + return self._failed_build_result( + "PyLucene backend requires exactly one index configuration" + ) + + index_config = indexes[0] + build_params = index_config.build_param + index_path = Path(index_config.file) + try: + _validate_build_params(build_params) + codec_name = _configured_codec(build_params, self.config) + except Exception as exc: + return self._failed_build_result( + str(exc), str(index_path), build_params + ) + + if dry_run: + return self._dry_run_build_result( + index_path, codec_name, build_params + ) + + existing_index_decision = self._decide_existing_index( + dataset, index_path, codec_name, build_params, force + ) + if existing_index_decision.action is _ExistingIndexAction.BUILD: + return self._build_new_index( + dataset, index_path, codec_name, build_params + ) + + return existing_index_decision.completed_result() + + def _build_new_index( + self, + dataset: Dataset, + index_path: Path, + codec_name: str, + build_params: Dict[str, Any], + ) -> BuildResult: + """Build a validated index that is not eligible for reuse.""" + + created_for_build = False + try: + vectors = self._training_vectors_for_build(dataset, codec_name) + runtime = self._get_runtime() + # Preflight must succeed before an existing index is replaced. + resolved_codec = self._require_gpu_writer(runtime, codec_name) + + if index_path.exists(): + _safe_remove_index(index_path) + index_path.mkdir(parents=True) + created_for_build = True + + start = time.perf_counter() + telemetry = runtime.build_index( + index_path, vectors, resolved_codec + ) + build_time = time.perf_counter() - start + + metadata = { + "codec": codec_name, + "pylucene_version": runtime.pylucene_version, + "writer_path": resolved_codec.writer_path, + "writer_telemetry": telemetry, + } + metadata.update( + self._verify_and_persist_built_index_provenance( + runtime, index_path, codec_name, vectors + ) + ) + + return BuildResult( + index_path=str(index_path), + build_time_seconds=build_time, + index_size_bytes=_index_size(index_path), + algorithm=self.algo, + build_params=build_params, + metadata=metadata, + success=True, + ) + except Exception as exc: + cleanup_error = self._cleanup_partial_index( + index_path, created_for_build + ) + error_message = _exception_summary(exc) + if cleanup_error is not None: + error_message += ( + "; failed to remove partial index: " + f"{type(cleanup_error).__name__}: {cleanup_error}" + ) + return self._failed_build_result( + error_message, + str(index_path), + build_params, + ) + except BaseException as exc: + cleanup_error = self._cleanup_partial_index( + index_path, created_for_build + ) + if cleanup_error is not None: + exc.add_note( + "Failed to remove partial index: " + f"{type(cleanup_error).__name__}: {cleanup_error}" + ) + raise + + def _resolve_search_plan( + self, + index_config: IndexConfig, + *, + k: int, + batch_size: int, + mode: str, + search_threads: Optional[Union[int, str]], + ) -> _SearchPlan: + search_params = index_config.search_params or [{}] + _validate_build_params(index_config.build_param) + codec_name = _configured_codec(index_config.build_param, self.config) + if k < 1: + raise ValueError("k must be positive") + if batch_size < 1: + raise ValueError("batch_size must be positive") + if mode != "latency": + raise ValueError( + "PyLucene backend currently supports only latency mode" + ) + if search_threads not in (None, 1, "1"): + raise ValueError( + "PyLucene backend currently supports only one search thread" + ) + _validate_search_params(search_params) + if codec_name == _CAGRA_CODEC and k > 1024: + raise ValueError( + "CuVS2510GPUSearchCodec benchmarks support k <= 1024 " + "to avoid cuVS-Lucene search paths that can use GPU " + "brute-force search above that limit" + ) + return _SearchPlan( + index_path=Path(index_config.file), + codec_name=codec_name, + search_params=search_params, + k=k, + batch_size=batch_size, + mode=mode, + ) + + def _dry_run_search_result( + self, + plan: _SearchPlan, + ) -> SearchResult: + print( + f"[dry_run] Would search PyLucene index '{plan.index_path}' " + f"with codec={plan.codec_name}, k={plan.k}, " + f"batch_size={plan.batch_size}" + ) + return SearchResult( + neighbors=np.empty((0, plan.k), dtype=np.int64), + distances=np.empty((0, plan.k), dtype=np.float32), + search_time_ms=0.0, + queries_per_second=0.0, + recall=0.0, + algorithm=self.algo, + search_params=plan.search_params, + metadata={"codec": plan.codec_name}, + success=True, + ) + + @staticmethod + def _load_search_inputs(dataset: Dataset) -> _SearchInputs: + _validate_metric(dataset) + query_vectors = _validate_float32_matrix( + dataset.query_vectors, "query_vectors" + ) + expected_shape = _expected_training_shape(dataset) + if ( + expected_shape is not None + and expected_shape[1] != query_vectors.shape[1] + ): + raise ValueError( + "Query vector dimensions do not match the dataset: " + f"{query_vectors.shape[1]} != {expected_shape[1]}" + ) + return _SearchInputs( + query_vectors=query_vectors, + expected_training_shape=expected_shape, + ) + + def _execute_search( + self, dataset: Dataset, plan: _SearchPlan + ) -> SearchResult: + end_to_end_start = time.perf_counter() + inputs = self._load_search_inputs(dataset) + query_vectors = inputs.query_vectors + expected_shape = inputs.expected_training_shape + provenance, verification_metadata = self._verify_existing_index( + plan.index_path, + plan.codec_name, + expected_vector_count=( + expected_shape[0] if expected_shape else None + ), + expected_dimensions=int(query_vectors.shape[1]), + ) + runtime = self._get_runtime() + + runtime_result = runtime.search_index( + plan.index_path, + query_vectors, + plan.k, + plan.batch_size, + ) + end_to_end_time_ms = (time.perf_counter() - end_to_end_start) * 1000.0 + processed = _process_search_result( + runtime_result, + provenance, + query_count=int(query_vectors.shape[0]), + k=plan.k, + batch_size=plan.batch_size, + ) + metadata = { + "codec": plan.codec_name, + "pylucene_version": runtime.pylucene_version, + "index_dimensions": runtime_result.index_dimensions, + "document_count": runtime_result.document_count, + "batch_size": plan.batch_size, + "num_batches": processed.num_batches, + "mode": plan.mode, + "end_to_end_time_ms": end_to_end_time_ms, + "non_query_overhead_time_ms": max( + 0.0, end_to_end_time_ms - processed.search_time_ms + ), + } + metadata.update(verification_metadata) + + return SearchResult( + neighbors=processed.neighbors, + distances=processed.distances, + search_time_ms=processed.search_time_ms, + queries_per_second=processed.queries_per_second, + recall=0.0, + algorithm=self.algo, + search_params=plan.search_params, + latency_seconds=processed.latency_seconds, + latency_percentiles=processed.latency_percentiles, + metadata=metadata, + success=True, + ) + + def search( + self, + dataset: Dataset, + indexes: List[IndexConfig], + k: int, + batch_size: int = 10000, + mode: str = "latency", + force: bool = False, + search_threads: Optional[int] = None, + dry_run: bool = False, + ) -> SearchResult: + """Search one local Lucene index and report per-batch latency.""" + if len(indexes) != 1: + return self._failed_search_result( + k, "PyLucene backend requires exactly one index configuration" + ) + + index_config = indexes[0] + try: + plan = self._resolve_search_plan( + index_config, + k=k, + batch_size=batch_size, + mode=mode, + search_threads=search_threads, + ) + except Exception as exc: + return self._failed_search_result( + k, + str(exc), + index_config.search_params or [{}], + ) + + if dry_run: + return self._dry_run_search_result(plan) + + if not plan.index_path.is_dir(): + return self._failed_search_result( + k, + f"PyLucene index directory does not exist: {plan.index_path}", + plan.search_params, + ) + + try: + return self._execute_search(dataset, plan) + except Exception as exc: + return self._failed_search_result( + k, + _exception_summary(exc), + plan.search_params, + ) + + +__all__ = ["PyLuceneBackend", "PyLuceneConfigLoader"] diff --git a/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_cagra.yaml b/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_cagra.yaml new file mode 100644 index 0000000000..2e272895b7 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_cagra.yaml @@ -0,0 +1,10 @@ +name: pylucene_cuvs_cagra +groups: + base: + build: + codec: ["CuVS2510GPUSearchCodec"] + search: {} + test: + build: + codec: ["CuVS2510GPUSearchCodec"] + search: {} diff --git a/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_hnsw.yaml b/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_hnsw.yaml new file mode 100644 index 0000000000..a82700aeeb --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_hnsw.yaml @@ -0,0 +1,13 @@ +name: pylucene_cuvs_hnsw +groups: + base: + build: + codec: + - "Lucene101AcceleratedHNSWCodec" + - "Lucene101AcceleratedHNSWBaseLayerCodec" + - "Lucene101AcceleratedHNSWMultiLayerCodec" + search: {} + test: + build: + codec: ["Lucene101AcceleratedHNSWCodec"] + search: {} diff --git a/python/cuvs_bench/cuvs_bench/orchestrator/__init__.py b/python/cuvs_bench/cuvs_bench/orchestrator/__init__.py index 7600101439..36bc5539a7 100644 --- a/python/cuvs_bench/cuvs_bench/orchestrator/__init__.py +++ b/python/cuvs_bench/cuvs_bench/orchestrator/__init__.py @@ -1,10 +1,15 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # from .orchestrator import BenchmarkOrchestrator -from .config_loaders import ConfigLoader, BenchmarkConfig, DatasetConfig, CppGBenchConfigLoader +from .config_loaders import ( + ConfigLoader, + BenchmarkConfig, + DatasetConfig, + CppGBenchConfigLoader, +) from ..backends.registry import ( get_backend_class, list_backends, @@ -12,6 +17,7 @@ get_config_loader, ) from ..backends.opensearch import OpenSearchConfigLoader +from ..backends.pylucene import PyLuceneConfigLoader __all__ = [ # Main orchestrator @@ -22,6 +28,7 @@ "DatasetConfig", "CppGBenchConfigLoader", "OpenSearchConfigLoader", + "PyLuceneConfigLoader", # Registry functions "get_backend_class", "list_backends", @@ -34,10 +41,12 @@ # Register built-in config loaders # ============================================================================ + def _register_builtin_loaders(): """Register built-in config loaders.""" register_config_loader("cpp_gbench", CppGBenchConfigLoader) register_config_loader("opensearch", OpenSearchConfigLoader) + register_config_loader("pylucene", PyLuceneConfigLoader) # Auto-register when module is imported diff --git a/python/cuvs_bench/cuvs_bench/orchestrator/config_loaders.py b/python/cuvs_bench/cuvs_bench/orchestrator/config_loaders.py index 5e3a51b7f7..37be49f269 100644 --- a/python/cuvs_bench/cuvs_bench/orchestrator/config_loaders.py +++ b/python/cuvs_bench/cuvs_bench/orchestrator/config_loaders.py @@ -367,7 +367,7 @@ def gather_algorithm_configs( list A list of paths to the algorithm configuration files. """ - algos_conf_fs = os.listdir(os.path.join(config_path, "algos")) + algos_conf_fs = sorted(os.listdir(os.path.join(config_path, "algos"))) algos_conf_fs = [ os.path.join(config_path, "algos", f) for f in algos_conf_fs @@ -378,7 +378,7 @@ def gather_algorithm_configs( if os.path.isdir(algorithm_configuration): algos_conf_fs += [ os.path.join(algorithm_configuration, f) - for f in os.listdir(algorithm_configuration) + for f in sorted(os.listdir(algorithm_configuration)) if f.endswith((".yaml", ".yml")) ] elif os.path.isfile(algorithm_configuration): diff --git a/python/cuvs_bench/cuvs_bench/tests/_pylucene_test_utils.py b/python/cuvs_bench/cuvs_bench/tests/_pylucene_test_utils.py new file mode 100644 index 0000000000..16b40affc3 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/_pylucene_test_utils.py @@ -0,0 +1,224 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Shared fakes and data builders for PyLucene unit tests.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np + +import cuvs_bench.backends.pylucene as pylucene_backend +from cuvs_bench._bin_format import write_bin_header +from cuvs_bench.backends.base import Dataset +from cuvs_bench.backends.pylucene import ( + PyLuceneBackend, + _CagraIndexVerification, + _RuntimeSearchResult, + _SearchHit, +) +from cuvs_bench.orchestrator.config_loaders import IndexConfig + +_HNSW_CODEC = "Lucene101AcceleratedHNSWCodec" +_HNSW_BASE_LAYER_CODEC = "Lucene101AcceleratedHNSWBaseLayerCodec" +_CAGRA_CODEC = "CuVS2510GPUSearchCodec" + + +class _FakeRuntime: + pylucene_version = "test" + + def __init__( + self, + *, + writer_path: str = "gpu-hnsw", + index_dimensions: int = 4, + document_count: int = 10, + hits: list[list[_SearchHit]] | None = None, + batch_latencies_ms: list[float] | None = None, + validate_query_dimensions: bool = True, + ): + self.writer_path = writer_path + self.index_dimensions = index_dimensions + self.document_count = document_count + self.hits = hits + self.batch_latencies_ms = batch_latencies_ms + self.validate_query_dimensions = validate_query_dimensions + self.resolve_calls = [] + self.telemetry_calls = [] + self.build_calls = [] + self.search_calls = [] + self.verification_calls = [] + self.resolve_error = None + self.build_error = None + self.search_error = None + self.verification_error = None + + def resolve_codec(self, codec_name): + self.resolve_calls.append(codec_name) + if self.resolve_error is not None: + raise self.resolve_error + return codec_name + + def codec_telemetry(self, codec_name, java_codec): + self.telemetry_calls.append(codec_name) + assert java_codec == codec_name + return {"writerPath": self.writer_path, "codec": codec_name} + + def build_index(self, index_path, vectors, resolved_codec): + self.build_calls.append((index_path, vectors.copy(), resolved_codec)) + if self.build_error is not None: + (index_path / "partial").write_bytes(b"partial") + raise self.build_error + self.document_count = int(vectors.shape[0]) + (index_path / "segments_1").write_bytes(b"index") + return dict(resolved_codec.telemetry) + + def verify_cagra_index( + self, + index_path, + expected_vector_count=None, + expected_dimensions=None, + ): + self.verification_calls.append( + (index_path, expected_vector_count, expected_dimensions) + ) + if self.verification_error is not None: + raise self.verification_error + return _CagraIndexVerification( + segment_count=1, + field_count=1, + vector_count=( + expected_vector_count + if expected_vector_count is not None + else 10 + ), + dimensions=( + expected_dimensions if expected_dimensions is not None else 4 + ), + ) + + def search_index(self, index_path, query_vectors, k, batch_size): + self.search_calls.append( + (index_path, query_vectors.copy(), k, batch_size) + ) + if self.search_error is not None: + raise self.search_error + if ( + self.validate_query_dimensions + and query_vectors.shape[1] != self.index_dimensions + ): + raise ValueError( + "Query vector dimensions do not match the Lucene index" + ) + hits = self.hits + if hits is None: + hits = [ + [_SearchHit(document_id=0, score=1.0)] + for _ in range(query_vectors.shape[0]) + ] + batch_latencies_ms = self.batch_latencies_ms + if batch_latencies_ms is None: + batch_latencies_ms = [ + 1.0 for _ in range(0, query_vectors.shape[0], batch_size) + ] + return _RuntimeSearchResult( + hits=hits, + batch_latencies_ms=batch_latencies_ms, + index_dimensions=self.index_dimensions, + document_count=self.document_count, + ) + + +def _dataset( + *, + n_base: int = 10, + n_queries: int = 2, + dimensions: int = 4, + dtype=np.float32, + distance_metric: str = "euclidean", +) -> Dataset: + rng = np.random.default_rng(7) + base = rng.random((n_base, dimensions)).astype(dtype) + queries = rng.random((n_queries, dimensions)).astype(dtype) + return Dataset( + name="test", + training_vectors=base, + query_vectors=queries, + distance_metric=distance_metric, + ) + + +def _index( + index_path: Path, + *, + codec: str = _HNSW_CODEC, + search_params: list[dict] | None = None, +) -> IndexConfig: + return IndexConfig( + name="pylucene-test", + algo="pylucene_cuvs_hnsw", + build_param={"codec": codec}, + search_params=[{}] if search_params is None else search_params, + file=str(index_path), + ) + + +def _backend( + runtime: _FakeRuntime | None = None, + *, + codec: str = _HNSW_CODEC, +) -> PyLuceneBackend: + backend = PyLuceneBackend( + { + "name": "pylucene-test", + "algo": "pylucene_cuvs_hnsw", + "codec": codec, + } + ) + if runtime is not None: + backend._runtime = runtime + return backend + + +def _prepare_hnsw_index( + index_path: Path, + *, + codec: str = _HNSW_CODEC, + vector_count: int = 10, + dimensions: int = 4, +) -> Path: + index_path.mkdir(parents=True, exist_ok=True) + (index_path / "segments_1").write_bytes(b"index") + pylucene_backend._write_hnsw_provenance( + index_path, + codec, + vector_count=vector_count, + dimensions=dimensions, + ) + return index_path / pylucene_backend._HNSW_PROVENANCE_FILE + + +def _prepare_cagra_index( + index_path: Path, + *, + vector_count: int = 10, + dimensions: int = 4, +) -> Path: + index_path.mkdir(parents=True, exist_ok=True) + (index_path / "segments_1").write_bytes(b"index") + pylucene_backend._write_cagra_provenance( + index_path, + vector_count=vector_count, + dimensions=dimensions, + ) + return index_path / pylucene_backend._CAGRA_PROVENANCE_FILE + + +def _write_test_bin(path: Path, data: np.ndarray) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("wb") as file: + write_bin_header(file, data.shape[0], data.shape[1]) + np.ascontiguousarray(data).tofile(file) diff --git a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_backend.py b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_backend.py new file mode 100644 index 0000000000..a5b75654c5 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_backend.py @@ -0,0 +1,1038 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Build and search behavior tests for the PyLucene backend.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import pytest + +import cuvs_bench.backends.pylucene as pylucene_backend +from cuvs_bench.backends.pylucene import _SearchHit +from cuvs_bench.tests._pylucene_test_utils import ( + _CAGRA_CODEC, + _HNSW_BASE_LAYER_CODEC, + _HNSW_CODEC, + _FakeRuntime, + _backend, + _dataset, + _index, + _prepare_cagra_index, + _prepare_hnsw_index, +) + + +@pytest.mark.parametrize( + ("score", "distance"), + [(1.0, 0.0), (0.5, 1.0), (0.2, 4.0), (0.0, np.inf)], +) +def test_score_to_squared_euclidean(score, distance): + assert pylucene_backend._score_to_squared_euclidean( + score + ) == pytest.approx(distance) + + +def test_build_dry_run_does_not_initialize_pylucene(tmp_path): + backend = _backend() + result = backend.build( + _dataset(), [_index(tmp_path / "index")], dry_run=True + ) + + assert result.success + assert backend._runtime is None + assert result.metadata["codec"] == _HNSW_CODEC + + +def test_build_creates_index_and_reports_gpu_writer(tmp_path): + runtime = _FakeRuntime() + backend = _backend(runtime) + index_path = tmp_path / "index" + + result = backend.build(_dataset(), [_index(index_path)], force=True) + + assert result.success + assert result.index_size_bytes > len(b"index") + assert result.metadata["writer_path"] == "gpu-hnsw" + assert result.metadata["writer_telemetry"]["writerPath"] == "gpu-hnsw" + assert result.metadata["hnsw_verification"] == { + "status": "gpu-hnsw-provenance", + "schema_version": 1, + "codec": _HNSW_CODEC, + "writer_path": "gpu-hnsw", + "vector_count": 10, + "dimensions": 4, + "commit_file_count": 1, + } + assert (index_path / pylucene_backend._HNSW_PROVENANCE_FILE).is_file() + assert runtime.resolve_calls == [_HNSW_CODEC] + assert runtime.telemetry_calls == [_HNSW_CODEC] + assert len(runtime.build_calls) == 1 + assert runtime.verification_calls == [] + + +def test_build_verifies_persisted_cagra_index(tmp_path): + runtime = _FakeRuntime(writer_path="gpu-cagra") + index_path = tmp_path / "index" + + result = _backend(runtime, codec=_CAGRA_CODEC).build( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + force=True, + ) + + assert result.success + assert runtime.verification_calls == [(index_path, 10, 4)] + assert result.metadata["cagra_provenance"] == { + "status": "gpu-cagra-provenance", + "schema_version": 1, + "codec": _CAGRA_CODEC, + "writer_path": "gpu-cagra", + "vector_count": 10, + "dimensions": 4, + "commit_file_count": 1, + } + assert result.metadata["cagra_verification"] == { + "status": "cagra-only", + "segment_count": 1, + "field_count": 1, + "vector_count": 10, + "dimensions": 4, + } + assert (index_path / pylucene_backend._CAGRA_PROVENANCE_FILE).is_file() + + +def test_build_rejects_persisted_cagra_fallback_and_removes_index(tmp_path): + runtime = _FakeRuntime(writer_path="gpu-cagra") + runtime.verification_error = RuntimeError("persisted brute-force index") + index_path = tmp_path / "index" + + result = _backend(runtime, codec=_CAGRA_CODEC).build( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + force=True, + ) + + assert not result.success + assert "persisted brute-force index" in result.error_message + assert runtime.verification_calls == [(index_path, 10, 4)] + assert len(runtime.build_calls) == 1 + assert not index_path.exists() + + +def test_build_reuses_existing_index_without_runtime(tmp_path): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + runtime = _FakeRuntime() + backend = _backend(runtime) + + result = backend.build(_dataset(), [_index(index_path)]) + + assert result.success + assert result.metadata["skipped"] is True + assert result.metadata["hnsw_verification"]["status"] == ( + "gpu-hnsw-provenance" + ) + assert runtime.resolve_calls == [] + assert runtime.telemetry_calls == [] + assert runtime.build_calls == [] + assert runtime.verification_calls == [] + + +def test_build_reports_reused_index_sizing_failure(tmp_path, monkeypatch): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + runtime = _FakeRuntime() + + def fail_index_size(_index_path): + raise PermissionError("cannot read index size") + + monkeypatch.setattr(pylucene_backend, "_index_size", fail_index_size) + + result = _backend(runtime).build(_dataset(), [_index(index_path)]) + + assert not result.success + assert "cannot read index size" in result.error_message + assert runtime.telemetry_calls == [] + assert runtime.build_calls == [] + + +def test_build_rejects_reused_hnsw_index_without_provenance(tmp_path): + index_path = tmp_path / "index" + index_path.mkdir() + (index_path / "segments_1").write_bytes(b"existing") + runtime = _FakeRuntime() + + result = _backend(runtime).build(_dataset(), [_index(index_path)]) + + assert not result.success + assert "provenance manifest is missing" in result.error_message + assert runtime.telemetry_calls == [] + assert runtime.build_calls == [] + + +def test_build_verifies_reused_cagra_index(tmp_path): + index_path = tmp_path / "index" + _prepare_cagra_index(index_path) + runtime = _FakeRuntime(writer_path="gpu-cagra") + + result = _backend(runtime, codec=_CAGRA_CODEC).build( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + ) + + assert result.success + assert result.metadata["skipped"] is True + assert runtime.verification_calls == [(index_path, 10, 4)] + assert result.metadata["cagra_provenance"]["status"] == ( + "gpu-cagra-provenance" + ) + assert result.metadata["cagra_verification"] == { + "status": "cagra-only", + "segment_count": 1, + "field_count": 1, + "vector_count": 10, + "dimensions": 4, + } + assert runtime.telemetry_calls == [] + assert runtime.build_calls == [] + + +def test_build_rejects_reused_cagra_index_that_cannot_be_verified(tmp_path): + index_path = tmp_path / "index" + _prepare_cagra_index(index_path) + segments_file = index_path / "segments_1" + runtime = _FakeRuntime(writer_path="gpu-cagra") + runtime.verification_error = RuntimeError("persisted brute-force index") + + result = _backend(runtime, codec=_CAGRA_CODEC).build( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + ) + + assert not result.success + assert "persisted brute-force index" in result.error_message + assert runtime.verification_calls == [(index_path, 10, 4)] + assert segments_file.read_bytes() == b"index" + assert runtime.telemetry_calls == [] + assert runtime.build_calls == [] + + +def test_build_rejects_reused_cagra_index_without_provenance_before_runtime( + tmp_path, +): + index_path = tmp_path / "index" + index_path.mkdir() + (index_path / "segments_1").write_bytes(b"existing") + runtime = _FakeRuntime(writer_path="gpu-cagra") + + result = _backend(runtime, codec=_CAGRA_CODEC).build( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + ) + + assert not result.success + assert "CAGRA provenance manifest is missing" in result.error_message + assert runtime.verification_calls == [] + assert runtime.search_calls == [] + + +@pytest.mark.parametrize("force", [False, True]) +def test_build_rejects_unrelated_existing_directory(tmp_path, force): + index_path = tmp_path / "index" + index_path.mkdir() + unrelated_file = index_path / "unrelated" + unrelated_file.write_bytes(b"not a Lucene index") + runtime = _FakeRuntime() + + result = _backend(runtime).build( + _dataset(), [_index(index_path)], force=force + ) + + assert not result.success + assert "does not contain a Lucene segments file" in result.error_message + assert unrelated_file.exists() + assert runtime.telemetry_calls == [] + + +def test_build_rejects_existing_index_path_that_is_a_file(tmp_path): + index_path = tmp_path / "index" + index_path.write_bytes(b"not a directory") + + result = _backend(_FakeRuntime()).build( + _dataset(), [_index(index_path)], force=True + ) + + assert not result.success + assert "not a directory" in result.error_message + assert index_path.read_bytes() == b"not a directory" + + +def test_build_force_replaces_existing_index(tmp_path): + index_path = tmp_path / "index" + index_path.mkdir() + (index_path / "segments_1").write_bytes(b"existing") + old_file = index_path / "old" + old_file.write_bytes(b"old") + runtime = _FakeRuntime() + + result = _backend(runtime).build( + _dataset(), [_index(index_path)], force=True + ) + + assert result.success + assert not old_file.exists() + assert (index_path / "segments_1").exists() + + +def test_build_preserves_existing_index_if_codec_preflight_fails(tmp_path): + index_path = tmp_path / "index" + index_path.mkdir() + (index_path / "segments_1").write_bytes(b"existing") + old_file = index_path / "old" + old_file.write_bytes(b"old") + runtime = _FakeRuntime() + runtime.resolve_error = RuntimeError("codec unavailable") + + result = _backend(runtime).build( + _dataset(), [_index(index_path)], force=True + ) + + assert not result.success + assert "codec unavailable" in result.error_message + assert old_file.exists() + + +def test_build_rejects_silent_cpu_fallback_and_removes_index(tmp_path): + runtime = _FakeRuntime(writer_path="cpu-hnsw-fallback") + index_path = tmp_path / "index" + + result = _backend(runtime).build( + _dataset(), [_index(index_path)], force=True + ) + + assert not result.success + assert "silent CPU" in result.error_message + assert not index_path.exists() + assert runtime.build_calls == [] + + +def test_build_preserves_existing_index_on_cpu_fallback(tmp_path): + runtime = _FakeRuntime(writer_path="cpu-hnsw-fallback") + index_path = tmp_path / "index" + index_path.mkdir() + (index_path / "segments_1").write_bytes(b"existing") + old_file = index_path / "old" + old_file.write_bytes(b"old") + + result = _backend(runtime).build( + _dataset(), [_index(index_path)], force=True + ) + + assert not result.success + assert "silent CPU" in result.error_message + assert old_file.read_bytes() == b"old" + assert runtime.build_calls == [] + + +@pytest.mark.parametrize( + ("dataset", "error"), + [ + (_dataset(dtype=np.float64), "float32"), + (_dataset(distance_metric="cosine"), "Euclidean"), + ], +) +def test_build_rejects_unsupported_dataset(dataset, error, tmp_path): + result = _backend(_FakeRuntime()).build( + dataset, [_index(tmp_path / "index")], force=True + ) + + assert not result.success + assert error in result.error_message + + +def test_build_rejects_nonfinite_vectors(tmp_path): + dataset = _dataset() + dataset.training_vectors[0, 0] = np.nan + + result = _backend(_FakeRuntime()).build( + dataset, [_index(tmp_path / "index")], force=True + ) + + assert not result.success + assert "finite" in result.error_message + + +@pytest.mark.parametrize( + ("codec", "writer_path"), + [ + (_HNSW_CODEC, "gpu-hnsw"), + (_CAGRA_CODEC, "gpu-cagra"), + ], + ids=["hnsw", "cagra"], +) +def test_build_rejects_single_vector_cuvs_bypass(tmp_path, codec, writer_path): + runtime = _FakeRuntime(writer_path=writer_path) + + result = _backend(runtime, codec=codec).build( + _dataset(n_base=1), + [_index(tmp_path / "index", codec=codec)], + force=True, + ) + + assert not result.success + assert "at least two training vectors" in result.error_message + assert "does not invoke cuVS" in result.error_message + assert runtime.telemetry_calls == [] + assert not (tmp_path / "index").exists() + + +def test_build_rejects_unknown_codec(tmp_path): + result = _backend(_FakeRuntime()).build( + _dataset(), + [_index(tmp_path / "index", codec="UnknownCodec")], + force=True, + ) + + assert not result.success + assert "Unsupported PyLucene codec" in result.error_message + + +@pytest.mark.parametrize( + "codec", + [None, "", 0, False], + ids=["none", "empty", "zero", "false"], +) +def test_explicit_invalid_codec_does_not_fall_back_to_backend_config( + tmp_path, codec +): + index = _index(tmp_path / "index") + index.build_param["codec"] = codec + backend = _backend(codec=_HNSW_CODEC) + + build_result = backend.build(_dataset(), [index], dry_run=True) + search_result = backend.search(_dataset(), [index], k=3, dry_run=True) + + for result in (build_result, search_result): + assert not result.success + assert f"Unsupported PyLucene codec {codec!r}" in result.error_message + + +def test_build_rejects_unsupported_parameter(tmp_path): + index = _index(tmp_path / "index") + index.build_param["ignored"] = 1 + + result = _backend(_FakeRuntime()).build(_dataset(), [index]) + + assert not result.success + assert "Unsupported PyLucene build parameter" in result.error_message + + +@pytest.mark.parametrize("index_count", [0, 2]) +def test_build_requires_exactly_one_index(index_count, tmp_path): + indexes = [ + _index(tmp_path / f"index-{index_id}") + for index_id in range(index_count) + ] + + result = _backend(_FakeRuntime()).build(_dataset(), indexes, force=True) + + assert not result.success + assert "exactly one" in result.error_message + + +@pytest.mark.parametrize( + ("vectors", "error"), + [ + (np.empty((0, 4), dtype=np.float32), "at least one vector"), + (np.ones(4, dtype=np.float32), "two-dimensional"), + (np.ones((1, 4097), dtype=np.float32), "4096"), + ], +) +def test_build_rejects_invalid_vector_shape(vectors, error, tmp_path): + dataset = _dataset() + dataset.training_vectors = vectors + + result = _backend(_FakeRuntime()).build( + dataset, [_index(tmp_path / "index")], force=True + ) + + assert not result.success + assert error in result.error_message + + +def test_build_removes_partial_index_after_runtime_failure(tmp_path): + runtime = _FakeRuntime() + runtime.build_error = RuntimeError("Java build failed") + index_path = tmp_path / "index" + + result = _backend(runtime).build( + _dataset(), [_index(index_path)], force=True + ) + + assert not result.success + assert "Java build failed" in result.error_message + assert not index_path.exists() + + +def test_build_reports_chained_runtime_failure(tmp_path): + runtime = _FakeRuntime() + runtime.build_error = RuntimeError("rollback failed") + runtime.build_error.__cause__ = RuntimeError("add failed") + index_path = tmp_path / "index" + + result = _backend(runtime).build( + _dataset(), [_index(index_path)], force=True + ) + + assert not result.success + assert "RuntimeError: rollback failed" in result.error_message + assert "caused by RuntimeError: add failed" in result.error_message + assert not index_path.exists() + + +def test_build_removes_partial_index_and_reraises_interrupt(tmp_path): + runtime = _FakeRuntime() + runtime.build_error = KeyboardInterrupt() + index_path = tmp_path / "index" + + with pytest.raises(KeyboardInterrupt): + _backend(runtime).build(_dataset(), [_index(index_path)], force=True) + + assert not index_path.exists() + + +def test_build_preserves_interrupt_when_partial_index_cleanup_fails( + tmp_path, monkeypatch +): + runtime = _FakeRuntime() + runtime.build_error = KeyboardInterrupt() + index_path = tmp_path / "index" + + def fail_cleanup(_index_path): + raise PermissionError("cleanup denied") + + monkeypatch.setattr(pylucene_backend, "_safe_remove_index", fail_cleanup) + with pytest.raises(KeyboardInterrupt) as exc_info: + _backend(runtime).build(_dataset(), [_index(index_path)], force=True) + + assert exc_info.value.__notes__ == [ + "Failed to remove partial index: PermissionError: cleanup denied" + ] + assert index_path.exists() + + +def test_build_reports_partial_index_cleanup_failure(tmp_path, monkeypatch): + runtime = _FakeRuntime() + runtime.build_error = RuntimeError("Java build failed") + index_path = tmp_path / "index" + + def fail_cleanup(_index_path): + raise PermissionError("cleanup denied") + + monkeypatch.setattr(pylucene_backend, "_safe_remove_index", fail_cleanup) + result = _backend(runtime).build( + _dataset(), [_index(index_path)], force=True + ) + + assert not result.success + assert "Java build failed" in result.error_message + assert "failed to remove partial index" in result.error_message + assert "cleanup denied" in result.error_message + assert index_path.exists() + + +def test_search_dry_run_does_not_initialize_pylucene(tmp_path): + backend = _backend() + + result = backend.search( + _dataset(), [_index(tmp_path / "index")], k=3, dry_run=True + ) + + assert result.success + assert backend._runtime is None + assert result.neighbors.shape == (0, 3) + + +def test_search_converts_hits_scores_and_padding(tmp_path): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + runtime = _FakeRuntime( + hits=[ + [ + _SearchHit(document_id=3, score=1.0), + _SearchHit(document_id=7, score=0.5), + ], + [_SearchHit(document_id=2, score=0.2)], + ] + ) + + result = _backend(runtime).search( + _dataset(), [_index(index_path)], k=3, batch_size=1 + ) + + assert result.success + np.testing.assert_array_equal( + result.neighbors, + np.array([[3, 7, -1], [2, -1, -1]], dtype=np.int64), + ) + np.testing.assert_allclose(result.distances[0, :2], [0.0, 1.0]) + assert result.distances[1, 0] == pytest.approx(4.0) + assert np.isinf(result.distances[:, -1]).all() + assert result.search_time_ms == pytest.approx(2.0) + assert result.latency_seconds == pytest.approx(0.001) + assert result.queries_per_second == pytest.approx(1000.0) + assert result.latency_percentiles == { + "p50": 1.0, + "p95": 1.0, + "p99": 1.0, + } + assert result.metadata["num_batches"] == 2 + assert result.metadata["hnsw_verification"]["status"] == ( + "gpu-hnsw-provenance" + ) + assert "per_search_param_results" not in result.metadata + assert runtime.search_calls[0][3] == 1 + assert runtime.verification_calls == [] + + +def test_search_end_to_end_timing_starts_before_input_verification( + tmp_path, monkeypatch +): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + backend = _backend(_FakeRuntime()) + events = [] + times = iter((10.0, 10.012)) + load_search_inputs = backend._load_search_inputs + + def record_time(): + events.append("time") + return next(times) + + def record_input_loading(dataset): + events.append("inputs") + return load_search_inputs(dataset) + + monkeypatch.setattr(pylucene_backend.time, "perf_counter", record_time) + monkeypatch.setattr(backend, "_load_search_inputs", record_input_loading) + + result = backend.search( + _dataset(), [_index(index_path)], k=3, batch_size=2 + ) + + assert result.success + assert events == ["time", "inputs", "time"] + assert result.metadata["end_to_end_time_ms"] == pytest.approx(12.0) + assert result.metadata["non_query_overhead_time_ms"] == pytest.approx(11.0) + + +@pytest.mark.parametrize( + ("runtime", "error"), + [ + ( + _FakeRuntime(document_count=9), + "document count does not match index provenance", + ), + ( + _FakeRuntime( + hits=[ + [ + _SearchHit(document_id=3, score=1.0), + _SearchHit(document_id=3, score=0.5), + ], + [_SearchHit(document_id=2, score=1.0)], + ] + ), + "duplicate stored ID", + ), + ( + _FakeRuntime( + hits=[ + [_SearchHit(document_id=10, score=1.0)], + [_SearchHit(document_id=2, score=1.0)], + ] + ), + "out-of-range stored ID", + ), + ( + _FakeRuntime(hits=[[_SearchHit(document_id=0, score=1.0)]]), + "unexpected number of query results", + ), + ( + _FakeRuntime( + hits=[ + [_SearchHit(document_id=0, score=float("nan"))], + [_SearchHit(document_id=1, score=1.0)], + ] + ), + "non-finite score", + ), + ( + _FakeRuntime( + hits=[ + [_SearchHit(document_id=0, score=1.1)], + [_SearchHit(document_id=1, score=1.0)], + ] + ), + "outside the Euclidean range", + ), + ( + _FakeRuntime( + hits=[ + [ + _SearchHit(document_id=0, score=1.0), + _SearchHit(document_id=1, score=0.9), + _SearchHit(document_id=2, score=0.8), + _SearchHit(document_id=3, score=0.7), + ], + [_SearchHit(document_id=1, score=1.0)], + ] + ), + "too many hits", + ), + ( + _FakeRuntime(batch_latencies_ms=[1.0, 2.0]), + "unexpected number of batch latencies", + ), + ( + _FakeRuntime(batch_latencies_ms=[float("nan")]), + "invalid batch latency", + ), + ], + ids=[ + "document-count", + "duplicate-id", + "out-of-range-id", + "query-count", + "score", + "score-range", + "hit-count", + "latency-count", + "batch-latency", + ], +) +def test_search_rejects_invalid_runtime_results(tmp_path, runtime, error): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + + result = _backend(runtime).search( + _dataset(), [_index(index_path)], k=3, batch_size=2 + ) + + assert not result.success + assert error in result.error_message + + +def test_search_verifies_cagra_index_before_querying(tmp_path): + index_path = tmp_path / "index" + _prepare_cagra_index(index_path) + runtime = _FakeRuntime(writer_path="gpu-cagra") + + result = _backend(runtime, codec=_CAGRA_CODEC).search( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + k=3, + ) + + assert result.success + assert runtime.verification_calls == [(index_path, 10, 4)] + assert len(runtime.search_calls) == 1 + assert result.metadata["cagra_provenance"]["status"] == ( + "gpu-cagra-provenance" + ) + assert result.metadata["cagra_verification"] == { + "status": "cagra-only", + "segment_count": 1, + "field_count": 1, + "vector_count": 10, + "dimensions": 4, + } + + +def test_search_rejects_unverified_cagra_index_before_querying(tmp_path): + index_path = tmp_path / "index" + _prepare_cagra_index(index_path) + runtime = _FakeRuntime(writer_path="gpu-cagra") + runtime.verification_error = RuntimeError("persisted brute-force index") + + result = _backend(runtime, codec=_CAGRA_CODEC).search( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + k=3, + ) + + assert not result.success + assert "persisted brute-force index" in result.error_message + assert runtime.verification_calls == [(index_path, 10, 4)] + assert runtime.search_calls == [] + + +@pytest.mark.parametrize("manifest_state", ["missing", "stale"]) +def test_search_rejects_invalid_cagra_provenance_before_runtime( + tmp_path, manifest_state +): + index_path = tmp_path / "index" + manifest_path = _prepare_cagra_index(index_path) + if manifest_state == "missing": + manifest_path.unlink() + else: + (index_path / "segments_1").write_bytes(b"changed commit") + runtime = _FakeRuntime(writer_path="gpu-cagra") + + result = _backend(runtime, codec=_CAGRA_CODEC).search( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + k=3, + ) + + assert not result.success + assert "CAGRA provenance" in result.error_message + assert runtime.verification_calls == [] + assert runtime.search_calls == [] + + +@pytest.mark.parametrize( + ("index", "k", "batch_size", "search_threads", "mode", "error"), + [ + ( + _index(Path("/tmp/index")), + 0, + 10000, + None, + "latency", + "k must be positive", + ), + ( + _index(Path("/tmp/index")), + 3, + 10000, + 2, + "latency", + "one search thread", + ), + ( + _index(Path("/tmp/index")), + 3, + 10000, + None, + "throughput", + "only latency mode", + ), + ( + _index(Path("/tmp/index")), + 3, + 0, + None, + "latency", + "batch_size must be positive", + ), + ( + _index(Path("/tmp/index"), search_params=[{"search_width": 16}]), + 3, + 10000, + None, + "latency", + "does not expose", + ), + ( + _index(Path("/tmp/index"), codec=_CAGRA_CODEC), + 1025, + 10000, + None, + "latency", + "GPU brute-force", + ), + ], +) +def test_search_rejects_unsupported_options( + index, k, batch_size, search_threads, mode, error +): + result = _backend(_FakeRuntime()).search( + _dataset(), + [index], + k=k, + batch_size=batch_size, + mode=mode, + search_threads=search_threads, + dry_run=True, + ) + + assert not result.success + assert error in result.error_message + + +def test_resolve_search_plan_normalizes_the_complete_request(tmp_path): + index = _index(tmp_path / "index", search_params=[]) + + plan = _backend()._resolve_search_plan( + index, + k=3, + batch_size=7, + mode="latency", + search_threads="1", + ) + + assert plan == pylucene_backend._SearchPlan( + index_path=tmp_path / "index", + codec_name=_HNSW_CODEC, + search_params=[{}], + k=3, + batch_size=7, + mode="latency", + ) + + +def test_search_validation_preserves_first_error_precedence(tmp_path): + result = _backend().search( + _dataset(), + [_index(tmp_path / "index", search_params=[{"unsupported": True}])], + k=0, + batch_size=0, + mode="throughput", + search_threads=2, + dry_run=True, + ) + + assert not result.success + assert result.error_message == "k must be positive" + + +@pytest.mark.parametrize("index_count", [0, 2]) +def test_search_requires_exactly_one_index(index_count, tmp_path): + indexes = [ + _index(tmp_path / f"index-{index_id}") + for index_id in range(index_count) + ] + + result = _backend(_FakeRuntime()).search(_dataset(), indexes, k=3) + + assert not result.success + assert "exactly one" in result.error_message + + +def test_search_rejects_missing_index(tmp_path): + result = _backend(_FakeRuntime()).search( + _dataset(), [_index(tmp_path / "missing")], k=3 + ) + + assert not result.success + assert "does not exist" in result.error_message + + +def test_search_rejects_empty_query_vectors(tmp_path): + index_path = tmp_path / "index" + index_path.mkdir() + dataset = _dataset() + dataset.query_vectors = np.empty((0, 4), dtype=np.float32) + + result = _backend(_FakeRuntime()).search( + dataset, [_index(index_path)], k=3 + ) + + assert not result.success + assert "at least one vector" in result.error_message + + +def test_search_reports_runtime_query_dimension_mismatch(tmp_path): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + runtime = _FakeRuntime(index_dimensions=8) + + result = _backend(runtime).search( + _dataset(dimensions=4), [_index(index_path)], k=3 + ) + + assert not result.success + assert "dimensions do not match the Lucene index" in result.error_message + + +def test_search_rejects_runtime_dimensions_that_disagree_with_provenance( + tmp_path, +): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path, dimensions=4) + runtime = _FakeRuntime( + index_dimensions=8, + validate_query_dimensions=False, + ) + + result = _backend(runtime).search( + _dataset(dimensions=4), [_index(index_path)], k=3 + ) + + assert not result.success + assert ( + "Lucene index dimensions do not match index provenance" + in result.error_message + ) + assert len(runtime.search_calls) == 1 + + +def test_search_rejects_query_dimensions_that_disagree_with_dataset(tmp_path): + index_path = tmp_path / "index" + index_path.mkdir() + dataset = _dataset(dimensions=4) + dataset.query_vectors = np.zeros((2, 8), dtype=np.float32) + runtime = _FakeRuntime() + + result = _backend(runtime).search(dataset, [_index(index_path)], k=3) + + assert not result.success + assert "Query vector dimensions do not match the dataset" in ( + result.error_message + ) + assert runtime.search_calls == [] + + +def test_search_reports_runtime_failure(tmp_path): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + runtime = _FakeRuntime() + runtime.search_error = RuntimeError("Java search failed") + + result = _backend(runtime).search(_dataset(), [_index(index_path)], k=3) + + assert not result.success + assert "Java search failed" in result.error_message + + +@pytest.mark.parametrize( + "manifest_state", + ["missing", "corrupt", "stale", "wrong-codec"], +) +def test_search_rejects_invalid_hnsw_provenance_before_runtime( + tmp_path, manifest_state +): + index_path = tmp_path / "index" + manifest_path = _prepare_hnsw_index(index_path) + if manifest_state == "missing": + manifest_path.unlink() + elif manifest_state == "corrupt": + manifest_path.write_text("{") + elif manifest_state == "stale": + (index_path / "segments_1").write_bytes(b"changed commit") + else: + payload = json.loads(manifest_path.read_text()) + payload["codec"] = _HNSW_BASE_LAYER_CODEC + manifest_path.write_text(json.dumps(payload)) + + runtime = _FakeRuntime() + result = _backend(runtime).search(_dataset(), [_index(index_path)], k=3) + + assert not result.success + assert "provenance" in result.error_message.lower() + assert runtime.search_calls == [] + + +def test_cleanup_releases_runtime_reference(): + backend = _backend(_FakeRuntime()) + + backend.cleanup() + + assert backend._runtime is None diff --git a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_cagra_verifier.py b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_cagra_verifier.py new file mode 100644 index 0000000000..fa320fdde4 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_cagra_verifier.py @@ -0,0 +1,815 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Persisted CAGRA verifier unit tests.""" + +from __future__ import annotations + +import os +from types import SimpleNamespace + +import pytest + +import cuvs_bench.backends.pylucene as pylucene_backend + + +class _FakeCagraMetadataInput: + def __init__(self, integers, variable_longs): + self._integers = iter(integers) + self._variable_longs = iter(variable_longs) + + def readInt(self): + return next(self._integers) + + def readVLong(self): + return next(self._variable_longs) + + +class _FakeChecksumInput(_FakeCagraMetadataInput): + def __init__(self, integers, variable_longs): + super().__init__(integers, variable_longs) + self.closed = False + + def close(self): + self.closed = True + + +class _FakeIndexInput: + def __init__(self, payload_start=57, payload_length=1000): + self.payload_start = payload_start + self.file_length = payload_start + payload_length + 16 + self.closed = False + + def getFilePointer(self): + return self.payload_start + + def length(self): + return self.file_length + + def close(self): + self.closed = True + + +class _FakeCodecUtil: + def __init__(self, error_at=None): + self.error_at = error_at + self.header_calls = [] + self.footer_calls = [] + self.checksum_calls = [] + self.retrieved_checksum_calls = [] + + def checkIndexHeader(self, *args): + self.header_calls.append(args) + if self.error_at == "header": + raise RuntimeError("unsupported metadata version") + return 0 + + def checkFooter(self, metadata_input): + self.footer_calls.append(metadata_input) + if self.error_at == "footer": + raise RuntimeError("invalid checksum footer") + + def footerLength(self): + return 16 + + def checksumEntireFile(self, data_input): + self.checksum_calls.append(data_input) + if self.error_at == "data-checksum": + raise RuntimeError("invalid data checksum") + return 0 + + def retrieveChecksum(self, data_input): + self.retrieved_checksum_calls.append(data_input) + return 0 + + +class _FakeFieldInfos: + def __init__(self, field_infos): + self._field_infos = { + field_info.number: field_info for field_info in field_infos + } + + def fieldInfo(self, field_number): + return self._field_infos.get(field_number) + + def __iter__(self): + return iter(self._field_infos.values()) + + +def _fake_cagra_index_verifier( + metadata_input, + *, + metadata_files=("_0.vemc",), + codec_util=None, + data_payload_length=1000, + field_name="vector", + field_dimensions=32, + field_updates=False, + max_documents=4, + deletion_count=0, + soft_deletion_count=0, + extra_vector_field=False, + data_error=None, +): + data_input = _FakeIndexInput(payload_length=data_payload_length) + vector_encoding = object() + vector_similarity = object() + field_info = SimpleNamespace( + number=1, + getName=lambda: field_name, + getVectorDimension=lambda: field_dimensions, + getVectorEncoding=lambda: vector_encoding, + getVectorSimilarityFunction=lambda: vector_similarity, + ) + all_field_infos = [field_info] + if extra_vector_field: + all_field_infos.append( + SimpleNamespace( + number=2, + getName=lambda: "extra-vector", + getVectorDimension=lambda: field_dimensions, + getVectorEncoding=lambda: vector_encoding, + getVectorSimilarityFunction=lambda: vector_similarity, + ) + ) + field_infos = _FakeFieldInfos(all_field_infos) + field_infos_format = SimpleNamespace( + read=lambda _directory, _info, _suffix, _context: field_infos + ) + + def open_data(_file_name, _context): + if data_error is not None: + raise data_error + return data_input + + directory = SimpleNamespace( + openChecksumInput=lambda _file_name: metadata_input, + openInput=open_data, + close=lambda: None, + ) + segment_info = SimpleNamespace( + info=SimpleNamespace( + name="_0", + getId=lambda: b"segment-id", + getCodec=lambda: SimpleNamespace( + fieldInfosFormat=lambda: field_infos_format + ), + maxDoc=lambda: max_documents, + ), + files=lambda: metadata_files, + hasFieldUpdates=lambda: field_updates, + hasDeletions=lambda: deletion_count > 0, + getDelCount=lambda: deletion_count, + getSoftDelCount=lambda: soft_deletion_count, + ) + verifier = pylucene_backend._CagraIndexVerifier( + attach_current_thread=lambda: None, + paths=SimpleNamespace(get=lambda path: path), + codec_util=codec_util or _FakeCodecUtil(), + field_info=SimpleNamespace( + cast_=lambda raw_field_info: raw_field_info + ), + segment_commit_info=SimpleNamespace( + cast_=lambda raw_segment_info: raw_segment_info + ), + segment_infos=SimpleNamespace( + readLatestCommit=lambda _directory: [segment_info] + ), + vector_encoding=SimpleNamespace(FLOAT32=vector_encoding), + vector_similarity_function=SimpleNamespace( + EUCLIDEAN=vector_similarity + ), + fs_directory=SimpleNamespace(open=lambda _path: directory), + io_context=SimpleNamespace(READONCE=object()), + ) + verifier._test_data_input = data_input + verifier._test_directory = directory + verifier._test_segment_info = segment_info + return verifier + + +def _cagra_data_context(verifier, index_path): + return pylucene_backend._CagraDataFileContext( + index_path=index_path, + directory=verifier._test_directory, + segment_info=verifier._test_segment_info, + suffix="", + metadata_file="_0.vemc", + ) + + +@pytest.mark.parametrize( + ("metadata_file", "expected_suffix"), + [("_0.vemc", ""), ("_0_CuVS_0.vemc", "CuVS_0")], +) +def test_cagra_segment_suffix(metadata_file, expected_suffix): + assert ( + pylucene_backend._CagraIndexVerifier._segment_suffix( + "_0", metadata_file + ) + == expected_suffix + ) + + +def test_cagra_segment_suffix_rejects_unrelated_metadata_file(): + with pytest.raises(RuntimeError, match="does not match segment"): + pylucene_backend._CagraIndexVerifier._segment_suffix("_0", "_1.vemc") + + +def test_read_cagra_fields_accepts_only_persisted_cagra_data(): + metadata_input = _FakeCagraMetadataInput( + integers=[ + 1, + 1, + 0, + 32, + 4, + 2, + 1, + 0, + 32, + 0, + -1, + ], + variable_longs=[57, 1000, 1057, 0, 0, 0, 0, 0], + ) + + assert pylucene_backend._CagraIndexVerifier._read_cagra_fields( + metadata_input, "_0.vemc" + ) == [ + pylucene_backend._CagraFieldMetadata( + field_number=1, + dimensions=32, + vector_count=4, + cagra_offset=57, + cagra_length=1000, + ) + ] + + +@pytest.mark.parametrize( + ("cagra_length", "brute_force_length", "error"), + [ + (0, 512, "persisted brute-force"), + (1000, 512, "persisted brute-force"), + ], +) +def test_read_cagra_fields_rejects_non_cagra_only_data( + cagra_length, brute_force_length, error +): + metadata_input = _FakeCagraMetadataInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[ + 57, + cagra_length, + 57 + cagra_length, + brute_force_length, + ], + ) + + with pytest.raises(RuntimeError, match=error): + pylucene_backend._CagraIndexVerifier._read_cagra_fields( + metadata_input, "_0.vemc" + ) + + +def test_read_cagra_fields_rejects_data_for_empty_field(): + metadata_input = _FakeCagraMetadataInput( + integers=[1, 1, 0, 32, 0, -1], + variable_longs=[57, 1, 58, 0], + ) + + with pytest.raises(RuntimeError, match="empty field"): + pylucene_backend._CagraIndexVerifier._read_cagra_fields( + metadata_input, "_0.vemc" + ) + + +@pytest.mark.parametrize( + ("encoding", "similarity", "dimensions", "error"), + [ + (0, 0, 32, "encoding ordinal 0"), + (1, 1, 32, "similarity ordinal 1"), + (1, 0, 0, "invalid field metadata"), + (1, 0, 4097, "invalid field metadata"), + ], +) +def test_read_cagra_fields_rejects_unsupported_vector_semantics( + encoding, similarity, dimensions, error +): + metadata_input = _FakeCagraMetadataInput( + integers=[1, encoding, similarity, dimensions, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ) + + with pytest.raises(RuntimeError, match=error): + pylucene_backend._CagraIndexVerifier._read_cagra_fields( + metadata_input, "_0.vemc" + ) + + +def test_verify_cagra_index_validates_header_footer_and_expected_count( + tmp_path, +): + metadata_input = _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ) + codec_util = _FakeCodecUtil() + verifier = _fake_cagra_index_verifier( + metadata_input, codec_util=codec_util + ) + + verification = verifier.verify_index( + tmp_path, expected_vector_count=4, expected_dimensions=32 + ) + + assert verification.to_metadata() == { + "status": "cagra-only", + "segment_count": 1, + "field_count": 1, + "vector_count": 4, + "dimensions": 32, + } + assert codec_util.header_calls == [ + ( + metadata_input, + "Lucene102CuVSVectorsFormatMeta", + 0, + 0, + b"segment-id", + "", + ), + ( + verifier._test_data_input, + "Lucene102CuVSVectorsFormatIndex", + 0, + 0, + b"segment-id", + "", + ), + ] + assert codec_util.footer_calls == [metadata_input] + assert codec_util.checksum_calls == [verifier._test_data_input] + assert metadata_input.closed is True + assert verifier._test_data_input.closed is True + + +def test_verify_cagra_index_traverses_segments_and_closes_each_input( + tmp_path, +): + events = [] + metadata_inputs = { + name: _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ) + for name in ("_0.vemc", "_1.vemc") + } + data_inputs = {name: _FakeIndexInput() for name in ("_0.vcag", "_1.vcag")} + verifier = _fake_cagra_index_verifier(metadata_inputs["_0.vemc"]) + + def segment(name): + field_info = SimpleNamespace( + number=1, + getName=lambda: "vector", + getVectorDimension=lambda: 32, + getVectorEncoding=lambda: verifier.VectorEncoding.FLOAT32, + getVectorSimilarityFunction=( + lambda: verifier.VectorSimilarityFunction.EUCLIDEAN + ), + ) + field_infos_format = SimpleNamespace( + read=lambda _directory, _info, _suffix, _context: ( + _FakeFieldInfos([field_info]) + ) + ) + return SimpleNamespace( + info=SimpleNamespace( + name=name, + getId=lambda: name.encode(), + getCodec=lambda: SimpleNamespace( + fieldInfosFormat=lambda: field_infos_format + ), + maxDoc=lambda: 4, + ), + files=lambda: (f"{name}.vemc",), + hasFieldUpdates=lambda: False, + hasDeletions=lambda: False, + getDelCount=lambda: 0, + getSoftDelCount=lambda: 0, + ) + + def open_metadata(file_name): + events.append(f"open {file_name}") + metadata_input = metadata_inputs[file_name] + + def close_metadata(): + metadata_input.closed = True + events.append(f"close {file_name}") + + metadata_input.close = close_metadata + return metadata_input + + def open_data(file_name, _context): + events.append(f"open {file_name}") + data_input = data_inputs[file_name] + + def close_data(): + data_input.closed = True + events.append(f"close {file_name}") + + data_input.close = close_data + return data_input + + verifier.SegmentInfos.readLatestCommit = lambda _directory: [ + segment("_0"), + segment("_1"), + ] + verifier._test_directory.openChecksumInput = open_metadata + verifier._test_directory.openInput = open_data + verifier._test_directory.close = lambda: events.append("close directory") + + verification = verifier.verify_index(tmp_path) + + assert verification == pylucene_backend._CagraIndexVerification( + segment_count=2, + field_count=2, + vector_count=8, + dimensions=32, + ) + assert events == [ + "open _0.vemc", + "close _0.vemc", + "open _0.vcag", + "close _0.vcag", + "open _1.vemc", + "close _1.vemc", + "open _1.vcag", + "close _1.vcag", + "close directory", + ] + assert all(item.closed for item in metadata_inputs.values()) + assert all(item.closed for item in data_inputs.values()) + + +@pytest.mark.parametrize("error_at", ["header", "footer"]) +def test_verify_cagra_index_fails_closed_on_invalid_format(error_at, tmp_path): + metadata_input = _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ) + verifier = _fake_cagra_index_verifier( + metadata_input, codec_util=_FakeCodecUtil(error_at=error_at) + ) + + with pytest.raises(RuntimeError, match="metadata format v0"): + verifier.verify_index(tmp_path) + + assert metadata_input.closed is True + + +def test_verify_cagra_index_rejects_missing_segment_metadata(tmp_path): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput([], []), metadata_files=() + ) + + with pytest.raises(RuntimeError, match="no .vemc metadata"): + verifier.verify_index(tmp_path) + + +def test_verify_cagra_index_rejects_only_empty_fields(tmp_path): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 0, -1], + variable_longs=[0, 0, 0, 0], + ), + data_payload_length=0, + ) + + with pytest.raises(RuntimeError, match="without matching CAGRA metadata"): + verifier.verify_index(tmp_path) + + +def test_verify_cagra_index_rejects_vector_count_mismatch(tmp_path): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ) + ) + + with pytest.raises(RuntimeError, match="4 vectors; expected 5"): + verifier.verify_index(tmp_path, expected_vector_count=5) + + +def test_verify_cagra_index_rejects_dimension_mismatch(tmp_path): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ) + ) + + with pytest.raises(RuntimeError, match="32 dimensions; expected 16"): + verifier.verify_index(tmp_path, expected_dimensions=16) + + +@pytest.mark.parametrize( + ("runtime_kwargs", "expected_counts"), + [ + ({"deletion_count": 1}, "deleted=1, soft_deleted=0"), + ({"soft_deletion_count": 1}, "deleted=0, soft_deleted=1"), + ], + ids=["hard-deletion", "soft-deletion"], +) +def test_verify_cagra_index_rejects_committed_deletions( + tmp_path, runtime_kwargs, expected_counts +): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ), + **runtime_kwargs, + ) + + with pytest.raises(RuntimeError, match=expected_counts): + verifier.verify_index(tmp_path) + + +def test_verify_cagra_index_rejects_segment_document_count_mismatch(tmp_path): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ), + max_documents=5, + ) + + with pytest.raises(RuntimeError, match="4 vectors for 5 documents"): + verifier.verify_index(tmp_path) + + +def test_verify_cagra_index_rejects_unaccounted_vector_field(tmp_path): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ), + extra_vector_field=True, + ) + + with pytest.raises(RuntimeError, match=r"metadata=\[1\], Lucene=\[1, 2\]"): + verifier.verify_index(tmp_path) + + +def test_verify_cagra_index_rejects_duplicate_field_across_metadata_files( + tmp_path, +): + metadata_inputs = { + name: _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ) + for name in ("_0.vemc", "_0_CuVS_0.vemc") + } + verifier = _fake_cagra_index_verifier( + metadata_inputs["_0.vemc"], + metadata_files=tuple(metadata_inputs), + ) + verifier._test_directory.openChecksumInput = metadata_inputs.__getitem__ + verifier._test_directory.openInput = ( + lambda _file_name, _context: _FakeIndexInput() + ) + + with pytest.raises(RuntimeError, match="duplicate field 1"): + verifier.verify_index(tmp_path) + + +def test_verify_cagra_index_rejects_inconsistent_dimensions_across_segments( + tmp_path, +): + metadata_inputs = { + "_0.vemc": _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ), + "_1.vemc": _FakeChecksumInput( + integers=[1, 1, 0, 16, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ), + } + verifier = _fake_cagra_index_verifier(metadata_inputs["_0.vemc"]) + + def segment(name, dimensions): + field_info = SimpleNamespace( + number=1, + getName=lambda: "vector", + getVectorDimension=lambda: dimensions, + getVectorEncoding=lambda: verifier.VectorEncoding.FLOAT32, + getVectorSimilarityFunction=( + lambda: verifier.VectorSimilarityFunction.EUCLIDEAN + ), + ) + field_infos = _FakeFieldInfos([field_info]) + field_infos_format = SimpleNamespace( + read=lambda _directory, _info, _suffix, _context: field_infos + ) + return SimpleNamespace( + info=SimpleNamespace( + name=name, + getId=lambda: name.encode(), + getCodec=lambda: SimpleNamespace( + fieldInfosFormat=lambda: field_infos_format + ), + maxDoc=lambda: 4, + ), + files=lambda: (f"{name}.vemc",), + hasFieldUpdates=lambda: False, + hasDeletions=lambda: False, + getDelCount=lambda: 0, + getSoftDelCount=lambda: 0, + ) + + segments = [segment("_0", 32), segment("_1", 16)] + verifier.SegmentInfos.readLatestCommit = lambda _directory: segments + verifier._test_directory.openChecksumInput = metadata_inputs.__getitem__ + verifier._test_directory.openInput = ( + lambda _file_name, _context: _FakeIndexInput() + ) + + with pytest.raises(RuntimeError, match=r"dimensions: \[16, 32\]"): + verifier.verify_index(tmp_path) + + +@pytest.mark.parametrize( + ("runtime_kwargs", "error"), + [ + ({"field_name": "other"}, "unexpected field"), + ({"field_dimensions": 16}, "field metadata"), + ({"field_updates": True}, "field-info updates"), + ({"data_payload_length": 999}, "do not exactly cover"), + ], +) +def test_verify_cagra_index_rejects_foreign_or_inconsistent_index( + tmp_path, runtime_kwargs, error +): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ), + **runtime_kwargs, + ) + + with pytest.raises(RuntimeError, match=error): + verifier.verify_index(tmp_path) + + +def test_verify_cagra_index_rejects_corrupt_data_file(tmp_path): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ), + codec_util=_FakeCodecUtil(error_at="data-checksum"), + ) + + with pytest.raises(RuntimeError, match="invalid data checksum"): + verifier.verify_index(tmp_path) + + assert verifier._test_data_input.closed is True + + +def test_verify_cagra_index_rejects_missing_data_file(tmp_path): + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput( + integers=[1, 1, 0, 32, 4, -1], + variable_longs=[57, 1000, 1057, 0], + ), + data_error=FileNotFoundError("_0.vcag"), + ) + + with pytest.raises(RuntimeError, match="cannot read '_0.vcag'"): + verifier.verify_index(tmp_path) + + +def test_file_signature_exposes_named_stat_fields(tmp_path): + data_path = tmp_path / "data.vcag" + data_path.write_bytes(b"data") + file_stat = data_path.stat() + + signature = pylucene_backend._file_signature(data_path) + + assert signature == pylucene_backend._FileSignature( + resolved_path=str(data_path.resolve()), + device=file_stat.st_dev, + inode=file_stat.st_ino, + size=file_stat.st_size, + modified_at_ns=file_stat.st_mtime_ns, + changed_at_ns=file_stat.st_ctime_ns, + ) + + +def test_cagra_data_checksum_is_cached_for_unchanged_file( + tmp_path, monkeypatch +): + codec_util = _FakeCodecUtil() + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput([], []), codec_util=codec_util + ) + (tmp_path / "_0.vcag").write_bytes(b"data") + data_ctime_ns = (tmp_path / "_0.vcag").stat().st_ctime_ns + monkeypatch.setattr( + pylucene_backend.time, + "time_ns", + lambda: ( + data_ctime_ns + pylucene_backend._CAGRA_CACHE_MIN_FILE_AGE_NS + 1 + ), + ) + fields = [ + pylucene_backend._CagraFieldMetadata( + field_number=1, + dimensions=32, + vector_count=4, + cagra_offset=57, + cagra_length=1000, + ) + ] + context = _cagra_data_context(verifier, tmp_path) + + for _ in range(2): + verifier._verify_cagra_data_file(context, fields) + + assert codec_util.checksum_calls == [verifier._test_data_input] + assert codec_util.retrieved_checksum_calls == [verifier._test_data_input] + + +def test_cagra_data_checksum_cache_invalidates_on_size_change(tmp_path): + codec_util = _FakeCodecUtil() + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput([], []), codec_util=codec_util + ) + data_path = tmp_path / "_0.vcag" + data_path.write_bytes(b"data") + fields = [ + pylucene_backend._CagraFieldMetadata( + field_number=1, + dimensions=32, + vector_count=4, + cagra_offset=57, + cagra_length=1000, + ) + ] + context = _cagra_data_context(verifier, tmp_path) + + verifier._verify_cagra_data_file(context, fields) + data_path.write_bytes(b"changed-size") + verifier._verify_cagra_data_file(context, fields) + + assert codec_util.checksum_calls == [ + verifier._test_data_input, + verifier._test_data_input, + ] + assert codec_util.retrieved_checksum_calls == [] + + +def test_cagra_data_checksum_cache_invalidates_same_size_restored_mtime( + tmp_path, +): + codec_util = _FakeCodecUtil() + verifier = _fake_cagra_index_verifier( + _FakeChecksumInput([], []), codec_util=codec_util + ) + data_path = tmp_path / "_0.vcag" + data_path.write_bytes(b"data") + fields = [ + pylucene_backend._CagraFieldMetadata( + field_number=1, + dimensions=32, + vector_count=4, + cagra_offset=57, + cagra_length=1000, + ) + ] + context = _cagra_data_context(verifier, tmp_path) + + verifier._verify_cagra_data_file(context, fields) + original_stat = data_path.stat() + data_path.write_bytes(b"evil") + os.utime( + data_path, + ns=(original_stat.st_atime_ns, original_stat.st_mtime_ns), + ) + assert data_path.stat().st_mtime_ns == original_stat.st_mtime_ns + + verifier._verify_cagra_data_file(context, fields) + + assert codec_util.checksum_calls == [ + verifier._test_data_input, + verifier._test_data_input, + ] + assert codec_util.retrieved_checksum_calls == [] diff --git a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_cli_config.py b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_cli_config.py new file mode 100644 index 0000000000..e3c2ddedc0 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_cli_config.py @@ -0,0 +1,899 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""CLI and configuration tests for the PyLucene backend.""" + +from __future__ import annotations + +import csv +import json +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pytest +from click.testing import CliRunner + +import cuvs_bench.backends.pylucene as pylucene_backend +from cuvs_bench.backends import CppGoogleBenchmarkBackend, get_registry +from cuvs_bench.backends.pylucene import ( + PyLuceneBackend, + PyLuceneConfigLoader, + _SearchHit, +) +from cuvs_bench.backends.registry import list_config_loaders +from cuvs_bench.tests._pylucene_test_utils import ( + _CAGRA_CODEC, + _HNSW_CODEC, + _FakeRuntime, + _write_test_bin, +) + + +def _prepare_cli_dataset(dataset_path: Path) -> None: + training_vectors = np.asarray( + [ + [0.0, 0.0, 0.0, 0.0], + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + ], + dtype=np.float32, + ) + query_vectors = training_vectors[:2].copy() + groundtruth_neighbors = np.asarray([[0], [1]], dtype=np.int32) + dataset_dir = dataset_path / "test-dataset" + _write_test_bin(dataset_dir / "base.fbin", training_vectors) + _write_test_bin(dataset_dir / "query.fbin", query_vectors) + _write_test_bin( + dataset_dir / "groundtruth.neighbors.ibin", groundtruth_neighbors + ) + + +def _pylucene_cli_args( + backend_config: Path, + config_dir: Path, + dataset_path: Path, + *, + dry_run: bool = False, +) -> list[str]: + args = [ + "--backend-config", + str(backend_config), + "--dataset-configuration", + str(config_dir / "datasets" / "datasets.yaml"), + "--configuration", + str(config_dir / "algos" / "pylucene_test.yaml"), + "--dataset", + "test-dataset", + "--dataset-path", + str(dataset_path), + "--algorithms", + "pylucene_test", + "--groups", + "test", + "--batch-size", + "2", + "-k", + "1", + "-m", + "latency", + "--build", + "--search", + "--force", + ] + if dry_run: + args.append("--dry-run") + return args + + +@pytest.fixture +def config_dir(tmp_path): + (tmp_path / "datasets").mkdir() + (tmp_path / "datasets" / "datasets.yaml").write_text( + """\ +- name: test-dataset + base_file: test-dataset/base.fbin + query_file: test-dataset/query.fbin + groundtruth_neighbors_file: test-dataset/groundtruth.neighbors.ibin + distance: euclidean + dims: 4 +""" + ) + (tmp_path / "algos").mkdir() + (tmp_path / "algos" / "pylucene_test.yaml").write_text( + f"""\ +name: pylucene_test +groups: + base: + build: + codec: [{_HNSW_CODEC}, {_CAGRA_CODEC}] + search: {{}} + test: + build: + codec: [{_HNSW_CODEC}] + search: {{}} +""" + ) + (tmp_path / "algos" / "unrelated.yaml").write_text( + """\ +name: unrelated +groups: + base: + build: {} + search: {} +""" + ) + return tmp_path + + +def test_backend_and_loader_are_registered(): + assert get_registry().is_registered("pylucene") + assert list_config_loaders()["pylucene"] is PyLuceneConfigLoader + + +def test_result_file_ownership_is_explicit(): + assert PyLuceneBackend.orchestrator_persists_results is True + assert CppGoogleBenchmarkBackend.orchestrator_persists_results is False + + +def test_import_does_not_load_pylucene(): + subprocess.run( + [ + sys.executable, + "-c", + ( + "import sys; import cuvs_bench.backends; " + "assert 'lucene' not in sys.modules" + ), + ], + check=True, + ) + + +def test_cli_backend_config_supports_pylucene_dry_run(config_dir, tmp_path): + from cuvs_bench.run.__main__ import main as run_main + + backend_config = tmp_path / "backend.yaml" + backend_config.write_text("backend: pylucene\n") + dataset_path = tmp_path / "runtime-datasets" + result_path = dataset_path / "test-dataset" / "result" + assert not result_path.exists() + + result = CliRunner().invoke( + run_main, + _pylucene_cli_args( + backend_config, + config_dir, + dataset_path, + dry_run=True, + ), + ) + + assert result.exit_code == 0, result.output + assert "Would build PyLucene index" in result.output + assert "Would search PyLucene index" in result.output + assert not result_path.exists() + + +def test_cli_build_search_persists_metrics_and_build_join( + config_dir, tmp_path, monkeypatch +): + from cuvs_bench.run.__main__ import main as run_main + + backend_config = tmp_path / "backend.yaml" + backend_config.write_text("backend: pylucene\n") + dataset_path = tmp_path / "runtime-datasets" + _prepare_cli_dataset(dataset_path) + result_path = dataset_path / "test-dataset" / "result" + assert not result_path.exists() + + runtime = _FakeRuntime( + hits=[ + [_SearchHit(document_id=0, score=1.0)], + [_SearchHit(document_id=1, score=1.0)], + ] + ) + monkeypatch.setattr( + pylucene_backend._PyLuceneRuntime, + "create", + staticmethod(lambda _config: runtime), + ) + + result = CliRunner().invoke( + run_main, + _pylucene_cli_args(backend_config, config_dir, dataset_path), + ) + + assert result.exit_code == 0, result.output + index_name = f"pylucene_test_test.codec{_HNSW_CODEC}" + index_path = dataset_path / "test-dataset" / "index" / index_name + assert (index_path / "segments_1").is_file() + assert (index_path / pylucene_backend._HNSW_PROVENANCE_FILE).is_file() + expected_index_size = sum( + path.stat().st_size for path in index_path.rglob("*") if path.is_file() + ) + + build_json = result_path / "build" / "pylucene_test,test.json" + search_json = result_path / "search" / "pylucene_test,test,k1,bs2.json" + build_rows = json.loads(build_json.read_text())["benchmarks"] + search_rows = json.loads(search_json.read_text())["benchmarks"] + assert len(build_rows) == 1 + assert len(search_rows) == 1 + + build_row = build_rows[0] + assert build_row["name"] == f"{index_name}/build" + assert build_row["real_time"] > 0.0 + assert build_row["index_size"] == expected_index_size + assert build_row["codec"] == _HNSW_CODEC + assert build_row["writer_path"] == "gpu-hnsw" + + search_row = search_rows[0] + assert search_row["name"] == f"{index_name}/search" + assert search_row["Recall"] == pytest.approx(1.0) + assert search_row["items_per_second"] == pytest.approx(2000.0) + assert search_row["search_time_ms"] == pytest.approx(1.0) + assert search_row["real_time"] == pytest.approx(1.0) + assert search_row["Latency"] == pytest.approx(0.001) + assert search_row["p50"] == pytest.approx(1.0) + + raw_csv = result_path / "search" / "pylucene_test,test,k1,bs2,raw.csv" + with raw_csv.open(newline="") as file: + csv_rows = list(csv.DictReader(file)) + assert len(csv_rows) == 1 + csv_row = csv_rows[0] + assert csv_row["index_name"] == index_name + assert float(csv_row["recall"]) == pytest.approx(1.0) + assert float(csv_row["throughput"]) == pytest.approx(2000.0) + assert float(csv_row["latency"]) == pytest.approx(0.001) + assert float(csv_row["build time"]) > 0.0 + assert (raw_csv.parent / "pylucene_test,test,k1,bs2,latency.csv").is_file() + assert ( + raw_csv.parent / "pylucene_test,test,k1,bs2,throughput.csv" + ).is_file() + + +def test_cli_returns_nonzero_and_persists_build_failure( + config_dir, tmp_path, monkeypatch +): + from cuvs_bench.run.__main__ import main as run_main + + backend_config = tmp_path / "backend.yaml" + backend_config.write_text("backend: pylucene\n") + dataset_path = tmp_path / "runtime-datasets" + _prepare_cli_dataset(dataset_path) + runtime = _FakeRuntime() + runtime.build_error = RuntimeError("intentional build failure") + monkeypatch.setattr( + pylucene_backend._PyLuceneRuntime, + "create", + staticmethod(lambda _config: runtime), + ) + + result = CliRunner().invoke( + run_main, + _pylucene_cli_args(backend_config, config_dir, dataset_path), + ) + + assert result.exit_code != 0 + assert "intentional build failure" in result.output + build_json = ( + dataset_path + / "test-dataset" + / "result" + / "build" + / "pylucene_test,test.json" + ) + build_rows = json.loads(build_json.read_text())["benchmarks"] + assert len(build_rows) == 1 + assert build_rows[0]["success"] is False + assert "intentional build failure" in build_rows[0]["error_message"] + assert not (dataset_path / "test-dataset" / "result" / "search").exists() + + +def test_config_loader_expands_codecs_and_forwards_runtime_config(config_dir): + loader = PyLuceneConfigLoader(config_path=config_dir) + dataset_config, configs = loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_test", + groups="base", + cuvs_java_jar="/artifacts/cuvs-java.jar", + cuvs_lucene_jar="/artifacts/cuvs-lucene.jar", + java_library_path="/native", + jvm_args=["-Xms1g"], + ) + + assert dataset_config.distance == "euclidean" + assert len(configs) == 2 + assert {config.indexes[0].build_param["codec"] for config in configs} == { + _HNSW_CODEC, + _CAGRA_CODEC, + } + for config in configs: + assert config.indexes[0].search_params == [{}] + assert config.backend_config["cuvs_java_jar"].endswith("cuvs-java.jar") + assert config.backend_config["cuvs_lucene_jar"].endswith( + "cuvs-lucene.jar" + ) + assert config.backend_config["java_library_path"] == "/native" + assert config.backend_config["jvm_args"] == ["-Xms1g"] + assert config.backend_config["requires_gpu"] is True + assert config.output_filename == ( + "pylucene_test,base", + "pylucene_test,base,k10,bs10000", + ) + + +@pytest.mark.parametrize( + ("build_yaml", "search_yaml", "error"), + [ + ( + f"codec: [{_HNSW_CODEC}]\n ignored: [1]", + "{}", + "Unsupported PyLucene build parameter.*ignored", + ), + ( + f"codec: [{_HNSW_CODEC}]", + "ignored: [1]", + "does not expose cuVS-specific search parameters", + ), + ], +) +def test_config_loader_rejects_unsupported_parameters( + config_dir, tmp_path, build_yaml, search_yaml, error +): + custom_config = tmp_path / "unsupported-parameters.yaml" + custom_config.write_text( + f"""\ +backend: pylucene +name: pylucene_unsupported +groups: + test: + build: + {build_yaml} + search: + {search_yaml} +""" + ) + + with pytest.raises(ValueError, match=error): + PyLuceneConfigLoader(config_path=config_dir).load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_unsupported", + groups="test", + algorithm_configuration=str(custom_config), + ) + + +def test_config_loader_scopes_artifact_identity_by_dataset_subset(config_dir): + loader = PyLuceneConfigLoader(config_path=config_dir) + identities = [] + + for subset_size, subset_scope in ( + (None, None), + (2, "subset2"), + (3, "subset3"), + ): + dataset_config, configs = loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_test", + groups="test", + subset_size=subset_size, + count=1, + batch_size=2, + ) + + assert dataset_config.subset_size == subset_size + assert len(configs) == 1 + config = configs[0] + index_prefix = "pylucene_test_test" + result_stem = "pylucene_test,test" + if subset_scope is not None: + index_prefix = f"{index_prefix}.{subset_scope}" + result_stem = f"{result_stem},{subset_scope}" + index_name = f"{index_prefix}.codec{_HNSW_CODEC}" + assert config.index_name == index_name + assert config.index_path == ( + Path("/datasets/test-dataset/index") / index_name + ) + assert config.output_filename == ( + result_stem, + f"{result_stem},k1,bs2", + ) + assert all(Path(stem).name == stem for stem in config.output_filename) + identities.append( + (config.index_name, config.index_path, config.output_filename) + ) + + assert len(set(identities)) == 3 + + +@pytest.mark.parametrize("subset_size", [0, -1, True, "../other"]) +def test_config_loader_rejects_unsafe_subset_identity(config_dir, subset_size): + loader = PyLuceneConfigLoader(config_path=config_dir) + + with pytest.raises(ValueError, match="positive integer"): + loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_test", + groups="test", + subset_size=subset_size, + ) + + +@pytest.mark.parametrize( + ("backend_declaration", "algorithm_name"), + [ + ("", "pylucene_custom"), + ("backend: pylucene\n", "custom_lucene"), + ], + ids=["parsed-name", "declared-backend"], +) +def test_config_loader_discovers_custom_filename_by_parsed_identity( + config_dir, + tmp_path, + backend_declaration, + algorithm_name, +): + custom_config = tmp_path / "arbitrary-custom-name.yaml" + custom_config.write_text( + f"""\ +{backend_declaration}name: {algorithm_name} +groups: + test: + build: + codec: [{_HNSW_CODEC}] + search: {{}} +""" + ) + loader = PyLuceneConfigLoader(config_path=config_dir) + + _, configs = loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms=algorithm_name, + groups="test", + algorithm_configuration=str(custom_config), + ) + + assert len(configs) == 1 + assert configs[0].indexes[0].algo == algorithm_name + assert configs[0].index_name.startswith(f"{algorithm_name}_test.") + + +def test_config_loader_excludes_explicit_other_backend(config_dir, tmp_path): + custom_config = tmp_path / "misleading-pylucene-name.yaml" + custom_config.write_text( + f"""\ +backend: cpp_gbench +name: pylucene_other_backend +groups: + test: + build: + codec: [{_HNSW_CODEC}] + search: {{}} +""" + ) + loader = PyLuceneConfigLoader(config_path=config_dir) + + with pytest.raises( + ValueError, + match="Unknown PyLucene algorithm selector.*pylucene_other_backend", + ): + loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_other_backend", + groups="test", + algorithm_configuration=str(custom_config), + ) + + +def test_config_loader_honors_algorithm_and_group_filters(config_dir): + loader = PyLuceneConfigLoader(config_path=config_dir) + + _, configs = loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_test", + groups="test", + ) + assert len(configs) == 1 + assert configs[0].index_name.startswith("pylucene_test_test") + + +def test_algorithm_config_discovery_is_deterministic(tmp_path): + config_path = tmp_path / "config" + bundled_algorithms = config_path / "algos" + bundled_algorithms.mkdir(parents=True) + bundled_zeta = bundled_algorithms / "zeta.yaml" + bundled_alpha = bundled_algorithms / "alpha.yml" + bundled_zeta.touch() + bundled_alpha.touch() + (bundled_algorithms / "ignored.txt").touch() + + custom_algorithms = tmp_path / "custom" + custom_algorithms.mkdir() + custom_zeta = custom_algorithms / "zeta.yml" + custom_alpha = custom_algorithms / "alpha.yaml" + custom_zeta.touch() + custom_alpha.touch() + + files = PyLuceneConfigLoader( + config_path=config_path + ).gather_algorithm_configs(config_path, str(custom_algorithms)) + + assert files == [ + str(bundled_alpha), + str(bundled_zeta), + str(custom_alpha), + str(custom_zeta), + ] + + +def test_duplicate_algorithm_config_uses_last_definition_and_position( + config_dir, monkeypatch +): + loader = PyLuceneConfigLoader(config_path=config_dir) + config_files = ["alpha-original", "beta", "alpha-override"] + configs_by_file = { + "alpha-original": { + "name": "pylucene_alpha", + "groups": {"original": {}}, + }, + "beta": { + "name": "pylucene_beta", + "groups": {"base": {}}, + }, + "alpha-override": { + "name": "pylucene_alpha", + "groups": {"override": {}}, + }, + } + monkeypatch.setattr(loader, "load_yaml_file", configs_by_file.__getitem__) + + algorithm_configs = loader._load_algorithm_configs(config_files) + + assert list(algorithm_configs) == ["pylucene_beta", "pylucene_alpha"] + assert algorithm_configs["pylucene_alpha"] == {"override": {}} + + +def test_config_loader_unions_global_and_algorithm_specific_groups(config_dir): + loader = PyLuceneConfigLoader(config_path=config_dir) + + _, configs = loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algorithms="pylucene_test", + groups="base", + algo_groups="pylucene_test.test", + ) + + assert len(configs) == 3 + assert ( + sum( + config.index_name.startswith("pylucene_test_test.") + for config in configs + ) + == 1 + ) + + +def test_config_loader_selects_only_explicit_algorithm_group(config_dir): + loader = PyLuceneConfigLoader(config_path=config_dir) + + _, configs = loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algo_groups="pylucene_test.test", + ) + + assert len(configs) == 1 + assert configs[0].index_name.startswith("pylucene_test_test.") + + +@pytest.mark.parametrize( + ("selectors", "expected_groups"), + [ + ( + {}, + ["pylucene_alpha,base", "pylucene_beta,base"], + ), + ( + {"algorithms": "pylucene_beta,pylucene_alpha"}, + ["pylucene_alpha,base", "pylucene_beta,base"], + ), + ( + {"groups": "shared"}, + ["pylucene_alpha,shared", "pylucene_beta,shared"], + ), + ( + { + "algorithms": "pylucene_alpha,pylucene_beta", + "groups": "alpha_only,shared", + }, + [ + "pylucene_alpha,shared", + "pylucene_alpha,alpha_only", + "pylucene_beta,shared", + ], + ), + ( + {"algo_groups": "pylucene_beta.beta_only"}, + ["pylucene_beta,beta_only"], + ), + ( + { + "algorithms": "pylucene_alpha", + "algo_groups": "pylucene_beta.beta_only", + }, + ["pylucene_alpha,base", "pylucene_beta,beta_only"], + ), + ( + { + "groups": "alpha_only", + "algo_groups": "pylucene_beta.beta_only", + }, + ["pylucene_alpha,alpha_only", "pylucene_beta,beta_only"], + ), + ( + { + "algorithms": "pylucene_beta", + "groups": "shared", + "algo_groups": "pylucene_alpha.alpha_only", + }, + ["pylucene_alpha,alpha_only", "pylucene_beta,shared"], + ), + ( + { + "algorithms": "pylucene_alpha", + "groups": "shared", + "algo_groups": "pylucene_alpha.shared", + }, + ["pylucene_alpha,shared"], + ), + ( + { + "algo_groups": ( + "pylucene_beta.beta_only,pylucene_alpha.alpha_only" + ) + }, + ["pylucene_alpha,alpha_only", "pylucene_beta,beta_only"], + ), + ( + { + "algorithms": " pylucene_beta, pylucene_beta ", + "groups": " base,base ", + }, + ["pylucene_beta,base"], + ), + ], + ids=[ + "defaults", + "algorithm", + "group", + "algorithm-and-group", + "explicit-pair", + "algorithm-plus-explicit-pair", + "group-plus-explicit-pair", + "all-selectors", + "overlapping-selectors", + "explicit-pair-order-follows-config", + "duplicate-selectors", + ], +) +def test_config_loader_combines_global_and_explicit_selectors( + config_dir, monkeypatch, selectors, expected_groups +): + alpha_config = config_dir / "algos" / "pylucene_alpha.yaml" + alpha_config.write_text( + f"""\ +name: pylucene_alpha +groups: + base: + build: + codec: [{_HNSW_CODEC}] + search: {{}} + shared: + build: + codec: [{_HNSW_CODEC}] + search: {{}} + alpha_only: + build: + codec: [{_HNSW_CODEC}] + search: {{}} +""" + ) + beta_config = config_dir / "algos" / "pylucene_beta.yaml" + beta_config.write_text( + f"""\ +name: pylucene_beta +groups: + base: + build: + codec: [{_HNSW_CODEC}] + search: {{}} + shared: + build: + codec: [{_HNSW_CODEC}] + search: {{}} + beta_only: + build: + codec: [{_HNSW_CODEC}] + search: {{}} +""" + ) + loader = PyLuceneConfigLoader(config_path=config_dir) + monkeypatch.setattr( + loader, + "gather_algorithm_configs", + lambda *_args: [str(alpha_config), str(beta_config)], + ) + + _, configs = loader.load( + dataset="test-dataset", + dataset_path="/datasets", + **selectors, + ) + + selected_groups = [config.output_filename[0] for config in configs] + assert selected_groups == expected_groups + + +def test_algorithm_selection_resolution_preserves_inputs_and_source_order(): + alpha_base = {"build": {"codec": [_HNSW_CODEC]}} + beta_base = {"build": {"codec": [_HNSW_CODEC]}} + algorithm_configs = { + "pylucene_beta": {"base": beta_base}, + "pylucene_alpha": {"base": alpha_base}, + } + original_items = [ + (algorithm, list(groups.items())) + for algorithm, groups in algorithm_configs.items() + ] + selection = pylucene_backend._AlgorithmSelection( + algorithms=frozenset({"pylucene_alpha", "pylucene_beta"}), + groups=None, + explicit_groups=frozenset(), + ) + + selected = selection.resolve(algorithm_configs) + + assert [(group.algorithm, group.group) for group in selected] == [ + ("pylucene_beta", "base"), + ("pylucene_alpha", "base"), + ] + assert selected[0].configuration is beta_base + assert selected[1].configuration is alpha_base + assert [ + (algorithm, list(groups.items())) + for algorithm, groups in algorithm_configs.items() + ] == original_items + + +def test_algorithm_selection_rejects_group_outside_selected_algorithms(): + selection = pylucene_backend._AlgorithmSelection( + algorithms=frozenset({"pylucene_alpha"}), + groups=frozenset({"beta_only"}), + explicit_groups=frozenset(), + ) + algorithm_configs = { + "pylucene_alpha": {"base": {}}, + "pylucene_beta": {"beta_only": {}}, + } + + with pytest.raises( + ValueError, + match="Unknown PyLucene group selector\\(s\\): beta_only", + ): + selection.resolve(algorithm_configs) + + +def test_algorithm_selection_reports_explicit_errors_deterministically(): + selection = pylucene_backend._AlgorithmSelection( + algorithms=None, + groups=None, + explicit_groups=frozenset( + { + ("pylucene_zeta", "base"), + ("pylucene_beta", "base"), + } + ), + ) + + with pytest.raises(ValueError) as error: + selection.resolve({"pylucene_alpha": {"base": {}}}) + + assert str(error.value) == ( + "Unknown PyLucene algorithm in --algo-groups: pylucene_beta" + ) + + +def test_config_loader_rejects_malformed_algorithm_group(config_dir): + loader = PyLuceneConfigLoader(config_path=config_dir) + + with pytest.raises(ValueError, match="."): + loader.load( + dataset="test-dataset", + dataset_path="/datasets", + algo_groups="pylucene_test", + ) + + +@pytest.mark.parametrize( + ("selectors", "error"), + [ + ( + {"algorithms": "not-pylucene", "groups": "base"}, + "Unknown PyLucene algorithm selector.*not-pylucene", + ), + ( + {"algorithms": "pylucene_test", "groups": "missing"}, + "Unknown PyLucene group selector.*missing", + ), + ( + { + "algorithms": "pylucene_test", + "groups": "base", + "algo_groups": "not-pylucene.test", + }, + "Unknown PyLucene algorithm in --algo-groups.*not-pylucene", + ), + ( + { + "algorithms": "pylucene_test", + "groups": "base", + "algo_groups": "pylucene_test.missing", + }, + "Unknown PyLucene group for pylucene_test.*missing", + ), + ], + ids=[ + "algorithm", + "global-group", + "algo-group-algorithm", + "algo-group-group", + ], +) +def test_config_loader_rejects_unknown_selectors(config_dir, selectors, error): + loader = PyLuceneConfigLoader(config_path=config_dir) + + with pytest.raises(ValueError, match=error): + loader.load( + dataset="test-dataset", + dataset_path="/datasets", + **selectors, + ) + + +@pytest.mark.parametrize( + ("algorithm", "expected_codecs"), + [ + ( + "pylucene_cuvs_hnsw", + { + "Lucene101AcceleratedHNSWCodec", + "Lucene101AcceleratedHNSWBaseLayerCodec", + "Lucene101AcceleratedHNSWMultiLayerCodec", + }, + ), + ("pylucene_cuvs_cagra", {_CAGRA_CODEC}), + ], +) +def test_shipped_algorithm_configs_load(algorithm, expected_codecs): + _, configs = PyLuceneConfigLoader().load( + dataset="test-data", + dataset_path="/datasets", + algorithms=algorithm, + groups="base", + ) + + assert { + config.indexes[0].build_param["codec"] for config in configs + } == expected_codecs diff --git a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_provenance.py b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_provenance.py new file mode 100644 index 0000000000..475aaba9c1 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_provenance.py @@ -0,0 +1,270 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Index provenance and training-shape tests for PyLucene.""" + +from __future__ import annotations + +import json +import stat + +import numpy as np +import pytest + +import cuvs_bench.backends.pylucene as pylucene_backend +from cuvs_bench.backends.base import Dataset +from cuvs_bench.tests._pylucene_test_utils import ( + _CAGRA_CODEC, + _HNSW_CODEC, + _FakeRuntime, + _backend, + _dataset, + _index, + _prepare_cagra_index, + _prepare_hnsw_index, + _write_test_bin, +) + + +def test_hnsw_provenance_round_trip(tmp_path): + index_path = tmp_path / "index" + manifest_path = _prepare_hnsw_index(index_path) + + verification = pylucene_backend._verify_hnsw_provenance( + index_path, + _HNSW_CODEC, + expected_vector_count=10, + expected_dimensions=4, + ) + + assert verification.to_metadata() == { + "status": "gpu-hnsw-provenance", + "schema_version": 1, + "codec": _HNSW_CODEC, + "writer_path": "gpu-hnsw", + "vector_count": 10, + "dimensions": 4, + "commit_file_count": 1, + } + payload = json.loads(manifest_path.read_text()) + assert payload["commit_fingerprints"] == [ + { + "name": "segments_1", + "sha256": ( + "1bc04b5291c26a46d918139138b992d2de976d6851d0893b" + "0476b85bfbdfc6e6" + ), + } + ] + assert stat.S_IMODE(manifest_path.stat().st_mode) == 0o644 + assert list(index_path.glob(f"{manifest_path.name}.*.tmp")) == [] + + +def test_cagra_provenance_round_trip(tmp_path): + index_path = tmp_path / "index" + manifest_path = _prepare_cagra_index(index_path) + + verification = pylucene_backend._verify_cagra_provenance( + index_path, + expected_vector_count=10, + expected_dimensions=4, + ) + + assert verification.to_metadata() == { + "status": "gpu-cagra-provenance", + "schema_version": 1, + "codec": _CAGRA_CODEC, + "writer_path": "gpu-cagra", + "vector_count": 10, + "dimensions": 4, + "commit_file_count": 1, + } + assert stat.S_IMODE(manifest_path.stat().st_mode) == 0o644 + + +@pytest.mark.parametrize( + ("expected_vector_count", "expected_dimensions", "error"), + [ + (11, 4, "vector count"), + (10, 8, "dimensions"), + ], +) +def test_hnsw_provenance_rejects_dataset_shape_mismatch( + tmp_path, expected_vector_count, expected_dimensions, error +): + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + + with pytest.raises(RuntimeError, match=error): + pylucene_backend._verify_hnsw_provenance( + index_path, + _HNSW_CODEC, + expected_vector_count=expected_vector_count, + expected_dimensions=expected_dimensions, + ) + + +@pytest.mark.parametrize( + ("manifest_state", "error"), + [ + ("wrong-writer", "required GPU writer path"), + ("boolean-count", "positive integer"), + ("malformed-fingerprint", "malformed Lucene commit fingerprint"), + ("duplicate-fingerprint", "must be unique and sorted"), + ], +) +def test_hnsw_provenance_rejects_malformed_manifest_fields( + tmp_path, manifest_state, error +): + index_path = tmp_path / "index" + manifest_path = _prepare_hnsw_index(index_path) + payload = json.loads(manifest_path.read_text()) + + if manifest_state == "wrong-writer": + payload["writer_path"] = "cpu-hnsw" + elif manifest_state == "boolean-count": + payload["vector_count"] = True + elif manifest_state == "malformed-fingerprint": + payload["commit_fingerprints"] = [{"name": "segments_1"}] + else: + payload["commit_fingerprints"] *= 2 + manifest_path.write_text(json.dumps(payload)) + + with pytest.raises(RuntimeError, match=error): + pylucene_backend._verify_hnsw_provenance( + index_path, + _HNSW_CODEC, + ) + + +@pytest.mark.parametrize( + ("field_name", "field_value", "error"), + [ + ("name", 1, "commit filename must be a string"), + ("name", "commit_1", "must start with 'segments_'"), + ("name", "segments_1/child", "must not contain a path"), + ("sha256", 1, "commit SHA-256 must be a string"), + ("sha256", "0" * 63, "must contain 64 characters"), + ("sha256", "A" * 64, "must be lowercase hexadecimal"), + ], +) +def test_hnsw_provenance_reports_invalid_commit_fingerprint_field( + tmp_path, field_name, field_value, error +): + index_path = tmp_path / "index" + manifest_path = _prepare_hnsw_index(index_path) + payload = json.loads(manifest_path.read_text()) + payload["commit_fingerprints"][0][field_name] = field_value + manifest_path.write_text(json.dumps(payload)) + + with pytest.raises(RuntimeError, match=error): + pylucene_backend._verify_hnsw_provenance(index_path, _HNSW_CODEC) + + +def test_expected_training_shape_prefers_explicit_vectors_over_base_file( + tmp_path, +): + base_file = tmp_path / "different.fbin" + _write_test_bin(base_file, np.zeros((3, 8), dtype=np.float32)) + dataset = _dataset(n_base=10, dimensions=4) + dataset.base_file = str(base_file) + + assert pylucene_backend._expected_training_shape(dataset) == (10, 4) + + +def test_expected_training_shape_does_not_validate_unused_file_metadata( + tmp_path, +): + dataset = _dataset(n_base=10, dimensions=4) + dataset.base_file = str(tmp_path / "missing.ibin") + dataset.metadata["subset_size"] = True + + assert pylucene_backend._expected_training_shape(dataset) == (10, 4) + + +def test_expected_training_shape_returns_none_without_training_source(): + dataset = Dataset( + name="query-only", + query_vectors=np.zeros((2, 4), dtype=np.float32), + distance_metric="euclidean", + ) + + assert pylucene_backend._expected_training_shape(dataset) is None + + +def test_expected_training_shape_clamps_file_to_valid_subset(tmp_path): + base_file = tmp_path / "base.fbin" + _write_test_bin(base_file, np.zeros((10, 4), dtype=np.float32)) + dataset = Dataset( + name="file-backed", + query_vectors=np.zeros((2, 4), dtype=np.float32), + base_file=str(base_file), + distance_metric="euclidean", + metadata={"subset_size": 3}, + ) + + assert pylucene_backend._expected_training_shape(dataset) == (3, 4) + assert dataset.loaded_training_vectors is None + + +def test_build_reuses_file_backed_subset_without_loading_vectors(tmp_path): + base_file = tmp_path / "base.fbin" + _write_test_bin(base_file, np.zeros((10, 4), dtype=np.float32)) + dataset = Dataset( + name="file-backed", + query_vectors=np.zeros((2, 4), dtype=np.float32), + base_file=str(base_file), + distance_metric="euclidean", + metadata={"subset_size": 3}, + ) + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path, vector_count=3, dimensions=4) + backend = _backend() + + result = backend.build(dataset, [_index(index_path)], force=False) + + assert result.success + assert result.metadata["skipped"] is True + assert result.metadata["hnsw_verification"]["vector_count"] == 3 + assert dataset.loaded_training_vectors is None + assert backend._runtime is None + + +@pytest.mark.parametrize("subset_size", [0, -1, True, "3"]) +def test_expected_training_shape_rejects_invalid_file_subset( + tmp_path, subset_size +): + base_file = tmp_path / "base.fbin" + _write_test_bin(base_file, np.zeros((10, 4), dtype=np.float32)) + dataset = Dataset( + name="file-backed", + query_vectors=np.zeros((2, 4), dtype=np.float32), + base_file=str(base_file), + distance_metric="euclidean", + metadata={"subset_size": subset_size}, + ) + + with pytest.raises(ValueError, match="positive integer"): + pylucene_backend._expected_training_shape(dataset) + + +def test_reused_index_rejects_non_float32_base_file(tmp_path): + base_file = tmp_path / "base.ibin" + _write_test_bin(base_file, np.zeros((10, 4), dtype=np.int32)) + dataset = Dataset( + name="integer-data", + query_vectors=np.zeros((2, 4), dtype=np.float32), + base_file=str(base_file), + distance_metric="euclidean", + ) + index_path = tmp_path / "index" + _prepare_hnsw_index(index_path) + + result = _backend(_FakeRuntime()).build( + dataset, [_index(index_path)], force=False + ) + + assert not result.success + assert "must use float32 values, got int32" in result.error_message diff --git a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_runtime.py b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_runtime.py new file mode 100644 index 0000000000..da62e13421 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_runtime.py @@ -0,0 +1,602 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""PyLucene JVM and Java-boundary unit tests.""" + +from __future__ import annotations + +import zipfile +from types import SimpleNamespace + +import numpy as np +import pytest + +import cuvs_bench.backends.pylucene as pylucene_backend +from cuvs_bench.backends.pylucene import _ResolvedCodec +from cuvs_bench.tests._pylucene_test_utils import _HNSW_CODEC + + +def _resolved_codec(java_codec=None, writer_path="gpu-hnsw") -> _ResolvedCodec: + return _ResolvedCodec( + codec_name=_HNSW_CODEC, + java_codec=java_codec if java_codec is not None else object(), + telemetry={"writerPath": writer_path}, + writer_path=writer_path, + ) + + +class _FakeVMEnvironment: + def __init__(self): + self.attach_count = 0 + + def attachCurrentThread(self): + self.attach_count += 1 + + +class _FakeLuceneModule: + CLASSPATH = "/pylucene/lucene-core.jar" + VERSION = "test" + + def __init__(self): + self.environment = None + self.init_calls = [] + + def getVMEnv(self): + return self.environment + + def initVM(self, **kwargs): + self.init_calls.append(kwargs) + self.environment = _FakeVMEnvironment() + return self.environment + + +class _FakeIndexWriter: + def __init__(self, error_at=None): + self.error_at = error_at + self.documents = [] + self.committed = False + self.rollback_called = False + self.close_called = False + + def addDocument(self, document): + if self.error_at in {"interrupt", "interrupt-rollback"}: + raise KeyboardInterrupt + if self.error_at in {"add", "rollback"}: + raise RuntimeError("add failed") + self.documents.append(document) + + def commit(self): + if self.error_at == "commit": + raise RuntimeError("commit failed") + self.committed = True + + def rollback(self): + self.rollback_called = True + if self.error_at in {"rollback", "interrupt-rollback"}: + raise RuntimeError("rollback failed") + + def close(self): + self.close_called = True + if self.error_at == "close": + raise RuntimeError("close failed") + + +class _FakeIndexWriterConfig: + OpenMode = SimpleNamespace(CREATE=object()) + + def setOpenMode(self, _mode): + pass + + def setCodec(self, _codec): + pass + + def setUseCompoundFile(self, _enabled): + pass + + def setMergeScheduler(self, _scheduler): + pass + + +def _fake_index_writer_runtime(error_at=None, directory_close_error=None): + writer = _FakeIndexWriter(error_at=error_at) + directory = SimpleNamespace(closed=False) + + def close_directory(): + directory.closed = True + if directory_close_error is not None: + raise directory_close_error + + directory.close = close_directory + runtime = pylucene_backend._PyLuceneRuntime.__new__( + pylucene_backend._PyLuceneRuntime + ) + runtime.attach_current_thread = lambda: None + runtime.resolve_codec = lambda _codec_name: object() + runtime.Paths = SimpleNamespace(get=lambda path: path) + runtime.FSDirectory = SimpleNamespace(open=lambda _path: directory) + runtime.IndexWriterConfig = _FakeIndexWriterConfig + runtime.SerialMergeScheduler = lambda: object() + runtime.IndexWriter = lambda _directory, _config: writer + runtime._vector_document = lambda document_id, vector: ( + document_id, + vector.copy(), + ) + runtime._test_writer = writer + runtime._test_directory = directory + runtime._test_codec = object() + return runtime + + +@pytest.fixture(autouse=True) +def _reset_jvm_tracking(monkeypatch): + monkeypatch.setattr(pylucene_backend, "_INITIALIZED_CLASSPATH", None) + monkeypatch.setattr(pylucene_backend, "_INITIALIZED_VMARGS", None) + + +def test_initialize_pylucene_uses_verified_classpath_and_vmargs( + tmp_path, monkeypatch +): + cuvs_java = tmp_path / "cuvs-java.jar" + cuvs_lucene = tmp_path / "cuvs-lucene.jar" + cuvs_java.touch() + cuvs_lucene.touch() + fake_lucene = _FakeLuceneModule() + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda name: fake_lucene, + ) + + returned = pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": cuvs_java, + "cuvs_lucene_jar": cuvs_lucene, + "java_library_path": "/native", + "jvm_args": ["-Xms1g"], + } + ) + + assert returned is fake_lucene + assert len(fake_lucene.init_calls) == 1 + init_call = fake_lucene.init_calls[0] + assert init_call["classpath"].split(":") == [ + str(cuvs_java), + str(cuvs_lucene), + fake_lucene.CLASSPATH, + ] + assert init_call["vmargs"] == [ + "--enable-native-access=ALL-UNNAMED", + "--add-modules=jdk.incubator.vector", + "-Djava.library.path=/native", + "-Xms1g", + ] + assert fake_lucene.environment.attach_count == 1 + + pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": cuvs_java, + "cuvs_lucene_jar": cuvs_lucene, + "java_library_path": "/native", + "jvm_args": ["-Xms1g"], + } + ) + assert len(fake_lucene.init_calls) == 1 + assert fake_lucene.environment.attach_count == 2 + + +@pytest.mark.parametrize( + ("library_paths", "expected_library_path"), + [ + ( + { + "JAVA_LIBRARY_PATH": "/java-native", + "LD_LIBRARY_PATH": "/ld-native", + }, + "/java-native", + ), + ({"LD_LIBRARY_PATH": "/ld-native"}, "/ld-native"), + ], + ids=["java-library-path-precedence", "ld-library-path-fallback"], +) +def test_initialize_pylucene_uses_environment_runtime_config( + library_paths, + expected_library_path, + tmp_path, + monkeypatch, +): + cuvs_java = tmp_path / "cuvs-java.jar" + cuvs_lucene = tmp_path / "cuvs-lucene.jar" + cuvs_java.touch() + cuvs_lucene.touch() + fake_lucene = _FakeLuceneModule() + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda _name: fake_lucene, + ) + monkeypatch.setenv("CUVS_LUCENE_CUVS_JAVA_JAR", str(cuvs_java)) + monkeypatch.setenv("CUVS_LUCENE_JAR", str(cuvs_lucene)) + monkeypatch.delenv("JAVA_LIBRARY_PATH", raising=False) + monkeypatch.delenv("LD_LIBRARY_PATH", raising=False) + for name, value in library_paths.items(): + monkeypatch.setenv(name, value) + + pylucene_backend._initialize_pylucene({}) + + init_call = fake_lucene.init_calls[0] + assert init_call["classpath"].split(":") == [ + str(cuvs_java), + str(cuvs_lucene), + fake_lucene.CLASSPATH, + ] + assert init_call["vmargs"] == [ + "--enable-native-access=ALL-UNNAMED", + "--add-modules=jdk.incubator.vector", + f"-Djava.library.path={expected_library_path}", + ] + + +def test_initialize_pylucene_rejects_fat_cuvs_lucene_jar_before_init( + tmp_path, monkeypatch +): + cuvs_java = tmp_path / "cuvs-java.jar" + cuvs_lucene = tmp_path / "cuvs-lucene-jar-with-dependencies.jar" + cuvs_java.touch() + with zipfile.ZipFile(cuvs_lucene, "w") as archive: + archive.writestr( + pylucene_backend._LUCENE_CORE_CLASS, + b"bundled Lucene bytecode", + ) + fake_lucene = _FakeLuceneModule() + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda _name: fake_lucene, + ) + + with pytest.raises(RuntimeError, match="standard thin cuvs-lucene JAR"): + pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": cuvs_java, + "cuvs_lucene_jar": cuvs_lucene, + } + ) + + assert fake_lucene.init_calls == [] + + +def test_initialize_pylucene_rejects_externally_started_jvm( + tmp_path, monkeypatch +): + cuvs_java = tmp_path / "cuvs-java.jar" + cuvs_lucene = tmp_path / "cuvs-lucene.jar" + cuvs_java.touch() + cuvs_lucene.touch() + fake_lucene = _FakeLuceneModule() + fake_lucene.environment = _FakeVMEnvironment() + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda name: fake_lucene, + ) + + with pytest.raises(RuntimeError, match="initialized before"): + pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": cuvs_java, + "cuvs_lucene_jar": cuvs_lucene, + } + ) + + +def test_initialize_pylucene_rejects_different_jars_after_start( + tmp_path, monkeypatch +): + first_java = tmp_path / "first-java.jar" + first_lucene = tmp_path / "first-lucene.jar" + second_java = tmp_path / "second-java.jar" + for path in (first_java, first_lucene, second_java): + path.touch() + fake_lucene = _FakeLuceneModule() + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda name: fake_lucene, + ) + + pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": first_java, + "cuvs_lucene_jar": first_lucene, + } + ) + + with pytest.raises(RuntimeError, match="different"): + pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": second_java, + "cuvs_lucene_jar": first_lucene, + } + ) + + +def test_initialize_pylucene_rejects_different_vmargs_after_start( + tmp_path, monkeypatch +): + cuvs_java = tmp_path / "cuvs-java.jar" + cuvs_lucene = tmp_path / "cuvs-lucene.jar" + cuvs_java.touch() + cuvs_lucene.touch() + fake_lucene = _FakeLuceneModule() + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda name: fake_lucene, + ) + + pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": cuvs_java, + "cuvs_lucene_jar": cuvs_lucene, + "java_library_path": "/first", + } + ) + + with pytest.raises(RuntimeError, match="different JVM arguments"): + pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": cuvs_java, + "cuvs_lucene_jar": cuvs_lucene, + "java_library_path": "/second", + } + ) + + +def test_initialize_pylucene_reports_missing_binding(monkeypatch): + def missing_import(_name): + raise ImportError("missing") + + monkeypatch.setattr( + pylucene_backend.importlib, "import_module", missing_import + ) + + with pytest.raises(ImportError, match="must be built"): + pylucene_backend._initialize_pylucene({}) + + +def test_initialize_pylucene_reports_missing_jar(monkeypatch): + monkeypatch.delenv("CUVS_LUCENE_CUVS_JAVA_JAR", raising=False) + monkeypatch.delenv("CUVS_LUCENE_JAR", raising=False) + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda _name: _FakeLuceneModule(), + ) + + with pytest.raises(RuntimeError, match="cuvs_java_jar"): + pylucene_backend._initialize_pylucene({}) + + +@pytest.mark.parametrize("jvm_args", ["-Xmx1g", ["-Xmx1g", 1]]) +def test_initialize_pylucene_rejects_invalid_jvm_args( + jvm_args, tmp_path, monkeypatch +): + cuvs_java = tmp_path / "cuvs-java.jar" + cuvs_lucene = tmp_path / "cuvs-lucene.jar" + cuvs_java.touch() + cuvs_lucene.touch() + monkeypatch.setattr( + pylucene_backend.importlib, + "import_module", + lambda _name: _FakeLuceneModule(), + ) + + with pytest.raises(TypeError, match="jvm_args"): + pylucene_backend._initialize_pylucene( + { + "cuvs_java_jar": cuvs_java, + "cuvs_lucene_jar": cuvs_lucene, + "jvm_args": jvm_args, + } + ) + + +@pytest.mark.parametrize( + ("description", "expected"), + [ + ( + "Format(writerPath=gpu-hnsw;hnswLayers=1)", + {"writerPath": "gpu-hnsw", "hnswLayers": "1"}, + ), + ("Format(writerPath=gpu-cagra)", {"writerPath": "gpu-cagra"}), + ], +) +def test_parse_writer_telemetry(description, expected): + assert pylucene_backend._parse_writer_telemetry(description) == expected + + +@pytest.mark.parametrize( + "description", + ["Format", "Format(writerPath)", "Format(=gpu-hnsw)"], +) +def test_parse_writer_telemetry_rejects_malformed_description(description): + with pytest.raises(RuntimeError, match="Malformed"): + pylucene_backend._parse_writer_telemetry(description) + + +def test_resolved_codec_copies_and_freezes_writer_telemetry(): + telemetry = {"writerPath": "gpu-hnsw"} + + resolved_codec = _ResolvedCodec( + codec_name=_HNSW_CODEC, + java_codec=object(), + telemetry=telemetry, + writer_path="gpu-hnsw", + ) + telemetry["writerPath"] = "changed" + + assert dict(resolved_codec.telemetry) == {"writerPath": "gpu-hnsw"} + with pytest.raises(TypeError): + resolved_codec.telemetry["writerPath"] = "changed" + + +def test_runtime_build_index_commits_and_closes_writer_and_directory(tmp_path): + runtime = _fake_index_writer_runtime() + vectors = np.zeros((2, 4), dtype=np.float32) + + telemetry = runtime.build_index( + tmp_path, + vectors, + _resolved_codec(runtime._test_codec), + ) + + assert telemetry == {"writerPath": "gpu-hnsw"} + assert len(runtime._test_writer.documents) == 2 + assert runtime._test_writer.committed is True + assert runtime._test_writer.rollback_called is False + assert runtime._test_writer.close_called is True + assert runtime._test_directory.closed is True + + +@pytest.mark.parametrize( + ("error_at", "expected_error", "expect_rollback", "expect_close"), + [ + ("add", "add failed", True, False), + ("commit", "commit failed", True, False), + ("rollback", "rollback failed", True, False), + ("close", "close failed", False, True), + ], +) +def test_runtime_build_index_preserves_transaction_cleanup( + tmp_path, error_at, expected_error, expect_rollback, expect_close +): + runtime = _fake_index_writer_runtime(error_at=error_at) + vectors = np.zeros((2, 4), dtype=np.float32) + + with pytest.raises(RuntimeError, match=expected_error): + runtime.build_index( + tmp_path, + vectors, + _resolved_codec(runtime._test_codec), + ) + + assert runtime._test_writer.rollback_called is expect_rollback + assert runtime._test_writer.close_called is expect_close + assert runtime._test_directory.closed is True + + +def test_runtime_build_index_rolls_back_on_interrupt(tmp_path): + runtime = _fake_index_writer_runtime(error_at="interrupt") + vectors = np.zeros((2, 4), dtype=np.float32) + + with pytest.raises(KeyboardInterrupt): + runtime.build_index( + tmp_path, + vectors, + _resolved_codec(runtime._test_codec), + ) + + assert runtime._test_writer.rollback_called is True + assert runtime._test_writer.close_called is False + assert runtime._test_directory.closed is True + + +def test_runtime_build_index_preserves_interrupt_when_rollback_fails(tmp_path): + runtime = _fake_index_writer_runtime(error_at="interrupt-rollback") + vectors = np.zeros((2, 4), dtype=np.float32) + + with pytest.raises(KeyboardInterrupt) as exc_info: + runtime.build_index( + tmp_path, + vectors, + _resolved_codec(runtime._test_codec), + ) + + assert exc_info.value.__notes__ == [ + "IndexWriter rollback also failed: RuntimeError: rollback failed" + ] + assert runtime._test_writer.rollback_called is True + assert runtime._test_directory.closed is True + + +def test_runtime_build_index_preserves_interrupt_when_directory_close_fails( + tmp_path, +): + runtime = _fake_index_writer_runtime( + error_at="interrupt", + directory_close_error=RuntimeError("directory close failed"), + ) + vectors = np.zeros((2, 4), dtype=np.float32) + + with pytest.raises(KeyboardInterrupt) as exc_info: + runtime.build_index( + tmp_path, + vectors, + _resolved_codec(runtime._test_codec), + ) + + assert exc_info.value.__notes__ == [ + "Failed to close Lucene directory: " + "RuntimeError: directory close failed" + ] + assert runtime._test_writer.rollback_called is True + assert runtime._test_directory.closed is True + + +def test_search_cleanup_attempts_directory_after_reader_close_fails(): + close_calls = [] + + def fail_reader_close(): + close_calls.append("reader") + raise RuntimeError("reader close failed") + + def fail_directory_close(): + close_calls.append("directory") + raise RuntimeError("directory close failed") + + reader = SimpleNamespace(close=fail_reader_close) + directory = SimpleNamespace(close=fail_directory_close) + + with pytest.raises(RuntimeError, match="reader close failed") as exc_info: + with pylucene_backend._CleanupStack() as cleanups: + cleanups.add("close Lucene directory", directory.close) + cleanups.add("close Lucene index reader", reader.close) + + assert close_calls == ["reader", "directory"] + assert exc_info.value.__notes__ == [ + "Failed to close Lucene directory: " + "RuntimeError: directory close failed" + ] + + +def test_cleanup_stack_ignores_unrelated_handled_exception(): + def fail_close(): + raise RuntimeError("close failed") + + unrelated_error = None + try: + raise ValueError("unrelated") + except ValueError as exc: + unrelated_error = exc + with pytest.raises(RuntimeError, match="close failed"): + with pylucene_backend._CleanupStack() as cleanups: + cleanups.add("close test resource", fail_close) + + assert getattr(unrelated_error, "__notes__", []) == [] + + +def test_cleanup_stack_prioritizes_cleanup_interrupt_over_operation_error(): + def interrupt_close(): + raise KeyboardInterrupt("stop") + + with pytest.raises(KeyboardInterrupt, match="stop") as exc_info: + with pylucene_backend._CleanupStack() as cleanups: + cleanups.add("close test resource", interrupt_close) + raise RuntimeError("operation failed") + + assert exc_info.value.__notes__ == [ + "Raised while attempting to close test resource; prior failure: " + "RuntimeError: operation failed" + ] diff --git a/python/cuvs_bench/cuvs_bench/tests/test_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.""" From 0483c2ba3efa7fb0aba97238e39bd3c5838fda6d Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Fri, 31 Jul 2026 13:29:59 +0000 Subject: [PATCH 04/12] Add live PyLucene integration coverage Exercise real JVM and cuVS-Lucene HNSW and CAGRA build/search paths behind an opt-in pytest marker. Cover persisted GPU formats, index reuse, CLI execution, fallback rejection, and integrity failures. --- .../tests/test_pylucene_integration.py | 704 ++++++++++++++++++ python/cuvs_bench/pyproject.toml | 1 + 2 files changed, 705 insertions(+) create mode 100644 python/cuvs_bench/cuvs_bench/tests/test_pylucene_integration.py diff --git a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_integration.py b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_integration.py new file mode 100644 index 0000000000..af886641f7 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_integration.py @@ -0,0 +1,704 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +"""Opt-in integration coverage for the PyLucene benchmark backend.""" + +import csv +import importlib +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pytest + +from cuvs_bench._bin_format import write_bin_header +from cuvs_bench.backends._utils import compute_recall +from cuvs_bench.backends.base import Dataset +from cuvs_bench.backends.pylucene import ( + _CAGRA_PROVENANCE_FILE, + PyLuceneBackend, + _HNSW_PROVENANCE_FILE, + _PyLuceneRuntime, + _ResolvedCodec, +) +from cuvs_bench.orchestrator.config_loaders import IndexConfig + +pytestmark = [ + pytest.mark.pylucene, + pytest.mark.filterwarnings( + "ignore:builtin type .* has no __module__ attribute:DeprecationWarning" + ), +] + +_OPT_IN_ENV = "CUVS_BENCH_PYLUCENE_INTEGRATION" +_CUVS_JAVA_JAR_ENV = "CUVS_LUCENE_CUVS_JAVA_JAR" +_CUVS_LUCENE_JAR_ENV = "CUVS_LUCENE_JAR" +_CAGRA_CODEC = "CuVS2510GPUSearchCodec" +_HNSW_CODEC = "Lucene101AcceleratedHNSWCodec" + + +def _write_test_bin(path: Path, data: np.ndarray) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("wb") as file: + write_bin_header(file, data.shape[0], data.shape[1]) + np.ascontiguousarray(data).tofile(file) + + +def _required_jar(env_name: str) -> Path: + configured = os.environ.get(env_name) + if not configured: + pytest.fail(f"{env_name} must point to the required runtime jar") + + jar = Path(configured).expanduser() + if not jar.is_file(): + pytest.fail(f"{env_name} does not point to an existing file: {jar}") + return jar.resolve() + + +def _assert_cagra_verification( + metadata, expected_vector_count, expected_dimensions +): + verification = metadata["cagra_verification"] + assert verification["status"] == "cagra-only" + assert verification["segment_count"] >= 1 + assert verification["field_count"] >= 1 + assert verification["vector_count"] == expected_vector_count + assert verification["dimensions"] == expected_dimensions + + +def _assert_cagra_provenance( + metadata, expected_vector_count, expected_dimensions +): + assert metadata["cagra_provenance"] == { + "status": "gpu-cagra-provenance", + "schema_version": 1, + "codec": _CAGRA_CODEC, + "writer_path": "gpu-cagra", + "vector_count": expected_vector_count, + "dimensions": expected_dimensions, + "commit_file_count": 1, + } + + +def _commit_one_deletion(backend, index_path): + runtime = backend._get_runtime() + runtime.attach_current_thread() + directory = runtime.FSDirectory.open(runtime.Paths.get(str(index_path))) + writer = None + reader = None + try: + writer_config = runtime.IndexWriterConfig() + writer_config.setOpenMode(runtime.IndexWriterConfig.OpenMode.APPEND) + writer = runtime.IndexWriter(directory, writer_config) + reader = runtime.DirectoryReader.open(writer) + assert int(writer.tryDeleteDocument(reader, 0)) >= 0 + writer.commit() + finally: + try: + if reader is not None: + reader.close() + finally: + try: + if writer is not None: + writer.close() + finally: + directory.close() + + +def _assert_hnsw_verification( + metadata, codec, expected_vector_count, expected_dimensions +): + verification = metadata["hnsw_verification"] + assert verification == { + "status": "gpu-hnsw-provenance", + "schema_version": 1, + "codec": codec, + "writer_path": "gpu-hnsw", + "vector_count": expected_vector_count, + "dimensions": expected_dimensions, + "commit_file_count": 1, + } + + +@pytest.fixture(scope="module") +def pylucene_runtime_config(): + """Resolve the explicitly configured PyLucene/cuVS runtime.""" + if os.environ.get(_OPT_IN_ENV) != "1": + pytest.skip(f"set {_OPT_IN_ENV}=1 to run PyLucene integration tests") + + cuvs_java_jar = _required_jar(_CUVS_JAVA_JAR_ENV) + cuvs_lucene_jar = _required_jar(_CUVS_LUCENE_JAR_ENV) + + java_library_path = os.environ.get("JAVA_LIBRARY_PATH") or os.environ.get( + "LD_LIBRARY_PATH" + ) + if not java_library_path: + pytest.fail( + "JAVA_LIBRARY_PATH or LD_LIBRARY_PATH must provide the native " + "cuVS runtime libraries" + ) + + try: + importlib.import_module("lucene") + except ImportError as exc: + pytest.fail(f"PyLucene is not importable: {exc}") + + return { + "cuvs_java_jar": str(cuvs_java_jar), + "cuvs_lucene_jar": str(cuvs_lucene_jar), + "java_library_path": java_library_path, + } + + +@pytest.mark.parametrize( + ( + "algo", + "codec", + "expected_writer_telemetry", + "expected_index_suffix", + ), + [ + pytest.param( + "pylucene_cuvs_hnsw", + "Lucene101AcceleratedHNSWCodec", + { + "writerPath": "gpu-hnsw", + "hnswLayers": "1", + "cagraGraphBuildAlgo": "NN_DESCENT", + "cagraGraphDegree": "64", + "cagraIntermediateGraphDegree": "128", + }, + ".vex", + id="accelerated-hnsw", + ), + pytest.param( + "pylucene_cuvs_hnsw", + "Lucene101AcceleratedHNSWBaseLayerCodec", + { + "writerPath": "gpu-hnsw", + "hnswLayers": "1", + "cagraGraphBuildAlgo": "NN_DESCENT", + "cagraGraphDegree": "32", + "cagraIntermediateGraphDegree": "64", + }, + ".vex", + id="accelerated-hnsw-base-layer", + ), + pytest.param( + "pylucene_cuvs_hnsw", + "Lucene101AcceleratedHNSWMultiLayerCodec", + { + "writerPath": "gpu-hnsw", + "hnswLayers": "3", + "cagraGraphBuildAlgo": "NN_DESCENT", + "cagraGraphDegree": "32", + "cagraIntermediateGraphDegree": "64", + }, + ".vex", + id="accelerated-hnsw-multi-layer", + ), + pytest.param( + "pylucene_cuvs_cagra", + "CuVS2510GPUSearchCodec", + { + "writerPath": "gpu-cagra", + "cagraStrategy": "HEURISTIC", + "cagraGraphBuildAlgo": "NN_DESCENT", + "cagraGraphDegree": "64", + "cagraIntermediateGraphDegree": "128", + }, + ".vcag", + id="cagra", + ), + ], +) +def test_build_and_search_with_real_pylucene_runtime( + tmp_path, + pylucene_runtime_config, + algo, + codec, + expected_writer_telemetry, + expected_index_suffix, +): + """Build and search through Lucene's public codec and query SPI.""" + rng = np.random.default_rng(1907) + training_vectors = rng.standard_normal((512, 32)).astype(np.float32) + query_ids = np.asarray([0, 137, 259, 511], dtype=np.int64) + query_vectors = training_vectors[query_ids].copy() + k = 5 + + squared_distances = np.sum( + (query_vectors[:, np.newaxis, :] - training_vectors[np.newaxis, :, :]) + ** 2, + axis=2, + ) + groundtruth_neighbors = np.argsort(squared_distances, axis=1)[:, :k] + groundtruth_distances = np.take_along_axis( + squared_distances, groundtruth_neighbors, axis=1 + ) + dataset = Dataset( + name="pylucene-integration", + training_vectors=training_vectors, + query_vectors=query_vectors, + groundtruth_neighbors=groundtruth_neighbors.astype(np.int32), + groundtruth_distances=groundtruth_distances.astype(np.float32), + distance_metric="euclidean", + ) + + index_path = tmp_path / f"{algo}-index" + index = IndexConfig( + name=f"{algo}-integration", + algo=algo, + build_param={"codec": codec}, + search_params=[{}], + file=str(index_path), + ) + backend = PyLuceneBackend( + { + "name": index.name, + "algo": algo, + "codec": codec, + **pylucene_runtime_config, + } + ) + + try: + build_result = backend.build(dataset, [index], force=True) + assert build_result.success, build_result.error_message + assert build_result.build_time_seconds > 0 + assert build_result.index_size_bytes > 0 + assert build_result.metadata["codec"] == codec + assert build_result.metadata["pylucene_version"] != "unknown" + assert ( + build_result.metadata["writer_path"] + == expected_writer_telemetry["writerPath"] + ) + assert ( + build_result.metadata["writer_telemetry"] + == expected_writer_telemetry + ) + assert any( + path.suffix == expected_index_suffix + for path in index_path.iterdir() + ) + if codec == _CAGRA_CODEC: + _assert_cagra_provenance( + build_result.metadata, + training_vectors.shape[0], + training_vectors.shape[1], + ) + _assert_cagra_verification( + build_result.metadata, + training_vectors.shape[0], + training_vectors.shape[1], + ) + assert (index_path / _CAGRA_PROVENANCE_FILE).is_file() + else: + _assert_hnsw_verification( + build_result.metadata, + codec, + training_vectors.shape[0], + training_vectors.shape[1], + ) + assert (index_path / _HNSW_PROVENANCE_FILE).is_file() + + reused_result = backend.build(dataset, [index], force=False) + assert reused_result.success, reused_result.error_message + assert reused_result.metadata["skipped"] is True + if codec == _CAGRA_CODEC: + _assert_cagra_provenance( + reused_result.metadata, + training_vectors.shape[0], + training_vectors.shape[1], + ) + _assert_cagra_verification( + reused_result.metadata, + training_vectors.shape[0], + training_vectors.shape[1], + ) + else: + _assert_hnsw_verification( + reused_result.metadata, + codec, + training_vectors.shape[0], + training_vectors.shape[1], + ) + + search_result = backend.search(dataset, [index], k=k, batch_size=2) + assert search_result.success, search_result.error_message + assert search_result.neighbors.shape == (query_vectors.shape[0], k) + assert search_result.distances.shape == (query_vectors.shape[0], k) + np.testing.assert_array_equal(search_result.neighbors[:, 0], query_ids) + assert np.all(np.isfinite(search_result.distances)) + recall = compute_recall( + search_result.neighbors, groundtruth_neighbors, k + ) + assert recall >= 0.75 + returned_distances = np.take_along_axis( + squared_distances, search_result.neighbors, axis=1 + ) + np.testing.assert_allclose( + search_result.distances, + returned_distances, + rtol=1e-5, + atol=1e-5, + ) + assert np.all(np.diff(search_result.distances, axis=1) >= -1e-6) + assert search_result.search_time_ms > 0 + assert search_result.latency_seconds is not None + assert search_result.latency_seconds > 0 + assert search_result.queries_per_second > 0 + assert search_result.metadata["codec"] == codec + assert search_result.metadata["pylucene_version"] != "unknown" + assert search_result.metadata["num_batches"] == 2 + assert search_result.metadata["mode"] == "latency" + assert set(search_result.latency_percentiles) == {"p50", "p95", "p99"} + if codec == _CAGRA_CODEC: + _assert_cagra_provenance( + search_result.metadata, + training_vectors.shape[0], + training_vectors.shape[1], + ) + _assert_cagra_verification( + search_result.metadata, + training_vectors.shape[0], + training_vectors.shape[1], + ) + + wrong_count_dataset = Dataset( + name="pylucene-integration-wrong-count", + training_vectors=training_vectors[:-1], + query_vectors=query_vectors, + distance_metric="euclidean", + ) + rejected = backend.build(wrong_count_dataset, [index], force=False) + assert not rejected.success + assert "vector count does not match the dataset: 512 != 511" in ( + rejected.error_message + ) + rejected = backend.search(wrong_count_dataset, [index], k=k) + assert not rejected.success + assert "vector count does not match the dataset: 512 != 511" in ( + rejected.error_message + ) + + data_path = next(index_path.glob("*.vcag")) + original_data = data_path.read_bytes() + + data_path.unlink() + rejected = backend.search(dataset, [index], k=k) + assert not rejected.success + assert "cannot read" in rejected.error_message + data_path.write_bytes(original_data) + + data_path.write_bytes(original_data[:-1]) + rejected = backend.search(dataset, [index], k=k) + assert not rejected.success + assert "do not exactly cover" in rejected.error_message + data_path.write_bytes(original_data) + + corrupted_data = bytearray(original_data) + corrupted_data[len(corrupted_data) // 2] ^= 0xFF + original_stat = data_path.stat() + data_path.write_bytes(corrupted_data) + os.utime( + data_path, + ns=(original_stat.st_atime_ns, original_stat.st_mtime_ns), + ) + rejected = backend.search(dataset, [index], k=k) + assert not rejected.success + assert "checksum" in rejected.error_message.lower() + data_path.write_bytes(original_data) + + deleted_index_path = tmp_path / f"{algo}-deleted-index" + shutil.copytree(index_path, deleted_index_path) + deleted_index = IndexConfig( + name=f"{algo}-deleted-integration", + algo=algo, + build_param={"codec": codec}, + search_params=[{}], + file=str(deleted_index_path), + ) + _commit_one_deletion(backend, deleted_index_path) + with pytest.raises(RuntimeError, match="committed deletions"): + backend._get_runtime().verify_cagra_index(deleted_index_path) + rejected = backend.search(dataset, [deleted_index], k=k) + assert not rejected.success + assert "does not match the current Lucene commit" in ( + rejected.error_message + ) + + metadata_path = next(index_path.glob("*.vemc")) + contents = bytearray(metadata_path.read_bytes()) + contents[-1] ^= 0xFF + metadata_path.write_bytes(contents) + with pytest.raises(RuntimeError, match="metadata format v0"): + backend._get_runtime().verify_cagra_index(index_path) + else: + assert "cagra_verification" not in search_result.metadata + _assert_hnsw_verification( + search_result.metadata, + codec, + training_vectors.shape[0], + training_vectors.shape[1], + ) + if codec == _HNSW_CODEC: + manifest_path = index_path / _HNSW_PROVENANCE_FILE + original_manifest = manifest_path.read_bytes() + segment_path = next(index_path.glob("segments_*")) + original_segment = segment_path.read_bytes() + + manifest_path.unlink() + rejected = backend.search(dataset, [index], k=k) + assert not rejected.success + assert "provenance manifest is missing" in ( + rejected.error_message + ) + manifest_path.write_bytes(original_manifest) + + manifest_path.write_text("{") + rejected = backend.search(dataset, [index], k=k) + assert not rejected.success + assert "provenance manifest cannot be read" in ( + rejected.error_message + ) + manifest_path.write_bytes(original_manifest) + + segment_path.write_bytes(original_segment + b"stale") + rejected = backend.search(dataset, [index], k=k) + assert not rejected.success + assert "does not match the current Lucene commit" in ( + rejected.error_message + ) + segment_path.write_bytes(original_segment) + + payload = json.loads(original_manifest) + payload["codec"] = "Lucene101AcceleratedHNSWBaseLayerCodec" + manifest_path.write_text(json.dumps(payload)) + rejected = backend.search(dataset, [index], k=k) + assert not rejected.success + assert "does not match the requested codec" in ( + rejected.error_message + ) + manifest_path.write_bytes(original_manifest) + finally: + backend.cleanup() + + +def test_real_verifier_rejects_cagra_brute_force_fallback( + tmp_path, pylucene_runtime_config +): + """Reject cuVS-Lucene's real one-vector fallback before search.""" + runtime = _PyLuceneRuntime.create(pylucene_runtime_config) + vectors = np.ones((1, 32), dtype=np.float32) + index_path = tmp_path / "cagra-fallback-index" + index_path.mkdir() + java_codec = runtime.resolve_codec(_CAGRA_CODEC) + telemetry = runtime.codec_telemetry(_CAGRA_CODEC, java_codec) + assert telemetry["writerPath"] == "gpu-cagra" + runtime.build_index( + index_path, + vectors, + _ResolvedCodec( + codec_name=_CAGRA_CODEC, + java_codec=java_codec, + telemetry=telemetry, + writer_path=telemetry["writerPath"], + ), + ) + + with pytest.raises(RuntimeError, match="persisted brute-force"): + runtime.verify_cagra_index(index_path, expected_vector_count=1) + + dataset = Dataset( + name="pylucene-cagra-fallback", + training_vectors=vectors, + query_vectors=vectors.copy(), + distance_metric="euclidean", + ) + index = IndexConfig( + name="pylucene-cagra-fallback", + algo="pylucene_cuvs_cagra", + build_param={"codec": _CAGRA_CODEC}, + search_params=[{}], + file=str(index_path), + ) + backend = PyLuceneBackend( + { + "name": index.name, + "algo": index.algo, + "codec": _CAGRA_CODEC, + **pylucene_runtime_config, + } + ) + backend._runtime = runtime + try: + result = backend.search(dataset, [index], k=1) + assert not result.success + assert "CAGRA provenance manifest is missing" in result.error_message + finally: + backend.cleanup() + + +def test_cli_build_and_search_with_real_pylucene_runtime( + tmp_path, pylucene_runtime_config +): + """Exercise the documented module CLI in a fresh PyLucene process.""" + rng = np.random.default_rng(174) + training_vectors = rng.standard_normal((512, 32)).astype(np.float32) + query_ids = np.asarray([0, 137, 259, 511], dtype=np.int64) + query_vectors = training_vectors[query_ids].copy() + k = 5 + squared_distances = np.sum( + (query_vectors[:, np.newaxis, :] - training_vectors[np.newaxis, :, :]) + ** 2, + axis=2, + ) + groundtruth_neighbors = np.argsort(squared_distances, axis=1)[:, :k] + groundtruth_distances = np.take_along_axis( + squared_distances, groundtruth_neighbors, axis=1 + ) + + dataset_name = "pylucene-cli-integration" + dataset_path = tmp_path / "datasets" + dataset_dir = dataset_path / dataset_name + _write_test_bin(dataset_dir / "base.fbin", training_vectors) + _write_test_bin(dataset_dir / "query.fbin", query_vectors) + _write_test_bin( + dataset_dir / "groundtruth.neighbors.ibin", + groundtruth_neighbors.astype(np.int32), + ) + _write_test_bin( + dataset_dir / "groundtruth.distances.fbin", + groundtruth_distances.astype(np.float32), + ) + + dataset_config = tmp_path / "datasets.yaml" + dataset_config.write_text( + json.dumps( + [ + { + "name": dataset_name, + "base_file": f"{dataset_name}/base.fbin", + "query_file": f"{dataset_name}/query.fbin", + "groundtruth_neighbors_file": ( + f"{dataset_name}/groundtruth.neighbors.ibin" + ), + "groundtruth_distances_file": ( + f"{dataset_name}/groundtruth.distances.fbin" + ), + "distance": "euclidean", + "dims": training_vectors.shape[1], + } + ] + ) + ) + backend_config = tmp_path / "pylucene-backend.yaml" + backend_config.write_text( + json.dumps({"backend": "pylucene", **pylucene_runtime_config}) + ) + + environment = os.environ.copy() + source_root = str(Path(__file__).resolve().parents[2]) + environment["PYTHONPATH"] = os.pathsep.join( + value + for value in (source_root, environment.get("PYTHONPATH")) + if value + ) + completed = subprocess.run( + [ + sys.executable, + "-m", + "cuvs_bench.run", + "--backend-config", + str(backend_config), + "--dataset-configuration", + str(dataset_config), + "--dataset", + dataset_name, + "--dataset-path", + str(dataset_path), + "--algorithms", + "pylucene_cuvs_hnsw", + "--groups", + "test", + "--batch-size", + "2", + "-k", + str(k), + "-m", + "latency", + "--build", + "--search", + "--force", + ], + cwd=tmp_path, + env=environment, + capture_output=True, + text=True, + timeout=600, + check=False, + ) + assert completed.returncode == 0, ( + f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}" + ) + + codec = "Lucene101AcceleratedHNSWCodec" + index_name = f"pylucene_cuvs_hnsw_test.codec{codec}" + index_path = dataset_dir / "index" / index_name + assert any(path.suffix == ".vex" for path in index_path.iterdir()) + assert (index_path / _HNSW_PROVENANCE_FILE).is_file() + + result_path = dataset_dir / "result" + build_json = result_path / "build" / "pylucene_cuvs_hnsw,test.json" + search_stem = f"pylucene_cuvs_hnsw,test,k{k},bs2" + search_json = result_path / "search" / f"{search_stem}.json" + build_rows = json.loads(build_json.read_text())["benchmarks"] + search_rows = json.loads(search_json.read_text())["benchmarks"] + assert len(build_rows) == 1 + assert len(search_rows) == 1 + + build_row = build_rows[0] + assert build_row["name"] == f"{index_name}/build" + assert build_row["real_time"] > 0 + assert build_row["index_size"] > 0 + assert build_row["codec"] == codec + assert build_row["writer_path"] == "gpu-hnsw" + _assert_hnsw_verification( + build_row, + codec, + training_vectors.shape[0], + training_vectors.shape[1], + ) + + search_row = search_rows[0] + assert search_row["name"] == f"{index_name}/search" + assert search_row["Recall"] >= 0.75 + assert search_row["items_per_second"] > 0 + assert search_row["Latency"] > 0 + assert search_row["search_time_ms"] > 0 + _assert_hnsw_verification( + search_row, + codec, + training_vectors.shape[0], + training_vectors.shape[1], + ) + + raw_csv = result_path / "search" / f"{search_stem},raw.csv" + with raw_csv.open(newline="") as file: + csv_rows = list(csv.DictReader(file)) + assert len(csv_rows) == 1 + csv_row = csv_rows[0] + assert csv_row["index_name"] == index_name + assert float(csv_row["recall"]) >= 0.75 + assert float(csv_row["throughput"]) > 0 + assert float(csv_row["latency"]) > 0 + assert float(csv_row["build time"]) > 0 + assert (result_path / "search" / f"{search_stem},latency.csv").is_file() + assert (result_path / "search" / f"{search_stem},throughput.csv").is_file() diff --git a/python/cuvs_bench/pyproject.toml b/python/cuvs_bench/pyproject.toml index 162564fcf4..9f582bb313 100644 --- a/python/cuvs_bench/pyproject.toml +++ b/python/cuvs_bench/pyproject.toml @@ -58,6 +58,7 @@ Homepage = "https://github.com/rapidsai/cuvs" [tool.pytest.ini_options] markers = [ "opensearch: tests that require a live OpenSearch node (run with '-m opensearch')", + "pylucene: real PyLucene/JVM/cuVS tests (set CUVS_BENCH_PYLUCENE_INTEGRATION=1 and run with '-m pylucene')", ] [tool.isort] From 57787eb4125b190a99c3775d7e049321ef0a7325 Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Fri, 31 Jul 2026 13:30:23 +0000 Subject: [PATCH 05/12] Document PyLucene backend setup and usage Document the verified dependency build, runtime configuration, supported codecs and limits, manual smoke workflow, index reuse behavior, and benchmark result semantics. --- fern/pages/cuvs_bench/install.md | 56 +++++++++++++++++ fern/pages/cuvs_bench/pluggable_backend.md | 4 ++ fern/pages/cuvs_bench/running.md | 71 +++++++++++++++++++++- 3 files changed, 129 insertions(+), 2 deletions(-) diff --git a/fern/pages/cuvs_bench/install.md b/fern/pages/cuvs_bench/install.md index 9c525b4973..bec57d8d29 100644 --- a/fern/pages/cuvs_bench/install.md +++ b/fern/pages/cuvs_bench/install.md @@ -49,6 +49,62 @@ Exact tags are listed on Docker Hub: **Note:** GPU containers use the CUDA toolkit inside the container. The host only needs a compatible driver, so CUDA 12 containers can run on systems with CUDA 13.x-capable drivers. GPU access also requires the NVIDIA Docker runtime from the [NVIDIA Container Toolkit](https://github.com/NVIDIA/nvidia-docker). +## PyLucene backend prerequisites + +The optional `pylucene` backend requires components that cuVS Bench does not install automatically: + +- An NVIDIA GPU supported by cuVS, plus matching CUDA and cuVS native libraries. +- JDK 22, `pytest`, and a [source-built PyLucene installation](https://lucene.apache.org/pylucene/install.html) compatible with the Lucene APIs used by the selected cuVS-Lucene checkout. The verified environment used PyLucene 10.0.0; the selected cuVS-Lucene POM compiles against Lucene 10.2.0, so validate the exact pair below. +- [Maven 3.9.6 or newer](https://maven.apache.org/download.cgi) to build cuVS-Lucene. +- The base `cuvs-java` JAR, standard `cuvs-lucene` JAR, and native libraries built from the pinned compatible revisions below. + +The required PyLucene service-provider and codec support is currently proposed in [NVIDIA/cuvs-lucene#174](https://github.com/NVIDIA/cuvs-lucene/pull/174). Until that work is merged and released, build cuVS-Lucene from that PR's branch; a release or `main` checkout without those changes is insufficient. + +Source compatibility is revision-specific while that PR is under review: the verified PR head `7d70d2f` does not compile against cuVS `main` at `f72199e3`. The commands below pin a source-compatible pair rather than a moving PR head. + +Build the dependency in a separate checkout so this does not change your cuVS Bench working tree: + +```bash +git clone https://github.com/NVIDIA/cuvs.git cuvs-pylucene-deps +cd cuvs-pylucene-deps +git switch --detach a59e2445 +./build.sh libcuvs java +cd .. +``` + +If matching native cuVS libraries are already built and installed, `./build.sh java` is sufficient. The Java build installs the base and native-classifier JARs into the local Maven repository; see the [cuVS Java build guide](https://github.com/NVIDIA/cuvs/blob/main/java/README.md). + +```bash +git clone https://github.com/NVIDIA/cuvs-lucene.git +cd cuvs-lucene +git fetch origin pull/174/head +git switch --detach 7d70d2f +mvn clean package -DskipTests +``` + +After the build, the conventional JAR paths are: + +```text +~/.m2/repository/com/nvidia/cuvs/cuvs-java//cuvs-java-.jar +/target/cuvs-lucene-.jar +``` + +Use the base `cuvs-java` JAR, not a native-classifier JAR. Use the standard cuVS-Lucene JAR, not its `-jar-with-dependencies`, sources, or Javadoc variants. Native-library paths must resolve `libcuvs.so`, `libcuvs_c.so`, their dependencies, and the CUDA runtime libraries from the pinned cuVS build. Then validate the artifacts from the cuVS-Lucene checkout: + +Use a clean environment without another cuVS native installation on its library path; otherwise, the JVM can load the other `libcuvs_c.so` first and reject the Java/native version mismatch. + +```bash +python -m pip install pytest + +CUVS_NATIVE_BUILD="$(cd ../cuvs-pylucene-deps/cpp/build && pwd)" +export JAVA_LIBRARY_PATH="$CUVS_NATIVE_BUILD:$CUVS_NATIVE_BUILD/c:/usr/local/cuda/lib64" +export LD_LIBRARY_PATH="$JAVA_LIBRARY_PATH${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" + +./test_pylucene.sh --full-e2e +``` + +See [Running the PyLucene backend](/user-guide/benchmarking-guide/cu-vs-bench-tool/usage#running-the-pylucene-backend) for a complete smoke workflow. + ## Build from Source Build cuVS Bench from source when you need local benchmark executables that match a development checkout, include custom algorithm targets, or use dependencies that are not available in the pre-built packages. diff --git a/fern/pages/cuvs_bench/pluggable_backend.md b/fern/pages/cuvs_bench/pluggable_backend.md index ce5eb2ec51..a72620350b 100644 --- a/fern/pages/cuvs_bench/pluggable_backend.md +++ b/fern/pages/cuvs_bench/pluggable_backend.md @@ -222,6 +222,10 @@ get_registry().register("elasticsearch", ElasticsearchBackend) | `BenchmarkBackend` | Abstract class whose `build(...)` method returns `BuildResult` and whose `search(...)` method returns `SearchResult`. Register with `BackendRegistry.register(name, backend_class)`. | | `BackendRegistry` | Singleton registry returned by `get_registry()`. It maps backend type names to backend classes. | +## PyLucene Backend + +The built-in `pylucene` loader expands algorithm YAML groups into one Lucene index per selected codec. The backend initializes PyLucene's process-global JVM and resolves cuVS-Lucene codecs through Lucene's service-provider interface. See [Installation](/user-guide/benchmarking-guide/cu-vs-bench-tool/installation#pylucene-backend-prerequisites) and [Usage](/user-guide/benchmarking-guide/cu-vs-bench-tool/usage#running-the-pylucene-backend) for the provisional dependency, configuration, and runtime limits. + ## C++ Backend The built-in `CppGoogleBenchmarkBackend` uses `backend_type="cpp_gbench"`. Its config loader reads YAML under `config/datasets` and `config/algos`, expands parameter combinations, and validates constraints. Its backend runs the C++ benchmark executables and merges their results. diff --git a/fern/pages/cuvs_bench/running.md b/fern/pages/cuvs_bench/running.md index 94606bb214..2dec4c8485 100644 --- a/fern/pages/cuvs_bench/running.md +++ b/fern/pages/cuvs_bench/running.md @@ -64,6 +64,7 @@ Create a custom YAML file with a `base` group to override the default benchmark | HNSWLIB | `hnswlib` | | DiskANN | `diskann_memory`, `diskann_ssd` | | NVIDIA cuVS | `cuvs_brute_force`, `cuvs_cagra`, `cuvs_ivf_flat`, `cuvs_ivf_pq`, `cuvs_cagra_hnswlib`, `cuvs_vamana` | +| PyLucene/cuVS | `pylucene_cuvs_hnsw`, `pylucene_cuvs_cagra` | ### Multi-GPU algorithms @@ -75,6 +76,70 @@ cuVS Bench includes single-node multi-GPU versions of IVF-Flat, IVF-PQ, and CAGR | IVF-PQ | `cuvs_mg_ivf_pq` | | CAGRA | `cuvs_mg_cagra` | +## Running the PyLucene backend + +The PyLucene backend type is `pylucene`. It runs through Python's embedded JVM, not the `*_ANN_BENCH` executables or `--executable-dir`. It accepts only FLOAT32 datasets with Euclidean distance and requires the prerequisites described in [Installation](/user-guide/benchmarking-guide/cu-vs-bench-tool/installation#pylucene-backend-prerequisites). + +Install the development checkout into the activated PyLucene environment: + +```bash +cd +python -m pip install -e ./python/cuvs_bench +``` + +Create `pylucene-backend.yaml` with absolute paths to the base `cuvs-java` JAR, the standard `cuvs-lucene` JAR, and a path-separated list of directories containing the matching cuVS and CUDA native libraries: + +```yaml +backend: pylucene +cuvs_java_jar: /home/user/.m2/repository/com/nvidia/cuvs/cuvs-java/VERSION/cuvs-java-VERSION.jar +cuvs_lucene_jar: /absolute/path/to/cuvs-lucene/target/cuvs-lucene-VERSION.jar +java_library_path: /work/cuvs-pylucene-deps/cpp/build:/work/cuvs-pylucene-deps/cpp/build/c:/usr/local/cuda/lib64 +``` + +Configuration sources are: + +| Backend configuration key | Environment variable | +| --- | --- | +| `cuvs_java_jar` | `CUVS_LUCENE_CUVS_JAVA_JAR` | +| `cuvs_lucene_jar` | `CUVS_LUCENE_JAR` | +| `java_library_path` | `JAVA_LIBRARY_PATH` or `LD_LIBRARY_PATH` | +| `jvm_args` | No environment alias; YAML list of additional JVM arguments. | + +Use one explicit dataset root for both preparation and execution. This small synthetic run is a functional smoke test: + +```bash +export DATASET_ROOT=/absolute/path/to/cuvs-bench-data + +python -m cuvs_bench.get_dataset \ + --dataset test-data \ + --dataset-path "$DATASET_ROOT" \ + --test-data-n-train 512 \ + --test-data-n-test 4 \ + --test-data-k 5 + +python -m cuvs_bench.run \ + --backend-config pylucene-backend.yaml \ + --dataset test-data \ + --dataset-path "$DATASET_ROOT" \ + --algorithms pylucene_cuvs_hnsw \ + --groups test \ + --batch-size 2 -k 5 \ + -m latency \ + --build --search --force +``` + +Do not use this tiny synthetic smoke run as performance or quality evidence; it exists only to verify the workflow. For a representative recall run, prepare `deep-image-96-angular` with `--normalize` into the same `DATASET_ROOT`, then run the backend against `deep-image-96-inner` with the same `--dataset-path`. + +The HNSW `base` group sweeps `Lucene101AcceleratedHNSWCodec`, `Lucene101AcceleratedHNSWBaseLayerCodec`, and `Lucene101AcceleratedHNSWMultiLayerCodec`; its `test` group uses `Lucene101AcceleratedHNSWCodec`. These codecs build the HNSW graph on the GPU with cuVS and use Lucene's CPU HNSW search. Use `--algorithms pylucene_cuvs_cagra` to run `CuVS2510GPUSearchCodec`, which builds and searches on the GPU. + +The supplied configurations use Lucene's public `KnnFloatVectorQuery` and do not expose search-time tuning parameters. The backend accepts FLOAT32 Euclidean datasets with at most 4096 dimensions. All supplied codecs require at least two indexed vectors because cuVS-Lucene bypasses cuVS for a single-vector index. CAGRA requires a GPU for build and search. Although cuVS-Lucene uses an effective `lucene_k` of `min(k, document_count)`, the backend conservatively requires `k <= 1024` to avoid cuVS-Lucene paths that can use brute-force search above that limit. HNSW requires a GPU for graph construction, and new builds reject cuVS-Lucene's CPU fallback. + +New HNSW and CAGRA builds atomically write commit-bound provenance manifests named `.cuvs-bench-pylucene-hnsw.json` and `.cuvs-bench-pylucene-cagra.json`, respectively. Reuse and search fail if the applicable manifest is missing, malformed, stale, or names a different codec. Indexes created outside this backend without the applicable manifest must be rebuilt with `--force`. CAGRA indexes also undergo structural and checksum-integrity verification. + +The backend currently supports latency mode with one search thread. `--batch-size` groups queries for measurement, and the reported latency percentiles are per batch. Throughput mode and multiple search threads are not implemented. + +PyLucene's JVM is process-global and can be initialized only once. Set the JAR, native-library locations, and `jvm_args` before the first PyLucene benchmark, and start a new Python process to change any of them. + ## Smaller-scale benchmarks (<1M to 10M vectors) Use `cuvs_bench.get_dataset` to prepare a built-in dataset. By default, datasets are stored under `RAPIDS_DATASET_ROOT_DIR` when that environment variable is set, or under a local `datasets` directory otherwise. @@ -192,7 +257,9 @@ Containers can also run in detached mode. ## Evaluating results -Build benchmarks report: +The tables below describe fields emitted by the default C++ Google Benchmark backend. Other backends report the fields that apply to their execution model, so not every field is present for every backend. In particular, PyLucene reports `index_size` as the total on-disk index size in bytes. + +C++ build benchmarks report: | Name | Description | | --- | --- | @@ -203,7 +270,7 @@ Build benchmarks report: | GPU | GPU time spent building. | | index_size | Number of vectors used to train the index. | -Search benchmarks report: +C++ search benchmarks report: | Name | Description | | --- | --- | From 1f206beb33a8e68dbb06afd458b29a87bdb4a37c Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Fri, 14 Aug 2026 19:46:55 +0000 Subject: [PATCH 06/12] Align Python backend results with current cuvs-bench Use the shared direct CSV exporter for in-process backends and remove the redundant JSON persistence layer. Preserve scoped result identities, latency percentiles, safe artifact paths, stale-result cleanup, and accurate CLI exit behavior. Signed-off-by: nvzm123 --- python/cuvs_bench/cuvs_bench/_validation.py | 61 ++++ python/cuvs_bench/cuvs_bench/backends/base.py | 65 +--- .../cuvs_bench/orchestrator/config_loaders.py | 5 - .../cuvs_bench/orchestrator/orchestrator.py | 16 +- .../cuvs_bench/orchestrator/result_files.py | 232 ------------ python/cuvs_bench/cuvs_bench/run/__main__.py | 67 +++- .../cuvs_bench/cuvs_bench/run/data_export.py | 193 +++++----- .../cuvs_bench/cuvs_bench/tests/test_cli.py | 162 ++++++--- .../cuvs_bench/tests/test_data_export.py | 205 ++++++++++- .../cuvs_bench/tests/test_registry.py | 338 ------------------ 10 files changed, 559 insertions(+), 785 deletions(-) create mode 100644 python/cuvs_bench/cuvs_bench/_validation.py delete mode 100644 python/cuvs_bench/cuvs_bench/orchestrator/result_files.py 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/base.py b/python/cuvs_bench/cuvs_bench/backends/base.py index 5f70586a6a..1c9ebfeb55 100644 --- a/python/cuvs_bench/cuvs_bench/backends/base.py +++ b/python/cuvs_bench/cuvs_bench/backends/base.py @@ -250,23 +250,15 @@ def to_json(self) -> Dict[str, Any]: Dict[str, Any] Dictionary with benchmark results """ - result = { + return { + "name": f"{self.algorithm}/build", + "real_time": self.build_time_seconds, + "time_unit": "s", + "index_size": self.index_size_bytes, + "success": self.success, **self.build_params, **self.metadata, } - result.pop("error_message", None) - result.update( - { - "name": f"{self.algorithm}/build", - "real_time": self.build_time_seconds, - "time_unit": "s", - "index_size": self.index_size_bytes, - "success": self.success, - } - ) - if self.error_message is not None: - result["error_message"] = self.error_message - return result @dataclass @@ -303,10 +295,6 @@ class SearchResult: Whether the search succeeded error_message : Optional[str] Error message if search failed - latency_seconds : Optional[float] - Backend-reported mean latency in seconds. For batched backends, this - is mean batch latency. This field is last to preserve - compatibility with existing positional construction. """ neighbors: np.ndarray @@ -322,7 +310,6 @@ class SearchResult: metadata: Dict[str, Any] = field(default_factory=dict) success: bool = True error_message: Optional[str] = None - latency_seconds: Optional[float] = None def to_json(self) -> Dict[str, Any]: """ @@ -333,39 +320,26 @@ def to_json(self) -> Dict[str, Any]: Dict[str, Any] Dictionary with benchmark results """ - result = dict(self.metadata) - result.pop("error_message", None) + result = { + "name": f"{self.algorithm}/search", + "real_time": self.search_time_ms, + "time_unit": "ms", + "items_per_second": self.queries_per_second, + "Recall": self.recall, + "success": self.success, + "search_params": self.search_params, + **self.metadata, + } + if self.latency_percentiles: result.update(self.latency_percentiles) - real_time_ms = ( - self.latency_seconds * 1000.0 - if self.latency_seconds is not None - else self.search_time_ms - ) - result.update( - { - "name": f"{self.algorithm}/search", - "real_time": real_time_ms, - "time_unit": "ms", - "search_time_ms": self.search_time_ms, - "items_per_second": self.queries_per_second, - "Recall": self.recall, - "success": self.success, - "search_params": self.search_params, - } - ) - if self.latency_seconds is not None: - result["Latency"] = self.latency_seconds if self.gpu_time_seconds is not None: result["GPU"] = self.gpu_time_seconds if self.cpu_time_seconds is not None: result["cpu_time"] = self.cpu_time_seconds - if self.error_message is not None: - result["error_message"] = self.error_message - return result @@ -401,11 +375,6 @@ class BenchmarkBackend(ABC): - requires_network : bool - Whether backend requires network (default: False) """ - # Opt in when this backend's sweep results should be persisted by the - # orchestrator. Tune-mode and existing backends retain their own - # result-file behavior. - orchestrator_persists_results = False - def __init__(self, config: Dict[str, Any]): """Initialize backend with configuration.""" self.config = config diff --git a/python/cuvs_bench/cuvs_bench/orchestrator/config_loaders.py b/python/cuvs_bench/cuvs_bench/orchestrator/config_loaders.py index 37be49f269..7d6375ce96 100644 --- a/python/cuvs_bench/cuvs_bench/orchestrator/config_loaders.py +++ b/python/cuvs_bench/cuvs_bench/orchestrator/config_loaders.py @@ -91,11 +91,6 @@ def search_params_list(self) -> List[Dict[str, Any]]: def index_path(self) -> Path: return Path(self.indexes[0].file) if self.indexes else Path("") - @property - def output_filename(self) -> Tuple[str, str]: - """Return the build and search result filename stems.""" - return self.backend_config.get("output_filename", ("", "")) - @dataclass class DatasetConfig: diff --git a/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py b/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py index c070d22fa9..abffad0217 100644 --- a/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py +++ b/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py @@ -243,6 +243,7 @@ def _run_sweep( dry_run=dry_run, ) results.append(build_result) + if not build_result.success: print( f"Build failed for {config.index_name}: {build_result.error_message}" @@ -506,17 +507,10 @@ def objective(trial) -> float: print(f"Best params: {study.best_params}") print(f"Best {optimize_metric}: {study.best_value:.4f}") except ValueError: - if any(result.success for result in all_results): - print( - f"\n⚠️ All {n_trials} trials were pruned " - "(constraints not met)" - ) - print(f" Consider relaxing constraints: {hard_constraints}") - else: - print( - f"\n⚠️ All {n_trials} benchmark trials failed before " - "producing valid metrics" - ) + print( + f"\n⚠️ All {n_trials} trials were pruned (constraints not met)" + ) + print(f" Consider relaxing constraints: {hard_constraints}") return all_results diff --git a/python/cuvs_bench/cuvs_bench/orchestrator/result_files.py b/python/cuvs_bench/cuvs_bench/orchestrator/result_files.py deleted file mode 100644 index 8985aa0051..0000000000 --- a/python/cuvs_bench/cuvs_bench/orchestrator/result_files.py +++ /dev/null @@ -1,232 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# - -"""Persistence for benchmark backends that return in-process results.""" - -from __future__ import annotations - -import json -import re -import stat -import tempfile -from collections import defaultdict -from dataclasses import dataclass -from pathlib import Path -from typing import Any, Dict, Iterable, List, Tuple, Union - -from ..backends.base import BuildResult, SearchResult - -BenchmarkResult = Union[BuildResult, SearchResult] -_SEARCH_SUFFIX_PATTERN = re.compile(r"k[1-9]\d*,bs[1-9]\d*") - - -@dataclass(frozen=True) -class ResultRecord: - """Associate a backend result with its benchmark artifact identity.""" - - result: BenchmarkResult - index_name: str - output_filename: str - - @property - def phase(self) -> str: - return "build" if isinstance(self.result, BuildResult) else "search" - - def to_json(self) -> Dict[str, Any]: - row = self.result.to_json() - row["name"] = f"{self.index_name}/{self.phase}" - return row - - -def _validate_path_component(value: object, description: str) -> str: - if not isinstance(value, str): - raise ValueError(f"Invalid {description}: {value!r}") - path = Path(value) - if not value or path.name != value or value in {".", ".."}: - raise ValueError(f"Invalid {description}: {value!r}") - return value - - -def validate_dataset_name(dataset: object) -> str: - """Validate a dataset name before using it as a path component.""" - return _validate_path_component(dataset, "benchmark dataset name") - - -def _validate_output_filename(output_filename: str) -> None: - _validate_path_component(output_filename, "benchmark result filename") - - -def validate_sweep_output_filenames( - output_filenames: object, -) -> Tuple[str, str]: - """Validate and normalize opt-in sweep result filename stems.""" - if ( - not isinstance(output_filenames, (list, tuple)) - or len(output_filenames) != 2 - ): - raise ValueError( - "Opt-in sweep persistence requires exactly two result filename " - "stems: (build_stem, search_stem)" - ) - - build_stem, search_stem = output_filenames - _validate_output_filename(build_stem) - _validate_output_filename(search_stem) - - build_parts = build_stem.split(",") - if len(build_parts) < 2 or any(not part for part in build_parts): - raise ValueError( - "Build result filename stem must contain at least " - "','" - ) - - expected_prefix = f"{build_stem}," - search_suffix = ( - search_stem[len(expected_prefix) :] - if search_stem.startswith(expected_prefix) - else "" - ) - if _SEARCH_SUFFIX_PATTERN.fullmatch(search_suffix) is None: - raise ValueError( - "Search result filename stem must be " - "',k,bs'" - ) - - return build_stem, search_stem - - -def _read_benchmarks(path: Path) -> List[Dict[str, Any]]: - if not path.exists(): - return [] - with path.open(encoding="utf-8") as file: - payload = json.load(file) - benchmarks = ( - payload.get("benchmarks") if isinstance(payload, dict) else None - ) - if not isinstance(benchmarks, list): - raise ValueError( - f"Benchmark result file must contain a 'benchmarks' list: {path}" - ) - seen_names = set() - for position, row in enumerate(benchmarks): - if not isinstance(row, dict) or not isinstance(row.get("name"), str): - raise ValueError( - f"Benchmark row {position} must be an object with a string " - f"'name': {path}" - ) - if row["name"] in seen_names: - raise ValueError( - f"Benchmark result file contains duplicate name " - f"{row['name']!r}: {path}" - ) - seen_names.add(row["name"]) - return benchmarks - - -def _merge_build_rows( - path: Path, records: List[ResultRecord] -) -> List[Dict[str, Any]]: - rows_by_name = {row["name"]: row for row in _read_benchmarks(path)} - for record in records: - row = record.to_json() - previous = rows_by_name.get(row["name"]) - if ( - record.result.success - and record.result.metadata.get("skipped") - and previous is not None - and previous.get("success", True) is True - ): - continue - rows_by_name[row["name"]] = row - return list(rows_by_name.values()) - - -def _write_payload(path: Path, rows: List[Dict[str, Any]]) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - file_mode = stat.S_IMODE(path.stat().st_mode) if path.exists() else 0o644 - temp_path = None - try: - with tempfile.NamedTemporaryFile( - mode="w", - encoding="utf-8", - dir=path.parent, - prefix=f".{path.name}.", - suffix=".tmp", - delete=False, - ) as file: - temp_path = Path(file.name) - json.dump({"benchmarks": rows}, file, indent=2, allow_nan=False) - file.write("\n") - temp_path.chmod(file_mode) - temp_path.replace(path) - finally: - if temp_path is not None: - temp_path.unlink(missing_ok=True) - - -def _invalidate_derived_csvs(path: Path, phase: str) -> None: - suffixes = ( - (".csv",) - if phase == "build" - else (",raw.csv", ",throughput.csv", ",latency.csv") - ) - stem = path.with_suffix("") - for suffix in suffixes: - stem.with_name(f"{stem.name}{suffix}").unlink(missing_ok=True) - - -def _remove_result_artifacts(path: Path, phase: str) -> None: - path.unlink(missing_ok=True) - _invalidate_derived_csvs(path, phase) - - -def write_result_files( - records: List[ResultRecord], - dataset: str, - dataset_path: str, - *, - stale_search_filenames: Iterable[str] = (), -) -> None: - """Write grouped Google Benchmark-compatible JSON result files.""" - validate_dataset_name(dataset) - grouped: Dict[Tuple[str, str], List[ResultRecord]] = defaultdict(list) - for record in records: - _validate_output_filename(record.output_filename) - grouped[(record.phase, record.output_filename)].append(record) - - stale_search_filenames = set(stale_search_filenames) - for output_filename in stale_search_filenames: - _validate_output_filename(output_filename) - - result_root = Path(dataset_path) / dataset / "result" - current_search_filenames = { - output_filename - for phase, output_filename in grouped - if phase == "search" - } - for output_filename in stale_search_filenames - current_search_filenames: - _remove_result_artifacts( - result_root / "search" / f"{output_filename}.json", - "search", - ) - - for (phase, output_filename), file_records in grouped.items(): - output_path = result_root / phase / f"{output_filename}.json" - rows = ( - _merge_build_rows(output_path, file_records) - if phase == "build" - else [record.to_json() for record in file_records] - ) - if rows: - _write_payload(output_path, rows) - _invalidate_derived_csvs(output_path, phase) - - -__all__ = [ - "ResultRecord", - "validate_dataset_name", - "validate_sweep_output_filenames", - "write_result_files", -] diff --git a/python/cuvs_bench/cuvs_bench/run/__main__.py b/python/cuvs_bench/cuvs_bench/run/__main__.py index 897eb4edb1..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 " @@ -311,24 +346,11 @@ def main( **backend_kwargs, ) - failures = [result for result in results if not result.success] - if mode == "sweep": - run_failed = not results or bool(failures) - else: - run_failed = not any(result.success for result in results) - if run_failed: - details = ( - "; ".join( - f"{result.algorithm}: " - f"{result.error_message or 'benchmark failed'}" - for result in failures - ) - if failures - else f"{mode} mode produced no benchmark results" - ) - raise click.ClickException(details) - if dry_run: + # 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": @@ -337,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 6a3be8350a..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( @@ -32,16 +37,7 @@ ) skip_search_cols = ( - set( - [ - "recall", - "qps", - "latency", - "items_per_second", - "Recall", - "Latency", - ] - ) + set(["recall", "qps", "latency", "items_per_second", "Recall", "Latency"]) | skip_build_cols ) @@ -62,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") @@ -73,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: @@ -107,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") @@ -140,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: @@ -157,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"], @@ -188,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( @@ -203,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") @@ -212,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() @@ -221,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. @@ -240,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 @@ -276,15 +348,6 @@ def clean_algo_name(algo_name): return name.removesuffix(".json") -def _remove_derived_csvs(file, suffixes): - stem = os.path.splitext(file)[0] - for suffix in suffixes: - try: - os.remove(f"{stem}{suffix}") - except FileNotFoundError: - pass - - def write_csv(file, algo_name, df, extra_columns=None, skip_cols=None): """ Write a DataFrame to CSV with specified columns skipped. @@ -335,9 +398,6 @@ def convert_json_to_csv_build(dataset, dataset_path): """ for file, algo_name, df in read_json_files(dataset, dataset_path, "build"): try: - if df.empty: - _remove_derived_csvs(file, (".csv",)) - continue algo_name = clean_algo_name(algo_name) write_csv(file, algo_name, df, skip_cols=skip_build_cols) except Exception as e: @@ -360,26 +420,12 @@ def convert_json_to_csv_search(dataset, dataset_path): dataset, dataset_path, "search" ): try: - if df.empty: - _remove_derived_csvs( - file, (",raw.csv", ",throughput.csv", ",latency.csv") - ) - continue - search_stem = os.path.splitext(os.path.basename(file))[0] - # Search stems end in k and batch-size components. Preserve every - # preceding component, including an optional dataset subset. - search_stem_parts = search_stem.rsplit(",", maxsplit=2) - build_stem = ( - search_stem_parts[0] - if len(search_stem_parts) == 3 - else ",".join(algo_name) - ) build_file = os.path.join( dataset_path, dataset, "result", "build", - f"{build_stem}.csv", + f"{','.join(algo_name)}.csv", ) algo_name = clean_algo_name(algo_name) df["name"] = df["name"].str.split("/").str[0] @@ -401,46 +447,29 @@ def convert_json_to_csv_search(dataset, dataset_path): ] if os.path.exists(build_file): build_df = pd.read_csv(build_file) + write_ncols = len(write.columns) write["build time"] = None write["build threads"] = None write["build cpu_time"] = None - build_columns = { - "time": "build time", - "threads": "build threads", - "cpu_time": "build cpu_time", - } + start_idx = 5 if "GPU" in build_df.columns: + start_idx = 6 write["build GPU"] = None - build_columns["GPU"] = "build GPU" - - ignored_columns = {"algo_name", "index_name"} - for col_name in build_df.columns: - if ( - col_name in ignored_columns - or col_name in build_columns - ): - continue - target_name = ( - "build_num_threads" - if col_name == "num_threads" - else col_name - ) - if target_name not in write.columns: - write[target_name] = None - build_columns[col_name] = target_name - + for col_idx in range(start_idx, len(build_df.columns)): + col_name = build_df.columns[col_idx] + write[col_name] = None + if col_name == "num_threads": + write["build_num_threads"] = None for s_index, search_row in write.iterrows(): - for _, build_row in build_df.iterrows(): + for b_index, build_row in build_df.iterrows(): if search_row["index_name"] == build_row["index_name"]: - for ( - source_name, - target_name, - ) in build_columns.items(): - if source_name in build_df.columns: - write.at[s_index, target_name] = build_row[ - source_name - ] + write.iloc[s_index, write_ncols] = build_df.iloc[ + b_index, 2 + ] + write.iloc[s_index, write_ncols + 1 :] = ( + build_df.iloc[b_index, 3:] + ) break # Write search data and compute frontiers write.to_csv(file.replace(".json", ",raw.csv"), index=False) diff --git a/python/cuvs_bench/cuvs_bench/tests/test_cli.py b/python/cuvs_bench/cuvs_bench/tests/test_cli.py index a1525d234d..a9dd27e76d 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_cli.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_cli.py @@ -9,6 +9,7 @@ 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 @@ -525,7 +526,15 @@ def test_plot_command_creates_png_files(temp_datasets_dir: Path): # without requiring native benchmark execution. -def _invoke_run_with_results(monkeypatch, tmp_path, mode, results): +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( @@ -539,47 +548,66 @@ def _invoke_run_with_results(monkeypatch, tmp_path, mode, results): exported = [] monkeypatch.setattr( run_module, - "convert_json_to_csv_build", - lambda dataset, dataset_path: exported.append( - ("build", dataset, dataset_path) - ), - ) - monkeypatch.setattr( - run_module, - "convert_json_to_csv_search", - lambda dataset, dataset_path: exported.append( - ("search", dataset, dataset_path) - ), + "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, - [ - "--dataset", - "test-data", - "--dataset-path", - str(tmp_path), - "--algorithms", - "fake", - "--groups", - "base", - "--batch-size", - "1", - "-k", - "1", - "-m", - "latency", - "--mode", - mode, - ], + args, ) return result, exported -def _benchmark_result(success, error_message=None): - return SimpleNamespace( +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, ) @@ -590,22 +618,21 @@ def test_tune_mixed_trial_results_exit_zero_and_export(monkeypatch, tmp_path): tmp_path, "tune", [ - _benchmark_result(False, "trial failed"), - _benchmark_result(True), + _search_result(False, "trial failed"), + _search_result(), ], ) assert result.exit_code == 0, result.output - assert exported == [ - ("build", "test-data", str(tmp_path)), - ("search", "test-data", str(tmp_path)), - ] + assert len(exported) == 1 + assert exported[0][1]["search_requested"] is True @pytest.mark.parametrize( ("results", "expected_message"), [ - ([_benchmark_result(False, "trial failed")], "trial failed"), + ([_search_result(False, "trial failed")], "trial failed"), + ([_build_result()], "no successful search result"), ([], "tune mode produced no benchmark results"), ], ) @@ -618,7 +645,7 @@ def test_tune_without_successful_results_exits_nonzero( assert result.exit_code != 0 assert expected_message in result.output - assert exported == [] + assert len(exported) == 1 def test_tune_all_constraint_pruned_measurements_exit_zero_and_export( @@ -630,14 +657,11 @@ def test_tune_all_constraint_pruned_measurements_exit_zero_and_export( monkeypatch, tmp_path, "tune", - [_benchmark_result(True), _benchmark_result(True)], + [_search_result(), _search_result()], ) assert result.exit_code == 0, result.output - assert exported == [ - ("build", "test-data", str(tmp_path)), - ("search", "test-data", str(tmp_path)), - ] + assert len(exported) == 1 def test_sweep_mixed_results_exit_nonzero(monkeypatch, tmp_path): @@ -646,14 +670,14 @@ def test_sweep_mixed_results_exit_nonzero(monkeypatch, tmp_path): tmp_path, "sweep", [ - _benchmark_result(False, "benchmark failed"), - _benchmark_result(True), + _build_result(False, "benchmark failed"), + _search_result(), ], ) assert result.exit_code != 0 assert "benchmark failed" in result.output - assert exported == [] + assert len(exported) == 1 def test_sweep_without_results_exits_nonzero(monkeypatch, tmp_path): @@ -663,6 +687,48 @@ def test_sweep_without_results_exits_nonzero(monkeypatch, tmp_path): 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 == [] 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_registry.py b/python/cuvs_bench/cuvs_bench/tests/test_registry.py index 8c79223c4f..eadf804db3 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_registry.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_registry.py @@ -6,8 +6,6 @@ Unit tests for the backend registry system. """ -import json - import pytest import numpy as np @@ -21,12 +19,6 @@ register_backend, get_backend, ) -from cuvs_bench.orchestrator import ( - BenchmarkConfig, - BenchmarkOrchestrator, - DatasetConfig, -) -from cuvs_bench.orchestrator.config_loaders import IndexConfig class DummyBackend(BenchmarkBackend): @@ -209,28 +201,6 @@ def test_build_result_to_json(self): assert json_result["nlist"] == 1024 assert json_result["gpu_time"] == 4.2 - def test_build_result_core_fields_cannot_be_overridden(self): - result = BuildResult( - index_path="/path/to/index", - build_time_seconds=5.5, - index_size_bytes=1024000, - algorithm="test_algo", - build_params={"name": "wrong", "real_time": -1}, - metadata={ - "time_unit": "ms", - "success": False, - "error_message": "wrong", - }, - ) - - json_result = result.to_json() - - assert json_result["name"] == "test_algo/build" - assert json_result["real_time"] == 5.5 - assert json_result["time_unit"] == "s" - assert json_result["success"] is True - assert "error_message" not in json_result - class TestSearchResult: """Tests for SearchResult dataclass.""" @@ -254,36 +224,6 @@ def test_search_result_creation(self): assert result.queries_per_second == 20.0 assert result.success is True - def test_search_result_preserves_legacy_positional_arguments(self): - neighbors = np.array([[1]]) - distances = np.array([[0.1]]) - latency_percentiles = {"p50": 1.0} - metadata = {"source": "legacy"} - - result = SearchResult( - neighbors, - distances, - 2.0, - 500.0, - 0.9, - "test_algo", - [{}], - latency_percentiles, - 0.5, - 0.75, - metadata, - False, - "failed", - ) - - assert result.latency_percentiles is latency_percentiles - assert result.gpu_time_seconds == 0.5 - assert result.cpu_time_seconds == 0.75 - assert result.metadata is metadata - assert result.success is False - assert result.error_message == "failed" - assert result.latency_seconds is None - def test_search_result_to_json(self): """Test conversion to JSON format.""" neighbors = np.array([[1, 2, 3]]) @@ -310,284 +250,6 @@ def test_search_result_to_json(self): assert json_result["p50"] == 50.0 assert json_result["p95"] == 95.0 - def test_search_result_core_fields_cannot_be_overridden(self): - result = SearchResult( - neighbors=np.array([[1]]), - distances=np.array([[0.1]]), - search_time_ms=6.0, - queries_per_second=500.0, - recall=0.95, - algorithm="test_algo", - search_params=[{}], - latency_seconds=0.003, - gpu_time_seconds=0.004, - cpu_time_seconds=0.005, - metadata={ - "name": "wrong", - "real_time": -1, - "Latency": -1, - "GPU": -1, - "cpu_time": -1, - "Recall": -1, - "success": False, - "error_message": "wrong", - }, - ) - - json_result = result.to_json() - - assert json_result["name"] == "test_algo/search" - assert json_result["real_time"] == 3.0 - assert json_result["Latency"] == 0.003 - assert json_result["GPU"] == 0.004 - assert json_result["cpu_time"] == 0.005 - assert json_result["Recall"] == 0.95 - assert json_result["success"] is True - assert "error_message" not in json_result - - def test_search_result_optional_metadata_fields_are_fallbacks(self): - result = SearchResult( - neighbors=np.array([[1]]), - distances=np.array([[0.1]]), - search_time_ms=6.0, - queries_per_second=500.0, - recall=0.95, - algorithm="test_algo", - search_params=[{}], - metadata={ - "Latency": 0.006, - "GPU": 0.004, - "cpu_time": 0.005, - }, - ) - - json_result = result.to_json() - - assert json_result["Latency"] == 0.006 - assert json_result["GPU"] == 0.004 - assert json_result["cpu_time"] == 0.005 - - -def test_backend_without_result_stems_runs_without_generic_persistence( - tmp_path, monkeypatch -): - class DummyConfigLoader: - def load(self, **_kwargs): - return DatasetConfig(name="dummy"), [ - BenchmarkConfig( - indexes=[ - IndexConfig( - name="dummy-index", - algo="dummy", - build_param={}, - search_params=[{}], - file=str(tmp_path / "dummy-index"), - ) - ], - backend_config={"name": "dummy-index"}, - ) - ] - - monkeypatch.setattr( - "cuvs_bench.orchestrator.orchestrator.get_backend_class", - lambda _backend_type: DummyBackend, - ) - monkeypatch.setattr( - "cuvs_bench.orchestrator.orchestrator.get_config_loader", - lambda _backend_type: DummyConfigLoader, - ) - orchestrator = BenchmarkOrchestrator(backend_type="dummy") - - results = orchestrator.run_benchmark( - build=True, - search=False, - dataset="dummy", - dataset_path=str(tmp_path), - ) - - assert len(results) == 1 - assert results[0].success is True - assert not (tmp_path / "dummy" / "result").exists() - - -def test_failed_opt_in_build_clears_only_matching_stale_search_artifacts( - tmp_path, monkeypatch -): - build_stem = "dummy,base" - search_stem = f"{build_stem},k1,bs1" - - class DummyConfigLoader: - def load(self, **_kwargs): - return DatasetConfig(name="dummy"), [ - BenchmarkConfig( - indexes=[ - IndexConfig( - name="dummy-index", - algo="dummy", - build_param={}, - search_params=[{}], - file=str(tmp_path / "dummy-index"), - ) - ], - backend_config={ - "name": "dummy-index", - "output_filename": (build_stem, search_stem), - }, - ) - ] - - class FailedBuildBackend(DummyBackend): - orchestrator_persists_results = True - - def build(self, dataset, indexes, force=False, dry_run=False): - return BuildResult( - index_path=indexes[0].file, - build_time_seconds=0.0, - index_size_bytes=0, - algorithm=self.algo, - build_params={}, - success=False, - error_message="expected build failure", - ) - - def search(self, *args, **kwargs): - raise AssertionError("search must not run after a failed build") - - result_root = tmp_path / "dummy" / "result" - search_dir = result_root / "search" - search_dir.mkdir(parents=True) - stale_paths = [ - search_dir / f"{search_stem}.json", - search_dir / f"{search_stem},raw.csv", - search_dir / f"{search_stem},throughput.csv", - search_dir / f"{search_stem},latency.csv", - ] - unrelated_paths = [ - search_dir / "other,base,k1,bs1.json", - search_dir / "other,base,k1,bs1,raw.csv", - ] - for path in [*stale_paths, *unrelated_paths]: - path.write_text("stale\n", encoding="utf-8") - - monkeypatch.setattr( - "cuvs_bench.orchestrator.orchestrator.get_backend_class", - lambda _backend_type: FailedBuildBackend, - ) - monkeypatch.setattr( - "cuvs_bench.orchestrator.orchestrator.get_config_loader", - lambda _backend_type: DummyConfigLoader, - ) - - results = BenchmarkOrchestrator(backend_type="dummy").run_benchmark( - build=True, - search=True, - count=1, - batch_size=1, - dataset="dummy", - dataset_path=str(tmp_path), - ) - - assert len(results) == 1 - assert results[0].success is False - assert not any(path.exists() for path in stale_paths) - assert all( - path.read_text(encoding="utf-8") == "stale\n" - for path in unrelated_paths - ) - - build_path = result_root / "build" / f"{build_stem}.json" - build_rows = json.loads(build_path.read_text(encoding="utf-8"))[ - "benchmarks" - ] - assert len(build_rows) == 1 - assert build_rows[0]["success"] is False - assert build_rows[0]["error_message"] == "expected build failure" - - -@pytest.mark.parametrize( - ("index_count", "output_filenames", "expected_message"), - [ - ( - 2, - ("dummy,base", "dummy,base,k1,bs1"), - "exactly one index", - ), - ( - 1, - ("dummy", "dummy,k1,bs1"), - "Build result filename stem", - ), - ( - 1, - ("dummy,base", "other,base,k1,bs1"), - "Search result filename stem", - ), - ( - 1, - ("dummy,base",), - "exactly two result filename stems", - ), - ( - 1, - ("../dummy,base", "../dummy,base,k1,bs1"), - "Invalid benchmark result filename", - ), - ], -) -def test_opt_in_sweep_persistence_validates_artifact_identity( - tmp_path, - monkeypatch, - index_count, - output_filenames, - expected_message, -): - indexes = [ - IndexConfig( - name=f"dummy-index-{position}", - algo="dummy", - build_param={}, - search_params=[{}], - file=str(tmp_path / f"dummy-index-{position}"), - ) - for position in range(index_count) - ] - - class DummyConfigLoader: - def load(self, **_kwargs): - return DatasetConfig(name="dummy"), [ - BenchmarkConfig( - indexes=indexes, - backend_config={ - "name": "dummy-index", - "output_filename": output_filenames, - }, - ) - ] - - class PersistedDummyBackend(DummyBackend): - orchestrator_persists_results = True - - monkeypatch.setattr( - "cuvs_bench.orchestrator.orchestrator.get_backend_class", - lambda _backend_type: PersistedDummyBackend, - ) - monkeypatch.setattr( - "cuvs_bench.orchestrator.orchestrator.get_config_loader", - lambda _backend_type: DummyConfigLoader, - ) - - with pytest.raises(ValueError, match=expected_message): - BenchmarkOrchestrator(backend_type="dummy").run_benchmark( - build=True, - search=False, - count=1, - batch_size=1, - dataset="dummy", - dataset_path=str(tmp_path), - ) - - assert not (tmp_path / "dummy" / "result").exists() - class TestBackendRegistry: """Tests for BackendRegistry.""" From 2ab45a9444f681b3fa63ce96ae262ac424fa866e Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Fri, 14 Aug 2026 19:49:00 +0000 Subject: [PATCH 07/12] Update PyLucene backend for current cuVS-Lucene codecs Target the production HNSW and CAGRA codecs exposed by the current cuVS-Lucene PR. Preserve the intentional HNSW CPU fallback, keep fail-closed CAGRA integrity checks, and verify actual HNSW writer selection through a downstream test-only codec adapter. Signed-off-by: nvzm123 --- .../cuvs_bench/backends/pylucene.py | 278 +++++++++--------- .../config/algos/pylucene_cuvs_hnsw.yaml | 2 - .../cuvs_bench/tests/_pylucene_test_utils.py | 14 +- .../cuvs_bench/tests/test_pylucene_backend.py | 265 ++++++++++------- .../tests/test_pylucene_cli_config.py | 161 ++++++---- .../tests/test_pylucene_integration.py | 263 ++++++++++------- .../tests/test_pylucene_provenance.py | 16 +- .../cuvs_bench/tests/test_pylucene_runtime.py | 134 +++++---- .../bench/PyLuceneWriterSelectionCodec.java | 63 ++++ 9 files changed, 695 insertions(+), 501 deletions(-) create mode 100644 python/cuvs_bench/tests/java/com/nvidia/cuvs/bench/PyLuceneWriterSelectionCodec.java diff --git a/python/cuvs_bench/cuvs_bench/backends/pylucene.py b/python/cuvs_bench/cuvs_bench/backends/pylucene.py index 8a10fcdfbe..9e5f620510 100644 --- a/python/cuvs_bench/cuvs_bench/backends/pylucene.py +++ b/python/cuvs_bench/cuvs_bench/backends/pylucene.py @@ -19,13 +19,11 @@ from dataclasses import dataclass from enum import Enum from pathlib import Path -from types import MappingProxyType from typing import ( Any, Callable, Dict, List, - Mapping, NamedTuple, Optional, Tuple, @@ -35,6 +33,7 @@ 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, @@ -48,18 +47,12 @@ _VECTOR_FIELD = "vector" _MAX_DIMENSIONS = 4096 -_HNSW_CODECS = frozenset( - { - "Lucene101AcceleratedHNSWCodec", - "Lucene101AcceleratedHNSWBaseLayerCodec", - "Lucene101AcceleratedHNSWMultiLayerCodec", - } -) +_HNSW_CODEC = "Lucene101AcceleratedHNSWCodec" _CAGRA_CODEC = "CuVS2510GPUSearchCodec" -_SUPPORTED_CODECS = _HNSW_CODECS | {_CAGRA_CODEC} +_SUPPORTED_CODECS = frozenset({_HNSW_CODEC, _CAGRA_CODEC}) _SUPPORTED_BUILD_KEYS = frozenset({"codec"}) -_EXPECTED_WRITER_PATH = { - **{codec: "gpu-hnsw" for codec in _HNSW_CODECS}, +_EXPECTED_WRITER_POLICY = { + _HNSW_CODEC: "gpu-with-cpu-fallback", _CAGRA_CODEC: "gpu-cagra", } _CAGRA_META_EXTENSION = ".vemc" @@ -76,12 +69,12 @@ _HNSW_PROVENANCE_FILE = ".cuvs-bench-pylucene-hnsw.json" _CAGRA_PROVENANCE_FILE = ".cuvs-bench-pylucene-cagra.json" -_PROVENANCE_SCHEMA_VERSION = 1 +_PROVENANCE_SCHEMA_VERSION = 2 _PROVENANCE_KEYS = frozenset( { "schema_version", "codec", - "writer_path", + "writer_policy", "vector_count", "dimensions", "commit_fingerprints", @@ -294,26 +287,6 @@ def _initialize_pylucene(config: Dict[str, Any]) -> Any: return lucene -def _parse_writer_telemetry(description: str) -> Dict[str, str]: - """Parse cuVS-Lucene vector-format diagnostics.""" - _, separator, payload = description.partition("(") - if not separator or not payload.endswith(")"): - raise RuntimeError( - f"Malformed cuVS-Lucene writer telemetry: {description!r}" - ) - - telemetry: Dict[str, str] = {} - for item in payload[:-1].split(";"): - key, item_separator, value = item.partition("=") - if not item_separator or not key: - raise RuntimeError( - "Malformed cuVS-Lucene writer telemetry item " - f"{item!r} in {description!r}" - ) - telemetry[key] = value - return telemetry - - def _score_to_squared_euclidean(score: float) -> float: """Convert Lucene's Euclidean score to squared Euclidean distance.""" if score <= 0.0: @@ -390,7 +363,7 @@ def _validate_search_params(search_params: Any) -> List[Dict[str, Any]]: return search_params -def _safe_remove_index(index_path: Path) -> None: +def _safe_remove_index(index_path: Path, trusted_index_root: Path) -> None: resolved = index_path.resolve() forbidden = { Path(resolved.anchor), @@ -399,6 +372,11 @@ def _safe_remove_index(index_path: Path) -> None: } if resolved in forbidden: raise ValueError(f"Refusing to remove unsafe index path: {resolved}") + if resolved.parent != trusted_index_root.resolve(): + raise ValueError( + "Refusing to remove PyLucene index outside its configured root: " + f"{resolved}" + ) if index_path.is_symlink() or not index_path.is_dir(): raise ValueError( f"PyLucene index path must be a directory: {index_path}" @@ -474,17 +452,17 @@ def _commit_fingerprints(index_path: Path) -> List[Dict[str, str]]: @dataclass(frozen=True) class _IndexProvenanceVerification: codec: str - writer_path: str + writer_policy: str vector_count: int dimensions: int commit_file_count: int def to_metadata(self) -> Dict[str, Union[str, int]]: return { - "status": f"{self.writer_path}-provenance", + "status": f"{self.writer_policy}-provenance", "schema_version": _PROVENANCE_SCHEMA_VERSION, "codec": self.codec, - "writer_path": self.writer_path, + "writer_policy": self.writer_policy, "vector_count": self.vector_count, "dimensions": self.dimensions, "commit_file_count": self.commit_file_count, @@ -514,7 +492,7 @@ def _write_index_provenance( payload = { "schema_version": _PROVENANCE_SCHEMA_VERSION, "codec": codec, - "writer_path": _EXPECTED_WRITER_PATH[codec], + "writer_policy": _EXPECTED_WRITER_POLICY[codec], "vector_count": int(vector_count), "dimensions": int(dimensions), "commit_fingerprints": _commit_fingerprints(index_path), @@ -628,12 +606,12 @@ def _validate_provenance_identity( f"{payload['codec']!r} != {expected_codec!r}" ) - writer_path = payload["writer_path"] - if writer_path != _EXPECTED_WRITER_PATH[expected_codec]: + writer_policy = payload["writer_policy"] + if writer_policy != _EXPECTED_WRITER_POLICY[expected_codec]: raise _IndexProvenanceError( - f"{provenance_label} does not record the required GPU writer path" + f"{provenance_label} does not record the expected writer policy" ) - return writer_path + return writer_policy def _validate_provenance_shape( @@ -746,7 +724,7 @@ def _verify_index_provenance( payload = _read_index_provenance( index_path / expectation.manifest_name, expectation.label ) - writer_path = _validate_provenance_identity( + writer_policy = _validate_provenance_identity( payload, expectation.codec, expectation.label ) vector_count, dimensions = _validate_provenance_shape( @@ -771,7 +749,7 @@ def _verify_index_provenance( return _IndexProvenanceVerification( codec=expectation.codec, - writer_path=writer_path, + writer_policy=writer_policy, vector_count=vector_count, dimensions=dimensions, commit_file_count=len(current_fingerprints), @@ -908,16 +886,9 @@ class _SearchPlan: @dataclass(frozen=True) -class _ResolvedCodec: - codec_name: str +class _BuildCodec: java_codec: Any - telemetry: Mapping[str, str] - writer_path: str - - def __post_init__(self) -> None: - object.__setattr__( - self, "telemetry", MappingProxyType(dict(self.telemetry)) - ) + writer_policy: str class _ExistingIndexAction(Enum): @@ -1847,17 +1818,12 @@ def resolve_codec(self, codec_name: str) -> Any: raise RuntimeError( f"Requested codec {codec_name}, got {codec.getName()}" ) - self._codec_cache[codec_name] = codec - return codec - - def codec_telemetry(self, codec_name: str, codec: Any) -> Dict[str, str]: - """Inspect the selected writer path before indexing starts.""" - vectors_format = codec.knnVectorsFormat() - if vectors_format is None: + if codec.knnVectorsFormat() is None: raise RuntimeError( f"{codec_name} did not initialize a Lucene vector format" ) - return _parse_writer_telemetry(str(vectors_format)) + self._codec_cache[codec_name] = codec + return codec def _java_float_array(self, vector: np.ndarray) -> Any: return self.lucene.JArray("float")( @@ -1880,20 +1846,19 @@ def build_index( self, index_path: Path, vectors: np.ndarray, - resolved_codec: _ResolvedCodec, - ) -> Dict[str, str]: + build_codec: _BuildCodec, + ) -> None: self.attach_current_thread() directory = self.FSDirectory.open(self.Paths.get(str(index_path))) with _CleanupStack() as cleanups: cleanups.add("close Lucene directory", directory.close) writer_config = self.IndexWriterConfig() writer_config.setOpenMode(self.IndexWriterConfig.OpenMode.CREATE) - writer_config.setCodec(resolved_codec.java_codec) + writer_config.setCodec(build_codec.java_codec) writer_config.setUseCompoundFile(False) writer_config.setMergeScheduler(self.SerialMergeScheduler()) writer = self.IndexWriter(directory, writer_config) self._write_and_close_index(writer, vectors) - return dict(resolved_codec.telemetry) def _write_and_close_index(self, writer: Any, vectors: np.ndarray) -> None: try: @@ -2163,8 +2128,6 @@ class _BenchmarkConfigContext: dataset_path: str subset_scope: Optional[str] runtime_config: Dict[str, Any] - neighbor_count: int - batch_size: int class PyLuceneConfigLoader(ConfigLoader): @@ -2320,13 +2283,11 @@ def _index_label( subset_scope: Optional[str], build_params: Dict[str, Any], ) -> str: - prefix = algorithm if group == "base" else f"{algorithm}_{group}" - label_parts = [prefix] - if subset_scope is not None: - label_parts.append(subset_scope) - for key, value in build_params.items(): - label_parts.append(f"{key}{value}") - return ".".join(label_parts) + identity = format_artifact_identity( + algorithm, group, subset_scope, description="PyLucene" + ) + codec = _validate_codec(build_params.get("codec")) + return f"{identity}[codec={codec}]" @classmethod def _benchmark_config( @@ -2337,36 +2298,30 @@ def _benchmark_config( build_params: Dict[str, Any], search_params: List[Dict[str, Any]], ) -> BenchmarkConfig: + dataset = validate_path_component( + context.dataset, "PyLucene dataset name" + ) index_label = cls._index_label( algorithm, group, context.subset_scope, build_params ) - index_path = os.path.join( - context.dataset_path, - context.dataset, - "index", - index_label, - ) - result_stem = f"{algorithm},{group}" - if context.subset_scope is not None: - result_stem = f"{result_stem},{context.subset_scope}" - + index_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=index_path, + 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.get("codec"), - "requires_gpu": True, - "output_filename": ( - result_stem, - f"{result_stem},k{context.neighbor_count}," - f"bs{context.batch_size}", - ), + "requires_gpu": build_params.get("codec") == _CAGRA_CODEC, **context.runtime_config, } return BenchmarkConfig( @@ -2412,8 +2367,6 @@ def _build_benchmark_configs( dataset_path=dataset_path, subset_scope=self._subset_scope(dataset_config.subset_size), runtime_config=self._runtime_config(kwargs), - neighbor_count=kwargs.get("count", 10), - batch_size=kwargs.get("batch_size", 10000), ) tune_mode = kwargs.get("_tune_mode", False) tune_build_params = kwargs.get("_tune_build_params") @@ -2451,8 +2404,6 @@ def _build_benchmark_configs( class PyLuceneBackend(BenchmarkBackend): """Build and search cuVS-Lucene indexes through PyLucene.""" - orchestrator_persists_results = True - def __init__(self, config: Dict[str, Any]): super().__init__(config) self._runtime: Optional[_PyLuceneRuntime] = None @@ -2461,6 +2412,18 @@ def __init__(self, config: Dict[str, Any]): 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 cleanup(self) -> None: """Release Python references; the process-wide JVM remains active.""" self._runtime = None @@ -2470,6 +2433,27 @@ def _get_runtime(self) -> _PyLuceneRuntime: 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, @@ -2482,6 +2466,7 @@ def _failed_build_result( index_size_bytes=0, algorithm=self.algo, build_params=build_params or {}, + metadata=self._result_identity(), success=False, error_message=error_message, ) @@ -2501,6 +2486,7 @@ def _failed_search_result( recall=0.0, algorithm=self.algo, search_params=search_params or [], + metadata=self._result_identity(), success=False, error_message=error_message, ) @@ -2555,7 +2541,7 @@ def _dry_run_build_result( index_size_bytes=0, algorithm=self.algo, build_params=build_params, - metadata={"codec": codec_name}, + metadata={**self._result_identity(), "codec": codec_name}, success=True, ) @@ -2590,6 +2576,7 @@ def _decide_existing_index( return _ExistingIndexDecision.build() metadata: Dict[str, Any] = { + **self._result_identity(), "codec": codec_name, "skipped": True, } @@ -2643,24 +2630,12 @@ def _training_vectors_for_build( return vectors @staticmethod - def _require_gpu_writer( + def _resolve_build_codec( runtime: _PyLuceneRuntime, codec_name: str - ) -> _ResolvedCodec: - java_codec = runtime.resolve_codec(codec_name) - telemetry = runtime.codec_telemetry(codec_name, java_codec) - writer_path = telemetry.get("writerPath") - expected_writer_path = _EXPECTED_WRITER_PATH[codec_name] - if writer_path != expected_writer_path: - raise RuntimeError( - f"{codec_name} uses writerPath={writer_path!r}; " - f"expected {expected_writer_path!r}. Refusing to run " - "a silent CPU or alternate-index fallback as a cuVS build." - ) - return _ResolvedCodec( - codec_name=codec_name, - java_codec=java_codec, - telemetry=telemetry, - writer_path=writer_path, + ) -> _BuildCodec: + return _BuildCodec( + java_codec=runtime.resolve_codec(codec_name), + writer_policy=_EXPECTED_WRITER_POLICY[codec_name], ) @staticmethod @@ -2707,14 +2682,13 @@ def _verify_and_persist_built_index_provenance( ) return {"hnsw_verification": verification.to_metadata()} - @staticmethod def _cleanup_partial_index( - index_path: Path, created_for_build: bool + self, index_path: Path, created_for_build: bool ) -> Optional[Exception]: if not created_for_build or not index_path.exists(): return None try: - _safe_remove_index(index_path) + self._remove_index(index_path) except Exception as exc: return exc return None @@ -2738,6 +2712,7 @@ def build( try: _validate_build_params(build_params) codec_name = _configured_codec(build_params, self.config) + self._validate_index_location(index_path) except Exception as exc: return self._failed_build_result( str(exc), str(index_path), build_params @@ -2772,24 +2747,22 @@ def _build_new_index( vectors = self._training_vectors_for_build(dataset, codec_name) runtime = self._get_runtime() # Preflight must succeed before an existing index is replaced. - resolved_codec = self._require_gpu_writer(runtime, codec_name) + build_codec = self._resolve_build_codec(runtime, codec_name) if index_path.exists(): - _safe_remove_index(index_path) + self._remove_index(index_path) index_path.mkdir(parents=True) created_for_build = True start = time.perf_counter() - telemetry = runtime.build_index( - index_path, vectors, resolved_codec - ) + runtime.build_index(index_path, vectors, build_codec) build_time = time.perf_counter() - start metadata = { + **self._result_identity(), "codec": codec_name, "pylucene_version": runtime.pylucene_version, - "writer_path": resolved_codec.writer_path, - "writer_telemetry": telemetry, + "writer_policy": build_codec.writer_policy, } metadata.update( self._verify_and_persist_built_index_provenance( @@ -2889,7 +2862,7 @@ def _dry_run_search_result( recall=0.0, algorithm=self.algo, search_params=plan.search_params, - metadata={"codec": plan.codec_name}, + metadata={**self._result_identity(), "codec": plan.codec_name}, success=True, ) @@ -2945,6 +2918,7 @@ def _execute_search( batch_size=plan.batch_size, ) metadata = { + **self._result_identity(), "codec": plan.codec_name, "pylucene_version": runtime.pylucene_version, "index_dimensions": runtime_result.index_dimensions, @@ -2952,6 +2926,7 @@ def _execute_search( "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 @@ -2967,7 +2942,6 @@ def _execute_search( recall=0.0, algorithm=self.algo, search_params=plan.search_params, - latency_seconds=processed.latency_seconds, latency_percentiles=processed.latency_percentiles, metadata=metadata, success=True, @@ -2983,12 +2957,15 @@ def search( force: bool = False, search_threads: Optional[int] = None, dry_run: bool = False, - ) -> SearchResult: + ) -> 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" - ) + return [ + self._failed_search_result( + k, + "PyLucene backend requires exactly one index configuration", + ) + ] index_config = indexes[0] try: @@ -3000,30 +2977,37 @@ def search( search_threads=search_threads, ) except Exception as exc: - return self._failed_search_result( - k, - str(exc), - index_config.search_params or [{}], - ) + return [ + self._failed_search_result( + k, + str(exc), + index_config.search_params or [{}], + ) + ] if dry_run: - return self._dry_run_search_result(plan) + return [self._dry_run_search_result(plan)] if not plan.index_path.is_dir(): - return self._failed_search_result( - k, - f"PyLucene index directory does not exist: {plan.index_path}", - plan.search_params, - ) + return [ + self._failed_search_result( + k, + "PyLucene index directory does not exist: " + f"{plan.index_path}", + plan.search_params, + ) + ] try: - return self._execute_search(dataset, plan) + return [self._execute_search(dataset, plan)] except Exception as exc: - return self._failed_search_result( - k, - _exception_summary(exc), - plan.search_params, - ) + return [ + self._failed_search_result( + k, + _exception_summary(exc), + plan.search_params, + ) + ] __all__ = ["PyLuceneBackend", "PyLuceneConfigLoader"] diff --git a/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_hnsw.yaml b/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_hnsw.yaml index a82700aeeb..bd8ab2849c 100644 --- a/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_hnsw.yaml +++ b/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_hnsw.yaml @@ -4,8 +4,6 @@ groups: build: codec: - "Lucene101AcceleratedHNSWCodec" - - "Lucene101AcceleratedHNSWBaseLayerCodec" - - "Lucene101AcceleratedHNSWMultiLayerCodec" search: {} test: build: diff --git a/python/cuvs_bench/cuvs_bench/tests/_pylucene_test_utils.py b/python/cuvs_bench/cuvs_bench/tests/_pylucene_test_utils.py index 16b40affc3..1c3d946c9f 100644 --- a/python/cuvs_bench/cuvs_bench/tests/_pylucene_test_utils.py +++ b/python/cuvs_bench/cuvs_bench/tests/_pylucene_test_utils.py @@ -23,7 +23,6 @@ from cuvs_bench.orchestrator.config_loaders import IndexConfig _HNSW_CODEC = "Lucene101AcceleratedHNSWCodec" -_HNSW_BASE_LAYER_CODEC = "Lucene101AcceleratedHNSWBaseLayerCodec" _CAGRA_CODEC = "CuVS2510GPUSearchCodec" @@ -33,21 +32,18 @@ class _FakeRuntime: def __init__( self, *, - writer_path: str = "gpu-hnsw", index_dimensions: int = 4, document_count: int = 10, hits: list[list[_SearchHit]] | None = None, batch_latencies_ms: list[float] | None = None, validate_query_dimensions: bool = True, ): - self.writer_path = writer_path self.index_dimensions = index_dimensions self.document_count = document_count self.hits = hits self.batch_latencies_ms = batch_latencies_ms self.validate_query_dimensions = validate_query_dimensions self.resolve_calls = [] - self.telemetry_calls = [] self.build_calls = [] self.search_calls = [] self.verification_calls = [] @@ -62,19 +58,13 @@ def resolve_codec(self, codec_name): raise self.resolve_error return codec_name - def codec_telemetry(self, codec_name, java_codec): - self.telemetry_calls.append(codec_name) - assert java_codec == codec_name - return {"writerPath": self.writer_path, "codec": codec_name} - - def build_index(self, index_path, vectors, resolved_codec): - self.build_calls.append((index_path, vectors.copy(), resolved_codec)) + 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") - return dict(resolved_codec.telemetry) def verify_cagra_index( self, diff --git a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_backend.py b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_backend.py index a5b75654c5..416268566e 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_backend.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_backend.py @@ -17,7 +17,6 @@ from cuvs_bench.backends.pylucene import _SearchHit from cuvs_bench.tests._pylucene_test_utils import ( _CAGRA_CODEC, - _HNSW_BASE_LAYER_CODEC, _HNSW_CODEC, _FakeRuntime, _backend, @@ -28,6 +27,11 @@ ) +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)], @@ -49,7 +53,7 @@ def test_build_dry_run_does_not_initialize_pylucene(tmp_path): assert result.metadata["codec"] == _HNSW_CODEC -def test_build_creates_index_and_reports_gpu_writer(tmp_path): +def test_build_creates_index_and_records_hnsw_writer_policy(tmp_path): runtime = _FakeRuntime() backend = _backend(runtime) index_path = tmp_path / "index" @@ -58,26 +62,25 @@ def test_build_creates_index_and_reports_gpu_writer(tmp_path): assert result.success assert result.index_size_bytes > len(b"index") - assert result.metadata["writer_path"] == "gpu-hnsw" - assert result.metadata["writer_telemetry"]["writerPath"] == "gpu-hnsw" + assert result.metadata["writer_policy"] == "gpu-with-cpu-fallback" assert result.metadata["hnsw_verification"] == { - "status": "gpu-hnsw-provenance", - "schema_version": 1, + "status": "gpu-with-cpu-fallback-provenance", + "schema_version": 2, "codec": _HNSW_CODEC, - "writer_path": "gpu-hnsw", + "writer_policy": "gpu-with-cpu-fallback", "vector_count": 10, "dimensions": 4, "commit_file_count": 1, } assert (index_path / pylucene_backend._HNSW_PROVENANCE_FILE).is_file() assert runtime.resolve_calls == [_HNSW_CODEC] - assert runtime.telemetry_calls == [_HNSW_CODEC] assert len(runtime.build_calls) == 1 + assert runtime.build_calls[0][2].writer_policy == ("gpu-with-cpu-fallback") assert runtime.verification_calls == [] def test_build_verifies_persisted_cagra_index(tmp_path): - runtime = _FakeRuntime(writer_path="gpu-cagra") + runtime = _FakeRuntime() index_path = tmp_path / "index" result = _backend(runtime, codec=_CAGRA_CODEC).build( @@ -90,9 +93,9 @@ def test_build_verifies_persisted_cagra_index(tmp_path): assert runtime.verification_calls == [(index_path, 10, 4)] assert result.metadata["cagra_provenance"] == { "status": "gpu-cagra-provenance", - "schema_version": 1, + "schema_version": 2, "codec": _CAGRA_CODEC, - "writer_path": "gpu-cagra", + "writer_policy": "gpu-cagra", "vector_count": 10, "dimensions": 4, "commit_file_count": 1, @@ -108,7 +111,7 @@ def test_build_verifies_persisted_cagra_index(tmp_path): def test_build_rejects_persisted_cagra_fallback_and_removes_index(tmp_path): - runtime = _FakeRuntime(writer_path="gpu-cagra") + runtime = _FakeRuntime() runtime.verification_error = RuntimeError("persisted brute-force index") index_path = tmp_path / "index" @@ -136,10 +139,9 @@ def test_build_reuses_existing_index_without_runtime(tmp_path): assert result.success assert result.metadata["skipped"] is True assert result.metadata["hnsw_verification"]["status"] == ( - "gpu-hnsw-provenance" + "gpu-with-cpu-fallback-provenance" ) assert runtime.resolve_calls == [] - assert runtime.telemetry_calls == [] assert runtime.build_calls == [] assert runtime.verification_calls == [] @@ -158,7 +160,6 @@ def fail_index_size(_index_path): assert not result.success assert "cannot read index size" in result.error_message - assert runtime.telemetry_calls == [] assert runtime.build_calls == [] @@ -172,14 +173,13 @@ def test_build_rejects_reused_hnsw_index_without_provenance(tmp_path): assert not result.success assert "provenance manifest is missing" in result.error_message - assert runtime.telemetry_calls == [] assert runtime.build_calls == [] def test_build_verifies_reused_cagra_index(tmp_path): index_path = tmp_path / "index" _prepare_cagra_index(index_path) - runtime = _FakeRuntime(writer_path="gpu-cagra") + runtime = _FakeRuntime() result = _backend(runtime, codec=_CAGRA_CODEC).build( _dataset(), @@ -199,7 +199,6 @@ def test_build_verifies_reused_cagra_index(tmp_path): "vector_count": 10, "dimensions": 4, } - assert runtime.telemetry_calls == [] assert runtime.build_calls == [] @@ -207,7 +206,7 @@ def test_build_rejects_reused_cagra_index_that_cannot_be_verified(tmp_path): index_path = tmp_path / "index" _prepare_cagra_index(index_path) segments_file = index_path / "segments_1" - runtime = _FakeRuntime(writer_path="gpu-cagra") + runtime = _FakeRuntime() runtime.verification_error = RuntimeError("persisted brute-force index") result = _backend(runtime, codec=_CAGRA_CODEC).build( @@ -219,7 +218,6 @@ def test_build_rejects_reused_cagra_index_that_cannot_be_verified(tmp_path): assert "persisted brute-force index" in result.error_message assert runtime.verification_calls == [(index_path, 10, 4)] assert segments_file.read_bytes() == b"index" - assert runtime.telemetry_calls == [] assert runtime.build_calls == [] @@ -229,7 +227,7 @@ def test_build_rejects_reused_cagra_index_without_provenance_before_runtime( index_path = tmp_path / "index" index_path.mkdir() (index_path / "segments_1").write_bytes(b"existing") - runtime = _FakeRuntime(writer_path="gpu-cagra") + runtime = _FakeRuntime() result = _backend(runtime, codec=_CAGRA_CODEC).build( _dataset(), @@ -257,7 +255,6 @@ def test_build_rejects_unrelated_existing_directory(tmp_path, force): assert not result.success assert "does not contain a Lucene segments file" in result.error_message assert unrelated_file.exists() - assert runtime.telemetry_calls == [] def test_build_rejects_existing_index_path_that_is_a_file(tmp_path): @@ -290,6 +287,39 @@ def test_build_force_replaces_existing_index(tmp_path): 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() @@ -308,36 +338,17 @@ def test_build_preserves_existing_index_if_codec_preflight_fails(tmp_path): assert old_file.exists() -def test_build_rejects_silent_cpu_fallback_and_removes_index(tmp_path): - runtime = _FakeRuntime(writer_path="cpu-hnsw-fallback") - index_path = tmp_path / "index" - - result = _backend(runtime).build( - _dataset(), [_index(index_path)], force=True - ) - - assert not result.success - assert "silent CPU" in result.error_message - assert not index_path.exists() - assert runtime.build_calls == [] - - -def test_build_preserves_existing_index_on_cpu_fallback(tmp_path): - runtime = _FakeRuntime(writer_path="cpu-hnsw-fallback") +def test_build_accepts_production_hnsw_fallback_policy(tmp_path): + runtime = _FakeRuntime() index_path = tmp_path / "index" - index_path.mkdir() - (index_path / "segments_1").write_bytes(b"existing") - old_file = index_path / "old" - old_file.write_bytes(b"old") result = _backend(runtime).build( _dataset(), [_index(index_path)], force=True ) - assert not result.success - assert "silent CPU" in result.error_message - assert old_file.read_bytes() == b"old" - assert runtime.build_calls == [] + 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( @@ -369,15 +380,15 @@ def test_build_rejects_nonfinite_vectors(tmp_path): @pytest.mark.parametrize( - ("codec", "writer_path"), + "codec", [ - (_HNSW_CODEC, "gpu-hnsw"), - (_CAGRA_CODEC, "gpu-cagra"), + _HNSW_CODEC, + _CAGRA_CODEC, ], ids=["hnsw", "cagra"], ) -def test_build_rejects_single_vector_cuvs_bypass(tmp_path, codec, writer_path): - runtime = _FakeRuntime(writer_path=writer_path) +def test_build_rejects_single_vector_cuvs_bypass(tmp_path, codec): + runtime = _FakeRuntime() result = _backend(runtime, codec=codec).build( _dataset(n_base=1), @@ -388,7 +399,6 @@ def test_build_rejects_single_vector_cuvs_bypass(tmp_path, codec, writer_path): assert not result.success assert "at least two training vectors" in result.error_message assert "does not invoke cuVS" in result.error_message - assert runtime.telemetry_calls == [] assert not (tmp_path / "index").exists() @@ -416,7 +426,9 @@ def test_explicit_invalid_codec_does_not_fall_back_to_backend_config( backend = _backend(codec=_HNSW_CODEC) build_result = backend.build(_dataset(), [index], dry_run=True) - search_result = backend.search(_dataset(), [index], k=3, dry_run=True) + 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 @@ -514,7 +526,7 @@ def test_build_preserves_interrupt_when_partial_index_cleanup_fails( runtime.build_error = KeyboardInterrupt() index_path = tmp_path / "index" - def fail_cleanup(_index_path): + def fail_cleanup(_index_path, trusted_index_root=None): raise PermissionError("cleanup denied") monkeypatch.setattr(pylucene_backend, "_safe_remove_index", fail_cleanup) @@ -532,7 +544,7 @@ def test_build_reports_partial_index_cleanup_failure(tmp_path, monkeypatch): runtime.build_error = RuntimeError("Java build failed") index_path = tmp_path / "index" - def fail_cleanup(_index_path): + def fail_cleanup(_index_path, trusted_index_root=None): raise PermissionError("cleanup denied") monkeypatch.setattr(pylucene_backend, "_safe_remove_index", fail_cleanup) @@ -550,8 +562,10 @@ def fail_cleanup(_index_path): def test_search_dry_run_does_not_initialize_pylucene(tmp_path): backend = _backend() - result = backend.search( - _dataset(), [_index(tmp_path / "index")], k=3, dry_run=True + result = _only_search_result( + backend.search( + _dataset(), [_index(tmp_path / "index")], k=3, dry_run=True + ) ) assert result.success @@ -572,8 +586,10 @@ def test_search_converts_hits_scores_and_padding(tmp_path): ] ) - result = _backend(runtime).search( - _dataset(), [_index(index_path)], k=3, batch_size=1 + result = _only_search_result( + _backend(runtime).search( + _dataset(), [_index(index_path)], k=3, batch_size=1 + ) ) assert result.success @@ -585,7 +601,7 @@ def test_search_converts_hits_scores_and_padding(tmp_path): assert result.distances[1, 0] == pytest.approx(4.0) assert np.isinf(result.distances[:, -1]).all() assert result.search_time_ms == pytest.approx(2.0) - assert result.latency_seconds == pytest.approx(0.001) + assert result.metadata["latency_seconds"] == pytest.approx(0.001) assert result.queries_per_second == pytest.approx(1000.0) assert result.latency_percentiles == { "p50": 1.0, @@ -594,7 +610,7 @@ def test_search_converts_hits_scores_and_padding(tmp_path): } assert result.metadata["num_batches"] == 2 assert result.metadata["hnsw_verification"]["status"] == ( - "gpu-hnsw-provenance" + "gpu-with-cpu-fallback-provenance" ) assert "per_search_param_results" not in result.metadata assert runtime.search_calls[0][3] == 1 @@ -622,8 +638,8 @@ def record_input_loading(dataset): monkeypatch.setattr(pylucene_backend.time, "perf_counter", record_time) monkeypatch.setattr(backend, "_load_search_inputs", record_input_loading) - result = backend.search( - _dataset(), [_index(index_path)], k=3, batch_size=2 + result = _only_search_result( + backend.search(_dataset(), [_index(index_path)], k=3, batch_size=2) ) assert result.success @@ -721,8 +737,10 @@ def test_search_rejects_invalid_runtime_results(tmp_path, runtime, error): index_path = tmp_path / "index" _prepare_hnsw_index(index_path) - result = _backend(runtime).search( - _dataset(), [_index(index_path)], k=3, batch_size=2 + result = _only_search_result( + _backend(runtime).search( + _dataset(), [_index(index_path)], k=3, batch_size=2 + ) ) assert not result.success @@ -732,12 +750,14 @@ def test_search_rejects_invalid_runtime_results(tmp_path, runtime, error): def test_search_verifies_cagra_index_before_querying(tmp_path): index_path = tmp_path / "index" _prepare_cagra_index(index_path) - runtime = _FakeRuntime(writer_path="gpu-cagra") + runtime = _FakeRuntime() - result = _backend(runtime, codec=_CAGRA_CODEC).search( - _dataset(), - [_index(index_path, codec=_CAGRA_CODEC)], - k=3, + result = _only_search_result( + _backend(runtime, codec=_CAGRA_CODEC).search( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + k=3, + ) ) assert result.success @@ -758,13 +778,15 @@ def test_search_verifies_cagra_index_before_querying(tmp_path): def test_search_rejects_unverified_cagra_index_before_querying(tmp_path): index_path = tmp_path / "index" _prepare_cagra_index(index_path) - runtime = _FakeRuntime(writer_path="gpu-cagra") + runtime = _FakeRuntime() runtime.verification_error = RuntimeError("persisted brute-force index") - result = _backend(runtime, codec=_CAGRA_CODEC).search( - _dataset(), - [_index(index_path, codec=_CAGRA_CODEC)], - k=3, + result = _only_search_result( + _backend(runtime, codec=_CAGRA_CODEC).search( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + k=3, + ) ) assert not result.success @@ -783,12 +805,14 @@ def test_search_rejects_invalid_cagra_provenance_before_runtime( manifest_path.unlink() else: (index_path / "segments_1").write_bytes(b"changed commit") - runtime = _FakeRuntime(writer_path="gpu-cagra") + runtime = _FakeRuntime() - result = _backend(runtime, codec=_CAGRA_CODEC).search( - _dataset(), - [_index(index_path, codec=_CAGRA_CODEC)], - k=3, + result = _only_search_result( + _backend(runtime, codec=_CAGRA_CODEC).search( + _dataset(), + [_index(index_path, codec=_CAGRA_CODEC)], + k=3, + ) ) assert not result.success @@ -853,14 +877,16 @@ def test_search_rejects_invalid_cagra_provenance_before_runtime( def test_search_rejects_unsupported_options( index, k, batch_size, search_threads, mode, error ): - result = _backend(_FakeRuntime()).search( - _dataset(), - [index], - k=k, - batch_size=batch_size, - mode=mode, - search_threads=search_threads, - dry_run=True, + 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 @@ -889,14 +915,21 @@ def test_resolve_search_plan_normalizes_the_complete_request(tmp_path): def test_search_validation_preserves_first_error_precedence(tmp_path): - result = _backend().search( - _dataset(), - [_index(tmp_path / "index", search_params=[{"unsupported": True}])], - k=0, - batch_size=0, - mode="throughput", - search_threads=2, - dry_run=True, + 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 @@ -910,15 +943,19 @@ def test_search_requires_exactly_one_index(index_count, tmp_path): for index_id in range(index_count) ] - result = _backend(_FakeRuntime()).search(_dataset(), indexes, k=3) + 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 = _backend(_FakeRuntime()).search( - _dataset(), [_index(tmp_path / "missing")], k=3 + result = _only_search_result( + _backend(_FakeRuntime()).search( + _dataset(), [_index(tmp_path / "missing")], k=3 + ) ) assert not result.success @@ -931,8 +968,8 @@ def test_search_rejects_empty_query_vectors(tmp_path): dataset = _dataset() dataset.query_vectors = np.empty((0, 4), dtype=np.float32) - result = _backend(_FakeRuntime()).search( - dataset, [_index(index_path)], k=3 + result = _only_search_result( + _backend(_FakeRuntime()).search(dataset, [_index(index_path)], k=3) ) assert not result.success @@ -944,8 +981,10 @@ def test_search_reports_runtime_query_dimension_mismatch(tmp_path): _prepare_hnsw_index(index_path) runtime = _FakeRuntime(index_dimensions=8) - result = _backend(runtime).search( - _dataset(dimensions=4), [_index(index_path)], k=3 + result = _only_search_result( + _backend(runtime).search( + _dataset(dimensions=4), [_index(index_path)], k=3 + ) ) assert not result.success @@ -962,8 +1001,10 @@ def test_search_rejects_runtime_dimensions_that_disagree_with_provenance( validate_query_dimensions=False, ) - result = _backend(runtime).search( - _dataset(dimensions=4), [_index(index_path)], k=3 + result = _only_search_result( + _backend(runtime).search( + _dataset(dimensions=4), [_index(index_path)], k=3 + ) ) assert not result.success @@ -981,7 +1022,9 @@ def test_search_rejects_query_dimensions_that_disagree_with_dataset(tmp_path): dataset.query_vectors = np.zeros((2, 8), dtype=np.float32) runtime = _FakeRuntime() - result = _backend(runtime).search(dataset, [_index(index_path)], k=3) + 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 ( @@ -996,7 +1039,9 @@ def test_search_reports_runtime_failure(tmp_path): runtime = _FakeRuntime() runtime.search_error = RuntimeError("Java search failed") - result = _backend(runtime).search(_dataset(), [_index(index_path)], k=3) + 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 @@ -1019,11 +1064,13 @@ def test_search_rejects_invalid_hnsw_provenance_before_runtime( (index_path / "segments_1").write_bytes(b"changed commit") else: payload = json.loads(manifest_path.read_text()) - payload["codec"] = _HNSW_BASE_LAYER_CODEC + payload["codec"] = "DifferentCodec" manifest_path.write_text(json.dumps(payload)) runtime = _FakeRuntime() - result = _backend(runtime).search(_dataset(), [_index(index_path)], k=3) + result = _only_search_result( + _backend(runtime).search(_dataset(), [_index(index_path)], k=3) + ) assert not result.success assert "provenance" in result.error_message.lower() 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 index e3c2ddedc0..f24c4f3cf7 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_cli_config.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_cli_config.py @@ -8,7 +8,6 @@ from __future__ import annotations import csv -import json import subprocess import sys from pathlib import Path @@ -18,9 +17,8 @@ from click.testing import CliRunner import cuvs_bench.backends.pylucene as pylucene_backend -from cuvs_bench.backends import CppGoogleBenchmarkBackend, get_registry +from cuvs_bench.backends import get_registry from cuvs_bench.backends.pylucene import ( - PyLuceneBackend, PyLuceneConfigLoader, _SearchHit, ) @@ -134,11 +132,6 @@ def test_backend_and_loader_are_registered(): assert list_config_loaders()["pylucene"] is PyLuceneConfigLoader -def test_result_file_ownership_is_explicit(): - assert PyLuceneBackend.orchestrator_persists_results is True - assert CppGoogleBenchmarkBackend.orchestrator_persists_results is False - - def test_import_does_not_load_pylucene(): subprocess.run( [ @@ -208,36 +201,19 @@ def test_cli_build_search_persists_metrics_and_build_join( ) assert result.exit_code == 0, result.output - index_name = f"pylucene_test_test.codec{_HNSW_CODEC}" + index_name = f"pylucene_test[group=test][codec={_HNSW_CODEC}]" index_path = dataset_path / "test-dataset" / "index" / index_name assert (index_path / "segments_1").is_file() assert (index_path / pylucene_backend._HNSW_PROVENANCE_FILE).is_file() - expected_index_size = sum( - path.stat().st_size for path in index_path.rglob("*") if path.is_file() - ) - - build_json = result_path / "build" / "pylucene_test,test.json" - search_json = result_path / "search" / "pylucene_test,test,k1,bs2.json" - build_rows = json.loads(build_json.read_text())["benchmarks"] - search_rows = json.loads(search_json.read_text())["benchmarks"] + 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 - assert len(search_rows) == 1 - build_row = build_rows[0] - assert build_row["name"] == f"{index_name}/build" - assert build_row["real_time"] > 0.0 - assert build_row["index_size"] == expected_index_size + assert build_row["index_name"] == index_name + assert float(build_row["time"]) > 0.0 assert build_row["codec"] == _HNSW_CODEC - assert build_row["writer_path"] == "gpu-hnsw" - - search_row = search_rows[0] - assert search_row["name"] == f"{index_name}/search" - assert search_row["Recall"] == pytest.approx(1.0) - assert search_row["items_per_second"] == pytest.approx(2000.0) - assert search_row["search_time_ms"] == pytest.approx(1.0) - assert search_row["real_time"] == pytest.approx(1.0) - assert search_row["Latency"] == pytest.approx(0.001) - assert search_row["p50"] == pytest.approx(1.0) + 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: @@ -249,13 +225,14 @@ def test_cli_build_search_persists_metrics_and_build_join( 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_and_persists_build_failure( +def test_cli_returns_nonzero_without_persisting_failed_measurement( config_dir, tmp_path, monkeypatch ): from cuvs_bench.run.__main__ import main as run_main @@ -279,17 +256,14 @@ def test_cli_returns_nonzero_and_persists_build_failure( assert result.exit_code != 0 assert "intentional build failure" in result.output - build_json = ( + build_csv = ( dataset_path / "test-dataset" / "result" / "build" - / "pylucene_test,test.json" + / "pylucene_test,test.csv" ) - build_rows = json.loads(build_json.read_text())["benchmarks"] - assert len(build_rows) == 1 - assert build_rows[0]["success"] is False - assert "intentional build failure" in build_rows[0]["error_message"] + assert not build_csv.exists() assert not (dataset_path / "test-dataset" / "result" / "search").exists() @@ -320,11 +294,14 @@ def test_config_loader_expands_codecs_and_forwards_runtime_config(config_dir): ) assert config.backend_config["java_library_path"] == "/native" assert config.backend_config["jvm_args"] == ["-Xms1g"] - assert config.backend_config["requires_gpu"] is True - assert config.output_filename == ( - "pylucene_test,base", - "pylucene_test,base,k10,bs10000", + 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 @pytest.mark.parametrize( @@ -369,6 +346,40 @@ def test_config_loader_rejects_unsupported_parameters( ) +@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 = [] @@ -391,28 +402,48 @@ def test_config_loader_scopes_artifact_identity_by_dataset_subset(config_dir): assert dataset_config.subset_size == subset_size assert len(configs) == 1 config = configs[0] - index_prefix = "pylucene_test_test" - result_stem = "pylucene_test,test" + index_prefix = "pylucene_test[group=test]" if subset_scope is not None: - index_prefix = f"{index_prefix}.{subset_scope}" - result_stem = f"{result_stem},{subset_scope}" - index_name = f"{index_prefix}.codec{_HNSW_CODEC}" + index_prefix = f"{index_prefix}[scope={subset_scope}]" + index_name = f"{index_prefix}[codec={_HNSW_CODEC}]" assert config.index_name == index_name assert config.index_path == ( Path("/datasets/test-dataset/index") / index_name ) - assert config.output_filename == ( - result_stem, - f"{result_stem},k1,bs2", - ) - assert all(Path(stem).name == stem for stem in config.output_filename) + 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.output_filename) + ( + 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 = {"codec": _HNSW_CODEC} + + 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) @@ -464,7 +495,9 @@ def test_config_loader_discovers_custom_filename_by_parsed_identity( assert len(configs) == 1 assert configs[0].indexes[0].algo == algorithm_name - assert configs[0].index_name.startswith(f"{algorithm_name}_test.") + assert configs[0].index_name == ( + f"{algorithm_name}[group=test][codec={_HNSW_CODEC}]" + ) def test_config_loader_excludes_explicit_other_backend(config_dir, tmp_path): @@ -505,7 +538,9 @@ def test_config_loader_honors_algorithm_and_group_filters(config_dir): groups="test", ) assert len(configs) == 1 - assert configs[0].index_name.startswith("pylucene_test_test") + assert configs[0].index_name == ( + f"pylucene_test[group=test][codec={_HNSW_CODEC}]" + ) def test_algorithm_config_discovery_is_deterministic(tmp_path): @@ -578,7 +613,7 @@ def test_config_loader_unions_global_and_algorithm_specific_groups(config_dir): assert len(configs) == 3 assert ( sum( - config.index_name.startswith("pylucene_test_test.") + config.index_name.startswith("pylucene_test[group=test]") for config in configs ) == 1 @@ -595,7 +630,7 @@ def test_config_loader_selects_only_explicit_algorithm_group(config_dir): ) assert len(configs) == 1 - assert configs[0].index_name.startswith("pylucene_test_test.") + assert configs[0].index_name.startswith("pylucene_test[group=test]") @pytest.mark.parametrize( @@ -742,7 +777,9 @@ def test_config_loader_combines_global_and_explicit_selectors( **selectors, ) - selected_groups = [config.output_filename[0] for config in configs] + selected_groups = [ + f"{config.algo},{config.backend_config['group']}" for config in configs + ] assert selected_groups == expected_groups @@ -877,11 +914,7 @@ def test_config_loader_rejects_unknown_selectors(config_dir, selectors, error): [ ( "pylucene_cuvs_hnsw", - { - "Lucene101AcceleratedHNSWCodec", - "Lucene101AcceleratedHNSWBaseLayerCodec", - "Lucene101AcceleratedHNSWMultiLayerCodec", - }, + {_HNSW_CODEC}, ), ("pylucene_cuvs_cagra", {_CAGRA_CODEC}), ], diff --git a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_integration.py b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_integration.py index af886641f7..e5f6494376 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_integration.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_integration.py @@ -21,11 +21,11 @@ from cuvs_bench.backends._utils import compute_recall from cuvs_bench.backends.base import Dataset from cuvs_bench.backends.pylucene import ( + _BuildCodec, _CAGRA_PROVENANCE_FILE, PyLuceneBackend, _HNSW_PROVENANCE_FILE, _PyLuceneRuntime, - _ResolvedCodec, ) from cuvs_bench.orchestrator.config_loaders import IndexConfig @@ -41,6 +41,21 @@ _CUVS_LUCENE_JAR_ENV = "CUVS_LUCENE_JAR" _CAGRA_CODEC = "CuVS2510GPUSearchCodec" _HNSW_CODEC = "Lucene101AcceleratedHNSWCodec" +_HNSW_WRITER_POLICY = "gpu-with-cpu-fallback" +_GPU_HNSW_WRITER = ( + "com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsWriter" +) +_WRITER_SELECTION_CODEC = "com.nvidia.cuvs.bench.PyLuceneWriterSelectionCodec" +_WRITER_SELECTION_SOURCE = ( + Path(__file__).resolve().parents[2] + / "tests" + / "java" + / "com" + / "nvidia" + / "cuvs" + / "bench" + / "PyLuceneWriterSelectionCodec.java" +) def _write_test_bin(path: Path, data: np.ndarray) -> None: @@ -61,6 +76,55 @@ def _required_jar(env_name: str) -> Path: return jar.resolve() +def _single_search_result(results): + assert len(results) == 1 + return results[0] + + +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 ): @@ -77,9 +141,9 @@ def _assert_cagra_provenance( ): assert metadata["cagra_provenance"] == { "status": "gpu-cagra-provenance", - "schema_version": 1, + "schema_version": 2, "codec": _CAGRA_CODEC, - "writer_path": "gpu-cagra", + "writer_policy": "gpu-cagra", "vector_count": expected_vector_count, "dimensions": expected_dimensions, "commit_file_count": 1, @@ -116,10 +180,10 @@ def _assert_hnsw_verification( ): verification = metadata["hnsw_verification"] assert verification == { - "status": "gpu-hnsw-provenance", - "schema_version": 1, + "status": "gpu-with-cpu-fallback-provenance", + "schema_version": 2, "codec": codec, - "writer_path": "gpu-hnsw", + "writer_policy": _HNSW_WRITER_POLICY, "vector_count": expected_vector_count, "dimensions": expected_dimensions, "commit_file_count": 1, @@ -127,7 +191,7 @@ def _assert_hnsw_verification( @pytest.fixture(scope="module") -def pylucene_runtime_config(): +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") @@ -145,10 +209,19 @@ def pylucene_runtime_config(): ) try: - importlib.import_module("lucene") + lucene = importlib.import_module("lucene") except ImportError as exc: pytest.fail(f"PyLucene is not importable: {exc}") + 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 { "cuvs_java_jar": str(cuvs_java_jar), "cuvs_lucene_jar": str(cuvs_lucene_jar), @@ -160,59 +233,21 @@ def pylucene_runtime_config(): ( "algo", "codec", - "expected_writer_telemetry", + "expected_writer_policy", "expected_index_suffix", ), [ pytest.param( "pylucene_cuvs_hnsw", "Lucene101AcceleratedHNSWCodec", - { - "writerPath": "gpu-hnsw", - "hnswLayers": "1", - "cagraGraphBuildAlgo": "NN_DESCENT", - "cagraGraphDegree": "64", - "cagraIntermediateGraphDegree": "128", - }, + _HNSW_WRITER_POLICY, ".vex", id="accelerated-hnsw", ), - pytest.param( - "pylucene_cuvs_hnsw", - "Lucene101AcceleratedHNSWBaseLayerCodec", - { - "writerPath": "gpu-hnsw", - "hnswLayers": "1", - "cagraGraphBuildAlgo": "NN_DESCENT", - "cagraGraphDegree": "32", - "cagraIntermediateGraphDegree": "64", - }, - ".vex", - id="accelerated-hnsw-base-layer", - ), - pytest.param( - "pylucene_cuvs_hnsw", - "Lucene101AcceleratedHNSWMultiLayerCodec", - { - "writerPath": "gpu-hnsw", - "hnswLayers": "3", - "cagraGraphBuildAlgo": "NN_DESCENT", - "cagraGraphDegree": "32", - "cagraIntermediateGraphDegree": "64", - }, - ".vex", - id="accelerated-hnsw-multi-layer", - ), pytest.param( "pylucene_cuvs_cagra", "CuVS2510GPUSearchCodec", - { - "writerPath": "gpu-cagra", - "cagraStrategy": "HEURISTIC", - "cagraGraphBuildAlgo": "NN_DESCENT", - "cagraGraphDegree": "64", - "cagraIntermediateGraphDegree": "128", - }, + "gpu-cagra", ".vcag", id="cagra", ), @@ -223,7 +258,7 @@ def test_build_and_search_with_real_pylucene_runtime( pylucene_runtime_config, algo, codec, - expected_writer_telemetry, + expected_writer_policy, expected_index_suffix, ): """Build and search through Lucene's public codec and query SPI.""" @@ -275,14 +310,7 @@ def test_build_and_search_with_real_pylucene_runtime( assert build_result.index_size_bytes > 0 assert build_result.metadata["codec"] == codec assert build_result.metadata["pylucene_version"] != "unknown" - assert ( - build_result.metadata["writer_path"] - == expected_writer_telemetry["writerPath"] - ) - assert ( - build_result.metadata["writer_telemetry"] - == expected_writer_telemetry - ) + assert build_result.metadata["writer_policy"] == expected_writer_policy assert any( path.suffix == expected_index_suffix for path in index_path.iterdir() @@ -330,7 +358,9 @@ def test_build_and_search_with_real_pylucene_runtime( training_vectors.shape[1], ) - search_result = backend.search(dataset, [index], k=k, batch_size=2) + search_result = _single_search_result( + backend.search(dataset, [index], k=k, batch_size=2) + ) assert search_result.success, search_result.error_message assert search_result.neighbors.shape == (query_vectors.shape[0], k) assert search_result.distances.shape == (query_vectors.shape[0], k) @@ -351,8 +381,7 @@ def test_build_and_search_with_real_pylucene_runtime( ) assert np.all(np.diff(search_result.distances, axis=1) >= -1e-6) assert search_result.search_time_ms > 0 - assert search_result.latency_seconds is not None - assert search_result.latency_seconds > 0 + assert search_result.metadata["latency_seconds"] > 0 assert search_result.queries_per_second > 0 assert search_result.metadata["codec"] == codec assert search_result.metadata["pylucene_version"] != "unknown" @@ -382,7 +411,9 @@ def test_build_and_search_with_real_pylucene_runtime( assert "vector count does not match the dataset: 512 != 511" in ( rejected.error_message ) - rejected = backend.search(wrong_count_dataset, [index], k=k) + rejected = _single_search_result( + backend.search(wrong_count_dataset, [index], k=k) + ) assert not rejected.success assert "vector count does not match the dataset: 512 != 511" in ( rejected.error_message @@ -392,13 +423,17 @@ def test_build_and_search_with_real_pylucene_runtime( original_data = data_path.read_bytes() data_path.unlink() - rejected = backend.search(dataset, [index], k=k) + rejected = _single_search_result( + backend.search(dataset, [index], k=k) + ) assert not rejected.success assert "cannot read" in rejected.error_message data_path.write_bytes(original_data) data_path.write_bytes(original_data[:-1]) - rejected = backend.search(dataset, [index], k=k) + rejected = _single_search_result( + backend.search(dataset, [index], k=k) + ) assert not rejected.success assert "do not exactly cover" in rejected.error_message data_path.write_bytes(original_data) @@ -411,7 +446,9 @@ def test_build_and_search_with_real_pylucene_runtime( data_path, ns=(original_stat.st_atime_ns, original_stat.st_mtime_ns), ) - rejected = backend.search(dataset, [index], k=k) + rejected = _single_search_result( + backend.search(dataset, [index], k=k) + ) assert not rejected.success assert "checksum" in rejected.error_message.lower() data_path.write_bytes(original_data) @@ -428,7 +465,9 @@ def test_build_and_search_with_real_pylucene_runtime( _commit_one_deletion(backend, deleted_index_path) with pytest.raises(RuntimeError, match="committed deletions"): backend._get_runtime().verify_cagra_index(deleted_index_path) - rejected = backend.search(dataset, [deleted_index], k=k) + rejected = _single_search_result( + backend.search(dataset, [deleted_index], k=k) + ) assert not rejected.success assert "does not match the current Lucene commit" in ( rejected.error_message @@ -455,7 +494,9 @@ def test_build_and_search_with_real_pylucene_runtime( original_segment = segment_path.read_bytes() manifest_path.unlink() - rejected = backend.search(dataset, [index], k=k) + rejected = _single_search_result( + backend.search(dataset, [index], k=k) + ) assert not rejected.success assert "provenance manifest is missing" in ( rejected.error_message @@ -463,7 +504,9 @@ def test_build_and_search_with_real_pylucene_runtime( manifest_path.write_bytes(original_manifest) manifest_path.write_text("{") - rejected = backend.search(dataset, [index], k=k) + rejected = _single_search_result( + backend.search(dataset, [index], k=k) + ) assert not rejected.success assert "provenance manifest cannot be read" in ( rejected.error_message @@ -471,7 +514,9 @@ def test_build_and_search_with_real_pylucene_runtime( manifest_path.write_bytes(original_manifest) segment_path.write_bytes(original_segment + b"stale") - rejected = backend.search(dataset, [index], k=k) + rejected = _single_search_result( + backend.search(dataset, [index], k=k) + ) assert not rejected.success assert "does not match the current Lucene commit" in ( rejected.error_message @@ -479,9 +524,11 @@ def test_build_and_search_with_real_pylucene_runtime( segment_path.write_bytes(original_segment) payload = json.loads(original_manifest) - payload["codec"] = "Lucene101AcceleratedHNSWBaseLayerCodec" + payload["codec"] = _CAGRA_CODEC manifest_path.write_text(json.dumps(payload)) - rejected = backend.search(dataset, [index], k=k) + rejected = _single_search_result( + backend.search(dataset, [index], k=k) + ) assert not rejected.success assert "does not match the requested codec" in ( rejected.error_message @@ -491,6 +538,30 @@ def test_build_and_search_with_real_pylucene_runtime( backend.cleanup() +def test_real_hnsw_codec_selects_gpu_writer(tmp_path, pylucene_runtime_config): + """Prove the production HNSW codec selects its GPU writer on this host.""" + runtime = _PyLuceneRuntime.create(pylucene_runtime_config) + runtime.attach_current_thread() + reflected_codec = runtime.Class.forName( + _WRITER_SELECTION_CODEC + ).newInstance() + java_codec = runtime.Codec.cast_(reflected_codec) + index_path = tmp_path / "hnsw-writer-selection-index" + index_path.mkdir() + + runtime.build_index( + index_path, + np.ones((2, 32), dtype=np.float32), + _BuildCodec( + java_codec=java_codec, + writer_policy=_HNSW_WRITER_POLICY, + ), + ) + + diagnostics = str(java_codec.knnVectorsFormat()) + assert f"writerClass={_GPU_HNSW_WRITER}" in diagnostics + + def test_real_verifier_rejects_cagra_brute_force_fallback( tmp_path, pylucene_runtime_config ): @@ -500,16 +571,12 @@ def test_real_verifier_rejects_cagra_brute_force_fallback( index_path = tmp_path / "cagra-fallback-index" index_path.mkdir() java_codec = runtime.resolve_codec(_CAGRA_CODEC) - telemetry = runtime.codec_telemetry(_CAGRA_CODEC, java_codec) - assert telemetry["writerPath"] == "gpu-cagra" runtime.build_index( index_path, vectors, - _ResolvedCodec( - codec_name=_CAGRA_CODEC, + _BuildCodec( java_codec=java_codec, - telemetry=telemetry, - writer_path=telemetry["writerPath"], + writer_policy="gpu-cagra", ), ) @@ -539,7 +606,7 @@ def test_real_verifier_rejects_cagra_brute_force_fallback( ) backend._runtime = runtime try: - result = backend.search(dataset, [index], k=1) + 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: @@ -650,45 +717,23 @@ def test_cli_build_and_search_with_real_pylucene_runtime( ) codec = "Lucene101AcceleratedHNSWCodec" - index_name = f"pylucene_cuvs_hnsw_test.codec{codec}" + index_name = f"pylucene_cuvs_hnsw[group=test][codec={codec}]" index_path = dataset_dir / "index" / index_name assert any(path.suffix == ".vex" for path in index_path.iterdir()) assert (index_path / _HNSW_PROVENANCE_FILE).is_file() result_path = dataset_dir / "result" - build_json = result_path / "build" / "pylucene_cuvs_hnsw,test.json" + build_csv = result_path / "build" / "pylucene_cuvs_hnsw,test.csv" search_stem = f"pylucene_cuvs_hnsw,test,k{k},bs2" - search_json = result_path / "search" / f"{search_stem}.json" - build_rows = json.loads(build_json.read_text())["benchmarks"] - search_rows = json.loads(search_json.read_text())["benchmarks"] + with build_csv.open(newline="") as file: + build_rows = list(csv.DictReader(file)) assert len(build_rows) == 1 - assert len(search_rows) == 1 build_row = build_rows[0] - assert build_row["name"] == f"{index_name}/build" - assert build_row["real_time"] > 0 - assert build_row["index_size"] > 0 + assert build_row["index_name"] == index_name + assert float(build_row["time"]) > 0 assert build_row["codec"] == codec - assert build_row["writer_path"] == "gpu-hnsw" - _assert_hnsw_verification( - build_row, - codec, - training_vectors.shape[0], - training_vectors.shape[1], - ) - - search_row = search_rows[0] - assert search_row["name"] == f"{index_name}/search" - assert search_row["Recall"] >= 0.75 - assert search_row["items_per_second"] > 0 - assert search_row["Latency"] > 0 - assert search_row["search_time_ms"] > 0 - _assert_hnsw_verification( - search_row, - codec, - training_vectors.shape[0], - training_vectors.shape[1], - ) + assert build_row["writer_policy"] == _HNSW_WRITER_POLICY raw_csv = result_path / "search" / f"{search_stem},raw.csv" with raw_csv.open(newline="") as file: @@ -700,5 +745,9 @@ def test_cli_build_and_search_with_real_pylucene_runtime( 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 (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 index 475aaba9c1..f8f8332c9d 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_provenance.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_provenance.py @@ -40,10 +40,10 @@ def test_hnsw_provenance_round_trip(tmp_path): ) assert verification.to_metadata() == { - "status": "gpu-hnsw-provenance", - "schema_version": 1, + "status": "gpu-with-cpu-fallback-provenance", + "schema_version": 2, "codec": _HNSW_CODEC, - "writer_path": "gpu-hnsw", + "writer_policy": "gpu-with-cpu-fallback", "vector_count": 10, "dimensions": 4, "commit_file_count": 1, @@ -74,9 +74,9 @@ def test_cagra_provenance_round_trip(tmp_path): assert verification.to_metadata() == { "status": "gpu-cagra-provenance", - "schema_version": 1, + "schema_version": 2, "codec": _CAGRA_CODEC, - "writer_path": "gpu-cagra", + "writer_policy": "gpu-cagra", "vector_count": 10, "dimensions": 4, "commit_file_count": 1, @@ -109,7 +109,7 @@ def test_hnsw_provenance_rejects_dataset_shape_mismatch( @pytest.mark.parametrize( ("manifest_state", "error"), [ - ("wrong-writer", "required GPU writer path"), + ("wrong-policy", "expected writer policy"), ("boolean-count", "positive integer"), ("malformed-fingerprint", "malformed Lucene commit fingerprint"), ("duplicate-fingerprint", "must be unique and sorted"), @@ -122,8 +122,8 @@ def test_hnsw_provenance_rejects_malformed_manifest_fields( manifest_path = _prepare_hnsw_index(index_path) payload = json.loads(manifest_path.read_text()) - if manifest_state == "wrong-writer": - payload["writer_path"] = "cpu-hnsw" + if manifest_state == "wrong-policy": + payload["writer_policy"] = "gpu-cagra" elif manifest_state == "boolean-count": payload["vector_count"] = True elif manifest_state == "malformed-fingerprint": diff --git a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_runtime.py b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_runtime.py index da62e13421..b0dcb25e69 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_runtime.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_runtime.py @@ -14,16 +14,14 @@ import pytest import cuvs_bench.backends.pylucene as pylucene_backend -from cuvs_bench.backends.pylucene import _ResolvedCodec +from cuvs_bench.backends.pylucene import _BuildCodec from cuvs_bench.tests._pylucene_test_utils import _HNSW_CODEC -def _resolved_codec(java_codec=None, writer_path="gpu-hnsw") -> _ResolvedCodec: - return _ResolvedCodec( - codec_name=_HNSW_CODEC, +def _build_codec(java_codec=None) -> _BuildCodec: + return _BuildCodec( java_codec=java_codec if java_codec is not None else object(), - telemetry={"writerPath": writer_path}, - writer_path=writer_path, + writer_policy="gpu-with-cpu-fallback", ) @@ -99,6 +97,37 @@ def setMergeScheduler(self, _scheduler): pass +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_index_writer_runtime(error_at=None, directory_close_error=None): writer = _FakeIndexWriter(error_at=error_at) directory = SimpleNamespace(closed=False) @@ -379,6 +408,46 @@ def test_initialize_pylucene_reports_missing_jar(monkeypatch): 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) + + @pytest.mark.parametrize("jvm_args", ["-Xmx1g", ["-Xmx1g", 1]]) def test_initialize_pylucene_rejects_invalid_jvm_args( jvm_args, tmp_path, monkeypatch @@ -403,56 +472,17 @@ def test_initialize_pylucene_rejects_invalid_jvm_args( ) -@pytest.mark.parametrize( - ("description", "expected"), - [ - ( - "Format(writerPath=gpu-hnsw;hnswLayers=1)", - {"writerPath": "gpu-hnsw", "hnswLayers": "1"}, - ), - ("Format(writerPath=gpu-cagra)", {"writerPath": "gpu-cagra"}), - ], -) -def test_parse_writer_telemetry(description, expected): - assert pylucene_backend._parse_writer_telemetry(description) == expected - - -@pytest.mark.parametrize( - "description", - ["Format", "Format(writerPath)", "Format(=gpu-hnsw)"], -) -def test_parse_writer_telemetry_rejects_malformed_description(description): - with pytest.raises(RuntimeError, match="Malformed"): - pylucene_backend._parse_writer_telemetry(description) - - -def test_resolved_codec_copies_and_freezes_writer_telemetry(): - telemetry = {"writerPath": "gpu-hnsw"} - - resolved_codec = _ResolvedCodec( - codec_name=_HNSW_CODEC, - java_codec=object(), - telemetry=telemetry, - writer_path="gpu-hnsw", - ) - telemetry["writerPath"] = "changed" - - assert dict(resolved_codec.telemetry) == {"writerPath": "gpu-hnsw"} - with pytest.raises(TypeError): - resolved_codec.telemetry["writerPath"] = "changed" - - def test_runtime_build_index_commits_and_closes_writer_and_directory(tmp_path): runtime = _fake_index_writer_runtime() vectors = np.zeros((2, 4), dtype=np.float32) - telemetry = runtime.build_index( + result = runtime.build_index( tmp_path, vectors, - _resolved_codec(runtime._test_codec), + _build_codec(runtime._test_codec), ) - assert telemetry == {"writerPath": "gpu-hnsw"} + assert result is None assert len(runtime._test_writer.documents) == 2 assert runtime._test_writer.committed is True assert runtime._test_writer.rollback_called is False @@ -479,7 +509,7 @@ def test_runtime_build_index_preserves_transaction_cleanup( runtime.build_index( tmp_path, vectors, - _resolved_codec(runtime._test_codec), + _build_codec(runtime._test_codec), ) assert runtime._test_writer.rollback_called is expect_rollback @@ -495,7 +525,7 @@ def test_runtime_build_index_rolls_back_on_interrupt(tmp_path): runtime.build_index( tmp_path, vectors, - _resolved_codec(runtime._test_codec), + _build_codec(runtime._test_codec), ) assert runtime._test_writer.rollback_called is True @@ -511,7 +541,7 @@ def test_runtime_build_index_preserves_interrupt_when_rollback_fails(tmp_path): runtime.build_index( tmp_path, vectors, - _resolved_codec(runtime._test_codec), + _build_codec(runtime._test_codec), ) assert exc_info.value.__notes__ == [ @@ -534,7 +564,7 @@ def test_runtime_build_index_preserves_interrupt_when_directory_close_fails( runtime.build_index( tmp_path, vectors, - _resolved_codec(runtime._test_codec), + _build_codec(runtime._test_codec), ) assert exc_info.value.__notes__ == [ 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..2251449a0c --- /dev/null +++ b/python/cuvs_bench/tests/java/com/nvidia/cuvs/bench/PyLuceneWriterSelectionCodec.java @@ -0,0 +1,63 @@ +/* + * 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.lucene.Lucene101AcceleratedHNSWCodec; +import java.io.IOException; +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 production HNSW codec. */ +public final class PyLuceneWriterSelectionCodec extends Lucene101AcceleratedHNSWCodec { + + public PyLuceneWriterSelectionCodec() throws Exception { + super(); + setKnnFormat(new WriterSelectionFormat(super.knnVectorsFormat())); + } + + private static final class WriterSelectionFormat extends KnnVectorsFormat { + + private static final String NOT_SELECTED = "not-selected"; + + private final KnnVectorsFormat delegate; + private volatile String writerClass = NOT_SELECTED; + + private WriterSelectionFormat(KnnVectorsFormat delegate) { + super(delegate.getName()); + this.delegate = delegate; + } + + @Override + public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + KnnVectorsWriter writer = delegate.fieldsWriter(state); + String selectedClass = writer.getClass().getName(); + String previousClass = writerClass; + if (!NOT_SELECTED.equals(previousClass) && !previousClass.equals(selectedClass)) { + throw new AssertionError( + "Vector writer selection changed from " + previousClass + " to " + selectedClass); + } + writerClass = selectedClass; + return writer; + } + + @Override + public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { + return delegate.fieldsReader(state); + } + + @Override + public int getMaxDimensions(String fieldName) { + return delegate.getMaxDimensions(fieldName); + } + + @Override + public String toString() { + return getName() + "(writerClass=" + writerClass + ")"; + } + } +} From cfea7d8bc4c9009e4f0839fa2485a91c7b74d0b1 Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Fri, 14 Aug 2026 19:49:13 +0000 Subject: [PATCH 08/12] Refresh PyLucene setup and usage documentation Document the current JDK, PyLucene, cuVS Java, and cuVS-Lucene requirements. Describe the two supported production codecs, HNSW fallback behavior, CAGRA validation contract, and the current manual validation workflow. Signed-off-by: nvzm123 --- fern/pages/cuvs_bench/install.md | 19 +++++++++---------- fern/pages/cuvs_bench/pluggable_backend.md | 8 ++++---- fern/pages/cuvs_bench/running.md | 8 ++++---- 3 files changed, 17 insertions(+), 18 deletions(-) diff --git a/fern/pages/cuvs_bench/install.md b/fern/pages/cuvs_bench/install.md index bec57d8d29..5e3ebfa199 100644 --- a/fern/pages/cuvs_bench/install.md +++ b/fern/pages/cuvs_bench/install.md @@ -53,21 +53,18 @@ Exact tags are listed on Docker Hub: The optional `pylucene` backend requires components that cuVS Bench does not install automatically: -- An NVIDIA GPU supported by cuVS, plus matching CUDA and cuVS native libraries. -- JDK 22, `pytest`, and a [source-built PyLucene installation](https://lucene.apache.org/pylucene/install.html) compatible with the Lucene APIs used by the selected cuVS-Lucene checkout. The verified environment used PyLucene 10.0.0; the selected cuVS-Lucene POM compiles against Lucene 10.2.0, so validate the exact pair below. +- 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, `pytest`, and a [source-built PyLucene 10.2.0 installation](https://lucene.apache.org/pylucene/install.html), matching the Lucene 10.2.0 APIs used by the current cuVS-Lucene PR. - [Maven 3.9.6 or newer](https://maven.apache.org/download.cgi) to build cuVS-Lucene. -- The base `cuvs-java` JAR, standard `cuvs-lucene` JAR, and native libraries built from the pinned compatible revisions below. +- The base `cuvs-java` JAR, standard `cuvs-lucene` JAR, and, for GPU execution, native libraries from the same cuVS version. The current PR targets cuVS Java 26.10.0. The required PyLucene service-provider and codec support is currently proposed in [NVIDIA/cuvs-lucene#174](https://github.com/NVIDIA/cuvs-lucene/pull/174). Until that work is merged and released, build cuVS-Lucene from that PR's branch; a release or `main` checkout without those changes is insufficient. -Source compatibility is revision-specific while that PR is under review: the verified PR head `7d70d2f` does not compile against cuVS `main` at `f72199e3`. The commands below pin a source-compatible pair rather than a moving PR head. - -Build the dependency in a separate checkout so this does not change your cuVS Bench working tree: +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 a59e2445 ./build.sh libcuvs java cd .. ``` @@ -78,7 +75,7 @@ If matching native cuVS libraries are already built and installed, `./build.sh j git clone https://github.com/NVIDIA/cuvs-lucene.git cd cuvs-lucene git fetch origin pull/174/head -git switch --detach 7d70d2f +git switch --detach FETCH_HEAD mvn clean package -DskipTests ``` @@ -89,10 +86,12 @@ After the build, the conventional JAR paths are: /target/cuvs-lucene-.jar ``` -Use the base `cuvs-java` JAR, not a native-classifier JAR. Use the standard cuVS-Lucene JAR, not its `-jar-with-dependencies`, sources, or Javadoc variants. Native-library paths must resolve `libcuvs.so`, `libcuvs_c.so`, their dependencies, and the CUDA runtime libraries from the pinned cuVS build. Then validate the artifacts from the cuVS-Lucene checkout: +Use 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. +Then validate the artifacts from the 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 @@ -100,7 +99,7 @@ CUVS_NATIVE_BUILD="$(cd ../cuvs-pylucene-deps/cpp/build && pwd)" export JAVA_LIBRARY_PATH="$CUVS_NATIVE_BUILD:$CUVS_NATIVE_BUILD/c:/usr/local/cuda/lib64" export LD_LIBRARY_PATH="$JAVA_LIBRARY_PATH${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" -./test_pylucene.sh --full-e2e +python3 -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 complete smoke workflow. diff --git a/fern/pages/cuvs_bench/pluggable_backend.md b/fern/pages/cuvs_bench/pluggable_backend.md index a72620350b..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,12 +219,12 @@ 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 cuVS-Lucene codecs through Lucene's service-provider interface. See [Installation](/user-guide/benchmarking-guide/cu-vs-bench-tool/installation#pylucene-backend-prerequisites) and [Usage](/user-guide/benchmarking-guide/cu-vs-bench-tool/usage#running-the-pylucene-backend) for the provisional dependency, configuration, and runtime limits. +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 diff --git a/fern/pages/cuvs_bench/running.md b/fern/pages/cuvs_bench/running.md index 2dec4c8485..136d45f1e8 100644 --- a/fern/pages/cuvs_bench/running.md +++ b/fern/pages/cuvs_bench/running.md @@ -130,11 +130,11 @@ python -m cuvs_bench.run \ Do not use this tiny synthetic smoke run as performance or quality evidence; it exists only to verify the workflow. For a representative recall run, prepare `deep-image-96-angular` with `--normalize` into the same `DATASET_ROOT`, then run the backend against `deep-image-96-inner` with the same `--dataset-path`. -The HNSW `base` group sweeps `Lucene101AcceleratedHNSWCodec`, `Lucene101AcceleratedHNSWBaseLayerCodec`, and `Lucene101AcceleratedHNSWMultiLayerCodec`; its `test` group uses `Lucene101AcceleratedHNSWCodec`. These codecs build the HNSW graph on the GPU with cuVS and use Lucene's CPU HNSW search. Use `--algorithms pylucene_cuvs_cagra` to run `CuVS2510GPUSearchCodec`, which builds and searches on the GPU. +The 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 supplied configurations use Lucene's public `KnnFloatVectorQuery` and do not expose search-time tuning parameters. The backend accepts FLOAT32 Euclidean datasets with at most 4096 dimensions. All supplied codecs require at least two indexed vectors because cuVS-Lucene bypasses cuVS for a single-vector index. CAGRA requires a GPU for build and search. Although cuVS-Lucene uses an effective `lucene_k` of `min(k, document_count)`, the backend conservatively requires `k <= 1024` to avoid cuVS-Lucene paths that can use brute-force search above that limit. HNSW requires a GPU for graph construction, and new builds reject cuVS-Lucene's CPU fallback. +The supplied configurations use Lucene's public `KnnFloatVectorQuery` and do not expose search-time tuning parameters. The backend accepts FLOAT32 Euclidean datasets with at most 4096 dimensions 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` to avoid cuVS-Lucene 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 a different codec. Indexes created outside this backend without the applicable manifest must be rebuilt with `--force`. CAGRA indexes also undergo structural and checksum-integrity verification. +New HNSW and CAGRA builds atomically write commit-bound provenance manifests named `.cuvs-bench-pylucene-hnsw.json` and `.cuvs-bench-pylucene-cagra.json`, respectively. Reuse and search fail if the applicable manifest is missing, malformed, stale, or names a different codec or writer 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. The backend currently supports latency mode with one search thread. `--batch-size` groups queries for measurement, and the reported latency percentiles are per batch. Throughput mode and multiple search threads are not implemented. @@ -257,7 +257,7 @@ Containers can also run in detached mode. ## Evaluating results -The tables below describe fields emitted by the default C++ Google Benchmark backend. Other backends report the fields that apply to their execution model, so not every field is present for every backend. In particular, PyLucene reports `index_size` as the total on-disk index size in bytes. +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: From 3155a67f865eb58567a7a706475de776d51130c6 Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Mon, 17 Aug 2026 05:16:14 +0000 Subject: [PATCH 09/12] Align PyLucene with Lucene 10.2 Validate the generated PyLucene ABI before JVM startup and follow the current cuVS-Lucene codec and writer-selection contracts. Preserve Lucene defaults for HNSW while keeping CAGRA segment files directly verifiable across flushes and merges. Extend provenance and live coverage for GPU selection, CPU fallback, and merged CAGRA indexes. Signed-off-by: nvzm123 --- .../cuvs_bench/backends/pylucene.py | 79 +- .../tests/pylucene_cpu_fallback_probe.py | 83 ++ .../cuvs_bench/tests/test_pylucene_backend.py | 9 +- .../tests/test_pylucene_integration.py | 929 ++++++++++++------ .../tests/test_pylucene_provenance.py | 9 +- .../cuvs_bench/tests/test_pylucene_runtime.py | 95 +- .../bench/PyLuceneWriterSelectionCodec.java | 25 +- 7 files changed, 882 insertions(+), 347 deletions(-) create mode 100644 python/cuvs_bench/cuvs_bench/tests/pylucene_cpu_fallback_probe.py diff --git a/python/cuvs_bench/cuvs_bench/backends/pylucene.py b/python/cuvs_bench/cuvs_bench/backends/pylucene.py index 9e5f620510..0d72bd8956 100644 --- a/python/cuvs_bench/cuvs_bench/backends/pylucene.py +++ b/python/cuvs_bench/cuvs_bench/backends/pylucene.py @@ -46,6 +46,7 @@ _ID_FIELD = "id" _VECTOR_FIELD = "vector" _MAX_DIMENSIONS = 4096 +_REQUIRED_PYLUCENE_VERSION = "10.2.0" _HNSW_CODEC = "Lucene101AcceleratedHNSWCodec" _CAGRA_CODEC = "CuVS2510GPUSearchCodec" @@ -55,6 +56,10 @@ _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" @@ -69,12 +74,13 @@ _HNSW_PROVENANCE_FILE = ".cuvs-bench-pylucene-hnsw.json" _CAGRA_PROVENANCE_FILE = ".cuvs-bench-pylucene-cagra.json" -_PROVENANCE_SCHEMA_VERSION = 2 +_PROVENANCE_SCHEMA_VERSION = 3 _PROVENANCE_KEYS = frozenset( { "schema_version", "codec", "writer_policy", + "compound_file_policy", "vector_count", "dimensions", "commit_fingerprints", @@ -193,6 +199,17 @@ def _load_pylucene() -> Any: ) 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" @@ -281,6 +298,7 @@ def _validate_pylucene_jvm_configuration( 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) @@ -453,6 +471,7 @@ def _commit_fingerprints(index_path: Path) -> List[Dict[str, str]]: class _IndexProvenanceVerification: codec: str writer_policy: str + compound_file_policy: str vector_count: int dimensions: int commit_file_count: int @@ -463,6 +482,7 @@ def to_metadata(self) -> Dict[str, Union[str, int]]: "schema_version": _PROVENANCE_SCHEMA_VERSION, "codec": self.codec, "writer_policy": self.writer_policy, + "compound_file_policy": self.compound_file_policy, "vector_count": self.vector_count, "dimensions": self.dimensions, "commit_file_count": self.commit_file_count, @@ -493,6 +513,7 @@ def _write_index_provenance( "schema_version": _PROVENANCE_SCHEMA_VERSION, "codec": codec, "writer_policy": _EXPECTED_WRITER_POLICY[codec], + "compound_file_policy": _COMPOUND_FILE_POLICY[codec], "vector_count": int(vector_count), "dimensions": int(dimensions), "commit_fingerprints": _commit_fingerprints(index_path), @@ -592,7 +613,7 @@ def _validate_provenance_identity( payload: Dict[str, Any], expected_codec: str, provenance_label: str, -) -> str: +) -> Tuple[str, str]: if ( type(payload["schema_version"]) is not int or payload["schema_version"] != _PROVENANCE_SCHEMA_VERSION @@ -611,7 +632,13 @@ def _validate_provenance_identity( raise _IndexProvenanceError( f"{provenance_label} does not record the expected writer policy" ) - return 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 writer_policy, compound_file_policy def _validate_provenance_shape( @@ -724,7 +751,7 @@ def _verify_index_provenance( payload = _read_index_provenance( index_path / expectation.manifest_name, expectation.label ) - writer_policy = _validate_provenance_identity( + writer_policy, compound_file_policy = _validate_provenance_identity( payload, expectation.codec, expectation.label ) vector_count, dimensions = _validate_provenance_shape( @@ -750,6 +777,7 @@ def _verify_index_provenance( return _IndexProvenanceVerification( codec=expectation.codec, writer_policy=writer_policy, + compound_file_policy=compound_file_policy, vector_count=vector_count, dimensions=dimensions, commit_file_count=len(current_fingerprints), @@ -887,6 +915,7 @@ class _SearchPlan: @dataclass(frozen=True) class _BuildCodec: + codec_name: str java_codec: Any writer_policy: str @@ -1742,7 +1771,6 @@ def __init__(self, lucene: Any): IndexWriterConfig, SegmentCommitInfo, SegmentInfos, - SerialMergeScheduler, VectorEncoding, VectorSimilarityFunction, ) @@ -1762,7 +1790,6 @@ def __init__(self, lucene: Any): self.DirectoryReader = DirectoryReader self.IndexWriter = IndexWriter self.IndexWriterConfig = IndexWriterConfig - self.SerialMergeScheduler = SerialMergeScheduler self.VectorSimilarityFunction = VectorSimilarityFunction self.IndexSearcher = IndexSearcher self.KnnFloatVectorQuery = KnnFloatVectorQuery @@ -1852,14 +1879,23 @@ def build_index( directory = self.FSDirectory.open(self.Paths.get(str(index_path))) with _CleanupStack() as cleanups: cleanups.add("close Lucene directory", directory.close) - writer_config = self.IndexWriterConfig() - writer_config.setOpenMode(self.IndexWriterConfig.OpenMode.CREATE) - writer_config.setCodec(build_codec.java_codec) - writer_config.setUseCompoundFile(False) - writer_config.setMergeScheduler(self.SerialMergeScheduler()) + writer_config = self._new_index_writer_config(build_codec) writer = self.IndexWriter(directory, writer_config) self._write_and_close_index(writer, vectors) + def _new_index_writer_config( + self, build_codec: _BuildCodec + ) -> 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) + return writer_config + def _write_and_close_index(self, writer: Any, vectors: np.ndarray) -> None: try: for document_id, vector in enumerate(vectors): @@ -2424,6 +2460,13 @@ def _result_identity(self) -> Dict[str, Any]: identity["result_scope"] = result_scope return identity + def _index_metadata(self, codec_name: str) -> Dict[str, Any]: + return { + **self._result_identity(), + "codec": codec_name, + "compound_file_policy": _COMPOUND_FILE_POLICY[codec_name], + } + def cleanup(self) -> None: """Release Python references; the process-wide JVM remains active.""" self._runtime = None @@ -2541,7 +2584,7 @@ def _dry_run_build_result( index_size_bytes=0, algorithm=self.algo, build_params=build_params, - metadata={**self._result_identity(), "codec": codec_name}, + metadata=self._index_metadata(codec_name), success=True, ) @@ -2576,8 +2619,7 @@ def _decide_existing_index( return _ExistingIndexDecision.build() metadata: Dict[str, Any] = { - **self._result_identity(), - "codec": codec_name, + **self._index_metadata(codec_name), "skipped": True, } try: @@ -2634,6 +2676,7 @@ def _resolve_build_codec( runtime: _PyLuceneRuntime, codec_name: str ) -> _BuildCodec: return _BuildCodec( + codec_name=codec_name, java_codec=runtime.resolve_codec(codec_name), writer_policy=_EXPECTED_WRITER_POLICY[codec_name], ) @@ -2759,8 +2802,7 @@ def _build_new_index( build_time = time.perf_counter() - start metadata = { - **self._result_identity(), - "codec": codec_name, + **self._index_metadata(codec_name), "pylucene_version": runtime.pylucene_version, "writer_policy": build_codec.writer_policy, } @@ -2862,7 +2904,7 @@ def _dry_run_search_result( recall=0.0, algorithm=self.algo, search_params=plan.search_params, - metadata={**self._result_identity(), "codec": plan.codec_name}, + metadata=self._index_metadata(plan.codec_name), success=True, ) @@ -2918,8 +2960,7 @@ def _execute_search( batch_size=plan.batch_size, ) metadata = { - **self._result_identity(), - "codec": plan.codec_name, + **self._index_metadata(plan.codec_name), "pylucene_version": runtime.pylucene_version, "index_dimensions": runtime_result.index_dimensions, "document_count": runtime_result.document_count, 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..10351f89b2 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/pylucene_cpu_fallback_probe.py @@ -0,0 +1,83 @@ +# +# 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 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" + + +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"], + } + ) + 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", + ), + ) + 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, + ) + 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_pylucene_backend.py b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_backend.py index 416268566e..65eeafe457 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_backend.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_backend.py @@ -51,6 +51,7 @@ def test_build_dry_run_does_not_initialize_pylucene(tmp_path): 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): @@ -63,11 +64,13 @@ def test_build_creates_index_and_records_hnsw_writer_policy(tmp_path): 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": 2, + "schema_version": 3, "codec": _HNSW_CODEC, "writer_policy": "gpu-with-cpu-fallback", + "compound_file_policy": "lucene-default", "vector_count": 10, "dimensions": 4, "commit_file_count": 1, @@ -90,12 +93,14 @@ def test_build_verifies_persisted_cagra_index(tmp_path): ) 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": 2, + "schema_version": 3, "codec": _CAGRA_CODEC, "writer_policy": "gpu-cagra", + "compound_file_policy": "disabled", "vector_count": 10, "dimensions": 4, "commit_file_count": 1, diff --git a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_integration.py b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_integration.py index e5f6494376..720db47570 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_integration.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_integration.py @@ -9,9 +9,11 @@ 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 @@ -19,13 +21,14 @@ from cuvs_bench._bin_format import write_bin_header from cuvs_bench.backends._utils import compute_recall -from cuvs_bench.backends.base import Dataset +from cuvs_bench.backends.base import BuildResult, Dataset from cuvs_bench.backends.pylucene import ( _BuildCodec, _CAGRA_PROVENANCE_FILE, PyLuceneBackend, _HNSW_PROVENANCE_FILE, _PyLuceneRuntime, + _validate_pylucene_version, ) from cuvs_bench.orchestrator.config_loaders import IndexConfig @@ -42,9 +45,12 @@ _CAGRA_CODEC = "CuVS2510GPUSearchCodec" _HNSW_CODEC = "Lucene101AcceleratedHNSWCodec" _HNSW_WRITER_POLICY = "gpu-with-cpu-fallback" -_GPU_HNSW_WRITER = ( - "com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsWriter" -) +_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" _WRITER_SELECTION_SOURCE = ( Path(__file__).resolve().parents[2] @@ -56,6 +62,33 @@ / "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 + 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: @@ -81,6 +114,55 @@ def _single_search_result(results): 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 _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={"codec": baseline.codec}, + 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, @@ -125,9 +207,7 @@ def _compile_writer_selection_codec( ) -def _assert_cagra_verification( - metadata, expected_vector_count, expected_dimensions -): +def _assert_cagra_verification(metadata, expected_vector_count, expected_dimensions): verification = metadata["cagra_verification"] assert verification["status"] == "cagra-only" assert verification["segment_count"] >= 1 @@ -136,14 +216,13 @@ def _assert_cagra_verification( assert verification["dimensions"] == expected_dimensions -def _assert_cagra_provenance( - metadata, expected_vector_count, expected_dimensions -): +def _assert_cagra_provenance(metadata, expected_vector_count, expected_dimensions): assert metadata["cagra_provenance"] == { "status": "gpu-cagra-provenance", - "schema_version": 2, + "schema_version": 3, "codec": _CAGRA_CODEC, "writer_policy": "gpu-cagra", + "compound_file_policy": "disabled", "vector_count": expected_vector_count, "dimensions": expected_dimensions, "commit_file_count": 1, @@ -181,9 +260,10 @@ def _assert_hnsw_verification( verification = metadata["hnsw_verification"] assert verification == { "status": "gpu-with-cpu-fallback-provenance", - "schema_version": 2, + "schema_version": 3, "codec": codec, "writer_policy": _HNSW_WRITER_POLICY, + "compound_file_policy": "lucene-default", "vector_count": expected_vector_count, "dimensions": expected_dimensions, "commit_file_count": 1, @@ -212,6 +292,7 @@ def pylucene_runtime_config(tmp_path_factory): 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( @@ -222,46 +303,18 @@ def pylucene_runtime_config(tmp_path_factory): ) lucene.CLASSPATH = os.pathsep.join((str(test_classes), lucene.CLASSPATH)) - return { - "cuvs_java_jar": str(cuvs_java_jar), - "cuvs_lucene_jar": str(cuvs_lucene_jar), - "java_library_path": java_library_path, - } + 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.mark.parametrize( - ( - "algo", - "codec", - "expected_writer_policy", - "expected_index_suffix", - ), - [ - pytest.param( - "pylucene_cuvs_hnsw", - "Lucene101AcceleratedHNSWCodec", - _HNSW_WRITER_POLICY, - ".vex", - id="accelerated-hnsw", - ), - pytest.param( - "pylucene_cuvs_cagra", - "CuVS2510GPUSearchCodec", - "gpu-cagra", - ".vcag", - id="cagra", - ), - ], -) -def test_build_and_search_with_real_pylucene_runtime( - tmp_path, - pylucene_runtime_config, - algo, - codec, - expected_writer_policy, - expected_index_suffix, -): - """Build and search through Lucene's public codec and query SPI.""" +@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) @@ -269,24 +322,38 @@ def test_build_and_search_with_real_pylucene_runtime( k = 5 squared_distances = np.sum( - (query_vectors[:, np.newaxis, :] - training_vectors[np.newaxis, :, :]) - ** 2, + (query_vectors[:, np.newaxis, :] - training_vectors[np.newaxis, :, :]) ** 2, axis=2, ) groundtruth_neighbors = np.argsort(squared_distances, axis=1)[:, :k] groundtruth_distances = np.take_along_axis( squared_distances, groundtruth_neighbors, axis=1 ) - dataset = Dataset( - name="pylucene-integration", - training_vectors=training_vectors, - query_vectors=query_vectors, - groundtruth_neighbors=groundtruth_neighbors.astype(np.int32), - groundtruth_distances=groundtruth_distances.astype(np.float32), - distance_metric="euclidean", + 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, ) - index_path = tmp_path / f"{algo}-index" + +def _build_baseline( + tmp_path_factory, + runtime_fixture, + dataset_case, + *, + algo, + codec, + writer_policy, +): + index_path = tmp_path_factory.mktemp(f"{algo}-baseline") / "index" index = IndexConfig( name=f"{algo}-integration", algo=algo, @@ -294,279 +361,522 @@ def test_build_and_search_with_real_pylucene_runtime( search_params=[{}], file=str(index_path), ) - backend = PyLuceneBackend( - { - "name": index.name, - "algo": algo, - "codec": codec, - **pylucene_runtime_config, - } + 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, + index_path=index_path, + index=index, + build_result=build_result, + dataset_case=dataset_case, + runtime_config=runtime_fixture.backend_config, ) - try: - build_result = backend.build(dataset, [index], force=True) - assert build_result.success, build_result.error_message - assert build_result.build_time_seconds > 0 - assert build_result.index_size_bytes > 0 - assert build_result.metadata["codec"] == codec - assert build_result.metadata["pylucene_version"] != "unknown" - assert build_result.metadata["writer_policy"] == expected_writer_policy - assert any( - path.suffix == expected_index_suffix - for path in index_path.iterdir() - ) - if codec == _CAGRA_CODEC: - _assert_cagra_provenance( - build_result.metadata, - training_vectors.shape[0], - training_vectors.shape[1], - ) - _assert_cagra_verification( - build_result.metadata, - training_vectors.shape[0], - training_vectors.shape[1], - ) - assert (index_path / _CAGRA_PROVENANCE_FILE).is_file() - else: - _assert_hnsw_verification( - build_result.metadata, - codec, - training_vectors.shape[0], - training_vectors.shape[1], - ) - assert (index_path / _HNSW_PROVENANCE_FILE).is_file() - - reused_result = backend.build(dataset, [index], force=False) - assert reused_result.success, reused_result.error_message - assert reused_result.metadata["skipped"] is True - if codec == _CAGRA_CODEC: - _assert_cagra_provenance( - reused_result.metadata, - training_vectors.shape[0], - training_vectors.shape[1], - ) - _assert_cagra_verification( - reused_result.metadata, - training_vectors.shape[0], - training_vectors.shape[1], - ) - else: - _assert_hnsw_verification( - reused_result.metadata, - codec, - training_vectors.shape[0], - training_vectors.shape[1], - ) - search_result = _single_search_result( - backend.search(dataset, [index], k=k, batch_size=2) +@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.metadata["codec"] == baseline.codec + assert result.metadata["pylucene_version"] != "unknown" + assert result.metadata["writer_policy"] == baseline.writer_policy + 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], ) - assert search_result.success, search_result.error_message - assert search_result.neighbors.shape == (query_vectors.shape[0], k) - assert search_result.distances.shape == (query_vectors.shape[0], k) - np.testing.assert_array_equal(search_result.neighbors[:, 0], query_ids) - assert np.all(np.isfinite(search_result.distances)) - recall = compute_recall( - search_result.neighbors, groundtruth_neighbors, k + assert (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["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_cagra_provenance( + result.metadata, + training_vectors.shape[0], + training_vectors.shape[1], ) - assert recall >= 0.75 - returned_distances = np.take_along_axis( - squared_distances, search_result.neighbors, axis=1 + _assert_cagra_verification( + result.metadata, + training_vectors.shape[0], + training_vectors.shape[1], ) - np.testing.assert_allclose( - search_result.distances, - returned_distances, - rtol=1e-5, - atol=1e-5, + else: + assert "cagra_verification" not in result.metadata + _assert_hnsw_verification( + result.metadata, + baseline.codec, + training_vectors.shape[0], + training_vectors.shape[1], ) - assert np.all(np.diff(search_result.distances, axis=1) >= -1e-6) - assert search_result.search_time_ms > 0 - assert search_result.metadata["latency_seconds"] > 0 - assert search_result.queries_per_second > 0 - assert search_result.metadata["codec"] == codec - assert search_result.metadata["pylucene_version"] != "unknown" - assert search_result.metadata["num_batches"] == 2 - assert search_result.metadata["mode"] == "latency" - assert set(search_result.latency_percentiles) == {"p50", "p95", "p99"} - if codec == _CAGRA_CODEC: - _assert_cagra_provenance( - search_result.metadata, - training_vectors.shape[0], - training_vectors.shape[1], - ) - _assert_cagra_verification( - search_result.metadata, - training_vectors.shape[0], - training_vectors.shape[1], - ) - wrong_count_dataset = Dataset( - name="pylucene-integration-wrong-count", - training_vectors=training_vectors[:-1], - query_vectors=query_vectors, - distance_metric="euclidean", - ) - rejected = backend.build(wrong_count_dataset, [index], force=False) - assert not rejected.success - assert "vector count does not match the dataset: 512 != 511" in ( - rejected.error_message - ) - rejected = _single_search_result( - backend.search(wrong_count_dataset, [index], k=k) - ) - assert not rejected.success - assert "vector count does not match the dataset: 512 != 511" in ( - rejected.error_message - ) - data_path = next(index_path.glob("*.vcag")) - original_data = data_path.read_bytes() +def test_build_hnsw_with_real_pylucene_runtime(hnsw_baseline): + _assert_build_contract(hnsw_baseline) - data_path.unlink() - rejected = _single_search_result( - backend.search(dataset, [index], k=k) - ) - assert not rejected.success - assert "cannot read" in rejected.error_message - data_path.write_bytes(original_data) - data_path.write_bytes(original_data[:-1]) - rejected = _single_search_result( - backend.search(dataset, [index], k=k) - ) - assert not rejected.success - assert "do not exactly cover" in rejected.error_message - data_path.write_bytes(original_data) - - corrupted_data = bytearray(original_data) - corrupted_data[len(corrupted_data) // 2] ^= 0xFF - original_stat = data_path.stat() - data_path.write_bytes(corrupted_data) - os.utime( - data_path, - ns=(original_stat.st_atime_ns, original_stat.st_mtime_ns), - ) - rejected = _single_search_result( - backend.search(dataset, [index], k=k) - ) - assert not rejected.success - assert "checksum" in rejected.error_message.lower() - data_path.write_bytes(original_data) - - deleted_index_path = tmp_path / f"{algo}-deleted-index" - shutil.copytree(index_path, deleted_index_path) - deleted_index = IndexConfig( - name=f"{algo}-deleted-integration", - algo=algo, - build_param={"codec": codec}, - search_params=[{}], - file=str(deleted_index_path), - ) - _commit_one_deletion(backend, deleted_index_path) - with pytest.raises(RuntimeError, match="committed deletions"): - backend._get_runtime().verify_cagra_index(deleted_index_path) - rejected = _single_search_result( - backend.search(dataset, [deleted_index], k=k) - ) - assert not rejected.success - assert "does not match the current Lucene commit" in ( - rejected.error_message +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.metadata["skipped"] is True + 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], + ) + + +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 + ) + - metadata_path = next(index_path.glob("*.vemc")) - contents = bytearray(metadata_path.read_bytes()) - contents[-1] ^= 0xFF - metadata_path.write_bytes(contents) - with pytest.raises(RuntimeError, match="metadata format v0"): - backend._get_runtime().verify_cagra_index(index_path) - else: - assert "cagra_verification" not in search_result.metadata - _assert_hnsw_verification( - search_result.metadata, - codec, - training_vectors.shape[0], - training_vectors.shape[1], +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, ) - if codec == _HNSW_CODEC: - manifest_path = index_path / _HNSW_PROVENANCE_FILE - original_manifest = manifest_path.read_bytes() - segment_path = next(index_path.glob("segments_*")) - original_segment = segment_path.read_bytes() - - manifest_path.unlink() - rejected = _single_search_result( - backend.search(dataset, [index], k=k) - ) - assert not rejected.success - assert "provenance manifest is missing" in ( - rejected.error_message - ) - manifest_path.write_bytes(original_manifest) - - manifest_path.write_text("{") - rejected = _single_search_result( - backend.search(dataset, [index], k=k) - ) - assert not rejected.success - assert "provenance manifest cannot be read" in ( - rejected.error_message - ) - manifest_path.write_bytes(original_manifest) - - segment_path.write_bytes(original_segment + b"stale") - rejected = _single_search_result( - backend.search(dataset, [index], k=k) - ) - assert not rejected.success - assert "does not match the current Lucene commit" in ( - rejected.error_message - ) - segment_path.write_bytes(original_segment) - - payload = json.loads(original_manifest) - payload["codec"] = _CAGRA_CODEC - manifest_path.write_text(json.dumps(payload)) - rejected = _single_search_result( - backend.search(dataset, [index], k=k) - ) - assert not rejected.success - assert "does not match the requested codec" in ( - rejected.error_message - ) - manifest_path.write_bytes(original_manifest) + ) finally: backend.cleanup() + 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_real_hnsw_codec_selects_gpu_writer(tmp_path, pylucene_runtime_config): - """Prove the production HNSW codec selects its GPU writer on this host.""" - runtime = _PyLuceneRuntime.create(pylucene_runtime_config) + """Exercise GPU writer selection during default-scheduler merges.""" + runtime = _PyLuceneRuntime.create(pylucene_runtime_config.backend_config) runtime.attach_current_thread() - reflected_codec = runtime.Class.forName( - _WRITER_SELECTION_CODEC - ).newInstance() + reflected_codec = runtime.Class.forName(_WRITER_SELECTION_CODEC).newInstance() java_codec = runtime.Codec.cast_(reflected_codec) index_path = tmp_path / "hnsw-writer-selection-index" index_path.mkdir() - - runtime.build_index( - index_path, - np.ones((2, 32), dtype=np.float32), + 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, - ), + ) ) + 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) - diagnostics = str(java_codec.knnVectorsFormat()) - assert f"writerClass={_GPU_HNSW_WRITER}" in diagnostics + 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) + assert search.hits[0][0].document_id == 3 + + +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", + ) + ) + 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) + 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) + 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() @@ -575,6 +885,7 @@ def test_real_verifier_rejects_cagra_brute_force_fallback( index_path, vectors, _BuildCodec( + codec_name=_CAGRA_CODEC, java_codec=java_codec, writer_policy="gpu-cagra", ), @@ -601,7 +912,7 @@ def test_real_verifier_rejects_cagra_brute_force_fallback( "name": index.name, "algo": index.algo, "codec": _CAGRA_CODEC, - **pylucene_runtime_config, + **pylucene_runtime_config.backend_config, } ) backend._runtime = runtime @@ -623,8 +934,7 @@ def test_cli_build_and_search_with_real_pylucene_runtime( query_vectors = training_vectors[query_ids].copy() k = 5 squared_distances = np.sum( - (query_vectors[:, np.newaxis, :] - training_vectors[np.newaxis, :, :]) - ** 2, + (query_vectors[:, np.newaxis, :] - training_vectors[np.newaxis, :, :]) ** 2, axis=2, ) groundtruth_neighbors = np.argsort(squared_distances, axis=1)[:, :k] @@ -668,15 +978,18 @@ def test_cli_build_and_search_with_real_pylucene_runtime( ) backend_config = tmp_path / "pylucene-backend.yaml" backend_config.write_text( - json.dumps({"backend": "pylucene", **pylucene_runtime_config}) + json.dumps( + { + "backend": "pylucene", + **pylucene_runtime_config.backend_config, + } + ) ) environment = os.environ.copy() source_root = str(Path(__file__).resolve().parents[2]) environment["PYTHONPATH"] = os.pathsep.join( - value - for value in (source_root, environment.get("PYTHONPATH")) - if value + value for value in (source_root, environment.get("PYTHONPATH")) if value ) completed = subprocess.run( [ @@ -719,7 +1032,7 @@ def test_cli_build_and_search_with_real_pylucene_runtime( codec = "Lucene101AcceleratedHNSWCodec" index_name = f"pylucene_cuvs_hnsw[group=test][codec={codec}]" index_path = dataset_dir / "index" / index_name - assert any(path.suffix == ".vex" for path in index_path.iterdir()) + assert any(index_path.glob("segments_*")) assert (index_path / _HNSW_PROVENANCE_FILE).is_file() result_path = dataset_dir / "result" @@ -734,6 +1047,7 @@ def test_cli_build_and_search_with_real_pylucene_runtime( 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: @@ -749,5 +1063,6 @@ def test_cli_build_and_search_with_real_pylucene_runtime( 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 index f8f8332c9d..f50a8a8c97 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_provenance.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_provenance.py @@ -41,9 +41,10 @@ def test_hnsw_provenance_round_trip(tmp_path): assert verification.to_metadata() == { "status": "gpu-with-cpu-fallback-provenance", - "schema_version": 2, + "schema_version": 3, "codec": _HNSW_CODEC, "writer_policy": "gpu-with-cpu-fallback", + "compound_file_policy": "lucene-default", "vector_count": 10, "dimensions": 4, "commit_file_count": 1, @@ -74,9 +75,10 @@ def test_cagra_provenance_round_trip(tmp_path): assert verification.to_metadata() == { "status": "gpu-cagra-provenance", - "schema_version": 2, + "schema_version": 3, "codec": _CAGRA_CODEC, "writer_policy": "gpu-cagra", + "compound_file_policy": "disabled", "vector_count": 10, "dimensions": 4, "commit_file_count": 1, @@ -110,6 +112,7 @@ def test_hnsw_provenance_rejects_dataset_shape_mismatch( ("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"), @@ -124,6 +127,8 @@ def test_hnsw_provenance_rejects_malformed_manifest_fields( 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": diff --git a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_runtime.py b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_runtime.py index b0dcb25e69..1a5f6b57cb 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_runtime.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_runtime.py @@ -15,13 +15,18 @@ import cuvs_bench.backends.pylucene as pylucene_backend from cuvs_bench.backends.pylucene import _BuildCodec -from cuvs_bench.tests._pylucene_test_utils import _HNSW_CODEC +from cuvs_bench.tests._pylucene_test_utils import _CAGRA_CODEC, _HNSW_CODEC -def _build_codec(java_codec=None) -> _BuildCodec: +def _build_codec(java_codec=None, codec_name=_HNSW_CODEC) -> _BuildCodec: return _BuildCodec( + codec_name=codec_name, java_codec=java_codec if java_codec is not None else object(), - writer_policy="gpu-with-cpu-fallback", + writer_policy=( + "gpu-cagra" + if codec_name == _CAGRA_CODEC + else "gpu-with-cpu-fallback" + ), ) @@ -35,7 +40,7 @@ def attachCurrentThread(self): class _FakeLuceneModule: CLASSPATH = "/pylucene/lucene-core.jar" - VERSION = "test" + VERSION = pylucene_backend._REQUIRED_PYLUCENE_VERSION def __init__(self): self.environment = None @@ -81,20 +86,38 @@ def close(self): raise RuntimeError("close failed") +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()) + def __init__(self): + self.use_compound_file = None + self.merge_policy = _FakeMergePolicy() + def setOpenMode(self, _mode): pass def setCodec(self, _codec): pass - def setUseCompoundFile(self, _enabled): - pass + def setUseCompoundFile(self, enabled): + self.use_compound_file = enabled + + def getMergePolicy(self): + return self.merge_policy + + def setMergePolicy(self, _merge_policy): + raise AssertionError("PyLucene must use Lucene's default merge policy") def setMergeScheduler(self, _scheduler): - pass + raise AssertionError("PyLucene must use Lucene's default scheduler") class _FakeAvailableCodecs: @@ -146,8 +169,12 @@ def close_directory(): runtime.Paths = SimpleNamespace(get=lambda path: path) runtime.FSDirectory = SimpleNamespace(open=lambda _path: directory) runtime.IndexWriterConfig = _FakeIndexWriterConfig - runtime.SerialMergeScheduler = lambda: object() - runtime.IndexWriter = lambda _directory, _config: writer + + 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(), @@ -215,6 +242,34 @@ def test_initialize_pylucene_uses_verified_classpath_and_vmargs( assert fake_lucene.environment.attach_count == 2 +@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"), [ @@ -488,6 +543,28 @@ def test_runtime_build_index_commits_and_closes_writer_and_directory(tmp_path): 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 + + +@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) + ) + + assert config.use_compound_file is use_compound_file + assert config.merge_policy.no_cfs_ratio == no_cfs_ratio @pytest.mark.parametrize( 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 index 2251449a0c..7d0e0d4e1f 100644 --- a/python/cuvs_bench/tests/java/com/nvidia/cuvs/bench/PyLuceneWriterSelectionCodec.java +++ b/python/cuvs_bench/tests/java/com/nvidia/cuvs/bench/PyLuceneWriterSelectionCodec.java @@ -25,7 +25,8 @@ private static final class WriterSelectionFormat extends KnnVectorsFormat { private static final String NOT_SELECTED = "not-selected"; private final KnnVectorsFormat delegate; - private volatile String writerClass = NOT_SELECTED; + private String writerClass = NOT_SELECTED; + private int fieldsWriterCalls; private WriterSelectionFormat(KnnVectorsFormat delegate) { super(delegate.getName()); @@ -35,14 +36,17 @@ private WriterSelectionFormat(KnnVectorsFormat delegate) { @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { KnnVectorsWriter writer = delegate.fieldsWriter(state); - String selectedClass = writer.getClass().getName(); - String previousClass = writerClass; - if (!NOT_SELECTED.equals(previousClass) && !previousClass.equals(selectedClass)) { + 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 " + previousClass + " to " + selectedClass); + "Vector writer selection changed from " + writerClass + " to " + selectedClass); } writerClass = selectedClass; - return writer; + fieldsWriterCalls++; } @Override @@ -56,8 +60,13 @@ public int getMaxDimensions(String fieldName) { } @Override - public String toString() { - return getName() + "(writerClass=" + writerClass + ")"; + public synchronized String toString() { + return getName() + + "(writerClass=" + + writerClass + + ", fieldsWriterCalls=" + + fieldsWriterCalls + + ")"; } } } From 4110093afe863821c27169fff8e6f999dac8061b Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Mon, 17 Aug 2026 05:16:31 +0000 Subject: [PATCH 10/12] Refresh PyLucene 10.2 setup guidance Document the custom Lucene 10.2 wrapper requirement, temporary PR 174 artifact workflow, codec-specific compound-file policies, fallback behavior, and latency units. Signed-off-by: nvzm123 --- fern/pages/cuvs_bench/install.md | 18 +++++++++--------- fern/pages/cuvs_bench/running.md | 10 ++++++---- 2 files changed, 15 insertions(+), 13 deletions(-) diff --git a/fern/pages/cuvs_bench/install.md b/fern/pages/cuvs_bench/install.md index 5e3ebfa199..2434198e55 100644 --- a/fern/pages/cuvs_bench/install.md +++ b/fern/pages/cuvs_bench/install.md @@ -54,11 +54,11 @@ Exact tags are listed on Docker Hub: 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, `pytest`, and a [source-built PyLucene 10.2.0 installation](https://lucene.apache.org/pylucene/install.html), matching the Lucene 10.2.0 APIs used by the current cuVS-Lucene PR. +- JDK 22 and a custom PyLucene wrapper generated against Lucene 10.2.0. 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 current PR targets cuVS Java 26.10.0. +- 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. -The required PyLucene service-provider and codec support is currently proposed in [NVIDIA/cuvs-lucene#174](https://github.com/NVIDIA/cuvs-lucene/pull/174). Until that work is merged and released, build cuVS-Lucene from that PR's branch; a release or `main` checkout without those changes is insufficient. +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. Until those changes are ported and merged in cuVS, build the thin cuVS-Lucene JAR from that PR in a temporary checkout; a release or `main` checkout without the PR changes is insufficient. 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. @@ -72,8 +72,8 @@ cd .. If matching native cuVS libraries are already built and installed, `./build.sh java` is sufficient. The Java build installs the base and native-classifier JARs into the local Maven repository; see the [cuVS Java build guide](https://github.com/NVIDIA/cuvs/blob/main/java/README.md). ```bash -git clone https://github.com/NVIDIA/cuvs-lucene.git -cd cuvs-lucene +git clone https://github.com/NVIDIA/cuvs-lucene.git cuvs-lucene-pr174 +cd cuvs-lucene-pr174 git fetch origin pull/174/head git switch --detach FETCH_HEAD mvn clean package -DskipTests @@ -83,14 +83,14 @@ After the build, the conventional JAR paths are: ```text ~/.m2/repository/com/nvidia/cuvs/cuvs-java//cuvs-java-.jar -/target/cuvs-lucene-.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. -Then validate the artifacts from the 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. +The backend checks that `lucene.VERSION` is exactly `10.2.0` before starting the process-wide JVM. Then validate the 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 @@ -99,10 +99,10 @@ 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}" -python3 -m pytest -q -s src/test/python/test_pylucene_end_to_end.py +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 complete smoke workflow. +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 diff --git a/fern/pages/cuvs_bench/running.md b/fern/pages/cuvs_bench/running.md index 136d45f1e8..4188f0d290 100644 --- a/fern/pages/cuvs_bench/running.md +++ b/fern/pages/cuvs_bench/running.md @@ -92,7 +92,7 @@ Create `pylucene-backend.yaml` with absolute paths to the base `cuvs-java` JAR, ```yaml backend: pylucene cuvs_java_jar: /home/user/.m2/repository/com/nvidia/cuvs/cuvs-java/VERSION/cuvs-java-VERSION.jar -cuvs_lucene_jar: /absolute/path/to/cuvs-lucene/target/cuvs-lucene-VERSION.jar +cuvs_lucene_jar: /absolute/path/to/cuvs-lucene-pr174/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 ``` @@ -134,11 +134,13 @@ The HNSW `base` and `test` groups use `Lucene101AcceleratedHNSWCodec`. It uses c The supplied configurations use Lucene's public `KnnFloatVectorQuery` and do not expose search-time tuning parameters. The backend accepts FLOAT32 Euclidean datasets with at most 4096 dimensions 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` to avoid cuVS-Lucene 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 a different codec or writer 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. +New HNSW and CAGRA builds atomically write commit-bound provenance manifests named `.cuvs-bench-pylucene-hnsw.json` and `.cuvs-bench-pylucene-cagra.json`, respectively. Reuse and search fail if the applicable manifest is missing, malformed, stale, or names a different codec, 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. -The backend currently supports latency mode with one search thread. `--batch-size` groups queries for measurement, and the reported latency percentiles are per batch. Throughput mode and multiple search threads are not implemented. +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. Both codecs retain Lucene's default index-writer scheduling behavior. -PyLucene's JVM is process-global and can be initialized only once. Set the JAR, native-library locations, and `jvm_args` before the first PyLucene benchmark, and start a new Python process to change any of them. +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) From aee920e56f2662b11573cc7e738c9f30a1f407e2 Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Mon, 17 Aug 2026 15:44:04 +0000 Subject: [PATCH 11/12] Expose PyLucene HNSW benchmark parameters Pass m and ef_construction through a PyLucene-compatible codec adapter, expose num_candidates sweeps, and support verified direct single-segment builds. Persist the complete build identity and extend unit and live coverage for writer selection, fallback, topology, tuning, and reuse. Signed-off-by: nvzm123 --- .../_java/PyLuceneConfiguredHnswCodec.java | 71 ++ .../cuvs_bench/backends/_java/__init__.py | 6 + .../cuvs_bench/backends/_pylucene_java.py | 154 ++++ .../cuvs_bench/backends/pylucene.py | 660 ++++++++++++++---- .../cuvs_bench/backends/search_spaces.py | 16 + .../config/algos/pylucene_cuvs_hnsw.yaml | 9 + .../cuvs_bench/orchestrator/orchestrator.py | 33 +- .../cuvs_bench/tests/_pylucene_test_utils.py | 36 +- .../tests/pylucene_cpu_fallback_probe.py | 38 +- .../cuvs_bench/tests/test_pylucene_backend.py | 330 ++++++++- .../tests/test_pylucene_cli_config.py | 298 +++++++- .../tests/test_pylucene_integration.py | 371 ++++++++-- .../tests/test_pylucene_provenance.py | 85 ++- .../cuvs_bench/tests/test_pylucene_runtime.py | 372 +++++++++- .../bench/PyLuceneWriterSelectionCodec.java | 37 +- 15 files changed, 2245 insertions(+), 271 deletions(-) create mode 100644 python/cuvs_bench/cuvs_bench/backends/_java/PyLuceneConfiguredHnswCodec.java create mode 100644 python/cuvs_bench/cuvs_bench/backends/_java/__init__.py create mode 100644 python/cuvs_bench/cuvs_bench/backends/_pylucene_java.py 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/pylucene.py b/python/cuvs_bench/cuvs_bench/backends/pylucene.py index 0d72bd8956..e62a478c20 100644 --- a/python/cuvs_bench/cuvs_bench/backends/pylucene.py +++ b/python/cuvs_bench/cuvs_bench/backends/pylucene.py @@ -41,6 +41,12 @@ 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" @@ -51,7 +57,15 @@ _HNSW_CODEC = "Lucene101AcceleratedHNSWCodec" _CAGRA_CODEC = "CuVS2510GPUSearchCodec" _SUPPORTED_CODECS = frozenset({_HNSW_CODEC, _CAGRA_CODEC}) -_SUPPORTED_BUILD_KEYS = frozenset({"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", @@ -74,15 +88,17 @@ _HNSW_PROVENANCE_FILE = ".cuvs-bench-pylucene-hnsw.json" _CAGRA_PROVENANCE_FILE = ".cuvs-bench-pylucene-cagra.json" -_PROVENANCE_SCHEMA_VERSION = 3 +_PROVENANCE_SCHEMA_VERSION = 4 _PROVENANCE_KEYS = frozenset( { "schema_version", "codec", + "build_parameters", "writer_policy", "compound_file_policy", "vector_count", "dimensions", + "segment_count", "commit_fingerprints", } ) @@ -90,6 +106,7 @@ _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 @@ -113,8 +130,7 @@ def _attempt_cleanup( ) return cleanup_error primary_error.add_note( - f"Failed to {description}: " - f"{type(cleanup_error).__name__}: {cleanup_error}" + f"Failed to {description}: {type(cleanup_error).__name__}: {cleanup_error}" ) return primary_error @@ -154,6 +170,13 @@ def _exception_summary(error: Exception) -> str: 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: @@ -218,8 +241,18 @@ def _pylucene_classpath(config: Dict[str, Any], lucene: Any) -> str: 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(cuvs_java_jar), str(cuvs_lucene_jar), lucene.CLASSPATH] + [ + str(adapter_classes), + str(cuvs_java_jar), + str(cuvs_lucene_jar), + str(lucene.CLASSPATH), + ] ) @@ -322,8 +355,7 @@ def _validate_float32_matrix(vectors: np.ndarray, name: str) -> np.ndarray: raise TypeError(f"{name} must use float32 values, got {array.dtype}") if array.shape[1] > _MAX_DIMENSIONS: raise ValueError( - f"{name} dimensions must not exceed {_MAX_DIMENSIONS}, " - f"got {array.shape[1]}" + 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") @@ -342,8 +374,7 @@ def _validate_codec(codec_name: Any) -> str: if not isinstance(codec_name, str) or codec_name not in _SUPPORTED_CODECS: available = ", ".join(sorted(_SUPPORTED_CODECS)) raise ValueError( - f"Unsupported PyLucene codec {codec_name!r}. " - f"Supported codecs: {available}" + f"Unsupported PyLucene codec {codec_name!r}. Supported codecs: {available}" ) return codec_name @@ -358,27 +389,107 @@ def _configured_codec( return _validate_codec(codec_name) -def _validate_build_params(build_params: Any) -> Dict[str, Any]: +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}") - return build_params + 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 _validate_search_params(search_params: Any) -> List[Dict[str, Any]]: - if ( - not isinstance(search_params, list) - or len(search_params) != 1 - or search_params[0] != {} - ): + +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's public Lucene query API does not expose " - "cuVS-specific search parameters" + "PyLucene search parameters must be a non-empty list of mappings" ) - return search_params + + 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: @@ -392,8 +503,7 @@ def _safe_remove_index(index_path: Path, trusted_index_root: Path) -> None: raise ValueError(f"Refusing to remove unsafe index path: {resolved}") if resolved.parent != trusted_index_root.resolve(): raise ValueError( - "Refusing to remove PyLucene index outside its configured root: " - f"{resolved}" + f"Refusing to remove PyLucene index outside its configured root: {resolved}" ) if index_path.is_symlink() or not index_path.is_dir(): raise ValueError( @@ -460,8 +570,7 @@ def _commit_fingerprints(index_path: Path) -> List[Dict[str, str]]: for path in commit_files: if path.is_symlink() or not path.is_file(): raise RuntimeError( - "Lucene commit fingerprint target must be a regular file: " - f"{path}" + f"Lucene commit fingerprint target must be a regular file: {path}" ) fingerprints.append({"name": path.name, "sha256": _sha256_file(path)}) return fingerprints @@ -470,21 +579,25 @@ def _commit_fingerprints(index_path: Path) -> List[Dict[str, str]]: @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, Union[str, 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, } @@ -492,6 +605,7 @@ def to_metadata(self) -> Dict[str, Union[str, int]]: @dataclass(frozen=True) class _ProvenanceExpectation: codec: str + build_parameters: Dict[str, Any] manifest_name: str label: str vector_count: Optional[int] @@ -505,17 +619,21 @@ class _IndexProvenanceError(RuntimeError): 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 @@ -552,12 +670,19 @@ def _write_hnsw_provenance( 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, ) @@ -566,12 +691,19 @@ 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, ) @@ -581,8 +713,7 @@ def _require_positive_int( ) -> int: if type(value) is not int or value < 1: raise _IndexProvenanceError( - f"{provenance_label} field {field_name!r} must be a " - "positive integer" + f"{provenance_label} field {field_name!r} must be a positive integer" ) return value @@ -598,8 +729,7 @@ def _read_index_provenance( ) from exc except (OSError, UnicodeError, json.JSONDecodeError) as exc: raise _IndexProvenanceError( - f"{provenance_label} manifest cannot be read: " - f"{manifest_path}: {exc}" + f"{provenance_label} manifest cannot be read: {manifest_path}: {exc}" ) from exc if not isinstance(payload, dict) or set(payload) != _PROVENANCE_KEYS: @@ -611,9 +741,9 @@ def _read_index_provenance( def _validate_provenance_identity( payload: Dict[str, Any], - expected_codec: str, + expected_build_parameters: Dict[str, Any], provenance_label: str, -) -> Tuple[str, str]: +) -> Tuple[Dict[str, Any], str, str]: if ( type(payload["schema_version"]) is not int or payload["schema_version"] != _PROVENANCE_SCHEMA_VERSION @@ -621,11 +751,29 @@ def _validate_provenance_identity( 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]: @@ -635,10 +783,9 @@ def _validate_provenance_identity( 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" + f"{provenance_label} does not record the expected compound-file policy" ) - return writer_policy, compound_file_policy + return stored_build_parameters, writer_policy, compound_file_policy def _validate_provenance_shape( @@ -678,6 +825,14 @@ def _validate_provenance_shape( 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]]: @@ -751,8 +906,12 @@ def _verify_index_provenance( payload = _read_index_provenance( index_path / expectation.manifest_name, expectation.label ) - writer_policy, compound_file_policy = _validate_provenance_identity( - payload, expectation.codec, 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, @@ -760,6 +919,14 @@ def _verify_index_provenance( 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 ) @@ -776,10 +943,12 @@ def _verify_index_provenance( 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), ) @@ -788,13 +957,18 @@ 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, @@ -806,13 +980,18 @@ def _verify_hnsw_provenance( 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, @@ -841,8 +1020,7 @@ def _expected_training_shape(dataset: Dataset) -> Optional[Tuple[int, int]]: ) if vectors.dtype != np.float32: raise TypeError( - "training_vectors must use float32 values, " - f"got {vectors.dtype}" + f"training_vectors must use float32 values, got {vectors.dtype}" ) rows, dimensions = vectors.shape elif dataset.base_file: @@ -907,7 +1085,9 @@ class _SearchInputs: class _SearchPlan: index_path: Path codec_name: str - search_params: List[Dict[str, Any]] + build_parameters: Dict[str, Any] + search_parameters: Dict[str, Any] + num_candidates: int k: int batch_size: int mode: str @@ -918,6 +1098,56 @@ 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): @@ -993,8 +1223,7 @@ def _validate_search_hit( ) if not np.isfinite(hit.score): raise RuntimeError( - "PyLucene returned a non-finite score for " - f"query {query_id}: {hit.score}" + f"PyLucene returned a non-finite score for query {query_id}: {hit.score}" ) if not 0.0 <= hit.score <= 1.0: raise RuntimeError( @@ -1031,8 +1260,7 @@ def _convert_query_hits( ) -> Tuple[List[int], List[float]]: if len(hits) > min(k, document_count): raise RuntimeError( - f"PyLucene returned too many hits for query {query_id}: " - f"{len(hits)}" + f"PyLucene returned too many hits for query {query_id}: {len(hits)}" ) neighbors = [] @@ -1382,15 +1610,11 @@ def _verify_cagra_payload_coverage( if not intervals: if payload_start != payload_end: raise _CagraVerificationError( - "CAGRA data file contains unreferenced payload: " - f"{data_file!r}" + f"CAGRA data file contains unreferenced payload: {data_file!r}" ) return - coverage_error = ( - "CAGRA metadata ranges do not exactly cover data " - f"file {data_file!r}" - ) + 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: @@ -1481,8 +1705,7 @@ def _verify_cagra_data_file( raise except Exception as exc: raise _CagraVerificationError( - "CAGRA-only verification cannot read " - f"{context.data_file!r}: {exc}" + f"CAGRA-only verification cannot read {context.data_file!r}: {exc}" ) from exc def _verify_field_against_lucene_metadata( @@ -1756,7 +1979,7 @@ class _PyLuceneRuntime: """Own generated PyLucene/Lucene bindings and index operations.""" def __init__(self, lucene: Any): - from java.lang import Class + 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 ( @@ -1769,6 +1992,7 @@ def __init__(self, lucene: Any): FieldInfo, IndexWriter, IndexWriterConfig, + NoMergePolicy, SegmentCommitInfo, SegmentInfos, VectorEncoding, @@ -1782,6 +2006,7 @@ def __init__(self, lucene: Any): self.lucene = lucene self.Class = Class + self.System = System self.Paths = Paths self.Codec = Codec self.Document = Document @@ -1790,6 +2015,7 @@ def __init__(self, lucene: Any): self.DirectoryReader = DirectoryReader self.IndexWriter = IndexWriter self.IndexWriterConfig = IndexWriterConfig + self.NoMergePolicy = NoMergePolicy self.VectorSimilarityFunction = VectorSimilarityFunction self.IndexSearcher = IndexSearcher self.KnnFloatVectorQuery = KnnFloatVectorQuery @@ -1841,6 +2067,45 @@ def resolve_codec(self, codec_name: str) -> Any: 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()}" @@ -1849,8 +2114,6 @@ def resolve_codec(self, codec_name: str) -> Any: raise RuntimeError( f"{codec_name} did not initialize a Lucene vector format" ) - self._codec_cache[codec_name] = codec - return codec def _java_float_array(self, vector: np.ndarray) -> Any: return self.lucene.JArray("float")( @@ -1874,17 +2137,25 @@ def build_index( index_path: Path, vectors: np.ndarray, build_codec: _BuildCodec, - ) -> None: + ) -> _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) + 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 + self, build_codec: _BuildCodec, *, vector_count: int ) -> Any: writer_config = self.IndexWriterConfig() writer_config.setOpenMode(self.IndexWriterConfig.OpenMode.CREATE) @@ -1894,8 +2165,35 @@ def _new_index_writer_config( # 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): @@ -1945,8 +2243,7 @@ def _index_vector_dimensions(reader: Any) -> int: ) if len(dimensions) != 1: raise RuntimeError( - "Lucene index has inconsistent vector dimensions: " - f"{dimensions}" + f"Lucene index has inconsistent vector dimensions: {dimensions}" ) return dimensions.pop() @@ -1956,11 +2253,12 @@ def _search_vector( stored_fields: Any, vector: np.ndarray, k: int, + num_candidates: int, ) -> List[_SearchHit]: query = self.KnnFloatVectorQuery( _VECTOR_FIELD, self._java_float_array(vector), - k, + num_candidates, ) score_docs = searcher.search(query, k).scoreDocs hits = [] @@ -1984,6 +2282,7 @@ def search_index( 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))) @@ -2002,6 +2301,10 @@ def search_index( 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() @@ -2018,6 +2321,7 @@ def search_index( stored_fields, query_vectors[batch_start:batch_end], lucene_k, + lucene_num_candidates, ) ) batch_latencies_ms.append( @@ -2037,11 +2341,18 @@ def _search_vectors( 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) + self._search_vector( + searcher, + stored_fields, + vector, + k, + num_candidates, + ) ) return hits @@ -2309,8 +2620,14 @@ def _effective_parameters( 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 [tune_build_params], search_params + return [build_params], search_params @staticmethod def _index_label( @@ -2322,8 +2639,18 @@ def _index_label( identity = format_artifact_identity( algorithm, group, subset_scope, description="PyLucene" ) - codec = _validate_codec(build_params.get("codec")) - return f"{identity}[codec={codec}]" + 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( @@ -2334,6 +2661,10 @@ def _benchmark_config( 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" ) @@ -2356,8 +2687,8 @@ def _benchmark_config( "index_name": index_label, "index_root": str(index_root), "result_scope": context.subset_scope, - "codec": build_params.get("codec"), - "requires_gpu": build_params.get("codec") == _CAGRA_CODEC, + "codec": build_params["codec"], + "requires_gpu": build_params["codec"] == _CAGRA_CODEC, **context.runtime_config, } return BenchmarkConfig( @@ -2374,16 +2705,15 @@ def _group_benchmark_configs( build_params: List[Dict[str, Any]], search_params: List[Dict[str, Any]], ) -> List[BenchmarkConfig]: - _validate_search_params(search_params) benchmark_configs = [] for params in build_params: - _validate_build_params(params) + effective_params = _normalize_build_params(params) benchmark_configs.append( cls._benchmark_config( context, algorithm, group, - params, + effective_params, search_params, ) ) @@ -2460,11 +2790,14 @@ def _result_identity(self) -> Dict[str, Any]: identity["result_scope"] = result_scope return identity - def _index_metadata(self, codec_name: str) -> Dict[str, Any]: + def _index_metadata( + self, build_parameters: Dict[str, Any] + ) -> Dict[str, Any]: + codec_name = build_parameters["codec"] return { **self._result_identity(), - "codec": codec_name, "compound_file_policy": _COMPOUND_FILE_POLICY[codec_name], + **build_parameters, } def cleanup(self) -> None: @@ -2537,15 +2870,17 @@ def _failed_search_result( def _verify_existing_index( self, index_path: Path, - codec_name: str, + 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, ) @@ -2562,6 +2897,7 @@ def _verify_existing_index( provenance = _verify_hnsw_provenance( index_path, codec_name, + expected_build_parameters=build_parameters, expected_vector_count=expected_vector_count, expected_dimensions=expected_dimensions, ) @@ -2571,20 +2907,19 @@ def _verify_existing_index( def _dry_run_build_result( self, index_path: Path, - codec_name: str, - build_params: Dict[str, Any], + build_parameters: Dict[str, Any], ) -> BuildResult: print( f"[dry_run] Would build PyLucene index '{index_path}' " - f"with codec={codec_name}" + 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_params, - metadata=self._index_metadata(codec_name), + build_params=build_parameters, + metadata=self._index_metadata(build_parameters), success=True, ) @@ -2592,8 +2927,7 @@ def _decide_existing_index( self, dataset: Dataset, index_path: Path, - codec_name: str, - build_params: Dict[str, Any], + build_parameters: Dict[str, Any], force: bool, ) -> _ExistingIndexDecision: if not index_path.exists(): @@ -2603,7 +2937,7 @@ def _decide_existing_index( self._failed_build_result( f"PyLucene index path is not a directory: {index_path}", str(index_path), - build_params, + build_parameters, ) ) if not _has_lucene_segments(index_path): @@ -2612,22 +2946,22 @@ def _decide_existing_index( "Existing PyLucene index directory does not contain " f"a Lucene segments file: {index_path}", str(index_path), - build_params, + build_parameters, ) ) if force: return _ExistingIndexDecision.build() metadata: Dict[str, Any] = { - **self._index_metadata(codec_name), + **self._index_metadata(build_parameters), "skipped": True, } try: _validate_metric(dataset) expected_shape = _expected_training_shape(dataset) - _, verification_metadata = self._verify_existing_index( + provenance, verification_metadata = self._verify_existing_index( index_path, - codec_name, + build_parameters, expected_vector_count=( expected_shape[0] if expected_shape else None ), @@ -2635,13 +2969,14 @@ def _decide_existing_index( 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_params, + build_params=build_parameters, metadata=metadata, success=True, ) @@ -2650,7 +2985,7 @@ def _decide_existing_index( self._failed_build_result( _exception_summary(exc), str(index_path), - build_params, + build_parameters, ) ) @@ -2673,21 +3008,32 @@ def _training_vectors_for_build( @staticmethod def _resolve_build_codec( - runtime: _PyLuceneRuntime, codec_name: str + 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=runtime.resolve_codec(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, - codec_name: str, + 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: @@ -2700,9 +3046,12 @@ def _verify_and_persist_built_index_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, ) @@ -2716,10 +3065,13 @@ def _verify_and_persist_built_index_provenance( 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, ) @@ -2750,29 +3102,26 @@ def build( ) index_config = indexes[0] - build_params = index_config.build_param + requested_build_params = index_config.build_param index_path = Path(index_config.file) try: - _validate_build_params(build_params) - codec_name = _configured_codec(build_params, self.config) + 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), build_params + str(exc), str(index_path), requested_build_params ) if dry_run: - return self._dry_run_build_result( - index_path, codec_name, build_params - ) + return self._dry_run_build_result(index_path, build_parameters) existing_index_decision = self._decide_existing_index( - dataset, index_path, codec_name, build_params, force + dataset, index_path, build_parameters, force ) if existing_index_decision.action is _ExistingIndexAction.BUILD: - return self._build_new_index( - dataset, index_path, codec_name, build_params - ) + return self._build_new_index(dataset, index_path, build_parameters) return existing_index_decision.completed_result() @@ -2780,17 +3129,17 @@ def _build_new_index( self, dataset: Dataset, index_path: Path, - codec_name: str, - build_params: Dict[str, Any], + 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, codec_name) + build_codec = self._resolve_build_codec(runtime, build_parameters) if index_path.exists(): self._remove_index(index_path) @@ -2798,17 +3147,22 @@ def _build_new_index( created_for_build = True start = time.perf_counter() - runtime.build_index(index_path, vectors, build_codec) + topology = runtime.build_index(index_path, vectors, build_codec) build_time = time.perf_counter() - start metadata = { - **self._index_metadata(codec_name), + **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, codec_name, vectors + runtime, + index_path, + build_parameters, + vectors, + topology, ) ) @@ -2817,7 +3171,7 @@ def _build_new_index( build_time_seconds=build_time, index_size_bytes=_index_size(index_path), algorithm=self.algo, - build_params=build_params, + build_params=build_parameters, metadata=metadata, success=True, ) @@ -2834,7 +3188,7 @@ def _build_new_index( return self._failed_build_result( error_message, str(index_path), - build_params, + build_parameters, ) except BaseException as exc: cleanup_error = self._cleanup_partial_index( @@ -2847,7 +3201,7 @@ def _build_new_index( ) raise - def _resolve_search_plan( + def _resolve_search_plans( self, index_config: IndexConfig, *, @@ -2855,10 +3209,11 @@ def _resolve_search_plan( batch_size: int, mode: str, search_threads: Optional[Union[int, str]], - ) -> _SearchPlan: - search_params = index_config.search_params or [{}] - _validate_build_params(index_config.build_param) - codec_name = _configured_codec(index_config.build_param, self.config) + ) -> 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: @@ -2871,21 +3226,28 @@ def _resolve_search_plan( raise ValueError( "PyLucene backend currently supports only one search thread" ) - _validate_search_params(search_params) + 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, - search_params=search_params, - k=k, - batch_size=batch_size, - mode=mode, - ) + 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, @@ -2894,6 +3256,7 @@ def _dry_run_search_result( 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( @@ -2903,8 +3266,14 @@ def _dry_run_search_result( queries_per_second=0.0, recall=0.0, algorithm=self.algo, - search_params=plan.search_params, - metadata=self._index_metadata(plan.codec_name), + 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, ) @@ -2937,7 +3306,7 @@ def _execute_search( expected_shape = inputs.expected_training_shape provenance, verification_metadata = self._verify_existing_index( plan.index_path, - plan.codec_name, + plan.build_parameters, expected_vector_count=( expected_shape[0] if expected_shape else None ), @@ -2950,6 +3319,7 @@ def _execute_search( 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( @@ -2960,10 +3330,13 @@ def _execute_search( batch_size=plan.batch_size, ) metadata = { - **self._index_metadata(plan.codec_name), + **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, @@ -2982,7 +3355,7 @@ def _execute_search( queries_per_second=processed.queries_per_second, recall=0.0, algorithm=self.algo, - search_params=plan.search_params, + search_params=[plan.search_parameters], latency_percentiles=processed.latency_percentiles, metadata=metadata, success=True, @@ -3010,7 +3383,7 @@ def search( index_config = indexes[0] try: - plan = self._resolve_search_plan( + plans = self._resolve_search_plans( index_config, k=k, batch_size=batch_size, @@ -3027,28 +3400,31 @@ def search( ] if dry_run: - return [self._dry_run_search_result(plan)] + return [self._dry_run_search_result(plan) for plan in plans] - if not plan.index_path.is_dir(): + index_path = plans[0].index_path + if not index_path.is_dir(): return [ self._failed_search_result( k, - "PyLucene index directory does not exist: " - f"{plan.index_path}", - plan.search_params, + f"PyLucene index directory does not exist: {index_path}", + [plan.search_parameters for plan in plans], ) ] - try: - return [self._execute_search(dataset, plan)] - except Exception as exc: - return [ - self._failed_search_result( - k, - _exception_summary(exc), - plan.search_params, + 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_hnsw.yaml b/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_hnsw.yaml index bd8ab2849c..16c5bab34b 100644 --- a/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_hnsw.yaml +++ b/python/cuvs_bench/cuvs_bench/config/algos/pylucene_cuvs_hnsw.yaml @@ -4,8 +4,17 @@ groups: 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/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/tests/_pylucene_test_utils.py b/python/cuvs_bench/cuvs_bench/tests/_pylucene_test_utils.py index 1c3d946c9f..0b9cf6dc90 100644 --- a/python/cuvs_bench/cuvs_bench/tests/_pylucene_test_utils.py +++ b/python/cuvs_bench/cuvs_bench/tests/_pylucene_test_utils.py @@ -17,6 +17,7 @@ from cuvs_bench.backends.pylucene import ( PyLuceneBackend, _CagraIndexVerification, + _IndexTopology, _RuntimeSearchResult, _SearchHit, ) @@ -44,11 +45,13 @@ def __init__( 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 @@ -58,6 +61,12 @@ def resolve_codec(self, codec_name): 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: @@ -65,6 +74,15 @@ def build_index(self, index_path, vectors, build_codec): 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, @@ -90,9 +108,17 @@ def verify_cagra_index( ), ) - def search_index(self, index_path, query_vectors, k, batch_size): + 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) + ( + index_path, + query_vectors.copy(), + k, + batch_size, + num_candidates, + ) ) if self.search_error is not None: raise self.search_error @@ -145,12 +171,16 @@ 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={"codec": codec}, + build_param=parameters, search_params=[{}] if search_params is None else search_params, file=str(index_path), ) 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 index 10351f89b2..647e4797fa 100644 --- a/python/cuvs_bench/cuvs_bench/tests/pylucene_cpu_fallback_probe.py +++ b/python/cuvs_bench/cuvs_bench/tests/pylucene_cpu_fallback_probe.py @@ -13,20 +13,34 @@ 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" +_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) + 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)) @@ -44,9 +58,15 @@ def main(): "java_library_path": os.environ["JAVA_LIBRARY_PATH"], } ) - reflected_codec = runtime.Class.forName(_WRITER_SELECTION_CODEC).newInstance() + 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) + 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) @@ -57,17 +77,25 @@ def main(): 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}") + 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 diff --git a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_backend.py b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_backend.py index 65eeafe457..5f1184ccea 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_backend.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_backend.py @@ -67,18 +67,35 @@ def test_build_creates_index_and_records_hnsw_writer_policy(tmp_path): assert result.metadata["compound_file_policy"] == "lucene-default" assert result.metadata["hnsw_verification"] == { "status": "gpu-with-cpu-fallback-provenance", - "schema_version": 3, + "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.resolve_calls == [_HNSW_CODEC] + 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 == [] @@ -97,12 +114,14 @@ def test_build_verifies_persisted_cagra_index(tmp_path): assert runtime.verification_calls == [(index_path, 10, 4)] assert result.metadata["cagra_provenance"] == { "status": "gpu-cagra-provenance", - "schema_version": 3, + "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"] == { @@ -450,6 +469,96 @@ def test_build_rejects_unsupported_parameter(tmp_path): 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 = [ @@ -497,6 +606,32 @@ def test_build_removes_partial_index_after_runtime_failure(tmp_path): 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") @@ -576,6 +711,73 @@ def test_search_dry_run_does_not_initialize_pylucene(tmp_path): 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): @@ -614,11 +816,13 @@ def test_search_converts_hits_scores_and_padding(tmp_path): "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][3] == 1 + assert runtime.search_calls[0][2:] == (3, 1, 3) assert runtime.verification_calls == [] @@ -867,7 +1071,7 @@ def test_search_rejects_invalid_cagra_provenance_before_runtime( 10000, None, "latency", - "does not expose", + "Unsupported PyLucene search parameter", ), ( _index(Path("/tmp/index"), codec=_CAGRA_CODEC), @@ -898,10 +1102,97 @@ def test_search_rejects_unsupported_options( assert error in result.error_message -def test_resolve_search_plan_normalizes_the_complete_request(tmp_path): +@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=[]) - plan = _backend()._resolve_search_plan( + plans = _backend()._resolve_search_plans( index, k=3, batch_size=7, @@ -909,14 +1200,23 @@ def test_resolve_search_plan_normalizes_the_complete_request(tmp_path): search_threads="1", ) - assert plan == pylucene_backend._SearchPlan( - index_path=tmp_path / "index", - codec_name=_HNSW_CODEC, - search_params=[{}], - k=3, - batch_size=7, - mode="latency", - ) + 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): 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 index f24c4f3cf7..129f4758a1 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_cli_config.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_cli_config.py @@ -11,6 +11,7 @@ import subprocess import sys from pathlib import Path +from types import SimpleNamespace import numpy as np import pytest @@ -18,11 +19,14 @@ 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, @@ -31,6 +35,68 @@ ) +_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( [ @@ -201,7 +267,7 @@ def test_cli_build_search_persists_metrics_and_build_join( ) assert result.exit_code == 0, result.output - index_name = f"pylucene_test[group=test][codec={_HNSW_CODEC}]" + 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() @@ -304,6 +370,217 @@ def test_config_loader_expands_codecs_and_forwards_runtime_config(config_dir): 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"), [ @@ -315,7 +592,7 @@ def test_config_loader_expands_codecs_and_forwards_runtime_config(config_dir): ( f"codec: [{_HNSW_CODEC}]", "ignored: [1]", - "does not expose cuVS-specific search parameters", + "Unsupported PyLucene search parameter.*ignored", ), ], ) @@ -402,10 +679,9 @@ def test_config_loader_scopes_artifact_identity_by_dataset_subset(config_dir): assert dataset_config.subset_size == subset_size assert len(configs) == 1 config = configs[0] - index_prefix = "pylucene_test[group=test]" - if subset_scope is not None: - index_prefix = f"{index_prefix}[scope={subset_scope}]" - index_name = f"{index_prefix}[codec={_HNSW_CODEC}]" + 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 @@ -432,7 +708,7 @@ def test_config_loader_scopes_artifact_identity_by_dataset_subset(config_dir): ], ) def test_index_labels_are_injective(first, second): - build_params = {"codec": _HNSW_CODEC} + build_params = _DEFAULT_HNSW_BUILD_PARAMETERS first_label = PyLuceneConfigLoader._index_label( first[0], first[1], first[2], build_params @@ -495,9 +771,7 @@ def test_config_loader_discovers_custom_filename_by_parsed_identity( assert len(configs) == 1 assert configs[0].indexes[0].algo == algorithm_name - assert configs[0].index_name == ( - f"{algorithm_name}[group=test][codec={_HNSW_CODEC}]" - ) + assert configs[0].index_name == _hnsw_index_name(algorithm_name, "test") def test_config_loader_excludes_explicit_other_backend(config_dir, tmp_path): @@ -538,9 +812,7 @@ def test_config_loader_honors_algorithm_and_group_filters(config_dir): groups="test", ) assert len(configs) == 1 - assert configs[0].index_name == ( - f"pylucene_test[group=test][codec={_HNSW_CODEC}]" - ) + assert configs[0].index_name == _hnsw_index_name("pylucene_test", "test") def test_algorithm_config_discovery_is_deterministic(tmp_path): diff --git a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_integration.py b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_integration.py index 720db47570..fdfc4b04fb 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_integration.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_integration.py @@ -20,6 +20,10 @@ 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 ( @@ -28,6 +32,7 @@ PyLuceneBackend, _HNSW_PROVENANCE_FILE, _PyLuceneRuntime, + _restore_java_property, _validate_pylucene_version, ) from cuvs_bench.orchestrator.config_loaders import IndexConfig @@ -49,9 +54,20 @@ _HNSW_CODEC: "lucene-default", _CAGRA_CODEC: "disabled", } -_GPU_HNSW_WRITER = "com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsWriter" -_CPU_HNSW_WRITER = "org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsWriter" +_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" @@ -62,7 +78,9 @@ / "bench" / "PyLuceneWriterSelectionCodec.java" ) -_CPU_FALLBACK_PROBE = Path(__file__).with_name("pylucene_cpu_fallback_probe.py") +_CPU_FALLBACK_PROBE = Path(__file__).with_name( + "pylucene_cpu_fallback_probe.py" +) @dataclass(frozen=True) @@ -84,6 +102,7 @@ class _BaselineIndex: algo: str codec: str writer_policy: str + build_parameters: dict[str, object] index_path: Path index: IndexConfig build_result: BuildResult @@ -116,11 +135,31 @@ def _single_search_result(results): def _writer_diagnostics(java_codec): diagnostics = str(java_codec.knnVectorsFormat()) - match = re.search(r"writerClass=([^,)]+), fieldsWriterCalls=(\d+)", diagnostics) + 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( { @@ -136,7 +175,7 @@ def _index_at(baseline, index_path): return IndexConfig( name=baseline.index.name, algo=baseline.algo, - build_param={"codec": baseline.codec}, + build_param=baseline.build_parameters, search_params=[{}], file=str(index_path), ) @@ -150,7 +189,9 @@ def _copy_baseline(baseline, tmp_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 + search_dataset = ( + baseline.dataset_case.dataset if dataset is None else dataset + ) try: return _single_search_result( backend.search( @@ -207,7 +248,9 @@ def _compile_writer_selection_codec( ) -def _assert_cagra_verification(metadata, expected_vector_count, expected_dimensions): +def _assert_cagra_verification( + metadata, expected_vector_count, expected_dimensions +): verification = metadata["cagra_verification"] assert verification["status"] == "cagra-only" assert verification["segment_count"] >= 1 @@ -216,15 +259,21 @@ def _assert_cagra_verification(metadata, expected_vector_count, expected_dimensi assert verification["dimensions"] == expected_dimensions -def _assert_cagra_provenance(metadata, expected_vector_count, expected_dimensions): +def _assert_cagra_provenance( + metadata, + expected_vector_count, + expected_dimensions, +): assert metadata["cagra_provenance"] == { "status": "gpu-cagra-provenance", - "schema_version": 3, + "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, } @@ -255,17 +304,23 @@ def _commit_one_deletion(backend, index_path): def _assert_hnsw_verification( - metadata, codec, expected_vector_count, expected_dimensions + metadata, + codec, + expected_vector_count, + expected_dimensions, + expected_build_parameters, ): verification = metadata["hnsw_verification"] assert verification == { "status": "gpu-with-cpu-fallback-provenance", - "schema_version": 3, + "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, } @@ -322,7 +377,8 @@ def integration_dataset_case(): k = 5 squared_distances = np.sum( - (query_vectors[:, np.newaxis, :] - training_vectors[np.newaxis, :, :]) ** 2, + (query_vectors[:, np.newaxis, :] - training_vectors[np.newaxis, :, :]) + ** 2, axis=2, ) groundtruth_neighbors = np.argsort(squared_distances, axis=1)[:, :k] @@ -353,11 +409,16 @@ def _build_baseline( 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={"codec": codec}, + build_param=build_parameters, search_params=[{}], file=str(index_path), ) @@ -371,6 +432,7 @@ def _build_baseline( algo=algo, codec=codec, writer_policy=writer_policy, + build_parameters=build_parameters, index_path=index_path, index=index, build_result=build_result, @@ -380,7 +442,9 @@ def _build_baseline( @pytest.fixture(scope="module") -def hnsw_baseline(tmp_path_factory, pylucene_runtime_config, integration_dataset_case): +def hnsw_baseline( + tmp_path_factory, pylucene_runtime_config, integration_dataset_case +): return _build_baseline( tmp_path_factory, pylucene_runtime_config, @@ -392,7 +456,9 @@ def hnsw_baseline(tmp_path_factory, pylucene_runtime_config, integration_dataset @pytest.fixture(scope="module") -def cagra_baseline(tmp_path_factory, pylucene_runtime_config, integration_dataset_case): +def cagra_baseline( + tmp_path_factory, pylucene_runtime_config, integration_dataset_case +): return _build_baseline( tmp_path_factory, pylucene_runtime_config, @@ -409,17 +475,28 @@ def _assert_build_contract(baseline): 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] + 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 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( @@ -427,6 +504,7 @@ def _assert_build_contract(baseline): baseline.codec, vectors.shape[0], vectors.shape[1], + baseline.build_parameters, ) assert (baseline.index_path / _HNSW_PROVENANCE_FILE).is_file() @@ -441,7 +519,9 @@ def _assert_search_contract(baseline, result): 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) + 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 @@ -454,14 +534,17 @@ def _assert_search_contract(baseline, result): 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] + 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], @@ -473,12 +556,15 @@ def _assert_search_contract(baseline, result): 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, ) @@ -505,21 +591,28 @@ def _assert_reuse_contract(baseline_index): 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]) + _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, ) @@ -742,7 +835,9 @@ def _committed_segment_compound_flags(runtime, index_path): segment_infos = verifier.SegmentInfos.readLatestCommit(directory) return [ bool( - verifier.SegmentCommitInfo.cast_(raw_segment).info.getUseCompoundFile() + verifier.SegmentCommitInfo.cast_( + raw_segment + ).info.getUseCompoundFile() ) for raw_segment in segment_infos ] @@ -750,29 +845,46 @@ def _committed_segment_compound_flags(runtime, index_path): directory.close() -def test_real_hnsw_codec_selects_gpu_writer(tmp_path, pylucene_runtime_config): - """Exercise GPU writer selection during default-scheduler merges.""" +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() - reflected_codec = runtime.Class.forName(_WRITER_SELECTION_CODEC).newInstance() - java_codec = runtime.Codec.cast_(reflected_codec) + 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) + 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 float( + writer_config.getMergePolicy().getNoCFSRatio() + ) == pytest.approx(float(default_config.getMergePolicy().getNoCFSRatio())) _assert_default_merge_scheduler(writer_config) writer_config.setMaxBufferedDocs(2) @@ -785,10 +897,86 @@ def test_real_hnsw_codec_selects_gpu_writer(tmp_path, pylucene_runtime_config): 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) + 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 ): @@ -803,7 +991,9 @@ def test_real_cagra_codec_merges_without_compound_files( 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 @@ -827,7 +1017,13 @@ def test_real_cagra_codec_merges_without_compound_files( expected_dimensions=vectors.shape[1], ) assert verification.segment_count == 1 - search = runtime.search_index(index_path, vectors[[137]], k=1, batch_size=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 @@ -853,7 +1049,9 @@ def test_real_hnsw_codec_falls_back_to_cpu_in_fresh_process( pylucene_runtime_config.writer_selection_classes ), "PYTHONPATH": os.pathsep.join( - value for value in (source_root, environment.get("PYTHONPATH")) if value + value + for value in (source_root, environment.get("PYTHONPATH")) + if value ), } ) @@ -888,6 +1086,7 @@ def test_real_verifier_rejects_cagra_brute_force_fallback( codec_name=_CAGRA_CODEC, java_codec=java_codec, writer_policy="gpu-cagra", + build_parameters=_CAGRA_BUILD_PARAMETERS, ), ) @@ -927,14 +1126,15 @@ def test_real_verifier_rejects_cagra_brute_force_fallback( def test_cli_build_and_search_with_real_pylucene_runtime( tmp_path, pylucene_runtime_config ): - """Exercise the documented module CLI in a fresh PyLucene process.""" + """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, + (query_vectors[:, np.newaxis, :] - training_vectors[np.newaxis, :, :]) + ** 2, axis=2, ) groundtruth_neighbors = np.argsort(squared_distances, axis=1)[:, :k] @@ -985,11 +1185,32 @@ def test_cli_build_and_search_with_real_pylucene_runtime( } ) ) + 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 + value + for value in (source_root, environment.get("PYTHONPATH")) + if value ) completed = subprocess.run( [ @@ -1000,6 +1221,8 @@ def test_cli_build_and_search_with_real_pylucene_runtime( str(backend_config), "--dataset-configuration", str(dataset_config), + "--configuration", + str(algorithm_config), "--dataset", dataset_name, "--dataset-path", @@ -1029,40 +1252,56 @@ def test_cli_build_and_search_with_real_pylucene_runtime( f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}" ) - codec = "Lucene101AcceleratedHNSWCodec" - index_name = f"pylucene_cuvs_hnsw[group=test][codec={codec}]" - index_path = dataset_dir / "index" / index_name - assert any(index_path.glob("segments_*")) - assert (index_path / _HNSW_PROVENANCE_FILE).is_file() + 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) == 1 - - build_row = build_rows[0] - assert build_row["index_name"] == index_name - 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" + 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) == 1 - csv_row = csv_rows[0] - assert csv_row["index_name"] == index_name - assert float(csv_row["recall"]) >= 0.75 - assert float(csv_row["throughput"]) > 0 - assert float(csv_row["latency"]) > 0 - assert float(csv_row["build time"]) > 0 - assert 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 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 index f50a8a8c97..8d29a54d1c 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_provenance.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_provenance.py @@ -28,6 +28,14 @@ ) +_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) @@ -41,21 +49,23 @@ def test_hnsw_provenance_round_trip(tmp_path): assert verification.to_metadata() == { "status": "gpu-with-cpu-fallback-provenance", - "schema_version": 3, + "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": ( - "1bc04b5291c26a46d918139138b992d2de976d6851d0893b" - "0476b85bfbdfc6e6" + "1bc04b5291c26a46d918139138b992d2de976d6851d0893b0476b85bfbdfc6e6" ), } ] @@ -75,17 +85,84 @@ def test_cagra_provenance_round_trip(tmp_path): assert verification.to_metadata() == { "status": "gpu-cagra-provenance", - "schema_version": 3, + "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"), [ diff --git a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_runtime.py b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_runtime.py index 1a5f6b57cb..ddd9daec7b 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_pylucene_runtime.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_pylucene_runtime.py @@ -13,12 +13,29 @@ 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 +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) -> _BuildCodec: +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(), @@ -27,6 +44,7 @@ def _build_codec(java_codec=None, codec_name=_HNSW_CODEC) -> _BuildCodec: if codec_name == _CAGRA_CODEC else "gpu-with-cpu-fallback" ), + build_parameters=build_parameters, ) @@ -62,6 +80,7 @@ def __init__(self, error_at=None): 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"}: @@ -85,6 +104,9 @@ def close(self): 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): @@ -96,10 +118,13 @@ def setNoCFSRatio(self, 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 @@ -113,8 +138,14 @@ def setUseCompoundFile(self, enabled): def getMergePolicy(self): return self.merge_policy - def setMergePolicy(self, _merge_policy): - raise AssertionError("PyLucene must use Lucene's default 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") @@ -151,6 +182,56 @@ def for_name(codec_name): 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) @@ -169,6 +250,7 @@ def close_directory(): 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 @@ -182,6 +264,10 @@ def create_writer(_directory, config): 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 @@ -189,6 +275,11 @@ def create_writer(_directory, config): 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( @@ -218,6 +309,7 @@ def test_initialize_pylucene_uses_verified_classpath_and_vmargs( 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, @@ -242,6 +334,28 @@ def test_initialize_pylucene_uses_verified_classpath_and_vmargs( 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()], @@ -311,6 +425,7 @@ def test_initialize_pylucene_uses_environment_runtime_config( init_call = fake_lucene.init_calls[0] assert init_call["classpath"].split(":") == [ + "/configured-codec-classes", str(cuvs_java), str(cuvs_lucene), fake_lucene.CLASSPATH, @@ -503,6 +618,46 @@ def test_resolve_codec_rejects_unusable_codec( 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 @@ -537,7 +692,10 @@ def test_runtime_build_index_commits_and_closes_writer_and_directory(tmp_path): _build_codec(runtime._test_codec), ) - assert result is None + 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 @@ -545,6 +703,7 @@ def test_runtime_build_index_commits_and_closes_writer_and_directory(tmp_path): 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( @@ -560,13 +719,172 @@ def test_runtime_writer_config_applies_codec_compound_file_policy( runtime = _fake_index_writer_runtime() config = runtime._new_index_writer_config( - _build_codec(runtime._test_codec, codec_name) + _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"), [ @@ -645,13 +963,48 @@ def test_runtime_build_index_preserves_interrupt_when_directory_close_fails( ) assert exc_info.value.__notes__ == [ - "Failed to close Lucene directory: " - "RuntimeError: directory close failed" + "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 = [] @@ -673,8 +1026,7 @@ def fail_directory_close(): assert close_calls == ["reader", "directory"] assert exc_info.value.__notes__ == [ - "Failed to close Lucene directory: " - "RuntimeError: directory close failed" + "Failed to close Lucene directory: RuntimeError: directory close failed" ] 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 index 7d0e0d4e1f..934238b73d 100644 --- a/python/cuvs_bench/tests/java/com/nvidia/cuvs/bench/PyLuceneWriterSelectionCodec.java +++ b/python/cuvs_bench/tests/java/com/nvidia/cuvs/bench/PyLuceneWriterSelectionCodec.java @@ -4,20 +4,47 @@ */ package com.nvidia.cuvs.bench; -import com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWCodec; 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 production HNSW codec. */ -public final class PyLuceneWriterSelectionCodec extends Lucene101AcceleratedHNSWCodec { +/** 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 { - super(); - setKnnFormat(new WriterSelectionFormat(super.knnVectorsFormat())); + 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 { From 06b9a8f9b794f1958db3f4bcffd0582748b476b5 Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Mon, 17 Aug 2026 15:44:22 +0000 Subject: [PATCH 12/12] Document configurable PyLucene HNSW benchmarks Describe the validated cuVS-Lucene source combination and adapter requirements. Document supported HNSW parameters, tuning ranges, single-segment constraints, and a manual sweep workflow. Signed-off-by: nvzm123 --- fern/pages/cuvs_bench/install.md | 19 +++++++---- fern/pages/cuvs_bench/running.md | 57 +++++++++++++++++++++++++++++--- 2 files changed, 65 insertions(+), 11 deletions(-) diff --git a/fern/pages/cuvs_bench/install.md b/fern/pages/cuvs_bench/install.md index 2434198e55..919d8b3031 100644 --- a/fern/pages/cuvs_bench/install.md +++ b/fern/pages/cuvs_bench/install.md @@ -54,17 +54,18 @@ Exact tags are listed on Docker Hub: 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 and a custom PyLucene wrapper generated against Lucene 10.2.0. 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. +- 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. Until those changes are ported and merged in cuVS, build the thin cuVS-Lucene JAR from that PR in a temporary checkout; a release or `main` checkout without the PR changes is insufficient. +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 .. ``` @@ -72,10 +73,14 @@ 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-pr174 -cd cuvs-lucene-pr174 +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 FETCH_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 ``` @@ -83,14 +88,14 @@ After the build, the conventional JAR paths are: ```text ~/.m2/repository/com/nvidia/cuvs/cuvs-java//cuvs-java-.jar -/target/cuvs-lucene-.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. Then validate the 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. +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 diff --git a/fern/pages/cuvs_bench/running.md b/fern/pages/cuvs_bench/running.md index 4188f0d290..e215a005e0 100644 --- a/fern/pages/cuvs_bench/running.md +++ b/fern/pages/cuvs_bench/running.md @@ -92,7 +92,7 @@ Create `pylucene-backend.yaml` with absolute paths to the base `cuvs-java` JAR, ```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-pr174/target/cuvs-lucene-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 ``` @@ -132,11 +132,60 @@ Do not use this tiny synthetic smoke run as performance or quality evidence; it 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 supplied configurations use Lucene's public `KnnFloatVectorQuery` and do not expose search-time tuning parameters. The backend accepts FLOAT32 Euclidean datasets with at most 4096 dimensions 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` to avoid cuVS-Lucene paths that can use brute-force search above that limit. +The HNSW configuration maps benchmark parameters to cuVS-Lucene as follows: -New HNSW and CAGRA builds atomically write commit-bound provenance manifests named `.cuvs-bench-pylucene-hnsw.json` and `.cuvs-bench-pylucene-cagra.json`, respectively. Reuse and search fail if the applicable manifest is missing, malformed, stale, or names a different codec, 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. +| 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. | -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. Both codecs retain Lucene's default index-writer scheduling behavior. +`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.