From 964abea0fee437463585ffdfa5a8862ec7431f18 Mon Sep 17 00:00:00 2001 From: Erik Sundell Date: Wed, 10 Jan 2024 11:43:19 +0100 Subject: [PATCH] Delegate wait_for_workers to cluster instances only when implemented (#8441) --- distributed/client.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/distributed/client.py b/distributed/client.py index 0ca5447b34..a61a9f72d8 100644 --- a/distributed/client.py +++ b/distributed/client.py @@ -1454,10 +1454,10 @@ 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.cluster.wait_for_workers(n_workers, timeout) + return self.sync(self._wait_for_workers, n_workers, timeout=timeout) def _heartbeat(self): # Don't send heartbeat if scheduler comm or cluster are already closed