Delegate wait_for_workers to cluster instances only when implemented

dask · Jan 9, 2024 · d0bc4ca · d0bc4ca
1 parent 7562f9c
commit d0bc4ca
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/distributed/client.py b/distributed/client.py
@@ -1454,10 +1454,11 @@ def wait_for_workers(self, n_workers: int, timeout: float | None = None) -> None
                 f"`n_workers` must be a positive integer. Instead got {n_workers}."
             )
 
-        if self.cluster is None:
-            return self.sync(self._wait_for_workers, n_workers, timeout=timeout)
+        if self.cluster and hasattr(self.cluster, "wait_for_workers"):
+            return self.cluster.wait_for_workers(n_workers, timeout)
+
+        return self.sync(self._wait_for_workers, n_workers, timeout=timeout)
 
-        return self.cluster.wait_for_workers(n_workers, timeout)
 
     def _heartbeat(self):
         # Don't send heartbeat if scheduler comm or cluster are already closed