Repository: dask/distributed

Issue:
Add support for wait(FIRST_COMPLETED)
Concurrent futures [wait](https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.wait) supports FIRST_COMPLETED, FIRST_EXCEPTION, and ALL_COMPLETED. It appears Dask only supports ALL_COMPLETED at the moment. There are many applications that are much easier to write with the wait(FIRST_COMPLETED) primitive as opposed to other primitives provided.


Discussion:
Agreed.  This would be a good thing for someone to implement if they have the time.  One approach might be to use `as_completed` to help solve this problem.

Resolving patch:
diff --git a/distributed/client.py b/distributed/client.py
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -63,7 +63,7 @@
 from .utils import (All, sync, funcname, ignoring, queue_to_iterator,
                     tokey, log_errors, str_graph, key_split, format_bytes, asciitable,
                     thread_state, no_default, PeriodicCallback, LoopRunner,
-                    parse_timedelta, shutting_down)
+                    parse_timedelta, shutting_down, Any)
 from .versions import get_versions
 
 
@@ -3637,35 +3637,43 @@ def CompatibleExecutor(*args, **kwargs):
     raise Exception("This has been moved to the Client.get_executor() method")
 
 
+ALL_COMPLETED = 'ALL_COMPLETED'
+FIRST_COMPLETED = 'FIRST_COMPLETED'
+
+
 @gen.coroutine
-def _wait(fs, timeout=None, return_when='ALL_COMPLETED'):
+def _wait(fs, timeout=None, return_when=ALL_COMPLETED):
     if timeout is not None and not isinstance(timeout, Number):
         raise TypeError("timeout= keyword received a non-numeric value.\n"
                         "Beware that wait expects a list of values\n"
                         "  Bad:  wait(x, y, z)\n"
                         "  Good: wait([x, y, z])")
     fs = futures_of(fs)
-    if return_when == 'ALL_COMPLETED':
-        future = All({f._state.wait() for f in fs})
-        if timeout is not None:
-            future = gen.with_timeout(timedelta(seconds=timeout), future)
-        yield future
-        done, not_done = set(fs), set()
-        cancelled = [f.key for f in done
-                     if f.status == 'cancelled']
-        if cancelled:
-            raise CancelledError(cancelled)
+    if return_when == ALL_COMPLETED:
+        wait_for = All
+    elif return_when == FIRST_COMPLETED:
+        wait_for = Any
     else:
-        raise NotImplementedError("Only return_when='ALL_COMPLETED' supported")
+        raise NotImplementedError("Only return_when='ALL_COMPLETED' and 'FIRST_COMPLETED' are "
+                                  "supported")
 
-    raise gen.Return(DoneAndNotDoneFutures(done, not_done))
+    future = wait_for({f._state.wait() for f in fs})
+    if timeout is not None:
+        future = gen.with_timeout(timedelta(seconds=timeout), future)
+    yield future
 
+    done, not_done = ({fu for fu in fs if fu.status != 'pending'},
+                      {fu for fu in fs if fu.status == 'pending'})
+    cancelled = [f.key for f in done
+                 if f.status == 'cancelled']
+    if cancelled:
+        raise CancelledError(cancelled)
 
-ALL_COMPLETED = 'ALL_COMPLETED'
+    raise gen.Return(DoneAndNotDoneFutures(done, not_done))
 
 
-def wait(fs, timeout=None, return_when='ALL_COMPLETED'):
-    """ Wait until all futures are complete
+def wait(fs, timeout=None, return_when=ALL_COMPLETED):
+    """ Wait until all/any futures are finished
 
     Parameters
     ----------
diff --git a/distributed/utils.py b/distributed/utils.py
--- a/distributed/utils.py
+++ b/distributed/utils.py
@@ -233,6 +233,45 @@ def quiet():
     raise gen.Return(results)
 
 
+@gen.coroutine
+def Any(args, quiet_exceptions=()):
+    """ Wait on many tasks at the same time and return when any is finished
+
+    Err once any of the tasks err.
+
+    Parameters
+    ----------
+    args: futures to wait for
+    quiet_exceptions: tuple, Exception
+        Exception types to avoid logging if they fail
+    """
+    tasks = gen.WaitIterator(*args)
+    results = [None for _ in args]
+    while not tasks.done():
+        try:
+            result = yield tasks.next()
+        except Exception:
+            @gen.coroutine
+            def quiet():
+                """ Watch unfinished tasks
+
+                Otherwise if they err they get logged in a way that is hard to
+                control.  They need some other task to watch them so that they
+                are not orphaned
+                """
+                for task in list(tasks._unfinished):
+                    try:
+                        yield task
+                    except quiet_exceptions:
+                        pass
+            quiet()
+            raise
+
+        results[tasks.current_index] = result
+        break
+    raise gen.Return(results)
+
+
 def sync(loop, func, *args, **kwargs):
     """
     Run coroutine in loop running in separate thread.


Test patch:
diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
--- a/distributed/tests/test_client.py
+++ b/distributed/tests/test_client.py
@@ -609,6 +609,21 @@ def test_wait(c, s, a, b):
     assert x.status == y.status == 'finished'
 
 
+@gen_cluster(client=True)
+def test_wait_first_completed(c, s, a, b):
+    x = c.submit(slowinc, 1)
+    y = c.submit(slowinc, 1)
+    z = c.submit(inc, 2)
+
+    done, not_done = yield wait([x, y, z], return_when='FIRST_COMPLETED')
+
+    assert done == {z}
+    assert not_done == {x, y}
+    assert z.status == 'finished'
+    assert x.status == 'pending'
+    assert y.status == 'pending'
+
+
 @gen_cluster(client=True, timeout=2)
 def test_wait_timeout(c, s, a, b):
     future = c.submit(sleep, 0.3)