Repository: dask/distributed

Issue:
GPU Utilization chart doesn't show in vanilla dask
I have a vanilla Dask worker on a machine with a GPU.  I've verified that pynvml is initialized

```python
from distributed.diagnostics import nvml
client.run(nvml.device_get_count)
```
```
{'tls://10.0.51.50:34287': 1,
 'tls://10.0.51.93:35409': 1,
 'tls://10.0.51.96:34807': 1,
 'tls://10.0.55.35:37955': 1,
 'tls://10.0.56.189:33125': 1,
 'tls://10.0.58.50:33941': 1,
 'tls://10.0.59.25:40323': 1,
 'tls://10.0.60.189:36467': 1,
 'tls://10.0.61.129:40929': 1,
 'tls://10.0.62.14:37877': 1}
```

But my chart is still blank

<img width="264" alt="Screenshot 2024-03-09 at 9 55 29 AM" src="https://github.com/dask/distributed/assets/306380/c11c02fc-f3c2-4cb5-a375-ac0edc1c7f4d">

My *guess* is that this chart now only works when something is triggered explicitly, probably something that happens in dask-cuda.  @quasiben @rjzamora @jacobtomlinson I don't suppose you have much experience here?


Discussion:
Here's an MRE for folks.

```bash
conda create -n dask-gpu-dash-test -c conda-forge python ipython dask nvidia-ml-py jupyterlab dask-labextension -y

conda activate dask-gpu-dash-test

dask scheduler --jupyter &

dask worker tcp://127.0.0.1:8786
```

The GPU plots are empty in both the regular dashboard and the Jupyter Lab extension.

<img width="1199" alt="image" src="https://github.com/dask/distributed/assets/1610850/39ef82d6-cca0-4113-b660-342e51c25084">

```python
In [1]: from dask.distributed import Client

In [2]: client = Client("tcp://localhost:8786")

In [3]: from distributed.diagnostics import nvml
   ...: client.run(nvml.device_get_count)
Out[3]: {'tcp://127.0.0.1:43475': 2}
```
Downgrading to `dask=2023.11.0` fixes things (I only found this because the `defaults` channel gave me `2023.11.0` when I accidentally missed `-c conda-forge` off).

```bash
conda create -n dask-gpu-dash-test -c conda-forge python ipython dask=2023.11.0 nvidia-ml-py jupyterlab dask-labextension -y
```

<img width="1163" alt="image" src="https://github.com/dask/distributed/assets/1610850/57263b71-8066-464a-95d5-881c608e7b69">

cc @charlesbluca @pentschev @wence- who have been in this code more recently.
Still triaging -- looks like it may have broken in https://github.com/dask/distributed/pull/8399 .  The dashboard relied on the worker having `gpu` in `extra`:

https://github.com/dask/distributed/blob/043876806d42880db0bf1405f16c514a19def838/distributed/dashboard/components/nvml.py#L126-L130

Resolving patch:
diff --git a/distributed/dashboard/components/nvml.py b/distributed/dashboard/components/nvml.py
--- a/distributed/dashboard/components/nvml.py
+++ b/distributed/dashboard/components/nvml.py
@@ -125,13 +125,11 @@ def update(self):
 
         for idx, ws in enumerate(workers):
             try:
-                info = ws.extra["gpu"]
+                mem_used = ws.metrics["gpu_memory_used"]
+                mem_total = ws.metrics["gpu-memory-total"]
+                u = ws.metrics["gpu_utilization"]
             except KeyError:
                 continue
-            metrics = ws.metrics["gpu"]
-            u = metrics["utilization"]
-            mem_used = metrics["memory-used"]
-            mem_total = info["memory-total"]
             memory_max = max(memory_max, mem_total)
             memory_total += mem_total
             utilization.append(int(u))
diff --git a/distributed/system_monitor.py b/distributed/system_monitor.py
--- a/distributed/system_monitor.py
+++ b/distributed/system_monitor.py
@@ -129,6 +129,7 @@ def __init__(
             gpu_extra = nvml.one_time()
             self.gpu_name = gpu_extra["name"]
             self.gpu_memory_total = gpu_extra["memory-total"]
+            self.quantities["gpu-memory-total"] = deque(maxlen=1)
             self.quantities["gpu_utilization"] = deque(maxlen=maxlen)
             self.quantities["gpu_memory_used"] = deque(maxlen=maxlen)
         else:
@@ -207,6 +208,7 @@ def update(self) -> dict[str, Any]:
 
         if self.gpu_name:
             gpu_metrics = nvml.real_time()
+            result["gpu-memory-total"] = self.gpu_memory_total
             result["gpu_utilization"] = gpu_metrics["utilization"]
             result["gpu_memory_used"] = gpu_metrics["memory-used"]
 


Test patch: