diff --git a/distributed/client.py b/distributed/client.py index e31b808abcf..fa6d9b5e387 100644 --- a/distributed/client.py +++ b/distributed/client.py @@ -1053,6 +1053,14 @@ def _update_scheduler_info(self): except EnvironmentError: logger.debug("Not able to query scheduler for identity") + @gen.coroutine + def _wait_until_n_workers(self, n): + while n and len(self.cluster.scheduler.workers) < n: + yield gen.sleep(0.01) + + def wait_until_n_workers(self, n): + return self.sync(self._wait_until_n_workers, n) + def _heartbeat(self): if self.scheduler_comm: self.scheduler_comm.send({"op": "heartbeat-client"})