[FEA] Add missing pytests for cuvs-bench algorithms

cpp/bench/ann/src/diskann/diskann_benchmark.cpp

@@ -26,7 +26,7 @@ void parse_build_param(const nlohmann::json& conf,
 {
   param.R       = conf.at("R");
   param.L_build = conf.at("L_build");
-  if (conf.contains("alpha")) { param.num_threads = conf.at("alpha"); }
+  if (conf.contains("alpha")) { param.alpha = conf.at("alpha"); }
   if (conf.contains("num_threads")) { param.num_threads = conf.at("num_threads"); }
 }
 

cpp/src/neighbors/detail/vamana/vamana_serialize.cuh

@@ -65,13 +65,20 @@ void serialize_dataset(raft::resources const& res,
     const auto* strided_dataset =
       dynamic_cast<const cuvs::neighbors::strided_dataset<T, int64_t>*>(dataset);
     if (strided_dataset) {
-      auto h_dataset =
-        raft::make_host_matrix<T, int64_t>(strided_dataset->n_rows(), strided_dataset->dim());
-      raft::copy(res,
-                 raft::make_host_vector_view(h_dataset.data_handle(),
-                                             strided_dataset->n_rows() * strided_dataset->dim()),
-                 raft::make_device_vector_view(strided_dataset->view().data_handle(),
-                                               strided_dataset->n_rows() * strided_dataset->dim()));
+      auto nrows     = strided_dataset->n_rows();
+      auto dim       = strided_dataset->dim();
+      auto stride    = strided_dataset->stride();
+      auto d_data    = strided_dataset->view();
+      auto h_dataset = raft::make_host_matrix<T, int64_t>(nrows, dim);
+      RAFT_CUDA_TRY(cudaMemcpy2DAsync(h_dataset.data_handle(),
+                                      sizeof(T) * dim,
+                                      d_data.data_handle(),
+                                      sizeof(T) * stride,
+                                      sizeof(T) * dim,
+                                      nrows,
+                                      cudaMemcpyDefault,
+                                      raft::resource::get_cuda_stream(res)));
+      raft::resource::sync_stream(res);
       to_file(dataset_base_file, h_dataset);
     } else {
       RAFT_LOG_DEBUG("dynamic_cast to strided_dataset failed");
@@ -91,6 +98,7 @@ void serialize_dataset(raft::resources const& res,
   try {
     auto h_dataset = raft::make_host_matrix<T, int64_t>(dataset.extent(0), dataset.extent(1));
     raft::copy(res, h_dataset.view(), dataset);
+    raft::resource::sync_stream(res);
     to_file(dataset_base_file, h_dataset);
   } catch (std::bad_alloc& e) {
     RAFT_LOG_INFO("Failed to serialize dataset");

python/cuvs_bench/cuvs_bench/config/algos/cuvs_brute_force.yaml

@@ -3,3 +3,6 @@ groups:
   base:
     build:
     search:
+  test:
+    build:
+    search:

python/cuvs_bench/cuvs_bench/config/algos/cuvs_cagra_diskann.yaml

@@ -9,3 +9,10 @@ groups:
       graph_build_algo: ["NN_DESCENT"]
     search:
       L_search: [10, 20, 30, 40, 50, 100, 200, 300]
+  test:
+    build:
+      graph_degree: [32]
+      intermediate_graph_degree: [32]
+      graph_build_algo: ["NN_DESCENT"]
+    search:
+      L_search: [10]

python/cuvs_bench/cuvs_bench/config/algos/cuvs_vamana.yaml

@@ -8,3 +8,10 @@ groups:
     search:
       L_search: [10, 20, 30, 40, 50, 100, 200, 300]
       num_threads: [32]
+  test:
+    build:
+      graph_degree: [32]
+      visited_size: [128]
+      alpha: [1.2]
+    search:
+      L_search: [10]

python/cuvs_bench/cuvs_bench/config/algos/diskann_memory.yaml

@@ -10,3 +10,11 @@ groups:
       num_threads: [32]
     search:
       L_search: [10, 20, 30, 40, 50, 100, 200, 300]
+  test:
+    build:
+      R: [32]
+      L_build: [128]
+      alpha: [1.2]
+      num_threads: [1]
+    search:
+      L_search: [10]

python/cuvs_bench/cuvs_bench/run/data_export.py

@@ -1,5 +1,5 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 #
 
@@ -213,9 +213,10 @@ def convert_json_to_csv_search(dataset, dataset_path):
                     write["build GPU"] = None
                 for col_idx in range(start_idx, len(build_df.columns)):
                     col_name = build_df.columns[col_idx]
-                    write[col_name] = None
                     if col_name == "num_threads":
                         write["build_num_threads"] = None
+                    else:
+                        write[col_name] = None
                 for s_index, search_row in write.iterrows():
                     for b_index, build_row in build_df.iterrows():
                         if search_row["index_name"] == build_row["index_name"]:

python/cuvs_bench/cuvs_bench/tests/test_cli.py

@@ -1,8 +1,9 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 #
 
+import platform
 from pathlib import Path
 
 import pandas as pd
@@ -11,6 +12,11 @@
 from cuvs_bench.get_dataset.__main__ import main
 
 
+def is_arm_cpu():
+    """Check if running on ARM Linux architecture."""
+    return platform.machine() in ("aarch64", "arm64")
+
+
 @pytest.fixture(scope="session")
 def temp_datasets_dir(tmp_path_factory):
     return tmp_path_factory.mktemp("datasets")
@@ -51,7 +57,8 @@ def test_run_command_creates_results(temp_datasets_dir: Path):
 
         python -m cuvs_bench.run --dataset test-data --dataset-path datasets/ \
             --algorithms faiss_gpu_ivf_flat,faiss_gpu_ivf_sq,cuvs_ivf_flat,\
-            cuvs_cagra,ggnn,cuvs_cagra_hnswlib,cuvs_ivf_pq \
+            cuvs_cagra,ggnn,cuvs_cagra_hnswlib,cuvs_ivf_pq,cuvs_vamana,\
+            diskann_memory \
             --batch-size 100 -k 10 --groups test -m latency --force
 
     It then verifies that the set of expected result files
@@ -63,14 +70,19 @@ def test_run_command_creates_results(temp_datasets_dir: Path):
 
     from cuvs_bench.run.__main__ import main as run_main
 
+    skip_diskann = is_arm_cpu()
+    algorithms = "faiss_gpu_ivf_flat,faiss_gpu_ivf_sq,cuvs_ivf_flat,cuvs_cagra,ggnn,cuvs_cagra_hnswlib,cuvs_ivf_pq"
+    if not skip_diskann:
+        algorithms += ",cuvs_vamana,diskann_memory"
+
     runner = CliRunner()
     run_args = [
         "--dataset",
         "test-data",
         "--dataset-path",
         dataset_path_arg,
         "--algorithms",
-        "faiss_gpu_ivf_flat,faiss_gpu_ivf_sq,cuvs_ivf_flat,cuvs_cagra,ggnn,cuvs_cagra_hnswlib,cuvs_ivf_pq",  # noqa: E501
+        algorithms,
         "--batch-size",
         "100",
         "-k",
@@ -105,6 +117,8 @@ def test_run_command_creates_results(temp_datasets_dir: Path):
     ]
 
     # --- Verify that the expected result files exist and are not empty ---
+    skip_diskann = is_arm_cpu()
+
     expected_files = {
         # Build files:
         "test-data/result/build/cuvs_ivf_flat,test.csv": {
@@ -423,6 +437,143 @@ def test_run_command_creates_results(temp_datasets_dir: Path):
         },
     }
 
+    if not skip_diskann:
+        expected_files.update(
+            {
+                "test-data/result/build/cuvs_vamana,test.csv": {
+                    "header": common_build_header
+                    + [
+                        "GPU",
+                        "alpha",
+                        "graph_degree",
+                        "visited_size",
+                    ],
+                    "rows": 1,
+                },
+                "test-data/result/search/cuvs_vamana,test,k10,bs100,raw.csv": {
+                    "header": common_search_header
+                    + [
+                        "GPU",
+                        "L_search",
+                        "end_to_end",
+                        "k",
+                        "n_queries",
+                        "total_queries",
+                        "build time",
+                        "build threads",
+                        "build cpu_time",
+                        "build GPU",
+                        "alpha",
+                        "graph_degree",
+                        "visited_size",
+                    ],
+                    "rows": 1,
+                },
+                "test-data/result/search/cuvs_vamana,test,k10,bs100,latency.csv": {
+                    "header": common_search_header
+                    + [
+                        "GPU",
+                        "L_search",
+                        "end_to_end",
+                        "k",
+                        "n_queries",
+                        "total_queries",
+                        "build time",
+                        "build threads",
+                        "build cpu_time",
+                        "build GPU",
+                        "alpha",
+                        "graph_degree",
+                        "visited_size",
+                    ],
+                    "rows": 1,
+                },
+                "test-data/result/search/cuvs_vamana,test,k10,bs100,throughput.csv": {
+                    "header": common_search_header
+                    + [
+                        "GPU",
+                        "L_search",
+                        "end_to_end",
+                        "k",
+                        "n_queries",
+                        "total_queries",
+                        "build time",
+                        "build threads",
+                        "build cpu_time",
+                        "build GPU",
+                        "alpha",
+                        "graph_degree",
+                        "visited_size",
+                    ],
+                    "rows": 1,
+                },
+                "test-data/result/build/diskann_memory,test.csv": {
+                    "header": common_build_header
+                    + [
+                        "L_build",
+                        "R",
+                        "alpha",
+                        "num_threads",
+                    ],
+                    "rows": 1,
+                },
+                "test-data/result/search/diskann_memory,test,k10,bs100,raw.csv": {
+                    "header": common_search_header
+                    + [
+                        "L_search",
+                        "end_to_end",
+                        "k",
+                        "n_queries",
+                        "total_queries",
+                        "build time",
+                        "build threads",
+                        "build cpu_time",
+                        "L_build",
+                        "R",
+                        "alpha",
+                        "build_num_threads",
+                    ],
+                    "rows": 1,
+                },
+                "test-data/result/search/diskann_memory,test,k10,bs100,latency.csv": {
+                    "header": common_search_header
+                    + [
+                        "L_search",
+                        "end_to_end",
+                        "k",
+                        "n_queries",
+                        "total_queries",
+                        "build time",
+                        "build threads",
+                        "build cpu_time",
+                        "L_build",
+                        "R",
+                        "alpha",
+                        "build_num_threads",
+                    ],
+                    "rows": 1,
+                },
+                "test-data/result/search/diskann_memory,test,k10,bs100,throughput.csv": {
+                    "header": common_search_header
+                    + [
+                        "L_search",
+                        "end_to_end",
+                        "k",
+                        "n_queries",
+                        "total_queries",
+                        "build time",
+                        "build threads",
+                        "build cpu_time",
+                        "L_build",
+                        "R",
+                        "alpha",
+                        "build_num_threads",
+                    ],
+                    "rows": 1,
+                },
+            }
+        )
+
     for rel_path, expectations in expected_files.items():
         file_path = temp_datasets_dir / rel_path
         assert file_path.exists(), f"Expected file {file_path} does not exist."
@@ -448,7 +599,8 @@ def test_plot_command_creates_png_files(temp_datasets_dir: Path):
 
       python -m cuvs_bench.plot --dataset test-data --dataset-path datasets/ \
           --algorithms faiss_gpu_ivf_flat,faiss_gpu_ivf_sq, \
-          cuvs_ivf_flat,cuvs_cagra,ggnn,cuvs_cagra_hnswlib,cuvs_ivf_pq \
+          cuvs_ivf_flat,cuvs_cagra,ggnn,cuvs_cagra_hnswlib,cuvs_ivf_pq,\
+          cuvs_vamana,diskann_memory \
           --batch-size 100 -k 10 --groups test -m latency
 
     and then verifies that the following files are produced in the
@@ -463,6 +615,11 @@ def test_plot_command_creates_png_files(temp_datasets_dir: Path):
 
     from cuvs_bench.plot.__main__ import main as plot_main
 
+    skip_diskann = is_arm_cpu()
+    algorithms = "faiss_gpu_ivf_flat,faiss_gpu_ivf_sq,cuvs_ivf_flat,cuvs_cagra,ggnn,cuvs_cagra_hnswlib,cuvs_ivf_pq"
+    if not skip_diskann:
+        algorithms += ",cuvs_vamana,diskann_memory"
+
     runner = CliRunner()
     args = [
         "--dataset",
@@ -472,7 +629,7 @@ def test_plot_command_creates_png_files(temp_datasets_dir: Path):
         "--output-filepath",
         dataset_path_arg,
         "--algorithms",
-        "faiss_gpu_ivf_flat,faiss_gpu_ivf_sq,cuvs_ivf_flat,cuvs_cagra,ggnn,cuvs_cagra_hnswlib,cuvs_ivf_pq",  # noqa: E501
+        algorithms,
         "--batch-size",
         "100",
         "-k",