diff --git a/distributed/deploy/adaptive.py b/distributed/deploy/adaptive.py
index 9b1d8511045..2efc18dfe0c 100644
--- a/distributed/deploy/adaptive.py
+++ b/distributed/deploy/adaptive.py
@@ -1,16 +1,14 @@
-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.
@@ -23,19 +21,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.
@@ -84,45 +76,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)
+
+ @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, **kwargs):
+ async def workers_to_close(self, target: int):
"""
Determine which, if any, workers should potentially be removed from
the cluster.
@@ -140,114 +134,30 @@ 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
-
- 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
+ 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 _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,
- )
+ async def scale_down(self, workers):
if not workers:
- raise gen.Return(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
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)
- workers = set(
- await self.workers_to_close(
- key=pickle.dumps(self.worker_key) if self.worker_key else None,
- minimum=self.minimum,
- )
- )
- 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",
- }
-
- elif n > current:
- self.close_counts.clear()
- return {"status": "up", "n": n}
-
- elif workers:
- d = {}
- to_close = []
- for w, c in self.close_counts.items():
- if w in workers:
- if c >= self.wait_count:
- to_close.append(w)
- else:
- d[w] = c
-
- for w in workers:
- 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
- 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
new file mode 100644
index 00000000000..6732bb20284
--- /dev/null
+++ b/distributed/deploy/adaptive_core.py
@@ -0,0 +1,197 @@
+import collections
+import math
+
+from tornado.ioloop import IOLoop
+import toolz
+
+from ..metrics import time
+from ..utils import parse_timedelta, PeriodicCallback
+
+
+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: int = 0,
+ maximum: int = 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
+
+ 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 stop(self):
+ if self.periodic_callback:
+ self.periodic_callback.stop()
+ 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:
+ """
+ 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:
+ """ Used internally, like target, but respects minimum/maximum """
+ 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:
+ """
+ Make scale up/down recommendations based on current state and target
+ """
+ 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": list(firmly_close)}
+ else:
+ 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
+
+ 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)
+ except OSError:
+ self.stop()
+ finally:
+ self._adapting = False
diff --git a/distributed/deploy/spec.py b/distributed/deploy/spec.py
index c9e1aca1a87..feb0dfe63b5 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
@@ -376,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()
@@ -411,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
@@ -473,6 +474,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)
@@ -492,13 +494,13 @@ def _widget_status(self):
}
- | Workers | %d |
+ | Workers | %s |
| Cores | %d |
| Memory | %s |
""" % (
- workers,
+ workers if workers == requested else "%d / %d" % (workers, requested),
cores,
memory,
)
@@ -547,6 +549,7 @@ def _widget(self):
def adapt_cb(b):
self.adapt(minimum=minimum.value, maximum=maximum.value)
+ update()
adapt.on_click(adapt_cb)
@@ -556,6 +559,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_adaptive.py b/distributed/deploy/tests/test_adaptive.py
index 2d3d2235e21..261b4355251 100644
--- a/distributed/deploy/tests/test_adaptive.py
+++ b/distributed/deploy/tests/test_adaptive.py
@@ -1,7 +1,6 @@
from time import sleep
import pytest
-from toolz import frequencies, pluck
from tornado import gen
from tornado.ioloop import IOLoop
@@ -68,51 +67,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):
@@ -164,7 +158,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)
@@ -177,7 +171,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
@@ -185,41 +179,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)
@@ -238,7 +226,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)
@@ -246,10 +234,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
@@ -362,7 +350,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
@@ -372,21 +360,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},
@@ -394,9 +382,7 @@ def test_worker_keys():
"b-2": {"cls": Worker},
},
asynchronous=True,
- )
-
- try:
+ ) as cluster:
def key(ws):
return ws.name.split("-")[0]
@@ -404,12 +390,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
new file mode 100644
index 00000000000..a073314223d
--- /dev/null
+++ b/distributed/deploy/tests/test_adaptive_core.py
@@ -0,0 +1,90 @@
+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 = 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()
+ await asyncio.sleep(0.050)
+
+ adapt._target = 10
+ await asyncio.sleep(0.020)
+ assert len(adapt.plan) == 1 # last value from before, unchanged
diff --git a/distributed/deploy/tests/test_slow_adaptive.py b/distributed/deploy/tests/test_slow_adaptive.py
new file mode 100644
index 00000000000..4f565a78289
--- /dev/null
+++ b/distributed/deploy/tests/test_slow_adaptive.py
@@ -0,0 +1,98 @@
+import asyncio
+import pytest
+
+from dask.distributed import Worker, Scheduler, SpecCluster, Client
+from distributed.utils_test import slowinc, 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)
+
+
+@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
+
+
+@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"]
diff --git a/distributed/deploy/tests/test_spec_cluster.py b/distributed/deploy/tests/test_spec_cluster.py
index 58bbbaef44d..64633428a38 100644
--- a/distributed/deploy/tests/test_spec_cluster.py
+++ b/distributed/deploy/tests/test_spec_cluster.py
@@ -252,7 +252,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
@@ -262,3 +265,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()
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/scheduler.py b/distributed/scheduler.py
index acabce22c63..e9fdbe60b51 100644
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -2898,7 +2898,14 @@ 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,
+ attribute="address",
):
"""
Find workers that we can close with low cost
@@ -2925,6 +2932,11 @@ 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
+ attribute : str
+ The attribute of the WorkerState object to return, like "address"
+ or "name". Defaults to "address".
Examples
--------
@@ -2952,6 +2964,13 @@ def workers_to_close(
--------
Scheduler.retire_workers
"""
+ if target is not None and n is None:
+ n = len(self.workers) - target
+ if n is not None:
+ if n < 0:
+ n = 0
+ target = len(self.workers) - n
+
if n is None and memory_ratio is None:
memory_ratio = 2
@@ -2976,12 +2995,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)
@@ -2996,7 +3015,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)
@@ -3005,22 +3024,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
@@ -3042,6 +3069,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:
@@ -3052,17 +3084,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:
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)
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"
################