Merge branch 'branch-25.02' into branch-25.04

.github/workflows/build.yaml

@@ -47,6 +47,19 @@ jobs:
       node_type: "gpu-l4-latest-1"
       run_script: "ci/build_rust.sh"
       sha: ${{ inputs.sha }}
+  go-build:
+    needs: cpp-build
+    secrets: inherit
+    uses: rapidsai/shared-workflows/.github/workflows/[email protected]
+    with:
+      build_type: ${{ inputs.build_type || 'branch' }}
+      branch: ${{ inputs.branch }}
+      arch: "amd64"
+      date: ${{ inputs.date }}
+      container_image: "rapidsai/ci-conda:latest"
+      node_type: "gpu-v100-latest-1"
+      run_script: "ci/build_go.sh"
+      sha: ${{ inputs.sha }}
   python-build:
     needs: [cpp-build]
     secrets: inherit

.github/workflows/pr.yaml

@@ -22,16 +22,28 @@ jobs:
       - conda-python-tests
       - docs-build
       - rust-build
+      - go-build
       - wheel-build-libcuvs
       - wheel-build-cuvs
       - wheel-tests-cuvs
       - devcontainer
+      - telemetry-setup
     secrets: inherit
     uses: rapidsai/shared-workflows/.github/workflows/[email protected]
     if: always()
     with:
       needs: ${{ toJSON(needs) }}
+  telemetry-setup:
+    continue-on-error: true
+    runs-on: ubuntu-latest
+    env:
+      OTEL_SERVICE_NAME: 'pr-cuvs'
+    steps:
+      - name: Telemetry setup
+        if: ${{ vars.TELEMETRY_ENABLED == 'true' }}
+        uses: rapidsai/shared-actions/telemetry-dispatch-stash-base-env-vars@main
   check-nightly-ci:
+    needs: telemetry-setup
     # Switch to ubuntu-latest once it defaults to a version of Ubuntu that
     # provides at least Python 3.11 (see
     # https://docs.python.org/3/library/datetime.html#datetime.date.fromisoformat)
@@ -44,6 +56,7 @@ jobs:
         with:
           repo: cuvs
   changed-files:
+    needs: telemetry-setup
     secrets: inherit
     uses: rapidsai/shared-workflows/.github/workflows/[email protected]
     with:
@@ -58,13 +71,15 @@ jobs:
           - '!notebooks/**'
           - '!python/**'
           - '!rust/**'
+          - '!go/**'
           - '!thirdparty/LICENSES/**'
         test_notebooks:
           - '**'
           - '!.devcontainer/**'
           - '!.pre-commit-config.yaml'
           - '!README.md'
           - '!rust/**'
+          - '!go/**'
           - '!thirdparty/LICENSES/**'
         test_python:
           - '**'
@@ -75,12 +90,15 @@ jobs:
           - '!img/**'
           - '!notebooks/**'
           - '!rust/**'
+          - '!go/**'
           - '!thirdparty/LICENSES/**'
   checks:
+    needs: telemetry-setup
     secrets: inherit
     uses: rapidsai/shared-workflows/.github/workflows/[email protected]
     with:
       enable_check_generated_files: false
+      ignored_pr_jobs: "telemetry-summarize"
   conda-cpp-build:
     needs: checks
     secrets: inherit
@@ -136,6 +154,16 @@ jobs:
       arch: "amd64"
       container_image: "rapidsai/ci-conda:latest"
       run_script: "ci/build_rust.sh"
+  go-build:
+    needs: conda-cpp-build
+    secrets: inherit
+    uses: rapidsai/shared-workflows/.github/workflows/[email protected]
+    with:
+      build_type: pull-request
+      node_type: "gpu-v100-latest-1"
+      arch: "amd64"
+      container_image: "rapidsai/ci-conda:latest"
+      run_script: "ci/build_go.sh"
   wheel-build-libcuvs:
     needs: checks
     secrets: inherit
@@ -162,6 +190,7 @@ jobs:
       script: ci/test_wheel_cuvs.sh
   devcontainer:
     secrets: inherit
+    needs: telemetry-setup
     uses: rapidsai/shared-workflows/.github/workflows/[email protected]
     with:
       arch: '["amd64"]'
@@ -170,3 +199,13 @@ jobs:
         sccache -z;
         build-all --verbose;
         sccache -s;
+
+  telemetry-summarize:
+    # This job must use a self-hosted runner to record telemetry traces.
+    runs-on: linux-amd64-cpu4
+    needs: pr-builder
+    if: ${{ vars.TELEMETRY_ENABLED == 'true' && !cancelled() }}
+    continue-on-error: true
+    steps:
+      - name: Telemetry summarize
+        uses: rapidsai/shared-actions/telemetry-dispatch-summarize@main

build.sh

@@ -18,14 +18,15 @@ ARGS=$*
 # scripts, and that this script resides in the repo dir!
 REPODIR=$(cd $(dirname $0); pwd)
 
-VALIDARGS="clean libcuvs python rust java docs tests bench-ann examples --uninstall  -v -g -n --compile-static-lib --allgpuarch --no-mg --no-cpu --cpu-only --no-shared-libs --no-nvtx --show_depr_warn --incl-cache-stats --time -h"
+VALIDARGS="clean libcuvs python rust go java docs tests bench-ann examples --uninstall  -v -g -n --compile-static-lib --allgpuarch --no-mg --no-cpu --cpu-only --no-shared-libs --no-nvtx --show_depr_warn --incl-cache-stats --time -h"
 HELP="$0 [<target> ...] [<flag> ...] [--cmake-args=\"<args>\"] [--cache-tool=<tool>] [--limit-tests=<targets>] [--limit-bench-ann=<targets>] [--build-metrics=<filename>]
  where <target> is:
    clean            - remove all existing build artifacts and configuration (start over)
    libcuvs          - build the cuvs C++ code only. Also builds the C-wrapper library
                       around the C++ code.
    python           - build the cuvs Python package
    rust             - build the cuvs Rust bindings
+   go               - build the cuvs Go bindings
    java             - build the cuvs Java bindings
    docs             - build the documentation
    tests            - build the tests
@@ -447,6 +448,13 @@ if (( ${NUMARGS} == 0 )) || hasArg rust; then
     cargo test
 fi
 
+# Build the cuvs Go bindings
+if (( ${NUMARGS} == 0 )) || hasArg go; then
+    cd ${REPODIR}/go
+    go build ./...
+    go test ./...
+fi
+
 # Build the cuvs Java bindings
 if (( ${NUMARGS} == 0 )) || hasArg java; then
     if ! hasArg libcuvs; then

ci/build_go.sh

@@ -0,0 +1,39 @@
+#!/bin/bash
+# Copyright (c) 2024, NVIDIA CORPORATION.
+
+set -euo pipefail
+
+rapids-logger "Create test conda environment"
+. /opt/conda/etc/profile.d/conda.sh
+
+RAPIDS_VERSION="$(rapids-version)"
+
+rapids-dependency-file-generator \
+  --output conda \
+  --file-key go \
+  --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch);py=${RAPIDS_PY_VERSION}" | tee env.yaml
+
+rapids-mamba-retry env create --yes -f env.yaml -n go
+
+# seeing failures on activating the environment here on unbound locals
+# apply workaround from https://github.com/conda/conda/issues/8186#issuecomment-532874667
+set +eu
+conda activate go
+set -eu
+
+rapids-print-env
+
+export CGO_CFLAGS="-I${CONDA_PREFIX}/include"
+export CGO_LDFLAGS="-L${CONDA_PREFIX}/lib -lcudart -lcuvs -lcuvs_c"
+export LD_LIBRARY_PATH="$CONDA_PREFIX/lib:$LD_LIBRARY_PATH"
+export CC=clang
+
+rapids-logger "Downloading artifacts from previous jobs"
+CPP_CHANNEL=$(rapids-download-conda-from-s3 cpp)
+
+rapids-mamba-retry install \
+  --channel "${CPP_CHANNEL}" \
+  "libcuvs=${RAPIDS_VERSION}" \
+  "libraft=${RAPIDS_VERSION}"
+
+bash ./build.sh go

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}

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

@@ -23,6 +23,7 @@ dependencies:
 - dlpack>=0.8,<1.0
 - doxygen>=1.8.20
 - gcc_linux-aarch64=11.*
+- go
 - graphviz
 - ipython
 - libclang==16.0.6

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

@@ -23,6 +23,7 @@ dependencies:
 - dlpack>=0.8,<1.0
 - doxygen>=1.8.20
 - gcc_linux-64=11.*
+- go
 - graphviz
 - ipython
 - libclang==16.0.6

conda/environments/all_cuda-128_arch-aarch64.yaml

@@ -24,6 +24,7 @@ dependencies:
 - dlpack>=0.8,<1.0
 - doxygen>=1.8.20
 - gcc_linux-aarch64=13.*
+- go
 - graphviz
 - ipython
 - libclang==16.0.6

conda/environments/all_cuda-128_arch-x86_64.yaml

@@ -24,6 +24,7 @@ dependencies:
 - dlpack>=0.8,<1.0
 - doxygen>=1.8.20
 - gcc_linux-64=13.*
+- go
 - graphviz
 - ipython
 - libclang==16.0.6

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

@@ -37,7 +37,7 @@ dependencies:
 - libcusparse=11.7.5.86
 - libcuvs==25.4.*,>=0.0.0a0
 - librmm==25.4.*,>=0.0.0a0
-- matplotlib
+- matplotlib-base
 - nccl>=2.19
 - ninja
 - nlohmann_json>=3.11.2
@@ -47,6 +47,8 @@ dependencies:
 - pylibraft==25.4.*,>=0.0.0a0
 - pyyaml
 - rapids-build-backend>=0.3.0,<0.4.0.dev0
+- requests
+- scikit-learn
 - setuptools
 - sysroot_linux-aarch64==2.28
 - wheel

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

@@ -37,7 +37,7 @@ dependencies:
 - libcusparse=11.7.5.86
 - libcuvs==25.4.*,>=0.0.0a0
 - librmm==25.4.*,>=0.0.0a0
-- matplotlib
+- matplotlib-base
 - nccl>=2.19
 - ninja
 - nlohmann_json>=3.11.2
@@ -47,6 +47,8 @@ dependencies:
 - pylibraft==25.4.*,>=0.0.0a0
 - pyyaml
 - rapids-build-backend>=0.3.0,<0.4.0.dev0
+- requests
+- scikit-learn
 - setuptools
 - sysroot_linux-64==2.28
 - wheel

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

@@ -34,7 +34,7 @@ dependencies:
 - libcusparse-dev
 - libcuvs==25.4.*,>=0.0.0a0
 - librmm==25.4.*,>=0.0.0a0
-- matplotlib
+- matplotlib-base
 - nccl>=2.19
 - ninja
 - nlohmann_json>=3.11.2
@@ -43,6 +43,8 @@ dependencies:
 - pylibraft==25.4.*,>=0.0.0a0
 - pyyaml
 - rapids-build-backend>=0.3.0,<0.4.0.dev0
+- requests
+- scikit-learn
 - setuptools
 - sysroot_linux-aarch64==2.28
 - wheel

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

@@ -34,7 +34,7 @@ dependencies:
 - libcusparse-dev
 - libcuvs==25.4.*,>=0.0.0a0
 - librmm==25.4.*,>=0.0.0a0
-- matplotlib
+- matplotlib-base
 - nccl>=2.19
 - ninja
 - nlohmann_json>=3.11.2
@@ -43,6 +43,8 @@ dependencies:
 - pylibraft==25.4.*,>=0.0.0a0
 - pyyaml
 - rapids-build-backend>=0.3.0,<0.4.0.dev0
+- requests
+- scikit-learn
 - setuptools
 - sysroot_linux-64==2.28
 - wheel

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

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/

cpp/CMakeLists.txt

@@ -396,6 +396,10 @@ if(BUILD_SHARED_LIBS)
     src/neighbors/cagra_serialize_half.cu
     src/neighbors/cagra_serialize_int8.cu
     src/neighbors/cagra_serialize_uint8.cu
+    src/neighbors/cagra_merge_float.cu
+    src/neighbors/cagra_merge_half.cu
+    src/neighbors/cagra_merge_int8.cu
+    src/neighbors/cagra_merge_uint8.cu
     src/neighbors/iface/iface_cagra_float_uint32_t.cu
     src/neighbors/iface/iface_cagra_half_uint32_t.cu
     src/neighbors/iface/iface_cagra_int8_t_uint32_t.cu
@@ -693,6 +697,7 @@ target_compile_definitions(cuvs::cuvs INTERFACE $<$<BOOL:${CUVS_NVTX}>:NVTX_ENAB
       $<$<BOOL:${BUILD_CAGRA_HNSWLIB}>:src/neighbors/hnsw_c.cpp>
       src/neighbors/nn_descent_c.cpp
       src/neighbors/refine/refine_c.cpp
+      src/preprocessing/quantize/binary_c.cpp
       src/preprocessing/quantize/scalar_c.cpp
       src/distance/pairwise_distance_c.cpp
     )

cpp/bench/ann/src/common/blob.hpp

@@ -453,7 +453,8 @@ struct blob_mmap {
           size_t size = data_end - data_start;
           mmap_owner owner{size, flags};
           std::fseek(file_.descriptor().value(), data_start, SEEK_SET);
-          size_t n_elems = file_.rows_limit() * file_.n_cols();
+          auto n_elems =
+            static_cast<size_t>(file_.rows_limit()) * static_cast<size_t>(file_.n_cols());
           if (std::fread(owner.data(), sizeof(T), n_elems, file_.descriptor().value()) != n_elems) {
             throw std::runtime_error{"cuvs::bench::blob_mmap() fread " + file_.path() + " failed"};
           }