From 21e943da00b53063afccc010f856c14eedf63ee6 Mon Sep 17 00:00:00 2001
From: Kyle Edwards <kyedwards@nvidia.com>
Date: Wed, 12 Feb 2025 13:53:06 -0500
Subject: [PATCH] Create Conda CI test env in one step (#2580)

Issue: https://github.com/rapidsai/build-planning/issues/22

Authors:
  - Kyle Edwards (https://github.com/KyleFromNVIDIA)

Approvers:
  - Mike Sarahan (https://github.com/msarahan)

URL: https://github.com/rapidsai/raft/pull/2580
---
 ci/release/update-version.sh |  3 +++
 ci/test_cpp.sh               | 18 ++++++------------
 ci/test_python.sh            | 21 +++++++--------------
 dependencies.yaml            | 30 ++++++++++++++++++++++++++++++
 4 files changed, 46 insertions(+), 26 deletions(-)

diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh
index 244f66e99a..75a096f673 100755
--- a/ci/release/update-version.sh
+++ b/ci/release/update-version.sh
@@ -47,6 +47,9 @@ DEPENDENCIES=(
   pylibraft
   rmm
   rapids-dask-dependency
+  libraft-headers
+  raft-dask
+  libraft-tests
 )
 UCXX_DEPENDENCIES=(
   ucx-py
diff --git a/ci/test_cpp.sh b/ci/test_cpp.sh
index 64400858ec..851af74716 100755
--- a/ci/test_cpp.sh
+++ b/ci/test_cpp.sh
@@ -8,13 +8,17 @@ cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../
 
 . /opt/conda/etc/profile.d/conda.sh
 
-RAPIDS_VERSION="$(rapids-version)"
+CPP_CHANNEL=$(rapids-download-conda-from-s3 cpp)
+RAPIDS_TESTS_DIR=${RAPIDS_TESTS_DIR:-"${PWD}/test-results"}/
+mkdir -p "${RAPIDS_TESTS_DIR}"
 
 rapids-logger "Generate C++ testing dependencies"
 rapids-dependency-file-generator \
   --output conda \
   --file-key test_cpp \
-  --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch)" | tee env.yaml
+  --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch)" \
+  --prepend-channel "${CPP_CHANNEL}" \
+  | tee env.yaml
 
 rapids-mamba-retry env create --yes -f env.yaml -n test
 
@@ -23,18 +27,8 @@ set +u
 conda activate test
 set -u
 
-CPP_CHANNEL=$(rapids-download-conda-from-s3 cpp)
-RAPIDS_TESTS_DIR=${RAPIDS_TESTS_DIR:-"${PWD}/test-results"}/
-mkdir -p "${RAPIDS_TESTS_DIR}"
-
 rapids-print-env
 
-rapids-mamba-retry install \
-  --channel "${CPP_CHANNEL}" \
-  "libraft-headers=${RAPIDS_VERSION}" \
-  "libraft=${RAPIDS_VERSION}" \
-  "libraft-tests=${RAPIDS_VERSION}"
-
 rapids-logger "Check GPU usage"
 nvidia-smi
 
diff --git a/ci/test_python.sh b/ci/test_python.sh
index af93d2e04b..0b57173b03 100755
--- a/ci/test_python.sh
+++ b/ci/test_python.sh
@@ -8,13 +8,18 @@ cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../
 
 . /opt/conda/etc/profile.d/conda.sh
 
-RAPIDS_VERSION="$(rapids-version)"
+rapids-logger "Downloading artifacts from previous jobs"
+CPP_CHANNEL=$(rapids-download-conda-from-s3 cpp)
+PYTHON_CHANNEL=$(rapids-download-conda-from-s3 python)
 
 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}" | tee env.yaml
+  --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch);py=${RAPIDS_PY_VERSION}" \
+  --prepend-channel "${CPP_CHANNEL}" \
+  --prepend-channel "${PYTHON_CHANNEL}" \
+  | tee env.yaml
 
 rapids-mamba-retry env create --yes -f env.yaml -n test
 
@@ -23,24 +28,12 @@ set +u
 conda activate test
 set -u
 
-rapids-logger "Downloading artifacts from previous jobs"
-CPP_CHANNEL=$(rapids-download-conda-from-s3 cpp)
-PYTHON_CHANNEL=$(rapids-download-conda-from-s3 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}"
 
 rapids-print-env
 
-rapids-mamba-retry install \
-  --channel "${CPP_CHANNEL}" \
-  --channel "${PYTHON_CHANNEL}" \
-  "libraft=${RAPIDS_VERSION}" \
-  "libraft-headers=${RAPIDS_VERSION}" \
-  "pylibraft=${RAPIDS_VERSION}" \
-  "raft-dask=${RAPIDS_VERSION}"
-
 rapids-logger "Check GPU usage"
 nvidia-smi
 
diff --git a/dependencies.yaml b/dependencies.yaml
index 225103391f..ab2431a1c4 100644
--- a/dependencies.yaml
+++ b/dependencies.yaml
@@ -29,6 +29,9 @@ files:
     includes:
       - cuda_version
       - test_libraft
+      - depends_on_libraft_headers
+      - depends_on_libraft
+      - depends_on_libraft_tests
   test_python:
     output: none
     includes:
@@ -37,6 +40,10 @@ files:
       - py_version
       - test_pylibraft
       - test_python_common
+      - depends_on_libraft
+      - depends_on_libraft_headers
+      - depends_on_pylibraft
+      - depends_on_raft_dask
   checks:
     output: none
     includes:
@@ -545,6 +552,9 @@ dependencies:
           # pip recognizes the index as a global option for the requirements.txt file
           - --extra-index-url=https://pypi.nvidia.com
           - --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple
+      - output_types: conda
+        packages:
+          - libraft==25.4.*,>=0.0.0a0
     specific:
       - output_types: [requirements, pyproject]
         matrices:
@@ -561,6 +571,26 @@ dependencies:
           - matrix:
             packages:
               - libraft==25.4.*,>=0.0.0a0
+  depends_on_libraft_headers:
+    common:
+      - output_types: conda
+        packages:
+          - libraft-headers==25.4.*,>=0.0.0a0
+  depends_on_pylibraft:
+    common:
+      - output_types: conda
+        packages:
+          - pylibraft==25.4.*,>=0.0.0a0
+  depends_on_raft_dask:
+    common:
+      - output_types: conda
+        packages:
+          - raft-dask==25.4.*,>=0.0.0a0
+  depends_on_libraft_tests:
+    common:
+      - output_types: conda
+        packages:
+          - libraft-tests==25.4.*,>=0.0.0a0
   depends_on_librmm:
     common:
       - output_types: conda