Replace gen.with_timeout with asyncio.wait_for

distributed/client.py

@@ -5,7 +5,6 @@
 from concurrent.futures._base import DoneAndNotDoneFutures
 from contextlib import contextmanager
 import copy
-from datetime import timedelta
 import errno
 from functools import partial
 import html
@@ -761,7 +760,7 @@ def sync(self, func, *args, asynchronous=None, callback_timeout=None, **kwargs):
         ):
             future = func(*args, **kwargs)
             if callback_timeout is not None:
-                future = gen.with_timeout(timedelta(seconds=callback_timeout), future)
+                future = asyncio.wait_for(future, callback_timeout)
             return future
         else:
             return sync(
@@ -1043,9 +1042,7 @@ async def _ensure_connected(self, timeout=None):
             )
             comm.name = "Client->Scheduler"
             if timeout is not None:
-                await gen.with_timeout(
-                    timedelta(seconds=timeout), self._update_scheduler_info()
-                )
+                await asyncio.wait_for(self._update_scheduler_info(), timeout)
             else:
                 await self._update_scheduler_info()
             await comm.write(
@@ -1064,7 +1061,7 @@ async def _ensure_connected(self, timeout=None):
         finally:
             self._connecting_to_scheduler = False
         if timeout is not None:
-            msg = await gen.with_timeout(timedelta(seconds=timeout), comm.read())
+            msg = await asyncio.wait_for(comm.read(), timeout)
         else:
             msg = await comm.read()
         assert len(msg) == 1
@@ -1268,11 +1265,9 @@ async def _close(self, fast=False):
 
             # Give the scheduler 'stream-closed' message 100ms to come through
             # This makes the shutdown slightly smoother and quieter
-            with ignoring(AttributeError, gen.TimeoutError):
-                await gen.with_timeout(
-                    timedelta(milliseconds=100),
-                    self._handle_scheduler_coroutine,
-                    quiet_exceptions=(CancelledError,),
+            with ignoring(AttributeError, CancelledError, asyncio.TimeoutError):
+                await asyncio.wait_for(
+                    asyncio.shield(self._handle_scheduler_coroutine), 0.1
                 )
 
             if (
@@ -1308,7 +1303,7 @@ async def _close(self, fast=False):
 
             if not fast:
                 with ignoring(TimeoutError):
-                    await gen.with_timeout(timedelta(seconds=2), list(coroutines))
+                    await asyncio.wait_for(asyncio.gather(*coroutines), 2)
 
             with ignoring(AttributeError):
                 await self.scheduler.close_rpc()
@@ -1344,7 +1339,7 @@ def close(self, timeout=no_default):
         if self.asynchronous:
             future = self._close()
             if timeout:
-                future = gen.with_timeout(timedelta(seconds=timeout), future)
+                future = asyncio.wait_for(future, timeout)
             return future
 
         if self._start_arg is None:
@@ -4077,7 +4072,7 @@ async def _wait(fs, timeout=None, return_when=ALL_COMPLETED):
 
     future = wait_for({f._state.wait() for f in fs})
     if timeout is not None:
-        future = gen.with_timeout(timedelta(seconds=timeout), future)
+        future = asyncio.wait_for(future, timeout)
     await future
 
     done, not_done = (

distributed/deploy/cluster.py

@@ -1,10 +1,8 @@
 import asyncio
-from datetime import timedelta
 import logging
 import threading
 
 from dask.utils import format_bytes
-from tornado import gen
 
 from .adaptive import Adaptive
 
@@ -156,7 +154,7 @@ def sync(self, func, *args, asynchronous=None, callback_timeout=None, **kwargs):
         if asynchronous:
             future = func(*args, **kwargs)
             if callback_timeout is not None:
-                future = gen.with_timeout(timedelta(seconds=callback_timeout), future)
+                future = asyncio.wait_for(future, callback_timeout)
             return future
         else:
             return sync(self.loop, func, *args, **kwargs)

distributed/deploy/tests/test_local.py

@@ -523,7 +523,7 @@ def test_memory_nanny(loop, n_workers):
 
 
 def test_death_timeout_raises(loop):
-    with pytest.raises(gen.TimeoutError):
+    with pytest.raises(asyncio.TimeoutError):
         with LocalCluster(
             scheduler_port=0,
             silence_logs=False,

distributed/lock.py

@@ -1,9 +1,8 @@
 from collections import defaultdict, deque
-from datetime import timedelta
 import logging
 import uuid
+import asyncio
 
-from tornado import gen
 import tornado.locks
 
 from .client import _get_global_client
@@ -45,10 +44,10 @@ async def acquire(self, stream=None, name=None, id=None, timeout=None):
                     self.events[name].append(event)
                     future = event.wait()
                     if timeout is not None:
-                        future = gen.with_timeout(timedelta(seconds=timeout), future)
+                        future = asyncio.wait_for(future, timeout)
                     try:
                         await future
-                    except gen.TimeoutError:
+                    except asyncio.TimeoutError:
                         result = False
                         break
                     else:

distributed/nanny.py

@@ -1,5 +1,4 @@
 import asyncio
-from datetime import timedelta
 import logging
 from multiprocessing.queues import Empty
 import os
@@ -31,6 +30,7 @@
     json_load_robust,
     PeriodicCallback,
     parse_timedelta,
+    ignoring,
 )
 from .worker import run, parse_memory_limit, Worker
 
@@ -219,14 +219,10 @@ async def _unregister(self, timeout=10):
             EnvironmentError,
             RPCClosed,
         )
-        try:
-            await gen.with_timeout(
-                timedelta(seconds=timeout),
-                self.scheduler.unregister(address=self.worker_address),
-                quiet_exceptions=allowed_errors,
+        with ignoring(allowed_errors):
+            await asyncio.wait_for(
+                self.scheduler.unregister(address=self.worker_address), timeout
             )
-        except allowed_errors:
-            pass
 
     @property
     def worker_address(self):
@@ -318,8 +314,8 @@ async def instantiate(self, comm=None):
         self.auto_restart = True
         if self.death_timeout:
             try:
-                result = await gen.with_timeout(
-                    timedelta(seconds=self.death_timeout), self.process.start()
+                result = await asyncio.wait_for(
+                    self.process.start(), self.death_timeout
                 )
             except gen.TimeoutError:
                 await self.close(timeout=self.death_timeout)
@@ -343,8 +339,8 @@ async def _():
                 await self.instantiate()
 
         try:
-            await gen.with_timeout(timedelta(seconds=timeout), _())
-        except gen.TimeoutError:
+            await asyncio.wait_for(_(), timeout)
+        except asyncio.TimeoutError:
             logger.error("Restart timed out, returning before finished")
             return "timed out"
         else:

distributed/node.py

@@ -173,7 +173,7 @@ async def wait_for(future, timeout=None):
                         await asyncio.wait_for(future, timeout=timeout)
                     except Exception:
                         await self.close(timeout=1)
-                        raise gen.TimeoutError(
+                        raise asyncio.TimeoutError(
                             "{} failed to start in {} seconds".format(
                                 type(self).__name__, timeout
                             )

distributed/process.py

@@ -1,11 +1,11 @@
 import atexit
-from datetime import timedelta
 import logging
 import os
 from queue import Queue as PyQueue
 import re
 import threading
 import weakref
+import asyncio
 import dask
 
 from .utils import mp_context
@@ -282,7 +282,7 @@ def join(self, timeout=None):
             yield self._exit_future
         else:
             try:
-                yield gen.with_timeout(timedelta(seconds=timeout), self._exit_future)
+                yield asyncio.wait_for(self._exit_future, timeout)
             except gen.TimeoutError:
                 pass
 

distributed/scheduler.py

@@ -2878,7 +2878,7 @@ async def restart(self, client=None, timeout=3):
                         for nanny in nannies
                     ]
                 )
-                resps = await gen.with_timeout(timedelta(seconds=timeout), resps)
+                resps = await asyncio.wait_for(resps, timeout)
                 if not all(resp == "OK" for resp in resps):
                     logger.error(
                         "Not all workers responded positively: %s", resps, exc_info=True

distributed/tests/test_batched.py

@@ -1,10 +1,8 @@
 import asyncio
-from datetime import timedelta
 import random
 
 import pytest
 from toolz import assoc
-from tornado import gen
 
 from distributed.batched import BatchedSend
 from distributed.core import listen, connect, CommClosedError
@@ -170,7 +168,7 @@ async def send():
 
         async def recv():
             while True:
-                result = await gen.with_timeout(timedelta(seconds=1), comm.read())
+                result = await asyncio.wait_for(comm.read(), 1)
                 L.extend(result)
                 if result[-1] == 9999:
                     break
@@ -205,7 +203,7 @@ async def run_traffic_jam(nsends, nbytes):
             # If this times out then I think it's a backpressure issue
             # Somehow we're able to flood the socket so that the receiving end
             # loses some of our messages
-            L = await gen.with_timeout(timedelta(seconds=5), comm.read())
+            L = await asyncio.wait_for(comm.read(), 5)
             count += 1
             results.extend(r["i"] for r in L)
 
@@ -254,5 +252,5 @@ async def test_serializers():
         msg = await comm.read()
         assert list(msg) == [{"x": 123}, {"x": "hello"}]
 
-        with pytest.raises(gen.TimeoutError):
-            msg = await gen.with_timeout(timedelta(milliseconds=100), comm.read())
+        with pytest.raises(asyncio.TimeoutError):
+            msg = await asyncio.wait_for(comm.read(), 0.1)

distributed/tests/test_client.py

@@ -496,7 +496,7 @@ def test_thread(c):
     assert x.result() == 2
 
     x = c.submit(slowinc, 1, delay=0.3)
-    with pytest.raises(gen.TimeoutError):
+    with pytest.raises((gen.TimeoutError, asyncio.TimeoutError)):
         x.result(timeout=0.01)
     assert x.result() == 2
 
@@ -681,7 +681,7 @@ def test_wait_first_completed(c, s, a, b):
 @gen_cluster(client=True, timeout=2)
 def test_wait_timeout(c, s, a, b):
     future = c.submit(sleep, 0.3)
-    with pytest.raises(gen.TimeoutError):
+    with pytest.raises(asyncio.TimeoutError):
         yield wait(future, timeout=0.01)
 
 
@@ -695,7 +695,7 @@ def test_wait_sync(c):
     assert x.status == y.status == "finished"
 
     future = c.submit(sleep, 0.3)
-    with pytest.raises(gen.TimeoutError):
+    with pytest.raises(asyncio.TimeoutError):
         wait(future, timeout=0.01)
 
 
@@ -5279,7 +5279,7 @@ def test_client_active_bad_port():
     http_server.listen(8080)
     with dask.config.set({"distributed.comm.timeouts.connect": "10ms"}):
         c = Client("127.0.0.1:8080", asynchronous=True)
-        with pytest.raises((TimeoutError, IOError)):
+        with pytest.raises((asyncio.TimeoutError, IOError)):
             yield c
         yield c._close(fast=True)
     http_server.stop()

distributed/tests/test_client_executor.py

@@ -14,7 +14,16 @@
 import pytest
 from toolz import take
 
-from distributed.utils_test import slowinc, slowadd, slowdec, inc, throws, varying
+from distributed import Client
+from distributed.utils_test import (
+    slowinc,
+    slowadd,
+    slowdec,
+    inc,
+    throws,
+    varying,
+    cluster,
+)
 from distributed.utils_test import client, cluster_fixture, loop, s, a, b  # noqa: F401
 
 
@@ -218,30 +227,32 @@ def test_retries(client):
         exc_info.match("one")
 
 
-def test_shutdown(client):
-    # shutdown(wait=True) waits for pending tasks to finish
-    e = client.get_executor()
-    fut = e.submit(time.sleep, 1.0)
-    t1 = time.time()
-    e.shutdown()
-    dt = time.time() - t1
-    assert 0.5 <= dt <= 2.0
-    time.sleep(0.1)  # wait for future outcome to propagate
-    assert fut.done()
-    fut.result()  # doesn't raise
-
-    with pytest.raises(RuntimeError):
-        e.submit(time.sleep, 1.0)
-
-    # shutdown(wait=False) cancels pending tasks
-    e = client.get_executor()
-    fut = e.submit(time.sleep, 2.0)
-    t1 = time.time()
-    e.shutdown(wait=False)
-    dt = time.time() - t1
-    assert dt < 0.5
-    time.sleep(0.1)  # wait for future outcome to propagate
-    assert fut.cancelled()
-
-    with pytest.raises(RuntimeError):
-        e.submit(time.sleep, 1.0)
+def test_shutdown(loop):
+    with cluster(disconnect_timeout=10) as (s, [a, b]):
+        with Client(s["address"], loop=loop) as client:
+            # shutdown(wait=True) waits for pending tasks to finish
+            e = client.get_executor()
+            fut = e.submit(time.sleep, 1.0)
+            t1 = time.time()
+            e.shutdown()
+            dt = time.time() - t1
+            assert 0.5 <= dt <= 2.0
+            time.sleep(0.1)  # wait for future outcome to propagate
+            assert fut.done()
+            fut.result()  # doesn't raise
+
+            with pytest.raises(RuntimeError):
+                e.submit(time.sleep, 1.0)
+
+            # shutdown(wait=False) cancels pending tasks
+            e = client.get_executor()
+            fut = e.submit(time.sleep, 2.0)
+            t1 = time.time()
+            e.shutdown(wait=False)
+            dt = time.time() - t1
+            assert dt < 0.5
+            time.sleep(0.1)  # wait for future outcome to propagate
+            assert fut.cancelled()
+
+            with pytest.raises(RuntimeError):
+                e.submit(time.sleep, 1.0)

distributed/tests/test_failed_workers.py

@@ -26,7 +26,7 @@
 
 
 def test_submit_after_failed_worker_sync(loop):
-    with cluster(active_rpc_timeout=10) as (s, [a, b]):
+    with cluster(active_rpc_timeout=10, disconnect_timeout=10) as (s, [a, b]):
         with Client(s["address"], loop=loop) as c:
             L = c.map(inc, range(10))
             wait(L)
@@ -64,7 +64,7 @@ def test_submit_after_failed_worker(c, s, a, b):
 
 
 def test_gather_after_failed_worker(loop):
-    with cluster(active_rpc_timeout=10) as (s, [a, b]):
+    with cluster(active_rpc_timeout=10, disconnect_timeout=10) as (s, [a, b]):
         with Client(s["address"], loop=loop) as c:
             L = c.map(inc, range(10))
             wait(L)