Repository: dask/dask
Pull request: Support `traverse` keyword in persist/optimize
Adds support for `traverse` keyword to `persist`/`optimize`, and changes how this is implemented for `compute`.

The general gist is that we wrap all collection transforms in `unpack`/`repack` calls. The pattern is:

```
collections, repack = unpack(*args)  # extract all collections from arguments
results = transform(collections)     # do something with the collections
return repack(results)               # repackage the results in python collections
```

Supersedes #3404, fixes dask/distributed#1918.

Problem statement:
Progressbar on list of delayed objects 
I have a list of delayed objects in "scores". Following the documentation, I should be able to do 
```
from dask import compute, persist
from dask.distributed import progress
...
x = persist(scores)  # start computation in the background
progress(x)  # watch progress
scores = compute(x[0])[0]
```

The problem is that the progress bar in the terminal shows a 100% bar immediately. I can observe with the dashboard that the computation is still going for about 10 minutes. 

Am I doing something wrong here? The docs are not explicit for the list of delayed objects.

Hints from issue discussion:
Perhaps they are already in memory?  If you call x[0].compute() does it
return immediately?

On Fri, Apr 13, 2018 at 12:57 PM, Alvaro Ulloa <notifications@github.com>
wrote:

> I have a list of delayed objects in "scores". Following the documentation,
> I should be able to do
>
> from dask import compute, persist
> from dask.distributed import progress
> ...
> x = persist(scores)  # start computation in the background
> progress(x)  # watch progress
> scores = compute(x[0])[0]
>
> The problem is that the progress bar in the terminal shows a 100% bar
> immediately. I can observe with the dashboard that the computation is still
> going for about 10 minutes.
>
> Am I doing something wrong here? The docs are not explicit for the list of
> delayed objects.
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/distributed/issues/1918>, or mute the thread
> <https://github.com/notifications/unsubscribe-auth/AASszCoj2BZz0coFSGjGA7_6zFhYVz8nks5toOcegaJpZM4TT2Gt>
> .
>

no, it takes about 10 minutes to return. 
Can I ask you for a reproducible example?  Just something that someone can
copy-paste to reproduce the problem.  Perhaps you could construct a
function that uses time.sleep rather than your actual computation.

On Fri, Apr 13, 2018 at 2:23 PM, Alvaro Ulloa <notifications@github.com>
wrote:

> no, it takes about 10 minutes to return.
>
> —
> You are receiving this because you commented.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/distributed/issues/1918#issuecomment-381237678>,
> or mute the thread
> <https://github.com/notifications/unsubscribe-auth/AASszBT6ntwPDCg69z_jiUm-MnZvuMcwks5toPtGgaJpZM4TT2Gt>
> .
>

```
from dask import compute, persist, delayed
from dask.distributed import progress, Client
from time import sleep


@delayed
def task(delay):
    sleep(delay)
    return delay


if __name__ == '__main__':
    client = Client('20670:8786')  # Change accordingly
    scores = [task(x) for x in range(1, 20)]
    x = persist(scores)  # start computation in the background
    progress(x)  # watch progress
    scores = compute(x)  # wait until completed
    print('Done')

```

Just for the record, this seems to work but it is ugly

```
from dask import compute, persist, delayed
from dask.distributed import progress, Client
from time import sleep


@delayed
def task(delay):
    sleep(delay)
    return delay


if __name__ == '__main__':
    client = Client('x20670:8786')  # Change accordingly
    scores = [task(x) for x in range(1, 20)]

    x = delayed(compute)(scores)
    x = persist(x)  # start computation in the background
    progress(x)  # watch progress
    scores = compute(x)

    print(scores[0][0][0])
    print('Done')

```

is there a better way?
In this case I think you want the following:

    x = persist(*scores)  # start computation in the background

Persist takes many arguments.  It currently doesn't dive into each element of a list, though perhaps it should.
@martindurant if you're around an interested this would be useful to solve.  The change probably needs to be made in `distributed/client.py:Client.persist`.  We should verify that this *would* work with the local scheduler as well (I suspect that it does).
@alvarouc thank you for the reproducible example by the way!  It was very helpful to quickly identify the problem.
The error appears to be in `dask.persist`; if you call `persist(scores)`, this does not call `client.persist` at all, but `persist(*scores)` does. `client.persist` works with either `scores` or `*scores`.

Commit message: Add traverse keyword to optimize

Also support task literals in both graph and results

Patch to review:
diff --git a/dask/base.py b/dask/base.py
index e129fadfe07..015c9caabff 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -4,6 +4,7 @@
 from collections import OrderedDict, Iterator
 from functools import partial
 from hashlib import md5
+from operator import getitem
 import inspect
 import pickle
 import os
@@ -16,7 +17,7 @@
 
 from .compatibility import long, unicode
 from .context import _globals, thread_state
-from .core import flatten
+from .core import flatten, quote, get as simple_get
 from .hashing import hash_buffer_hex
 from .utils import Dispatch, ensure_dict
 
@@ -125,7 +126,7 @@ def persist(self, **kwargs):
         --------
         dask.base.persist
         """
-        (result,) = persist(self, **kwargs)
+        (result,) = persist(self, traverse=False, **kwargs)
         return result
 
     def compute(self, **kwargs):
@@ -281,6 +282,74 @@ def _extract_graph_and_keys(vals):
     return dsk, keys
 
 
+def unpack_collections(*args, **kwargs):
+    """Extract collections in preparation for compute/persist/etc...
+
+    Intended use is to find all collections in a set of (possibly nested)
+    python objects, do something to them (compute, etc...), then repackage them
+    in equivalent python objects.
+
+    Parameters
+    ----------
+    *args
+        Any number of objects. If it is a dask collection, it's extracted and
+        added to the list of collections returned. By default, python builtin
+        collections are also traversed to look for dask collections (for more
+        information see the ``traverse`` keyword).
+    traverse : bool, optional
+        If True (default), builtin python collections are traversed looking for
+        any dask collections they might contain.
+
+    Returns
+    -------
+    collections : list
+        A list of all dask collections contained in ``args``
+    repack : callable
+        A function to call on the transformed collections to repackage them as
+        they were in the original ``args``.
+    """
+    traverse = kwargs.pop('traverse', True)
+
+    collections = []
+    repack_dsk = {}
+
+    def _unpack(expr):
+        if is_dask_collection(expr):
+            tok = tokenize(expr)
+            if tok not in repack_dsk:
+                repack_dsk[tok] = (getitem, 'collections', len(collections))
+                collections.append(expr)
+            return tok
+
+        tok = uuid.uuid4().hex
+        if not traverse:
+            tsk = quote(expr)
+        else:
+            # Treat iterators like lists
+            typ = list if isinstance(expr, Iterator) else type(expr)
+
+            if typ in (list, tuple, set):
+                tsk = (typ, [_unpack(i) for i in expr])
+            elif typ is dict:
+                tsk = (dict, [[_unpack(k), _unpack(v)]
+                              for k, v in expr.items()])
+            else:
+                return expr
+
+        repack_dsk[tok] = tsk
+        return tok
+
+    out = uuid.uuid4().hex
+    repack_dsk[out] = (tuple, [_unpack(i) for i in args])
+
+    def repack(results):
+        dsk = repack_dsk.copy()
+        dsk['collections'] = quote(results)
+        return simple_get(dsk, out)
+
+    return collections, repack
+
+
 def optimize(*args, **kwargs):
     """Optimize several dask collections at once.
 
@@ -298,6 +367,11 @@ def optimize(*args, **kwargs):
         merged with all those of all other dask objects before returning an
         equivalent dask collection. Non-dask arguments are passed through
         unchanged.
+    traverse : bool, optional
+        By default dask traverses builtin python collections looking for dask
+        objects passed to ``optimize``. For large collections this can be
+        expensive. If none of the arguments contain any dask objects, set
+        ``traverse=False`` to avoid doing this traversal.
     optimizations : list of callables, optional
         Additional optimization passes to perform.
     **kwargs
@@ -315,15 +389,20 @@ def optimize(*args, **kwargs):
     >>> b2.compute() == b.compute()
     True
     """
-    variables = [a for a in args if is_dask_collection(a)]
-    if not variables:
+    collections, repack = unpack_collections(*args, **kwargs)
+    if not collections:
         return args
 
-    dsk = collections_to_dsk(variables, **kwargs)
+    dsk = collections_to_dsk(collections, **kwargs)
     postpersists = [a.__dask_postpersist__() if is_dask_collection(a)
                     else (None, a) for a in args]
 
-    return tuple(a if f is None else f(dsk, *a) for f, a in postpersists)
+    keys, postpersists = [], []
+    for a in collections:
+        keys.extend(flatten(a.__dask_keys__()))
+        postpersists.append(a.__dask_postpersist__())
+
+    return repack([r(dsk, *s) for r, s in postpersists])
 
 
 # TODO: remove after deprecation cycle of `dask.optimize` module completes
@@ -371,40 +450,31 @@ def compute(*args, **kwargs):
     >>> compute({'a': a, 'b': b, 'c': 1})  # doctest: +SKIP
     ({'a': 45, 'b': 4.5, 'c': 1},)
     """
-    from dask.delayed import delayed
     traverse = kwargs.pop('traverse', True)
-    if traverse:
-        args = tuple(delayed(a)
-                     if isinstance(a, (list, set, tuple, dict, Iterator))
-                     else a for a in args)
-
     optimize_graph = kwargs.pop('optimize_graph', True)
-    variables = [a for a in args if is_dask_collection(a)]
-    if not variables:
-        return args
-
     get = kwargs.pop('get', None) or _globals['get']
 
+    collections, repack = unpack_collections(*args, traverse=traverse)
+    if not collections:
+        return args
+
     if get is None and getattr(thread_state, 'key', False):
         from distributed.worker import get_worker
         get = get_worker().client.get
 
     if not get:
-        get = variables[0].__dask_scheduler__
-        if not all(a.__dask_scheduler__ == get for a in variables):
+        get = collections[0].__dask_scheduler__
+        if not all(a.__dask_scheduler__ == get for a in collections):
             raise ValueError("Compute called on multiple collections with "
                              "differing default schedulers. Please specify a "
                              "scheduler `get` function using either "
                              "the `get` kwarg or globally with `set_options`.")
 
-    dsk = collections_to_dsk(variables, optimize_graph, **kwargs)
-    keys = [var.__dask_keys__() for var in variables]
-    postcomputes = [a.__dask_postcompute__() if is_dask_collection(a)
-                    else (None, a) for a in args]
+    dsk = collections_to_dsk(collections, optimize_graph, **kwargs)
+    keys = [x.__dask_keys__() for x in collections]
+    postcomputes = [x.__dask_postcompute__() for x in collections]
     results = get(dsk, keys, **kwargs)
-    results_iter = iter(results)
-    return tuple(a if f is None else f(next(results_iter), *a)
-                 for f, a in postcomputes)
+    return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
 
 
 def visualize(*args, **kwargs):
@@ -535,6 +605,11 @@ def persist(*args, **kwargs):
         A scheduler ``get`` function to use. If not provided, the default
         is to check the global settings first, and then fall back to
         the collection defaults.
+    traverse : bool, optional
+        By default dask traverses builtin python collections looking for dask
+        objects passed to ``persist``. For large collections this can be
+        expensive. If none of the arguments contain any dask objects, set
+        ``traverse=False`` to avoid doing this traversal.
     optimize_graph : bool, optional
         If True [default], the graph is optimized before computation.
         Otherwise the graph is run as is. This can be useful for debugging.
@@ -545,12 +620,14 @@ def persist(*args, **kwargs):
     -------
     New dask collections backed by in-memory data
     """
-    collections = [a for a in args if is_dask_collection(a)]
+    traverse = kwargs.pop('traverse', True)
+    optimize_graph = kwargs.pop('optimize_graph', True)
+    get = kwargs.pop('get', None) or _globals['get']
+
+    collections, repack = unpack_collections(*args, traverse=traverse)
     if not collections:
         return args
 
-    get = kwargs.pop('get', None) or _globals['get']
-
     if get is None and getattr(thread_state, 'key', False):
         from distributed.worker import get_worker
         get = get_worker().client.get
@@ -577,8 +654,6 @@ def persist(*args, **kwargs):
                                  else next(results_iter)
                                  for a in args)
 
-    optimize_graph = kwargs.pop('optimize_graph', True)
-
     if not get:
         get = collections[0].__dask_scheduler__
         if not all(a.__dask_scheduler__ == get for a in collections):
@@ -588,21 +663,17 @@ def persist(*args, **kwargs):
                              "the `get` kwarg or globally with `set_options`.")
 
     dsk = collections_to_dsk(collections, optimize_graph, **kwargs)
-
     keys, postpersists = [], []
-    for a in args:
-        if is_dask_collection(a):
-            a_keys = list(flatten(a.__dask_keys__()))
-            rebuild, state = a.__dask_postpersist__()
-            keys.extend(a_keys)
-            postpersists.append((rebuild, a_keys, state))
-        else:
-            postpersists.append((None, None, a))
+    for a in collections:
+        a_keys = list(flatten(a.__dask_keys__()))
+        rebuild, state = a.__dask_postpersist__()
+        keys.extend(a_keys)
+        postpersists.append((rebuild, a_keys, state))
 
     results = get(dsk, keys, **kwargs)
     d = dict(zip(keys, results))
-    return tuple(s if r is None else r({k: d[k] for k in ks}, *s)
-                 for r, ks, s in postpersists)
+    results2 = [r({k: d[k] for k in ks}, *s) for r, ks, s in postpersists]
+    return repack(results2)
 
 
 ############
diff --git a/dask/tests/test_base.py b/dask/tests/test_base.py
index c0a21b86304..efc22d193e6 100644
--- a/dask/tests/test_base.py
+++ b/dask/tests/test_base.py
@@ -11,7 +11,7 @@
 from dask import delayed
 from dask.base import (compute, tokenize, normalize_token, normalize_function,
                        visualize, persist, function_cache, is_dask_collection,
-                       DaskMethodsMixin, optimize)
+                       DaskMethodsMixin, optimize, unpack_collections)
 from dask.delayed import Delayed
 from dask.utils import tmpdir, tmpfile, ignoring
 from dask.utils_test import inc, dec
@@ -342,6 +342,51 @@ def __dask_graph__(self):
     assert not is_dask_collection(DummyCollection)
 
 
+def test_unpack_collections():
+    a = delayed(1) + 5
+    b = a + 1
+    c = a + 2
+
+    def build(a, b, c, iterator):
+        return (a, b,               # Top-level collections
+                {'a': a,            # dict
+                 a: b,              # collections as keys
+                 'b': [1, 2, b],    # list
+                 'c': 10,           # other builtins pass through unchanged
+                 'd': (c, 2),       # tuple
+                 'e': {a, 2, 3}},   # set
+                iterator)           # Iterator
+
+    args = build(a, b, c, (i for i in [a, b, c]))
+
+    collections, repack = unpack_collections(*args)
+    assert len(collections) == 3
+
+    # Replace collections with `'~a'` strings
+    result = repack(['~a', '~b', '~c'])
+    sol = build('~a', '~b', '~c', ['~a', '~b', '~c'])
+    assert result == sol
+
+    # traverse=False
+    collections, repack = unpack_collections(*args, traverse=False)
+    assert len(collections) == 2  # just a and b
+    assert repack(collections) == args
+
+    # No collections
+    collections, repack = unpack_collections(1, 2, {'a': 3})
+    assert not collections
+    assert repack(collections) == (1, 2, {'a': 3})
+
+    # Result that looks like a task
+    def fail(*args):
+        raise ValueError("Shouldn't have been called")
+
+    collections, repack = unpack_collections(a, (fail, 1), [(fail, 2, 3)],
+                                             traverse=False)
+    repack(collections)  # Smoketest task literals
+    repack([(fail, 1)])  # Smoketest results that look like tasks
+
+
 class Tuple(DaskMethodsMixin):
     __slots__ = ('_dask', '_keys')
     __dask_scheduler__ = staticmethod(dask.threaded.get)
@@ -623,6 +668,29 @@ def test_optimize():
         assert dict(a.dask) == dict(b.dask)
 
 
+def test_optimize_nested():
+    a = dask.delayed(inc)(1)
+    b = dask.delayed(inc)(a)
+    c = a + b
+
+    result = optimize({'a': a, 'b': [1, 2, b]}, (c, 2))
+
+    a2 = result[0]['a']
+    b2 = result[0]['b'][2]
+    c2 = result[1][0]
+
+    assert isinstance(a2, Delayed)
+    assert isinstance(b2, Delayed)
+    assert isinstance(c2, Delayed)
+    assert dict(a2.dask) == dict(b2.dask) == dict(c2.dask)
+    assert compute(*result) == ({'a': 2, 'b': [1, 2, 3]}, (5, 2))
+
+    res = optimize([a, b], c, traverse=False)
+    assert res[0][0] is a
+    assert res[0][1] is b
+    assert res[1].compute() == 5
+
+
 # TODO: remove after deprecation cycle of `dask.optimize` module is completed
 def test_optimize_has_deprecated_module_functions_as_attributes():
     import dask.optimize as deprecated_optimize
@@ -657,6 +725,22 @@ def test_persist_literals():
     assert persist(1, 2, 3) == (1, 2, 3)
 
 
+def test_persist_nested():
+    a = delayed(1) + 5
+    b = a + 1
+    c = a + 2
+    result = persist({'a': a, 'b': [1, 2, b]}, (c, 2))
+    assert isinstance(result[0]['a'], Delayed)
+    assert isinstance(result[0]['b'][2], Delayed)
+    assert isinstance(result[1][0], Delayed)
+    assert compute(*result) == ({'a': 6, 'b': [1, 2, 7]}, (8, 2))
+
+    res = persist([a, b], c, traverse=False)
+    assert res[0][0] is a
+    assert res[0][1] is b
+    assert res[1].compute() == 8
+
+
 def test_persist_delayed():
     x1 = delayed(1)
     x2 = delayed(inc)(x1)


Code review comments:
diff_hunk: @@ -342,6 +342,51 @@ def __dask_graph__(self):
     assert not is_dask_collection(DummyCollection)
 
 
+def test_unpack_collections():
+    a = delayed(1) + 5
+    b = a + 1
+    c = a + 2
+
+    def build(a, b, c, iterator):
+        return (a, b,               # Top-level collections
+                {'a': a,            # dict
+                 a: b,              # collections as keys
+                 'b': [1, 2, b],    # list
path: dask/tests/test_base.py
text: @user1:
Can I ask for a nested list, perhaps just by subbing `b -> [b]`

@author:
Done.