Repository: dask/distributed

Issue:
Should Cluster.scale be an async method?
Currently when using a cluster with `asynchronous=True`, `Cluster.scale` is not an async method, which means that any work that it does either blocks, or is scheduled to run as background task. In contrast, `dask_gateway.GatewayCluster.scale` currently is an async method when running with `asynchronous=True`. Note that it doesn't wait for the workers to be running, but it does wait for them to be submitted.

```python
cluster = await ClusterClass(asynchronous=True)

# Most cluster managers don't need an await here
# Any submissions either block, or run as background tasks
cluster.scale(2)

# dask-gateway currently needs an await here
# This returns after the scale request has been sent to the gateway server,
# not when the workers have been received.
await cluster.scale(2)
```

This was discovered when trying to make dask-gateway work with dask-labextension. It's easy to make labextension work with both, or we could make dask-gateway use a background task here to send the request. Happy with either.

cc @lesteve, @jacobtomlinson, @mrocklin @guillaumeeb, @ian-r-rose for thoughts.



Discussion:
One pro of making it an async method is that in asynchronous mode all standard api methods are then also async methods, making this a more uniform interface.

```python
cluster = await ClusterClass()
await cluster.scale(2)
await cluster.close()
```
I would be very much in favor of this. I have a memory of talking to @mrocklin about implementing async scaling when we were working on `dask-labextension`, but there was some difficulty that I can no longer remember, so we punted on the question. Unfortunately, I couldn't find any record of that -- do you remember @mrocklin?

```python
# Any submissions either block, or run as background tasks
cluster.scale(2)
```

Right, so the current behavior is that scale is immediate, and is expected to take very little time.  Typically it sets some target and is quickly done.

> I have a memory of talking to @mrocklin about implementing async scaling when we were working on dask-labextension

To be clear, asynchronous scaling is already doable.  


> One pro of making it an async method is that in asynchronous mode all standard api methods are then also async methods, making this a more uniform interface.

```python
cluster = await ClusterClass()
await cluster.scale(2)
await cluster.close()
```

Some are and some aren't, but yes, most are.  Counter examples include user-facing functions like `submit`, `map`, and `persist`, all of which operate in this fire-and-forget mode.  

If we change the behavior of `scale` (which I'm totally open to) we may want to figure out how to make things backwards compatible.  The `scale` method is commonly used, and it has always been assumed immediate.  (However, async programming is rare, which may save us here).
It might be easy to support either.  If current scale methods returned an empty awaitable object then things would be fine.  Because most cluster managers depend on SpecCluster this might be an easy change to make.

I would be interested in hearing from @jacobtomlinson as to whether or not this would affect dask-kubernetes work positively.
Looking briefly through the `Cluster` API, the other method for which this is relevant is `adapt`.
> Right, so the current behavior is that scale is immediate, and is expected to take very little time. Typically it sets some target and is quickly done.

The `scale` method for dask-gateway also sets a target, but that target is remote on a different server - the `await` here is only to wait for the target to be set. If we were to change this in dask-gateway we'd launch a background task that update the set point in the background. This is less clean than the current implementation, but still doable.

> If current scale methods returned an empty awaitable object then things would be fine.

This seems reasonable-ish. Note that this is compatible only one way. Other cluster managers can be used with or without the `await` on `scale` - `dask-gateway` can only be used with an `await`. I'd worry about code that consumes the cluster interface relying on the former, making the latter not work properly.

---

Anyway, I'd be happy either way, just want to make sure we have a group decision before I make any changes.
Right, if we make this decision we'll have to change the adaptive code and dask-labextension. Anything else?
I think this is a good idea. I've already found myself typing `await cluster.scale(n)` out of habit when in async mode, so this feels like it would be natural.

In the `SpecCluster` rewrite of dask-kubernetes we are only implementing the scale method in order to insert the cluster max shim from #3054. 
Ok. I'll push a fix to dask-labextension making async support for scale optional (at least for now), and will push a PR to distributed making scale support async calls (but not require them?).

Resolving patch:
diff --git a/distributed/deploy/spec.py b/distributed/deploy/spec.py
--- a/distributed/deploy/spec.py
+++ b/distributed/deploy/spec.py
@@ -77,6 +77,19 @@ async def __aexit__(self, *args, **kwargs):
         await self.close()
 
 
+class NoOpAwaitable(object):
+    """An awaitable object that always returns None.
+
+    Useful to return from a method that can be called in both asynchronous and
+    synchronous contexts"""
+
+    def __await__(self):
+        async def f():
+            return None
+
+        return f().__await__()
+
+
 class SpecCluster(Cluster):
     """ Cluster that requires a full specification of workers
 
@@ -411,15 +424,15 @@ def scale(self, n=0, memory=None, cores=None):
         while len(self.worker_spec) > n:
             self.worker_spec.popitem()
 
-        if self.status in ("closing", "closed"):
-            self.loop.add_callback(self._correct_state)
-            return
-
-        while len(self.worker_spec) < n:
-            self.worker_spec.update(self.new_worker_spec())
+        if self.status not in ("closing", "closed"):
+            while len(self.worker_spec) < n:
+                self.worker_spec.update(self.new_worker_spec())
 
         self.loop.add_callback(self._correct_state)
 
+        if self.asynchronous:
+            return NoOpAwaitable()
+
     def new_worker_spec(self):
         """ Return name and spec for the next worker
 


Test patch:
diff --git a/distributed/deploy/tests/test_spec_cluster.py b/distributed/deploy/tests/test_spec_cluster.py
--- a/distributed/deploy/tests/test_spec_cluster.py
+++ b/distributed/deploy/tests/test_spec_cluster.py
@@ -125,6 +125,11 @@ async def test_scale(cleanup):
         await cluster
         assert len(cluster.workers) == 1
 
+        # Can use with await
+        await cluster.scale(2)
+        await cluster
+        assert len(cluster.workers) == 2
+
 
 @pytest.mark.asyncio
 async def test_unexpected_closed_worker(cleanup):