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" ################