Cluster wait

distributed/client.py

@@ -1332,7 +1332,10 @@ def wait_for_workers(self, n_workers=0, timeout=None):
             Time in seconds after which to raise a
             ``dask.distributed.TimeoutError``
         """
-        return self.sync(self._wait_for_workers, n_workers, timeout=timeout)
+        if self.cluster is not None:
+            return self.cluster.wait_for_workers(n_workers, timeout)
+        else:
+            return self.sync(self._wait_for_workers, n_workers, timeout=timeout)
 
     def _heartbeat(self):
         if self.scheduler_comm:

distributed/deploy/cluster.py

@@ -17,6 +17,7 @@
 
 from distributed.core import Status
 from distributed.deploy.adaptive import Adaptive
+from distributed.metrics import time
 from distributed.objects import SchedulerInfo
 from distributed.utils import (
     Log,
@@ -531,3 +532,44 @@ def __eq__(self, other):
 
     def __hash__(self):
         return id(self)
+
+    async def _wait_for_workers(self, n_workers=0, timeout=None):
+        info = self.scheduler.identity()
+        self._scheduler_identity = SchedulerInfo(info)
+        if timeout:
+            deadline = time() + parse_timedelta(timeout)
+        else:
+            deadline = None
+
+        def running_workers(info):
+            return len(
+                [
+                    ws
+                    for ws in info["workers"].values()
+                    if ws["status"] == Status.running.name
+                ]
+            )
+
+        while n_workers and running_workers(info) < n_workers:
+            if deadline and time() > deadline:
+                raise TimeoutError(
+                    "Only %d/%d workers arrived after %s"
+                    % (running_workers(info), n_workers, timeout)
+                )
+            await asyncio.sleep(0.1)
+
+            info = self.scheduler.identity()
+            self._scheduler_identity = SchedulerInfo(info)
+
+    def wait_for_workers(self, n_workers=0, timeout=None):
+        """Blocking call to wait for n workers before continuing
+
+        Parameters
+        ----------
+        n_workers : int
+            The number of workers
+        timeout : number, optional
+            Time in seconds after which to raise a
+            ``dask.distributed.TimeoutError``
+        """
+        return self.sync(self._wait_for_workers, n_workers, timeout=timeout)

distributed/deploy/tests/test_cluster.py

@@ -2,6 +2,7 @@
 
 import pytest
 
+from distributed import LocalCluster
 from distributed.deploy.cluster import Cluster
 from distributed.utils_test import gen_test
 
@@ -48,3 +49,24 @@ def __init__(self):
     assert "foo" in cluster._cluster_info  # exists before start() called
     with cluster:  # start and stop the cluster to avoid a resource warning
         pass
+
+
+@gen_test()
+async def test_cluster_wait_for_worker(loop):
+    with LocalCluster(n_workers=3, loop=loop) as cluster:
+        assert all(
+            [
+                worker.status.name == "running"
+                for _, worker in cluster.scheduler.workers.items()
+            ]
+        )
+        assert len(cluster.scheduler.workers) == 3
+        cluster.scale(10)
+        cluster.wait_for_workers(10)
+        assert all(
+            [
+                worker.status.name == "running"
+                for _, worker in cluster.scheduler.workers.items()
+            ]
+        )
+        assert len(cluster.scheduler.workers) == 10