Create Conda test environment in one go (#1101)

Partially addresses issue: rapidsai/build-planning#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: #1101

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

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
 

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.*

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]