Merge branch 'branch-0.35' into benchmark-gil-contention

.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.04
+    uses: rapidsai/shared-action-workflows/.github/workflows/conda-python-build.yaml@cuda-120-arm
     with:
       build_type: ${{ inputs.build_type || 'branch' }}
       branch: ${{ inputs.branch }}
@@ -37,32 +37,25 @@ jobs:
   upload-conda:
     needs: [conda-python-build]
     secrets: inherit
-    uses: rapidsai/shared-action-workflows/.github/workflows/conda-upload-packages.yaml@branch-23.04
+    uses: rapidsai/shared-action-workflows/.github/workflows/conda-upload-packages.yaml@cuda-120-arm
     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-manylinux-build.[email protected]
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-build.yaml@cuda-120-arm
     with:
       build_type: ${{ inputs.build_type || 'branch' }}
       branch: ${{ inputs.branch }}
       sha: ${{ inputs.sha }}
       date: ${{ inputs.date }}
-      package-name: ucx_py
-      package-dir: .
-      # copy custom auditwheel script to known location with '{package}' placeholder
-      # in the auditwheel repair command, the '{package}' placeholder can't be used
-      cibw-before-build: 'cp {package}/.github/build-scripts/patch-wheel.sh /patch-wheel.sh'
-      post-repair-amd64: "/patch-wheel.sh {dest_dir}/*.whl"
-      post-repair-arm64: "/patch-wheel.sh {dest_dir}/*.whl"
-      uses-setup-env-vars: false
+      script: ci/build_wheel.sh
   wheel-publish:
     needs: wheel-build
     secrets: inherit
-    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-manylinux-publish.[email protected]
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-publish.yaml@cuda-120-arm
     with:
       build_type: ${{ inputs.build_type || 'branch' }}
       branch: ${{ inputs.branch }}

.github/workflows/pr.yaml

@@ -18,43 +18,34 @@ jobs:
       - wheel-build
       - wheel-tests
     secrets: inherit
-    uses: rapidsai/shared-action-workflows/.github/workflows/pr-builder.yaml@branch-23.04
+    uses: rapidsai/shared-action-workflows/.github/workflows/pr-builder.yaml@cuda-120-arm
   checks:
     secrets: inherit
-    uses: rapidsai/shared-action-workflows/.github/workflows/checks.yaml@branch-23.04
+    uses: rapidsai/shared-action-workflows/.github/workflows/checks.yaml@cuda-120-arm
   conda-python-build:
     needs: checks
     secrets: inherit
-    uses: rapidsai/shared-action-workflows/.github/workflows/conda-python-build.yaml@branch-23.04
+    uses: rapidsai/shared-action-workflows/.github/workflows/conda-python-build.yaml@cuda-120-arm
     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.04
+    uses: rapidsai/shared-action-workflows/.github/workflows/conda-python-tests.yaml@cuda-120-arm
     with:
       build_type: pull-request
   wheel-build:
     needs: checks
     secrets: inherit
-    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-manylinux-build.[email protected]
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-build.yaml@cuda-120-arm
     with:
       build_type: pull-request
-      package-name: ucx_py
-      package-dir: .
-      cibw-before-build: 'cp {package}/.github/build-scripts/patch-wheel.sh /patch-wheel.sh'
-      post-repair-amd64: "/patch-wheel.sh {dest_dir}/*.whl"
-      post-repair-arm64: "/patch-wheel.sh {dest_dir}/*.whl"
-      uses-setup-env-vars: false
+      script: ci/build_wheel.sh
   wheel-tests:
     needs: wheel-build
     secrets: inherit
-    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-manylinux-test.[email protected]
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-test.yaml@cuda-120-arm
     with:
       build_type: pull-request
-      package-name: ucx_py
+      script: ci/test_wheel.sh
       test-docker-options: "--cap-add CAP_SYS_PTRACE --shm-size=8g"
-      test-before-arm64: "pip install cupy-cuda11x -f https://pip.cupy.dev/aarch64"
-      # skipped test context: https://github.com/rapidsai/ucx-py/pull/909
-      test-unittest: "pytest -k 'not test_send_recv_am' --cache-clear -vs ./ucp/_libs/tests/ && pytest --cache-clear -vs ./tests"
-      test-smoketest: "pytest -k 'not test_send_recv_am' --cache-clear -vs ./ucp/_libs/tests/ && pytest --cache-clear -vs ./tests"

.github/workflows/test.yaml

@@ -16,22 +16,19 @@ on:
 jobs:
   conda-python-tests:
     secrets: inherit
-    uses: rapidsai/shared-action-workflows/.github/workflows/conda-python-tests.yaml@branch-23.04
+    uses: rapidsai/shared-action-workflows/.github/workflows/conda-python-tests.yaml@cuda-120-arm
     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-manylinux-test.[email protected]
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-test.yaml@cuda-120-arm
     with:
       build_type: nightly
       branch: ${{ inputs.branch }}
       date: ${{ inputs.date }}
       sha: ${{ inputs.sha }}
-      package-name: ucx_py
+      script: ci/test_wheel.sh
       test-docker-options: "--cap-add CAP_SYS_PTRACE --shm-size=8g"
-      test-before-arm64: "pip install cupy-cuda11x -f https://pip.cupy.dev/aarch64"
-      # skipped test context: https://github.com/rapidsai/ucx-py/pull/909
-      test-unittest: "pytest -k 'not test_send_recv_am' --cache-clear -vs ./ucp/_libs/tests/ && pytest --cache-clear -vs ./tests"

.pre-commit-config.yaml

@@ -23,5 +23,10 @@ repos:
                 name: flake8-cython
                 args: ["--config=.flake8.cython"]
                 types: [cython]
+      - repo: https://github.com/rapidsai/dependency-file-generator
+        rev: v1.5.2
+        hooks:
+              - id: rapids-dependency-file-generator
+                args: ["--clean"]
 default_language_version:
       python: python3

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

@@ -0,0 +1,125 @@
+#!/bin/bash
+# Copyright (c) 2023, NVIDIA CORPORATION.
+
+set -euo pipefail
+
+package_name="ucx-py"
+underscore_package_name=$(echo "${package_name}" | tr "-" "_")
+
+source rapids-configure-sccache
+source rapids-date-string
+
+# Use gha-tools rapids-pip-wheel-version to generate wheel version then
+# update the necessary files
+version_override="$(rapids-pip-wheel-version ${RAPIDS_DATE_STRING})"
+
+RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})"
+
+# 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
+
+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) }')
+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/${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
+cp -P /usr/lib/ucx/* .
+
+# 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 --replace-needed libnuma.so.1 $LIBNUMA $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

ci/release/update-version.sh

@@ -24,8 +24,11 @@ NEXT_MINOR=$(echo $NEXT_FULL_TAG | awk '{split($0, a, "."); print a[2]}')
 NEXT_SHORT_TAG=${NEXT_MAJOR}.${NEXT_MINOR}
 
 # Get RAPIDS version associated w/ ucx-py version
-NEXT_RAPIDS_VERSION="$(curl -sL https://version.gpuci.io/ucx-py/${NEXT_SHORT_TAG})"
+NEXT_RAPIDS_SHORT_TAG="$(curl -sL https://version.gpuci.io/ucx-py/${NEXT_SHORT_TAG})"
 
+# 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"
 
@@ -34,8 +37,18 @@ function sed_runner() {
     sed -i.bak ''"$1"'' $2 && rm -f ${2}.bak
 }
 
-sed_runner "s/cudf=.*/cudf=${NEXT_RAPIDS_VERSION}/g" dependencies.yaml
+DEPENDENCIES=(
+  cudf
+)
+for FILE in dependencies.yaml; do
+  for DEP in "${DEPENDENCIES[@]}"; do
+    sed_runner "/-.* ${DEP}==/ s/==.*/==${NEXT_RAPIDS_SHORT_TAG_PEP440}\.*/g" ${FILE};
+  done
+done
 
 for FILE in .github/workflows/*.yaml; do
-  sed_runner "/shared-action-workflows/ s/@.*/@branch-${NEXT_RAPIDS_VERSION}/g" "${FILE}"
+  sed_runner "/shared-action-workflows/ s/@.*/@branch-${NEXT_RAPIDS_SHORT_TAG}/g" "${FILE}"
 done
+
+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