Create Conda test environment in one go

Instead of creating the Conda test environment and then installing the
test dependencies into the environment, create the Conda environment
once with the test dependencies.

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,6 +213,11 @@ 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]