From d0bc4ca47e7bcbee35420a104afec673910c31b5 Mon Sep 17 00:00:00 2001 From: Erik Sundell Date: Thu, 4 Jan 2024 11:52:55 +0100 Subject: [PATCH] Delegate wait_for_workers to cluster instances only when implemented --- distributed/client.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/distributed/client.py b/distributed/client.py index 0ca5447b341..c7f724854fa 100644 --- 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