from __future__ import annotations import asyncio import logging import random from collections import defaultdict from collections.abc import Callable, Collection, Coroutine, Mapping from functools import partial from itertools import cycle from typing import Any, TypeVar from tlz import drop, groupby, merge import dask.config from dask._task_spec import TaskRef from dask.optimization import SubgraphCallable from dask.typing import Key from dask.utils import is_namedtuple_instance, parse_timedelta from distributed.core import ConnectionPool, rpc from distributed.utils import All logger = logging.getLogger(__name__) # Backwards compat WrappedKey = TaskRef async def gather_from_workers( who_has: Mapping[Key, Collection[str]], rpc: ConnectionPool, *, serializers: list[str] | None = None, who: str | None = None, ) -> tuple[dict[Key, object], list[Key], list[Key], list[str]]: """Gather data directly from peers Parameters ---------- who_has: mapping from keys to worker addresses rpc: RPC channel to use Returns ------- Tuple: - Successfully retrieved: ``{key: value, ...}`` - Keys that were not available on any worker: ``[key, ...]`` - Keys that raised exception; e.g. failed to deserialize: ``[key, ...]`` - Workers that failed to respond: ``[address, ...]`` See Also -------- gather _gather Scheduler.get_who_has """ from distributed.worker import get_data_from_worker to_gather = {k: set(v) for k, v in who_has.items()} data: dict[Key, object] = {} failed_keys: list[Key] = [] missing_workers: set[str] = set() busy_workers: set[str] = set() while to_gather: d = defaultdict(list) for key, addresses in to_gather.items(): addresses -= missing_workers ready_addresses = addresses - busy_workers if ready_addresses: d[random.choice(list(ready_addresses))].append(key) if not d: if busy_workers: await asyncio.sleep(0.15) busy_workers.clear() continue return data, list(to_gather), failed_keys, list(missing_workers) tasks = [ asyncio.create_task( retry_operation( partial( get_data_from_worker, rpc, keys, address, who=who, serializers=serializers, ), operation="get_data_from_worker", ), name=f"get-data-from-{address}", ) for address, keys in d.items() ] results = await asyncio.gather(*tasks, return_exceptions=True) for address, r in zip(d, results): if isinstance(r, OSError): missing_workers.add(address) elif isinstance(r, Exception): # For example, deserialization error logger.error( "Unexpected error while collecting tasks %s from %s", d[address], address, exc_info=r, ) for key in d[address]: failed_keys.append(key) del to_gather[key] elif isinstance(r, BaseException): # pragma: nocover # for example, asyncio.CancelledError raise r else: assert isinstance(r, dict), r if r["status"] == "busy": busy_workers.add(address) continue assert r["status"] == "OK" for key in d[address]: if key in r["data"]: data[key] = r["data"][key] del to_gather[key] else: to_gather[key].remove(address) return data, [], failed_keys, list(missing_workers) _round_robin_counter = [0] async def scatter_to_workers(workers, data, rpc=rpc): """Scatter data directly to workers This distributes data in a round-robin fashion to a set of workers. See scatter for parameter docstring """ assert isinstance(data, dict) workers = sorted(workers) names, data = list(zip(*data.items())) worker_iter = drop(_round_robin_counter[0] % len(workers), cycle(workers)) _round_robin_counter[0] += len(data) L = list(zip(worker_iter, names, data)) d = groupby(0, L) d = {worker: {key: value for _, key, value in v} for worker, v in d.items()} rpcs = {addr: rpc(addr) for addr in d} try: out = await All([rpcs[address].update_data(data=v) for address, v in d.items()]) finally: for r in rpcs.values(): await r.close_rpc() nbytes = merge(o["nbytes"] for o in out) who_has = {k: [w for w, _, _ in v] for k, v in groupby(1, L).items()} return (names, who_has, nbytes) collection_types = (tuple, list, set, frozenset) def _namedtuple_packing(o: Any, handler: Callable[..., Any]) -> Any: """Special pack/unpack dispatcher for namedtuples respecting their potential constructors.""" assert is_namedtuple_instance(o) typ = type(o) if hasattr(o, "__getnewargs_ex__"): args, kwargs = o.__getnewargs_ex__() handled_args = [handler(item) for item in args] handled_kwargs = {k: handler(v) for k, v in kwargs.items()} return typ(*handled_args, **handled_kwargs) args = o.__getnewargs__() if hasattr(typ, "__getnewargs__") else o handled_args = [handler(item) for item in args] return typ(*handled_args) def _unpack_remotedata_inner( o: Any, byte_keys: bool, found_futures: set[TaskRef] ) -> Any: """Inner implementation of `unpack_remotedata` that adds found wrapped keys to `found_futures`""" typ = type(o) if typ is tuple: if not o: return o if type(o[0]) is SubgraphCallable: # Unpack futures within the arguments of the subgraph callable futures: set[TaskRef] = set() args = tuple(_unpack_remotedata_inner(i, byte_keys, futures) for i in o[1:]) found_futures.update(futures) # Unpack futures within the subgraph callable itself sc: SubgraphCallable = o[0] futures = set() dsk = { k: _unpack_remotedata_inner(v, byte_keys, futures) for k, v in sc.dsk.items() } future_keys: tuple = () if futures: # If no futures is in the subgraph, we just use `sc` as-is found_futures.update(futures) future_keys = ( tuple(f.key for f in futures) if byte_keys else tuple(f.key for f in futures) ) inkeys = tuple(sc.inkeys) + future_keys sc = SubgraphCallable(dsk, sc.outkey, inkeys, sc.name) return (sc,) + args + future_keys else: return tuple( _unpack_remotedata_inner(item, byte_keys, found_futures) for item in o ) elif is_namedtuple_instance(o): return _namedtuple_packing( o, partial( _unpack_remotedata_inner, byte_keys=byte_keys, found_futures=found_futures, ), ) if typ in collection_types: if not o: return o outs = [_unpack_remotedata_inner(item, byte_keys, found_futures) for item in o] return typ(outs) elif typ is dict: if o: return { k: _unpack_remotedata_inner(v, byte_keys, found_futures) for k, v in o.items() } else: return o elif issubclass(typ, TaskRef): # TODO use type is Future k = o.key found_futures.add(o) return k else: return o def unpack_remotedata(o: Any, byte_keys: bool = False) -> tuple[Any, set]: """Unpack TaskRef objects from collection Returns original collection and set of all found TaskRef objects Examples -------- >>> rd = TaskRef('mykey') >>> unpack_remotedata(1) (1, set()) >>> unpack_remotedata(()) ((), set()) >>> unpack_remotedata(rd) ('mykey', {TaskRef('mykey')}) >>> unpack_remotedata([1, rd]) ([1, 'mykey'], {TaskRef('mykey')}) >>> unpack_remotedata({1: rd}) ({1: 'mykey'}, {TaskRef('mykey')}) >>> unpack_remotedata({1: [rd]}) ({1: ['mykey']}, {TaskRef('mykey')}) Use the ``byte_keys=True`` keyword to force string keys >>> rd = TaskRef(('x', 1)) >>> unpack_remotedata(rd, byte_keys=True) ("('x', 1)", {TaskRef('('x', 1)')}) """ found_futures: set[TaskRef] = set() return _unpack_remotedata_inner(o, byte_keys, found_futures), found_futures def pack_data(o, d, key_types=object): """Merge known data into tuple or dict Parameters ---------- o core data structures containing literals and keys d : dict mapping of keys to data Examples -------- >>> data = {'x': 1} >>> pack_data(('x', 'y'), data) (1, 'y') >>> pack_data({'a': 'x', 'b': 'y'}, data) # doctest: +SKIP {'a': 1, 'b': 'y'} >>> pack_data({'a': ['x'], 'b': 'y'}, data) # doctest: +SKIP {'a': [1], 'b': 'y'} """ typ = type(o) try: if isinstance(o, key_types) and o in d: return d[o] except TypeError: pass if typ in collection_types: return typ([pack_data(x, d, key_types=key_types) for x in o]) elif typ is dict: return {k: pack_data(v, d, key_types=key_types) for k, v in o.items()} elif is_namedtuple_instance(o): return _namedtuple_packing(o, partial(pack_data, d=d, key_types=key_types)) else: return o def subs_multiple(o, d): """Perform substitutions on a tasks Parameters ---------- o Core data structures containing literals and keys d : dict Mapping of keys to values Examples -------- >>> dsk = {"a": (sum, ["x", 2])} >>> data = {"x": 1} >>> subs_multiple(dsk, data) # doctest: +SKIP {'a': (sum, [1, 2])} """ typ = type(o) if typ is tuple and o and callable(o[0]): # istask(o) return (o[0],) + tuple(subs_multiple(i, d) for i in o[1:]) elif typ is list: return [subs_multiple(i, d) for i in o] elif typ is dict: return {k: subs_multiple(v, d) for (k, v) in o.items()} else: try: return d.get(o, o) except TypeError: return o T = TypeVar("T") async def retry( coro: Callable[[], Coroutine[Any, Any, T]], count: int, delay_min: float, delay_max: float, jitter_fraction: float = 0.1, retry_on_exceptions: type[BaseException] | tuple[type[BaseException], ...] = ( EnvironmentError, IOError, ), operation: str | None = None, ) -> T: """ Return the result of ``await coro()``, re-trying in case of exceptions The delay between attempts is ``delay_min * (2 ** i - 1)`` where ``i`` enumerates the attempt that just failed (starting at 0), but never larger than ``delay_max``. This yields no delay between the first and second attempt, then ``delay_min``, ``3 * delay_min``, etc. (The reason to re-try with no delay is that in most cases this is sufficient and will thus recover faster from a communication failure). Parameters ---------- coro The coroutine function to call and await count The maximum number of re-tries before giving up. 0 means no re-try; must be >= 0. delay_min The base factor for the delay (in seconds); this is the first non-zero delay between re-tries. delay_max The maximum delay (in seconds) between consecutive re-tries (without jitter) jitter_fraction The maximum jitter to add to the delay, as fraction of the total delay. No jitter is added if this value is <= 0. Using a non-zero value here avoids "herd effects" of many operations re-tried at the same time retry_on_exceptions A tuple of exception classes to retry. Other exceptions are not caught and re-tried, but propagate immediately. operation A human-readable description of the operation attempted; used only for logging failures Returns ------- Any Whatever `await coro()` returned """ # this loop is a no-op in case max_retries<=0 for i_try in range(count): try: return await coro() except retry_on_exceptions as ex: operation = operation or str(coro) logger.info( f"Retrying {operation} after exception in attempt {i_try}/{count}: {ex}" ) delay = min(delay_min * (2**i_try - 1), delay_max) if jitter_fraction > 0: delay *= 1 + random.random() * jitter_fraction await asyncio.sleep(delay) return await coro() # ParamSpec is not supported here due to the additional "operation" kwarg async def retry_operation( coro: Callable[..., Coroutine[Any, Any, T]], *args: object, operation: str | None = None, **kwargs: object, ) -> T: """ Retry an operation using the configuration values for the retry parameters """ retry_count = dask.config.get("distributed.comm.retry.count") retry_delay_min = parse_timedelta( dask.config.get("distributed.comm.retry.delay.min"), default="s" ) retry_delay_max = parse_timedelta( dask.config.get("distributed.comm.retry.delay.max"), default="s" ) return await retry( partial(coro, *args, **kwargs), count=retry_count, delay_min=retry_delay_min, delay_max=retry_delay_max, operation=operation, )