Repository: dask/dask

Issue:
Caching option for dask.DataFrame
As part of the analytical monitoring related workflow, we find querying the same type of data over and over again and there is currently no way to persist the dataframes in-memory for a near-realtime and interactive analysis.  This may also lead into  "aggressive caching" bit that i know @mrocklin is interested in solving for as well.



Discussion:
For single node operation the solution is to back a dask.dataframe off of a pandas.dataframe in the following way.

```
df = dd.from_pandas(df.compute(), npartition=10)
```

However from our talks I believe that @dataisle is more interested in the distributed memory case, where indeed no such solution exists.  

This is more a question about the distributed scheduler and how we can choose to keep results in memory on the workers.  I see this problem as two parts
1.  How the client tells the scheduler that we want to keep data on the workers
2.  A quick optimization that `culls` incoming graphs of tasks whose results are already in worker data.

The second form of caching that we discussed was about opportunistically caching intermediate results without the user asking for it.  A brief writeup for what I was thinking about previously exists [here](https://github.com/mrocklin/sbe/blob/master/caching.md).  This topic feels somewhat orthogonal though.

cc @cowlicks, who has been thinking about distributed scheduling recently.

Thanks. Opportunistic caching seems like a really awesome end-state. i am in the middle of creating a computation setup with medium size data (around 60GB) using dask.distributed and i think i could potentially play around the scheduler API to let the workers know to cache a dataframe.

How are you thinking about the optimization that culls the precomputed elements of the graphs?

It's in `dask.optimize.cull`
- [Docs](http://dask.pydata.org/en/latest/optimize.html)
- [Code](https://github.com/ContinuumIO/dask/blob/master/dask/optimize.py#L9-L33)

Cool. thanks.

So when a scheduler [receives a graph from the client](https://github.com/ContinuumIO/dask/blob/master/dask/distributed/scheduler.py#L573) it can cull the graph based on the [keys found in the workers](https://github.com/ContinuumIO/dask/blob/master/dask/distributed/scheduler.py#L73-L76)

I can find some time to do this tomorrow or maybe later this afternoon.

Assuming you don't beat me too it that is :)

haha, i am quite slow, just started looking into the code base today.


Resolving patch:
diff --git a/dask/async.py b/dask/async.py
--- a/dask/async.py
+++ b/dask/async.py
@@ -170,13 +170,16 @@ def start_state_from_dask(dsk, cache=None, sortkey=None):
             cache[k] = v
             data_keys.add(k)
 
-    dependencies = dict((k, get_dependencies(dsk, k)) for k in dsk)
+    dsk2 = dsk.copy()
+    dsk2.update(cache)
+
+    dependencies = dict((k, get_dependencies(dsk2, k)) for k in dsk)
     waiting = dict((k, v.copy()) for k, v in dependencies.items()
                                  if k not in data_keys)
 
     dependents = reverse_dict(dependencies)
-    for a in data_keys:
-        for b in dependents[a]:
+    for a in cache:
+        for b in dependents.get(a, ()):
             waiting[b].remove(a)
     waiting_data = dict((k, v.copy()) for k, v in dependents.items() if v)
 
diff --git a/dask/distributed/client.py b/dask/distributed/client.py
--- a/dask/distributed/client.py
+++ b/dask/distributed/client.py
@@ -37,10 +37,10 @@ def __init__(self, scheduler, address=None):
         self.socket.connect(self.address_to_scheduler)
         self.register_client()
 
-    def get(self, dsk, keys):
+    def get(self, dsk, keys, keep_results=False):
         header = {'function': 'schedule',
                   'jobid': next(jobids)}
-        payload = {'dask': dsk, 'keys': keys}
+        payload = {'dask': dsk, 'keys': keys, 'keep_results': keep_results}
 
         self.send_to_scheduler(header, payload)
         header2, payload2 = self.recv_from_scheduler()
diff --git a/dask/distributed/scheduler.py b/dask/distributed/scheduler.py
--- a/dask/distributed/scheduler.py
+++ b/dask/distributed/scheduler.py
@@ -6,6 +6,7 @@
 import traceback
 import sys
 import random
+from functools import partial
 from collections import defaultdict
 from multiprocessing.pool import ThreadPool
 from datetime import datetime
@@ -22,6 +23,7 @@
     import pickle
 
 from ..core import get_dependencies, flatten
+from ..optimize import cull
 from .. import core
 from ..async import (sortkey, finish_task,
         start_state_from_dask as dag_state_from_dask)
@@ -293,8 +295,8 @@ def send_to_worker(self, address, header, payload):
         header['timestamp'] = datetime.utcnow()
 
         self.send_to_workers_queue.put([address,
-                                            pickle.dumps(header),
-                                            dumps(payload)])
+                                        pickle.dumps(header),
+                                        dumps(payload)])
         self.send_to_workers_send.send(b'')
 
     def send_to_client(self, address, header, result):
@@ -311,20 +313,19 @@ def send_to_client(self, address, header, result):
                                             pickle.dumps(header),
                                             dumps(result)])
 
-    def trigger_task(self, dsk, key, queue):
+    def trigger_task(self, key, task, deps, queue):
         """ Send a single task to the next available worker
 
         See also:
             Scheduler.schedule
             Scheduler.worker_finished_task
         """
-        deps = get_dependencies(dsk, key)
         worker = self.available_workers.get()
         locations = dict((dep, self.who_has[dep]) for dep in deps)
 
         header = {'function': 'compute', 'jobid': key,
                   'dumps': dill.dumps, 'loads': dill.loads}
-        payload = {'key': key, 'task': dsk[key], 'locations': locations,
+        payload = {'key': key, 'task': task, 'locations': locations,
                    'queue': queue}
         self.send_to_worker(worker, header, payload)
 
@@ -570,7 +571,7 @@ def close(self):
         self.block()
         self.context.destroy(linger=3)
 
-    def schedule(self, dsk, result, **kwargs):
+    def schedule(self, dsk, result, keep_results=False, **kwargs):
         """ Execute dask graph against workers
 
         Parameters
@@ -602,15 +603,26 @@ def schedule(self, dsk, result, **kwargs):
                 result_flat = set([result])
             results = set(result_flat)
 
-            cache = dict()
+            for k in self.who_has:  # remove keys that we already know about
+                if self.who_has[k]:
+                    del dsk[k]
+                    if k in results:
+                        results.remove(k)
+
+            dsk = cull(dsk, results)
+
+            preexisting_data = set(k for k, v in self.who_has.items() if v)
+            cache = dict((k, None) for k in preexisting_data)
             dag_state = dag_state_from_dask(dsk, cache=cache)
-            if cache:
-                self.scatter(cache.items())  # send data in dask up to workers
+            del dag_state['cache']
 
-            tick = [0]
+            new_data = dict((k, v) for k, v in cache.items()
+                                   if not (k in self.who_has and
+                                           self.who_has[k]))
+            if new_data:
+                self.scatter(new_data.items())  # send data in dask up to workers
 
-            if dag_state['waiting'] and not dag_state['ready']:
-                raise ValueError("Found no accessible jobs in dask graph")
+            tick = [0]
 
             event_queue = Queue()
             qkey = str(uuid.uuid1())
@@ -624,7 +636,8 @@ def fire_task():
                 dag_state['ready-set'].remove(key)
                 dag_state['running'].add(key)
 
-                self.trigger_task(dsk, key, qkey)  # Fire
+                self.trigger_task(key, dsk[key],
+                        dag_state['dependencies'][key], qkey)  # Fire
 
             try:
                 worker = self.available_workers.get(timeout=20)
@@ -637,6 +650,7 @@ def fire_task():
                 fire_task()
 
             # Main loop, wait on tasks to finish, insert new ones
+            release_data = partial(self._release_data, protected=preexisting_data)
             while dag_state['waiting'] or dag_state['ready'] or dag_state['running']:
                 payload = event_queue.get()
 
@@ -645,14 +659,20 @@ def fire_task():
 
                 key = payload['key']
                 finish_task(dsk, key, dag_state, results, sortkey,
-                            release_data=self._release_data)
+                            release_data=release_data,
+                            delete=key not in preexisting_data)
 
                 while dag_state['ready'] and self.available_workers.qsize() > 0:
                     fire_task()
 
             result2 = self.gather(result)
-            for key in flatten(result):  # release result data from workers
-                self.release_key(key)
+            if not keep_results:  # release result data from workers
+                for key in flatten(result):
+                    if key not in preexisting_data:
+                        self.release_key(key)
+
+            self.cull_redundant_data(3)
+
         return result2
 
     def _schedule_from_client(self, header, payload):
@@ -668,11 +688,12 @@ def _schedule_from_client(self, header, payload):
             address = header['address']
             dsk = payload['dask']
             keys = payload['keys']
+            keep_results = payload.get('keep_results', False)
 
             header2 = {'jobid': header.get('jobid'),
                        'function': 'schedule-ack'}
             try:
-                result = self.schedule(dsk, keys)
+                result = self.schedule(dsk, keys, keep_results)
                 header2['status'] = 'OK'
             except Exception as e:
                 result = e
@@ -681,7 +702,7 @@ def _schedule_from_client(self, header, payload):
             payload2 = {'keys': keys, 'result': result}
             self.send_to_client(address, header2, payload2)
 
-    def _release_data(self, key, state, delete=True):
+    def _release_data(self, key, state, delete=True, protected=()):
         """ Remove data from temporary storage during scheduling run
 
         See Also
@@ -694,7 +715,7 @@ def _release_data(self, key, state, delete=True):
 
         state['released'].add(key)
 
-        if delete:
+        if delete and key not in protected:
             self.release_key(key)
 
     def _set_collection(self, header, payload):
@@ -758,3 +779,23 @@ def prune_and_notify(self, timeout=20):
                 header = {'function': 'worker-death'}
                 payload = {'removed': removed}
                 self.send_to_worker(w_address, header, payload)
+
+    def cull_redundant_data(self, k):
+        """ Remove highly redundant data from workers
+
+        Finds all keys that are replicated more than k times across all workers
+        and releases data from a randomly chosen subset until there are only k
+        workers left holding this data.
+
+        Operates asynchronously and returns quickly.  Scheduler metadata is
+        updated synchronously.
+        """
+        with logerrors():
+            for key, v in self.who_has.items():
+                while len(v) > k:
+                    worker = random.choice(list(v))
+                    header = {'function': 'delitem', 'jobid': key}
+                    payload = {'key': key}
+                    self.send_to_worker(worker, header, payload)
+                    self.who_has[key].remove(worker)
+                    self.worker_has[worker].remove(key)


Test patch:
diff --git a/dask/distributed/tests/test_client.py b/dask/distributed/tests/test_client.py
--- a/dask/distributed/tests/test_client.py
+++ b/dask/distributed/tests/test_client.py
@@ -138,3 +138,12 @@ def test_get_workers():
 
         s.close_workers()
         assert c.get_registered_workers() == {}
+
+
+def test_keep_results():
+    with scheduler_and_workers() as (s, (a, b)):
+        c = Client(s.address_to_clients)
+
+        assert c.get({'x': (inc, 1)}, 'x', keep_results=True) == 2
+
+        assert 'x' in a.data or 'x' in b.data
diff --git a/dask/distributed/tests/test_scheduler.py b/dask/distributed/tests/test_scheduler.py
--- a/dask/distributed/tests/test_scheduler.py
+++ b/dask/distributed/tests/test_scheduler.py
@@ -110,7 +110,7 @@ def test_compute_cycle():
         assert s.available_workers.qsize() == 2
 
         dsk = {'a': (add, 1, 2), 'b': (inc, 'a')}
-        s.trigger_task(dsk, 'a', 'queue-key')
+        s.trigger_task('a', dsk['a'], set([]), 'queue-key')
         sleep(0.1)
 
         assert 'a' in s.who_has
@@ -119,7 +119,7 @@ def test_compute_cycle():
         assert a.address in s.worker_has or b.address in s.worker_has
         assert s.available_workers.qsize() == 2
 
-        s.trigger_task(dsk, 'b', 'queue-key')
+        s.trigger_task('b', dsk['b'], set(['a']), 'queue-key')
         sleep(0.1)
 
         assert 'b' in s.who_has
@@ -307,3 +307,44 @@ def test_monitor_workers():
         while w2.address in s.workers:  # wait to be removed
             sleep(1e-6)
         assert w2.address not in s.workers
+
+
+def test_scheduler_reuses_worker_state():
+    with scheduler_and_workers() as (s, (a, b)):
+        assert s.schedule({'x': (inc, 1)}, 'x') == 2
+
+        s.send_data('x', 10, address=a.address)  # send x to a worker
+        assert a.data['x'] == 10
+        assert s.schedule({'x': (inc, 1)}, 'x') == 10  # recall, don't compute
+        sleep(0.1)
+        assert a.data['x'] == 10
+        assert s.schedule({'x': (inc, 1), 'y': (lambda x: x, 'x')}, 'y') == 10
+
+def test_scheduler_keeps_preexisting_data():
+    with scheduler_and_workers() as (s, (a, b)):
+        s.send_data('x', 10, address=a.address)
+        dsk = {'x': (inc, 1), 'y': (lambda x: x, 'x')}
+        assert s.schedule(dsk, 'y') == 10
+        sleep(0.01)
+        assert a.data['x'] == 10
+
+
+def test_keep_results():
+    with scheduler_and_workers() as (s, (a, b)):
+        assert s.schedule({'x': (inc, 1)}, 'x', keep_results=True) == 2
+        sleep(0.1)
+        assert a.data.get('x') == 2 or b.data.get('x') == 2
+
+
+def test_cull_redundant_data():
+    with scheduler_and_workers() as (s, (a, b)):
+        s.send_data('x', 10, address=a.address)  # send x to a worker
+        s.send_data('x', 10, address=b.address)  # send x to a worker
+
+        s.cull_redundant_data(1)
+
+        while 'x' in a.data and 'x' in b.data:
+            sleep(0.01)
+
+        assert ('x' in a.data and 'x' not in b.data or
+                'x' in b.data and 'x' not in a.data)
diff --git a/dask/tests/test_async.py b/dask/tests/test_async.py
--- a/dask/tests/test_async.py
+++ b/dask/tests/test_async.py
@@ -33,6 +33,13 @@ def test_start_state():
                                 'y': set(['w']),
                                 'z': set(['w'])}}
 
+def test_start_state_looks_at_cache():
+    dsk = {'b': (inc, 'a')}
+    cache = {'a': 1}
+    result = start_state_from_dask(dsk, cache)
+    assert result['dependencies']['b'] == set(['a'])
+    assert result['ready'] == ['b']
+
 
 def test_start_state_with_redirects():
     dsk = {'x': 1, 'y': 'x', 'z': (inc, 'y')}