From 211deb58a380aa385b439e1dae2c8ad15d018452 Mon Sep 17 00:00:00 2001
From: Kyle Edwards <kyedwards@nvidia.com>
Date: Thu, 13 Feb 2025 09:30:55 -0500
Subject: [PATCH] Create Conda CI test env in one step (#684)

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

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

Approvers:
  - James Lamb (https://github.com/jameslamb)

URL: https://github.com/rapidsai/cuvs/pull/684
---
 ci/release/update-version.sh |  2 ++
 ci/test_cpp.sh               | 12 ++++--------
 ci/test_python.sh            | 20 +++++++-------------
 dependencies.yaml            | 20 ++++++++++++++++++++
 4 files changed, 33 insertions(+), 21 deletions(-)

diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh
index 9bc916a87..058e0e46a 100755
--- a/ci/release/update-version.sh
+++ b/ci/release/update-version.sh
@@ -44,7 +44,9 @@ echo "${NEXT_FULL_TAG}" > VERSION
 DEPENDENCIES=(
   dask-cuda
   cuvs
+  cuvs-bench
   libcuvs
+  libcuvs-tests
   libraft
   librmm
   pylibraft
diff --git a/ci/test_cpp.sh b/ci/test_cpp.sh
index 134dc4421..34098175c 100755
--- a/ci/test_cpp.sh
+++ b/ci/test_cpp.sh
@@ -5,13 +5,15 @@ set -euo pipefail
 
 . /opt/conda/etc/profile.d/conda.sh
 
-RAPIDS_VERSION="$(rapids-version)"
+CPP_CHANNEL=$(rapids-download-conda-from-s3 cpp)
 
 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
 
@@ -20,17 +22,11 @@ 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}" \
-  "libcuvs=${RAPIDS_VERSION}" \
-  "libcuvs-tests=${RAPIDS_VERSION}"
-
 rapids-logger "Check GPU usage"
 nvidia-smi
 
diff --git a/ci/test_python.sh b/ci/test_python.sh
index 6ba36bc5a..c040416b0 100755
--- a/ci/test_python.sh
+++ b/ci/test_python.sh
@@ -5,13 +5,18 @@ set -euo pipefail
 
 . /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
 
@@ -20,23 +25,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}" \
-  "libcuvs=${RAPIDS_VERSION}" \
-  "cuvs=${RAPIDS_VERSION}" \
-  "cuvs-bench=${RAPIDS_VERSION}"
-
 rapids-logger "Check GPU usage"
 nvidia-smi
 
diff --git a/dependencies.yaml b/dependencies.yaml
index 483e3fe76..474e629a6 100644
--- a/dependencies.yaml
+++ b/dependencies.yaml
@@ -49,6 +49,8 @@ files:
     includes:
       - cuda_version
       - test_libcuvs
+      - depends_on_libcuvs
+      - depends_on_libcuvs_tests
   test_python:
     output: none
     includes:
@@ -57,6 +59,9 @@ files:
       - py_version
       - test_python_common
       - test_py_cuvs
+      - depends_on_libcuvs
+      - depends_on_cuvs
+      - depends_on_cuvs_bench
   checks:
     output: none
     includes:
@@ -566,6 +571,16 @@ dependencies:
       - output_types: [requirements, pyproject]
         packages:
           - matplotlib
+  depends_on_cuvs:
+    common:
+      - output_types: conda
+        packages:
+          - cuvs==25.4.*,>=0.0.0a0
+  depends_on_cuvs_bench:
+    common:
+      - output_types: conda
+        packages:
+          - cuvs-bench==25.4.*,>=0.0.0a0
   depends_on_libcuvs:
     common:
       - output_types: conda
@@ -590,6 +605,11 @@ dependencies:
             packages:
               - libcuvs-cu11==25.4.*,>=0.0.0a0
           - {matrix: null, packages: [*libcuvs_unsuffixed]}
+  depends_on_libcuvs_tests:
+    common:
+      - output_types: conda
+        packages:
+          - libcuvs-tests==25.4.*,>=0.0.0a0
   depends_on_libraft:
     common:
       - output_types: conda