From 632ea74974eb103f328a922898d821adebeb741d Mon Sep 17 00:00:00 2001
From: jakirkham <jakirkham@gmail.com>
Date: Fri, 10 Jan 2025 07:01:28 -0800
Subject: [PATCH] Create Conda test environment in one go (#1101)

Partially addresses issue: https://github.com/rapidsai/build-planning/issues/22

Instead of creating the Conda test environment and then installing `ucx-py` into it, just create the Conda environment with the `ucx-py` package to start with. Thus the environment is constructed only once with everything it needs before testing begins.

Authors:
  - https://github.com/jakirkham

Approvers:
  - Mike Sarahan (https://github.com/msarahan)
  - Peter Andreas Entschev (https://github.com/pentschev)

URL: https://github.com/rapidsai/ucx-py/pull/1101
---
 ci/release/update-version.sh | 17 ++++++++++++-----
 ci/test_python.sh            | 16 ++++++----------
 dependencies.yaml            |  8 +++++++-
 pyproject.toml               |  2 +-
 4 files changed, 26 insertions(+), 17 deletions(-)

diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh
index b6fa1a34..ed22e1a9 100755
--- a/ci/release/update-version.sh
+++ b/ci/release/update-version.sh
@@ -27,8 +27,9 @@ NEXT_SHORT_TAG=${NEXT_MAJOR}.${NEXT_MINOR}
 NEXT_RAPIDS_SHORT_TAG="$(curl -sL https://version.gpuci.io/ucx-py/${NEXT_SHORT_TAG})"
 
 # Need to distutils-normalize the versions for some use cases
+NEXT_SHORT_TAG_PEP440=$(python -c "from packaging.version import Version; print(Version('${NEXT_SHORT_TAG}'))")
+NEXT_FULL_TAG_PEP440=$(python -c "from packaging.version import Version; print(Version('${NEXT_FULL_TAG}'))")
 NEXT_RAPIDS_SHORT_TAG_PEP440=$(python -c "from packaging.version import Version; print(Version('${NEXT_RAPIDS_SHORT_TAG}'))")
-NEXT_RAPIDS_FULL_TAG_PEP440=$(python -c "from packaging.version import Version; print(Version('${NEXT_FULL_TAG}'))")
 
 echo "Preparing release $CURRENT_TAG => $NEXT_FULL_TAG"
 
@@ -39,18 +40,24 @@ function sed_runner() {
 
 DEPENDENCIES=(
   cudf
+  rapids-dask-dependency
 )
-for DEP in "${DEPENDENCIES[@]}"; do
-  for FILE in dependencies.yaml conda/environments/*.yml; do
+UCX_PY_DEPENDENCIES=(
+  ucx-py
+)
+for FILE in dependencies.yaml conda/environments/*.yml; do
+  for DEP in "${DEPENDENCIES[@]}"; do
     sed_runner "/-.* ${DEP}\(-cu[[:digit:]]\{2\}\)\{0,1\}==/ s/==.*/==${NEXT_RAPIDS_SHORT_TAG_PEP440}.*,>=0.0.0a0/g" "${FILE}"
   done
-  sed_runner "/\"${DEP}==/ s/==.*\"/==${NEXT_RAPIDS_SHORT_TAG_PEP440}.*,>=0.0.0a0\"/g" pyproject.toml
+  for DEP in "${UCX_PY_DEPENDENCIES[@]}"; do
+    sed_runner "/-.* ${DEP}\(-cu[[:digit:]]\{2\}\)\{0,1\}==/ s/==.*/==${NEXT_SHORT_TAG_PEP440}.*,>=0.0.0a0/g" "${FILE}"
+  done
 done
 
 for FILE in .github/workflows/*.yaml; do
   sed_runner "/shared-workflows/ s/@.*/@branch-${NEXT_RAPIDS_SHORT_TAG}/g" "${FILE}"
 done
 
-echo "${NEXT_RAPIDS_FULL_TAG_PEP440}" > VERSION
+echo "${NEXT_FULL_TAG_PEP440}" > VERSION
 
 sed_runner "s/--rapids-version=[[:digit:]]\{2\}.[[:digit:]]\{2\}/--rapids-version=${NEXT_RAPIDS_SHORT_TAG}/g" .pre-commit-config.yaml
diff --git a/ci/test_python.sh b/ci/test_python.sh
index c3e7acdf..a74f95b3 100755
--- a/ci/test_python.sh
+++ b/ci/test_python.sh
@@ -6,17 +6,20 @@ set -euo pipefail
 # Support invoking test_python.sh outside the script directory
 cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../
 
-rapids-logger "Create test conda environment"
+rapids-logger "Create test conda environment using artifacts from previous job"
 . /opt/conda/etc/profile.d/conda.sh
 
 UCX_PY_VERSION="$(head -1 ./VERSION)"
+PYTHON_CHANNEL=$(rapids-download-conda-from-s3 python)
 
 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
+rapids-mamba-retry env create -yq -f env.yaml -n test
 conda activate test
 
 rapids-print-env
@@ -47,13 +50,6 @@ run_tests() {
   ./ci/run_benchmark_pytests.sh
 }
 
-rapids-logger "Downloading artifacts from previous jobs"
-PYTHON_CHANNEL=$(rapids-download-conda-from-s3 python)
-
-rapids-mamba-retry install \
-  --channel "${PYTHON_CHANNEL}" \
-  "ucx-py=${UCX_PY_VERSION}"
-
 rapids-logger "Run tests with conda package"
 run_tests
 
diff --git a/dependencies.yaml b/dependencies.yaml
index f9e59df4..8411a292 100644
--- a/dependencies.yaml
+++ b/dependencies.yaml
@@ -17,6 +17,7 @@ files:
       - cuda_version
       - depends_on_cupy
       - py_version
+      - test_ucx_py
       - test_python
   checks:
     output: none
@@ -212,12 +213,17 @@ dependencies:
           # without CUDA, e.g. 'pip install .'
           - matrix: null
             packages: null
+  test_ucx_py:
+    common:
+      - output_types: conda
+        packages:
+          - ucx-py==0.42.*,>=0.0.0a0
   test_python:
     common:
       - output_types: [conda, requirements, pyproject]
         packages:
           - cloudpickle
-          - dask
+          - rapids-dask-dependency==25.2.*,>=0.0.0a0
           - distributed
           - numba>=0.57
           - pytest==7.*
diff --git a/pyproject.toml b/pyproject.toml
index 94174b58..1ebe1d33 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -50,12 +50,12 @@ test = [
     "cloudpickle",
     "cudf==25.2.*,>=0.0.0a0",
     "cupy-cuda12x>=12.0.0",
-    "dask",
     "distributed",
     "numba>=0.57",
     "pytest-asyncio",
     "pytest-rerunfailures",
     "pytest==7.*",
+    "rapids-dask-dependency==25.2.*,>=0.0.0a0",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit dependencies.yaml and run `rapids-dependency-file-generator`.
 
 [project.urls]