diff --git a/python/cuvs_bench/cuvs_bench/backends/base.py b/python/cuvs_bench/cuvs_bench/backends/base.py index 15209decb6..bfdef6b367 100644 --- a/python/cuvs_bench/cuvs_bench/backends/base.py +++ b/python/cuvs_bench/cuvs_bench/backends/base.py @@ -11,7 +11,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass, field -from typing import Dict, List, Any, Optional, TYPE_CHECKING +from typing import Dict, List, Any, Optional, Sequence, TYPE_CHECKING import numpy as np if TYPE_CHECKING: @@ -368,6 +368,10 @@ class BenchmarkBackend(ABC): - requires_network : bool - Whether backend requires network (default: False) """ + # Registration exposes this optional class-level default to the generic + # CLI. + default_algorithm: Optional[str] = None + def __init__(self, config: Dict[str, Any]): """Initialize backend with configuration.""" self.config = config @@ -441,7 +445,7 @@ def search( batch_size : int, optional Number of queries to process at once (default: 10000) mode : str, optional - "latency" (measure individual query latency with percentiles) or + "latency" (measure batch latency with percentiles) or "throughput" (measure overall QPS) (default: "latency") force : bool, optional Whether to force the execution regardless of existing results (default: False) @@ -495,6 +499,13 @@ def cleanup(self) -> None: """ pass + @classmethod + def result_failure_message( + cls, results: Sequence[BuildResult | SearchResult] + ) -> Optional[str]: + """Return a fatal result summary when the backend requires one.""" + return None + def _check_gpu_available(self) -> bool: """ Check if GPU is available (same logic as run.py rmm_present()). diff --git a/python/cuvs_bench/cuvs_bench/backends/cpp_gbench.py b/python/cuvs_bench/cuvs_bench/backends/cpp_gbench.py index 1e2d861456..78004c374b 100644 --- a/python/cuvs_bench/cuvs_bench/backends/cpp_gbench.py +++ b/python/cuvs_bench/cuvs_bench/backends/cpp_gbench.py @@ -51,6 +51,8 @@ class CppGoogleBenchmarkBackend(BenchmarkBackend): >>> print(backend.algo) # "cuvs_ivf_flat" (from config) """ + default_algorithm = "cuvs_cagra" + def __init__(self, config: Dict[str, Any]): """Initialize C++ benchmark backend.""" super().__init__(config) diff --git a/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py b/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py index abffad0217..ef7f9a4c68 100644 --- a/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py +++ b/python/cuvs_bench/cuvs_bench/orchestrator/orchestrator.py @@ -87,6 +87,12 @@ def __init__(self, backend_type: str = "cpp_gbench"): # Instantiate config loader self.config_loader = loader_class() + def result_failure_message( + self, results: List[Union[BuildResult, SearchResult]] + ) -> Optional[str]: + """Return the selected backend's fatal result summary, if any.""" + return self.backend_class.result_failure_message(results) + def run_benchmark( self, mode: str = "sweep", diff --git a/python/cuvs_bench/cuvs_bench/run/__main__.py b/python/cuvs_bench/cuvs_bench/run/__main__.py index 225a3d3f9c..efd0a72085 100644 --- a/python/cuvs_bench/cuvs_bench/run/__main__.py +++ b/python/cuvs_bench/cuvs_bench/run/__main__.py @@ -6,17 +6,64 @@ import json import os from pathlib import Path -from typing import Optional +from typing import Any, Optional import click import yaml +from ..backends.registry import get_backend_class +from ..orchestrator import BenchmarkOrchestrator from .data_export import ( convert_json_to_csv_build, convert_json_to_csv_search, write_results_to_csv, ) -from ..orchestrator import BenchmarkOrchestrator + + +_DEFAULT_BACKEND = "cpp_gbench" +_BACKEND_CONFIG_CACHE_KEY = "cuvs_bench.backend_config" + + +def _read_backend_config(path: str) -> dict[str, Any]: + with open(path, "r") as stream: + config = yaml.safe_load(stream) + if not isinstance(config, dict): + raise ValueError( + "--backend-config must parse to a mapping, " + f"got {type(config).__name__}" + ) + if "backend" not in config: + raise ValueError("--backend-config must include a 'backend' field") + return dict(config) + + +def _backend_config_snapshot(path: str) -> dict[str, Any]: + """Return one backend configuration snapshot per CLI invocation.""" + context = click.get_current_context() + cached_config = context.meta.get(_BACKEND_CONFIG_CACHE_KEY) + if cached_config is None: + cached_config = _read_backend_config(path) + context.meta[_BACKEND_CONFIG_CACHE_KEY] = cached_config + return dict(cached_config) + + +def _default_algorithm() -> str: + """Choose the displayed prompt default from the selected backend.""" + context = click.get_current_context(silent=True) + backend = _DEFAULT_BACKEND + if context is not None and ( + backend_config := context.params.get("backend_config") + ): + backend = _backend_config_snapshot(backend_config)["backend"] + backend_class = get_backend_class(str(backend)) + if backend_class.default_algorithm is not None: + return backend_class.default_algorithm + default_backend_class = get_backend_class(_DEFAULT_BACKEND) + if default_backend_class.default_algorithm is None: + raise RuntimeError( + "The default backend must define a default algorithm" + ) + return default_backend_class.default_algorithm @click.command() @@ -88,12 +135,13 @@ @click.option("--search", is_flag=True, help="Perform the search") @click.option( "--algorithms", - default="cuvs_cagra", + default=_default_algorithm, show_default=True, prompt="Enter the comma separated list of named algorithms to run", help="Run only comma separated list of named algorithms. If parameters " "`groups` and `algo-groups` are both undefined, then group `base` " - "is run by default.", + "is run by default. A backend may provide its own prompt default; " + "otherwise the CLI default is retained.", ) @click.option( "--groups", @@ -251,10 +299,10 @@ def main( n_trials : Optional[int] Number of Optuna trials for tune mode. backend_config : Optional[str] - Path to YAML config for non-C++ backends. If not provided, - defaults to the C++ Google Benchmark backend. The YAML file - must contain a 'backend' field (e.g., 'opensearch', 'elastic') - and any backend-specific connection parameters (host, port, etc.). + Path to YAML config for non-C++ backends. If not provided, defaults + to the C++ Google Benchmark backend. The YAML file must contain a + 'backend' field (e.g., 'opensearch', 'elastic') and any + backend-specific connection parameters (host, port, etc.). """ if data_export: @@ -271,18 +319,10 @@ def main( if not build and not search: build = search = True - backend_type = "cpp_gbench" + backend_type = _DEFAULT_BACKEND 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") + cfg = _backend_config_snapshot(backend_config) backend_type = cfg.pop("backend") backend_kwargs = cfg @@ -321,6 +361,8 @@ def main( convert_json_to_csv_search(dataset, dataset_path) else: write_results_to_csv(results, dataset, dataset_path, count, batch_size) + if failure_message := orchestrator.result_failure_message(results): + raise click.ClickException(failure_message) if __name__ == "__main__": diff --git a/python/cuvs_bench/cuvs_bench/run/data_export.py b/python/cuvs_bench/cuvs_bench/run/data_export.py index 50afa7a5ea..69e146d291 100644 --- a/python/cuvs_bench/cuvs_bench/run/data_export.py +++ b/python/cuvs_bench/cuvs_bench/run/data_export.py @@ -52,11 +52,15 @@ }, } +_SEARCH_PARAMS_IDENTITY_COLUMN = "_search_params_identity" + 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: + if result.metadata.get("dry_run"): + continue group = result.metadata.get("group") index_name = result.metadata.get("index_name") if group is None or ( @@ -84,13 +88,15 @@ def write_results_to_csv(results, 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"): + if ( + not result.success + or result.metadata.get("skipped") + or result.metadata.get("dry_run") + ): continue metadata = _scalar_metadata(result.metadata) rows.append( @@ -98,11 +104,19 @@ def _write_build_results(results, algorithm, group, dataset, dataset_path): **result.build_params, **metadata, "algo_name": algo_name, - "index_name": result.index_path, + "index_name": result.metadata.get( + "index_name", result.index_path + ), "time": result.build_time_seconds, } ) + if not rows: + # A failed, skipped, or dry run must not replace prior measurements. + return + + output_dir = os.path.join(dataset_path, dataset, "result", "build") + os.makedirs(output_dir, exist_ok=True) columns = ["algo_name", "index_name", "time"] dataframe = pd.DataFrame(rows) build_file = os.path.join(output_dir, f"{algorithm},{group}.csv") @@ -110,6 +124,7 @@ def _write_build_results(results, algorithm, group, dataset, dataset_path): complete_run = all( result.success and not result.metadata.get("skipped") for result in results + if not result.metadata.get("dry_run") ) if not complete_run and os.path.exists(build_file): dataframe = pd.concat( @@ -118,11 +133,6 @@ def _write_build_results(results, algorithm, group, dataset, dataset_path): 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] @@ -133,13 +143,15 @@ def _write_build_results(results, algorithm, group, dataset, dataset_path): 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: + if ( + not result.success + or result.metadata.get("skipped") + or result.metadata.get("dry_run") + ): continue metadata = _scalar_metadata(result.metadata) search_params = ( @@ -156,9 +168,18 @@ def _write_search_results( "latency": result.metadata.get( "latency_seconds", result.search_time_ms / 1000.0 ), + _SEARCH_PARAMS_IDENTITY_COLUMN: _search_params_identity( + result.search_params + ), } ) + if not rows: + # Preserve prior evidence when every attempted search failed or was dry. + return + + output_dir = os.path.join(dataset_path, dataset, "result", "search") + os.makedirs(output_dir, exist_ok=True) columns = [ "algo_name", "index_name", @@ -167,12 +188,9 @@ def _write_search_results( "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] - ] + dataframe = dataframe[ + columns + [name for name in dataframe if name not in columns] + ] build_file = os.path.join( dataset_path, @@ -196,10 +214,53 @@ def _write_search_results( stem = f"{algorithm},{group},k{count},bs{batch_size}" raw_file = os.path.join(output_dir, f"{stem},raw.csv") + complete_run = all( + result.success and not result.metadata.get("skipped") + for result in results + ) + if not complete_run and os.path.exists(raw_file): + existing = _with_legacy_search_param_identities( + pd.read_csv(raw_file) + ) + identity_columns = ["index_name", _SEARCH_PARAMS_IDENTITY_COLUMN] + dataframe = dataframe.drop_duplicates( + subset=identity_columns, keep="last" + ) + existing_identities = pd.MultiIndex.from_frame( + existing[identity_columns] + ) + replacement_identities = pd.MultiIndex.from_frame( + dataframe[identity_columns] + ) + existing = existing.loc[ + ~existing_identities.isin(replacement_identities) + ] + dataframe = pd.concat( + [existing, dataframe], ignore_index=True, sort=False + ) 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") + frontier_data = dataframe.drop(columns=[_SEARCH_PARAMS_IDENTITY_COLUMN]) + write_frontier(frontier_file, frontier_data, "throughput") + write_frontier(frontier_file, frontier_data, "latency") + + +def _search_params_identity(search_params): + """Return a stable key for the parameter plans represented by one row.""" + return json.dumps(search_params, sort_keys=True, separators=(",", ":")) + + +def _with_legacy_search_param_identities(dataframe): + """Give legacy rows opaque keys so a partial run cannot erase them.""" + if _SEARCH_PARAMS_IDENTITY_COLUMN not in dataframe: + dataframe[_SEARCH_PARAMS_IDENTITY_COLUMN] = pd.NA + + missing_identity = dataframe[_SEARCH_PARAMS_IDENTITY_COLUMN].isna() + dataframe.loc[missing_identity, _SEARCH_PARAMS_IDENTITY_COLUMN] = [ + f"legacy:{row_number}" + for row_number in dataframe.index[missing_identity] + ] + return dataframe def _scalar_metadata(metadata): diff --git a/python/cuvs_bench/cuvs_bench/tests/test_backend_cli.py b/python/cuvs_bench/cuvs_bench/tests/test_backend_cli.py new file mode 100644 index 0000000000..68f329f345 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_backend_cli.py @@ -0,0 +1,266 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CLI contracts shared by pluggable cuVS Bench backends.""" + +from pathlib import Path +from unittest.mock import patch + +from click.testing import CliRunner + +from cuvs_bench.backends.base import BenchmarkBackend, BuildResult +from cuvs_bench.orchestrator import BenchmarkOrchestrator +from cuvs_bench.run.__main__ import main as run_main + + +class _DefaultBackend: + default_algorithm = "cuvs_cagra" + + @classmethod + def result_failure_message(cls, results): + return None + + +class _PluginBackend: + default_algorithm = "plugin_default" + + @classmethod + def result_failure_message(cls, results): + failures = [result for result in results if not result.success] + if not failures: + return None + return "; ".join(result.error_message for result in failures) + + +def _write_backend_config( + tmp_path: Path, backend: str, *settings: str +) -> Path: + backend_config = tmp_path / "backend.yaml" + lines = [f"backend: {backend}", *settings] + backend_config.write_text("\n".join(lines) + "\n", encoding="utf-8") + return backend_config + + +def _invoke_run( + tmp_path: Path, + *extra_args: str, + input_text: str = "", + run_results=None, + dry_run: bool = True, + plugin_backend=_PluginBackend, +): + captured = {} + backend_classes = { + "cpp_gbench": _DefaultBackend, + "opensearch": plugin_backend, + "third_party": plugin_backend, + } + + class RecordingOrchestrator: + def __init__(self, backend_type): + captured["backend_type"] = backend_type + self.backend_class = backend_classes[backend_type] + + def run_benchmark(self, **kwargs): + captured["run_kwargs"] = kwargs + return [] if run_results is None else run_results + + def result_failure_message(self, results): + return self.backend_class.result_failure_message(results) + + arguments = [ + "--dataset", + "test-data", + "--dataset-path", + str(tmp_path), + "--batch-size", + "10", + "-k", + "10", + "--groups", + "test", + "-m", + "latency", + *extra_args, + ] + if dry_run: + arguments.append("--dry-run") + + with ( + patch( + "cuvs_bench.run.__main__.get_backend_class", + side_effect=backend_classes.__getitem__, + ), + patch( + "cuvs_bench.run.__main__.BenchmarkOrchestrator", + RecordingOrchestrator, + ), + patch("cuvs_bench.run.__main__.write_results_to_csv") as write_results, + ): + result = CliRunner().invoke(run_main, arguments, input=input_text) + + return result, captured, write_results + + +def test_cli_keeps_the_cpp_backend_and_algorithm_defaults(tmp_path: Path): + result, captured, _write_results = _invoke_run( + tmp_path, input_text="\n" + ) + + assert result.exit_code == 0, result.output + assert captured["backend_type"] == "cpp_gbench" + assert captured["run_kwargs"]["algorithms"] == "cuvs_cagra" + + +def test_registered_backend_controls_its_prompt_default(tmp_path: Path): + backend_config = _write_backend_config(tmp_path, "third_party") + + result, captured, _write_results = _invoke_run( + tmp_path, + "--backend-config", + str(backend_config), + input_text="\n", + ) + + assert result.exit_code == 0, result.output + assert captured["backend_type"] == "third_party" + assert captured["run_kwargs"]["algorithms"] == "plugin_default" + + +def test_backend_default_and_dispatch_share_one_config_snapshot( + tmp_path: Path, +): + backend_config = _write_backend_config(tmp_path, "third_party") + changing_snapshots = ( + {"backend": "third_party"}, + {"backend": "opensearch"}, + {"backend": "cpp_gbench"}, + ) + + with patch( + "cuvs_bench.run.__main__._read_backend_config", + side_effect=changing_snapshots, + ) as read_backend_config: + result, captured, _write_results = _invoke_run( + tmp_path, + "--backend-config", + str(backend_config), + input_text="\n", + ) + + assert result.exit_code == 0, result.output + assert read_backend_config.call_count == 1 + assert captured["backend_type"] == "third_party" + assert captured["run_kwargs"]["algorithms"] == "plugin_default" + + +def test_backend_without_a_default_keeps_the_cli_default(tmp_path: Path): + class BackendWithoutDefault(_PluginBackend): + default_algorithm = None + + backend_config = _write_backend_config(tmp_path, "third_party") + + result, captured, _write_results = _invoke_run( + tmp_path, + "--backend-config", + str(backend_config), + input_text="\n", + plugin_backend=BackendWithoutDefault, + ) + + assert result.exit_code == 0, result.output + assert captured["run_kwargs"]["algorithms"] == "cuvs_cagra" + + +def test_cli_preserves_an_explicit_algorithm(tmp_path: Path): + backend_config = _write_backend_config(tmp_path, "third_party") + + result, captured, _write_results = _invoke_run( + tmp_path, + "--backend-config", + str(backend_config), + "--algorithms", + "user_selected_algorithm", + ) + + assert result.exit_code == 0, result.output + assert captured["run_kwargs"]["algorithms"] == "user_selected_algorithm" + + +def test_backend_config_keeps_selecting_the_backend(tmp_path: Path): + backend_config = _write_backend_config( + tmp_path, "opensearch", "host: search.example" + ) + + result, captured, _write_results = _invoke_run( + tmp_path, + "--backend-config", + str(backend_config), + "--algorithms", + "plugin_algorithm", + ) + + assert result.exit_code == 0, result.output + assert captured["backend_type"] == "opensearch" + assert captured["run_kwargs"]["host"] == "search.example" + + +def test_cli_invokes_export_before_reporting_backend_failure(tmp_path: Path): + failed = BuildResult( + index_path="", + build_time_seconds=0.0, + index_size_bytes=0, + algorithm="plugin_algorithm", + build_params={}, + success=False, + error_message="backend operation failed", + ) + + backend_config = _write_backend_config(tmp_path, "third_party") + + result, _captured, write_results = _invoke_run( + tmp_path, + "--backend-config", + str(backend_config), + "--algorithms", + "plugin_algorithm", + run_results=[failed], + dry_run=False, + ) + + assert result.exit_code == 1 + assert "backend operation failed" in result.output + write_results.assert_called_once() + + +def test_backend_failure_policy_is_opt_in(): + failed = BuildResult( + index_path="", + build_time_seconds=0.0, + index_size_bytes=0, + algorithm="plugin_algorithm", + build_params={}, + success=False, + error_message="backend operation failed", + ) + + assert BenchmarkBackend.result_failure_message([failed]) is None + + +def test_orchestrator_delegates_failure_policy_to_selected_backend(): + orchestrator = BenchmarkOrchestrator.__new__(BenchmarkOrchestrator) + orchestrator.backend_class = _PluginBackend + failed = BuildResult( + index_path="", + build_time_seconds=0.0, + index_size_bytes=0, + algorithm="plugin_algorithm", + build_params={}, + success=False, + error_message="backend operation failed", + ) + + assert orchestrator.result_failure_message([failed]) == ( + "backend operation failed" + ) diff --git a/python/cuvs_bench/cuvs_bench/tests/test_data_export.py b/python/cuvs_bench/cuvs_bench/tests/test_data_export.py index 721d03c5ac..e2af5f5d0b 100644 --- a/python/cuvs_bench/cuvs_bench/tests/test_data_export.py +++ b/python/cuvs_bench/cuvs_bench/tests/test_data_export.py @@ -5,6 +5,7 @@ import numpy as np import pandas as pd +import pytest from cuvs_bench.backends.base import BuildResult, SearchResult from cuvs_bench.orchestrator.config_loaders import ( @@ -14,7 +15,35 @@ ) from cuvs_bench.orchestrator.orchestrator import BenchmarkOrchestrator from cuvs_bench.plot.__main__ import load_all_results -from cuvs_bench.run.data_export import write_results_to_csv +from cuvs_bench.run.data_export import ( + _SEARCH_PARAMS_IDENTITY_COLUMN, + write_results_to_csv, +) + + +def _search_result( + algorithm, + index_name, + search_params, + recall, + *, + outcome="completed", +): + metadata = {"group": "base", "index_name": index_name} + if outcome == "skipped": + metadata["skipped"] = True + return 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=100.0, + recall=recall, + algorithm=algorithm, + search_params=[search_params], + metadata=metadata, + success=outcome != "failed", + error_message="search failed" if outcome == "failed" else None, + ) def test_python_backend_csv_is_plot_compatible(tmp_path): @@ -23,12 +52,12 @@ def test_python_backend_csv_is_plot_compatible(tmp_path): index_name = "test-index" results = [ BuildResult( - index_path=index_name, + index_path="temporary-build-path", build_time_seconds=1.5, index_size_bytes=1024, algorithm=algorithm, build_params={"m": 16}, - metadata={"group": "base"}, + metadata={"group": "base", "index_name": index_name}, ), SearchResult( neighbors=np.empty((0, 2), dtype=np.int64), @@ -115,6 +144,331 @@ def test_python_backend_csv_is_plot_compatible(tmp_path): assert plotted[algorithm] +def test_failed_and_dry_results_preserve_existing_measurements(tmp_path): + dataset = "test-dataset" + algorithm = "test-backend-algorithm" + index_name = "test-index" + build = BuildResult( + index_path=index_name, + build_time_seconds=1.5, + index_size_bytes=1024, + algorithm=algorithm, + build_params={}, + metadata={"group": "base"}, + ) + search = 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=100.0, + recall=1.0, + algorithm=algorithm, + search_params=[{}], + metadata={ + "group": "base", + "index_name": index_name, + "latency_seconds": 0.01, + }, + ) + write_results_to_csv( + [build, search], dataset, str(tmp_path), count=2, batch_size=2 + ) + + result_root = tmp_path / dataset / "result" + exported_files = tuple(result_root.rglob("*.csv")) + original_contents = {path: path.read_bytes() for path in exported_files} + dry_build = BuildResult( + index_path=index_name, + build_time_seconds=0.0, + index_size_bytes=0, + algorithm=algorithm, + build_params={}, + metadata={"group": "base", "dry_run": True}, + ) + failed_build = BuildResult( + index_path=index_name, + build_time_seconds=0.0, + index_size_bytes=0, + algorithm=algorithm, + build_params={}, + metadata={"group": "base"}, + success=False, + error_message="build failed", + ) + dry_search = SearchResult( + neighbors=np.empty((0, 2), dtype=np.int64), + distances=np.empty((0, 2), dtype=np.float32), + search_time_ms=0.0, + queries_per_second=0.0, + recall=0.0, + algorithm=algorithm, + search_params=[{}], + metadata={ + "group": "base", + "index_name": index_name, + "dry_run": True, + }, + ) + failed_search = SearchResult( + neighbors=np.empty((0, 2), dtype=np.int64), + distances=np.empty((0, 2), dtype=np.float32), + search_time_ms=0.0, + queries_per_second=0.0, + recall=0.0, + algorithm=algorithm, + search_params=[{}], + metadata={"group": "base", "index_name": index_name}, + success=False, + error_message="search failed", + ) + skipped_search = SearchResult( + neighbors=np.empty((0, 2), dtype=np.int64), + distances=np.empty((0, 2), dtype=np.float32), + search_time_ms=0.0, + queries_per_second=0.0, + recall=0.0, + algorithm=algorithm, + search_params=[{}], + metadata={ + "group": "base", + "index_name": index_name, + "skipped": True, + }, + ) + + write_results_to_csv( + [ + dry_build, + failed_build, + dry_search, + failed_search, + skipped_search, + ], + dataset, + str(tmp_path), + count=2, + batch_size=2, + ) + + assert {path: path.read_bytes() for path in exported_files} == ( + original_contents + ) + + +@pytest.mark.parametrize( + "incomplete_outcome", + ["failed", "skipped"], +) +def test_partial_search_run_replaces_only_completed_measurements( + tmp_path, incomplete_outcome +): + dataset = "test-dataset" + algorithm = "test-backend-algorithm" + index_name = "test-index" + retained_index_name = "retained-index" + + write_results_to_csv( + [ + _search_result( + algorithm, + index_name, + {"ef_search": 50, "search_width": 16}, + 0.5, + ), + _search_result( + algorithm, + index_name, + {"ef_search": 100, "search_width": 32}, + 1.0, + ), + _search_result( + algorithm, + retained_index_name, + {"ef_search": 50, "search_width": 16}, + 0.6, + ), + ], + dataset, + str(tmp_path), + count=2, + batch_size=2, + ) + write_results_to_csv( + [ + _search_result( + algorithm, + index_name, + {"search_width": 16, "ef_search": 50}, + 0.75, + ), + _search_result( + algorithm, + index_name, + {"search_width": 32, "ef_search": 100}, + 0.0, + outcome=incomplete_outcome, + ), + ], + dataset, + str(tmp_path), + count=2, + batch_size=2, + ) + + raw_file = ( + tmp_path + / dataset + / "result" + / "search" + / f"{algorithm},base,k2,bs2,raw.csv" + ) + measurements = pd.read_csv(raw_file) + assert len(measurements) == 3 + assert not measurements.duplicated( + subset=["index_name", "ef_search"] + ).any() + measurements_by_identity = measurements.set_index( + ["index_name", "ef_search"] + )["recall"].to_dict() + assert measurements_by_identity == { + (index_name, 50): 0.75, + (index_name, 100): 1.0, + (retained_index_name, 50): 0.6, + } + + +def test_failed_search_with_new_parameter_preserves_existing_results( + tmp_path, +): + dataset = "test-dataset" + algorithm = "test-backend-algorithm" + index_name = "test-index" + + write_results_to_csv( + [_search_result(algorithm, index_name, {}, 0.5)], + dataset, + str(tmp_path), + count=2, + batch_size=2, + ) + write_results_to_csv( + [ + _search_result(algorithm, index_name, {}, 0.75), + _search_result( + algorithm, + index_name, + {"ef_search": 100}, + 0.0, + outcome="failed", + ), + ], + dataset, + str(tmp_path), + count=2, + batch_size=2, + ) + + raw_file = ( + tmp_path + / dataset + / "result" + / "search" + / f"{algorithm},base,k2,bs2,raw.csv" + ) + measurements = pd.read_csv(raw_file) + assert measurements["recall"].tolist() == [0.75] + assert "ef_search" not in measurements + + +@pytest.mark.parametrize( + "legacy_csv", [False, True], ids=["current", "legacy"] +) +@pytest.mark.parametrize( + "completed_index_name", + ["historical-index", "new-index"], + ids=["same-index", "different-index"], +) +def test_partial_search_run_preserves_unmatched_parameter_plans( + tmp_path, legacy_csv, completed_index_name +): + dataset = "test-dataset" + algorithm = "test-backend-algorithm" + historical_index_name = "historical-index" + + write_results_to_csv( + [ + _search_result( + algorithm, + historical_index_name, + {"ef_search": 50}, + 0.5, + ), + _search_result( + algorithm, + historical_index_name, + {"ef_search": 100}, + 1.0, + ), + ], + dataset, + str(tmp_path), + count=2, + batch_size=2, + ) + raw_file = ( + tmp_path + / dataset + / "result" + / "search" + / f"{algorithm},base,k2,bs2,raw.csv" + ) + if legacy_csv: + legacy_measurements = pd.read_csv(raw_file).drop( + columns=[_SEARCH_PARAMS_IDENTITY_COLUMN] + ) + legacy_measurements.to_csv(raw_file, index=False) + + write_results_to_csv( + [ + _search_result( + algorithm, completed_index_name, {}, 0.75 + ), + _search_result( + algorithm, + "failed-index", + {}, + 0.0, + outcome="failed", + ), + ], + dataset, + str(tmp_path), + count=2, + batch_size=2, + ) + + measurements = pd.read_csv(raw_file) + historical_measurements = measurements.loc[ + measurements["ef_search"].notna() + ] + assert historical_measurements.set_index("ef_search")[ + "recall" + ].to_dict() == {50: 0.5, 100: 1.0} + completed_measurements = measurements.loc[ + (measurements["index_name"] == completed_index_name) + & measurements["ef_search"].isna() + ] + assert completed_measurements["recall"].tolist() == [0.75] + assert len(measurements) == 3 + + latency_frontier = raw_file.with_name( + raw_file.name.replace(",raw.csv", ",latency.csv") + ) + assert _SEARCH_PARAMS_IDENTITY_COLUMN not in pd.read_csv( + latency_frontier + ).columns + + def test_tune_trial_retains_build_and_search_results(): algorithm = "opensearch_faiss_hnsw" index = IndexConfig(