Run cuvs-bench pytests and end-to-end tests in CI (#574)

PR does the following: - [x] Modifies CI to run pytest and e2e test of cuvs-bench - [x] We need to test the additional time needed to run the tests. They should be fast, but if they are not, then we can add an additional job to run them in parallel. - [x] Adds synthetic test-data generation so the CI jobs don't depend on downloading datasets, and users can have easy testing locally. - [ ] Few improvements to be done to docs, yaml and other things to make it easy for users. - [x] Check in some additional pytests that hadn't been checked in before. Authors: - Dante Gama Dessavre (https://github.com/dantegd) - Corey J. Nolet (https://github.com/cjnolet) - Micka (https://github.com/lowener) Approvers: - James Lamb (https://github.com/jameslamb) - Corey J. Nolet (https://github.com/cjnolet) URL: #574
rapidsai · Feb 7, 2025 · d760d85 · d760d85
1 parent ab597ec
commit d760d85
Show file tree

Hide file tree

Showing 48 changed files with 869 additions and 11,428 deletions.
diff --git a/ci/test_python.sh b/ci/test_python.sh
@@ -1,5 +1,5 @@
 #!/bin/bash
-# Copyright (c) 2022-2024, NVIDIA CORPORATION.
+# Copyright (c) 2022-2025, NVIDIA CORPORATION.
 
 set -euo pipefail
 
@@ -34,7 +34,8 @@ rapids-mamba-retry install \
   --channel "${CPP_CHANNEL}" \
   --channel "${PYTHON_CHANNEL}" \
   "libcuvs=${RAPIDS_VERSION}" \
-  "cuvs=${RAPIDS_VERSION}"
+  "cuvs=${RAPIDS_VERSION}" \
+  "cuvs-bench=${RAPIDS_VERSION}"
 
 rapids-logger "Check GPU usage"
 nvidia-smi
@@ -54,5 +55,17 @@ pytest \
  --cov-report=term \
  tests
 
+rapids-logger "pytest cuvs-bench"
+popd
+pushd python/cuvs_bench/cuvs_bench
+pytest \
+ --cache-clear \
+ --junitxml="${RAPIDS_TESTS_DIR}/junit-cuvs.xml" \
+ --cov-config=../.coveragerc \
+ --cov=cuvs \
+ --cov-report=xml:"${RAPIDS_COVERAGE_DIR}/cuvs-bench-coverage.xml" \
+ --cov-report=term \
+ tests
+
 rapids-logger "Test script exiting with value: $EXITCODE"
 exit ${EXITCODE}
diff --git a/conda/environments/bench_ann_cuda-118_arch-aarch64.yaml b/conda/environments/bench_ann_cuda-118_arch-aarch64.yaml
@@ -37,7 +37,7 @@ dependencies:
 - libcusparse=11.7.5.86
 - libcuvs==25.2.*,>=0.0.0a0
 - librmm==25.2.*,>=0.0.0a0
-- matplotlib
+- matplotlib-base
 - nccl>=2.19
 - ninja
 - nlohmann_json>=3.11.2
@@ -47,6 +47,8 @@ dependencies:
 - pylibraft==25.2.*,>=0.0.0a0
 - pyyaml
 - rapids-build-backend>=0.3.0,<0.4.0.dev0
+- requests
+- scikit-learn
 - setuptools
 - sysroot_linux-aarch64==2.28
 - wheel

diff --git a/conda/environments/bench_ann_cuda-118_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-118_arch-x86_64.yaml
@@ -37,7 +37,7 @@ dependencies:
 - libcusparse=11.7.5.86
 - libcuvs==25.2.*,>=0.0.0a0
 - librmm==25.2.*,>=0.0.0a0
-- matplotlib
+- matplotlib-base
 - nccl>=2.19
 - ninja
 - nlohmann_json>=3.11.2
@@ -47,6 +47,8 @@ dependencies:
 - pylibraft==25.2.*,>=0.0.0a0
 - pyyaml
 - rapids-build-backend>=0.3.0,<0.4.0.dev0
+- requests
+- scikit-learn
 - setuptools
 - sysroot_linux-64==2.28
 - wheel

diff --git a/conda/environments/bench_ann_cuda-128_arch-aarch64.yaml b/conda/environments/bench_ann_cuda-128_arch-aarch64.yaml
@@ -34,7 +34,7 @@ dependencies:
 - libcusparse-dev
 - libcuvs==25.2.*,>=0.0.0a0
 - librmm==25.2.*,>=0.0.0a0
-- matplotlib
+- matplotlib-base
 - nccl>=2.19
 - ninja
 - nlohmann_json>=3.11.2
@@ -43,6 +43,8 @@ dependencies:
 - pylibraft==25.2.*,>=0.0.0a0
 - pyyaml
 - rapids-build-backend>=0.3.0,<0.4.0.dev0
+- requests
+- scikit-learn
 - setuptools
 - sysroot_linux-aarch64==2.28
 - wheel

diff --git a/conda/environments/bench_ann_cuda-128_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-128_arch-x86_64.yaml
@@ -34,7 +34,7 @@ dependencies:
 - libcusparse-dev
 - libcuvs==25.2.*,>=0.0.0a0
 - librmm==25.2.*,>=0.0.0a0
-- matplotlib
+- matplotlib-base
 - nccl>=2.19
 - ninja
 - nlohmann_json>=3.11.2
@@ -43,6 +43,8 @@ dependencies:
 - pylibraft==25.2.*,>=0.0.0a0
 - pyyaml
 - rapids-build-backend>=0.3.0,<0.4.0.dev0
+- requests
+- scikit-learn
 - setuptools
 - sysroot_linux-64==2.28
 - wheel

diff --git a/conda/recipes/cuvs-bench-cpu/meta.yaml b/conda/recipes/cuvs-bench-cpu/meta.yaml
@@ -58,11 +58,12 @@ requirements:
     - click
     - glog {{ glog_version }}
     - h5py {{ h5py_version }}
-    - matplotlib
+    - matplotlib-base
     - numpy >=1.23,<3.0a0
     - pandas
     - pyyaml
     - python
+    - requests
 about:
   home: https://rapids.ai/
   license: Apache-2.0

diff --git a/conda/recipes/cuvs-bench/meta.yaml b/conda/recipes/cuvs-bench/meta.yaml
@@ -93,12 +93,13 @@ requirements:
     - glog {{ glog_version }}
     - cuvs {{ version }}
     - h5py {{ h5py_version }}
-    - matplotlib
+    - matplotlib-base
     - pandas
     - pyyaml
     # rmm is needed to determine if package is gpu-enabled
     - pylibraft ={{ minor_version }}
     - python
+    - requests
     - rmm ={{ minor_version }}
 about:
   home: https://rapids.ai/

diff --git a/dependencies.yaml b/dependencies.yaml
@@ -555,9 +555,16 @@ dependencies:
         packages:
           - click
           - cuvs==25.2.*,>=0.0.0a0
-          - matplotlib
           - pandas
           - pyyaml
+          - requests
+          - scikit-learn
+      - output_types: [conda]
+        packages:
+          - matplotlib-base
+      - output_types: [requirements, pyproject]
+        packages:
+          - matplotlib
   depends_on_libcuvs:
     common:
       - output_types: conda

diff --git a/python/cuvs_bench/.coveragerc b/python/cuvs_bench/.coveragerc
@@ -0,0 +1,3 @@
+# Configuration file for Python coverage tests
+[run]
+source = cuvs_bench
diff --git a/python/cuvs_bench/cuvs_bench/config/algos/cuvs_cagra.yaml b/python/cuvs_bench/cuvs_bench/config/algos/cuvs_cagra.yaml
@@ -11,3 +11,11 @@ groups:
     search:
       itopk: [32, 64, 128, 256, 512]
       search_width: [1, 2, 4, 8, 16, 32, 64]
+  test:
+    build:
+      graph_degree: [32]
+      intermediate_graph_degree: [32]
+      graph_build_algo: ["NN_DESCENT"]
+    search:
+      itopk: [32]
+      search_width: [1, 2]
diff --git a/python/cuvs_bench/cuvs_bench/config/algos/cuvs_cagra_hnswlib.yaml b/python/cuvs_bench/cuvs_bench/config/algos/cuvs_cagra_hnswlib.yaml
@@ -12,3 +12,12 @@ groups:
       ef_construction: [64, 128, 256, 512]
     search:
       ef: [10, 20, 40, 60, 80, 120, 200, 400, 600, 800]
+  test:
+    build:
+      graph_degree: [32]
+      intermediate_graph_degree: [32]
+      graph_build_algo: ["NN_DESCENT"]
+      hierarchy: ["none", "cpu"]
+      ef_construction: [64]
+    search:
+      ef: [10,20]
diff --git a/python/cuvs_bench/cuvs_bench/config/algos/cuvs_ivf_flat.yaml b/python/cuvs_bench/cuvs_bench/config/algos/cuvs_ivf_flat.yaml
@@ -7,3 +7,10 @@ groups:
       niter: [20, 25]
     search:
       nprobe: [1, 5, 10, 50, 100, 200, 500, 1000, 2000]
+  test:
+    build:
+      nlist: [1024]
+      ratio: [1]
+      niter: [20]
+    search:
+      nprobe: [1, 5]
diff --git a/python/cuvs_bench/cuvs_bench/config/algos/cuvs_ivf_pq.yaml b/python/cuvs_bench/cuvs_bench/config/algos/cuvs_ivf_pq.yaml
@@ -39,3 +39,15 @@ groups:
       internalDistanceDtype: ["float"]
       smemLutDtype: ["float", "fp8", "half"]
       refine_ratio: [1]
+  test:
+    build:
+      nlist: [1024]
+      pq_dim: [64]
+      pq_bits: [8]
+      ratio: [10]
+      niter: [25]
+    search:
+      nprobe: [1]
+      internalDistanceDtype: ["float"]
+      smemLutDtype: ["float", "fp8", "half"]
+      refine_ratio: [1]
diff --git a/python/cuvs_bench/cuvs_bench/config/algos/cuvs_mg_cagra.yaml b/python/cuvs_bench/cuvs_bench/config/algos/cuvs_mg_cagra.yaml
@@ -11,3 +11,11 @@ groups:
     search:
       itopk: [32, 64, 128, 256, 512]
       search_width: [1, 2, 4, 8, 16, 32, 64]
+  test:
+    build:
+      graph_degree: [32]
+      intermediate_graph_degree: [32]
+      graph_build_algo: ["NN_DESCENT"]
+    search:
+      itopk: [32]
+      search_width: [1, 2]
diff --git a/python/cuvs_bench/cuvs_bench/config/algos/cuvs_mg_ivf_flat.yaml b/python/cuvs_bench/cuvs_bench/config/algos/cuvs_mg_ivf_flat.yaml
@@ -7,3 +7,10 @@ groups:
       niter: [20, 25]
     search:
       nprobe: [1, 5, 10, 50, 100, 200, 500, 1000, 2000]
+  test:
+    build:
+      nlist: [1024]
+      ratio: [1]
+      niter: [20]
+    search:
+      nprobe: [1, 5]
diff --git a/python/cuvs_bench/cuvs_bench/config/algos/cuvs_mg_ivf_pq.yaml b/python/cuvs_bench/cuvs_bench/config/algos/cuvs_mg_ivf_pq.yaml
@@ -39,3 +39,15 @@ groups:
       internalDistanceDtype: ["float"]
       smemLutDtype: ["float", "fp8", "half"]
       refine_ratio: [1]
+  test:
+    build:
+      nlist: [1024]
+      pq_dim: [64]
+      pq_bits: [8]
+      ratio: [10]
+      niter: [25]
+    search:
+      nprobe: [1]
+      internalDistanceDtype: ["float"]
+      smemLutDtype: ["float", "fp8", "half"]
+      refine_ratio: [1]
diff --git a/python/cuvs_bench/cuvs_bench/config/algos/faiss_cpu_ivf_flat.yaml b/python/cuvs_bench/cuvs_bench/config/algos/faiss_cpu_ivf_flat.yaml
@@ -8,3 +8,11 @@ groups:
     search:
       nprobe: [1, 5, 10, 50, 100, 200]
       refine_ratio: [1]
+  test:
+    build:
+      nlist: [2048]
+      ratio: [10]
+      useFloat16: [False]
+    search:
+      nprobe: [1, 5]
+      refine_ratio: [1]
diff --git a/python/cuvs_bench/cuvs_bench/config/algos/faiss_cpu_ivf_pq.yaml b/python/cuvs_bench/cuvs_bench/config/algos/faiss_cpu_ivf_pq.yaml
@@ -16,3 +16,11 @@ groups:
       bitsPerCode: [8, 6, 5, 4]
     search:
       nprobe: [20, 30, 40, 50, 100, 200, 500, 1000]
+  test:
+    build:
+      nlist: [1024]
+      M: [48]
+      ratio: [10]
+      bitsPerCode: [8]
+    search:
+      nprobe: [1, 5]
diff --git a/python/cuvs_bench/cuvs_bench/config/algos/faiss_gpu_cagra.yaml b/python/cuvs_bench/cuvs_bench/config/algos/faiss_gpu_cagra.yaml
@@ -11,3 +11,11 @@ groups:
     search:
       itopk: [32, 64, 128, 256, 512]
       search_width: [1, 2, 4, 8, 16, 32, 64]
+  test:
+    build:
+      graph_degree: [32]
+      intermediate_graph_degree: [32]
+      graph_build_algo: ["NN_DESCENT"]
+    search:
+      itopk: [32]
+      search_width: [1, 2]
diff --git a/python/cuvs_bench/cuvs_bench/config/algos/faiss_gpu_ivf_flat.yaml b/python/cuvs_bench/cuvs_bench/config/algos/faiss_gpu_ivf_flat.yaml
@@ -28,3 +28,10 @@ groups:
       use_cuvs: [True]
     search:
       nprobe: [10, 20, 30, 50, 100, 200, 500, 1000]
+  test:
+    build:
+      nlist: [1024]
+      ratio: [4]
+      use_cuvs: [False]
+    search:
+      nprobe: [1, 5]
diff --git a/python/cuvs_bench/cuvs_bench/config/algos/faiss_gpu_ivf_pq.yaml b/python/cuvs_bench/cuvs_bench/config/algos/faiss_gpu_ivf_pq.yaml
@@ -74,3 +74,15 @@ groups:
     search:
       nprobe: [20, 30, 40, 50, 100, 200, 500, 1000]
       refine_ratio: [1, 2, 4]
+  test:
+    build:
+      nlist: [1024]
+      M: [96]
+      ratio: [4]
+      usePrecomputed: [True]
+      useFloat16: [True]
+      use_cuvs: [False]
+      bitsPerCode: [8]
+    search:
+      nprobe: [10]
+      refine_ratio: [1, 2]
diff --git a/python/cuvs_bench/cuvs_bench/config/algos/hnswlib.yaml b/python/cuvs_bench/cuvs_bench/config/algos/hnswlib.yaml
@@ -8,3 +8,9 @@ groups:
       efConstruction: [64, 128, 256, 512]
     search:
       ef: [10, 20, 40, 60, 80, 120, 200, 400, 600, 800]
+  test:
+    build:
+      M: [12]
+      efConstruction: [64]
+    search:
+      ef: [10, 20]