Repository: dask/dask

Issue:
ProgressBar doesn't exit cleanly on error
When call to get terminates abruptly due to an error in one of the tasks the progressbar continues counting off time.  Presumably we should call the finalize callback when we encounter an error.



Discussion:
Hmm, that is a problem. Calling `finish` on error seems wrong to me though, as there's no way to distinguish between an error and the actual end of dask. I see three solutions to this:
- Add yet another callback to get called on errors. Could potentially call it `cleanup` or something, and have it called at the end as well. I imagine this taking no parameters.
- Handle this in the `__exit__` method of `Profiler`. Unfortunately this will fail if the callbacks are set globally instead of contextually.
- Modify `finish` to take a third parameter that's either `None` or the traceback. This would work, but destroys the symmetry between `start` and `finish`.

I'd prefer either the 1st or 3rd option.

In current code is there anything that would break by calling `finish` on error?

None at all. I would want to add a parameter though to note whether it was a clean finish or not. Could be boolean, or include the traceback (thinking about it's uses, I think I'd actual prefer the parameter to be `clean_finish` or `errored`, and boolean).

I trust your judgement here more than my own.  


Resolving patch:
diff --git a/dask/async.py b/dask/async.py
--- a/dask/async.py
+++ b/dask/async.py
@@ -467,6 +467,8 @@ def fire_task():
     while state['waiting'] or state['ready'] or state['running']:
         key, res, tb, worker_id = queue.get()
         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))
@@ -495,7 +497,7 @@ def fire_task():
         key, res, tb, worker_id = queue.get()
 
     for f in finish_cbs:
-        f(dsk, state)
+        f(dsk, state, False)
 
     return nested_get(result, state['cache'])
 
diff --git a/dask/callbacks.py b/dask/callbacks.py
--- a/dask/callbacks.py
+++ b/dask/callbacks.py
@@ -1,4 +1,3 @@
-from contextlib import contextmanager
 from collections import namedtuple
 
 from .context import _globals
@@ -18,8 +17,7 @@ def unpack_callbacks(cbs):
         return [(), (), (), ()]
 
 
-@contextmanager
-def add_callbacks(*args):
+class add_callbacks(object):
     """Context manager for callbacks.
 
     Takes several callback tuples and applies them only in the enclosed
@@ -33,7 +31,12 @@ def add_callbacks(*args):
     >>> with add_callbacks(callbacks):    # doctest: +SKIP
     ...     res.compute()
     """
-    old = _globals['callbacks'].copy()
-    _globals['callbacks'].update(args)
-    yield
-    _globals['callbacks'] = old
+    def __init__(self, *args):
+        self.old = _globals['callbacks'].copy()
+        _globals['callbacks'].update(args)
+
+    def __enter__(self):
+        return
+
+    def __exit__(self, type, value, traceback):
+        _globals['callbacks'] = self.old
diff --git a/dask/diagnostics/progress.py b/dask/diagnostics/progress.py
--- a/dask/diagnostics/progress.py
+++ b/dask/diagnostics/progress.py
@@ -55,7 +55,7 @@ def _start(self, dsk, state):
     def _posttask(self, key, value, dsk, state, id):
         self._ndone += 1
 
-    def _finish(self, dsk, state):
+    def _finish(self, dsk, state, errored):
         self._running = False
         self._timer.join()
         self._finalize_bar()


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
@@ -11,7 +11,7 @@
        'e': (mul, 'c', 'd')}
 
 
-def test_profiler(capsys):
+def test_progressbar(capsys):
     with ProgressBar():
         out = get(dsk, 'e')
     assert out == 6
@@ -27,6 +27,17 @@ def test_profiler(capsys):
     assert percent == "100% Completed"
 
 
+def test_clean_exit():
+    dsk = {'a': (lambda: 1/0,)}
+    try:
+        with ProgressBar() as pbar:
+            get(dsk, 'a')
+    except:
+        pass
+    assert not pbar._running
+    assert not pbar._timer.is_alive()
+
+
 def test_format_time():
     assert format_time(1.4) == ' 1.4s'
     assert format_time(10.4) == '10.4s'