Repository: dask/dask

Issue:
Dask doesn't play well with interrupts
This is on `Python 3.5.2`, `dask 0.10.1`, and `jupyter-notebook 4.2.1`, all installed via conda on a 64-bit Ubuntu machine.

If you run the following piece of code:

``` Python
import json

from dask.diagnostics import ProgressBar
import dask.bag as db

j = json.dumps({"a": 1, "b": 1})

for i in range(8):
    data = [j for _ in range(10 ** i)]
    bag = db.from_sequence(data, npartitions=4).map(json.loads)
    with ProgressBar():
        db.zip(bag.pluck("a"), bag.pluck('b')).count().compute()
```

And try to do a keyboard interrupt in the middle, sometimes dask will flip out and just refuse to exit (no matter how many `CTRL-C`s and `CTRL-D`s you press. See https://asciinema.org/a/2yqxbhn1patwenwy024m2316l for a video.

This is especially annoying on a Jupyter notebook. Sometimes, dask will mysteriously keep printing out the progress bar from the last session, even after restarting the kernel (!).

_EDIT_: When this happens, dask will also leave around `python` processes still running. See https://asciinema.org/a/8x1liiyuiw7910mugj9tscher for an example (particularly the extra `python` process at the end).



Discussion:
Notably, I can only reproduce this in the python repl (and ipython). When I run it as a script (`python test.py`), 3 `KeyboardInterrupt`s are enough to exit cleanly without leaving phantom python processes anywhere (see https://asciinema.org/a/eitz5acubx6ot7a9cmm1sdz2l)

...contrary to my previous assumption, this doesn't have anything to do with the progressbar. I just thought it did because the phantom progressbar on a jupyter notebook was the most obvious indicator something was wrong.

Everything above can be replicated without the `with ProgressBar():` line.

Sorry for the noise.

I noticed that behavior when using `dask` in a jupyter notebook: the progressbar running in background after a stop and also orphan processes even after killing all kernels.

FWIW, I used `ProgressBar().register()` at the beginning of the notebook.

I've so far been unable to reproduce this without a progressbar.

I have been able to reproduce this with the synchronous single-threaded scheduler


Resolving patch:
diff --git a/dask/async.py b/dask/async.py
--- a/dask/async.py
+++ b/dask/async.py
@@ -405,7 +405,7 @@ def get_async(apply_async, num_workers, dsk, result, cache=None,
         Whether to rerun failing tasks in local process to enable debugging
         (False by default)
     callbacks : tuple or list of tuples, optional
-        Callbacks are passed in as tuples of length 4. Multiple sets of
+        Callbacks are passed in as tuples of length 5. Multiple sets of
         callbacks may be passed in as a list of tuples. For more information,
         see the dask.diagnostics documentation.
 
@@ -418,7 +418,7 @@ def get_async(apply_async, num_workers, dsk, result, cache=None,
 
     if callbacks is None:
         callbacks = _globals['callbacks']
-    start_cbs, start_state_cbs, pretask_cbs, posttask_cbs, finish_cbs = unpack_callbacks(callbacks)
+    _, _, pretask_cbs, posttask_cbs, _ = unpack_callbacks(callbacks)
 
     if isinstance(result, list):
         result_flat = set(flatten(result))
@@ -427,74 +427,81 @@ def get_async(apply_async, num_workers, dsk, result, cache=None,
     results = set(result_flat)
 
     dsk = dsk.copy()
-    for f in start_cbs:
-        f(dsk)
-
-    dsk, dependencies = cull(dsk, list(results))
+    started_cbs = []
+    try:
+        for cb in callbacks:
+            if cb[0]:
+                cb[0](dsk)
+            started_cbs.append(cb)
 
-    keyorder = order(dsk)
+        dsk, dependencies = cull(dsk, list(results))
 
-    state = start_state_from_dask(dsk, cache=cache, sortkey=keyorder.get)
+        keyorder = order(dsk)
 
-    for f in start_state_cbs:
-        f(dsk, state)
+        state = start_state_from_dask(dsk, cache=cache, sortkey=keyorder.get)
 
-    if rerun_exceptions_locally is None:
-        rerun_exceptions_locally = _globals.get('rerun_exceptions_locally', False)
+        for _, start_state, _, _, _ in callbacks:
+            if start_state:
+                start_state(dsk, state)
 
-    if state['waiting'] and not state['ready']:
-        raise ValueError("Found no accessible jobs in dask")
+        if rerun_exceptions_locally is None:
+            rerun_exceptions_locally = _globals.get('rerun_exceptions_locally', False)
 
-    def fire_task():
-        """ Fire off a task to the thread pool """
-        # Choose a good task to compute
-        key = state['ready'].pop()
-        state['running'].add(key)
-        for f in pretask_cbs:
-            f(key, dsk, state)
+        if state['waiting'] and not state['ready']:
+            raise ValueError("Found no accessible jobs in dask")
 
-        # Prep data to send
-        data = dict((dep, state['cache'][dep])
-                    for dep in get_dependencies(dsk, key))
-        # Submit
-        apply_async(execute_task, args=[key, dsk[key], data, queue,
-                                        get_id, raise_on_exception])
+        def fire_task():
+            """ Fire off a task to the thread pool """
+            # Choose a good task to compute
+            key = state['ready'].pop()
+            state['running'].add(key)
+            for f in pretask_cbs:
+                f(key, dsk, state)
 
-    # Seed initial tasks into the thread pool
-    while state['ready'] and len(state['running']) < num_workers:
-        fire_task()
+            # Prep data to send
+            data = dict((dep, state['cache'][dep])
+                        for dep in get_dependencies(dsk, key))
+            # Submit
+            apply_async(execute_task, args=[key, dsk[key], data, queue,
+                                            get_id, raise_on_exception])
 
-    # Main loop, wait on tasks to finish, insert new ones
-    while state['waiting'] or state['ready'] or state['running']:
-        try:
-            key, res, tb, worker_id = queue.get()
-        except KeyboardInterrupt:
-            for f in finish_cbs:
-                f(dsk, state, True)
-            raise
-        if isinstance(res, Exception):
-            for f in finish_cbs:
-                f(dsk, state, True)
-            if rerun_exceptions_locally:
-                data = dict((dep, state['cache'][dep])
-                            for dep in get_dependencies(dsk, key))
-                task = dsk[key]
-                _execute_task(task, data)  # Re-execute locally
-            else:
-                raise(remote_exception(res, tb))
-        state['cache'][key] = res
-        finish_task(dsk, key, state, results, keyorder.get)
-        for f in posttask_cbs:
-            f(key, res, dsk, state, worker_id)
+        # Seed initial tasks into the thread pool
         while state['ready'] and len(state['running']) < num_workers:
             fire_task()
 
+        # Main loop, wait on tasks to finish, insert new ones
+        while state['waiting'] or state['ready'] or state['running']:
+            key, res, tb, worker_id = queue.get()
+            if isinstance(res, Exception):
+                for _, _, _, _, finish in callbacks:
+                    if finish:
+                        finish(dsk, state, True)
+                if rerun_exceptions_locally:
+                    data = dict((dep, state['cache'][dep])
+                                for dep in get_dependencies(dsk, key))
+                    task = dsk[key]
+                    _execute_task(task, data)  # Re-execute locally
+                else:
+                    raise(remote_exception(res, tb))
+            state['cache'][key] = res
+            finish_task(dsk, key, state, results, keyorder.get)
+            for f in posttask_cbs:
+                f(key, res, dsk, state, worker_id)
+            while state['ready'] and len(state['running']) < num_workers:
+                fire_task()
+    except KeyboardInterrupt:
+        for cb in started_cbs:
+            if cb[-1]:
+                cb[-1](dsk, state, True)
+        raise
+
     # Final reporting
     while state['running'] or not queue.empty():
         key, res, tb, worker_id = queue.get()
 
-    for f in finish_cbs:
-        f(dsk, state, False)
+    for _, _, _, _, finish in started_cbs:
+        if finish:
+            finish(dsk, state, False)
 
     return nested_get(result, state['cache'])
 


Test patch: