test with rmm CI artifacts

rapidsai · Apr 26, 2024 · 122c8eb · 122c8eb
1 parent ad80928
commit 122c8eb
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 1 deletion.
diff --git a/ci/build_cpp.sh b/ci/build_cpp.sh
@@ -13,10 +13,14 @@ export CMAKE_GENERATOR=Ninja
 
 rapids-print-env
 
+LIBRMM_CHANNEL=$(rapids-get-pr-conda-artifact rmm 1544 cpp)
+
 version=$(rapids-generate-version)
 
 rapids-logger "Begin cpp build"
 
-RAPIDS_PACKAGE_VERSION=${version} rapids-conda-retry mambabuild conda/recipes/libraft
+RAPIDS_PACKAGE_VERSION=${version} rapids-conda-retry \
+    --channel "${LIBRMM_CHANNEL}" \
+    mambabuild conda/recipes/libraft
 
 rapids-upload-conda-to-s3 cpp
diff --git a/ci/build_python.sh b/ci/build_python.sh
@@ -16,6 +16,8 @@ rapids-print-env
 rapids-logger "Begin py build"
 
 CPP_CHANNEL=$(rapids-download-conda-from-s3 cpp)
+LIBRMM_CHANNEL=$(rapids-get-pr-conda-artifact rmm 1544 cpp)
+RMM_CHANNEL=$(rapids-get-pr-conda-artifact rmm 1544 python)
 
 version=$(rapids-generate-version)
 git_commit=$(git rev-parse HEAD)
@@ -33,19 +35,25 @@ done
 rapids-conda-retry mambabuild \
   --no-test \
   --channel "${CPP_CHANNEL}" \
+  --channel "${LIBRMM_CHANNEL}" \
+  --channel "${RMM_CHANNEL}" \
   conda/recipes/pylibraft
 
 rapids-conda-retry mambabuild \
   --no-test \
   --channel "${CPP_CHANNEL}" \
   --channel "${RAPIDS_CONDA_BLD_OUTPUT_DIR}" \
+  --channel "${LIBRMM_CHANNEL}" \
+  --channel "${RMM_CHANNEL}" \
   conda/recipes/raft-dask
 
 # Build ann-bench for each cuda and python version
 rapids-conda-retry mambabuild \
   --no-test \
   --channel "${CPP_CHANNEL}" \
   --channel "${RAPIDS_CONDA_BLD_OUTPUT_DIR}" \
+  --channel "${LIBRMM_CHANNEL}" \
+  --channel "${RMM_CHANNEL}" \
   conda/recipes/raft-ann-bench
 
 # Build ann-bench-cpu only in CUDA 11 jobs since it only depends on python
@@ -56,6 +64,8 @@ if [[ ${RAPIDS_CUDA_MAJOR} == "11" ]]; then
   --no-test \
   --channel "${CPP_CHANNEL}" \
   --channel "${RAPIDS_CONDA_BLD_OUTPUT_DIR}" \
+  --channel "${LIBRMM_CHANNEL}" \
+  --channel "${RMM_CHANNEL}" \
   conda/recipes/raft-ann-bench-cpu
 fi
 

diff --git a/ci/test_cpp.sh b/ci/test_cpp.sh
@@ -22,13 +22,18 @@ conda activate test
 set -u
 
 CPP_CHANNEL=$(rapids-download-conda-from-s3 cpp)
+LIBRMM_CHANNEL=$(rapids-get-pr-conda-artifact rmm 1544 cpp)
+RMM_CHANNEL=$(rapids-get-pr-conda-artifact rmm 1544 python)
+
 RAPIDS_TESTS_DIR=${RAPIDS_TESTS_DIR:-"${PWD}/test-results"}/
 mkdir -p "${RAPIDS_TESTS_DIR}"
 
 rapids-print-env
 
 rapids-mamba-retry install \
   --channel "${CPP_CHANNEL}" \
+  --channel "${LIBRMM_CHANNEL}" \
+  --channel "${RMM_CHANNEL}" \
   libraft-headers libraft libraft-tests
 
 rapids-logger "Check GPU usage"

diff --git a/ci/test_python.sh b/ci/test_python.sh
@@ -25,6 +25,9 @@ rapids-logger "Downloading artifacts from previous jobs"
 CPP_CHANNEL=$(rapids-download-conda-from-s3 cpp)
 PYTHON_CHANNEL=$(rapids-download-conda-from-s3 python)
 
+LIBRMM_CHANNEL=$(rapids-get-pr-conda-artifact rmm 1544 cpp)
+RMM_CHANNEL=$(rapids-get-pr-conda-artifact rmm 1544 python)
+
 RAPIDS_TESTS_DIR=${RAPIDS_TESTS_DIR:-"${PWD}/test-results"}
 RAPIDS_COVERAGE_DIR=${RAPIDS_COVERAGE_DIR:-"${PWD}/coverage-results"}
 mkdir -p "${RAPIDS_TESTS_DIR}" "${RAPIDS_COVERAGE_DIR}"
@@ -34,6 +37,8 @@ rapids-print-env
 rapids-mamba-retry install \
   --channel "${CPP_CHANNEL}" \
   --channel "${PYTHON_CHANNEL}" \
+  --channel "${LIBRMM_CHANNEL}" \
+  --channel "${RMM_CHANNEL}" \
   libraft libraft-headers pylibraft raft-dask
 
 rapids-logger "Check GPU usage"