From 9e65a042be44761a21abd2f07e10b0d886ffaa22 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Wed, 1 May 2024 04:33:33 -0700 Subject: [PATCH] Switch to `ucx_wheel` package --- ci/build_wheel.sh | 11 +++++++++-- ci/test_wheel_raft_dask.sh | 6 ++++-- dependencies.yaml | 8 ++++++++ python/raft-dask/pyproject.toml | 2 ++ python/raft-dask/raft_dask/__init__.py | 12 +++++++++++- 5 files changed, 34 insertions(+), 5 deletions(-) diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh index 146f4844e8..81baf9661b 100755 --- a/ci/build_wheel.sh +++ b/ci/build_wheel.sh @@ -7,6 +7,10 @@ package_name=$1 package_dir=$2 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/libuc* + source rapids-configure-sccache source rapids-date-string @@ -15,6 +19,8 @@ git_commit=$(git rev-parse HEAD) RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})" +libucx_wheelhouse=$(RAPIDS_PY_WHEEL_NAME="${RAPIDS_PY_CUDA_SUFFIX}" rapids-get-pr-wheel-artifact ucx-wheels 1 cpp) + # 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}" @@ -38,6 +44,7 @@ fi if [[ ${package_name} == "raft-dask" ]]; then sed -r -i "s/pylibraft==(.*)\"/pylibraft${PACKAGE_CUDA_SUFFIX}==\1${alpha_spec}\"/g" ${pyproject_file} + sed -r -i "s/libucx(.*)\"/libucx${PACKAGE_CUDA_SUFFIX}\1${alpha_spec}\"/g" ${pyproject_file} sed -r -i "s/ucx-py==(.*)\"/ucx-py${PACKAGE_CUDA_SUFFIX}==\1${alpha_spec}\"/g" ${pyproject_file} sed -r -i "s/rapids-dask-dependency==(.*)\"/rapids-dask-dependency==\1${alpha_spec}\"/g" ${pyproject_file} sed -r -i "s/dask-cuda==(.*)\"/dask-cuda==\1${alpha_spec}\"/g" ${pyproject_file} @@ -54,9 +61,9 @@ fi cd "${package_dir}" # Hardcode the output dir -python -m pip wheel . -w dist -vvv --no-deps --disable-pip-version-check +python -m pip wheel . -w dist -vvv --no-deps --disable-pip-version-check --find-links ${libucx_wheelhouse} mkdir -p final_dist -python -m auditwheel repair -w final_dist dist/* +python -m auditwheel repair -w final_dist --exclude "libucp.so.0" dist/* RAPIDS_PY_WHEEL_NAME="${underscore_package_name}_${RAPIDS_PY_CUDA_SUFFIX}" rapids-upload-wheels-to-s3 final_dist diff --git a/ci/test_wheel_raft_dask.sh b/ci/test_wheel_raft_dask.sh index b4ebee8008..9c2621e38c 100755 --- a/ci/test_wheel_raft_dask.sh +++ b/ci/test_wheel_raft_dask.sh @@ -6,13 +6,15 @@ set -euo pipefail mkdir -p ./dist RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})" RAPIDS_PY_WHEEL_NAME="raft_dask_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 ./dist +libucx_wheelhouse=$(RAPIDS_PY_WHEEL_NAME="${RAPIDS_PY_CUDA_SUFFIX}" rapids-get-pr-wheel-artifact ucx-wheels 1 cpp) +ucxx_wheelhouse=$(RAPIDS_PY_WHEEL_NAME="ucxx_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 ./local-raft-dask-dep) + # Download the pylibraft built in the previous step RAPIDS_PY_WHEEL_NAME="pylibraft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 ./local-pylibraft-dep python -m pip install --no-deps ./local-pylibraft-dep/pylibraft*.whl -# echo to expand wildcard before adding `[extra]` requires for pip -python -m pip install $(echo ./dist/raft_dask*.whl)[test] +python -m pip install "raft_dask-${RAPIDS_PY_CUDA_SUFFIX}[test]>=0.0.0a0" --find-links dist/ --find-links "${libucx_wheelhouse}" --find-links "${ucxx_wheelhouse}" python -m pytest ./python/raft-dask/raft_dask/test python -m pytest ./python/raft-dask/raft_dask/test --run_ucx diff --git a/dependencies.yaml b/dependencies.yaml index 7d815a8dfe..c00e9bb8d9 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -92,6 +92,7 @@ files: table: build-system includes: - build + - ucx_wheel py_run_raft_dask: output: pyproject pyproject_dir: python/raft-dask @@ -107,6 +108,7 @@ files: key: test includes: - test_python_common + - ucx_wheel - depends_on_distributed_ucxx py_build_raft_ann_bench: output: pyproject @@ -503,3 +505,9 @@ dependencies: packages: - distributed-ucxx-cu11==0.38.* - {matrix: null, packages: [*distributed_ucxx_conda]} + ucx_wheel: + common: + - output_types: [requirements, pyproject] + packages: + # TODO: Update version + - libucx==1.14.1 diff --git a/python/raft-dask/pyproject.toml b/python/raft-dask/pyproject.toml index a7b77a1139..cfb82c16c9 100644 --- a/python/raft-dask/pyproject.toml +++ b/python/raft-dask/pyproject.toml @@ -18,6 +18,7 @@ build-backend = "scikit_build_core.build" requires = [ "cmake>=3.26.4", "cython>=3.0.0", + "libucx==1.14.1", "ninja", "scikit-build-core[pyproject]>=0.7.0", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. @@ -52,6 +53,7 @@ classifiers = [ [project.optional-dependencies] test = [ "distributed-ucxx==0.38.*", + "libucx==1.14.1", "pytest-cov", "pytest==7.*", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. diff --git a/python/raft-dask/raft_dask/__init__.py b/python/raft-dask/raft_dask/__init__.py index fbbaee4118..19a037ae75 100644 --- a/python/raft-dask/raft_dask/__init__.py +++ b/python/raft-dask/raft_dask/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2020-2023, NVIDIA CORPORATION. +# Copyright (c) 2020-2024, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,3 +14,13 @@ # from raft_dask._version import __git_commit__, __version__ + +# 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