Add GPU executor if GPU is present

distributed/tests/test_worker.py

@@ -20,6 +20,7 @@
 from dask import delayed
 from dask.system import CPU_COUNT
 
+import distributed
 from distributed import (
     Client,
     Nanny,
@@ -33,6 +34,7 @@
 from distributed.comm.tcp import TCPBackend
 from distributed.compatibility import LINUX, WINDOWS
 from distributed.core import CommClosedError, Status, rpc
+from distributed.diagnostics import nvml
 from distributed.diagnostics.plugin import PipInstall
 from distributed.metrics import time
 from distributed.scheduler import Scheduler
@@ -1979,16 +1981,16 @@ def get_thread_name():
     async with Worker(
         s.address,
         nthreads=2,
-        executor={"GPU": ThreadPoolExecutor(1, thread_name_prefix="Dask-GPU-Threads")},
+        executor={"foo": ThreadPoolExecutor(1, thread_name_prefix="Dask-Foo-Threads")},
     ):
         futures = []
         with dask.annotate(executor="default"):
             futures.append(c.submit(get_thread_name, pure=False))
-        with dask.annotate(executor="GPU"):
+        with dask.annotate(executor="foo"):
             futures.append(c.submit(get_thread_name, pure=False))
         default_result, gpu_result = await c.gather(futures)
         assert "Dask-Default-Threads" in default_result
-        assert "Dask-GPU-Threads" in gpu_result
+        assert "Dask-Foo-Threads" in gpu_result
 
 
 @gen_cluster(client=True)
@@ -2055,6 +2057,17 @@ async def test_process_executor_raise_exception(c, s, a, b):
             await future
 
 
+@pytest.mark.gpu
+@gen_cluster(client=True, nthreads=[("127.0.0.1", 1)])
+async def test_gpu_executor(c, s, w):
+    if nvml.device_get_count() > 0:
+        e = w.executors["gpu"]
+        assert isinstance(e, distributed.threadpoolexecutor.ThreadPoolExecutor)
+        assert e._max_workers == 1
+    else:
+        assert "gpu" not in w.executors
+
+
 def assert_task_states_on_worker(expected, worker):
     for dep_key, expected_state in expected.items():
         assert dep_key in worker.tasks, (worker.name, dep_key, worker.tasks)

distributed/worker.py

@@ -635,6 +635,10 @@ def __init__(
             "offload": utils._offload_executor,
             "actor": ThreadPoolExecutor(1, thread_name_prefix="Dask-Actor-Threads"),
         }
+        if nvml.device_get_count() > 0:
+            self.executors["gpu"] = ThreadPoolExecutor(
+                1, thread_name_prefix="Dask-GPU-Threads"
+            )
 
         # Find the default executor
         if executor == "offload":