diff --git a/.circleci/config.yml b/.circleci/config.yml index cf70c4355..3ca6a2611 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -39,9 +39,6 @@ jobs: - run: command: | sudo kubectl get deployment - - restore_cache: - keys: - - miniconda-v1-{{ checksum "ci/environment-3.7.yml" }} - run: name: install miniconda command: | @@ -64,10 +61,6 @@ jobs: fi conda env list conda list ${ENV_NAME} - - save_cache: - key: miniconda-v1-{{ checksum "ci/environment-3.7.yml" }} - paths: - - "/home/circleci/miniconda" - run: name: build docker command: | diff --git a/dask_kubernetes/adaptive.py b/dask_kubernetes/adaptive.py new file mode 100644 index 000000000..903268a23 --- /dev/null +++ b/dask_kubernetes/adaptive.py @@ -0,0 +1,255 @@ +# Taken from commit e1e36e437a53a937d14519f06fced6b645cff361 of dask/distributed +from collections import deque +import logging + +from tornado import gen + +from distributed.metrics import time +from distributed.utils import log_errors, PeriodicCallback, parse_timedelta +from distributed.protocol import pickle + +logger = logging.getLogger(__name__) + + +class Adaptive(object): + ''' + Adaptively allocate workers based on scheduler load. A superclass. + + Contains logic to dynamically resize a Dask cluster based on current use. + This class needs to be paired with a system that can create and destroy + Dask workers using a cluster resource manager. Typically it is built into + already existing solutions, rather than used directly by users. + It is most commonly used from the ``.adapt(...)`` method of various Dask + cluster classes. + + 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 + 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. + worker_key: Callable[WorkerState] + Function to group workers together when scaling down + See Scheduler.workers_to_close for more information + minimum: int + Minimum number of workers to keep around + maximum: int + Maximum number of workers to keep around + **kwargs: + Extra parameters to pass to Scheduler.workers_to_close + + Examples + -------- + + This is commonly used from existing Dask classes, like KubeCluster + + >>> from dask_kubernetes import KubeCluster + >>> cluster = KubeCluster() + >>> cluster.adapt(minimum=10, maximum=100) + + Alternatively you can use it from your own Cluster class by subclassing + from Dask's Cluster superclass + + >>> from distributed.deploy import Cluster + >>> class MyCluster(Cluster): + ... def scale_up(self, n): + ... """ Bring worker count up to n """ + ... def scale_down(self, workers): + ... """ Remove worker addresses from cluster """ + + >>> cluster = MyCluster() + >>> cluster.adapt(minimum=10, maximum=100) + + Notes + ----- + Subclasses can override :meth:`Adaptive.should_scale_up` and + :meth:`Adaptive.workers_to_close` to control when the cluster should be + resized. The default implementation checks if there are too many tasks + per worker or too little memory available (see :meth:`Adaptive.needs_cpu` + and :meth:`Adaptive.needs_memory`). + ''' + + def __init__( + self, + cluster=None, + interval="1s", + startup_cost="1s", + scale_factor=2, + minimum=0, + maximum=None, + 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._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) + + @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 + + async def workers_to_close(self, **kwargs): + """ + Determine which, if any, workers should potentially be removed from + the cluster. + + Notes + ----- + ``Adaptive.workers_to_close`` dispatches to Scheduler.workers_to_close(), + but may be overridden in subclasses. + + Returns + ------- + List of worker addresses to close, if any + + See Also + -------- + 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 + + 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) + with log_errors(): + await self.scheduler.retire_workers( + workers=workers, remove=True, close_workers=True + ) + + 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) + + @property + def loop(self): + return self.cluster.loop diff --git a/dask_kubernetes/core.py b/dask_kubernetes/core.py index 5a444f864..485a75a68 100644 --- a/dask_kubernetes/core.py +++ b/dask_kubernetes/core.py @@ -22,6 +22,7 @@ from .objects import make_pod_from_dict, clean_pod_template from .auth import ClusterAuth from .logs import Log, Logs +from .adaptive import Adaptive logger = logging.getLogger(__name__) @@ -378,6 +379,9 @@ def logs(self, pod=None): ) ) + def adapt(self, Adaptive=Adaptive, **kwargs): + return super().adapt(Adaptive=Adaptive, **kwargs) + def scale(self, n): """ Scale cluster to n workers @@ -537,6 +541,14 @@ def __exit__(self, type, value, traceback): def scheduler_comm(self): return self.cluster.scheduler_comm + @property + def scheduler_info(self): + return self.cluster.scheduler_info + + @property + def periodic_callbacks(self): + return self.cluster.periodic_callbacks + def _cleanup_pods(namespace, labels): """ Remove all pods with these labels in this namespace """ diff --git a/dask_kubernetes/tests/test_core.py b/dask_kubernetes/tests/test_core.py index 3612305e4..a0383bf4e 100644 --- a/dask_kubernetes/tests/test_core.py +++ b/dask_kubernetes/tests/test_core.py @@ -580,7 +580,7 @@ def test_escape_name(pod_spec, loop, ns): def test_maximum(cluster): - with dask.config.set(**{"kubernetes.count.max": 1}): + with dask.config.set({"kubernetes.count.max": 1}): with captured_logger("dask_kubernetes") as logger: cluster.scale(10)