Repository: dask/distributed

Issue:
add `cluster.wait_for_workers` (called by `client.wait_for_workers` if the `client.client` exists)
Currently users call `client.wait` to wait for a certain number of workers.

Currently a common pattern is:

```
cluster.scale(100)
client.wait_for_workers(100)
```

Cluster managers (like Coiled) would like to include custom handling to give users nice information about what's happening (including especially failures) while waiting for workers. At the moment, this isn't possible.

Proposal to make this possible:

- add `cluster.wait_for_workers` (with logic similar to the existing `client.wait_for_workers`)
- `client.wait_for_workers` should call `cluster.wait_for_workers` if there's a cluster object.

Then cluster managers can override `cluster.wait_for_workers` to provide custom handling, which will be called whether the user calls `cluster.wait_for_workers` or `client.wait_for_workers`.

This suggestion arose from discussion with @jacobtomlinson in https://github.com/dask/distributed/pull/6041.

Discussion:
GPUtester: Can one of the admins verify this patch?
github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/distributed/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

       15 files  ±0         15 suites  ±0   6h 12m 59s [:stopwatch:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "duration of all tests") + 9m 26s
  2 951 tests +1    2 863 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "passed tests")  - 2       85 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "skipped / disabled tests") ±0  3 [:x:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "failed tests") +3 
21 879 runs  +7  20 849 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "passed tests")  - 5  1 023 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "skipped / disabled tests") +5  7 [:x:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "failed tests") +7 

For more details on these failures, see [this check](https://github.com/dask/distributed/runs/7265174618).

Results for commit 4ee2739b. ± Comparison against base commit 6765e6e9.

[test-results]:data:application/gzip;base64,H4sIAK+5yWIC/03NTQ7CIBCG4asY1i4YoGXwMqalQ0Lsj6GwMt7dabXF5fsM+XiJEEdaxe0CzfUi1hJzraGkLsdl5lZKW8fE17zdlWvgyPtavN8MW13tEZ9s2JwSujiy1CeU0pJYJEsq8z4LuH+z5Tkr0VT7zoJU+qDfrj36f9Yv0xQzhzBEymrXOzNYHYB0a3qUAxGAwtCRQYvKkxTvD7X4DPUTAQAA

quasiben: add to allowlist
idorrington92: > add to allowlist

@quasiben Sorry, what does this mean?

Also, I've got 3 failing tests in test_ssh. Is it possible to run these locally? They just fail for me, saying cluster failed to start. This happens on main branch as well. I'm not really sure how to debug them otherwise... 
quasiben: @idorrington92 the comment allows us to run your PR against a GPU CI system:
https://docs.dask.org/en/stable/develop.html#gpu-ci

Distributed has been flaky in the past and recently there has been a concerted effort to stabilize the CI test:
https://dask.github.io/distributed/test_report.html

As you can see in the linked report, `test_dashboard_port_zero` is a known flaky test

I would suggest @jacobtomlinson review this PR but he is currently at SciPy (as are other dask maintainers) so it may be a little bit before you hear back
idorrington92: Doesn't it say listing/pre-commit hooks passed ok? 

I agree that test failure is related but I can't run it locally as I get an error saying cluster failed to start. 

Is there something I need to do in order to run it locally? Otherwise the only thing I've got to go on is the automated tests here... 
jacobtomlinson: Apologies I was looking at `Linting / Unit Test Results (pull_request)` which actually has nothing to do with linting, sorry for the noise.

Can you run `ssh localhost` without having to enter a password?
idorrington92: Thanks @jacobtomlinson, that fixed the error I was getting. I would have taken ages to think of that.

I had a little look at the bug that was coming up in the tests, but it's not obvious to me why it's happening. I'll have proper look later in the week.
idorrington92: @jacobtomlinson can I get some help please? 

I managed to fix those failing tests by calling client.scale in client.wait_for_workers. But when I merged in master, the test I wrote to check the cluster.wait_for_worker method started failing. It waits forever for the workers because they never change their status to running, they just get stuck on init. 

Do you have any ideas why this would be happening? No worries if not, I'll keep digging. 
idorrington92: I think it's all working properly now. I had to create a function to get the worker status and update the scheduler_info. As it was, scheduler_info was only updated when the workers were added, which was why it looked like they were stuck with status 'init'. The new function updates the scheduler_info with the latest worker status. 

One test is failing on python 3.8, but I think that's nothing to do with my changes.
idorrington92: Hi @jacobtomlinson and @fjetter, just pinging you to make sure this isn't forgotten about :) 

I've merged main branch into this branch, and copied the new behaviour for when there are no workers
idorrington92: Can we ping someone from Coiled to review this? It's been months now and would be nice to merge this
mrocklin: Hi @idorrington92 sorry that this hasn't gotten much love.  I suspect that most folks are out during the holidays.  I'm holding down the fort for the moment.  

If I were reviewing this initially I probably would have said "the try-except logic around `cluster.wait_for_workers` seems super simple and good to me 👍 .  However, the implementation for `Cluster.wait_for_workers` and the changes to `update_scheduler_info` seem also *probably* good, but like an area that might be error prone.  Do we actually want/need this for some reason?  If not I'd be inclined to skip this and stay with status quo.  It seems to me like the upstream issue was mostly asking for the ability of *downstream* clusters to have their `wait_for_worker` methods respected, not tha we implement one for all `Cluster` classes.  

My preference is to drop the implementation (even if this means dropping the test), and just stick with the try-except logic that you have in `client.py` (which seems super straightforward).  However, I appreciate that you've also gone through several rounds of review, and having expectations change on you after so long probably isn't fun. 

Do you have any thoughts on the above?  Is there a solid reason to implement `wait_for_workers` and `update_scheduler_info` that I'm not seeing here?
mrocklin: Said a different way, I'm more than happy to merge the client.py changes immediately.  I'd want to look over the other changes more thoroughly, which would take me some time.  Or somoene like @jacobtomlinson who has already looked over them could merge if he's feeling confident.
idorrington92: > Hi @idorrington92 sorry that this hasn't gotten much love. I suspect that most folks are out during the holidays. I'm holding down the fort for the moment.
> 
> If I were reviewing this initially I probably would have said "the try-except logic around `cluster.wait_for_workers` seems super simple and good to me +1 . However, the implementation for `Cluster.wait_for_workers` and the changes to `update_scheduler_info` seem also _probably_ good, but like an area that might be error prone. Do we actually want/need this for some reason? If not I'd be inclined to skip this and stay with status quo. It seems to me like the upstream issue was mostly asking for the ability of _downstream_ clusters to have their `wait_for_worker` methods respected, not tha we implement one for all `Cluster` classes.
> 
> My preference is to drop the implementation (even if this means dropping the test), and just stick with the try-except logic that you have in `client.py` (which seems super straightforward). However, I appreciate that you've also gone through several rounds of review, and having expectations change on you after so long probably isn't fun.
> 
> Do you have any thoughts on the above? Is there a solid reason to implement `wait_for_workers` and `update_scheduler_info` that I'm not seeing here?

Hi @mrocklin,
Thank you for your quick reply :)
I see what you mean. I took the issue quite literally and just implemented the changes that were requested. But perhaps a simple try-catch in client.py would allow clusters to do everything the issue really needs. I've not used clusters much, so don't know from a users perspective whether cluster.wait_for_workers is very useful.

FYI - Following the comment from @fjetter above, I've replaced that try-catch with an if-else, so that will need to be reverted if we do go down that route. 

I'm happy either way, I learned a lot from working on this issue, and would rather remove most of my changes and contribute something useful than add a load of code that'll cause problems later on :) 

I think it's better if people with more user knowledge than me decide which route we go down though :) 
idorrington92: I've handled the merge conflicts. I'm getting some failing tests, but I'm seeing these test failures in other PRs as well so don't think it's to do with my changes... 
jacobtomlinson: It looks like `distributed/deploy/tests/test_cluster.py::test_cluster_wait_for_worker` is failing a bunch here so it's certainly related to these changes. Would you mind taking another look?
idorrington92: It is only failing in Python 3.8, because of threads leaking. Is this 

Patch:
diff --git a/distributed/client.py b/distributed/client.py
index 264fe232..8c897e96 100644
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -1458,7 +1458,11 @@ class Client(SyncMethodMixin):
             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
diff --git a/distributed/deploy/cluster.py b/distributed/deploy/cluster.py
index bbba149d..c9170c1c 100644
--- a/distributed/deploy/cluster.py
+++ b/distributed/deploy/cluster.py
@@ -19,6 +19,7 @@ from dask.widgets import get_template
 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 @@ from distributed.utils import (
 logger = logging.getLogger(__name__)
 
 
+no_default = "__no_default__"
+
+
 class Cluster(SyncMethodMixin):
     """Superclass for cluster objects
 
@@ -582,6 +586,57 @@ class Cluster(SyncMethodMixin):
     def __hash__(self):
         return id(self)
 
+    async def _wait_for_workers(self, n_workers=0, timeout=None):
+        self.scheduler_info = SchedulerInfo(await self.scheduler_comm.identity())
+        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)
+
+            self.scheduler_info = SchedulerInfo(await self.scheduler_comm.identity())
+
+    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)
+
 
 def _exponential_backoff(
     attempt: int, multiplier: float, exponential_base: float, max_interval: float


Test patch:
diff --git a/distributed/deploy/tests/test_cluster.py b/distributed/deploy/tests/test_cluster.py
index cb0b189e..55339160 100644
--- a/distributed/deploy/tests/test_cluster.py
+++ b/distributed/deploy/tests/test_cluster.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 import pytest
 from tornado.ioloop import IOLoop
 
+from distributed import LocalCluster, Status
 from distributed.deploy.cluster import Cluster, _exponential_backoff
 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():
+    async with LocalCluster(n_workers=2, asynchronous=True) as cluster:
+        assert len(cluster.scheduler.workers) == 2
+        cluster.scale(4)
+        await cluster.wait_for_workers(4)
+        assert all(
+            [
+                worker["status"] == Status.running.name
+                for _, worker in cluster.scheduler_info["workers"].items()
+            ]
+        )
+        assert len(cluster.scheduler.workers) == 4
+
+
 @gen_test()
 async def test_deprecated_loop_properties():
     class ExampleCluster(Cluster):