Use `nvml` to detect CUDA driver version

[charliermarsh]

See: #12070 (comment)

[charliermarsh]

I think I can use https://crates.io/crates/nvml-wrapper for this.

[geofft]

(cc @DEKHTIARJonathan)

So, for the specific issue that spawned this, it's because nvidia-smi --query-gpu=driver_version --format=csv is producing multiple lines of output if you have multiple GPUs.

I suspect that this is just an artifact of CSV being a tabular format, and it's printing one line per GPU and fills in the driver_version column with the same value. I can't really imagine how you would have multiple versions of the same driver loaded.

The Python code using NVML is calling nvmlSystemGetDriverVersion() once, a function documented as not being GPU-specific.

If this really is the problem, then let's just read multiple lines of input and check that they're equal to the first one (and throw an error saying to report a bug if they're not).

---

My hesitation around NVML is that you don't need to have it installed in a way that uv can find. You can have it installed on bin/python's rpath, or on a nonstandard LD_LIBRARY_PATH, or with some sort of magical Python library loader calling ctypes, or whatever. You can use a uv with an incompatible glibc—in particular I don't think this change is going to work with our "distroless" uv Docker image which uses musl but is capable of installing manylinux wheels if it sees a glibc present on the system, because uv itself won't be able to load the library.

This is also largely my hesitation around nvidia-smi (which doesn't need to be installed at all, or if it's installed it doesn't need to be on your PATH, or it could be misconfigured in a way where it can't load NVML itself, etc.), which is why we try parsing out of /proc and /sys first. I guess the problem is that WSL does something more magical where /dev/nvidia* relays things to the actual Windows driver in the NT kernel driving the actual hardware, and there's no nvidia driver in the WSL kernel.

I could buy the argument that we ought to swap out nvidia-smi for NVML, or have one be a fallback for the other.

If we could somehow bundle and statically link our own NVML, then that seems like it would be fine. I assume licensing makes that hard.

But—the ioctl that nvmlSystemGetDriverVersion() wraps is documented in NVIDIA/open-gpu-kernel-modules... what if we just called the ioctl? Is that a stable/public interface?

How about changing the logic to try the ioctl first, then try /sys and /proc (which are helpful if for whatever reason you don't have permissions on /dev/nvidia*, or you're in a container and forgot to share the devices into the container but that's a problem you can fix later once you get an error from PyTorch, or whatever), then finally try nvidia-smi and dedup the lines?

[DEKHTIARJonathan]

What do you think @vyasr? Is NVML not a guarantee to be on LD_LIBRARY PATH? I thought it was?

[geofft]

Well, I can say from prior experience that I had a setup where we installed libnvidia-ml.so.1 on a directory added to LD_LIBRARY_PATH by a Python startup wrapper script, and not globally :) IIRC this is one of the drivers that's in the "forward compatibility" package, along with libcuda.so.1. It looks like the CUDA runtime shipped on PyPI doesn't currently contain the forward compatibility stuff, but I think it'd be nice, because it means everything can be in a virtualenv (or virtualenv-like thing).

[geofft]

Looks like NVIDIA/open-gpu-kernel-modules#530 has some links to sample code that speaks the NVIDIA ioctl API if we want to go that route. (To be clear, I'm not sure if we do....)

Getting the driver version out of the ioctl API seems easy; if we also want to get compute capabilities, which it sounds like we should from discussion in the other issue, that's a bit more involved.

[DEKHTIARJonathan]

A long time ago we used to compile and execute some C code to get both libcuda version and SM. Though if nvml is not a guarantee on the system - we will have a problem with https://github.com/wheelnext/nvidia-variant-provider. So let's find the best approach and we can both leverage it.

[DEKHTIARJonathan]

Reading the doc: https://docs.nvidia.com/deploy/pdf/NVML_API_Reference_Guide.pdf

The NVIDIA Management Library (NVML) is a C-based programmatic interface for
monitoring and managing various states within NVIDIA Tesla™ GPUs. It is intended to be a 
platform for building 3rd party applications, and is also the underlying library for the 
NVIDIA-supported nvidia-smi tool. NVML is thread-safe so it is safe to make simultaneous 
NVML calls from multiple threads.

Well, I can say from prior experience that I had a setup where we installed libnvidia-ml.so.1 on a directory added to LD_LIBRARY_PATH by a Python startup wrapper script, and not globally

I do not see how uv can fix it. Both nvml and nvidia-smi that you rely on right now depend on the "visible libcuda" through LD_LIBRARY_PATH ... If someone modifies it at runtime but not install time ... Well ... I really don't see how you can "guess what the user will do at exec time". That will also be true with Wheel Variant.

Currently you are using nvidia-smi - nvml is the low level library underneath. I don't see how using nvidia-smi can do what you want but not nvml

[geofft]

OK, mostly out of curiosity, I implemented #14743 to query the driver version via the ioctl. I am pretty confident in the stability of that particular ABI, as there are comments in the open kernel module saying that it needs to be stable for the sake of older/newer clients, so they can properly figure out that they're incompatible.

We ought to be able to get compute capability out of the same interface, but it will be more complicated and we have to think about ABI stability. (I just learned that https://github.com/google/gvisor/tree/master/pkg/sentry/devices/nvproxy tracks changes in the NVIDIA kernel ABI, because it's used to emulate/relay ioctls from the guest to a host that might be running one of many possible driver versions. I haven't looked in detail but maybe it's helpful if we do go down this road, which, again, I'm not sure we should.)

Currently you are using nvidia-smi - nvml is the low level library underneath. I don't see how using nvidia-smi can do what you want but not nvml

Well, the one case that comes to mind is the our "distroless" uv binary, which (if I remember correctly) is statically linked against musl and therefore cannot open shared libraries. But it can run subprocesses just fine, and we do support using it to create manylinux (glibc-based) venvs. I find this a little weird, to be fair, but it seems to work fine and there are arguments in favor of it.

Also, it doesn't seem that wacky to have a setup where the NVIDIA libraries are accessible within your Python environment and for the execution of nvidia-smi but not for other commands on the system. For something like pip that is running in the context of the Python environment that is about to run PyTorch (or whatever), it makes sense to expect it to be possible to load NVML, but you don't have that guarantee for uv.

(I think it would be sound to move some of this logic into Python interpreter discovery, and say that if Python can't load libnvidia-ml.so, then really nothing is going to work. But that would make things slower, I think.)

[DEKHTIARJonathan]

I'm probably missing some context and depth on uv. I really struggle to understand why that statement

For something like pip that is running in the context of the Python environment that is about to run PyTorch (or whatever), it makes sense to expect it to be possible to load NVML, but you don't have that guarantee for uv.

As I said it's probably lacking some context. I don't understand how the environment is able to load nvml in Python or 'nvidia-smi' but uv can not. That's puzzling to me...

[geofft]

After a bunch of digging I have learned that the ioctl interface does not actually answer what compute capability a device is, that's something NVML presumably figures out in userspace given the device's hardware version from the NV2080_CTRL_CMD_MC_GET_ARCH_INFO operation.

On the one hand, that's a strong argument in favor of using NVML (directly or via nvidia-smi).

On the other hand, the fact that this answer is potentially versioned/variable with respect to the current userspace NVML installation (certainly an older version can fail to tell you what the compute capability is for a newer GPU, even if it's not the case that two different versions of NVML will give different successful answers) worries me, since again I've worked with systems where the installation in /usr is unavailable or out of date (usually thanks to old container images) and you really need the compat libraries from the CUDA installation to make things work.

So at this point I really am leaning towards having a giant mapping from PCI IDs to associated compute capabilities.... although that still leaves the question of whether that will work right on WSL2 if the whole deal is that there is no PCI device on the WSL2 kernel. (But there is a pretty clear ioctl to ask the NVIDIA driver for the cards' PCI IDs, so maybe that will work.)

(Alternatively, it would be nice if NVIDIA could publish some MIT-licensed source code to do the parts of NVML that we and the variant provider need.)

[DEKHTIARJonathan]

You keep saying there's a problem and I still don't understand why you can't use the very library designed to do exactly this. I'm just confused by your messages.

A few clear scenarios where using nvml would not work and why would be useful.

[geofft]

Here are a few scenarios I have in mind:

1. You're following the instructions at https://docs.astral.sh/uv/guides/integration/docker/ to install uv via COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/, which gets you a statically-musl-linked binary. This binary is unable to load any libraries because it is statically linked. Even if it were dynamically linked against musl, it would not be able to load libnvidia-ml.so, which has a dependency on glibc. However, this build of uv is still perfectly capable of either using a system-installed Python linking glibc or downloading one from python-build-standalone, and creating manylinux virtualenvs, where NVML, CUDA, and the Torch wheels will all work fine.
2. You're running a generic container image like debian on a host, where the host NVIDIA drivers are installed normally on the host, but the userspace part of the drivers are not installed in the image in /usr. You can still expect actually doing GPU stuff in Python to work as long as Python code can open libcudart.so.12 (e.g.) and libcuda.so.1. Some specific scenarios where this might be the case:
   1. The container manager (e.g. with the help of nvidia-docker or something) is bind-mounting /usr/local/nvidia/ from the host into the image, and something is setting LD_LIBRARY_PATH at runtime for Python startup to point there, but it is not set in general in the container.
   2. Something is setting LD_LIBRARY_PATH in general, but something is unsetting it before uv gets to run. This type of thing has happened to me much more than I would like.
   3. You have the CUDA runtime and cuda-compat (which includes libcuda.so.1 but not libnvidia-ml.so.1) installed either globally or directly into the Python environment, but you do not have NVML or any other components of the normal userspace drivers installed. You do not actually need NVML installed to do CUDA stuff and such a container image will work fine.

Here's a way to convince yourself of that latter point without building a container image:

$ uvx --with torch,numpy python
Installed 26 packages in 351ms
Python 3.12.3 (main, Jun 18 2025, 17:59:45) [GCC 13.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.cuda.is_available()
True
>>> torch.tensor([[1.]], device='cuda')
tensor([[1.]], device='cuda:0')
>>> 'nvidia-ml' in open('/proc/self/maps').read()
True
>>> exit()
$ echo 'nope' > libnvidia-ml.so.1
$ LD_LIBRARY_PATH=. uvx --with torch,numpy python
Python 3.12.3 (main, Jun 18 2025, 17:59:45) [GCC 13.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.cuda.is_available()
True
>>> torch.tensor([[1.]], device='cuda')
tensor([[1.]], device='cuda:0')
>>> 'nvidia-ml' in open('/proc/self/maps').read()
False

While PyTorch uses NVML if available to query what devices are on the system etc., it will fall back to asking those questions to libcudart, which does not seem to need to load libnvidia-ml to answer those questions.

(By "installed directly into the Python environment" I mean with again either a wrapper setting LD_LIBRARY_PATH or something involving rpaths. The Torch wheels on PyPI use rpaths to find libcudart from the nvidia-cuda-runtime-cu12 etc. wheels, and do not find the global libcudart. While the CUDA runtime wheels on PyPI do not include the cuda-compat libraries, shipping them into Python environments is definitely a thing I have done on my own—specifically to handle the case of wanting to be able to use a newer Python environment + newer physical GPU without having to update the container image. If it is helpful I can build such a container image to demonstrate what it looks like.)

In both of these cases, something like pip, being Python code running inside the Python environment in question, will behave more closely to the behavior of libraries like PyTorch—though, even then, e.g. ctypes.CDLL("libcudart.so.12") will fail (or find the wrong library) because libtorch_cuda.so relies on rpaths to find it, and will also fail if you have not yet installed the CUDA runtime into the environment. But uv is not running inside the Python environment at all, so its behavior can be very different. (For the same reason, "what if it's statically linked against musl and cannot load libraries at all" is something we have to worry about but pip doesn't—if pip is running inside a Python interpreter that is itself statically linked against musl, that interpreter cannot load compiled extension modules like PyTorch, let alone CUDA, so pip doesn't need to answer any questions about GPUs.)

We could decide to move some of the discovery logic into the set of things we query the Python interpreter for. But

• This still has the downside that you need to install some set of userspace library that might not be currently installed to get the answer. Again, it is wholly possible to set up a self-contained virtual environment for GPU-enabled PyTorch that runs on a container image with no NVIDIA libraries installed in the container image, and it's often preferable to do so, so that everything is versioned along with the venv and not dependent on the system (the version of the container image). The only dependency is on the kernel driver and the physical hardware.
• As noted above, even if they are installed, they may not be implicitly importable, and we'd have to add some heuristics around how we think an actual library like libtorch_cuda is going to look for things.
• For uv more so than pip, calling into Python is slow. We do already do some interpreter discovery by calling into Python, but (as I understand it) we try to avoid doing it as much as possible, and we cache it as a function of the Python install, where here it's a function of the system CUDA install, and either of these can change without the other changing.