Allow timeout strings in distributed.wait (#7081)

dask · Sep 29, 2022 · 482941e · 482941e
1 parent 8f36aa5
commit 482941e
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/distributed/client.py b/distributed/client.py
@@ -4741,16 +4741,18 @@ def wait(fs, timeout=None, return_when=ALL_COMPLETED):
     Parameters
     ----------
     fs : List[Future]
-    timeout : number, optional
-        Time in seconds after which to raise a
-        ``dask.distributed.TimeoutError``
+    timeout : number, string, optional
+        Time after which to raise a ``dask.distributed.TimeoutError``.
+        Can be a string like ``"10 minutes"`` or a number of seconds to wait.
     return_when : str, optional
         One of `ALL_COMPLETED` or `FIRST_COMPLETED`
 
     Returns
     -------
     Named tuple of completed, not completed
     """
+    if timeout is not None and isinstance(timeout, (Number, str)):
+        timeout = parse_timedelta(timeout, default="s")
     client = default_client()
     result = client.sync(_wait, fs, timeout=timeout, return_when=return_when)
     return result

diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
@@ -749,6 +749,11 @@ async def test_wait_timeout(c, s, a, b):
     with pytest.raises(TimeoutError):
         await wait(future, timeout=0.01)
 
+    # Ensure timeout can be a string
+    future = c.submit(sleep, 0.3)
+    with pytest.raises(TimeoutError):
+        await wait(future, timeout="0.01 s")
+
 
 def test_wait_sync(c):
     x = c.submit(inc, 1)