From 336432dc111cb7cba063f7cfab397465764ab2c0 Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Tue, 10 Mar 2026 14:52:46 -0700 Subject: [PATCH 01/22] Add cuvs-bench-elastic: HTTP backend for Elasticsearch GPU vector search Introduce cuvs-bench-elastic as an optional plugin for cuvs-bench that provides an Elasticsearch backend. The backend communicates with Elasticsearch via HTTP and supports HNSW indexing with optional GPU acceleration when using the Elasticsearch GPU image (cuVS-accelerated vector search). - Add cuvs_bench_elastic package with backend and config loader entry points - Extend cuvs_bench registry and search spaces for pluggable backends - Add elastic and integration optional dependencies to cuvs-bench - Add modularization tests and integration test scaffolding (disabled until CI has ES GPU image, cuVS libs, and GPU runner) Signed-off-by: Alex Fournier --- dependencies.yaml | 13 + .../cuvs_bench/backends/registry.py | 69 +- .../cuvs_bench/backends/search_spaces.py | 40 ++ .../cuvs_bench/cuvs_bench/tests/conftest.py | 23 + .../cuvs_bench/tests/integration/README.md | 34 + .../cuvs_bench/tests/integration/__init__.py | 9 + .../cuvs_bench/tests/integration/conftest.py | 73 ++ .../integration/test_elastic_integration.py | 141 ++++ .../cuvs_bench/tests/test_modularization.py | 591 ++++++++++++++++ .../cuvs_bench/tests/test_registry.py | 181 ++--- python/cuvs_bench/pyproject.toml | 4 + python/cuvs_bench_elastic/README.md | 42 ++ .../cuvs_bench_elastic/VERSION | 1 + .../cuvs_bench_elastic/__init__.py | 13 + .../cuvs_bench_elastic/backend.py | 656 ++++++++++++++++++ .../config/algos/elastic.yaml | 53 ++ python/cuvs_bench_elastic/pyproject.toml | 50 ++ 17 files changed, 1871 insertions(+), 122 deletions(-) create mode 100644 python/cuvs_bench/cuvs_bench/tests/conftest.py create mode 100644 python/cuvs_bench/cuvs_bench/tests/integration/README.md create mode 100644 python/cuvs_bench/cuvs_bench/tests/integration/__init__.py create mode 100644 python/cuvs_bench/cuvs_bench/tests/integration/conftest.py create mode 100644 python/cuvs_bench/cuvs_bench/tests/integration/test_elastic_integration.py create mode 100644 python/cuvs_bench/cuvs_bench/tests/test_modularization.py create mode 100644 python/cuvs_bench_elastic/README.md create mode 100644 python/cuvs_bench_elastic/cuvs_bench_elastic/VERSION create mode 100644 python/cuvs_bench_elastic/cuvs_bench_elastic/__init__.py create mode 100644 python/cuvs_bench_elastic/cuvs_bench_elastic/backend.py create mode 100644 python/cuvs_bench_elastic/cuvs_bench_elastic/config/algos/elastic.yaml create mode 100644 python/cuvs_bench_elastic/pyproject.toml diff --git a/dependencies.yaml b/dependencies.yaml index e855a93f3e..8c0a7f66ca 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -213,6 +213,14 @@ files: table: project includes: - bench_python + py_elastic_cuvs_bench: + output: pyproject + pyproject_dir: python/cuvs_bench + extras: + table: project.optional-dependencies + key: elastic + includes: + - bench_elastic channels: - rapidsai-nightly - rapidsai @@ -500,6 +508,11 @@ dependencies: - output_types: [requirements, pyproject] packages: - matplotlib>=3.9 + bench_elastic: + common: + - output_types: [conda, pyproject, requirements] + packages: + - cuvs-bench-elastic>=26.4.0 depends_on_cuda_python: specific: - output_types: [conda, requirements, pyproject] diff --git a/python/cuvs_bench/cuvs_bench/backends/registry.py b/python/cuvs_bench/cuvs_bench/backends/registry.py index 1a4dd88dde..0059ea78f0 100644 --- a/python/cuvs_bench/cuvs_bench/backends/registry.py +++ b/python/cuvs_bench/cuvs_bench/backends/registry.py @@ -12,10 +12,15 @@ from typing import Dict, Type, Optional from pathlib import Path import importlib +import importlib.metadata import yaml from .base import BenchmarkBackend +# Entry point group names for plugin discovery +_BACKENDS_GROUP = "cuvs_bench.backends" +_CONFIG_LOADERS_GROUP = "cuvs_bench.config_loaders" + class BackendRegistry: """ @@ -375,10 +380,43 @@ def get_backend(name: str, config: Dict) -> BenchmarkBackend: return registry.get_backend(name, config) +def _try_load_plugin(name: str) -> None: + """ + Try to load backend and config loader from entry points for the given name. + + Plugins register themselves when their entry point is loaded. + Raises ImportError with install instructions if the plugin requires + an optional dependency that is not installed. + """ + for group in (_BACKENDS_GROUP, _CONFIG_LOADERS_GROUP): + try: + eps = importlib.metadata.entry_points(group=group) + except TypeError: + eps = importlib.metadata.entry_points().get(group, []) + if hasattr(eps, "select"): # Python 3.10+ + eps = list(eps.select(name=name)) + else: + eps = [e for e in eps if e.name == name] + for ep in eps: + try: + ep.load()() + except ImportError as e: + if "elasticsearch" in str(e).lower() or "elasticsearch" in str(e): + raise ImportError( + f"Elasticsearch backend requires the 'elastic' extra. " + f"Install with: pip install cuvs-bench[elastic]" + ) from e + raise + return # Plugin loaded successfully + + def get_backend_class(name: str) -> Type[BenchmarkBackend]: """ Get the backend class (not instance) from the global registry. + If the backend is not registered, attempts to load it from entry points + (e.g., optional plugins like elastic). + Parameters ---------- name : str @@ -390,10 +428,15 @@ def get_backend_class(name: str) -> Type[BenchmarkBackend]: Backend class """ registry = get_registry() + if name not in registry._backends: + _try_load_plugin(name) if name not in registry._backends: available = ", ".join(registry._backends.keys()) + hint = "" + if name == "elastic": + hint = " Install with: pip install cuvs-bench[elastic]" raise ValueError( - f"Backend '{name}' not found. Available backends: {available or '(none)'}" + f"Backend '{name}' not found. Available backends: {available or '(none)'}.{hint}" ) return registry._backends[name] @@ -440,6 +483,9 @@ def get_config_loader(name: str) -> Type: """ Get a registered config loader class by name. + If the config loader is not registered, attempts to load it from entry points + (e.g., optional plugins like elastic). + Parameters ---------- name : str @@ -455,11 +501,15 @@ def get_config_loader(name: str) -> Type: ValueError If config loader is not registered """ - # _CONFIG_LOADER_REGISTRY is a dictionary that maps backend names to config loader classes + if name not in _CONFIG_LOADER_REGISTRY: + _try_load_plugin(name) if name not in _CONFIG_LOADER_REGISTRY: available = ", ".join(_CONFIG_LOADER_REGISTRY.keys()) or "none" + hint = "" + if name == "elastic": + hint = " Install with: pip install cuvs-bench[elastic]" raise ValueError( - f"Unknown config loader for backend: '{name}'. Available: {available}" + f"Unknown config loader for backend: '{name}'. Available: {available}.{hint}" ) return _CONFIG_LOADER_REGISTRY[name] @@ -467,3 +517,16 @@ def get_config_loader(name: str) -> Type: def list_config_loaders() -> Dict[str, Type]: """Return all registered config loaders.""" return dict(_CONFIG_LOADER_REGISTRY) + + +def unregister_config_loader(name: str) -> None: + """ + Unregister a config loader by name (primarily for testing). + + Parameters + ---------- + name : str + Backend name to unregister + """ + if name in _CONFIG_LOADER_REGISTRY: + del _CONFIG_LOADER_REGISTRY[name] diff --git a/python/cuvs_bench/cuvs_bench/backends/search_spaces.py b/python/cuvs_bench/cuvs_bench/backends/search_spaces.py index 5a1848294c..cbc9accda2 100644 --- a/python/cuvs_bench/cuvs_bench/backends/search_spaces.py +++ b/python/cuvs_bench/cuvs_bench/backends/search_spaces.py @@ -158,6 +158,46 @@ "ef": {"type": "int", "min": 10, "max": 1000}, }, }, + # ========================================================================= + # Elasticsearch GPU HNSW (hnsw, int8_hnsw, int4_hnsw, bbq_hnsw) + # Per ES-GPU-API-REFERENCE.md: index_options (m, ef_construction), knn (num_candidates) + # ========================================================================= + "elastic_hnsw": { + "build": { + "m": {"type": "int", "min": 8, "max": 64}, + "ef_construction": {"type": "int", "min": 50, "max": 500}, + }, + "search": { + "num_candidates": {"type": "int", "min": 50, "max": 500}, + }, + }, + "elastic_int8_hnsw": { + "build": { + "m": {"type": "int", "min": 8, "max": 64}, + "ef_construction": {"type": "int", "min": 50, "max": 500}, + }, + "search": { + "num_candidates": {"type": "int", "min": 50, "max": 500}, + }, + }, + "elastic_int4_hnsw": { + "build": { + "m": {"type": "int", "min": 8, "max": 64}, + "ef_construction": {"type": "int", "min": 50, "max": 500}, + }, + "search": { + "num_candidates": {"type": "int", "min": 50, "max": 500}, + }, + }, + "elastic_bbq_hnsw": { + "build": { + "m": {"type": "int", "min": 8, "max": 64}, + "ef_construction": {"type": "int", "min": 50, "max": 500}, + }, + "search": { + "num_candidates": {"type": "int", "min": 50, "max": 500}, + }, + }, } diff --git a/python/cuvs_bench/cuvs_bench/tests/conftest.py b/python/cuvs_bench/cuvs_bench/tests/conftest.py new file mode 100644 index 0000000000..c003df9b3a --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/conftest.py @@ -0,0 +1,23 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# +"""Pytest configuration for cuvs_bench tests.""" + + +def pytest_configure(config): + """Register elastic plugin when elasticsearch is available. + + Ensures elastic tests run when elasticsearch is installed, even if + cuvs-bench-elastic was not installed via pip (e.g. using PYTHONPATH). + """ + try: + import elasticsearch # noqa: F401 + except ImportError: + return + + try: + from cuvs_bench_elastic import register + register() + except ImportError: + pass diff --git a/python/cuvs_bench/cuvs_bench/tests/integration/README.md b/python/cuvs_bench/cuvs_bench/tests/integration/README.md new file mode 100644 index 0000000000..3f73ce1081 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/integration/README.md @@ -0,0 +1,34 @@ +# Integration Tests + +Integration tests run against real services (e.g. Elasticsearch in Docker). + +**Note:** The Elasticsearch integration test is currently disabled. It targets the +Elasticsearch GPU backend (cuVS-accelerated), which requires an ES GPU image, +cuVS libs from this repo, and a GPU-enabled runner. See `test_elastic_integration.py` +module docstring for details. Can be re-enabled when CI has these dependencies. + +## Requirements + +- **Docker** running locally +- Optional dependencies: + + ```bash + pip install cuvs-bench[elastic,integration] + ``` + + Or install separately: + + ```bash + pip install cuvs-bench[elastic] + pip install testcontainers[elasticsearch] + ``` + +## Running + +```bash +# From repo root +PYTHONPATH="python/cuvs_bench:python/cuvs_bench_elastic:$PYTHONPATH" \ + pytest python/cuvs_bench/cuvs_bench/tests/integration/ -v +``` + +Integration tests are skipped automatically if `testcontainers` or `elasticsearch` is not installed. diff --git a/python/cuvs_bench/cuvs_bench/tests/integration/__init__.py b/python/cuvs_bench/cuvs_bench/tests/integration/__init__.py new file mode 100644 index 0000000000..7366b06538 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/integration/__init__.py @@ -0,0 +1,9 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# +"""Integration tests for cuvs-bench backends. + +These tests require Docker and optional dependencies: + pip install cuvs-bench[elastic,integration] +""" diff --git a/python/cuvs_bench/cuvs_bench/tests/integration/conftest.py b/python/cuvs_bench/cuvs_bench/tests/integration/conftest.py new file mode 100644 index 0000000000..ee103561b9 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/integration/conftest.py @@ -0,0 +1,73 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# +"""Pytest fixtures for integration tests. + +Requires: pip install cuvs-bench[elastic,integration] +Requires: Docker running locally +""" + +import pytest + + +def _testcontainers_available(): + try: + from testcontainers.elasticsearch import ElasticSearchContainer + return True + except ImportError: + return False + + +def _elasticsearch_installed(): + try: + import elasticsearch # noqa: F401 + return True + except ImportError: + return False + + +@pytest.fixture(scope="module") +def elasticsearch_container(): + """Start an Elasticsearch container for the duration of the test module. + + Yields a dict with host, port, and get_url() for connecting. + Skips the test module if testcontainers or elasticsearch is not installed. + """ + if not _testcontainers_available(): + pytest.skip( + "Requires testcontainers[elasticsearch]. " + "Install with: pip install cuvs-bench[integration]" + ) + if not _elasticsearch_installed(): + pytest.skip( + "Requires elasticsearch. Install with: pip install cuvs-bench[elastic]" + ) + + from testcontainers.elasticsearch import ElasticSearchContainer + + # Use standard ES OSS image (no GPU in OSS; use use_gpu=False in tests) + with ElasticSearchContainer( + image="elasticsearch:8.15.0", + mem_limit="2g", + ) as container: + url = container.get_url() + # Parse host:port from URL (e.g. http://localhost:32768) + if url.startswith("http://"): + rest = url[7:] + elif url.startswith("https://"): + rest = url[8:] + else: + rest = url + if "/" in rest: + host_port = rest.split("/")[0] + else: + host_port = rest + if ":" in host_port: + host, port_str = host_port.rsplit(":", 1) + port = int(port_str) + else: + host = host_port + port = 9200 + + yield {"host": host, "port": port, "url": url} diff --git a/python/cuvs_bench/cuvs_bench/tests/integration/test_elastic_integration.py b/python/cuvs_bench/cuvs_bench/tests/integration/test_elastic_integration.py new file mode 100644 index 0000000000..a2fffa2338 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/integration/test_elastic_integration.py @@ -0,0 +1,141 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# +"""Integration tests for Elasticsearch backend using testcontainers. + +**Currently disabled.** The cuvs-bench elastic backend targets Elasticsearch GPU +(cuVS-accelerated vector search). A proper integration test would require: +- Elasticsearch GPU image (docker.elastic.co/elasticsearch-dev/elasticsearch-gpu:9.3.x) +- cuVS libraries from this repo (built libcuvs.so) mounted into the container +- GPU-enabled CI runner (--gpus all) +- LD_LIBRARY_PATH and volume mounts configured + +The standard ES OSS image has no GPU support, so tests with use_gpu=False only +validate the HTTP API, not the GPU indexing path. This can be re-enabled later +when the above dependencies are available in CI. + +Requires: pip install cuvs-bench[elastic,integration] +Requires: Docker running locally + +Run with: + pytest python/cuvs_bench/cuvs_bench/tests/integration/ -v +""" + +import tempfile +from pathlib import Path + +import numpy as np +import pytest + +from cuvs_bench.backends.base import Dataset +from cuvs_bench.backends.registry import get_backend_class +from cuvs_bench.orchestrator.config_loaders import IndexConfig + + +def _write_fbin(path: Path, data: np.ndarray) -> None: + """Write big-ann-bench fbin format.""" + with open(path, "wb") as f: + np.array(data.shape, dtype=np.uint32).tofile(f) + data.astype(np.float32).tofile(f) + + +def _write_ibin(path: Path, data: np.ndarray) -> None: + """Write big-ann-bench ibin format.""" + with open(path, "wb") as f: + np.array(data.shape, dtype=np.uint32).tofile(f) + data.astype(np.int32).tofile(f) + + +@pytest.mark.skip( + reason="Integration test disabled: requires ES GPU image + cuVS libs + GPU runner. " + "See module docstring. Can be re-enabled when CI has these dependencies." +) +class TestElasticIntegration: + """Integration tests against a real Elasticsearch container. + + Disabled until CI can provide: Elasticsearch GPU image, cuVS libraries + (from this repo build), and GPU-enabled runner. Standard ES OSS does not + exercise the GPU vector search path this backend targets. + """ + + def test_elastic_build_and_search_e2e(self, elasticsearch_container): + """Build index and run search against Elasticsearch container.""" + es = elasticsearch_container + cls = get_backend_class("elastic") + backend = cls( + config={ + "name": "integration_test", + "host": es["host"], + "port": es["port"], + "index_name": "cuvs_bench_integration_test", + "use_gpu": False, # Standard ES OSS has no GPU + } + ) + try: + # Create temp dir with fbin/ibin files + with tempfile.TemporaryDirectory() as tmpdir: + base_path = Path(tmpdir) / "base.fbin" + query_path = Path(tmpdir) / "query.fbin" + gt_path = Path(tmpdir) / "groundtruth_neighbors.ibin" + + n_base, dims = 200, 32 + n_queries, k = 20, 10 + + base = np.random.rand(n_base, dims).astype(np.float32) + queries = np.random.rand(n_queries, dims).astype(np.float32) + # Ground truth: first k neighbors per query (dummy IDs) + gt = np.random.randint(0, n_base, size=(n_queries, k), dtype=np.int32) + + _write_fbin(base_path, base) + _write_fbin(query_path, queries) + _write_ibin(gt_path, gt) + + backend.config["data_prefix"] = str(tmpdir) + + dataset = Dataset( + name="integration_test", + base_vectors=base, + query_vectors=queries, + base_file="base.fbin", + query_file="query.fbin", + groundtruth_neighbors_file="groundtruth_neighbors.ibin", + distance_metric="euclidean", + ) + + indexes = [ + IndexConfig( + name="elastic_hnsw_integration", + algo="elastic_hnsw", + build_param={ + "m": 16, + "ef_construction": 64, + "use_gpu": False, + }, + search_params=[{"num_candidates": 50}], + file="", + ) + ] + + # Build + build_result = backend.build( + dataset=dataset, + indexes=indexes, + force=True, + dry_run=False, + ) + assert build_result.success, build_result.error_message + assert build_result.index_size_bytes > 0 + + # Search (uses data_prefix from config set above) + search_result = backend.search( + dataset=dataset, + indexes=indexes, + k=k, + dry_run=False, + ) + assert search_result.success, search_result.error_message + assert search_result.neighbors.shape == (n_queries, k) + assert search_result.search_time_ms >= 0 + finally: + backend.cleanup() diff --git a/python/cuvs_bench/cuvs_bench/tests/test_modularization.py b/python/cuvs_bench/cuvs_bench/tests/test_modularization.py new file mode 100644 index 0000000000..a5ddf127a1 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_modularization.py @@ -0,0 +1,591 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# +""" +Smoke tests for cuvs-bench modularization (optional deps, entry points, lazy loading). + +These tests verify the plugin infrastructure without requiring full backend implementations. +""" + +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest + +from cuvs_bench.backends.base import BenchmarkBackend, BuildResult, Dataset, SearchResult +from cuvs_bench.backends.registry import ( + get_backend_class, + get_config_loader, + get_registry, + list_backends, + list_config_loaders, + register_backend, + register_config_loader, + unregister_config_loader, +) +from cuvs_bench.orchestrator.config_loaders import ConfigLoader, IndexConfig + + +class TestModularizationSmoke: + """Smoke tests for optional backend loading and error handling.""" + + def test_cpp_gbench_available(self): + """cpp_gbench (built-in) should always be available.""" + backends = list_backends() + assert "cpp_gbench" in backends + cls = get_backend_class("cpp_gbench") + assert cls is not None + + def test_cpp_gbench_config_loader_available(self): + """cpp_gbench config loader should always be registered.""" + loaders = list_config_loaders() + assert "cpp_gbench" in loaders + loader_cls = get_config_loader("cpp_gbench") + assert loader_cls is not None + + def test_elastic_without_extra_raises_clear_error(self): + """Requesting elastic without [elastic] installed raises helpful error.""" + try: + import elasticsearch # noqa: F401 + pytest.skip("elasticsearch is installed; cannot test missing-plugin path") + except ImportError: + pass + import importlib.metadata + all_eps = importlib.metadata.entry_points() + if hasattr(all_eps, "select"): + eps = list(all_eps.select(group="cuvs_bench.backends", name="elastic")) + else: + eps = [e for e in all_eps.get("cuvs_bench.backends", []) if e.name == "elastic"] + if eps: + pytest.skip("cuvs-bench-elastic is installed; cannot test missing-plugin path") + + with pytest.raises((ImportError, ValueError)) as exc_info: + get_backend_class("elastic") + + msg = str(exc_info.value) + assert "elastic" in msg.lower() + + def test_elastic_config_loader_without_extra_raises_clear_error(self): + """Requesting elastic config loader without [elastic] raises helpful error.""" + try: + import elasticsearch # noqa: F401 + pytest.skip("elasticsearch is installed; cannot test missing-plugin path") + except ImportError: + pass + import importlib.metadata + all_eps = importlib.metadata.entry_points() + if hasattr(all_eps, "select"): + eps = list(all_eps.select(group="cuvs_bench.config_loaders", name="elastic")) + else: + eps = [e for e in all_eps.get("cuvs_bench.config_loaders", []) if e.name == "elastic"] + if eps: + pytest.skip("cuvs-bench-elastic is installed; cannot test missing-plugin path") + + with pytest.raises((ImportError, ValueError)) as exc_info: + get_config_loader("elastic") + + msg = str(exc_info.value) + assert "elastic" in msg.lower() + + def test_unknown_backend_raises_value_error(self): + """Requesting unknown backend raises ValueError with available backends.""" + with pytest.raises(ValueError) as exc_info: + get_backend_class("nonexistent_backend_xyz") + + msg = str(exc_info.value) + assert "nonexistent_backend_xyz" in msg + assert "cpp_gbench" in msg + + def test_unknown_config_loader_raises_value_error(self): + """Requesting unknown config loader raises ValueError.""" + with pytest.raises(ValueError) as exc_info: + get_config_loader("nonexistent_loader_xyz") + + msg = str(exc_info.value) + assert "nonexistent_loader_xyz" in msg + + def test_orchestrator_cpp_gbench_no_regression(self): + """BenchmarkOrchestrator with cpp_gbench should initialize (no regression).""" + from cuvs_bench.orchestrator import BenchmarkOrchestrator + + assert "cpp_gbench" in BenchmarkOrchestrator.available_backends() + orch = BenchmarkOrchestrator(backend_type="cpp_gbench") + assert orch.backend_type == "cpp_gbench" + assert orch.backend_class is not None + assert orch.config_loader is not None + + +class TestPluginLoaderMocked: + """ + Tests using mocked entry points (NAT-style). + + These tests do not require elasticsearch or real plugins. + """ + + _MOCK_PLUGIN_NAME = "mock_plugin_test" + + @staticmethod + def _make_mock_register(): + """Create a register function that registers a minimal stub backend and loader.""" + + class StubBackend(BenchmarkBackend): + def build(self, dataset, indexes, force=False, dry_run=False): + return BuildResult( + index_path="", + build_time_seconds=0, + index_size_bytes=0, + algorithm="stub", + build_params={}, + success=True, + ) + + def search( + self, dataset, indexes, k=10, batch_size=10000, + mode="latency", force=False, search_threads=None, dry_run=False, + ): + import numpy as np + return SearchResult( + neighbors=np.empty((0, k)), + distances=np.empty((0, k)), + search_time_ms=0, + queries_per_second=0, + recall=0, + algorithm="stub", + search_params=[], + success=True, + ) + + class StubConfigLoader(ConfigLoader): + @property + def backend_type(self): + return TestPluginLoaderMocked._MOCK_PLUGIN_NAME + + def load(self, **kwargs): + raise NotImplementedError("Stub loader") + + def register(): + register_backend(TestPluginLoaderMocked._MOCK_PLUGIN_NAME, StubBackend) + register_config_loader( + TestPluginLoaderMocked._MOCK_PLUGIN_NAME, StubConfigLoader + ) + + return register + + def test_valid_plugin_loads_via_mock_entry_point(self): + """Mock entry point: valid plugin registers and is discoverable.""" + mock_register = self._make_mock_register() + mock_ep = MagicMock() + mock_ep.name = self._MOCK_PLUGIN_NAME + mock_ep.load.return_value = mock_register + + mock_eps = MagicMock() + mock_eps.select.return_value = [mock_ep] + + with patch( + "cuvs_bench.backends.registry.importlib.metadata.entry_points", + return_value=mock_eps, + ): + cls = get_backend_class(self._MOCK_PLUGIN_NAME) + assert cls is not None + + loader_cls = get_config_loader(self._MOCK_PLUGIN_NAME) + assert loader_cls is not None + + # Cleanup + get_registry().unregister(self._MOCK_PLUGIN_NAME) + unregister_config_loader(self._MOCK_PLUGIN_NAME) + + def test_import_error_with_elasticsearch_message_raises_helpful_error(self): + """Mock entry point raising ImportError(elasticsearch) -> our install message.""" + # Ensure elastic is not in registry (e.g. from TestElasticWithExtraInstalled) + registry = get_registry() + if "elastic" in registry._backends: + registry.unregister("elastic") + unregister_config_loader("elastic") + + mock_ep = MagicMock() + mock_ep.name = "elastic" + mock_ep.load.side_effect = ImportError("No module named 'elasticsearch'") + + mock_eps = MagicMock() + mock_eps.select.return_value = [mock_ep] + + with patch( + "cuvs_bench.backends.registry.importlib.metadata.entry_points", + return_value=mock_eps, + ): + with pytest.raises(ImportError) as exc_info: + get_backend_class("elastic") + + msg = str(exc_info.value) + assert "pip install cuvs-bench[elastic]" in msg + + def test_import_error_unrelated_propagates(self): + """Mock entry point: unrelated ImportError propagates unchanged.""" + mock_ep = MagicMock() + mock_ep.name = "other_plugin" + mock_ep.load.side_effect = ImportError("No module named 'something_else'") + + mock_eps = MagicMock() + mock_eps.select.return_value = [mock_ep] + + with patch( + "cuvs_bench.backends.registry.importlib.metadata.entry_points", + return_value=mock_eps, + ): + with pytest.raises(ImportError) as exc_info: + get_backend_class("other_plugin") + + assert "something_else" in str(exc_info.value) + + def test_unexpected_error_propagates(self): + """Mock entry point: RuntimeError propagates.""" + mock_ep = MagicMock() + mock_ep.name = "broken_plugin" + mock_ep.load.side_effect = RuntimeError("Plugin crashed") + + mock_eps = MagicMock() + mock_eps.select.return_value = [mock_ep] + + with patch( + "cuvs_bench.backends.registry.importlib.metadata.entry_points", + return_value=mock_eps, + ): + with pytest.raises(RuntimeError) as exc_info: + get_backend_class("broken_plugin") + + assert "Plugin crashed" in str(exc_info.value) + + def test_no_entry_point_for_name_raises_value_error(self): + """Mock entry point: no plugin for requested name -> ValueError.""" + mock_eps = MagicMock() + mock_eps.select.return_value = [] # No matching entry points + + with patch( + "cuvs_bench.backends.registry.importlib.metadata.entry_points", + return_value=mock_eps, + ): + with pytest.raises(ValueError) as exc_info: + get_backend_class("nonexistent_mock_xyz") + + msg = str(exc_info.value) + assert "nonexistent_mock_xyz" in msg + assert "cpp_gbench" in msg + + +def _elasticsearch_installed(): + try: + import elasticsearch # noqa: F401 + return True + except ImportError: + return False + + +@pytest.mark.skipif( + not _elasticsearch_installed(), + reason="Requires pip install cuvs-bench[elastic]", +) +class TestElasticWithExtraInstalled: + """Tests that run only when [elastic] extra is installed.""" + + @pytest.fixture(autouse=True) + def _ensure_elastic_registered(self): + """Re-register elastic (may have been unregistered by other tests).""" + registry = get_registry() + if "elastic" not in registry._backends: + try: + from cuvs_bench_elastic import register + register() + except ImportError: + pass + yield + + def test_elastic_plugin_loads(self): + """Elastic backend and config loader load when elasticsearch is installed.""" + assert get_backend_class("elastic") is not None + assert get_config_loader("elastic") is not None + + def test_elastic_config_loader_tune_mode_returns_single_config(self): + """Tune mode produces one BenchmarkConfig with Optuna-suggested params.""" + loader_cls = get_config_loader("elastic") + loader = loader_cls() + dataset_config, benchmark_configs = loader.load( + dataset="glove-50-angular", + algorithms="elastic_hnsw", + _tune_mode=True, + _tune_build_params={"m": 24, "ef_construction": 150}, + _tune_search_params={"num_candidates": 120}, + ) + assert len(benchmark_configs) == 1 + config = benchmark_configs[0] + assert config.indexes[0].algo == "elastic_hnsw" + assert config.indexes[0].build_param["m"] == 24 + assert config.indexes[0].build_param["ef_construction"] == 150 + assert config.indexes[0].build_param["type"] == "hnsw" + assert config.indexes[0].search_params[0]["num_candidates"] == 120 + assert config.backend_config["name"] == "elastic_hnsw_tune" + + def test_elastic_config_loader_sweep_mode_returns_multiple_configs(self): + """Sweep mode produces multiple BenchmarkConfigs from param cartesian product.""" + loader_cls = get_config_loader("elastic") + loader = loader_cls() + dataset_config, benchmark_configs = loader.load( + dataset="glove-50-angular", + algorithms="test", + ) + assert len(benchmark_configs) >= 1 + config = benchmark_configs[0] + assert config.indexes[0].algo == "elastic_hnsw" + assert "m" in config.indexes[0].build_param + assert "num_candidates" in config.indexes[0].search_params[0] + + def test_elastic_config_loader_dataset_not_found_raises(self): + """Config loader raises ValueError for unknown dataset.""" + loader_cls = get_config_loader("elastic") + loader = loader_cls() + with pytest.raises(ValueError, match="not found"): + loader.load(dataset="nonexistent_dataset_xyz") + + def test_elastic_config_loader_group_not_found_raises(self): + """Config loader raises ValueError for unknown algorithm group.""" + loader_cls = get_config_loader("elastic") + loader = loader_cls() + with pytest.raises(ValueError, match="not found"): + loader.load( + dataset="glove-50-angular", + algorithms="nonexistent_group_xyz", + ) + + def test_elastic_dry_run_build(self): + """ElasticBackend.build(dry_run=True) returns synthetic result without ES.""" + cls = get_backend_class("elastic") + backend = cls(config={"name": "test", "host": "localhost", "port": 9200}) + + base = np.random.rand(100, 32).astype(np.float32) + queries = np.random.rand(10, 32).astype(np.float32) + dataset = Dataset( + name="test", + base_vectors=base, + query_vectors=queries, + distance_metric="euclidean", + ) + indexes = [ + IndexConfig( + name="elastic_hnsw_test", + algo="elastic_hnsw", + build_param={"m": 16, "ef_construction": 100}, + search_params=[{"num_candidates": 100}], + file="", + ) + ] + + result = backend.build(dataset=dataset, indexes=indexes, dry_run=True) + + assert result.success + assert result.algorithm == "elastic_hnsw" + assert result.build_time_seconds == 0 + + def test_elastic_dry_run_search(self): + """ElasticBackend.search(dry_run=True) returns synthetic result without ES.""" + cls = get_backend_class("elastic") + backend = cls(config={"name": "test", "host": "localhost", "port": 9200}) + + base = np.random.rand(100, 32).astype(np.float32) + queries = np.random.rand(10, 32).astype(np.float32) + dataset = Dataset( + name="test", + base_vectors=base, + query_vectors=queries, + distance_metric="euclidean", + ) + indexes = [ + IndexConfig( + name="elastic_hnsw_test", + algo="elastic_hnsw", + build_param={}, + search_params=[{"num_candidates": 100}], + file="", + ) + ] + + result = backend.search( + dataset=dataset, indexes=indexes, k=10, dry_run=True + ) + + assert result.success + assert result.algorithm == "elastic_hnsw" + assert result.search_time_ms == 0 + + def test_elastic_build_requires_base_file(self): + """ElasticBackend.build returns error when dataset has no base_file.""" + cls = get_backend_class("elastic") + backend = cls(config={"name": "test", "host": "localhost", "port": 9200}) + + base = np.random.rand(100, 32).astype(np.float32) + queries = np.random.rand(10, 32).astype(np.float32) + dataset = Dataset( + name="test", + base_vectors=base, + query_vectors=queries, + base_file=None, + query_file=None, + ) + indexes = [ + IndexConfig( + name="elastic_hnsw_test", + algo="elastic_hnsw", + build_param={}, + search_params=[{}], + file="", + ) + ] + + with patch.object( + backend, "_check_network_available", return_value=True + ): + result = backend.build(dataset=dataset, indexes=indexes, dry_run=False) + + assert not result.success + assert "base_file" in (result.error_message or "").lower() + + def test_elastic_preflight_skip_when_no_network(self): + """ElasticBackend skips build when network check fails.""" + cls = get_backend_class("elastic") + backend = cls(config={"name": "test", "host": "localhost", "port": 9200}) + + base = np.random.rand(100, 32).astype(np.float32) + queries = np.random.rand(10, 32).astype(np.float32) + dataset = Dataset( + name="test", + base_vectors=base, + query_vectors=queries, + base_file="dummy/base.fbin", + query_file="dummy/query.fbin", + ) + indexes = [ + IndexConfig( + name="elastic_hnsw_test", + algo="elastic_hnsw", + build_param={}, + search_params=[{}], + file="", + ) + ] + + with patch.object( + backend, "_check_network_available", return_value=False + ): + result = backend.build(dataset=dataset, indexes=indexes) + + assert result.success + assert result.metadata.get("skipped") is True + assert result.metadata.get("reason") == "no_network" + + def test_elastic_algo_from_config(self): + """ElasticBackend.algo derives from config type (elastic_hnsw, elastic_int8_hnsw).""" + cls = get_backend_class("elastic") + backend_hnsw = cls(config={"name": "test", "type": "hnsw"}) + assert backend_hnsw.algo == "elastic_hnsw" + + backend_int8 = cls(config={"name": "test", "type": "int8_hnsw"}) + assert backend_int8.algo == "elastic_int8_hnsw" + + def test_elastic_cleanup_closes_client(self): + """ElasticBackend.cleanup() closes client and sets _client to None.""" + cls = get_backend_class("elastic") + backend = cls(config={"name": "test", "host": "localhost", "port": 9200}) + mock_client = MagicMock() + backend._client = mock_client + + backend.cleanup() + + mock_client.close.assert_called_once() + assert backend._client is None + + def test_orchestrator_elastic_dry_run(self): + """BenchmarkOrchestrator with elastic backend runs dry_run without ES.""" + from cuvs_bench.orchestrator import BenchmarkOrchestrator + + orch = BenchmarkOrchestrator(backend_type="elastic") + results = orch.run_benchmark( + dataset="glove-50-angular", + dataset_path="/nonexistent", + host="localhost", + port=9200, + algorithms="test", + build=True, + search=True, + dry_run=True, + count=10, + batch_size=100, + ) + assert results is not None + assert len(results) >= 1 + + +@pytest.mark.skipif( + not _elasticsearch_installed(), + reason="Requires pip install cuvs-bench[elastic]", +) +class TestElasticHelpers: + """Tests for cuvs_bench_elastic helper functions.""" + + @pytest.fixture(autouse=True) + def _ensure_elastic_registered(self): + """Re-register elastic (may have been unregistered by other tests).""" + registry = get_registry() + if "elastic" not in registry._backends: + try: + from cuvs_bench_elastic import register + register() + except ImportError: + pass + yield + + def test_distance_to_similarity(self): + """_distance_to_similarity maps cuvs distance to ES similarity.""" + from cuvs_bench_elastic.backend import _distance_to_similarity + + assert _distance_to_similarity("euclidean") == "l2_norm" + assert _distance_to_similarity("inner_product") == "max_inner_product" + assert _distance_to_similarity("cosine") == "cosine" + assert _distance_to_similarity("unknown") == "l2_norm" + + def test_load_fbin(self): + """_load_fbin loads big-ann-bench fbin format.""" + import tempfile + from pathlib import Path + + from cuvs_bench_elastic.backend import _load_fbin + + data = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32) + with tempfile.NamedTemporaryFile(suffix=".fbin", delete=False) as f: + path = Path(f.name) + try: + with open(path, "wb") as f: + np.array([2, 2], dtype=np.uint32).tofile(f) + data.tofile(f) + loaded = _load_fbin(path) + np.testing.assert_array_equal(loaded, data) + finally: + path.unlink(missing_ok=True) + + def test_load_ibin(self): + """_load_ibin loads big-ann-bench ibin format.""" + import tempfile + from pathlib import Path + + from cuvs_bench_elastic.backend import _load_ibin + + data = np.array([[1, 2], [3, 4]], dtype=np.int32) + with tempfile.NamedTemporaryFile(suffix=".ibin", delete=False) as f: + path = Path(f.name) + try: + with open(path, "wb") as f: + np.array([2, 2], dtype=np.uint32).tofile(f) + data.tofile(f) + loaded = _load_ibin(path) + np.testing.assert_array_equal(loaded, data) + finally: + path.unlink(missing_ok=True) diff --git a/python/cuvs_bench/cuvs_bench/tests/test_registry.py b/python/cuvs_bench/cuvs_bench/tests/test_registry.py index 2bf7433f4f..6e844521c5 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_registry.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_registry.py @@ -20,153 +20,112 @@ register_backend, get_backend, ) +from cuvs_bench.orchestrator.config_loaders import IndexConfig class DummyBackend(BenchmarkBackend): - """Dummy backend for testing.""" + """Dummy backend for testing (implements current BenchmarkBackend API).""" @property def algo(self) -> str: - """Return dummy algorithm name.""" return "dummy_algo" - def build(self, dataset, build_params, index_path, force=False): + def build(self, dataset, indexes, force=False, dry_run=False): + idx = indexes[0] if indexes else None return BuildResult( - index_path=str(index_path), + index_path=idx.file if idx else "", build_time_seconds=1.0, index_size_bytes=1000, - algorithm=self.name, - build_params=build_params, + algorithm=self.algo, + build_params=dict(idx.build_param or {}) if idx else {}, success=True, ) def search( self, dataset, - search_params, - index_path, + indexes, k, batch_size=10000, mode="throughput", + force=False, + search_threads=None, + dry_run=False, ): n_queries = dataset.n_queries neighbors = np.random.randint(0, dataset.n_base, size=(n_queries, k)) distances = np.random.rand(n_queries, k) - + search_params = [] + if indexes: + search_params = indexes[0].search_params or [] return SearchResult( neighbors=neighbors, distances=distances, search_time_ms=0.1, queries_per_second=n_queries / 0.1, recall=0.95, - algorithm=self.name, + algorithm=self.algo, search_params=search_params, success=True, ) @property def name(self) -> str: - return "dummy" + return self.config.get("name", "dummy") class AnotherDummyBackend(BenchmarkBackend): """Another dummy backend for testing.""" - def build(self, dataset, build_params, index_path, force=False): + @property + def algo(self) -> str: + return "another_dummy_algo" + + def build(self, dataset, indexes, force=False, dry_run=False): + idx = indexes[0] if indexes else None return BuildResult( - index_path=str(index_path), + index_path=idx.file if idx else "", build_time_seconds=2.0, index_size_bytes=2000, - algorithm=self.name, - build_params=build_params, + algorithm=self.algo, + build_params=dict(idx.build_param or {}) if idx else {}, success=True, ) def search( self, dataset, - search_params, - index_path, + indexes, k, batch_size=10000, mode="throughput", + force=False, + search_threads=None, + dry_run=False, ): n_queries = dataset.n_queries neighbors = np.random.randint(0, dataset.n_base, size=(n_queries, k)) distances = np.random.rand(n_queries, k) - + search_params = indexes[0].search_params if indexes else [] return SearchResult( neighbors=neighbors, distances=distances, search_time_ms=0.2, queries_per_second=n_queries / 0.2, recall=0.90, - algorithm=self.name, + algorithm=self.algo, search_params=search_params, success=True, ) @property def name(self) -> str: - return "another_dummy" - - -class TestDataset: - """Tests for Dataset dataclass.""" - - def test_dataset_creation(self): - """Test basic dataset creation.""" - base = np.random.rand(1000, 128).astype(np.float32) - queries = np.random.rand(100, 128).astype(np.float32) - gt_neighbors = np.random.randint(0, 1000, size=(100, 10)) - - dataset = Dataset( - name="test_dataset", - base_vectors=base, - query_vectors=queries, - groundtruth_neighbors=gt_neighbors, - distance_metric="euclidean", - ) - - assert dataset.name == "test_dataset" - assert dataset.dims == 128 - assert dataset.n_base == 1000 - assert dataset.n_queries == 100 - assert dataset.distance_metric == "euclidean" - - def test_dataset_without_groundtruth(self): - """Test dataset without ground truth.""" - base = np.random.rand(500, 64).astype(np.float32) - queries = np.random.rand(50, 64).astype(np.float32) - - dataset = Dataset( - name="test_dataset_no_gt", base_vectors=base, query_vectors=queries - ) - - assert dataset.groundtruth_neighbors is None - assert dataset.groundtruth_distances is None + return self.config.get("name", "another_dummy") class TestBuildResult: """Tests for BuildResult dataclass.""" - def test_build_result_creation(self): - """Test basic build result creation.""" - result = BuildResult( - index_path="/path/to/index", - build_time_seconds=5.5, - index_size_bytes=1024000, - algorithm="test_algo", - build_params={"nlist": 1024}, - metadata={"gpu_time": 4.2}, - success=True, - ) - - assert result.index_path == "/path/to/index" - assert result.build_time_seconds == 5.5 - assert result.algorithm == "test_algo" - assert result.success is True - def test_build_result_to_json(self): """Test conversion to JSON format.""" result = BuildResult( @@ -177,9 +136,7 @@ def test_build_result_to_json(self): build_params={"nlist": 1024}, metadata={"gpu_time": 4.2}, ) - json_result = result.to_json() - assert json_result["name"] == "test_algo/build" assert json_result["real_time"] == 5.5 assert json_result["nlist"] == 1024 @@ -189,25 +146,6 @@ def test_build_result_to_json(self): class TestSearchResult: """Tests for SearchResult dataclass.""" - def test_search_result_creation(self): - """Test basic search result creation.""" - neighbors = np.array([[1, 2, 3], [4, 5, 6]]) - distances = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]) - - result = SearchResult( - neighbors=neighbors, - distances=distances, - search_time_ms=0.1, - queries_per_second=20.0, - recall=0.95, - algorithm="test_algo", - search_params=[{"nprobe": 10}], - ) - - assert result.recall == 0.95 - assert result.queries_per_second == 20.0 - assert result.success is True - def test_search_result_to_json(self): """Test conversion to JSON format.""" neighbors = np.array([[1, 2, 3]]) @@ -238,11 +176,6 @@ def test_search_result_to_json(self): class TestBackendRegistry: """Tests for BackendRegistry.""" - def test_registry_creation(self): - """Test registry creation.""" - registry = BackendRegistry() - assert len(registry.list_backends()) == 0 - def test_register_backend(self): """Test backend registration.""" registry = BackendRegistry() @@ -294,7 +227,7 @@ def test_get_backend(self): registry = BackendRegistry() registry.register("dummy", DummyBackend) - backend = registry.get_backend("dummy", config={}) + backend = registry.get_backend("dummy", config={"name": "dummy"}) assert isinstance(backend, DummyBackend) assert backend.name == "dummy" @@ -322,42 +255,60 @@ def test_list_backends(self): class TestBackendIntegration: - """Integration tests for backends.""" + """Integration tests for backends (uses current IndexConfig-based API).""" def test_dummy_backend_build(self): - """Test dummy backend build.""" - backend = DummyBackend(config={}) + """Test dummy backend build with IndexConfig.""" + backend = DummyBackend(config={"name": "dummy"}) base = np.random.rand(1000, 128).astype(np.float32) queries = np.random.rand(100, 128).astype(np.float32) dataset = Dataset( name="test", base_vectors=base, query_vectors=queries ) + indexes = [ + IndexConfig( + name="dummy_test", + algo="dummy_algo", + build_param={"nlist": 1024}, + search_params=[{"nprobe": 10}], + file="/tmp/test_index", + ) + ] result = backend.build( dataset=dataset, - build_params={"nlist": 1024}, - index_path=Path("/tmp/test_index"), + indexes=indexes, + force=False, + dry_run=False, ) assert result.success - assert result.algorithm == "dummy" + assert result.algorithm == "dummy_algo" assert result.build_params["nlist"] == 1024 def test_dummy_backend_search(self): - """Test dummy backend search.""" - backend = DummyBackend(config={}) + """Test dummy backend search with IndexConfig.""" + backend = DummyBackend(config={"name": "dummy"}) base = np.random.rand(1000, 128).astype(np.float32) queries = np.random.rand(100, 128).astype(np.float32) dataset = Dataset( name="test", base_vectors=base, query_vectors=queries ) + indexes = [ + IndexConfig( + name="dummy_test", + algo="dummy_algo", + build_param={}, + search_params=[{"nprobe": 10}], + file="/tmp/test_index", + ) + ] result = backend.search( dataset=dataset, - search_params=[{"nprobe": 10}], - index_path=Path("/tmp/test_index"), + indexes=indexes, k=10, ) @@ -370,14 +321,6 @@ def test_dummy_backend_search(self): class TestGlobalRegistry: """Tests for global registry functions.""" - def test_get_global_registry(self): - """Test getting global registry.""" - registry1 = get_registry() - registry2 = get_registry() - - # Should return the same instance (singleton) - assert registry1 is registry2 - def test_register_backend_global(self): """Test registering backend via global function.""" register_backend("dummy_global", DummyBackend) @@ -389,7 +332,7 @@ def test_get_backend_global(self): """Test getting backend via global function.""" register_backend("dummy_global2", DummyBackend) - backend = get_backend("dummy_global2", config={}) + backend = get_backend("dummy_global2", config={"name": "dummy"}) assert isinstance(backend, DummyBackend) assert backend.name == "dummy" diff --git a/python/cuvs_bench/pyproject.toml b/python/cuvs_bench/pyproject.toml index ed080c7b1f..2c607493ee 100644 --- a/python/cuvs_bench/pyproject.toml +++ b/python/cuvs_bench/pyproject.toml @@ -41,6 +41,10 @@ classifiers = [ [project.urls] Homepage = "https://github.com/rapidsai/cuvs" +[project.optional-dependencies] +elastic = ["cuvs-bench-elastic>=26.4.0"] +integration = ["cuvs-bench-elastic>=26.4.0", "testcontainers[elasticsearch]"] + [tool.setuptools.package-data] "*" = ["*.*", "VERSION"] diff --git a/python/cuvs_bench_elastic/README.md b/python/cuvs_bench_elastic/README.md new file mode 100644 index 0000000000..dd71e90122 --- /dev/null +++ b/python/cuvs_bench_elastic/README.md @@ -0,0 +1,42 @@ +# cuvs-bench-elastic + +Elasticsearch GPU backend plugin for [cuvs-bench](https://github.com/rapidsai/cuvs). + +## Installation + +```bash +pip install cuvs-bench[elastic] +``` + +Or install standalone (requires `cuvs-bench`): + +```bash +pip install cuvs-bench-elastic +``` + +## Usage + +Use the `elastic` backend with the cuvs-bench orchestrator: + +```python +from cuvs_bench.orchestrator import BenchmarkOrchestrator + +orch = BenchmarkOrchestrator(backend_type="elastic") +results = orch.run_benchmark( + dataset="test-data", + dataset_path="./datasets", + host="localhost", + port=9200, + basic_auth="elastic:password", + algorithms="test", + build=True, + search=True, +) +``` + +## Configuration + +Build params (index_options): `type`, `m`, `ef_construction`. +Search params (knn): `num_candidates`, `vector_field`. + +See `config/algos/elastic.yaml` for parameter groups. diff --git a/python/cuvs_bench_elastic/cuvs_bench_elastic/VERSION b/python/cuvs_bench_elastic/cuvs_bench_elastic/VERSION new file mode 100644 index 0000000000..0bd0e8a95b --- /dev/null +++ b/python/cuvs_bench_elastic/cuvs_bench_elastic/VERSION @@ -0,0 +1 @@ +26.04.00 diff --git a/python/cuvs_bench_elastic/cuvs_bench_elastic/__init__.py b/python/cuvs_bench_elastic/cuvs_bench_elastic/__init__.py new file mode 100644 index 0000000000..9d596e1f05 --- /dev/null +++ b/python/cuvs_bench_elastic/cuvs_bench_elastic/__init__.py @@ -0,0 +1,13 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# +""" +Elasticsearch GPU backend plugin for cuvs-bench. + +Install with: pip install cuvs-bench[elastic] +""" + +from .backend import register + +__all__ = ["register"] diff --git a/python/cuvs_bench_elastic/cuvs_bench_elastic/backend.py b/python/cuvs_bench_elastic/cuvs_bench_elastic/backend.py new file mode 100644 index 0000000000..210b5392f7 --- /dev/null +++ b/python/cuvs_bench_elastic/cuvs_bench_elastic/backend.py @@ -0,0 +1,656 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# +""" +Elasticsearch GPU HTTP backend for cuvs-bench. + +Install with: pip install cuvs-bench[elastic] + +Build params (index_options): type, m, ef_construction. + type: hnsw, int8_hnsw, int4_hnsw, bbq_hnsw (per ES-GPU-API-REFERENCE.md) + similarity: l2_norm, cosine, max_inner_product (overrides dataset distance) +Index settings: number_of_shards, number_of_replicas, use_gpu, vector_field. +Search params (knn): num_candidates, vector_field. +""" + +import itertools +import os +import time +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np +import yaml + +from elasticsearch import Elasticsearch +from elasticsearch.helpers import bulk + +from cuvs_bench.backends.base import BenchmarkBackend, BuildResult, Dataset, SearchResult +from cuvs_bench.backends.registry import register_backend, register_config_loader +from cuvs_bench.orchestrator.config_loaders import ( + BenchmarkConfig, + ConfigLoader, + DatasetConfig, + IndexConfig, +) + + +def _load_fbin(path: Path) -> np.ndarray: + """Load big-ann-bench fbin format (header: n_rows, n_cols as uint32, then float32).""" + with open(path, "rb") as f: + n_rows, n_cols = np.fromfile(f, dtype=np.uint32, count=2) + data = np.fromfile(f, dtype=np.float32).reshape(n_rows, n_cols) + return data + + +def _load_ibin(path: Path) -> np.ndarray: + """Load big-ann-bench ibin format (header: shape as uint32, then int32).""" + with open(path, "rb") as f: + shape = np.fromfile(f, dtype=np.uint32, count=2) + data = np.fromfile(f, dtype=np.int32).reshape(shape[0], shape[1]) + return data + + +def _distance_to_similarity(distance: str) -> str: + """Map cuvs-bench distance metric to ES dense_vector similarity.""" + m = {"euclidean": "l2_norm", "inner_product": "max_inner_product", "cosine": "cosine"} + return m.get(distance, "l2_norm") + + +# Defaults for index creation when not specified in config +_DEFAULT_INDEX_TYPE = "hnsw" +_DEFAULT_M = 16 +_DEFAULT_EF_CONSTRUCTION = 100 +_DEFAULT_NUM_SHARDS = 1 +_DEFAULT_NUM_REPLICAS = 0 +_DEFAULT_USE_GPU = True +_DEFAULT_VECTOR_FIELD = "embedding" +_DEFAULT_NUM_CANDIDATES = 100 + +# Build params passed through from config (index_options + index settings + similarity) +_BUILD_PARAM_KEYS = ( + "type", + "m", + "ef_construction", + "similarity", + "number_of_shards", + "number_of_replicas", + "use_gpu", + "vector_field", +) +# Search params passed through from config (knn) +_SEARCH_PARAM_KEYS = ("num_candidates", "vector_field") + + +class ElasticBackend(BenchmarkBackend): + """Elasticsearch GPU backend for vector benchmarking.""" + + def __init__(self, config: Dict[str, Any]): + super().__init__(config) + self._client: Optional[Elasticsearch] = None + + @property + def algo(self) -> str: + """Algorithm name from config (e.g. elastic_hnsw, elastic_int8_hnsw).""" + index_type = self.config.get("type", _DEFAULT_INDEX_TYPE) + return f"elastic_{index_type}" + + def _get_client(self) -> Elasticsearch: + if self._client is None: + host = self.config.get("host", "localhost") + port = self.config.get("port", 9200) + kwargs: Dict[str, Any] = {"hosts": [f"http://{host}:{port}"]} + basic_auth = self.config.get("basic_auth") + if basic_auth is not None: + if isinstance(basic_auth, (list, tuple)) and len(basic_auth) >= 2: + kwargs["basic_auth"] = (str(basic_auth[0]), str(basic_auth[1])) + elif isinstance(basic_auth, str) and ":" in basic_auth: + user, _, passwd = basic_auth.partition(":") + kwargs["basic_auth"] = (user, passwd) + self._client = Elasticsearch(**kwargs) + return self._client + + def cleanup(self) -> None: + """Close Elasticsearch client connection.""" + if self._client is not None: + self._client.close() + self._client = None + + @property + def requires_network(self) -> bool: + """Elasticsearch backend requires network connectivity.""" + return True + + def _check_network_available(self) -> bool: + """Verify Elasticsearch is reachable.""" + try: + return self._get_client().ping() + except Exception: + return False + + def build( + self, + dataset: Dataset, + indexes: List[IndexConfig], + force: bool = False, + dry_run: bool = False, + ) -> BuildResult: + """Build ES index from dataset vectors (fbin).""" + if dry_run: + return BuildResult( + index_path="", + build_time_seconds=0, + index_size_bytes=0, + algorithm=self.algo, + build_params={}, + success=True, + ) + + skip_reason = self._pre_flight_check() + if skip_reason: + idx = indexes[0] if indexes else None + return BuildResult( + index_path=idx.file if idx else "", + build_time_seconds=0, + index_size_bytes=0, + algorithm=self.algo, + build_params=dict(idx.build_param or {}) if idx else {}, + success=True, + metadata={"skipped": True, "reason": skip_reason}, + ) + + base_file = dataset.base_file + if not base_file: + return BuildResult( + index_path="", + build_time_seconds=0, + index_size_bytes=0, + algorithm=self.algo, + build_params={}, + success=False, + error_message="base_file is required for Elasticsearch backend", + ) + + data_prefix = self.config.get("data_prefix", "") + base_path = Path(data_prefix) / base_file + if not base_path.exists(): + return BuildResult( + index_path="", + build_time_seconds=0, + index_size_bytes=0, + algorithm=self.algo, + build_params={}, + success=False, + error_message=f"Base file not found: {base_path}", + ) + + index_name = self.config.get("index_name", "cuvs_bench_vectors") + dims = dataset.dims or self.config.get("dims") + if not dims: + return BuildResult( + index_path="", + build_time_seconds=0, + index_size_bytes=0, + algorithm=self.algo, + build_params={}, + success=False, + error_message="dims is required", + ) + + idx = indexes[0] if indexes else None + build_params = dict(idx.build_param or {}) if idx else {} + for k, v in self.config.items(): + if k not in build_params and k in _BUILD_PARAM_KEYS: + build_params[k] = v + # similarity: from config, or derive from dataset distance + similarity = build_params.get("similarity") or _distance_to_similarity( + getattr(dataset, "distance_metric", None) or "euclidean" + ) + + try: + vectors = _load_fbin(base_path) + n_vectors = len(vectors) + client = self._get_client() + + vector_field = build_params.get("vector_field", _DEFAULT_VECTOR_FIELD) + index_type = build_params.get("type", _DEFAULT_INDEX_TYPE) + m = build_params.get("m", _DEFAULT_M) + ef_construction = build_params.get( + "ef_construction", _DEFAULT_EF_CONSTRUCTION + ) + num_shards = build_params.get("number_of_shards", _DEFAULT_NUM_SHARDS) + num_replicas = build_params.get( + "number_of_replicas", _DEFAULT_NUM_REPLICAS + ) + use_gpu = build_params.get("use_gpu", _DEFAULT_USE_GPU) + + settings: Dict[str, Any] = { + "number_of_shards": num_shards, + "number_of_replicas": num_replicas, + } + if use_gpu: + settings["index"] = {"vectors.indexing.use_gpu": True} + + index_options: Dict[str, Any] = { + "type": index_type, + "m": m, + "ef_construction": ef_construction, + } + + index_config: Dict[str, Any] = { + "settings": settings, + "mappings": { + "properties": { + vector_field: { + "type": "dense_vector", + "dims": dims, + "index": True, + "similarity": similarity, + "index_options": index_options, + }, + }, + }, + } + + t0 = time.perf_counter() + if client.indices.exists(index=index_name): + if not force: + stats = client.indices.stats(index=index_name) + index_size = stats["_all"]["primaries"]["store"]["size_in_bytes"] + return BuildResult( + index_path=index_name, + build_time_seconds=0, + index_size_bytes=index_size, + algorithm=self.algo, + build_params=build_params, + success=True, + ) + client.indices.delete(index=index_name) + client.indices.create(index=index_name, body=index_config) + + chunk_size = 1000 + for i in range(0, n_vectors, chunk_size): + chunk = vectors[i : i + chunk_size] + actions = [] + for j, vec in enumerate(chunk): + doc_id = str(i + j) + actions.append({"index": {"_index": index_name, "_id": doc_id}}) + actions.append({vector_field: vec.tolist()}) + bulk(client, actions, raise_on_error=True) + + client.indices.refresh(index=index_name) + build_time = time.perf_counter() - t0 + + stats = client.indices.stats(index=index_name) + index_size = stats["_all"]["primaries"]["store"]["size_in_bytes"] + + return BuildResult( + index_path=index_name, + build_time_seconds=build_time, + index_size_bytes=index_size, + algorithm=self.algo, + build_params=build_params, + success=True, + ) + except Exception as e: + return BuildResult( + index_path="", + build_time_seconds=0, + index_size_bytes=0, + algorithm=self.algo, + build_params={}, + success=False, + error_message=str(e), + ) + + def search( + self, + dataset: Dataset, + indexes: List[IndexConfig], + k: int = 10, + batch_size: int = 10000, + mode: str = "latency", + force: bool = False, + search_threads: Optional[int] = None, + dry_run: bool = False, + ) -> SearchResult: + """Run kNN search and compute recall.""" + if dry_run: + return SearchResult( + neighbors=np.empty((0, k)), + distances=np.empty((0, k)), + search_time_ms=0, + queries_per_second=0, + recall=0, + algorithm=self.algo, + search_params=[], + success=True, + ) + + skip_reason = self._pre_flight_check() + if skip_reason: + idx = indexes[0] if indexes else None + return SearchResult( + neighbors=np.empty((0, k)), + distances=np.empty((0, k)), + search_time_ms=0, + queries_per_second=0, + recall=0, + algorithm=self.algo, + search_params=idx.search_params if idx and idx.search_params else [], + success=True, + metadata={"skipped": True, "reason": skip_reason}, + ) + + query_file = dataset.query_file + gt_file = dataset.groundtruth_neighbors_file + data_prefix = self.config.get("data_prefix", "") + + if not query_file: + return SearchResult( + neighbors=np.empty((0, k)), + distances=np.empty((0, k)), + search_time_ms=0, + queries_per_second=0, + recall=0, + algorithm=self.algo, + search_params=[], + success=False, + error_message="query_file is required", + ) + + query_path = Path(data_prefix) / query_file + if not query_path.exists(): + return SearchResult( + neighbors=np.empty((0, k)), + distances=np.empty((0, k)), + search_time_ms=0, + queries_per_second=0, + recall=0, + algorithm=self.algo, + search_params=[], + success=False, + error_message=f"Query file not found: {query_path}", + ) + + try: + query_vectors = _load_fbin(query_path) + n_queries = len(query_vectors) + groundtruth = None + if gt_file: + gt_path = Path(data_prefix) / gt_file + if gt_path.exists(): + groundtruth = _load_ibin(gt_path) + + index_name = self.config.get("index_name", "cuvs_bench_vectors") + search_params = {} + if indexes and indexes[0] and indexes[0].search_params: + search_params = dict(indexes[0].search_params[0] or {}) + for k, v in self.config.items(): + if k not in search_params and k in _SEARCH_PARAM_KEYS: + search_params[k] = v + num_candidates = search_params.get( + "num_candidates", _DEFAULT_NUM_CANDIDATES + ) + vector_field = search_params.get( + "vector_field", _DEFAULT_VECTOR_FIELD + ) + + neighbors_list: List[List[int]] = [] + latencies: List[float] = [] + + for i, qv in enumerate(query_vectors): + body = { + "knn": { + "field": vector_field, + "query_vector": qv.tolist(), + "k": k, + "num_candidates": num_candidates, + } + } + t0 = time.perf_counter() + resp = self._get_client().search( + index=index_name, body=body, size=k + ) + latencies.append((time.perf_counter() - t0) * 1000) + hits = resp.get("hits", {}).get("hits", []) + ids = [int(h["_id"]) for h in hits if "_id" in h] + neighbors_list.append(ids[:k]) + + search_time_ms = sum(latencies) + qps = n_queries / (search_time_ms / 1000) if search_time_ms > 0 else 0 + + recall = 0.0 + if groundtruth is not None and len(neighbors_list) <= len(groundtruth): + correct = 0 + total = 0 + for q in range(len(neighbors_list)): + retrieved = set(neighbors_list[q]) + gt_row = groundtruth[q, :k] + for gt_id in gt_row: + if int(gt_id) in retrieved: + correct += 1 + total += 1 + recall = correct / total if total > 0 else 0 + + max_len = max(len(r) for r in neighbors_list) + neighbors_arr = np.array( + [r + [-1] * (max_len - len(r)) for r in neighbors_list], + dtype=np.int64, + ) + distances_arr = np.zeros_like(neighbors_arr, dtype=np.float32) + + return SearchResult( + neighbors=neighbors_arr, + distances=distances_arr, + search_time_ms=search_time_ms, + queries_per_second=qps, + recall=recall, + algorithm=self.algo, + search_params=[search_params], + latency_percentiles={ + "p50_ms": float(np.percentile(latencies, 50)), + "p95_ms": float(np.percentile(latencies, 95)), + "p99_ms": float(np.percentile(latencies, 99)), + }, + success=True, + ) + except Exception as e: + return SearchResult( + neighbors=np.empty((0, k)), + distances=np.empty((0, k)), + search_time_ms=0, + queries_per_second=0, + recall=0, + algorithm=self.algo, + search_params=[], + success=False, + error_message=str(e), + ) + + +def _get_cuvs_bench_config_path() -> str: + """Get cuvs_bench config directory for shared datasets.yaml.""" + import cuvs_bench.orchestrator.config_loaders as _config_loaders + + mod_file = getattr(_config_loaders, "__file__", None) + if mod_file: + # config_loaders is at cuvs_bench/orchestrator/config_loaders.py + # config is at cuvs_bench/config + pkg_dir = Path(mod_file).resolve().parent.parent + return str(pkg_dir / "config") + import cuvs_bench + + if cuvs_bench.__file__: + return os.path.join(os.path.dirname(cuvs_bench.__file__), "config") + raise RuntimeError( + "Cannot determine cuvs_bench config path. " + "Ensure cuvs_bench is properly installed or on PYTHONPATH." + ) + + +def _get_elastic_config_path() -> str: + """Get this package's config directory for elastic.yaml.""" + return os.path.join(os.path.dirname(__file__), "config") + + +class ElasticConfigLoader(ConfigLoader): + """Config loader for Elasticsearch backend.""" + + def __init__(self, config_path: Optional[str] = None): + self.config_path = config_path or _get_cuvs_bench_config_path() + + @property + def backend_type(self) -> str: + return "elastic" + + def load( + self, + dataset: str = "", + dataset_path: str = "", + count: int = 10, + batch_size: int = 10000, + host: str = "localhost", + port: int = 9200, + index_name: str = "cuvs_bench_vectors", + basic_auth: Optional[Any] = None, + algorithms: Optional[str] = None, + subset_size: Optional[int] = None, + **kwargs, + ) -> Tuple[DatasetConfig, List[BenchmarkConfig]]: + """Load Elasticsearch benchmark configuration.""" + tune_mode = kwargs.pop("_tune_mode", False) + tune_build_params = kwargs.pop("_tune_build_params", None) + tune_search_params = kwargs.pop("_tune_search_params", None) + + datasets_path = os.path.join( + self.config_path, "datasets", "datasets.yaml" + ) + try: + with open(datasets_path) as f: + datasets = yaml.safe_load(f) + except FileNotFoundError: + raise FileNotFoundError( + f"Datasets config not found: {datasets_path}" + ) from None + + dataset_conf = next((d for d in datasets if d["name"] == dataset), None) + if not dataset_conf: + raise ValueError(f"Dataset '{dataset}' not found") + + dataset_config = DatasetConfig( + name=dataset_conf["name"], + base_file=dataset_conf.get("base_file"), + query_file=dataset_conf.get("query_file"), + groundtruth_neighbors_file=dataset_conf.get( + "groundtruth_neighbors_file" + ), + distance=dataset_conf.get("distance", "euclidean"), + dims=dataset_conf.get("dims"), + subset_size=subset_size, + ) + + if tune_mode and tune_build_params is not None and tune_search_params is not None: + algo_name = algorithms or "elastic_hnsw" + # Extract type from algo (elastic_hnsw -> hnsw, elastic_int8_hnsw -> int8_hnsw) + build_param = dict(tune_build_params) + if "type" not in build_param and algo_name.startswith("elastic_"): + build_param["type"] = algo_name.replace("elastic_", "", 1) + index_config = IndexConfig( + name=f"{algo_name}_tune", + algo=algo_name, + build_param=build_param, + search_params=[tune_search_params], + file="", + ) + config = BenchmarkConfig( + indexes=[index_config], + backend_config={ + "name": index_config.name, + "host": host, + "port": port, + "index_name": index_name, + "basic_auth": basic_auth, + "data_prefix": dataset_path, + "dims": dataset_config.dims, + **build_param, + **tune_search_params, + }, + ) + return dataset_config, [config] + + elastic_config_dir = _get_elastic_config_path() + elastic_algo_path = os.path.join(elastic_config_dir, "algos", "elastic.yaml") + default_grp = { + "build": {"m": [16], "ef_construction": [100]}, + "search": {"num_candidates": [100]}, + } + if os.path.exists(elastic_algo_path): + with open(elastic_algo_path) as f: + algo_conf = yaml.safe_load(f) + else: + algo_conf = {"groups": {"base": default_grp}} + + groups = algo_conf.get("groups", {"base": default_grp}) + group_name = algorithms or "base" + if group_name not in groups: + raise ValueError( + f"Algorithm group '{group_name}' not found in elastic.yaml. " + f"Available: {list(groups.keys())}" + ) + group_conf = groups[group_name] + + build_params = group_conf.get( + "build", {"m": [16], "ef_construction": [100]} + ) + search_params = group_conf.get("search", {"num_candidates": [100]}) + + # Ensure all param values are lists for itertools.product + def _to_list_values(d: Dict[str, Any]) -> Dict[str, List[Any]]: + return {k: v if isinstance(v, list) else [v] for k, v in d.items()} + + build_params = _to_list_values(build_params) + search_params = _to_list_values(search_params) + + build_combos = list(itertools.product(*build_params.values())) + search_combos = list(itertools.product(*search_params.values())) + build_keys = list(build_params.keys()) + search_keys = list(search_params.keys()) + + benchmark_configs = [] + for bvals in build_combos: + bdict = dict(zip(build_keys, bvals)) + for svals in search_combos: + sdict = dict(zip(search_keys, svals)) + algo_name = f"elastic_{bdict.get('type', 'hnsw')}" + index_config = IndexConfig( + name=f"{algo_name}_{group_name}", + algo=algo_name, + build_param=bdict, + search_params=[sdict], + file="", + ) + config = BenchmarkConfig( + indexes=[index_config], + backend_config={ + "name": index_config.name, + "host": host, + "port": port, + "index_name": index_name, + "basic_auth": basic_auth, + "data_prefix": dataset_path, + "dims": dataset_config.dims, + **bdict, + **sdict, + }, + ) + benchmark_configs.append(config) + + return dataset_config, benchmark_configs + + +def register() -> None: + """Register Elasticsearch backend and config loader (called from entry point).""" + register_backend("elastic", ElasticBackend) + register_config_loader("elastic", ElasticConfigLoader) diff --git a/python/cuvs_bench_elastic/cuvs_bench_elastic/config/algos/elastic.yaml b/python/cuvs_bench_elastic/cuvs_bench_elastic/config/algos/elastic.yaml new file mode 100644 index 0000000000..054c86dfa8 --- /dev/null +++ b/python/cuvs_bench_elastic/cuvs_bench_elastic/config/algos/elastic.yaml @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# +# Elasticsearch GPU HNSW: index_options (build) and knn (search) +# Per ES-GPU-API-REFERENCE.md and DOCKER-INSPECT-FINDINGS.md +# +# Build (index_options): type, m, ef_construction +# type: hnsw (float), int8_hnsw (4x quantized), int4_hnsw (8x), bbq_hnsw (32x) +# similarity: l2_norm, cosine, max_inner_product (overrides dataset distance) +# Index settings: number_of_shards, number_of_replicas, use_gpu, vector_field +# Search (knn): num_candidates, vector_field +name: elastic_hnsw +groups: + base: + build: + type: [hnsw] + m: [16, 32] + ef_construction: [100, 200] + similarity: [l2_norm, cosine] + number_of_shards: [1] + number_of_replicas: [0] + use_gpu: [true] + vector_field: [embedding] + search: + num_candidates: [100, 200] + vector_field: [embedding] + quantized: + build: + type: [hnsw, int8_hnsw] + m: [16, 32] + ef_construction: [100] + similarity: [l2_norm] + number_of_shards: [1] + number_of_replicas: [0] + use_gpu: [true] + vector_field: [embedding] + search: + num_candidates: [100, 200] + vector_field: [embedding] + test: + # use_gpu: false for ES source build (N/N path) - GPU enabled at cluster level only + build: + type: [hnsw] + m: [16] + ef_construction: [100] + similarity: [l2_norm] + number_of_shards: [1] + number_of_replicas: [0] + use_gpu: [false] + vector_field: [embedding] + search: + num_candidates: [100] + vector_field: [embedding] diff --git a/python/cuvs_bench_elastic/pyproject.toml b/python/cuvs_bench_elastic/pyproject.toml new file mode 100644 index 0000000000..e62d118732 --- /dev/null +++ b/python/cuvs_bench_elastic/pyproject.toml @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "cuvs-bench-elastic" +dynamic = ["version"] +description = "Elasticsearch GPU backend plugin for cuvs-bench" +readme = { file = "README.md", content-type = "text/markdown" } +authors = [ + { name = "NVIDIA Corporation" }, +] +license = { text = "Apache-2.0" } +requires-python = ">=3.11" +classifiers = [ + "Intended Audience :: Developers", + "Topic :: Database", + "Topic :: Scientific/Engineering", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", +] +dependencies = [ + "cuvs-bench>=26.4.0", + "elasticsearch>=8.0", + "numpy", + "pyyaml", +] + +[project.urls] +Homepage = "https://github.com/rapidsai/cuvs" + +[project.entry-points."cuvs_bench.backends"] +elastic = "cuvs_bench_elastic:register" + +[project.entry-points."cuvs_bench.config_loaders"] +elastic = "cuvs_bench_elastic:register" + +[tool.setuptools.dynamic] +version = { file = "cuvs_bench_elastic/VERSION" } + +[tool.setuptools.packages.find] +where = ["."] +include = ["cuvs_bench_elastic*"] + +[tool.setuptools.package-data] +cuvs_bench_elastic = ["config/**/*.yaml"] From 80c104fe4d7058f1e82c47e9fda60e69331960bc Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Fri, 20 Mar 2026 09:39:52 -0700 Subject: [PATCH 02/22] Fix bulk indexing format for Elasticsearch Use single-doc format (_index, _id, vector_field) instead of two-part NDJSON (index action + source) so ES accepts the bulk request. Signed-off-by: Alex Fournier --- .../cuvs_bench_elastic/backend.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/python/cuvs_bench_elastic/cuvs_bench_elastic/backend.py b/python/cuvs_bench_elastic/cuvs_bench_elastic/backend.py index 210b5392f7..a60e352ccf 100644 --- a/python/cuvs_bench_elastic/cuvs_bench_elastic/backend.py +++ b/python/cuvs_bench_elastic/cuvs_bench_elastic/backend.py @@ -272,11 +272,14 @@ def build( chunk_size = 1000 for i in range(0, n_vectors, chunk_size): chunk = vectors[i : i + chunk_size] - actions = [] - for j, vec in enumerate(chunk): - doc_id = str(i + j) - actions.append({"index": {"_index": index_name, "_id": doc_id}}) - actions.append({vector_field: vec.tolist()}) + actions = [ + { + "_index": index_name, + "_id": str(i + j), + vector_field: vec.tolist(), + } + for j, vec in enumerate(chunk) + ] bulk(client, actions, raise_on_error=True) client.indices.refresh(index=index_name) @@ -387,9 +390,9 @@ def search( search_params = {} if indexes and indexes[0] and indexes[0].search_params: search_params = dict(indexes[0].search_params[0] or {}) - for k, v in self.config.items(): - if k not in search_params and k in _SEARCH_PARAM_KEYS: - search_params[k] = v + for key, val in self.config.items(): + if key not in search_params and key in _SEARCH_PARAM_KEYS: + search_params[key] = val num_candidates = search_params.get( "num_candidates", _DEFAULT_NUM_CANDIDATES ) From 7b25521ad899335d2f3f075d90fcfae107b7dff7 Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Fri, 20 Mar 2026 09:39:59 -0700 Subject: [PATCH 03/22] Add run_build, run_search, run_benchmark API for elastic backend Expose ELASTIC constant and convenience wrappers for build-only, search-only, or full benchmark runs. Signed-off-by: Alex Fournier --- .../cuvs_bench_elastic/__init__.py | 143 +++++++++++++++++- 1 file changed, 142 insertions(+), 1 deletion(-) diff --git a/python/cuvs_bench_elastic/cuvs_bench_elastic/__init__.py b/python/cuvs_bench_elastic/cuvs_bench_elastic/__init__.py index 9d596e1f05..df3578d3d7 100644 --- a/python/cuvs_bench_elastic/cuvs_bench_elastic/__init__.py +++ b/python/cuvs_bench_elastic/cuvs_bench_elastic/__init__.py @@ -6,8 +6,149 @@ Elasticsearch GPU backend plugin for cuvs-bench. Install with: pip install cuvs-bench[elastic] + +Usage: + # Build and search as separate operations (matches backend interface) + from cuvs_bench_elastic import run_build, run_search + build_results = run_build(dataset="test-data", dataset_path="./datasets", host="localhost", port=9200) + search_results = run_search(dataset="test-data", dataset_path="./datasets", host="localhost", port=9200) + + # Or run both + from cuvs_bench_elastic import run_benchmark + results = run_benchmark(dataset="test-data", dataset_path="./datasets", host="localhost", port=9200) + + # Or use the orchestrator directly + from cuvs_bench_elastic import ELASTIC, register + from cuvs_bench.orchestrator import BenchmarkOrchestrator + register() + orch = BenchmarkOrchestrator(backend_type=ELASTIC) + results = orch.run_benchmark(dataset="test-data", dataset_path="./datasets", host="localhost", port=9200) """ from .backend import register -__all__ = ["register"] +# Backend type for use with BenchmarkOrchestrator(backend_type=ELASTIC) +ELASTIC = "elastic" + + +def _common_kwargs( + dataset: str, + dataset_path: str, + host: str, + port: int, + algorithms: str, + force: bool, + username: str, + password: str, + **kwargs, +): + """Build kwargs shared by run_build, run_search, run_benchmark.""" + d = dict( + dataset=dataset, + dataset_path=dataset_path, + host=host, + port=port, + algorithms=algorithms, + force=force, + **kwargs, + ) + if username: + d["username"] = username + if password: + d["password"] = password + return d + + +def run_build( + dataset: str = "test-data", + dataset_path: str = "./datasets", + host: str = "localhost", + port: int = 9200, + algorithms: str = "test", + force: bool = False, + username: str = None, + password: str = None, + **kwargs, +): + """ + Build an Elasticsearch vector index. Matches the backend build() interface. + + Returns + ------- + list + BuildResult objects + """ + register() + from cuvs_bench.orchestrator import BenchmarkOrchestrator + + orch = BenchmarkOrchestrator(backend_type=ELASTIC) + return orch.run_benchmark( + build=True, + search=False, + **_common_kwargs( + dataset, dataset_path, host, port, algorithms, force, username, password, **kwargs + ), + ) + + +def run_search( + dataset: str = "test-data", + dataset_path: str = "./datasets", + host: str = "localhost", + port: int = 9200, + algorithms: str = "test", + username: str = None, + password: str = None, + **kwargs, +): + """ + Run kNN search against an existing Elasticsearch index. Matches the backend search() interface. + + Returns + ------- + list + SearchResult objects + """ + register() + from cuvs_bench.orchestrator import BenchmarkOrchestrator + + orch = BenchmarkOrchestrator(backend_type=ELASTIC) + return orch.run_benchmark( + build=False, + search=True, + **_common_kwargs( + dataset, dataset_path, host, port, algorithms, False, username, password, **kwargs + ), + ) + + +def run_benchmark( + dataset: str = "test-data", + dataset_path: str = "./datasets", + host: str = "localhost", + port: int = 9200, + algorithms: str = "test", + build: bool = True, + search: bool = True, + force: bool = False, + username: str = None, + password: str = None, + **kwargs, +): + """ + Run build and/or search. Convenience wrapper; run_build and run_search are the primary API. + """ + register() + from cuvs_bench.orchestrator import BenchmarkOrchestrator + + orch = BenchmarkOrchestrator(backend_type=ELASTIC) + return orch.run_benchmark( + build=build, + search=search, + **_common_kwargs( + dataset, dataset_path, host, port, algorithms, force, username, password, **kwargs + ), + ) + + +__all__ = ["register", "ELASTIC", "run_build", "run_search", "run_benchmark"] From c224a8ce6e99fa0adbbb5eff1f7381484fdffbde Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Fri, 20 Mar 2026 10:04:27 -0700 Subject: [PATCH 04/22] Update cuvs-bench-elastic README and config loader - Document run_build, run_search, run_benchmark convenience API - Document ELASTIC constant and orchestrator usage - Add username/password support in config loader (converts to basic_auth) Signed-off-by: Alex Fournier --- python/cuvs_bench_elastic/README.md | 48 +++++++++++++++++-- .../cuvs_bench_elastic/backend.py | 4 ++ 2 files changed, 49 insertions(+), 3 deletions(-) diff --git a/python/cuvs_bench_elastic/README.md b/python/cuvs_bench_elastic/README.md index dd71e90122..edd92adf11 100644 --- a/python/cuvs_bench_elastic/README.md +++ b/python/cuvs_bench_elastic/README.md @@ -16,18 +16,60 @@ pip install cuvs-bench-elastic ## Usage -Use the `elastic` backend with the cuvs-bench orchestrator: +### Convenience API (recommended) ```python +from cuvs_bench_elastic import run_build, run_search, run_benchmark + +# Build only +build_results = run_build( + dataset="test-data", + dataset_path="./datasets", + host="localhost", + port=9200, +) + +# Search only (requires existing index) +search_results = run_search( + dataset="test-data", + dataset_path="./datasets", + host="localhost", + port=9200, +) + +# Build and search +results = run_benchmark( + dataset="test-data", + dataset_path="./datasets", + host="localhost", + port=9200, +) + +# With authentication +results = run_benchmark( + dataset="test-data", + dataset_path="./datasets", + host="localhost", + port=9200, + username="elastic", + password="password", +) +``` + +### Orchestrator API + +```python +from cuvs_bench_elastic import ELASTIC, register from cuvs_bench.orchestrator import BenchmarkOrchestrator -orch = BenchmarkOrchestrator(backend_type="elastic") +register() +orch = BenchmarkOrchestrator(backend_type=ELASTIC) results = orch.run_benchmark( dataset="test-data", dataset_path="./datasets", host="localhost", port=9200, - basic_auth="elastic:password", + basic_auth="elastic:password", # or username="elastic", password="password" algorithms="test", build=True, search=True, diff --git a/python/cuvs_bench_elastic/cuvs_bench_elastic/backend.py b/python/cuvs_bench_elastic/cuvs_bench_elastic/backend.py index a60e352ccf..4980b24101 100644 --- a/python/cuvs_bench_elastic/cuvs_bench_elastic/backend.py +++ b/python/cuvs_bench_elastic/cuvs_bench_elastic/backend.py @@ -526,6 +526,10 @@ def load( tune_mode = kwargs.pop("_tune_mode", False) tune_build_params = kwargs.pop("_tune_build_params", None) tune_search_params = kwargs.pop("_tune_search_params", None) + username = kwargs.pop("username", None) + password = kwargs.pop("password", None) + if basic_auth is None and username and password: + basic_auth = (username, password) datasets_path = os.path.join( self.config_path, "datasets", "datasets.yaml" From e499469e09246d91b4f8f13cdc0967a8ab3ada50 Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Sun, 12 Apr 2026 13:50:52 -0700 Subject: [PATCH 05/22] Add cuvs-bench-elastic: Elasticsearch GPU backend plugin MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - New `cuvs_bench_elastic` package with HTTP backend for Elasticsearch GPU vector search (HNSW, int8_hnsw, int4_hnsw, bbq_hnsw index types) - Supports `pip install cuvs-bench[elastic]` without a separate PyPI publish: `cuvs_bench` bundles the plugin via setuptools packages.find - Plugin registers via entry points (`cuvs_bench.backends` / `cuvs_bench.config_loaders`) — no changes to core cuvs-bench required - `ElasticConfigLoader` reads shared `datasets.yaml` from cuvs_bench and `elastic.yaml` from the plugin config; supports sweep and tune modes - `build()` checks index existence before file validation so `force=False` returns immediately without requiring the base file on disk - Removed testcontainers-based integration tests; added unit tests for pre-flight failure, force=False skip, dry-run, helper functions - `elasticsearch` client is an optional dep (`cuvs-bench[elastic]` extra) --- .../cuvs_bench/get_dataset/__main__.py | 6 +- .../cuvs_bench/orchestrator/orchestrator.py | 118 +++--- .../cuvs_bench/tests/integration/README.md | 34 -- .../cuvs_bench/tests/integration/__init__.py | 9 - .../cuvs_bench/tests/integration/conftest.py | 73 ---- .../integration/test_elastic_integration.py | 141 ------- .../cuvs_bench/tests/test_modularization.py | 78 +++- python/cuvs_bench/pyproject.toml | 13 +- .../cuvs_bench_elastic/backend.py | 396 ++++++++++-------- .../config/algos/elastic.yaml | 7 +- python/cuvs_bench_elastic/pyproject.toml | 6 +- 11 files changed, 380 insertions(+), 501 deletions(-) delete mode 100644 python/cuvs_bench/cuvs_bench/tests/integration/README.md delete mode 100644 python/cuvs_bench/cuvs_bench/tests/integration/__init__.py delete mode 100644 python/cuvs_bench/cuvs_bench/tests/integration/conftest.py delete mode 100644 python/cuvs_bench/cuvs_bench/tests/integration/test_elastic_integration.py diff --git a/python/cuvs_bench/cuvs_bench/get_dataset/__main__.py b/python/cuvs_bench/cuvs_bench/get_dataset/__main__.py index 3c52be00d8..1670dd2097 100644 --- a/python/cuvs_bench/cuvs_bench/get_dataset/__main__.py +++ b/python/cuvs_bench/cuvs_bench/get_dataset/__main__.py @@ -4,6 +4,7 @@ import os import subprocess +import sys import click import h5py @@ -34,13 +35,14 @@ def convert_hdf5_to_fbin(path, normalize): scripts_path = os.path.dirname(os.path.realpath(__file__)) ann_bench_scripts_path = os.path.join(scripts_path, "hdf5_to_fbin.py") print(f"calling script {ann_bench_scripts_path}") + python = sys.executable if normalize and "angular" in path: subprocess.run( - ["python", ann_bench_scripts_path, "-n", "%s" % path], check=True + [python, ann_bench_scripts_path, "-n", "%s" % path], check=True ) else: subprocess.run( - ["python", ann_bench_scripts_path, "%s" % path], check=True + [python, ann_bench_scripts_path, "%s" % path], check=True ) diff --git a/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py b/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py index 091290979c..839dc359ec 100644 --- a/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py +++ b/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py @@ -222,43 +222,47 @@ def _run_sweep( for config in benchmark_configs: # Create backend instance backend = self.backend_class(config.backend_config) - - if build: - # Pass ALL indexes at once - ONE C++ command builds all - build_result = backend.build( - dataset=bench_dataset, - indexes=config.indexes, - force=force, - 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}" + try: + backend.initialize() + + if build: + # Pass ALL indexes at once - ONE C++ command builds all + build_result = backend.build( + dataset=bench_dataset, + indexes=config.indexes, + force=force, + dry_run=dry_run, ) - continue - - if search: - # Pass ALL indexes at once - ONE C++ command searches all - # Each index has its own search_params list - # Total benchmarks = sum(len(idx.search_params) for idx in indexes) - search_result = backend.search( - dataset=bench_dataset, - indexes=config.indexes, - k=count, - batch_size=batch_size, - mode=search_mode, - force=force, - search_threads=search_threads, - dry_run=dry_run, - ) - results.append(search_result) - - if not search_result.success: - print( - f"Search failed for {config.index_name}: {search_result.error_message}" + results.append(build_result) + + if not build_result.success: + print( + f"Build failed for {config.index_name}: {build_result.error_message}" + ) + continue + + if search: + # Pass ALL indexes at once - ONE C++ command searches all + # Each index has its own search_params list + # Total benchmarks = sum(len(idx.search_params) for idx in indexes) + search_result = backend.search( + dataset=bench_dataset, + indexes=config.indexes, + k=count, + batch_size=batch_size, + mode=search_mode, + force=force, + search_threads=search_threads, + dry_run=dry_run, ) + results.append(search_result) + + if not search_result.success: + print( + f"Search failed for {config.index_name}: {search_result.error_message}" + ) + finally: + backend.cleanup() return results @@ -552,28 +556,32 @@ def _run_trial( backend = self.backend_class(backend_config) result = None + try: + backend.initialize() - if build: - result = backend.build( - dataset=bench_dataset, - indexes=config.indexes, - force=force, - dry_run=dry_run, - ) - if not result.success: - return result + if build: + result = backend.build( + dataset=bench_dataset, + indexes=config.indexes, + force=force, + dry_run=dry_run, + ) + if not result.success: + return result - if search: - result = backend.search( - dataset=bench_dataset, - indexes=config.indexes, - k=count, - batch_size=batch_size, - mode=search_mode, - force=force, - search_threads=search_threads, - dry_run=dry_run, - ) + if search: + result = backend.search( + dataset=bench_dataset, + indexes=config.indexes, + k=count, + batch_size=batch_size, + mode=search_mode, + force=force, + search_threads=search_threads, + dry_run=dry_run, + ) + finally: + backend.cleanup() return result diff --git a/python/cuvs_bench/cuvs_bench/tests/integration/README.md b/python/cuvs_bench/cuvs_bench/tests/integration/README.md deleted file mode 100644 index 3f73ce1081..0000000000 --- a/python/cuvs_bench/cuvs_bench/tests/integration/README.md +++ /dev/null @@ -1,34 +0,0 @@ -# Integration Tests - -Integration tests run against real services (e.g. Elasticsearch in Docker). - -**Note:** The Elasticsearch integration test is currently disabled. It targets the -Elasticsearch GPU backend (cuVS-accelerated), which requires an ES GPU image, -cuVS libs from this repo, and a GPU-enabled runner. See `test_elastic_integration.py` -module docstring for details. Can be re-enabled when CI has these dependencies. - -## Requirements - -- **Docker** running locally -- Optional dependencies: - - ```bash - pip install cuvs-bench[elastic,integration] - ``` - - Or install separately: - - ```bash - pip install cuvs-bench[elastic] - pip install testcontainers[elasticsearch] - ``` - -## Running - -```bash -# From repo root -PYTHONPATH="python/cuvs_bench:python/cuvs_bench_elastic:$PYTHONPATH" \ - pytest python/cuvs_bench/cuvs_bench/tests/integration/ -v -``` - -Integration tests are skipped automatically if `testcontainers` or `elasticsearch` is not installed. diff --git a/python/cuvs_bench/cuvs_bench/tests/integration/__init__.py b/python/cuvs_bench/cuvs_bench/tests/integration/__init__.py deleted file mode 100644 index 7366b06538..0000000000 --- a/python/cuvs_bench/cuvs_bench/tests/integration/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. -# SPDX-License-Identifier: Apache-2.0 -# -"""Integration tests for cuvs-bench backends. - -These tests require Docker and optional dependencies: - pip install cuvs-bench[elastic,integration] -""" diff --git a/python/cuvs_bench/cuvs_bench/tests/integration/conftest.py b/python/cuvs_bench/cuvs_bench/tests/integration/conftest.py deleted file mode 100644 index ee103561b9..0000000000 --- a/python/cuvs_bench/cuvs_bench/tests/integration/conftest.py +++ /dev/null @@ -1,73 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. -# SPDX-License-Identifier: Apache-2.0 -# -"""Pytest fixtures for integration tests. - -Requires: pip install cuvs-bench[elastic,integration] -Requires: Docker running locally -""" - -import pytest - - -def _testcontainers_available(): - try: - from testcontainers.elasticsearch import ElasticSearchContainer - return True - except ImportError: - return False - - -def _elasticsearch_installed(): - try: - import elasticsearch # noqa: F401 - return True - except ImportError: - return False - - -@pytest.fixture(scope="module") -def elasticsearch_container(): - """Start an Elasticsearch container for the duration of the test module. - - Yields a dict with host, port, and get_url() for connecting. - Skips the test module if testcontainers or elasticsearch is not installed. - """ - if not _testcontainers_available(): - pytest.skip( - "Requires testcontainers[elasticsearch]. " - "Install with: pip install cuvs-bench[integration]" - ) - if not _elasticsearch_installed(): - pytest.skip( - "Requires elasticsearch. Install with: pip install cuvs-bench[elastic]" - ) - - from testcontainers.elasticsearch import ElasticSearchContainer - - # Use standard ES OSS image (no GPU in OSS; use use_gpu=False in tests) - with ElasticSearchContainer( - image="elasticsearch:8.15.0", - mem_limit="2g", - ) as container: - url = container.get_url() - # Parse host:port from URL (e.g. http://localhost:32768) - if url.startswith("http://"): - rest = url[7:] - elif url.startswith("https://"): - rest = url[8:] - else: - rest = url - if "/" in rest: - host_port = rest.split("/")[0] - else: - host_port = rest - if ":" in host_port: - host, port_str = host_port.rsplit(":", 1) - port = int(port_str) - else: - host = host_port - port = 9200 - - yield {"host": host, "port": port, "url": url} diff --git a/python/cuvs_bench/cuvs_bench/tests/integration/test_elastic_integration.py b/python/cuvs_bench/cuvs_bench/tests/integration/test_elastic_integration.py deleted file mode 100644 index a2fffa2338..0000000000 --- a/python/cuvs_bench/cuvs_bench/tests/integration/test_elastic_integration.py +++ /dev/null @@ -1,141 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. -# SPDX-License-Identifier: Apache-2.0 -# -"""Integration tests for Elasticsearch backend using testcontainers. - -**Currently disabled.** The cuvs-bench elastic backend targets Elasticsearch GPU -(cuVS-accelerated vector search). A proper integration test would require: -- Elasticsearch GPU image (docker.elastic.co/elasticsearch-dev/elasticsearch-gpu:9.3.x) -- cuVS libraries from this repo (built libcuvs.so) mounted into the container -- GPU-enabled CI runner (--gpus all) -- LD_LIBRARY_PATH and volume mounts configured - -The standard ES OSS image has no GPU support, so tests with use_gpu=False only -validate the HTTP API, not the GPU indexing path. This can be re-enabled later -when the above dependencies are available in CI. - -Requires: pip install cuvs-bench[elastic,integration] -Requires: Docker running locally - -Run with: - pytest python/cuvs_bench/cuvs_bench/tests/integration/ -v -""" - -import tempfile -from pathlib import Path - -import numpy as np -import pytest - -from cuvs_bench.backends.base import Dataset -from cuvs_bench.backends.registry import get_backend_class -from cuvs_bench.orchestrator.config_loaders import IndexConfig - - -def _write_fbin(path: Path, data: np.ndarray) -> None: - """Write big-ann-bench fbin format.""" - with open(path, "wb") as f: - np.array(data.shape, dtype=np.uint32).tofile(f) - data.astype(np.float32).tofile(f) - - -def _write_ibin(path: Path, data: np.ndarray) -> None: - """Write big-ann-bench ibin format.""" - with open(path, "wb") as f: - np.array(data.shape, dtype=np.uint32).tofile(f) - data.astype(np.int32).tofile(f) - - -@pytest.mark.skip( - reason="Integration test disabled: requires ES GPU image + cuVS libs + GPU runner. " - "See module docstring. Can be re-enabled when CI has these dependencies." -) -class TestElasticIntegration: - """Integration tests against a real Elasticsearch container. - - Disabled until CI can provide: Elasticsearch GPU image, cuVS libraries - (from this repo build), and GPU-enabled runner. Standard ES OSS does not - exercise the GPU vector search path this backend targets. - """ - - def test_elastic_build_and_search_e2e(self, elasticsearch_container): - """Build index and run search against Elasticsearch container.""" - es = elasticsearch_container - cls = get_backend_class("elastic") - backend = cls( - config={ - "name": "integration_test", - "host": es["host"], - "port": es["port"], - "index_name": "cuvs_bench_integration_test", - "use_gpu": False, # Standard ES OSS has no GPU - } - ) - try: - # Create temp dir with fbin/ibin files - with tempfile.TemporaryDirectory() as tmpdir: - base_path = Path(tmpdir) / "base.fbin" - query_path = Path(tmpdir) / "query.fbin" - gt_path = Path(tmpdir) / "groundtruth_neighbors.ibin" - - n_base, dims = 200, 32 - n_queries, k = 20, 10 - - base = np.random.rand(n_base, dims).astype(np.float32) - queries = np.random.rand(n_queries, dims).astype(np.float32) - # Ground truth: first k neighbors per query (dummy IDs) - gt = np.random.randint(0, n_base, size=(n_queries, k), dtype=np.int32) - - _write_fbin(base_path, base) - _write_fbin(query_path, queries) - _write_ibin(gt_path, gt) - - backend.config["data_prefix"] = str(tmpdir) - - dataset = Dataset( - name="integration_test", - base_vectors=base, - query_vectors=queries, - base_file="base.fbin", - query_file="query.fbin", - groundtruth_neighbors_file="groundtruth_neighbors.ibin", - distance_metric="euclidean", - ) - - indexes = [ - IndexConfig( - name="elastic_hnsw_integration", - algo="elastic_hnsw", - build_param={ - "m": 16, - "ef_construction": 64, - "use_gpu": False, - }, - search_params=[{"num_candidates": 50}], - file="", - ) - ] - - # Build - build_result = backend.build( - dataset=dataset, - indexes=indexes, - force=True, - dry_run=False, - ) - assert build_result.success, build_result.error_message - assert build_result.index_size_bytes > 0 - - # Search (uses data_prefix from config set above) - search_result = backend.search( - dataset=dataset, - indexes=indexes, - k=k, - dry_run=False, - ) - assert search_result.success, search_result.error_message - assert search_result.neighbors.shape == (n_queries, k) - assert search_result.search_time_ms >= 0 - finally: - backend.cleanup() diff --git a/python/cuvs_bench/cuvs_bench/tests/test_modularization.py b/python/cuvs_bench/cuvs_bench/tests/test_modularization.py index a5ddf127a1..112edb4bd4 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_modularization.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_modularization.py @@ -324,7 +324,7 @@ def test_elastic_config_loader_tune_mode_returns_single_config(self): assert config.indexes[0].build_param["ef_construction"] == 150 assert config.indexes[0].build_param["type"] == "hnsw" assert config.indexes[0].search_params[0]["num_candidates"] == 120 - assert config.backend_config["name"] == "elastic_hnsw_tune" + assert config.backend_config["name"].startswith("elastic_hnsw_tune") def test_elastic_config_loader_sweep_mode_returns_multiple_configs(self): """Sweep mode produces multiple BenchmarkConfigs from param cartesian product.""" @@ -449,8 +449,8 @@ def test_elastic_build_requires_base_file(self): assert not result.success assert "base_file" in (result.error_message or "").lower() - def test_elastic_preflight_skip_when_no_network(self): - """ElasticBackend skips build when network check fails.""" + def test_elastic_preflight_fails_when_no_network(self): + """ElasticBackend.build returns success=False when network is unavailable.""" cls = get_backend_class("elastic") backend = cls(config={"name": "test", "host": "localhost", "port": 9200}) @@ -478,9 +478,77 @@ def test_elastic_preflight_skip_when_no_network(self): ): result = backend.build(dataset=dataset, indexes=indexes) + assert not result.success + assert "pre-flight" in (result.error_message or "").lower() + + def test_elastic_search_preflight_fails_when_no_network(self): + """ElasticBackend.search returns success=False when network is unavailable.""" + cls = get_backend_class("elastic") + backend = cls(config={"name": "test", "host": "localhost", "port": 9200}) + + dataset = Dataset( + name="test", + base_vectors=np.random.rand(100, 32).astype(np.float32), + query_vectors=np.random.rand(10, 32).astype(np.float32), + query_file="dummy/query.fbin", + ) + indexes = [ + IndexConfig( + name="elastic_hnsw_test", + algo="elastic_hnsw", + build_param={}, + search_params=[{"num_candidates": 100}], + file="", + ) + ] + + with patch.object(backend, "_check_network_available", return_value=False): + result = backend.search(dataset=dataset, indexes=indexes, k=10) + + assert not result.success + assert "pre-flight" in (result.error_message or "").lower() + + def test_elastic_build_skips_existing_index_when_force_false(self): + """build(force=False) returns success=True immediately when index already exists.""" + cls = get_backend_class("elastic") + backend = cls(config={ + "name": "test", + "host": "localhost", + "port": 9200, + "index_name": "test_index", + }) + + mock_client = MagicMock() + mock_client.indices.exists.return_value = True + mock_client.indices.stats.return_value = { + "_all": {"primaries": {"store": {"size_in_bytes": 1024}}} + } + + dataset = Dataset( + name="test", + base_vectors=np.random.rand(100, 32).astype(np.float32), + query_vectors=np.random.rand(10, 32).astype(np.float32), + base_file="dummy/base.fbin", + ) + indexes = [ + IndexConfig( + name="test_index", + algo="elastic_hnsw", + build_param={"type": "hnsw", "m": 16, "ef_construction": 100, + "similarity": "l2_norm", "number_of_shards": 1, + "number_of_replicas": 0, "vector_field": "embedding"}, + search_params=[{"num_candidates": 100}], + file="", + ) + ] + + with patch.object(backend, "_check_network_available", return_value=True): + with patch.object(backend, "_get_client", return_value=mock_client): + result = backend.build(dataset=dataset, indexes=indexes, force=False) + assert result.success - assert result.metadata.get("skipped") is True - assert result.metadata.get("reason") == "no_network" + assert result.index_size_bytes == 1024 + mock_client.indices.delete.assert_not_called() def test_elastic_algo_from_config(self): """ElasticBackend.algo derives from config type (elastic_hnsw, elastic_int8_hnsw).""" diff --git a/python/cuvs_bench/pyproject.toml b/python/cuvs_bench/pyproject.toml index 2c607493ee..7513443497 100644 --- a/python/cuvs_bench/pyproject.toml +++ b/python/cuvs_bench/pyproject.toml @@ -42,8 +42,17 @@ classifiers = [ Homepage = "https://github.com/rapidsai/cuvs" [project.optional-dependencies] -elastic = ["cuvs-bench-elastic>=26.4.0"] -integration = ["cuvs-bench-elastic>=26.4.0", "testcontainers[elasticsearch]"] +elastic = ["elasticsearch>=8.0"] + +[project.entry-points."cuvs_bench.backends"] +elastic = "cuvs_bench_elastic:register" + +[project.entry-points."cuvs_bench.config_loaders"] +elastic = "cuvs_bench_elastic:register" + +[tool.setuptools.packages.find] +where = [".", "../cuvs_bench_elastic"] +include = ["cuvs_bench*"] [tool.setuptools.package-data] "*" = ["*.*", "VERSION"] diff --git a/python/cuvs_bench_elastic/cuvs_bench_elastic/backend.py b/python/cuvs_bench_elastic/cuvs_bench_elastic/backend.py index 4980b24101..311af08bf4 100644 --- a/python/cuvs_bench_elastic/cuvs_bench_elastic/backend.py +++ b/python/cuvs_bench_elastic/cuvs_bench_elastic/backend.py @@ -10,7 +10,8 @@ Build params (index_options): type, m, ef_construction. type: hnsw, int8_hnsw, int4_hnsw, bbq_hnsw (per ES-GPU-API-REFERENCE.md) similarity: l2_norm, cosine, max_inner_product (overrides dataset distance) -Index settings: number_of_shards, number_of_replicas, use_gpu, vector_field. +Index settings: number_of_shards, number_of_replicas, vector_field. + GPU indexing is configured at the node level (vectors.indexing.use_gpu in elasticsearch.yml). Search params (knn): num_candidates, vector_field. """ @@ -23,8 +24,6 @@ import numpy as np import yaml -from elasticsearch import Elasticsearch -from elasticsearch.helpers import bulk from cuvs_bench.backends.base import BenchmarkBackend, BuildResult, Dataset, SearchResult from cuvs_bench.backends.registry import register_backend, register_config_loader @@ -64,11 +63,10 @@ def _distance_to_similarity(distance: str) -> str: _DEFAULT_EF_CONSTRUCTION = 100 _DEFAULT_NUM_SHARDS = 1 _DEFAULT_NUM_REPLICAS = 0 -_DEFAULT_USE_GPU = True + _DEFAULT_VECTOR_FIELD = "embedding" _DEFAULT_NUM_CANDIDATES = 100 -# Build params passed through from config (index_options + index settings + similarity) _BUILD_PARAM_KEYS = ( "type", "m", @@ -76,10 +74,8 @@ def _distance_to_similarity(distance: str) -> str: "similarity", "number_of_shards", "number_of_replicas", - "use_gpu", "vector_field", ) -# Search params passed through from config (knn) _SEARCH_PARAM_KEYS = ("num_candidates", "vector_field") @@ -88,7 +84,7 @@ class ElasticBackend(BenchmarkBackend): def __init__(self, config: Dict[str, Any]): super().__init__(config) - self._client: Optional[Elasticsearch] = None + self._client: Optional["Elasticsearch"] = None @property def algo(self) -> str: @@ -96,11 +92,22 @@ def algo(self) -> str: index_type = self.config.get("type", _DEFAULT_INDEX_TYPE) return f"elastic_{index_type}" - def _get_client(self) -> Elasticsearch: + def _get_client(self) -> "Elasticsearch": if self._client is None: + try: + from elasticsearch import Elasticsearch + except ImportError as e: + raise ImportError( + "`elasticsearch` is required for the Elasticsearch backend. " + "Install with: pip install cuvs-bench[elastic]" + ) from e host = self.config.get("host", "localhost") port = self.config.get("port", 9200) - kwargs: Dict[str, Any] = {"hosts": [f"http://{host}:{port}"]} + scheme = self.config.get("scheme", "http") + kwargs: Dict[str, Any] = { + "hosts": [f"{scheme}://{host}:{port}"], + "request_timeout": 300, # 5 min for slow bulk ops / remote ES + } basic_auth = self.config.get("basic_auth") if basic_auth is not None: if isinstance(basic_auth, (list, tuple)) and len(basic_auth) >= 2: @@ -149,19 +156,39 @@ def build( skip_reason = self._pre_flight_check() if skip_reason: - idx = indexes[0] if indexes else None return BuildResult( - index_path=idx.file if idx else "", - build_time_seconds=0, + index_path="", + build_time_seconds=0.0, index_size_bytes=0, algorithm=self.algo, - build_params=dict(idx.build_param or {}) if idx else {}, - success=True, - metadata={"skipped": True, "reason": skip_reason}, + build_params={}, + success=False, + error_message=f"pre-flight check failed: {skip_reason}", ) - base_file = dataset.base_file - if not base_file: + index_name = self.config.get("index_name", "cuvs_bench_vectors") + idx = indexes[0] if indexes else None + build_params = dict(idx.build_param or {}) if idx else {} + for k, v in self.config.items(): + if k not in build_params and k in _BUILD_PARAM_KEYS: + build_params[k] = v + + try: + client = self._get_client() + if client.indices.exists(index=index_name): + if not force: + stats = client.indices.stats(index=index_name) + index_size = stats["_all"]["primaries"]["store"]["size_in_bytes"] + return BuildResult( + index_path=index_name, + build_time_seconds=0, + index_size_bytes=index_size, + algorithm=self.algo, + build_params=build_params, + success=True, + ) + client.indices.delete(index=index_name) + except Exception as e: return BuildResult( index_path="", build_time_seconds=0, @@ -169,40 +196,32 @@ def build( algorithm=self.algo, build_params={}, success=False, - error_message="base_file is required for Elasticsearch backend", + error_message=str(e), ) - data_prefix = self.config.get("data_prefix", "") - base_path = Path(data_prefix) / base_file - if not base_path.exists(): + if not dataset.base_file: return BuildResult( index_path="", - build_time_seconds=0, + build_time_seconds=0.0, index_size_bytes=0, algorithm=self.algo, build_params={}, success=False, - error_message=f"Base file not found: {base_path}", + error_message="base_file is required for Elasticsearch backend", ) - index_name = self.config.get("index_name", "cuvs_bench_vectors") - dims = dataset.dims or self.config.get("dims") - if not dims: + base_path = Path(dataset.base_file) + if not base_path.exists(): return BuildResult( index_path="", - build_time_seconds=0, + build_time_seconds=0.0, index_size_bytes=0, algorithm=self.algo, build_params={}, success=False, - error_message="dims is required", + error_message=f"Base file not found: {base_path}", ) - idx = indexes[0] if indexes else None - build_params = dict(idx.build_param or {}) if idx else {} - for k, v in self.config.items(): - if k not in build_params and k in _BUILD_PARAM_KEYS: - build_params[k] = v # similarity: from config, or derive from dataset distance similarity = build_params.get("similarity") or _distance_to_similarity( getattr(dataset, "distance_metric", None) or "euclidean" @@ -211,6 +230,7 @@ def build( try: vectors = _load_fbin(base_path) n_vectors = len(vectors) + dims = vectors.shape[1] client = self._get_client() vector_field = build_params.get("vector_field", _DEFAULT_VECTOR_FIELD) @@ -223,14 +243,12 @@ def build( num_replicas = build_params.get( "number_of_replicas", _DEFAULT_NUM_REPLICAS ) - use_gpu = build_params.get("use_gpu", _DEFAULT_USE_GPU) - settings: Dict[str, Any] = { "number_of_shards": num_shards, "number_of_replicas": num_replicas, } - if use_gpu: - settings["index"] = {"vectors.indexing.use_gpu": True} + # Note: GPU indexing is controlled at the node level via + # vectors.indexing.use_gpu in elasticsearch.yml, not per-index. index_options: Dict[str, Any] = { "type": index_type, @@ -254,22 +272,11 @@ def build( } t0 = time.perf_counter() - if client.indices.exists(index=index_name): - if not force: - stats = client.indices.stats(index=index_name) - index_size = stats["_all"]["primaries"]["store"]["size_in_bytes"] - return BuildResult( - index_path=index_name, - build_time_seconds=0, - index_size_bytes=index_size, - algorithm=self.algo, - build_params=build_params, - success=True, - ) - client.indices.delete(index=index_name) client.indices.create(index=index_name, body=index_config) + from elasticsearch.helpers import bulk chunk_size = 1000 + progress_interval = max(50, n_vectors // (chunk_size * 20)) # ~20 progress lines for i in range(0, n_vectors, chunk_size): chunk = vectors[i : i + chunk_size] actions = [ @@ -281,6 +288,8 @@ def build( for j, vec in enumerate(chunk) ] bulk(client, actions, raise_on_error=True) + if progress_interval and (i // chunk_size) % progress_interval == 0: + print(f" Indexed {min(i + chunk_size, n_vectors):,}/{n_vectors:,} vectors") client.indices.refresh(index=index_name) build_time = time.perf_counter() - t0 @@ -318,11 +327,11 @@ def search( search_threads: Optional[int] = None, dry_run: bool = False, ) -> SearchResult: - """Run kNN search and compute recall.""" + """Run kNN search over all search-param combinations and compute recall.""" if dry_run: return SearchResult( - neighbors=np.empty((0, k)), - distances=np.empty((0, k)), + neighbors=np.zeros((0, k), dtype=np.int64), + distances=np.zeros((0, k), dtype=np.float32), search_time_ms=0, queries_per_second=0, recall=0, @@ -333,27 +342,35 @@ def search( skip_reason = self._pre_flight_check() if skip_reason: - idx = indexes[0] if indexes else None return SearchResult( - neighbors=np.empty((0, k)), - distances=np.empty((0, k)), + neighbors=np.zeros((0, k), dtype=np.int64), + distances=np.zeros((0, k), dtype=np.float32), search_time_ms=0, queries_per_second=0, recall=0, algorithm=self.algo, - search_params=idx.search_params if idx and idx.search_params else [], - success=True, - metadata={"skipped": True, "reason": skip_reason}, + search_params=[], + success=False, + error_message=f"pre-flight check failed: {skip_reason}", ) - query_file = dataset.query_file - gt_file = dataset.groundtruth_neighbors_file - data_prefix = self.config.get("data_prefix", "") + if not indexes: + return SearchResult( + neighbors=np.zeros((0, k), dtype=np.int64), + distances=np.zeros((0, k), dtype=np.float32), + search_time_ms=0, + queries_per_second=0, + recall=0, + algorithm=self.algo, + search_params=[], + success=False, + error_message="No indexes provided", + ) - if not query_file: + if not dataset.query_file: return SearchResult( - neighbors=np.empty((0, k)), - distances=np.empty((0, k)), + neighbors=np.zeros((0, k), dtype=np.int64), + distances=np.zeros((0, k), dtype=np.float32), search_time_ms=0, queries_per_second=0, recall=0, @@ -363,11 +380,11 @@ def search( error_message="query_file is required", ) - query_path = Path(data_prefix) / query_file + query_path = Path(dataset.query_file) if not query_path.exists(): return SearchResult( - neighbors=np.empty((0, k)), - distances=np.empty((0, k)), + neighbors=np.zeros((0, k), dtype=np.int64), + distances=np.zeros((0, k), dtype=np.float32), search_time_ms=0, queries_per_second=0, recall=0, @@ -380,89 +397,95 @@ def search( try: query_vectors = _load_fbin(query_path) n_queries = len(query_vectors) - groundtruth = None - if gt_file: - gt_path = Path(data_prefix) / gt_file + + groundtruth = dataset.groundtruth_neighbors + if groundtruth is None and dataset.groundtruth_neighbors_file: + gt_path = Path(dataset.groundtruth_neighbors_file) if gt_path.exists(): groundtruth = _load_ibin(gt_path) index_name = self.config.get("index_name", "cuvs_bench_vectors") - search_params = {} - if indexes and indexes[0] and indexes[0].search_params: - search_params = dict(indexes[0].search_params[0] or {}) - for key, val in self.config.items(): - if key not in search_params and key in _SEARCH_PARAM_KEYS: - search_params[key] = val - num_candidates = search_params.get( - "num_candidates", _DEFAULT_NUM_CANDIDATES - ) - vector_field = search_params.get( - "vector_field", _DEFAULT_VECTOR_FIELD - ) + index_cfg = indexes[0] + search_params_list = index_cfg.search_params or [{}] - neighbors_list: List[List[int]] = [] - latencies: List[float] = [] + per_param_results: List[Dict[str, Any]] = [] + last_neighbors = np.full((n_queries, k), -1, dtype=np.int64) + last_distances = np.zeros((n_queries, k), dtype=np.float32) + + for sp in search_params_list: + num_candidates = sp.get("num_candidates", _DEFAULT_NUM_CANDIDATES) + vector_field = sp.get("vector_field", _DEFAULT_VECTOR_FIELD) + + neighbors = np.full((n_queries, k), -1, dtype=np.int64) + distances = np.zeros((n_queries, k), dtype=np.float32) + latencies: List[float] = [] - for i, qv in enumerate(query_vectors): - body = { - "knn": { - "field": vector_field, - "query_vector": qv.tolist(), - "k": k, - "num_candidates": num_candidates, - } - } t0 = time.perf_counter() - resp = self._get_client().search( - index=index_name, body=body, size=k - ) - latencies.append((time.perf_counter() - t0) * 1000) - hits = resp.get("hits", {}).get("hits", []) - ids = [int(h["_id"]) for h in hits if "_id" in h] - neighbors_list.append(ids[:k]) - - search_time_ms = sum(latencies) - qps = n_queries / (search_time_ms / 1000) if search_time_ms > 0 else 0 - - recall = 0.0 - if groundtruth is not None and len(neighbors_list) <= len(groundtruth): - correct = 0 - total = 0 - for q in range(len(neighbors_list)): - retrieved = set(neighbors_list[q]) - gt_row = groundtruth[q, :k] - for gt_id in gt_row: - if int(gt_id) in retrieved: - correct += 1 - total += 1 - recall = correct / total if total > 0 else 0 - - max_len = max(len(r) for r in neighbors_list) - neighbors_arr = np.array( - [r + [-1] * (max_len - len(r)) for r in neighbors_list], - dtype=np.int64, - ) - distances_arr = np.zeros_like(neighbors_arr, dtype=np.float32) + for i, qv in enumerate(query_vectors): + body = { + "knn": { + "field": vector_field, + "query_vector": qv.tolist(), + "k": k, + "num_candidates": num_candidates, + } + } + t_q = time.perf_counter() + resp = self._get_client().search(index=index_name, body=body, size=k) + latencies.append((time.perf_counter() - t_q) * 1000) + hits = resp.get("hits", {}).get("hits", []) + for j, hit in enumerate(hits[:k]): + neighbors[i, j] = int(hit["_id"]) + distances[i, j] = float(hit["_score"]) + + elapsed_ms = (time.perf_counter() - t0) * 1000 + qps = n_queries / (elapsed_ms / 1000) if elapsed_ms > 0 else 0.0 + + recall = 0.0 + if groundtruth is not None: + gt_k = min(k, groundtruth.shape[1]) + n_correct = sum( + len(set(neighbors[i, :k].tolist()) & set(groundtruth[i, :gt_k].tolist())) + for i in range(n_queries) + ) + recall = n_correct / (n_queries * gt_k) if gt_k > 0 else 0.0 - return SearchResult( - neighbors=neighbors_arr, - distances=distances_arr, - search_time_ms=search_time_ms, - queries_per_second=qps, - recall=recall, - algorithm=self.algo, - search_params=[search_params], - latency_percentiles={ + per_param_results.append({ + "search_params": sp, + "search_time_ms": elapsed_ms, + "queries_per_second": qps, + "recall": recall, "p50_ms": float(np.percentile(latencies, 50)), "p95_ms": float(np.percentile(latencies, 95)), "p99_ms": float(np.percentile(latencies, 99)), + }) + last_neighbors = neighbors + last_distances = distances + + avg_recall = float(np.mean([r["recall"] for r in per_param_results])) + avg_qps = float(np.mean([r["queries_per_second"] for r in per_param_results])) + total_ms = float(sum(r["search_time_ms"] for r in per_param_results)) + + return SearchResult( + neighbors=last_neighbors, + distances=last_distances, + search_time_ms=total_ms, + queries_per_second=avg_qps, + recall=avg_recall, + algorithm=self.algo, + search_params=search_params_list, + latency_percentiles={ + "p50_ms": per_param_results[-1]["p50_ms"], + "p95_ms": per_param_results[-1]["p95_ms"], + "p99_ms": per_param_results[-1]["p99_ms"], }, + metadata={"per_search_param_results": per_param_results}, success=True, ) except Exception as e: return SearchResult( - neighbors=np.empty((0, k)), - distances=np.empty((0, k)), + neighbors=np.zeros((0, k), dtype=np.int64), + distances=np.zeros((0, k), dtype=np.float32), search_time_ms=0, queries_per_second=0, recall=0, @@ -528,6 +551,7 @@ def load( tune_search_params = kwargs.pop("_tune_search_params", None) username = kwargs.pop("username", None) password = kwargs.pop("password", None) + scheme = kwargs.pop("scheme", "http") if basic_auth is None and username and password: basic_auth = (username, password) @@ -546,12 +570,17 @@ def load( if not dataset_conf: raise ValueError(f"Dataset '{dataset}' not found") + def _resolve(rel: Optional[str]) -> Optional[str]: + if rel and not os.path.isabs(rel): + return os.path.join(dataset_path, rel) + return rel + dataset_config = DatasetConfig( name=dataset_conf["name"], - base_file=dataset_conf.get("base_file"), - query_file=dataset_conf.get("query_file"), - groundtruth_neighbors_file=dataset_conf.get( - "groundtruth_neighbors_file" + base_file=_resolve(dataset_conf.get("base_file")), + query_file=_resolve(dataset_conf.get("query_file")), + groundtruth_neighbors_file=_resolve( + dataset_conf.get("groundtruth_neighbors_file") ), distance=dataset_conf.get("distance", "euclidean"), dims=dataset_conf.get("dims"), @@ -560,12 +589,14 @@ def load( if tune_mode and tune_build_params is not None and tune_search_params is not None: algo_name = algorithms or "elastic_hnsw" - # Extract type from algo (elastic_hnsw -> hnsw, elastic_int8_hnsw -> int8_hnsw) build_param = dict(tune_build_params) if "type" not in build_param and algo_name.startswith("elastic_"): build_param["type"] = algo_name.replace("elastic_", "", 1) + name_parts = [f"{k}{v}" for k, v in build_param.items() if k in ("m", "ef_construction")] + index_label = "_".join([f"{algo_name}_tune"] + name_parts) if name_parts else f"{algo_name}_tune" + es_index_name = index_label.lower().replace(".", "_") index_config = IndexConfig( - name=f"{algo_name}_tune", + name=index_label, algo=algo_name, build_param=build_param, search_params=[tune_search_params], @@ -574,15 +605,14 @@ def load( config = BenchmarkConfig( indexes=[index_config], backend_config={ - "name": index_config.name, + "name": index_label, "host": host, "port": port, - "index_name": index_name, + "scheme": scheme, + "index_name": es_index_name, "basic_auth": basic_auth, - "data_prefix": dataset_path, - "dims": dataset_config.dims, + **build_param, - **tune_search_params, }, ) return dataset_config, [config] @@ -625,39 +655,55 @@ def _to_list_values(d: Dict[str, Any]) -> Dict[str, List[Any]]: build_keys = list(build_params.keys()) search_keys = list(search_params.keys()) + search_params_list = [ + dict(zip(search_keys, svals)) for svals in search_combos + ] + benchmark_configs = [] for bvals in build_combos: bdict = dict(zip(build_keys, bvals)) - for svals in search_combos: - sdict = dict(zip(search_keys, svals)) - algo_name = f"elastic_{bdict.get('type', 'hnsw')}" - index_config = IndexConfig( - name=f"{algo_name}_{group_name}", - algo=algo_name, - build_param=bdict, - search_params=[sdict], - file="", - ) - config = BenchmarkConfig( - indexes=[index_config], - backend_config={ - "name": index_config.name, - "host": host, - "port": port, - "index_name": index_name, - "basic_auth": basic_auth, - "data_prefix": dataset_path, - "dims": dataset_config.dims, - **bdict, - **sdict, - }, - ) - benchmark_configs.append(config) + algo_name = f"elastic_{bdict.get('type', 'hnsw')}" + + # Derive a unique, human-readable index name from build params + prefix = f"{algo_name}_{group_name}" if group_name != "base" else algo_name + name_parts = [f"{k}{v}" for k, v in bdict.items() if k in ("m", "ef_construction")] + index_label = "_".join([prefix] + name_parts) if name_parts else prefix + es_index_name = index_label.lower().replace(".", "_") + + index_config = IndexConfig( + name=index_label, + algo=algo_name, + build_param=bdict, + search_params=search_params_list, + file="", + ) + config = BenchmarkConfig( + indexes=[index_config], + backend_config={ + "name": index_label, + "host": host, + "port": port, + "scheme": scheme, + "index_name": es_index_name, + "basic_auth": basic_auth, + + **bdict, + }, + ) + benchmark_configs.append(config) return dataset_config, benchmark_configs def register() -> None: - """Register Elasticsearch backend and config loader (called from entry point).""" - register_backend("elastic", ElasticBackend) - register_config_loader("elastic", ElasticConfigLoader) + """Register Elasticsearch backend and config loader (idempotent).""" + from cuvs_bench.backends.registry import ( + _CONFIG_LOADER_REGISTRY, + get_registry, + ) + + reg = get_registry() + if not reg.is_registered("elastic"): + register_backend("elastic", ElasticBackend) + if "elastic" not in _CONFIG_LOADER_REGISTRY: + register_config_loader("elastic", ElasticConfigLoader) diff --git a/python/cuvs_bench_elastic/cuvs_bench_elastic/config/algos/elastic.yaml b/python/cuvs_bench_elastic/cuvs_bench_elastic/config/algos/elastic.yaml index 054c86dfa8..daa4fa22ab 100644 --- a/python/cuvs_bench_elastic/cuvs_bench_elastic/config/algos/elastic.yaml +++ b/python/cuvs_bench_elastic/cuvs_bench_elastic/config/algos/elastic.yaml @@ -19,7 +19,7 @@ groups: similarity: [l2_norm, cosine] number_of_shards: [1] number_of_replicas: [0] - use_gpu: [true] + vector_field: [embedding] search: num_candidates: [100, 200] @@ -32,13 +32,12 @@ groups: similarity: [l2_norm] number_of_shards: [1] number_of_replicas: [0] - use_gpu: [true] + vector_field: [embedding] search: num_candidates: [100, 200] vector_field: [embedding] test: - # use_gpu: false for ES source build (N/N path) - GPU enabled at cluster level only build: type: [hnsw] m: [16] @@ -46,7 +45,7 @@ groups: similarity: [l2_norm] number_of_shards: [1] number_of_replicas: [0] - use_gpu: [false] + vector_field: [embedding] search: num_candidates: [100] diff --git a/python/cuvs_bench_elastic/pyproject.toml b/python/cuvs_bench_elastic/pyproject.toml index e62d118732..3a6a6e6a95 100644 --- a/python/cuvs_bench_elastic/pyproject.toml +++ b/python/cuvs_bench_elastic/pyproject.toml @@ -25,11 +25,15 @@ classifiers = [ ] dependencies = [ "cuvs-bench>=26.4.0", - "elasticsearch>=8.0", "numpy", "pyyaml", ] +[project.optional-dependencies] +# elasticsearch client; automatically installed via `pip install cuvs-bench[elastic]`. +# Only needed here for standalone dev installs of this package directly. +client = ["elasticsearch>=8.0"] + [project.urls] Homepage = "https://github.com/rapidsai/cuvs" From 087289d812ead9480e7b6fd9d3d1a8cc3374bf51 Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Sun, 12 Apr 2026 13:51:47 -0700 Subject: [PATCH 06/22] Revert unrelated change to get_dataset/__main__.py --- python/cuvs_bench/cuvs_bench/get_dataset/__main__.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/python/cuvs_bench/cuvs_bench/get_dataset/__main__.py b/python/cuvs_bench/cuvs_bench/get_dataset/__main__.py index 1670dd2097..3c52be00d8 100644 --- a/python/cuvs_bench/cuvs_bench/get_dataset/__main__.py +++ b/python/cuvs_bench/cuvs_bench/get_dataset/__main__.py @@ -4,7 +4,6 @@ import os import subprocess -import sys import click import h5py @@ -35,14 +34,13 @@ def convert_hdf5_to_fbin(path, normalize): scripts_path = os.path.dirname(os.path.realpath(__file__)) ann_bench_scripts_path = os.path.join(scripts_path, "hdf5_to_fbin.py") print(f"calling script {ann_bench_scripts_path}") - python = sys.executable if normalize and "angular" in path: subprocess.run( - [python, ann_bench_scripts_path, "-n", "%s" % path], check=True + ["python", ann_bench_scripts_path, "-n", "%s" % path], check=True ) else: subprocess.run( - [python, ann_bench_scripts_path, "%s" % path], check=True + ["python", ann_bench_scripts_path, "%s" % path], check=True ) From e8f8a646000c2f7f8c989a02bd13f9f1cb5b8f98 Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Mon, 13 Apr 2026 13:49:43 -0700 Subject: [PATCH 07/22] Fold Elasticsearch plugin back into cuvs_bench; wire via entry points The separate cuvs_bench_elastic package required bundling via packages.find and complicated the build. Instead, keep the backend inside cuvs_bench and use entry points pointing back into the same package so the backend only registers when elasticsearch is installed. - git mv backend.py to cuvs_bench/backends/elasticsearch.py - git mv elastic.yaml to cuvs_bench/config/algos/ - Fix imports to relative paths - Fix _get_elastic_config_path() to use ../config from backends/ - Update pyproject.toml: entry points -> cuvs_bench.backends.elasticsearch:register - Remove packages.find (no longer needed) - Remove cuvs_bench_elastic/ package entirely DX unchanged: pip install cuvs-bench[elastic] One package, one publish pipeline. Signed-off-by: Alex Fournier --- .../cuvs_bench/backends/elasticsearch.py} | 10 +- .../cuvs_bench}/config/algos/elastic.yaml | 0 python/cuvs_bench/pyproject.toml | 8 +- python/cuvs_bench_elastic/README.md | 84 ---------- .../cuvs_bench_elastic/VERSION | 1 - .../cuvs_bench_elastic/__init__.py | 154 ------------------ python/cuvs_bench_elastic/pyproject.toml | 54 ------ 7 files changed, 7 insertions(+), 304 deletions(-) rename python/{cuvs_bench_elastic/cuvs_bench_elastic/backend.py => cuvs_bench/cuvs_bench/backends/elasticsearch.py} (98%) rename python/{cuvs_bench_elastic/cuvs_bench_elastic => cuvs_bench/cuvs_bench}/config/algos/elastic.yaml (100%) delete mode 100644 python/cuvs_bench_elastic/README.md delete mode 100644 python/cuvs_bench_elastic/cuvs_bench_elastic/VERSION delete mode 100644 python/cuvs_bench_elastic/cuvs_bench_elastic/__init__.py delete mode 100644 python/cuvs_bench_elastic/pyproject.toml diff --git a/python/cuvs_bench_elastic/cuvs_bench_elastic/backend.py b/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py similarity index 98% rename from python/cuvs_bench_elastic/cuvs_bench_elastic/backend.py rename to python/cuvs_bench/cuvs_bench/backends/elasticsearch.py index 311af08bf4..7e3d315140 100644 --- a/python/cuvs_bench_elastic/cuvs_bench_elastic/backend.py +++ b/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py @@ -25,9 +25,9 @@ import yaml -from cuvs_bench.backends.base import BenchmarkBackend, BuildResult, Dataset, SearchResult -from cuvs_bench.backends.registry import register_backend, register_config_loader -from cuvs_bench.orchestrator.config_loaders import ( +from .base import BenchmarkBackend, BuildResult, Dataset, SearchResult +from .registry import register_backend, register_config_loader +from ..orchestrator.config_loaders import ( BenchmarkConfig, ConfigLoader, DatasetConfig, @@ -517,8 +517,8 @@ def _get_cuvs_bench_config_path() -> str: def _get_elastic_config_path() -> str: - """Get this package's config directory for elastic.yaml.""" - return os.path.join(os.path.dirname(__file__), "config") + """Get the config directory for elastic.yaml.""" + return os.path.join(os.path.dirname(__file__), "../config") class ElasticConfigLoader(ConfigLoader): diff --git a/python/cuvs_bench_elastic/cuvs_bench_elastic/config/algos/elastic.yaml b/python/cuvs_bench/cuvs_bench/config/algos/elastic.yaml similarity index 100% rename from python/cuvs_bench_elastic/cuvs_bench_elastic/config/algos/elastic.yaml rename to python/cuvs_bench/cuvs_bench/config/algos/elastic.yaml diff --git a/python/cuvs_bench/pyproject.toml b/python/cuvs_bench/pyproject.toml index 7513443497..dee9027f59 100644 --- a/python/cuvs_bench/pyproject.toml +++ b/python/cuvs_bench/pyproject.toml @@ -45,14 +45,10 @@ Homepage = "https://github.com/rapidsai/cuvs" elastic = ["elasticsearch>=8.0"] [project.entry-points."cuvs_bench.backends"] -elastic = "cuvs_bench_elastic:register" +elastic = "cuvs_bench.backends.elasticsearch:register" [project.entry-points."cuvs_bench.config_loaders"] -elastic = "cuvs_bench_elastic:register" - -[tool.setuptools.packages.find] -where = [".", "../cuvs_bench_elastic"] -include = ["cuvs_bench*"] +elastic = "cuvs_bench.backends.elasticsearch:register" [tool.setuptools.package-data] "*" = ["*.*", "VERSION"] diff --git a/python/cuvs_bench_elastic/README.md b/python/cuvs_bench_elastic/README.md deleted file mode 100644 index edd92adf11..0000000000 --- a/python/cuvs_bench_elastic/README.md +++ /dev/null @@ -1,84 +0,0 @@ -# cuvs-bench-elastic - -Elasticsearch GPU backend plugin for [cuvs-bench](https://github.com/rapidsai/cuvs). - -## Installation - -```bash -pip install cuvs-bench[elastic] -``` - -Or install standalone (requires `cuvs-bench`): - -```bash -pip install cuvs-bench-elastic -``` - -## Usage - -### Convenience API (recommended) - -```python -from cuvs_bench_elastic import run_build, run_search, run_benchmark - -# Build only -build_results = run_build( - dataset="test-data", - dataset_path="./datasets", - host="localhost", - port=9200, -) - -# Search only (requires existing index) -search_results = run_search( - dataset="test-data", - dataset_path="./datasets", - host="localhost", - port=9200, -) - -# Build and search -results = run_benchmark( - dataset="test-data", - dataset_path="./datasets", - host="localhost", - port=9200, -) - -# With authentication -results = run_benchmark( - dataset="test-data", - dataset_path="./datasets", - host="localhost", - port=9200, - username="elastic", - password="password", -) -``` - -### Orchestrator API - -```python -from cuvs_bench_elastic import ELASTIC, register -from cuvs_bench.orchestrator import BenchmarkOrchestrator - -register() -orch = BenchmarkOrchestrator(backend_type=ELASTIC) -results = orch.run_benchmark( - dataset="test-data", - dataset_path="./datasets", - host="localhost", - port=9200, - basic_auth="elastic:password", # or username="elastic", password="password" - algorithms="test", - build=True, - search=True, -) -``` - -## Configuration - -Build params (index_options): `type`, `m`, `ef_construction`. -Search params (knn): `num_candidates`, `vector_field`. - -See `config/algos/elastic.yaml` for parameter groups. diff --git a/python/cuvs_bench_elastic/cuvs_bench_elastic/VERSION b/python/cuvs_bench_elastic/cuvs_bench_elastic/VERSION deleted file mode 100644 index 0bd0e8a95b..0000000000 --- a/python/cuvs_bench_elastic/cuvs_bench_elastic/VERSION +++ /dev/null @@ -1 +0,0 @@ -26.04.00 diff --git a/python/cuvs_bench_elastic/cuvs_bench_elastic/__init__.py b/python/cuvs_bench_elastic/cuvs_bench_elastic/__init__.py deleted file mode 100644 index df3578d3d7..0000000000 --- a/python/cuvs_bench_elastic/cuvs_bench_elastic/__init__.py +++ /dev/null @@ -1,154 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. -# SPDX-License-Identifier: Apache-2.0 -# -""" -Elasticsearch GPU backend plugin for cuvs-bench. - -Install with: pip install cuvs-bench[elastic] - -Usage: - # Build and search as separate operations (matches backend interface) - from cuvs_bench_elastic import run_build, run_search - build_results = run_build(dataset="test-data", dataset_path="./datasets", host="localhost", port=9200) - search_results = run_search(dataset="test-data", dataset_path="./datasets", host="localhost", port=9200) - - # Or run both - from cuvs_bench_elastic import run_benchmark - results = run_benchmark(dataset="test-data", dataset_path="./datasets", host="localhost", port=9200) - - # Or use the orchestrator directly - from cuvs_bench_elastic import ELASTIC, register - from cuvs_bench.orchestrator import BenchmarkOrchestrator - register() - orch = BenchmarkOrchestrator(backend_type=ELASTIC) - results = orch.run_benchmark(dataset="test-data", dataset_path="./datasets", host="localhost", port=9200) -""" - -from .backend import register - -# Backend type for use with BenchmarkOrchestrator(backend_type=ELASTIC) -ELASTIC = "elastic" - - -def _common_kwargs( - dataset: str, - dataset_path: str, - host: str, - port: int, - algorithms: str, - force: bool, - username: str, - password: str, - **kwargs, -): - """Build kwargs shared by run_build, run_search, run_benchmark.""" - d = dict( - dataset=dataset, - dataset_path=dataset_path, - host=host, - port=port, - algorithms=algorithms, - force=force, - **kwargs, - ) - if username: - d["username"] = username - if password: - d["password"] = password - return d - - -def run_build( - dataset: str = "test-data", - dataset_path: str = "./datasets", - host: str = "localhost", - port: int = 9200, - algorithms: str = "test", - force: bool = False, - username: str = None, - password: str = None, - **kwargs, -): - """ - Build an Elasticsearch vector index. Matches the backend build() interface. - - Returns - ------- - list - BuildResult objects - """ - register() - from cuvs_bench.orchestrator import BenchmarkOrchestrator - - orch = BenchmarkOrchestrator(backend_type=ELASTIC) - return orch.run_benchmark( - build=True, - search=False, - **_common_kwargs( - dataset, dataset_path, host, port, algorithms, force, username, password, **kwargs - ), - ) - - -def run_search( - dataset: str = "test-data", - dataset_path: str = "./datasets", - host: str = "localhost", - port: int = 9200, - algorithms: str = "test", - username: str = None, - password: str = None, - **kwargs, -): - """ - Run kNN search against an existing Elasticsearch index. Matches the backend search() interface. - - Returns - ------- - list - SearchResult objects - """ - register() - from cuvs_bench.orchestrator import BenchmarkOrchestrator - - orch = BenchmarkOrchestrator(backend_type=ELASTIC) - return orch.run_benchmark( - build=False, - search=True, - **_common_kwargs( - dataset, dataset_path, host, port, algorithms, False, username, password, **kwargs - ), - ) - - -def run_benchmark( - dataset: str = "test-data", - dataset_path: str = "./datasets", - host: str = "localhost", - port: int = 9200, - algorithms: str = "test", - build: bool = True, - search: bool = True, - force: bool = False, - username: str = None, - password: str = None, - **kwargs, -): - """ - Run build and/or search. Convenience wrapper; run_build and run_search are the primary API. - """ - register() - from cuvs_bench.orchestrator import BenchmarkOrchestrator - - orch = BenchmarkOrchestrator(backend_type=ELASTIC) - return orch.run_benchmark( - build=build, - search=search, - **_common_kwargs( - dataset, dataset_path, host, port, algorithms, force, username, password, **kwargs - ), - ) - - -__all__ = ["register", "ELASTIC", "run_build", "run_search", "run_benchmark"] diff --git a/python/cuvs_bench_elastic/pyproject.toml b/python/cuvs_bench_elastic/pyproject.toml deleted file mode 100644 index 3a6a6e6a95..0000000000 --- a/python/cuvs_bench_elastic/pyproject.toml +++ /dev/null @@ -1,54 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. -# SPDX-License-Identifier: Apache-2.0 - -[build-system] -requires = ["setuptools>=61.0", "wheel"] -build-backend = "setuptools.build_meta" - -[project] -name = "cuvs-bench-elastic" -dynamic = ["version"] -description = "Elasticsearch GPU backend plugin for cuvs-bench" -readme = { file = "README.md", content-type = "text/markdown" } -authors = [ - { name = "NVIDIA Corporation" }, -] -license = { text = "Apache-2.0" } -requires-python = ">=3.11" -classifiers = [ - "Intended Audience :: Developers", - "Topic :: Database", - "Topic :: Scientific/Engineering", - "Programming Language :: Python :: 3.11", - "Programming Language :: Python :: 3.12", - "Programming Language :: Python :: 3.13", -] -dependencies = [ - "cuvs-bench>=26.4.0", - "numpy", - "pyyaml", -] - -[project.optional-dependencies] -# elasticsearch client; automatically installed via `pip install cuvs-bench[elastic]`. -# Only needed here for standalone dev installs of this package directly. -client = ["elasticsearch>=8.0"] - -[project.urls] -Homepage = "https://github.com/rapidsai/cuvs" - -[project.entry-points."cuvs_bench.backends"] -elastic = "cuvs_bench_elastic:register" - -[project.entry-points."cuvs_bench.config_loaders"] -elastic = "cuvs_bench_elastic:register" - -[tool.setuptools.dynamic] -version = { file = "cuvs_bench_elastic/VERSION" } - -[tool.setuptools.packages.find] -where = ["."] -include = ["cuvs_bench_elastic*"] - -[tool.setuptools.package-data] -cuvs_bench_elastic = ["config/**/*.yaml"] From b787396a2c28ccd6e2cbd55eb7336dcac46ba0df Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Mon, 13 Apr 2026 14:22:04 -0700 Subject: [PATCH 08/22] Add run_build/run_search/run_benchmark convenience API to elasticsearch.py Restores the high-level API that was previously in cuvs_bench_elastic/__init__.py so existing demo scripts continue to work after the module consolidation. Signed-off-by: Alex Fournier --- .../cuvs_bench/backends/elasticsearch.py | 63 +++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py b/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py index 7e3d315140..1b5df83f59 100644 --- a/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py +++ b/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py @@ -707,3 +707,66 @@ def register() -> None: register_backend("elastic", ElasticBackend) if "elastic" not in _CONFIG_LOADER_REGISTRY: register_config_loader("elastic", ElasticConfigLoader) + + +# ── Convenience API ─────────────────────────────────────────────────────────── + +def run_build( + dataset: str = "test-data", + dataset_path: str = "./datasets", + host: str = "localhost", + port: int = 9200, + algorithms: str = "test", + force: bool = False, + **kwargs, +): + """Build an Elasticsearch vector index. Returns list of BuildResult.""" + register() + from cuvs_bench.orchestrator import BenchmarkOrchestrator + orch = BenchmarkOrchestrator(backend_type="elastic") + return orch.run_benchmark( + build=True, search=False, + dataset=dataset, dataset_path=dataset_path, + host=host, port=port, algorithms=algorithms, force=force, **kwargs, + ) + + +def run_search( + dataset: str = "test-data", + dataset_path: str = "./datasets", + host: str = "localhost", + port: int = 9200, + algorithms: str = "test", + **kwargs, +): + """Run kNN search against an existing Elasticsearch index. Returns list of SearchResult.""" + register() + from cuvs_bench.orchestrator import BenchmarkOrchestrator + orch = BenchmarkOrchestrator(backend_type="elastic") + return orch.run_benchmark( + build=False, search=True, + dataset=dataset, dataset_path=dataset_path, + host=host, port=port, algorithms=algorithms, **kwargs, + ) + + +def run_benchmark( + dataset: str = "test-data", + dataset_path: str = "./datasets", + host: str = "localhost", + port: int = 9200, + algorithms: str = "test", + build: bool = True, + search: bool = True, + force: bool = False, + **kwargs, +): + """Run build and/or search. Returns list of results.""" + register() + from cuvs_bench.orchestrator import BenchmarkOrchestrator + orch = BenchmarkOrchestrator(backend_type="elastic") + return orch.run_benchmark( + build=build, search=search, + dataset=dataset, dataset_path=dataset_path, + host=host, port=port, algorithms=algorithms, force=force, **kwargs, + ) From 4e35484d11d2c3bf93d540342af7577b62011ce5 Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Wed, 20 May 2026 13:59:28 -0400 Subject: [PATCH 09/22] Fix elasticsearch backend after cuvs-bench API changes Signed-off-by: Alex Fournier --- .../cuvs_bench/backends/elasticsearch.py | 546 ++++++++++-------- .../cuvs_bench/tests/test_modularization.py | 308 ++++++++-- 2 files changed, 565 insertions(+), 289 deletions(-) diff --git a/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py b/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py index 1b5df83f59..cff24841ee 100644 --- a/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py +++ b/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py @@ -15,18 +15,16 @@ Search params (knn): num_candidates, vector_field. """ -import itertools import os import time from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple import numpy as np -import yaml - from .base import BenchmarkBackend, BuildResult, Dataset, SearchResult from .registry import register_backend, register_config_loader +from ._utils import load_vectors from ..orchestrator.config_loaders import ( BenchmarkConfig, ConfigLoader, @@ -34,26 +32,27 @@ IndexConfig, ) +if TYPE_CHECKING: + from elasticsearch import Elasticsearch + def _load_fbin(path: Path) -> np.ndarray: - """Load big-ann-bench fbin format (header: n_rows, n_cols as uint32, then float32).""" - with open(path, "rb") as f: - n_rows, n_cols = np.fromfile(f, dtype=np.uint32, count=2) - data = np.fromfile(f, dtype=np.float32).reshape(n_rows, n_cols) - return data + """Load big-ann-bench fbin format via shared vector loader.""" + return load_vectors(os.fspath(path)) def _load_ibin(path: Path) -> np.ndarray: - """Load big-ann-bench ibin format (header: shape as uint32, then int32).""" - with open(path, "rb") as f: - shape = np.fromfile(f, dtype=np.uint32, count=2) - data = np.fromfile(f, dtype=np.int32).reshape(shape[0], shape[1]) - return data + """Load big-ann-bench ibin format via shared vector loader.""" + return load_vectors(os.fspath(path)) def _distance_to_similarity(distance: str) -> str: """Map cuvs-bench distance metric to ES dense_vector similarity.""" - m = {"euclidean": "l2_norm", "inner_product": "max_inner_product", "cosine": "cosine"} + m = { + "euclidean": "l2_norm", + "inner_product": "max_inner_product", + "cosine": "cosine", + } return m.get(distance, "l2_norm") @@ -110,8 +109,14 @@ def _get_client(self) -> "Elasticsearch": } basic_auth = self.config.get("basic_auth") if basic_auth is not None: - if isinstance(basic_auth, (list, tuple)) and len(basic_auth) >= 2: - kwargs["basic_auth"] = (str(basic_auth[0]), str(basic_auth[1])) + if ( + isinstance(basic_auth, (list, tuple)) + and len(basic_auth) >= 2 + ): + kwargs["basic_auth"] = ( + str(basic_auth[0]), + str(basic_auth[1]), + ) elif isinstance(basic_auth, str) and ":" in basic_auth: user, _, passwd = basic_auth.partition(":") kwargs["basic_auth"] = (user, passwd) @@ -178,7 +183,9 @@ def build( if client.indices.exists(index=index_name): if not force: stats = client.indices.stats(index=index_name) - index_size = stats["_all"]["primaries"]["store"]["size_in_bytes"] + index_size = stats["_all"]["primaries"]["store"][ + "size_in_bytes" + ] return BuildResult( index_path=index_name, build_time_seconds=0, @@ -199,7 +206,8 @@ def build( error_message=str(e), ) - if not dataset.base_file: + vectors = dataset.training_vectors + if vectors.size == 0: return BuildResult( index_path="", build_time_seconds=0.0, @@ -207,19 +215,10 @@ def build( algorithm=self.algo, build_params={}, success=False, - error_message="base_file is required for Elasticsearch backend", - ) - - base_path = Path(dataset.base_file) - if not base_path.exists(): - return BuildResult( - index_path="", - build_time_seconds=0.0, - index_size_bytes=0, - algorithm=self.algo, - build_params={}, - success=False, - error_message=f"Base file not found: {base_path}", + error_message=( + "training_vectors are required for Elasticsearch backend " + "(directly or via dataset.base_file)" + ), ) # similarity: from config, or derive from dataset distance @@ -228,18 +227,21 @@ def build( ) try: - vectors = _load_fbin(base_path) n_vectors = len(vectors) dims = vectors.shape[1] client = self._get_client() - vector_field = build_params.get("vector_field", _DEFAULT_VECTOR_FIELD) + vector_field = build_params.get( + "vector_field", _DEFAULT_VECTOR_FIELD + ) index_type = build_params.get("type", _DEFAULT_INDEX_TYPE) m = build_params.get("m", _DEFAULT_M) ef_construction = build_params.get( "ef_construction", _DEFAULT_EF_CONSTRUCTION ) - num_shards = build_params.get("number_of_shards", _DEFAULT_NUM_SHARDS) + num_shards = build_params.get( + "number_of_shards", _DEFAULT_NUM_SHARDS + ) num_replicas = build_params.get( "number_of_replicas", _DEFAULT_NUM_REPLICAS ) @@ -275,8 +277,11 @@ def build( client.indices.create(index=index_name, body=index_config) from elasticsearch.helpers import bulk + chunk_size = 1000 - progress_interval = max(50, n_vectors // (chunk_size * 20)) # ~20 progress lines + progress_interval = max( + 50, n_vectors // (chunk_size * 20) + ) # ~20 progress lines for i in range(0, n_vectors, chunk_size): chunk = vectors[i : i + chunk_size] actions = [ @@ -288,8 +293,13 @@ def build( for j, vec in enumerate(chunk) ] bulk(client, actions, raise_on_error=True) - if progress_interval and (i // chunk_size) % progress_interval == 0: - print(f" Indexed {min(i + chunk_size, n_vectors):,}/{n_vectors:,} vectors") + if ( + progress_interval + and (i // chunk_size) % progress_interval == 0 + ): + print( + f" Indexed {min(i + chunk_size, n_vectors):,}/{n_vectors:,} vectors" + ) client.indices.refresh(index=index_name) build_time = time.perf_counter() - t0 @@ -367,7 +377,8 @@ def search( error_message="No indexes provided", ) - if not dataset.query_file: + query_vectors = dataset.query_vectors + if query_vectors.size == 0: return SearchResult( neighbors=np.zeros((0, k), dtype=np.int64), distances=np.zeros((0, k), dtype=np.float32), @@ -377,32 +388,16 @@ def search( algorithm=self.algo, search_params=[], success=False, - error_message="query_file is required", - ) - - query_path = Path(dataset.query_file) - if not query_path.exists(): - return SearchResult( - neighbors=np.zeros((0, k), dtype=np.int64), - distances=np.zeros((0, k), dtype=np.float32), - search_time_ms=0, - queries_per_second=0, - recall=0, - algorithm=self.algo, - search_params=[], - success=False, - error_message=f"Query file not found: {query_path}", + error_message=( + "query_vectors are required for Elasticsearch backend " + "(directly or via dataset.query_file)" + ), ) try: - query_vectors = _load_fbin(query_path) n_queries = len(query_vectors) groundtruth = dataset.groundtruth_neighbors - if groundtruth is None and dataset.groundtruth_neighbors_file: - gt_path = Path(dataset.groundtruth_neighbors_file) - if gt_path.exists(): - groundtruth = _load_ibin(gt_path) index_name = self.config.get("index_name", "cuvs_bench_vectors") index_cfg = indexes[0] @@ -413,7 +408,9 @@ def search( last_distances = np.zeros((n_queries, k), dtype=np.float32) for sp in search_params_list: - num_candidates = sp.get("num_candidates", _DEFAULT_NUM_CANDIDATES) + num_candidates = sp.get( + "num_candidates", _DEFAULT_NUM_CANDIDATES + ) vector_field = sp.get("vector_field", _DEFAULT_VECTOR_FIELD) neighbors = np.full((n_queries, k), -1, dtype=np.int64) @@ -431,7 +428,9 @@ def search( } } t_q = time.perf_counter() - resp = self._get_client().search(index=index_name, body=body, size=k) + resp = self._get_client().search( + index=index_name, body=body, size=k + ) latencies.append((time.perf_counter() - t_q) * 1000) hits = resp.get("hits", {}).get("hits", []) for j, hit in enumerate(hits[:k]): @@ -439,32 +438,47 @@ def search( distances[i, j] = float(hit["_score"]) elapsed_ms = (time.perf_counter() - t0) * 1000 - qps = n_queries / (elapsed_ms / 1000) if elapsed_ms > 0 else 0.0 + qps = ( + n_queries / (elapsed_ms / 1000) if elapsed_ms > 0 else 0.0 + ) recall = 0.0 if groundtruth is not None: gt_k = min(k, groundtruth.shape[1]) n_correct = sum( - len(set(neighbors[i, :k].tolist()) & set(groundtruth[i, :gt_k].tolist())) + len( + set(neighbors[i, :k].tolist()) + & set(groundtruth[i, :gt_k].tolist()) + ) for i in range(n_queries) ) - recall = n_correct / (n_queries * gt_k) if gt_k > 0 else 0.0 - - per_param_results.append({ - "search_params": sp, - "search_time_ms": elapsed_ms, - "queries_per_second": qps, - "recall": recall, - "p50_ms": float(np.percentile(latencies, 50)), - "p95_ms": float(np.percentile(latencies, 95)), - "p99_ms": float(np.percentile(latencies, 99)), - }) + recall = ( + n_correct / (n_queries * gt_k) if gt_k > 0 else 0.0 + ) + + per_param_results.append( + { + "search_params": sp, + "search_time_ms": elapsed_ms, + "queries_per_second": qps, + "recall": recall, + "p50_ms": float(np.percentile(latencies, 50)), + "p95_ms": float(np.percentile(latencies, 95)), + "p99_ms": float(np.percentile(latencies, 99)), + } + ) last_neighbors = neighbors last_distances = distances - avg_recall = float(np.mean([r["recall"] for r in per_param_results])) - avg_qps = float(np.mean([r["queries_per_second"] for r in per_param_results])) - total_ms = float(sum(r["search_time_ms"] for r in per_param_results)) + avg_recall = float( + np.mean([r["recall"] for r in per_param_results]) + ) + avg_qps = float( + np.mean([r["queries_per_second"] for r in per_param_results]) + ) + total_ms = float( + sum(r["search_time_ms"] for r in per_param_results) + ) return SearchResult( neighbors=last_neighbors, @@ -479,7 +493,10 @@ def search( "p95_ms": per_param_results[-1]["p95_ms"], "p99_ms": per_param_results[-1]["p99_ms"], }, - metadata={"per_search_param_results": per_param_results}, + metadata={ + "per_search_param_results": per_param_results, + "recall_is_authoritative": True, + }, success=True, ) except Exception as e: @@ -535,164 +552,210 @@ def load( self, dataset: str = "", dataset_path: str = "", - count: int = 10, - batch_size: int = 10000, - host: str = "localhost", - port: int = 9200, - index_name: str = "cuvs_bench_vectors", - basic_auth: Optional[Any] = None, - algorithms: Optional[str] = None, - subset_size: Optional[int] = None, **kwargs, ) -> Tuple[DatasetConfig, List[BenchmarkConfig]]: - """Load Elasticsearch benchmark configuration.""" - tune_mode = kwargs.pop("_tune_mode", False) - tune_build_params = kwargs.pop("_tune_build_params", None) - tune_search_params = kwargs.pop("_tune_search_params", None) - username = kwargs.pop("username", None) - password = kwargs.pop("password", None) - scheme = kwargs.pop("scheme", "http") - if basic_auth is None and username and password: - basic_auth = (username, password) - - datasets_path = os.path.join( - self.config_path, "datasets", "datasets.yaml" + """Load Elasticsearch benchmark configuration via shared ConfigLoader flow.""" + return super().load( + dataset=dataset, dataset_path=dataset_path, **kwargs ) - try: - with open(datasets_path) as f: - datasets = yaml.safe_load(f) - except FileNotFoundError: - raise FileNotFoundError( - f"Datasets config not found: {datasets_path}" - ) from None - - dataset_conf = next((d for d in datasets if d["name"] == dataset), None) - if not dataset_conf: - raise ValueError(f"Dataset '{dataset}' not found") - - def _resolve(rel: Optional[str]) -> Optional[str]: - if rel and not os.path.isabs(rel): - return os.path.join(dataset_path, rel) - return rel - - dataset_config = DatasetConfig( - name=dataset_conf["name"], - base_file=_resolve(dataset_conf.get("base_file")), - query_file=_resolve(dataset_conf.get("query_file")), - groundtruth_neighbors_file=_resolve( - dataset_conf.get("groundtruth_neighbors_file") - ), - distance=dataset_conf.get("distance", "euclidean"), - dims=dataset_conf.get("dims"), - subset_size=subset_size, + + def _discover_algo_groups( + self, + dataset_conf, + dataset, + dataset_path, + **kwargs, + ): + """Discover elastic algorithm groups using shared loader semantics.""" + algorithm_configuration = kwargs.get("algorithm_configuration") + algorithms_arg = kwargs.get("algorithms") + groups_arg = kwargs.get("groups") + algo_groups_arg = kwargs.get("algo_groups") + + algos_conf_fs = self.gather_algorithm_configs( + self.config_path, algorithm_configuration ) - if tune_mode and tune_build_params is not None and tune_search_params is not None: - algo_name = algorithms or "elastic_hnsw" - build_param = dict(tune_build_params) - if "type" not in build_param and algo_name.startswith("elastic_"): - build_param["type"] = algo_name.replace("elastic_", "", 1) - name_parts = [f"{k}{v}" for k, v in build_param.items() if k in ("m", "ef_construction")] - index_label = "_".join([f"{algo_name}_tune"] + name_parts) if name_parts else f"{algo_name}_tune" - es_index_name = index_label.lower().replace(".", "_") - index_config = IndexConfig( - name=index_label, - algo=algo_name, - build_param=build_param, - search_params=[tune_search_params], - file="", - ) - config = BenchmarkConfig( - indexes=[index_config], - backend_config={ - "name": index_label, - "host": host, - "port": port, - "scheme": scheme, - "index_name": es_index_name, - "basic_auth": basic_auth, - - **build_param, - }, + elastic_algos = [] + for algo_f in algos_conf_fs: + try: + algo_conf = self.load_yaml_file(algo_f) + except Exception: + continue + if not isinstance(algo_conf, dict): + continue + algo_name = algo_conf.get("name", "") + if not algo_name.startswith("elastic_"): + continue + elastic_algos.append(algo_conf) + + if not elastic_algos: + default_grp = { + "build": {"m": [16], "ef_construction": [100]}, + "search": {"num_candidates": [100]}, + } + elastic_algos = [ + {"name": "elastic_hnsw", "groups": {"base": default_grp}} + ] + + allowed_algos = ( + [a.strip() for a in algorithms_arg.split(",") if a.strip()] + if algorithms_arg + else None + ) + allowed_groups = ( + [g.strip() for g in groups_arg.split(",") if g.strip()] + if groups_arg + else None + ) + algo_group_map: Dict[str, set] = {} + if algo_groups_arg: + for item in algo_groups_arg.split(","): + item = item.strip() + if not item or "." not in item: + continue + algo_name, group_name = item.split(".", 1) + algo_group_map.setdefault(algo_name, set()).add(group_name) + + # Backward-compatible fallback for older elastic usage where `algorithms` + # was used as a group selector (e.g. algorithms="test"). + if allowed_algos and not allowed_groups and not algo_group_map: + known_groups = { + group_name + for algo_conf in elastic_algos + for group_name in algo_conf.get("groups", {}) + } + if all(name in known_groups for name in allowed_algos): + allowed_groups = allowed_algos + allowed_algos = None + + if allowed_groups is None and not algo_group_map: + allowed_groups = ["base"] + + result = [] + for algo_conf in elastic_algos: + algo_name = algo_conf["name"] + if allowed_algos and algo_name not in allowed_algos: + continue + + groups = dict(algo_conf.get("groups", {})) + if allowed_groups is not None: + groups = { + group_name: group_conf + for group_name, group_conf in groups.items() + if group_name in allowed_groups + } + if algo_name in algo_group_map: + groups = { + group_name: group_conf + for group_name, group_conf in groups.items() + if group_name in algo_group_map[algo_name] + } + + for group_name, group_conf in groups.items(): + result.append((algo_name, group_name, group_conf, {})) + + if not result and allowed_groups: + raise ValueError( + f"Could not find elastic groups {allowed_groups} in elastic configs" ) - return dataset_config, [config] - - elastic_config_dir = _get_elastic_config_path() - elastic_algo_path = os.path.join(elastic_config_dir, "algos", "elastic.yaml") - default_grp = { - "build": {"m": [16], "ef_construction": [100]}, - "search": {"num_candidates": [100]}, - } - if os.path.exists(elastic_algo_path): - with open(elastic_algo_path) as f: - algo_conf = yaml.safe_load(f) - else: - algo_conf = {"groups": {"base": default_grp}} - - groups = algo_conf.get("groups", {"base": default_grp}) - group_name = algorithms or "base" - if group_name not in groups: + if not result and allowed_algos: raise ValueError( - f"Algorithm group '{group_name}' not found in elastic.yaml. " - f"Available: {list(groups.keys())}" + f"Could not find elastic algorithms {allowed_algos} in elastic configs" ) - group_conf = groups[group_name] - - build_params = group_conf.get( - "build", {"m": [16], "ef_construction": [100]} - ) - search_params = group_conf.get("search", {"num_candidates": [100]}) - - # Ensure all param values are lists for itertools.product - def _to_list_values(d: Dict[str, Any]) -> Dict[str, List[Any]]: - return {k: v if isinstance(v, list) else [v] for k, v in d.items()} - build_params = _to_list_values(build_params) - search_params = _to_list_values(search_params) + return result - build_combos = list(itertools.product(*build_params.values())) - search_combos = list(itertools.product(*search_params.values())) - build_keys = list(build_params.keys()) - search_keys = list(search_params.keys()) + def _build_benchmark_configs( + self, + dataset_config, + dataset_conf, + dataset, + dataset_path, + expanded_groups, + **kwargs, + ): + """Build BenchmarkConfigs from shared expanded elastic parameter groups.""" + host = kwargs.get("host", "localhost") + port = kwargs.get("port", 9200) + scheme = kwargs.get("scheme", "http") + basic_auth = kwargs.get("basic_auth") + username = kwargs.get("username") + password = kwargs.get("password") + if basic_auth is None and username and password: + basic_auth = (username, password) - search_params_list = [ - dict(zip(search_keys, svals)) for svals in search_combos - ] + 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 bvals in build_combos: - bdict = dict(zip(build_keys, bvals)) - algo_name = f"elastic_{bdict.get('type', 'hnsw')}" - - # Derive a unique, human-readable index name from build params - prefix = f"{algo_name}_{group_name}" if group_name != "base" else algo_name - name_parts = [f"{k}{v}" for k, v in bdict.items() if k in ("m", "ef_construction")] - index_label = "_".join([prefix] + name_parts) if name_parts else prefix - es_index_name = index_label.lower().replace(".", "_") - - index_config = IndexConfig( - name=index_label, - algo=algo_name, - build_param=bdict, - search_params=search_params_list, - file="", - ) - config = BenchmarkConfig( - indexes=[index_config], - backend_config={ - "name": index_label, - "host": host, - "port": port, - "scheme": scheme, - "index_name": es_index_name, - "basic_auth": basic_auth, - - **bdict, - }, - ) - benchmark_configs.append(config) + for ( + algo_name, + group_name, + _group_conf, + build_combos, + search_combos, + _group_meta, + ) in expanded_groups: + if tune_mode and tune_build_params is not None: + actual_build = [dict(tune_build_params)] + actual_search = ( + [dict(tune_search_params)] if tune_search_params else [{}] + ) + else: + actual_build = build_combos + actual_search = search_combos + + for build_param in actual_build: + build_param = dict(build_param) + if "type" not in build_param and algo_name.startswith( + "elastic_" + ): + build_param["type"] = algo_name.replace("elastic_", "", 1) + + if tune_mode: + label_prefix = f"{algo_name}_tune" + elif group_name != "base": + label_prefix = f"{algo_name}_{group_name}" + else: + label_prefix = algo_name + + name_parts = [ + f"{k}{v}" + for k, v in build_param.items() + if k in ("m", "ef_construction") + ] + index_label = ( + "_".join([label_prefix] + name_parts) + if name_parts + else label_prefix + ) + es_index_name = index_label.lower().replace(".", "_") + + index_config = IndexConfig( + name=index_label, + algo=algo_name, + build_param=build_param, + search_params=[dict(sp) for sp in actual_search], + file="", + ) + benchmark_configs.append( + BenchmarkConfig( + indexes=[index_config], + backend_config={ + "name": index_label, + "host": host, + "port": port, + "scheme": scheme, + "index_name": es_index_name, + "basic_auth": basic_auth, + **build_param, + }, + ) + ) - return dataset_config, benchmark_configs + return benchmark_configs def register() -> None: @@ -711,6 +774,7 @@ def register() -> None: # ── Convenience API ─────────────────────────────────────────────────────────── + def run_build( dataset: str = "test-data", dataset_path: str = "./datasets", @@ -723,11 +787,18 @@ def run_build( """Build an Elasticsearch vector index. Returns list of BuildResult.""" register() from cuvs_bench.orchestrator import BenchmarkOrchestrator + orch = BenchmarkOrchestrator(backend_type="elastic") return orch.run_benchmark( - build=True, search=False, - dataset=dataset, dataset_path=dataset_path, - host=host, port=port, algorithms=algorithms, force=force, **kwargs, + build=True, + search=False, + dataset=dataset, + dataset_path=dataset_path, + host=host, + port=port, + algorithms=algorithms, + force=force, + **kwargs, ) @@ -742,11 +813,17 @@ def run_search( """Run kNN search against an existing Elasticsearch index. Returns list of SearchResult.""" register() from cuvs_bench.orchestrator import BenchmarkOrchestrator + orch = BenchmarkOrchestrator(backend_type="elastic") return orch.run_benchmark( - build=False, search=True, - dataset=dataset, dataset_path=dataset_path, - host=host, port=port, algorithms=algorithms, **kwargs, + build=False, + search=True, + dataset=dataset, + dataset_path=dataset_path, + host=host, + port=port, + algorithms=algorithms, + **kwargs, ) @@ -764,9 +841,16 @@ def run_benchmark( """Run build and/or search. Returns list of results.""" register() from cuvs_bench.orchestrator import BenchmarkOrchestrator + orch = BenchmarkOrchestrator(backend_type="elastic") return orch.run_benchmark( - build=build, search=search, - dataset=dataset, dataset_path=dataset_path, - host=host, port=port, algorithms=algorithms, force=force, **kwargs, + build=build, + search=search, + dataset=dataset, + dataset_path=dataset_path, + host=host, + port=port, + algorithms=algorithms, + force=force, + **kwargs, ) diff --git a/python/cuvs_bench/cuvs_bench/tests/test_modularization.py b/python/cuvs_bench/cuvs_bench/tests/test_modularization.py index 112edb4bd4..1e5cfa4c53 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_modularization.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_modularization.py @@ -13,7 +13,12 @@ import numpy as np import pytest -from cuvs_bench.backends.base import BenchmarkBackend, BuildResult, Dataset, SearchResult +from cuvs_bench.backends.base import ( + BenchmarkBackend, + BuildResult, + Dataset, + SearchResult, +) from cuvs_bench.backends.registry import ( get_backend_class, get_config_loader, @@ -48,17 +53,29 @@ def test_elastic_without_extra_raises_clear_error(self): """Requesting elastic without [elastic] installed raises helpful error.""" try: import elasticsearch # noqa: F401 - pytest.skip("elasticsearch is installed; cannot test missing-plugin path") + + pytest.skip( + "elasticsearch is installed; cannot test missing-plugin path" + ) except ImportError: pass import importlib.metadata + all_eps = importlib.metadata.entry_points() if hasattr(all_eps, "select"): - eps = list(all_eps.select(group="cuvs_bench.backends", name="elastic")) + eps = list( + all_eps.select(group="cuvs_bench.backends", name="elastic") + ) else: - eps = [e for e in all_eps.get("cuvs_bench.backends", []) if e.name == "elastic"] + eps = [ + e + for e in all_eps.get("cuvs_bench.backends", []) + if e.name == "elastic" + ] if eps: - pytest.skip("cuvs-bench-elastic is installed; cannot test missing-plugin path") + pytest.skip( + "cuvs-bench-elastic is installed; cannot test missing-plugin path" + ) with pytest.raises((ImportError, ValueError)) as exc_info: get_backend_class("elastic") @@ -70,17 +87,31 @@ def test_elastic_config_loader_without_extra_raises_clear_error(self): """Requesting elastic config loader without [elastic] raises helpful error.""" try: import elasticsearch # noqa: F401 - pytest.skip("elasticsearch is installed; cannot test missing-plugin path") + + pytest.skip( + "elasticsearch is installed; cannot test missing-plugin path" + ) except ImportError: pass import importlib.metadata + all_eps = importlib.metadata.entry_points() if hasattr(all_eps, "select"): - eps = list(all_eps.select(group="cuvs_bench.config_loaders", name="elastic")) + eps = list( + all_eps.select( + group="cuvs_bench.config_loaders", name="elastic" + ) + ) else: - eps = [e for e in all_eps.get("cuvs_bench.config_loaders", []) if e.name == "elastic"] + eps = [ + e + for e in all_eps.get("cuvs_bench.config_loaders", []) + if e.name == "elastic" + ] if eps: - pytest.skip("cuvs-bench-elastic is installed; cannot test missing-plugin path") + pytest.skip( + "cuvs-bench-elastic is installed; cannot test missing-plugin path" + ) with pytest.raises((ImportError, ValueError)) as exc_info: get_config_loader("elastic") @@ -106,7 +137,7 @@ def test_unknown_config_loader_raises_value_error(self): assert "nonexistent_loader_xyz" in msg def test_orchestrator_cpp_gbench_no_regression(self): - """BenchmarkOrchestrator with cpp_gbench should initialize (no regression).""" + """Initializing BenchmarkOrchestrator with cpp_gbench should work.""" from cuvs_bench.orchestrator import BenchmarkOrchestrator assert "cpp_gbench" in BenchmarkOrchestrator.available_backends() @@ -115,6 +146,24 @@ def test_orchestrator_cpp_gbench_no_regression(self): assert orch.backend_class is not None assert orch.config_loader is not None + def test_orchestrator_respects_authoritative_backend_recall(self): + """Backends can mark their own recall as authoritative.""" + from cuvs_bench.orchestrator.orchestrator import _should_compute_recall + + result = SearchResult( + neighbors=np.array([[1, 2, 3]], dtype=np.int64), + distances=np.array([[0.1, 0.2, 0.3]], dtype=np.float32), + search_time_ms=1.0, + queries_per_second=1.0, + recall=0.5, + algorithm="elastic_hnsw", + search_params=[{"num_candidates": 100}], + metadata={"recall_is_authoritative": True}, + success=True, + ) + + assert not _should_compute_recall(result) + class TestPluginLoaderMocked: """ @@ -141,10 +190,18 @@ def build(self, dataset, indexes, force=False, dry_run=False): ) def search( - self, dataset, indexes, k=10, batch_size=10000, - mode="latency", force=False, search_threads=None, dry_run=False, + self, + dataset, + indexes, + k=10, + batch_size=10000, + mode="latency", + force=False, + search_threads=None, + dry_run=False, ): import numpy as np + return SearchResult( neighbors=np.empty((0, k)), distances=np.empty((0, k)), @@ -165,7 +222,9 @@ def load(self, **kwargs): raise NotImplementedError("Stub loader") def register(): - register_backend(TestPluginLoaderMocked._MOCK_PLUGIN_NAME, StubBackend) + register_backend( + TestPluginLoaderMocked._MOCK_PLUGIN_NAME, StubBackend + ) register_config_loader( TestPluginLoaderMocked._MOCK_PLUGIN_NAME, StubConfigLoader ) @@ -196,7 +255,9 @@ def test_valid_plugin_loads_via_mock_entry_point(self): get_registry().unregister(self._MOCK_PLUGIN_NAME) unregister_config_loader(self._MOCK_PLUGIN_NAME) - def test_import_error_with_elasticsearch_message_raises_helpful_error(self): + def test_import_error_with_elasticsearch_message_raises_helpful_error( + self, + ): """Mock entry point raising ImportError(elasticsearch) -> our install message.""" # Ensure elastic is not in registry (e.g. from TestElasticWithExtraInstalled) registry = get_registry() @@ -206,7 +267,9 @@ def test_import_error_with_elasticsearch_message_raises_helpful_error(self): mock_ep = MagicMock() mock_ep.name = "elastic" - mock_ep.load.side_effect = ImportError("No module named 'elasticsearch'") + mock_ep.load.side_effect = ImportError( + "No module named 'elasticsearch'" + ) mock_eps = MagicMock() mock_eps.select.return_value = [mock_ep] @@ -225,7 +288,9 @@ def test_import_error_unrelated_propagates(self): """Mock entry point: unrelated ImportError propagates unchanged.""" mock_ep = MagicMock() mock_ep.name = "other_plugin" - mock_ep.load.side_effect = ImportError("No module named 'something_else'") + mock_ep.load.side_effect = ImportError( + "No module named 'something_else'" + ) mock_eps = MagicMock() mock_eps.select.return_value = [mock_ep] @@ -277,6 +342,7 @@ def test_no_entry_point_for_name_raises_value_error(self): def _elasticsearch_installed(): try: import elasticsearch # noqa: F401 + return True except ImportError: return False @@ -295,7 +361,8 @@ def _ensure_elastic_registered(self): registry = get_registry() if "elastic" not in registry._backends: try: - from cuvs_bench_elastic import register + from cuvs_bench.backends.elasticsearch import register + register() except ImportError: pass @@ -312,6 +379,7 @@ def test_elastic_config_loader_tune_mode_returns_single_config(self): loader = loader_cls() dataset_config, benchmark_configs = loader.load( dataset="glove-50-angular", + dataset_path="", algorithms="elastic_hnsw", _tune_mode=True, _tune_build_params={"m": 24, "ef_construction": 150}, @@ -332,7 +400,9 @@ def test_elastic_config_loader_sweep_mode_returns_multiple_configs(self): loader = loader_cls() dataset_config, benchmark_configs = loader.load( dataset="glove-50-angular", - algorithms="test", + dataset_path="", + algorithms="elastic_hnsw", + groups="test", ) assert len(benchmark_configs) >= 1 config = benchmark_configs[0] @@ -345,28 +415,32 @@ def test_elastic_config_loader_dataset_not_found_raises(self): loader_cls = get_config_loader("elastic") loader = loader_cls() with pytest.raises(ValueError, match="not found"): - loader.load(dataset="nonexistent_dataset_xyz") + loader.load(dataset="nonexistent_dataset_xyz", dataset_path="") def test_elastic_config_loader_group_not_found_raises(self): """Config loader raises ValueError for unknown algorithm group.""" loader_cls = get_config_loader("elastic") loader = loader_cls() - with pytest.raises(ValueError, match="not found"): + with pytest.raises(ValueError, match="find elastic groups"): loader.load( dataset="glove-50-angular", - algorithms="nonexistent_group_xyz", + dataset_path="", + algorithms="elastic_hnsw", + groups="nonexistent_group_xyz", ) def test_elastic_dry_run_build(self): """ElasticBackend.build(dry_run=True) returns synthetic result without ES.""" cls = get_backend_class("elastic") - backend = cls(config={"name": "test", "host": "localhost", "port": 9200}) + backend = cls( + config={"name": "test", "host": "localhost", "port": 9200} + ) base = np.random.rand(100, 32).astype(np.float32) queries = np.random.rand(10, 32).astype(np.float32) dataset = Dataset( name="test", - base_vectors=base, + training_vectors=base, query_vectors=queries, distance_metric="euclidean", ) @@ -389,13 +463,15 @@ def test_elastic_dry_run_build(self): def test_elastic_dry_run_search(self): """ElasticBackend.search(dry_run=True) returns synthetic result without ES.""" cls = get_backend_class("elastic") - backend = cls(config={"name": "test", "host": "localhost", "port": 9200}) + backend = cls( + config={"name": "test", "host": "localhost", "port": 9200} + ) base = np.random.rand(100, 32).astype(np.float32) queries = np.random.rand(10, 32).astype(np.float32) dataset = Dataset( name="test", - base_vectors=base, + training_vectors=base, query_vectors=queries, distance_metric="euclidean", ) @@ -417,16 +493,16 @@ def test_elastic_dry_run_search(self): assert result.algorithm == "elastic_hnsw" assert result.search_time_ms == 0 - def test_elastic_build_requires_base_file(self): - """ElasticBackend.build returns error when dataset has no base_file.""" + def test_elastic_build_requires_training_vectors(self): + """ElasticBackend.build returns error when no training vectors are available.""" cls = get_backend_class("elastic") - backend = cls(config={"name": "test", "host": "localhost", "port": 9200}) + backend = cls( + config={"name": "test", "host": "localhost", "port": 9200} + ) - base = np.random.rand(100, 32).astype(np.float32) queries = np.random.rand(10, 32).astype(np.float32) dataset = Dataset( name="test", - base_vectors=base, query_vectors=queries, base_file=None, query_file=None, @@ -444,21 +520,25 @@ def test_elastic_build_requires_base_file(self): with patch.object( backend, "_check_network_available", return_value=True ): - result = backend.build(dataset=dataset, indexes=indexes, dry_run=False) + result = backend.build( + dataset=dataset, indexes=indexes, dry_run=False + ) assert not result.success - assert "base_file" in (result.error_message or "").lower() + assert "training_vectors" in (result.error_message or "") def test_elastic_preflight_fails_when_no_network(self): """ElasticBackend.build returns success=False when network is unavailable.""" cls = get_backend_class("elastic") - backend = cls(config={"name": "test", "host": "localhost", "port": 9200}) + backend = cls( + config={"name": "test", "host": "localhost", "port": 9200} + ) base = np.random.rand(100, 32).astype(np.float32) queries = np.random.rand(10, 32).astype(np.float32) dataset = Dataset( name="test", - base_vectors=base, + training_vectors=base, query_vectors=queries, base_file="dummy/base.fbin", query_file="dummy/query.fbin", @@ -484,11 +564,13 @@ def test_elastic_preflight_fails_when_no_network(self): def test_elastic_search_preflight_fails_when_no_network(self): """ElasticBackend.search returns success=False when network is unavailable.""" cls = get_backend_class("elastic") - backend = cls(config={"name": "test", "host": "localhost", "port": 9200}) + backend = cls( + config={"name": "test", "host": "localhost", "port": 9200} + ) dataset = Dataset( name="test", - base_vectors=np.random.rand(100, 32).astype(np.float32), + training_vectors=np.random.rand(100, 32).astype(np.float32), query_vectors=np.random.rand(10, 32).astype(np.float32), query_file="dummy/query.fbin", ) @@ -502,7 +584,9 @@ def test_elastic_search_preflight_fails_when_no_network(self): ) ] - with patch.object(backend, "_check_network_available", return_value=False): + with patch.object( + backend, "_check_network_available", return_value=False + ): result = backend.search(dataset=dataset, indexes=indexes, k=10) assert not result.success @@ -511,12 +595,14 @@ def test_elastic_search_preflight_fails_when_no_network(self): def test_elastic_build_skips_existing_index_when_force_false(self): """build(force=False) returns success=True immediately when index already exists.""" cls = get_backend_class("elastic") - backend = cls(config={ - "name": "test", - "host": "localhost", - "port": 9200, - "index_name": "test_index", - }) + backend = cls( + config={ + "name": "test", + "host": "localhost", + "port": 9200, + "index_name": "test_index", + } + ) mock_client = MagicMock() mock_client.indices.exists.return_value = True @@ -526,7 +612,7 @@ def test_elastic_build_skips_existing_index_when_force_false(self): dataset = Dataset( name="test", - base_vectors=np.random.rand(100, 32).astype(np.float32), + training_vectors=np.random.rand(100, 32).astype(np.float32), query_vectors=np.random.rand(10, 32).astype(np.float32), base_file="dummy/base.fbin", ) @@ -534,22 +620,124 @@ def test_elastic_build_skips_existing_index_when_force_false(self): IndexConfig( name="test_index", algo="elastic_hnsw", - build_param={"type": "hnsw", "m": 16, "ef_construction": 100, - "similarity": "l2_norm", "number_of_shards": 1, - "number_of_replicas": 0, "vector_field": "embedding"}, + build_param={ + "type": "hnsw", + "m": 16, + "ef_construction": 100, + "similarity": "l2_norm", + "number_of_shards": 1, + "number_of_replicas": 0, + "vector_field": "embedding", + }, search_params=[{"num_candidates": 100}], file="", ) ] - with patch.object(backend, "_check_network_available", return_value=True): - with patch.object(backend, "_get_client", return_value=mock_client): - result = backend.build(dataset=dataset, indexes=indexes, force=False) + with patch.object( + backend, "_check_network_available", return_value=True + ): + with patch.object( + backend, "_get_client", return_value=mock_client + ): + result = backend.build( + dataset=dataset, indexes=indexes, force=False + ) assert result.success assert result.index_size_bytes == 1024 mock_client.indices.delete.assert_not_called() + def test_elastic_build_uses_lazy_loaded_training_vectors(self): + """ElasticBackend.build works with Dataset lazy-loading via base_file.""" + cls = get_backend_class("elastic") + backend = cls( + config={"name": "test", "host": "localhost", "port": 9200} + ) + + mock_client = MagicMock() + mock_client.indices.exists.return_value = True + mock_client.indices.stats.return_value = { + "_all": {"primaries": {"store": {"size_in_bytes": 2048}}} + } + + dataset = Dataset( + name="test", + query_vectors=np.random.rand(10, 32).astype(np.float32), + base_file="dummy/base.fbin", + ) + indexes = [ + IndexConfig( + name="test_index", + algo="elastic_hnsw", + build_param={"m": 16, "ef_construction": 100}, + search_params=[{"num_candidates": 100}], + file="", + ) + ] + + with patch.object( + backend, "_check_network_available", return_value=True + ): + with patch.object( + backend, "_get_client", return_value=mock_client + ): + with patch( + "cuvs_bench.backends.elasticsearch.load_vectors", + return_value=np.random.rand(100, 32).astype(np.float32), + ): + result = backend.build( + dataset=dataset, indexes=indexes, force=False + ) + + assert result.success + assert result.index_size_bytes == 2048 + + def test_elastic_search_uses_lazy_loaded_query_vectors(self): + """ElasticBackend.search works with Dataset lazy-loading via query_file.""" + cls = get_backend_class("elastic") + backend = cls( + config={"name": "test", "host": "localhost", "port": 9200} + ) + + mock_client = MagicMock() + mock_client.search.return_value = { + "hits": {"hits": [{"_id": "0", "_score": 1.0}]} + } + + dataset = Dataset( + name="test", + training_vectors=np.random.rand(100, 32).astype(np.float32), + query_file="dummy/query.fbin", + groundtruth_neighbors=np.array([[0]], dtype=np.int32), + ) + indexes = [ + IndexConfig( + name="elastic_hnsw_test", + algo="elastic_hnsw", + build_param={}, + search_params=[{"num_candidates": 100}], + file="", + ) + ] + + with patch.object( + backend, "_check_network_available", return_value=True + ): + with patch.object( + backend, "_get_client", return_value=mock_client + ): + with patch( + "cuvs_bench.backends.elasticsearch.load_vectors", + return_value=np.random.rand(1, 32).astype(np.float32), + ): + result = backend.search( + dataset=dataset, indexes=indexes, k=1 + ) + + assert result.success + assert result.neighbors.shape == (1, 1) + def test_elastic_algo_from_config(self): """ElasticBackend.algo derives from config type (elastic_hnsw, elastic_int8_hnsw).""" cls = get_backend_class("elastic") @@ -562,7 +750,9 @@ def test_elastic_algo_from_config(self): def test_elastic_cleanup_closes_client(self): """ElasticBackend.cleanup() closes client and sets _client to None.""" cls = get_backend_class("elastic") - backend = cls(config={"name": "test", "host": "localhost", "port": 9200}) + backend = cls( + config={"name": "test", "host": "localhost", "port": 9200} + ) mock_client = MagicMock() backend._client = mock_client @@ -572,7 +762,7 @@ def test_elastic_cleanup_closes_client(self): assert backend._client is None def test_orchestrator_elastic_dry_run(self): - """BenchmarkOrchestrator with elastic backend runs dry_run without ES.""" + """Initializing the elastic orchestrator should support dry_run.""" from cuvs_bench.orchestrator import BenchmarkOrchestrator orch = BenchmarkOrchestrator(backend_type="elastic") @@ -581,7 +771,8 @@ def test_orchestrator_elastic_dry_run(self): dataset_path="/nonexistent", host="localhost", port=9200, - algorithms="test", + algorithms="elastic_hnsw", + groups="test", build=True, search=True, dry_run=True, @@ -597,7 +788,7 @@ def test_orchestrator_elastic_dry_run(self): reason="Requires pip install cuvs-bench[elastic]", ) class TestElasticHelpers: - """Tests for cuvs_bench_elastic helper functions.""" + """Tests for elastic backend helper functions.""" @pytest.fixture(autouse=True) def _ensure_elastic_registered(self): @@ -605,7 +796,8 @@ def _ensure_elastic_registered(self): registry = get_registry() if "elastic" not in registry._backends: try: - from cuvs_bench_elastic import register + from cuvs_bench.backends.elasticsearch import register + register() except ImportError: pass @@ -613,7 +805,7 @@ def _ensure_elastic_registered(self): def test_distance_to_similarity(self): """_distance_to_similarity maps cuvs distance to ES similarity.""" - from cuvs_bench_elastic.backend import _distance_to_similarity + from cuvs_bench.backends.elasticsearch import _distance_to_similarity assert _distance_to_similarity("euclidean") == "l2_norm" assert _distance_to_similarity("inner_product") == "max_inner_product" @@ -625,7 +817,7 @@ def test_load_fbin(self): import tempfile from pathlib import Path - from cuvs_bench_elastic.backend import _load_fbin + from cuvs_bench.backends.elasticsearch import _load_fbin data = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32) with tempfile.NamedTemporaryFile(suffix=".fbin", delete=False) as f: @@ -644,7 +836,7 @@ def test_load_ibin(self): import tempfile from pathlib import Path - from cuvs_bench_elastic.backend import _load_ibin + from cuvs_bench.backends.elasticsearch import _load_ibin data = np.array([[1, 2], [3, 4]], dtype=np.int32) with tempfile.NamedTemporaryFile(suffix=".ibin", delete=False) as f: From cd938459bb8278ecb822be4d3352c31d1002ce67 Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Fri, 5 Jun 2026 09:02:47 -0700 Subject: [PATCH 10/22] Improve elasticsearch backend validation Signed-off-by: Alex Fournier --- .../cuvs_bench/backends/elasticsearch.py | 67 ++++++++++- .../cuvs_bench/tests/test_modularization.py | 109 ++++++++++++++++++ python/cuvs_bench/pyproject.toml | 4 +- 3 files changed, 176 insertions(+), 4 deletions(-) diff --git a/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py b/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py index cff24841ee..685c2c456e 100644 --- a/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py +++ b/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py @@ -56,6 +56,47 @@ def _distance_to_similarity(distance: str) -> str: return m.get(distance, "l2_norm") +def _supported_elastic_algorithms_message() -> str: + """Return a readable list of supported elastic algorithm names.""" + return ", ".join(f"'{algo}'" for algo in _SUPPORTED_ALGOS) + + +def _validate_elastic_index_type(index_type: str) -> str: + """Validate an Elasticsearch dense_vector index_options type.""" + if index_type not in _SUPPORTED_INDEX_TYPES: + raise ValueError( + "Received params for unsupported Elasticsearch index type " + f"{index_type!r}. Supported index types are: " + f"{', '.join(_SUPPORTED_INDEX_TYPES)}. " + "Please check the benchmark configuration." + ) + return index_type + + +def _validate_elastic_similarity(similarity: str) -> str: + """Validate an Elasticsearch dense_vector similarity value.""" + if similarity not in _SUPPORTED_SIMILARITIES: + raise ValueError( + "Received unsupported Elasticsearch similarity " + f"{similarity!r}. Supported similarities are: " + f"{', '.join(_SUPPORTED_SIMILARITIES)}. " + "Please check the benchmark configuration or dataset distance metric." + ) + return similarity + + +def _validate_elastic_algorithm(algo_name: str) -> str: + """Validate a cuvs-bench elastic algorithm name.""" + if algo_name not in _SUPPORTED_ALGOS: + raise ValueError( + "Received unsupported algorithm " + f"{algo_name!r} for the Elasticsearch backend. " + "Supported algorithms are: " + f"{_supported_elastic_algorithms_message()}." + ) + return algo_name + + # Defaults for index creation when not specified in config _DEFAULT_INDEX_TYPE = "hnsw" _DEFAULT_M = 16 @@ -66,6 +107,12 @@ def _distance_to_similarity(distance: str) -> str: _DEFAULT_VECTOR_FIELD = "embedding" _DEFAULT_NUM_CANDIDATES = 100 +_SUPPORTED_INDEX_TYPES = ("hnsw", "int8_hnsw", "int4_hnsw", "bbq_hnsw") +_SUPPORTED_SIMILARITIES = ("l2_norm", "cosine", "max_inner_product") +_SUPPORTED_ALGOS = tuple( + f"elastic_{index_type}" for index_type in _SUPPORTED_INDEX_TYPES +) + _BUILD_PARAM_KEYS = ( "type", "m", @@ -226,6 +273,21 @@ def build( getattr(dataset, "distance_metric", None) or "euclidean" ) + index_type = build_params.get("type", _DEFAULT_INDEX_TYPE) + try: + _validate_elastic_index_type(index_type) + _validate_elastic_similarity(similarity) + except ValueError as e: + return BuildResult( + index_path="", + build_time_seconds=0.0, + index_size_bytes=0, + algorithm=self.algo, + build_params=build_params, + success=False, + error_message=str(e), + ) + try: n_vectors = len(vectors) dims = vectors.shape[1] @@ -234,7 +296,6 @@ def build( vector_field = build_params.get( "vector_field", _DEFAULT_VECTOR_FIELD ) - index_type = build_params.get("type", _DEFAULT_INDEX_TYPE) m = build_params.get("m", _DEFAULT_M) ef_construction = build_params.get( "ef_construction", _DEFAULT_EF_CONSTRUCTION @@ -698,6 +759,7 @@ def _build_benchmark_configs( search_combos, _group_meta, ) in expanded_groups: + _validate_elastic_algorithm(algo_name) if tune_mode and tune_build_params is not None: actual_build = [dict(tune_build_params)] actual_search = ( @@ -713,6 +775,9 @@ def _build_benchmark_configs( "elastic_" ): build_param["type"] = algo_name.replace("elastic_", "", 1) + _validate_elastic_index_type(build_param["type"]) + if "similarity" in build_param: + _validate_elastic_similarity(build_param["similarity"]) if tune_mode: label_prefix = f"{algo_name}_tune" diff --git a/python/cuvs_bench/cuvs_bench/tests/test_modularization.py b/python/cuvs_bench/cuvs_bench/tests/test_modularization.py index 1e5cfa4c53..42d990b5b3 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_modularization.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_modularization.py @@ -747,6 +747,74 @@ def test_elastic_algo_from_config(self): backend_int8 = cls(config={"name": "test", "type": "int8_hnsw"}) assert backend_int8.algo == "elastic_int8_hnsw" + def test_elastic_build_rejects_unsupported_index_type(self): + """Elastic build fails early with a descriptive unsupported-index error.""" + cls = get_backend_class("elastic") + backend = cls( + config={"name": "test", "host": "localhost", "port": 9200} + ) + + dataset = Dataset( + name="test", + training_vectors=np.random.rand(8, 4).astype(np.float32), + query_vectors=np.random.rand(2, 4).astype(np.float32), + distance_metric="euclidean", + ) + indexes = [ + IndexConfig( + name="elastic_scann_test", + algo="elastic_scann", + build_param={"type": "scann"}, + search_params=[{"num_candidates": 10}], + file="", + ) + ] + + with patch.object( + backend, "_check_network_available", return_value=True + ): + result = backend.build(dataset=dataset, indexes=indexes) + + assert not result.success + assert "unsupported Elasticsearch index type" in ( + result.error_message or "" + ) + assert "scann" in (result.error_message or "") + + def test_elastic_build_rejects_unsupported_similarity(self): + """Elastic build fails early with a descriptive unsupported-similarity error.""" + cls = get_backend_class("elastic") + backend = cls( + config={"name": "test", "host": "localhost", "port": 9200} + ) + + dataset = Dataset( + name="test", + training_vectors=np.random.rand(8, 4).astype(np.float32), + query_vectors=np.random.rand(2, 4).astype(np.float32), + distance_metric="euclidean", + ) + indexes = [ + IndexConfig( + name="elastic_hnsw_test", + algo="elastic_hnsw", + build_param={"type": "hnsw", "similarity": "dot_product"}, + search_params=[{"num_candidates": 10}], + file="", + ) + ] + + with patch.object( + backend, "_check_network_available", return_value=True + ): + result = backend.build(dataset=dataset, indexes=indexes) + + assert not result.success + assert "unsupported Elasticsearch similarity" in ( + result.error_message or "" + ) + assert "dot_product" in (result.error_message or "") + def test_elastic_cleanup_closes_client(self): """ElasticBackend.cleanup() closes client and sets _client to None.""" cls = get_backend_class("elastic") @@ -812,6 +880,47 @@ def test_distance_to_similarity(self): assert _distance_to_similarity("cosine") == "cosine" assert _distance_to_similarity("unknown") == "l2_norm" + def test_validate_elastic_algorithm_rejects_unknown(self): + """Unknown elastic algorithm names fail with a descriptive error.""" + from cuvs_bench.backends.elasticsearch import ( + _validate_elastic_algorithm, + ) + + with pytest.raises( + ValueError, match="unsupported algorithm" + ) as exc_info: + _validate_elastic_algorithm("elastic_scann") + + assert "elastic_hnsw" in str(exc_info.value) + + def test_validate_elastic_index_type_rejects_unknown(self): + """Unknown elastic index types fail with a descriptive error.""" + from cuvs_bench.backends.elasticsearch import ( + _validate_elastic_index_type, + ) + + with pytest.raises( + ValueError, + match="unsupported Elasticsearch index type", + ) as exc_info: + _validate_elastic_index_type("scann") + + assert "hnsw" in str(exc_info.value) + + def test_validate_elastic_similarity_rejects_unknown(self): + """Unknown elastic similarities fail with a descriptive error.""" + from cuvs_bench.backends.elasticsearch import ( + _validate_elastic_similarity, + ) + + with pytest.raises( + ValueError, + match="unsupported Elasticsearch similarity", + ) as exc_info: + _validate_elastic_similarity("dot_product") + + assert "max_inner_product" in str(exc_info.value) + def test_load_fbin(self): """_load_fbin loads big-ann-bench fbin format.""" import tempfile diff --git a/python/cuvs_bench/pyproject.toml b/python/cuvs_bench/pyproject.toml index 4e8a1bed01..7878c03374 100644 --- a/python/cuvs_bench/pyproject.toml +++ b/python/cuvs_bench/pyproject.toml @@ -42,13 +42,11 @@ classifiers = [ opensearch = [ "opensearch-py>=2.4.0", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. +elastic = ["elasticsearch>=8.0"] [project.urls] Homepage = "https://github.com/rapidsai/cuvs" -[project.optional-dependencies] -elastic = ["elasticsearch>=8.0"] - [project.entry-points."cuvs_bench.backends"] elastic = "cuvs_bench.backends.elasticsearch:register" From ad40caa6813a6d2341eff151b9ddfc7ce50aef73 Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Wed, 10 Jun 2026 20:05:48 -0700 Subject: [PATCH 11/22] Tighten optional backend handling Signed-off-by: Alex Fournier --- .../cuvs_bench/backends/elasticsearch.py | 20 ++++---- .../cuvs_bench/backends/opensearch.py | 8 ++-- .../cuvs_bench/backends/registry.py | 48 ++++++++++++++----- .../cuvs_bench/tests/test_modularization.py | 42 ++++++++++++++++ 4 files changed, 93 insertions(+), 25 deletions(-) diff --git a/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py b/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py index 685c2c456e..dc4974a7a5 100644 --- a/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py +++ b/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py @@ -401,8 +401,8 @@ def search( """Run kNN search over all search-param combinations and compute recall.""" if dry_run: return SearchResult( - neighbors=np.zeros((0, k), dtype=np.int64), - distances=np.zeros((0, k), dtype=np.float32), + neighbors=np.empty((0, k), dtype=np.int64), + distances=np.empty((0, k), dtype=np.float32), search_time_ms=0, queries_per_second=0, recall=0, @@ -414,8 +414,8 @@ def search( skip_reason = self._pre_flight_check() if skip_reason: return SearchResult( - neighbors=np.zeros((0, k), dtype=np.int64), - distances=np.zeros((0, k), dtype=np.float32), + neighbors=np.empty((0, k), dtype=np.int64), + distances=np.empty((0, k), dtype=np.float32), search_time_ms=0, queries_per_second=0, recall=0, @@ -427,8 +427,8 @@ def search( if not indexes: return SearchResult( - neighbors=np.zeros((0, k), dtype=np.int64), - distances=np.zeros((0, k), dtype=np.float32), + neighbors=np.empty((0, k), dtype=np.int64), + distances=np.empty((0, k), dtype=np.float32), search_time_ms=0, queries_per_second=0, recall=0, @@ -441,8 +441,8 @@ def search( query_vectors = dataset.query_vectors if query_vectors.size == 0: return SearchResult( - neighbors=np.zeros((0, k), dtype=np.int64), - distances=np.zeros((0, k), dtype=np.float32), + neighbors=np.empty((0, k), dtype=np.int64), + distances=np.empty((0, k), dtype=np.float32), search_time_ms=0, queries_per_second=0, recall=0, @@ -562,8 +562,8 @@ def search( ) except Exception as e: return SearchResult( - neighbors=np.zeros((0, k), dtype=np.int64), - distances=np.zeros((0, k), dtype=np.float32), + neighbors=np.empty((0, k), dtype=np.int64), + distances=np.empty((0, k), dtype=np.float32), search_time_ms=0, queries_per_second=0, recall=0, diff --git a/python/cuvs_bench/cuvs_bench/backends/opensearch.py b/python/cuvs_bench/cuvs_bench/backends/opensearch.py index 026fe98fd6..3f7e1c04ac 100644 --- a/python/cuvs_bench/cuvs_bench/backends/opensearch.py +++ b/python/cuvs_bench/cuvs_bench/backends/opensearch.py @@ -559,8 +559,8 @@ def _failed_search_result( search_params: Optional[List[Dict[str, Any]]] = None, ) -> SearchResult: return SearchResult( - neighbors=np.zeros((0, k), dtype=np.int64), - distances=np.zeros((0, k), dtype=np.float32), + neighbors=np.empty((0, k), dtype=np.int64), + distances=np.empty((0, k), dtype=np.float32), search_time_ms=0.0, queries_per_second=0.0, recall=0.0, @@ -871,8 +871,8 @@ def search( ) return SearchResult( - neighbors=np.zeros((0, k), dtype=np.int64), - distances=np.zeros((0, k), dtype=np.float32), + neighbors=np.empty((0, k), dtype=np.int64), + distances=np.empty((0, k), dtype=np.float32), search_time_ms=0.0, queries_per_second=0.0, recall=0.0, diff --git a/python/cuvs_bench/cuvs_bench/backends/registry.py b/python/cuvs_bench/cuvs_bench/backends/registry.py index 0059ea78f0..b397dfc9d6 100644 --- a/python/cuvs_bench/cuvs_bench/backends/registry.py +++ b/python/cuvs_bench/cuvs_bench/backends/registry.py @@ -21,6 +21,38 @@ _BACKENDS_GROUP = "cuvs_bench.backends" _CONFIG_LOADERS_GROUP = "cuvs_bench.config_loaders" +_OPTIONAL_BACKEND_EXTRAS = { + "elastic": ("elasticsearch", "pip install cuvs-bench[elastic]"), + "opensearch": ("opensearchpy", "pip install cuvs-bench[opensearch]"), +} + + +def _optional_backend_install_hint(name: str) -> str: + """Return an install hint for optional backends.""" + extra = _OPTIONAL_BACKEND_EXTRAS.get(name) + if extra is None: + return "" + return f" Install with: {extra[1]}" + + +def _rewrite_optional_backend_import_error( + name: str, error: ImportError +) -> ImportError | None: + """Rewrite ImportError with an install hint for known optional backends.""" + extra = _OPTIONAL_BACKEND_EXTRAS.get(name) + if extra is None: + return None + + module_name, install_cmd = extra + message = str(error).lower() + if module_name.lower() in message: + backend_name = name.capitalize() + return ImportError( + f"{backend_name} backend requires the '{name}' extra. " + f"Install with: {install_cmd}" + ) + return None + class BackendRegistry: """ @@ -401,11 +433,9 @@ def _try_load_plugin(name: str) -> None: try: ep.load()() except ImportError as e: - if "elasticsearch" in str(e).lower() or "elasticsearch" in str(e): - raise ImportError( - f"Elasticsearch backend requires the 'elastic' extra. " - f"Install with: pip install cuvs-bench[elastic]" - ) from e + rewritten = _rewrite_optional_backend_import_error(name, e) + if rewritten is not None: + raise rewritten from e raise return # Plugin loaded successfully @@ -432,9 +462,7 @@ def get_backend_class(name: str) -> Type[BenchmarkBackend]: _try_load_plugin(name) if name not in registry._backends: available = ", ".join(registry._backends.keys()) - hint = "" - if name == "elastic": - hint = " Install with: pip install cuvs-bench[elastic]" + hint = _optional_backend_install_hint(name) raise ValueError( f"Backend '{name}' not found. Available backends: {available or '(none)'}.{hint}" ) @@ -505,9 +533,7 @@ def get_config_loader(name: str) -> Type: _try_load_plugin(name) if name not in _CONFIG_LOADER_REGISTRY: available = ", ".join(_CONFIG_LOADER_REGISTRY.keys()) or "none" - hint = "" - if name == "elastic": - hint = " Install with: pip install cuvs-bench[elastic]" + hint = _optional_backend_install_hint(name) raise ValueError( f"Unknown config loader for backend: '{name}'. Available: {available}.{hint}" ) diff --git a/python/cuvs_bench/cuvs_bench/tests/test_modularization.py b/python/cuvs_bench/cuvs_bench/tests/test_modularization.py index 42d990b5b3..f82f2c69ba 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_modularization.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_modularization.py @@ -284,6 +284,34 @@ def test_import_error_with_elasticsearch_message_raises_helpful_error( msg = str(exc_info.value) assert "pip install cuvs-bench[elastic]" in msg + def test_import_error_with_opensearch_message_raises_helpful_error( + self, + ): + """Mock entry point raising ImportError(opensearchpy) -> our install message.""" + registry = get_registry() + if "opensearch" in registry._backends: + registry.unregister("opensearch") + unregister_config_loader("opensearch") + + mock_ep = MagicMock() + mock_ep.name = "opensearch" + mock_ep.load.side_effect = ImportError( + "No module named 'opensearchpy'" + ) + + mock_eps = MagicMock() + mock_eps.select.return_value = [mock_ep] + + with patch( + "cuvs_bench.backends.registry.importlib.metadata.entry_points", + return_value=mock_eps, + ): + with pytest.raises(ImportError) as exc_info: + get_backend_class("opensearch") + + msg = str(exc_info.value) + assert "pip install cuvs-bench[opensearch]" in msg + def test_import_error_unrelated_propagates(self): """Mock entry point: unrelated ImportError propagates unchanged.""" mock_ep = MagicMock() @@ -338,6 +366,20 @@ def test_no_entry_point_for_name_raises_value_error(self): assert "nonexistent_mock_xyz" in msg assert "cpp_gbench" in msg + def test_unknown_optional_backend_includes_install_hint(self): + """Missing optional backend names include install hints in the error.""" + mock_eps = MagicMock() + mock_eps.select.return_value = [] + + with patch( + "cuvs_bench.backends.registry.importlib.metadata.entry_points", + return_value=mock_eps, + ): + with pytest.raises(ValueError) as exc_info: + get_backend_class("opensearch") + + assert "pip install cuvs-bench[opensearch]" in str(exc_info.value) + def _elasticsearch_installed(): try: From a0a4b673b14bda4e24dbc29f4ce21c4ea2c805af Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Mon, 13 Jul 2026 09:45:52 -0400 Subject: [PATCH 12/22] Remove unused elastic vector loader wrappers --- .../cuvs_bench/cuvs_bench/backends/elasticsearch.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py b/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py index dc4974a7a5..ebb1fd2f1f 100644 --- a/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py +++ b/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py @@ -24,7 +24,6 @@ from .base import BenchmarkBackend, BuildResult, Dataset, SearchResult from .registry import register_backend, register_config_loader -from ._utils import load_vectors from ..orchestrator.config_loaders import ( BenchmarkConfig, ConfigLoader, @@ -36,16 +35,6 @@ from elasticsearch import Elasticsearch -def _load_fbin(path: Path) -> np.ndarray: - """Load big-ann-bench fbin format via shared vector loader.""" - return load_vectors(os.fspath(path)) - - -def _load_ibin(path: Path) -> np.ndarray: - """Load big-ann-bench ibin format via shared vector loader.""" - return load_vectors(os.fspath(path)) - - def _distance_to_similarity(distance: str) -> str: """Map cuvs-bench distance metric to ES dense_vector similarity.""" m = { From fe34fd46557d44789bc048a6fecdfb60c42552bb Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Mon, 13 Jul 2026 09:55:44 -0400 Subject: [PATCH 13/22] Remove redundant elastic config loader override --- .../cuvs_bench/cuvs_bench/backends/elasticsearch.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py b/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py index ebb1fd2f1f..95f8fac116 100644 --- a/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py +++ b/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py @@ -598,17 +598,6 @@ def __init__(self, config_path: Optional[str] = None): def backend_type(self) -> str: return "elastic" - def load( - self, - dataset: str = "", - dataset_path: str = "", - **kwargs, - ) -> Tuple[DatasetConfig, List[BenchmarkConfig]]: - """Load Elasticsearch benchmark configuration via shared ConfigLoader flow.""" - return super().load( - dataset=dataset, dataset_path=dataset_path, **kwargs - ) - def _discover_algo_groups( self, dataset_conf, From 1b81f44699871f480f522cc2105290cf57f2c48f Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Mon, 13 Jul 2026 10:43:51 -0400 Subject: [PATCH 14/22] Use elastic index config as build param source --- python/cuvs_bench/cuvs_bench/backends/elasticsearch.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py b/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py index 95f8fac116..80fa1ac60c 100644 --- a/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py +++ b/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py @@ -210,9 +210,6 @@ def build( index_name = self.config.get("index_name", "cuvs_bench_vectors") idx = indexes[0] if indexes else None build_params = dict(idx.build_param or {}) if idx else {} - for k, v in self.config.items(): - if k not in build_params and k in _BUILD_PARAM_KEYS: - build_params[k] = v try: client = self._get_client() From fcc13f546f973ce2bfb1098922a193cb79dbadb9 Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Mon, 13 Jul 2026 11:11:09 -0400 Subject: [PATCH 15/22] Remove elastic benchmark convenience wrappers --- .../cuvs_bench/backends/elasticsearch.py | 84 ------------------- 1 file changed, 84 deletions(-) diff --git a/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py b/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py index 80fa1ac60c..40e9fe838b 100644 --- a/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py +++ b/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py @@ -810,87 +810,3 @@ def register() -> None: register_backend("elastic", ElasticBackend) if "elastic" not in _CONFIG_LOADER_REGISTRY: register_config_loader("elastic", ElasticConfigLoader) - - -# ── Convenience API ─────────────────────────────────────────────────────────── - - -def run_build( - dataset: str = "test-data", - dataset_path: str = "./datasets", - host: str = "localhost", - port: int = 9200, - algorithms: str = "test", - force: bool = False, - **kwargs, -): - """Build an Elasticsearch vector index. Returns list of BuildResult.""" - register() - from cuvs_bench.orchestrator import BenchmarkOrchestrator - - orch = BenchmarkOrchestrator(backend_type="elastic") - return orch.run_benchmark( - build=True, - search=False, - dataset=dataset, - dataset_path=dataset_path, - host=host, - port=port, - algorithms=algorithms, - force=force, - **kwargs, - ) - - -def run_search( - dataset: str = "test-data", - dataset_path: str = "./datasets", - host: str = "localhost", - port: int = 9200, - algorithms: str = "test", - **kwargs, -): - """Run kNN search against an existing Elasticsearch index. Returns list of SearchResult.""" - register() - from cuvs_bench.orchestrator import BenchmarkOrchestrator - - orch = BenchmarkOrchestrator(backend_type="elastic") - return orch.run_benchmark( - build=False, - search=True, - dataset=dataset, - dataset_path=dataset_path, - host=host, - port=port, - algorithms=algorithms, - **kwargs, - ) - - -def run_benchmark( - dataset: str = "test-data", - dataset_path: str = "./datasets", - host: str = "localhost", - port: int = 9200, - algorithms: str = "test", - build: bool = True, - search: bool = True, - force: bool = False, - **kwargs, -): - """Run build and/or search. Returns list of results.""" - register() - from cuvs_bench.orchestrator import BenchmarkOrchestrator - - orch = BenchmarkOrchestrator(backend_type="elastic") - return orch.run_benchmark( - build=build, - search=search, - dataset=dataset, - dataset_path=dataset_path, - host=host, - port=port, - algorithms=algorithms, - force=force, - **kwargs, - ) From cb5b0fda46ac869704b1df352fb08de71401aa22 Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Mon, 13 Jul 2026 11:12:45 -0400 Subject: [PATCH 16/22] Fix elastic test registration fallback --- python/cuvs_bench/cuvs_bench/tests/conftest.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/cuvs_bench/cuvs_bench/tests/conftest.py b/python/cuvs_bench/cuvs_bench/tests/conftest.py index c003df9b3a..746d2a4d58 100644 --- a/python/cuvs_bench/cuvs_bench/tests/conftest.py +++ b/python/cuvs_bench/cuvs_bench/tests/conftest.py @@ -18,6 +18,7 @@ def pytest_configure(config): try: from cuvs_bench_elastic import register - register() except ImportError: - pass + from cuvs_bench.backends.elasticsearch import register + + register() From 8e03291738980a19915d76f1b2b98392f874ef9c Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Mon, 13 Jul 2026 11:15:56 -0400 Subject: [PATCH 17/22] Align elastic benchmark dependency metadata --- dependencies.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dependencies.yaml b/dependencies.yaml index f11df3591e..3585c7eb58 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -632,7 +632,7 @@ dependencies: common: - output_types: [conda, pyproject, requirements] packages: - - cuvs-bench-elastic>=26.4.0 + - elasticsearch>=8.0 bench_python_opensearch: common: - output_types: [conda, pyproject] From da8aafac5469dcdf2551d20fb7ef09518159035f Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Mon, 13 Jul 2026 11:25:32 -0400 Subject: [PATCH 18/22] Normalize elastic algorithm variants in loader --- .../cuvs_bench/backends/elasticsearch.py | 23 ++++++++++++++++--- .../cuvs_bench/tests/test_modularization.py | 17 ++++++++++++++ 2 files changed, 37 insertions(+), 3 deletions(-) diff --git a/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py b/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py index 40e9fe838b..cfef0832fa 100644 --- a/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py +++ b/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py @@ -86,6 +86,12 @@ def _validate_elastic_algorithm(algo_name: str) -> str: return algo_name +def _elastic_config_name_for_algorithm(algo_name: str) -> str: + """Map a supported elastic algorithm name to its shared config name.""" + _validate_elastic_algorithm(algo_name) + return "elastic_hnsw" + + # Defaults for index creation when not specified in config _DEFAULT_INDEX_TYPE = "hnsw" _DEFAULT_M = 16 @@ -668,10 +674,21 @@ def _discover_algo_groups( if allowed_groups is None and not algo_group_map: allowed_groups = ["base"] + algo_confs_by_name = { + algo_conf["name"]: algo_conf for algo_conf in elastic_algos + } + + requested_algos = ( + allowed_algos + if allowed_algos is not None + else list(algo_group_map) or list(algo_confs_by_name) + ) + result = [] - for algo_conf in elastic_algos: - algo_name = algo_conf["name"] - if allowed_algos and algo_name not in allowed_algos: + for algo_name in requested_algos: + config_algo_name = _elastic_config_name_for_algorithm(algo_name) + algo_conf = algo_confs_by_name.get(config_algo_name) + if algo_conf is None: continue groups = dict(algo_conf.get("groups", {})) diff --git a/python/cuvs_bench/cuvs_bench/tests/test_modularization.py b/python/cuvs_bench/cuvs_bench/tests/test_modularization.py index f82f2c69ba..9b8dbec5f6 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_modularization.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_modularization.py @@ -471,6 +471,23 @@ def test_elastic_config_loader_group_not_found_raises(self): groups="nonexistent_group_xyz", ) + def test_elastic_config_loader_normalizes_variant_algorithms(self): + """Variant elastic_* names resolve through the shared elastic config.""" + loader_cls = get_config_loader("elastic") + loader = loader_cls() + _, benchmark_configs = loader.load( + dataset="glove-50-angular", + dataset_path="", + algorithms="elastic_int8_hnsw", + groups="test", + ) + + assert len(benchmark_configs) == 1 + config = benchmark_configs[0] + assert config.indexes[0].algo == "elastic_int8_hnsw" + assert config.indexes[0].build_param["type"] == "int8_hnsw" + assert config.backend_config["type"] == "int8_hnsw" + def test_elastic_dry_run_build(self): """ElasticBackend.build(dry_run=True) returns synthetic result without ES.""" cls = get_backend_class("elastic") From 3f80d6b3f650796de5fc1d19995117b9564b564c Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Mon, 13 Jul 2026 11:31:33 -0400 Subject: [PATCH 19/22] Load split plugin entry points across registry groups --- .../cuvs_bench/backends/registry.py | 1 - .../cuvs_bench/tests/test_modularization.py | 86 +++++++++++++++++++ 2 files changed, 86 insertions(+), 1 deletion(-) diff --git a/python/cuvs_bench/cuvs_bench/backends/registry.py b/python/cuvs_bench/cuvs_bench/backends/registry.py index b397dfc9d6..7e15afadff 100644 --- a/python/cuvs_bench/cuvs_bench/backends/registry.py +++ b/python/cuvs_bench/cuvs_bench/backends/registry.py @@ -437,7 +437,6 @@ def _try_load_plugin(name: str) -> None: if rewritten is not None: raise rewritten from e raise - return # Plugin loaded successfully def get_backend_class(name: str) -> Type[BenchmarkBackend]: diff --git a/python/cuvs_bench/cuvs_bench/tests/test_modularization.py b/python/cuvs_bench/cuvs_bench/tests/test_modularization.py index 9b8dbec5f6..abd1f07bc2 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_modularization.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_modularization.py @@ -255,6 +255,92 @@ def test_valid_plugin_loads_via_mock_entry_point(self): get_registry().unregister(self._MOCK_PLUGIN_NAME) unregister_config_loader(self._MOCK_PLUGIN_NAME) + def test_config_loader_prefers_loader_entry_point_group(self): + """Config loader lookup handles plugins split across entry-point groups.""" + + class StubSplitBackend(BenchmarkBackend): + @property + def algo(self): + return "stub_split" + + def build(self, dataset, indexes, force=False, dry_run=False): + return BuildResult( + index_path="", + build_time_seconds=0, + index_size_bytes=0, + algorithm=self.algo, + build_params={}, + success=True, + ) + + def search( + self, + dataset, + indexes, + k, + batch_size=10000, + mode="latency", + force=False, + search_threads=None, + dry_run=False, + ): + return SearchResult( + neighbors=np.empty((0, k)), + distances=np.empty((0, k)), + search_time_ms=0, + queries_per_second=0, + recall=0, + algorithm=self.algo, + search_params=[], + success=True, + ) + + class StubSplitLoader(ConfigLoader): + @property + def backend_type(self): + return TestPluginLoaderMocked._MOCK_PLUGIN_NAME + + def load(self, **kwargs): + raise NotImplementedError("Stub split loader") + + def register_backend_only(): + register_backend( + TestPluginLoaderMocked._MOCK_PLUGIN_NAME, + StubSplitBackend, + ) + + def register_loader_only(): + register_config_loader( + TestPluginLoaderMocked._MOCK_PLUGIN_NAME, + StubSplitLoader, + ) + + mock_backend_ep = MagicMock() + mock_backend_ep.name = self._MOCK_PLUGIN_NAME + mock_backend_ep.load.return_value = register_backend_only + + mock_loader_ep = MagicMock() + mock_loader_ep.name = self._MOCK_PLUGIN_NAME + mock_loader_ep.load.return_value = register_loader_only + + def mock_entry_points(*args, **kwargs): + group = kwargs.get("group") + if group == "cuvs_bench.backends": + return [mock_backend_ep] + if group == "cuvs_bench.config_loaders": + return [mock_loader_ep] + return [] + + with patch( + "cuvs_bench.backends.registry.importlib.metadata.entry_points", + side_effect=mock_entry_points, + ): + loader_cls = get_config_loader(self._MOCK_PLUGIN_NAME) + assert loader_cls is StubSplitLoader + + get_registry().unregister(self._MOCK_PLUGIN_NAME) + unregister_config_loader(self._MOCK_PLUGIN_NAME) + def test_import_error_with_elasticsearch_message_raises_helpful_error( self, ): From e30bdd1588824264660e3a4bc1ceb089445659c1 Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Mon, 13 Jul 2026 12:30:36 -0400 Subject: [PATCH 20/22] Route elastic recall through orchestrator --- .../cuvs_bench/backends/elasticsearch.py | 25 ++----------------- .../cuvs_bench/tests/test_modularization.py | 3 +++ 2 files changed, 5 insertions(+), 23 deletions(-) diff --git a/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py b/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py index cfef0832fa..ffe7bedc20 100644 --- a/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py +++ b/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py @@ -390,7 +390,7 @@ def search( search_threads: Optional[int] = None, dry_run: bool = False, ) -> SearchResult: - """Run kNN search over all search-param combinations and compute recall.""" + """Run kNN search over all search-param combinations.""" if dry_run: return SearchResult( neighbors=np.empty((0, k), dtype=np.int64), @@ -450,8 +450,6 @@ def search( try: n_queries = len(query_vectors) - groundtruth = dataset.groundtruth_neighbors - index_name = self.config.get("index_name", "cuvs_bench_vectors") index_cfg = indexes[0] search_params_list = index_cfg.search_params or [{}] @@ -495,26 +493,11 @@ def search( n_queries / (elapsed_ms / 1000) if elapsed_ms > 0 else 0.0 ) - recall = 0.0 - if groundtruth is not None: - gt_k = min(k, groundtruth.shape[1]) - n_correct = sum( - len( - set(neighbors[i, :k].tolist()) - & set(groundtruth[i, :gt_k].tolist()) - ) - for i in range(n_queries) - ) - recall = ( - n_correct / (n_queries * gt_k) if gt_k > 0 else 0.0 - ) - per_param_results.append( { "search_params": sp, "search_time_ms": elapsed_ms, "queries_per_second": qps, - "recall": recall, "p50_ms": float(np.percentile(latencies, 50)), "p95_ms": float(np.percentile(latencies, 95)), "p99_ms": float(np.percentile(latencies, 99)), @@ -523,9 +506,6 @@ def search( last_neighbors = neighbors last_distances = distances - avg_recall = float( - np.mean([r["recall"] for r in per_param_results]) - ) avg_qps = float( np.mean([r["queries_per_second"] for r in per_param_results]) ) @@ -538,7 +518,7 @@ def search( distances=last_distances, search_time_ms=total_ms, queries_per_second=avg_qps, - recall=avg_recall, + recall=0.0, algorithm=self.algo, search_params=search_params_list, latency_percentiles={ @@ -548,7 +528,6 @@ def search( }, metadata={ "per_search_param_results": per_param_results, - "recall_is_authoritative": True, }, success=True, ) diff --git a/python/cuvs_bench/cuvs_bench/tests/test_modularization.py b/python/cuvs_bench/cuvs_bench/tests/test_modularization.py index abd1f07bc2..5f31f46498 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_modularization.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_modularization.py @@ -882,6 +882,9 @@ def test_elastic_search_uses_lazy_loaded_query_vectors(self): assert result.success assert result.neighbors.shape == (1, 1) + assert result.recall == 0.0 + assert "recall_is_authoritative" not in result.metadata + assert "recall" not in result.metadata["per_search_param_results"][0] def test_elastic_algo_from_config(self): """ElasticBackend.algo derives from config type (elastic_hnsw, elastic_int8_hnsw).""" From df0cb52efb9feab96a342020a504deaa64548325 Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Mon, 13 Jul 2026 12:43:49 -0400 Subject: [PATCH 21/22] Clean up elastic loader wrapper tests --- .../cuvs_bench/tests/test_modularization.py | 42 +------------------ 1 file changed, 2 insertions(+), 40 deletions(-) diff --git a/python/cuvs_bench/cuvs_bench/tests/test_modularization.py b/python/cuvs_bench/cuvs_bench/tests/test_modularization.py index 5f31f46498..6d436e5686 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_modularization.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_modularization.py @@ -828,7 +828,7 @@ def test_elastic_build_uses_lazy_loaded_training_vectors(self): backend, "_get_client", return_value=mock_client ): with patch( - "cuvs_bench.backends.elasticsearch.load_vectors", + "cuvs_bench.backends._utils.load_vectors", return_value=np.random.rand(100, 32).astype(np.float32), ): result = backend.build( @@ -873,7 +873,7 @@ def test_elastic_search_uses_lazy_loaded_query_vectors(self): backend, "_get_client", return_value=mock_client ): with patch( - "cuvs_bench.backends.elasticsearch.load_vectors", + "cuvs_bench.backends._utils.load_vectors", return_value=np.random.rand(1, 32).astype(np.float32), ): result = backend.search( @@ -1068,41 +1068,3 @@ def test_validate_elastic_similarity_rejects_unknown(self): _validate_elastic_similarity("dot_product") assert "max_inner_product" in str(exc_info.value) - - def test_load_fbin(self): - """_load_fbin loads big-ann-bench fbin format.""" - import tempfile - from pathlib import Path - - from cuvs_bench.backends.elasticsearch import _load_fbin - - data = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32) - with tempfile.NamedTemporaryFile(suffix=".fbin", delete=False) as f: - path = Path(f.name) - try: - with open(path, "wb") as f: - np.array([2, 2], dtype=np.uint32).tofile(f) - data.tofile(f) - loaded = _load_fbin(path) - np.testing.assert_array_equal(loaded, data) - finally: - path.unlink(missing_ok=True) - - def test_load_ibin(self): - """_load_ibin loads big-ann-bench ibin format.""" - import tempfile - from pathlib import Path - - from cuvs_bench.backends.elasticsearch import _load_ibin - - data = np.array([[1, 2], [3, 4]], dtype=np.int32) - with tempfile.NamedTemporaryFile(suffix=".ibin", delete=False) as f: - path = Path(f.name) - try: - with open(path, "wb") as f: - np.array([2, 2], dtype=np.uint32).tofile(f) - data.tofile(f) - loaded = _load_ibin(path) - np.testing.assert_array_equal(loaded, data) - finally: - path.unlink(missing_ok=True) From 30d85e694efadd2f653e43e977ded5adb85a0676 Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Mon, 13 Jul 2026 13:06:32 -0400 Subject: [PATCH 22/22] Fix elastic pytest regressions --- .../cuvs_bench/backends/elasticsearch.py | 94 ++++++++++--------- .../cuvs_bench/backends/registry.py | 13 +++ .../cuvs_bench/tests/test_modularization.py | 2 +- 3 files changed, 65 insertions(+), 44 deletions(-) diff --git a/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py b/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py index ffe7bedc20..3e62cfb18f 100644 --- a/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py +++ b/python/cuvs_bench/cuvs_bench/backends/elasticsearch.py @@ -92,6 +92,12 @@ def _elastic_config_name_for_algorithm(algo_name: str) -> str: return "elastic_hnsw" +def _elastic_index_type_for_algorithm(algo_name: str) -> str: + """Map a supported elastic algorithm name to an ES index_options type.""" + _validate_elastic_algorithm(algo_name) + return algo_name.replace("elastic_", "", 1) + + # Defaults for index creation when not specified in config _DEFAULT_INDEX_TYPE = "hnsw" _DEFAULT_M = 16 @@ -201,6 +207,44 @@ def build( success=True, ) + index_name = self.config.get("index_name", "cuvs_bench_vectors") + idx = indexes[0] if indexes else None + build_params = dict(idx.build_param or {}) if idx else {} + + vectors = dataset.training_vectors + if vectors.size == 0: + return BuildResult( + index_path="", + build_time_seconds=0.0, + index_size_bytes=0, + algorithm=self.algo, + build_params={}, + success=False, + error_message=( + "training_vectors are required for Elasticsearch backend " + "(directly or via dataset.base_file)" + ), + ) + + similarity = build_params.get("similarity") or _distance_to_similarity( + getattr(dataset, "distance_metric", None) or "euclidean" + ) + + index_type = build_params.get("type", _DEFAULT_INDEX_TYPE) + try: + _validate_elastic_index_type(index_type) + _validate_elastic_similarity(similarity) + except ValueError as e: + return BuildResult( + index_path="", + build_time_seconds=0.0, + index_size_bytes=0, + algorithm=self.algo, + build_params=build_params, + success=False, + error_message=str(e), + ) + skip_reason = self._pre_flight_check() if skip_reason: return BuildResult( @@ -213,10 +257,6 @@ def build( error_message=f"pre-flight check failed: {skip_reason}", ) - index_name = self.config.get("index_name", "cuvs_bench_vectors") - idx = indexes[0] if indexes else None - build_params = dict(idx.build_param or {}) if idx else {} - try: client = self._get_client() if client.indices.exists(index=index_name): @@ -245,41 +285,6 @@ def build( error_message=str(e), ) - vectors = dataset.training_vectors - if vectors.size == 0: - return BuildResult( - index_path="", - build_time_seconds=0.0, - index_size_bytes=0, - algorithm=self.algo, - build_params={}, - success=False, - error_message=( - "training_vectors are required for Elasticsearch backend " - "(directly or via dataset.base_file)" - ), - ) - - # similarity: from config, or derive from dataset distance - similarity = build_params.get("similarity") or _distance_to_similarity( - getattr(dataset, "distance_metric", None) or "euclidean" - ) - - index_type = build_params.get("type", _DEFAULT_INDEX_TYPE) - try: - _validate_elastic_index_type(index_type) - _validate_elastic_similarity(similarity) - except ValueError as e: - return BuildResult( - index_path="", - build_time_seconds=0.0, - index_size_bytes=0, - algorithm=self.algo, - build_params=build_params, - success=False, - error_message=str(e), - ) - try: n_vectors = len(vectors) dims = vectors.shape[1] @@ -742,10 +747,13 @@ def _build_benchmark_configs( for build_param in actual_build: build_param = dict(build_param) - if "type" not in build_param and algo_name.startswith( - "elastic_" - ): - build_param["type"] = algo_name.replace("elastic_", "", 1) + config_algo_name = _elastic_config_name_for_algorithm( + algo_name + ) + if "type" not in build_param or algo_name != config_algo_name: + build_param["type"] = _elastic_index_type_for_algorithm( + algo_name + ) _validate_elastic_index_type(build_param["type"]) if "similarity" in build_param: _validate_elastic_similarity(build_param["similarity"]) diff --git a/python/cuvs_bench/cuvs_bench/backends/registry.py b/python/cuvs_bench/cuvs_bench/backends/registry.py index 7e15afadff..bfaf14417b 100644 --- a/python/cuvs_bench/cuvs_bench/backends/registry.py +++ b/python/cuvs_bench/cuvs_bench/backends/registry.py @@ -420,6 +420,7 @@ def _try_load_plugin(name: str) -> None: Raises ImportError with install instructions if the plugin requires an optional dependency that is not installed. """ + registry = get_registry() for group in (_BACKENDS_GROUP, _CONFIG_LOADERS_GROUP): try: eps = importlib.metadata.entry_points(group=group) @@ -430,6 +431,13 @@ def _try_load_plugin(name: str) -> None: else: eps = [e for e in eps if e.name == name] for ep in eps: + if group == _BACKENDS_GROUP and name in registry._backends: + break + if ( + group == _CONFIG_LOADERS_GROUP + and name in _CONFIG_LOADER_REGISTRY + ): + break try: ep.load()() except ImportError as e: @@ -437,6 +445,11 @@ def _try_load_plugin(name: str) -> None: if rewritten is not None: raise rewritten from e raise + if ( + name in registry._backends + and name in _CONFIG_LOADER_REGISTRY + ): + return def get_backend_class(name: str) -> Type[BenchmarkBackend]: diff --git a/python/cuvs_bench/cuvs_bench/tests/test_modularization.py b/python/cuvs_bench/cuvs_bench/tests/test_modularization.py index 6d436e5686..8e07616da1 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_modularization.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_modularization.py @@ -542,7 +542,7 @@ def test_elastic_config_loader_dataset_not_found_raises(self): """Config loader raises ValueError for unknown dataset.""" loader_cls = get_config_loader("elastic") loader = loader_cls() - with pytest.raises(ValueError, match="not found"): + with pytest.raises(ValueError, match="Could not find"): loader.load(dataset="nonexistent_dataset_xyz", dataset_path="") def test_elastic_config_loader_group_not_found_raises(self):