Create thread state within dask/base.py

dask/distributed will now use this thread_state by default

dask/__init__.py

@@ -8,7 +8,7 @@
 except ImportError:
     pass
 try:
-    from .base import visualize, compute, persist
+    from .base import visualize, compute, persist, thread_state
 except ImportError:
     pass
 

dask/base.py

@@ -24,6 +24,9 @@
 __all__ = ("Base", "compute", "normalize_token", "tokenize", "visualize")
 
 
+thread_state = threading.local()
+
+
 class Base(object):
     """Base class for dask collections"""
     __slots__ = ()
@@ -193,7 +196,7 @@ def compute(*args, **kwargs):
 
     get = kwargs.pop('get', None) or _globals['get']
 
-    if get is None and _globals.get('distributed_worker'):
+    if get is None and getattr(thread_state, 'key', False):
         from distributed.worker import get_worker
         get = get_worker().client.get
 
@@ -558,7 +561,7 @@ def persist(*args, **kwargs):
 
     get = kwargs.pop('get', None) or _globals['get']
 
-    if get is None and _globals.get('distributed_worker'):
+    if get is None and getattr(thread_state, 'key', False):
         from distributed.worker import get_worker
         get = get_worker().client.get
 

docs/source/changelog.rst

@@ -24,6 +24,8 @@ Core
 ++++
 
 -  Allow tuples as sharedict keys (:pr:`2763`)
+-  Calling compute within a dask.distributed task defaults to distributed
+   scheduler (:pr:`2762`)
 
 0.15.4 / 2017-10-06
 -------------------