Add cuvs_bench.run python code and build (#279)

Builds on top of #275

Authors:
  - Dante Gama Dessavre (https://github.com/dantegd)
  - Corey J. Nolet (https://github.com/cjnolet)
  - Divye Gala (https://github.com/divyegala)
  - rhdong (https://github.com/rhdong)
  - Bradley Dice (https://github.com/bdice)

Approvers:
  - Divye Gala (https://github.com/divyegala)
  - Corey J. Nolet (https://github.com/cjnolet)
  - Bradley Dice (https://github.com/bdice)

URL: #279

.gitignore

@@ -79,3 +79,6 @@ cagra_index
 ivf_flat_index
 ivf_pq_index
 
+# cuvs_bench
+datasets/
+/*.json

build.sh

@@ -275,7 +275,7 @@ if hasArg tests || (( ${NUMARGS} == 0 )); then
 fi
 
 if hasArg bench-ann || (( ${NUMARGS} == 0 )); then
-    BUILD_ANN_BENCH=ON
+    BUILD_CUVS_BENCH=ON
     CMAKE_TARGET="${CMAKE_TARGET};${ANN_BENCH_TARGETS}"
 fi
 
@@ -351,7 +351,7 @@ if (( ${NUMARGS} == 0 )) || hasArg libcuvs || hasArg docs || hasArg tests || has
           -DDISABLE_DEPRECATION_WARNINGS=${DISABLE_DEPRECATION_WARNINGS} \
           -DBUILD_TESTS=${BUILD_TESTS} \
           -DBUILD_C_TESTS=${BUILD_TESTS} \
-          -DBUILD_ANN_BENCH=${BUILD_ANN_BENCH} \
+          -DBUILD_CUVS_BENCH=${BUILD_CUVS_BENCH} \
           -DBUILD_CPU_ONLY=${BUILD_CPU_ONLY} \
           -DCMAKE_MESSAGE_LOG_LEVEL=${CMAKE_LOG_LEVEL} \
           ${CACHE_ARGS} \
@@ -419,6 +419,11 @@ if (( ${NUMARGS} == 0 )) || hasArg python; then
         python -m pip install --no-build-isolation --no-deps --config-settings rapidsai.disable-cuda=true ${REPODIR}/python/cuvs
 fi
 
+# Build and (optionally) install the cuvs_bench Python package
+if (( ${NUMARGS} == 0 )) || hasArg bench-ann; then
+    python -m pip install --no-build-isolation --no-deps --config-settings rapidsai.disable-cuda=true ${REPODIR}/python/cuvs_bench
+fi
+
 # Build the cuvs Rust bindings
 if (( ${NUMARGS} == 0 )) || hasArg rust; then
     cd ${REPODIR}/rust

conda/environments/bench_ann_cuda-118_arch-aarch64.yaml

@@ -42,5 +42,8 @@ dependencies:
 - pandas
 - pylibraft==24.10.*,>=0.0.0a0
 - pyyaml
+- rapids-build-backend>=0.3.0,<0.4.0.dev0
+- setuptools
 - sysroot_linux-aarch64==2.17
+- wheel
 name: bench_ann_cuda-118_arch-aarch64

conda/environments/bench_ann_cuda-118_arch-x86_64.yaml

@@ -42,5 +42,8 @@ dependencies:
 - pandas
 - pylibraft==24.10.*,>=0.0.0a0
 - pyyaml
+- rapids-build-backend>=0.3.0,<0.4.0.dev0
+- setuptools
 - sysroot_linux-64==2.17
+- wheel
 name: bench_ann_cuda-118_arch-x86_64

conda/environments/bench_ann_cuda-125_arch-aarch64.yaml

@@ -38,5 +38,8 @@ dependencies:
 - pandas
 - pylibraft==24.10.*,>=0.0.0a0
 - pyyaml
+- rapids-build-backend>=0.3.0,<0.4.0.dev0
+- setuptools
 - sysroot_linux-aarch64==2.17
+- wheel
 name: bench_ann_cuda-125_arch-aarch64

conda/environments/bench_ann_cuda-125_arch-x86_64.yaml

@@ -38,5 +38,8 @@ dependencies:
 - pandas
 - pylibraft==24.10.*,>=0.0.0a0
 - pyyaml
+- rapids-build-backend>=0.3.0,<0.4.0.dev0
+- setuptools
 - sysroot_linux-64==2.17
+- wheel
 name: bench_ann_cuda-125_arch-x86_64

conda/recipes/libcuvs/build_libcuvs_tests.sh

@@ -1,5 +1,5 @@
 #!/usr/bin/env bash
 # Copyright (c) 2022-2024, NVIDIA CORPORATION.
 
-./build.sh tests bench-ann --allgpuarch --no-nvtx --build-metrics=tests_bench --incl-cache-stats
+./build.sh tests --allgpuarch --no-nvtx --build-metrics=tests --incl-cache-stats
 cmake --install cpp/build --component testing

cpp/CMakeLists.txt

@@ -55,7 +55,7 @@ option(BUILD_SHARED_LIBS "Build cuvs shared libraries" ON)
 option(BUILD_TESTS "Build cuvs unit-tests" ON)
 option(BUILD_C_LIBRARY "Build cuVS C API library" OFF)
 option(BUILD_C_TESTS "Build cuVS C API tests" OFF)
-option(BUILD_ANN_BENCH "Build cuVS ann benchmarks" OFF)
+option(BUILD_CUVS_BENCH "Build cuVS ann benchmarks" OFF)
 option(BUILD_CAGRA_HNSWLIB "Build CAGRA+hnswlib interface" ON)
 option(CUDA_ENABLE_KERNELINFO "Enable kernel resource usage info" OFF)
 option(CUDA_ENABLE_LINEINFO
@@ -96,7 +96,7 @@ include(CMakeDependentOption)
 
 message(VERBOSE "cuVS: Build cuVS unit-tests: ${BUILD_TESTS}")
 message(VERBOSE "cuVS: Build CPU only components: ${BUILD_CPU_ONLY}")
-message(VERBOSE "cuVS: Build ANN benchmarks: ${BUILD_ANN_BENCH}")
+message(VERBOSE "cuVS: Build ANN benchmarks: ${BUILD_CUVS_BENCH}")
 message(VERBOSE "cuVS: Enable detection of conda environment for dependencies: ${DETECT_CONDA_ENV}")
 message(VERBOSE "cuVS: Disable depreaction warnings " ${DISABLE_DEPRECATION_WARNINGS})
 message(VERBOSE "cuVS: Disable OpenMP: ${DISABLE_OPENMP}")
@@ -188,7 +188,7 @@ endif()
 
 include(cmake/thirdparty/get_cutlass.cmake)
 
-if(BUILD_ANN_BENCH)
+if(BUILD_CUVS_BENCH)
   include(${rapids-cmake-dir}/cpm/gbench.cmake)
   rapids_cpm_gbench(BUILD_STATIC)
 endif()
@@ -651,6 +651,6 @@ endif()
 # ##################################################################################################
 # * build ann benchmark executable -----------------------------------------------
 
-if(BUILD_ANN_BENCH)
+if(BUILD_CUVS_BENCH)
   add_subdirectory(bench/ann/)
 endif()

cpp/bench/ann/CMakeLists.txt

@@ -199,30 +199,19 @@ if(NOT TARGET CUVS_ANN_BENCH_ALL)
 endif()
 
 if(CUVS_ANN_BENCH_USE_HNSWLIB)
-  ConfigureAnnBench(
-    NAME HNSWLIB PATH src/hnswlib/hnswlib_benchmark.cpp LINKS hnswlib::hnswlib
-  )
+  ConfigureAnnBench(NAME HNSWLIB PATH src/hnswlib/hnswlib_benchmark.cpp LINKS hnswlib::hnswlib)
 
 endif()
 
 if(CUVS_ANN_BENCH_USE_CUVS_IVF_PQ)
   ConfigureAnnBench(
-    NAME CUVS_IVF_PQ
-    PATH
-    src/cuvs/cuvs_benchmark.cu
-    src/cuvs/cuvs_ivf_pq.cu
-    LINKS cuvs
+    NAME CUVS_IVF_PQ PATH src/cuvs/cuvs_benchmark.cu src/cuvs/cuvs_ivf_pq.cu LINKS cuvs
   )
 endif()
 
 if(CUVS_ANN_BENCH_USE_CUVS_IVF_FLAT)
   ConfigureAnnBench(
-    NAME CUVS_IVF_FLAT
-    PATH
-    src/cuvs/cuvs_benchmark.cu
-    src/cuvs/cuvs_ivf_flat.cu
-    LINKS
-    cuvs
+    NAME CUVS_IVF_FLAT PATH src/cuvs/cuvs_benchmark.cu src/cuvs/cuvs_ivf_flat.cu LINKS cuvs
   )
 endif()
 
@@ -232,12 +221,8 @@ endif()
 
 if(CUVS_KNN_BENCH_USE_CUVS_BRUTE_FORCE)
   ConfigureAnnBench(
-    NAME
-    CUVS_KNN_BRUTE_FORCE
-    PATH
-    $<$<BOOL:${CUVS_KNN_BENCH_USE_CUVS_BRUTE_FORCE}>:src/cuvs/cuvs_brute_force_knn.cu>
-    LINKS
-    cuvs
+    NAME CUVS_KNN_BRUTE_FORCE PATH
+    $<$<BOOL:${CUVS_KNN_BENCH_USE_CUVS_BRUTE_FORCE}>:src/cuvs/cuvs_brute_force_knn.cu> LINKS cuvs
   )
 endif()
 
@@ -258,45 +243,39 @@ endif()
 
 if(CUVS_ANN_BENCH_USE_CUVS_CAGRA_HNSWLIB)
   ConfigureAnnBench(
-    NAME CUVS_CAGRA_HNSWLIB PATH src/cuvs/cuvs_cagra_hnswlib.cu LINKS cuvs
-    hnswlib::hnswlib
+    NAME CUVS_CAGRA_HNSWLIB PATH src/cuvs/cuvs_cagra_hnswlib.cu LINKS cuvs hnswlib::hnswlib
   )
 endif()
 
 message("CUVS_FAISS_TARGETS: ${CUVS_FAISS_TARGETS}")
 message("CUDAToolkit_LIBRARY_DIR: ${CUDAToolkit_LIBRARY_DIR}")
 if(CUVS_ANN_BENCH_USE_FAISS_CPU_FLAT)
   ConfigureAnnBench(
-    NAME FAISS_CPU_FLAT PATH src/faiss/faiss_cpu_benchmark.cpp LINKS
-    ${CUVS_FAISS_TARGETS}
+    NAME FAISS_CPU_FLAT PATH src/faiss/faiss_cpu_benchmark.cpp LINKS ${CUVS_FAISS_TARGETS}
   )
 endif()
 
 if(CUVS_ANN_BENCH_USE_FAISS_CPU_IVF_FLAT)
   ConfigureAnnBench(
-    NAME FAISS_CPU_IVF_FLAT PATH src/faiss/faiss_cpu_benchmark.cpp LINKS
-    ${CUVS_FAISS_TARGETS}
+    NAME FAISS_CPU_IVF_FLAT PATH src/faiss/faiss_cpu_benchmark.cpp LINKS ${CUVS_FAISS_TARGETS}
   )
 endif()
 
 if(CUVS_ANN_BENCH_USE_FAISS_CPU_IVF_PQ)
   ConfigureAnnBench(
-    NAME FAISS_CPU_IVF_PQ PATH src/faiss/faiss_cpu_benchmark.cpp LINKS
-    ${CUVS_FAISS_TARGETS}
+    NAME FAISS_CPU_IVF_PQ PATH src/faiss/faiss_cpu_benchmark.cpp LINKS ${CUVS_FAISS_TARGETS}
   )
 endif()
 
 if(CUVS_ANN_BENCH_USE_FAISS_GPU_IVF_FLAT AND CUVS_FAISS_ENABLE_GPU)
   ConfigureAnnBench(
-    NAME FAISS_GPU_IVF_FLAT PATH src/faiss/faiss_gpu_benchmark.cu LINKS
-    ${CUVS_FAISS_TARGETS}
+    NAME FAISS_GPU_IVF_FLAT PATH src/faiss/faiss_gpu_benchmark.cu LINKS ${CUVS_FAISS_TARGETS}
   )
 endif()
 
 if(CUVS_ANN_BENCH_USE_FAISS_GPU_IVF_PQ AND CUVS_FAISS_ENABLE_GPU)
   ConfigureAnnBench(
-    NAME FAISS_GPU_IVF_PQ PATH src/faiss/faiss_gpu_benchmark.cu LINKS
-    ${CUVS_FAISS_TARGETS}
+    NAME FAISS_GPU_IVF_PQ PATH src/faiss/faiss_gpu_benchmark.cu LINKS ${CUVS_FAISS_TARGETS}
   )
 endif()
 
@@ -322,13 +301,8 @@ if(CUVS_ANN_BENCH_SINGLE_EXE)
 
   target_link_libraries(
     ANN_BENCH
-    PRIVATE raft::raft
-            nlohmann_json::nlohmann_json
-            benchmark::benchmark
-            dl
-            fmt::fmt-header-only
-            spdlog::spdlog_header_only
-            $<$<BOOL:${NVTX3_HEADERS_FOUND}>:CUDA::nvtx3>
+    PRIVATE raft::raft nlohmann_json::nlohmann_json benchmark::benchmark dl fmt::fmt-header-only
+            spdlog::spdlog_header_only $<$<BOOL:${NVTX3_HEADERS_FOUND}>:CUDA::nvtx3>
   )
   set_target_properties(
     ANN_BENCH

dependencies.yaml

@@ -38,6 +38,7 @@ files:
       - develop
       - bench
       - bench_python
+      - rapids_build_setuptools
   test_cpp:
     output: none
     includes:
@@ -115,6 +116,13 @@ files:
       table: build-system
     includes:
       - rapids_build_setuptools
+  py_rapids_build_py_cuvs_bench:
+    output: pyproject
+    pyproject_dir: python/cuvs_bench
+    extras:
+      table: tool.rapids-build-backend
+      key: requires
+    includes: []
   py_run_cuvs_bench:
     output: pyproject
     pyproject_dir: python/cuvs_bench
@@ -187,7 +195,7 @@ dependencies:
 
   rapids_build_setuptools:
     common:
-      - output_types: [requirements, pyproject]
+      - output_types: [conda, requirements, pyproject]
         packages:
           - &rapids_build_backend rapids-build-backend>=0.3.0,<0.4.0.dev0
           - setuptools
@@ -469,13 +477,12 @@ dependencies:
           - openblas
   bench_python:
     common:
-      - output_types: [conda]
+      - output_types: [conda, pyproject, requirements]
         packages:
+          - click
           - matplotlib
           - pandas
           - pyyaml
-          - pandas
-          - click
   depends_on_librmm:
     common:
       - output_types: conda

python/cuvs/CMakeLists.txt

@@ -83,14 +83,22 @@ if(NOT cuvs_FOUND)
 
   if(NOT CUDA_STATIC_MATH_LIBRARIES AND USE_CUDA_MATH_WHEELS)
     set(rpaths
-      "$ORIGIN/../nvidia/cublas/lib"
-      "$ORIGIN/../nvidia/curand/lib"
-      "$ORIGIN/../nvidia/cusolver/lib"
-      "$ORIGIN/../nvidia/cusparse/lib"
-      "$ORIGIN/../nvidia/nvjitlink/lib"
+        "$ORIGIN/../nvidia/cublas/lib"
+        "$ORIGIN/../nvidia/curand/lib"
+        "$ORIGIN/../nvidia/cusolver/lib"
+        "$ORIGIN/../nvidia/cusparse/lib"
+        "$ORIGIN/../nvidia/nvjitlink/lib"
+    )
+    set_property(
+      TARGET cuvs
+      PROPERTY INSTALL_RPATH ${rpaths}
+      APPEND
+    )
+    set_property(
+      TARGET cuvs_c
+      PROPERTY INSTALL_RPATH ${rpaths}
+      APPEND
     )
-    set_property(TARGET cuvs PROPERTY INSTALL_RPATH ${rpaths} APPEND)
-    set_property(TARGET cuvs_c PROPERTY INSTALL_RPATH ${rpaths} APPEND)
   endif()
 
   set(cython_lib_dir cuvs)

python/cuvs/cuvs/neighbors/filters/CMakeLists.txt

@@ -20,6 +20,5 @@ set(linked_libraries cuvs::cuvs cuvs::c_api)
 rapids_cython_create_modules(
   CXX
   SOURCE_FILES "${cython_sources}"
-  LINKED_LIBRARIES "${linked_libraries}" ASSOCIATED_TARGETS cuvs MODULE_PREFIX
-                   neighbors_prefilter_
+  LINKED_LIBRARIES "${linked_libraries}" ASSOCIATED_TARGETS cuvs MODULE_PREFIX neighbors_prefilter_
 )

python/cuvs_bench/cuvs_bench/config/__init__.py

@@ -0,0 +1,17 @@
+#
+# Copyright (c) 2024, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from .algos.constraints import *

python/cuvs_bench/cuvs_bench/config/algorithms.yaml

@@ -0,0 +1,42 @@
+faiss_gpu_flat:
+  executable: FAISS_GPU_FLAT_ANN_BENCH
+  requires_gpu: true
+faiss_gpu_ivf_flat:
+  executable: FAISS_GPU_IVF_FLAT_ANN_BENCH
+  requires_gpu: true
+faiss_gpu_ivf_pq:
+  executable: FAISS_GPU_IVF_PQ_ANN_BENCH
+  requires_gpu: true
+faiss_gpu_ivf_sq:
+  executable: FAISS_GPU_IVF_PQ_ANN_BENCH
+  requires_gpu: true
+faiss_cpu_flat:
+  executable: FAISS_CPU_FLAT_ANN_BENCH
+  requires_gpu: false
+faiss_cpu_ivf_flat:
+  executable: FAISS_CPU_IVF_FLAT_ANN_BENCH
+  requires_gpu: false
+faiss_cpu_ivf_pq:
+  executable: FAISS_CPU_IVF_PQ_ANN_BENCH
+  requires_gpu: false
+cuvs_ivf_flat:
+  executable: CUVS_IVF_FLAT_ANN_BENCH
+  requires_gpu: true
+cuvs_ivf_pq:
+  executable: CUVS_IVF_PQ_ANN_BENCH
+  requires_gpu: true
+cuvs_cagra:
+  executable: CUVS_CAGRA_ANN_BENCH
+  requires_gpu: true
+cuvs_brute_force:
+  executable: CUVS_BRUTE_FORCE_ANN_BENCH
+  requires_gpu: true
+ggnn:
+  executable: GGNN_ANN_BENCH
+  requires_gpu: true
+hnswlib:
+  executable: HNSWLIB_ANN_BENCH
+  requires_gpu: false
+cuvs_cagra_hnswlib:
+  executable: CUVS_CAGRA_HNSWLIB_ANN_BENCH
+  requires_gpu: true

python/cuvs_bench/cuvs_bench/config/constraints.py → python/cuvs_bench/cuvs_bench/config/algos/constraints/__init__.py

@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2023-2024, NVIDIA CORPORATION.
+# Copyright (c) 2024, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.