from __future__ import annotations import asyncio import logging import uuid from collections import defaultdict from contextlib import suppress from tlz import merge from dask.utils import parse_timedelta from distributed.client import Future from distributed.metrics import time from distributed.utils import Deadline, TimeoutError, log_errors, wait_for from distributed.worker import get_client logger = logging.getLogger(__name__) class VariableExtension: """An extension for the scheduler to manage Variables This adds the following routes to the scheduler * variable-set * variable-get * variable-delete """ def __init__(self, scheduler): self.scheduler = scheduler self.variables = dict() self.waiting = defaultdict(set) self.waiting_conditions = defaultdict(asyncio.Condition) self.started = asyncio.Condition() self.scheduler.handlers.update( {"variable_set": self.set, "variable_get": self.get} ) self.scheduler.stream_handlers["variable-future-received-confirm"] = ( self.future_received_confirm ) self.scheduler.stream_handlers["variable_delete"] = self.delete async def set(self, name=None, key=None, data=None, client=None, timeout=None): deadline = Deadline.after(parse_timedelta(timeout)) if key is not None: record = {"type": "Future", "value": key} while key not in self.scheduler.tasks: await asyncio.sleep(0.01) if deadline.expired: raise TimeoutError(f"Task {key} unknown to scheduler.") self.scheduler.client_desires_keys(keys=[key], client="variable-%s" % name) else: record = {"type": "msgpack", "value": data} try: old = self.variables[name] except KeyError: pass else: if old["type"] == "Future" and old["value"] != key: asyncio.ensure_future(self.release(old["value"], name)) if name not in self.variables: async with self.started: self.started.notify_all() self.variables[name] = record async def release(self, key, name): while self.waiting[key, name]: async with self.waiting_conditions[name]: await self.waiting_conditions[name].wait() self.scheduler.client_releases_keys(keys=[key], client="variable-%s" % name) del self.waiting[key, name] async def future_received_confirm( self, name=None, key=None, token=None, client=None ): self.scheduler.client_desires_keys([key], client) self.waiting[key, name].remove(token) if not self.waiting[key, name]: async with self.waiting_conditions[name]: self.waiting_conditions[name].notify_all() async def get(self, name=None, client=None, timeout=None): start = time() while name not in self.variables: if timeout is not None: left = timeout - (time() - start) else: left = None if left and left < 0: raise TimeoutError() try: async def _(): # Python 3.6 is odd and requires special help here await self.started.acquire() await self.started.wait() await wait_for(_(), timeout=left) finally: self.started.release() record = self.variables[name] if record["type"] == "Future": key = record["value"] token = uuid.uuid4().hex ts = self.scheduler.tasks.get(key) state = ts.state if ts is not None else "lost" msg = {"token": token, "state": state} if state == "erred": msg["exception"] = ts.exception_blame.exception msg["traceback"] = ts.exception_blame.traceback record = merge(record, msg) self.waiting[key, name].add(token) return record @log_errors async def delete(self, name=None, client=None): try: old = self.variables[name] except KeyError: pass else: if old["type"] == "Future": await self.release(old["value"], name) with suppress(KeyError): del self.waiting_conditions[name] with suppress(KeyError): del self.variables[name] self.scheduler.remove_client("variable-%s" % name) class Variable: """Distributed Global Variable This allows multiple clients to share futures and data between each other with a single mutable variable. All metadata is sequentialized through the scheduler. Race conditions can occur. Values must be either Futures or msgpack-encodable data (ints, lists, strings, etc..) All data will be kept and sent through the scheduler, so it is wise not to send too much. If you want to share a large amount of data then ``scatter`` it and share the future instead. Parameters ---------- name: string (optional) Name used by other clients and the scheduler to identify the variable. If not given, a random name will be generated. client: Client (optional) Client used for communication with the scheduler. If not given, the default global client will be used. Examples -------- >>> from dask.distributed import Client, Variable # doctest: +SKIP >>> client = Client() # doctest: +SKIP >>> x = Variable('x') # doctest: +SKIP >>> x.set(123) # docttest: +SKIP >>> x.get() # docttest: +SKIP 123 >>> future = client.submit(f, x) # doctest: +SKIP >>> x.set(future) # doctest: +SKIP See Also -------- Queue: shared multi-producer/multi-consumer queue between clients """ def __init__(self, name=None, client=None): self._client = client self.name = name or "variable-" + uuid.uuid4().hex @property def client(self): if not self._client: try: self._client = get_client() except ValueError: pass return self._client def _verify_running(self): if not self.client: raise RuntimeError( f"{type(self)} object not properly initialized. This can happen" " if the object is being deserialized outside of the context of" " a Client or Worker." ) async def _set(self, value, timeout): if isinstance(value, Future): await self.client.scheduler.variable_set( key=value.key, name=self.name, timeout=timeout ) else: await self.client.scheduler.variable_set( data=value, name=self.name, timeout=timeout ) def set(self, value, timeout="30 s", **kwargs): """Set the value of this variable Parameters ---------- value : Future or object Must be either a Future or a msgpack-encodable value """ self._verify_running() return self.client.sync(self._set, value, timeout=timeout, **kwargs) async def _get(self, timeout=None): d = await self.client.scheduler.variable_get( timeout=timeout, name=self.name, client=self.client.id ) if d["type"] == "Future": value = Future(d["value"], self.client, state=d["state"]) if d["state"] == "erred": value._state.set_error(d["exception"], d["traceback"]) self.client._send_to_scheduler( { "op": "variable-future-received-confirm", "name": self.name, "key": d["value"], "token": d["token"], } ) else: value = d["value"] return value def get(self, timeout=None, **kwargs): """Get the value of this variable Parameters ---------- timeout : number or string or timedelta, optional Time in seconds to wait before timing out. Instead of number of seconds, it is also possible to specify a timedelta in string format, e.g. "200ms". """ self._verify_running() timeout = parse_timedelta(timeout) return self.client.sync(self._get, timeout=timeout, **kwargs) def delete(self): """Delete this variable Caution, this affects all clients currently pointing to this variable. """ self._verify_running() if self.client.status == "running": # TODO: can leave zombie futures self.client._send_to_scheduler({"op": "variable_delete", "name": self.name}) def __reduce__(self): return Variable, (self.name,)