Add docstrings to MockWorker class

rapidsai · Jun 7, 2021 · 5a175f0 · 5a175f0
1 parent 9bc5c09
commit 5a175f0
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/dask_cuda/utils.py b/dask_cuda/utils.py
@@ -533,6 +533,13 @@ def parse_device_memory_limit(device_memory_limit, device_index=0):
 
 
 class MockWorker(Worker):
+    """Mock Worker class preventing NVML from getting used by SystemMonitor.
+
+    By preventing the Worker from initializing NVML in the SystemMonitor, we can
+    mock test multiple devices in `CUDA_VISIBLE_DEVICES` behavior with single-GPU
+    machines.
+    """
+
     def __init__(self, *args, **kwargs):
         import distributed