Use CUDA wheels to avoid statically linking CUDA components in our wheels

[vyasr]

In order to achieve manylinux compliance, RAPIDS wheels currently statically link all components of the CTK that they consume. This leads to heavily bloated binaries, especially when the effect is compounded across many packages. Since NVIDIA now publishes wheels containing the CUDA libraries and these libraries have been stress tested by the wheels for various deep learning frameworks (e.g. pytorch now depends on the CUDA wheels), RAPIDS should now do the same to reduce our wheel sizes. This work is a companion to #33 that should probably be tackled afterwards since #33 will reduce the scope of these changes to just the resulting C++ wheels, a meaningful reduction since multiple RAPIDS repos produce multiple wheels. While the goals of this are aligned with #33 and the approach is similar, there are some notable differences because of the way the CUDA wheels are structured. In particular, they are not really designed to be compiled against, only run against. They do generally seem to contain both includes and libraries, which is helpful, but they do not contain any CMake or other packaging metadata, nor do they contain the multiple symlinked copies of libraries (e.g. linker name->soname->library name). The latter is a fundamental limitation of wheels not supporting symlinks, but could cause issues for library discovery using standardized solutions like CMake's FindCUDAToolkit or pkg-config that rely on a specific version of those files existing (AFAICT only the SONAME is present). We should stage work on this in a way that minimizes conflicts with #31 and #33, both of which should facilitate this change. I propose the following, but all of it is open for discussion:

1. Test dynamically linking in a build, then manually installing the CUDA wheels at runtime - Our first attempt should simply be to verify that we are able to interchange the libraries as expected. To achieve this, we will want to do the following:
   1. Pick a repo. raft is probably the best choice here since it's the main entry point for a lot of math libraries in RAPIDS (cuml, cugraph, and cuopt all use it that way) and because it only depends on rmm as a header-only library so there's minimal conflict with the ongoing work to introduce wheel interdependencies.
   2. Turn off static linking in the build, then configure auditwheel to exclude the relevant CUDA math libraries from inclusion. This can be done with the --exclude flag. The resulting wheel should be inspected to verify that all CUDA math libraries have been excluded from the build. Note that (at least for now) we want to continue statically linking the CUDA runtime. This change will likely require some CMake work to decouple static linking of cudart from the static linking of other CUDA libraries.
   3. We will then want to try installing these wheels into a new environment without the necessary CUDA libraries installed. This could be done using a container with a different CUDA version, or on a machine with CUDA drivers installed but relying on e.g. conda for installing the CUDA runtime and libraries. Attempting to import the wheel should give a linker error.
   4. We will then want to try installing the CUDA wheels and verify that we can make things work. The easiest choice at this stage will probably be to just set LD_LIBRARY_PATH
2. Build against the CUDA wheels: In the long run, we would like to be able to build against the CUDA wheels to ensure that we see a consistent build and runtime layout of CUDA files. At present, this likely to be challenging due to some of the layout issues mentioned above. Concretely, I think that we will achieve the most benefit here if we attempt to make things work with the current layout, but do so in a way that makes it manifestly clear why the current layout is difficult to work with. We can then have a more productive discussion with the CUDA wheels teams about changes that we'd like to see (I have already started some of those discussions, but I think it'll be a lot easier to make headway when we have something concrete to discuss). With that in mind, I would suggest that at this stage we focus on writing custom CMake find modules for the CUDA libraries that work when we're building wheels. This will allow us to determine what shortcomings there are with the existing layouts.
3. Layer on top of the C++ wheels - In the long run, RAPIDS Python packages should never need to deal directly with the CUDA wheels. Instead, they should be getting all their CUDA dependencies transitively via the C++ wheels. To achieve this goal, once we've reached this point with these wheels we should rework the above changes on top of the ongoing work to create separate C++ wheels.
4. Figure out a suitable CUDA library loading strategy - The easiest way to make our wheels work with CUDA wheels at runtime is by setting the RPATHs to do a relative lookup of the libraries in the CUDA packages. Ideally I think we would want to push for the CUDA packages to instead manage the libraries via dynamic loading (the way I've set up the RAPIDS C++ wheels) to insulate consumers from the file layout of the wheel, the use of layered environments, etc, but that's probably not going to be an option in the near to intermediate term. Therefore, our options will likely be to set the RPATHs of our binaries directly, or to load the libraries in Python ourselves. The latter is a bit more flexible in that it would allow the potential for coexistence with system-installed CUDA libraries if desired, so for the purpose of e.g. DLFW containers we may still want to go that route. This would be the stage where we try to figure out in general the degree to which we want to support system vs pip-installed CUDA libraries when using our pip wheels.
5. Roll out to all the libraries - Once we reach this point, we can make analogous changes to other RAPIDS packages.

DLFW / devcontainers adjustments

Beta Give feedback

1. update matrix selectors (Use CUDA math wheels raft#2415 (comment))

[jameslamb]

Putting up some of my notes from poking around at this today.

On an x86_64 machine, ran the following.

docker run \
    --rm \
    -it rapidsai/ci-wheel:cuda12.2.2-rockylinux8-py3.10 \
    bash

Looked for the default python interpreter's location for storing architecture-specific libraries:

python -c "import sysconfig; print(sysconfig.get_path('platlib'))"
# /pyenv/versions/3.10.14/lib/python3.10/site-packages

Installed some stuff.

python -m pip install \
    --extra-index-url https://pypi.nvidia.com \
        nvidia-cusparse-cu12 \
        nvidia-cublas-cu12 \
        nvidia-cufft-cu12 \
        'pylibraft-cu12==24.4.*'

Ok, so where did it put all those CUDA libraries?

find \
    /pyenv/versions/3.10.14/lib/python3.10/site-packages \
    -type f \
    -name 'libcu*.so*'

# /pyenv/versions/3.10.14/lib/python3.10/site-packages/nvidia/cublas/lib/libcublasLt.so.12
# /pyenv/versions/3.10.14/lib/python3.10/site-packages/nvidia/cublas/lib/libcublas.so.12
# /pyenv/versions/3.10.14/lib/python3.10/site-packages/nvidia/cufft/lib/libcufft.so.11
# /pyenv/versions/3.10.14/lib/python3.10/site-packages/nvidia/cufft/lib/libcufftw.so.11
# /pyenv/versions/3.10.14/lib/python3.10/site-packages/nvidia/cusparse/lib/libcusparse.so.12

And what about libraft.so?

find \
    /pyenv/versions/3.10.14/lib/python3.10/site-packages \
    -type f \
    -name 'libraft*.so*'

# /pyenv/versions/3.10.14/lib/python3.10/site-packages/pylibraft/libraft.so

So by default, it looks like those CUDA libraries will be at this path relative to libraft.so:

$ORIGIN/../nvidia/{project}/lib/

Where {project} is the top-level name with no lib prefix, e.g. "cublas", "cufft", etc.

I saw at least one example where two libraries that depend on each other are installed together, and one has RUNPATH pointing to the other.

SITE=" /pyenv/versions/3.10.14/lib/python3.10/site-packages"

ldd ${SITE}/nvidia/cublas/lib/libcublas.so.12
# libcublasLt.so.12 => /pyenv/versions/3.10.14/lib/python3.10/site-packages/nvidia/cublas/lib/libcublasLt.so.12 (0x00007f1c23f31000

readelf -d ${SITE}/nvidia/cublas/lib/libcublas.so.12 | grep PATH
#  0x000000000000001d (RUNPATH)            Library runpath: [$ORIGIN]

Tomorrow, I'll try a wheel build of raft customized as described at the top of this issue, and see what I learn.

---

Dumping some other links I've been consulting to get a better understanding of the difference between what happens to be true in the places I'm testing and what we can assume to be true about installation layouts generally.

• "Understanding site-packages directories" (Python Discourse)
• pip install docs (link)
• "Python, Debian, and the install locations" (Filipe Lains blog)
• "How does setuptools determine the appropriate site-packages directory?" (pypa/setuptools#2813)
• wheel file contents spec (packagin.python.org docs)

[bdice]

We’ve mostly explored this for CUDA math library components thus far but we should do the same for nvcomp: https://pypi.org/project/nvidia-nvcomp-cu12/

RAPIDS should shift to using nvcomp wheels as a dependency of our own wheel builds of cudf and kvikio so we do not redistribute nvcomp libraries as part of the cudf and kvikio wheels.