from __future__ import annotations import logging from collections.abc import Hashable from datetime import timedelta from inspect import isawaitable from typing import TYPE_CHECKING, Any, Callable, Literal, cast from tornado.ioloop import IOLoop import dask.config from dask.utils import parse_timedelta from distributed.compatibility import PeriodicCallback from distributed.core import Status from distributed.deploy.adaptive_core import AdaptiveCore from distributed.protocol import pickle from distributed.utils import log_errors if TYPE_CHECKING: from typing_extensions import TypeAlias import distributed from distributed.deploy.cluster import Cluster logger = logging.getLogger(__name__) AdaptiveStateState: TypeAlias = Literal[ "starting", "running", "stopped", "inactive", ] class Adaptive(AdaptiveCore): ''' 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 ---------- cluster: object 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. 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.target` 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:`distributed.Scheduler.adaptive_target`). The values for interval, min, max, wait_count and target_duration can be specified in the dask config under the distributed.adaptive key. ''' interval: float | None periodic_callback: PeriodicCallback | None #: Whether this adaptive strategy is periodically adapting state: AdaptiveStateState def __init__( self, cluster: Cluster, interval: str | float | timedelta | None = None, minimum: int | None = None, maximum: int | float | None = None, wait_count: int | None = None, target_duration: str | float | timedelta | None = None, worker_key: ( Callable[[distributed.scheduler.WorkerState], Hashable] | None ) = None, **kwargs: Any, ): self.cluster = cluster self.worker_key = worker_key self._workers_to_close_kwargs = kwargs if interval is None: interval = dask.config.get("distributed.adaptive.interval") if minimum is None: minimum = cast(int, dask.config.get("distributed.adaptive.minimum")) if maximum is None: maximum = cast(float, dask.config.get("distributed.adaptive.maximum")) if wait_count is None: wait_count = cast(int, dask.config.get("distributed.adaptive.wait-count")) if target_duration is None: target_duration = cast( str, dask.config.get("distributed.adaptive.target-duration") ) super().__init__(minimum=minimum, maximum=maximum, wait_count=wait_count) self.interval = parse_timedelta(interval, "seconds") self.periodic_callback = None if self.interval and self.cluster: import weakref self_ref = weakref.ref(self) async def _adapt(): adaptive = self_ref() if not adaptive or adaptive.state != "running": return if adaptive.cluster.status != Status.running: adaptive.stop(reason="cluster-not-running") return try: await adaptive.adapt() except Exception: logger.warning( "Adaptive encountered an error while adapting", exc_info=True ) self.periodic_callback = PeriodicCallback(_adapt, self.interval * 1000) self.state = "starting" self.loop.add_callback(self._start) else: self.state = "inactive" self.target_duration = parse_timedelta(target_duration) def _start(self) -> None: if self.state != "starting": return assert self.periodic_callback is not None self.periodic_callback.start() self.state = "running" logger.info( "Adaptive scaling started: minimum=%s maximum=%s", self.minimum, self.maximum, ) def stop(self, reason: str = "unknown") -> None: if self.state in ("inactive", "stopped"): return if self.state == "running": assert self.periodic_callback is not None self.periodic_callback.stop() logger.info( "Adaptive scaling stopped: minimum=%s maximum=%s. Reason: %s", self.minimum, self.maximum, reason, ) self.periodic_callback = None self.state = "stopped" @property def scheduler(self): return self.cluster.scheduler_comm @property def plan(self): return self.cluster.plan @property def requested(self): return self.cluster.requested @property def observed(self): return self.cluster.observed async def target(self): """ Determine target number of workers that should exist. Notes ----- ``Adaptive.target`` dispatches to Scheduler.adaptive_target(), but may be overridden in subclasses. Returns ------- Target number of workers See Also -------- Scheduler.adaptive_target """ return await self.scheduler.adaptive_target( target_duration=self.target_duration ) async def recommendations(self, target: int) -> dict: if len(self.plan) != len(self.requested): # Ensure that the number of planned and requested workers # are in sync before making recommendations. await self.cluster return await super().recommendations(target) async def workers_to_close(self, target: int) -> list[str]: """ 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 names to close, if any See Also -------- Scheduler.workers_to_close """ 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, ) @log_errors async def scale_down(self, workers): if not workers: return logger.info("Retiring workers %s", workers) # Ask scheduler to cleanly retire workers await self.scheduler.retire_workers( names=workers, remove=True, close_workers=True, ) # close workers more forcefully f = self.cluster.scale_down(workers) if isawaitable(f): await f async def scale_up(self, n): f = self.cluster.scale(n) if isawaitable(f): await f @property def loop(self) -> IOLoop: """Override Adaptive.loop""" if self.cluster: return self.cluster.loop # type: ignore[return-value] else: return IOLoop.current() def __del__(self): self.stop(reason="adaptive-deleted")