From c7f2a8b11b1609cc23bb645de6f008b6cdb195a4 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Wed, 15 Apr 2026 13:02:07 -0500 Subject: [PATCH 1/3] Fix device_uvector construction for new RMM resource model --- cpp/src/cluster/detail/kmeans_batched.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/cluster/detail/kmeans_batched.cuh b/cpp/src/cluster/detail/kmeans_batched.cuh index e2fc8d334f..93888490a4 100644 --- a/cpp/src/cluster/detail/kmeans_batched.cuh +++ b/cpp/src/cluster/detail/kmeans_batched.cuh @@ -173,7 +173,7 @@ void accumulate_batch_centroids( cudaStream_t stream = raft::resource::get_cuda_stream(handle); auto workspace = rmm::device_uvector( - batch_data.extent(0), stream, raft::resource::get_workspace_resource(handle)); + batch_data.extent(0), stream, raft::resource::get_workspace_resource_ref(handle)); cuvs::cluster::kmeans::detail::KeyValueIndexOp conversion_op; thrust::transform_iterator, From ae767d44062262ed22cf4771ca655709e20dabbb Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Thu, 16 Apr 2026 11:00:59 -0500 Subject: [PATCH 2/3] Replace deprecated set_current_device_resource_ref with set_current_device_resource --- c/src/core/c_api.cpp | 4 ++-- cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/c/src/core/c_api.cpp b/c/src/core/c_api.cpp index 6c9da46da0..dc1d93379f 100644 --- a/c/src/core/c_api.cpp +++ b/c/src/core/c_api.cpp @@ -168,14 +168,14 @@ extern "C" cuvsError_t cuvsRMMPoolMemoryResourceEnable(int initial_pool_size_per pool_mr.emplace(pool_upstream, initial_size, max_size); - rmm::mr::set_current_device_resource_ref(*pool_mr); + rmm::mr::set_current_device_resource(*pool_mr); }); } extern "C" cuvsError_t cuvsRMMMemoryResourceReset() { return cuvs::core::translate_exceptions([=] { - rmm::mr::reset_current_device_resource_ref(); + rmm::mr::reset_current_device_resource(); pool_mr.reset(); }); } diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h index c0966c0391..537f0f0aa7 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h @@ -76,7 +76,7 @@ class shared_raft_resources { : pool_resource_(rmm::mr::get_current_device_resource_ref(), 1024 * 1024 * 1024ull), resource_(pool_resource_, rmm_oom_callback, nullptr), large_mr_() { - orig_resource_ = rmm::mr::set_current_device_resource_ref(resource_); + orig_resource_ = rmm::mr::set_current_device_resource(resource_); } catch (const std::exception& e) { auto cuda_status = cudaGetLastError(); size_t free = 0; @@ -96,7 +96,7 @@ class shared_raft_resources { shared_raft_resources(const shared_raft_resources& res) = delete; auto operator=(const shared_raft_resources& other) -> shared_raft_resources& = delete; - ~shared_raft_resources() noexcept { rmm::mr::set_current_device_resource_ref(orig_resource_); } + ~shared_raft_resources() noexcept { rmm::mr::set_current_device_resource(orig_resource_); } auto get_large_memory_resource() noexcept -> rmm::device_async_resource_ref { return large_mr_; } From 4df8847d7568edb72f052b27bc2054f379b43270 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Thu, 16 Apr 2026 08:51:13 -0500 Subject: [PATCH 3/3] TEMP: use CI artifacts from RMM/dependency PRs for testing --- ci/build_cpp.sh | 3 +- ci/build_docs.sh | 2 ++ ci/build_go.sh | 2 ++ ci/build_java.sh | 4 ++- ci/build_python.sh | 1 + ci/build_rust.sh | 4 ++- ci/build_standalone_c.sh | 1 + ci/build_wheel_cuvs.sh | 3 +- ci/build_wheel_libcuvs.sh | 3 +- ci/test_cpp.sh | 2 ++ ci/test_python.sh | 2 ++ ci/test_wheel_cuvs.sh | 3 +- ci/use_conda_packages_from_prs.sh | 24 ++++++++++++++ ci/use_wheels_from_prs.sh | 25 +++++++++++++++ .../neighbors/ivf_flat/ivf_flat_search.cuh | 31 ++++++++++--------- 15 files changed, 89 insertions(+), 21 deletions(-) create mode 100755 ci/use_conda_packages_from_prs.sh create mode 100755 ci/use_wheels_from_prs.sh diff --git a/ci/build_cpp.sh b/ci/build_cpp.sh index b4195323ba..5eb3f8d808 100755 --- a/ci/build_cpp.sh +++ b/ci/build_cpp.sh @@ -1,8 +1,9 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail +source ./ci/use_conda_packages_from_prs.sh source rapids-configure-sccache source rapids-date-string diff --git a/ci/build_docs.sh b/ci/build_docs.sh index 65c8f29f8a..2489393edc 100755 --- a/ci/build_docs.sh +++ b/ci/build_docs.sh @@ -3,6 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 set -euo pipefail +source ./ci/use_conda_packages_from_prs.sh rapids-logger "Downloading artifacts from previous jobs" CPP_CHANNEL=$(rapids-download-conda-from-github cpp) @@ -18,6 +19,7 @@ RAPIDS_VERSION_MAJOR_MINOR="$(rapids-version-major-minor)" export RAPIDS_VERSION_MAJOR_MINOR rapids-dependency-file-generator \ + "${RAPIDS_EXTRA_CONDA_CHANNEL_ARGS[@]}" \ --output conda \ --file-key docs \ --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch);py=${RAPIDS_PY_VERSION}" \ diff --git a/ci/build_go.sh b/ci/build_go.sh index 80370048ff..127190491b 100755 --- a/ci/build_go.sh +++ b/ci/build_go.sh @@ -3,6 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 set -euo pipefail +source ./ci/use_conda_packages_from_prs.sh rapids-logger "Downloading artifacts from previous jobs" CPP_CHANNEL=$(rapids-download-conda-from-github cpp) @@ -14,6 +15,7 @@ rapids-logger "Configuring conda strict channel priority" conda config --set channel_priority strict rapids-dependency-file-generator \ + "${RAPIDS_EXTRA_CONDA_CHANNEL_ARGS[@]}" \ --output conda \ --file-key go \ --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch);py=${RAPIDS_PY_VERSION}" \ diff --git a/ci/build_java.sh b/ci/build_java.sh index daf814884d..c958657155 100755 --- a/ci/build_java.sh +++ b/ci/build_java.sh @@ -1,8 +1,9 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail +source ./ci/use_conda_packages_from_prs.sh source rapids-configure-sccache @@ -33,6 +34,7 @@ rapids-logger "Generate Java testing dependencies" ENV_YAML_DIR="$(mktemp -d)" rapids-dependency-file-generator \ + "${RAPIDS_EXTRA_CONDA_CHANNEL_ARGS[@]}" \ --output conda \ --file-key java \ --prepend-channel "${CPP_CHANNEL}" \ diff --git a/ci/build_python.sh b/ci/build_python.sh index c153bdd948..cce29bd41c 100755 --- a/ci/build_python.sh +++ b/ci/build_python.sh @@ -3,6 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 set -euo pipefail +source ./ci/use_conda_packages_from_prs.sh source rapids-configure-sccache source rapids-date-string diff --git a/ci/build_rust.sh b/ci/build_rust.sh index edce9447f2..166eac86cc 100755 --- a/ci/build_rust.sh +++ b/ci/build_rust.sh @@ -1,8 +1,9 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail +source ./ci/use_conda_packages_from_prs.sh rapids-logger "Downloading artifacts from previous jobs" CPP_CHANNEL=$(rapids-download-conda-from-github cpp) @@ -14,6 +15,7 @@ rapids-logger "Configuring conda strict channel priority" conda config --set channel_priority strict rapids-dependency-file-generator \ + "${RAPIDS_EXTRA_CONDA_CHANNEL_ARGS[@]}" \ --output conda \ --file-key rust \ --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch);py=${RAPIDS_PY_VERSION}" \ diff --git a/ci/build_standalone_c.sh b/ci/build_standalone_c.sh index 33123cd79c..1b09b35907 100755 --- a/ci/build_standalone_c.sh +++ b/ci/build_standalone_c.sh @@ -3,6 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 set -euo pipefail +source ./ci/use_conda_packages_from_prs.sh TOOLSET_VERSION=14 NINJA_VERSION=v1.13.1 diff --git a/ci/build_wheel_cuvs.sh b/ci/build_wheel_cuvs.sh index 75be8a6dff..2cdcc5460a 100755 --- a/ci/build_wheel_cuvs.sh +++ b/ci/build_wheel_cuvs.sh @@ -3,9 +3,10 @@ # SPDX-License-Identifier: Apache-2.0 set -euo pipefail - source rapids-init-pip +source ./ci/use_wheels_from_prs.sh + package_dir="python/cuvs" RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen "${RAPIDS_CUDA_VERSION}")" diff --git a/ci/build_wheel_libcuvs.sh b/ci/build_wheel_libcuvs.sh index 6eb6bfa5ab..bd2694c31b 100755 --- a/ci/build_wheel_libcuvs.sh +++ b/ci/build_wheel_libcuvs.sh @@ -3,9 +3,10 @@ # SPDX-License-Identifier: Apache-2.0 set -euo pipefail - source rapids-init-pip +source ./ci/use_wheels_from_prs.sh + package_name="libcuvs" package_dir="python/libcuvs" diff --git a/ci/test_cpp.sh b/ci/test_cpp.sh index 5b7a393f31..3444658256 100755 --- a/ci/test_cpp.sh +++ b/ci/test_cpp.sh @@ -3,6 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 set -euo pipefail +source ./ci/use_conda_packages_from_prs.sh . /opt/conda/etc/profile.d/conda.sh @@ -13,6 +14,7 @@ CPP_CHANNEL=$(rapids-download-conda-from-github cpp) rapids-logger "Generate C++ testing dependencies" rapids-dependency-file-generator \ + "${RAPIDS_EXTRA_CONDA_CHANNEL_ARGS[@]}" \ --output conda \ --file-key test_cpp \ --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch)" \ diff --git a/ci/test_python.sh b/ci/test_python.sh index ebb6ab2762..e12bcb03d2 100755 --- a/ci/test_python.sh +++ b/ci/test_python.sh @@ -3,6 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 set -euo pipefail +source ./ci/use_conda_packages_from_prs.sh . /opt/conda/etc/profile.d/conda.sh @@ -15,6 +16,7 @@ PYTHON_CHANNEL=$(rapids-download-from-github "$(rapids-package-name "conda_pytho rapids-logger "Generate Python testing dependencies" rapids-dependency-file-generator \ + "${RAPIDS_EXTRA_CONDA_CHANNEL_ARGS[@]}" \ --output conda \ --file-key test_python \ --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch);py=${RAPIDS_PY_VERSION}" \ diff --git a/ci/test_wheel_cuvs.sh b/ci/test_wheel_cuvs.sh index c0758da20b..dd85d0a886 100755 --- a/ci/test_wheel_cuvs.sh +++ b/ci/test_wheel_cuvs.sh @@ -3,9 +3,10 @@ # SPDX-License-Identifier: Apache-2.0 set -euo pipefail - source rapids-init-pip +source ./ci/use_wheels_from_prs.sh + # Delete system libnccl.so to ensure the wheel is used rm -rf /usr/lib64/libnccl* diff --git a/ci/use_conda_packages_from_prs.sh b/ci/use_conda_packages_from_prs.sh new file mode 100755 index 0000000000..97892ea0dd --- /dev/null +++ b/ci/use_conda_packages_from_prs.sh @@ -0,0 +1,24 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +# download CI artifacts +LIBRMM_CHANNEL=$(rapids-get-pr-artifact rmm 2361 cpp conda) +RMM_CHANNEL=$(rapids-get-pr-artifact rmm 2361 python conda --stable) +LIBRAFT_CHANNEL=$(rapids-get-pr-artifact raft 2996 cpp conda) +RAFT_CHANNEL=$(rapids-get-pr-artifact raft 2996 python conda --stable) + +RAPIDS_PREPENDED_CONDA_CHANNELS=( + "${LIBRMM_CHANNEL}" + "${RMM_CHANNEL}" + "${LIBRAFT_CHANNEL}" + "${RAFT_CHANNEL}" +) +export RAPIDS_PREPENDED_CONDA_CHANNELS + +RAPIDS_EXTRA_CONDA_CHANNEL_ARGS=() +for _channel in "${RAPIDS_PREPENDED_CONDA_CHANNELS[@]}" +do + conda config --system --add channels "${_channel}" + RAPIDS_EXTRA_CONDA_CHANNEL_ARGS+=(--prepend-channel "${_channel}") +done diff --git a/ci/use_wheels_from_prs.sh b/ci/use_wheels_from_prs.sh new file mode 100755 index 0000000000..a49b027e4b --- /dev/null +++ b/ci/use_wheels_from_prs.sh @@ -0,0 +1,25 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +RAPIDS_PY_CUDA_SUFFIX=$(rapids-wheel-ctk-name-gen "${RAPIDS_CUDA_VERSION}") + +LIBRMM_WHEELHOUSE=$( + RAPIDS_PY_WHEEL_NAME="librmm_${RAPIDS_PY_CUDA_SUFFIX}" rapids-get-pr-artifact rmm 2361 cpp wheel +) +RMM_WHEELHOUSE=$( + rapids-get-pr-artifact rmm 2361 python wheel --stable +) +LIBRAFT_WHEELHOUSE=$( + RAPIDS_PY_WHEEL_NAME="libraft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-get-pr-artifact raft 2996 cpp wheel +) +RAFT_WHEELHOUSE=$( + rapids-get-pr-artifact raft 2996 python wheel --stable --pkg_name pylibraft +) + +cat >> "${PIP_CONSTRAINT}" <(handle, - index, - effective_metric, - params, - queries + offset_q * index.dim(), - queries_batch, - offset_q, - k, - n_probes, - max_samples, - cuvs::distance::is_min_close(effective_metric), - neighbors + offset_q * k, - distances + offset_q * k, - raft::resource::get_workspace_resource_ref(handle), - sample_filter); + search_impl( + handle, + index, + effective_metric, + params, + queries + offset_q * index.dim(), + queries_batch, + offset_q, + k, + n_probes, + max_samples, + cuvs::distance::is_min_close(effective_metric), + neighbors + offset_q * k, + distances + offset_q * k, + raft::resource::get_workspace_resource_ref(handle), + sample_filter); } }