Repository: dask/dask

Issue:
ProgressBar keeps running after exception
Using dask 0.12.0 the progress bar keeps running after an exception occurs. Using [this census_sample.h5](http://assets.holoviews.org/census_sample.h5) file, the issue occurs when running:

```python
pbar = ProgressBar()
pbar.register()
df = dd.read_hdf('census_sample.h5', key='census')
dd.io.to_castra(df, './census.castra')
```


Discussion:
Here's a self-contained example needing no sample file:

``` python
import dask.dataframe as dd
import numpy as np

from dask.diagnostics import ProgressBar
ProgressBar().register()

!touch dummy.castra
x = np.zeros(100)
df=dd.from_array(x)
dd.io.to_castra(df, 'dummy.castra')
```

Just put this in a cell in a Jupyter notebook and execute that cell.  (Run All doesn't do it, only cell by cell.)


Resolving patch:
diff --git a/dask/async.py b/dask/async.py
--- a/dask/async.py
+++ b/dask/async.py
@@ -268,14 +268,14 @@ def execute_task(key, task_info, dumps, loads, get_id, raise_on_exception=False)
         result = _execute_task(task, data)
         id = get_id()
         result = dumps((result, None, id))
-    except Exception as e:
+    except BaseException as e:
         if raise_on_exception:
             raise
         exc_type, exc_value, exc_traceback = sys.exc_info()
         tb = ''.join(traceback.format_tb(exc_traceback))
         try:
             result = dumps((e, tb, None))
-        except Exception as e:
+        except BaseException as e:
             if raise_on_exception:
                 raise
             exc_type, exc_value, exc_traceback = sys.exc_info()
@@ -482,17 +482,8 @@ def fire_task():
         # Main loop, wait on tasks to finish, insert new ones
         while state['waiting'] or state['ready'] or state['running']:
             key, res_info = queue.get()
-            try:
-                res, tb, worker_id = loads(res_info)
-            except Exception:
-                for _, _, _, _, finish in callbacks:
-                    if finish:
-                        finish(dsk, state, True)
-                raise
-            if isinstance(res, Exception):
-                for _, _, _, _, finish in callbacks:
-                    if finish:
-                        finish(dsk, state, True)
+            res, tb, worker_id = loads(res_info)
+            if isinstance(res, BaseException):
                 if rerun_exceptions_locally:
                     data = dict((dep, state['cache'][dep])
                                 for dep in get_dependencies(dsk, key))
@@ -508,16 +499,12 @@ def fire_task():
             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)
+    except BaseException:
+        for _, _, _, _, finish in started_cbs:
+            if finish:
+                finish(dsk, state, True)
         raise
 
-    # Final reporting
-    while state['running'] or not queue.empty():
-        key, res, tb, worker_id = queue.get()
-
     for _, _, _, _, finish in started_cbs:
         if finish:
             finish(dsk, state, False)


Test patch:
diff --git a/dask/diagnostics/tests/test_progress.py b/dask/diagnostics/tests/test_progress.py
--- a/dask/diagnostics/tests/test_progress.py
+++ b/dask/diagnostics/tests/test_progress.py
@@ -1,8 +1,11 @@
 from operator import add, mul
+
 import pytest
+
+from dask.async import get_sync
 from dask.diagnostics import ProgressBar
 from dask.diagnostics.progress import format_time
-from dask.threaded import get
+from dask.threaded import get as get_threaded
 from dask.context import _globals
 
 
@@ -22,26 +25,27 @@ def check_bar_completed(capsys, width=40):
 
 def test_progressbar(capsys):
     with ProgressBar():
-        out = get(dsk, 'e')
+        out = get_threaded(dsk, 'e')
     assert out == 6
     check_bar_completed(capsys)
     with ProgressBar(width=20):
-        out = get(dsk, 'e')
+        out = get_threaded(dsk, 'e')
     check_bar_completed(capsys, 20)
 
 
 def test_minimum_time(capsys):
     with ProgressBar(1.0):
-        out = get(dsk, 'e')
+        out = get_threaded(dsk, 'e')
     out, err = capsys.readouterr()
     assert out == '' and err == ''
 
 
-def test_clean_exit():
+@pytest.mark.parametrize('get', [get_threaded, get_sync])
+def test_clean_exit(get):
     dsk = {'a': (lambda: 1 / 0, )}
     try:
         with ProgressBar() as pbar:
-            get(dsk, 'a')
+            get_threaded(dsk, 'a')
     except ZeroDivisionError:
         pass
     assert not pbar._running
@@ -63,7 +67,7 @@ def test_register(capsys):
 
         assert _globals['callbacks']
 
-        get(dsk, 'e')
+        get_threaded(dsk, 'e')
         check_bar_completed(capsys)
 
         p.unregister()
@@ -75,7 +79,7 @@ def test_register(capsys):
 
 def test_no_tasks(capsys):
     with ProgressBar():
-        get({'x': 1}, 'x')
+        get_threaded({'x': 1}, 'x')
     check_bar_completed(capsys)
 
 
@@ -87,13 +91,13 @@ def test_with_cache(capsys):
 
     with cc:
         with ProgressBar():
-            assert get({'x': (mul, 1, 2)}, 'x') == 2
+            assert get_threaded({'x': (mul, 1, 2)}, 'x') == 2
     check_bar_completed(capsys)
     assert c.data['x'] == 2
 
     with cc:
         with ProgressBar():
-            assert get({'x': (mul, 1, 2), 'y': (mul, 'x', 3)}, 'y') == 6
+            assert get_threaded({'x': (mul, 1, 2), 'y': (mul, 'x', 3)}, 'y') == 6
     check_bar_completed(capsys)
 
 
@@ -105,13 +109,13 @@ def test_with_alias(capsys):
            'e': 'd',
            'f': (mul, 'e', 'c')}
     with ProgressBar():
-        get(dsk, 'f')
+        get_threaded(dsk, 'f')
     check_bar_completed(capsys)
 
 
 def test_store_time():
     p = ProgressBar()
     with p:
-        get({'x': 1}, 'x')
+        get_threaded({'x': 1}, 'x')
 
     assert isinstance(p.last_duration, float)
diff --git a/dask/tests/test_async.py b/dask/tests/test_async.py
--- a/dask/tests/test_async.py
+++ b/dask/tests/test_async.py
@@ -155,7 +155,7 @@ def test_order_of_startstate():
     assert result['ready'] == ['b', 'y']
 
 
-def test_nonstandard_exceptions_propagate():
+def test_exceptions_propagate():
     class MyException(Exception):
         def __init__(self, a, b):
             self.a = a
diff --git a/dask/tests/test_callbacks.py b/dask/tests/test_callbacks.py
--- a/dask/tests/test_callbacks.py
+++ b/dask/tests/test_callbacks.py
@@ -1,4 +1,5 @@
 from dask.async import get_sync
+from dask.threaded import get as get_threaded
 from dask.callbacks import Callback
 
 
@@ -28,3 +29,44 @@ def _start_state(self, dsk, state):
         get_sync({'x': 1}, 'x')
 
     assert flag[0] is True
+
+
+def test_finish_always_called():
+    flag = [False]
+
+    class MyCallback(Callback):
+        def _finish(self, dsk, state, errored):
+            flag[0] = True
+            assert errored
+
+    dsk = {'x': (lambda: 1 / 0,)}
+
+    # `raise_on_exception=True`
+    try:
+        with MyCallback():
+            get_sync(dsk, 'x')
+    except Exception as e:
+        assert isinstance(e, ZeroDivisionError)
+    assert flag[0]
+
+    # `raise_on_exception=False`
+    flag[0] = False
+    try:
+        with MyCallback():
+            get_threaded(dsk, 'x')
+    except Exception as e:
+        assert isinstance(e, ZeroDivisionError)
+    assert flag[0]
+
+    # KeyboardInterrupt
+    def raise_keyboard():
+        raise KeyboardInterrupt()
+
+    dsk = {'x': (raise_keyboard,)}
+    flag[0] = False
+    try:
+        with MyCallback():
+            get_sync(dsk, 'x')
+    except BaseException as e:
+        assert isinstance(e, KeyboardInterrupt)
+    assert flag[0]