Add GPU executor if GPU is present

See dask#5084
mrocklin · Jul 26, 2021 · b4a2363 · b4a2363
1 parent 145e8aa
commit b4a2363
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/distributed/worker.py b/distributed/worker.py
@@ -635,6 +635,15 @@ def __init__(
             "offload": utils._offload_executor,
             "actor": ThreadPoolExecutor(1, thread_name_prefix="Dask-Actor-Threads"),
         }
+        try:
+            import pynvml
+        except ImportError:
+            pass
+        else:
+            if pynvml.nvmlDeviceGetCount():
+                self.executors["gpu"] = ThreadPoolExecutor(
+                    1, thread_name_prefix="Dask-GPU-Threads"
+                )
 
         # Find the default executor
         if executor == "offload":