From d2e960385920b245fab8422f8c0480daa6f092f3 Mon Sep 17 00:00:00 2001
From: Matthew Rocklin <mrocklin@gmail.com>
Date: Sun, 28 Jul 2019 18:25:09 -0700
Subject: [PATCH 01/12] Support slowly arriving workers in SpecCluster

Previously SpecCluster waited until all workers had checked in with the
scheduler.  This made sense for LocalCluster or SSHCluster because
there isn't really a significant delay in starting things that we can't
control.  However, for other systems like dask-jobqueue or
dask-kubernetes workers might not ever start, so we need a different
system.

Now, SpecCluster still awaits the Worker object that it is passed, but
doesn't require that the worker has started in the scheduler.  We now
expect awaiting to mean

*"We have successfully handed control of starting the worker to some other robust system"*

Our job at this point is done.  We hope that the worker arrives, but
from our perspective this local Worker object is awaited and "running".

This commit also includes a minimal example of a Worker class,
`SlowWorker`, that serves as a nice minimal example for what SpecCluster
expects.
---
 distributed/deploy/spec.py                    |  3 --
 .../deploy/tests/test_slow_adaptive.py        | 52 +++++++++++++++++++
 distributed/worker.py                         |  2 +
 3 files changed, 54 insertions(+), 3 deletions(-)
 create mode 100644 distributed/deploy/tests/test_slow_adaptive.py

diff --git a/distributed/deploy/spec.py b/distributed/deploy/spec.py
index 87228d5693e..7d0ef417449 100644
--- a/distributed/deploy/spec.py
+++ b/distributed/deploy/spec.py
@@ -192,7 +192,6 @@ def __init__(
             self._loop_runner.start()
             self.sync(self._start)
             self.sync(self._correct_state)
-            self.sync(self._wait_for_workers)
 
     async def _start(self):
         while self.status == "starting":
@@ -306,7 +305,6 @@ async def _():
             await self._correct_state()
             if self.workers:
                 await asyncio.wait(list(self.workers.values()))  # maybe there are more
-            await self._wait_for_workers()
             return self
 
         return _().__await__()
@@ -367,7 +365,6 @@ def __del__(self):
 
     def __enter__(self):
         self.sync(self._correct_state)
-        self.sync(self._wait_for_workers)
         assert self.status == "running"
         return self
 
diff --git a/distributed/deploy/tests/test_slow_adaptive.py b/distributed/deploy/tests/test_slow_adaptive.py
new file mode 100644
index 00000000000..d2d861a0b5c
--- /dev/null
+++ b/distributed/deploy/tests/test_slow_adaptive.py
@@ -0,0 +1,52 @@
+import pytest
+
+from dask.distributed import Worker, Scheduler, SpecCluster, Client
+from distributed.utils_test import cleanup  # noqa: F401
+from distributed.metrics import time
+
+
+class SlowWorker(object):
+    def __init__(self, *args, delay=0, **kwargs):
+        self.worker = Worker(*args, **kwargs)
+        self.delay = delay
+        self.status = None
+
+    @property
+    def address(self):
+        return self.worker.address
+
+    def __await__(self):
+        async def now():
+            if self.status != "running":
+                self.worker.loop.call_later(self.delay, self.worker.start)
+                self.status = "running"
+            return self
+
+        return now().__await__()
+
+    async def close(self):
+        await self.worker.close()
+        self.status = "closed"
+
+
+scheduler = {"cls": Scheduler, "options": {"port": 0}}
+
+
+@pytest.mark.asyncio
+async def test_startup(cleanup):
+    start = time()
+    async with SpecCluster(
+        scheduler=scheduler,
+        workers={
+            0: {"cls": Worker, "options": {}},
+            1: {"cls": SlowWorker, "options": {"delay": 5}},
+            2: {"cls": SlowWorker, "options": {"delay": 0}},
+        },
+        asynchronous=True,
+    ) as cluster:
+        assert len(cluster.workers) == len(cluster.worker_spec) == 3
+        assert time() < start + 5
+        assert 1 <= len(cluster.scheduler_info["workers"]) <= 2
+
+        async with Client(cluster, asynchronous=True) as client:
+            await client.wait_for_workers(n_workers=2)
diff --git a/distributed/worker.py b/distributed/worker.py
index 08927eb741d..d8492aeaeaa 100644
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -906,6 +906,8 @@ async def gather(self, comm=None, who_has=None):
     #############
 
     async def start(self):
+        if self.status and self.status.startswith("clos"):
+            return
         assert self.status is None, self.status
 
         enable_gc_diagnosis()

From e6985eeb19d4c12b89ab105f9333ea631c566fc2 Mon Sep 17 00:00:00 2001
From: Matthew Rocklin <mrocklin@gmail.com>
Date: Sun, 28 Jul 2019 19:16:30 -0700
Subject: [PATCH 02/12] Scale down non-existent workers first

---
 distributed/deploy/spec.py                    |  7 +++++++
 .../deploy/tests/test_slow_adaptive.py        | 21 +++++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/distributed/deploy/spec.py b/distributed/deploy/spec.py
index 7d0ef417449..afb39276b58 100644
--- a/distributed/deploy/spec.py
+++ b/distributed/deploy/spec.py
@@ -373,6 +373,13 @@ def __exit__(self, typ, value, traceback):
         self._loop_runner.stop()
 
     def scale(self, n):
+        if len(self.worker_spec) > n:
+            not_yet_launched = set(self.worker_spec) - {
+                v["name"] for v in self.scheduler_info["workers"].values()
+            }
+            while len(self.worker_spec) > n and not_yet_launched:
+                del self.worker_spec[not_yet_launched.pop()]
+
         while len(self.worker_spec) > n:
             self.worker_spec.popitem()
 
diff --git a/distributed/deploy/tests/test_slow_adaptive.py b/distributed/deploy/tests/test_slow_adaptive.py
index d2d861a0b5c..749bac39554 100644
--- a/distributed/deploy/tests/test_slow_adaptive.py
+++ b/distributed/deploy/tests/test_slow_adaptive.py
@@ -50,3 +50,24 @@ async def test_startup(cleanup):
 
         async with Client(cluster, asynchronous=True) as client:
             await client.wait_for_workers(n_workers=2)
+
+
+@pytest.mark.asyncio
+async def test_scale_up_down(cleanup):
+    start = time()
+    async with SpecCluster(
+        scheduler=scheduler,
+        workers={
+            "slow": {"cls": SlowWorker, "options": {"delay": 5}},
+            "fast": {"cls": Worker, "options": {}},
+        },
+        asynchronous=True,
+    ) as cluster:
+        cluster.scale(1)  # remove a worker, hopefully the one we don't have
+        await cluster
+
+        assert list(cluster.worker_spec) == ["fast"]
+
+        cluster.scale(0)
+        await cluster
+        assert not cluster.worker_spec

From 36d0cab69704e0d3612ac2579842bb4f355db332 Mon Sep 17 00:00:00 2001
From: Matthew Rocklin <mrocklin@gmail.com>
Date: Sun, 28 Jul 2019 20:07:55 -0700
Subject: [PATCH 03/12] wip

---
 distributed/deploy/adaptive.py                | 35 +++++++++----------
 .../deploy/tests/test_slow_adaptive.py        | 27 +++++++++++++-
 2 files changed, 42 insertions(+), 20 deletions(-)

diff --git a/distributed/deploy/adaptive.py b/distributed/deploy/adaptive.py
index 761a7d300ee..4358f8b7ff3 100644
--- a/distributed/deploy/adaptive.py
+++ b/distributed/deploy/adaptive.py
@@ -151,6 +151,7 @@ async def workers_to_close(self, **kwargs):
         if self.maximum is not None and len(self.cluster.workers) > self.maximum:
             kw["n"] = len(self.cluster.workers) - self.maximum
 
+        breakpoint()
         L = await self.scheduler.workers_to_close(**kw)
         if len(self.cluster.workers) - len(L) < self.minimum:
             L = L[: len(self.cluster.workers) - self.minimum]
@@ -183,48 +184,44 @@ async def recommendations(self, comm=None):
             n = min(self.maximum, n)
         if self.minimum is not None:
             n = max(self.minimum, n)
-        workers = set(
-            await self.workers_to_close(
-                key=pickle.dumps(self.worker_key) if self.worker_key else None,
-                minimum=self.minimum,
-            )
-        )
+
+        else:
+            workers_to_close = []
         try:
             current = len(self.cluster.worker_spec)
         except AttributeError:
             current = len(self.cluster.workers)
-        if n > current and workers:
-            logger.info("Attempting to scale up and scale down simultaneously.")
-            self.close_counts.clear()
-            return {
-                "status": "error",
-                "msg": "Trying to scale up and down simultaneously",
-            }
 
+        if n == current:
+            self.close_counts.clear()
+            return None
         elif n > current:
             self.close_counts.clear()
             return {"status": "up", "n": n}
 
-        elif workers:
+        else:
+            workers_to_close = set(
+                await self.workers_to_close(
+                    key=pickle.dumps(self.worker_key) if self.worker_key else None,
+                    minimum=self.minimum,
+                )
+            )
             d = {}
             to_close = []
             for w, c in self.close_counts.items():
-                if w in workers:
+                if w in workers_to_close:
                     if c >= self.wait_count:
                         to_close.append(w)
                     else:
                         d[w] = c
 
-            for w in workers:
+            for w in workers_to_close:
                 d[w] = d.get(w, 0) + 1
 
             self.close_counts = d
 
             if to_close:
                 return {"status": "down", "workers": to_close}
-        else:
-            self.close_counts.clear()
-            return None
 
     async def _adapt(self):
         if self._adapting:  # Semaphore to avoid overlapping adapt calls
diff --git a/distributed/deploy/tests/test_slow_adaptive.py b/distributed/deploy/tests/test_slow_adaptive.py
index 749bac39554..4f565a78289 100644
--- a/distributed/deploy/tests/test_slow_adaptive.py
+++ b/distributed/deploy/tests/test_slow_adaptive.py
@@ -1,7 +1,8 @@
+import asyncio
 import pytest
 
 from dask.distributed import Worker, Scheduler, SpecCluster, Client
-from distributed.utils_test import cleanup  # noqa: F401
+from distributed.utils_test import slowinc, cleanup  # noqa: F401
 from distributed.metrics import time
 
 
@@ -71,3 +72,27 @@ async def test_scale_up_down(cleanup):
         cluster.scale(0)
         await cluster
         assert not cluster.worker_spec
+
+
+@pytest.mark.asyncio
+async def test_adaptive(cleanup):
+    start = time()
+    async with SpecCluster(
+        scheduler=scheduler,
+        workers={"fast": {"cls": Worker, "options": {}}},
+        worker={"cls": SlowWorker, "options": {"delay": 5}},
+        asynchronous=True,
+    ) as cluster:
+        cluster.adapt(minimum=1, maximum=4, target_duration="1s", interval="20ms")
+        async with Client(cluster, asynchronous=True) as client:
+            futures = client.map(slowinc, range(200), delay=0.1)
+
+            while len(cluster.worker_spec) <= 1:
+                await asyncio.sleep(0.05)
+
+            del futures
+
+            while len(cluster.worker_spec) > 1:
+                await asyncio.sleep(0.05)
+
+            assert list(cluster.worker_spec) == ["fast"]

From 473a034ff7177009c132f48d147302df7de70524 Mon Sep 17 00:00:00 2001
From: Matthew Rocklin <mrocklin@gmail.com>
Date: Mon, 29 Jul 2019 11:20:27 -0700
Subject: [PATCH 04/12] Close workers more gracefully

This commit does two things:

1.  We wait to shutdown the executor a little longer in case it is still
    in use
2.  The worker no longer asks the Nanny to terminate it.  Instead it
    asks the nanny to shutdown gracefully after it is gone, and then
    continues closing itself as normal.
---
 distributed/nanny.py            |  3 +++
 distributed/tests/test_nanny.py | 17 ++++++++++++++++-
 distributed/worker.py           | 26 ++++++++++++--------------
 3 files changed, 31 insertions(+), 15 deletions(-)

diff --git a/distributed/nanny.py b/distributed/nanny.py
index 228e37c2839..49522217e86 100644
--- a/distributed/nanny.py
+++ b/distributed/nanny.py
@@ -382,6 +382,9 @@ async def _on_exit(self, exitcode):
                     if self.auto_restart:
                         logger.warning("Restarting worker")
                         await self.instantiate()
+                elif self.status == "closing-gracefully":
+                    await self.close()
+
             except Exception:
                 logger.error(
                     "Failed to restart worker after its process exited", exc_info=True
diff --git a/distributed/tests/test_nanny.py b/distributed/tests/test_nanny.py
index 579af8dbc2c..2f8b01f13d7 100644
--- a/distributed/tests/test_nanny.py
+++ b/distributed/tests/test_nanny.py
@@ -15,7 +15,7 @@
 from tornado.ioloop import IOLoop
 
 import dask
-from distributed import Nanny, rpc, Scheduler, Worker
+from distributed import Nanny, rpc, Scheduler, Worker, Client
 from distributed.core import CommClosedError
 from distributed.metrics import time
 from distributed.protocol.pickle import dumps
@@ -398,3 +398,18 @@ async def test_nanny_closes_cleanly(cleanup):
         assert not n.process
         assert not proc.is_alive()
         assert proc.exitcode == 0
+
+
+@pytest.mark.asyncio
+async def test_nanny_closes_cleanly(cleanup):
+    async with Scheduler() as s:
+        async with Nanny(s.address) as n:
+            async with Client(s.address, asynchronous=True) as client:
+                with client.rpc(n.worker_address) as w:
+                    IOLoop.current().add_callback(w.terminate)
+                    start = time()
+                    while n.status != "closed":
+                        await gen.sleep(0.01)
+                        assert time() < start + 5
+
+                    assert n.status == "closed"
diff --git a/distributed/worker.py b/distributed/worker.py
index ae7ee648ff9..4a562015f96 100644
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -577,7 +577,7 @@ def __init__(
             "get_data": self.get_data,
             "update_data": self.update_data,
             "delete_data": self.delete_data,
-            "terminate": self.terminate,
+            "terminate": self.close,
             "ping": pingpong,
             "upload_file": self.upload_file,
             "start_ipython": self.start_ipython,
@@ -832,7 +832,7 @@ async def handle_scheduler(self, comm):
             logger.exception(e)
             raise
         finally:
-            if self.reconnect:
+            if self.reconnect and self.status == "running":
                 logger.info("Connection to scheduler broken.  Reconnecting...")
                 self.loop.add_callback(self._register_with_scheduler)
             else:
@@ -998,13 +998,6 @@ async def close(self, report=True, timeout=10, nanny=True, executor_wait=True):
                         self.scheduler.unregister(address=self.contact_address),
                     )
             self.scheduler.close_rpc()
-            self.actor_executor._work_queue.queue.clear()
-            if isinstance(self.executor, ThreadPoolExecutor):
-                self.executor._work_queue.queue.clear()
-                self.executor.shutdown(wait=executor_wait, timeout=timeout)
-            else:
-                self.executor.shutdown(wait=False)
-            self.actor_executor.shutdown(wait=executor_wait, timeout=timeout)
             self._workdir.release()
 
             for k, v in self.services.items():
@@ -1016,9 +1009,13 @@ async def close(self, report=True, timeout=10, nanny=True, executor_wait=True):
             if self.batched_stream:
                 self.batched_stream.close()
 
-            if nanny and self.nanny:
-                with self.rpc(self.nanny) as r:
-                    await r.terminate()
+            self.actor_executor._work_queue.queue.clear()
+            if isinstance(self.executor, ThreadPoolExecutor):
+                self.executor._work_queue.queue.clear()
+                self.executor.shutdown(wait=executor_wait, timeout=timeout)
+            else:
+                self.executor.shutdown(wait=False)
+            self.actor_executor.shutdown(wait=executor_wait, timeout=timeout)
 
             self.stop()
             self.rpc.close()
@@ -1028,9 +1025,10 @@ async def close(self, report=True, timeout=10, nanny=True, executor_wait=True):
             await ServerNode.close(self)
 
             setproctitle("dask-worker [closed]")
+        return "OK"
 
-    async def terminate(self, comm, report=True):
-        await self.close(report=report)
+    async def terminate(self, comm, report=True, **kwargs):
+        await self.close(report=report, **kwargs)
         return "OK"
 
     async def wait_until_closed(self):

From e3278759c2cbbf004ef9dd791591c8999f2a0206 Mon Sep 17 00:00:00 2001
From: Matthew Rocklin <mrocklin@gmail.com>
Date: Mon, 29 Jul 2019 18:31:57 -0700
Subject: [PATCH 05/12] Add AdaptiveCore class

---
 distributed/deploy/adaptive_core.py           | 127 ++++++++++++++++++
 .../deploy/tests/test_adaptive_core.py        |  89 ++++++++++++
 2 files changed, 216 insertions(+)
 create mode 100644 distributed/deploy/adaptive_core.py
 create mode 100644 distributed/deploy/tests/test_adaptive_core.py

diff --git a/distributed/deploy/adaptive_core.py b/distributed/deploy/adaptive_core.py
new file mode 100644
index 00000000000..6e6a11f726a
--- /dev/null
+++ b/distributed/deploy/adaptive_core.py
@@ -0,0 +1,127 @@
+import collections
+import math
+
+import toolz
+
+from ..metrics import time
+from ..utils import parse_timedelta, PeriodicCallback
+
+
+class AdaptiveCore:
+    def __init__(
+        self,
+        minimum: float = 0,
+        maximum: float = math.inf,
+        wait_count: int = 3,
+        interval: str = "1s",
+    ):
+        self.minimum = minimum
+        self.maximum = maximum
+        self.wait_count = wait_count
+        self.interval = parse_timedelta(interval, "seconds") if interval else interval
+        self.periodic_callback = None
+
+        self.plan = set()
+        self.requested = set()
+        self.observed = set()
+
+        # internal state
+        self.close_counts = collections.defaultdict(int)
+        self._adapting = False
+        self.log = collections.deque(maxlen=10000)
+
+    def __await__(self):
+        async def f():
+            if self.interval:
+                self.periodic_callback = PeriodicCallback(
+                    self.adapt, self.interval * 1000
+                )
+                self.periodic_callback.start()
+            return self
+
+        return f().__await__()
+
+    def stop(self):
+        if self.periodic_callback:
+            self.periodic_callback.stop()
+            self.periodic_callback = None
+
+    async def target(self) -> int:
+        pass
+
+    async def _workers_to_close(self, target: int) -> list:
+        """
+        Give a list of workers to close that brings us down to target workers
+        """
+        # TODO, improve me with something that thinks about current load
+        return list(self.observed)[target:]
+
+    async def safe_target(self) -> int:
+        n = await self.target()
+        if n > self.maximum:
+            n = self.maximum
+
+        if n < self.minimum:
+            n = self.minimum
+
+        return n
+
+    async def recommendations(self, target: int) -> dict:
+        plan = self.plan
+        requested = self.requested
+        observed = self.observed
+
+        if target == len(plan):
+            self.close_counts.clear()
+            return {"status": "same"}
+
+        elif target > len(plan):
+            self.close_counts.clear()
+            return {"status": "up", "n": target}
+
+        elif target < len(plan):
+            not_yet_arrived = requested - observed
+            to_close = set()
+            if not_yet_arrived:
+                to_close.update((toolz.take(len(plan) - target, not_yet_arrived)))
+
+            if target < len(plan) - len(to_close):
+                L = await self._workers_to_close(target=target)
+                to_close.update(L)
+
+            firmly_close = set()
+            for w in to_close:
+                self.close_counts[w] += 1
+                if self.close_counts[w] >= self.wait_count:
+                    firmly_close.add(w)
+
+            for k in list(self.close_counts):  # clear out unseen keys
+                if k in firmly_close or k not in to_close:
+                    del self.close_counts[k]
+
+            if firmly_close:
+                return {"status": "down", "workers": firmly_close}
+            else:
+                return {"status": "same"}
+
+    async def adapt(self) -> None:
+        if self._adapting:  # Semaphore to avoid overlapping adapt calls
+            return
+        self._adapting = True
+
+        try:
+            target = await self.safe_target()
+            recommendations = await self.recommendations(target)
+
+            if recommendations["status"] != "same":
+                self.log.append((time(), dict(recommendations)))
+
+            status = recommendations.pop("status")
+            if status == "same":
+                return
+            if status == "up":
+                await self.scale_up(**recommendations)
+            if status == "down":
+                await self.scale_down(**recommendations)
+        finally:
+            self._adapting = False
diff --git a/distributed/deploy/tests/test_adaptive_core.py b/distributed/deploy/tests/test_adaptive_core.py
new file mode 100644
index 00000000000..e01977700a9
--- /dev/null
+++ b/distributed/deploy/tests/test_adaptive_core.py
@@ -0,0 +1,89 @@
+import asyncio
+import pytest
+
+from distributed.deploy.adaptive_core import AdaptiveCore
+from distributed.metrics import time
+
+
+class MyAdaptive(AdaptiveCore):
+    def __init__(self, *args, interval=None, **kwargs):
+        super().__init__(*args, interval=interval, **kwargs)
+        self._target = 0
+        self._log = []
+
+    async def target(self):
+        return self._target
+
+    async def scale_up(self, n=0):
+        self.plan = self.requested = set(range(n))
+
+    async def scale_down(self, workers=()):
+        for collection in [self.plan, self.requested, self.observed]:
+            for w in workers:
+                collection.discard(w)
+
+
+@pytest.mark.asyncio
+async def test_safe_target():
+    adapt = MyAdaptive(minimum=1, maximum=4)
+    assert await adapt.safe_target() == 1
+    adapt._target = 10
+    assert await adapt.safe_target() == 4
+
+
+@pytest.mark.asyncio
+async def test_scale_up():
+    adapt = MyAdaptive(minimum=1, maximum=4)
+    await adapt.adapt()
+    assert adapt.log[-1][1] == {"status": "up", "n": 1}
+    assert adapt.plan == {0}
+
+    adapt._target = 10
+    await adapt.adapt()
+    assert adapt.log[-1][1] == {"status": "up", "n": 4}
+    assert adapt.plan == {0, 1, 2, 3}
+
+
+@pytest.mark.asyncio
+async def test_scale_down():
+    adapt = MyAdaptive(minimum=1, maximum=4, wait_count=2)
+    adapt._target = 10
+    await adapt.adapt()
+    assert len(adapt.log) == 1
+
+    adapt.observed = {0, 1, 3}  # all but 2 have arrived
+
+    adapt._target = 2
+    await adapt.adapt()
+    assert len(adapt.log) == 1  # no change after only one call
+    await adapt.adapt()
+    assert len(adapt.log) == 2  # no change after only one call
+    assert adapt.log[-1][1]["status"] == "down"
+    assert 2 in adapt.log[-1][1]["workers"]
+    assert len(adapt.log[-1][1]["workers"]) == 2
+
+    old = list(adapt.log)
+    await adapt.adapt()
+    await adapt.adapt()
+    await adapt.adapt()
+    await adapt.adapt()
+    assert list(adapt.log) == old
+
+
+@pytest.mark.asyncio
+async def test_interval():
+    adapt = await MyAdaptive(interval="5 ms")
+    assert not adapt.plan
+
+    for i in [0, 3, 1]:
+        start = time()
+        adapt._target = i
+        while len(adapt.plan) != i:
+            await asyncio.sleep(0.001)
+            assert time() < start + 2
+
+    adapt.stop()
+
+    adapt._target = 10
+    await asyncio.sleep(0.020)
+    assert len(adapt.plan) == 1  # last value from before, unchanged

From 38786320174cbb5353281889d5caf7c721b4af84 Mon Sep 17 00:00:00 2001
From: Matthew Rocklin <mrocklin@gmail.com>
Date: Mon, 29 Jul 2019 19:58:12 -0700
Subject: [PATCH 06/12] Back Adaptive with AdaptiveCore

---
 distributed/deploy/adaptive.py                | 172 +++++-------------
 distributed/deploy/adaptive_core.py           |  41 +++--
 distributed/deploy/tests/test_adaptive.py     | 120 ++++++------
 .../deploy/tests/test_adaptive_core.py        |   2 +-
 distributed/scheduler.py                      |  17 +-
 5 files changed, 130 insertions(+), 222 deletions(-)

diff --git a/distributed/deploy/adaptive.py b/distributed/deploy/adaptive.py
index 4358f8b7ff3..2d7641d3c16 100644
--- a/distributed/deploy/adaptive.py
+++ b/distributed/deploy/adaptive.py
@@ -1,18 +1,16 @@
 from __future__ import print_function, division, absolute_import
 
-from collections import deque
 import logging
+import math
 
-from tornado import gen
-
-from ..metrics import time
-from ..utils import log_errors, PeriodicCallback, parse_timedelta
+from .adaptive_core import AdaptiveCore
+from ..utils import log_errors, parse_timedelta
 from ..protocol import pickle
 
 logger = logging.getLogger(__name__)
 
 
-class Adaptive(object):
+class Adaptive(AdaptiveCore):
     '''
     Adaptively allocate workers based on scheduler load.  A superclass.
 
@@ -25,19 +23,13 @@ class Adaptive(object):
 
     Parameters
     ----------
-    scheduler: distributed.Scheduler
     cluster: object
-        Must have scale_up and scale_down methods/coroutines
-    startup_cost : timedelta or str, default "1s"
-        Estimate of the number of seconds for nnFactor representing how costly it is to start an additional worker.
-        Affects quickly to adapt to high tasks per worker loads
+        Must have scale and scale_down methods/coroutines
     interval : timedelta or str, default "1000 ms"
         Milliseconds between checks
     wait_count: int, default 3
         Number of consecutive times that a worker should be suggested for
         removal before we remove it.
-    scale_factor : int, default 2
-        Factor to scale by when it's determined additional workers are needed
     target_duration: timedelta or str, default "5s"
         Amount of time we want a computation to take.
         This affects how aggressively we scale up.
@@ -86,45 +78,47 @@ def __init__(
         self,
         cluster=None,
         interval="1s",
-        startup_cost="1s",
-        scale_factor=2,
         minimum=0,
-        maximum=None,
+        maximum=math.inf,
         wait_count=3,
         target_duration="5s",
         worker_key=None,
         **kwargs
     ):
-        interval = parse_timedelta(interval, default="ms")
-        self.worker_key = worker_key
         self.cluster = cluster
-        self.startup_cost = parse_timedelta(startup_cost, default="s")
-        self.scale_factor = scale_factor
-        if self.cluster:
-            self._adapt_callback = PeriodicCallback(
-                self._adapt, interval * 1000, io_loop=self.loop
-            )
-            self.loop.add_callback(self._adapt_callback.start)
-        self._adapting = False
+        self.worker_key = worker_key
         self._workers_to_close_kwargs = kwargs
-        self.minimum = minimum
-        self.maximum = maximum
-        self.log = deque(maxlen=1000)
-        self.close_counts = {}
-        self.wait_count = wait_count
         self.target_duration = parse_timedelta(target_duration)
 
+        super().__init__(
+            minimum=minimum, maximum=maximum, wait_count=wait_count, interval=interval
+        )
+
     @property
     def scheduler(self):
         return self.cluster.scheduler_comm
 
-    def stop(self):
-        if self.cluster:
-            self._adapt_callback.stop()
-            self._adapt_callback = None
-            del self._adapt_callback
+    @property
+    def plan(self):
+        try:
+            return set(self.cluster.worker_spec)
+        except AttributeError:
+            return set(self.cluster.workers)
 
-    async def workers_to_close(self, **kwargs):
+    @property
+    def requested(self):
+        return set(self.cluster.workers)
+
+    @property
+    def observed(self):
+        return {d["name"] for d in self.cluster.scheduler_info["workers"].values()}
+
+    async def target(self):
+        return await self.scheduler.adaptive_target(
+            target_duration=self.target_duration
+        )
+
+    async def workers_to_close(self, target: int):
         """
         Determine which, if any, workers should potentially be removed from
         the cluster.
@@ -142,111 +136,27 @@ async def workers_to_close(self, **kwargs):
         --------
         Scheduler.workers_to_close
         """
-        if len(self.cluster.workers) <= self.minimum:
-            return []
-
-        kw = dict(self._workers_to_close_kwargs)
-        kw.update(kwargs)
-
-        if self.maximum is not None and len(self.cluster.workers) > self.maximum:
-            kw["n"] = len(self.cluster.workers) - self.maximum
+        return await self.scheduler.workers_to_close(
+            target=target,
+            key=pickle.dumps(self.worker_key) if self.worker_key else None,
+            **self._workers_to_close_kwargs
+        )
 
-        breakpoint()
-        L = await self.scheduler.workers_to_close(**kw)
-        if len(self.cluster.workers) - len(L) < self.minimum:
-            L = L[: len(self.cluster.workers) - self.minimum]
-
-        return L
-
-    async def _retire_workers(self, workers=None):
-        if workers is None:
-            workers = await self.workers_to_close(
-                key=pickle.dumps(self.worker_key) if self.worker_key else None,
-                minimum=self.minimum,
-            )
-        if not workers:
-            raise gen.Return(workers)
+    async def scale_down(self, workers):
         with log_errors():
+            # Ask scheduler to cleanly retire workers
             await self.scheduler.retire_workers(
                 workers=workers, remove=True, close_workers=True
             )
 
+            # close workers more forcefully
             logger.info("Retiring workers %s", workers)
             f = self.cluster.scale_down(workers)
             if hasattr(f, "__await__"):
                 await f
 
-            return workers
-
-    async def recommendations(self, comm=None):
-        n = await self.scheduler.adaptive_target(target_duration=self.target_duration)
-        if self.maximum is not None:
-            n = min(self.maximum, n)
-        if self.minimum is not None:
-            n = max(self.minimum, n)
-
-        else:
-            workers_to_close = []
-        try:
-            current = len(self.cluster.worker_spec)
-        except AttributeError:
-            current = len(self.cluster.workers)
-
-        if n == current:
-            self.close_counts.clear()
-            return None
-        elif n > current:
-            self.close_counts.clear()
-            return {"status": "up", "n": n}
-
-        else:
-            workers_to_close = set(
-                await self.workers_to_close(
-                    key=pickle.dumps(self.worker_key) if self.worker_key else None,
-                    minimum=self.minimum,
-                )
-            )
-            d = {}
-            to_close = []
-            for w, c in self.close_counts.items():
-                if w in workers_to_close:
-                    if c >= self.wait_count:
-                        to_close.append(w)
-                    else:
-                        d[w] = c
-
-            for w in workers_to_close:
-                d[w] = d.get(w, 0) + 1
-
-            self.close_counts = d
-
-            if to_close:
-                return {"status": "down", "workers": to_close}
-
-    async def _adapt(self):
-        if self._adapting:  # Semaphore to avoid overlapping adapt calls
-            return
-
-        self._adapting = True
-        try:
-            recommendations = await self.recommendations()
-            if not recommendations:
-                return
-            status = recommendations.pop("status")
-            if status == "up":
-                f = self.cluster.scale(**recommendations)
-                self.log.append((time(), "up", recommendations))
-                if hasattr(f, "__await__"):
-                    await f
-
-            elif status == "down":
-                self.log.append((time(), "down", recommendations["workers"]))
-                workers = await self._retire_workers(workers=recommendations["workers"])
-        finally:
-            self._adapting = False
-
-    def adapt(self):
-        self.loop.add_callback(self._adapt)
+    async def scale_up(self, n):
+        self.cluster.scale(n)
 
     @property
     def loop(self):
diff --git a/distributed/deploy/adaptive_core.py b/distributed/deploy/adaptive_core.py
index 6e6a11f726a..b1a6424833c 100644
--- a/distributed/deploy/adaptive_core.py
+++ b/distributed/deploy/adaptive_core.py
@@ -1,6 +1,7 @@
 import collections
 import math
 
+from tornado.ioloop import IOLoop
 import toolz
 
 from ..metrics import time
@@ -21,35 +22,37 @@ def __init__(
         self.interval = parse_timedelta(interval, "seconds") if interval else interval
         self.periodic_callback = None
 
-        self.plan = set()
-        self.requested = set()
-        self.observed = set()
+        def f():
+            self.periodic_callback = PeriodicCallback(self.adapt, self.interval * 1000)
+            self.periodic_callback.start()
+
+        if self.interval:
+            try:
+                self.loop.add_callback(f)
+            except AttributeError:
+                IOLoop.current().add_callback(f)
+
+        try:
+            self.plan = set()
+            self.requested = set()
+            self.observed = set()
+        except Exception:
+            pass
 
         # internal state
         self.close_counts = collections.defaultdict(int)
         self._adapting = False
         self.log = collections.deque(maxlen=10000)
 
-    def __await__(self):
-        async def f():
-            if self.interval:
-                self.periodic_callback = PeriodicCallback(
-                    self.adapt, self.interval * 1000
-                )
-                self.periodic_callback.start()
-            return self
-
-        return f().__await__()
-
     def stop(self):
         if self.periodic_callback:
             self.periodic_callback.stop()
             self.periodic_callback = None
 
     async def target(self) -> int:
-        pass
+        raise NotImplementedError()
 
-    async def _workers_to_close(self, target: int) -> list:
+    async def workers_to_close(self, target: int) -> list:
         """
         Give a list of workers to close that brings us down to target workers
         """
@@ -86,7 +89,7 @@ async def recommendations(self, target: int) -> dict:
                 to_close.update((toolz.take(len(plan) - target, not_yet_arrived)))
 
             if target < len(plan) - len(to_close):
-                L = await self._workers_to_close(target=target)
+                L = await self.workers_to_close(target=target)
                 to_close.update(L)
 
             firmly_close = set()
@@ -100,7 +103,7 @@ async def recommendations(self, target: int) -> dict:
                     del self.close_counts[k]
 
             if firmly_close:
-                return {"status": "down", "workers": firmly_close}
+                return {"status": "down", "workers": list(firmly_close)}
             else:
                 return {"status": "same"}
 
@@ -123,5 +126,7 @@ async def adapt(self) -> None:
                 await self.scale_up(**recommendations)
             if status == "down":
                 await self.scale_down(**recommendations)
+        except OSError:
+            self.stop()
         finally:
             self._adapting = False
diff --git a/distributed/deploy/tests/test_adaptive.py b/distributed/deploy/tests/test_adaptive.py
index e0478a9cbdb..375683372be 100644
--- a/distributed/deploy/tests/test_adaptive.py
+++ b/distributed/deploy/tests/test_adaptive.py
@@ -3,7 +3,6 @@
 from time import sleep
 
 import pytest
-from toolz import frequencies, pluck
 from tornado import gen
 from tornado.ioloop import IOLoop
 
@@ -70,51 +69,46 @@ def test_adaptive_local_cluster(loop):
             assert not c.nthreads()
 
 
-@nodebug
-@gen_test(timeout=30)
-def test_adaptive_local_cluster_multi_workers():
-    cluster = yield LocalCluster(
+@pytest.mark.asyncio
+async def test_adaptive_local_cluster_multi_workers(cleanup):
+    async with LocalCluster(
         0,
         scheduler_port=0,
         silence_logs=False,
         processes=False,
         dashboard_address=None,
         asynchronous=True,
-    )
-    try:
-        cluster.scheduler.allowed_failures = 1000
-        alc = cluster.adapt(interval=100)
-        c = yield Client(cluster, asynchronous=True)
-
-        futures = c.map(slowinc, range(100), delay=0.01)
+    ) as cluster:
 
-        start = time()
-        while not cluster.scheduler.workers:
-            yield gen.sleep(0.01)
-            assert time() < start + 15, alc.log
+        cluster.scheduler.allowed_failures = 1000
+        adapt = cluster.adapt(interval="100 ms")
+        async with Client(cluster, asynchronous=True) as c:
+            futures = c.map(slowinc, range(100), delay=0.01)
 
-        yield c.gather(futures)
-        del futures
+            start = time()
+            while not cluster.scheduler.workers:
+                await gen.sleep(0.01)
+                assert time() < start + 15, adapt.log
 
-        start = time()
-        # while cluster.workers:
-        while cluster.scheduler.workers:
-            yield gen.sleep(0.01)
-            assert time() < start + 15, alc.log
+            await c.gather(futures)
+            del futures
 
-        # no workers for a while
-        for i in range(10):
-            assert not cluster.scheduler.workers
-            yield gen.sleep(0.05)
+            start = time()
+            # while cluster.workers:
+            while cluster.scheduler.workers:
+                await gen.sleep(0.01)
+                assert time() < start + 15, adapt.log
 
-        futures = c.map(slowinc, range(100), delay=0.01)
-        yield c.gather(futures)
+            # no workers for a while
+            for i in range(10):
+                assert not cluster.scheduler.workers
+                await gen.sleep(0.05)
 
-    finally:
-        yield c.close()
-        yield cluster.close()
+            futures = c.map(slowinc, range(100), delay=0.01)
+            await c.gather(futures)
 
 
+@pytest.mark.xfail(reason="changed API")
 @pytest.mark.asyncio
 async def test_adaptive_scale_down_override(cleanup):
     class TestAdaptive(Adaptive):
@@ -166,7 +160,7 @@ def test_min_max():
 
         yield gen.sleep(0.2)
         assert len(cluster.scheduler.workers) == 1
-        assert frequencies(pluck(1, adapt.log)) == {"up": 1}
+        assert len(adapt.log) == 1 and adapt.log[-1][1] == {"status": "up", "n": 1}
 
         futures = c.map(slowinc, range(100), delay=0.1)
 
@@ -179,7 +173,7 @@ def test_min_max():
         yield gen.sleep(0.5)
         assert len(cluster.scheduler.workers) == 2
         assert len(cluster.workers) == 2
-        assert frequencies(pluck(1, adapt.log)) == {"up": 2}
+        assert len(adapt.log) == 2 and all(d["status"] == "up" for _, d in adapt.log)
 
         del futures
 
@@ -187,41 +181,35 @@ def test_min_max():
         while len(cluster.scheduler.workers) != 1:
             yield gen.sleep(0.01)
             assert time() < start + 2
-        assert frequencies(pluck(1, adapt.log)) == {"up": 2, "down": 1}
+        assert adapt.log[-1][1]["status"] == "down"
     finally:
         yield c.close()
         yield cluster.close()
 
 
-@gen_test()
-def test_avoid_churn():
+@pytest.mark.asyncio
+async def test_avoid_churn(cleanup):
     """ We want to avoid creating and deleting workers frequently
 
     Instead we want to wait a few beats before removing a worker in case the
     user is taking a brief pause between work
     """
-    cluster = yield LocalCluster(
+    async with LocalCluster(
         0,
         asynchronous=True,
         processes=False,
         scheduler_port=0,
         silence_logs=False,
         dashboard_address=None,
-    )
-    client = yield Client(cluster, asynchronous=True)
-    try:
-        adapt = cluster.adapt(interval="20 ms", wait_count=5)
+    ) as cluster:
+        async with Client(cluster, asynchronous=True) as client:
+            adapt = cluster.adapt(interval="20 ms", wait_count=5)
 
-        for i in range(10):
-            yield client.submit(slowinc, i, delay=0.040)
-            yield gen.sleep(0.040)
+            for i in range(10):
+                await client.submit(slowinc, i, delay=0.040)
+                await gen.sleep(0.040)
 
-        from toolz.curried import pipe, unique, pluck, frequencies
-
-        assert pipe(adapt.log, unique(key=str), pluck(1), frequencies) == {"up": 1}
-    finally:
-        yield client.close()
-        yield cluster.close()
+            assert len(adapt.log) == 1
 
 
 @gen_test(timeout=None)
@@ -240,7 +228,7 @@ def test_adapt_quickly():
         dashboard_address=None,
     )
     client = yield Client(cluster, asynchronous=True)
-    adapt = cluster.adapt(interval=20, wait_count=5, maximum=10)
+    adapt = cluster.adapt(interval="20 ms", wait_count=5, maximum=10)
     try:
         future = client.submit(slowinc, 1, delay=0.100)
         yield wait(future)
@@ -248,10 +236,10 @@ def test_adapt_quickly():
 
         # Scale up when there is plenty of available work
         futures = client.map(slowinc, range(1000), delay=0.100)
-        while frequencies(pluck(1, adapt.log)) == {"up": 1}:
+        while len(adapt.log) == 1:
             yield gen.sleep(0.01)
         assert len(adapt.log) == 2
-        assert "up" in adapt.log[-1]
+        assert adapt.log[-1][1]["status"] == "up"
         d = [x for x in adapt.log[-1] if isinstance(x, dict)][0]
         assert 2 < d["n"] <= adapt.maximum
 
@@ -364,7 +352,7 @@ def test_target_duration():
         dashboard_address=None,
     )
     client = yield Client(cluster, asynchronous=True)
-    adaptive = cluster.adapt(interval="20ms", minimum=2, target_duration="5s")
+    adapt = cluster.adapt(interval="20ms", minimum=2, target_duration="5s")
 
     cluster.scheduler.task_duration["slowinc"] = 1
 
@@ -374,21 +362,21 @@ def test_target_duration():
 
         futures = client.map(slowinc, range(100), delay=0.3)
 
-        while len(adaptive.log) < 2:
+        while len(adapt.log) < 2:
             yield gen.sleep(0.01)
 
-        assert adaptive.log[0][1:] == ("up", {"n": 2})
-        assert adaptive.log[1][1:] == ("up", {"n": 20})
+        assert adapt.log[0][1] == {"status": "up", "n": 2}
+        assert adapt.log[1][1] == {"status": "up", "n": 20}
 
     finally:
         yield client.close()
         yield cluster.close()
 
 
-@gen_test(timeout=None)
-def test_worker_keys():
+@pytest.mark.asyncio
+async def test_worker_keys(cleanup):
     """ Ensure that redefining adapt with a lower maximum removes workers """
-    cluster = yield SpecCluster(
+    async with SpecCluster(
         workers={
             "a-1": {"cls": Worker},
             "a-2": {"cls": Worker},
@@ -396,9 +384,7 @@ def test_worker_keys():
             "b-2": {"cls": Worker},
         },
         asynchronous=True,
-    )
-
-    try:
+    ) as cluster:
 
         def key(ws):
             return ws.name.split("-")[0]
@@ -406,12 +392,10 @@ def key(ws):
         cluster._adaptive_options = {"worker_key": key}
 
         adaptive = cluster.adapt(minimum=1)
-        yield adaptive._adapt()
+        await adaptive.adapt()
 
         while len(cluster.scheduler.workers) == 4:
-            yield gen.sleep(0.01)
+            await gen.sleep(0.01)
 
         names = {ws.name for ws in cluster.scheduler.workers.values()}
         assert names == {"a-1", "a-2"} or names == {"b-1", "b-2"}
-    finally:
-        yield cluster.close()
diff --git a/distributed/deploy/tests/test_adaptive_core.py b/distributed/deploy/tests/test_adaptive_core.py
index e01977700a9..5229f0ef65a 100644
--- a/distributed/deploy/tests/test_adaptive_core.py
+++ b/distributed/deploy/tests/test_adaptive_core.py
@@ -72,7 +72,7 @@ async def test_scale_down():
 
 @pytest.mark.asyncio
 async def test_interval():
-    adapt = await MyAdaptive(interval="5 ms")
+    adapt = MyAdaptive(interval="5 ms")
     assert not adapt.plan
 
     for i in [0, 3, 1]:
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
index b0db6653d2b..695a47e578a 100644
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -2904,7 +2904,7 @@ def replicate(
         )
 
     def workers_to_close(
-        self, comm=None, memory_ratio=None, n=None, key=None, minimum=None
+        self, comm=None, memory_ratio=None, n=None, key=None, minimum=None, target=None
     ):
         """
         Find workers that we can close with low cost
@@ -2931,6 +2931,8 @@ def workers_to_close(
             An optional callable mapping a WorkerState object to a group
             affiliation.  Groups will be closed together.  This is useful when
             closing workers must be done collectively, such as by hostname.
+        target: int
+            Target number of workers to have after we close
 
         Examples
         --------
@@ -2958,6 +2960,13 @@ def workers_to_close(
         --------
         Scheduler.retire_workers
         """
+        if target is not None and n is None:
+            n = len(self.workers) - target
+        if n:
+            if n < 0:
+                n = 0
+            target = len(self.workers) - n
+
         if n is None and memory_ratio is None:
             memory_ratio = 2
 
@@ -2982,12 +2991,12 @@ def workers_to_close(
             limit = sum(limit_bytes.values())
             total = sum(group_bytes.values())
 
-            def key(group):
+            def _key(group):
                 is_idle = not any(ws.processing for ws in groups[group])
                 bytes = -group_bytes[group]
                 return (is_idle, bytes)
 
-            idle = sorted(groups, key=key)
+            idle = sorted(groups, key=_key)
 
             to_close = []
             n_remain = len(self.workers)
@@ -3002,7 +3011,7 @@ def key(group):
 
                 limit -= limit_bytes[group]
 
-                if (n is not None and len(to_close) < n) or (
+                if (n is not None and n_remain - len(groups[group]) >= target) or (
                     memory_ratio is not None and limit >= memory_ratio * total
                 ):
                     to_close.append(group)

From 8c8ce8443f5a306d0aeef12605bb92bd64de2895 Mon Sep 17 00:00:00 2001
From: Matthew Rocklin <mrocklin@gmail.com>
Date: Tue, 30 Jul 2019 10:31:16 -0700
Subject: [PATCH 07/12] Include requested workers in widget

---
 distributed/deploy/spec.py                    | 7 +++++--
 distributed/deploy/tests/test_spec_cluster.py | 8 +++++++-
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/distributed/deploy/spec.py b/distributed/deploy/spec.py
index afb39276b58..63620487bf1 100644
--- a/distributed/deploy/spec.py
+++ b/distributed/deploy/spec.py
@@ -477,6 +477,7 @@ def dashboard_link(self):
 
     def _widget_status(self):
         workers = len(self.scheduler_info["workers"])
+        requested = len(self.worker_spec)
         cores = sum(v["nthreads"] for v in self.scheduler_info["workers"].values())
         memory = sum(v["memory_limit"] for v in self.scheduler_info["workers"].values())
         memory = format_bytes(memory)
@@ -496,13 +497,13 @@ def _widget_status(self):
     }
   </style>
   <table style="text-align: right;">
-    <tr><th>Workers</th> <td>%d</td></tr>
+    <tr><th>Workers</th> <td>%s</td></tr>
     <tr><th>Cores</th> <td>%d</td></tr>
     <tr><th>Memory</th> <td>%s</td></tr>
   </table>
 </div>
 """ % (
-            workers,
+            workers if workers == requested else "%d / %d" % (workers, requested),
             cores,
             memory,
         )
@@ -551,6 +552,7 @@ def _widget(self):
 
             def adapt_cb(b):
                 self.adapt(minimum=minimum.value, maximum=maximum.value)
+                update()
 
             adapt.on_click(adapt_cb)
 
@@ -560,6 +562,7 @@ def scale_cb(b):
                     with ignoring(AttributeError):
                         self._adaptive.stop()
                     self.scale(n)
+                    update()
 
             scale.on_click(scale_cb)
         else:
diff --git a/distributed/deploy/tests/test_spec_cluster.py b/distributed/deploy/tests/test_spec_cluster.py
index 84c868b2585..c755a628879 100644
--- a/distributed/deploy/tests/test_spec_cluster.py
+++ b/distributed/deploy/tests/test_spec_cluster.py
@@ -250,7 +250,10 @@ async def test_dashboard_link(cleanup):
 @pytest.mark.asyncio
 async def test_widget(cleanup):
     async with SpecCluster(
-        workers=worker_spec, scheduler=scheduler, asynchronous=True
+        workers=worker_spec,
+        scheduler=scheduler,
+        asynchronous=True,
+        worker={"cls": Worker, "options": {"nthreads": 1}},
     ) as cluster:
 
         start = time()  # wait for all workers
@@ -260,3 +263,6 @@ async def test_widget(cleanup):
 
         assert "3" in cluster._widget_status()
         assert "GB" in cluster._widget_status()
+
+        cluster.scale(5)
+        assert "3 / 5" in cluster._widget_status()

From 2be16dcf61e4f1d8772b4cc5d7fcad07f6f598e9 Mon Sep 17 00:00:00 2001
From: Matthew Rocklin <mrocklin@gmail.com>
Date: Tue, 30 Jul 2019 11:20:52 -0700
Subject: [PATCH 08/12] Use worker names throughout adaptive

---
 distributed/deploy/adaptive.py |  5 +++-
 distributed/deploy/spec.py     |  9 +++----
 distributed/scheduler.py       | 48 +++++++++++++++++++++++++---------
 3 files changed, 42 insertions(+), 20 deletions(-)

diff --git a/distributed/deploy/adaptive.py b/distributed/deploy/adaptive.py
index 2d7641d3c16..70167401ab4 100644
--- a/distributed/deploy/adaptive.py
+++ b/distributed/deploy/adaptive.py
@@ -139,14 +139,17 @@ async def workers_to_close(self, target: int):
         return await self.scheduler.workers_to_close(
             target=target,
             key=pickle.dumps(self.worker_key) if self.worker_key else None,
+            attribute="name",
             **self._workers_to_close_kwargs
         )
 
     async def scale_down(self, workers):
+        if not workers:
+            return
         with log_errors():
             # Ask scheduler to cleanly retire workers
             await self.scheduler.retire_workers(
-                workers=workers, remove=True, close_workers=True
+                names=workers, remove=True, close_workers=True
             )
 
             # close workers more forcefully
diff --git a/distributed/deploy/spec.py b/distributed/deploy/spec.py
index 63620487bf1..bd2146c5513 100644
--- a/distributed/deploy/spec.py
+++ b/distributed/deploy/spec.py
@@ -415,12 +415,9 @@ def _supports_scaling(self):
         return not not self.new_spec
 
     async def scale_down(self, workers):
-        workers = set(workers)
-
-        for k, v in self.workers.items():
-            if getattr(v, "worker_address", v.address) in workers:
-                del self.worker_spec[k]
-
+        for w in workers:
+            if w in self.worker_spec:
+                del self.worker_spec[w]
         await self
 
     scale_up = scale  # backwards compatibility
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
index 695a47e578a..33fc9012c54 100644
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -2904,7 +2904,14 @@ def replicate(
         )
 
     def workers_to_close(
-        self, comm=None, memory_ratio=None, n=None, key=None, minimum=None, target=None
+        self,
+        comm=None,
+        memory_ratio=None,
+        n=None,
+        key=None,
+        minimum=None,
+        target=None,
+        attribute="address",
     ):
         """
         Find workers that we can close with low cost
@@ -2933,6 +2940,9 @@ def workers_to_close(
             closing workers must be done collectively, such as by hostname.
         target: int
             Target number of workers to have after we close
+        attribute : str
+            The attribute of the WorkerState object to return, like "address"
+            or "name".  Defaults to "address".
 
         Examples
         --------
@@ -2962,7 +2972,7 @@ def workers_to_close(
         """
         if target is not None and n is None:
             n = len(self.workers) - target
-        if n:
+        if n is not None:
             if n < 0:
                 n = 0
             target = len(self.workers) - n
@@ -3020,22 +3030,30 @@ def _key(group):
                 else:
                     break
 
-            result = [ws.address for g in to_close for ws in groups[g]]
+            result = [getattr(ws, attribute) for g in to_close for ws in groups[g]]
             if result:
                 logger.info("Suggest closing workers: %s", result)
 
             return result
 
     async def retire_workers(
-        self, comm=None, workers=None, remove=True, close_workers=False, **kwargs
+        self,
+        comm=None,
+        workers=None,
+        remove=True,
+        close_workers=False,
+        names=None,
+        **kwargs
     ):
         """ Gracefully retire workers from cluster
 
         Parameters
         ----------
         workers: list (optional)
-            List of worker IDs to retire.
+            List of worker addresses to retire.
             If not provided we call ``workers_to_close`` which finds a good set
+        workers_names: list (optional)
+            List of worker names to retire.
         remove: bool (defaults to True)
             Whether or not to remove the worker metadata immediately or else
             wait for the worker to contact us
@@ -3057,6 +3075,11 @@ async def retire_workers(
         Scheduler.workers_to_close
         """
         with log_errors():
+            if names is not None:
+                names = set(names)
+                workers = [
+                    ws.address for ws in self.workers.values() if ws.name in names
+                ]
             if workers is None:
                 while True:
                     try:
@@ -3067,17 +3090,16 @@ async def retire_workers(
                                 remove=remove,
                                 close_workers=close_workers,
                             )
-                        raise gen.Return(workers)
+                        return workers
                     except KeyError:  # keys left during replicate
                         pass
-
             workers = {self.workers[w] for w in workers if w in self.workers}
-            if len(workers) > 0:
-                # Keys orphaned by retiring those workers
-                keys = set.union(*[w.has_what for w in workers])
-                keys = {ts.key for ts in keys if ts.who_has.issubset(workers)}
-            else:
-                keys = set()
+            if not workers:
+                return []
+
+            # Keys orphaned by retiring those workers
+            keys = set.union(*[w.has_what for w in workers])
+            keys = {ts.key for ts in keys if ts.who_has.issubset(workers)}
 
             other_workers = set(self.workers.values()) - workers
             if keys:

From 31d79d8aa27ea89cd2c479302a2924d8bcea1f55 Mon Sep 17 00:00:00 2001
From: Matthew Rocklin <mrocklin@gmail.com>
Date: Tue, 30 Jul 2019 13:25:26 -0700
Subject: [PATCH 09/12] fixup test_interval

---
 distributed/deploy/tests/test_adaptive_core.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/distributed/deploy/tests/test_adaptive_core.py b/distributed/deploy/tests/test_adaptive_core.py
index 5229f0ef65a..a073314223d 100644
--- a/distributed/deploy/tests/test_adaptive_core.py
+++ b/distributed/deploy/tests/test_adaptive_core.py
@@ -83,6 +83,7 @@ async def test_interval():
             assert time() < start + 2
 
     adapt.stop()
+    await asyncio.sleep(0.050)
 
     adapt._target = 10
     await asyncio.sleep(0.020)

From 8ff1ed18193909c1bca5a59ff8557735915ccfda Mon Sep 17 00:00:00 2001
From: Matthew Rocklin <mrocklin@gmail.com>
Date: Tue, 30 Jul 2019 14:42:43 -0700
Subject: [PATCH 10/12] Add docstrings to AdaptiveCore

---
 distributed/deploy/adaptive_core.py | 69 ++++++++++++++++++++++++++++-
 1 file changed, 67 insertions(+), 2 deletions(-)

diff --git a/distributed/deploy/adaptive_core.py b/distributed/deploy/adaptive_core.py
index b1a6424833c..6732bb20284 100644
--- a/distributed/deploy/adaptive_core.py
+++ b/distributed/deploy/adaptive_core.py
@@ -9,10 +9,65 @@
 
 
 class AdaptiveCore:
+    """
+    The core logic for adaptive deployments, with none of the cluster details
+
+    This class controls our adaptive scaling behavior.  It is intended to be
+    sued as a super-class or mixin.  It expects the following state and methods:
+
+    **State**
+
+    plan: set
+        A set of workers that we think should exist.
+        Here and below worker is just a token, often an address or name string
+
+    requested: set
+        A set of workers that the cluster class has successfully requested from
+        the resource manager.  We expect that resource manager to work to make
+        these exist.
+
+    observed: set
+        A set of workers that have successfully checked in with the scheduler
+
+    These sets are not necessarily equivalent.  Often plan and requested will
+    be very similar (requesting is usually fast) but there may be a large delay
+    between requested and observed (often resource managers don't give us what
+    we want).
+
+    **Functions**
+
+    target : -> int
+        Returns the target number of workers that should exist.
+        This is often obtained by querying the scheduler
+
+    workers_to_close : int -> Set[worker]
+        Given a target number of workers,
+        returns a set of workers that we should close when we're scaling down
+
+    scale_up : int -> None
+        Scales the cluster up to a target number of workers, presumably
+        changing at least ``plan`` and hopefully eventually also ``requested``
+
+   scale_down : Set[worker] -> None
+        Closes the provided set of workers
+
+    Parameters
+    ----------
+    minimum: int
+        The minimum number of allowed workers
+    maximum: int
+        The maximum number of allowed workers
+    wait_count: int
+        The number of scale-down requests we should receive before actually
+        scaling down
+    interval: str
+        The amount of time, like ``"1s"`` between checks
+    """
+
     def __init__(
         self,
-        minimum: float = 0,
-        maximum: float = math.inf,
+        minimum: int = 0,
+        maximum: int = math.inf,
         wait_count: int = 3,
         interval: str = "1s",
     ):
@@ -50,6 +105,7 @@ def stop(self):
             self.periodic_callback = None
 
     async def target(self) -> int:
+        """ The target number of workers that should exist """
         raise NotImplementedError()
 
     async def workers_to_close(self, target: int) -> list:
@@ -60,6 +116,7 @@ async def workers_to_close(self, target: int) -> list:
         return list(self.observed)[target:]
 
     async def safe_target(self) -> int:
+        """ Used internally, like target, but respects minimum/maximum """
         n = await self.target()
         if n > self.maximum:
             n = self.maximum
@@ -70,6 +127,9 @@ async def safe_target(self) -> int:
         return n
 
     async def recommendations(self, target: int) -> dict:
+        """
+        Make scale up/down recommendations based on current state and target
+        """
         plan = self.plan
         requested = self.requested
         observed = self.observed
@@ -108,6 +168,11 @@ async def recommendations(self, target: int) -> dict:
                 return {"status": "same"}
 
     async def adapt(self) -> None:
+        """
+        Check the current state, make recommendations, call scale
+
+        This is the main event of the system
+        """
         if self._adapting:  # Semaphore to avoid overlapping adapt calls
             return
         self._adapting = True

From 8c84dcefb4f8eb93fa0e5d86b3975f72d239690a Mon Sep 17 00:00:00 2001
From: Matthew Rocklin <mrocklin@gmail.com>
Date: Thu, 1 Aug 2019 12:48:38 -0700
Subject: [PATCH 11/12] deprecate Worker.wait_until_closed in favor of finished

---
 distributed/nanny.py  | 2 +-
 distributed/worker.py | 6 ++----
 2 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/distributed/nanny.py b/distributed/nanny.py
index 155fde98158..b6d8dadbf9a 100644
--- a/distributed/nanny.py
+++ b/distributed/nanny.py
@@ -690,7 +690,7 @@ async def run():
                         }
                     )
                     init_result_q.close()
-                    await worker.wait_until_closed()
+                    await worker.finished()
                     logger.info("Worker closed")
 
         try:
diff --git a/distributed/worker.py b/distributed/worker.py
index 672358d4460..290128c84e1 100644
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -26,7 +26,6 @@
     from toolz import pluck, partial, merge, first
 from tornado import gen
 from tornado.ioloop import IOLoop
-from tornado.locks import Event
 
 from . import profile, comm
 from .batched import BatchedSend
@@ -546,7 +545,6 @@ def __init__(
         self.actors = {}
         self.loop = loop or IOLoop.current()
         self.status = None
-        self._closed = Event()
         self.reconnect = reconnect
         self.executor = executor or ThreadPoolExecutor(
             self.nthreads, thread_name_prefix="Dask-Worker-Threads'"
@@ -1054,7 +1052,6 @@ async def close(
 
             self.stop()
             self.rpc.close()
-            self._closed.set()
 
             self.status = "closed"
             await ServerNode.close(self)
@@ -1084,7 +1081,8 @@ async def terminate(self, comm, report=True, **kwargs):
         return "OK"
 
     async def wait_until_closed(self):
-        await self._closed.wait()
+        warnings.warn("wait_until_closed has moved to finished()")
+        await self.finished()
         assert self.status == "closed"
 
     ################

From f973b30208ff02e90fd370392c28b28199ca4ef5 Mon Sep 17 00:00:00 2001
From: Matthew Rocklin <mrocklin@gmail.com>
Date: Thu, 1 Aug 2019 12:49:04 -0700
Subject: [PATCH 12/12] ignore leaking thread in test_log_tasks_during_restart

---
 distributed/tests/test_scheduler.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/distributed/tests/test_scheduler.py b/distributed/tests/test_scheduler.py
index 80cc04c81b3..0331ac0a972 100644
--- a/distributed/tests/test_scheduler.py
+++ b/distributed/tests/test_scheduler.py
@@ -1231,7 +1231,9 @@ def test_cancel_fire_and_forget(c, s, a, b):
     assert not s.tasks
 
 
-@gen_cluster(client=True, Worker=Nanny, clean_kwargs={"processes": False})
+@gen_cluster(
+    client=True, Worker=Nanny, clean_kwargs={"processes": False, "threads": False}
+)
 def test_log_tasks_during_restart(c, s, a, b):
     future = c.submit(sys.exit, 0)
     yield wait(future)