Cluster wait

distributed/client.py

@@ -1438,7 +1438,11 @@ def wait_for_workers(
             raise ValueError(
                 f"`n_workers` must be a positive integer. Instead got {n_workers}."
             )
-        return self.sync(self._wait_for_workers, n_workers, timeout=timeout)
+
+        if self.cluster is None:
+            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

distributed/deploy/cluster.py

@@ -19,6 +19,7 @@
 from distributed.compatibility import PeriodicCallback
 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,
@@ -33,6 +34,9 @@
 logger = logging.getLogger(__name__)
 
 
+no_default = "__no_default__"
+
+
 class Cluster(SyncMethodMixin):
     """Superclass for cluster objects
 
@@ -73,6 +77,7 @@ def __init__(
         self.periodic_callbacks = {}
         self._watch_worker_status_comm = None
         self._watch_worker_status_task = None
+        self._scheduler_info_comm = None
         self._cluster_manager_logs = []
         self.quiet = quiet
         self.scheduler_comm = None
@@ -120,15 +125,17 @@ def name(self, name):
         self._cluster_info["name"] = name
 
     async def _start(self):
-        comm = await self.scheduler_comm.live_comm()
-        comm.name = "Cluster worker status"
-        await comm.write({"op": "subscribe_worker_status"})
-        self.scheduler_info = SchedulerInfo(await comm.read())
-        self._watch_worker_status_comm = comm
+        self._watch_worker_status_comm = await self.scheduler_comm.live_comm()
+        self._watch_worker_status_comm.name = "Cluster worker status"
+        await self._watch_worker_status_comm.write({"op": "subscribe_worker_status"})
+        self.scheduler_info = SchedulerInfo(await self._watch_worker_status_comm.read())
         self._watch_worker_status_task = asyncio.ensure_future(
-            self._watch_worker_status(comm)
+            self._watch_worker_status(self._watch_worker_status_comm)
         )
 
+        self._scheduler_info_comm = await self.scheduler_comm.live_comm()
+        self._scheduler_info_comm.name = "Scheduler info"
+
         info = await self.scheduler_comm.get_metadata(
             keys=["cluster-manager-info"], default={}
         )
@@ -187,6 +194,8 @@ async def _close(self):
         with suppress(AttributeError):
             self._adaptive.stop()
 
+        if self._scheduler_info_comm:
+            await self._scheduler_info_comm.close()
         if self._watch_worker_status_comm:
             await self._watch_worker_status_comm.close()
         if self._watch_worker_status_task:
@@ -580,3 +589,59 @@ def __eq__(self, other):
 
     def __hash__(self):
         return id(self)
+
+    async def _wait_for_workers(self, n_workers=0, timeout=None):
+        await self.update_scheduler_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(self.scheduler_info) < n_workers:
+            if deadline and time() > deadline:
+                raise TimeoutError(
+                    "Only %d/%d workers arrived after %s"
+                    % (running_workers(self.scheduler_info), n_workers, timeout)
+                )
+            await asyncio.sleep(0.1)
+
+            await self.update_scheduler_info()
+
+    async def update_scheduler_info(self) -> None:
+        """Send comm to scheduler requesting information and update scheduler_info accordingly"""
+        await self._scheduler_info_comm.write({"op": "identity"})
+        self.scheduler_info = SchedulerInfo(await self._scheduler_info_comm.read())
+
+    def wait_for_workers(
+        self, n_workers: int | str = no_default, timeout: float | None = None
+    ) -> 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``
+        """
+        if n_workers is no_default:
+            warnings.warn(
+                "Please specify the `n_workers` argument when using `Client.wait_for_workers`. Not specifying `n_workers` will no longer be supported in future versions.",
+                FutureWarning,
+            )
+            n_workers = 0
+        elif not isinstance(n_workers, int) or n_workers < 1:
+            raise ValueError(
+                f"`n_workers` must be a positive integer. Instead got {n_workers}."
+            )
+        return self.sync(self._wait_for_workers, n_workers, timeout=timeout)

distributed/deploy/tests/test_cluster.py

@@ -3,6 +3,7 @@
 import pytest
 from tornado.ioloop import IOLoop
 
+from distributed import LocalCluster, Status
 from distributed.deploy.cluster import Cluster
 from distributed.utils_test import gen_test
 
@@ -38,6 +39,21 @@ async def test_logs_deprecated():
             cluster.logs()
 
 
+@gen_test()
+async def test_cluster_wait_for_worker(loop):
+    with LocalCluster(n_workers=3, loop=loop) as cluster:
+        assert len(cluster.scheduler.workers) == 3
+        cluster.scale(10)
+        cluster.wait_for_workers(10)
+        assert all(
+            [
+                worker["status"] == Status.running.name
+                for _, worker in cluster.scheduler_info["workers"].items()
+            ]
+        )
+        assert len(cluster.scheduler.workers) == 10
+
+
 @gen_test()
 async def test_deprecated_loop_properties():
     class ExampleCluster(Cluster):