Repository: dask/dask
Pull request: Add colors to represent high level layer types
- [ ] Closes #7919
- [x] Tests passed
- [x] Passes `black dask` / `flake8 dask` / `isort dask`


In this PR, I have added an extra attribute to the graphviz output of High-Level graphs - node fill color. Nodes are colored on the basis of their `layer_type`.

This was achieved by using DOT's `fillcolor` attribute.

Currently, this option is made optional to the users. They need to provide an additional kwarg while calling `dask.visualize()` with `color="layer_type"`.

## Demo

```py
import dask
import dask.dataframe as dd

df = dask.datasets.timeseries()
df2 = df[df.y > 0]
df3 = df2.groupby('name').x.std()
```

### `c.dask.visualize(color="layer_type")`
![7974-AA](https://user-images.githubusercontent.com/62539811/129655948-004df074-7449-4725-8fe6-811ee050289e.png)

### `c.dask.visualize(color="layer_type")`
![image](https://user-images.githubusercontent.com/62539811/129747767-145a5092-8ee7-4414-a87f-18e50a8e7c2a.png)


## Color Scheme

```py
layer_colors = {
    "DataFrameIOLayer": "purple",
    "ShuffleLayer": "rose",
    "SimpleShuffleLayer": "rose",
    "ArrayOverlayLayer": "pink",
    "BroadcastJoinLayer": "blue",
    "Blockwise": "green",
    "BlockwiseLayer": "green",
    "BlockwiseCreateArray": "green",
    "MaterializedLayer": "gray",
}
```



![final](https://user-images.githubusercontent.com/62539811/129032590-422ecfd4-1007-43fe-8fb6-d257bb62b95f.png)

## Explanation

- `DataFrameIOLayer`: inefficient;
- `ShuffleLayer`, `SimpleShuffleLayer`: inefficient;
- `ArrayOverlayLayer`: inefficient;
- `BroadcastJoinLayer`: (?);
- `Blockwise`, `BlockwiseLayer`, `BlockwiseCreateArray`: efficient; easy to parallize;
- `MaterializedLayer`: inefficient; better to materialize as late as possible;

#### Key Points

- Blockwise Layers are more efficient than others and can be readily parallelized. As a result, they are **green** (a color used to signify something that is right, something that is correct). When users see a green layer, they may be certain that it is the most efficient method to accomplish things and that no optimization is required.
- The **gray** color (which signifies neutrality and balance) is used for Materialized layers to indicate that the layer should be materialized as late as feasible. Since we are not sure of the optimal way to optimize , we use the **gray** color to indicate that we are not sure.
- All of the other layers are inefficient in some way. As a result, they are colored brighter and stronger to attract the user's attention. Colors such as **purple**, **pink**, and **blue** are used to indicate to the user that something needs to be optimized. Users should utilize as little of these levels as feasible.


Problem statement:
Categorise high level layers for display
There are currently a concrete number of subclasses of the base highlevelgraph.Layer. Some of these have specific contexts or collection linkage (array Vs dataframe), others do not. For the sake of the work being done by @freyam , it would be nice to create some categories for the purpose of being shown in `.visualize()`.

Layers allow for attaching attributes at instantiation. I suggests there might also be class attributes giving information about the layer type, which will be true for all instances.

Example: `DataFrameIOLayer` is IO by operation type and dataframe by collection.  It would be reasonable for these to be among the default annotations of all instances.

Hints from issue discussion:
cc @GenevieveBuckley 
This is amazing! I will get started on finding more about this and share my findings here! :rocket: 
I suggest starting with a concrete list of the current known implementations of Layer.
# List of Layers

## General
```py
# High level graph layer
class Layer(collections.abc.Mapping)

# Fully materialized layer of `Layer`
class MaterializedLayer(Layer)

# Tensor Operation
class Blockwise(Layer)
```

## Array Layers

```py
# Specialized Blockwise Layer for array creation routines
class BlockwiseCreateArray(Blockwise)

# Simple HighLevelGraph array overlap layer
class ArrayOverlapLayer(Layer)
```

## Dataframe Layers

```py
# DataFrame-based HighLevelGraph Layer
class DataFrameLayer(Layer)

# High-level graph layer for a simple shuffle operation in which each output partition depends on all input partitions
class SimpleShuffleLayer(DataFrameLayer)

# High-level graph layer corresponding to a single stage of a multi-stage inter-partition shuffle operation.
class ShuffleLayer(SimpleShuffleLayer)

# High-level graph layer for a join operation requiring the smaller collection to be broadcasted to every partition of the larger collection.
class BroadcastJoinLayer(DataFrameLayer)

# DataFrame-based Blockwise Layer with IO
class DataFrameIOLayer(Blockwise, DataFrameLayer)
```

Copying over my comment from the Slack thread earlier today:

> I think figuring out what categories should be included will be the tricky thing, and trying to make those categories be something users will care about. It seems like it'd be more useful if we can use them to highlight bottlenecks. I assume that's why you suggest IO and shuffle. We seem a bit vague on things beyond that though, I don't know that "all other CPU computations" is likely to be very useful to display (and will probably be most cases)
So I think I'm more advocating that the Layer classes have some default annotations - and that's all. That way, we only need to look at these annotations, we don't need to do any `isinstance` checks. Of course, we still *might* want to look into the layer contents, to make the difference between GPU/CPU.

Totally agree that "CPU" should not have any associated decoration, as it will be the default, most common box.
@martindurant 
What would the default annotations be for all the layers? I can think of `layer_id`.

Meanwhile, I have come up with the visualization for the HLGs!

## Design
Group the listed layer_types into 4 categories

### IO (DataFrame)
1. `DataFrameIOLayer`

### Shuffle (DataFrame)
1. `SimpleShuffleLayer`
2. `ShuffleLayer`

### Blockwise (Array mostly + DataFrame)
1. `Blockwise`
2. `BlockwiseCreateArray`

### Materialized (Array  + DataFrame)
1. `MaterializedLayer`

The remaining layers:
1. `ArrayOverlapLayer`
2. `DataFrameLayer`
3. `BroadcastJoinLayer`
can be just left as it is.

### Visualization
 
#### Graphviz (`Layer.dask.visualize()`)
We can show the different groups with different colors.

EDIT: Updated Graph Reprentation
![image](https://user-images.githubusercontent.com/62539811/126983749-f07450fd-b36c-437f-a5a7-8256c68af0bf.png)


### HTML Repr (`Layer.dask`)
Using the same colors used in the Graphviz, we change the tiny circles next to each layer to reflect its `layer_type`.

EDIT: Updated HTML Reprentation
![image](https://user-images.githubusercontent.com/62539811/126983834-cbe441c2-100d-4f54-9e5d-e0fa2501c987.png)


(I have also modified the `MaterializedLayer`'s existing color to be a little lighter. I feel it looks nicer and easier on the eyes)
It might seem too colorful, but that's because I am showing all the different types of layers at the same time. This is highly unlikely to been observed in real Dask Graphs.


/cc @jacobtomlinson I would love to hear your opinion as well on the colors. You have done an amazing job adding all the existing colors :sparkles: 
> What would the default annotations be for all the layers

Empty? If we can't think of anything useful, it's better not to complicate the visuals.
/cc @mrocklin What do you think about the use of colors?
(colors are open for discussion)
I hadn't thought of changing the colours in the dots next to the layer titles in the HTML representation, that's a nice touch.

Personally, I'm not sure we're gaining much by having multiple shades of green/blue/etc. for different types of layers in the same larger category. I think that adds more confusion.

We'll need some effort spent on:
1. Explaining what information the user should gain from this (eg: shuffles are inefficient so you don't want to see too many of them, blockwise layers are likely to be nicely parallelized, etc.). You'll need to be able to explain this very clearly for all of them.
2. How we are going to communicate that information to the user (a legend for graphviz? added information in the docs? something in the HTML repr?)

Commit message: update doc

Patch to review:
diff --git a/dask/base.py b/dask/base.py
index d321269fd7a..9bc6996ecb6 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -571,7 +571,7 @@ def compute(*args, **kwargs):
 
 def visualize(*args, **kwargs):
     """
-    Visualize several dask graphs at once.
+    Visualize a dask low level graph.
 
     Requires ``graphviz`` to be installed. All options that are not the dask
     graph(s) should be passed as keyword arguments.
@@ -579,7 +579,7 @@ def visualize(*args, **kwargs):
     Parameters
     ----------
     dsk : dict(s) or collection(s)
-        The dask graph(s) to visualize.
+        The dask low level graph(s) to visualize.
     filename : str or None, optional
         The name of the file to write to disk. If the provided `filename`
         doesn't include an extension, '.png' will be used by default.
@@ -593,6 +593,8 @@ def visualize(*args, **kwargs):
     color : {None, 'order'}, optional
         Options to color nodes.  Provide ``cmap=`` keyword for additional
         colormap
+        - None, the default, no colors.
+        - 'order', colors the nodes' border based on the order they appear in the graph
     collapse_outputs : bool, optional
         Whether to collapse output boxes, which often have empty labels.
         Default is False.
diff --git a/dask/highlevelgraph.py b/dask/highlevelgraph.py
index 05ebb718f82..3ea52e93963 100644
--- a/dask/highlevelgraph.py
+++ b/dask/highlevelgraph.py
@@ -870,6 +870,49 @@ def merge(cls, *graphs):
         return cls(layers, dependencies)
 
     def visualize(self, filename="dask-hlg.svg", format=None, **kwargs):
+        """
+        Visualize a dask high level graph.
+
+        Requires ``graphviz`` to be installed. All options that are not the dask
+        graph(s) should be passed as keyword arguments.
+
+        Parameters
+        ----------
+        filename : str or None, optional
+            The name of the file to write to disk. If the provided `filename`
+            doesn't include an extension, '.svg' will be used by default.
+            If `filename` is None, no file will be written, and we communicate
+            with dot using only pipes.
+        format : {'png', 'pdf', 'dot', 'svg', 'jpeg', 'jpg'}, optional
+            Format in which to write output file.  Default is 'svg'.
+        color : {None, 'layer_type'}, optional
+            Options to color nodes.
+            - None, the default, no colors.
+            - 'layer_type', colored nodes based on the layer type.
+        **kwargs
+           Additional keyword arguments to forward to ``to_graphviz``.
+
+        Examples
+        --------
+        >>> x.dask.visualize(filename='dask.svg')  # doctest: +SKIP
+        >>> x.dask.visualize(filename='dask.svg', color='layer_type')  # doctest: +SKIP
+
+        Returns
+        -------
+        result : IPython.diplay.Image, IPython.display.SVG, or None
+            See dask.dot.dot_graph for more information.
+
+        See Also
+        --------
+        dask.dot.dot_graph
+        dask.base.visualize # low level variant
+
+        Notes
+        -----
+        For more information on optimization see here:
+
+        https://docs.dask.org/en/latest/optimize.html
+        """
         from .dot import graphviz_to_file
 
         g = to_graphviz(self, **kwargs)
@@ -1228,18 +1271,34 @@ def to_graphviz(
     for layer in hg.dependencies:
         n_tasks[layer] = len(hg.layers[layer])
 
-    mn = min(n_tasks.values())
-    mx = max(n_tasks.values())
+    min_tasks = min(n_tasks.values())
+    max_tasks = max(n_tasks.values())
 
     cache = {}
 
+    color = kwargs.get("color")
+    if color == "layer_type":
+        layer_colors = {
+            "DataFrameIOLayer": ["#CCC7F9", False],  # purple
+            "ShuffleLayer": ["#F9CCC7", False],  # rose
+            "SimpleShuffleLayer": ["#F9CCC7", False],  # rose
+            "ArrayOverlayLayer": ["#FFD9F2", False],  # pink
+            "BroadcastJoinLayer": ["#D9F2FF", False],  # blue
+            "Blockwise": ["#D9FFE6", False],  # green
+            "BlockwiseLayer": ["#D9FFE6", False],  # green
+            "BlockwiseCreateArray": ["#D9FFE6", False],  # green
+            "MaterializedLayer": ["#DBDEE5", False],  # gray
+        }
+
     for layer in hg.dependencies:
         layer_name = name(layer)
         attrs = data_attributes.get(layer, {})
 
         node_label = label(layer, cache=cache)
         node_size = (
-            20 if mx == mn else int(20 + ((n_tasks[layer] - mn) / (mx - mn)) * 20)
+            20
+            if max_tasks == min_tasks
+            else int(20 + ((n_tasks[layer] - min_tasks) / (max_tasks - min_tasks)) * 20)
         )
 
         layer_type = str(type(hg.layers[layer]).__name__)
@@ -1271,6 +1330,13 @@ def to_graphviz(
         attrs.setdefault("fontsize", str(node_size))
         attrs.setdefault("tooltip", str(node_tooltips))
 
+        if color == "layer_type":
+            node_color = layer_colors.get(layer_type)[0]
+            layer_colors.get(layer_type)[1] = True
+
+            attrs.setdefault("fillcolor", str(node_color))
+            attrs.setdefault("style", "filled")
+
         g.node(layer_name, **attrs)
 
     for layer, deps in hg.dependencies.items():
@@ -1278,6 +1344,28 @@ def to_graphviz(
         for dep in deps:
             dep_name = name(dep)
             g.edge(dep_name, layer_name)
+
+    if color == "layer_type":
+        legend_title = "Key"
+
+        legend_label = (
+            '<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" CELLPADDING="5">'
+            "<TR><TD><B>Legend: Layer types</B></TD></TR>"
+        )
+
+        for layer_type, color in layer_colors.items():
+            if color[1]:
+                legend_label += f'<TR><TD BGCOLOR="{color[0]}">{layer_type}</TD></TR>'
+
+        legend_label += "</TABLE>>"
+
+        attrs = data_attributes.get(legend_title, {})
+        attrs.setdefault("label", str(legend_label))
+        attrs.setdefault("fontsize", "20")
+        attrs.setdefault("margin", "0")
+
+        g.node(legend_title, **attrs)
+
     return g
 
 


Code review comments:
diff_hunk: @@ -571,15 +571,15 @@ def compute(*args, **kwargs):
 
 def visualize(*args, **kwargs):
     """
-    Visualize several dask graphs at once.
+    Visualize a dask low level graph.
path: dask/base.py
text: @user1:
**THIS** function does indeed take as many graphs as you like, the original docstring was correct. It's only for the `self.visualize` that the graph in question is already implied and there is no way to show more.

@author:
Oh! Didn't know that :heavy_check_mark:

diff_hunk: @@ -870,6 +870,49 @@ def merge(cls, *graphs):
         return cls(layers, dependencies)
 
     def visualize(self, filename="dask-hlg.svg", format=None, **kwargs):
+        """
+        Visualize a dask high level graph.
+
+        Requires ``graphviz`` to be installed. All options that are not the dask
+        graph(s) should be passed as keyword arguments.
+
+        Parameters
+        ----------
+        filename : str or None, optional
+            The name of the file to write to disk. If the provided `filename`
+            doesn't include an extension, '.svg' will be used by default.
+            If `filename` is None, no file will be written, and we communicate
path: dask/highlevelgraph.py
text: @user1:
Use `None` to avoid writing a file (e.g., for output in a jupyter notebook)

diff_hunk: @@ -571,15 +571,15 @@ def compute(*args, **kwargs):
 
 def visualize(*args, **kwargs):
     """
-    Visualize several dask graphs at once.
+    Visualize a dask low level graph.
 
     Requires ``graphviz`` to be installed. All options that are not the dask
     graph(s) should be passed as keyword arguments.
 
     Parameters
     ----------
     dsk : dict(s) or collection(s)
path: dask/base.py
text: @user1:
This parameter is, however, wrong. `*args` is the set of dask graphs to show, one or more of them.

diff_hunk: @@ -870,6 +870,49 @@ def merge(cls, *graphs):
         return cls(layers, dependencies)
 
     def visualize(self, filename="dask-hlg.svg", format=None, **kwargs):
+        """
+        Visualize a dask high level graph.
+
+        Requires ``graphviz`` to be installed. All options that are not the dask
+        graph(s) should be passed as keyword arguments.
+
+        Parameters
+        ----------
+        filename : str or None, optional
+            The name of the file to write to disk. If the provided `filename`
+            doesn't include an extension, '.svg' will be used by default.
path: dask/highlevelgraph.py
text: @user1:
Doesn't the extension you get depend on the `format` argument?

@author:
It also depends on `filename`.

https://github.com/dask/dask/blob/88a13f77f374a11425961c074f4e579f3befe721/dask/dot.py#L282-L283

but, `format` would overwrite the format in the `filename`.
`x.dask.visualize(filename="thishouldbesvg.svg", format="png")`
would `thishouldbesvg.svg.png`!

@user1:
I don't know what the code produces (you can test), but the wording here implies that when you give a filename without any extension at all, ".svg" is appended - which I suspect is not necessarily true.

@author:
Oh, you are indeed right!

@author:
```py
import dask.array as da

x = da.ones((10, 10), chunks=(5, 5))
x = x + 100
x = x * 100

x.dask.visualize(filename="isthissvg")
```
produces `isthissvg.png`

@author:
Wow, I found some tiny flaws in the filename/format system.
I will create an issue on it once I confirm them.

@author:
For now, I have just changed 
`'.svg' will be used by default.`
to 
`'.png' will be used by default.`
which is true.

_I will make this doc even better in another PR which fixes the filename/format situation for HLGs._

Is this fine for now?


@user1:
>  '.png' will be used by default.

Is it true regardless of the `format` argument?

> I will make this doc even better in another PR

yes, that's fine.

@author:
> Is it true regardless of the format argument?

yes! 

@author:
EDIT: format would overwrite the filename (if the filename has an extension attached or not)

```py
x.dask.visualize(filename="whatextension.png", format='pdf') # whatextension.png.pdf
x.dask.visualize(filename="whatextension", format=None) # whatextension.png
```

diff_hunk: @@ -593,6 +593,8 @@ def visualize(*args, **kwargs):
     color : {None, 'order'}, optional
         Options to color nodes.  Provide ``cmap=`` keyword for additional
path: dask/base.py
text: @user1:
If we're being thorough, then `cmap` should be a separate kwargs in this list with a better description than given here; or just remove it, since we don't say anywhere what kind of object cmap can be (and it is probably not tested).

@author:
Right!

diff_hunk: @@ -870,6 +870,49 @@ def merge(cls, *graphs):
         return cls(layers, dependencies)
 
     def visualize(self, filename="dask-hlg.svg", format=None, **kwargs):
+        """
+        Visualize a dask high level graph.
+
+        Requires ``graphviz`` to be installed. All options that are not the dask
+        graph(s) should be passed as keyword arguments.
path: dask/highlevelgraph.py
text: @user1:
Still refers to "dask graphs", but we are not passing any of those.