Repository: dask/dask

Issue:
Another task ordering issue
This example shows poor memory management:

```python
import dask.array as da
from colorize import visualize_colors

A, B = 20, 99 
x = da.random.normal(size=(A, B), chunks=(1, None))
for _ in range(2):
    y = (x[:, None, :] * x[:, :, None]).cumsum(axis=0)
    x = x.cumsum(axis=0)
w = (y * x[:, None]).sum(axis=(1, 2))

visualize_colors(w.__dask_graph__(), filename='dask.pdf')
```

![dask](https://user-images.githubusercontent.com/306380/34782973-e48cc3ba-f5ef-11e7-894d-bde8c2aa5d1c.png)



Discussion:
Found it, we're doing string comparison

```python
In [18]: key((k, 2))
Out[18]: (1, "('sum-aggregate-bf54015949fd59d3b23045ced546ed38', 2)")

In [19]: key((k, 10))
Out[19]: (1, "('sum-aggregate-bf54015949fd59d3b23045ced546ed38', 10)")

In [20]: _18 < _19
Out[20]: False
```

We would like to be able to compare the keys directly and then fall back on string comparison.  

Resolving patch:
diff --git a/dask/order.py b/dask/order.py
--- a/dask/order.py
+++ b/dask/order.py
@@ -1,13 +1,13 @@
 """ Static order of nodes in dask graph
 
-We can make decisions on what tasks to run next both
+Dask makes decisions on what tasks to prioritize both
 
 *  Dynamically at runtime
 *  Statically before runtime
 
-Dask's async scheduler runs dynamically and prefers to run tasks that were just
-made available.  However when several tasks become available at the same time
-we have an opportunity to break ties in an intelligent way
+Dynamically we prefer to run tasks that were just made available.  However when
+several tasks become available at the same time we have an opportunity to break
+ties in an intelligent way
 
         d
         |
@@ -15,8 +15,9 @@
      \ /
       a
 
-E.g. when we run ``a`` we can choose to run either ``b`` or ``c`` next.  In
-this case we may choose to start with ``c``, because it has other dependencies.
+For example after we finish ``a`` we can choose to run either ``b`` or ``c``
+next.  In this case we may prefer to start with ``c``, because it has other
+dependents.
 
 This is particularly important at the beginning of the computation when we
 often dump hundreds of leaf nodes onto the scheduler at once.  The order in
@@ -37,22 +38,23 @@
 There are several ways in which we might order our keys.  In practice we have
 found the following objectives important:
 
-1.  **Finish subtrees before starting new subtrees:** Often our computation
-    consists of many independent subtrees (e.g. reductions in an array).  We
-    want to work on and finish individual subtrees before moving on to others
-    in order to keep a low memory footprint.
-2.  **Run heavily depended-on tasks first**: Some tasks produce data that is
-    required by many other tasks, either in a deep linear chain (critical path)
-    or in a shallow but broad nexus (critical point).  By preferring these we
-    allow other computations to flow to completion more easily.
-
-
-Approach: Depth First Search with Intelligent Tie-Breaking
-----------------------------------------------------------
-
-To satisfy concern (1) we perform a depth first search (``dfs``).  To satisfy
-concern (2) we prefer to traverse down children in the order of which child has
-the descendent on whose result the most tasks depend.
+1.  **Depth first**:  By traversing a tree deeply before broadly we encourage
+    the completion of cohesive parts of the computation before starting new
+    parts.  This helps to reduce memory footprint.
+2.  **Lean**: We avoid tasks that have multiple dependents.  These tend to
+    start new branches of computation that we would prefer to avoid until we
+    have finished our current work.
+3.  **Heavy-first**: When deciding between two possible dependencies we choose
+    the dependency that is likely to take the longest first.  That way after it
+    finishes it has to stay in memory only a short time while waiting for its
+    co-dependent task to finish
+
+So we perform a depth first search where we choose to traverse down children
+with the following priority:
+
+1.  Number of dependents (smaller is better)
+2.  Number of total downstream dependencies (larger is better)
+3.  Name of the task itself, as a tie breaker
 """
 from __future__ import absolute_import, division, print_function
 
@@ -63,24 +65,24 @@
 def order(dsk, dependencies=None):
     """ Order nodes in dask graph
 
-    The ordering will be a toposort but will also have other convenient
-    properties
-
-    1.  Depth first search
-    2.  DFS prefers nodes that enable the most data
+    The ordering will be a topological sort but will also tend to produce
+    computations that have a small memory footprint.
 
+    Examples
+    --------
     >>> dsk = {'a': 1, 'b': 2, 'c': (inc, 'a'), 'd': (add, 'b', 'c')}
     >>> order(dsk)
-    {'a': 3, 'c': 2, 'b': 1, 'd': 0}
+    {'a': 2, 'c': 1, 'b': 3, 'd': 0}
     """
     if dependencies is None:
         dependencies = {k: get_dependencies(dsk, k) for k in dsk}
     dependents = reverse_dict(dependencies)
 
-    ndeps = ndependents(dependencies, dependents)
+    ndepts = {k: len(dependents[k]) for k in dependents}
+    ndepends = ndependencies(dependencies, dependents)
 
     def key(x):
-        return ndeps.get(x, 0), str(x)
+        return StrComparable((ndepts[x], -ndepends.get(x, 0), x))
 
     return dfs(dependencies, dependents, key=key)
 
@@ -102,8 +104,8 @@ def ndependents(dependencies, dependents):
     [('a', 3), ('b', 2), ('c', 1)]
     """
     result = dict()
-    num_needed = dict((k, len(v)) for k, v in dependents.items())
-    current = set(k for k, v in num_needed.items() if v == 0)
+    num_needed = {k: len(v) for k, v in dependents.items()}
+    current = {k for k, v in num_needed.items() if v == 0}
     while current:
         key = current.pop()
         result[key] = 1 + sum(result[parent] for parent in dependents[key])
@@ -114,36 +116,26 @@ def ndependents(dependencies, dependents):
     return result
 
 
-def child_max(dependencies, dependents, scores):
-    """ Maximum-ish of scores of children
-
-    This takes a dictionary of scores per key and returns a new set of scores
-    per key that is the maximum of the scores of all children of that node plus
-    its own score.  In some sense this ranks each node by the maximum
-    importance of their children plus their own value.
+def ndependencies(dependencies, dependents):
+    """ Number of total data elements on which this key depends
 
-    This is generally fed the result from ``ndependents``
+    For each key we return the number of tasks that must be run for us to run
+    this task.
 
     Examples
     --------
 
-    >>> dsk = {'a': 1, 'b': 2, 'c': (inc, 'a'), 'd': (add, 'b', 'c')}
-    >>> scores = {'a': 3, 'b': 2, 'c': 2, 'd': 1}
+    >>> dsk = {'a': 1, 'b': (inc, 'a'), 'c': (inc, 'b')}
     >>> dependencies, dependents = get_deps(dsk)
-
-    >>> sorted(child_max(dependencies, dependents, scores).items())
-    [('a', 3), ('b', 2), ('c', 5), ('d', 6)]
+    >>> sorted(ndependencies(dependencies, dependents).items())
+    [('a', 1), ('b', 2), ('c', 3)]
     """
     result = dict()
-    num_needed = dict((k, len(v)) for k, v in dependencies.items())
-    current = set(k for k, v in num_needed.items() if v == 0)
+    num_needed = {k: len(v) for k, v in dependencies.items()}
+    current = {k for k, v in num_needed.items() if v == 0}
     while current:
         key = current.pop()
-        score = scores[key]
-        children = dependencies[key]
-        if children:
-            score += max(result[child] for child in children)
-        result[key] = score
+        result[key] = 1 + sum(result[child] for child in dependencies[key])
         for parent in dependents[key]:
             num_needed[parent] -= 1
             if num_needed[parent] == 0:
@@ -163,7 +155,6 @@ def dfs(dependencies, dependents, key=lambda x: x):
 
     Examples
     --------
-
     >>> dsk = {'a': 1, 'b': 2, 'c': (inc, 'a'), 'd': (add, 'b', 'c')}
     >>> dependencies, dependents = get_deps(dsk)
 
@@ -192,3 +183,31 @@ def dfs(dependencies, dependents, key=lambda x: x):
         i += 1
 
     return result
+
+
+class StrComparable(object):
+    """ Wrap object so that it defaults to string comparison
+
+    When comparing two objects of different types Python fails
+
+    >>> 'a' < 1  # doctest: +SKIP
+    Traceback (most recent call last):
+        ...
+    TypeError: '<' not supported between instances of 'str' and 'int'
+
+    This class wraps the object so that, when this would occur it instead
+    compares the string representation
+
+    >>> StrComparable('a') < StrComparable(1)
+    False
+    """
+    __slots__ = ('obj',)
+
+    def __init__(self, obj):
+        self.obj = obj
+
+    def __lt__(self, other):
+        try:
+            return self.obj < other.obj
+        except Exception:
+            return str(self.obj) < str(other.obj)


Test patch:
diff --git a/dask/tests/test_order.py b/dask/tests/test_order.py
--- a/dask/tests/test_order.py
+++ b/dask/tests/test_order.py
@@ -1,10 +1,13 @@
-from itertools import chain
+import pytest
 
-from dask.order import child_max, ndependents, order
+from dask.order import ndependencies, order
 from dask.core import get_deps
 from dask.utils_test import add, inc
 
-a, b, c, d, e = 'abcde'
+
+@pytest.fixture(params=['abcde', 'edcba'])
+def abcde(request):
+    return request.param
 
 
 def issorted(L, reverse=False):
@@ -15,7 +18,8 @@ def f(*args):
     pass
 
 
-def test_ordering_keeps_groups_together():
+def test_ordering_keeps_groups_together(abcde):
+    a, b, c, d, e = abcde
     d = dict(((a, i), (f,)) for i in range(4))
     d.update({(b, 0): (f, (a, 0), (a, 1)),
               (b, 1): (f, (a, 2), (a, 3))})
@@ -33,16 +37,16 @@ def test_ordering_keeps_groups_together():
     assert abs(o[(a, 1)] - o[(a, 3)]) == 1
 
 
-def test_avoid_broker_nodes():
+def test_avoid_broker_nodes(abcde):
     """
 
     b0    b1  b2
     |      \  /
     a0      a1
 
-    a1 should be run before a0
+    a0 should be run before a1
     """
-    a, b, c = 'abc'
+    a, b, c, d, e = abcde
     dsk = {(a, 0): (f,), (a, 1): (f,),
            (b, 0): (f, (a, 0)), (b, 1): (f, (a, 1)), (b, 2): (f, (a, 1))}
 
@@ -59,7 +63,7 @@ def test_avoid_broker_nodes():
     assert o[(a, 0)] > o[(a, 1)]
 
 
-def test_base_of_reduce_preferred():
+def test_base_of_reduce_preferred(abcde):
     """
                a3
               /|
@@ -75,9 +79,10 @@ def test_base_of_reduce_preferred():
 
     We really want to run b0 quickly
     """
-    dsk = dict(((a, i), (f, (a, i - 1), (b, i))) for i in [1, 2, 3])
+    a, b, c, d, e = abcde
+    dsk = {(a, i): (f, (a, i - 1), (b, i)) for i in [1, 2, 3]}
     dsk[(a, 0)] = (f, (b, 0))
-    dsk.update(dict(((b, i), (f, c, 1)) for i in [0, 1, 2, 3]))
+    dsk.update({(b, i): (f, c, 1) for i in [0, 1, 2, 3]})
     dsk[c] = 1
 
     o = order(dsk)
@@ -96,7 +101,7 @@ def test_base_of_reduce_preferred():
     assert min([(b, i) for i in [0, 1, 2, 3]], key=o.get) == (b, 0)
 
 
-def test_avoid_upwards_branching():
+def test_avoid_upwards_branching(abcde):
     """
          a1
          |
@@ -112,6 +117,7 @@ def test_avoid_upwards_branching():
 
     Prefer b1 over c1 because it won't stick around waiting for d1 to complete
     """
+    a, b, c, d, e = abcde
     dsk = {(a, 1): (f, (a, 2)),
            (a, 2): (f, (a, 3)),
            (a, 3): (f, (b, 1), (c, 1)),
@@ -125,7 +131,7 @@ def test_avoid_upwards_branching():
     assert o[(b, 1)] < o[(c, 1)]
 
 
-def test_avoid_upwards_branching_complex():
+def test_avoid_upwards_branching_complex(abcde):
     """
          a1
          |
@@ -139,8 +145,10 @@ def test_avoid_upwards_branching_complex():
             |
             c3
 
-    Prefer b1 over c1 because it will be easier to finish its dependents
+    Prefer c1 over b1 because c1 will stay in memory less long while b1
+    computes
     """
+    a, b, c, d, e = abcde
     dsk = {(a, 1): (f, (a, 2)),
            (a, 2): (f, (a, 3)),
            (a, 3): (f, (b, 1), (c, 1)),
@@ -155,16 +163,17 @@ def test_avoid_upwards_branching_complex():
 
     o = order(dsk)
 
-    assert o[(b, 1)] < o[(c, 1)]
+    assert o[(c, 1)] < o[(b, 1)]
 
 
-def test_deep_bases_win_over_dependents():
+@pytest.mark.xfail(reason="this case is ambiguous")
+def test_deep_bases_win_over_dependents(abcde):
     """
     It's not clear who should run first, e or d
 
     1.  d is nicer because it exposes parallelism
     2.  e is nicer (hypothetically) because it will be sooner released
-        (though in this we need d to run first regardless)
+        (though in this case we need d to run first regardless)
 
             a
           / | \   .
@@ -172,6 +181,7 @@ def test_deep_bases_win_over_dependents():
         / \ | /
        e    d
     """
+    a, b, c, d, e = abcde
     dsk = {a: (f, b, c, d), b: (f, d, e), c: (f, d), d: 1, e: 2}
 
     o = order(dsk)
@@ -179,7 +189,7 @@ def test_deep_bases_win_over_dependents():
     assert o[d] < o[b] or o[d] < o[c]
 
 
-def test_prefer_deep():
+def test_prefer_deep(abcde):
     """
         c
         |
@@ -189,6 +199,7 @@ def test_prefer_deep():
 
     Prefer longer chains first so we should start with c
     """
+    a, b, c, d, e = abcde
     dsk = {a: 1, b: (f, a), c: (f, b),
            'x': 1, 'y': (f, 'x')}
 
@@ -196,38 +207,14 @@ def test_prefer_deep():
     assert o == {c: 0, b: 1, a: 2, 'y': 3, 'x': 4}
 
 
-def test_stacklimit():
+def test_stacklimit(abcde):
     dsk = dict(('x%s' % (i + 1), (inc, 'x%s' % i)) for i in range(10000))
     dependencies, dependents = get_deps(dsk)
-    scores = dict.fromkeys(dsk, 1)
-    child_max(dependencies, dependents, scores)
-    ndependents(dependencies, dependents)
-
-
-def test_ndependents():
-    a, b, c = 'abc'
-    dsk = dict(chain((((a, i), i * 2) for i in range(5)),
-                     (((b, i), (add, i, (a, i))) for i in range(5)),
-                     (((c, i), (add, i, (b, i))) for i in range(5))))
-    result = ndependents(*get_deps(dsk))
-    expected = dict(chain((((a, i), 3) for i in range(5)),
-                          (((b, i), 2) for i in range(5)),
-                          (((c, i), 1) for i in range(5))))
-    assert result == expected
-
-    dsk = {a: 1, b: 1}
-    deps = get_deps(dsk)
-    assert ndependents(*deps) == dsk
-
-    dsk = {a: 1, b: (add, a, 1), c: (add, b, a)}
-    assert ndependents(*get_deps(dsk)) == {a: 4, b: 2, c: 1}
-
-    dsk = {a: 1, b: a, c: b}
-    deps = get_deps(dsk)
-    assert ndependents(*deps) == {a: 3, b: 2, c: 1}
+    ndependencies(dependencies, dependents)
 
 
-def test_break_ties_by_str():
+def test_break_ties_by_str(abcde):
+    a, b, c, d, e = abcde
     dsk = {('x', i): (inc, i) for i in range(10)}
     x_keys = sorted(dsk)
     dsk['y'] = list(x_keys)
@@ -239,7 +226,28 @@ def test_break_ties_by_str():
     assert o == expected
 
 
-def test_order_doesnt_fail_on_mixed_type_keys():
+def test_order_doesnt_fail_on_mixed_type_keys(abcde):
     order({'x': (inc, 1),
            ('y', 0): (inc, 2),
            'z': (add, 'x', ('y', 0))})
+
+
+def test_gh_3055():
+    da = pytest.importorskip('dask.array')
+    A, B = 20, 99
+    x = da.random.normal(size=(A, B), chunks=(1, None))
+    for _ in range(2):
+        y = (x[:, None, :] * x[:, :, None]).cumsum(axis=0)
+        x = x.cumsum(axis=0)
+    w = (y * x[:, None]).sum(axis=(1,2))
+
+    dsk = dict(w.__dask_graph__())
+    o = order(dsk)
+    L = [o[k] for k in w.__dask_keys__()]
+    assert sorted(L) == L[::-1]
+
+
+def test_type_comparisions_ok(abcde):
+    a, b, c, d, e = abcde
+    dsk = {a: 1, (a, 1): 2, (a, b, 1): 3}
+    order(dsk)  # this doesn't err