Merge pull request #993 from rapidsai/branch-0.34

[RELEASE] ucx-py v0.34

.github/copy-pr-bot.yaml

@@ -0,0 +1,4 @@
+# Configuration file for `copy-pr-bot` GitHub App
+# https://docs.gha-runners.nvidia.com/apps/copy-pr-bot/
+
+enabled: true

.github/ops-bot.yaml

@@ -5,5 +5,4 @@ auto_merger: true
 branch_checker: false
 label_checker: false
 release_drafter: false
-copy_prs: true
 recently_updated: true

.github/workflows/build.yaml

@@ -28,7 +28,7 @@ concurrency:
 jobs:
   conda-python-build:
     secrets: inherit
-    uses: rapidsai/shared-action-workflows/.github/workflows/conda-python-build.yaml@branch-23.08
+    uses: rapidsai/shared-action-workflows/.github/workflows/conda-python-build.yaml@branch-23.10
     with:
       build_type: ${{ inputs.build_type || 'branch' }}
       branch: ${{ inputs.branch }}
@@ -37,15 +37,15 @@ jobs:
   upload-conda:
     needs: [conda-python-build]
     secrets: inherit
-    uses: rapidsai/shared-action-workflows/.github/workflows/conda-upload-packages.yaml@branch-23.08
+    uses: rapidsai/shared-action-workflows/.github/workflows/conda-upload-packages.yaml@branch-23.10
     with:
       build_type: ${{ inputs.build_type || 'branch' }}
       branch: ${{ inputs.branch }}
       date: ${{ inputs.date }}
       sha: ${{ inputs.sha }}
   wheel-build:
     secrets: inherit
-    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-build.yaml@branch-23.08
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-build.yaml@branch-23.10
     with:
       build_type: ${{ inputs.build_type || 'branch' }}
       branch: ${{ inputs.branch }}
@@ -55,7 +55,7 @@ jobs:
   wheel-publish:
     needs: wheel-build
     secrets: inherit
-    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-publish.yaml@branch-23.08
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-publish.yaml@branch-23.10
     with:
       build_type: ${{ inputs.build_type || 'branch' }}
       branch: ${{ inputs.branch }}

.github/workflows/pr.yaml

@@ -18,33 +18,33 @@ jobs:
       - wheel-build
       - wheel-tests
     secrets: inherit
-    uses: rapidsai/shared-action-workflows/.github/workflows/pr-builder.yaml@branch-23.08
+    uses: rapidsai/shared-action-workflows/.github/workflows/pr-builder.yaml@branch-23.10
   checks:
     secrets: inherit
-    uses: rapidsai/shared-action-workflows/.github/workflows/checks.yaml@branch-23.08
+    uses: rapidsai/shared-action-workflows/.github/workflows/checks.yaml@branch-23.10
   conda-python-build:
     needs: checks
     secrets: inherit
-    uses: rapidsai/shared-action-workflows/.github/workflows/conda-python-build.yaml@branch-23.08
+    uses: rapidsai/shared-action-workflows/.github/workflows/conda-python-build.yaml@branch-23.10
     with:
       build_type: pull-request
   conda-python-tests:
     needs: conda-python-build
     secrets: inherit
-    uses: rapidsai/shared-action-workflows/.github/workflows/conda-python-tests.yaml@branch-23.08
+    uses: rapidsai/shared-action-workflows/.github/workflows/conda-python-tests.yaml@branch-23.10
     with:
       build_type: pull-request
   wheel-build:
     needs: checks
     secrets: inherit
-    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-build.yaml@branch-23.08
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-build.yaml@branch-23.10
     with:
       build_type: pull-request
       script: ci/build_wheel.sh
   wheel-tests:
     needs: wheel-build
     secrets: inherit
-    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-test.yaml@branch-23.08
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-test.yaml@branch-23.10
     with:
       build_type: pull-request
       script: ci/test_wheel.sh

.github/workflows/test.yaml

@@ -16,15 +16,15 @@ on:
 jobs:
   conda-python-tests:
     secrets: inherit
-    uses: rapidsai/shared-action-workflows/.github/workflows/conda-python-tests.yaml@branch-23.08
+    uses: rapidsai/shared-action-workflows/.github/workflows/conda-python-tests.yaml@branch-23.10
     with:
       build_type: nightly
       branch: ${{ inputs.branch }}
       date: ${{ inputs.date }}
       sha: ${{ inputs.sha }}
   wheel-tests:
     secrets: inherit
-    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-test.yaml@branch-23.08
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-test.yaml@branch-23.10
     with:
       build_type: nightly
       branch: ${{ inputs.branch }}

ci/build_python.sh

@@ -9,7 +9,7 @@ rapids-print-env
 
 rapids-logger "Begin py build"
 
-rapids-mamba-retry mambabuild \
+rapids-conda-retry mambabuild \
   conda/recipes/ucx-py
 
 rapids-upload-conda-to-s3 python

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

ci/release/update-version.sh

@@ -28,6 +28,7 @@ NEXT_RAPIDS_SHORT_TAG="$(curl -sL https://version.gpuci.io/ucx-py/${NEXT_SHORT_T
 
 # Need to distutils-normalize the versions for some use cases
 NEXT_RAPIDS_SHORT_TAG_PEP440=$(python -c "from setuptools.extern import packaging; print(packaging.version.Version('${NEXT_RAPIDS_SHORT_TAG}'))")
+NEXT_RAPIDS_FULL_TAG_PEP440=$(python -c "from setuptools.extern import packaging; print(packaging.version.Version('${NEXT_FULL_TAG}'))")
 
 echo "Preparing release $CURRENT_TAG => $NEXT_FULL_TAG"
 
@@ -49,5 +50,5 @@ for FILE in .github/workflows/*.yaml; do
   sed_runner "/shared-action-workflows/ s/@.*/@branch-${NEXT_RAPIDS_SHORT_TAG}/g" "${FILE}"
 done
 
-sed_runner "s/^version = .*/version = \"${NEXT_FULL_TAG}\"/g" pyproject.toml
-sed_runner "s/^__version__ = .*/__version__ = \"${NEXT_FULL_TAG}\"/g" ucp/__init__.py
+sed_runner "s/^version = .*/version = \"${NEXT_RAPIDS_FULL_TAG_PEP440}\"/g" pyproject.toml
+sed_runner "s/^__version__ = .*/__version__ = \"${NEXT_RAPIDS_FULL_TAG_PEP440}\"/g" ucp/__init__.py

dependencies.yaml

@@ -100,23 +100,23 @@ dependencies:
       - output_types: [conda, requirements, pyproject]
         packages:
           - setuptools>=64.0.0
-          - cython>=0.29.14,<0.30.0a0
+          - cython>=3.0.0
           - tomli  # Not needed for Python 3.11+
   run:
     common:
       - output_types: [conda, requirements, pyproject]
         packages:
           - numpy>=1.21
           - pynvml>=11.4.1
-      - output_types: [conda]
+      - output_types: conda
         packages:
           - ucx
   test_python:
     common:
       - output_types: [conda, requirements, pyproject]
         packages:
           - cloudpickle
-          - cudf==23.8.*
+          - cudf==23.10.*
           - dask
           - distributed
           - numba>=0.57

pyproject.toml

@@ -15,14 +15,14 @@
 [build-system]
 build-backend = "setuptools.build_meta"
 requires = [
-    "cython>=0.29.14,<0.30.0a0",
+    "cython>=3.0.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`.
 
 [project]
 name = "ucx-py"
-version = "0.33.0"
+version = "0.34.0"
 description = "Python Bindings for the Unified Communication X library (UCX)"
 readme = { file = "README.md", content-type = "text/markdown" }
 authors = [
@@ -49,7 +49,7 @@ classifiers = [
 [project.optional-dependencies]
 test = [
     "cloudpickle",
-    "cudf==23.8.*",
+    "cudf==23.10.*",
     "cupy-cuda11x>=12.0.0",
     "dask",
     "distributed",

tests/utils.py

@@ -76,9 +76,7 @@ async def send(ep, frames):
 
     await ep.send(np.array([len(frames)], dtype=np.uint64))
     await ep.send(
-        np.array(
-            [hasattr(f, "__cuda_array_interface__") for f in frames], dtype=np.bool
-        )
+        np.array([hasattr(f, "__cuda_array_interface__") for f in frames], dtype=bool)
     )
     await ep.send(np.array([nbytes(f) for f in frames], dtype=np.uint64))
     # Send frames
@@ -96,7 +94,7 @@ async def recv(ep):
         # Recv meta data
         nframes = np.empty(1, dtype=np.uint64)
         await ep.recv(nframes)
-        is_cudas = np.empty(nframes[0], dtype=np.bool)
+        is_cudas = np.empty(nframes[0], dtype=bool)
         await ep.recv(is_cudas)
         sizes = np.empty(nframes[0], dtype=np.uint64)
         await ep.recv(sizes)

ucp/__init__.py

@@ -64,9 +64,14 @@ def _is_mig_device(handle):
             if _is_mig_device(handle):
                 continue
 
-            total_memory = pynvml.nvmlDeviceGetMemoryInfo(handle).total
-            bar1_total = pynvml.nvmlDeviceGetBAR1MemoryInfo(handle).bar1Total
+            try:
+                bar1_total = pynvml.nvmlDeviceGetBAR1MemoryInfo(handle).bar1Total
+            except pynvml.nvml.NVMLError_NotSupported:
+                # Bar1 access not supported on this device, set it to
+                # zero (always lower than device memory).
+                bar1_total = 0
 
+            total_memory = pynvml.nvmlDeviceGetMemoryInfo(handle).total
             if total_memory <= bar1_total:
                 large_bar1[dev_idx] = True
 
@@ -85,7 +90,7 @@ def _is_mig_device(handle):
     os.environ["UCX_MAX_RNDV_RAILS"] = "1"
 
 
-__version__ = "0.33.0"
+__version__ = "0.34.0"
 __ucx_version__ = "%d.%d.%d" % get_ucx_version()
 
 if get_ucx_version() < (1, 11, 1):

ucp/_libs/ucx_worker.pyx

@@ -47,7 +47,7 @@ cdef _drain_worker_tag_recv(ucp_worker_h handle):
 
         buf = malloc(info.length)
         status = ucp_tag_msg_recv_nb(
-            handle, buf, info.length, ucp_dt_make_contig(1), message, _tag_recv_callback
+            handle, buf, info.length, ucp_dt_make_contig(1), message, _tag_recv_cb
         )
 
         try:
@@ -130,7 +130,7 @@ cdef class UCXWorker(UCXObject):
                 UCP_AM_HANDLER_PARAM_FIELD_ARG
             )
             am_handler_param.id = AM_MSG_ID
-            am_handler_param.cb = _am_recv_callback
+            am_handler_param.cb = <ucp_am_recv_callback_t>_am_recv_callback
             am_handler_param.arg = <void *>self
             status = ucp_worker_set_am_recv_handler(self._handle, &am_handler_param)