Repository: dask/distributed

Issue:
Incorrect GPU Memory Totals 
On the `/individual-gpu-memory` page, the GPU memory title value far exceeds what is used and available:

<img width="1470" alt="Screen Shot 2020-05-17 at 10 50 09 PM" src="https://user-images.githubusercontent.com/1403768/82170031-4f84bb00-9891-11ea-997d-b065d4d213bb.png">

This was executed on a DGX1 with 256GBs of 
GPU memory, not 2TBs.  Perhaps a counting issue within https://github.com/dask/distributed/blob/1bcbaeee4a92d57fb2d399c9580fce6480a6d7c9/distributed/dashboard/components/nvml.py#L136-L153


Discussion:
Out of curiosity how much system memory does that DGX1 have? Is that where the 2TB number is coming from or is it more random than that?
> bzaitlen@dgx12:~$ free -m
              total        used        free      shared  buff/cache   available
Mem:        1031977       14079      291489          90      726408     1011670
Swap:             0           0           0

1TB 
It looks like there is some repetition.  My guess is that every worker is
looking at all of the GPUs on the machine, and so you're getting 8x8
records.  They should probably filter by CUDA_VISIBLE_DEVICES or something.

On Mon, May 18, 2020 at 5:06 AM Benjamin Zaitlen <notifications@github.com>
wrote:

> bzaitlen@dgx12:~$ free -m
> total used free shared buff/cache available
> Mem: 1031977 14079 291489 90 726408 1011670
> Swap: 0 0 0
>
> 1TB
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/distributed/issues/3808#issuecomment-630137208>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AACKZTAI25ZMEGDBH7AIBYDRSEQEDANCNFSM4NDV5RLQ>
> .
>


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
@@ -131,29 +131,23 @@ def update(self):
             memory_total = 0
             memory_max = 0
             worker = []
-            i = 0
 
-            for ws in workers:
+            for idx, ws in enumerate(workers):
                 try:
                     info = ws.extra["gpu"]
                 except KeyError:
                     continue
                 metrics = ws.metrics["gpu"]
-                for j, (u, mem_used, mem_total) in enumerate(
-                    zip(
-                        metrics["utilization"],
-                        metrics["memory-used"],
-                        info["memory-total"],
-                    )
-                ):
-                    memory_max = max(memory_max, mem_total)
-                    memory_total += mem_total
-                    utilization.append(int(u))
-                    memory.append(mem_used)
-                    worker.append(ws.address)
-                    gpu_index.append(j)
-                    y.append(i)
-                    i += 1
+                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))
+                memory.append(mem_used)
+                worker.append(ws.address)
+                gpu_index.append(idx)
+                y.append(idx)
 
             memory_text = [format_bytes(m) for m in memory]
 
diff --git a/distributed/diagnostics/nvml.py b/distributed/diagnostics/nvml.py
--- a/distributed/diagnostics/nvml.py
+++ b/distributed/diagnostics/nvml.py
@@ -1,28 +1,48 @@
+import os
 import pynvml
 
-handles = None
+nvmlInit = None
+
+
+def init_once():
+    global nvmlInit
+    if nvmlInit is not None:
+        return
+
+    from pynvml import nvmlInit as _nvmlInit
+
+    nvmlInit = _nvmlInit
+    nvmlInit()
 
 
 def _pynvml_handles():
-    global handles
-    if handles is None:
-        pynvml.nvmlInit()
-        count = pynvml.nvmlDeviceGetCount()
-        handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in range(count)]
-    return handles
+    count = pynvml.nvmlDeviceGetCount()
+    try:
+        cuda_visible_devices = [
+            int(idx) for idx in os.environ.get("CUDA_VISIBLE_DEVICES", "").split(",")
+        ]
+    except ValueError:
+        # CUDA_VISIBLE_DEVICES is not set
+        cuda_visible_devices = False
+    if not cuda_visible_devices:
+        cuda_visible_devices = list(range(count))
+    gpu_idx = cuda_visible_devices[0]
+    return pynvml.nvmlDeviceGetHandleByIndex(gpu_idx)
 
 
 def real_time():
-    handles = _pynvml_handles()
+    init_once()
+    h = _pynvml_handles()
     return {
-        "utilization": [pynvml.nvmlDeviceGetUtilizationRates(h).gpu for h in handles],
-        "memory-used": [pynvml.nvmlDeviceGetMemoryInfo(h).used for h in handles],
+        "utilization": pynvml.nvmlDeviceGetUtilizationRates(h).gpu,
+        "memory-used": pynvml.nvmlDeviceGetMemoryInfo(h).used,
     }
 
 
 def one_time():
-    handles = _pynvml_handles()
+    init_once()
+    h = _pynvml_handles()
     return {
-        "memory-total": [pynvml.nvmlDeviceGetMemoryInfo(h).total for h in handles],
-        "name": [pynvml.nvmlDeviceGetName(h).decode() for h in handles],
+        "memory-total": pynvml.nvmlDeviceGetMemoryInfo(h).total,
+        "name": pynvml.nvmlDeviceGetName(h).decode(),
     }


Test patch:
diff --git a/distributed/diagnostics/tests/test_nvml.py b/distributed/diagnostics/tests/test_nvml.py
new file mode 100644
--- /dev/null
+++ b/distributed/diagnostics/tests/test_nvml.py
@@ -0,0 +1,34 @@
+import pytest
+import os
+
+pynvml = pytest.importorskip("pynvml")
+
+from distributed.diagnostics import nvml
+
+
+def test_one_time():
+    output = nvml.one_time()
+    assert "memory-total" in output
+    assert "name" in output
+
+    assert len(output["name"]) > 0
+
+
+def test_1_visible_devices():
+    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+    output = nvml.one_time()
+    assert len(output["memory-total"]) == 1
+
+
+@pytest.mark.parametrize("CVD", ["1,0", "0,1"])
+def test_2_visible_devices(CVD):
+    os.environ["CUDA_VISIBLE_DEVICES"] = CVD
+    idx = int(CVD.split(",")[0])
+
+    h = nvml._pynvml_handles()
+    h2 = pynvml.nvmlDeviceGetHandleByIndex(idx)
+
+    s = pynvml.nvmlDeviceGetSerial(h)
+    s2 = pynvml.nvmlDeviceGetSerial(h2)
+
+    assert s == s2