Repository: dask/distributed

Issue:
Direct function evaluation
For diagnostic purposes you sometimes want to run a function outside of the task scheduling framework and get an immediate result.  We can use the scheduler's broadcast functionality to manage this.



Discussion:
@quasiben 


Resolving patch:
diff --git a/distributed/executor.py b/distributed/executor.py
--- a/distributed/executor.py
+++ b/distributed/executor.py
@@ -142,9 +142,9 @@ def exception(self):
         """
         return sync(self.executor.loop, self._exception)
 
-    def cancel(self, block=False):
+    def cancel(self):
         """ Returns True if the future has been cancelled """
-        return self.executor.cancel([self], block=False)
+        return self.executor.cancel([self])
 
     def cancelled(self):
         """ Returns True if the future has been cancelled """
@@ -828,16 +828,14 @@ def scatter(self, data, workers=None, broadcast=False):
             return sync(self.loop, self._scatter, data, workers=workers,
                         broadcast=broadcast)
     @gen.coroutine
-    def _cancel(self, futures, block=False):
+    def _cancel(self, futures):
         keys = {f.key for f in futures_of(futures)}
-        f = self.scheduler.cancel(keys=list(keys), client=self.id)
-        if block:
-            yield f
+        f = yield self.scheduler.cancel(keys=list(keys), client=self.id)
         for k in keys:
             with ignoring(KeyError):
                 del self.futures[k]
 
-    def cancel(self, futures, block=False):
+    def cancel(self, futures):
         """
         Cancel running futures
 
@@ -849,7 +847,52 @@ def cancel(self, futures, block=False):
         ----------
         futures: list of Futures
         """
-        return sync(self.loop, self._cancel, futures, block=False)
+        return sync(self.loop, self._cancel, futures)
+
+    @gen.coroutine
+    def _run(self, function, *args, **kwargs):
+        workers = kwargs.pop('workers', None)
+        result = yield self.scheduler.broadcast(msg=dict(op='run',
+                                                function=dumps(function),
+                                                args=dumps(args),
+                                                kwargs=dumps(kwargs)),
+                                                workers=workers)
+        raise Return(result)
+
+    def run(self, function, *args, **kwargs):
+        """
+        Run a function on all workers outside of task scheduling system
+
+        This calls a function on all currently known workers immediately,
+        blocks until those results come back, and returns the results
+        asynchronously as a dictionary keyed by worker address.  This method
+        if generally used for side effects, such and collecting diagnostic
+        information or installing libraries.
+
+        Parameters
+        ----------
+        function: callable
+        *args: arguments for remote function
+        **kwargs: keyword arguments for remote function
+        workers: list
+            Workers on which to run the function. Defaults to all known workers.
+
+        Examples
+        --------
+        >>> e.run(os.getpid)  # doctest: +SKIP
+        {'192.168.0.100:9000': 1234,
+         '192.168.0.101:9000': 4321,
+         '192.168.0.102:9000': 5555}
+
+        Restrict computation to particular workers with the ``workers=``
+        keyword argument.
+
+        >>> e.run(os.getpid, workers=['192.168.0.100:9000',
+        ...                           '192.168.0.101:9000'])  # doctest: +SKIP
+        {'192.168.0.100:9000': 1234,
+         '192.168.0.101:9000': 4321}
+        """
+        return sync(self.loop, self._run, function, *args, **kwargs)
 
     @gen.coroutine
     def _get(self, dsk, keys, restrictions=None, raise_on_error=True):
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -1187,9 +1187,10 @@ def get_ncores(self, stream, addresses=None):
             return self.ncores
 
     @gen.coroutine
-    def broadcast(self, stream, msg=None):
+    def broadcast(self, stream=None, msg=None, workers=None):
         """ Broadcast message to workers, return all results """
-        workers = list(self.ncores)
+        if workers is None:
+            workers = list(self.ncores)
         results = yield All([send_recv(arg=address, close=True, **msg)
                              for address in workers])
         raise Return(dict(zip(workers, results)))
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -115,6 +115,7 @@ def __init__(self, center_ip, center_port, ip=None, ncores=None,
             sys.path.insert(0, self.local_dir)
 
         handlers = {'compute': self.compute,
+                    'run': self.run,
                     'get_data': self.get_data,
                     'update_data': self.update_data,
                     'delete_data': self.delete_data,
@@ -297,6 +298,17 @@ def compute(self, stream, function=None, key=None, args=(), kwargs={},
             self.active.remove(key)
         raise Return(out)
 
+    @gen.coroutine
+    def run(self, stream, function=None, args=(), kwargs={}):
+        function = loads(function)
+        if args:
+            args = loads(args)
+        if kwargs:
+            kwargs = loads(kwargs)
+
+        result = function(*args, **kwargs)
+        raise Return(result)
+
     @gen.coroutine
     def update_data(self, stream, data=None, report=True):
         data = valmap(loads, data)


Test patch:
diff --git a/distributed/tests/test_executor.py b/distributed/tests/test_executor.py
--- a/distributed/tests/test_executor.py
+++ b/distributed/tests/test_executor.py
@@ -1733,7 +1733,7 @@ def test__cancel(e, s, a, b):
     while y.key not in s.tasks:
         yield gen.sleep(0.01)
 
-    yield e._cancel([x], block=True)
+    yield e._cancel([x])
 
     assert x.cancelled()
     assert 'cancel' in str(x)
@@ -1760,12 +1760,15 @@ def test__cancel_multi_client(s, a, b):
 
     assert x.key == y.key
 
-    yield e._cancel([x], block=True)
+    yield e._cancel([x])
 
     assert x.cancelled()
     assert not y.cancelled()
 
-    assert y.key in s.tasks
+    start = time()
+    while y.key not in s.tasks:
+        yield gen.sleep(0.01)
+        assert time() < start + 5
 
     out = yield y._result()
     assert out == 2
@@ -1797,7 +1800,7 @@ def test_cancel(loop):
             y = e.submit(slowinc, x)
             z = e.submit(slowinc, y)
 
-            e.cancel([y], block=True)
+            e.cancel([y])
 
             start = time()
             while not z.cancelled():
@@ -1806,7 +1809,7 @@ def test_cancel(loop):
 
             assert x.result() == 2
 
-            z.cancel(block=True)
+            z.cancel()
             assert z.cancelled()
 
 
@@ -2042,3 +2045,29 @@ def test_balance_tasks_by_stacks(e, s, a, b):
     yield _wait(y)
 
     assert len(a.data) == len(b.data) == 1
+
+
+@gen_cluster(executor=True)
+def test_run(e, s, a, b):
+    results = yield e._run(inc, 1)
+    assert results == {a.address: 2, b.address: 2}
+
+    results = yield e._run(inc, 1, workers=[a.address])
+    assert results == {a.address: 2}
+
+    results = yield e._run(inc, 1, workers=[])
+    assert results == {}
+
+
+def test_run_sync(loop):
+    def func(x, y=10):
+        return x + y
+
+    with cluster() as (s, [a, b]):
+        with Executor(('127.0.0.1', s['port']), loop=loop) as e:
+            result = e.run(func, 1, y=2)
+            assert result == {'127.0.0.1:%d' % a['port']: 3,
+                              '127.0.0.1:%d' % b['port']: 3}
+
+            result = e.run(func, 1, y=2, workers=['127.0.0.1:%d' % a['port']])
+            assert result == {'127.0.0.1:%d' % a['port']: 3}
diff --git a/distributed/tests/test_scheduler.py b/distributed/tests/test_scheduler.py
--- a/distributed/tests/test_scheduler.py
+++ b/distributed/tests/test_scheduler.py
@@ -916,3 +916,12 @@ def test_str_graph():
         skeys = [str(k) for k in keys]
         assert all(isinstance(k, (str, bytes)) for k in sdsk)
         assert dask.get(dsk, keys) == dask.get(sdsk, skeys)
+
+
+@gen_cluster()
+def test_ready_add_worker(s, a, b):
+    result = yield s.broadcast(msg={'op': 'ping'})
+    assert result == {a.address: b'pong', b.address: b'pong'}
+
+    result = yield s.broadcast(msg={'op': 'ping'}, workers=[a.address])
+    assert result == {a.address: b'pong'}