diff --git a/conda/recipes/cuvs-bench-cpu/recipe.yaml b/conda/recipes/cuvs-bench-cpu/recipe.yaml index 27ae190305..20838876db 100644 --- a/conda/recipes/cuvs-bench-cpu/recipe.yaml +++ b/conda/recipes/cuvs-bench-cpu/recipe.yaml @@ -130,9 +130,7 @@ tests: pip_check: false - script: content: | - python -c "from importlib.metadata import entry_points; assert any(item.name == 'lucene' for item in entry_points(group='cuvs_bench.backends'))" - python -c "from importlib.metadata import entry_points; assert any(item.name == 'lucene' for item in entry_points(group='cuvs_bench.config_loaders'))" - python -c "from cuvs_bench.orchestrator import BenchmarkOrchestrator; assert BenchmarkOrchestrator('lucene').backend_type == 'lucene'" + python -m cuvs_bench.tests.test_lucene_scaffold about: homepage: ${{ load_from_file("python/cuvs_bench/pyproject.toml").project.urls.Homepage }} diff --git a/python/cuvs_bench/tools/pylucene/build_pylucene_10_2.sh b/conda/recipes/cuvs-bench/build_pylucene_10_2.sh similarity index 100% rename from python/cuvs_bench/tools/pylucene/build_pylucene_10_2.sh rename to conda/recipes/cuvs-bench/build_pylucene_10_2.sh diff --git a/python/cuvs_bench/tools/pylucene/pylucene-10.2.0.patch b/conda/recipes/cuvs-bench/pylucene-10.2.0.patch similarity index 100% rename from python/cuvs_bench/tools/pylucene/pylucene-10.2.0.patch rename to conda/recipes/cuvs-bench/pylucene-10.2.0.patch diff --git a/conda/recipes/cuvs-bench/recipe.yaml b/conda/recipes/cuvs-bench/recipe.yaml index 199733caf2..4dc7bab8c6 100644 --- a/conda/recipes/cuvs-bench/recipe.yaml +++ b/conda/recipes/cuvs-bench/recipe.yaml @@ -80,9 +80,7 @@ tests: pip_check: false - script: content: | - python -c "from importlib.metadata import entry_points; assert any(item.name == 'lucene' for item in entry_points(group='cuvs_bench.backends'))" - python -c "from importlib.metadata import entry_points; assert any(item.name == 'lucene' for item in entry_points(group='cuvs_bench.config_loaders'))" - python -c "from cuvs_bench.orchestrator import BenchmarkOrchestrator; assert BenchmarkOrchestrator('lucene').backend_type == 'lucene'" + python -m cuvs_bench.tests.test_lucene_scaffold about: homepage: ${{ load_from_file("python/cuvs_bench/pyproject.toml").project.urls.Homepage }} diff --git a/fern/pages/cuvs_bench/lucene_backend.md b/fern/pages/cuvs_bench/lucene_backend.md index 3be47ec9ac..e18a405795 100644 --- a/fern/pages/cuvs_bench/lucene_backend.md +++ b/fern/pages/cuvs_bench/lucene_backend.md @@ -107,7 +107,7 @@ least 3 GB of disk space. ```bash export PYLUCENE_BUILD_ROOT="$HOME/.local/share/cuvs/pylucene-10.2.0" -python/cuvs_bench/tools/pylucene/build_pylucene_10_2.sh \ +conda/recipes/cuvs-bench/build_pylucene_10_2.sh \ --python python3 \ --build-root "$PYLUCENE_BUILD_ROOT" diff --git a/python/cuvs_bench/cuvs_bench/backends/_lucene_runtime.py b/python/cuvs_bench/cuvs_bench/backends/_lucene_runtime.py index 65bd18ec9a..ba933e4bcd 100644 --- a/python/cuvs_bench/cuvs_bench/backends/_lucene_runtime.py +++ b/python/cuvs_bench/cuvs_bench/backends/_lucene_runtime.py @@ -26,14 +26,6 @@ CAGRA_CODEC = "CuVS2510GPUSearchCodec" MAX_CAGRA_TOP_K = 1024 REQUIRED_PYLUCENE_VERSION = "10.2.0" -_PYLUCENE_SETUP_GUIDANCE = ( - "Setup documentation: fern/pages/cuvs_bench/lucene_backend.md " - "(published at https://docs.nvidia.com/cuvs/user-guide/benchmarking-guide/" - "cu-vs-bench-tool/lucene-backend). Source-checkout helper: " - "python/cuvs_bench/tools/pylucene/build_pylucene_10_2.sh " - "(run it with --build-root , then source " - "/activate.sh)." -) _ID_FIELD = "id" _VECTOR_FIELD = "vector" @@ -45,8 +37,22 @@ _CAGRA_FORMAT_VERSION = 0 _FLOAT32_ENCODING_ORDINAL = 1 _EUCLIDEAN_SIMILARITY_ORDINAL = 0 +_CUVS_JAVA_PACKAGE = "com.nvidia.cuvs" +_CUVS_LUCENE_PACKAGE = f"{_CUVS_JAVA_PACKAGE}.lucene" +_CUVS_JAVA_PACKAGE_DIRECTORY = _CUVS_JAVA_PACKAGE.replace(".", "/") +_CUVS_LUCENE_PACKAGE_DIRECTORY = _CUVS_LUCENE_PACKAGE.replace(".", "/") +_MAVEN_METADATA_DIRECTORY = "META-INF/maven" +_CUVS_JAVA_MAVEN_PROPERTIES = ( + f"{_MAVEN_METADATA_DIRECTORY}/{_CUVS_JAVA_PACKAGE}/" + "cuvs-java/pom.properties" +) +_CUVS_LUCENE_MAVEN_PROPERTIES = ( + f"{_MAVEN_METADATA_DIRECTORY}/{_CUVS_LUCENE_PACKAGE}/" + "cuvs-lucene/pom.properties" +) +_CODEC_SERVICE_DESCRIPTOR = "META-INF/services/org.apache.lucene.codecs.Codec" _INDEX_SEARCHER_TIMING_BRIDGE = ( - "com.nvidia.cuvs.lucene.IndexSearcherTimingBridge" + f"{_CUVS_LUCENE_PACKAGE}.IndexSearcherTimingBridge" ) _SEARCHER_REQUEST_KEY = "searcher" _QUERY_REQUEST_KEY = "query" @@ -124,9 +130,9 @@ def _rollback_writer(writer: Any, error: BaseException) -> None: def _read_jar(path: Path, label: str) -> tuple[set[str], dict[str, bytes]]: inspected = { "META-INF/MANIFEST.MF", - "META-INF/maven/com.nvidia.cuvs/cuvs-java/pom.properties", - "META-INF/maven/com.nvidia.cuvs.lucene/cuvs-lucene/pom.properties", - "META-INF/services/org.apache.lucene.codecs.Codec", + _CUVS_JAVA_MAVEN_PROPERTIES, + _CUVS_LUCENE_MAVEN_PROPERTIES, + _CODEC_SERVICE_DESCRIPTOR, } try: with zipfile.ZipFile(path) as archive: @@ -247,9 +253,10 @@ def _validate_artifacts( } java_entries, java_contents = _read_jar(java_jar, "cuvs_java_jar") java_required = { - "com/nvidia/cuvs/CagraIndex.class", - "com/nvidia/cuvs/CuVSResources.class", - "META-INF/versions/22/com/nvidia/cuvs/spi/JDKProvider.class", + f"{_CUVS_JAVA_PACKAGE_DIRECTORY}/CagraIndex.class", + f"{_CUVS_JAVA_PACKAGE_DIRECTORY}/CuVSResources.class", + "META-INF/versions/22/" + f"{_CUVS_JAVA_PACKAGE_DIRECTORY}/spi/JDKProvider.class", } missing = sorted(java_required - java_entries) if missing: @@ -273,11 +280,12 @@ def _validate_artifacts( lucene_entries, lucene_contents = _read_jar(lucene_jar, "cuvs_lucene_jar") lucene_required = { - "com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.class", - "com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.class", - "com/nvidia/cuvs/lucene/IndexSearcherTimingBridge.class", - "com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.class", - "META-INF/services/org.apache.lucene.codecs.Codec", + f"{_CUVS_LUCENE_PACKAGE_DIRECTORY}/CuVS2510GPUVectorsFormat.class", + f"{_CUVS_LUCENE_PACKAGE_DIRECTORY}/CuVS2510GPUSearchCodec.class", + f"{_CUVS_LUCENE_PACKAGE_DIRECTORY}/IndexSearcherTimingBridge.class", + f"{_CUVS_LUCENE_PACKAGE_DIRECTORY}/" + "Lucene101AcceleratedHNSWCodec.class", + _CODEC_SERVICE_DESCRIPTOR, } missing = sorted(lucene_required - lucene_entries) if missing: @@ -304,36 +312,37 @@ def _validate_artifacts( ) providers = { line.partition("#")[0].strip() - for line in lucene_contents[ - "META-INF/services/org.apache.lucene.codecs.Codec" - ] + for line in lucene_contents[_CODEC_SERVICE_DESCRIPTOR] .decode("utf-8") .splitlines() if line.partition("#")[0].strip() } - if "com.nvidia.cuvs.lucene.CuVS2510GPUSearchCodec" not in providers: + if f"{_CUVS_LUCENE_PACKAGE}.CuVS2510GPUSearchCodec" not in providers: raise RuntimeError( "cuvs_lucene_jar does not advertise CuVS2510GPUSearchCodec" ) - if "com.nvidia.cuvs.lucene.Lucene101AcceleratedHNSWCodec" not in providers: + if ( + f"{_CUVS_LUCENE_PACKAGE}.Lucene101AcceleratedHNSWCodec" + not in providers + ): raise RuntimeError( "cuvs_lucene_jar does not advertise Lucene101AcceleratedHNSWCodec" ) java_coordinates = _maven_coordinates( java_contents, - "META-INF/maven/com.nvidia.cuvs/cuvs-java/pom.properties", + _CUVS_JAVA_MAVEN_PROPERTIES, ) lucene_coordinates = _maven_coordinates( lucene_contents, - "META-INF/maven/com.nvidia.cuvs.lucene/cuvs-lucene/pom.properties", + _CUVS_LUCENE_MAVEN_PROPERTIES, ) - if java_coordinates[:2] != ("com.nvidia.cuvs", "cuvs-java"): + if java_coordinates[:2] != (_CUVS_JAVA_PACKAGE, "cuvs-java"): raise RuntimeError( f"Unexpected cuvs-java coordinates: {java_coordinates[:2]}" ) if lucene_coordinates[:2] != ( - "com.nvidia.cuvs.lucene", + _CUVS_LUCENE_PACKAGE, "cuvs-lucene", ): raise RuntimeError( @@ -375,8 +384,9 @@ def _load_pylucene() -> Any: except ImportError as error: raise ImportError( "The Lucene backend requires the custom PyLucene 10.2.0 runtime, " - "which is not included in cuVS Bench packages. " - f"{_PYLUCENE_SETUP_GUIDANCE} PyLucene import failed: {error}" + "which is not included in cuVS Bench packages. Install and " + "activate the optional Lucene runtime before selecting this " + f"backend. PyLucene import failed: {error}" ) from error @@ -432,8 +442,9 @@ def initialize_pylucene( if actual_version != REQUIRED_PYLUCENE_VERSION: raise RuntimeError( "PyLucene must match cuvs-lucene's Lucene version: expected " - f"{REQUIRED_PYLUCENE_VERSION}, found {actual_version}. " - f"{_PYLUCENE_SETUP_GUIDANCE}" + f"{REQUIRED_PYLUCENE_VERSION}, found {actual_version}. Install " + "and activate a compatible PyLucene runtime before selecting " + "this backend." ) vmargs = _vmargs(config) with _JVM_LOCK: diff --git a/python/cuvs_bench/cuvs_bench/backends/_lucene_runtime_config.py b/python/cuvs_bench/cuvs_bench/backends/_lucene_runtime_config.py index 3c2d0770a9..e89cf8ee7d 100644 --- a/python/cuvs_bench/cuvs_bench/backends/_lucene_runtime_config.py +++ b/python/cuvs_bench/cuvs_bench/backends/_lucene_runtime_config.py @@ -20,6 +20,8 @@ _PACKAGE_ROOT = Path(__file__).resolve().parents[1] _REPOSITORY_ROOT = Path(__file__).resolve().parents[4] _VERSION_FILE = _PACKAGE_ROOT / "VERSION" +_CUVS_MAVEN_DIRECTORY = Path("com/nvidia/cuvs") +_CUVS_LUCENE_MAVEN_DIRECTORY = _CUVS_MAVEN_DIRECTORY / "lucene" _CUDA_TARGET_BY_MACHINE = { "aarch64": "sbsa-linux", @@ -57,9 +59,7 @@ def _artifact_candidates(kind: str) -> tuple[Path, ...]: / "target" / f"cuvs-java-{version}.jar", repository - / "com" - / "nvidia" - / "cuvs" + / _CUVS_MAVEN_DIRECTORY / "cuvs-java" / version / f"cuvs-java-{version}.jar", @@ -72,10 +72,7 @@ def _artifact_candidates(kind: str) -> tuple[Path, ...]: / "target" / f"cuvs-lucene-{version}.jar", repository - / "com" - / "nvidia" - / "cuvs" - / "lucene" + / _CUVS_LUCENE_MAVEN_DIRECTORY / "cuvs-lucene" / version / f"cuvs-lucene-{version}.jar", diff --git a/python/cuvs_bench/cuvs_bench/backends/base.py b/python/cuvs_bench/cuvs_bench/backends/base.py index bc04408ce9..bfdef6b367 100644 --- a/python/cuvs_bench/cuvs_bench/backends/base.py +++ b/python/cuvs_bench/cuvs_bench/backends/base.py @@ -11,7 +11,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass, field -from typing import Dict, List, Any, Optional, TYPE_CHECKING +from typing import Dict, List, Any, Optional, Sequence, TYPE_CHECKING import numpy as np if TYPE_CHECKING: @@ -368,6 +368,10 @@ class BenchmarkBackend(ABC): - requires_network : bool - Whether backend requires network (default: False) """ + # Registration exposes this optional class-level default to the generic + # CLI. + default_algorithm: Optional[str] = None + def __init__(self, config: Dict[str, Any]): """Initialize backend with configuration.""" self.config = config @@ -495,6 +499,13 @@ def cleanup(self) -> None: """ pass + @classmethod + def result_failure_message( + cls, results: Sequence[BuildResult | SearchResult] + ) -> Optional[str]: + """Return a fatal result summary when the backend requires one.""" + return None + def _check_gpu_available(self) -> bool: """ Check if GPU is available (same logic as run.py rmm_present()). diff --git a/python/cuvs_bench/cuvs_bench/backends/cpp_gbench.py b/python/cuvs_bench/cuvs_bench/backends/cpp_gbench.py index 1e2d861456..78004c374b 100644 --- a/python/cuvs_bench/cuvs_bench/backends/cpp_gbench.py +++ b/python/cuvs_bench/cuvs_bench/backends/cpp_gbench.py @@ -51,6 +51,8 @@ class CppGoogleBenchmarkBackend(BenchmarkBackend): >>> print(backend.algo) # "cuvs_ivf_flat" (from config) """ + default_algorithm = "cuvs_cagra" + def __init__(self, config: Dict[str, Any]): """Initialize C++ benchmark backend.""" super().__init__(config) diff --git a/python/cuvs_bench/cuvs_bench/backends/lucene.py b/python/cuvs_bench/cuvs_bench/backends/lucene.py index 0dd7feac27..af8fd2ad2a 100644 --- a/python/cuvs_bench/cuvs_bench/backends/lucene.py +++ b/python/cuvs_bench/cuvs_bench/backends/lucene.py @@ -19,7 +19,7 @@ from dataclasses import dataclass from numbers import Integral from pathlib import Path -from typing import Any, Callable, Mapping, Optional +from typing import Any, Callable, Mapping, Optional, Sequence import numpy as np @@ -855,6 +855,8 @@ def _build_benchmark_configs( class LuceneBackend(BenchmarkBackend): """Build and query one immutable Lucene vector index per configuration.""" + default_algorithm = CAGRA_ALGORITHM + def __init__( self, config: dict[str, Any], @@ -874,6 +876,19 @@ def __init__( self._runtime_factory = runtime_factory self._runtime: LuceneRuntime | None = None + @classmethod + def result_failure_message( + cls, results: Sequence[BuildResult | SearchResult] + ) -> str | None: + """Fail the public command when any Lucene operation failed.""" + failures = [result for result in results if not result.success] + if not failures: + return None + return "; ".join( + result.error_message or "unknown Lucene backend failure" + for result in failures + ) + @property def algo(self) -> str: """Return the selected cuVS Bench algorithm name.""" diff --git a/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py b/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py index abffad0217..ef7f9a4c68 100644 --- a/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py +++ b/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py @@ -87,6 +87,12 @@ def __init__(self, backend_type: str = "cpp_gbench"): # Instantiate config loader self.config_loader = loader_class() + def result_failure_message( + self, results: List[Union[BuildResult, SearchResult]] + ) -> Optional[str]: + """Return the selected backend's fatal result summary, if any.""" + return self.backend_class.result_failure_message(results) + def run_benchmark( self, mode: str = "sweep", diff --git a/python/cuvs_bench/cuvs_bench/run/__main__.py b/python/cuvs_bench/cuvs_bench/run/__main__.py index 0c65a8f91e..7956562f5a 100644 --- a/python/cuvs_bench/cuvs_bench/run/__main__.py +++ b/python/cuvs_bench/cuvs_bench/run/__main__.py @@ -12,17 +12,16 @@ import yaml from click.core import ParameterSource +from ..backends.registry import get_backend_class +from ..orchestrator import BenchmarkOrchestrator from .data_export import ( convert_json_to_csv_build, convert_json_to_csv_search, write_results_to_csv, ) -from ..orchestrator import BenchmarkOrchestrator _DEFAULT_BACKEND = "cpp_gbench" -_DEFAULT_ALGORITHM = "cuvs_cagra" -_LUCENE_DEFAULT_ALGORITHM = "lucene_cuvs_cagra" def _read_backend_config(path: str) -> dict[str, Any]: @@ -41,16 +40,20 @@ def _read_backend_config(path: str) -> dict[str, Any]: def _default_algorithm() -> str: """Choose the displayed prompt default from the selected backend.""" context = click.get_current_context(silent=True) - if context is None: - return _DEFAULT_ALGORITHM - backend = context.params.get("backend", _DEFAULT_BACKEND) - if backend_config := context.params.get("backend_config"): - backend = _read_backend_config(backend_config)["backend"] - return ( - _LUCENE_DEFAULT_ALGORITHM - if backend == "lucene" - else _DEFAULT_ALGORITHM - ) + backend = _DEFAULT_BACKEND + if context is not None: + backend = context.params.get("backend", _DEFAULT_BACKEND) + if backend_config := context.params.get("backend_config"): + backend = _read_backend_config(backend_config)["backend"] + backend_class = get_backend_class(str(backend)) + if backend_class.default_algorithm is not None: + return backend_class.default_algorithm + default_backend_class = get_backend_class(_DEFAULT_BACKEND) + if default_backend_class.default_algorithm is None: + raise RuntimeError( + "The default backend must define a default algorithm" + ) + return default_backend_class.default_algorithm @click.command() @@ -127,8 +130,8 @@ def _default_algorithm() -> str: prompt="Enter the comma separated list of named algorithms to run", help="Run only comma separated list of named algorithms. If parameters " "`groups` and `algo-groups` are both undefined, then group `base` " - "is run by default. The prompt defaults to `lucene_cuvs_cagra` for the " - "Lucene backend and `cuvs_cagra` otherwise.", + "is run by default. A backend may provide its own prompt default; " + "otherwise the CLI default is retained.", ) @click.option( "--groups", @@ -212,15 +215,14 @@ def _default_algorithm() -> str: "--backend", default=_DEFAULT_BACKEND, show_default=True, - help="Backend type to run. The default preserves the C++ benchmark " - "workflow; select 'lucene' to opt in to the Lucene backend.", + help="Backend type to run.", ) @click.option( "--backend-config", default=None, help="Path to YAML configuration file for non-C++ backends. " "The file must contain a 'backend' field specifying the backend " - "type (e.g., 'lucene', 'opensearch', 'elastic'). All other fields are " + "type (e.g., 'opensearch', 'elastic'). All other fields are " "passed as backend-specific parameters. If --backend is also provided, " "the values must match.", ) @@ -298,7 +300,7 @@ def main( Backend type to run. Defaults to the C++ Google Benchmark backend. backend_config : Optional[str] Path to YAML config for non-C++ backends. The YAML file must contain - a 'backend' field (e.g., 'lucene', 'opensearch', 'elastic') and any + a 'backend' field (e.g., 'opensearch', 'elastic') and any backend-specific connection parameters (host, port, etc.). If ``--backend`` is also provided, the values must match. @@ -371,14 +373,8 @@ def main( convert_json_to_csv_search(dataset, dataset_path) else: write_results_to_csv(results, dataset, dataset_path, count, batch_size) - if backend_type == "lucene": - failures = [result for result in results if not result.success] - if failures: - details = "; ".join( - result.error_message or "unknown Lucene backend failure" - for result in failures - ) - raise click.ClickException(details) + if failure_message := orchestrator.result_failure_message(results): + raise click.ClickException(failure_message) if __name__ == "__main__": diff --git a/python/cuvs_bench/cuvs_bench/tests/test_lucene_runtime.py b/python/cuvs_bench/cuvs_bench/tests/test_lucene_runtime.py index a31818ff7e..50287522a9 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_lucene_runtime.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_lucene_runtime.py @@ -370,13 +370,10 @@ def missing_module(_name: str): message = str(error.value) assert "requires the custom PyLucene 10.2.0 runtime" in message - assert "fern/pages/cuvs_bench/lucene_backend.md" in message - assert "python/cuvs_bench/tools/pylucene/build_pylucene_10_2.sh" in message - assert "/activate.sh" in message - assert ( - "https://docs.nvidia.com/cuvs/user-guide/benchmarking-guide/" - "cu-vs-bench-tool/lucene-backend" - ) in message + assert "Install and activate the optional Lucene runtime" in message + assert "http" not in message + assert ".md" not in message + assert ".sh" not in message assert "PyLucene import failed: no lucene module" in message @@ -393,13 +390,10 @@ def test_incompatible_pylucene_version_fails_before_vm_initialization( message = str(error.value) assert "expected 10.2.0, found 10.1.0" in message - assert "fern/pages/cuvs_bench/lucene_backend.md" in message - assert "python/cuvs_bench/tools/pylucene/build_pylucene_10_2.sh" in message - assert "/activate.sh" in message - assert ( - "https://docs.nvidia.com/cuvs/user-guide/benchmarking-guide/" - "cu-vs-bench-tool/lucene-backend" - ) in message + assert "Install and activate a compatible PyLucene runtime" in message + assert "http" not in message + assert ".md" not in message + assert ".sh" not in message assert fake_lucene.initializations == [] diff --git a/python/cuvs_bench/cuvs_bench/tests/test_lucene_scaffold.py b/python/cuvs_bench/cuvs_bench/tests/test_lucene_scaffold.py index 5b95ef3af0..18098c088c 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_lucene_scaffold.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_lucene_scaffold.py @@ -7,19 +7,49 @@ import subprocess import tomllib from importlib import resources +from importlib.metadata import entry_points from pathlib import Path from unittest.mock import patch import yaml from click.testing import CliRunner +from cuvs_bench.backends.base import BuildResult +from cuvs_bench.backends.lucene import ( + CAGRA_ALGORITHM, + LuceneBackend, + register as register_lucene, +) +from cuvs_bench.orchestrator import BenchmarkOrchestrator from cuvs_bench.run.__main__ import main as run_main _PROJECT_ROOT = Path(__file__).parents[2] +_REPOSITORY_ROOT = Path(__file__).parents[4] +_LUCENE_ENTRY_POINT = "cuvs_bench.backends.lucene:register" +_PYLUCENE_RECIPE_DIRECTORY = ( + _REPOSITORY_ROOT / "conda" / "recipes" / "cuvs-bench" +) +_PYLUCENE_BUILD_HELPER = ( + _PYLUCENE_RECIPE_DIRECTORY / "build_pylucene_10_2.sh" +) +_PYLUCENE_PATCH = _PYLUCENE_RECIPE_DIRECTORY / "pylucene-10.2.0.patch" + + +def _verify_installed_lucene_plugin() -> None: + """Exercise the installed entry points without importing PyLucene.""" + for group in ("cuvs_bench.backends", "cuvs_bench.config_loaders"): + assert any( + item.name == "lucene" and item.value == _LUCENE_ENTRY_POINT + for item in entry_points(group=group) + ), f"installed cuvs-bench is missing the {group!r} Lucene entry point" + + orchestrator = BenchmarkOrchestrator("lucene") + assert orchestrator.backend_type == "lucene" def _invoke_run(tmp_path, *extra_args, input_text=""): + register_lucene() captured = {} class RecordingOrchestrator: @@ -72,6 +102,51 @@ def test_lucene_backend_selects_cagra_default(tmp_path): assert captured["run_kwargs"]["algorithms"] == "lucene_cuvs_cagra" +def test_registered_backend_class_controls_the_prompt_default(tmp_path): + class PluginBackend: + default_algorithm = "plugin_default" + + with patch( + "cuvs_bench.run.__main__.get_backend_class", + return_value=PluginBackend, + ): + result, captured = _invoke_run( + tmp_path, + "--backend", + "third_party", + input_text="\n", + ) + + assert result.exit_code == 0, result.output + assert captured["backend_type"] == "third_party" + assert captured["run_kwargs"]["algorithms"] == "plugin_default" + + +def test_plugin_without_a_default_keeps_the_cli_default(tmp_path): + class PluginBackend: + default_algorithm = None + + class DefaultBackend: + default_algorithm = "cuvs_cagra" + + def backend_class(name): + return PluginBackend if name == "third_party" else DefaultBackend + + with patch( + "cuvs_bench.run.__main__.get_backend_class", + side_effect=backend_class, + ): + result, captured = _invoke_run( + tmp_path, + "--backend", + "third_party", + input_text="\n", + ) + + assert result.exit_code == 0, result.output + assert captured["run_kwargs"]["algorithms"] == "cuvs_cagra" + + def test_lucene_backend_preserves_explicit_cpu_algorithm(tmp_path): result, captured = _invoke_run( tmp_path, @@ -167,9 +242,83 @@ def test_lucene_plugin_uses_lazy_entry_points(): pyproject = tomllib.loads((_PROJECT_ROOT / "pyproject.toml").read_text()) entry_points = pyproject["project"]["entry-points"] - target = "cuvs_bench.backends.lucene:register" - assert entry_points["cuvs_bench.backends"]["lucene"] == target - assert entry_points["cuvs_bench.config_loaders"]["lucene"] == target + assert ( + entry_points["cuvs_bench.backends"]["lucene"] + == _LUCENE_ENTRY_POINT + ) + assert ( + entry_points["cuvs_bench.config_loaders"]["lucene"] + == _LUCENE_ENTRY_POINT + ) + + +def test_lucene_backend_makes_failed_results_fatal() -> None: + successful = BuildResult( + index_path="", + build_time_seconds=0.0, + index_size_bytes=0, + algorithm=CAGRA_ALGORITHM, + build_params={}, + ) + failed = BuildResult( + index_path="", + build_time_seconds=0.0, + index_size_bytes=0, + algorithm=CAGRA_ALGORITHM, + build_params={}, + success=False, + error_message="GPU path was unavailable", + ) + + assert LuceneBackend.result_failure_message([successful]) is None + assert LuceneBackend.result_failure_message([successful, failed]) == ( + "GPU path was unavailable" + ) + + +def test_lucene_cli_fails_after_recording_failed_results(tmp_path: Path): + failed = BuildResult( + index_path="", + build_time_seconds=0.0, + index_size_bytes=0, + algorithm=CAGRA_ALGORITHM, + build_params={}, + success=False, + error_message="GPU path was unavailable", + ) + arguments = [ + "--dataset", + "test-data", + "--dataset-path", + str(tmp_path), + "--batch-size", + "10", + "-k", + "10", + "--groups", + "test", + "-m", + "latency", + "--backend", + "lucene", + "--algorithms", + CAGRA_ALGORITHM, + ] + register_lucene() + + with ( + patch.object( + BenchmarkOrchestrator, + "run_benchmark", + return_value=[failed], + ), + patch("cuvs_bench.run.__main__.write_results_to_csv") as write_results, + ): + result = CliRunner().invoke(run_main, arguments) + + assert result.exit_code == 1 + assert "GPU path was unavailable" in result.output + write_results.assert_called_once() def test_lucene_algorithm_configs_are_packaged_resources(): @@ -194,19 +343,15 @@ def test_lucene_algorithm_configs_are_packaged_resources(): def test_pylucene_source_builder_is_self_describing() -> None: - tool_directory = _PROJECT_ROOT / "tools" / "pylucene" - helper = tool_directory / "build_pylucene_10_2.sh" - compatibility_patch = tool_directory / "pylucene-10.2.0.patch" - completed = subprocess.run( - ["bash", str(helper), "--help"], + ["bash", str(_PYLUCENE_BUILD_HELPER), "--help"], check=False, capture_output=True, text=True, ) - assert os.access(helper, os.X_OK) - assert compatibility_patch.is_file() + assert os.access(_PYLUCENE_BUILD_HELPER, os.X_OK) + assert _PYLUCENE_PATCH.is_file() assert completed.returncode == 0, completed.stderr assert "Build an isolated PyLucene 10.2.0 environment" in completed.stdout @@ -214,14 +359,12 @@ def test_pylucene_source_builder_is_self_describing() -> None: def test_pylucene_source_builder_rejects_unsafe_build_roots( tmp_path: Path, ) -> None: - helper = _PROJECT_ROOT / "tools" / "pylucene" / "build_pylucene_10_2.sh" - for unsafe_character in (" ", ":", ";", "$", "`", "&", "#", "|"): build_root = tmp_path / f"unsafe{unsafe_character}root" completed = subprocess.run( [ "bash", - str(helper), + str(_PYLUCENE_BUILD_HELPER), "--build-root", str(build_root), "--prepare-only", @@ -239,7 +382,6 @@ def test_pylucene_source_builder_rejects_unsafe_build_roots( def test_pylucene_source_builder_rejects_unsafe_resolved_build_root( tmp_path: Path, ) -> None: - helper = _PROJECT_ROOT / "tools" / "pylucene" / "build_pylucene_10_2.sh" unsafe_target = tmp_path / "unsafe;target" unsafe_target.mkdir() build_root = tmp_path / "safe-link" @@ -248,7 +390,7 @@ def test_pylucene_source_builder_rejects_unsafe_resolved_build_root( completed = subprocess.run( [ "bash", - str(helper), + str(_PYLUCENE_BUILD_HELPER), "--build-root", str(build_root), "--prepare-only", @@ -261,3 +403,7 @@ def test_pylucene_source_builder_rejects_unsafe_resolved_build_root( assert completed.returncode != 0 assert "--build-root may contain only" in completed.stderr assert not (unsafe_target / ".build.lock").exists() + + +if __name__ == "__main__": + _verify_installed_lucene_plugin()