Repository: dask/distributed

Issue:
`has_cuda_context` fails in WSL
<!-- Please include a self-contained copy-pastable example that generates the issue if possible.

Please be concise with code posted. See guidelines below on how to provide a good bug report:

- Craft Minimal Bug Reports http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
- Minimal Complete Verifiable Examples https://stackoverflow.com/help/mcve

Bug reports that follow these guidelines are easier to diagnose, and so are often handled much more quickly.
-->

**What happened**:
When attempting to check if a CUDA context has already been created using `has_cuda_context` in WSL, we see failures at the NVML level:

```python
NVMLError_Unknown                         Traceback (most recent call last)
<ipython-input-2-086bc6eeaba8> in <module>
----> 1 has_cuda_context()

~/dev/rapids-dask/distributed/distributed/diagnostics/nvml.py in has_cuda_context()
     76         handle = pynvml.nvmlDeviceGetHandleByIndex(index)
     77         if hasattr(pynvml, "nvmlDeviceGetComputeRunningProcesses_v2"):
---> 78             running_processes = pynvml.nvmlDeviceGetComputeRunningProcesses_v2(handle)
     79         else:
     80             running_processes = pynvml.nvmlDeviceGetComputeRunningProcesses(handle)

~/miniconda3/envs/rapids-dask/lib/python3.8/site-packages/pynvml/nvml.py in nvmlDeviceGetComputeRunningProcesses_v2(handle)
   2139     else:
   2140         # error case
-> 2141         raise NVMLError(ret)
   2142 
   2143 def nvmlDeviceGetComputeRunningProcesses(handle):
```

This is actually expected, as [querying active processes is not yet supported in WSL](https://docs.nvidia.com/cuda/wsl-user-guide/index.html#features-not-yet-supported).

**What you expected to happen**:
We would want the function to work as desired, i.e. returning false or the index of the device for which there's a CUDA context.

**Minimal Complete Verifiable Example**:

```python
from distributed.diagnostics.nvml import has_cuda_context

has_cuda_context()
```

The error also occurs running the bare PyNVML code:

```python
from pynvml import *

nvmlInit()

h = nvmlDeviceGetHandleByIndex(0)
nvmlDeviceGetComputeRunningProcesses_v2(h)
```

**Anything else we need to know?**:
For context, this causes issues when attempting to create a CUDA context in dask-cuda, which is currently being tracked in https://github.com/rapidsai/dask-cuda/issues/816.

I think a potential solution to this issue, in the same vein as https://github.com/dask/distributed/pull/5343, would be to make a special case for WSL in `has_cuda_context`. This could either be done implicitly with a try/except block to catch the NVML error, or explicitly by checking if the OS is WSL and doing something different in that case:

```python
# https://www.scivision.dev/python-detect-wsl/
from platform import uname

def in_wsl() -> bool:
    return 'microsoft-standard' in uname().release
```

**Environment**:

- Dask version: latest `main`
- Python version: 3.8
- Operating System: ubuntu 20.04 (WSL2)
- Install method (conda, pip, source): source

cc @pentschev 


Discussion:
> I think a potential solution to this issue, in the same vein as #5343, would be to make a special case for WSL in `has_cuda_context`. This could either be done implicitly with a try/except block to catch the NVML error, or explicitly by checking if the OS is WSL and doing something different in that case:

I would any day prefer an explicit check. Implicitly disabling NVML based on an error makes any future non-WSL-related issues very hard to debug.

> 
> ```python
> # https://www.scivision.dev/python-detect-wsl/
> from platform import uname
> 
> def in_wsl() -> bool:
>     return 'microsoft-standard' in uname().release
> ```

This looks like a good enough check for me, +1 to just verifying this way and returning `False` at all times. A better alternative mixing both implicit and explicit checks would be to implicitly try/except, and return `False` in the except block when system is WSL.
> A better alternative mixing both implicit and explicit checks would be to implicitly try/except, and return False in the except block when system is WSL.

Good idea! Playing around with the fix in a branch now and now that we are past the failures in `has_cuda_context`, we run into issues with the real-time NVML monitoring:

```python
tornado.application - ERROR - Exception in callback <bound method SystemMonitor.update of <SystemMonitor: cpu: 6 memory: 338 MB fds: 30>>
Traceback (most recent call last):
  File "/home/charlesb/miniconda3/envs/rapids-dask/lib/python3.8/site-packages/tornado/ioloop.py", line 905, in _run
    return self.callback()
  File "/home/charlesb/dev/rapids-dask/distributed/distributed/system_monitor.py", line 132, in update
    gpu_metrics = nvml.real_time()
  File "/home/charlesb/dev/rapids-dask/distributed/distributed/diagnostics/nvml.py", line 129, in real_time
    h = _pynvml_handles()
  File "/home/charlesb/dev/rapids-dask/distributed/distributed/diagnostics/nvml.py", line 64, in _pynvml_handles
    return pynvml.nvmlDeviceGetHandleByIndex(gpu_idx)
  File "/home/charlesb/miniconda3/envs/rapids-dask/lib/python3.8/site-packages/pynvml/nvml.py", line 1576, in nvmlDeviceGetHandleByIndex
    _nvmlCheckReturn(ret)
  File "/home/charlesb/miniconda3/envs/rapids-dask/lib/python3.8/site-packages/pynvml/nvml.py", line 743, in _nvmlCheckReturn
    raise NVMLError(ret)
pynvml.nvml.NVMLError_Unknown: Unknown Error
```

Is this expected given the workaround we're doing here?
I'm guessing this is now the general issue with NVML in WSL. You'll probably need to forcefully disable it with `DASK_DISTRIBUTED__DIAGNOSTICS__NVML=False`. Could you try that? If it works then I believe the CUDA context issue can be considered resolved (given current NVML limitations with WSL).

Resolving patch:
diff --git a/distributed/diagnostics/nvml.py b/distributed/diagnostics/nvml.py
--- a/distributed/diagnostics/nvml.py
+++ b/distributed/diagnostics/nvml.py
@@ -1,4 +1,5 @@
 import os
+from platform import uname
 
 import dask
 
@@ -12,10 +13,17 @@
 nvmlOwnerPID = None
 
 
+def _in_wsl():
+    """Check if we are in Windows Subsystem for Linux; some PyNVML queries are not supported there.
+    Taken from https://www.scivision.dev/python-detect-wsl/
+    """
+    return "microsoft-standard" in uname().release
+
+
 def init_once():
     global nvmlInitialized, nvmlLibraryNotFound, nvmlOwnerPID
 
-    if dask.config.get("distributed.diagnostics.nvml") is False:
+    if dask.config.get("distributed.diagnostics.nvml") is False or _in_wsl():
         nvmlInitialized = False
         return
 
@@ -72,8 +80,13 @@ def has_cuda_context():
     index of the device for which there's a CUDA context.
     """
     init_once()
+    if nvmlLibraryNotFound or not nvmlInitialized:
+        return False
     for index in range(device_get_count()):
         handle = pynvml.nvmlDeviceGetHandleByIndex(index)
+        # TODO: WSL doesn't support this NVML query yet; when NVML monitoring is enabled
+        # there we may need to wrap this in a try/except block.
+        # See https://github.com/dask/distributed/pull/5568
         if hasattr(pynvml, "nvmlDeviceGetComputeRunningProcesses_v2"):
             running_processes = pynvml.nvmlDeviceGetComputeRunningProcesses_v2(handle)
         else:


Test patch: