Add GPU and CPU packages for ANN benchmarks (#1773)

Builds on top of #1769 

- [x] Removes `libraft-ann-bench` C++ based package
- [x] Creates `raft-ann-bench` packages that includes C++ tests as well as Python scripts
    - [x] `raft-ann-bench` package includes all tests for CPU and GPU
    - [x] `raft-ann-bench-cpu` package that does not depend on CUDA or RAFT GPU code
- [x] Update docs
- [x] Test artifacts and scripts in CI  
- [x] Minor code cleaning

Some changes include:

- Use `RAPIDS_DATASET_ROOT_DIR` env variable to indicate location of datasets (optional) consistent with other repos: https://docs.rapids.ai/maintainers/datasets/
- CPU and GPU packages are built in the existing GPU build GHA. Only the CUDA 12 jobs build the CPU packages. 
- Small change for invocation of scripts, for example: `python bench/ann/run.py --dataset deep-image-96-inner` is now `python -m raft-ann-bench.run --dataset deep-image-96-inner`, but still scripts meant to be invoked from the command line. 

Future improvements: 

- Remove use of popen python scripts from python scripts. 
- Improve printing and logging
- Allow functions of package to be called from python scripts. 

Closes #1744

Authors:
  - Dante Gama Dessavre (https://github.com/dantegd)
  - Artem M. Chirkin (https://github.com/achirkin)
  - Divye Gala (https://github.com/divyegala)
  - Corey J. Nolet (https://github.com/cjnolet)

Approvers:
  - Corey J. Nolet (https://github.com/cjnolet)
  - Divye Gala (https://github.com/divyegala)
  - AJ Schmidt (https://github.com/ajschmidt8)
  - Artem M. Chirkin (https://github.com/achirkin)

URL: #1773

build.sh

@@ -18,7 +18,7 @@ ARGS=$*
 # scripts, and that this script resides in the repo dir!
 REPODIR=$(cd $(dirname $0); pwd)
 
-VALIDARGS="clean libraft pylibraft raft-dask docs tests template bench-prims bench-ann clean --uninstall  -v -g -n --compile-lib --compile-static-lib --allgpuarch --no-nvtx --show_depr_warn --incl-cache-stats --time -h"
+VALIDARGS="clean libraft pylibraft raft-dask docs tests template bench-prims bench-ann clean --uninstall  -v -g -n --compile-lib --compile-static-lib --allgpuarch --no-nvtx --cpu-only --show_depr_warn --incl-cache-stats --time -h"
 HELP="$0 [<target> ...] [<flag> ...] [--cmake-args=\"<args>\"] [--cache-tool=<tool>] [--limit-tests=<targets>] [--limit-bench-prims=<targets>] [--limit-bench-ann=<targets>] [--build-metrics=<filename>]
  where <target> is:
    clean            - remove all existing build artifacts and configuration (start over)
@@ -39,6 +39,7 @@ HELP="$0 [<target> ...] [<flag> ...] [--cmake-args=\"<args>\"] [--cache-tool=<to
    --uninstall                 - uninstall files for specified targets which were built and installed prior
    --compile-lib               - compile shared library for all components
    --compile-static-lib        - compile static library for all components
+   --cpu-only                  - build CPU only components without CUDA. Applies to bench-ann only currently.
    --limit-tests               - semicolon-separated list of test executables to compile (e.g. NEIGHBORS_TEST;CLUSTER_TEST)
    --limit-bench-prims         - semicolon-separated list of prims benchmark executables to compute (e.g. NEIGHBORS_PRIMS_BENCH;CLUSTER_PRIMS_BENCH)
    --limit-bench-ann           - semicolon-separated list of ann benchmark executables to compute (e.g. HNSWLIB_ANN_BENCH;RAFT_IVF_PQ_ANN_BENCH)
@@ -71,6 +72,7 @@ BUILD_TESTS=OFF
 BUILD_TYPE=Release
 BUILD_PRIMS_BENCH=OFF
 BUILD_ANN_BENCH=OFF
+BUILD_CPU_ONLY=OFF
 COMPILE_LIBRARY=OFF
 INSTALL_TARGET=install
 BUILD_REPORT_METRICS=""
@@ -152,7 +154,7 @@ function limitTests {
             # Remove the full LIMIT_TEST_TARGETS argument from list of args so that it passes validArgs function
             ARGS=${ARGS//--limit-tests=$LIMIT_TEST_TARGETS/}
             TEST_TARGETS=${LIMIT_TEST_TARGETS}
-	    echo "Limiting tests to $TEST_TARGETS"
+            echo "Limiting tests to $TEST_TARGETS"
         fi
     fi
 }
@@ -347,7 +349,13 @@ fi
 if hasArg bench-ann || (( ${NUMARGS} == 0 )); then
     BUILD_ANN_BENCH=ON
     CMAKE_TARGET="${CMAKE_TARGET};${ANN_BENCH_TARGETS}"
-    COMPILE_LIBRARY=ON
+    if hasArg --cpu-only; then
+        COMPILE_LIBRARY=OFF
+        BUILD_CPU_ONLY=ON
+        NVTX=OFF
+    else
+        COMPILE_LIBRARY=ON
+    fi
 fi
 
 if hasArg --no-nvtx; then
@@ -420,6 +428,7 @@ if (( ${NUMARGS} == 0 )) || hasArg libraft || hasArg docs || hasArg tests || has
           -DBUILD_TESTS=${BUILD_TESTS} \
           -DBUILD_PRIMS_BENCH=${BUILD_PRIMS_BENCH} \
           -DBUILD_ANN_BENCH=${BUILD_ANN_BENCH} \
+          -DBUILD_CPU_ONLY=${BUILD_CPU_ONLY} \
           -DCMAKE_MESSAGE_LOG_LEVEL=${CMAKE_LOG_LEVEL} \
           ${CACHE_ARGS} \
           ${EXTRA_CMAKE_ARGS}
@@ -493,6 +502,10 @@ if (( ${NUMARGS} == 0 )) || hasArg raft-dask; then
         python -m pip install --no-build-isolation --no-deps ${REPODIR}/python/raft-dask
 fi
 
+# Build and (optionally) install the raft-ann-bench Python package
+if (( ${NUMARGS} == 0 )) || hasArg bench-ann; then
+    python -m pip install --no-build-isolation --no-deps ${REPODIR}/python/raft-ann-bench -vvv
+fi
 
 if hasArg docs; then
     set -x

ci/build_python.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-# Copyright (c) 2022, NVIDIA CORPORATION.
+# Copyright (c) 2022-2023, NVIDIA CORPORATION.
 
 set -euo pipefail
 
@@ -26,4 +26,22 @@ rapids-mamba-retry mambabuild \
   --channel "${RAPIDS_CONDA_BLD_OUTPUT_DIR}" \
   conda/recipes/raft-dask
 
+# Build ann-bench for each cuda and python version
+rapids-mamba-retry mambabuild \
+--no-test \
+--channel "${CPP_CHANNEL}" \
+--channel "${RAPIDS_CONDA_BLD_OUTPUT_DIR}" \
+conda/recipes/raft-ann-bench
+
+# Build ann-bench-cpu only in CUDA 11 jobs since it only depends on python
+# version
+RAPIDS_CUDA_MAJOR="${RAPIDS_CUDA_VERSION%%.*}"
+if [[ ${RAPIDS_CUDA_MAJOR} == "11" ]]; then
+  rapids-mamba-retry mambabuild \
+  --no-test \
+  --channel "${CPP_CHANNEL}" \
+  --channel "${RAPIDS_CONDA_BLD_OUTPUT_DIR}" \
+  conda/recipes/raft-ann-bench-cpu
+fi
+
 rapids-upload-conda-to-s3 python

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

@@ -31,12 +31,9 @@ dependencies:
 - libcusparse-dev=11.7.5.86
 - libcusparse=11.7.5.86
 - libfaiss>=1.7.1
-- matplotlib
 - nccl>=2.9.9
 - ninja
 - nlohmann_json>=3.11.2
-- pandas
-- pyyaml
 - scikit-build>=0.13.1
 - sysroot_linux-64==2.17
 name: bench_ann_cuda-118_arch-x86_64

conda/recipes/libraft/build_libraft.sh

@@ -1,4 +1,4 @@
 #!/usr/bin/env bash
 # Copyright (c) 2022-2023, NVIDIA CORPORATION.
 
-./build.sh libraft --allgpuarch --compile-lib --build-metrics=compile_lib --incl-cache-stats --no-nvtx
+./build.sh libraft -v --allgpuarch --compile-lib --build-metrics=compile_lib --incl-cache-stats --no-nvtx

conda/recipes/libraft/build_libraft_headers.sh

@@ -1,4 +1,4 @@
 #!/usr/bin/env bash
 # Copyright (c) 2022-2023, NVIDIA CORPORATION.
 
-./build.sh libraft --allgpuarch --no-nvtx
+./build.sh libraft -v --allgpuarch --no-nvtx

conda/recipes/libraft/build_libraft_template.sh

@@ -2,4 +2,4 @@
 # Copyright (c) 2022-2023, NVIDIA CORPORATION.
 
 # Just building template so we verify it uses libraft.so and fail if it doesn't build
-./build.sh template
+./build.sh template -v

conda/recipes/libraft/build_libraft_tests.sh

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

conda/recipes/libraft/meta.yaml

@@ -320,62 +320,3 @@ outputs:
       home: https://rapids.ai/
       license: Apache-2.0
       summary: libraft template
-  - name: libraft-ann-bench
-    version: {{ version }}
-    script: build_libraft_nn_bench.sh
-    build:
-      script_env: *script_env
-      number: {{ GIT_DESCRIBE_NUMBER }}
-      string: cuda{{ cuda_major }}_{{ date_string }}_{{ GIT_DESCRIBE_HASH }}_{{ GIT_DESCRIBE_NUMBER }}
-      ignore_run_exports_from:
-        {% if cuda_major == "11" %}
-        - {{ compiler('cuda11') }}
-        {% endif %}
-    requirements:
-      build:
-        - {{ compiler('c') }}
-        - {{ compiler('cxx') }}
-        {% if cuda_major == "11" %}
-        - {{ compiler('cuda11') }} ={{ cuda_version }}
-        {% else %}
-        - {{ compiler('cuda') }}
-        {% endif %}
-        - cuda-version ={{ cuda_version }}
-        - cmake {{ cmake_version }}
-        - ninja
-        - sysroot_{{ target_platform }} {{ sysroot_version }}
-      host:
-        - {{ pin_subpackage('libraft', exact=True) }}
-        - cuda-version ={{ cuda_version }}
-        {% if cuda_major == "11" %}
-        - cuda-profiler-api {{ cuda11_cuda_profiler_api_run_version }}
-        - libcublas {{ cuda11_libcublas_host_version }}
-        - libcublas-dev {{ cuda11_libcublas_host_version }}
-        {% else %}
-        - cuda-profiler-api
-        - libcublas-dev
-        {% endif %}
-        - glog {{ glog_version }}
-        - nlohmann_json {{ nlohmann_json_version }}
-        # Temporarily ignore faiss benchmarks on CUDA 12 because packages do not exist yet
-        {% if cuda_major == "11" %}
-        - faiss-proc=*=cuda
-        - libfaiss {{ faiss_version }}
-        {% endif %}
-      run:
-        - {{ pin_subpackage('libraft', exact=True) }}
-        - {{ pin_compatible('cuda-version', max_pin='x', min_pin='x') }}
-        {% if cuda_major == "11" %}
-        - cudatoolkit
-        {% endif %}
-        - glog {{ glog_version }}
-        # Temporarily ignore faiss benchmarks on CUDA 12 because packages do not exist yet
-        {% if cuda_major == "11" %}
-        - faiss-proc=*=cuda
-        - libfaiss {{ faiss_version }}
-        {% endif %}
-        - h5py {{ h5py_version }}
-    about:
-      home: https://rapids.ai/
-      license: Apache-2.0
-      summary: libraft ann bench

conda/recipes/libraft/build_libraft_nn_bench.sh → conda/recipes/raft-ann-bench-cpu/build.sh

@@ -1,5 +1,5 @@
 #!/usr/bin/env bash
 # Copyright (c) 2023, NVIDIA CORPORATION.
 
-./build.sh bench-ann --allgpuarch --no-nvtx --build-metrics=bench_ann --incl-cache-stats
+./build.sh bench-ann --cpu-only --no-nvtx --build-metrics=bench_ann_cpu --incl-cache-stats
 cmake --install cpp/build --component ann_bench

conda/recipes/raft-ann-bench-cpu/conda_build_config.yaml

@@ -0,0 +1,20 @@
+c_compiler_version:
+  - 11
+
+cxx_compiler_version:
+  - 11
+
+sysroot_version:
+  - "2.17"
+
+cmake_version:
+  - ">=3.26.4"
+
+glog_version:
+  - ">=0.6.0"
+
+h5py_version:
+  - ">=3.8.0"
+
+nlohmann_json_version:
+  - ">=3.11.2"

conda/recipes/raft-ann-bench-cpu/meta.yaml

@@ -0,0 +1,64 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION.
+
+# Usage:
+#   conda build . -c conda-forge -c nvidia -c rapidsai
+{% set version = environ.get('GIT_DESCRIBE_TAG', '0.0.0.dev').lstrip('v') %}
+{% set minor_version = version.split('.')[0] + '.' + version.split('.')[1] %}
+{% set py_version = environ['CONDA_PY'] %}
+{% set cuda_version = '.'.join(environ['RAPIDS_CUDA_VERSION'].split('.')[:2]) %}
+{% set date_string = environ['RAPIDS_DATE_STRING'] %}
+
+package:
+  name: raft-ann-bench-cpu
+  version: {{ version }}
+  script: build.sh
+
+source:
+  git_url: ../../..
+
+build:
+  script_env:
+    - AWS_ACCESS_KEY_ID
+    - AWS_SECRET_ACCESS_KEY
+    - AWS_SESSION_TOKEN
+    - CMAKE_C_COMPILER_LAUNCHER
+    - CMAKE_CUDA_COMPILER_LAUNCHER
+    - CMAKE_CXX_COMPILER_LAUNCHER
+    - CMAKE_GENERATOR
+    - PARALLEL_LEVEL
+    - RAPIDS_ARTIFACTS_DIR
+    - SCCACHE_BUCKET
+    - SCCACHE_IDLE_TIMEOUT
+    - SCCACHE_REGION
+    - SCCACHE_S3_KEY_PREFIX=libraft-aarch64 # [aarch64]
+    - SCCACHE_S3_KEY_PREFIX=libraft-linux64 # [linux64]
+    - SCCACHE_S3_USE_SSL
+  number: {{ GIT_DESCRIBE_NUMBER }}
+  string: py{{ py_version }}_{{ date_string }}_{{ GIT_DESCRIBE_HASH }}_{{ GIT_DESCRIBE_NUMBER }}
+
+requirements:
+  build:
+    - {{ compiler('c') }}
+    - {{ compiler('cxx') }}
+    - cmake {{ cmake_version }}
+    - ninja
+    - sysroot_{{ target_platform }} {{ sysroot_version }}
+
+  host:
+    - glog {{ glog_version }}
+    - matplotlib
+    - nlohmann_json {{ nlohmann_json_version }}
+    - python
+    - pyyaml
+
+  run:
+    - glog {{ glog_version }}
+    - h5py {{ h5py_version }}
+    - matplotlib
+    - python
+    - pyyaml
+
+about:
+  home: https://rapids.ai/
+  license: Apache-2.0
+  summary: libraft ann bench

conda/recipes/raft-ann-bench/build.sh

@@ -0,0 +1,5 @@
+#!/usr/bin/env bash
+# Copyright (c) 2023, NVIDIA CORPORATION.
+
+./build.sh bench-ann -v --allgpuarch --no-nvtx --build-metrics=bench_ann --incl-cache-stats
+cmake --install cpp/build --component ann_bench

conda/recipes/raft-ann-bench/conda_build_config.yaml

@@ -0,0 +1,73 @@
+c_compiler_version:
+  - 11
+
+cxx_compiler_version:
+  - 11
+
+cuda_compiler:
+  - cuda-nvcc
+
+cuda11_compiler:
+  - nvcc
+
+sysroot_version:
+  - "2.17"
+
+cmake_version:
+  - ">=3.26.4"
+
+nccl_version:
+  - ">=2.9.9"
+
+gtest_version:
+  - ">=1.13.0"
+
+glog_version:
+  - ">=0.6.0"
+
+faiss_version:
+  - ">=1.7.1"
+
+h5py_version:
+  - ">=3.8.0"
+
+nlohmann_json_version:
+  - ">=3.11.2"
+
+# The CTK libraries below are missing from the conda-forge::cudatoolkit package
+# for CUDA 11. The "*_host_*" version specifiers correspond to `11.8` packages
+# and the "*_run_*" version specifiers correspond to `11.x` packages.
+
+cuda11_libcublas_host_version:
+  - "=11.11.3.6"
+
+cuda11_libcublas_run_version:
+  - ">=11.5.2.43,<12.0.0"
+
+cuda11_libcurand_host_version:
+  - "=10.3.0.86"
+
+cuda11_libcurand_run_version:
+  - ">=10.2.5.43,<10.3.1"
+
+cuda11_libcusolver_host_version:
+  - "=11.4.1.48"
+
+cuda11_libcusolver_run_version:
+  - ">=11.2.0.43,<11.4.2"
+
+cuda11_libcusparse_host_version:
+  - "=11.7.5.86"
+
+cuda11_libcusparse_run_version:
+  - ">=11.6.0.43,<12.0.0"
+
+# `cuda-profiler-api` only has `11.8.0` and `12.0.0` packages for all
+# architectures. The "*_host_*" version specifiers correspond to `11.8` packages and the
+# "*_run_*" version specifiers correspond to `11.x` packages.
+
+cuda11_cuda_profiler_api_host_version:
+  - "=11.8.86"
+
+cuda11_cuda_profiler_api_run_version:
+  - ">=11.4.240,<12"