Repository: dask/dask

Issue:
Profiler fails with KeyError with dask.bag + dask.do
Here's a snippet that reproduces the issue

``` python
import dask.bag as db
from dask import do
from dask.diagnostics import Profiler


def inc(x):
    return x + 1


@do
def dec(seq):
    return [x - 1 for x in seq]

bag = dec(db.from_sequence([1, 2, 3]).map(inc))

result = bag.compute()
assert result == [1, 2, 3]  # works

with Profiler() as prof:
    result = bag.compute()  # somehow fails because of the profiler
```

I suspect this is related to using both thread and process based schedulers



Discussion:
This is also fairly strange behavior.  You're wrapping one dask thing with another.

I suggest that dask calling dask is something of an anti-pattern.  You really want to merge graphs, not nest them.

happy to chalk it up to user error, should i close? 

I'm curious what the behavior was that lead to this error.  The fact that you got to this state is an error in documentation.

The reason the profiler errors is because it's recording all computations performed by the scheduler, but then only working with the outermost dask. Definitely user error, not sure how we'd deal with this otherwise.

In the future I plan to make imperative play nice with the collections, so then doing it this way should work as you expected. Until then though we should update the docs to say not to mix collections.


Resolving patch:
diff --git a/dask/bag/core.py b/dask/bag/core.py
--- a/dask/bag/core.py
+++ b/dask/bag/core.py
@@ -174,7 +174,7 @@ def to_textfiles(b, path, name_function=str):
 
 def finalize(bag, results):
     if isinstance(bag, Item):
-        return results
+        return results[0]
     if isinstance(results, Iterator):
         results = list(results)
     if isinstance(results[0], Iterable) and not isinstance(results[0], str):
@@ -194,7 +194,7 @@ def __init__(self, dsk, key):
         self.key = key
 
     def _keys(self):
-        return self.key
+        return [self.key]
 
     def apply(self, func):
         name = next(names)
diff --git a/dask/imperative.py b/dask/imperative.py
--- a/dask/imperative.py
+++ b/dask/imperative.py
@@ -6,7 +6,8 @@
 from toolz import merge, unique, curry
 
 from .optimize import cull, fuse
-from .context import _globals
+from .utils import concrete
+from . import base
 from .compatibility import apply
 from . import threaded
 
@@ -63,6 +64,12 @@ def to_task_dasks(expr):
     """
     if isinstance(expr, Value):
         return expr.key, expr._dasks
+    elif isinstance(expr, base.Base):
+        name = tokenize(str(expr), True)
+        keys = expr._keys()
+        dsk = expr._optimize(expr.dask, keys)
+        dsk[name] = (expr._finalize, expr, (concrete, keys))
+        return name, [dsk]
     elif isinstance(expr, (Iterator, list, tuple, set)):
         args, dasks = unzip(map(to_task_dasks, expr), 2)
         args = list(args)
@@ -176,16 +183,13 @@ def optimize(dsk, keys):
     return fuse(dsk2)
 
 
-def get(dsk, keys, get=None, **kwargs):
-    """Specialized get function"""
-    get = get or _globals['get'] or threaded.get
-    dsk2 = optimize(dsk, keys)
-    return get(dsk2, keys, **kwargs)
-
-
 def compute(*args, **kwargs):
     """Evaluate several ``Value``s at once.
 
+    Note that the only difference between this function and
+    ``dask.base.compute`` is that this implicitly converts python objects to
+    ``Value``s, allowing for collections of dask objects to be computed.
+
     Examples
     --------
     >>> a = value(1)
@@ -193,13 +197,11 @@ def compute(*args, **kwargs):
     >>> c = a + 3
     >>> compute(b, c)  # Compute both simultaneously
     (3, 4)
-    >>> compute(a, [b, c])
+    >>> compute(a, [b, c])  # Works for lists of Values
     (1, [3, 4])
     """
     args = [value(a) for a in args]
-    dsk = merge(*[arg.dask for arg in args])
-    keys = [arg.key for arg in args]
-    return tuple(get(dsk, keys, **kwargs))
+    return base.compute(*args, **kwargs)
 
 
 def right(method):
@@ -209,21 +211,20 @@ def _inner(self, other):
     return _inner
 
 
-class Value(object):
+class Value(base.Base):
     """Represents a value to be computed by dask.
 
     Equivalent to the output from a single key in a dask graph.
     """
     __slots__ = ('_key', '_dasks')
+    _optimize = staticmethod(optimize)
+    _finalize = staticmethod(lambda a, r: r[0])
+    _default_get = staticmethod(threaded.get)
 
     def __init__(self, name, dasks):
         object.__setattr__(self, '_key', name)
         object.__setattr__(self, '_dasks', dasks)
 
-    def compute(self, **kwargs):
-        """Compute the result."""
-        return get(self.dask, self.key, **kwargs)
-
     @property
     def dask(self):
         return merge(*self._dasks)
@@ -232,13 +233,8 @@ def dask(self):
     def key(self):
         return self._key
 
-    def visualize(self, optimize_graph=False, **kwargs):
-        """Visualize the dask as a graph"""
-        from dask.dot import dot_graph
-        if optimize_graph:
-            return dot_graph(optimize(self.dask, self.key), **kwargs)
-        else:
-            return dot_graph(self.dask, **kwargs)
+    def _keys(self):
+        return [self.key]
 
     def __repr__(self):
         return "Value({0})".format(repr(self.key))


Test patch:
diff --git a/dask/tests/test_imperative.py b/dask/tests/test_imperative.py
--- a/dask/tests/test_imperative.py
+++ b/dask/tests/test_imperative.py
@@ -2,7 +2,9 @@
 from collections import Iterator
 from random import random
 
-from dask.imperative import value, do, to_task_dasks, compute
+import pytest
+
+from dask.imperative import value, do, to_task_dasks, compute, Value
 from dask.utils import raises
 
 
@@ -154,3 +156,39 @@ def test_pure():
 
     myrand = do(random)
     assert myrand().key != myrand().key
+
+
+da = pytest.importorskip('dask.array')
+import numpy as np
+
+
+def test_array_imperative():
+    arr = np.arange(100).reshape((10, 10))
+    darr = da.from_array(arr, chunks=(5, 5))
+    val = do(sum)([arr, darr, 1])
+    assert isinstance(val, Value)
+    assert np.allclose(val.compute(), arr + arr + 1)
+    assert val.sum().compute() == (arr + arr + 1).sum()
+    assert val[0, 0].compute() == (arr + arr + 1)[0, 0]
+
+    task, dasks = to_task_dasks(darr)
+    assert len(dasks) == 1
+    orig = set(darr.dask)
+    final = set(dasks[0])
+    assert orig.issubset(final)
+    diff = final.difference(orig)
+    assert len(diff) == 1
+
+
+db = pytest.importorskip('dask.bag')
+
+
+def test_array_bag_imperative():
+    arr1 = np.arange(100).reshape((10, 10))
+    arr2 = arr1.dot(arr1.T)
+    darr1 = da.from_array(arr1, chunks=(5, 5))
+    darr2 = da.from_array(arr2, chunks=(5, 5))
+    b = db.from_sequence([1, 2, 3])
+    seq = [arr1, arr2, darr1, darr2, b]
+    out = do(sum)([i.sum() for i in seq])
+    assert out.compute() == 2*arr1.sum() + 2*arr2.sum() + sum([1, 2, 3])