Repository: dask/dask

Issue:
Rerun erroring task to provide pdb support
It is difficult to debug erroring code in a parallel context.  Functions run in separate threads/processes/machines don't provide the ability to run `pdb` when errors occur in user functions.  Currently we send back the error and the traceback, but this often isn't as good as actually investigating the errored state interactively.

Fortunately we have the troublemaking function and arguments, we could rerun the task in the user's thread and drop them into the error that originally occurred on the remote resource.

This is a bit tricky if the function has side effects.  We might want users to opt-in to this. 



Discussion:

Resolving patch:
diff --git a/dask/async.py b/dask/async.py
--- a/dask/async.py
+++ b/dask/async.py
@@ -380,7 +380,9 @@ def default_get_id():
 
 def get_async(apply_async, num_workers, dsk, result, cache=None,
               queue=None, get_id=default_get_id, raise_on_exception=False,
-              start_callback=None, end_callback=None, **kwargs):
+              start_callback=None, end_callback=None,
+              rerun_exceptions_locally=None,
+              **kwargs):
     """ Asynchronous get function
 
     This is a general version of various asynchronous schedulers for dask.  It
@@ -408,6 +410,9 @@ def get_async(apply_async, num_workers, dsk, result, cache=None,
     get_id : callable, optional
         Function to return the worker id, takes no arguments. Examples are
         `threading.current_thread` and `multiprocessing.current_process`.
+    rerun_exceptions_locally : bool, optional
+        Whether to rerun failing tasks in local process to enable debugging
+        (False by default)
     start_callback : function, optional
         Callback run every time a new task is started. Receives the key of the
         task to be run, the dask, and the scheduler state. At the end of
@@ -437,6 +442,9 @@ def get_async(apply_async, num_workers, dsk, result, cache=None,
 
     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)
+
     if state['waiting'] and not state['ready']:
         raise ValueError("Found no accessible jobs in dask")
 
@@ -464,8 +472,22 @@ def fire_task():
     while state['waiting'] or state['ready'] or state['running']:
         key, res, tb, worker_id = queue.get()
         if isinstance(res, Exception):
-            raise type(res)(" Exception in remote process\n\n"
-                + str(res) + "\n\nTraceback:\n" + tb)
+            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 type(res)(
+                "Exception occurred in remote worker.\n\n"
+                "Something you've asked dask to compute raised an exception.\n"
+                "That exception and the traceback are copied below.\n"
+                "To use pdb, rerun the computation with the keyword argument\n"
+                "    dask.set_options(rerun_exceptions_locally=True)\n"
+                "    or\n"
+                "    dataset.compute(rerun_exceptions_locally=True)\n\n"
+                "The original exception and traceback follow below:\n\n"
+                    + str(res) + "\n\nTraceback:\n" + tb)
         state['cache'][key] = res
         finish_task(dsk, key, state, results, keyorder.get)
         if end_callback:
diff --git a/dask/context.py b/dask/context.py
--- a/dask/context.py
+++ b/dask/context.py
@@ -21,6 +21,7 @@ class set_options(object):
         cache - Cache to use for intermediate results
         func_loads/func_dumps - loads/dumps functions for serialization of data
             likely to contain functions.  Defaults to dill.loads/dill.dumps
+        rerun_exceptions_locally - rerun failed tasks in master process
 
     Example
     -------


Test patch:
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
@@ -147,3 +147,27 @@ def test_order_of_startstate():
     result = start_state_from_dask(dsk)
 
     assert result['ready'] == ['b', 'y']
+
+
+def test_rerun_exceptions_locally():
+    counter = [0]
+    def f():
+        counter[0] += 1
+        raise Exception('TOKEN')
+
+    from dask.threaded import get
+    try:
+        get({'x': (f,)}, 'x')
+    except Exception as e:
+        assert 'remote' in str(e).lower()
+
+    try:
+        get({'x': (f,)}, 'x', rerun_exceptions_locally=True)
+    except Exception as e:
+        assert 'remote' not in str(e).lower()
+
+    try:
+        with dask.set_options(rerun_exceptions_locally=True):
+            get({'x': (f,)}, 'x')
+    except Exception as e:
+        assert 'remote' not in str(e).lower()