Simplify wheel build scripts and allow alphas of RAPIDS dependencies (#…

…986)

This PR:

1. Removes `ci/apply_wheel_modifications.sh` and uses it inline in wheel build scripts
2. Allows for specifying alpha versioned dependencies of RAPIDS projects

Authors:
  - Divye Gala (https://github.com/divyegala)
  - Vyas Ramasubramani (https://github.com/vyasr)

Approvers:
  - Vyas Ramasubramani (https://github.com/vyasr)
  - AJ Schmidt (https://github.com/ajschmidt8)

URL: #986

ci/build_wheel.sh

@@ -3,6 +3,9 @@
 
 set -euo pipefail
 
+package_name="ucx-py"
+underscore_package_name=$(echo "${package_name}" | tr "-" "_")
+
 source rapids-configure-sccache
 source rapids-date-string
 
@@ -12,9 +15,31 @@ version_override="$(rapids-pip-wheel-version ${RAPIDS_DATE_STRING})"
 
 RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})"
 
-bash ci/release/apply_wheel_modifications.sh ${version_override} "-${RAPIDS_PY_CUDA_SUFFIX}"
-echo "The package name and/or version was modified in the package source. The git diff is:"
-git diff
+# This is the version of the suffix with a preceding hyphen. It's used
+# everywhere except in the final wheel name.
+PACKAGE_CUDA_SUFFIX="-${RAPIDS_PY_CUDA_SUFFIX}"
+
+# Patch project metadata files to include the CUDA version suffix and version override.
+pyproject_file="pyproject.toml"
+
+sed -i "s/^version = .*/version = \"${version_override}\"/g" ${pyproject_file}
+sed -i "s/name = \"${package_name}\"/name = \"${package_name}${PACKAGE_CUDA_SUFFIX}\"/g" ${pyproject_file}
+
+# For nightlies we want to ensure that we're pulling in alphas as well. The
+# easiest way to do so is to augment the spec with a constraint containing a
+# min alpha version that doesn't affect the version bounds but does allow usage
+# of alpha versions for that dependency without --pre
+alpha_spec=''
+if ! rapids-is-release-build; then
+    alpha_spec=',>=0.0.0a0'
+fi
+
+sed -r -i "s/cudf==(.*)\"/cudf${PACKAGE_CUDA_SUFFIX}==\1${alpha_spec}\"/g" ${pyproject_file}
+
+if [[ $PACKAGE_CUDA_SUFFIX == "-cu12" ]]; then
+    sed -i "s/cupy-cuda11x/cupy-cuda12x/g" ${pyproject_file}
+fi
+
 
 python -m pip wheel . -w dist -vvv --no-deps --disable-pip-version-check
 
@@ -27,7 +52,7 @@ python -m auditwheel repair -w final_dist dist/*
 # what these libraries are, we mimic the behaviour of auditwheel by using the
 # same hash-based uniqueness scheme and rewriting the link paths.
 
-WHL=$(realpath final_dist/ucx_py*manylinux*.whl)
+WHL=$(realpath final_dist/${underscore_package_name}*manylinux*.whl)
 
 # first grab the auditwheel hashes for libuc{tms}
 LIBUCM=$(unzip -l $WHL | awk 'match($4, /libucm-[^\.]+\./) { print substr($4, RSTART) }')
@@ -36,11 +61,11 @@ LIBUCS=$(unzip -l $WHL | awk 'match($4, /libucs-[^\.]+\./) { print substr($4, RS
 LIBNUMA=$(unzip -l $WHL | awk 'match($4, /libnuma-[^\.]+\./) { print substr($4, RSTART) }')
 
 # Extract the libraries that have already been patched in by auditwheel
-mkdir -p repair_dist/ucx_py_${RAPIDS_PY_CUDA_SUFFIX}.libs/ucx
-unzip $WHL "ucx_py_${RAPIDS_PY_CUDA_SUFFIX}.libs/*.so*" -d repair_dist/
+mkdir -p repair_dist/${underscore_package_name}_${RAPIDS_PY_CUDA_SUFFIX}.libs/ucx
+unzip $WHL "${underscore_package_name}_${RAPIDS_PY_CUDA_SUFFIX}.libs/*.so*" -d repair_dist/
 
 # Patch the RPATH to include ORIGIN for each library
-pushd repair_dist/ucx_py_${RAPIDS_PY_CUDA_SUFFIX}.libs
+pushd repair_dist/${underscore_package_name}_${RAPIDS_PY_CUDA_SUFFIX}.libs
 for f in libu*.so*
 do
     if [[ -f $f ]]; then
@@ -51,10 +76,10 @@ done
 popd
 
 # Now copy in all the extra libraries that are only ever loaded at runtime
-pushd repair_dist/ucx_py_${RAPIDS_PY_CUDA_SUFFIX}.libs/ucx
+pushd repair_dist/${underscore_package_name}_${RAPIDS_PY_CUDA_SUFFIX}.libs/ucx
 cp -P /usr/lib/ucx/* .
 
-# we link against <python>/lib/site-packages/ucx_py_${RAPIDS_PY_CUDA_SUFFIX}.lib/libuc{ptsm}
+# we link against <python>/lib/site-packages/${underscore_package_name}_${RAPIDS_PY_CUDA_SUFFIX}.lib/libuc{ptsm}
 # we also amend the rpath to search one directory above to *find* libuc{tsm}
 for f in libu*.so*
 do
@@ -94,7 +119,7 @@ patchelf --add-rpath '$ORIGIN' libuct_cuda.so
 popd
 
 pushd repair_dist
-zip -r $WHL ucx_py_${RAPIDS_PY_CUDA_SUFFIX}.libs/
+zip -r $WHL ${underscore_package_name}_${RAPIDS_PY_CUDA_SUFFIX}.libs/
 popd
 
-RAPIDS_PY_WHEEL_NAME="ucx_py_${RAPIDS_PY_CUDA_SUFFIX}" rapids-upload-wheels-to-s3 final_dist
+RAPIDS_PY_WHEEL_NAME="${underscore_package_name}_${RAPIDS_PY_CUDA_SUFFIX}" rapids-upload-wheels-to-s3 final_dist

conda/recipes/ucx-py/meta.yaml

@@ -5,6 +5,7 @@
 {% set version = environ.get('GIT_DESCRIBE_TAG', '0.0.0.dev').lstrip('v') %}
 {% set py_version = environ['CONDA_PY'] %}
 {% set cuda_version = '.'.join(environ['RAPIDS_CUDA_VERSION'].split('.')[:2]) %}
+{% set cuda_major = cuda_version.split('.')[0] %}
 {% set date_string = environ['RAPIDS_DATE_STRING'] %}
 
 package:
@@ -27,7 +28,13 @@ requirements:
   host:
     - python
     - pip
+    {% if cuda_major == "11" %}
     - ucx
+    {% else %}
+    # Temporary pin because the subsequent build (h64cca9d_3) is missing proper
+    # CUDA support
+    - ucx==1.14.1=h195a15c_3
+    {% endif %}
     {% for r in data.get("build-system", {}).get("requires", []) %}
     - {{ r }}
     {% endfor %}
@@ -36,6 +43,14 @@ requirements:
     {% for r in data.get("project", {}).get("dependencies", []) %}
     - {{ r }}
     {% endfor %}
+    # Something seems wrong with the run export as well now
+    {% if cuda_major == "11" %}
+    - ucx
+    {% else %}
+    # Temporary pin because the subsequent build (h64cca9d_3) is missing proper
+    # CUDA support
+    - ucx==1.14.1=h195a15c_3
+    {% endif %}
 
 test:
   imports:

dependencies.yaml

@@ -108,9 +108,19 @@ dependencies:
         packages:
           - numpy>=1.21
           - pynvml>=11.4.1
-      - output_types: [conda]
-        packages:
-          - ucx
+    specific:
+      - output_types: conda
+        matrices:
+          - matrix:
+              cuda: "12.0"
+            packages:
+              # Temporary pin because the subsequent build (h64cca9d_3) is
+              # missing proper CUDA support
+              - ucx==1.14.1=h195a15c_3
+          - matrix:
+              cuda: "11.8"
+            packages:
+              - ucx
   test_python:
     common:
       - output_types: [conda, requirements, pyproject]