Repository: dask/distributed

Issue:
future.result() freezes in add_done_callback()
Hi,

below is a simplified example of how calling `future.result()` in `add_done_callback` freezes.

```
c = Client()

def sleep():
    time.sleep(2)
    return 1

def cb(f):
    print(f)
    print('Result should be {}'.format(sleep()))
    print('Result was', f.result())

d = delayed(sleep)()
d.compute() # 1

c.compute(d).result() # 1

f2 = c.compute(d)
f2.add_done_callback(cb)
# <Future: status: finished, type: int, key: finalize-e24717a1ea49f34025d9f69a8d4b9e3e>
# Result should be 1
```

You never get the `f.result()` in `cb`. Execution freezes.


Discussion:
`add_done_callback` schedules the callback to run in the event loop, so you cannot use `f.result()`, which will block synchronously. Instead you must the coroutine API `f._result()`:

```
@gen.coroutine
def cb(f):
    print(f)
    res = yield f._result()
    print('Result was', res)

# [...]
f2.add_done_callback(cb)
```

Ok, thank you!

By the way, would `del f` be safe in the callback `cb`?

`del f` would probably be safe, but would also not achieve anything particular.

http://stackoverflow.com/questions/40477518/how-to-get-the-result-of-a-future-in-a-callback/40477862#40477862

So how would one be able to mark that the scheduler can free the data pointed by the future in the callback?

@mrocklin I see this was already asked :)

http://distributed.readthedocs.io/en/latest/memory.html

On Thu, Nov 10, 2016 at 9:01 AM, Jarno Lintusaari notifications@github.com
wrote:

> So how would one be able to mark that the scheduler can free the data
> pointed by the future in the callback?
> 
> @mrocklin https://github.com/mrocklin I see it was already asked :)
> 
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> https://github.com/dask/distributed/issues/647#issuecomment-259696801,
> or mute the thread
> https://github.com/notifications/unsubscribe-auth/AASszDU6xhYjm4eDpdzWqSJ_t9zVHpdBks5q8yOkgaJpZM4Kug9N
> .

Yes, `del f` is what I wish to do in the callback, I was just confused by:

> pitrou commented 3 hours ago
> del f would probably be safe, but would also not achieve anything particular.

So I wasn't sure if `del f` would work in the callback since e.g. `f.result()` would not.


Resolving patch:
diff --git a/distributed/client.py b/distributed/client.py
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -1,6 +1,7 @@
 from __future__ import print_function, division, absolute_import
 
 from collections import defaultdict, Iterator, Iterable
+from concurrent.futures import ThreadPoolExecutor
 from concurrent.futures._base import DoneAndNotDoneFutures, CancelledError
 from contextlib import contextmanager
 import copy
@@ -72,17 +73,21 @@ class Future(WrappedKey):
     --------
     Client:  Creates futures
     """
+    _cb_executor = None
+    _cb_executor_pid = None
+
     def __init__(self, key, client):
         self.key = key
+        tkey = tokey(key)
         self.client = client
-        self.client._inc_ref(tokey(key))
+        self.client._inc_ref(tkey)
         self._generation = self.client.generation
         self._cleared = False
 
-        tkey = tokey(key)
-
-        if tkey not in client.futures:
-            client.futures[tkey] = {'event': Event(), 'status': 'pending'}
+        if tkey in client.futures:
+            self._state = client.futures[tkey]
+        else:
+            self._state = client.futures[tkey] = FutureState(Event())
 
     @property
     def executor(self):
@@ -90,14 +95,11 @@ def executor(self):
 
     @property
     def status(self):
-        try:
-            return self.client.futures[tokey(self.key)]['status']
-        except KeyError:
-            return 'cancelled'
+        return self._state.status
 
     @property
     def event(self):
-        return self.client.futures[tokey(self.key)]['event']
+        return self._state.event
 
     def done(self):
         """ Is the computation complete? """
@@ -108,28 +110,27 @@ def result(self):
         result = sync(self.client.loop, self._result, raiseit=False)
         if self.status == 'error':
             six.reraise(*result)
-        if self.status == 'cancelled':
+        elif self.status == 'cancelled':
             raise result
         else:
             return result
 
     @gen.coroutine
     def _result(self, raiseit=True):
-        try:
-            d = self.client.futures[tokey(self.key)]
-        except KeyError:
+        yield self._state.event.wait()
+        if self.status == 'error':
+            exc = clean_exception(self._state.exception,
+                                  self._state.traceback)
+            if raiseit:
+                six.reraise(*exc)
+            else:
+                raise Return(exc)
+        elif self.status == 'cancelled':
             exception = CancelledError(self.key)
             if raiseit:
                 raise exception
             else:
                 raise gen.Return(exception)
-
-        yield d['event'].wait()
-        if self.status == 'error':
-            if raiseit:
-                six.reraise(*clean_exception(**d))
-            else:
-                raise Return(clean_exception(**d))
         else:
             result = yield self.client._gather([self])
             raise gen.Return(result[0])
@@ -138,8 +139,7 @@ def _result(self, raiseit=True):
     def _exception(self):
         yield self.event.wait()
         if self.status == 'error':
-            exception = self.client.futures[self.key]['exception']
-            raise Return(exception)
+            raise Return(self._state.exception)
         else:
             raise Return(None)
 
@@ -158,8 +158,22 @@ def add_done_callback(self, fn):
         The callback ``fn`` should take the future as its only argument.  This
         will be called regardless of if the future completes successfully,
         errs, or is cancelled
+
+        The callback is executed in a separate thread.
         """
-        self.client.loop.add_callback(done_callback, self, fn)
+        cls = Future
+        if cls._cb_executor is None or cls._cb_executor_pid != os.getpid():
+            cls._cb_executor = ThreadPoolExecutor(1)
+            cls._cb_executor_pid = os.getpid()
+
+        def execute_callback(fut):
+            try:
+                fn(fut)
+            except BaseException:
+                logger.exception("Error in callback %s of %s:", fn, fut)
+
+        self.client.loop.add_callback(done_callback, self,
+                                      partial(cls._cb_executor.submit, execute_callback))
 
     def cancel(self):
         """ Returns True if the future has been cancelled """
@@ -167,13 +181,13 @@ def cancel(self):
 
     def cancelled(self):
         """ Returns True if the future has been cancelled """
-        return tokey(self.key) not in self.client.futures
+        return self._state.status == 'cancelled'
 
     @gen.coroutine
     def _traceback(self):
         yield self.event.wait()
         if self.status == 'error':
-            raise Return(self.client.futures[tokey(self.key)]['traceback'])
+            raise Return(self._state.traceback)
         else:
             raise Return(None)
 
@@ -199,10 +213,7 @@ def traceback(self):
 
     @property
     def type(self):
-        try:
-            return self.client.futures[tokey(self.key)]['type']
-        except KeyError:
-            return None
+        return self._state.type
 
     def release(self):
         if not self._cleared and self.client.generation == self._generation:
@@ -235,9 +246,42 @@ def __str__(self):
     __repr__ = __str__
 
 
+class FutureState(object):
+    """A Future's internal state.
+
+    This is shared between all Futures with the same key and client.
+    """
+    __slots__ = ('event', 'status', 'type', 'exception', 'traceback')
+
+    def __init__(self, event):
+        self.event = event
+        self.status = 'pending'
+        self.type = None
+
+    def cancel(self):
+        self.status = 'cancelled'
+        self.event.set()
+
+    def finish(self, type=None):
+        self.status = 'finished'
+        self.event.set()
+        if type is not None:
+            self.type = type
+
+    def lose(self):
+        self.status = 'lost'
+        self.event.clear()
+
+    def set_error(self, exception, traceback):
+        self.status = 'error'
+        self.exception = exception
+        self.traceback = traceback
+        self.event.set()
+
+
 @gen.coroutine
 def done_callback(future, callback):
-    """ Wait on future, then call callback """
+    """ Coroutine that waits on future, then calls callback """
     while future.status == 'pending':
         yield future.event.wait()
     callback(future)
@@ -248,6 +292,11 @@ def normalize_future(f):
     return [f.key, type(f)]
 
 
+class AllExit(Exception):
+    """Custom exception class to exit All(...) early.
+    """
+
+
 class Client(object):
     """ Drive computations on a distributed cluster
 
@@ -314,11 +363,13 @@ def __init__(self, address=None, start=True, loop=None, timeout=3,
     def __str__(self):
         if hasattr(self, '_loop_thread'):
             n = sync(self.loop, self.scheduler.ncores)
-            return '<Client: scheduler="%s:%d" processes=%d cores=%d>' % (
+            return '<%s: scheduler="%s:%d" processes=%d cores=%d>' % (
+                    self.__class__.__name__,
                     self.scheduler.ip, self.scheduler.port, len(n),
                     sum(n.values()))
         else:
-            return '<Client: scheduler="%s:%d">' % (
+            return '<%s: scheduler="%s:%d">' % (
+                    self.__class__.__name__,
                     self.scheduler.ip, self.scheduler.port)
 
     __repr__ = __str__
@@ -376,10 +427,9 @@ def reconnect(self, timeout=0.1):
             self.status = 'connecting'
             self.scheduler_stream = None
 
-            events = [d['event'] for d in self.futures.values()]
+            for st in self.futures.values():
+                st.cancel()
             self.futures.clear()
-            for e in events:
-                e.set()
 
             while self.status == 'connecting':
                 try:
@@ -443,9 +493,9 @@ def _dec_ref(self, key):
     def _release_key(self, key):
         """ Release key from distributed memory """
         logger.debug("Release key %s", key)
-        if key in self.futures:
-            self.futures[key]['event'].clear()
-            del self.futures[key]
+        st = self.futures.pop(key, None)
+        if st is not None:
+            st.cancel()
         if self.status != 'closed':
             self._send_to_scheduler({'op': 'client-releases-keys',
                                      'keys': [key],
@@ -480,37 +530,38 @@ def _handle_report(self):
                         breakout = True
                         break
                     elif msg['op'] == 'key-in-memory':
-                        if msg['key'] in self.futures:
-                            self.futures[msg['key']]['status'] = 'finished'
-                            self.futures[msg['key']]['event'].set()
-                            if (msg.get('type') and
-                                not self.futures[msg['key']].get('type')):
-                                self.futures[msg['key']]['type'] = loads(msg['type'])
+                        st = self.futures.get(msg['key'])
+                        if st is not None:
+                            typ = (loads(msg['type'])
+                                   if (msg.get('type') and not st.type)
+                                   else None)
+                            st.finish(typ)
                     elif msg['op'] == 'lost-data':
-                        if msg['key'] in self.futures:
-                            self.futures[msg['key']]['status'] = 'lost'
-                            self.futures[msg['key']]['event'].clear()
+                        st = self.futures.get(msg['key'])
+                        if st is not None:
+                            st.lose()
                     elif msg['op'] == 'cancelled-key':
-                        if msg['key'] in self.futures:
-                            self.futures[msg['key']]['event'].set()
-                            del self.futures[msg['key']]
+                        st = self.futures.pop(msg['key'], None)
+                        if st is not None:
+                            st.cancel()
                     elif msg['op'] == 'task-erred':
-                        if msg['key'] in self.futures:
-                            self.futures[msg['key']]['status'] = 'error'
+                        st = self.futures.get(msg['key'])
+                        if st is not None:
                             try:
-                                self.futures[msg['key']]['exception'] = loads(msg['exception'])
+                                exception = loads(msg['exception'])
                             except TypeError:
-                                self.futures[msg['key']]['exception'] = \
-                                    Exception('Undeserializable exception', msg['exception'])
-                            self.futures[msg['key']]['traceback'] = (loads(msg['traceback'])
-                                                                     if msg['traceback'] else None)
-                            self.futures[msg['key']]['event'].set()
+                                exception = Exception(
+                                    'Undeserializable exception',
+                                    msg['exception'])
+                            st.set_error(exception,
+                                         loads(msg['traceback'])
+                                         if msg['traceback']
+                                         else None)
                     elif msg['op'] == 'restart':
                         logger.info("Receive restart signal from scheduler")
-                        events = [d['event'] for d in self.futures.values()]
+                        for st in self.futures.values():
+                            st.cancel()
                         self.futures.clear()
-                        for e in events:
-                            e.set()
                         with ignoring(AttributeError):
                             self._restart_event.set()
                     elif 'error' in msg['op']:
@@ -778,28 +829,29 @@ def _gather(self, futures, errors='raise'):
         @gen.coroutine
         def wait(k):
             """ Want to stop the All(...) early if we find an error """
-            yield self.futures[k]['event'].wait()
-            if self.futures[k]['status'] != 'finished':
-                raise Exception()
+            st = self.futures[k]
+            yield st.event.wait()
+            if st.status != 'finished':
+                raise AllExit()
 
         while True:
             logger.debug("Waiting on futures to clear before gather")
 
-            with ignoring(Exception):
+            with ignoring(AllExit):
                 yield All([wait(key) for key in keys if key in self.futures])
 
             exceptions = set()
             bad_keys = set()
             for key in keys:
                 if (key not in self.futures or
-                    self.futures[key]['status'] == 'error'):
+                    self.futures[key].status == 'error'):
                     exceptions.add(key)
                     if errors == 'raise':
                         try:
-                            d = self.futures[key]
-                            six.reraise(type(d['exception']),
-                                        d['exception'],
-                                        d['traceback'])
+                            st = self.futures[key]
+                            six.reraise(type(st.exception),
+                                        st.exception,
+                                        st.traceback)
                         except KeyError:
                             six.reraise(CancelledError,
                                         CancelledError(key),
@@ -818,7 +870,7 @@ def wait(k):
                 self._send_to_scheduler({'op': 'missing-data',
                                          'keys': response['keys']})
                 for key in response['keys']:
-                    self.futures[key]['event'].clear()
+                    self.futures[key].event.clear()
             else:
                 break
 
@@ -916,15 +968,14 @@ def _scatter(self, data, workers=None, broadcast=False):
                     "Input to scatter must be a list, iterator, or queue")
 
         for key in keys:
-            self.futures[key]['status'] = 'finished'
-            self.futures[key]['event'].set()
+            self.futures[key].finish(type=None)
 
         if isinstance(types, list):
             for key, typ in zip(keys, types):
-                self.futures[key]['type'] = typ
+                self.futures[key].type = typ
         elif isinstance(types, dict):
             for key in keys:
-                self.futures[key]['type'] = types[key]
+                self.futures[key].type = types[key]
 
         raise gen.Return(out)
 
@@ -1027,8 +1078,9 @@ def _cancel(self, futures):
         keys = {tokey(f.key) for f in futures_of(futures)}
         yield self.scheduler.cancel(keys=list(keys), client=self.id)
         for k in keys:
-            with ignoring(KeyError):
-                del self.futures[k]
+            st = self.futures.pop(k, None)
+            if st is not None:
+                st.cancel()
 
     def cancel(self, futures):
         """
@@ -2110,12 +2162,12 @@ def _wait(fs, timeout=None, return_when='ALL_COMPLETED'):
     if timeout is not None:
         raise NotImplementedError("Timeouts not yet supported")
     if return_when == 'ALL_COMPLETED':
-        try:
-            yield All({f.event.wait() for f in fs})
-        except KeyError:
-            raise CancelledError([f.key for f in fs
-                                        if f.key not in f.client.futures])
+        yield All({f.event.wait() for f in fs})
         done, not_done = set(fs), set()
+        cancelled = [f.key for f in done
+                     if f.status == 'cancelled']
+        if cancelled:
+            raise CancelledError(cancelled)
     else:
         raise NotImplementedError("Only return_when='ALL_COMPLETED' supported")
 
diff --git a/distributed/worker_client.py b/distributed/worker_client.py
--- a/distributed/worker_client.py
+++ b/distributed/worker_client.py
@@ -9,10 +9,10 @@
 from dask.base import tokenize
 from tornado import gen
 
-from .client import Client, Future, pack_data, unpack_remotedata
+from .client import AllExit, Client, Future, pack_data, unpack_remotedata
 from .sizeof import sizeof
 from .threadpoolexecutor import secede
-from .utils import log_errors, sync, tokey, ignoring
+from .utils import All, log_errors, sync, tokey, ignoring
 from .worker import thread_state
 
 
@@ -122,8 +122,7 @@ def _scatter(self, data, workers=None, broadcast=False):
                         "Input to scatter must be a list or dict")
 
             for key in keys:
-                self.futures[key]['status'] = 'finished'
-                self.futures[key]['event'].set()
+                self.futures[key].finish(type=None)
 
             raise gen.Return(out)
 
@@ -143,11 +142,11 @@ def _gather(self, futures, errors='raise'):
         @gen.coroutine
         def wait(k):
             """ Want to stop the All(...) early if we find an error """
-            yield self.futures[k]['event'].wait()
-            if self.futures[k]['status'] != 'finished':
-                raise Exception()
+            yield self.futures[k].event.wait()
+            if self.futures[k].status != 'finished':
+                raise AllExit()
 
-        with ignoring(Exception):
+        with ignoring(AllExit):
             yield All([wait(key) for key in keys if key in self.futures])
 
         while True:  # not threadsafe, so try until success


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
@@ -549,7 +549,7 @@ def test_tokenize_on_futures(c, s, a, b):
     assert tokenize(x) == tokenize(x)
     assert tokenize(x) == tokenize(y)
 
-    c.futures[x.key]['status'] = 'finished'
+    c.futures[x.key].finish()
 
     assert tok == tokenize(y)
 
@@ -2039,8 +2039,8 @@ def test_long_traceback(c, s, a, b):
     try:
         x = c.submit(deep, 1000)
         yield _wait([x])
-        assert len(dumps(c.futures[x.key]['traceback'])) < 10000
-        assert isinstance(c.futures[x.key]['exception'], RuntimeError)
+        assert len(dumps(c.futures[x.key].traceback)) < 10000
+        assert isinstance(c.futures[x.key].exception, RuntimeError)
     finally:
         sys.setrecursionlimit(n)
 
@@ -2915,7 +2915,7 @@ def test_shutdown_idempotent(loop):
 @gen_cluster(client=True)
 def test_get_returns_early(c, s, a, b):
     start = time()
-    with ignoring(Exception):
+    with ignoring(RuntimeError):
         result = yield c._get({'x': (throws, 1), 'y': (sleep, 1)}, ['x', 'y'])
     assert time() < start + 0.5
     assert not c.futures
@@ -2928,7 +2928,7 @@ def test_get_returns_early(c, s, a, b):
     x = c.submit(inc, 1)
     yield x._result()
 
-    with ignoring(Exception):
+    with ignoring(RuntimeError):
         result = yield c._get({'x': (throws, 1),
                                x.key: (inc, 1)}, ['x', x.key])
     assert x.key in s.tasks
@@ -3694,18 +3694,31 @@ def f(x, y=0):
 
 @gen_cluster(client=True)
 def test_add_done_callback(c, s, a, b):
-    x = c.submit(inc, 1)
+    S = set()
 
-    L = []
     def f(future):
-        L.append(future.key)
-        L.append(future.status)
+        future.add_done_callback(g)
+
+    def g(future):
+        S.add((future.key, future.status))
+
+    u = c.submit(inc, 1, key='u')
+    v = c.submit(throws, "hello", key='v')
+    w = c.submit(slowinc, 2, delay=0.3, key='w')
+    x = c.submit(inc, 3, key='x')
+    u.add_done_callback(f)
+    v.add_done_callback(f)
+    w.add_done_callback(f)
+
+    yield _wait((u, v, w, x))
 
     x.add_done_callback(f)
 
-    yield _wait(x)
+    t = time()
+    while len(S) < 4 and time() - t < 2.0:
+        yield gen.sleep(0.01)
 
-    assert L == [x.key, x.status]
+    assert S == {(f.key, f.status) for f in (u, v, w, x)}
 
 
 @gen_cluster(client=True)
diff --git a/distributed/tests/test_worker.py b/distributed/tests/test_worker.py
--- a/distributed/tests/test_worker.py
+++ b/distributed/tests/test_worker.py
@@ -492,7 +492,7 @@ def test_active_holds_tasks(e, s, w):
     assert future.key not in w.active
 
     future = e.submit(throws, 1)
-    with ignoring(Exception):
+    with ignoring(RuntimeError):
         yield _wait([future])
     assert not w.active
 
diff --git a/distributed/utils_test.py b/distributed/utils_test.py
--- a/distributed/utils_test.py
+++ b/distributed/utils_test.py
@@ -119,7 +119,7 @@ def deep(n):
 
 
 def throws(x):
-    raise Exception('hello!')
+    raise RuntimeError('hello!')
 
 
 def double(x):
@@ -255,7 +255,7 @@ def cluster(nworkers=2, nanny=False, worker_kwargs={}):
                 loop.run_sync(lambda: disconnect('127.0.0.1', port),
                               timeout=0.5)
         for proc in [w['proc'] for w in workers]:
-            with ignoring(Exception):
+            with ignoring(EnvironmentError):
                 proc.terminate()
                 proc.join(timeout=2)
         for q in [w['queue'] for w in workers]: