Repository: dask/distributed

Issue:
Publish collections/futures
From one client/executor I may want to create a dataset and then share it with my colleagues working from other clients/executors.  Currently I can pickle my dask collection, send it to them, they unpickle it, and through the creation of Futures they get registered as clients to the data.  However this pickling-and-sending process is inconvenient.  We could publish collections to the scheduler itself
### Client 1

``` python
e.publish_dataset(df, 'analysis')
```
### Client 2

``` python
df = e.collect_dataset('analysis')
```
### Scheduler

``` python
>>> s.published_datasets
{'analysis': {'thing-to-send-to-client': ...,
              'relevant-keys': {set of keys}}
```

This is completely orthogonal functionality for the scheduler, but fairly trivial to build.  It should be a couple of handlers and a dictionary on the scheduler, a couple of methods on the Executor, and a sane way to collect futures and things that wrap futures into something that we can store in the dictionary.

This would be an interesting and moderately challenging task for someone like @jcrist or @martindurant if they have interest.  It would touch major pieces of code peripherally but it isn't conceptually or algorithmically tricky.



Discussion:
Given that the user can already specify the key name associated with any future in their executor when submitting, isn't the only task here to be able to create a future pointing to a given key on the scheduler?

A collection might contain several futures.  

Under normal use of collections like dask.dataframe users usually don't specify keynames.  We would want to associate a name to both a set of futures and sufficient metadata to re-create the dask.dataframe that surrounds those futures on the new client.


Resolving patch:
diff --git a/distributed/executor.py b/distributed/executor.py
--- a/distributed/executor.py
+++ b/distributed/executor.py
@@ -2,12 +2,14 @@
 
 from collections import defaultdict, Iterator, Iterable
 from concurrent.futures._base import DoneAndNotDoneFutures, CancelledError
+from contextlib import contextmanager
 import copy
 from datetime import timedelta
 from functools import partial
 import logging
 import os
 import sys
+import pickle
 from time import sleep
 import uuid
 from threading import Thread
@@ -192,6 +194,15 @@ def release(self):
             self._cleared = True
             self.executor._dec_ref(tokey(self.key))
 
+    def __getstate__(self):
+        return self.key
+
+    def __setstate__(self, key):
+        e = default_executor()
+        Future.__init__(self, key, e)
+        e._send_to_scheduler({'op': 'update-graph', 'tasks': {},
+                              'keys': [tokey(self.key)], 'client': e.id})
+
     def __del__(self):
         self.release()
 
@@ -975,6 +986,62 @@ def cancel(self, futures):
         """
         return sync(self.loop, self._cancel, futures)
 
+    @gen.coroutine
+    def _publish_dataset(self, data, name):
+        keys = [tokey(f.key) for f in futures_of(data)]
+        out = yield self.scheduler.publish_dataset(keys=keys, name=tokey(name),
+                                             data=dumps(data),
+                                             client=self.id)
+        if out['status'] != "OK":
+            raise KeyError(name)
+
+    def publish_dataset(self, data, name):
+        """
+        Store a named reference to the data in the scheduler for other executors
+
+        Parameters
+        ----------
+        data: list of futures or dask collection
+        name: str or container
+            handle by which the stored data should be known.
+
+        Returns
+        -------
+        None
+        """
+        sync(self.loop, self._publish_dataset, data, name)
+
+    def unpublish_dataset(self, name):
+        """
+        Remove reference to data on the scheduler
+        """
+        sync(self.loop, self.scheduler.unpublish_dataset, name=name)
+
+    def list_datasets(self):
+        """ Returns list of named datasets referenced in the scheduler
+
+        See publish_dataset()
+        """
+        return sync(self.loop, self.scheduler.list_datasets)
+
+    @gen.coroutine
+    def _get_dataset(self, name):
+        out = yield self.scheduler.get_dataset(name=name, client=self.id)
+        data, keys = out['data'], out['keys']
+
+        if not data:
+            raise KeyError(name)
+        with temp_default_executor(self):
+            data = loads(data)
+        raise Return(data)
+
+    def get_dataset(self, name):
+        """ Fetch named dataset from the scheduler.
+
+        See publish_dataset()
+        """
+        return sync(self.loop, self._get_dataset, tokey(name))
+
     @gen.coroutine
     def _run(self, function, *args, **kwargs):
         workers = kwargs.pop('workers', None)
@@ -1872,3 +1939,20 @@ def futures_of(o, executor=None):
             raise CancelledError(bad)
 
     return list(futures)
+
+
+@contextmanager
+def temp_default_executor(e):
+    """ Set the default executor for the duration of the context
+
+    Parameters
+    ----------
+    e : Executor
+        This is what default_executor() will return within the with-block.
+    """
+    old_exec = default_executor()
+    _global_executor[0] = e
+    try:
+        yield
+    finally:
+        _global_executor[0] = old_exec
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -223,6 +223,7 @@ def __init__(self, center=None, loop=None,
         self.exceptions = dict()
         self.tracebacks = dict()
         self.exceptions_blame = dict()
+        self.datasets = dict()
         self.stealable = [set() for i in range(12)]
         self.stealable_unknown_durations = defaultdict(set)
 
@@ -261,7 +262,11 @@ def __init__(self, center=None, loop=None,
                          'add_keys': self.add_keys,
                          'rebalance': self.rebalance,
                          'replicate': self.replicate,
-                         'start_ipython': self.start_ipython}
+                         'start_ipython': self.start_ipython,
+                         'list_datasets': self.list_datasets,
+                         'get_dataset': self.get_dataset,
+                         'publish_dataset': self.publish_dataset,
+                         'unpublish_dataset': self.unpublish_dataset}
 
         self.services = {}
         for k, v in (services or {}).items():
@@ -932,7 +937,7 @@ def handle_messages(self, in_queue, report, client=None):
         """
         The master client coroutine.  Handles all inbound messages from clients.
 
-        This runs once per Cleint IOStream or Queue.
+        This runs once per Client IOStream or Queue.
 
         See Also
         --------
@@ -1590,6 +1595,24 @@ def get_nbytes(self, stream=None, keys=None, summary=True):
 
             return result
 
+    def publish_dataset(self, stream=None, keys=None, data=None, name=None,
+                        client=None):
+        if name in self.datasets:
+            return {'status': 'failed'}
+        self.client_wants_keys(keys, 'published-%s' % name)
+        self.datasets[name] = {'data': data, 'keys': keys}
+        return {'status':  'OK'}
+
+    def unpublish_dataset(self, stream=None, name=None):
+        out = self.datasets.pop(name, {'keys': []})
+        self.client_releases_keys(out['keys'], 'published-%s' % name)
+
+    def list_datasets(self, *args):
+        return list(sorted(self.datasets.keys()))
+
+    def get_dataset(self, stream, name=None, client=None):
+        return self.datasets.get(name, {'data': [], 'keys': []})
+
     #####################
     # State Transitions #
     #####################


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
@@ -7,6 +7,7 @@
 from datetime import timedelta
 import itertools
 from multiprocessing import Process
+import pickle
 from random import random, choice
 import sys
 from threading import Thread
@@ -28,7 +29,8 @@
 from distributed.client import WrappedKey
 from distributed.executor import (Executor, Future, CompatibleExecutor, _wait,
         wait, _as_completed, as_completed, tokenize, _global_executor,
-        default_executor, _first_completed, ensure_default_get, futures_of)
+        default_executor, _first_completed, ensure_default_get, futures_of,
+        temp_default_executor)
 from distributed.scheduler import Scheduler, KilledWorker
 from distributed.sizeof import sizeof
 from distributed.utils import sync, tmp_text, ignoring, tokey, All
@@ -3202,6 +3204,98 @@ def f():
                 future = e.submit(f)
                 assert future.result() == 1
 
+
+@gen_cluster(executor=False)
+def test_publish_simple(s, a, b):
+    e = Executor((s.ip, s.port), start=False)
+    yield e._start()
+    f = Executor((s.ip, s.port), start=False)
+    yield f._start()
+
+    data = yield e._scatter(range(3))
+    out = yield e._publish_dataset(data, 'data')
+    assert 'data' in s.datasets
+
+    with pytest.raises(KeyError):
+        out = yield e._publish_dataset(data, 'data')
+
+    result = yield e.scheduler.list_datasets()
+    assert result == ['data']
+
+    result = yield f.scheduler.list_datasets()
+    assert result == ['data']
+
+
+@gen_cluster(executor=False)
+def test_publish_roundtrip(s, a, b):
+    e = Executor((s.ip, s.port), start=False)
+    yield e._start()
+    f = Executor((s.ip, s.port), start=False)
+    yield f._start()
+
+    data = yield e._scatter([0, 1, 2])
+    yield e._publish_dataset(data, 'data')
+
+    assert 'published-data' in s.who_wants[data[0].key]
+    result = yield f._get_dataset(name='data')
+
+    assert len(result) == len(data)
+    out = yield f._gather(result)
+    assert out == [0, 1, 2]
+
+    with pytest.raises(KeyError):
+        result = yield f._get_dataset(name='nonexistent')
+
+
+@gen_cluster(executor=True)
+def test_unpublish(e, s, a, b):
+    data = yield e._scatter([0, 1, 2])
+    yield e._publish_dataset(data, 'data')
+
+    key = data[0].key
+    del data
+
+    yield e.scheduler.unpublish_dataset(name='data')
+
+    assert 'data' not in s.datasets
+
+    start = time()
+    while key in s.who_wants:
+        yield gen.sleep(0.01)
+        assert time() < start + 5
+
+    with pytest.raises(KeyError):
+        result = yield e._get_dataset(name='data')
+
+
+@gen_cluster(executor=False)
+def test_publish_bag(s, a, b):
+    db = pytest.importorskip('dask.bag')
+    e = Executor((s.ip, s.port), start=False)
+    yield e._start()
+    f = Executor((s.ip, s.port), start=False)
+    yield f._start()
+
+    bag = db.from_sequence([0, 1, 2])
+    bagp = e.persist(bag)
+
+    assert len(futures_of(bagp)) == 3
+    keys = {f.key for f in futures_of(bagp)}
+    assert keys == set(bag.dask)
+
+    yield e._publish_dataset(bagp, 'data')
+
+    # check that serialization didn't affect original bag's dask
+    assert len(futures_of(bagp)) == 3
+
+    result = yield f._get_dataset('data')
+    assert set(result.dask.keys()) == set(bagp.dask.keys())
+    assert {f.key for f in result.dask.values()} == {f.key for f in bagp.dask.values()}
+
+    out = yield f.compute(result)._result()
+    assert out == [0, 1, 2]
+
+
 @gen_cluster(executor=True)
 def test_lose_scattered_data(e, s, a, b):
     [x] = yield e._scatter([1], workers=a.address)
@@ -3317,3 +3411,36 @@ def test_scatter_compute_store_lose_processing(e, s, a, b):
     assert y.status == 'cancelled'
     assert z.status == 'cancelled'
 
+
+@gen_cluster(executor=False)
+def test_serialize_future(s, a, b):
+    e = Executor((s.ip, s.port), start=False)
+    yield e._start()
+    f = Executor((s.ip, s.port), start=False)
+    yield f._start()
+
+    future = e.submit(lambda: 1)
+    result = yield future._result()
+
+    with temp_default_executor(f):
+        future2 = pickle.loads(pickle.dumps(future))
+        assert future2.executor is f
+        assert tokey(future2.key) in f.futures
+        result2 = yield future2._result()
+        assert result == result2
+
+
+@gen_cluster(executor=False)
+def test_temp_executor(s, a, b):
+    e = Executor((s.ip, s.port), start=False)
+    yield e._start()
+    f = Executor((s.ip, s.port), start=False)
+    yield f._start()
+
+    with temp_default_executor(e):
+        assert default_executor() is e
+        assert default_executor(f) is f
+
+    with temp_default_executor(f):
+        assert default_executor() is f
+        assert default_executor(e) is e