Repository: dask/dask

Issue:
Dask array with non-highlevelgraph leads to KeyError
I have some [older code](https://github.com/ska-sa/katsdpcal/blob/5caa806649de33324977727a5ca1b5b76d61ba08/katsdpcal/scan.py#L86-L92) that assigns a plain dict directly into `array.dask`, as a workaround for an old dask bug. If that's no longer supported then feel free to close this bug (although it would be nice to add one or two sanity checks). Since dask 2.8.0 it raises a KeyError inside the optimizer.

Here's a minimum reproducing example:
```python
#!/usr/bin/env python3
import dask
import dask.array as da
import numpy as np

a = da.from_array(np.ones((4, 4), np.float64), chunks=(2, 4))
a.dask = dict(a.dask)
b = da.from_array(np.zeros((4, 4), np.float64), chunks=(2, 4))

x = a + b
da.compute(x)
```

```
Traceback (most recent call last):
  File "./bw.py", line 11, in <module>
    da.compute(x)
  File "/home/bmerry/work/sdp/env3/lib/python3.6/site-packages/dask/base.py", line 433, in compute
    dsk = collections_to_dsk(collections, optimize_graph, **kwargs)
  File "/home/bmerry/work/sdp/env3/lib/python3.6/site-packages/dask/base.py", line 219, in collections_to_dsk
    [opt(dsk, keys, **kwargs) for opt, (dsk, keys) in groups.items()],
  File "/home/bmerry/work/sdp/env3/lib/python3.6/site-packages/dask/base.py", line 219, in <listcomp>
    [opt(dsk, keys, **kwargs) for opt, (dsk, keys) in groups.items()],
  File "/home/bmerry/work/sdp/env3/lib/python3.6/site-packages/dask/array/optimization.py", line 43, in optimize
    dsk = fuse_roots(dsk, keys=keys)
  File "/home/bmerry/work/sdp/env3/lib/python3.6/site-packages/dask/blockwise.py", line 822, in fuse_roots
    and not any(dependencies[dep] for dep in deps)  # no need to fuse if 0 or 1
  File "/home/bmerry/work/sdp/env3/lib/python3.6/site-packages/dask/blockwise.py", line 822, in <genexpr>
    and not any(dependencies[dep] for dep in deps)  # no need to fuse if 0 or 1
KeyError: 'array-d6b7f7b33e2b14cd720de87a517147e8'
```
(run against dask 2.10.0, Python 3.6)

From what I can tell, the trouble starts when HighLevelGraph.from_collections is given a dependency that doesn't use a HighLevelGraph, and it names the resulting layer just using `id`: [here](https://github.com/dask/dask/blob/0274a5c6db618e7af8444562149ecc5774ae22a3/dask/highlevelgraph.py#L134). However, the Blockwise for the addition keys its `indices` based on the array names. Then `optimize_blockwise` uses those `indices` keys to construct the new dependencies, [here](https://github.com/dask/dask/blob/0274a5c6db618e7af8444562149ecc5774ae22a3/dask/blockwise.py#L552). That means the new dependencies don't line up with the layers. The final explosion happens in `fuse_roots`, but it could presumably happen anywhere that expected the dependencies to be consistent.


Discussion:

Actually I think this might happen even in cases where one doesn't assign directly to `Array.dask`: I haven't yet found a way to trigger the KeyError or reduce it to a minimal example, but if I sprinkle in some assertions that check for cases of a Blockwise whose indices aren't in the layers then I'm seeing some suspicious cases.


Ignore this - I tracked it down to a bug in some of my code (which constructed a HighLevelGraph but messed up the dependencies).
Thanks for raising @bmerry! 😄

Based on your last comment am trying to figure out if this is issue is still live. WDYT?
> Based on your last comment am trying to figure out if this is issue is still live.

It depends whether arrays whose `.dask` is a plain dict instead of a HighLevelGraph are considered valid. If so, then yes this is a bug, as per the reproduction above. If not, then I think I just need to change my own code to not construct such arrays.

Note that the `Array` constructor coerces its argument to a HighLevelGraph if it isn't currently one, so I think the only way to construct such a plain array is to assign to `.dask`.
Thanks for the pointer. I think it is less that `dict` is permissible and more that we need to use public APIs. If we want to get the underlying graph of an object, we need to use [`__dask_graph__`]( https://docs.dask.org/en/latest/custom-collections.html#__dask_graph__ ). If we are wanting to build an array from a graph, we should build a new `da.Array`. So the following should accomplish roughly what you are looking for...

```python
In [1]: import numpy as np

In [2]: import dask.array as da

In [3]: a = da.from_array(np.ones((4, 4), np.float64), chunks=(2, 4))

In [4]: b = da.from_array(np.zeros((4, 4), np.float64), chunks=(2, 4))

In [5]: c = da.Array(a.__dask_graph__(), a.name, a.chunks, meta=a)

In [6]: x = b + c

In [7]: da.compute(x)
Out[7]: 
(array([[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]]),)
```
> I think it is less that dict is permissible and more that we need to use public APIs.

Sure, but my difficulty is that it's unclear what is considered public API. What makes `a.dask` private but `a.name` and `a.chunks` public? I couldn't find any of them listed in the documentation for [class Array](https://docs.dask.org/en/latest/array-api.html#dask.array.Array).
I think trying to allow non-HLGs as `.dask` attributes on our collections isn't really feasible. Over the next few months I expect more and more places to be doing high-level graph optimization, and trying to handle plain dicts in there sounds painful.

If others agree, I'd suggest making `Array.dask` a property that validates that the value is a HLG. Likely the same for DataFrame and Bag.
It sounds sensible in the long term, but should there be a deprecation period during which assigning to `Array.dask` will automatically convert a dict to a HLG? Auto-converting a property on assignment has a mild code smell, but it would help with backwards compatibility and would be consistent with the behaviour of the constructor.
I'm not sure that you can easily create an HLG from a plain dict. You need additional information about the layers / dependencies (I think).
The Array constructor seems to manage: https://github.com/dask/dask/blob/03aca07be7f69964d290138fd86ce71c8797afe7/dask/array/core.py#L1013-L1014
Mmm, good to know. If that actually works then converting it in the `@dask.setter` seems OK. Are you interested in working on this?
Sure, I can have a go at it next week.
In principle it shouldn't be terribly hard to support non-HLGs.  

I think that the tricky thing here is having users mutate graphs of any sort.  I would rather users create new `Array` objects using the constructor.  

Over time I would also prefer that we remove the `.dask` attribute, preferring that people use the `.__dask_graph__()` method instead laid out in the protocol.  

So rather than do type checking in a `.dask` property, how about adding a deprecation warning?
Should add that I think you are right @bmerry. Documenting `.name` and `.chunks` seems reasonable. Though I could see an argument for a `.__dask_name__()` (thoughts from others?).
@mrocklin what would your preferred fix for this bug be? Regardless of whether we emit a deprecation warning, should Blockwise be made to work with collections build from non-HLG graphs?

> I think that the tricky thing here is having users mutate graphs of any sort.

Would that also mean deprecating the `out` parameter on all the elementwise and reduction functions that currently take them?
> what would your preferred fix for this bug be?

I haven't thought about this problem deeply enough to be able to define a fix.  I'm just emiting some thoughts.

> should Blockwise be made to work with collections build from non-HLG graphs?

Capability is often better than incapability, but ideally we would arrange things so that this situation wouldn't come up.  I'm also not sure exactly how you intend to use `Blockwise` with non-high-level-graphs.  A lot of the value of the Blockwise Mapping class comes from its use in HLGs.

> Would that also mean deprecating the out parameter on all the elementwise and reduction functions that currently take them?

No, because we control what happens here.

> Though I could see an argument for a .__dask_name__()

The closest analog is layers.  The name concept probably doesn't carry over to other collections like delayed or xarray.
> > Though I could see an argument for a .**dask_name**()
> 
> The closest analog is layers. The name concept probably doesn't carry over to other collections like delayed or xarray.

Fair enough. In that case we should document the `name` and `chunks` parameters as both are needed to build a new `Array`.
This is currently laid out here, but I agree that there is nothing in the docstring. 

https://docs.dask.org/en/latest/array-design.html
Good point. @bmerry, where were you looking for this info? The docstring? Somewhere else?
>  I'm also not sure exactly how you intend to use Blockwise with non-high-level-graphs.

I mean that at the moment the code I posted doesn't work, and it traces back to some interaction (which I don't fully understand yet) between Blockwise (implicitly created by `a + b`), the optimizer, and the non-HLG backing `a`. So the options I see are
(a) decide that the code is not supported (because it's mutating things it shouldn't, and/or because it's using a plain dict instead of a HLG); or
(b) make the code work by making assignment to `.dask` automatically create a HLG; or
(c) fixing the interaction so that the code works even through `a.dask` is not a HLG.

If I understand what you're saying, you'd like to do (c) now ("In principle it shouldn't be terribly hard to support non-HLGs"), but also deprecate external code assigning to `.dask`.

> > Would that also mean deprecating the out parameter on all the elementwise and reduction functions that currently take them?

> No, because we control what happens here.

If we're allowing users to mutate array objects indirectly, I would argue for preserving some direct method to mutate an array object. Otherwise it's possible but requires less efficient hacks like this (untested code):
```python
def mutate_array(dsk, name, chunks, dtype, out):
    """Modifies `out` in-place to have the given parameters."""
    add(dask.Array(dsk, name, chunks, dtype), 0, out=out)
```
> Good point. @bmerry, where were you looking for this info? The docstring? Somewhere else?

I was looking on the [API page](https://docs.dask.org/en/latest/array-api.html). I'd missed the references on the design page (unfortunately it's not easy to search for "chunks" and find the object attribute rather than parameter names or general prose). https://docs.dask.org/en/latest/array-design.html is titled "Internal design" and the "internal" might suggest that it is not describing public API. It also documents the `Array.dask` attribute, so if it is describing public API, then that attribute is also public API.

As an aside, that page should probably be updated for the HLGs. I'm prepared to have a go at that in a few weeks time.
Looking at the original case that crash, I suspect the problem is [this line](https://github.com/dask/dask/blob/0274a5c6db618e7af8444562149ecc5774ae22a3/dask/blockwise.py#L552). It is conflating key prefixes (from Blockwise.indices) with layer names. Is there any general expectation that they should be the same?

If I change the affected line to
```python3
            dependencies[layer] = (
                set().union(
                    *(full_graph.dependencies.get(l, set()) for l in blockwise_layers)
                )
                - blockwise_layers
            )
```
(i.e. make the dependencies of the rewritten layer be the union of the original layers' dependencies, but removing the layers themselves) then the case works. However, I'm not sure if that's going to cause performance regressions: `rewrite_blockwise` is a piece of black magic to me, and I don't know if it might be eliminating dependencies which this change will bring back.

If this approach sounds right to people I can expand it into a PR with a regression test.
> Looking at the original case that crash, I suspect the problem is this line. It is conflating key prefixes (from Blockwise.indices) with layer names. Is there any general expectation that they should be the same?

Digging some more, that's not the only place that assumes a layer with keys of the form `(name, indices...)` will be named `name`: see [this line](https://github.com/dask/dask/blob/5f61f7f79b861be5ffdbf59950889166628a8fa6/dask/blockwise.py#L494) (although I think that only depends on Blockwise layers being named like that).

@mrocklin is this assumption intended for HLGs (in which case [this line](https://github.com/dask/dask/blob/0274a5c6db618e7af8444562149ecc5774ae22a3/dask/highlevelgraph.py#L134) is the problem, and it should try to infer a layer name from the graph content), or should the code be able to deal with layer names that are unrelated to the keys in the layer?

Resolving patch:
diff --git a/dask/highlevelgraph.py b/dask/highlevelgraph.py
--- a/dask/highlevelgraph.py
+++ b/dask/highlevelgraph.py
@@ -131,7 +131,10 @@ def from_collections(cls, name, layer, dependencies=()):
                     with ignoring(AttributeError):
                         deps[name] |= set(collection.__dask_layers__())
                 else:
-                    key = id(graph)
+                    try:
+                        [key] = collection.__dask_layers__()
+                    except AttributeError:
+                        key = id(graph)
                     layers[key] = graph
                     deps[name].add(key)
                     deps[key] = set()


Test patch:
diff --git a/dask/array/tests/test_atop.py b/dask/array/tests/test_atop.py
--- a/dask/array/tests/test_atop.py
+++ b/dask/array/tests/test_atop.py
@@ -587,3 +587,12 @@ def test_atop_legacy():
     z = da.blockwise(inc, "i", x, "i", dtype=x.dtype)
     assert_eq(y, z)
     assert y.name == z.name
+
+
+def test_non_hlg():
+    # Regression test for https://github.com/dask/dask/issues/5850
+    a = da.from_array(np.ones(1, np.float64), chunks=(1,))
+    a.dask = dict(a.dask)  # Convert from HighLevelGraph to plain dict
+    b = da.from_array(np.zeros(1, np.float64), chunks=(1,))
+    x = a + b
+    assert_eq(x, a)