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.

Sign up for GitHub

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
merged 14 commits into from
Aug 22, 2024
Merged

Use CUDA math wheels #2415

Show file tree
Hide file tree
Changes from 13 commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
9ccf865
Use dynamic CUDA libraries
KyleFromNVIDIA Jul 9, 2024
92ee07c
Unconditionall exclude libucp
KyleFromNVIDIA Aug 19, 2024
d6b3714
Fix SKBUILD_CMAKE_ARGS
KyleFromNVIDIA Aug 19, 2024
5bcd36f
Review feedback
KyleFromNVIDIA Aug 19, 2024
f2d4e98
Oops
KyleFromNVIDIA Aug 19, 2024
e2073ff
Syntax
KyleFromNVIDIA Aug 19, 2024
d4e050f
Set rpath on correct target
KyleFromNVIDIA Aug 19, 2024
4adbc97
Remove cufft from rpath
KyleFromNVIDIA Aug 20, 2024
6d772a7
Add matrix entries
KyleFromNVIDIA Aug 21, 2024
d20ae1a
Make use_cuda_wheels the default
KyleFromNVIDIA Aug 21, 2024
896a394
Move cuda wheels into their own dependency set
KyleFromNVIDIA Aug 21, 2024
1e1fe91
Remove unneeded cuda_suffixed entry
KyleFromNVIDIA Aug 21, 2024
c68119a
Add explanatory comments to dependencies.yaml
KyleFromNVIDIA Aug 21, 2024
7f663a0
Add note about CUDA 11's lack of wheels
KyleFromNVIDIA Aug 21, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 20 additions & 1 deletion ci/build_wheel.sh
View file
Open in desktop
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
View file
Open in desktop
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
View file
Open in desktop
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
View file
Open in desktop
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
View file
Open in desktop
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.

rapidsai/devcontainers#382

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