Repository: dask/dask

Issue:
Use key names rather than function names in dask.visualize
When creating task graph images we currently inspect the task graph to find functions and use those function names in the graph.  Now that atop produces `subgraph_callables` all the time, this is less informative than it used to be.  We could either dig deeper here, or we could use the keynames instead.  This is what is done in most of the other diagnostics.  It might be a good time to switch over?

cc @jcrist 


Discussion:
A few ideas:
- Use the simplified keyname (the `"foo"` in `("foo", 0)`, I think this is what you're suggesting above). This is probably fine, as most keys are generated from function names, and single function names in the graph can be misleading.
- Set the `name` field in `SubgraphCallable` in `atop`. It may be easy to generate nice names, not sure.
- Leave the circles empty - the keys are already in the rectangle, this may provide adequate information, with the circles/edges just showing dependency information.

Resolving patch:
diff --git a/dask/dot.py b/dask/dot.py
--- a/dask/dot.py
+++ b/dask/dot.py
@@ -6,7 +6,7 @@
 
 from .compatibility import apply
 from .core import istask, get_dependencies, ishashable
-from .utils import funcname, import_required
+from .utils import funcname, import_required, key_split
 
 
 graphviz = import_required("graphviz", "Drawing dask graphs requires the "
@@ -104,6 +104,25 @@ def label(x, cache=None):
     return s
 
 
+def box_label(key):
+    """ Label boxes in graph by chunk index
+
+    >>> box_label(('x', 1, 2, 3))
+    '(1, 2, 3)'
+    >>> box_label(('x', 123))
+    '123'
+    >>> box_label('x')
+    ''
+    """
+    if isinstance(key, tuple):
+        key = key[1:]
+        if len(key) == 1:
+            [key] = key
+        return str(key)
+    else:
+        return ""
+
+
 def to_graphviz(dsk, data_attributes=None, function_attributes=None,
                 rankdir='BT', graph_attr={}, node_attr=None, edge_attr=None, **kwargs):
     if data_attributes is None:
@@ -119,14 +138,13 @@ def to_graphviz(dsk, data_attributes=None, function_attributes=None,
                          edge_attr=edge_attr)
 
     seen = set()
-    cache = {}
 
     for k, v in dsk.items():
         k_name = name(k)
         if k_name not in seen:
             seen.add(k_name)
             attrs = data_attributes.get(k, {})
-            attrs.setdefault('label', label(k, cache=cache))
+            attrs.setdefault('label', box_label(k))
             attrs.setdefault('shape', 'box')
             g.node(k_name, **attrs)
 
@@ -135,7 +153,7 @@ def to_graphviz(dsk, data_attributes=None, function_attributes=None,
             if func_name not in seen:
                 seen.add(func_name)
                 attrs = function_attributes.get(k, {})
-                attrs.setdefault('label', task_label(v))
+                attrs.setdefault('label', key_split(k))
                 attrs.setdefault('shape', 'circle')
                 g.node(func_name, **attrs)
             g.edge(func_name, k_name)
@@ -145,7 +163,7 @@ def to_graphviz(dsk, data_attributes=None, function_attributes=None,
                 if dep_name not in seen:
                     seen.add(dep_name)
                     attrs = data_attributes.get(dep, {})
-                    attrs.setdefault('label', label(dep, cache=cache))
+                    attrs.setdefault('label', box_label(dep))
                     attrs.setdefault('shape', 'box')
                     g.node(dep_name, **attrs)
                 g.edge(dep_name, func_name)


Test patch:
diff --git a/dask/tests/test_dot.py b/dask/tests/test_dot.py
--- a/dask/tests/test_dot.py
+++ b/dask/tests/test_dot.py
@@ -76,9 +76,7 @@ def test_to_graphviz():
     g = to_graphviz(dsk)
     labels = list(filter(None, map(get_label, g.body)))
     assert len(labels) == 10        # 10 nodes total
-    funcs = set(('add', 'sum', 'neg'))
-    assert set(labels).difference(dsk) == funcs
-    assert set(labels).difference(funcs) == set(dsk)
+    assert set(labels) == {'c', 'd', 'e', 'f', '""'}
     shapes = list(filter(None, map(get_shape, g.body)))
     assert set(shapes) == set(('box', 'circle'))
 
@@ -89,10 +87,8 @@ def test_to_graphviz_custom():
         data_attributes={'a': {'shape': 'square'}},
         function_attributes={'c': {'label': 'neg_c', 'shape': 'ellipse'}},
     )
-    labels = list(filter(None, map(get_label, g.body)))
-    funcs = set(('add', 'sum', 'neg', 'neg_c'))
-    assert set(labels).difference(dsk) == funcs
-    assert set(labels).difference(funcs) == set(dsk)
+    labels = set(filter(None, map(get_label, g.body)))
+    assert labels == {'neg_c', 'd', 'e', 'f', '""'}
     shapes = list(filter(None, map(get_shape, g.body)))
     assert set(shapes) == set(('box', 'circle', 'square', 'ellipse'))