From 1b7399a4bed50646fd400b34decc9f27f570ba70 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Mon, 6 May 2024 07:53:32 -0700
Subject: [PATCH 01/35] use libucx wheels

---
 .gitignore        |  3 +-
 ci/build_wheel.sh | 95 +++++------------------------------------------
 dependencies.yaml | 20 ++++++++++
 pyproject.toml    |  1 +
 ucp/__init__.py   | 11 ++++++
 5 files changed, 44 insertions(+), 86 deletions(-)

diff --git a/.gitignore b/.gitignore
index a155793c4..b0898274a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -10,8 +10,9 @@ dask-worker-space
 __pytestcache__
 __pycache__
 *.egg-info/
+final_dist/
 dist/
 .vscode
 
 *.sw[po]
-
+*.whl
diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh
index 9cc871e9a..ff73095d0 100755
--- a/ci/build_wheel.sh
+++ b/ci/build_wheel.sh
@@ -40,91 +40,16 @@ 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
-
-RAPIDS_PY_WHEEL_NAME="${underscore_package_name}_${RAPIDS_PY_CUDA_SUFFIX}" rapids-upload-wheels-to-s3 final_dist
+python -m auditwheel repair \
+    -w final_dist \
+    --exclude "libucm.so.0" \
+    --exclude "libucp.so.0" \
+    --exclude "libucx.so.0" \
+    --exclude "libucs.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
diff --git a/dependencies.yaml b/dependencies.yaml
index 1359a919c..b49eac0ec 100644
--- a/dependencies.yaml
+++ b/dependencies.yaml
@@ -34,6 +34,7 @@ files:
       table: project
     includes:
       - run
+      - depends_on_ucx_run
   py_optional_test:
     output: pyproject
     pyproject_dir: .
@@ -116,9 +117,28 @@ dependencies:
         packages:
           - numpy>=1.23,<2.0a0
           - pynvml>=11.4.1
+  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
+          - matrix: null
+            packages:
+              - libucx>=1.15.0,<1.16
   test_python:
     common:
       - output_types: [conda, requirements, pyproject]
diff --git a/pyproject.toml b/pyproject.toml
index 16a0dd89a..a8d760266 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -31,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`.
diff --git a/ucp/__init__.py b/ucp/__init__.py
index 791860b1d..b2bf80ba3 100644
--- a/ucp/__init__.py
+++ b/ucp/__init__.py
@@ -27,6 +27,17 @@
 except ImportError:
     pynvml = None
 
+# 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 ModuleNotFoundError:
+    pass
+else:
+    libucx.load_library()
+    del libucx
+
+
 # Setup UCX-Py logger
 logger = get_ucxpy_logger()
 

From cbb283ae300a730dab3253659d4f455f4b70ed69 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Mon, 6 May 2024 07:57:47 -0700
Subject: [PATCH 02/35] keep libucp

---
 ci/build_wheel.sh | 1 -
 1 file changed, 1 deletion(-)

diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh
index ff73095d0..f645c6d5b 100755
--- a/ci/build_wheel.sh
+++ b/ci/build_wheel.sh
@@ -46,7 +46,6 @@ mkdir -p final_dist
 python -m auditwheel repair \
     -w final_dist \
     --exclude "libucm.so.0" \
-    --exclude "libucp.so.0" \
     --exclude "libucx.so.0" \
     --exclude "libucs.so.0" \
     --exclude "libuct.so.0" \

From 875d5825d5c25fe5f323bd0581bf7fbecad9c3d7 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Mon, 6 May 2024 08:25:56 -0700
Subject: [PATCH 03/35] get all dependencies via conda

---
 conda/recipes/ucx-py/meta.yaml | 2 +-
 pyproject.toml                 | 1 -
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/conda/recipes/ucx-py/meta.yaml b/conda/recipes/ucx-py/meta.yaml
index f0a81f7f1..45689fa62 100644
--- a/conda/recipes/ucx-py/meta.yaml
+++ b/conda/recipes/ucx-py/meta.yaml
@@ -18,7 +18,7 @@ build:
   number: {{ GIT_DESCRIBE_NUMBER }}
   string: py{{ py_version }}_{{ date_string }}_{{ GIT_DESCRIBE_HASH }}_{{ GIT_DESCRIBE_NUMBER }}
   script:
-    - {{ PYTHON }} -m pip install . -vv
+    - {{ PYTHON }} -m pip install --no-build-isolation --no-deps . -vv
   ignore_run_exports_from:
     - ucx
 
diff --git a/pyproject.toml b/pyproject.toml
index a8d760266..16a0dd89a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -31,7 +31,6 @@ 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`.

From 14d023367ab0e5ae19042e874f302006639120e3 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Mon, 6 May 2024 08:38:47 -0700
Subject: [PATCH 04/35] ignore libucx dependency for conda packages

---
 conda/recipes/ucx-py/meta.yaml | 3 ++-
 pyproject.toml                 | 1 +
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/conda/recipes/ucx-py/meta.yaml b/conda/recipes/ucx-py/meta.yaml
index 45689fa62..f73504578 100644
--- a/conda/recipes/ucx-py/meta.yaml
+++ b/conda/recipes/ucx-py/meta.yaml
@@ -36,7 +36,8 @@ requirements:
   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 %}
 
diff --git a/pyproject.toml b/pyproject.toml
index 16a0dd89a..a8d760266 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -31,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`.

From 2f030df81d6e293ee4227a37780b9ed0ae629fc0 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Mon, 6 May 2024 10:48:42 -0700
Subject: [PATCH 05/35] upload artifacts

---
 ci/build_wheel.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh
index f645c6d5b..0ccd17c35 100755
--- a/ci/build_wheel.sh
+++ b/ci/build_wheel.sh
@@ -51,4 +51,4 @@ python -m auditwheel repair \
     --exclude "libuct.so.0" \
     dist/*
 
-#RAPIDS_PY_WHEEL_NAME="${underscore_package_name}_${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

From 5541ee16f61e4de5556921dd66f17b9018c5336c Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Mon, 6 May 2024 11:29:24 -0700
Subject: [PATCH 06/35] replace suffix for libucx

---
 ci/build_wheel.sh | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh
index 0ccd17c35..076b8a0eb 100755
--- a/ci/build_wheel.sh
+++ b/ci/build_wheel.sh
@@ -35,6 +35,7 @@ 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 "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}
@@ -51,4 +52,4 @@ python -m auditwheel repair \
     --exclude "libuct.so.0" \
     dist/*
 
-RAPIDS_PY_WHEEL_NAME="${underscore_package_name}_${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

From dc74edd411a886321b6ac630fc82526dc529a59e Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Mon, 6 May 2024 11:38:25 -0700
Subject: [PATCH 07/35] uncomment uploads again

---
 ci/build_wheel.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh
index 076b8a0eb..df43801df 100755
--- a/ci/build_wheel.sh
+++ b/ci/build_wheel.sh
@@ -52,4 +52,4 @@ python -m auditwheel repair \
     --exclude "libuct.so.0" \
     dist/*
 
-#RAPIDS_PY_WHEEL_NAME="${underscore_package_name}_${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

From 2b4b6e3c3ce0697b28411e14d44f9d2bd4513671 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Mon, 6 May 2024 11:58:31 -0700
Subject: [PATCH 08/35] maybe installation of the library itself can be skipped
 on readthedocs

---
 .readthedocs.yml | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/.readthedocs.yml b/.readthedocs.yml
index 937f721db..d72bb73ad 100644
--- a/.readthedocs.yml
+++ b/.readthedocs.yml
@@ -5,11 +5,6 @@ build:
   tools:
     python: "mambaforge-22.9"
 
-python:
-  install:
-    - method: pip
-      path: .
-
 conda:
   environment: conda/environments/builddocs.yml
 

From f7098fefaa32c0071c0a50483967eb76be79687f Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Mon, 6 May 2024 12:35:00 -0700
Subject: [PATCH 09/35] dlopen() earlier

---
 ucp/__init__.py | 21 +++++++++++----------
 1 file changed, 11 insertions(+), 10 deletions(-)

diff --git a/ucp/__init__.py b/ucp/__init__.py
index b2bf80ba3..7a51cecfb 100644
--- a/ucp/__init__.py
+++ b/ucp/__init__.py
@@ -16,16 +16,6 @@
     logger.debug("Setting env UCX_MEMTYPE_CACHE=n, which is required by UCX")
     os.environ["UCX_MEMTYPE_CACHE"] = "n"
 
-from .core import *  # noqa
-from .core import get_ucx_version  # noqa
-from .utils import get_ucxpy_logger  # noqa
-from ._libs.utils import get_address  # noqa
-from ._version import __git_commit__, __version__
-
-try:
-    import pynvml
-except ImportError:
-    pynvml = None
 
 # 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.
@@ -38,6 +28,17 @@
     del libucx
 
 
+from .core import *  # noqa
+from .core import get_ucx_version  # noqa
+from .utils import get_ucxpy_logger  # noqa
+from ._libs.utils import get_address  # noqa
+from ._version import __git_commit__, __version__
+
+try:
+    import pynvml
+except ImportError:
+    pynvml = None
+
 # Setup UCX-Py logger
 logger = get_ucxpy_logger()
 

From e424eb37c1dc1ee6dce8d4bf22951e062b9ce125 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Mon, 6 May 2024 13:14:14 -0700
Subject: [PATCH 10/35] add build dpeendency too, remove system installations

---
 ci/build_wheel.sh |  8 ++++++++
 dependencies.yaml | 23 +++++++++++++++++++++++
 pyproject.toml    |  1 +
 3 files changed, 32 insertions(+)

diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh
index df43801df..da0baee7b 100755
--- a/ci/build_wheel.sh
+++ b/ci/build_wheel.sh
@@ -6,6 +6,14 @@ 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.*
+
 source rapids-configure-sccache
 source rapids-date-string
 
diff --git a/dependencies.yaml b/dependencies.yaml
index b49eac0ec..0c414db3a 100644
--- a/dependencies.yaml
+++ b/dependencies.yaml
@@ -27,6 +27,7 @@ files:
       table: build-system
     includes:
       - build_python
+      - depends_on_ucx_build
   py_run:
     output: pyproject
     pyproject_dir: .
@@ -117,6 +118,28 @@ dependencies:
         packages:
           - numpy>=1.23,<2.0a0
           - pynvml>=11.4.1
+  depends_on_ucx_build:
+    common:
+      - output_types: conda
+        packages:
+          - &ucx_conda_build 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
+          - matrix: null
+            packages:
+              - libucx==1.15.0
   depends_on_ucx_run:
     common:
       - output_types: conda
diff --git a/pyproject.toml b/pyproject.toml
index a8d760266..07e135efa 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -16,6 +16,7 @@
 build-backend = "setuptools.build_meta"
 requires = [
     "cython>=3.0.0",
+    "libucx==1.15.0",
     "setuptools>=64.0.0",
     "tomli",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit dependencies.yaml and run `rapids-dependency-file-generator`.

From 0a397f02fcbac7001dadcb021bae6265dee382d1 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Mon, 6 May 2024 13:16:51 -0700
Subject: [PATCH 11/35] install is necessary for docs builds

---
 .readthedocs.yml | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/.readthedocs.yml b/.readthedocs.yml
index d72bb73ad..937f721db 100644
--- a/.readthedocs.yml
+++ b/.readthedocs.yml
@@ -5,6 +5,11 @@ build:
   tools:
     python: "mambaforge-22.9"
 
+python:
+  install:
+    - method: pip
+      path: .
+
 conda:
   environment: conda/environments/builddocs.yml
 

From c8d91ca594412fdc9f34750588d185fd91fdbe3c Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Mon, 6 May 2024 13:36:38 -0700
Subject: [PATCH 12/35] fix conda deps, try to fix readthedocs builds

---
 .readthedocs.yml               |  3 +++
 ci/build_docs_pre_install.sh   | 23 +++++++++++++++++++++++
 conda/recipes/ucx-py/meta.yaml |  3 ++-
 3 files changed, 28 insertions(+), 1 deletion(-)
 create mode 100755 ci/build_docs_pre_install.sh

diff --git a/.readthedocs.yml b/.readthedocs.yml
index 937f721db..3f466e746 100644
--- a/.readthedocs.yml
+++ b/.readthedocs.yml
@@ -4,6 +4,9 @@ build:
   os: "ubuntu-22.04"
   tools:
     python: "mambaforge-22.9"
+  jobs:
+    pre_install:
+      - bash ci/build_docs_pre_install.sh
 
 python:
   install:
diff --git a/ci/build_docs_pre_install.sh b/ci/build_docs_pre_install.sh
new file mode 100755
index 000000000..83f1187bf
--- /dev/null
+++ b/ci/build_docs_pre_install.sh
@@ -0,0 +1,23 @@
+#!/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
+#
+#   NOTE: This can go away if/when this project is cut over to rapids-build-backend.
+#         See https://github.com/rapidsai/build-planning/issues/31
+#
+
+set -euo pipefail
+
+# just remove libucx dependency from pyproject.toml... it's not necessary for docs builds,
+# and it's unsuffixed (e.g. no `-cu12`) in source control
+cat ./pyproject.toml | grep -v '"libucx' > pyproject.bak
+mv ./pyproject.bak ./pyproject.toml
diff --git a/conda/recipes/ucx-py/meta.yaml b/conda/recipes/ucx-py/meta.yaml
index f73504578..2ce70958f 100644
--- a/conda/recipes/ucx-py/meta.yaml
+++ b/conda/recipes/ucx-py/meta.yaml
@@ -30,7 +30,8 @@ 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:

From acecf8ee11c7d914cb2b8ac652609f7967ec0c3d Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Mon, 6 May 2024 13:54:33 -0700
Subject: [PATCH 13/35] fix libucx replacement in pyproject.toml

---
 ci/build_wheel.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh
index da0baee7b..88bbca8db 100755
--- a/ci/build_wheel.sh
+++ b/ci/build_wheel.sh
@@ -43,7 +43,7 @@ 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 "s/\"libucx>/\"libucx${PACKAGE_CUDA_SUFFIX}>/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}

From e6acd0c6539b74ca51d4b9f45696a39f7c0652a4 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Mon, 6 May 2024 15:15:43 -0700
Subject: [PATCH 14/35] fix lib paths (still not to update rpath/runpath)

---
 ci/build_wheel.sh |  1 +
 pyproject.toml    |  2 +-
 setup.py          | 18 +++++++++++++++++-
 3 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh
index 88bbca8db..6f23c2d03 100755
--- a/ci/build_wheel.sh
+++ b/ci/build_wheel.sh
@@ -55,6 +55,7 @@ mkdir -p final_dist
 python -m auditwheel repair \
     -w final_dist \
     --exclude "libucm.so.0" \
+    --exclude "libucp.so.0" \
     --exclude "libucx.so.0" \
     --exclude "libucs.so.0" \
     --exclude "libuct.so.0" \
diff --git a/pyproject.toml b/pyproject.toml
index 07e135efa..6b0c43836 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -22,7 +22,7 @@ requires = [
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit dependencies.yaml and run `rapids-dependency-file-generator`.
 
 [project]
-name = "ucx-py"
+name = "ucx-py-cu12"
 dynamic = ["version"]
 description = "Python Bindings for the Unified Communication X library (UCX)"
 readme = { file = "README.md", content-type = "text/markdown" }
diff --git a/setup.py b/setup.py
index 9a52b30c1..956de0fff 100644
--- a/setup.py
+++ b/setup.py
@@ -12,8 +12,24 @@
 from setuptools import setup
 from setuptools.extension import Extension
 
-include_dirs = [os.path.dirname(get_python_inc())]
 library_dirs = [get_config_var("LIBDIR")]
+
+# if the 'libucx' wheel was installed, add its library location to
+# library_dirs so it'll be linked against
+#
+# NOTE: doing this '.__file__' stuff because `sysconfig.get_path("platlib")` and similar
+#       can return paths to the main Python interpreters' installation locations...
+#       not where 'libucx' is going to be installed at build time when using
+#       build isolation
+try:
+    import libucx
+
+    libucx_libdir = os.path.join(os.path.dirname(libucx.__file__), "lib")
+    library_dirs.append(libucx_libdir)
+except ImportError:
+    pass
+
+include_dirs = [os.path.dirname(get_python_inc())]
 libraries = ["ucp", "uct", "ucm", "ucs"]
 extra_compile_args = ["-std=c99", "-Werror"]
 

From f8f25c9d8e222011177a730273bcdf01e08a7b6d Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Tue, 7 May 2024 08:43:43 -0700
Subject: [PATCH 15/35] remove debugging stuff, switch to ImportError

---
 ci/build_wheel.sh | 1 -
 pyproject.toml    | 2 +-
 setup.py          | 2 +-
 ucp/__init__.py   | 2 +-
 4 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh
index 6f23c2d03..4a404a378 100755
--- a/ci/build_wheel.sh
+++ b/ci/build_wheel.sh
@@ -56,7 +56,6 @@ python -m auditwheel repair \
     -w final_dist \
     --exclude "libucm.so.0" \
     --exclude "libucp.so.0" \
-    --exclude "libucx.so.0" \
     --exclude "libucs.so.0" \
     --exclude "libuct.so.0" \
     dist/*
diff --git a/pyproject.toml b/pyproject.toml
index 6b0c43836..07e135efa 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -22,7 +22,7 @@ requires = [
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit dependencies.yaml and run `rapids-dependency-file-generator`.
 
 [project]
-name = "ucx-py-cu12"
+name = "ucx-py"
 dynamic = ["version"]
 description = "Python Bindings for the Unified Communication X library (UCX)"
 readme = { file = "README.md", content-type = "text/markdown" }
diff --git a/setup.py b/setup.py
index 956de0fff..e5f4914d8 100644
--- a/setup.py
+++ b/setup.py
@@ -18,7 +18,7 @@
 # library_dirs so it'll be linked against
 #
 # NOTE: doing this '.__file__' stuff because `sysconfig.get_path("platlib")` and similar
-#       can return paths to the main Python interpreters' installation locations...
+#       can return paths to the main Python interpreter's installation locations...
 #       not where 'libucx' is going to be installed at build time when using
 #       build isolation
 try:
diff --git a/ucp/__init__.py b/ucp/__init__.py
index 7a51cecfb..390fbf459 100644
--- a/ucp/__init__.py
+++ b/ucp/__init__.py
@@ -21,7 +21,7 @@
 # Otherwise, we assume that the library was installed in a system path that ld can find.
 try:
     import libucx
-except ModuleNotFoundError:
+except ImportError:
     pass
 else:
     libucx.load_library()

From 9825f82450b856bcbb2b99638d85bd6440377d68 Mon Sep 17 00:00:00 2001
From: James Lamb <jaylamb20@gmail.com>
Date: Tue, 7 May 2024 15:12:48 -0500
Subject: [PATCH 16/35] Update dependencies.yaml

Co-authored-by: Vyas Ramasubramani <vyas.ramasubramani@gmail.com>
---
 dependencies.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/dependencies.yaml b/dependencies.yaml
index c3a012008..81475aff1 100644
--- a/dependencies.yaml
+++ b/dependencies.yaml
@@ -121,7 +121,7 @@ dependencies:
     common:
       - output_types: conda
         packages:
-          - &ucx_conda_build ucx==1.15.0
+          - ucx==1.15.0
       - output_types: requirements
         packages:
           # pip recognizes the index as a global option for the requirements.txt file

From 9bd984c3f5a56e8be205b0d4fc7198c8bcf9b214 Mon Sep 17 00:00:00 2001
From: James Lamb <jaylamb20@gmail.com>
Date: Tue, 7 May 2024 15:20:12 -0500
Subject: [PATCH 17/35] Update conda/recipes/ucx-py/meta.yaml

---
 conda/recipes/ucx-py/meta.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/conda/recipes/ucx-py/meta.yaml b/conda/recipes/ucx-py/meta.yaml
index 2ce70958f..1e6d6ab4c 100644
--- a/conda/recipes/ucx-py/meta.yaml
+++ b/conda/recipes/ucx-py/meta.yaml
@@ -18,7 +18,7 @@ build:
   number: {{ GIT_DESCRIBE_NUMBER }}
   string: py{{ py_version }}_{{ date_string }}_{{ GIT_DESCRIBE_HASH }}_{{ GIT_DESCRIBE_NUMBER }}
   script:
-    - {{ PYTHON }} -m pip install --no-build-isolation --no-deps . -vv
+    - {{ PYTHON }} -m pip install . -vv
   ignore_run_exports_from:
     - ucx
 

From 0523ba1905d9607962e45a04b5cd751859b3a503 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Wed, 8 May 2024 10:52:26 -0700
Subject: [PATCH 18/35] make build-time search for libucx insensitive to its
 internal layout

---
 setup.py | 21 +++++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/setup.py b/setup.py
index e5f4914d8..da6045981 100644
--- a/setup.py
+++ b/setup.py
@@ -22,10 +22,27 @@
 #       not where 'libucx' is going to be installed at build time when using
 #       build isolation
 try:
+    import glob
+
     import libucx
 
-    libucx_libdir = os.path.join(os.path.dirname(libucx.__file__), "lib")
-    library_dirs.append(libucx_libdir)
+    # 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})"
+        )
+
+    # add all the places 'libucx' stores libraries to that paths
+    # considered by the linker when compiling extensions
+    library_dirs.extend(lib_dirs)
+
 except ImportError:
     pass
 

From c60800f73fe327bbacc2d558d4e04ad6cb33b899 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Wed, 8 May 2024 11:40:42 -0700
Subject: [PATCH 19/35] grammar

---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index da6045981..7b716144d 100644
--- a/setup.py
+++ b/setup.py
@@ -39,7 +39,7 @@
             f"Did not find shared libraries in 'libucx' install location ({module_dir})"
         )
 
-    # add all the places 'libucx' stores libraries to that paths
+    # add all the places 'libucx' stores libraries to the paths
     # considered by the linker when compiling extensions
     library_dirs.extend(lib_dirs)
 

From 91cbdf5d919ff45a6154b16014b3e2bc88b1f9fe Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Wed, 8 May 2024 12:25:25 -0700
Subject: [PATCH 20/35] test with PIP_NO_DEPENDENCIES

---
 .readthedocs.yml             |  5 ++---
 ci/build_docs_pre_install.sh | 23 -----------------------
 2 files changed, 2 insertions(+), 26 deletions(-)
 delete mode 100755 ci/build_docs_pre_install.sh

diff --git a/.readthedocs.yml b/.readthedocs.yml
index 3f466e746..9590613c4 100644
--- a/.readthedocs.yml
+++ b/.readthedocs.yml
@@ -4,10 +4,9 @@ build:
   os: "ubuntu-22.04"
   tools:
     python: "mambaforge-22.9"
-  jobs:
-    pre_install:
-      - bash ci/build_docs_pre_install.sh
 
+# PIP_NO_DEPS="True" is set in the readthedocs build environment.
+# ref: https://docs.readthedocs.io/en/stable/guides/environment-variables.html
 python:
   install:
     - method: pip
diff --git a/ci/build_docs_pre_install.sh b/ci/build_docs_pre_install.sh
deleted file mode 100755
index 83f1187bf..000000000
--- a/ci/build_docs_pre_install.sh
+++ /dev/null
@@ -1,23 +0,0 @@
-#!/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
-#
-#   NOTE: This can go away if/when this project is cut over to rapids-build-backend.
-#         See https://github.com/rapidsai/build-planning/issues/31
-#
-
-set -euo pipefail
-
-# just remove libucx dependency from pyproject.toml... it's not necessary for docs builds,
-# and it's unsuffixed (e.g. no `-cu12`) in source control
-cat ./pyproject.toml | grep -v '"libucx' > pyproject.bak
-mv ./pyproject.bak ./pyproject.toml

From 95da3adf3dcd25911b3eb91123324050834c69ae Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Wed, 8 May 2024 12:38:20 -0700
Subject: [PATCH 21/35] try an empty commit... maybe RTD doesnt apply
 environment variable updates on rebuilds


From 327e93e6b4593ce373421b8e2a268cd1d8c5d46c Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Wed, 8 May 2024 12:57:40 -0700
Subject: [PATCH 22/35] try with PIP_NO_BUILD_ISOLATION

---
 .readthedocs.yml | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/.readthedocs.yml b/.readthedocs.yml
index 9590613c4..6bb568dd7 100644
--- a/.readthedocs.yml
+++ b/.readthedocs.yml
@@ -5,7 +5,9 @@ build:
   tools:
     python: "mambaforge-22.9"
 
-# PIP_NO_DEPS="True" is set in the readthedocs build environment.
+# PIP_NO_DEPS="True" and PIP_NO_BUILD_ISOLATION="True" are set in the readthedocs build environment.
+# Al of the project's build dependencies need to be provided via the conda environment
+# at conda/environments/builddocs.yml.
 # ref: https://docs.readthedocs.io/en/stable/guides/environment-variables.html
 python:
   install:

From 63976f2133661ddef0e6204d86aa9bc218a778b4 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Wed, 8 May 2024 13:37:36 -0700
Subject: [PATCH 23/35] just put a real dependency in pyproject.toml

---
 .readthedocs.yml  | 4 ----
 dependencies.yaml | 4 +++-
 pyproject.toml    | 2 +-
 3 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/.readthedocs.yml b/.readthedocs.yml
index 6bb568dd7..937f721db 100644
--- a/.readthedocs.yml
+++ b/.readthedocs.yml
@@ -5,10 +5,6 @@ build:
   tools:
     python: "mambaforge-22.9"
 
-# PIP_NO_DEPS="True" and PIP_NO_BUILD_ISOLATION="True" are set in the readthedocs build environment.
-# Al of the project's build dependencies need to be provided via the conda environment
-# at conda/environments/builddocs.yml.
-# ref: https://docs.readthedocs.io/en/stable/guides/environment-variables.html
 python:
   install:
     - method: pip
diff --git a/dependencies.yaml b/dependencies.yaml
index 81475aff1..6552689af 100644
--- a/dependencies.yaml
+++ b/dependencies.yaml
@@ -158,9 +158,11 @@ dependencies:
           - 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
+              - libucx-cu11>=1.15.0,<1.16
   test_python:
     common:
       - output_types: [conda, requirements, pyproject]
diff --git a/pyproject.toml b/pyproject.toml
index 0c3939f35..06701c74c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -31,7 +31,7 @@ authors = [
 license = { text = "BSD-3-Clause" }
 requires-python = ">=3.9"
 dependencies = [
-    "libucx>=1.15.0,<1.16",
+    "libucx-cu11>=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`.

From 72596ae0266906634f17773b010e3a50b4759841 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Wed, 8 May 2024 13:53:08 -0700
Subject: [PATCH 24/35] build and run

---
 dependencies.yaml | 4 +++-
 pyproject.toml    | 2 +-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/dependencies.yaml b/dependencies.yaml
index 6552689af..378c0be12 100644
--- a/dependencies.yaml
+++ b/dependencies.yaml
@@ -136,9 +136,11 @@ dependencies:
           - 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
+              - libucx-cu11==1.15.0
   depends_on_ucx_run:
     common:
       - output_types: conda
diff --git a/pyproject.toml b/pyproject.toml
index 06701c74c..917a32683 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -16,7 +16,7 @@
 build-backend = "setuptools.build_meta"
 requires = [
     "cython>=3.0.0",
-    "libucx==1.15.0",
+    "libucx-cu11==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`.
 

From 114623177271327dbeebbc68a01e4a5577635ef5 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Wed, 8 May 2024 14:04:33 -0700
Subject: [PATCH 25/35] please

---
 .readthedocs.yml                 | 5 -----
 conda/environments/builddocs.yml | 3 +++
 2 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/.readthedocs.yml b/.readthedocs.yml
index 937f721db..d72bb73ad 100644
--- a/.readthedocs.yml
+++ b/.readthedocs.yml
@@ -5,11 +5,6 @@ build:
   tools:
     python: "mambaforge-22.9"
 
-python:
-  install:
-    - method: pip
-      path: .
-
 conda:
   environment: conda/environments/builddocs.yml
 
diff --git a/conda/environments/builddocs.yml b/conda/environments/builddocs.yml
index 392ac63c0..b08e5d942 100644
--- a/conda/environments/builddocs.yml
+++ b/conda/environments/builddocs.yml
@@ -14,5 +14,8 @@ dependencies:
 - recommonmark
 - pandoc=<2.0.0
 - pip
+- pip:
+  - --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple
+  - ../../
 - ucx
 - cython

From 5046584938ef195c528dfe8107ba702a64ce38a3 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Wed, 8 May 2024 14:11:16 -0700
Subject: [PATCH 26/35] another empty commit to test removing env variables
 from the RTD UI


From 5c90c8ed095a08fc780068f09f0fe1a9d13a97e5 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Wed, 8 May 2024 14:17:59 -0700
Subject: [PATCH 27/35] catch more issues

---
 ucp/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ucp/__init__.py b/ucp/__init__.py
index 390fbf459..ce9764fa1 100644
--- a/ucp/__init__.py
+++ b/ucp/__init__.py
@@ -21,7 +21,7 @@
 # Otherwise, we assume that the library was installed in a system path that ld can find.
 try:
     import libucx
-except ImportError:
+except (ImportError, OSError, RuntimeError):
     pass
 else:
     libucx.load_library()

From f1f949dffcf4f24b2fbd9817ab40d67d76723968 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Wed, 8 May 2024 14:39:35 -0700
Subject: [PATCH 28/35] restore CPU support

---
 ucp/__init__.py | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/ucp/__init__.py b/ucp/__init__.py
index ce9764fa1..b8b5c67f6 100644
--- a/ucp/__init__.py
+++ b/ucp/__init__.py
@@ -21,11 +21,19 @@
 # Otherwise, we assume that the library was installed in a system path that ld can find.
 try:
     import libucx
-except (ImportError, OSError, RuntimeError):
+except ImportError:
     pass
 else:
-    libucx.load_library()
-    del libucx
+    try:
+        libucx.load_library()
+        del libucx
+    except RuntimeError as err:
+        msg = (
+            "Loading 'libucx' failed. Ignore this warning on CPU-only systems, "
+            "but on a system with GPUs this could indicate a more serious issue. "
+            f"Full error: {str(err)}"
+        )
+        logger.warning(msg)
 
 
 from .core import *  # noqa

From 6723baea944acd9e5c181df3819c1fcbae0993a0 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Fri, 10 May 2024 06:26:42 -0700
Subject: [PATCH 29/35] remove try-catch for no-GPU settings

---
 ucp/__init__.py | 12 ++----------
 1 file changed, 2 insertions(+), 10 deletions(-)

diff --git a/ucp/__init__.py b/ucp/__init__.py
index b8b5c67f6..390fbf459 100644
--- a/ucp/__init__.py
+++ b/ucp/__init__.py
@@ -24,16 +24,8 @@
 except ImportError:
     pass
 else:
-    try:
-        libucx.load_library()
-        del libucx
-    except RuntimeError as err:
-        msg = (
-            "Loading 'libucx' failed. Ignore this warning on CPU-only systems, "
-            "but on a system with GPUs this could indicate a more serious issue. "
-            f"Full error: {str(err)}"
-        )
-        logger.warning(msg)
+    libucx.load_library()
+    del libucx
 
 
 from .core import *  # noqa

From ca60927daa7009a350f9d2543cea87520c838e9d Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Fri, 10 May 2024 06:35:18 -0700
Subject: [PATCH 30/35] exclude libucs_signal too

---
 ci/build_wheel.sh | 1 +
 1 file changed, 1 insertion(+)

diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh
index 4a404a378..1b1d36892 100755
--- a/ci/build_wheel.sh
+++ b/ci/build_wheel.sh
@@ -57,6 +57,7 @@ python -m auditwheel repair \
     --exclude "libucm.so.0" \
     --exclude "libucp.so.0" \
     --exclude "libucs.so.0" \
+    --exclude "libucs_signal.so.0" \
     --exclude "libuct.so.0" \
     dist/*
 

From fb12aa34bfea86e041c0b80b870d5a4f3698fb8b Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Fri, 10 May 2024 07:52:37 -0700
Subject: [PATCH 31/35] remove ucx from docs conda env

---
 conda/environments/builddocs.yml | 1 -
 1 file changed, 1 deletion(-)

diff --git a/conda/environments/builddocs.yml b/conda/environments/builddocs.yml
index b08e5d942..cd9766abd 100644
--- a/conda/environments/builddocs.yml
+++ b/conda/environments/builddocs.yml
@@ -17,5 +17,4 @@ dependencies:
 - pip:
   - --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple
   - ../../
-- ucx
 - cython

From eba110fe2116978e9488b46b7d884ab6b5ffb876 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Fri, 10 May 2024 08:32:00 -0700
Subject: [PATCH 32/35] use the headers from the libucx wheels too

---
 ci/build_wheel.sh |  5 +++++
 setup.py          | 48 +++++++++++++++++++++++++++++++----------------
 2 files changed, 37 insertions(+), 16 deletions(-)

diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh
index 1b1d36892..5417f3090 100755
--- a/ci/build_wheel.sh
+++ b/ci/build_wheel.sh
@@ -14,6 +14,11 @@ 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
 
diff --git a/setup.py b/setup.py
index 7b716144d..2c238c3ca 100644
--- a/setup.py
+++ b/setup.py
@@ -5,6 +5,7 @@
 
 from __future__ import absolute_import, print_function
 
+import glob
 import os
 from distutils.sysconfig import get_config_var, get_python_inc
 
@@ -12,19 +13,21 @@
 from setuptools import setup
 from setuptools.extension import Extension
 
-library_dirs = [get_config_var("LIBDIR")]
 
-# if the 'libucx' wheel was installed, add its library location to
-# library_dirs so it'll be linked against
-#
-# NOTE: doing this '.__file__' stuff because `sysconfig.get_path("platlib")` and similar
-#       can return paths to the main Python interpreter's installation locations...
-#       not where 'libucx' is going to be installed at build time when using
-#       build isolation
-try:
-    import glob
+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.
 
-    import libucx
+    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__)
@@ -39,17 +42,30 @@
             f"Did not find shared libraries in 'libucx' install location ({module_dir})"
         )
 
-    # add all the places 'libucx' stores libraries to the paths
-    # considered by the linker when compiling extensions
-    library_dirs.extend(lib_dirs)
+    # 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)
 
-except ImportError:
-    pass
 
 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(

From fb022ee20592a2e9ef15cfed41a34b5c32917e18 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Fri, 10 May 2024 08:52:56 -0700
Subject: [PATCH 33/35] restore unsuffixed dependencies in pyproject.toml

---
 .readthedocs.yml             |  3 +++
 ci/build_docs_pre_install.sh | 17 +++++++++++++++++
 dependencies.yaml            |  4 ++--
 pyproject.toml               |  4 ++--
 4 files changed, 24 insertions(+), 4 deletions(-)
 create mode 100755 ci/build_docs_pre_install.sh

diff --git a/.readthedocs.yml b/.readthedocs.yml
index d72bb73ad..1326fa3e3 100644
--- a/.readthedocs.yml
+++ b/.readthedocs.yml
@@ -4,6 +4,9 @@ build:
   os: "ubuntu-22.04"
   tools:
     python: "mambaforge-22.9"
+  jobs:
+    pre_install:
+      - bash ci/build_docs_pre_install.sh
 
 conda:
   environment: conda/environments/builddocs.yml
diff --git a/ci/build_docs_pre_install.sh b/ci/build_docs_pre_install.sh
new file mode 100755
index 000000000..a422b1ac2
--- /dev/null
+++ b/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==\1\"/g" ./pyproject.toml
diff --git a/dependencies.yaml b/dependencies.yaml
index 378c0be12..9b7e3b562 100644
--- a/dependencies.yaml
+++ b/dependencies.yaml
@@ -140,7 +140,7 @@ dependencies:
           #       so 'pip install .' (e.g. as used in docs builds) will work
           - matrix: null
             packages:
-              - libucx-cu11==1.15.0
+              - libucx==1.15.0
   depends_on_ucx_run:
     common:
       - output_types: conda
@@ -164,7 +164,7 @@ dependencies:
           #       so "pip install ." (e.g. as used in docs builds) will work
           - matrix: null
             packages:
-              - libucx-cu11>=1.15.0,<1.16
+              - libucx>=1.15.0,<1.16
   test_python:
     common:
       - output_types: [conda, requirements, pyproject]
diff --git a/pyproject.toml b/pyproject.toml
index 917a32683..0c3939f35 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -16,7 +16,7 @@
 build-backend = "setuptools.build_meta"
 requires = [
     "cython>=3.0.0",
-    "libucx-cu11==1.15.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`.
 
@@ -31,7 +31,7 @@ authors = [
 license = { text = "BSD-3-Clause" }
 requires-python = ">=3.9"
 dependencies = [
-    "libucx-cu11>=1.15.0,<1.16",
+    "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`.

From 37dbad6a9fe920227d83f647e042c343186aeb6b Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Fri, 10 May 2024 08:57:26 -0700
Subject: [PATCH 34/35] run patching script earlier on docs builds

---
 .readthedocs.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.readthedocs.yml b/.readthedocs.yml
index 1326fa3e3..f672b74d8 100644
--- a/.readthedocs.yml
+++ b/.readthedocs.yml
@@ -5,7 +5,7 @@ build:
   tools:
     python: "mambaforge-22.9"
   jobs:
-    pre_install:
+    post_checkout:
       - bash ci/build_docs_pre_install.sh
 
 conda:

From e2cffe355769eb04d26b5e969dea20fad268a0a5 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Fri, 10 May 2024 09:06:16 -0700
Subject: [PATCH 35/35] catch run dependencies too

---
 ci/build_docs_pre_install.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ci/build_docs_pre_install.sh b/ci/build_docs_pre_install.sh
index a422b1ac2..95a38f724 100755
--- a/ci/build_docs_pre_install.sh
+++ b/ci/build_docs_pre_install.sh
@@ -14,4 +14,4 @@
 
 set -euo pipefail
 
-sed -r -i "s/libucx==(.*)\"/libucx-cu12==\1\"/g" ./pyproject.toml
+sed -r -i "s/\"libucx/\"libucx-cu12/g" ./pyproject.toml