From dc4e8f0960733f84623185a04ae2fbe83b4fb111 Mon Sep 17 00:00:00 2001
From: Sevag H <shanssian@nvidia.com>
Date: Tue, 31 Jan 2023 11:50:23 -0500
Subject: [PATCH] Pin wheel dependencies to same RAPIDS release (#1200)

Now that wheel CI is enabled, the nightly index will contain next release wheels, so we need to ensure those don't get installed while still building this release.

Also tests a new mechanism (`local-wheelhouse`) for cibuildwheel to prefer local wheels.

Authors:
  - Sevag H (https://github.com/sevagh)
  - Vyas Ramasubramani (https://github.com/vyasr)

Approvers:
  - AJ Schmidt (https://github.com/ajschmidt8)
  - Corey J. Nolet (https://github.com/cjnolet)

URL: https://github.com/rapidsai/raft/pull/1200
---
 .github/workflows/pr.yaml                 |  3 +--
 ci/release/update-version.sh              | 15 ++++++++++++++-
 python/pylibraft/_custom_build/backend.py |  7 +++----
 python/pylibraft/setup.py                 |  4 ++--
 python/raft-dask/setup.py                 |  4 ++--
 5 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml
index 61d976c605..b705557795 100644
--- a/.github/workflows/pr.yaml
+++ b/.github/workflows/pr.yaml
@@ -80,8 +80,7 @@ jobs:
       build_type: pull-request
       package-name: raft_dask
       package-dir: python/raft-dask
-      before-wheel: "RAPIDS_PY_WHEEL_NAME=pylibraft_cu11 rapids-download-wheels-from-s3 ./local-pylibraft-dep"
-      cibw-before-build: "pip install --no-deps ./local-pylibraft-dep/pylibraft*.whl"
+      before-wheel: "RAPIDS_PY_WHEEL_NAME=pylibraft_cu11 rapids-download-wheels-from-s3 ./local-wheelhouse"
       skbuild-configure-options: "-DRAFT_BUILD_WHEELS=ON -DDETECT_CONDA_ENV=OFF -DFIND_RAFT_CPP=OFF"
   wheel-tests-raft-dask:
     needs: wheel-build-raft-dask
diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh
index d1f849194e..00f6905032 100755
--- a/ci/release/update-version.sh
+++ b/ci/release/update-version.sh
@@ -17,12 +17,14 @@ CURRENT_MAJOR=$(echo $CURRENT_TAG | awk '{split($0, a, "."); print a[1]}')
 CURRENT_MINOR=$(echo $CURRENT_TAG | awk '{split($0, a, "."); print a[2]}')
 CURRENT_PATCH=$(echo $CURRENT_TAG | awk '{split($0, a, "."); print a[3]}')
 CURRENT_SHORT_TAG=${CURRENT_MAJOR}.${CURRENT_MINOR}
+CURRENT_UCX_PY_VERSION="$(curl -sL https://version.gpuci.io/rapids/${CURRENT_SHORT_TAG}).*"
 
 #Get <major>.<minor> for next version
 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}).*"
+NEXT_UCX_PY_SHORT_TAG="$(curl -sL https://version.gpuci.io/rapids/${NEXT_SHORT_TAG})"
+NEXT_UCX_PY_VERSION="${NEXT_UCX_PY_SHORT_TAG}.*"
 
 echo "Preparing release $CURRENT_TAG => $NEXT_FULL_TAG"
 
@@ -56,3 +58,14 @@ sed_runner "/^ucx_py_version:$/ {n;s/.*/  - \"${NEXT_UCX_PY_VERSION}\"/}" conda/
 
 # Wheel builds install dask-cuda from source, update its branch
 sed_runner "s/dask-cuda.git@branch-[^\"\s]\+/dask-cuda.git@branch-${NEXT_SHORT_TAG}/g" .github/workflows/*.yaml
+
+# Need to distutils-normalize the original version
+NEXT_SHORT_TAG_PEP440=$(python -c "from setuptools.extern import packaging; print(packaging.version.Version('${NEXT_SHORT_TAG}'))")
+NEXT_UCX_PY_SHORT_TAG_PEP440=$(python -c "from setuptools.extern import packaging; print(packaging.version.Version('${NEXT_UCX_PY_SHORT_TAG}'))")
+
+# Wheel builds install intra-RAPIDS dependencies from same release
+sed_runner "s/{cuda_suffix}[^\"].*\",/{cuda_suffix}==${NEXT_SHORT_TAG_PEP440}.*\",/g" python/pylibraft/setup.py
+sed_runner "s/{cuda_suffix}.*\"\]/{cuda_suffix}==${NEXT_SHORT_TAG_PEP440}.*\"\]/g" python/pylibraft/_custom_build/backend.py
+sed_runner "s/dask-cuda==.*\",/dask-cuda==${NEXT_SHORT_TAG_PEP440}.*\",/g" python/raft-dask/setup.py
+sed_runner "s/pylibraft{cuda_suffix}.*\",/pylibraft{cuda_suffix}==${NEXT_SHORT_TAG_PEP440}.*\",/g" python/raft-dask/setup.py
+sed_runner "s/ucx-py{cuda_suffix}.*\",/ucx-py{cuda_suffix}==${NEXT_UCX_PY_SHORT_TAG_PEP440}.*\",/g" python/raft-dask/setup.py
diff --git a/python/pylibraft/_custom_build/backend.py b/python/pylibraft/_custom_build/backend.py
index 7d1b334626..209e9e4b67 100644
--- a/python/pylibraft/_custom_build/backend.py
+++ b/python/pylibraft/_custom_build/backend.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2022, NVIDIA CORPORATION.
+# Copyright (c) 2022-2023, NVIDIA CORPORATION.
 
 """Custom build backend for pylibraft to get versioned requirements.
 
@@ -18,9 +18,8 @@ def replace_requirements(func):
     @wraps(func)
     def wrapper(config_settings=None):
         orig_list = getattr(_orig, func.__name__)(config_settings)
-        append_list = [
-            f"rmm{os.getenv('RAPIDS_PY_WHEEL_CUDA_SUFFIX', default='')}"
-        ]
+        cuda_suffix = os.getenv("RAPIDS_PY_WHEEL_CUDA_SUFFIX", default="")
+        append_list = [f"rmm{cuda_suffix}==23.2.*"]
         return orig_list + append_list
 
     return wrapper
diff --git a/python/pylibraft/setup.py b/python/pylibraft/setup.py
index 079825a30c..230b9127e3 100644
--- a/python/pylibraft/setup.py
+++ b/python/pylibraft/setup.py
@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2022, NVIDIA CORPORATION.
+# Copyright (c) 2022-2023, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -25,7 +25,7 @@
 install_requires = [
     "numpy",
     "cuda-python>=11.7.1,<12.0",
-    f"rmm{cuda_suffix}",
+    f"rmm{cuda_suffix}==23.2.*",
 ]
 
 extras_require = {
diff --git a/python/raft-dask/setup.py b/python/raft-dask/setup.py
index 0c65841341..3299ea9bc5 100644
--- a/python/raft-dask/setup.py
+++ b/python/raft-dask/setup.py
@@ -28,9 +28,9 @@
     "joblib>=0.11",
     "dask-cuda==23.2.*",
     "dask>=2022.12.0",
-    f"ucx-py{cuda_suffix}",
+    f"ucx-py{cuda_suffix}==0.30.*",
     "distributed>=2022.12.0",
-    f"pylibraft{cuda_suffix}",
+    f"pylibraft{cuda_suffix}==23.2.*",
 ]
 
 extras_require = {