add build wheel script and accompanying version info

.github/workflows/build.yaml

@@ -43,3 +43,22 @@ jobs:
       branch: ${{ inputs.branch }}
       date: ${{ inputs.date }}
       sha: ${{ inputs.sha }}
+  wheel-build:
+    secrets: inherit
+    uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-24.02
+    with:
+      build_type: ${{ inputs.build_type || 'branch' }}
+      branch: ${{ inputs.branch }}
+      sha: ${{ inputs.sha }}
+      date: ${{ inputs.date }}
+      script: ci/build_wheel.sh
+  wheel-publish:
+    needs: wheel-build
+    secrets: inherit
+    uses: rapidsai/shared-workflows/.github/workflows/wheels-publish.yaml@branch-24.02
+    with:
+      build_type: ${{ inputs.build_type || 'branch' }}
+      branch: ${{ inputs.branch }}
+      sha: ${{ inputs.sha }}
+      date: ${{ inputs.date }}
+      package-name: ucxx

.github/workflows/pr.yaml

@@ -16,6 +16,8 @@ jobs:
       - conda-cpp-build
       - conda-cpp-tests
       - conda-python-tests
+      - wheel-build
+      - wheel-tests
     secrets: inherit
     uses: rapidsai/shared-workflows/.github/workflows/pr-builder.yaml@branch-24.02
   checks:
@@ -43,3 +45,17 @@ jobs:
     with:
       build_type: pull-request
       container-options: "--cap-add CAP_SYS_PTRACE --shm-size=8g --ulimit=nofile=1000000:1000000"
+  wheel-build:
+    needs: checks
+    secrets: inherit
+    uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-24.02
+    with:
+      build_type: pull-request
+      script: ci/build_wheel.sh
+  wheel-tests:
+    needs: wheel-build
+    secrets: inherit
+    uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@branch-24.02
+    with:
+      build_type: pull-request
+      script: ci/test_wheel.sh

.github/workflows/test.yaml

@@ -32,3 +32,12 @@ jobs:
       date: ${{ inputs.date }}
       sha: ${{ inputs.sha }}
       container-options: "--cap-add CAP_SYS_PTRACE --shm-size=8g --ulimit=nofile=1000000:1000000"
+  wheel-tests:
+    secrets: inherit
+    uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@branch-24.02
+    with:
+      build_type: nightly
+      branch: ${{ inputs.branch }}
+      date: ${{ inputs.date }}
+      sha: ${{ inputs.sha }}
+      script: ci/test_wheel.sh

VERSION

@@ -0,0 +1 @@
+0.36.00

ci/build_wheel.sh

@@ -0,0 +1,52 @@
+#!/bin/bash
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.
+
+set -euo pipefail
+
+package_name="ucxx"
+package_dir="python"
+
+source rapids-configure-sccache
+source rapids-date-string
+
+version=$(rapids-generate-version)
+commit=$(git rev-parse HEAD)
+
+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="${package_dir}/pyproject.toml"
+
+sed -i -E "s/^name = \"${package_name}(.*)?\"$/name = \"${package_name}${PACKAGE_CUDA_SUFFIX}\"/g" ${pyproject_file}
+echo "${version}" > VERSION
+sed -i "/^__git_commit__/ s/= .*/= \"${commit}\"/g" "${package_dir}/${package_name//-/_}/_version.py"
+
+# 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
+
+# Add -cuXX to package name
+sed -r -i "s/rmm(.*)\"/rmm${PACKAGE_CUDA_SUFFIX}\1${alpha_spec}\"/g" ${pyproject_file}
+sed -r -i "s/cudf(.*)\"/cudf${PACKAGE_CUDA_SUFFIX}\1${alpha_spec}\"/g" ${pyproject_file}
+
+# Update cupy package name (different suffix from RAPIDS)
+if [[ $PACKAGE_CUDA_SUFFIX == "-cu12" ]]; then
+    sed -i "s/cupy-cuda11x/cupy-cuda12x/g" ${pyproject_file}
+fi
+
+SKBUILD_CMAKE_ARGS="-DUCXX_ENABLE_RMM=ON" \
+    python -m pip wheel "${package_dir}"/ -w "${package_dir}"/dist -vvv --no-deps --disable-pip-version-check
+
+python -m auditwheel repair -w ${package_dir}/final_dist ${package_dir}/dist/*
+
+RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})"
+RAPIDS_PY_WHEEL_NAME="ucxx_${RAPIDS_PY_CUDA_SUFFIX}" rapids-upload-wheels-to-s3 ${package_dir}/final_dist

ci/test_common.sh

@@ -0,0 +1,155 @@
+#!/bin/bash
+
+# SPDX-FileCopyrightText: Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: BSD-3-Clause
+
+set -euo pipefail
+
+source "$(dirname "$0")/test_utils.sh"
+
+
+##################################### C++ ######################################
+_SERVER_PORT=12345
+
+run_cpp_tests() {
+  RUNTIME_PATH=${CONDA_PREFIX:-./}
+  BINARY_PATH=${RUNTIME_PATH}/bin
+
+  CMD_LINE="timeout 10m ${BINARY_PATH}/gtests/libucxx/UCXX_TEST"
+
+  log_command "${CMD_LINE}"
+  UCX_TCP_CM_REUSEADDR=y ${CMD_LINE}
+}
+
+run_cpp_benchmark() {
+  SERVER_PORT=$1
+  PROGRESS_MODE=$2
+
+  RUNTIME_PATH=${CONDA_PREFIX:-./}
+  BINARY_PATH=${RUNTIME_PATH}/bin
+
+  CMD_LINE_SERVER="timeout 1m ${BINARY_PATH}/benchmarks/libucxx/ucxx_perftest -s 8388608 -r -n 20 -m ${PROGRESS_MODE} -p ${SERVER_PORT}"
+  CMD_LINE_CLIENT="timeout 1m ${BINARY_PATH}/benchmarks/libucxx/ucxx_perftest -s 8388608 -r -n 20 -m ${PROGRESS_MODE} -p ${SERVER_PORT} 127.0.0.1"
+
+  log_command "${CMD_LINE_SERVER}"
+  UCX_TCP_CM_REUSEADDR=y ${CMD_LINE_SERVER} &
+  sleep 1
+
+  log_command "${CMD_LINE_CLIENT}"
+  ${CMD_LINE_CLIENT}
+}
+
+run_cpp_example() {
+  SERVER_PORT=$1
+  PROGRESS_MODE=$2
+
+  RUNTIME_PATH=${CONDA_PREFIX:-./}
+  BINARY_PATH=${RUNTIME_PATH}/bin
+
+  CMD_LINE="timeout 1m ${BINARY_PATH}/examples/libucxx/ucxx_example_basic -m ${PROGRESS_MODE} -p ${SERVER_PORT}"
+
+  log_command "${CMD_LINE}"
+  UCX_TCP_CM_REUSEADDR=y ${CMD_LINE}
+}
+
+run_cpp_port_retry() {
+  MAX_ATTEMPTS=${1}
+  RUN_TYPE=${2}
+  PROGRESS_MODE=${3}
+
+  set +e
+  for attempt in $(seq 1 ${MAX_ATTEMPTS}); do
+    echo "Attempt ${attempt}/${MAX_ATTEMPTS} to run ${RUN_TYPE}"
+
+    _SERVER_PORT=$((_SERVER_PORT + 1))    # Use different ports every time to prevent `Device is busy`
+
+    if [[ "${RUN_TYPE}" == "benchmark" ]]; then
+      run_cpp_benchmark ${_SERVER_PORT} ${PROGRESS_MODE}
+    elif [[ "${RUN_TYPE}" == "example" ]]; then
+      run_cpp_example ${_SERVER_PORT} ${PROGRESS_MODE}
+    else
+      set -e
+      echo "Unknown test type "${RUN_TYPE}""
+      exit 1
+    fi
+
+    LAST_STATUS=$?
+    if [ ${LAST_STATUS} -eq 0 ]; then
+      break;
+    fi
+    sleep 1
+  done
+  set -e
+
+  if [ ${LAST_STATUS} -ne 0 ]; then
+    echo "Failure running benchmark client after ${MAX_ATTEMPTS} attempts"
+    exit $LAST_STATUS
+  fi
+}
+
+
+#################################### Python ####################################
+run_py_tests() {
+  CMD_LINE="timeout 2m pytest -vs python/ucxx/_lib/tests/"
+  log_command "${CMD_LINE}"
+  timeout 2m pytest -vs python/ucxx/_lib/tests/
+}
+
+run_py_tests_async() {
+  PROGRESS_MODE=$1
+  ENABLE_DELAYED_SUBMISSION=$2
+  ENABLE_PYTHON_FUTURE=$3
+  SKIP=$4
+
+  CMD_LINE="UCXPY_PROGRESS_MODE=${PROGRESS_MODE} UCXPY_ENABLE_DELAYED_SUBMISSION=${ENABLE_DELAYED_SUBMISSION} UCXPY_ENABLE_PYTHON_FUTURE=${ENABLE_PYTHON_FUTURE} timeout 20m pytest -vs python/ucxx/_lib_async/tests/ --durations=50"
+
+  if [ $SKIP -ne 0 ]; then
+    echo -e "\e[1;33mSkipping unstable test: ${CMD_LINE}\e[0m"
+  else
+    log_command "${CMD_LINE}"
+    UCXPY_PROGRESS_MODE=${PROGRESS_MODE} UCXPY_ENABLE_DELAYED_SUBMISSION=${ENABLE_DELAYED_SUBMISSION} UCXPY_ENABLE_PYTHON_FUTURE=${ENABLE_PYTHON_FUTURE} timeout 20m pytest -vs python/ucxx/_lib_async/tests/ --durations=50
+  fi
+}
+
+run_py_benchmark() {
+  BACKEND=$1
+  PROGRESS_MODE=$2
+  ASYNCIO_WAIT=$3
+  ENABLE_DELAYED_SUBMISSION=$4
+  ENABLE_PYTHON_FUTURE=$5
+  N_BUFFERS=$6
+  SLOW=$7
+
+  if [ $ASYNCIO_WAIT -ne 0 ]; then
+    ASYNCIO_WAIT="--asyncio-wait"
+  else
+    ASYNCIO_WAIT=""
+  fi
+
+  CMD_LINE="UCXPY_ENABLE_DELAYED_SUBMISSION=${ENABLE_DELAYED_SUBMISSION} UCXPY_ENABLE_PYTHON_FUTURE=${ENABLE_PYTHON_FUTURE} timeout 2m python -m ucxx.benchmarks.send_recv --backend ${BACKEND} -o cupy --reuse-alloc -n 8MiB --n-buffers $N_BUFFERS --progress-mode ${PROGRESS_MODE} ${ASYNCIO_WAIT}"
+
+  # Workaround for https://github.com/rapidsai/ucxx/issues/15
+  CMD_LINE="UCX_KEEPALIVE_INTERVAL=1ms ${CMD_LINE}"
+
+  log_command "${CMD_LINE}"
+  if [ $SLOW -ne 0 ]; then
+    echo -e "\e[1;33mSLOW BENCHMARK: it may seem like a deadlock but will eventually complete.\e[0m"
+  fi
+
+  UCX_KEEPALIVE_INTERVAL=1ms UCXPY_ENABLE_DELAYED_SUBMISSION=${ENABLE_DELAYED_SUBMISSION} UCXPY_ENABLE_PYTHON_FUTURE=${ENABLE_PYTHON_FUTURE} timeout 2m python -m ucxx.benchmarks.send_recv --backend ${BACKEND} -o cupy --reuse-alloc -n 8MiB --n-buffers $N_BUFFERS --progress-mode ${PROGRESS_MODE} ${ASYNCIO_WAIT}
+}
+
+################################## Distributed #################################
+run_distributed_ucxx_tests() {
+  PROGRESS_MODE=$1
+  ENABLE_DELAYED_SUBMISSION=$2
+  ENABLE_PYTHON_FUTURE=$3
+
+  CMD_LINE="UCXPY_PROGRESS_MODE=${PROGRESS_MODE} UCXPY_ENABLE_DELAYED_SUBMISSION=${ENABLE_DELAYED_SUBMISSION} UCXPY_ENABLE_PYTHON_FUTURE=${ENABLE_PYTHON_FUTURE} timeout 10m pytest -vs python/distributed-ucxx/distributed_ucxx/tests/"
+
+  # Workaround for https://github.com/rapidsai/ucxx/issues/15
+  # CMD_LINE="UCX_KEEPALIVE_INTERVAL=1ms ${CMD_LINE}"
+
+  log_command "${CMD_LINE}"
+  UCXPY_PROGRESS_MODE=${PROGRESS_MODE} UCXPY_ENABLE_DELAYED_SUBMISSION=${ENABLE_DELAYED_SUBMISSION} UCXPY_ENABLE_PYTHON_FUTURE=${ENABLE_PYTHON_FUTURE} timeout 10m pytest -vs python/distributed-ucxx/distributed_ucxx/tests/
+}

ci/test_cpp.sh

@@ -6,6 +6,7 @@
 set -euo pipefail
 
 source "$(dirname "$0")/test_utils.sh"
+source "$(dirname "$0")/test_common.sh"
 
 rapids-logger "Create test conda environment"
 . /opt/conda/etc/profile.d/conda.sh
@@ -24,75 +25,6 @@ print_system_stats
 
 BINARY_PATH=${CONDA_PREFIX}/bin
 
-_SERVER_PORT=12345
-
-run_tests() {
-  CMD_LINE="timeout 10m ${BINARY_PATH}/gtests/libucxx/UCXX_TEST"
-
-  log_command "${CMD_LINE}"
-  UCX_TCP_CM_REUSEADDR=y ${CMD_LINE}
-}
-
-run_benchmark() {
-  SERVER_PORT=$1
-  PROGRESS_MODE=$2
-
-  CMD_LINE_SERVER="timeout 1m ${BINARY_PATH}/benchmarks/libucxx/ucxx_perftest -s 8388608 -r -n 20 -m ${PROGRESS_MODE} -p ${SERVER_PORT}"
-  CMD_LINE_CLIENT="timeout 1m ${BINARY_PATH}/benchmarks/libucxx/ucxx_perftest -s 8388608 -r -n 20 -m ${PROGRESS_MODE} -p ${SERVER_PORT} 127.0.0.1"
-
-  log_command "${CMD_LINE_SERVER}"
-  UCX_TCP_CM_REUSEADDR=y ${CMD_LINE_SERVER} &
-  sleep 1
-
-  log_command "${CMD_LINE_CLIENT}"
-  ${CMD_LINE_CLIENT}
-}
-
-run_example() {
-  SERVER_PORT=$1
-  PROGRESS_MODE=$2
-
-  CMD_LINE="timeout 1m ${BINARY_PATH}/examples/libucxx/ucxx_example_basic -m ${PROGRESS_MODE} -p ${SERVER_PORT}"
-
-  log_command "${CMD_LINE}"
-  UCX_TCP_CM_REUSEADDR=y ${CMD_LINE}
-}
-
-run_port_retry() {
-  MAX_ATTEMPTS=${1}
-  RUN_TYPE=${2}
-  PROGRESS_MODE=${3}
-
-  set +e
-  for attempt in $(seq 1 ${MAX_ATTEMPTS}); do
-    echo "Attempt ${attempt}/${MAX_ATTEMPTS} to run ${RUN_TYPE}"
-
-    _SERVER_PORT=$((_SERVER_PORT + 1))    # Use different ports every time to prevent `Device is busy`
-
-    if [[ "${RUN_TYPE}" == "benchmark" ]]; then
-      run_benchmark ${_SERVER_PORT} ${PROGRESS_MODE}
-    elif [[ "${RUN_TYPE}" == "example" ]]; then
-      run_example ${_SERVER_PORT} ${PROGRESS_MODE}
-    else
-      set -e
-      echo "Unknown test type "${RUN_TYPE}""
-      exit 1
-    fi
-
-    LAST_STATUS=$?
-    if [ ${LAST_STATUS} -eq 0 ]; then
-      break;
-    fi
-    sleep 1
-  done
-  set -e
-
-  if [ ${LAST_STATUS} -ne 0 ]; then
-    echo "Failure running benchmark client after ${MAX_ATTEMPTS} attempts"
-    exit $LAST_STATUS
-  fi
-}
-
 rapids-logger "Downloading artifacts from previous jobs"
 CPP_CHANNEL=$(rapids-download-conda-from-s3 cpp)
 
@@ -104,20 +36,20 @@ print_ucx_config
 
 rapids-logger "Run tests with conda package"
 rapids-logger "C++ Tests"
-run_tests
+run_cpp_tests
 
 rapids-logger "C++ Benchmarks"
-# run_port_retry MAX_ATTEMPTS RUN_TYPE PROGRESS_MODE
-run_port_retry 10 "benchmark" "polling"
-run_port_retry 10 "benchmark" "blocking"
-run_port_retry 10 "benchmark" "thread-polling"
-run_port_retry 10 "benchmark" "thread-blocking"
-run_port_retry 10 "benchmark" "wait"
+# run_cpp_port_retry MAX_ATTEMPTS RUN_TYPE PROGRESS_MODE
+run_cpp_port_retry 10 "benchmark" "polling"
+run_cpp_port_retry 10 "benchmark" "blocking"
+run_cpp_port_retry 10 "benchmark" "thread-polling"
+run_cpp_port_retry 10 "benchmark" "thread-blocking"
+run_cpp_port_retry 10 "benchmark" "wait"
 
 rapids-logger "C++ Examples"
-# run_port_retry MAX_ATTEMPTS RUN_TYPE PROGRESS_MODE
-run_port_retry 10 "example" "polling"
-run_port_retry 10 "example" "blocking"
-run_port_retry 10 "example" "thread-polling"
-run_port_retry 10 "example" "thread-blocking"
-run_port_retry 10 "example" "wait"
+# run_cpp_port_retry MAX_ATTEMPTS RUN_TYPE PROGRESS_MODE
+run_cpp_port_retry 10 "example" "polling"
+run_cpp_port_retry 10 "example" "blocking"
+run_cpp_port_retry 10 "example" "thread-polling"
+run_cpp_port_retry 10 "example" "thread-blocking"
+run_cpp_port_retry 10 "example" "wait"