Repository: dask/distributed

Issue:
Disabling GPU diagnostics prevents GPU executor from getting created
After filing #8337, I ran across `distributed.diagnostics.nvml`, which seems to be supposed to serve the same purpose of toggling _diagnotics_ on GPUs. However, after taking a deeper look, it seems as though this setting has some unintended consequences and disables the GPU executor altogether:

In `distributed.diagnostics.nvml`, it is used within `init_once` to decide whether to initialize `pynvml` on a process:

https://github.com/dask/distributed/blob/0dc9e8800d6e49e77e2fd5cbc71f0033cccf7b3b/distributed/diagnostics/nvml.py#L89-L93

The initialization is then checked within `device_get_count()`, which returns 0 if the toggle is turned off:

https://github.com/dask/distributed/blob/0dc9e8800d6e49e77e2fd5cbc71f0033cccf7b3b/distributed/diagnostics/nvml.py#L125-L128

On the worker, this function is then used to decide whether to initialize the GPU executor:

https://github.com/dask/distributed/blob/0dc9e8800d6e49e77e2fd5cbc71f0033cccf7b3b/distributed/worker.py#L639-L642

I'm not an expert on this particular area of code, but I think that we should separate initialization of `pynvml` from the diagnostics toggle. We may also want to consider moving diagnostics toggle into the system monitor (akin to #8337) to have a conceptual grouping.


Discussion:
I can speak for the NVML diagnostics side, and we intentionally chose to decide whether to enable diagnostics or not when `device_get_count() == 0` which may indeed cause some confusion as it packs more information than its name suggests:

1. Is the NVML diagnostics toggle turned on?
2. Is `pynvml` installed?
3. Are NVML resources actually available (e.g., shared libraries)?
4. And finally, are there GPUs installed and visible?

We could separate all of those things but it adds a ton of boilerplate code that's easy to miss, even just having to test whether NVML is enabled and the device count is, from previous experience, easy enough to introduce unintended consequences. With that said, so far we assumed that if GPU diagnostics are disabled, everything GPU-related is intended to be disabled.

I wouldn't mind making changes so that we can still query whether GPUs are available even if the NVML toggle is off if there's use for you. I was not aware of the existence of `Dask-GPU-Threads` executor and I'm not familiar with its purpose, in Dask-CUDA we do not use it even though we rely heavily on the Distributed NVML diagnostics module. With that said, if having another path that can still keep `Dask-GPU-Threads` enabled despite NVML diagnostics switched off I'm ok with that, I would still like to retain the current behavior of `device_get_count() == 0` if NVML diagnostics are switched off though, and we can then perhaps introduce a new function or `bool` argument to `device_get_count` to still return the device count independent of the NVML diagnostics toggle.
@pentschev: Thanks for the additional context!

> I can speak for the NVML diagnostics side, and we intentionally chose to decide whether to enable diagnostics or not when `device_get_count() == 0`

That generally makes a lot of sense to me.

> With that said, so far we assumed that if GPU diagnostics are disabled, everything GPU-related is intended to be disabled.

This makes less sense to me. Shouldn't it instead be the other way around? If GPUs are generally enabled/disabled, then diagnostics should follow suit (possibly with the additional ability to turn off diagnostics either way).

> I was not aware of the existence of `Dask-GPU-Threads` executor and I'm not familiar with its purpose, in Dask-CUDA we do not use it even though we rely heavily on the Distributed NVML diagnostics module.

It was added a while ago (#5123), so I assumed that this was used by _something_.
Now I'm wondering if this is actually the case. If not, should we throw it out again? @mrocklin: Do you know more?

A more general question: Are NVML diagnostics independent of the system monitor? If not, I'm wondering if we should move the toggle there to be close to the other system monitor toggles `distributed.admin.system-monitor.{disk|host-cpu|gil}`.
> This makes less sense to me. Shouldn't it instead be the other way around? If GPUs are generally enabled/disabled, then diagnostics should follow suit (possibly with the additional ability to turn off diagnostics either way).

Yes, I would agree with that, but Distributed had no GPU code (at least to my knowledge) other than monitoring, thus we were making the assumption in the opposite direction (diagnostics disabled == intent to have no GPU support).

> A more general question: Are NVML diagnostics independent of the system monitor? If not, I'm wondering if we should move the toggle there to be close to the other system monitor toggles `distributed.admin.system-monitor.{disk|host-cpu|gil}`.

I don't mind moving the toggle, I think it was introduced in the NVML diagnostics module just to be the least intrusive as possible, if you see benefit in having it as part of system monitor I'm +1 for that.
> I don't mind moving the toggle, I think it was introduced in the NVML diagnostics module just to be the least intrusive as possible, if you see benefit in having it as part of system monitor I'm +1 for that.

Alright, I'll take another look at what distributed's NVML module does in detail to see if my proposal makes sense.
cc @dask/gpu for visibility
For context, Matt suggested in https://github.com/dask/distributed/issues/5084 to add an annotation to all RAPIDS layers such that the tasks are scheduled on this magical threadpool

https://github.com/dask/distributed/issues/5084#issuecomment-887051080
https://github.com/dask/distributed/issues/5084#issuecomment-887058906

Given this conversation, I suspect this never happened? Would anybody be upset if we removed that?

Resolving patch:
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -636,10 +636,6 @@ 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":


Test patch:
diff --git a/distributed/tests/test_worker.py b/distributed/tests/test_worker.py
--- a/distributed/tests/test_worker.py
+++ b/distributed/tests/test_worker.py
@@ -30,7 +30,6 @@
 from dask.system import CPU_COUNT
 from dask.utils import tmpfile
 
-import distributed
 from distributed import (
     Client,
     Event,
@@ -46,7 +45,6 @@
 from distributed.comm.utils import OFFLOAD_THRESHOLD
 from distributed.compatibility import LINUX, WINDOWS
 from distributed.core import CommClosedError, Status, rpc
-from distributed.diagnostics import nvml
 from distributed.diagnostics.plugin import ForwardOutput
 from distributed.metrics import time
 from distributed.protocol import pickle
@@ -2267,17 +2265,6 @@ 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
-
-
 async def assert_task_states_on_worker(
     expected: dict[str, str], worker: Worker
 ) -> None: