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.

Discussion:

Patch:
diff --git a/distributed/client.py b/distributed/client.py
index 08e8c252..413c41c2 100644
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -48,7 +48,7 @@ from dask.widgets import get_template
 
 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 @@ class as_completed:
         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 @@ class as_completed:
 
     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 @@ class as_completed:
         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():
diff --git a/distributed/shuffle/_arrow.py b/distributed/shuffle/_arrow.py
index 62193d21..6a0bc81e 100644
--- a/distributed/shuffle/_arrow.py
+++ b/distributed/shuffle/_arrow.py
@@ -45,7 +45,7 @@ def check_minimal_arrow_version() -> None:
         )
 
 
-def convert_partition(data: bytes) -> pa.Table:
+def convert_partition(data: bytes, meta: pd.DataFrame) -> pd.DataFrame:
     import pyarrow as pa
 
     file = BytesIO(data)
@@ -54,7 +54,9 @@ def convert_partition(data: bytes) -> pa.Table:
     while file.tell() < end:
         sr = pa.RecordBatchStreamReader(file)
         shards.append(sr.read_all())
-    return pa.concat_tables(shards)
+    table = pa.concat_tables(shards)
+    df = table.to_pandas(self_destruct=True)
+    return df.astype(meta.dtypes)
 
 
 def list_of_buffers_to_table(data: list[bytes]) -> pa.Table:
@@ -64,31 +66,6 @@ def list_of_buffers_to_table(data: list[bytes]) -> pa.Table:
     return pa.concat_tables(deserialize_table(buffer) for buffer in data)
 
 
-def deserialize_schema(data: bytes) -> pa.Schema:
-    """Deserialize an arrow schema
-
-    Examples
-    --------
-    >>> b = schema.serialize()  # doctest: +SKIP
-    >>> deserialize_schema(b)  # doctest: +SKIP
-
-    See also
-    --------
-    pa.Schema.serialize
-    """
-    import io
-
-    import pyarrow as pa
-
-    bio = io.BytesIO()
-    bio.write(data)
-    bio.seek(0)
-    sr = pa.RecordBatchStreamReader(bio)
-    table = sr.read_all()
-    bio.close()
-    return table.schema
-
-
 def serialize_table(table: pa.Table) -> bytes:
     import io
 
diff --git a/distributed/shuffle/_merge.py b/distributed/shuffle/_merge.py
index 40454a12..8a68b08a 100644
--- a/distributed/shuffle/_merge.py
+++ b/distributed/shuffle/_merge.py
@@ -108,9 +108,11 @@ def hash_join_p2p(
     join_layer = HashJoinP2PLayer(
         name=merge_name,
         name_input_left=lhs._name,
+        meta_input_left=lhs._meta,
         left_on=left_on,
         n_partitions_left=lhs.npartitions,
         name_input_right=rhs._name,
+        meta_input_right=rhs._meta,
         right_on=right_on,
         n_partitions_right=rhs.npartitions,
         meta_output=meta,
@@ -138,6 +140,7 @@ def merge_transfer(
     input_partition: int,
     npartitions: int,
     parts_out: set[int],
+    meta: pd.DataFrame,
 ):
     return shuffle_transfer(
         input=input,
@@ -146,6 +149,7 @@ def merge_transfer(
         npartitions=npartitions,
         column=_HASH_COLUMN_NAME,
         parts_out=parts_out,
+        meta=meta,
     )
 
 
@@ -164,12 +168,13 @@ def merge_unpack(
     from dask.dataframe.multi import merge_chunk
 
     ext = _get_worker_extension()
+    # If the partition is empty, it doesn't contain the hash column name
     left = ext.get_output_partition(
         shuffle_id_left, barrier_left, output_partition
-    ).drop(columns=_HASH_COLUMN_NAME)
+    ).drop(columns=_HASH_COLUMN_NAME, errors="ignore")
     right = ext.get_output_partition(
         shuffle_id_right, barrier_right, output_partition
-    ).drop(columns=_HASH_COLUMN_NAME)
+    ).drop(columns=_HASH_COLUMN_NAME, errors="ignore")
     return merge_chunk(
         left,
         right,
@@ -186,10 +191,12 @@ class HashJoinP2PLayer(Layer):
         self,
         name: str,
         name_input_left: str,
+        meta_input_left: pd.DataFrame,
         left_on,
         n_partitions_left: int,
         n_partitions_right: int,
         name_input_right: str,
+        meta_input_right: pd.DataFrame,
         right_on,
         meta_output: pd.DataFrame,
         left_index: bool,
@@ -203,8 +210,10 @@ class HashJoinP2PLayer(Layer):
     ) -> None:
         self.name = name
         self.name_input_left = name_input_left
+        self.meta_input_left = meta_input_left
         self.left_on = left_on
         self.name_input_right = name_input_right
+        self.meta_input_right = meta_input_right
         self.right_on = right_on
         self.how = how
         self.npartitions = npartitions
@@ -285,8 +294,10 @@ class HashJoinP2PLayer(Layer):
         return HashJoinP2PLayer(
             name=self.name,
             name_input_left=self.name_input_left,
+            meta_input_left=self.meta_input_left,
             left_on=self.left_on,
             name_input_right=self.name_input_right,
+            meta_input_right=self.meta_input_right,
             right_on=self.right_on,
             how=self.how,
             npartitions=self.npartitions,
@@ -344,6 +355,7 @@ class HashJoinP2PLayer(Layer):
                 i,
                 self.npartitions,
                 self.parts_out,
+                self.meta_input_left,
             )
         for i in range(self.n_partitions_right):
             transfer_keys_right.append((name_right, i))
@@ -354,6 +366,7 @@ class HashJoinP2PLayer(Layer):
                 i,
                 self.npartitions,
                 self.parts_out,
+                self.meta_input_right,
             )
 
         _barrier_key_left = barrier_key(ShuffleId(token_left))
diff --git a/distributed/shuffle/_scheduler_extension.py b/distributed/shuffle/_scheduler_extension.py
index 47dfe33c..ba511b63 100644
--- a/distributed/shuffle/_scheduler_extension.py
+++ b/distributed/shuffle/_scheduler_extension.py
@@ -12,6 +12,7 @@ from itertools import product
 from typing import TYPE_CHECKING, Any, ClassVar
 
 from distributed.diagnostics.plugin import SchedulerPlugin
+from distributed.protocol import to_serialize
 from distributed.shuffle._rechunk import ChunkedAxes, NIndex
 from distributed.shuffle._shuffle import (
     ShuffleId,
@@ -21,6 +22,8 @@ from distributed.shuffle._shuffle import (
 )
 
 if TYPE_CHECKING:
+    import pandas as pd
+
     from distributed.scheduler import (
         Recs,
         Scheduler,
@@ -50,7 +53,7 @@ class ShuffleState(abc.ABC):
 class DataFrameShuffleState(ShuffleState):
     type: ClassVar[ShuffleType] = ShuffleType.DATAFRAME
     worker_for: dict[int, str]
-    schema: bytes
+    meta: pd.DataFrame
     column: str
 
     def to_msg(self) -> dict[str, Any]:
@@ -60,7 +63,7 @@ class DataFrameShuffleState(ShuffleState):
             "run_id": self.run_id,
             "worker_for": self.worker_for,
             "column": self.column,
-            "schema": self.schema,
+            "meta": to_serialize(self.meta),
             "output_workers": self.output_workers,
         }
 
@@ -186,11 +189,11 @@ class ShuffleSchedulerExtension(SchedulerPlugin):
     def _create_dataframe_shuffle_state(
         self, id: ShuffleId, spec: dict[str, Any]
     ) -> DataFrameShuffleState:
-        schema = spec["schema"]
+        meta = spec["meta"]
         column = spec["column"]
         npartitions = spec["npartitions"]
         parts_out = spec["parts_out"]
-        assert schema is not None
+        assert meta is not None
         assert column is not None
         assert npartitions is not None
         assert parts_out is not None
@@ -204,7 +207,7 @@ class ShuffleSchedulerExtension(SchedulerPlugin):
             id=id,
             run_id=next(ShuffleState._run_id_iterator),
             worker_for=mapping,
-            schema=schema,
+            meta=meta,
             column=column,
             output_workers=output_workers,
             participating_workers=output_workers.copy(),
diff --git a/distributed/shuffle/_shuffle.py b/distributed/shuffle/_shuffle.py
index 59798b23..ad3dc08f 100644
--- a/distributed/shuffle/_shuffle.py
+++ b/distributed/shuffle/_shuffle.py
@@ -58,6 +58,7 @@ def shuffle_transfer(
     npartitions: int,
     column: str,
     parts_out: set[int],
+    meta: pd.DataFrame,
 ) -> int:
     try:
         return _get_worker_extension().add_partition(
@@ -68,6 +69,7 @@ def shuffle_transfer(
             npartitions=npartitions,
             column=column,
             parts_out=parts_out,
+            meta=meta,
         )
     except Exception as e:
         raise RuntimeError(f"shuffle_transfer failed during shuffle {id}") from e
@@ -100,13 +102,13 @@ def rearrange_by_column_p2p(
 ) -> DataFrame:
     from dask.dataframe import DataFrame
 
-    check_dtype_support(df._meta)
+    meta = df._meta
+    check_dtype_support(meta)
     npartitions = npartitions or df.npartitions
     token = tokenize(df, column, npartitions)
 
-    empty = df._meta.copy()
-    if any(not isinstance(c, str) for c in empty.columns):
-        unsupported = {c: type(c) for c in empty.columns if not isinstance(c, str)}
+    if any(not isinstance(c, str) for c in meta.columns):
+        unsupported = {c: type(c) for c in meta.columns if not isinstance(c, str)}
         raise TypeError(
             f"p2p requires all column names to be str, found: {unsupported}",
         )
@@ -118,11 +120,12 @@ def rearrange_by_column_p2p(
         npartitions,
         npartitions_input=df.npartitions,
         name_input=df._name,
+        meta_input=meta,
     )
     return DataFrame(
         HighLevelGraph.from_collections(name, layer, [df]),
         name,
-        empty,
+        meta,
         [None] * (npartitions + 1),
     )
 
@@ -139,6 +142,7 @@ class P2PShuffleLayer(Layer):
         npartitions: int,
         npartitions_input: int,
         name_input: str,
+        meta_input: pd.DataFrame,
         parts_out: Iterable | None = None,
         annotations: dict | None = None,
     ):
@@ -147,6 +151,7 @@ class P2PShuffleLayer(Layer):
         self.column = column
         self.npartitions = npartitions
         self.name_input = name_input
+        self.meta_input = meta_input
         if parts_out:
             self.parts_out = set(parts_out)
         else:
@@ -195,6 +200,7 @@ class P2PShuffleLayer(Layer):
             self.npartitions,
             self.npartitions_input,
             self.name_input,
+            self.meta_input,
             parts_out=parts_out,
         )
 
@@ -245,6 +251,7 @@ class P2PShuffleLayer(Layer):
                 self.npartitions,
                 self.column,
                 self.parts_out,
+                self.meta_input,
             )
 
         dsk[_barrier_key] = (shuffle_barrier, token, transfer_keys)
diff --git a/distributed/shuffle/_worker_extension.py b/distributed/shuffle/_worker_extension.py
index 30a031fc..105b40e2 100644
--- a/distributed/shuffle/_worker_extension.py
+++ b/distributed/shuffle/_worker_extension.py
@@ -23,7 +23,6 @@ from distributed.exceptions import Reschedule
 from distributed.protocol import to_serialize
 from distributed.shuffle._arrow import (
     convert_partition,
-    deserialize_schema,
     list_of_buffers_to_table,
     serialize_table,
 )
@@ -421,8 +420,8 @@ class DataFrameShuffleRun(ShuffleRun[int, int, "pd.DataFrame"]):
         A set of all participating worker (addresses).
     column:
         The data column we split the input partition by.
-    schema:
-        The schema of the payload data.
+    meta:
+        Empty metadata of the input.
     id:
         A unique `ShuffleID` this belongs to.
     run_id:
@@ -451,7 +450,7 @@ class DataFrameShuffleRun(ShuffleRun[int, int, "pd.DataFrame"]):
         worker_for: dict[int, str],
         output_workers: set,
         column: str,
-        schema: pa.Schema,
+        meta: pd.DataFrame,
         id: ShuffleId,
         run_id: int,
         local_address: str,
@@ -477,7 +476,7 @@ class DataFrameShuffleRun(ShuffleRun[int, int, "pd.DataFrame"]):
             memory_limiter_disk=memory_limiter_disk,
         )
         self.column = column
-        self.schema = schema
+        self.meta = meta
         partitions_of = defaultdict(list)
         for part, addr in worker_for.items():
             partitions_of[addr].append(part)
@@ -543,12 +542,11 @@ class DataFrameShuffleRun(ShuffleRun[int, int, "pd.DataFrame"]):
             data = self._read_from_disk((i,))
 
             def _() -> pd.DataFrame:
-                df = convert_partition(data)
-                return df.to_pandas()
+                return convert_partition(data, self.meta)
 
             out = await self.offload(_)
         except KeyError:
-            out = self.schema.empty_table().to_pandas()
+            out = self.meta.copy()
         return out
 
     def _get_assigned_worker(self, i: int) -> str:
@@ -649,8 +647,6 @@ class ShuffleWorkerExtension:
         type: ShuffleType,
         **kwargs: Any,
     ) -> int:
-        if type == ShuffleType.DATAFRAME:
-            kwargs["empty"] = data
         shuffle = self.get_or_create_shuffle(shuffle_id, type=type, **kwargs)
         return sync(
             self.worker.loop,
@@ -723,12 +719,8 @@ class ShuffleWorkerExtension:
         ----------
         shuffle_id
             Unique identifier of the shuffle
-        empty
-            Empty metadata of input collection
-        column
-            Column to be used to map rows to output partitions (by hashing)
-        npartitions
-            Number of output partitions
+        type:
+            Type of the shuffle operation
         """
         shuffle = self.shuffles.get(shuffle_id, None)
         if shuffle is None:
@@ -774,16 +766,12 @@ class ShuffleWorkerExtension:
                 worker=self.worker.address,
             )
         elif type == ShuffleType.DATAFRAME:
-            import pyarrow as pa
-
             assert kwargs is not None
             result = await self.worker.scheduler.shuffle_get_or_create(
                 id=shuffle_id,
                 type=type,
                 spec={
-                    "schema": pa.Schema.from_pandas(kwargs["empty"])
-                    .serialize()
-                    .to_pybytes(),
+                    "meta": to_serialize(kwargs["meta"]),
                     "npartitions": kwargs["npartitions"],
                     "column": kwargs["column"],
                     "parts_out": kwargs["parts_out"],
@@ -829,7 +817,7 @@ class ShuffleWorkerExtension:
                 column=result["column"],
                 worker_for=result["worker_for"],
                 output_workers=result["output_workers"],
-                schema=deserialize_schema(result["schema"]),
+                meta=result["meta"],
                 id=shuffle_id,
                 run_id=result["run_id"],
                 directory=os.path.join(


Test patch:
diff --git a/distributed/shuffle/tests/test_shuffle.py b/distributed/shuffle/tests/test_shuffle.py
index 79662c43..28f94aea 100644
--- a/distributed/shuffle/tests/test_shuffle.py
+++ b/distributed/shuffle/tests/test_shuffle.py
@@ -717,15 +717,18 @@ def test_processing_chain():
                 [pd.Timestamp.fromtimestamp(1641034800 + i) for i in range(100)],
                 dtype=pd.ArrowDtype(pa.timestamp("ms")),
             ),
-            # FIXME: distributed#7420
-            # f"col{next(counter)}": pd.array(
-            #     ["lorem ipsum"] * 100,
-            #     dtype="string[pyarrow]",
-            # ),
-            # f"col{next(counter)}": pd.array(
-            #     ["lorem ipsum"] * 100,
-            #     dtype=pd.StringDtype("pyarrow"),
-            # ),
+            f"col{next(counter)}": pd.array(
+                ["lorem ipsum"] * 100,
+                dtype="string[pyarrow]",
+            ),
+            f"col{next(counter)}": pd.array(
+                ["lorem ipsum"] * 100,
+                dtype=pd.StringDtype("pyarrow"),
+            ),
+            f"col{next(counter)}": pd.array(
+                ["lorem ipsum"] * 100,
+                dtype="string[python]",
+            ),
             # custom objects
             # FIXME: Serializing custom objects is not supported in P2P shuffling
             # f"col{next(counter)}": pd.array(
@@ -734,7 +737,6 @@ def test_processing_chain():
         }
     )
     df["_partitions"] = df.col4 % npartitions
-    schema = pa.Schema.from_pandas(df)
     worker_for = {i: random.choice(workers) for i in list(range(npartitions))}
     worker_for = pd.Series(worker_for, name="_worker").astype("category")
 
@@ -783,9 +785,9 @@ def test_processing_chain():
     out = {}
     for k, bio in filesystem.items():
         bio.seek(0)
-        out[k] = convert_partition(bio.read())
+        out[k] = convert_partition(bio.read(), df.head(0))
 
-    shuffled_df = pd.concat(table.to_pandas() for table in out.values())
+    shuffled_df = pd.concat(df for df in out.values())
     pd.testing.assert_frame_equal(
         df,
         shuffled_df,
@@ -1188,7 +1190,7 @@ class DataFrameShuffleTestPool(AbstractShuffleTestPool):
         self,
         name,
         worker_for_mapping,
-        schema,
+        meta,
         directory,
         loop,
         Shuffle=DataFrameShuffleRun,
@@ -1198,7 +1200,7 @@ class DataFrameShuffleTestPool(AbstractShuffleTestPool):
             worker_for=worker_for_mapping,
             # FIXME: Is output_workers redundant with worker_for?
             output_workers=set(worker_for_mapping.values()),
-            schema=schema,
+            meta=meta,
             directory=directory / name,
             id=ShuffleId(name),
             run_id=next(AbstractShuffleTestPool._shuffle_run_id_iterator),
@@ -1246,7 +1248,7 @@ async def test_basic_lowlevel_shuffle(
             npartitions, part, workers
         )
     assert len(set(worker_for_mapping.values())) == min(n_workers, npartitions)
-    schema = pa.Schema.from_pandas(dfs[0])
+    meta = dfs[0].head(0)
 
     with DataFrameShuffleTestPool() as local_shuffle_pool:
         shuffles = []
@@ -1255,7 +1257,7 @@ async def test_basic_lowlevel_shuffle(
                 local_shuffle_pool.new_shuffle(
                     name=workers[ix],
                     worker_for_mapping=worker_for_mapping,
-                    schema=schema,
+                    meta=meta,
                     directory=tmp_path,
                     loop=loop_in_thread,
                 )
@@ -1321,7 +1323,7 @@ async def test_error_offload(tmp_path, loop_in_thread):
             npartitions, part, workers
         )
         partitions_for_worker[w].append(part)
-    schema = pa.Schema.from_pandas(dfs[0])
+    meta = dfs[0].head(0)
 
     class ErrorOffload(DataFrameShuffleRun):
         async def offload(self, func, *args):
@@ -1331,7 +1333,7 @@ async def test_error_offload(tmp_path, loop_in_thread):
         sA = local_shuffle_pool.new_shuffle(
             name="A",
             worker_for_mapping=worker_for_mapping,
-            schema=schema,
+            meta=meta,
             directory=tmp_path,
             loop=loop_in_thread,
             Shuffle=ErrorOffload,
@@ -1339,7 +1341,7 @@ async def test_error_offload(tmp_path, loop_in_thread):
         sB = local_shuffle_pool.new_shuffle(
             name="B",
             worker_for_mapping=worker_for_mapping,
-            schema=schema,
+            meta=meta,
             directory=tmp_path,
             loop=loop_in_thread,
         )
@@ -1375,7 +1377,7 @@ async def test_error_send(tmp_path, loop_in_thread):
             npartitions, part, workers
         )
         partitions_for_worker[w].append(part)
-    schema = pa.Schema.from_pandas(dfs[0])
+    meta = dfs[0].head(0)
 
     class ErrorSend(DataFrameShuffleRun):
         async def send(self, *args: Any, **kwargs: Any) -> None:
@@ -1385,7 +1387,7 @@ async def test_error_send(tmp_path, loop_in_thread):
         sA = local_shuffle_pool.new_shuffle(
             name="A",
             worker_for_mapping=worker_for_mapping,
-            schema=schema,
+            meta=meta,
             directory=tmp_path,
             loop=loop_in_thread,
             Shuffle=ErrorSend,
@@ -1393,7 +1395,7 @@ async def test_error_send(tmp_path, loop_in_thread):
         sB = local_shuffle_pool.new_shuffle(
             name="B",
             worker_for_mapping=worker_for_mapping,
-            schema=schema,
+            meta=meta,
             directory=tmp_path,
             loop=loop_in_thread,
         )
@@ -1428,7 +1430,7 @@ async def test_error_receive(tmp_path, loop_in_thread):
             npartitions, part, workers
         )
         partitions_for_worker[w].append(part)
-    schema = pa.Schema.from_pandas(dfs[0])
+    meta = dfs[0].head(0)
 
     class ErrorReceive(DataFrameShuffleRun):
         async def receive(self, data: list[tuple[int, bytes]]) -> None:
@@ -1438,7 +1440,7 @@ async def test_error_receive(tmp_path, loop_in_thread):
         sA = local_shuffle_pool.new_shuffle(
             name="A",
             worker_for_mapping=worker_for_mapping,
-            schema=schema,
+            meta=meta,
             directory=tmp_path,
             loop=loop_in_thread,
             Shuffle=ErrorReceive,
@@ -1446,7 +1448,7 @@ async def test_error_receive(tmp_path, loop_in_thread):
         sB = local_shuffle_pool.new_shuffle(
             name="B",
             worker_for_mapping=worker_for_mapping,
-            schema=schema,
+            meta=meta,
             directory=tmp_path,
             loop=loop_in_thread,
         )
diff --git a/distributed/tests/test_as_completed.py b/distributed/tests/test_as_completed.py
index 45c37fbb..27b5b18c 100644
--- a/distributed/tests/test_as_completed.py
+++ b/distributed/tests/test_as_completed.py
@@ -9,10 +9,11 @@ from time import sleep
 
 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)