Repository: dask/distributed

Issue:
add a timeout to client.as_completed that mirrors concurrent.futures.as_completed's timeout
the timeout kwarg should work the same as https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.as_completed ie Futures that exceed the timeout are not cancelled, and

> The returned iterator raises a [TimeoutError](https://docs.python.org/3/library/exceptions.html#TimeoutError) if [__next__()](https://docs.python.org/3/library/stdtypes.html#iterator.__next__) is called and the result isn’t available after timeout seconds from the original call to [as_completed()](https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.as_completed). timeout can be an int or float. If timeout is not specified or None, there is no limit to the wait time.
add a timeout to client.as_completed that mirrors concurrent.futures.as_completed's timeout
the timeout kwarg should work the same as https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.as_completed ie Futures that exceed the timeout are not cancelled, and

> The returned iterator raises a [TimeoutError](https://docs.python.org/3/library/exceptions.html#TimeoutError) if [__next__()](https://docs.python.org/3/library/stdtypes.html#iterator.__next__) is called and the result isn’t available after timeout seconds from the original call to [as_completed()](https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.as_completed). timeout can be an int or float. If timeout is not specified or None, there is no limit to the wait time.


Discussion:

Resolving patch:
diff --git a/distributed/client.py b/distributed/client.py
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -48,7 +48,7 @@
 
 from distributed.core import ErrorMessage
 from distributed.protocol.serialize import _is_dumpable
-from distributed.utils import wait_for
+from distributed.utils import Deadline, wait_for
 
 try:
     from dask.delayed import single_key
@@ -5254,10 +5254,16 @@ class as_completed:
         complete
     with_results: bool (False)
         Whether to wait and include results of futures as well;
-        in this case `as_completed` yields a tuple of (future, result)
+        in this case ``as_completed`` yields a tuple of (future, result)
     raise_errors: bool (True)
         Whether we should raise when the result of a future raises an
-        exception; only affects behavior when `with_results=True`.
+        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
     --------
@@ -5294,7 +5300,15 @@ class as_completed:
     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)
@@ -5304,6 +5318,7 @@ def __init__(self, futures=None, loop=None, with_results=False, raise_errors=Tru
         self.thread_condition = threading.Condition()
         self.with_results = with_results
         self.raise_errors = raise_errors
+        self._deadline = Deadline.after(parse_timedelta(timeout))
 
         if futures:
             self.update(futures)
@@ -5402,6 +5417,8 @@ def _get_and_raise(self):
 
     def __next__(self):
         while self.queue.empty():
+            if self._deadline.expired:
+                raise TimeoutError()
             if self.is_empty():
                 raise StopIteration()
             with self.thread_condition:
@@ -5409,6 +5426,11 @@ def __next__(self):
         return self._get_and_raise()
 
     async def __anext__(self):
+        if not self._deadline.expires:
+            return await self._anext()
+        return await wait_for(self._anext(), self._deadline.remaining)
+
+    async def _anext(self):
         if not self.futures and self.queue.empty():
             raise StopAsyncIteration
         while self.queue.empty():


Test patch:
diff --git a/distributed/tests/test_as_completed.py b/distributed/tests/test_as_completed.py
--- a/distributed/tests/test_as_completed.py
+++ b/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)