diff --git a/distributed/client.py b/distributed/client.py index 7a39ec4b235..b578b4d16f1 100644 --- a/distributed/client.py +++ b/distributed/client.py @@ -1,3 +1,4 @@ +import asyncio import atexit from collections import defaultdict from collections.abc import Iterator @@ -8,6 +9,7 @@ import errno from functools import partial import html +from inspect import isawaitable import itertools import json import logging @@ -37,12 +39,7 @@ except ImportError: single_key = first from tornado import gen -from tornado.locks import Event, Condition, Semaphore from tornado.ioloop import IOLoop -from tornado.queues import Queue - -import asyncio -from asyncio import iscoroutine from .batched import BatchedSend from .utils_comm import ( @@ -431,7 +428,7 @@ def _get_event(self): # (https://github.com/tornadoweb/tornado/issues/2189) event = self._event if event is None: - event = self._event = Event() + event = self._event = asyncio.Event() return event def cancel(self): @@ -470,7 +467,7 @@ def reset(self): self._event.clear() async def wait(self, timeout=None): - await self._get_event().wait(timeout) + await asyncio.wait_for(self._get_event().wait(), timeout) def __repr__(self): return "<%s: %s>" % (self.__class__.__name__, self.status) @@ -630,10 +627,6 @@ def __init__( self._deserializers = deserializers self.direct_to_workers = direct_to_workers - self._gather_semaphore = Semaphore(5) - self._gather_keys = None - self._gather_future = None - # Communication self.scheduler_comm = None @@ -678,6 +671,10 @@ def __init__( self._loop_runner = LoopRunner(loop=loop, asynchronous=asynchronous) self.loop = self._loop_runner.loop + self._gather_semaphore = asyncio.Semaphore(5, loop=self.loop.asyncio_loop) + self._gather_keys = None + self._gather_future = None + if heartbeat_interval is None: heartbeat_interval = dask.config.get("distributed.client.heartbeat") heartbeat_interval = parse_timedelta(heartbeat_interval, default="ms") @@ -1182,7 +1179,9 @@ async def _handle_report(self): try: handler = self._stream_handlers[op] - handler(**msg) + result = handler(**msg) + if isawaitable(result): + await result except Exception as e: logger.exception(e) if breakout: @@ -1238,6 +1237,8 @@ def _handle_error(self, exception=None): async def _close(self, fast=False): """ Send close signal and wait until scheduler completes """ + if self.status == "closed": + return self.status = "closing" for pc in self._periodic_callbacks.values(): @@ -1252,8 +1253,6 @@ async def _close(self, fast=False): pass if self.get == dask.config.get("get", None): del dask.config.config["get"] - if self.status == "closed": - return if ( self.scheduler_comm @@ -1345,7 +1344,7 @@ def close(self, timeout=no_default): if self._start_arg is None: with ignoring(AttributeError): f = self.cluster.close() - if iscoroutine(f): + if asyncio.iscoroutine(f): async def _(): await f @@ -1365,6 +1364,7 @@ async def _shutdown(self): await self.cluster.close() else: with ignoring(CommClosedError): + self.status = "closing" await self.scheduler.terminate(close_workers=True) def shutdown(self): @@ -1800,12 +1800,11 @@ async def _gather_remote(self, direct, local_worker): few. In controls access using a Tornado semaphore, and picks up keys from other requests made recently. """ - await self._gather_semaphore.acquire() - keys = list(self._gather_keys) - self._gather_keys = None # clear state, these keys are being sent off - self._gather_future = None + async with self._gather_semaphore: + keys = list(self._gather_keys) + self._gather_keys = None # clear state, these keys are being sent off + self._gather_future = None - try: if direct or local_worker: # gather directly from workers who_has = await retry_operation(self.scheduler.who_has, keys=keys) data2, missing_keys, missing_workers = await gather_from_workers( @@ -1820,8 +1819,6 @@ async def _gather_remote(self, direct, local_worker): else: # ask scheduler to gather data for us response = await retry_operation(self.scheduler.gather, keys=keys) - finally: - self._gather_semaphore.release() return response @@ -2911,10 +2908,12 @@ async def _restart(self, timeout=no_default): if timeout == no_default: timeout = self._timeout * 2 self._send_to_scheduler({"op": "restart", "timeout": timeout}) - self._restart_event = Event() + self._restart_event = asyncio.Event() try: - await self._restart_event.wait(self.loop.time() + timeout) - except gen.TimeoutError: + await asyncio.wait_for( + self._restart_event.wait(), self.loop.time() + timeout + ) + except TimeoutError: logger.error("Restart timed out after %f seconds", timeout) pass self.generation += 1 @@ -4128,7 +4127,7 @@ async def _first_completed(futures): See Also: _as_completed """ - q = Queue() + q = asyncio.Queue() await _as_completed(futures, q) result = await q.get() return result @@ -4199,7 +4198,7 @@ def __init__(self, futures=None, loop=None, with_results=False, raise_errors=Tru self.queue = pyQueue() self.lock = threading.Lock() self.loop = loop or default_client().loop - self.condition = Condition() + self.condition = asyncio.Condition(loop=self.loop.asyncio_loop) self.thread_condition = threading.Condition() self.with_results = with_results self.raise_errors = raise_errors @@ -4207,11 +4206,6 @@ def __init__(self, futures=None, loop=None, with_results=False, raise_errors=Tru if futures: self.update(futures) - def _notify(self): - self.condition.notify() - with self.thread_condition: - self.thread_condition.notify() - async def _track_future(self, future): try: await _wait(future) @@ -4230,7 +4224,10 @@ async def _track_future(self, future): self.queue.put_nowait((future, result)) else: self.queue.put_nowait(future) - self._notify() + async with self.condition: + self.condition.notify() + with self.thread_condition: + self.thread_condition.notify() def update(self, futures): """ Add multiple futures to the collection. @@ -4297,7 +4294,8 @@ async def __anext__(self): while self.queue.empty(): if not self.futures: raise StopAsyncIteration - await self.condition.wait() + async with self.condition: + await self.condition.wait() return self._get_and_raise() diff --git a/distributed/comm/inproc.py b/distributed/comm/inproc.py index e46c2804ed1..09c08ddca67 100644 --- a/distributed/comm/inproc.py +++ b/distributed/comm/inproc.py @@ -7,7 +7,6 @@ import weakref import warnings -from tornado import locks from tornado.concurrent import Future from tornado.ioloop import IOLoop @@ -298,7 +297,7 @@ async def connect(self, address, deserialize=True, **connection_args): s2c_q=Queue(), c_loop=IOLoop.current(), c_addr=self.manager.new_address(), - conn_event=locks.Event(), + conn_event=asyncio.Event(), ) listener.connect_threadsafe(conn_req) # Wait for connection acknowledgement diff --git a/distributed/comm/tests/test_comms.py b/distributed/comm/tests/test_comms.py index 470c667b989..b9bfaba131a 100644 --- a/distributed/comm/tests/test_comms.py +++ b/distributed/comm/tests/test_comms.py @@ -9,7 +9,7 @@ import pkg_resources import pytest -from tornado import ioloop, locks, queues +from tornado import ioloop, queues from tornado.concurrent import Future import distributed @@ -901,7 +901,7 @@ async def handle_comm(comm): async def check_connector_deserialize(addr, deserialize, in_value, check_out): - done = locks.Event() + done = asyncio.Event() async def handle_comm(comm): await comm.write(in_value) diff --git a/distributed/core.py b/distributed/core.py index 81cd7adf8e4..c6cb72a9d14 100644 --- a/distributed/core.py +++ b/distributed/core.py @@ -13,7 +13,6 @@ from toolz import merge from tornado import gen from tornado.ioloop import IOLoop -from tornado.locks import Event from .comm import ( connect, @@ -134,7 +133,7 @@ def __init__( self.events = None self.event_counts = None self._ongoing_coroutines = weakref.WeakSet() - self._event_finished = Event() + self._event_finished = asyncio.Event() self.listeners = [] self.io_loop = io_loop or IOLoop.current() diff --git a/distributed/deploy/spec.py b/distributed/deploy/spec.py index 96279d15323..56a7740f74d 100644 --- a/distributed/deploy/spec.py +++ b/distributed/deploy/spec.py @@ -6,7 +6,6 @@ import weakref import dask -from tornado.locks import Event from tornado import gen from .adaptive import Adaptive @@ -42,7 +41,7 @@ def __init__(self, scheduler=None, name=None): self.external_address = None self.lock = asyncio.Lock() self.status = "created" - self._event_finished = Event() + self._event_finished = asyncio.Event() def __await__(self): async def _(): diff --git a/distributed/lock.py b/distributed/lock.py index c230a8e861c..218fcef8763 100644 --- a/distributed/lock.py +++ b/distributed/lock.py @@ -1,9 +1,7 @@ +import asyncio from collections import defaultdict, deque import logging import uuid -import asyncio - -import tornado.locks from .client import _get_global_client from .utils import log_errors, TimeoutError @@ -40,7 +38,7 @@ async def acquire(self, stream=None, name=None, id=None, timeout=None): result = True else: while name in self.ids: - event = tornado.locks.Event() + event = asyncio.Event() self.events[name].append(event) future = event.wait() if timeout is not None: diff --git a/distributed/nanny.py b/distributed/nanny.py index 945f33041d3..299b8bd2aff 100644 --- a/distributed/nanny.py +++ b/distributed/nanny.py @@ -12,7 +12,6 @@ import dask from dask.system import CPU_COUNT from tornado.ioloop import IOLoop -from tornado.locks import Event from tornado import gen from .comm import get_address_host, unparse_host_port @@ -507,8 +506,8 @@ async def start(self): ) self.process.daemon = dask.config.get("distributed.worker.daemon", default=True) self.process.set_exit_callback(self._on_exit) - self.running = Event() - self.stopped = Event() + self.running = asyncio.Event() + self.stopped = asyncio.Event() self.status = "starting" try: await self.process.start() diff --git a/distributed/pubsub.py b/distributed/pubsub.py index 9de133ddb47..889102dd7a2 100644 --- a/distributed/pubsub.py +++ b/distributed/pubsub.py @@ -1,14 +1,12 @@ +import asyncio from collections import defaultdict, deque -import datetime import logging import threading import weakref -import tornado.locks -from tornado import gen - from .core import CommClosedError -from .utils import sync, TimeoutError +from .metrics import time +from .utils import sync, TimeoutError, ignoring from .protocol.serialize import to_serialize logger = logging.getLogger(__name__) @@ -148,9 +146,9 @@ def remove_subscriber(self, name=None, address=None): def publish_scheduler(self, name=None, publish=None): self.publish_to_scheduler[name] = publish - def handle_message(self, name=None, msg=None): + async def handle_message(self, name=None, msg=None): for sub in self.subscribers.get(name, []): - sub._put(msg) + await sub._put(msg) def trigger_cleanup(self): self.worker.loop.add_callback(self.cleanup) @@ -180,9 +178,9 @@ def __init__(self, client): self.subscribers = defaultdict(weakref.WeakSet) self.client.extensions["pubsub"] = self # TODO: circular reference - def handle_message(self, name=None, msg=None): + async def handle_message(self, name=None, msg=None): for sub in self.subscribers[name]: - sub._put(msg) + await sub._put(msg) if not self.subscribers[name]: self.client.scheduler_comm.send( @@ -374,7 +372,7 @@ def __init__(self, name, worker=None, client=None): self.loop = self.client.loop self.name = name self.buffer = deque() - self.condition = tornado.locks.Condition() + self.condition = asyncio.Condition(loop=self.loop.asyncio_loop) if self.worker: pubsub = self.worker.extensions["pubsub"] @@ -393,20 +391,24 @@ def __init__(self, name, worker=None, client=None): weakref.finalize(self, pubsub.trigger_cleanup) async def _get(self, timeout=None): - if timeout is not None: - timeout = datetime.timedelta(seconds=timeout) - start = datetime.datetime.now() + start = time() while not self.buffer: if timeout is not None: - timeout2 = timeout - (datetime.datetime.now() - start) - if timeout2.total_seconds() < 0: + timeout2 = timeout - (time() - start) + if timeout2 < 0: raise TimeoutError() else: timeout2 = None + + async def _(): + await self.condition.acquire() + await self.condition.wait() + try: - await self.condition.wait(timeout=timeout2) - except gen.TimeoutError: - raise TimeoutError("Timed out waiting on Sub") + await asyncio.wait_for(_(), timeout2) + finally: + with ignoring(RuntimeError): # Python 3.6 fails here sometimes + self.condition.release() return self.buffer.popleft() @@ -431,9 +433,10 @@ def __iter__(self): def __aiter__(self): return self - def _put(self, msg): + async def _put(self, msg): self.buffer.append(msg) - self.condition.notify() + async with self.condition: + self.condition.notify() def __repr__(self): return "".format(self.name) diff --git a/distributed/queues.py b/distributed/queues.py index 6d1fc76571b..80fe30ac96d 100644 --- a/distributed/queues.py +++ b/distributed/queues.py @@ -1,14 +1,10 @@ +import asyncio from collections import defaultdict -import datetime import logging import uuid -import tornado.queues -from tornado.locks import Event -from tornado import gen - from .client import Future, _get_global_client, Client -from .utils import tokey, sync, thread_state, TimeoutError +from .utils import tokey, sync, thread_state from .worker import get_client logger = logging.getLogger(__name__) @@ -50,7 +46,7 @@ def __init__(self, scheduler): def create(self, stream=None, name=None, client=None, maxsize=0): logger.debug("Queue name: {}".format(name)) if name not in self.queues: - self.queues[name] = tornado.queues.Queue(maxsize=maxsize) + self.queues[name] = asyncio.Queue(maxsize=maxsize) self.client_refcount[name] = 1 else: self.client_refcount[name] += 1 @@ -77,12 +73,7 @@ async def put( self.scheduler.client_desires_keys(keys=[key], client="queue-%s" % name) else: record = {"type": "msgpack", "value": data} - if timeout is not None: - timeout = datetime.timedelta(seconds=timeout) - try: - await self.queues[name].put(record, timeout=timeout) - except gen.TimeoutError: - raise TimeoutError("Timed out waiting for Queue") + await asyncio.wait_for(self.queues[name].put(record), timeout=timeout) def future_release(self, name=None, key=None, client=None): self.future_refcount[name, key] -= 1 @@ -126,12 +117,7 @@ def process(record): out = [process(o) for o in out] return out else: - if timeout is not None: - timeout = datetime.timedelta(seconds=timeout) - try: - record = await self.queues[name].get(timeout=timeout) - except gen.TimeoutError: - raise TimeoutError("Timed out waiting for Queue") + record = await asyncio.wait_for(self.queues[name].get(), timeout=timeout) record = process(record) return record @@ -171,7 +157,7 @@ class Queue(object): def __init__(self, name=None, client=None, maxsize=0): self.client = client or _get_global_client() self.name = name or "queue-" + uuid.uuid4().hex - self._event_started = Event() + self._event_started = asyncio.Event() if self.client.asynchronous or getattr( thread_state, "on_event_loop_thread", False ): @@ -232,12 +218,9 @@ def qsize(self, **kwargs): return self.client.sync(self._qsize, **kwargs) async def _get(self, timeout=None, batch=False): - try: - resp = await self.client.scheduler.queue_get( - timeout=timeout, name=self.name, batch=batch - ) - except gen.TimeoutError: - raise TimeoutError("Timed out waiting for Queue") + resp = await self.client.scheduler.queue_get( + timeout=timeout, name=self.name, batch=batch + ) def process(d): if d["type"] == "Future": diff --git a/distributed/tests/test_nanny.py b/distributed/tests/test_nanny.py index 2ddc3b7e5db..0091a6126f1 100644 --- a/distributed/tests/test_nanny.py +++ b/distributed/tests/test_nanny.py @@ -12,7 +12,6 @@ from toolz import valmap, first from tornado import gen from tornado.ioloop import IOLoop -from tornado.locks import Event import dask from distributed.diagnostics import SchedulerPlugin @@ -453,7 +452,7 @@ async def test_nanny_closes_cleanly(cleanup): @pytest.mark.asyncio async def test_lifetime(cleanup): counter = 0 - event = Event() + event = asyncio.Event() class Plugin(SchedulerPlugin): def add_worker(self, **kwargs): diff --git a/distributed/tests/test_pubsub.py b/distributed/tests/test_pubsub.py index 847b0b88bf0..2e372dea88b 100644 --- a/distributed/tests/test_pubsub.py +++ b/distributed/tests/test_pubsub.py @@ -124,6 +124,8 @@ def test_timeouts(c, s, a, b): yield sub.get(timeout=0.1) stop = time() assert stop - start < 1 + with pytest.raises(TimeoutError): + yield sub.get(timeout=0.01) @gen_cluster(client=True) diff --git a/distributed/tests/test_variable.py b/distributed/tests/test_variable.py index 962b7a40e42..6e3b3bcdad6 100644 --- a/distributed/tests/test_variable.py +++ b/distributed/tests/test_variable.py @@ -5,9 +5,11 @@ import pytest from tornado import gen +from tornado.ioloop import IOLoop from distributed import Client, Variable, worker_client, Nanny, wait, TimeoutError from distributed.metrics import time +from distributed.compatibility import WINDOWS from distributed.utils_test import gen_cluster, inc, div from distributed.utils_test import client, cluster_fixture, loop # noqa: F401 @@ -83,20 +85,34 @@ def test_hold_futures(s, a, b): def test_timeout(c, s, a, b): v = Variable("v") - start = time() + start = IOLoop.current().time() + with pytest.raises(TimeoutError): + yield v.get(timeout=0.2) + stop = IOLoop.current().time() + + if WINDOWS: # timing is weird with asyncio and Windows + assert 0.1 < stop - start < 2.0 + else: + assert 0.2 < stop - start < 2.0 + with pytest.raises(TimeoutError): - yield v.get(timeout=0.1) - stop = time() - assert 0.1 < stop - start < 2.0 + yield v.get(timeout=0.01) def test_timeout_sync(client): v = Variable("v") - start = time() + start = IOLoop.current().time() + with pytest.raises(TimeoutError): + v.get(timeout=0.2) + stop = IOLoop.current().time() + + if WINDOWS: + assert 0.1 < stop - start < 2.0 + else: + assert 0.2 < stop - start < 2.0 + with pytest.raises(TimeoutError): - v.get(timeout=0.1) - stop = time() - assert 0.1 < stop - start < 2.0 + yield v.get(timeout=0.01) @gen_cluster(client=True) diff --git a/distributed/variable.py b/distributed/variable.py index 677e2997b32..975e5d48e8c 100644 --- a/distributed/variable.py +++ b/distributed/variable.py @@ -3,17 +3,13 @@ import logging import uuid -import tornado.locks -from tornado import gen - try: from cytoolz import merge except ImportError: from toolz import merge from .client import Future, _get_global_client, Client -from .metrics import time -from .utils import tokey, log_errors, TimeoutError +from .utils import tokey, log_errors, TimeoutError, ignoring from .worker import get_client logger = logging.getLogger(__name__) @@ -33,8 +29,8 @@ def __init__(self, scheduler): self.scheduler = scheduler self.variables = dict() self.waiting = defaultdict(set) - self.waiting_conditions = defaultdict(tornado.locks.Condition) - self.started = tornado.locks.Condition() + self.waiting_conditions = defaultdict(asyncio.Condition) + self.started = asyncio.Condition() self.scheduler.handlers.update( {"variable_set": self.set, "variable_get": self.get} @@ -45,7 +41,7 @@ def __init__(self, scheduler): self.scheduler.extensions["variables"] = self - def set(self, stream=None, name=None, key=None, data=None, client=None): + async def set(self, stream=None, name=None, key=None, data=None, client=None): if key is not None: record = {"type": "Future", "value": key} self.scheduler.client_desires_keys(keys=[key], client="variable-%s" % name) @@ -59,34 +55,44 @@ def set(self, stream=None, name=None, key=None, data=None, client=None): if old["type"] == "Future" and old["value"] != key: asyncio.ensure_future(self.release(old["value"], name)) if name not in self.variables: - self.started.notify_all() + async with self.started: + self.started.notify_all() self.variables[name] = record async def release(self, key, name): while self.waiting[key, name]: - await self.waiting_conditions[name].wait() + 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] - def future_release(self, name=None, key=None, token=None, client=None): + async def future_release(self, name=None, key=None, token=None, client=None): self.waiting[key, name].remove(token) if not self.waiting[key, name]: - self.waiting_conditions[name].notify_all() + async with self.waiting_conditions[name]: + self.waiting_conditions[name].notify_all() async def get(self, stream=None, name=None, client=None, timeout=None): - start = time() + start = self.scheduler.loop.time() while name not in self.variables: if timeout is not None: - left = timeout - (time() - start) + left = timeout - (self.scheduler.loop.time() - start) else: left = None if left and left < 0: raise TimeoutError() try: - await self.started.wait(timeout=left) - except gen.TimeoutError: - raise TimeoutError("Timed out waiting for Variable.get") + + async def _(): # Python 3.6 is odd and requires special help here + await self.started.acquire() + await self.started.wait() + + await asyncio.wait_for(_(), timeout=left) + finally: + with ignoring(RuntimeError): # Python 3.6 loses lock on finally clause + self.started.release() + record = self.variables[name] if record["type"] == "Future": key = record["value"] diff --git a/docs/source/actors.rst b/docs/source/actors.rst index b8bbebc743a..d3bf4f23c30 100644 --- a/docs/source/actors.rst +++ b/docs/source/actors.rst @@ -199,7 +199,7 @@ will run on the Worker's event loop thread rather than a separate thread. def Waiter(object): def __init__(self): - self.event = tornado.locks.Event() + self.event = asyncio.Event() async def set(self): self.event.set()