Repository: dask/distributed

Issue:
Asynchronous long-running tasks on workers
Dask is able to reason about and run long-running tasks on workers that themselves launch other tasks.  This enables some powerful and novel workflows.  http://distributed.readthedocs.io/en/latest/task-launch.html#extended-example

This system launches a separate thread for the long-running task to divorce it from the normal thread pool.  It then creates a local Client object and uses the synchronous API with that client.  

There may be some benefit to launching a long-running task inside of the main worker thread, running on the same IOLoop as the administrative worker code.  This has a few potential benefits:

1.  It might be less hacky than the current separate thread solution
2.  It would give the long-running task thread-safe access to the worker
3.  It would enable experimentation with asynchronous APIs in Python 3.5, which might bear long term fruit.

This would be useful for example in cases where we know the worker is sitting on some other data source, like a Kafka queue, and we want it to just pump that data into the system forever.  

cc @pitrou this sounds like the sort of thing that might interest you.


Discussion:
> There may be some benefit to launching a long-running task inside of the main worker thread, running on the same IOLoop as the administrative worker code

I suppose this assumes the long-running task is written as a coroutine, and yields execution at selected points to ensure reactivity to external requests. I agree it sounds useful to allow advanced users to run other "async" code on the same IO loop as the Worker.
Yes, the long running task would probably just be doing communication and administrative work.  We would have to encourage this strongly in documentation.
I'll try to tackle this soon.

Resolving patch:
diff --git a/distributed/client.py b/distributed/client.py
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -1250,7 +1250,7 @@ def _run(self, function, *args, **kwargs):
             if resp['status'] == 'OK':
                 results[key] = resp['result']
             elif resp['status'] == 'error':
-                raise loads(resp['exception'])
+                six.reraise(*clean_exception(**resp))
         raise Return(results)
 
     def run(self, function, *args, **kwargs):
@@ -1288,6 +1288,50 @@ def run(self, function, *args, **kwargs):
         """
         return sync(self.loop, self._run, function, *args, **kwargs)
 
+    @gen.coroutine
+    def _run_coroutine(self, function, *args, **kwargs):
+        workers = kwargs.pop('workers', None)
+        wait = kwargs.pop('wait', True)
+        responses = yield self.scheduler.broadcast(msg=dict(op='run_coroutine',
+                                                function=dumps(function),
+                                                args=dumps(args),
+                                                kwargs=dumps(kwargs),
+                                                wait=wait),
+                                                workers=workers)
+        if not wait:
+            raise Return(None)
+        else:
+            results = {}
+            for key, resp in responses.items():
+                if resp['status'] == 'OK':
+                    results[key] = resp['result']
+                elif resp['status'] == 'error':
+                    six.reraise(*clean_exception(**resp))
+            raise Return(results)
+
+    def run_coroutine(self, function, *args, **kwargs):
+        """
+        Spawn a coroutine on all workers.
+
+        This spaws a coroutine on all currently known workers and then waits
+        for the coroutine on each worker.  The coroutines' results are returned
+        as a dictionary keyed by worker address.
+
+        Parameters
+        ----------
+        function: a coroutine function
+            (typically a function wrapped in gen.coroutine or
+             a Python 3.5+ async function)
+        *args: arguments for remote function
+        **kwargs: keyword arguments for remote function
+        wait: boolean (default True)
+            Whether to wait for coroutines to end.
+        workers: list
+            Workers on which to run the function. Defaults to all known workers.
+
+        """
+        return sync(self.loop, self._run_coroutine, function, *args, **kwargs)
+
     def _graph_to_futures(self, dsk, keys, restrictions=None,
             loose_restrictions=None, allow_other_workers=True, priority=None,
             resources=None):
diff --git a/distributed/utils.py b/distributed/utils.py
--- a/distributed/utils.py
+++ b/distributed/utils.py
@@ -2,6 +2,7 @@
 
 from collections import Iterable
 from contextlib import contextmanager
+import inspect
 import logging
 import multiprocessing
 import os
@@ -48,6 +49,24 @@ def funcname(func):
         return str(func)
 
 
+def has_arg(func, argname):
+    """
+    Whether the function takes an argument with the given name.
+    """
+    while True:
+        try:
+            if 'dask_worker' in inspect.getargspec(func).args:
+                return True
+        except TypeError:
+            break
+        try:
+            # For Tornado coroutines and other decorated functions
+            func = func.__wrapped__
+        except AttributeError:
+            break
+    return False
+
+
 def get_fileno_limit():
     """
     Get the maximum number of open files per process.
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -37,8 +37,8 @@
 from .protocol.pickle import dumps, loads
 from .sizeof import sizeof
 from .threadpoolexecutor import ThreadPoolExecutor
-from .utils import (funcname, get_ip, _maybe_complex, log_errors, All,
-                    ignoring, validate_key, mp_context)
+from .utils import (funcname, get_ip, has_arg, _maybe_complex, log_errors,
+                    All, ignoring, validate_key, mp_context)
 
 _ncores = mp_context.cpu_count()
 
@@ -196,6 +196,7 @@ def __init__(self, scheduler_ip, scheduler_port, ip=None, ncores=None,
           'gather': self.gather,
           'compute-stream': self.compute_stream,
           'run': self.run,
+          'run_coroutine': self.run_coroutine,
           'get_data': self.get_data,
           'update_data': self.update_data,
           'delete_data': self.delete_data,
@@ -366,9 +367,13 @@ def executor_submit(self, key, function, *args, **kwargs):
         raise gen.Return(result)
 
 
-    def run(self, stream, function=None, args=(), kwargs={}):
+    def run(self, stream, function, args=(), kwargs={}):
         return run(self, stream, function=function, args=args, kwargs=kwargs)
 
+    def run_coroutine(self, stream, function, args=(), kwargs={}, wait=True):
+        return run(self, stream, function=function, args=args, kwargs=kwargs,
+                   is_coro=True, wait=wait)
+
     def update_data(self, stream=None, data=None, report=True):
         self.data.update(data)
         self.nbytes.update(valmap(sizeof, data))
@@ -717,20 +722,20 @@ def weight(k, v):
 
 
 @gen.coroutine
-def run(worker, stream, function=None, args=(), kwargs={}):
+def run(worker, stream, function, args=(), kwargs={}, is_coro=False, wait=True):
+    assert wait or is_coro, "Combination not supported"
     function = loads(function)
     if args:
         args = loads(args)
     if kwargs:
         kwargs = loads(kwargs)
-    try:
-        import inspect
-        if 'dask_worker' in inspect.getargspec(function).args:
-            kwargs['dask_worker'] = worker
-    except:
-        pass
+    if has_arg(function, 'dask_worker'):
+        kwargs['dask_worker'] = worker
+    logger.info("Run out-of-band function %r", funcname(function))
     try:
         result = function(*args, **kwargs)
+        if is_coro:
+            result = (yield result) if wait else None
     except Exception as e:
         logger.warn(" Run Failed\n"
             "Function: %s\n"


Test patch:
diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
--- a/distributed/tests/test_client.py
+++ b/distributed/tests/test_client.py
@@ -37,8 +37,8 @@
 from distributed.sizeof import sizeof
 from distributed.utils import sync, tmp_text, ignoring, tokey, All, mp_context
 from distributed.utils_test import (cluster, slow, slowinc, slowadd, slowdec,
-        randominc, loop, inc, dec, div, throws, gen_cluster, gen_test, double,
-        deep, popen)
+        randominc, loop, inc, dec, div, throws, geninc, asyncinc,
+        gen_cluster, gen_test, double, deep, popen)
 
 
 @gen_cluster(client=True, timeout=None)
@@ -2188,6 +2188,44 @@ def func(x, y=10):
             assert result == {'127.0.0.1:%d' % a['port']: 3}
 
 
+@gen_cluster(client=True)
+def test_run_coroutine(c, s, a, b):
+    results = yield c._run_coroutine(geninc, 1, delay=0.05)
+    assert results == {a.address: 2, b.address: 2}
+
+    results = yield c._run_coroutine(geninc, 1, delay=0.05, workers=[a.address])
+    assert results == {a.address: 2}
+
+    results = yield c._run_coroutine(geninc, 1, workers=[])
+    assert results == {}
+
+    with pytest.raises(RuntimeError) as exc_info:
+        yield c._run_coroutine(throws, 1)
+    exc_info.match("hello")
+
+    if sys.version_info >= (3, 5):
+        results = yield c._run_coroutine(asyncinc, 2, delay=0.01)
+        assert results == {a.address: 3, b.address: 3}
+
+
+def test_run_coroutine_sync(loop):
+    with cluster() as (s, [a, b]):
+        with Client(('127.0.0.1', s['port']), loop=loop) as c:
+            result = c.run_coroutine(geninc, 2, delay=0.01)
+            assert result == {'127.0.0.1:%d' % a['port']: 3,
+                              '127.0.0.1:%d' % b['port']: 3}
+
+            result = c.run_coroutine(geninc, 2,
+                                     workers=['127.0.0.1:%d' % a['port']])
+            assert result == {'127.0.0.1:%d' % a['port']: 3}
+
+            t1 = time()
+            result = c.run_coroutine(geninc, 2, delay=10, wait=False)
+            t2 = time()
+            assert result is None
+            assert t2 - t1 <= 1.0
+
+
 def test_run_exception(loop):
     def raise_exception(exc_type, exc_msg):
         raise exc_type(exc_msg)
diff --git a/distributed/tests/test_worker.py b/distributed/tests/test_worker.py
--- a/distributed/tests/test_worker.py
+++ b/distributed/tests/test_worker.py
@@ -13,6 +13,7 @@
 from dask import delayed
 import pytest
 from toolz import pluck, sliding_window
+import tornado
 from tornado import gen
 from tornado.ioloop import TimeoutError
 
@@ -364,6 +365,20 @@ def f(dask_worker=None):
     assert response == {a.address: a.id, b.address: b.id}
 
 
+@gen_cluster(client=True)
+def test_run_coroutine_dask_worker(c, s, a, b):
+    if sys.version_info < (3,) and tornado.version_info < (4, 5):
+        pytest.skip("test needs Tornado 4.5+ on Python 2.7")
+
+    @gen.coroutine
+    def f(dask_worker=None):
+        yield gen.sleep(0.001)
+        raise gen.Return(dask_worker.id)
+
+    response = yield c._run_coroutine(f)
+    assert response == {a.address: a.id, b.address: b.id}
+
+
 @gen_cluster(client=True, ncores=[])
 def test_Executor(c, s):
     with ThreadPoolExecutor(2) as e:
diff --git a/distributed/utils_test.py b/distributed/utils_test.py
--- a/distributed/utils_test.py
+++ b/distributed/utils_test.py
@@ -10,6 +10,7 @@
 import socket
 import subprocess
 import sys
+import textwrap
 from time import sleep
 import uuid
 
@@ -185,6 +186,31 @@ def slowidentity(*args, **kwargs):
     return args
 
 
+@gen.coroutine
+def geninc(x, delay=0.02):
+    yield gen.sleep(delay)
+    raise gen.Return(x + 1)
+
+
+def compile_snippet(code, dedent=True):
+    if dedent:
+        code = textwrap.dedent(code)
+    code = compile(code, '<dynamic>', 'exec')
+    ns = globals()
+    exec(code, ns, ns)
+
+
+if sys.version_info >= (3, 5):
+    compile_snippet("""
+        async def asyncinc(x, delay=0.02):
+            await gen.sleep(delay)
+            return x + 1
+        """)
+    assert asyncinc
+else:
+    asyncinc = None
+
+
 _readone_queues = {}
 
 @gen.coroutine