[DO NOT MERGE] Test RAFT 3052

ci/build_cpp.sh

@@ -23,6 +23,7 @@ mkdir -p "${RAPIDS_ARTIFACTS_DIR}"
 export RAPIDS_ARTIFACTS_DIR
 
 # populates `RATTLER_CHANNELS` array and `RATTLER_ARGS` array
+source ./ci/use_conda_packages_from_prs.sh
 source rapids-rattler-channel-string
 
 # --no-build-id allows for caching with `sccache`

ci/build_go.sh

@@ -13,10 +13,13 @@ rapids-logger "Create test conda environment"
 rapids-logger "Configuring conda strict channel priority"
 conda config --set channel_priority strict
 
+source ./ci/use_conda_packages_from_prs.sh
 rapids-dependency-file-generator \
   --output conda \
   --file-key go \
   --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch);py=${RAPIDS_PY_VERSION}" \
+  --prepend-channel "${LIBRAFT_CHANNEL}" \
+  --prepend-channel "${RAFT_CHANNEL}" \
   --prepend-channel "${CPP_CHANNEL}" \
   | tee env.yaml
 

ci/build_java.sh

@@ -32,9 +32,12 @@ rapids-logger "Generate Java testing dependencies"
 
 ENV_YAML_DIR="$(mktemp -d)"
 
+source ./ci/use_conda_packages_from_prs.sh
 rapids-dependency-file-generator \
   --output conda \
   --file-key java \
+  --prepend-channel "${LIBRAFT_CHANNEL}" \
+  --prepend-channel "${RAFT_CHANNEL}" \
   --prepend-channel "${CPP_CHANNEL}" \
   --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch)" | tee "${ENV_YAML_DIR}/env.yaml"
 

ci/build_python.sh

@@ -20,6 +20,7 @@ export RAPIDS_PACKAGE_VERSION=${version}
 echo "${version}" > VERSION
 
 # populates `RATTLER_CHANNELS` array and `RATTLER_ARGS` array
+source ./ci/use_conda_packages_from_prs.sh
 source rapids-rattler-channel-string
 
 rapids-logger "Prepending channel ${CPP_CHANNEL} to RATTLER_CHANNELS"

ci/build_rust.sh

@@ -13,10 +13,13 @@ rapids-logger "Create test conda environment"
 rapids-logger "Configuring conda strict channel priority"
 conda config --set channel_priority strict
 
+source ./ci/use_conda_packages_from_prs.sh
 rapids-dependency-file-generator \
   --output conda \
   --file-key rust \
   --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch);py=${RAPIDS_PY_VERSION}" \
+  --prepend-channel "${LIBRAFT_CHANNEL}" \
+  --prepend-channel "${RAFT_CHANNEL}" \
   --prepend-channel "${CPP_CHANNEL}" \
   | tee env.yaml
 

ci/build_wheel_cuvs.sh

@@ -15,6 +15,8 @@ RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen "${RAPIDS_CUDA_VERSION}")"
 #
 # env variable 'PIP_CONSTRAINT' is set up by rapids-init-pip. It constrains all subsequent
 # 'pip install', 'pip download', etc. calls (except those used in 'pip wheel', handled separately in build scripts)
+source ./ci/use_wheels_from_prs.sh
+
 LIBCUVS_WHEELHOUSE=$(rapids-download-from-github "$(rapids-artifact-name wheel_cpp libcuvs cuvs --cuda "$RAPIDS_CUDA_VERSION")")
 echo "libcuvs-${RAPIDS_PY_CUDA_SUFFIX} @ file://$(echo "${LIBCUVS_WHEELHOUSE}"/libcuvs_*.whl)" >> "${PIP_CONSTRAINT}"
 

ci/build_wheel_libcuvs.sh

@@ -19,6 +19,8 @@ rapids-dependency-file-generator \
   --matrix "${matrix_selectors}" \
 | tee /tmp/requirements-build.txt
 
+source ./ci/use_wheels_from_prs.sh
+
 rapids-logger "Installing build requirements"
 rapids-pip-retry install \
     -v \

ci/test_cpp.sh

@@ -11,11 +11,14 @@ conda config --set channel_priority strict
 
 CPP_CHANNEL=$(rapids-download-from-github "$(rapids-artifact-name conda_cpp libcuvs cuvs --cuda "$RAPIDS_CUDA_VERSION")")
 
+source ./ci/use_conda_packages_from_prs.sh
 rapids-logger "Generate C++ testing dependencies"
 rapids-dependency-file-generator \
   --output conda \
   --file-key test_cpp \
   --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch)" \
+  --prepend-channel "${LIBRAFT_CHANNEL}" \
+  --prepend-channel "${RAFT_CHANNEL}" \
   --prepend-channel "${CPP_CHANNEL}" \
   | tee env.yaml
 

ci/test_python.sh

@@ -13,11 +13,14 @@ rapids-logger "Downloading artifacts from previous jobs"
 CPP_CHANNEL=$(rapids-download-from-github "$(rapids-artifact-name conda_cpp libcuvs cuvs --cuda "$RAPIDS_CUDA_VERSION")")
 PYTHON_CHANNEL=$(rapids-download-from-github "$(rapids-artifact-name conda_python cuvs cuvs --stable --cuda "$RAPIDS_CUDA_VERSION")")
 
+source ./ci/use_conda_packages_from_prs.sh
 rapids-logger "Generate Python testing dependencies"
 rapids-dependency-file-generator \
   --output conda \
   --file-key test_python \
   --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch);py=${RAPIDS_PY_VERSION}" \
+  --prepend-channel "${LIBRAFT_CHANNEL}" \
+  --prepend-channel "${RAFT_CHANNEL}" \
   --prepend-channel "${CPP_CHANNEL}" \
   --prepend-channel "${PYTHON_CHANNEL}" \
   | tee env.yaml

ci/test_wheel_cuvs.sh

@@ -9,6 +9,8 @@ source rapids-init-pip
 # Delete system libnccl.so to ensure the wheel is used
 rm -rf /usr/lib64/libnccl*
 
+source ./ci/use_wheels_from_prs.sh
+
 LIBCUVS_WHEELHOUSE=$(rapids-download-from-github "$(rapids-artifact-name wheel_cpp libcuvs cuvs --cuda "$RAPIDS_CUDA_VERSION")")
 CUVS_WHEELHOUSE=$(rapids-download-from-github "$(rapids-artifact-name wheel_python cuvs cuvs --stable --cuda "$RAPIDS_CUDA_VERSION")")
 

ci/use_conda_packages_from_prs.sh

@@ -0,0 +1,30 @@
+#!/bin/bash
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+# SPDX-License-Identifier: Apache-2.0
+
+# download CI artifacts
+LIBRAFT_CHANNEL=$(rapids-get-pr-artifact raft 3052 cpp conda)
+RAFT_CHANNEL=$(rapids-get-pr-artifact raft 3052 python conda)
+
+# For `rattler` builds:
+#
+# Add these channels to the array checked by 'rapids-rattler-channel-string'.
+# This ensures that when conda packages are built with strict channel priority enabled,
+# the locally-downloaded packages will be preferred to remote packages (e.g. nightlies).
+#
+RAPIDS_PREPENDED_CONDA_CHANNELS=(
+    "${LIBRAFT_CHANNEL}"
+    "${RAFT_CHANNEL}"
+)
+export RAPIDS_PREPENDED_CONDA_CHANNELS
+
+# For tests and `conda-build` builds:
+#
+# Add these channels to the system-wide conda configuration.
+# This results in PREPENDING them to conda's channel list, so
+# these packages should be found first if strict channel priority is enabled.
+#
+for _channel in "${RAPIDS_PREPENDED_CONDA_CHANNELS[@]}"
+do
+   conda config --system --add channels "${_channel}"
+done

ci/use_wheels_from_prs.sh

@@ -0,0 +1,18 @@
+#!/bin/bash
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+# SPDX-License-Identifier: Apache-2.0
+
+# initialize PIP_CONSTRAINT
+source rapids-init-pip
+
+RAPIDS_PY_CUDA_SUFFIX=$(rapids-wheel-ctk-name-gen "${RAPIDS_CUDA_VERSION}")
+
+# download wheels, store the directories holding them in variables
+LIBRAFT_WHEELHOUSE=$(rapids-get-pr-artifact raft 3052 cpp wheel)
+PYLIBRAFT_WHEELHOUSE=$(rapids-get-pr-artifact raft 3052 python wheel --pkg_name pylibraft)
+
+# write a pip constraints file saying e.g. "whenever you encounter a requirement for 'librmm-cu12', use this wheel"
+cat > "${PIP_CONSTRAINT}" <<EOF
+libraft-${RAPIDS_PY_CUDA_SUFFIX} @ file://$(echo "${LIBRAFT_WHEELHOUSE}"/libraft_*.whl)
+raft-${RAPIDS_PY_CUDA_SUFFIX} @ file://$(echo "${PYLIBRAFT_WHEELHOUSE}"/pylibraft_*.whl)
+EOF