Adds as_completed utility for PrefectFuture

src/prefect/futures.py

@@ -292,6 +292,59 @@ def result(
             ) from exc
 
 
+def as_completed(futures: List[PrefectFuture], timeout=None) -> Iterator[PrefectFuture]:
+    """
+    An iterator over the given futures that yields each as it completes.
+
+    Args:
+        fs: The sequence of Futures (possibly created by different Executors) to
+            iterate over.
+        timeout: The maximum number of seconds to wait. If None, then there
+            is no limit on the wait time.
+
+    Returns:
+        An iterator that yields the given Futures as they complete (finished or
+        cancelled). If any given Futures are duplicated, they will be returned
+        once.
+
+    Raises:
+        TimeoutError: If the entire result iterator could not be generated
+            before the given timeout.
+
+    Examples:
+        ```python
+        @task
+        def sleep_task(seconds):
+            sleep(seconds)
+            return 42
+
+        @flow
+        def flow():
+            futures = random_task.map(range(10))
+            for future in as_completed(futures):
+                print(future.result())
+        ```
+    """
+    futures = set(futures)
+    total_futures = len(futures)
+    try:
+        with timeout_context(timeout):
+            done = {f for f in futures if f._final_state}
+            pending = futures - done
+            for future in done:
+                yield future
+
+            for future in pending.copy():
+                future.wait()
+                pending.remove(future)
+                yield future
+
+    except TimeoutError:
+        raise TimeoutError(
+            "%d (of %d) futures unfinished" % (len(pending), total_futures)
+        )
+
+
 DoneAndNotDoneFutures = collections.namedtuple("DoneAndNotDoneFutures", "done not_done")
 
 

tests/test_futures.py

@@ -15,6 +15,7 @@
     PrefectFuture,
     PrefectFutureList,
     PrefectWrappedFuture,
+    as_completed,
     resolve_futures_to_states,
     wait,
 )
@@ -55,6 +56,25 @@ def test_wait_with_timeout(self):
         futures = wait(mock_futures, timeout=0.01)
         assert futures.not_done == {mock_futures[-1]}
 
+    def test_as_completed(self):
+        mock_futures = [MockFuture(data=i) for i in range(5)]
+        for future in as_completed(mock_futures):
+            assert future.state.is_completed()
+
+    @pytest.mark.timeout(method="thread")
+    def test_as_completed_with_timeout(self):
+        mock_futures = [MockFuture(data=i) for i in range(5)]
+        hanging_future = Future()
+        mock_futures.append(PrefectConcurrentFuture(uuid.uuid4(), hanging_future))
+
+        with pytest.raises(TimeoutError) as exc_info:
+            for future in as_completed(mock_futures, timeout=0.01):
+                assert future.state.is_completed()
+
+        assert (
+            exc_info.value.args[0] == f"1 (of {len(mock_futures)}) futures unfinished"
+        )
+
 
 class TestPrefectConcurrentFuture:
     def test_wait_with_timeout(self):