From 756b69170c6fbe08c284c3a6308a004fbe93163b Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Tue, 28 Jul 2026 11:54:02 -0500 Subject: [PATCH 1/3] Initial commit Signed-off-by: James Bourbeau --- python/cuvs_bench/cuvs_bench/backends/base.py | 15 +- .../cuvs_bench/backends/cpp_gbench.py | 187 ++++++----- .../cuvs_bench/backends/opensearch.py | 142 ++++---- .../cuvs_bench/orchestrator/orchestrator.py | 129 ++++---- python/cuvs_bench/cuvs_bench/run/__main__.py | 124 ++++--- .../cuvs_bench/cuvs_bench/run/data_export.py | 166 +++++++++- .../cuvs_bench/tests/test_cpp_gbench.py | 6 +- .../cuvs_bench/tests/test_data_export.py | 304 ++++++++++++++++++ .../cuvs_bench/tests/test_opensearch.py | 67 +++- .../cuvs_bench/tests/test_registry.py | 48 +-- 10 files changed, 883 insertions(+), 305 deletions(-) create mode 100644 python/cuvs_bench/cuvs_bench/tests/test_data_export.py diff --git a/python/cuvs_bench/cuvs_bench/backends/base.py b/python/cuvs_bench/cuvs_bench/backends/base.py index bf802d2128..15209decb6 100644 --- a/python/cuvs_bench/cuvs_bench/backends/base.py +++ b/python/cuvs_bench/cuvs_bench/backends/base.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ @@ -274,8 +274,8 @@ class SearchResult: algorithm : str Algorithm name search_params : List[Dict[str, Any]] - List of search parameter combinations used (e.g., [{"nprobe": 1}, {"nprobe": 5}]) - All are batched into one C++ command (matches runners.py behavior) + Search parameter combinations represented by this result. Backends + normally return one result per combination. latency_percentiles : Optional[Dict[str, float]] Latency percentiles in milliseconds (p50, p95, p99) gpu_time_seconds : Optional[float] @@ -421,7 +421,7 @@ def search( force: bool = False, search_threads: Optional[int] = None, dry_run: bool = False, - ) -> SearchResult: + ) -> List[SearchResult]: """ Search for nearest neighbors using the built indexes. @@ -452,8 +452,11 @@ def search( Returns ------- - SearchResult - Search timing, results, and recall metrics + List[SearchResult] + One or more search result objects. Backends normally return one + result per independently measurable search point, but may return + an aggregate result when their native output remains authoritative + (for example, the C++ Google Benchmark backend). Raises ------ diff --git a/python/cuvs_bench/cuvs_bench/backends/cpp_gbench.py b/python/cuvs_bench/cuvs_bench/backends/cpp_gbench.py index a888e957e5..1e2d861456 100644 --- a/python/cuvs_bench/cuvs_bench/backends/cpp_gbench.py +++ b/python/cuvs_bench/cuvs_bench/backends/cpp_gbench.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ @@ -309,7 +309,7 @@ def search( force: bool = False, search_threads: Optional[int] = None, dry_run: bool = False, - ) -> SearchResult: + ) -> List[SearchResult]: """ Search using C++ Google Benchmark executable. @@ -339,40 +339,45 @@ def search( Returns ------- - SearchResult - Search timing, recall, and QPS (aggregated across all indexes) + List[SearchResult] + A single aggregate result containing search timing, recall, and + QPS across all indexes. """ if not indexes: - return SearchResult( - neighbors=np.array([]), - distances=np.array([]), - search_time_ms=0.0, - queries_per_second=0.0, - recall=0.0, - algorithm="", - search_params=[], - metadata={"skipped": True, "reason": "no_indexes"}, - success=True, - ) + return [ + SearchResult( + neighbors=np.array([]), + distances=np.array([]), + search_time_ms=0.0, + queries_per_second=0.0, + recall=0.0, + algorithm="", + search_params=[], + metadata={"skipped": True, "reason": "no_indexes"}, + success=True, + ) + ] first_index = indexes[0] # Pre-flight check (GPU, network, etc.) skip_reason = self._pre_flight_check() if skip_reason: - return SearchResult( - neighbors=np.array([]), - distances=np.array([]), - search_time_ms=0.0, - queries_per_second=0.0, - recall=0.0, - algorithm=first_index.algo, - search_params=first_index.search_params - if first_index.search_params - else [], - metadata={"skipped": True, "reason": skip_reason}, - success=True, - ) + return [ + SearchResult( + neighbors=np.array([]), + distances=np.array([]), + search_time_ms=0.0, + queries_per_second=0.0, + recall=0.0, + algorithm=first_index.algo, + search_params=first_index.search_params + if first_index.search_params + else [], + metadata={"skipped": True, "reason": skip_reason}, + success=True, + ) + ] # Note: runners.py doesn't validate and lets C++ fail. We validate here for # better Python-side error messages. @@ -470,21 +475,23 @@ def search( f"Benchmark command for {self.output_filename[1]}:\n{' '.join(cmd)}\n" ) Path(temp_config_path).unlink(missing_ok=True) - return SearchResult( - neighbors=np.array([]), - distances=np.array([]), - search_time_ms=0.0, - queries_per_second=0.0, - recall=0.0, - algorithm=first_index.algo, - search_params=first_index.search_params, - metadata={ - "dry_run": True, - "num_indexes": len(indexes), - "total_search_configs": total_search_configs, - }, - success=True, - ) + return [ + SearchResult( + neighbors=np.array([]), + distances=np.array([]), + search_time_ms=0.0, + queries_per_second=0.0, + recall=0.0, + algorithm=first_index.algo, + search_params=first_index.search_params, + metadata={ + "dry_run": True, + "num_indexes": len(indexes), + "total_search_configs": total_search_configs, + }, + success=True, + ) + ] # Execute subprocess start_time = time.perf_counter() @@ -544,54 +551,60 @@ def search( # Note: C++ Google Benchmark doesn't return actual neighbors/distances # This is a limitation of the current system - return SearchResult( - neighbors=np.array([]), # Not available from C++ benchmark - distances=np.array([]), # Not available from C++ benchmark - search_time_ms=total_search_time, - queries_per_second=avg_qps, - recall=avg_recall, - algorithm=first_index.algo, - search_params=first_index.search_params, - metadata={ - "num_indexes": len(indexes), - "num_benchmarks": len(benchmarks), - "elapsed_time": elapsed_time, - "latency_us": benchmarks[0].get("Latency") - if benchmarks - else None, - "end_to_end": benchmarks[0].get("end_to_end") - if benchmarks - else None, - "context": gbench_results.get("context", {}), - }, - success=True, - ) + return [ + SearchResult( + neighbors=np.array([]), # Not available from C++ benchmark + distances=np.array([]), # Not available from C++ benchmark + search_time_ms=total_search_time, + queries_per_second=avg_qps, + recall=avg_recall, + algorithm=first_index.algo, + search_params=first_index.search_params, + metadata={ + "num_indexes": len(indexes), + "num_benchmarks": len(benchmarks), + "elapsed_time": elapsed_time, + "latency_us": benchmarks[0].get("Latency") + if benchmarks + else None, + "end_to_end": benchmarks[0].get("end_to_end") + if benchmarks + else None, + "context": gbench_results.get("context", {}), + }, + success=True, + ) + ] except subprocess.CalledProcessError as e: - return SearchResult( - neighbors=np.array([]), - distances=np.array([]), - search_time_ms=time.perf_counter() - start_time, - queries_per_second=0.0, - recall=0.0, - algorithm=first_index.algo, - search_params=first_index.search_params, - success=False, - error_message=f"Search failed: {e.stderr}", - ) + return [ + SearchResult( + neighbors=np.array([]), + distances=np.array([]), + search_time_ms=time.perf_counter() - start_time, + queries_per_second=0.0, + recall=0.0, + algorithm=first_index.algo, + search_params=first_index.search_params, + success=False, + error_message=f"Search failed: {e.stderr}", + ) + ] except Exception as e: - return SearchResult( - neighbors=np.array([]), - distances=np.array([]), - search_time_ms=time.perf_counter() - start_time, - queries_per_second=0.0, - recall=0.0, - algorithm=first_index.algo, - search_params=first_index.search_params, - success=False, - error_message=f"Search error: {str(e)}", - ) + return [ + SearchResult( + neighbors=np.array([]), + distances=np.array([]), + search_time_ms=time.perf_counter() - start_time, + queries_per_second=0.0, + recall=0.0, + algorithm=first_index.algo, + search_params=first_index.search_params, + success=False, + error_message=f"Search error: {str(e)}", + ) + ] finally: # Cleanup temporary config diff --git a/python/cuvs_bench/cuvs_bench/backends/opensearch.py b/python/cuvs_bench/cuvs_bench/backends/opensearch.py index 026fe98fd6..84740f21c7 100644 --- a/python/cuvs_bench/cuvs_bench/backends/opensearch.py +++ b/python/cuvs_bench/cuvs_bench/backends/opensearch.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -190,6 +190,7 @@ def _build_benchmark_configs( backend_cfg: Dict[str, Any] = { "name": index_label, + "group": group_name, "host": host, "port": port, "index_name": os_index_name, @@ -264,6 +265,7 @@ class OpenSearchBackend(BenchmarkBackend): Required: - ``name`` – index label (e.g. ``"opensearch_faiss_hnsw.m16.ef_construction100"``) + - ``group`` – algorithm configuration group selected from YAML - ``index_name`` – OpenSearch index name (lowercase, no dots) - ``engine`` – ``"faiss"`` or ``"lucene"`` - ``algo`` – algorithm name (e.g. ``"opensearch_faiss_hnsw"``) @@ -543,11 +545,12 @@ def _failed_build_result( self, error_message: str, build_params: Optional[Dict[str, Any]] = None ) -> BuildResult: return BuildResult( - index_path="", + index_path=self.config.get("index_name", ""), build_time_seconds=0.0, index_size_bytes=0, algorithm=self.algo, build_params=build_params or {}, + metadata={"group": self.config["group"]}, success=False, error_message=error_message, ) @@ -566,6 +569,10 @@ def _failed_search_result( recall=0.0, algorithm=self.algo, search_params=search_params or [], + metadata={ + "group": self.config["group"], + "index_name": self.config.get("index_name", ""), + }, success=False, error_message=error_message, ) @@ -715,6 +722,7 @@ def build( index_size_bytes=0, algorithm=self.algo, build_params=build_param, + metadata={"group": self.config["group"]}, success=True, ) @@ -729,6 +737,10 @@ def build( index_size_bytes=0, algorithm=self.algo, build_params=build_param, + metadata={ + "group": self.config["group"], + "skipped": True, + }, success=True, ) @@ -790,6 +802,7 @@ def build( algorithm=self.algo, build_params=build_param, metadata={ + "group": self.config["group"], "engine": engine, "space_type": space_type, "remote_index_build": remote_index_build, @@ -807,21 +820,14 @@ def search( force: bool = False, search_threads: Optional[int] = None, dry_run: bool = False, - ) -> SearchResult: + ) -> List[SearchResult]: """ Search the OpenSearch k-NN index for nearest neighbors. Iterates over every search-parameter combination defined in the index config, updating the index-level ``ef_search`` setting between runs. - Metrics (QPS, latency) are collected per parameter set and stored in - ``SearchResult.metadata["per_search_param_results"]``. - - The *neighbors* and *distances* arrays in the returned result reflect - the **last** search-parameter combination (highest ef_search by - convention), while *queries_per_second* is the average across all - parameter combinations. This backend returns ``recall=0.0``; the - shared orchestrator path computes recall from the returned neighbors - and dataset ground truth. + Returns one result per parameter set so the orchestrator can compute + recall for each set independently. Parameters ---------- @@ -847,16 +853,18 @@ def search( Returns ------- - SearchResult + List[SearchResult] """ skip = self._pre_flight_check() if skip: - return self._failed_search_result( - k, f"pre-flight check failed: {skip}" - ) + return [ + self._failed_search_result( + k, f"pre-flight check failed: {skip}" + ) + ] if not indexes: - return self._failed_search_result(k, "No indexes provided") + return [self._failed_search_result(k, "No indexes provided")] index_cfg = indexes[0] index_name = self._resolve_index_name(index_cfg) @@ -870,35 +878,43 @@ def search( f"(k={k}, batch_size={batch_size})" ) - return SearchResult( - neighbors=np.zeros((0, k), dtype=np.int64), - distances=np.zeros((0, k), dtype=np.float32), - search_time_ms=0.0, - queries_per_second=0.0, - recall=0.0, - algorithm=self.algo, - search_params=search_params_list, - success=True, - ) + return [ + SearchResult( + neighbors=np.zeros((0, k), dtype=np.int64), + distances=np.zeros((0, k), dtype=np.float32), + search_time_ms=0.0, + queries_per_second=0.0, + recall=0.0, + algorithm=self.algo, + search_params=[search_params], + metadata={ + "group": self.config["group"], + "index_name": index_name, + "dry_run": True, + }, + success=True, + ) + for search_params in search_params_list + ] # Dataset handles lazy loading from query files when needed. query_vectors = dataset.query_vectors if query_vectors.size == 0: - return self._failed_search_result( - k, - "No query vectors available. Provide dataset.query_vectors " - "or a valid dataset.query_file path.", - search_params=search_params_list, - ) + return [ + self._failed_search_result( + k, + "No query vectors available. Provide " + "dataset.query_vectors or a valid dataset.query_file path.", + search_params=search_params_list, + ) + ] n_queries = query_vectors.shape[0] n_batches = (n_queries + batch_size - 1) // batch_size # Run search for each search-parameter combination - 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) + results: List[SearchResult] = [] for sp in search_params_list: ef_search = sp.get("ef_search", 100) @@ -957,39 +973,25 @@ def search( elapsed = time.perf_counter() - t0 qps = n_queries / elapsed if elapsed > 0 else 0.0 - per_param_results.append( - { - "search_params": sp, - "search_time_ms": elapsed * 1000.0, - "queries_per_second": qps, - "batch_size": batch_size, - "num_batches": n_batches, - } + results.append( + SearchResult( + neighbors=neighbors, + distances=distances, + search_time_ms=elapsed * 1000.0, + queries_per_second=qps, + recall=0.0, + algorithm=self.algo, + search_params=[sp], + metadata={ + "group": self.config["group"], + "index_name": index_name, + "engine": engine, + "batch_size": batch_size, + "num_batches": n_batches, + "latency_seconds": elapsed / n_batches, + }, + success=True, + ) ) - last_neighbors = neighbors - last_distances = distances - # Aggregate across all search-param combinations - avg_qps = float( - np.mean([r["queries_per_second"] for r in per_param_results]) - ) - total_search_time_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_search_time_ms, - queries_per_second=avg_qps, - recall=0.0, - algorithm=self.algo, - search_params=search_params_list, - metadata={ - "engine": engine, - "batch_size": batch_size, - "num_batches": n_batches, - "per_search_param_results": per_param_results, - }, - success=True, - ) + return results diff --git a/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py b/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py index 580291c2f0..0a50f1d8c9 100644 --- a/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py +++ b/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -143,9 +143,12 @@ def run_benchmark( Returns ------- List[Union[BuildResult, SearchResult]] - List of result objects, one per benchmark run: - - sweep mode: One result per IndexConfig (Cartesian product of params) - - tune mode: One result per Optuna trial (n_trials total) + All measurements produced by the benchmark: + - sweep mode: build results plus the search results returned for + each benchmark configuration. + - tune mode: build and search results produced by each trial. A + successful build-and-search trial normally contributes two + objects. Each SearchResult contains: recall, search_time_ms, queries_per_second, success, metadata, etc. @@ -246,7 +249,7 @@ def _run_sweep( # 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( + search_results = backend.search( dataset=bench_dataset, indexes=config.indexes, k=count, @@ -257,28 +260,17 @@ def _run_sweep( dry_run=dry_run, ) - # Compute recall for backends that return actual neighbors. - # The C++ backend computes recall in the subprocess and returns - # empty neighbors, so this is skipped for it. - # Empty neighbors or nonzero recall indicate that the backend - # already handled recall itself. - if ( - search_result.success - and search_result.neighbors.size > 0 - and search_result.recall == 0.0 - ): - gt = bench_dataset.groundtruth_neighbors - if gt is not None: - search_result.recall = compute_recall( - search_result.neighbors, gt, count - ) - - results.append(search_result) - - if not search_result.success: - print( - f"Search failed for {config.index_name}: {search_result.error_message}" + for search_result in search_results: + self._finalize_search_result( + search_result, bench_dataset, count ) + results.append(search_result) + + if not search_result.success: + print( + f"Search failed for {config.index_name}: " + f"{search_result.error_message}" + ) finally: backend.cleanup() @@ -430,7 +422,7 @@ def objective(trial) -> float: # Run single trial with these specific parameters # First trial (trial.number=0) overwrites, subsequent trials append - result = self._run_trial( + trial_results = self._run_trial( algorithm=algorithm, build_params=build_params, search_params=search_params_dict, @@ -446,13 +438,20 @@ def objective(trial) -> float: **loader_kwargs, ) - # Store result for pareto plot - all_results.append(result) + # Retain every measurement for export. The last result is the + # search result used as the Optuna objective on successful trials. + all_results.extend(trial_results) + result = trial_results[-1] # Check if trial failed if not result.success: raise optuna.TrialPruned() + if not isinstance(result, SearchResult): + raise RuntimeError( + "Successful tune trial did not produce a search result" + ) + # Build metrics dict from SearchResult attributes # No fallbacks - if metrics are missing, let it fail loudly so we can fix the root cause metrics = { @@ -525,7 +524,7 @@ def _run_trial( search_threads: Optional[int], append_results: bool = False, **loader_kwargs, - ) -> Union[BuildResult, SearchResult]: + ) -> List[Union[BuildResult, SearchResult]]: """ Run a single benchmark trial with specific parameters. @@ -558,12 +557,19 @@ def _run_trial( # Should have exactly one config for single trial if not benchmark_configs: - return SearchResult( - success=False, - error_message="No config generated for trial", - metrics={}, - search_params=[], - ) + return [ + SearchResult( + neighbors=np.empty((0, count), dtype=np.int64), + distances=np.empty((0, count), dtype=np.float32), + search_time_ms=0.0, + queries_per_second=0.0, + recall=0.0, + algorithm=algorithm, + search_params=[], + success=False, + error_message="No config generated for trial", + ) + ] config = benchmark_configs[0] # Pass append_results via config (backend-specific, not in base class) @@ -575,20 +581,21 @@ def _run_trial( try: backend.initialize() - result = None + trial_results: List[Union[BuildResult, SearchResult]] = [] if build: - result = backend.build( + build_result = backend.build( dataset=bench_dataset, indexes=config.indexes, force=force, dry_run=dry_run, ) - if not result.success: - return result + trial_results.append(build_result) + if not build_result.success: + return trial_results if search: - result = backend.search( + search_results = backend.search( dataset=bench_dataset, indexes=config.indexes, k=count, @@ -599,24 +606,36 @@ def _run_trial( dry_run=dry_run, ) - # Compute recall for backends that return actual neighbors. - # Empty neighbors or nonzero recall indicate that the backend - # already handled recall itself. - if ( - result.success - and result.neighbors.size > 0 - and result.recall == 0.0 - ): - gt = bench_dataset.groundtruth_neighbors - if gt is not None: - result.recall = compute_recall( - result.neighbors, gt, count - ) + if len(search_results) != 1: + raise RuntimeError( + "Tune mode expected one search-parameter result" + ) + search_result = search_results[0] + self._finalize_search_result( + search_result, bench_dataset, count + ) + trial_results.append(search_result) - return result + return trial_results finally: backend.cleanup() + @staticmethod + def _finalize_search_result( + result: SearchResult, dataset: Dataset, k: int + ) -> None: + """Compute recall for backends that return neighbor arrays.""" + if ( + result.success + and result.neighbors.size > 0 + and result.recall == 0.0 + ): + groundtruth = dataset.groundtruth_neighbors + if groundtruth is not None: + result.recall = compute_recall( + result.neighbors, groundtruth, k + ) + def _create_dataset(self, dataset_config: DatasetConfig) -> Dataset: """ Create a Dataset object from DatasetConfig. diff --git a/python/cuvs_bench/cuvs_bench/run/__main__.py b/python/cuvs_bench/cuvs_bench/run/__main__.py index 6950ff7202..225a3d3f9c 100644 --- a/python/cuvs_bench/cuvs_bench/run/__main__.py +++ b/python/cuvs_bench/cuvs_bench/run/__main__.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -11,7 +11,11 @@ import click import yaml -from .data_export import convert_json_to_csv_build, convert_json_to_csv_search +from .data_export import ( + convert_json_to_csv_build, + convert_json_to_csv_search, + write_results_to_csv, +) from ..orchestrator import BenchmarkOrchestrator @@ -144,12 +148,8 @@ @click.option( "--data-export", is_flag=True, - help="By default, the intermediate JSON outputs produced by " - "cuvs_bench.run to more easily readable CSV files is done " - "automatically, which are needed to build charts made by " - "cuvs_bench.plot. But if some of the benchmark runs failed or " - "were interrupted, use this option to convert those intermediate " - "files manually.", + help="Deprecated: convert existing C++ benchmark JSON files to CSV. " + "Benchmark runs now export CSV automatically.", ) @click.option( "--mode", @@ -243,7 +243,7 @@ def main( dry_run : bool Whether to perform a dry run without actual execution. data_export : bool - Whether to export intermediate JSON results to CSV. + Deprecated option for converting existing C++ JSON results to CSV. mode : str Benchmark mode: 'sweep' (exhaustive) or 'tune' (Optuna-based). constraints : Optional[str] @@ -257,52 +257,70 @@ def main( and any backend-specific connection parameters (host, port, etc.). """ - if not data_export: - # Determine backend type and extra kwargs from --backend-config - backend_type = "cpp_gbench" - backend_kwargs = {} - if backend_config: - with open(backend_config, "r") as f: - cfg = yaml.safe_load(f) - if not isinstance(cfg, dict): - raise ValueError( - f"--backend-config must parse to a mapping, " - f"got {type(cfg).__name__}" - ) - if "backend" not in cfg: - raise ValueError( - "--backend-config must include a 'backend' field" - ) - backend_type = cfg.pop("backend") - backend_kwargs = cfg - - orchestrator = BenchmarkOrchestrator(backend_type=backend_type) - orchestrator.run_benchmark( - mode=mode, - constraints=json.loads(constraints) if constraints else None, - n_trials=n_trials, - dataset=dataset, - dataset_path=dataset_path, - build=build, - search=search, - force=force, - dry_run=dry_run, - count=count, - batch_size=batch_size, - search_mode=search_mode, - search_threads=search_threads, - dataset_configuration=dataset_configuration, - algorithm_configuration=configuration, - algorithms=algorithms, - groups=groups, - algo_groups=algo_groups, - subset_size=subset_size, - executable_dir=executable_dir, - **backend_kwargs, + if data_export: + click.echo( + "Warning: --data-export is deprecated because benchmark runs now " + "export CSV automatically. Converting existing C++ JSON results.", + err=True, ) + convert_json_to_csv_build(dataset, dataset_path) + convert_json_to_csv_search(dataset, dataset_path) + return + + # The CLI historically runs both phases when neither flag is specified. + if not build and not search: + build = search = True + + backend_type = "cpp_gbench" + backend_kwargs = {} + if backend_config: + with open(backend_config, "r") as f: + cfg = yaml.safe_load(f) + if not isinstance(cfg, dict): + raise ValueError( + f"--backend-config must parse to a mapping, " + f"got {type(cfg).__name__}" + ) + if "backend" not in cfg: + raise ValueError("--backend-config must include a 'backend' field") + backend_type = cfg.pop("backend") + backend_kwargs = cfg + + orchestrator = BenchmarkOrchestrator(backend_type=backend_type) + results = orchestrator.run_benchmark( + mode=mode, + constraints=json.loads(constraints) if constraints else None, + n_trials=n_trials, + dataset=dataset, + dataset_path=dataset_path, + build=build, + search=search, + force=force, + dry_run=dry_run, + count=count, + batch_size=batch_size, + search_mode=search_mode, + search_threads=search_threads, + dataset_configuration=dataset_configuration, + algorithm_configuration=configuration, + algorithms=algorithms, + groups=groups, + algo_groups=algo_groups, + subset_size=subset_size, + executable_dir=executable_dir, + **backend_kwargs, + ) + + if dry_run: + return - convert_json_to_csv_build(dataset, dataset_path) - convert_json_to_csv_search(dataset, dataset_path) + if backend_type == "cpp_gbench": + if build: + convert_json_to_csv_build(dataset, dataset_path) + if search: + convert_json_to_csv_search(dataset, dataset_path) + else: + write_results_to_csv(results, dataset, dataset_path, count, batch_size) if __name__ == "__main__": diff --git a/python/cuvs_bench/cuvs_bench/run/data_export.py b/python/cuvs_bench/cuvs_bench/run/data_export.py index 707677a083..50afa7a5ea 100644 --- a/python/cuvs_bench/cuvs_bench/run/data_export.py +++ b/python/cuvs_bench/cuvs_bench/run/data_export.py @@ -1,14 +1,17 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # import json import os import traceback +from collections import defaultdict import pandas as pd +from ..backends.base import BuildResult, SearchResult + skip_build_cols = set( [ "algo_name", @@ -50,6 +53,165 @@ } +def write_results_to_csv(results, dataset, dataset_path, count, batch_size): + """Write Python-backend results using the existing plotting CSV schema.""" + grouped = defaultdict(list) + for result in results: + group = result.metadata.get("group") + index_name = result.metadata.get("index_name") + if group is None or ( + isinstance(result, SearchResult) and index_name is None + ): + continue + method = "build" if isinstance(result, BuildResult) else "search" + grouped[(method, result.algorithm, group)].append(result) + + for (method, algorithm, group), group_results in grouped.items(): + if method == "build": + _write_build_results( + group_results, algorithm, group, dataset, dataset_path + ) + else: + _write_search_results( + group_results, + algorithm, + group, + dataset, + dataset_path, + count, + batch_size, + ) + + +def _write_build_results(results, algorithm, group, dataset, dataset_path): + output_dir = os.path.join(dataset_path, dataset, "result", "build") + os.makedirs(output_dir, exist_ok=True) + algo_name = algorithm if group == "base" else f"{algorithm}_{group}" + + rows = [] + for result in results: + if not result.success or result.metadata.get("skipped"): + continue + metadata = _scalar_metadata(result.metadata) + rows.append( + { + **result.build_params, + **metadata, + "algo_name": algo_name, + "index_name": result.index_path, + "time": result.build_time_seconds, + } + ) + + columns = ["algo_name", "index_name", "time"] + dataframe = pd.DataFrame(rows) + build_file = os.path.join(output_dir, f"{algorithm},{group}.csv") + + complete_run = all( + result.success and not result.metadata.get("skipped") + for result in results + ) + if not complete_run and os.path.exists(build_file): + dataframe = pd.concat( + [pd.read_csv(build_file), dataframe], + ignore_index=True, + sort=False, + ) + + if dataframe.empty: + # Do not replace an existing measurement with a skipped or failed + # build, and do not create an empty result file. + return + + dataframe = dataframe.drop_duplicates(subset=["index_name"], keep="last") + dataframe = dataframe[ + columns + [name for name in dataframe if name not in columns] + ] + dataframe.to_csv(build_file, index=False) + + +def _write_search_results( + results, algorithm, group, dataset, dataset_path, count, batch_size +): + output_dir = os.path.join(dataset_path, dataset, "result", "search") + os.makedirs(output_dir, exist_ok=True) + algo_name = algorithm if group == "base" else f"{algorithm}_{group}" + + rows = [] + for result in results: + if not result.success: + continue + metadata = _scalar_metadata(result.metadata) + search_params = ( + result.search_params[0] if len(result.search_params) == 1 else {} + ) + rows.append( + { + **search_params, + **metadata, + "algo_name": algo_name, + "index_name": result.metadata["index_name"], + "recall": result.recall, + "throughput": result.queries_per_second, + "latency": result.metadata.get( + "latency_seconds", result.search_time_ms / 1000.0 + ), + } + ) + + columns = [ + "algo_name", + "index_name", + "recall", + "throughput", + "latency", + ] + dataframe = pd.DataFrame(rows) + if dataframe.empty: + dataframe = pd.DataFrame(columns=columns) + else: + dataframe = dataframe[ + columns + [name for name in dataframe if name not in columns] + ] + + build_file = os.path.join( + dataset_path, + dataset, + "result", + "build", + f"{algorithm},{group}.csv", + ) + if os.path.exists(build_file): + build = pd.read_csv(build_file).drop_duplicates( + subset=["index_name"], keep="last" + ) + if "time" in build: + dataframe = dataframe.merge( + build[["index_name", "time"]].rename( + columns={"time": "build time"} + ), + on="index_name", + how="left", + ) + + stem = f"{algorithm},{group},k{count},bs{batch_size}" + raw_file = os.path.join(output_dir, f"{stem},raw.csv") + dataframe.to_csv(raw_file, index=False) + frontier_file = os.path.join(output_dir, f"{stem}.json") + write_frontier(frontier_file, dataframe, "throughput") + write_frontier(frontier_file, dataframe, "latency") + + +def _scalar_metadata(metadata): + reserved = {"group", "index_name", "latency_seconds"} + return { + key: value + for key, value in metadata.items() + if key not in reserved + and isinstance(value, (str, int, float, bool, type(None))) + } + + def read_json_files(dataset, dataset_path, method): """ Yield file paths, algo names, and loaded JSON data as pandas DataFrames. @@ -70,6 +232,8 @@ def read_json_files(dataset, dataset_path, method): DataFrame of JSON content. """ dir_path = os.path.join(dataset_path, dataset, "result", method) + if not os.path.isdir(dir_path): + return for file in os.listdir(dir_path): if file.endswith(".json"): file_path = os.path.join(dir_path, file) diff --git a/python/cuvs_bench/cuvs_bench/tests/test_cpp_gbench.py b/python/cuvs_bench/cuvs_bench/tests/test_cpp_gbench.py index a25b4fbe04..02711c46c0 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_cpp_gbench.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_cpp_gbench.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ @@ -243,7 +243,7 @@ def test_search_with_no_indexes(self): # Search with empty indexes list result = backend.search( dataset=dataset, indexes=[], k=10, batch_size=1000 - ) + )[0] assert result.success is True assert result.metadata.get("skipped") is True @@ -351,7 +351,7 @@ def test_search_dry_run(self): k=10, batch_size=1000, dry_run=True, - ) + )[0] assert result.success is True assert result.metadata.get("dry_run") is True diff --git a/python/cuvs_bench/cuvs_bench/tests/test_data_export.py b/python/cuvs_bench/cuvs_bench/tests/test_data_export.py new file mode 100644 index 0000000000..2595d32095 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_data_export.py @@ -0,0 +1,304 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +import numpy as np +import pandas as pd +from click.testing import CliRunner + +from cuvs_bench.backends.base import BuildResult, SearchResult +from cuvs_bench.orchestrator.config_loaders import ( + BenchmarkConfig, + DatasetConfig, + IndexConfig, +) +from cuvs_bench.orchestrator.orchestrator import BenchmarkOrchestrator +from cuvs_bench.plot.__main__ import load_all_results +from cuvs_bench.run.__main__ import main +from cuvs_bench.run.data_export import write_results_to_csv + + +def test_python_backend_csv_is_plot_compatible(tmp_path): + dataset = "test-dataset" + algorithm = "opensearch_faiss_hnsw" + index_name = "test-index" + results = [ + BuildResult( + index_path=index_name, + build_time_seconds=1.5, + index_size_bytes=1024, + algorithm=algorithm, + build_params={"m": 16}, + metadata={"group": "base"}, + ), + SearchResult( + neighbors=np.empty((0, 2), dtype=np.int64), + distances=np.empty((0, 2), dtype=np.float32), + search_time_ms=20.0, + queries_per_second=100.0, + recall=0.5, + algorithm=algorithm, + search_params=[{"ef_search": 50}], + metadata={ + "group": "base", + "index_name": index_name, + "latency_seconds": 0.01, + }, + ), + SearchResult( + neighbors=np.empty((0, 2), dtype=np.int64), + distances=np.empty((0, 2), dtype=np.float32), + search_time_ms=10.0, + queries_per_second=200.0, + recall=1.0, + algorithm=algorithm, + search_params=[{"ef_search": 100}], + metadata={ + "group": "base", + "index_name": index_name, + "latency_seconds": 0.005, + }, + ), + ] + + write_results_to_csv( + results, dataset, str(tmp_path), count=2, batch_size=2 + ) + + result_path = tmp_path / dataset / "result" + raw_file = result_path / "search" / f"{algorithm},base,k2,bs2,raw.csv" + raw = pd.read_csv(raw_file) + assert raw.columns[:5].tolist() == [ + "algo_name", + "index_name", + "recall", + "throughput", + "latency", + ] + assert raw["ef_search"].tolist() == [50, 100] + assert raw["build time"].tolist() == [1.5, 1.5] + assert not list(result_path.rglob("*.json")) + + write_results_to_csv( + [ + BuildResult( + index_path=index_name, + build_time_seconds=0.0, + index_size_bytes=0, + algorithm=algorithm, + build_params={"m": 16}, + metadata={"group": "base", "skipped": True}, + ) + ], + dataset, + str(tmp_path), + count=2, + batch_size=2, + ) + build = pd.read_csv(result_path / "build" / f"{algorithm},base.csv") + assert build["time"].tolist() == [1.5] + + for mode in ("throughput", "latency"): + plotted = load_all_results( + str(result_path.parent), + algorithms=[algorithm], + groups=["base"], + algo_groups=[], + k=2, + batch_size=2, + method="search", + index_key="algo", + raw=False, + mode=mode, + time_unit="s", + ) + assert algorithm in plotted + assert plotted[algorithm] + + +def test_run_command_always_writes_python_backend_csv(tmp_path, monkeypatch): + algorithm = "opensearch_faiss_hnsw" + result = SearchResult( + neighbors=np.empty((0, 2), dtype=np.int64), + distances=np.empty((0, 2), dtype=np.float32), + search_time_ms=10.0, + queries_per_second=200.0, + recall=1.0, + algorithm=algorithm, + search_params=[{"ef_search": 100}], + metadata={ + "group": "base", + "index_name": "test-index", + "latency_seconds": 0.01, + }, + ) + + class FakeOrchestrator: + def __init__(self, backend_type): + assert backend_type == "opensearch" + + def run_benchmark(self, **kwargs): + return [result] + + monkeypatch.setattr( + "cuvs_bench.run.__main__.BenchmarkOrchestrator", FakeOrchestrator + ) + backend_config = tmp_path / "backend.yaml" + backend_config.write_text("backend: opensearch\n") + + cli_result = CliRunner().invoke( + main, + [ + "--dataset", + "test-dataset", + "--dataset-path", + str(tmp_path), + "--algorithms", + algorithm, + "--groups", + "base", + "--count", + "2", + "--batch-size", + "2", + "--search-mode", + "latency", + "--search", + "--backend-config", + str(backend_config), + ], + ) + + assert cli_result.exit_code == 0, cli_result.output + assert ( + tmp_path + / "test-dataset" + / "result" + / "search" + / f"{algorithm},base,k2,bs2,raw.csv" + ).exists() + help_output = CliRunner().invoke(main, ["--help"]).output + assert "--data-export" in help_output + assert "Deprecated" in help_output + + +def test_data_export_option_is_deprecated(tmp_path, monkeypatch): + converted = [] + monkeypatch.setattr( + "cuvs_bench.run.__main__.convert_json_to_csv_build", + lambda dataset, dataset_path: converted.append("build"), + ) + monkeypatch.setattr( + "cuvs_bench.run.__main__.convert_json_to_csv_search", + lambda dataset, dataset_path: converted.append("search"), + ) + + cli_result = CliRunner().invoke( + main, + [ + "--dataset", + "test-dataset", + "--dataset-path", + str(tmp_path), + "--algorithms", + "cuvs_cagra", + "--groups", + "base", + "--count", + "2", + "--batch-size", + "2", + "--search-mode", + "latency", + "--data-export", + ], + ) + + assert cli_result.exit_code == 0, cli_result.output + assert "--data-export is deprecated" in cli_result.output + assert converted == ["build", "search"] + + +def test_tune_trial_retains_build_and_search_results(): + algorithm = "opensearch_faiss_hnsw" + index = IndexConfig( + name="test-index", + algo=algorithm, + build_param={"m": 16}, + search_params=[{"ef_search": 100}], + file="test-index", + ) + dataset = DatasetConfig(name="test-dataset") + + class FakeLoader: + def load(self, **kwargs): + return dataset, [ + BenchmarkConfig( + indexes=[index], + backend_config={"name": index.name, "group": "base"}, + ) + ] + + class FakeBackend: + def __init__(self, config): + pass + + def initialize(self): + pass + + def cleanup(self): + pass + + def build(self, dataset, indexes, force, dry_run): + return BuildResult( + index_path=index.name, + build_time_seconds=1.5, + index_size_bytes=1024, + algorithm=algorithm, + build_params=index.build_param, + metadata={"group": "base"}, + ) + + def search(self, dataset, indexes, k, **kwargs): + return [ + SearchResult( + neighbors=np.array([[0, 1]], dtype=np.int64), + distances=np.zeros((1, k), dtype=np.float32), + search_time_ms=10.0, + queries_per_second=100.0, + recall=0.0, + algorithm=algorithm, + search_params=index.search_params, + metadata={ + "group": "base", + "index_name": index.name, + }, + ) + ] + + orchestrator = object.__new__(BenchmarkOrchestrator) + orchestrator.config_loader = FakeLoader() + orchestrator.backend_class = FakeBackend + orchestrator._create_dataset = lambda config: type( + "Dataset", + (), + {"groundtruth_neighbors": np.array([[0, 1]], dtype=np.int64)}, + )() + + results = orchestrator._run_trial( + algorithm=algorithm, + build_params=index.build_param, + search_params=index.search_params[0], + build=True, + search=True, + force=False, + dry_run=False, + count=2, + batch_size=1, + search_mode="latency", + search_threads=None, + ) + + assert [type(result) for result in results] == [BuildResult, SearchResult] + assert results[1].recall == 1.0 diff --git a/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py b/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py index d16eedf25d..67c02e1dd3 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # """ @@ -21,12 +21,14 @@ OpenSearchConfigLoader, ) from cuvs_bench.orchestrator.config_loaders import IndexConfig +from cuvs_bench.orchestrator.orchestrator import BenchmarkOrchestrator def _make_backend(config_overrides: dict = None) -> OpenSearchBackend: """Backend with no network requirement so pre-flight passes without a server.""" config = { "name": "test_index", + "group": "base", "index_name": "test_index", "engine": "faiss", "algo": "opensearch_faiss_hnsw", @@ -90,6 +92,7 @@ def live_backend(opensearch_url): backend = OpenSearchBackend( { "name": index_name, + "group": "base", "index_name": index_name, "engine": "faiss", "algo": "opensearch_faiss_hnsw", @@ -146,6 +149,7 @@ def test_load_produces_correct_configs(self, config_dir): assert len(benchmark_configs) == 4 bc = benchmark_configs[0] assert bc.backend_config["engine"] == "faiss" + assert bc.backend_config["group"] == "test" assert len(bc.indexes[0].search_params) == 2 # ef_search: [50, 100] def test_load_forwards_remote_build_kwargs(self, config_dir): @@ -176,11 +180,56 @@ def test_build_dry_run(self): assert result.index_path == backend.config["index_name"] def test_search_dry_run(self): - result = _make_backend().search( + results = _make_backend().search( _make_dataset(), [_make_index_cfg()], k=3, dry_run=True ) - assert result.success - assert len(result.search_params) == 2 + assert len(results) == 2 + assert all(result.success for result in results) + assert [result.search_params for result in results] == [ + [{"ef_search": 50}], + [{"ef_search": 100}], + ] + + def test_recall_is_computed_for_each_search_parameter(self): + class FakeIndices: + def __init__(self): + self.ef_search = None + + def put_settings(self, index, body): + self.ef_search = body["index.knn.algo_param.ef_search"] + + class FakeClient: + def __init__(self): + self.indices = FakeIndices() + + def msearch(self, index, body): + ids = [2, 3] if self.indices.ef_search == 50 else [0, 1] + response = { + "hits": { + "hits": [ + {"_id": str(neighbor), "_score": 1.0} + for neighbor in ids + ] + } + } + return {"responses": [response for _ in body[::2]]} + + dataset = Dataset( + name="test", + query_vectors=np.zeros((2, 4), dtype=np.float32), + groundtruth_neighbors=np.array([[0, 1], [0, 1]]), + ) + backend = _make_backend() + backend._OpenSearchBackend__client = FakeClient() + + results = backend.search( + dataset, [_make_index_cfg()], k=2, batch_size=2 + ) + for result in results: + BenchmarkOrchestrator._finalize_search_result(result, dataset, 2) + + assert [result.recall for result in results] == [0.0, 1.0] + assert all(result.neighbors.shape == (2, 2) for result in results) def test_remote_build_requires_faiss_engine(self): backend = _make_backend({"engine": "lucene"}) @@ -326,7 +375,7 @@ def test_search_fails_without_query_vectors(self): training_vectors=np.empty((0, 4), dtype=np.float32), query_vectors=np.empty((0, 4), dtype=np.float32), ) - result = _make_backend().search(dataset, [_make_index_cfg()], k=3) + result = _make_backend().search(dataset, [_make_index_cfg()], k=3)[0] assert not result.success assert "No query vectors" in result.error_message @@ -472,6 +521,7 @@ def live_remote_build_backend(opensearch_url, remote_build_env): backend = OpenSearchBackend( { "name": index_name, + "group": "base", "index_name": index_name, "engine": "faiss", "algo": "opensearch_faiss_hnsw", @@ -503,12 +553,11 @@ def test_build_and_search(self, live_backend): assert build_result.build_time_seconds > 0 assert build_result.index_size_bytes > 0 - search_result = live_backend.search(dataset, [idx], k=k) + search_result = live_backend.search(dataset, [idx], k=k)[0] assert search_result.success assert search_result.recall == 0.0 assert search_result.queries_per_second > 0 assert search_result.neighbors.shape == (10, k) - assert len(search_result.metadata["per_search_param_results"]) == 1 @pytest.mark.opensearch @@ -525,7 +574,9 @@ def test_remote_build_and_search(self, live_remote_build_backend): assert build_result.build_time_seconds > 0 assert build_result.metadata["remote_index_build"] is True - search_result = live_remote_build_backend.search(dataset, [idx], k=k) + search_result = live_remote_build_backend.search(dataset, [idx], k=k)[ + 0 + ] assert search_result.success assert search_result.recall == 0.0 assert search_result.queries_per_second > 0 diff --git a/python/cuvs_bench/cuvs_bench/tests/test_registry.py b/python/cuvs_bench/cuvs_bench/tests/test_registry.py index 960ea9cb25..eadf804db3 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_registry.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_registry.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ @@ -59,16 +59,18 @@ def search( distances = np.random.rand(n_queries, k) first = indexes[0] - return SearchResult( - neighbors=neighbors, - distances=distances, - search_time_ms=0.1, - queries_per_second=n_queries / 0.1, - recall=0.95, - algorithm=self.algo, - search_params=first.search_params, - success=True, - ) + return [ + SearchResult( + neighbors=neighbors, + distances=distances, + search_time_ms=0.1, + queries_per_second=n_queries / 0.1, + recall=0.95, + algorithm=self.algo, + search_params=first.search_params, + success=True, + ) + ] class AnotherDummyBackend(BenchmarkBackend): @@ -109,16 +111,18 @@ def search( distances = np.random.rand(n_queries, k) first = indexes[0] - return SearchResult( - neighbors=neighbors, - distances=distances, - search_time_ms=0.2, - queries_per_second=n_queries / 0.2, - recall=0.90, - algorithm=self.algo, - search_params=first.search_params if first else [], - success=True, - ) + return [ + SearchResult( + neighbors=neighbors, + distances=distances, + search_time_ms=0.2, + queries_per_second=n_queries / 0.2, + recall=0.90, + algorithm=self.algo, + search_params=first.search_params if first else [], + success=True, + ) + ] class TestDataset: @@ -386,7 +390,7 @@ def test_dummy_backend_search(self, tmp_path): ) ] - result = backend.search(dataset=dataset, indexes=indexes, k=10) + result = backend.search(dataset=dataset, indexes=indexes, k=10)[0] assert result.success assert result.recall == 0.95 From 34656e4d2471c29eae8a00200857c85833da81f6 Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Tue, 28 Jul 2026 13:29:47 -0500 Subject: [PATCH 2/3] Update Signed-off-by: James Bourbeau --- .../cuvs_bench/tests/test_data_export.py | 107 +----------------- 1 file changed, 1 insertion(+), 106 deletions(-) diff --git a/python/cuvs_bench/cuvs_bench/tests/test_data_export.py b/python/cuvs_bench/cuvs_bench/tests/test_data_export.py index 2595d32095..721d03c5ac 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_data_export.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_data_export.py @@ -5,7 +5,6 @@ import numpy as np import pandas as pd -from click.testing import CliRunner from cuvs_bench.backends.base import BuildResult, SearchResult from cuvs_bench.orchestrator.config_loaders import ( @@ -15,7 +14,6 @@ ) from cuvs_bench.orchestrator.orchestrator import BenchmarkOrchestrator from cuvs_bench.plot.__main__ import load_all_results -from cuvs_bench.run.__main__ import main from cuvs_bench.run.data_export import write_results_to_csv @@ -117,109 +115,6 @@ def test_python_backend_csv_is_plot_compatible(tmp_path): assert plotted[algorithm] -def test_run_command_always_writes_python_backend_csv(tmp_path, monkeypatch): - algorithm = "opensearch_faiss_hnsw" - result = SearchResult( - neighbors=np.empty((0, 2), dtype=np.int64), - distances=np.empty((0, 2), dtype=np.float32), - search_time_ms=10.0, - queries_per_second=200.0, - recall=1.0, - algorithm=algorithm, - search_params=[{"ef_search": 100}], - metadata={ - "group": "base", - "index_name": "test-index", - "latency_seconds": 0.01, - }, - ) - - class FakeOrchestrator: - def __init__(self, backend_type): - assert backend_type == "opensearch" - - def run_benchmark(self, **kwargs): - return [result] - - monkeypatch.setattr( - "cuvs_bench.run.__main__.BenchmarkOrchestrator", FakeOrchestrator - ) - backend_config = tmp_path / "backend.yaml" - backend_config.write_text("backend: opensearch\n") - - cli_result = CliRunner().invoke( - main, - [ - "--dataset", - "test-dataset", - "--dataset-path", - str(tmp_path), - "--algorithms", - algorithm, - "--groups", - "base", - "--count", - "2", - "--batch-size", - "2", - "--search-mode", - "latency", - "--search", - "--backend-config", - str(backend_config), - ], - ) - - assert cli_result.exit_code == 0, cli_result.output - assert ( - tmp_path - / "test-dataset" - / "result" - / "search" - / f"{algorithm},base,k2,bs2,raw.csv" - ).exists() - help_output = CliRunner().invoke(main, ["--help"]).output - assert "--data-export" in help_output - assert "Deprecated" in help_output - - -def test_data_export_option_is_deprecated(tmp_path, monkeypatch): - converted = [] - monkeypatch.setattr( - "cuvs_bench.run.__main__.convert_json_to_csv_build", - lambda dataset, dataset_path: converted.append("build"), - ) - monkeypatch.setattr( - "cuvs_bench.run.__main__.convert_json_to_csv_search", - lambda dataset, dataset_path: converted.append("search"), - ) - - cli_result = CliRunner().invoke( - main, - [ - "--dataset", - "test-dataset", - "--dataset-path", - str(tmp_path), - "--algorithms", - "cuvs_cagra", - "--groups", - "base", - "--count", - "2", - "--batch-size", - "2", - "--search-mode", - "latency", - "--data-export", - ], - ) - - assert cli_result.exit_code == 0, cli_result.output - assert "--data-export is deprecated" in cli_result.output - assert converted == ["build", "search"] - - def test_tune_trial_retains_build_and_search_results(): algorithm = "opensearch_faiss_hnsw" index = IndexConfig( @@ -277,7 +172,7 @@ def search(self, dataset, indexes, k, **kwargs): ) ] - orchestrator = object.__new__(BenchmarkOrchestrator) + orchestrator = BenchmarkOrchestrator(backend_type="opensearch") orchestrator.config_loader = FakeLoader() orchestrator.backend_class = FakeBackend orchestrator._create_dataset = lambda config: type( From e1baea2db3f59663787554a49d8bf6f6ed5ce617 Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Fri, 31 Jul 2026 10:17:08 -0500 Subject: [PATCH 3/3] Add ef_search fix Signed-off-by: James Bourbeau --- .../cuvs_bench/backends/opensearch.py | 15 ++-- .../cuvs_bench/tests/test_opensearch.py | 72 ++++++++----------- 2 files changed, 37 insertions(+), 50 deletions(-) diff --git a/python/cuvs_bench/cuvs_bench/backends/opensearch.py b/python/cuvs_bench/cuvs_bench/backends/opensearch.py index 84740f21c7..142942eee9 100644 --- a/python/cuvs_bench/cuvs_bench/backends/opensearch.py +++ b/python/cuvs_bench/cuvs_bench/backends/opensearch.py @@ -825,9 +825,9 @@ def search( Search the OpenSearch k-NN index for nearest neighbors. Iterates over every search-parameter combination defined in the index - config, updating the index-level ``ef_search`` setting between runs. - Returns one result per parameter set so the orchestrator can compute - recall for each set independently. + config, passing ``ef_search`` directly in every k-NN query. Returns one + result per parameter set so the orchestrator can compute recall for + each set independently. Parameters ---------- @@ -919,12 +919,6 @@ def search( for sp in search_params_list: ef_search = sp.get("ef_search", 100) - if engine == "faiss": - self._client.indices.put_settings( - index=index_name, - body={"index.knn.algo_param.ef_search": ef_search}, - ) - neighbors = np.full((n_queries, k), -1, dtype=np.int64) distances = np.zeros((n_queries, k), dtype=np.float32) @@ -942,6 +936,9 @@ def search( "vector": { "vector": q_vec.tolist(), "k": k, + "method_parameters": { + "ef_search": ef_search, + }, } } }, diff --git a/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py b/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py index 67c02e1dd3..124d6eacee 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_opensearch.py @@ -190,47 +190,6 @@ def test_search_dry_run(self): [{"ef_search": 100}], ] - def test_recall_is_computed_for_each_search_parameter(self): - class FakeIndices: - def __init__(self): - self.ef_search = None - - def put_settings(self, index, body): - self.ef_search = body["index.knn.algo_param.ef_search"] - - class FakeClient: - def __init__(self): - self.indices = FakeIndices() - - def msearch(self, index, body): - ids = [2, 3] if self.indices.ef_search == 50 else [0, 1] - response = { - "hits": { - "hits": [ - {"_id": str(neighbor), "_score": 1.0} - for neighbor in ids - ] - } - } - return {"responses": [response for _ in body[::2]]} - - dataset = Dataset( - name="test", - query_vectors=np.zeros((2, 4), dtype=np.float32), - groundtruth_neighbors=np.array([[0, 1], [0, 1]]), - ) - backend = _make_backend() - backend._OpenSearchBackend__client = FakeClient() - - results = backend.search( - dataset, [_make_index_cfg()], k=2, batch_size=2 - ) - for result in results: - BenchmarkOrchestrator._finalize_search_result(result, dataset, 2) - - assert [result.recall for result in results] == [0.0, 1.0] - assert all(result.neighbors.shape == (2, 2) for result in results) - def test_remote_build_requires_faiss_engine(self): backend = _make_backend({"engine": "lucene"}) with pytest.raises(ValueError, match="faiss engine"): @@ -559,6 +518,37 @@ def test_build_and_search(self, live_backend): assert search_result.queries_per_second > 0 assert search_result.neighbors.shape == (10, k) + def test_recall_is_computed_for_each_search_parameter(self, live_backend): + # Regression test for https://github.com/NVIDIA/cuvs/issues/2358 + k = 10 + dataset = _make_dataset( + n_base=5_000, + n_queries=100, + dims=16, + k=k, + ) + idx = IndexConfig( + name="test_index", + algo="opensearch_faiss_hnsw", + build_param={"m": 4, "ef_construction": 64}, + search_params=[{"ef_search": 10}, {"ef_search": 100}], + file="", + ) + + build_result = live_backend.build(dataset, [idx], force=True) + assert build_result.success + + results = live_backend.search(dataset, [idx], k=k) + for result in results: + BenchmarkOrchestrator._finalize_search_result(result, dataset, k) + + assert [result.search_params for result in results] == [ + [{"ef_search": 10}], + [{"ef_search": 100}], + ] + assert all(result.neighbors.shape == (100, k) for result in results) + assert results[0].recall < results[1].recall + @pytest.mark.opensearch class TestOpenSearchRemoteIndexBuildIntegration: