Repository: dask/distributed

Issue:
Wait for workers to join before continuing
_Moved from dask/dask-kubernetes#87_

It would be useful to have the `cluster.scale()` method block until the workers join. This was originally raised in dask-kubernetes but is applicable to dask-jobqueue and dask-yarn so would make more sense to implement here.

Suggestions from @mrocklin:

> One solution here might be to use SchedulerPlugins to explicitly trigger events whenever a new worker reaches the scheduler. Calling wait on the cluster object might register a tornado.locks.Event or tornado.locks.Condition that gets triggered when the right number of workers have arrived.

> @jhamman implemented a SchedulerPlugin for dask-jobqueue in order to improve adaptive operation (he wanted adaptive to be aware of pending jobs). This work happened in dask/dask-jobqueue#63 and might be the sort of thing we could think about generalizing. Any thoughts on this @jhamman?


Discussion:
Hi there. From dask-jobqueue perspective, I don't think it would be a good thing to wait until all workers join. Maybe block until one worker is there, but with job scheduler, waiting for all of them could take hours if not more. Maybe add a kwarg to scale() ?

From what I remember, using the plugin as suggested by @mrocklin looks simple enough.
You make a good point. I'm making a few assumptions here. 

The first is that I'm using a cluster which doesn't make me wait very long to give me workers (the cloud). I'm talking about interactive jobs rather than batch jobs.

I'm also assuming that I'm doing something brittle which needs all the workers before starting. The primary example of this for me right now is doing a live demo, I don't want it to run on one worker as it will be underwhelming. I want to scale and wait until I have a nice big cluster before plowing through my graph.

I appreciate that might not be a good enough reason to do this. I could write a little script to block until workers reaches n for live demos. 
I think it would make sense to have an explicit `wait` method or ` wait=` keyword to scale that would allow a user to opt-in to blocking semantics.  I would use this for benchmarking.
Perhaps we could have `wait=bool` and/or `wait_for=int` so you could specify a minimum number to wait for before starting. This may fit better with jobqueue.
Yep, I had something like this in mind. Being able to wait for a user defined number of workers to be online is something that would definitly be useful!
FWIW code like this has served this purpose quite well for us. Would be easy enough for someone to add it if they wish.

```python
while ((client.status == "running") and (len(client.scheduler_info()["workers"]) < nworkers)):
    sleep(1.0)
```
Could we use a non active wait solution for that?

Maybe using http://distributed.dask.org/en/latest/plugins.html#scheduler-plugins and some callback ?

Do you have suggestion on this @mrocklin ?
Adding some constraints here:

Currently the `scale` method is fast so that it can be reliably used within the event loop.  Typically scale starts some jobs and then returns immediately.  If we decide that it should block, then it has to block in an async-friendly way.  This means that `scale` has to become a coroutine if we're in the event loop (it has to return a future), and a blocking function if we're not.  This will require some finesse.  

We might solve this in a couple of ways:

1.  We could change the contract of scale across all of the projects.  This might be necessary at some point anyway.
2.  We could not add the `wait=` keyword to scale, and instead add a second method, `.wait()` which followed the coroutine semantics commonly used in the dask.distributed codebase.

    ```python
    async def _wait(self, n):
        while not condition():
            await gen.sleep(0.01)

    def wait(self, n):
        return self.sync(self._wait, n)
    ```
Just for my understanding, if we do not modify `scale` and instead just add a `wait` method, why would we want it to follow the common coroutine semantics?
In case someone wanted to use wait from within an asynchronous environment.

```python
async with SGECluster(...) as cluster:
    cluster.scale(10)
    await cluster.wait()
    x.compute()
```
As another data point, this is impacting my use of Dask in an HPC environment (and some thoughts on CLI integration below).

My jobs are started by a job scheduler (SLURM), and at the start of the job I want to spin up a scheduler and set of workers. For reasons not worth mentioning here, I'm doing this by calling `dask-scheduler` and `dask-worker` rather than `dask-ssh` or `dask-mpi`, but I believe the same issue appears no matter which route I go. Conceptually the setup looks like:

```
dask-scheduler &
srun ... dask-worker ... &
python my_dask_script.py
```

There is a race between the three steps. This is not a problem for *correctness* since the worker knows to retry if the scheduler isn't up yet, and the scheduler knows to wait for at least one worker to arrive before allowing the client to start running. But in my case this is a *performance* issue because I want to do accurate timing of execution, and if the client starts before all workers have arrived then the run isn't reflective of the steady state performance.

Currently I'm doing a fairly awful workaround based on parsing the log files from scheduler and counting the number of times it reports workers. You can see my script below. I'm not sure how exactly the API mentioned in the original post would integrate with my workflow; presumably this would need to be exposed via the CLI somehow. I would prefer not to have to write custom scripts to start Dask but I suppose I can do that if absolutely necessary.

https://github.com/elliottslaughter/task-bench/blob/dask/experiments/cori_metg_compute/metg_dask.sh#L18-L43
Could you create the client in the script, and then poll `client.ncores()` until you see as many workers are you are happy with?
I think could make that work. In my case I know how many nodes I'm planning to boot, so I can just tell that to the client.

Is there API documentation for this? I was trying to figure out if e.g. the `num_workers` keyword argument to compute was applicable to the distributed scheduler, but never found anything beyond this page (which only mentions multiprocessing and threaded schedulers): https://docs.dask.org/en/latest/scheduler-overview.html#configuring-the-schedulers
https://distributed.readthedocs.io/en/latest/api.html#distributed.Client.ncores

The method just gives information of the scheduler's view of the cluster.
You can also look at Client.scheduler_info, which is a decent all-purpose
method to get information about the general state.

On Fri, Apr 12, 2019 at 3:49 PM Martin Durant <notifications@github.com>
wrote:

>
> https://distributed.readthedocs.io/en/latest/api.html#distributed.Client.ncores
>
> The method just gives information of the scheduler's view of the cluster.
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/distributed/issues/2138#issuecomment-482718237>,
> or mute the thread
> <https://github.com/notifications/unsubscribe-auth/AASszLv2C8QIZ0BlSWL31ROErQ2A-CSfks5vgPFzgaJpZM4VaoCy>
> .
>

But, we should probably solve this problem properly at some point.  It's
well in-scope for the project.  If you or anyone you know would like to
work on it we'd be happy to help walk them through it.  Otherwise hopefully
it gets implemented within a few months.

On Fri, Apr 12, 2019 at 4:15 PM Matthew Rocklin <mrocklin@gmail.com> wrote:

> You can also look at Client.scheduler_info, which is a decent all-purpose
> method to get information about the general state.
>
> On Fri, Apr 12, 2019 at 3:49 PM Martin Durant <notifications@github.com>
> wrote:
>
>>
>> https://distributed.readthedocs.io/en/latest/api.html#distributed.Client.ncores
>>
>> The method just gives information of the scheduler's view of the cluster.
>>
>> —
>> You are receiving this because you were mentioned.
>> Reply to this email directly, view it on GitHub
>> <https://github.com/dask/distributed/issues/2138#issuecomment-482718237>,
>> or mute the thread
>> <https://github.com/notifications/unsubscribe-auth/AASszLv2C8QIZ0BlSWL31ROErQ2A-CSfks5vgPFzgaJpZM4VaoCy>
>> .
>>
>

@elliottslaughter I guess you using dask-jobqueue? Have you seen this PR: https://github.com/dask/dask-jobqueue/pull/223 by @danpf?

It is not properly finished, but may give you some insights.

Resolving patch:
diff --git a/distributed/client.py b/distributed/client.py
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -1053,6 +1053,17 @@ def _update_scheduler_info(self):
         except EnvironmentError:
             logger.debug("Not able to query scheduler for identity")
 
+    @gen.coroutine
+    def _wait_for_workers(self, n_workers=0):
+        info = yield self.scheduler.identity()
+        while n_workers and len(info["workers"]) < n_workers:
+            yield gen.sleep(0.1)
+            info = yield self.scheduler.identity()
+
+    def wait_for_workers(self, n_workers=0):
+        """Blocking call to wait for n workers before continuing"""
+        return self.sync(self._wait_for_workers, n_workers)
+
     def _heartbeat(self):
         if self.scheduler_comm:
             self.scheduler_comm.send({"op": "heartbeat-client"})


Test patch:
diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
--- a/distributed/tests/test_client.py
+++ b/distributed/tests/test_client.py
@@ -5712,5 +5712,18 @@ def test_instances(c, s, a, b):
     assert set(Worker._instances) == {a, b}
 
 
+@gen_cluster(client=True)
+def test_wait_for_workers(c, s, a, b):
+    future = c.wait_for_workers(n_workers=3)
+    yield gen.sleep(0.22)  # 2 chances
+    assert not future.done()
+
+    w = yield Worker(s.address)
+    start = time()
+    yield future
+    assert time() < start + 1
+    yield w.close()
+
+
 if sys.version_info >= (3, 5):
     from distributed.tests.py3_test_client import *  # noqa F401