Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use CUDA math wheels #2415

Merged
21 changes: 20 additions & 1 deletion ci/build_wheel.sh
Original file line number Diff line number Diff line change
Expand Up @@ -20,10 +20,29 @@ rapids-generate-version > VERSION

cd "${package_dir}"

case "${RAPIDS_CUDA_VERSION}" in
12.*)
EXCLUDE_ARGS=(
--exclude "libcublas.so.12"
--exclude "libcublasLt.so.12"
--exclude "libcurand.so.10"
--exclude "libcusolver.so.11"
--exclude "libcusparse.so.12"
--exclude "libnvJitLink.so.12"
--exclude "libucp.so.0"
)
;;
11.*)
EXCLUDE_ARGS=(
--exclude "libucp.so.0"
)
;;
esac

# Hardcode the output dir
python -m pip wheel . -w dist -vvv --no-deps --disable-pip-version-check

mkdir -p final_dist
python -m auditwheel repair -w final_dist --exclude "libucp.so.0" dist/*
python -m auditwheel repair -w final_dist "${EXCLUDE_ARGS[@]}" dist/*

RAPIDS_PY_WHEEL_NAME="${underscore_package_name}_${RAPIDS_PY_CUDA_SUFFIX}" rapids-upload-wheels-to-s3 final_dist
11 changes: 10 additions & 1 deletion ci/build_wheel_pylibraft.sh
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,16 @@

set -euo pipefail

case "${RAPIDS_CUDA_VERSION}" in
12.*)
EXTRA_CMAKE_ARGS=";-DUSE_CUDA_MATH_WHEELS=ON"
;;
11.*)
EXTRA_CMAKE_ARGS=";-DUSE_CUDA_MATH_WHEELS=OFF"
;;
esac

# Set up skbuild options. Enable sccache in skbuild config options
export SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF;-DFIND_RAFT_CPP=OFF"
export SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF;-DFIND_RAFT_CPP=OFF${EXTRA_CMAKE_ARGS}"

ci/build_wheel.sh pylibraft python/pylibraft
30 changes: 30 additions & 0 deletions dependencies.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,7 @@ files:
extras:
table: project
includes:
- cuda_wheels
- run_pylibraft
py_test_pylibraft:
output: pyproject
Expand Down Expand Up @@ -377,6 +378,35 @@ dependencies:
- *libcusolver114
- *libcusparse_dev114
- *libcusparse114
cuda_wheels:
specific:
- output_types: pyproject
matrices:
- matrix:
cuda: "12.*"
use_cuda_wheels: "true"
packages:
KyleFromNVIDIA marked this conversation as resolved.
Show resolved Hide resolved
- nvidia-cublas-cu12
- nvidia-curand-cu12
- nvidia-cusolver-cu12
- nvidia-cusparse-cu12
- matrix:
cuda: "11.*"
use_cuda_wheels: "true"
packages:
# if use_cuda_wheels=false is provided, do not add dependencies on any CUDA wheels
# (e.g. for DLFW and pip devcontainers)
- matrix:
use_cuda_wheels: "false"
packages:
# if no matching matrix selectors passed, list the unsuffixed packages
# (just as a source of documentation, as this populates pyproject.toml in source control)
- matrix:
packages:
- nvidia-cublas
- nvidia-curand
- nvidia-cusolver
- nvidia-cusparse

depends_on_cupy:
common:
Expand Down
18 changes: 18 additions & 0 deletions python/pylibraft/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ project(
option(FIND_RAFT_CPP "Search for existing RAFT C++ installations before defaulting to local files"
ON
)
option(USE_CUDA_MATH_WHEELS "Use the CUDA math wheels instead of the system libraries" OFF)

# If the user requested it we attempt to find RAFT.
if(FIND_RAFT_CPP)
Expand All @@ -48,15 +49,32 @@ endif()
include(rapids-cython-core)

if(NOT raft_FOUND)
find_package(CUDAToolkit REQUIRED)

set(BUILD_TESTS OFF)
set(BUILD_PRIMS_BENCH OFF)
set(BUILD_ANN_BENCH OFF)
set(RAFT_COMPILE_LIBRARY ON)
set(CUDA_STATIC_RUNTIME ON)
set(CUDA_STATIC_MATH_LIBRARIES ON)
if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.0)
KyleFromNVIDIA marked this conversation as resolved.
Show resolved Hide resolved
set(CUDA_STATIC_MATH_LIBRARIES OFF)
elseif(USE_CUDA_MATH_WHEELS)
message(FATAL_ERROR "Cannot use CUDA math wheels with CUDA < 12.0")
endif()

add_subdirectory(../../cpp raft-cpp EXCLUDE_FROM_ALL)

if(NOT CUDA_STATIC_MATH_LIBRARIES AND USE_CUDA_MATH_WHEELS)
set_property(TARGET raft_lib PROPERTY INSTALL_RPATH
"$ORIGIN/../nvidia/cublas/lib"
"$ORIGIN/../nvidia/curand/lib"
"$ORIGIN/../nvidia/cusolver/lib"
"$ORIGIN/../nvidia/cusparse/lib"
"$ORIGIN/../nvidia/nvjitlink/lib"
)
endif()

# When building the C++ libraries from source we must copy libraft.so alongside the
# pairwise_distance and random Cython libraries TODO: when we have a single 'compiled' raft
# library, we shouldn't need this
Expand Down
6 changes: 5 additions & 1 deletion python/pylibraft/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,10 @@ requires-python = ">=3.9"
dependencies = [
"cuda-python",
"numpy>=1.23,<2.0a0",
"nvidia-cublas",
"nvidia-curand",
"nvidia-cusolver",
"nvidia-cusparse",
"rmm==24.10.*,>=0.0.0a0",
] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
classifiers = [
Expand Down Expand Up @@ -124,7 +128,7 @@ requires = [
"rmm==24.10.*,>=0.0.0a0",
] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
dependencies-file = "../../dependencies.yaml"
matrix-entry = "cuda_suffixed=true"
matrix-entry = "cuda_suffixed=true;use_cuda_wheels=true"
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When we first talked about this idea, I was thinking usa_cuda_wheels should default to false, so that DLFW / pip devcontainers don't even have to think about it.

And then build scripts here in RAPIDS repos that prepare wheels for distribution would opt in to it by passing use_cuda_wheels=true through config settings.

As I type that, I've come around to the idea that that's not the right balance... one change in DLFW + one change in pip devcontainers is preferable to spreading this complexity around every build_wheel*.sh across RAPIDS.

So I'm good with this, but once we merge a PR with this pattern, please do go put up the corresponding changes in RAPIDS DLFW and devcontainers (devcontainers example: rapidsai/devcontainers#365).

Unused matrix selectors are harmless (I think), so use_cuda_wheels=false could be added unconditionally in those builds even if not all repos have matrices using that.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Will do

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, I added a tasklist with this in it to rapidsai/build-planning#35

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One thing I'd like to point out is that devcontainers generates a requirements.txt file, while this dependency set only has pyproject.toml output, so it won't actually get the wheels anyway. Still, I'll submit a PR just in case that ever changes.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

devcontainers generates a requirements.txt file, while this dependency set only has pyproject.toml output

I think your conclusion that the devcontainers PR isn't strictly necessary and is just defensive is correct. Everything below is for the benefit of others reading along or finding this from search, and doesn't change this PR.


The claim that this having only output_type: pyproject means we don't need to care about it for pip devcontainers isn't true by itself.

The runtime dependencies in the wheel metadata don't strictly have to be correct because in pip devcontainers, the package will be built with pip install --no-deps --no-build-isolation.

(devcontainers code link)

That --no-deps does mean that the runtime dependencies listed in the wheel's metadata don't have to be correct.

So it's output_type: pyproject and the fact that these CUDA wheels are only runtime dependencies that together make this change not disruptive for RAPIDS devcontainers.

HOWEVER... if the project metadata lists any build dependencies that don't already exist in the environment where that pip install runs, that build will fail.

So any list in dependencies.yaml which had only output_type: pyproject which was used to populate wheel build requirements could affect pip devcontainers, because those requirements would need to all be met by the pre-created build environment.


[tool.pytest.ini_options]
filterwarnings = [
Expand Down
Loading