Add configuration to enable/disable NVML diagnostics

distributed/diagnostics/nvml.py

@@ -1,5 +1,7 @@
 import os
 
+import dask
+
 try:
     import pynvml
 except ImportError:
@@ -12,6 +14,11 @@
 
 def init_once():
     global nvmlInitialized, nvmlLibraryNotFound, nvmlOwnerPID
+
+    if dask.config.get("distributed.diagnostics.nvml") is False:
+        nvmlInitialized = False
+        return
+
     if pynvml is None or (nvmlInitialized is True and nvmlOwnerPID == os.getpid()):
         return
 

distributed/diagnostics/tests/test_nvml.py

@@ -4,6 +4,8 @@
 
 pynvml = pytest.importorskip("pynvml")
 
+import dask
+
 from distributed.diagnostics import nvml
 from distributed.utils_test import gen_cluster
 
@@ -19,6 +21,23 @@ def test_one_time():
     assert len(output["name"]) > 0
 
 
+def test_enable_disable_nvml():
+    try:
+        pynvml.nvmlShutdown()
+    except pynvml.NVMLError_Uninitialized:
+        pass
+    else:
+        nvml.nvmlInitialized = False
+
+    with dask.config.set({"distributed.diagnostics.nvml": False}):
+        nvml.init_once()
+        assert nvml.nvmlInitialized is False
+
+    with dask.config.set({"distributed.diagnostics.nvml": True}):
+        nvml.init_once()
+        assert nvml.nvmlInitialized is True
+
+
 def test_1_visible_devices():
     if nvml.device_get_count() < 1:
         pytest.skip("No GPUs available")

distributed/distributed-schema.yaml

@@ -750,6 +750,17 @@ properties:
                       Alternatively, the key can be appended to the cert file
                       above, and this field left blank
 
+      diagnostics:
+        type: object
+        properties:
+          nvml:
+            type: boolean
+            description: |
+              If ``True``, enables GPU diagnostics with NVML. Generally leaving it enabled is
+              not a problem and will be automatically disabled if no GPUs are found in the
+              system, but in certain cases it may be desirable to completely disable NVML
+              diagnostics.
+
       dashboard:
         type: object
         properties:

distributed/distributed.yaml

@@ -195,6 +195,8 @@ distributed:
         key: null
         cert: null
 
+  diagnostics:
+    nvml: True
 
   ###################
   # Bokeh dashboard #