From a6dab457ec1fabd5bbf085c1544b67d0cb08d152 Mon Sep 17 00:00:00 2001
From: Jordan Jacobelli <jjacobelli@nvidia.com>
Date: Tue, 14 Dec 2021 18:14:06 +0100
Subject: [PATCH] Update ucx-py version on release using rvc

Signed-off-by: Jordan Jacobelli <jjacobelli@nvidia.com>
---
 ci/benchmark/build.sh                | 5 ++++-
 ci/cpu/build.sh                      | 3 +++
 ci/gpu/build.sh                      | 5 ++++-
 ci/release/update-version.sh         | 7 +++++++
 conda/recipes/cugraph/meta.yaml      | 5 +++--
 conda/recipes/pylibcugraph/meta.yaml | 3 ++-
 6 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/ci/benchmark/build.sh b/ci/benchmark/build.sh
index a25cb99e318..aaa378ea231 100644
--- a/ci/benchmark/build.sh
+++ b/ci/benchmark/build.sh
@@ -47,6 +47,9 @@ export GPUCI_CONDA_RETRY_SLEEP=30
 export DATASETS_DIR=${WORKSPACE}/datasets
 export BENCHMARKS_DIR=${WORKSPACE}/benchmarks
 
+# ucx-py version
+export UCX_PY_VERSION='0.24.*'
+
 ##########################################
 # Environment Setup                      #
 ##########################################
@@ -74,7 +77,7 @@ gpuci_conda_retry install -c nvidia -c rapidsai -c rapidsai-nightly -c conda-for
       "cudatoolkit=$CUDA_REL" \
       "dask-cudf=${MINOR_VERSION}" \
       "dask-cuda=${MINOR_VERSION}" \
-      "ucx-py=0.24.*" \
+      "ucx-py=${UCX_PY_VERSION}" \
       "ucx-proc=*=gpu" \
       "rapids-build-env=${MINOR_VERSION}" \
       rapids-pytest-benchmark
diff --git a/ci/cpu/build.sh b/ci/cpu/build.sh
index e8292522be3..7361b506e81 100755
--- a/ci/cpu/build.sh
+++ b/ci/cpu/build.sh
@@ -33,6 +33,9 @@ export GPUCI_CONDA_RETRY_SLEEP=30
 export CMAKE_GENERATOR="Ninja"
 export CONDA_BLD_DIR="${WORKSPACE}/.conda-bld"
 
+# ucx-py version
+export UCX_PY_VERSION='0.24.*'
+
 ################################################################################
 # SETUP - Check environment
 ################################################################################
diff --git a/ci/gpu/build.sh b/ci/gpu/build.sh
index c7680ed7c45..b59ba0b3968 100755
--- a/ci/gpu/build.sh
+++ b/ci/gpu/build.sh
@@ -39,6 +39,9 @@ cd $WORKSPACE
 export GIT_DESCRIBE_TAG=`git describe --tags`
 export MINOR_VERSION=`echo $GIT_DESCRIBE_TAG | grep -o -E '([0-9]+\.[0-9]+)'`
 
+# ucx-py version
+export UCX_PY_VERSION='0.24.*'
+
 ################################################################################
 # SETUP - Check environment
 ################################################################################
@@ -63,7 +66,7 @@ gpuci_mamba_retry install -y \
       "cudatoolkit=$CUDA_REL" \
       "dask-cudf=${MINOR_VERSION}" \
       "dask-cuda=${MINOR_VERSION}" \
-      "ucx-py=0.24.*" \
+      "ucx-py=${UCX_PY_VERSION}" \
       "ucx-proc=*=gpu" \
       "rapids-build-env=$MINOR_VERSION.*" \
       "rapids-notebook-env=$MINOR_VERSION.*" \
diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh
index 7ffe6bcae13..1dc87a922d7 100755
--- a/ci/release/update-version.sh
+++ b/ci/release/update-version.sh
@@ -30,6 +30,7 @@ CURRENT_SHORT_TAG=${CURRENT_MAJOR}.${CURRENT_MINOR}
 NEXT_MAJOR=$(echo $NEXT_FULL_TAG | awk '{split($0, a, "."); print a[1]}')
 NEXT_MINOR=$(echo $NEXT_FULL_TAG | awk '{split($0, a, "."); print a[2]}')
 NEXT_SHORT_TAG=${NEXT_MAJOR}.${NEXT_MINOR}
+NEXT_UCX_PY_VERSION="$(curl -sL https://version.gpuci.io/rapids/${NEXT_SHORT_TAG}).*"
 
 echo "Preparing release $CURRENT_TAG => $NEXT_FULL_TAG"
 
@@ -54,7 +55,13 @@ for FILE in conda/environments/*.yml; do
    sed_runner "s/dask-cuda=${CURRENT_SHORT_TAG}/dask-cuda=${NEXT_SHORT_TAG}/g" ${FILE};
    sed_runner "s/dask-cudf=${CURRENT_SHORT_TAG}/dask-cudf=${NEXT_SHORT_TAG}/g" ${FILE};
    sed_runner "s/cuxfilter=${CURRENT_SHORT_TAG}/cuxfilter=${NEXT_SHORT_TAG}/g" ${FILE};
+   sed_runner "s/ucx-py=.*/ucx-py=${NEXT_UCX_PY_VERSION}/g" ${FILE};
 done
 
 # Doxyfile update
 sed_runner "s|\(TAGFILES.*librmm/\).*|\1${NEXT_SHORT_TAG}|" cpp/doxygen/Doxyfile
+
+# ucx-py version
+sed_runner "s/export UCX_PY_VERSION=.*/export UCX_PY_VERSION='${NEXT_UCX_PY_VERSION}'/g" ci/benchmark/build.sh
+sed_runner "s/export UCX_PY_VERSION=.*/export UCX_PY_VERSION='${NEXT_UCX_PY_VERSION}'/g" ci/gpu/build.sh
+sed_runner "s/export UCX_PY_VERSION=.*/export UCX_PY_VERSION='${NEXT_UCX_PY_VERSION}'/g" ci/cpu/build.sh
diff --git a/conda/recipes/cugraph/meta.yaml b/conda/recipes/cugraph/meta.yaml
index 720f9855a30..df6199f4839 100644
--- a/conda/recipes/cugraph/meta.yaml
+++ b/conda/recipes/cugraph/meta.yaml
@@ -7,6 +7,7 @@
 {% set cuda_version='.'.join(environ.get('CUDA', 'unknown').split('.')[:2]) %}
 {% set cuda_major=cuda_version.split('.')[0] %}
 {% set py_version=environ.get('CONDA_PY', 36) %}
+{% set ucx_py_version=environ.get('UCX_PY_VERSION') %}
 package:
   name: cugraph
   version: {{ version }}
@@ -29,7 +30,7 @@ requirements:
     - cython>=0.29,<0.30
     - libcugraph={{ version }}
     - cudf={{ minor_version }}
-    - ucx-py 0.24
+    - ucx-py {{ ucx_py_version }}
     - ucx-proc=*=gpu
     - cudatoolkit {{ cuda_version }}.*
   run:
@@ -40,7 +41,7 @@ requirements:
     - dask-cuda {{ minor_version }}
     - dask>=2021.09.1,<=2021.11.2
     - distributed>=2021.09.1,<=2021.11.2
-    - ucx-py 0.24
+    - ucx-py {{ ucx_py_version }}
     - ucx-proc=*=gpu
     - {{ pin_compatible('cudatoolkit', max_pin='x', min_pin='x') }}
 
diff --git a/conda/recipes/pylibcugraph/meta.yaml b/conda/recipes/pylibcugraph/meta.yaml
index a8001344e1a..4ec3dd8a11d 100644
--- a/conda/recipes/pylibcugraph/meta.yaml
+++ b/conda/recipes/pylibcugraph/meta.yaml
@@ -7,6 +7,7 @@
 {% set cuda_version='.'.join(environ.get('CUDA', 'unknown').split('.')[:2]) %}
 {% set cuda_major=cuda_version.split('.')[0] %}
 {% set py_version=environ.get('CONDA_PY', 36) %}
+{% set ucx_py_version=environ.get('UCX_PY_VERSION') %}
 package:
   name: pylibcugraph
   version: {{ version }}
@@ -28,7 +29,7 @@ requirements:
     - python x.x
     - cython>=0.29,<0.30
     - libcugraph={{ version }}
-    - ucx-py 0.24
+    - ucx-py {{ ucx_py_version }}
     - ucx-proc=*=gpu
     - cudatoolkit {{ cuda_version }}.*
     - rmm {{ minor_version }}.*