Skip to content

Commit

Permalink
Delegate wait_for_workers to cluster instances only when implemented (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
consideRatio committed Jan 10, 2024
1 parent 7562f9c commit 964abea
Showing 1 changed file with 3 additions and 3 deletions.
6 changes: 3 additions & 3 deletions distributed/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 964abea

Please sign in to comment.