Repository: PrefectHQ/prefect
Issue category: Feature Request

Issue:
Add `wait` utility for `PrefectFuture`s
Add a `wait` utility that mirrors `concurrent.futures.wait`. Should accept an iterable of futures to wait for completion and a timeout. It should return a tuple of futures that have completed and those that have not yet completed.


Hints:

Patch:
diff --git a/src/prefect/futures.py b/src/prefect/futures.py
index d0df04ea9a5a..7379aa8c8b06 100644
--- a/src/prefect/futures.py
+++ b/src/prefect/futures.py
@@ -1,4 +1,5 @@
 import abc
+import collections
 import concurrent.futures
 import inspect
 import uuid
@@ -256,13 +257,7 @@ def wait(self, timeout: Optional[float] = None) -> None:
             timeout: The maximum number of seconds to wait for all futures to
                 complete. This method will not raise if the timeout is reached.
         """
-        try:
-            with timeout_context(timeout):
-                for future in self:
-                    future.wait()
-        except TimeoutError:
-            logger.debug("Timed out waiting for all futures to complete.")
-            return
+        wait(self, timeout=timeout)
 
     def result(
         self,
@@ -297,6 +292,57 @@ def result(
             ) from exc
 
 
+DoneAndNotDoneFutures = collections.namedtuple("DoneAndNotDoneFutures", "done not_done")
+
+
+def wait(futures: List[PrefectFuture], timeout=None) -> DoneAndNotDoneFutures:
+    """
+    Wait for the futures in the given sequence to complete.
+
+    Args:
+        futures: The sequence of Futures to wait upon.
+        timeout: The maximum number of seconds to wait. If None, then there
+            is no limit on the wait time.
+
+    Returns:
+        A named 2-tuple of sets. The first set, named 'done', contains the
+        futures that completed (is finished or cancelled) before the wait
+        completed. The second set, named 'not_done', contains uncompleted
+        futures. Duplicate futures given to *futures* are removed and will be
+        returned only once.
+
+    Examples:
+        ```python
+        @task
+        def sleep_task(seconds):
+            sleep(seconds)
+            return 42
+
+        @flow
+        def flow():
+            futures = random_task.map(range(10))
+            done, not_done = wait(futures, timeout=5)
+            print(f"Done: {len(done)}")
+            print(f"Not Done: {len(not_done)}")
+        ```
+    """
+    futures = set(futures)
+    done = {f for f in futures if f._final_state}
+    not_done = futures - done
+    if len(done) == len(futures):
+        return DoneAndNotDoneFutures(done, not_done)
+    try:
+        with timeout_context(timeout):
+            for future in not_done.copy():
+                future.wait()
+                done.add(future)
+                not_done.remove(future)
+            return DoneAndNotDoneFutures(done, not_done)
+    except TimeoutError:
+        logger.debug("Timed out waiting for all futures to complete.")
+        return DoneAndNotDoneFutures(done, not_done)
+
+
 def resolve_futures_to_states(
     expr: Union[PrefectFuture, Any],
 ) -> Union[State, Any]:


Test patch:
diff --git a/tests/test_futures.py b/tests/test_futures.py
index e009e0fa1716..43fcd2dc2460 100644
--- a/tests/test_futures.py
+++ b/tests/test_futures.py
@@ -16,6 +16,7 @@
     PrefectFutureList,
     PrefectWrappedFuture,
     resolve_futures_to_states,
+    wait,
 )
 from prefect.states import Completed, Failed
 from prefect.task_engine import run_task_async, run_task_sync
@@ -37,6 +38,24 @@ def result(
         return self._final_state.result()
 
 
+class TestUtilityFunctions:
+    def test_wait(self):
+        mock_futures = [MockFuture(data=i) for i in range(5)]
+        futures = wait(mock_futures)
+        assert futures.not_done == set()
+
+        for future in mock_futures:
+            assert future.state.is_completed()
+
+    @pytest.mark.timeout(method="thread")
+    def test_wait_with_timeout(self):
+        mock_futures = [MockFuture(data=i) for i in range(5)]
+        hanging_future = Future()
+        mock_futures.append(PrefectConcurrentFuture(uuid.uuid4(), hanging_future))
+        futures = wait(mock_futures, timeout=0.01)
+        assert futures.not_done == {mock_futures[-1]}
+
+
 class TestPrefectConcurrentFuture:
     def test_wait_with_timeout(self):
         wrapped_future = Future()