Replace nvgpu with nvidia-ml-py (#2160)

Resolves #2159

* As `nvgpu` is no longer maintained and uses `pynvml`, which directly
tells the user to use `nvidia-ml-py` instead at import

> The pynvml package is deprecated. Please install nvidia-ml-py instead.
If you did not install pynvml directly, please report this to the
maintainers of the package that installed pynvml for you."

drop `nvgpu` and replace its `nvgpu.gpu_info()` call with a single
function using `nvidia-ml-py` (which uses the `pynvml` namespace).
* Place a lower bound on `nvidia-ml-py` of `12.535.77`, which was the
first release to support `nvmlMemory_v2` which properly accounts for
system-reserved memory.
* Remove all mentions of `nvgpu` in other areas of the codebase and
replace them with `nvidia-ml-py`, except for `publications/` as this is
historical information.
- Do not add `nvidia-ml-py` to `dependabot.yml` as pinning this tightly
[will cause installation
issues](https://iscinumpy.dev/post/bound-version-constraints/),
especially with NVIDIA libraries.

---

Example:

```console
$ nvidia-smi --version
NVIDIA-SMI version  : 590.48.01
NVML version        : 590.48
DRIVER version      : 590.48.01
CUDA Version        : 13.1
```

* **On `main`** (2471d55)

```console
$ uv venv main
$ . main/bin/activate
$ uv pip install .
$ python
Python 3.13.7 (main, Sep 18 2025, 19:47:49) [Clang 20.1.4 ] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import desc
>>> desc.set_device("gpu")
/tmp/DESC/main/lib/python3.13/site-packages/nvgpu/__init__.py:8: SyntaxWarning: invalid escape sequence '\('
  gpu_infos = [re.match('GPU ([0-9]+): (.+?) \(UUID: ([^)]+)\)', gpu) for gpu in gpus]
>>> import os
>>> os.environ["CUDA_VISIBLE_DEVICES"]
'0'
>>> 
```
```console
$ python -c 'import nvgpu; print(nvgpu.gpu_info())'
[{'index': '0', 'type': 'NVIDIA GeForce RTX 4060 Laptop GPU', 'uuid': 'GPU-7fef9454-d8d1-86cf-c4b3-e2fd5e35e862', 'mem_used': 8, 'mem_total': 8188, 'mem_used_percent': 0.09770395701025891}]
```

* **This PR**

```console
$ uv venv
$ . .venv/bin/activate
$ uv pip install .
$ python
Python 3.13.7 (main, Sep 18 2025, 19:47:49) [Clang 20.1.4 ] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import desc
>>> desc.set_device("gpu")
>>> import os
>>> os.environ["CUDA_VISIBLE_DEVICES"]
'0'
>>> 
```

As I made it a guarded import I can't directly import it from `desc`,
but is the same code

```python
# _implementation.py
from pynvml import (
    nvmlDeviceGetCount,
    nvmlDeviceGetHandleByIndex,
    nvmlDeviceGetMemoryInfo,
    nvmlDeviceGetName,
    nvmlDeviceGetUUID,
    nvmlInit,
    nvmlMemory_v2,
    nvmlShutdown,
)


def _gpu_info():
    """Equivalent to nvgpu.gpu_info() using nvidia-ml-py."""
    nvmlInit()
    try:
        info = []
        for device_idx in range(nvmlDeviceGetCount()):
            handle = nvmlDeviceGetHandleByIndex(device_idx)
            mem = nvmlDeviceGetMemoryInfo(handle, version=nvmlMemory_v2)
            _bytes_to_mib = 1024 * 1024
            mem_used = mem.used // _bytes_to_mib
            mem_total = mem.total // _bytes_to_mib
            info.append(
                {
                    "index": str(device_idx),
                    "type": nvmlDeviceGetName(handle),
                    "uuid": nvmlDeviceGetUUID(handle),
                    "mem_used": mem_used,
                    "mem_total": mem_total,
                    "mem_used_percent": 100.0 * mem_used / mem_total,
                }
            )
        return info
    finally:
        nvmlShutdown()


if __name__ == "__main__":
    print(_gpu_info())
```

so

```console
$ python ./_implementation.py
[{'index': '0', 'type': 'NVIDIA GeForce RTX 4060 Laptop GPU', 'uuid': 'GPU-7fef9454-d8d1-86cf-c4b3-e2fd5e35e862', 'mem_used': 7, 'mem_total': 8188, 'mem_used_percent': 0.08549096238397655}]
```

So the memory consumption is effectively the same (good), and an
unmaintained dependency can be replaced with a maintained one.

---------

Co-authored-by: Daniel Dudt <33005725+ddudt@users.noreply.github.com>

Author: matthewfeickert

.github/dependabot.yml

@@ -29,7 +29,6 @@ updates:
       minor_packages:
         patterns:
           - "colorama"
-          - "nvgpu"
           - "psutil"
           - "pylatexenc"
           - "termcolor"

CHANGELOG.md

@@ -10,6 +10,8 @@ New Features
 - Method to compute bounce integrals in batches is now added to the public API ``Bounce2D.batch``.
 - Initiated deprecation of ``Bounce2D.compute_fieldline_length`` in favor of ``eq.compute("V_psi")``.
     - The quadrature resolution in ``Bounce2D.compute_fieldline_length`` now corresponds to the resolution over a single field period instead of the resolution over a toroidal transit.
+- Modernizes dependencies to use [``nvidia-ml-py``](https://pypi.org/project/nvidia-ml-py/) in place of [``nvgpu``](https://github.com/rossumai/nvgpu).
+  If you are updating an existing software environment uninstall ``pynvml`` first and then reinstall the dependencies to correctly get ``nvidia-ml-py``.
 
 Bug Fixes
 

desc/__init__.py

@@ -90,10 +90,46 @@ def set_device(kind="cpu", gpuid=None):
     if kind == "gpu":
         # Set CUDA_DEVICE_ORDER so the IDs assigned by CUDA match those from nvidia-smi
         os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
-        import nvgpu
+        # pynvml namespace is exposed through nvidia-ml-py
+        from pynvml import (
+            nvmlDeviceGetCount,
+            nvmlDeviceGetHandleByIndex,
+            nvmlDeviceGetMemoryInfo,
+            nvmlDeviceGetName,
+            nvmlDeviceGetUUID,
+            nvmlInit,
+            nvmlMemory_v2,
+            nvmlShutdown,
+        )
+
+        def _gpu_info():
+            """Equivalent to nvgpu.gpu_info() using nvidia-ml-py."""
+            nvmlInit()
+            try:
+                info = []
+                for device_idx in range(nvmlDeviceGetCount()):
+                    handle = nvmlDeviceGetHandleByIndex(device_idx)
+                    # Use nvmlMemory_v2 to account for system-reserved memory
+                    mem = nvmlDeviceGetMemoryInfo(handle, version=nvmlMemory_v2)
+                    _bytes_to_mib = 1024 * 1024
+                    mem_used = mem.used // _bytes_to_mib
+                    mem_total = mem.total // _bytes_to_mib
+                    info.append(
+                        {
+                            "index": str(device_idx),
+                            "type": nvmlDeviceGetName(handle),
+                            "uuid": nvmlDeviceGetUUID(handle),
+                            "mem_used": mem_used,
+                            "mem_total": mem_total,
+                            "mem_used_percent": 100.0 * mem_used / mem_total,
+                        }
+                    )
+                return info
+            finally:
+                nvmlShutdown()
 
         try:
-            devices = nvgpu.gpu_info()
+            devices = _gpu_info()
         except FileNotFoundError:
             devices = []
         if len(devices) == 0:

requirements.txt

@@ -10,7 +10,7 @@ matplotlib >= 3.7.3, <= 3.10.8
 mpmath >= 1.0.0, <= 1.4.1
 netcdf4 >= 1.5.4, !=1.7.4, <= 1.7.5
 numpy >= 1.20.0, <= 2.5
-nvgpu <= 0.10.0
+nvidia-ml-py >=12.535.77
 optax < 0.3.0
 orthax < 0.3
 plotly >= 5.16, <= 6.7.0