Repository: dask/dask

Issue:
diagnostics fail with nested dask schedulers
## Issue:
Example:
I have a custom graph which describes a flow of an algorithm with complex  inter dependencies.
Dask helps me efficiently pass data around.

Some subroutines of the algorithm use dask to parallelize the array computations.

I use the dask.get to handle the algorithm workflow so it's single threaded and there is no "daemon creating children" error.
I use dask.multiprocessing.get in the subroutines using dask (because there is some python bytecode).

When I try to use diagnostics:
```
with Profiler() as prof, ResourceProfiler(dt=0.25) as rprof, CacheProfiler() as cprof:
        get(dsk,'pull_m2')
    visualize([prof,rprof,cprof])
```
I get:
```
C:\Users\josh\Anaconda3\envs\compute\lib\site-packages\dask-0.14.3+28.gb0cf627-py3.5.egg\dask\diagnostics\profile.py in _finish(self, dsk, state, failed)
    329         t = default_timer()
    330         for k, (metric, start) in self._cache.items():
--> 331             self.results.append(CacheData(k, dsk[k], metric, start, t))
    332         self._cache.clear()
    333 

KeyError: ('sub-6b0bfecfb11b78fc1a21be54f83b65c1', 0, 0, 4)
```

This is a key from a subroutine. The routines them selves create the dask objects inside and do all the computing separately. The `with` context is trying to run diagnostics on all dask computations, when I only want it to work on the custom graph.

## Version:
dask-0.14.3+28.gb0cf627-py3.5
Windows 10 64 bit


Discussion:

Resolving patch:
diff --git a/dask/callbacks.py b/dask/callbacks.py
--- a/dask/callbacks.py
+++ b/dask/callbacks.py
@@ -1,5 +1,7 @@
 from __future__ import absolute_import, division, print_function
 
+from contextlib import contextmanager
+
 from .context import _globals
 
 __all__ = ['Callback', 'add_callbacks']
@@ -85,6 +87,22 @@ def unpack_callbacks(cbs):
         return [(), (), (), (), ()]
 
 
+@contextmanager
+def local_callbacks(callbacks=None):
+    """Allows callbacks to work with nested schedulers.
+
+    Callbacks will only be used by the first started scheduler they encounter.
+    This means that only the outermost scheduler will use global callbacks."""
+    global_callbacks = callbacks is None
+    if global_callbacks:
+        callbacks, _globals['callbacks'] = _globals['callbacks'], set()
+    try:
+        yield callbacks or ()
+    finally:
+        if global_callbacks:
+            _globals['callbacks'] = callbacks
+
+
 def normalize_callback(cb):
     """Normalizes a callback to a tuple"""
     if isinstance(cb, Callback):
diff --git a/dask/local.py b/dask/local.py
--- a/dask/local.py
+++ b/dask/local.py
@@ -122,7 +122,7 @@
                    has_tasks)
 from .context import _globals
 from .order import order
-from .callbacks import unpack_callbacks
+from .callbacks import unpack_callbacks, local_callbacks
 from .optimize import cull
 from .utils_test import add, inc  # noqa: F401
 
@@ -453,10 +453,6 @@ def get_async(apply_async, num_workers, dsk, result, cache=None,
     """
     queue = Queue()
 
-    if callbacks is None:
-        callbacks = _globals['callbacks']
-    _, _, pretask_cbs, posttask_cbs, _ = unpack_callbacks(callbacks)
-
     if isinstance(result, list):
         result_flat = set(flatten(result))
     else:
@@ -464,80 +460,80 @@ def get_async(apply_async, num_workers, dsk, result, cache=None,
     results = set(result_flat)
 
     dsk = dict(dsk)
-    started_cbs = []
-    try:
-        for cb in callbacks:
-            if cb[0]:
-                cb[0](dsk)
-            started_cbs.append(cb)
-
-        dsk, dependencies = cull(dsk, list(results))
-
-        keyorder = order(dsk)
-
-        state = start_state_from_dask(dsk, cache=cache, sortkey=keyorder.get)
-
-        for _, start_state, _, _, _ in callbacks:
-            if start_state:
-                start_state(dsk, state)
-
-        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")
-
-        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)
-
-            # Prep data to send
-            data = dict((dep, state['cache'][dep])
-                        for dep in get_dependencies(dsk, key))
-            # Submit
-            apply_async(execute_task,
-                        args=(key, dumps((dsk[key], data)),
-                              dumps, loads, get_id, pack_exception),
-                        callback=queue.put)
-
-        # 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_info, failed = queue_get(queue)
-            if failed:
-                exc, tb = loads(res_info)
-                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_exception(exc, tb)
-            res, worker_id = loads(res_info)
-            state['cache'][key] = res
-            finish_task(dsk, key, state, results, keyorder.get)
-            for f in posttask_cbs:
-                f(key, res, dsk, state, worker_id)
-
+    with local_callbacks(callbacks) as callbacks:
+        _, _, pretask_cbs, posttask_cbs, _ = unpack_callbacks(callbacks)
+        started_cbs = []
+        succeeded = False
+        try:
+            for cb in callbacks:
+                if cb[0]:
+                    cb[0](dsk)
+                started_cbs.append(cb)
+
+            dsk, dependencies = cull(dsk, list(results))
+
+            keyorder = order(dsk)
+
+            state = start_state_from_dask(dsk, cache=cache, sortkey=keyorder.get)
+
+            for _, start_state, _, _, _ in callbacks:
+                if start_state:
+                    start_state(dsk, state)
+
+            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")
+
+            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)
+
+                # Prep data to send
+                data = dict((dep, state['cache'][dep])
+                            for dep in get_dependencies(dsk, key))
+                # Submit
+                apply_async(execute_task,
+                            args=(key, dumps((dsk[key], data)),
+                                  dumps, loads, get_id, pack_exception),
+                            callback=queue.put)
+
+            # Seed initial tasks into the thread pool
             while state['ready'] and len(state['running']) < num_workers:
                 fire_task()
 
-    except BaseException:
-        for _, _, _, _, finish in started_cbs:
-            if finish:
-                finish(dsk, state, True)
-        raise
-
-    for _, _, _, _, finish in started_cbs:
-        if finish:
-            finish(dsk, state, False)
+            # Main loop, wait on tasks to finish, insert new ones
+            while state['waiting'] or state['ready'] or state['running']:
+                key, res_info, failed = queue_get(queue)
+                if failed:
+                    exc, tb = loads(res_info)
+                    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_exception(exc, tb)
+                res, worker_id = loads(res_info)
+                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()
+
+            succeeded = True
+
+        finally:
+            for _, _, _, _, finish in started_cbs:
+                if finish:
+                    finish(dsk, state, not succeeded)
 
     return nested_get(result, state['cache'])
 


Test patch:
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,6 +1,8 @@
 from dask.local import get_sync
+from dask.context import _globals
 from dask.threaded import get as get_threaded
 from dask.callbacks import Callback
+from dask.utils_test import add
 
 
 def test_start_callback():
@@ -70,3 +72,34 @@ def raise_keyboard():
     except BaseException as e:
         assert isinstance(e, KeyboardInterrupt)
     assert flag[0]
+
+
+def test_nested_schedulers():
+
+    class MyCallback(Callback):
+        def _start(self, dsk):
+            self.dsk = dsk
+
+        def _pretask(self, key, dsk, state):
+            assert key in self.dsk
+
+    inner_callback = MyCallback()
+    inner_dsk = {'x': (add, 1, 2),
+                 'y': (add, 'x', 3)}
+
+    def nested_call(x):
+        assert not _globals['callbacks']
+        with inner_callback:
+            return get_threaded(inner_dsk, 'y') + x
+
+    outer_callback = MyCallback()
+    outer_dsk = {'a': (nested_call, 1),
+                 'b': (add, 'a', 2)}
+
+    with outer_callback:
+        get_threaded(outer_dsk, 'b')
+
+    assert not _globals['callbacks']
+    assert outer_callback.dsk == outer_dsk
+    assert inner_callback.dsk == inner_dsk
+    assert not _globals['callbacks']