diff --git a/.github/build-scripts/patch-wheel.sh b/.github/build-scripts/patch-wheel.sh new file mode 100755 index 000000000..da36177b0 --- /dev/null +++ b/.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 /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/ diff --git a/.github/workflows/wheels.yml b/.github/workflows/wheels.yml new file mode 100644 index 000000000..d9e26fc30 --- /dev/null +++ b/.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 diff --git a/pyproject.toml b/pyproject.toml index c6fc063cf..1930a4dc1 100644 --- a/pyproject.toml +++ b/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] diff --git a/setup.py b/setup.py index 77f6800cb..0ffc1bbcf 100644 --- a/setup.py +++ b/setup.py @@ -6,7 +6,6 @@ from __future__ import absolute_import, print_function import os -import re from distutils.sysconfig import get_config_var, get_python_inc import versioneer @@ -14,22 +13,30 @@ 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(),