Repository: dask/dask

Issue:
KeyError when trying to compute/persist multiple keys
Hi, having this code:
```
import dask
from dask.bag import from_sequence, zip as bag_zip

a = from_sequence(range(10), npartitions=2)
b = from_sequence(range(10), npartitions=2)
c = bag_zip(a, b)
d1 = c.pluck(0)
d2 = c.pluck(1)
e = d1.map(lambda x: x+1)

dask.persist(d1,d2,e)
```
I'm getting this error:
```
Traceback (most recent call last):
  File "key_error.py", line 11, in <module>
    dask.compute(d1,d2,e)
  File "/home/cvb/src/dask/dask/base.py", line 398, in compute
    results = schedule(dsk, keys, **kwargs)
  File "/home/cvb/src/dask/dask/multiprocessing.py", line 173, in get
    dsk2, dependencies = cull(dsk, keys)
  File "/home/cvb/src/dask/dask/optimization.py", line 46, in cull
    for k in work]
  File "/home/cvb/src/dask/dask/optimization.py", line 46, in <listcomp>
    for k in work]
  File "/home/cvb/src/dask/dask/core.py", line 187, in get_dependencies
    arg = dsk[key]
KeyError: ('pluck-ddca17d615f33c398bce4e491736fb48', 0)
```
Can reproduce this with 1.1.0 version and with current master branch. Having same result when using `dask.compute` instead of `dask.persist`.


Discussion:
Thanks for the excellent bug report and reproducible example @cvb .  I'm able to reproduce it.

@eriknw any time and interest in looking into this?

Resolving patch:
diff --git a/dask/bag/core.py b/dask/bag/core.py
--- a/dask/bag/core.py
+++ b/dask/bag/core.py
@@ -33,7 +33,7 @@
 from ..bytes import open_files
 from ..compatibility import apply, urlopen, Iterable, Iterator
 from ..context import globalmethod
-from ..core import quote, istask, get_dependencies, reverse_dict
+from ..core import quote, istask, get_dependencies, reverse_dict, flatten
 from ..delayed import Delayed, unpack_collections
 from ..highlevelgraph import HighLevelGraph
 from ..multiprocessing import get as mpget
@@ -85,7 +85,7 @@ def lazify(dsk):
     return valmap(lazify_task, dsk)
 
 
-def inline_singleton_lists(dsk, dependencies=None):
+def inline_singleton_lists(dsk, keys, dependencies=None):
     """ Inline lists that are only used once.
 
     >>> d = {'b': (list, 'a'),
@@ -100,10 +100,11 @@ def inline_singleton_lists(dsk, dependencies=None):
                         for k, v in dsk.items()}
     dependents = reverse_dict(dependencies)
 
-    keys = [k for k, v in dsk.items()
-            if istask(v) and v and v[0] is list and len(dependents[k]) == 1]
-    dsk = inline(dsk, keys, inline_constants=False)
-    for k in keys:
+    inline_keys = {k for k, v in dsk.items()
+                   if istask(v) and v and v[0] is list and len(dependents[k]) == 1}
+    inline_keys.difference_update(flatten(keys))
+    dsk = inline(dsk, inline_keys, inline_constants=False)
+    for k in inline_keys:
         del dsk[k]
     return dsk
 
@@ -114,7 +115,7 @@ def optimize(dsk, keys, fuse_keys=None, rename_fused_keys=True, **kwargs):
     dsk2, dependencies = cull(dsk, keys)
     dsk3, dependencies = fuse(dsk2, keys + (fuse_keys or []), dependencies,
                               rename_keys=rename_fused_keys)
-    dsk4 = inline_singleton_lists(dsk3, dependencies)
+    dsk4 = inline_singleton_lists(dsk3, keys, dependencies)
     dsk5 = lazify(dsk4)
     return dsk5
 


Test patch:
diff --git a/dask/bag/tests/test_bag.py b/dask/bag/tests/test_bag.py
--- a/dask/bag/tests/test_bag.py
+++ b/dask/bag/tests/test_bag.py
@@ -502,18 +502,24 @@ def test_inline_singleton_lists():
     inp = {'b': (list, 'a'),
            'c': (f, 'b', 1)}
     out = {'c': (f, (list, 'a'), 1)}
-    assert inline_singleton_lists(inp) == out
+    assert inline_singleton_lists(inp, ['c']) == out
 
     out = {'c': (f, 'a', 1)}
     assert optimize(inp, ['c'], rename_fused_keys=False) == out
 
+    # If list is an output key, don't fuse it
+    assert inline_singleton_lists(inp, ['b', 'c']) == inp
+    assert optimize(inp, ['b', 'c'], rename_fused_keys=False) == inp
+
     inp = {'b': (list, 'a'),
            'c': (f, 'b', 1),
            'd': (f, 'b', 2)}
-    assert inline_singleton_lists(inp) == inp
+    assert inline_singleton_lists(inp, ['c', 'd']) == inp
 
-    inp = {'b': (4, 5)} # doesn't inline constants
-    assert inline_singleton_lists(inp) == inp
+    # Doesn't inline constants
+    inp = {'b': (4, 5),
+           'c': (f, 'b')}
+    assert inline_singleton_lists(inp, ['c']) == inp
 
 
 def test_take():