add a timeout to client.as_completed that mirrors concurrent.futures.as_completed's timeout

distributed/client.py

@@ -78,7 +78,7 @@
     WorkerPlugin,
     _get_plugin_name,
 )
-from distributed.metrics import time
+from distributed.metrics import monotonic, time
 from distributed.objects import HasWhat, SchedulerInfo, WhoHas
 from distributed.protocol import to_serialize
 from distributed.protocol.pickle import dumps, loads
@@ -5226,6 +5226,12 @@
     raise_errors: bool (True)
         Whether we should raise when the result of a future raises an
         exception; only affects behavior when `with_results=True`.
+    timeout: int (optional)
+        The returned iterator raises a ``dask.distributed.TimeoutError``
+        if `__next__()` or `__anext__()` is called and the result
+        isn't available after timeout seconds from the original call to
+        `as_completed()`. If timeout is not specified or None, there is no limit
+        to the wait time.
 
     Examples
     --------
@@ -5262,7 +5268,15 @@
     3
     """
 
-    def __init__(self, futures=None, loop=None, with_results=False, raise_errors=True):
+    def __init__(
+        self,
+        futures=None,
+        loop=None,
+        with_results=False,
+        raise_errors=True,
+        *,
+        timeout=None,
+    ):
         if futures is None:
             futures = []
         self.futures = defaultdict(lambda: 0)
@@ -5273,6 +5287,11 @@
         self.with_results = with_results
         self.raise_errors = raise_errors
 
+        if timeout is not None:
+            self._end_time = parse_timedelta(timeout) + monotonic()
+        else:
+            self._end_time = None
+
         if futures:
             self.update(futures)
 
@@ -5370,13 +5389,20 @@
 
     def __next__(self):
         while self.queue.empty():
+            if self._end_time is not None and (self._end_time - monotonic()) < 0:
+                raise TimeoutError()
             if self.is_empty():
                 raise StopIteration()
             with self.thread_condition:
                 self.thread_condition.wait(timeout=0.100)
         return self._get_and_raise()
 
     async def __anext__(self):
+        if self._end_time is None:
+            return await self._anext()
+        return await wait_for(self._anext(), self._end_time - monotonic())
+
+    async def _anext(self):
         if not self.futures and self.queue.empty():
             raise StopAsyncIteration
         while self.queue.empty():

distributed/tests/test_as_completed.py

@@ -9,10 +9,11 @@
 
 import pytest
 
+from distributed import TimeoutError
 from distributed.client import _as_completed, _first_completed, as_completed, wait
 from distributed.metrics import time
 from distributed.utils import CancelledError
-from distributed.utils_test import gen_cluster, inc, throws
+from distributed.utils_test import gen_cluster, inc, slowinc, throws
 
 
 @gen_cluster(client=True)
@@ -31,6 +32,31 @@ async def test_as_completed_async(c, s, a, b):
     assert result in [x, y, z]
 
 
+@gen_cluster(client=True)
+async def test_as_completed_timeout_async(c, s, a, b):
+    x = c.submit(inc, 1)
+    y = c.submit(inc, 1)
+    z = c.submit(slowinc, 2, delay=1)
+
+    seq = as_completed([x, y, z], timeout="0.5s")
+    await seq.__anext__()
+    await seq.__anext__()
+    with pytest.raises(TimeoutError):
+        await seq.__anext__()
+
+
+def test_as_completed_timeout_sync(client):
+    x = client.submit(inc, 1)
+    y = client.submit(inc, 1)
+    z = client.submit(slowinc, 2, delay=1)
+
+    seq = as_completed([x, y, z], timeout="0.5s")
+    next(seq)
+    next(seq)
+    with pytest.raises(TimeoutError):
+        next(seq)
+
+
 def test_as_completed_sync(client):
     x = client.submit(inc, 1)
     y = client.submit(inc, 2)