Add wheel builds (#901)

.github/build-scripts/patch-wheel.sh

@@ -0,0 +1,65 @@
+#!/usr/bin/env bash
+# Auditwheel rewrites dynamic libraries that are referenced at link time in the
+# package. However, UCX loads a number of sub-libraries at runtime via dlopen;
+# these are not picked up by auditwheel. Since we have a priori knowledge of
+# what these libraries are, we mimic the behaviour of auditwheel by using the
+# same hash-based uniqueness scheme and rewriting the link paths.
+
+set -Eeoxu pipefail
+
+mypyver=$(python --version)
+repair_dir="repair-${mypyver// /_}"
+
+mkdir -p "${repair_dir}" && cd "${repair_dir}"
+
+WHL=$1
+
+# first grab the auditwheel hashes for libuc{tms}
+LIBUCM=$(unzip -l $WHL | awk 'match($4, /libucm-[^\.]+\./) { print substr($4, RSTART) }')
+LIBUCT=$(unzip -l $WHL | awk 'match($4, /libuct-[^\.]+\./) { print substr($4, RSTART) }')
+LIBUCS=$(unzip -l $WHL | awk 'match($4, /libucs-[^\.]+\./) { print substr($4, RSTART) }')
+LIBNUMA=$(unzip -l $WHL | awk 'match($4, /libnuma-[^\.]+\./) { print substr($4, RSTART) }')
+
+# TODO: This directory is currently hardcoded, but it actually needs to take
+# another script argument to get the CUDA suffix used for the current build.
+mkdir -p ucx_py_cu11.libs/ucx
+cd ucx_py_cu11.libs/ucx
+cp -P /usr/lib/ucx/* .
+
+# we link against <python>/lib/site-packages/ucx_py_cu11.lib/libuc{ptsm}
+# we also amend the rpath to search one directory above to *find* libuc{tsm}
+for f in libu*.so*
+do
+  patchelf --replace-needed libuct.so.0 $LIBUCT $f
+  patchelf --replace-needed libucs.so.0 $LIBUCS $f
+  patchelf --replace-needed libucm.so.0 $LIBUCM $f
+  patchelf --replace-needed libnuma.so.1 $LIBNUMA $f
+  patchelf --add-rpath '$ORIGIN/..' $f
+done
+
+# Bring in cudart as well. To avoid symbol collision with other libraries e.g.
+# cupy we mimic auditwheel by renaming the libraries to include the hashes of
+# their names. Since there will typically be a chain of symlinks
+# libcudart.so->libcudart.so.X->libcudart.so.X.Y.Z we need to follow the chain
+# and rename all of them.
+
+find /usr/local/cuda/ -name "libcudart*.so*" | xargs cp -P -t .
+src=libcudart.so
+hash=$(sha256sum ${src} | awk '{print substr($1, 0, 8)}')
+target=$(basename $(readlink -f ${src}))
+
+mv ${target} ${target/libcudart/libcudart-${hash}}
+while readlink ${src} > /dev/null; do
+    target=$(readlink ${src})
+    ln -s ${target/libcudart/libcudart-${hash}} ${src/libcudart/libcudart-${hash}}
+    rm -f ${src}
+    src=${target}
+done
+
+to_rewrite=$(ldd libuct_cuda.so | awk '/libcudart/ { print $1 }')
+patchelf --replace-needed ${to_rewrite} libcudart-${hash}.so libuct_cuda.so
+patchelf --add-rpath '$ORIGIN' libuct_cuda.so
+
+cd -
+
+zip -r $WHL ucx_py_cu11.libs/

.github/workflows/wheels.yml

@@ -0,0 +1,60 @@
+name: ucx-py wheels
+
+on:
+  workflow_call:
+    inputs:
+      versioneer-override:
+        type: string
+        default: ''
+      build-tag:
+        type: string
+        default: ''
+      branch:
+        required: true
+        type: string
+      date:
+        required: true
+        type: string
+      sha:
+        required: true
+        type: string
+      build-type:
+        type: string
+        default: nightly
+
+concurrency:
+  group: "ucx-py-${{ github.workflow }}-${{ github.ref }}"
+  cancel-in-progress: true
+
+jobs:
+  ucx-py-wheels:
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-manylinux.yml@main
+    with:
+      repo: rapidsai/ucx-py
+
+      build-type: ${{ inputs.build-type }}
+      branch: ${{ inputs.branch }}
+      sha: ${{ inputs.sha }}
+      date: ${{ inputs.date }}
+
+      package-dir: .
+      package-name: ucx_py_cu11
+
+      python-package-versioneer-override: ${{ inputs.versioneer-override }}
+      python-package-build-tag: ${{ inputs.build-tag }}
+
+      # copy custom auditwheel script to known location with '{package}' placeholder
+      # in the auditwheel repair command, the '{package}' placeholder can't be used
+      # Also rewrite the name in the pyproject.toml file.
+      # TODO: Eventually this needs to be done more generically.
+      cibw-before-build: "cp {package}/.github/build-scripts/patch-wheel.sh /patch-wheel.sh && sed -i 's/name = \"ucx-py\"/name = \"ucx-py-cu11\"/g'  {package}/pyproject.toml"
+
+      auditwheel-repair-command-amd64: "auditwheel --verbose repair -w {dest_dir} --plat manylinux_2_17_x86_64 {wheel} && /patch-wheel.sh {dest_dir}/*.whl"
+      auditwheel-repair-command-arm64: "auditwheel --verbose repair -w {dest_dir} --plat manylinux_2_31_aarch64 {wheel} && /patch-wheel.sh {dest_dir}/*.whl"
+
+      test-docker-options: "--cap-add CAP_SYS_PTRACE --shm-size=8g"
+
+      test-before-arm64: "pip install cupy-cuda11x -f https://pip.cupy.dev/aarch64"
+      test-extras: test
+      test-unittest: "pytest -k 'not test_send_recv_am' --cache-clear -vs ./ucp/_libs/tests/ && pytest --cache-clear -vs ./tests"
+    secrets: inherit

pyproject.toml

@@ -1,4 +1,4 @@
-# Copyright (c) 2022, NVIDIA CORPORATION.
+# Copyright (c) 2021-2022, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -53,6 +53,10 @@ classifiers = [
 test = [
     "pytest",
     "pytest-asyncio",
+    "cloudpickle",
+    "dask",
+    "distributed",
+    "cupy-cuda11x",
 ]
 
 [project.urls]

setup.py

@@ -6,30 +6,37 @@
 from __future__ import absolute_import, print_function
 
 import os
-import re
 from distutils.sysconfig import get_config_var, get_python_inc
 
 import versioneer
 from Cython.Distutils.build_ext import new_build_ext as build_ext
 from setuptools import setup
 from setuptools.extension import Extension
 
-include_dirs = [os.path.dirname(get_python_inc())]
-library_dirs = [get_config_var("LIBDIR")]
-libraries = ["ucp", "uct", "ucm", "ucs"]
-extra_compile_args = ["-std=c99", "-Werror"]
+# Patch versioneer version for wheel builds.
+if "RAPIDS_PY_WHEEL_VERSIONEER_OVERRIDE" in os.environ:
+    orig_get_versions = versioneer.get_versions
+
+    version_override = os.environ["RAPIDS_PY_WHEEL_VERSIONEER_OVERRIDE"]
+    if not version_override:
+        raise RuntimeError(
+            "An empty RAPIDS_PY_WHEEL_VERSIONEER_OVERRIDE is not supported. "
+            "Either remove this variable from your environment or specify a "
+            "valid version override."
+        )
 
+    def get_versions():
+        data = orig_get_versions()
+        data["version"] = version_override
+        return data
 
-def get_ucp_version():
-    with open(include_dirs[0] + "/ucp/api/ucp_version.h") as f:
-        ftext = f.read()
-        major = re.findall("^#define.*UCP_API_MAJOR.*", ftext, re.MULTILINE)
-        minor = re.findall("^#define.*UCP_API_MINOR.*", ftext, re.MULTILINE)
+    versioneer.get_versions = get_versions
 
-        major = int(major[0].split()[-1])
-        minor = int(minor[0].split()[-1])
 
-        return (major, minor)
+include_dirs = [os.path.dirname(get_python_inc())]
+library_dirs = [get_config_var("LIBDIR")]
+libraries = ["ucp", "uct", "ucm", "ucs"]
+extra_compile_args = ["-std=c99", "-Werror"]
 
 
 ext_modules = [
@@ -56,6 +63,10 @@ def get_ucp_version():
 cmdclass = versioneer.get_cmdclass(cmdclass)
 
 setup(
+    # TODO: At present the ucx-py naming scheme is not dynamic and will not
+    # support different CUDA major versions if we need to build wheels. It is
+    # hardcoded in pyproject.toml, so overriding it here will not work.
+    # name="ucx-py" + os.getenv("RAPIDS_PY_WHEEL_CUDA_SUFFIX", default=""),
     ext_modules=ext_modules,
     cmdclass=cmdclass,
     version=versioneer.get_version(),