Repository: dask/distributed

Issue:
Cancel futures
There should be a way for the user to cancel a future or list of futures.  

For API the `concurrent.futures` API supports this for a single future, we should probably also have an `Executor.cancel(list_of_futures)` function for a list of futures.  This would cancel enqueued tasks, tasks that are not yet ready, and would also cancel all dependent tasks.  These tasks would likely be removed from memory and soon their metadata will also be removed from the scheduler.

For implementation this probably looks a lot like existing code for garbage collection.  It's unclear what to do when multiple clients want and don't want the same data.  My guess is that we still compute things while any client wants the result.



Discussion:

Resolving patch:
diff --git a/distributed/executor.py b/distributed/executor.py
--- a/distributed/executor.py
+++ b/distributed/executor.py
@@ -133,6 +133,10 @@ def exception(self):
         """
         return sync(self.executor.loop, self._exception)
 
+    def cancel(self):
+        """ Returns True if the future has been cancelled """
+        return self.executor.cancel([self])
+
     def cancelled(self):
         """ Returns True if the future has been cancelled """
         return self.key not in self.executor.futures
@@ -343,6 +347,8 @@ def _handle_report(self, start_event):
             except StreamClosedError:
                 break
 
+            logger.debug("Executor receives message %s", msg)
+
             if msg['op'] == 'stream-start':
                 start_event.set()
             if msg['op'] == 'close':
@@ -355,6 +361,10 @@ def _handle_report(self, start_event):
                 if msg['key'] in self.futures:
                     self.futures[msg['key']]['status'] = 'lost'
                     self.futures[msg['key']]['event'].clear()
+            if msg['op'] == 'cancelled-key':
+                if msg['key'] in self.futures:
+                    self.futures[msg['key']]['event'].set()
+                    del self.futures[msg['key']]
             if msg['op'] == 'task-erred':
                 if msg['key'] in self.futures:
                     self.futures[msg['key']]['status'] = 'error'
@@ -730,6 +740,27 @@ def _():
         else:
             return sync(self.loop, self._scatter, data, workers=workers,
                         broadcast=broadcast)
+    @gen.coroutine
+    def _cancel(self, futures):
+        keys = {f.key for f in flatten(futures)}
+        yield self.scheduler.cancel(keys=keys, client=self.id)
+        for k in keys:
+            with ignoring(KeyError):
+                del self.futures[k]
+
+    def cancel(self, futures):
+        """
+        Cancel running futures
+
+        This stops future tasks from being scheduled if they have not yet run
+        and deletes them if they have already run.  After calling, this result
+        and all dependent results will no longer be accessible
+
+        Parameters
+        ----------
+        futures: list of Futures
+        """
+        return sync(self.loop, self._cancel, futures)
 
     @gen.coroutine
     def _get(self, dsk, keys, restrictions=None, raise_on_error=True):
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -191,6 +191,7 @@ def __init__(self, center=None, loop=None,
                          'register': self.add_worker,
                          'unregister': self.remove_worker,
                          'gather': self.gather,
+                         'cancel': self.cancel,
                          'feed': self.feed,
                          'terminate': self.close,
                          'broadcast': self.broadcast,
@@ -385,6 +386,8 @@ def ensure_occupied(self, worker):
         while (self.stacks[worker] and
                self.ncores[worker] > len(self.processing[worker])):
             key = self.stacks[worker].pop()
+            if key not in self.dask:
+                continue
             self.processing[worker].add(key)
             logger.debug("Send job to worker: %s, %s, %s", worker, key, self.dask[key])
             self.worker_queues[worker].put_nowait(
@@ -416,7 +419,7 @@ def seed_ready_tasks(self, keys=None):
             if stack:
                 self.ensure_occupied(worker)
 
-    def update_data(self, who_has=None, nbytes=None):
+    def update_data(self, who_has=None, nbytes=None, client=None):
         """
         Learn that new data has entered the network from an external source
 
@@ -430,8 +433,9 @@ def update_data(self, who_has=None, nbytes=None):
 
         self.nbytes.update(nbytes)
 
-        for k in who_has:
-            self.who_wants[k] = set()
+        if client:
+            self.client_wants_keys(keys=list(who_has), client=client)
+
         self.in_play.update(who_has)
 
     def mark_task_erred(self, key, worker, exception, traceback):
@@ -635,6 +639,7 @@ def client_releases_keys(self, keys=None, client=None):
             if not self.who_wants[k]:
                 del self.who_wants[k]
                 self.release_held_data([k])
+                logger.debug("Delete who_wants[%s]", k)
 
     def client_wants_keys(self, keys=None, client=None):
         for k in keys:
@@ -686,6 +691,25 @@ def forget(self, key):
         if key in self.who_has:
             self.delete_data(keys=[key])
 
+    def cancel_key(self, key, client, retries=5):
+        if key not in self.who_wants:  # no key yet, lets try again in 500ms
+            if retries:
+                self.loop.add_future(gen.sleep(0.2),
+                        lambda _: self.cancel_key(key, client, retries - 1))
+            return
+        if self.who_wants[key] == {client}:  # no one else wants this key
+            for dep in list(self.dependents[key]):
+                self.cancel_key(dep, client)
+        logger.debug("Scheduler cancels key %s", key)
+        self.report({'op': 'cancelled-key', 'key': key})
+        self.client_releases_keys(keys=[key], client=client)
+
+    def cancel(self, stream, keys=None, client=None):
+        """ Stop execution on a list of keys """
+        logger.info("Client %s requests to cancel keys %s", client, keys)
+        for key in keys:
+            self.cancel_key(key, client)
+
     def heal_state(self):
         """ Recover from catastrophic change """
         logger.debug("Heal state")
@@ -717,7 +741,8 @@ def report(self, msg):
         for q in self.report_queues:
             q.put_nowait(msg)
         if 'key' in msg:
-            streams = {self.streams.get(c, ()) for c in self.who_wants[msg['key']]}
+            streams = {self.streams.get(c, ())
+                       for c in self.who_wants.get(msg['key'], ())}
         else:
             streams = self.streams.values()
         for s in streams:
@@ -1009,7 +1034,8 @@ def validate(self, allow_overlap=False, allow_bad_stacks=False):
         released = set(self.dask) - self.in_play
         validate_state(self.dependencies, self.dependents, self.waiting,
                 self.waiting_data, self.who_has, self.stacks,
-                self.processing, None, released, self.in_play,
+                self.processing, None, released, self.in_play, self.who_wants,
+                self.wants_what,
                 allow_overlap=allow_overlap, allow_bad_stacks=allow_bad_stacks)
         if not (set(self.ncores) == \
                 set(self.has_what) == \
@@ -1164,6 +1190,10 @@ def update_state(dsk, dependencies, dependents, who_wants, wants_what,
         if key not in dependents:
             dependents[key] = set()
 
+    for k in new_keys:
+        who_wants[k].add(client)
+        wants_what[client].add(k)
+
     exterior = keys_outside_frontier(dsk, dependencies, new_keys, in_play)
     in_play |= exterior
     for key in exterior:
@@ -1177,10 +1207,6 @@ def update_state(dsk, dependencies, dependents, who_wants, wants_what,
         if key not in waiting_data:
             waiting_data[key] = set()
 
-    for k in new_keys:
-        who_wants[k].add(client)
-        wants_what[client].add(k)
-
     return {'dsk': dsk,
             'dependencies': dependencies,
             'dependents': dependents,
@@ -1192,7 +1218,8 @@ def update_state(dsk, dependencies, dependents, who_wants, wants_what,
 
 def validate_state(dependencies, dependents, waiting, waiting_data,
         who_has, stacks, processing, finished_results, released, in_play,
-        allow_overlap=False, allow_bad_stacks=False, **kwargs):
+        who_wants, wants_what, allow_overlap=False, allow_bad_stacks=False,
+        **kwargs):
     """
     Validate a current runtime state
 
@@ -1243,6 +1270,11 @@ def check_key(key):
             if key in keys and who_has.get(key):
                 assert key in finished_results
 
+        for key, s in who_wants.items():
+            assert s, "empty who_wants"
+            for client in s:
+                assert key in wants_what[client]
+
         return True
 
     assert all(map(check_key, keys))
@@ -1376,7 +1408,7 @@ def unrunnable(key):
              'waiting': waiting, 'waiting_data': waiting_data,
              'who_has': who_has, 'processing': processing, 'stacks': stacks,
              'finished_results': finished_results, 'released': new_released,
-             'in_play': in_play}
+             'in_play': in_play, 'who_wants': {}, 'wants_what': {}}
     validate_state(**output)
     return output
 


Test patch:
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
@@ -1995,3 +1995,79 @@ def test_broadcast(loop):
             has_what = sync(e.loop, e.scheduler.has_what)
             assert has_what == {('127.0.0.1', a['port']): {x.key, y.key, z.key},
                                 ('127.0.0.1', b['port']): {x.key, y.key}}
+
+
+@gen_cluster()
+def test__cancel(s, a, b):
+    e = Executor((s.ip, s.port), start=False)
+    yield e._start()
+
+    x = e.submit(slowinc, 1)
+    y = e.submit(slowinc, x)
+
+    while y.key not in s.dask:
+        yield gen.sleep(0.01)
+
+    yield e._cancel([x])
+
+    assert x.cancelled()
+    s.validate()
+
+    start = time()
+    while not y.cancelled():
+        yield gen.sleep(0.01)
+        assert time() < start + 5
+
+    assert not s.dask
+    s.validate()
+
+    yield e._shutdown()
+
+
+@gen_cluster()
+def test__cancel_multi_client(s, a, b):
+    e = Executor((s.ip, s.port), start=False)
+    yield e._start()
+    f = Executor((s.ip, s.port), start=False)
+    yield f._start()
+
+    x = e.submit(slowinc, 1)
+    y = f.submit(slowinc, 1)
+
+    assert x.key == y.key
+
+    yield e._cancel([x])
+
+    assert x.cancelled()
+    assert not y.cancelled()
+
+    assert y.key in s.dask
+
+    out = yield y._result()
+    assert out == 2
+
+    with pytest.raises(CancelledError):
+        yield x._result()
+
+    yield e._shutdown()
+    yield f._shutdown()
+
+
+def test_cancel(loop):
+    with cluster() as (s, [a, b]):
+        with Executor(('127.0.0.1', s['port']), loop=loop) as e:
+            x = e.submit(slowinc, 1)
+            y = e.submit(slowinc, x)
+            z = e.submit(slowinc, y)
+
+            e.cancel([y])
+
+            start = time()
+            while not z.cancelled():
+                sleep(0.01)
+                assert time() < start + 5
+
+            assert x.result() == 2
+
+            z.cancel()
+            assert z.cancelled()
diff --git a/distributed/tests/test_scheduler.py b/distributed/tests/test_scheduler.py
--- a/distributed/tests/test_scheduler.py
+++ b/distributed/tests/test_scheduler.py
@@ -352,6 +352,8 @@ def test_validate_state():
     finished_results = set()
     released = set()
     in_play = {'x', 'y'}
+    who_wants = {'y': {'client'}}
+    wants_what = {'client': {'y'}}
 
     validate_state(**locals())
 
@@ -619,7 +621,7 @@ def test_add_worker(s, a, b):
     yield w._start(0)
 
     dsk = {('x', i): (inc, i) for i in range(10)}
-    s.update_graph(dsk=dsk, keys=list(dsk))
+    s.update_graph(dsk=dsk, keys=list(dsk), client='client')
 
     s.add_worker(address=w.address, keys=list(w.data),
                  ncores=w.ncores, services=s.services)
@@ -747,9 +749,9 @@ def test_delete_callback(s, a, b):
 @gen_cluster()
 def test_self_aliases(s, a, b):
     a.data['a'] = 1
-    s.update_data(who_has={'a': {a.address}}, nbytes={'a': 10})
+    s.update_data(who_has={'a': {a.address}}, nbytes={'a': 10}, client='client')
     s.update_graph(dsk={'a': 'a', 'b': (inc, 'a')},
-                   keys=['b'])
+                   keys=['b'], client='client')
 
     assert 'a' not in s.dask