Fix LocalCUDACluster's parsing of worker_class

rapidsai · Oct 24, 2023 · 30513fd · 30513fd
1 parent 6155732
commit 30513fd
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/dask_cuda/local_cuda_cluster.py b/dask_cuda/local_cuda_cluster.py
@@ -2,6 +2,7 @@
 import logging
 import os
 import warnings
+from functools import partial
 
 import dask
 from distributed import LocalCluster, Nanny, Worker
@@ -339,6 +340,8 @@ def __init__(
                 "`distributed.local_cuda_cluster.LoggedWorker`, and specify "
                 "`log_spilling=False`."
             )
+        if not isinstance(worker_class, Nanny):
+            worker_class = partial(Nanny, worker_class=worker_class)
 
         self.pre_import = pre_import