Repository: dask/distributed

Issue:
KeyboardInterrupt on `get` call doesn't cancel computations
Currently `get` doesn't stop computations when broken with a KeyboardInterrupt or when subsets of the  computations fail.



Discussion:
OK, this ends up being harder than expected.  The `_get` coroutine handles this all well in https://github.com/dask/distributed/pull/346 but the `KeyboardInterrupt` doesn't get passed down to the coroutine.  Instead it stops an `Event.wait` call.  I don't know how to dive into a particular running coroutine and send the `KeyboardInterrupt` there.


Resolving patch:
diff --git a/distributed/client.py b/distributed/client.py
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -380,17 +380,20 @@ def pack_data(o, d):
     {'a': [1], 'b': 'y'}
     """
     typ = type(o)
-    if typ is str or typ is bytes:
+    try:
         if o in d:
             return d[o]
-    else:
-        k = tokey(o)
-        if k in d:
-            return d[k]
+    except TypeError:
+        pass
 
     if typ in collection_types:
         return typ([pack_data(x, d) for x in o])
     elif typ is dict:
         return {k: pack_data(v, d) for k, v in o.items()}
     else:
+        try:
+            if o in d:
+                return d[o]
+        except TypeError:
+            pass
         return o
diff --git a/distributed/executor.py b/distributed/executor.py
--- a/distributed/executor.py
+++ b/distributed/executor.py
@@ -69,23 +69,23 @@ class Future(WrappedKey):
     def __init__(self, key, executor):
         self.key = key
         self.executor = executor
-        self.executor._inc_ref(key)
+        self.executor._inc_ref(tokey(key))
         self._generation = self.executor.generation
         self._cleared = False
 
         if key not in executor.futures:
-            executor.futures[key] = {'event': Event(), 'status': 'pending'}
+            executor.futures[tokey(key)] = {'event': Event(), 'status': 'pending'}
 
     @property
     def status(self):
         try:
-            return self.executor.futures[self.key]['status']
+            return self.executor.futures[tokey(self.key)]['status']
         except KeyError:
             return 'cancelled'
 
     @property
     def event(self):
-        return self.executor.futures[self.key]['event']
+        return self.executor.futures[tokey(self.key)]['event']
 
     def done(self):
         """ Is the computation complete? """
@@ -104,7 +104,7 @@ def result(self):
     @gen.coroutine
     def _result(self, raiseit=True):
         try:
-            d = self.executor.futures[self.key]
+            d = self.executor.futures[tokey(self.key)]
         except KeyError:
             exception = CancelledError(self.key)
             if raiseit:
@@ -146,13 +146,13 @@ def cancel(self):
 
     def cancelled(self):
         """ Returns True if the future has been cancelled """
-        return self.key not in self.executor.futures
+        return tokey(self.key) not in self.executor.futures
 
     @gen.coroutine
     def _traceback(self):
         yield self.event.wait()
         if self.status == 'error':
-            raise Return(self.executor.futures[self.key]['traceback'])
+            raise Return(self.executor.futures[tokey(self.key)]['traceback'])
         else:
             raise Return(None)
 
@@ -179,17 +179,17 @@ def traceback(self):
     @property
     def type(self):
         try:
-            return self.executor.futures[self.key]['type']
+            return self.executor.futures[tokey(self.key)]['type']
         except KeyError:
             return None
 
-    def _del(self):
+    def release(self):
         if not self._cleared and self.executor.generation == self._generation:
             self._cleared = True
-            self.executor._dec_ref(self.key)
+            self.executor._dec_ref(tokey(self.key))
 
     def __del__(self):
-        self._del()
+        self.release()
 
     def __str__(self):
         if self.type:
@@ -515,9 +515,7 @@ def submit(self, func, *args, **kwargs):
             else:
                 key = funcname(func) + '-' + str(uuid.uuid4())
 
-        key = tokey(key)
-
-        if key in self.futures:
+        if tokey(key) in self.futures:
             return Future(key, self)
 
         if allow_other_workers and workers is None:
@@ -532,30 +530,17 @@ def submit(self, func, *args, **kwargs):
             restrictions = {}
             loose_restrictions = []
 
-        args2, arg_dependencies = unpack_remotedata(args, byte_keys=True)
-        kwargs2, kwarg_dependencies = unpack_remotedata(kwargs, byte_keys=True)
-        dependencies = arg_dependencies | kwarg_dependencies
-
-        bad_deps = {d for d in dependencies if d not in self.futures}
-        if bad_deps:
-            raise CancelledError(*bad_deps)
+        if kwargs:
+            dsk = {key: (apply, func, list(args), kwargs)}
+        else:
+            dsk = {key: (func,) + tuple(args)}
 
-        task = {'function': dumps_function(func)}
-        if args2:
-            task['args'] = dumps(args2)
-        if kwargs2:
-            task['kwargs'] = dumps(kwargs2)
+        futures = self._graph_to_futures(dsk, [key], restrictions,
+                loose_restrictions)
 
         logger.debug("Submit %s(...), %s", funcname(func), key)
-        self._send_to_scheduler({'op': 'update-graph',
-                                 'tasks': {key: task},
-                                 'keys': [key],
-                                 'dependencies': {key: list(dependencies)},
-                                 'restrictions': valmap(list, restrictions),
-                                 'loose_restrictions': loose_restrictions,
-                                 'client': self.id})
 
-        return Future(key, self)
+        return futures[tokey(key)]
 
     def _threaded_map(self, q_out, func, qs_in, **kwargs):
         """ Internal function for mapping Queue """
@@ -637,8 +622,6 @@ def map(self, func, *iterables, **kwargs):
             keys = [funcname(func) + '-' + uid + '-' + str(i)
                     for i in range(min(map(len, iterables)))]
 
-        keys = [tokey(key) for key in keys]
-
         if not kwargs:
             dsk = {key: (func,) + args
                    for key, args in zip(keys, zip(*iterables))}
@@ -646,15 +629,6 @@ def map(self, func, *iterables, **kwargs):
             dsk = {key: (apply, func, (tuple, list(args)), kwargs)
                    for key, args in zip(keys, zip(*iterables))}
 
-        d = {key: unpack_remotedata(task, byte_keys=True) for key, task in dsk.items()}
-        dsk2 = str_graph({k: v[0] for k, v in d.items()})
-        d_futures = {k: v[1] for k, v in d.items()}
-        dependencies = {k: set(map(tokey, v[1])) for k, v in d.items()}
-
-        bad_deps = {k for v in d_futures.values() for k in v if k not in self.futures}
-        if bad_deps:
-            raise CancelledError(*bad_deps)
-
         if isinstance(workers, str):
             workers = [workers]
         if isinstance(workers, (list, set)):
@@ -676,16 +650,11 @@ def map(self, func, *iterables, **kwargs):
         else:
             loose_restrictions = set()
 
+        futures = self._graph_to_futures(dsk, keys, restrictions,
+                loose_restrictions)
         logger.debug("map(%s, ...)", funcname(func))
-        self._send_to_scheduler({'op': 'update-graph',
-                                 'tasks': valmap(dumps_task, dsk2),
-                                 'dependencies': valmap(list, dependencies),
-                                 'keys': keys,
-                                 'restrictions': valmap(list, restrictions),
-                                 'loose_restrictions': list(loose_restrictions),
-                                 'client': self.id})
 
-        return [Future(key, self) for key in keys]
+        return [futures[tokey(key)] for key in keys]
 
     @gen.coroutine
     def _gather(self, futures, errors='raise'):
@@ -987,25 +956,30 @@ def run(self, function, *args, **kwargs):
         """
         return sync(self.loop, self._run, function, *args, **kwargs)
 
-    @gen.coroutine
-    def _get(self, dsk, keys, restrictions=None, raise_on_error=True):
-        keyset = set(flatten([keys]))
-        flatkeys = list(map(tokey, flatten([keys])))
-        futures = {key: Future(key, self) for key in flatkeys}
+    def _graph_to_futures(self, dsk, keys, restrictions=None,
+                                loose_restrictions=None,
+                                allow_other_workers=True):
+        keyset = set(keys)
+        flatkeys = list(map(tokey, keys))
+        futures = {key: Future(key, self) for key in keyset}
 
         values = {k for k, v in dsk.items() if isinstance(v, Future)
                                             and k not in keyset}
         if values:
             dsk = dask.optimize.inline(dsk, keys=values)
 
-        d = {k: unpack_remotedata(v, byte_keys=True) for k, v in dsk.items()}
-        dsk2 = str_graph({k: v[0] for k, v in d.items()})
+        d = {k: unpack_remotedata(v) for k, v in dsk.items()}
+        extra_keys = set.union(*[v[1] for v in d.values()]) if d else set()
+        dsk2 = str_graph({k: v[0] for k, v in d.items()}, extra_keys)
         dsk3 = {k: v for k, v in dsk2.items() if k is not v}
 
         if restrictions:
             restrictions = keymap(tokey, restrictions)
             restrictions = valmap(list, restrictions)
 
+        if loose_restrictions is not None:
+            loose_restrictions = list(map(tokey, loose_restrictions))
+
         dependencies = {tokey(k): set(map(tokey, v[1])) for k, v in d.items()}
 
         for s in dependencies.values():
@@ -1021,8 +995,15 @@ def _get(self, dsk, keys, restrictions=None, raise_on_error=True):
                                  'dependencies': valmap(list, dependencies),
                                  'keys': list(flatkeys),
                                  'restrictions': restrictions or {},
+                                 'loose_restrictions': loose_restrictions,
                                  'client': self.id})
 
+        return futures
+
+    @gen.coroutine
+    def _get(self, dsk, keys, restrictions=None, raise_on_error=True):
+        futures = self._graph_to_futures(dsk, set(flatten([keys])), restrictions)
+
         packed = pack_data(keys, futures)
         try:
             result = yield self._gather(packed)
@@ -1034,12 +1015,12 @@ def _get(self, dsk, keys, restrictions=None, raise_on_error=True):
                 raise gen.Return(result)
         finally:
             for f in futures.values():
-                f._del()
+                f.release()
         if not raise_on_error:
             result = 'OK', result
         raise gen.Return(result)
 
-    def get(self, dsk, keys, **kwargs):
+    def get(self, dsk, keys, restrictions=None, loose_restrictions=None):
         """ Compute dask graph
 
         Parameters
@@ -1061,13 +1042,18 @@ def get(self, dsk, keys, **kwargs):
         --------
         Executor.compute: Compute asynchronous collections
         """
-        status, result = sync(self.loop, self._get, dsk, keys,
-                              raise_on_error=False, **kwargs)
+        futures = self._graph_to_futures(dsk, set(flatten([keys])),
+                restrictions, loose_restrictions)
 
-        if status == 'error':
-            raise result
-        else:
-            return result
+        try:
+            results = self.gather(futures)
+        except (KeyboardInterrupt, Exception) as e:
+            for f in futures.values():
+                f.release()
+            raise
+
+        results2 = pack_data(keys, results)
+        return results2
 
     def compute(self, args, sync=False):
         """ Compute dask collections on cluster
@@ -1117,27 +1103,16 @@ def compute(self, args, sync=False):
         dsk = merge([opt(merge([v.dask for v in val]),
                          [v._keys() for v in val])
                     for opt, val in groups.items()])
-        names = [tokey('finalize-%s' % tokenize(v)) for v in variables]
+        names = ['finalize-%s' % tokenize(v) for v in variables]
         dsk2 = {name: (v._finalize, v._keys()) for name, v in zip(names, variables)}
 
-        d = {k: unpack_remotedata(v, byte_keys=True) for k, v in merge(dsk, dsk2).items()}
-        dsk3 = str_graph({k: v[0] for k, v in d.items()})
-        dependencies = {tokey(k): set(map(tokey, v[1])) for k, v in d.items()}
-
-        for k, v in dsk3.items():
-            dependencies[k] |= set(_deps(dsk3, v))
-
-        self._send_to_scheduler({'op': 'update-graph',
-                                 'tasks': valmap(dumps_task, dsk3),
-                                 'dependencies': valmap(list, dependencies),
-                                 'keys': names,
-                                 'client': self.id})
+        futures_dict = self._graph_to_futures(merge(dsk2, dsk), names)
 
         i = 0
         futures = []
         for arg in args:
             if isinstance(arg, Base):
-                futures.append(Future(names[i], self))
+                futures.append(futures_dict[names[i]])
                 i += 1
             else:
                 futures.append(arg)
@@ -1190,22 +1165,12 @@ def persist(self, collections):
                          [v._keys() for v in val])
                     for opt, val in groups.items()])
 
-        d = {k: unpack_remotedata(v, byte_keys=True) for k, v in dsk.items()}
-        dsk2 = str_graph({k: v[0] for k, v in d.items()})
-        dependencies = {tokey(k): set(map(tokey, v[1])) for k, v in d.items()}
-
-        for k, v in dsk2.items():
-            dependencies[k] |= set(_deps(dsk2, v))
+        names = {k for c in collections for k in flatten(c._keys())}
 
-        names = list({tokey(k) for c in collections for k in flatten(c._keys())})
+        futures = self._graph_to_futures(dsk, names)
 
-        self._send_to_scheduler({'op': 'update-graph',
-                                 'tasks': valmap(dumps_task, dsk2),
-                                 'dependencies': valmap(list, dependencies),
-                                 'keys': names,
-                                 'client': self.id})
 
-        result = [redict_collection(c, {k: Future(tokey(k), self)
+        result = [redict_collection(c, {k: futures[k]
                                         for k in flatten(c._keys())})
                 for c in collections]
         if singleton:
diff --git a/distributed/utils.py b/distributed/utils.py
--- a/distributed/utils.py
+++ b/distributed/utils.py
@@ -319,7 +319,7 @@ def _maybe_complex(task):
             type(task) is dict and any(map(_maybe_complex, task.values())))
 
 
-def str_graph(dsk):
+def str_graph(dsk, extra_values=()):
     def convert(task):
         if type(task) is list:
             return [convert(v) for v in task]
@@ -328,7 +328,7 @@ def convert(task):
         if istask(task):
             return (task[0],) + tuple(map(convert, task[1:]))
         try:
-            if task in dsk:
+            if task in dsk or task in extra_values:
                 return tokey(task)
         except TypeError:
             pass


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
@@ -84,11 +84,6 @@ def test_pack_data():
     assert pack_data({'a': [b'x'], 'b': 'y'}, data) == {'a': [1], 'b': 'y'}
 
 
-def test_pack_data_with_key_mapping():
-    data = {tokey(('x', 1)): 1}
-    assert pack_data((('x', 1), 'y'), data) == (1, 'y')
-
-
 def test_gather_errors_voluminously(current_loop):
     with cluster_center() as (c, [a, b]):
         try:
diff --git a/distributed/tests/test_executor.py b/distributed/tests/test_executor.py
--- a/distributed/tests/test_executor.py
+++ b/distributed/tests/test_executor.py
@@ -15,6 +15,7 @@
 from toolz import (identity, isdistinct, first, concat, pluck, valmap,
         partition_all)
 from tornado import gen
+from tornado.ioloop import IOLoop
 
 from dask import delayed
 from dask.context import _globals
@@ -2177,7 +2178,6 @@ def test__persist(e, s, a, b):
     assert yy._keys() == y._keys()
 
     g, h = e.compute([y, yy])
-    yield gen.sleep(1)
     gg, hh = yield e._gather([g, h])
     assert (gg == hh).all()
 
@@ -3131,3 +3131,30 @@ def test_forgotten_futures_dont_clean_up_new_futures(e, s, a, b):
     import gc; gc.collect()
     yield gen.sleep(0.1)
     yield y._result()
+
+
+def test_get_stops_work_after_error(loop):
+    loop2 = IOLoop()
+    s = Scheduler(loop=loop2)
+    s.start(0)
+    w = Worker(s.ip, s.port, loop=loop2)
+    w.start(0)
+
+    t = Thread(target=loop2.start)
+    t.daemon = True
+    t.start()
+
+    with Executor(s.address, loop=loop) as e:
+        with pytest.raises(Exception):
+            e.get({'x': (throws, 1), 'y': (sleep, 1)}, ['x', 'y'])
+
+        start = time()
+        while len(s.tasks):
+            sleep(0.1)
+            assert time() < start + 5
+
+    loop2.add_callback(loop2.stop)
+    while loop2._running:
+        sleep(0.01)
+    loop2.close(all_fds=True)
+    t.join()