Replace tornado.locks with asyncio for Events/Locks/Conditions/Semaphore

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,19 +4198,14 @@ 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
 
         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()
 

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

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)

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()

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 _():

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:

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()