Repository: PrefectHQ/prefect
Pull request: Adds `wait` utility for `PrefectFutures`
<!-- 
Thanks for opening a pull request to Prefect! We've got a few requests to help us review contributions:

- Make sure that your title neatly summarizes the proposed changes.
- Provide a short overview of the change and the value it adds.
- Share an example to help us understand the change in user experience.
- Confirm that you've done common tasks so we can give a timely review.
- Review our contribution guidelines: https://docs.prefect.io/latest/contributing/overview/

Happy engineering!
-->
Adds a `wait`utility function that mirrors that of `concurrent.futures.wait`. 
Closes #14240


### Checklist
<!-- These boxes may be checked after opening the pull request. -->

- [ x ] This pull request includes a label categorizing the change e.g. `maintenance`, `fix`, `feature`, `enhancement`, `docs`.
- [ x ] This pull request references any related issue by including "closes `<link to issue>`"
  <!-- If you do not have permission to add a label, a maintainer will add one for you -->
  - If no issue exists and your change is not a small fix, please [create an issue](https://github.com/PrefectHQ/prefect/issues/new/choose) first.
- [ x ] If this pull request adds new functionality, it includes unit tests that cover the changes
- [ x ] If this pull request removes docs files, it includes redirect settings in `mint.json`.
- [ x ] If this pull request adds functions or classes, it includes helpful docstrings.


Problem statement:
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 from issue discussion:

Commit message: waits for sequence of futures

Patch to review:
diff --git a/src/prefect/futures.py b/src/prefect/futures.py
index d0df04ea9a5a..269cb040d39c 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
@@ -297,6 +298,41 @@ def result(
             ) from exc
 
 
+DoneAndNotDoneFutures = collections.namedtuple("DoneAndNotDoneFutures", "done not_done")
+
+
+def wait(pfs: List[PrefectFuture], timeout=None) -> DoneAndNotDoneFutures:
+    """
+    Wait for the futures in the given sequence to complete.
+
+    Args:
+        fs: 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 *fs* are removed and will be
+        returned only once.
+    """
+    pfs = set(pfs)
+    done = {f for f in pfs if f._final_state}
+    not_done = pfs - done
+    if len(done) == len(pfs):
+        return DoneAndNotDoneFutures(done, not_done)
+    try:
+        with timeout_context(timeout):
+            for f in not_done:
+                f.wait()
+                done.add(f)
+                not_done.remove(f)
+            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]:
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()


Code review comments:
diff_hunk: @@ -297,6 +298,41 @@ def result(
             ) from exc
 
 
+DoneAndNotDoneFutures = collections.namedtuple("DoneAndNotDoneFutures", "done not_done")
+
+
+def wait(pfs: List[PrefectFuture], timeout=None) -> DoneAndNotDoneFutures:
path: src/prefect/futures.py
text: @user1:
nit: I was confused about the purpose of `pfs` on first read.
```suggestion
def wait(futures: List[PrefectFuture], timeout=None) -> DoneAndNotDoneFutures:
```