From 4795210e7bf20cdac442ae10855d06ade0c6177c Mon Sep 17 00:00:00 2001
From: Kyle Edwards <kyedwards@nvidia.com>
Date: Thu, 13 Feb 2025 11:29:14 -0500
Subject: [PATCH] Create Conda CI test env in one step (#1448)

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

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

Approvers:
  - Gil Forsyth (https://github.com/gforsyth)

URL: https://github.com/rapidsai/dask-cuda/pull/1448
---
 ci/release/update-version.sh |  3 ++-
 ci/test_python.sh            | 16 ++++++----------
 dependencies.yaml            |  6 ++++++
 3 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh
index 5457a634f..823d57222 100755
--- a/ci/release/update-version.sh
+++ b/ci/release/update-version.sh
@@ -1,5 +1,5 @@
 #!/bin/bash
-# Copyright (c) 2020, NVIDIA CORPORATION.
+# Copyright (c) 2020-2025, NVIDIA CORPORATION.
 ################################################################################
 # dask-cuda version updater
 ################################################################################
@@ -37,6 +37,7 @@ sed_runner "s/ucxx==.*/ucxx==${NEXT_UCXPY_VERSION}.*,>=0.0.0a0/g" dependencies.y
 
 DEPENDENCIES=(
   cudf
+  dask-cuda
   dask-cudf
   kvikio
   rapids-dask-dependency
diff --git a/ci/test_python.sh b/ci/test_python.sh
index acecce435..6ca5e663e 100755
--- a/ci/test_python.sh
+++ b/ci/test_python.sh
@@ -1,17 +1,20 @@
 #!/bin/bash
-# Copyright (c) 2022-2023, NVIDIA CORPORATION.
+# Copyright (c) 2022-2025, NVIDIA CORPORATION.
 
 set -euo pipefail
 
 . /opt/conda/etc/profile.d/conda.sh
 
-RAPIDS_VERSION="$(rapids-version)"
+rapids-logger "Downloading artifacts from previous jobs"
+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 "${PYTHON_CHANNEL}" \
+  | tee env.yaml
 
 rapids-mamba-retry env create --yes -f env.yaml -n test
 
@@ -20,19 +23,12 @@ set +u
 conda activate test
 set -u
 
-rapids-logger "Downloading artifacts from previous jobs"
-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 "${PYTHON_CHANNEL}" \
-  "dask-cuda=${RAPIDS_VERSION}"
-
 rapids-logger "Check GPU usage"
 nvidia-smi
 
diff --git a/dependencies.yaml b/dependencies.yaml
index fef0a781c..f5954e36e 100644
--- a/dependencies.yaml
+++ b/dependencies.yaml
@@ -21,6 +21,7 @@ files:
       - cuda_version
       - py_version
       - test_python
+      - depends_on_dask_cuda
   checks:
     output: none
     includes:
@@ -206,3 +207,8 @@ dependencies:
               - *dask_cudf_unsuffixed
               - *kvikio_unsuffixed
               - *ucx_py_unsuffixed
+  depends_on_dask_cuda:
+    common:
+      - output_types: conda
+        packages:
+          - dask-cuda==25.4.*,>=0.0.0a0