use libucx wheels

.gitignore

@@ -10,8 +10,9 @@ dask-worker-space
 __pytestcache__
 __pycache__
 *.egg-info/
+final_dist/
 dist/
 .vscode
 
 *.sw[po]
-
+*.whl

.readthedocs.yml

@@ -4,11 +4,9 @@ build:
   os: "ubuntu-22.04"
   tools:
     python: "mambaforge-22.9"
-
-python:
-  install:
-    - method: pip
-      path: .
+  jobs:
+    post_checkout:
+      - bash ci/build_docs_pre_install.sh
 
 conda:
   environment: conda/environments/builddocs.yml

ci/build_docs_pre_install.sh

@@ -0,0 +1,17 @@
+#!/bin/bash
+# Copyright (c) 2024 NVIDIA CORPORATION.
+#
+# [description]
+#
+#   ucx-py's docs builds require installing the library.
+#
+#   It does that by running 'pip install .' from the root of the repo. This script
+#   is used to modify readthedocs' local checkout of this project's source code prior
+#   to that 'pip install' being run.
+#
+#   For more, see https://docs.readthedocs.io/en/stable/build-customization.html
+#
+
+set -euo pipefail
+
+sed -r -i "s/\"libucx/\"libucx-cu12/g" ./pyproject.toml

ci/build_wheel.sh

@@ -6,6 +6,19 @@ set -euo pipefail
 package_name="ucx-py"
 underscore_package_name=$(echo "${package_name}" | tr "-" "_")
 
+# Clear out system ucx files to ensure that we're getting ucx from the wheel.
+rm -rf /usr/lib64/ucx
+rm -rf /usr/lib64/libucm.*
+rm -rf /usr/lib64/libucp.*
+rm -rf /usr/lib64/libucs.*
+rm -rf /usr/lib64/libucs_signal.*
+rm -rf /usr/lib64/libuct.*
+
+rm -rf /usr/include/ucm
+rm -rf /usr/include/ucp
+rm -rf /usr/include/ucs
+rm -rf /usr/include/uct
+
 source rapids-configure-sccache
 source rapids-date-string
 
@@ -35,96 +48,22 @@ if ! rapids-is-release-build; then
 fi
 
 sed -r -i "s/cudf==(.*)\"/cudf${PACKAGE_CUDA_SUFFIX}==\1${alpha_spec}\"/g" ${pyproject_file}
+sed -r -i "/\"libucx([=><]+)/ s/\"libucx/\"libucx${PACKAGE_CUDA_SUFFIX}/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
 
 mkdir -p final_dist
-python -m auditwheel repair -w final_dist dist/*
-
-# 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.
-
-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) }')
-LIBUCT=$(unzip -l $WHL | awk 'match($4, /libuct-[^\.]+\./) { print substr($4, RSTART) }')
-LIBUCS=$(unzip -l $WHL | awk 'match($4, /libucs-[^\.]+\./) { print substr($4, RSTART) }')
-
-# Extract the libraries that have already been patched in by auditwheel
-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/${underscore_package_name}_${RAPIDS_PY_CUDA_SUFFIX}.libs
-for f in libu*.so*
-do
-    if [[ -f $f ]]; then
-        patchelf --add-rpath '$ORIGIN' $f
-    fi
-done
-
-popd
-
-# Now copy in all the extra libraries that are only ever loaded at runtime
-pushd repair_dist/${underscore_package_name}_${RAPIDS_PY_CUDA_SUFFIX}.libs/ucx
-if [[ -d /usr/lib64/ucx ]]; then
-    cp -P /usr/lib64/ucx/* .
-elif [[ -d /usr/lib/ucx ]]; then
-    cp -P /usr/lib/ucx/* .
-else
-    echo "Could not find ucx libraries"
-    exit 1
-fi
-
-# 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
-    # Avoid patching symlinks, which is redundant
-    if [[ ! -L $f ]]; then
-        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 --add-rpath '$ORIGIN/..' $f
-    fi
-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
-
-popd
-
-pushd repair_dist
-zip -r $WHL ${underscore_package_name}_${RAPIDS_PY_CUDA_SUFFIX}.libs/
-popd
+python -m auditwheel repair \
+    -w final_dist \
+    --exclude "libucm.so.0" \
+    --exclude "libucp.so.0" \
+    --exclude "libucs.so.0" \
+    --exclude "libucs_signal.so.0" \
+    --exclude "libuct.so.0" \
+    dist/*
 
 RAPIDS_PY_WHEEL_NAME="${underscore_package_name}_${RAPIDS_PY_CUDA_SUFFIX}" rapids-upload-wheels-to-s3 final_dist

conda/environments/builddocs.yml

@@ -14,5 +14,7 @@ dependencies:
 - recommonmark
 - pandoc=<2.0.0
 - pip
-- ucx
+- pip:
+  - --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple
+  - ../../
 - cython

conda/recipes/ucx-py/meta.yaml

@@ -30,13 +30,15 @@ requirements:
     - python
     - pip
     - ucx
-    {% for r in data.get("build-system", {}).get("requires", []) %}
+    # 'libucx' wheel dependency is unnecessary... the 'ucx' conda-forge package is used here instead
+    {% for r in data.get("build-system", {}).get("requires", []) if not r.startswith("libucx") %}
     - {{ r }}
     {% endfor %}
   run:
     - python
     - ucx >=1.15.0,<1.16.0
-    {% for r in data.get("project", {}).get("dependencies", []) %}
+    # 'libucx' wheel dependency is unnecessary... the 'ucx' conda-forge package is used here instead
+    {% for r in data.get("project", {}).get("dependencies", []) if not r.startswith("libucx") %}
     - {{ r }}
     {% endfor %}
 

dependencies.yaml

@@ -27,13 +27,15 @@ files:
       table: build-system
     includes:
       - build_python
+      - depends_on_ucx_build
   py_run:
     output: pyproject
     pyproject_dir: .
     extras:
       table: project
     includes:
       - run
+      - depends_on_ucx_run
   py_optional_test:
     output: pyproject
     pyproject_dir: .
@@ -115,9 +117,54 @@ dependencies:
         packages:
           - numpy>=1.23,<2.0a0
           - pynvml>=11.4.1
+  depends_on_ucx_build:
+    common:
+      - output_types: conda
+        packages:
+          - ucx==1.15.0
+      - output_types: requirements
+        packages:
+          # pip recognizes the index as a global option for the requirements.txt file
+          - --extra-index-url=https://pypi.nvidia.com
+          - --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple
+    specific:
+      - output_types: [requirements, pyproject]
+        matrices:
+          - matrix: {cuda: "12.*"}
+            packages:
+              - libucx-cu12==1.15.0
+          - matrix: {cuda: "11.*"}
+            packages:
+              - libucx-cu11==1.15.0
+          # NOTE: this fallback needs to be a real, suffixed version
+          #       so 'pip install .' (e.g. as used in docs builds) will work
+          - matrix: null
+            packages:
+              - libucx==1.15.0
+  depends_on_ucx_run:
+    common:
       - output_types: conda
         packages:
           - ucx>=1.15.0,<1.16
+      - output_types: requirements
+        packages:
+          # pip recognizes the index as a global option for the requirements.txt file
+          - --extra-index-url=https://pypi.nvidia.com
+          - --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple
+    specific:
+      - output_types: [requirements, pyproject]
+        matrices:
+          - matrix: {cuda: "12.*"}
+            packages:
+              - libucx-cu12>=1.15.0,<1.16
+          - matrix: {cuda: "11.*"}
+            packages:
+              - libucx-cu11>=1.15.0,<1.16
+          # NOTE: this fallback needs to be a real, suffixed version
+          #       so "pip install ." (e.g. as used in docs builds) will work
+          - matrix: null
+            packages:
+              - libucx>=1.15.0,<1.16
   test_python:
     common:
       - output_types: [conda, requirements, pyproject]

pyproject.toml

@@ -16,6 +16,7 @@
 build-backend = "setuptools.build_meta"
 requires = [
     "cython>=3.0.0",
+    "libucx==1.15.0",
     "setuptools>=64.0.0",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit dependencies.yaml and run `rapids-dependency-file-generator`.
 
@@ -30,6 +31,7 @@ authors = [
 license = { text = "BSD-3-Clause" }
 requires-python = ">=3.9"
 dependencies = [
+    "libucx>=1.15.0,<1.16",
     "numpy>=1.23,<2.0a0",
     "pynvml>=11.4.1",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit dependencies.yaml and run `rapids-dependency-file-generator`.

setup.py

@@ -5,18 +5,67 @@
 
 from __future__ import absolute_import, print_function
 
+import glob
 import os
 from distutils.sysconfig import get_config_var, get_python_inc
 
 from Cython.Distutils.build_ext import new_build_ext
 from setuptools import setup
 from setuptools.extension import Extension
 
+
+def _find_libucx_libs_and_headers():
+    """
+    If the 'libucx' wheel is not installed, returns a tuple of empty lists.
+    In that case, the project will be compiled against system installations
+    of the UCX libraries.
+
+    If 'libucx' is installed, returns lists of library and header paths to help
+    the compiler and linker find its contents. In that case, the project will
+    be compiled against those libucx-wheel-provided versions of the UCX libraries.
+    """
+    try:
+        import libucx
+    except ImportError:
+        return [], []
+
+    # find 'libucx'
+    module_dir = os.path.dirname(libucx.__file__)
+
+    # find where it stores files like 'libucm.so.0'
+    libs = glob.glob(f"{module_dir}/**/lib*.so*", recursive=True)
+
+    # deduplicate those library paths
+    lib_dirs = {os.path.dirname(f) for f in libs}
+    if not lib_dirs:
+        raise RuntimeError(
+            f"Did not find shared libraries in 'libucx' install location ({module_dir})"
+        )
+
+    # find where it stores headers
+    headers = glob.glob(f"{module_dir}/**/include", recursive=True)
+
+    # deduplicate those header paths (and ensure the list only includes directories)
+    header_dirs = {f for f in headers if os.path.isdir(f)}
+    if not header_dirs:
+        raise RuntimeError(
+            f"Did not find UCX headers 'libucx' install location ({module_dir})"
+        )
+
+    return list(lib_dirs), list(header_dirs)
+
+
 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"]
 
+# tell the compiler and linker where to find UCX libraries and their headers
+# provided by the 'libucx' wheel
+libucx_lib_dirs, libucx_header_dirs = _find_libucx_libs_and_headers()
+library_dirs.extend(libucx_lib_dirs)
+include_dirs.extend(libucx_header_dirs)
+
 
 ext_modules = [
     Extension(

ucp/__init__.py

@@ -16,6 +16,18 @@
     logger.debug("Setting env UCX_MEMTYPE_CACHE=n, which is required by UCX")
     os.environ["UCX_MEMTYPE_CACHE"] = "n"
 
+
+# If libucx was installed as a wheel, we must request it to load the library symbols.
+# Otherwise, we assume that the library was installed in a system path that ld can find.
+try:
+    import libucx
+except ImportError:
+    pass
+else:
+    libucx.load_library()
+    del libucx
+
+
 from .core import *  # noqa
 from .core import get_ucx_version  # noqa
 from .utils import get_ucxpy_logger  # noqa