Repository: dask/distributed

Issue:
as_completed using with_results option hangs getting next future if future errored
I ran into an issue where failed tasks were causing my code to hang waiting for `as_completed` to return the next result.  Here's an example of the scenario: 
```
from dask.distributed import Client
from dask.distributed import as_completed


client = Client()

def error_out():
    raise RuntimeError()

future = client.submit(error_out)
iterator = as_completed([future], with_results=True)
iterator.next()
```


Discussion:
Indeed.  Thank you for the excellent bug report and minimal reproducible example @rwedge !

This should probably be fixed here:

https://github.com/dask/distributed/blob/3715fd128eaa5dc6d027789470a5966b8d880253/distributed/client.py#L3465-L3484

and in particular the line 

```python
        if self.with_results:
            result = yield future._result()
```

This errs in the coroutine, which probably isn't right.  We should probably try-except around this and do *something* with the exception.  

What we should do is probably up for discussion.  We might return the exception as the result, or we might choose to raise the exception when the user iterates over that part of the `as_completed` iterator.

Figuring out what to do here might be a bit tricky, but the eventual solution should be approachable to a new contributor if anyone wants to pick this up.
We could also add an `errors='raise'/'return'` keyword to `as_completed`.  
@rwedge is this something that you would be interested in contributing?
I can take a shot at it once it's decided how to handle the exception
I think that the `errors=` keyword is likely a sensible choice.  We can leave this for a day to see if others come by with suggestions or concerns.  Also, do you have thoughts on the matter?
The `errors=` keyword seems reasonable.  This keyword determines how errors encountered while gathering results are handled, right?  I think the documentation should be clear about which errors the keyword is referring to.
Instead of a try-except, would it be sufficient to use raiseit=False and then raise the error in ``__next__`` or ``next_batch``?  
```
        if self.with_results:
                result = yield future._result(raiseit=False)
```
This could handle an error raised if there was an future with an 'error' status but an unexpected exception might still cause a hang.
Sounds reasonable to me.

One tricky bit here is that we would want to do this entirely within a coroutine, so that we knew that nothing changed between gathering the exception and checking the status afterwards like `if future.status == 'error': `.  If we do this in normal code, like in `next_batch` then these things might change in a subtle race condition (for example the worker holding the future might fail) or someone else might restart the task.

Resolving patch:
diff --git a/distributed/client.py b/distributed/client.py
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -3551,6 +3551,18 @@ class as_completed(object):
     Additionally, you can also add more futures to this object during
     computation with the ``.add`` method
 
+    Parameters
+    ----------
+    futures: Collection of futures
+        A list of Future objects to be iterated over in the order in which they
+        complete
+    with_results: bool (False)
+        Whether to wait and include results of futures as well;
+        in this case `as_completed` yields a tuple of (future, result)
+    raise_errors: bool (True)
+        Whether we should raise when the result of a future raises an exception;
+        only affects behavior when `with_results=True`.
+
     Examples
     --------
     >>> x, y, z = client.map(inc, [1, 2, 3])  # doctest: +SKIP
@@ -3586,7 +3598,7 @@ class as_completed(object):
     3
     """
 
-    def __init__(self, futures=None, loop=None, with_results=False):
+    def __init__(self, futures=None, loop=None, with_results=False, raise_errors=True):
         if futures is None:
             futures = []
         self.futures = defaultdict(lambda: 0)
@@ -3596,6 +3608,7 @@ def __init__(self, futures=None, loop=None, with_results=False):
         self.condition = Condition()
         self.thread_condition = threading.Condition()
         self.with_results = with_results
+        self.raise_errors = raise_errors
 
         if futures:
             self.update(futures)
@@ -3610,12 +3623,9 @@ def track_future(self, future):
         try:
             yield _wait(future)
         except CancelledError:
-            del self.futures[future]
-            if not self.futures:
-                self._notify()
-            return
+            pass
         if self.with_results:
-            result = yield future._result()
+            result = yield future._result(raiseit=False)
         with self.lock:
             self.futures[future] -= 1
             if not self.futures[future]:
@@ -3664,13 +3674,21 @@ def __iter__(self):
     def __aiter__(self):
         return self
 
+    def _get_and_raise(self):
+        res = self.queue.get()
+        if self.with_results:
+            future, result = res
+            if self.raise_errors and future.status == 'error':
+                six.reraise(*result)
+        return res
+
     def __next__(self):
         while self.queue.empty():
             if self.is_empty():
                 raise StopIteration()
             with self.thread_condition:
                 self.thread_condition.wait(timeout=0.100)
-        return self.queue.get()
+        return self._get_and_raise()
 
     @gen.coroutine
     def __anext__(self):
@@ -3681,7 +3699,7 @@ def __anext__(self):
                 raise StopAsyncIteration
             yield self.condition.wait()
 
-        raise gen.Return(self.queue.get())
+        raise gen.Return(self._get_and_raise())
 
     next = __next__
 


Test patch:
diff --git a/distributed/tests/py3_test_client.py b/distributed/tests/py3_test_client.py
--- a/distributed/tests/py3_test_client.py
+++ b/distributed/tests/py3_test_client.py
@@ -79,7 +79,7 @@ async def f():
 
     yield f()
 
-    assert L == [x]
+    assert L == [x, y]
 
 
 def test_async_with(loop):
diff --git a/distributed/tests/test_as_completed.py b/distributed/tests/test_as_completed.py
--- a/distributed/tests/test_as_completed.py
+++ b/distributed/tests/test_as_completed.py
@@ -1,4 +1,5 @@
 from collections import Iterator
+from concurrent.futures._base import CancelledError
 from operator import add
 import random
 from time import sleep
@@ -9,7 +10,7 @@
 from distributed import Client
 from distributed.client import _as_completed, as_completed, _first_completed
 from distributed.compatibility import Empty, StopAsyncIteration, Queue
-from distributed.utils_test import cluster, gen_cluster, inc
+from distributed.utils_test import cluster, gen_cluster, inc, throws
 from distributed.utils_test import loop  # noqa: F401
 
 
@@ -118,26 +119,30 @@ def test_as_completed_is_empty(loop):
 def test_as_completed_cancel(loop):
     with cluster() as (s, [a, b]):
         with Client(s['address'], loop=loop) as c:
-            x = c.submit(sleep, 1)
+            x = c.submit(inc, 1)
             y = c.submit(inc, 1)
 
             ac = as_completed([x, y])
             x.cancel()
 
-            assert next(ac) is y
+            assert next(ac) is x or y
+            assert next(ac) is y or x
 
             with pytest.raises(Empty):
                 ac.queue.get(timeout=0.1)
 
-            assert list(as_completed([x, y, x])) == [y]
+            res = list(as_completed([x, y, x]))
+            assert len(res) == 3
+            assert set(res) == {x, y}
+            assert res.count(x) == 2
 
 
 def test_as_completed_cancel_last(loop):
     with cluster() as (s, [a, b]):
         with Client(s['address'], loop=loop) as c:
-            w = c.submit(sleep, 0.3)
+            w = c.submit(inc, 0.3)
             x = c.submit(inc, 1)
-            y = c.submit(sleep, 0.3)
+            y = c.submit(inc, 0.3)
 
             @gen.coroutine
             def _():
@@ -148,9 +153,9 @@ def _():
             loop.add_callback(_)
 
             ac = as_completed([x, y])
-            result = list(ac)
+            result = set(ac)
 
-            assert result == [x]
+            assert result == {x, y}
 
 
 @gen_cluster(client=True)
@@ -168,3 +173,107 @@ def test_async_for_py2_equivalent(c, s, a, b):
 
     with pytest.raises(StopAsyncIteration):
         yield seq.__anext__()
+
+
+@gen_cluster(client=True)
+def test_as_completed_error_async(c, s, a, b):
+    x = c.submit(throws, 1)
+    y = c.submit(inc, 1)
+
+    ac = as_completed([x, y])
+    first = yield ac.__anext__()
+    second = yield ac.__anext__()
+    result = {first, second}
+
+    assert result == {x, y}
+    assert x.status == 'error'
+    assert y.status == 'finished'
+
+
+def test_as_completed_error(loop):
+    with cluster() as (s, [a, b]):
+        with Client(s['address'], loop=loop) as c:
+            x = c.submit(throws, 1)
+            y = c.submit(inc, 1)
+
+            ac = as_completed([x, y])
+            result = set(ac)
+
+            assert result == {x, y}
+            assert x.status == 'error'
+            assert y.status == 'finished'
+
+
+def test_as_completed_with_results(loop):
+    with cluster() as (s, [a, b]):
+        with Client(s['address'], loop=loop) as c:
+            x = c.submit(throws, 1)
+            y = c.submit(inc, 5)
+            z = c.submit(inc, 1)
+
+            ac = as_completed([x, y, z], with_results=True)
+            y.cancel()
+            with pytest.raises(RuntimeError) as exc:
+                res = list(ac)
+            assert str(exc.value) == 'hello!'
+
+
+@gen_cluster(client=True)
+def test_as_completed_with_results_async(c, s, a, b):
+    x = c.submit(throws, 1)
+    y = c.submit(inc, 5)
+    z = c.submit(inc, 1)
+
+    ac = as_completed([x, y, z], with_results=True)
+    y.cancel()
+    with pytest.raises(RuntimeError) as exc:
+        first = yield ac.__anext__()
+        second = yield ac.__anext__()
+        third = yield ac.__anext__()
+    assert str(exc.value) == 'hello!'
+
+
+def test_as_completed_with_results_no_raise(loop):
+    with cluster() as (s, [a, b]):
+        with Client(s['address'], loop=loop) as c:
+            x = c.submit(throws, 1)
+            y = c.submit(inc, 5)
+            z = c.submit(inc, 1)
+
+            ac = as_completed([x, y, z], with_results=True, raise_errors=False)
+            y.cancel()
+            res = list(ac)
+
+            dd = {r[0]: r[1:] for r in res}
+            assert set(dd.keys()) == {y, x, z}
+            assert x.status == 'error'
+            assert y.status == 'cancelled'
+            assert z.status == 'finished'
+
+            assert isinstance(dd[y][0], CancelledError)
+            assert isinstance(dd[x][0][1], RuntimeError)
+            assert dd[z][0] == 2
+
+
+@gen_cluster(client=True)
+def test_as_completed_with_results_no_raise_async(c, s, a, b):
+    x = c.submit(throws, 1)
+    y = c.submit(inc, 5)
+    z = c.submit(inc, 1)
+
+    ac = as_completed([x, y, z], with_results=True, raise_errors=False)
+    y.cancel()
+    first = yield ac.__anext__()
+    second = yield ac.__anext__()
+    third = yield ac.__anext__()
+    res = [first, second, third]
+
+    dd = {r[0]: r[1:] for r in res}
+    assert set(dd.keys()) == {y, x, z}
+    assert x.status == 'error'
+    assert y.status == 'cancelled'
+    assert z.status == 'finished'
+
+    assert isinstance(dd[y][0], CancelledError)
+    assert isinstance(dd[x][0][1], RuntimeError)
+    assert dd[z][0] == 2