Repository: dask/dask (Python)

Issue:
Ordering changes negatively affecting memory use
I will attempt to follow this report up with an example. I encountered this behaviour in a relatively complicated example and replicating it in a simpler case is eluding me. I am using dask (2.9.1/2.10.1, see remainder of post), distributed 2.10.0 and Python 3.6.9.

My specific problem involves reading large amounts of data from disk, processing it in independent chunks and then writing the results to disk again. Prior to dask 2.9.2 the memory usage of my code behaved as expected, exhibiting a typical heartbeat pattern as data was read - processed - written.

Post dask 2.9.2, the above is no longer the case and the maximum memory footprint has increased, often substantially. I believe that this is a result of changes to dask.order.order made [here](https://github.com/dask/dask/commit/c63dc149210d45cbf90bb3f9d5a56106339ac24c). 

As an example of what I believe is causing this behaviour, consider the following graph: 
```
   E          # Write operation.
   |   
   C D        # C: Additional computation. D: Write operation. 
   |/
   B          # Computation producing two outputs.
   |
   A          # Read operation.
```
dask.compute is then called on `[D, E]`. In my head (and in the past), this has worked to write two data products to disk per input chunk without growing extensively in memory. However, as of 2.9.2, it seems that the scheduler is no longer giving priority to node `D`. The result is that all those write operations are being kept in memory until every node `E` is completed.  This does not seem like desirable behaviour. 

There is one additional peculiarity - this memory growth is only evident when using the distributed scheduler. I cannot think of any explanation for this as its version was unchanged across my tests.

To boil it down - has anyone else experienced similar changes in behaviour? Is it plausible that changes to ordering could have this effect? Finally, does anyone have any advice for producing an MRE?


Discussion:
Thanks for the report.  `D` is definitely given priority in `order` (both before and after the change).  In fact, the new `order` will never leave a single task like `D` dangling before moving on (the previous version would sometimes do this).  I will need to see a more complete example to see whether `order` is behaving poorly.  It could indeed be something specific with the distributed scheduler and how it interacts with the order of tasks.

Have you tried running with the serial scheduler (`dask.get`)?
Thanks for the reply @eriknw. I have tried running with the single-threaded scheduler via compute, which I think is equivalent to `dask.get`. My experiments show that there is memory growth in the serial case too, which I had previously missed. 

I have also confirmed that it was [this commit](https://github.com/dask/dask/commit/c63dc149210d45cbf90bb3f9d5a56106339ac24c) that introduced the change - I installed from source and retested with HEAD at [this commit](https://github.com/dask/dask/commit/5db6dad7947f186a9764cdf409b8f50bf7c38f22), the one directly preceding the ordering changes. There was no memory growth at that point. 

In terms of a specific example, I am still trying. Unfortunately `order` for my code contains around 210 keys per input chunk. It makes it somewhat difficult to identify the precise degree of complexity which triggers my problem.
I have managed to get rid of the memory growth by removing some of the fields from the xarray.Datasets which are written to disk. This gives me a little more insight into where the problem is. Hopefully I can get a little closer to a reproducible example now.
I have been digging further into the problem. Unfortunately I feel like I may be further away from a reproducible example than before. As I mentioned in the original post, I call compute on the two write operations `[D, E]`. I have tested both `D` and `E` independently and neither exhibit the memory growth problem. Calling compute on `[D, E]` exposes the problem, which I think suggests that the dependencies/orderings are being combined sub-optimally.
I have saved the results of the call to `order` before and after the changes. The dictionary has 3502 keys. 
Writing an xarray.Dataset to disk produces keys which start with `_finalize_store`. Below are the values associated with all `_finalize_store` operations before and after the ordering change:

```
Before:
[ 212,  256,  420,  506,  592,  678,  764,  850,  936, 1025, 1111,
 1197, 1278, 1364, 1450, 1536, 1622, 1708, 1794, 1880, 1966, 2052,
 2138, 2224, 2310, 2396, 2482, 2492, 2502, 2590, 2678, 2845, 2929,
 2939, 3104, 3190, 3276, 3286, 3417, 3461]
After:
[ 858,  940,  988, 1036, 1084, 1132, 1180, 1228, 1276, 1324, 1372,
 1420, 1468, 1516, 1564, 1612, 1660, 1708, 1756, 1804, 1852, 1900,
 1948, 1996, 2044, 2092, 2140, 2542, 2590, 2638, 2686, 2734, 2782,
 2830, 2878, 2926, 2974, 3022, 3070, 3118]
```

If I am interpreting the above correctly, the new ordering waits a very long time before writing the results of the first chunk to disk. Another thing I found interesting is the non-uniform spacing between `_finalise_store` values (computed as `A[1:] - A[:-1]`):
```
Before:
[ 44, 164,  86,  86,  86,  86,  86,  86,  89,  86,  86,  81,  86,
 86,  86,  86,  86,  86,  86,  86,  86,  86,  86,  86,  86,  86,
 10,  10,  88,  88, 167,  84,  10, 165,  86,  86,  10, 131,  44]
After:
[ 82,  48,  48,  48,  48,  48,  48,  48,  48,  48,  48,  48,  48,
  48,  48,  48,  48,  48,  48,  48,  48,  48,  48,  48,  48,  48,
 402,  48,  48,  48,  48,  48,  48,  48,  48,  48,  48,  48,  48]
```
Prior to the ordering changes, the spacing is less uniform but is consistent with being more evenly spread throughout the ordering dict. With the new ordering spacing is very uniform but there seems to be one fairly extreme outlier.

This behaviour carries over to the other write operation too. 
```
Before:
      [  21,   22,   25,   27,   98,   99,  102,  104,  126,  141,  142,
        145,  147,  168,  272,  273,  276,  278,  334,  349,  350,  353,
        355,  376,  435,  436,  439,  441,  462,  521,  522,  525,  527,
        548,  607,  608,  611,  613,  634,  693,  694,  697,  699,  717,
        779,  780,  783,  785,  803,  865,  866,  869,  871,  892,  939,
        954,  955,  958,  960,  978, 1040, 1041, 1044, 1046, 1064, 1126,
       1127, 1130, 1132, 1150, 1207, 1208, 1211, 1213, 1231, 1293, 1294,
       1297, 1299, 1320, 1379, 1380, 1383, 1385, 1406, 1465, 1466, 1469,
       1471, 1489, 1551, 1552, 1555, 1557, 1575, 1637, 1638, 1641, 1643,
       1664, 1723, 1724, 1727, 1729, 1747, 1809, 1810, 1813, 1815, 1833,
       1895, 1896, 1899, 1901, 1919, 1981, 1982, 1985, 1987, 2008, 2067,
       2068, 2071, 2073, 2091, 2153, 2154, 2157, 2159, 2177, 2239, 2240,
       2243, 2245, 2263, 2325, 2326, 2329, 2331, 2352, 2411, 2412, 2415,
       2417, 2435, 2518, 2519, 2522, 2524, 2546, 2606, 2607, 2610, 2612,
       2668, 2700, 2701, 2704, 2706, 2774, 2775, 2778, 2780, 2801, 2848,
       2858, 2859, 2862, 2864, 2919, 2956, 2957, 2960, 2962, 3018, 3033,
       3034, 3037, 3039, 3060, 3119, 3120, 3123, 3125, 3146, 3205, 3206,
       3209, 3211, 3232, 3308, 3309, 3312, 3314, 3336, 3346, 3347, 3350,
       3352, 3373]
After:
      [  15,   35,   49,   63,   77,   91,  105,  119,  133,  147,  161,
        175,  189,  203,  217,  231,  245,  259,  273,  287,  301,  315,
        329,  343,  357,  371,  385,  402,  419,  436,  453,  470,  487,
        504,  521,  538,  555,  572,  589,  606,  623,  640,  657,  674,
        691,  708,  725,  742,  759,  776,  793,  810,  860,  862,  865,
        892, 2169, 2198, 2227, 2253, 2279, 2305, 2331, 2360, 2388, 2414,
       2440, 2468, 2494, 3120, 3122, 3125, 3127, 3129, 3132, 3134, 3136,
       3139, 3141, 3143, 3146, 3148, 3150, 3153, 3155, 3157, 3160, 3162,
       3164, 3167, 3169, 3171, 3174, 3176, 3178, 3181, 3183, 3185, 3188,
       3190, 3192, 3195, 3197, 3199, 3202, 3204, 3206, 3209, 3211, 3213,
       3216, 3218, 3220, 3223, 3225, 3227, 3230, 3232, 3234, 3237, 3239,
       3241, 3244, 3246, 3248, 3251, 3253, 3255, 3258, 3260, 3262, 3265,
       3267, 3269, 3272, 3274, 3276, 3279, 3281, 3283, 3286, 3288, 3290,
       3293, 3298, 3300, 3303, 3305, 3310, 3312, 3315, 3317, 3322, 3324,
       3327, 3329, 3334, 3336, 3339, 3341, 3346, 3348, 3351, 3353, 3358,
       3360, 3363, 3365, 3370, 3372, 3375, 3377, 3382, 3384, 3387, 3389,
       3394, 3396, 3399, 3401, 3406, 3408, 3411, 3413, 3418, 3420, 3423,
       3425, 3430, 3432, 3435, 3437, 3442, 3444, 3447, 3449, 3454, 3456,
       3459, 3461]
```
Before the ordering changes the writes were evenly spread across the ordering. Now they are highly concentrated. Once again computing `A[1:] - A[:-1]` for these writes:
```
Before:
      [  1,   3,   2,  71,   1,   3,   2,  22,  15,   1,   3,   2,  21,
       104,   1,   3,   2,  56,  15,   1,   3,   2,  21,  59,   1,   3,
         2,  21,  59,   1,   3,   2,  21,  59,   1,   3,   2,  21,  59,
         1,   3,   2,  18,  62,   1,   3,   2,  18,  62,   1,   3,   2,
        21,  47,  15,   1,   3,   2,  18,  62,   1,   3,   2,  18,  62,
         1,   3,   2,  18,  57,   1,   3,   2,  18,  62,   1,   3,   2,
        21,  59,   1,   3,   2,  21,  59,   1,   3,   2,  18,  62,   1,
         3,   2,  18,  62,   1,   3,   2,  21,  59,   1,   3,   2,  18,
        62,   1,   3,   2,  18,  62,   1,   3,   2,  18,  62,   1,   3,
         2,  21,  59,   1,   3,   2,  18,  62,   1,   3,   2,  18,  62,
         1,   3,   2,  18,  62,   1,   3,   2,  21,  59,   1,   3,   2,
        18,  83,   1,   3,   2,  22,  60,   1,   3,   2,  56,  32,   1,
         3,   2,  68,   1,   3,   2,  21,  47,  10,   1,   3,   2,  55,
        37,   1,   3,   2,  56,  15,   1,   3,   2,  21,  59,   1,   3,
         2,  21,  59,   1,   3,   2,  21,  76,   1,   3,   2,  22,  10,
         1,   3,   2,  21]
After:
      [  20,   14,   14,   14,   14,   14,   14,   14,   14,   14,   14,
         14,   14,   14,   14,   14,   14,   14,   14,   14,   14,   14,
         14,   14,   14,   14,   17,   17,   17,   17,   17,   17,   17,
         17,   17,   17,   17,   17,   17,   17,   17,   17,   17,   17,
         17,   17,   17,   17,   17,   17,   17,   50,    2,    3,   27,
       1277,   29,   29,   26,   26,   26,   26,   29,   28,   26,   26,
         28,   26,  626,    2,    3,    2,    2,    3,    2,    2,    3,
          2,    2,    3,    2,    2,    3,    2,    2,    3,    2,    2,
          3,    2,    2,    3,    2,    2,    3,    2,    2,    3,    2,
          2,    3,    2,    2,    3,    2,    2,    3,    2,    2,    3,
          2,    2,    3,    2,    2,    3,    2,    2,    3,    2,    2,
          3,    2,    2,    3,    2,    2,    3,    2,    2,    3,    2,
          2,    3,    2,    2,    3,    2,    2,    3,    2,    2,    3,
          5,    2,    3,    2,    5,    2,    3,    2,    5,    2,    3,
          2,    5,    2,    3,    2,    5,    2,    3,    2,    5,    2,
          3,    2,    5,    2,    3,    2,    5,    2,    3,    2,    5,
          2,    3,    2,    5,    2,    3,    2,    5,    2,    3,    2,
          5,    2,    3,    2,    5,    2,    3,    2,    5,    2,    3,
          2]
```
I find the extreme jumps in the post-change ordering worrying. 

Forgive the walls of numbers, but this seems non-trivial to debug/reproduce. I am relatively certain that it is a bug though.
Thanks for continuing to dig into this.  I am unsure how to interpret these numbers.  In order for me to look at the behavior of `order.order`, it would be sufficient to only have the dependencies of each task.  I don't need the actual tasks or the workflow that created them (although a workflow would be nice).  Could you find a way to post the results of `get_dependencies` on the dask graph that the scheduler gets (i.e., preferably after the optimization passes have occurred).  Thanks!
@eriknw Will try get something to you tomorrow. Thanks for the help! 

In terms of the numbers above, they just correspond to the integer values of the write operations as they appear in the output of `order.order`. I was just trying to establish how the changes have affected the ordering. 
You may also be interested in the `x.visualize("myfile.png", color="order")` option, which for modestly sized graphs (below 100-1000?) should create an image of the graph with the ordering overlayed as a colormap.  That can be useful when identifying ordering issues.
Thanks @mrocklin - I have done that for a subset of the problem. The graph is very complicated so interpreting the results is a little challenging. It does however look like the ordering change has resulted in a "breadthier" traversal. This is strange because each chunk of input data is processed independently - there shouldn't be any dependencies between data chunks, and the behaviour in the past was consistent with that.
@eriknw I have pickled the dictionary ([deps.txt](https://github.com/dask/dask/files/4157823/deps.txt)) produced [here](https://github.com/dask/dask/blob/03aca07be7f69964d290138fd86ce71c8797afe7/dask/order.py#L109). I believe that is what you asked for - please let me know if I am mistaken. The keys containing `_finialize_store` correspond to writing xarray.Datasets to disk. Keys starting with `C147_unflagged.MS-write` correspond to writing columns back to the input dataset. Let me know if I can be of any further assistance.




Just as a follow up to @mrocklin's suggestion, here are the ordering plots for a subset of my data.

Pre ordering change:
![pre](https://user-images.githubusercontent.com/6582745/73824915-acd58e80-4803-11ea-9a47-6de45b27be01.png)

Post ordering change:
![post](https://user-images.githubusercontent.com/6582745/73824998-dabad300-4803-11ea-8df4-3d1e840953ec.png)

To me it seems that the post change graph is being traversed in a sub-optimal, breadth-first fashion.

Excellent, thanks @JSKenyon!  I'm taking a look at deps.txt now.
Real, difficult cases like this are very useful, so thanks again for raising the issue @JSKenyon !

I've made a small variation to `order` to go from

![dsk_current-min](https://user-images.githubusercontent.com/2058401/73871701-ed3d0880-4813-11ea-8658-7e42254a60c3.png)

to

![dsk_lifo1-min](https://user-images.githubusercontent.com/2058401/73871928-56248080-4814-11ea-959e-b36f459446e1.png)

which I think is much better (and passes all the `order` tests).

fyi, I made these plots with this command `dask.visualize(dsk, color='order', cmap='autumn', filename='dsk_current.jpeg', node_attr={'penwidth': '10'})`.  Using a different colormap and larger penwidth makes it easier to see behaviors in large graphs.

I'll need more time to investigate the performance impact of this change and to run through other more complicated examples.  I'd like more time to investigate other possible solutions too.  I'll get an experimental branch up later today for you to test out.
@eriknw Thank you so much for the help! I look forward to trying out your branch when it appears. :-) And an extra heaping of gratitude for the visualize tips - that will prove super useful.
@eriknw I _think_ I run into the same bug in my work, where I have a very large intermediate representation that get's reduced by two different tasks. This gave me a nicely small MWE that might be of use:

```python
import numpy as np
import dask.array as da
import dask
from dask.distributed import Client, LocalCluster

cluster = LocalCluster(n_workers=1, threads_per_worker=2)
client = Client(cluster)

# Normally load actual data here and apply FFTs to do cross correlation
X = da.random.random((150,256,256), chunks=(50,256,256))
Corr = X[:, np.newaxis, ...] * X.conj()[np.newaxis, ...]
flatData = Corr.reshape(Corr.shape[:-2]+(-1,))
argmaxs = da.argmax(flatData, axis=-1)
maxima = da.max(flatData, axis=-1)
dask.visualize((maxima, argmaxs), optimize_graph='True', 
               color='order',
               cmap='viridis', 
               filename='dsk_MWE_argmaxs+maxima.png', 
               node_attr={'penwidth': '10'})
dask.visualize(argmaxs, verbose=True, optimize_graph='True', 
               color='order', 
               cmap='viridis', 
               filename='dsk_MWE_argmaxs.png', 
               node_attr={'penwidth': '10'})

da.compute(argmaxs) #fine 
da.compute(maxima) #fine
da.compute(argmaxs,maxima) #explodes
```
Obviously this is a somewhat awkward quadratic algorithm, but the first explodes memory use (especially if you increase the size somewhat more), while calculating either the argmaxs or maxima alone is perfectly fine, as the chunks of flatData don't remain in memory.

Argmaxs+maxima:
![dsk_MWE_argmaxs+maxima](https://user-images.githubusercontent.com/6816964/78362877-234ffc00-75bb-11ea-9594-764ff72a2506.png)
argmaxs only:
![dsk_MWE_argmaxs](https://user-images.githubusercontent.com/6816964/78362867-1fbc7500-75bb-11ea-9303-dd58380a6297.png)

I did not yet have time to test this MWE with #5872 , but it might be of use for checking and hope #5872 will actually resolve this.

(The full code at https://github.com/TAdeJong/LEEM-analysis/blob/master/2%20-%20Driftcorrection.ipynb used to be able to run with upto 1500 images in 32GB of memory, now it chokes on 500 :cry:  )
@TAdeJong, these are great example, thanks so much!

#5872 does handle this better, and I've already made an improvement to make it do even better :)
This was added as a test.  The new behavior in #5872 is:

![dask_new](https://user-images.githubusercontent.com/2058401/78406157-ac6d3000-75c7-11ea-8596-0a2838a03200.png)




Patch:
diff --git a/dask/order.py b/dask/order.py
index 8aad68d7b34..8ff1a260dc0 100644
--- a/dask/order.py
+++ b/dask/order.py
@@ -49,7 +49,8 @@
 
 Work towards *small goals* with *big steps*.
 
-1.  **Small goals**: prefer tasks whose final dependents have few dependencies.
+1.  **Small goals**: prefer tasks that have few total dependents and whose final
+    dependents have few total dependencies.
 
     We prefer to prioritize those tasks that help branches of computation that
     can terminate quickly.
@@ -74,6 +75,7 @@
     This relies on the regularity of graph constructors like dask.array to be a
     good proxy for ordering.  This is usually a good idea and a sane default.
 """
+from collections import defaultdict
 from math import log
 from .core import get_dependencies, reverse_dict, get_deps, getcycle  # noqa: F401
 from .utils_test import add, inc  # noqa: F401
@@ -154,19 +156,13 @@ def order(dsk, dependencies=None):
     def dependents_key(x):
         """ Choose a path from our starting task to our tactical goal
 
-        This path is connected to a large goal, but focuses on completing a small goal.
+        This path is connected to a large goal, but focuses on completing
+        a small goal and being memory efficient.
         """
-        num_dependents = len(dependents[x])
-        total_dependents, min_dependencies, _, min_heights, _ = metrics[x]
         return (
-            # tactically, finish small connected jobs first
-            min_dependencies,
-            num_dependents - min_heights,  # prefer tall and narrow
-            -total_dependents,  # take a big step
-            # try to be memory efficient
-            num_dependents - len(dependencies[x]) + num_needed[x],
-            num_dependents,
-            total_dependencies[x],  # already found work, so don't add more
+            # Focus on being memory-efficient
+            len(dependents[x]) - len(dependencies[x]) + num_needed[x],
+            -metrics[x][3],  # min_heights
             # tie-breaker
             StrComparable(x),
         )
@@ -174,7 +170,7 @@ def dependents_key(x):
     def dependencies_key(x):
         """ Choose which dependency to run as part of a reverse DFS
 
-        This is very similar to both ``initial_stack_key`` and ``dependents_key``.
+        This is very similar to both ``initial_stack_key``.
         """
         num_dependents = len(dependents[x])
         (
@@ -202,125 +198,268 @@ def dependencies_key(x):
             StrComparable(x),
         )
 
+    def finish_now_key(x):
+        """ Determine the order of dependents that are ready to run and be released"""
+        return (-len(dependencies[x]), StrComparable(x))
+
+    # Computing this for all keys can sometimes be relatively expensive :(
+    partition_keys = {
+        key: (
+            (min_dependencies - total_dependencies[key] + 1)
+            * (total_dependents - min_heights)
+        )
+        for key, (
+            total_dependents,
+            min_dependencies,
+            _,
+            min_heights,
+            _,
+        ) in metrics.items()
+    }
+
     result = {}
     i = 0
 
-    # Nodes in the current weakly connected subgraph that are not part of the current
-    # DFS along dependencies.  The next DFS will begin from these nodes.  `outer_stack`
-    # is populated from `init_stack` and the dependents of `outer_stack_seeds`.
+    # `inner_stask` is used to perform a DFS along dependencies.  Once emptied
+    # (when traversing dependencies), this continue down a path along dependents
+    # until a root node is reached.
+    #
+    # Sometimes, a better path along a dependent is discovered (i.e., something
+    # that is easier to compute and doesn't requiring holding too much in memory).
+    # In this case, the current `inner_stack` is appended to `inner_stacks` and
+    # we begin a new DFS from the better node.
+    #
+    # A "better path" is determined by comparing `partition_keys`.
+    inner_stacks = [[min(init_stack, key=initial_stack_key)]]
+    inner_stacks_append = inner_stacks.append
+    inner_stacks_extend = inner_stacks.extend
+    inner_stacks_pop = inner_stacks.pop
+
+    # Okay, now we get to the data structures used for fancy behavior.
+    #
+    # As we traverse nodes in the DFS along dependencies, we partition the dependents
+    # via `partition_key`.  A dependent goes to:
+    #    1) `inner_stack` if it's better than our current target,
+    #    2) `next_nodes` if the partition key is lower than it's parent,
+    #    3) `later_nodes` otherwise.
+    # When the inner stacks are depleted, we process `next_nodes`.  If `next_nodes` is
+    # empty (and `outer_stacks` is empty`), then we process `later_nodes` the same way.
+    # These dicts use `partition_keys` as keys.  We process them by placing the values
+    # in `outer_stack` so that the smallest keys will be processed first.
+    next_nodes = defaultdict(list)
+    later_nodes = defaultdict(list)
+
+    # `outer_stack` is used to populate `inner_stacks`.  From the time we partition the
+    # dependents of a node, we group them: one list per partition key per parent node.
+    # This likely results in many small lists.  We do this to avoid sorting many larger
+    # lists (i.e., to avoid n*log(n) behavior).  So, we have many small lists that we
+    # partitioned, and we keep them in the order that we saw them (we will process them
+    # in a FIFO manner).  By delaying sorting for as long as we can, we can first filter
+    # out nodes that have already been computed.  All this complexity is worth it!
     outer_stack = []
-
-    # Nodes get added to this as they complete.  We will populate `outer_stack` from
-    # dependents of these nodes.  We consider nodes using a FIFO policy (yes, we could
-    # have used collections.deque, but a list with the current index is faster).
-    # Note that this does not include nodes with no dependents!
-    outer_stack_seeds = []
-    outer_stack_seeds_index = 0
-
-    # Used to perform a DFS along dependencies.  Once emptied (when traversing dependencies),
-    # this continues (along dependents) along a path from the initial (leaf) node down to a
-    # root node (a tactical goal).  This ensures this root node will get calculated before
-    # we consider any other dependents from `outer_stack_seeds`.
-    inner_stack = []
-
-    # aliases for speed
-    outer_stack_append = outer_stack.append
-    outer_stack_pop = outer_stack.pop
     outer_stack_extend = outer_stack.extend
-    outer_stack_seeds_append = outer_stack_seeds.append
-    inner_stack_append = inner_stack.append
-    inner_stack_pop = inner_stack.pop
-    inner_stack_extend = inner_stack.extend
+    outer_stack_pop = outer_stack.pop
+
+    # Keep track of nodes that are in `inner_stack` or `inner_stacks` so we don't
+    # process them again.
+    seen = set()  # seen in an inner_stack (and has dependencies)
+    seen_update = seen.update
+    seen_add = seen.add
+
+    # alias for speed
     set_difference = set.difference
 
     is_init_sorted = False
-    item = min(init_stack, key=initial_stack_key)
     while True:
-        outer_stack_append(item)
-        while outer_stack:
-            item = outer_stack_pop()
-            if item not in result:
-                inner_stack_append(item)
-
-                while inner_stack:
-                    # Perform a DFS along dependencies until we complete our tactical goal
-                    item = inner_stack_pop()
-                    if item in result:
-                        continue
-
-                    if num_needed[item]:
-                        inner_stack_append(item)
-                        deps = set_difference(dependencies[item], result)
-                        if 1 < len(deps) < 1000:
-                            inner_stack_extend(
-                                sorted(deps, key=dependencies_key, reverse=True)
-                            )
-                        else:
-                            inner_stack_extend(deps)
-                        continue
-
-                    result[item] = i
-                    i += 1
-                    deps = dependents[item]
-
-                    if metrics[item][3] == 2:  # min_height
-                        # Don't leave any dangling single nodes!  Finish all dependents that are
-                        # ready and are also root nodes.  Doing this here also lets us continue
-                        # down to a different root node.
-                        finish_now = {
-                            dep
-                            for dep in deps
-                            if not dependents[dep] and num_needed[dep] == 1
-                        }
-                        if finish_now:
-                            deps -= finish_now
-                            if len(finish_now) > 1:
-                                finish_now = sorted(finish_now, key=dependents_key)
-                            for dep in finish_now:
-                                result[dep] = i
-                                i += 1
-
-                    if deps:
-                        for dep in deps:
-                            num_needed[dep] -= 1
-                        if not inner_stack:
-                            # Continue towards our tactical goal (an easy-to-compute root)
-                            if len(deps) == 1:
-                                inner_stack_extend(deps)
-                            else:
-                                # If there are many, many dependents, then calculating
-                                # `dependents_key` here could be relatively expensive.
-                                # However, this is still probably worth it tactically.
-                                inner_stack_append(min(deps, key=dependents_key))
-                                outer_stack_seeds_append(item)
-                        else:
-                            # Our next starting point will be "seeded" by a completed task in a
-                            # FIFO manner.  When our DFS with `inner_stack` is finished--which
-                            # means we computed our tactical goal--we will choose our next starting
-                            # point from the dependents of completed tasks.  However, it is too
-                            # expensive to consider all dependents of all completed tasks, so we
-                            # consider the dependents of tasks in the order they complete.
-                            outer_stack_seeds_append(item)
-
-            while not outer_stack and outer_stack_seeds_index < len(outer_stack_seeds):
-                # We wait for as long as possible to consider dependents of completed nodes:
-                # 1. some may have already completed, so waiting lets us sort fewer items, and
-                # 2. sorting via `dependents_key` depends on `num_needed`, which is dynamic.
-                deps = set_difference(
-                    dependents[outer_stack_seeds[outer_stack_seeds_index]], result
-                )
-                if deps:
+        while inner_stacks:
+            inner_stack = inner_stacks_pop()
+            inner_stack_pop = inner_stack.pop
+            while inner_stack:
+                # Perform a DFS along dependencies until we complete our tactical goal
+                item = inner_stack_pop()
+                if item in result:
+                    continue
+                if num_needed[item]:
+                    inner_stack.append(item)
+                    deps = set_difference(dependencies[item], result)
                     if 1 < len(deps) < 1000:
-                        outer_stack_extend(
-                            sorted(deps, key=dependents_key, reverse=True)
+                        inner_stack.extend(
+                            sorted(deps, key=dependencies_key, reverse=True)
                         )
                     else:
-                        outer_stack_extend(deps)
-                outer_stack_seeds_index += 1
+                        inner_stack.extend(deps)
+                    seen_update(deps)
+                    continue
+
+                result[item] = i
+                i += 1
+                deps = dependents[item]
+
+                # If inner_stack is empty, then we typically add the best dependent to it.
+                # However, we don't add to it if we complete a node early via "finish_now" below
+                # or if a dependent is already on an inner_stack.  In this case, we add the
+                # dependents (not in an inner_stack) to next_nodes or later_nodes to handle later.
+                # This serves three purposes:
+                #   1. shrink `deps` so that it can be processed faster,
+                #   2. make sure we don't process the same dependency repeatedly, and
+                #   3. make sure we don't accidentally continue down an expensive-to-compute path.
+                add_to_inner_stack = True
+                if metrics[item][3] == 1:  # min_height
+                    # Don't leave any dangling single nodes!  Finish all dependents that are
+                    # ready and are also root nodes.
+                    finish_now = {
+                        dep
+                        for dep in deps
+                        if not dependents[dep] and num_needed[dep] == 1
+                    }
+                    if finish_now:
+                        deps -= finish_now  # Safe to mutate
+                        if len(finish_now) > 1:
+                            finish_now = sorted(finish_now, key=finish_now_key)
+                        for dep in finish_now:
+                            result[dep] = i
+                            i += 1
+                        add_to_inner_stack = False
+
+                if deps:
+                    for dep in deps:
+                        num_needed[dep] -= 1
+
+                    already_seen = deps & seen
+                    if already_seen:
+                        if len(deps) == len(already_seen):
+                            continue
+                        add_to_inner_stack = False
+                        deps -= already_seen
+
+                    if len(deps) == 1:
+                        # Fast path!  We trim down `deps` above hoping to reach here.
+                        (dep,) = deps
+                        if not inner_stack:
+                            if add_to_inner_stack:
+                                inner_stack = [dep]
+                                inner_stack_pop = inner_stack.pop
+                                seen_add(dep)
+                                continue
+                            key = partition_keys[dep]
+                        else:
+                            key = partition_keys[dep]
+                            if key < partition_keys[inner_stack[0]]:
+                                # Run before `inner_stack` (change tactical goal!)
+                                inner_stacks_append(inner_stack)
+                                inner_stack = [dep]
+                                inner_stack_pop = inner_stack.pop
+                                seen_add(dep)
+                                continue
+                        if key < partition_keys[item]:
+                            next_nodes[key].append(deps)
+                        else:
+                            later_nodes[key].append(deps)
+                    else:
+                        # Slow path :(.  This requires grouping by partition_key.
+                        dep_pools = defaultdict(list)
+                        for dep in deps:
+                            dep_pools[partition_keys[dep]].append(dep)
+                        item_key = partition_keys[item]
+                        if inner_stack:
+                            # If we have an inner_stack, we need to look for a "better" path
+                            prev_key = partition_keys[inner_stack[0]]
+                            now_keys = []  # < inner_stack[0]
+                            for key, vals in dep_pools.items():
+                                if key < prev_key:
+                                    now_keys.append(key)
+                                elif key < item_key:
+                                    next_nodes[key].append(vals)
+                                else:
+                                    later_nodes[key].append(vals)
+                            if now_keys:
+                                # Run before `inner_stack` (change tactical goal!)
+                                inner_stacks_append(inner_stack)
+                                if 1 < len(now_keys):
+                                    now_keys.sort(reverse=True)
+                                for key in now_keys:
+                                    pool = dep_pools[key]
+                                    if 1 < len(pool) < 100:
+                                        pool.sort(key=dependents_key, reverse=True)
+                                    inner_stacks_extend([dep] for dep in pool)
+                                    seen_update(pool)
+                                inner_stack = inner_stacks_pop()
+                                inner_stack_pop = inner_stack.pop
+                        else:
+                            # If we don't have an inner_stack, then we don't need to look
+                            # for a "better" path, but we do need traverse along dependents.
+                            if add_to_inner_stack:
+                                min_key = min(dep_pools)
+                                min_pool = dep_pools.pop(min_key)
+                                if len(min_pool) == 1:
+                                    inner_stack = min_pool
+                                    seen_update(inner_stack)
+                                elif (
+                                    10 * item_key
+                                    > 11 * len(min_pool) * len(min_pool) * min_key
+                                ):
+                                    # Put all items in min_pool onto inner_stacks.
+                                    # I know this is a weird comparison.  Hear me out.
+                                    # Although it is often beneficial to put all of the items in `min_pool`
+                                    # onto `inner_stacks` to process next, it is very easy to be overzealous.
+                                    # Sometimes it is actually better to defer until `next_nodes` is handled.
+                                    # We should only put items onto `inner_stacks` that we're reasonably
+                                    # confident about.  The above formula is a best effort heuristic given
+                                    # what we have easily available.  It is obviously very specific to our
+                                    # choice of partition_key.  Dask tests take this route about 40%.
+                                    if len(min_pool) < 100:
+                                        min_pool.sort(key=dependents_key, reverse=True)
+                                    inner_stacks_extend([dep] for dep in min_pool)
+                                    inner_stack = inner_stacks_pop()
+                                    seen_update(min_pool)
+                                else:
+                                    # Put one item in min_pool onto inner_stack and the rest into next_nodes.
+                                    if len(min_pool) < 100:
+                                        inner_stack = [
+                                            min(min_pool, key=dependents_key)
+                                        ]
+                                    else:
+                                        inner_stack = [min_pool.pop()]
+                                    next_nodes[min_key].append(min_pool)
+                                    seen_update(inner_stack)
+
+                                inner_stack_pop = inner_stack.pop
+                            for key, vals in dep_pools.items():
+                                if key < item_key:
+                                    next_nodes[key].append(vals)
+                                else:
+                                    later_nodes[key].append(vals)
 
         if len(dependencies) == len(result):
             break  # all done!
 
+        if next_nodes:
+            for key in sorted(next_nodes, reverse=True):
+                # `outer_stacks` may not be empty here--it has data from previous `next_nodes`.
+                # Since we pop things off of it (onto `inner_nodes`), this means we handle
+                # multiple `next_nodes` in a LIFO manner.
+                outer_stack_extend(reversed(next_nodes[key]))
+            next_nodes = defaultdict(list)
+
+        while outer_stack:
+            # Try to add a few items to `inner_stacks`
+            deps = [x for x in outer_stack_pop() if x not in result]
+            if deps:
+                if 1 < len(deps) < 100:
+                    deps.sort(key=dependents_key, reverse=True)
+                inner_stacks_extend([dep] for dep in deps)
+                seen_update(deps)
+                break
+
+        if inner_stacks:
+            continue
+
+        if later_nodes:
+            # You know all those dependents with large keys we've been hanging onto to run "later"?
+            # Well, "later" has finally come.
+            next_nodes, later_nodes = later_nodes, next_nodes
+            continue
+
         # We just finished computing a connected group.
         # Let's choose the first `item` in the next group to compute.
         # If we have few large groups left, then it's best to find `item` by taking a minimum.
@@ -348,6 +487,7 @@ def dependencies_key(x):
         item = init_stack_pop()
         while item in result:
             item = init_stack_pop()
+        inner_stacks_append([item])
 
     return result
 
@@ -397,19 +537,19 @@ def graph_metrics(dependencies, dependents, total_dependencies):
 
     4.  The minimum height from a root node
 
-        1
+        0
         |
-        2   1
+        1   0
          \ /
-          2
+          1
 
     5.  The maximum height from a root node
 
-        1
+        0
         |
-        2   1
+        1   0
          \ /
-          3
+          2
 
     Examples
     --------
@@ -418,7 +558,7 @@ def graph_metrics(dependencies, dependents, total_dependencies):
     >>> _, total_dependencies = ndependencies(dependencies, dependents)
     >>> metrics = graph_metrics(dependencies, dependents, total_dependencies)
     >>> sorted(metrics.items())
-    [('a1', (4, 2, 3, 2, 3)), ('b1', (2, 3, 3, 2, 2)), ('b2', (1, 2, 2, 1, 1)), ('c1', (1, 3, 3, 1, 1))]
+    [('a1', (4, 2, 3, 1, 2)), ('b1', (2, 3, 3, 1, 1)), ('b2', (1, 2, 2, 0, 0)), ('c1', (1, 3, 3, 0, 0))]
 
     Returns
     -------
@@ -432,7 +572,7 @@ def graph_metrics(dependencies, dependents, total_dependencies):
     for key, deps in dependents.items():
         if not deps:
             val = total_dependencies[key]
-            result[key] = (1, val, val, 1, 1)
+            result[key] = (1, val, val, 0, 0)
             for child in dependencies[key]:
                 num_needed[child] -= 1
                 if not num_needed[child]:


Test patch:
diff --git a/dask/tests/test_order.py b/dask/tests/test_order.py
index 98e450f48e3..b3c77f225e4 100644
--- a/dask/tests/test_order.py
+++ b/dask/tests/test_order.py
@@ -656,3 +656,138 @@ def test_order_cycle():
 
 def test_order_empty():
     assert order({}) == {}
+
+
+def test_switching_dependents(abcde):
+    r"""
+
+    a7 a8  <-- do these last
+    | /
+    a6                e6
+    |                /
+    a5   c5    d5  e5
+    |    |    /   /
+    a4   c4 d4  e4
+    |  \ | /   /
+    a3   b3---/
+    |
+    a2
+    |
+    a1
+    |
+    a0  <-- start here
+
+    Test that we are able to switch to better dependents.
+    In this graph, we expect to start at a0.  To compute a4, we need to compute b3.
+    After computing b3, three "better" paths become available.
+    Confirm that we take the better paths before continuing down `a` path.
+
+    This test is pretty specific to how `order` is implemented
+    and is intended to increase code coverage.
+    """
+    a, b, c, d, e = abcde
+    dsk = {
+        (a, 0): 0,
+        (a, 1): (f, (a, 0)),
+        (a, 2): (f, (a, 1)),
+        (a, 3): (f, (a, 2)),
+        (a, 4): (f, (a, 3), (b, 3)),
+        (a, 5): (f, (a, 4)),
+        (a, 6): (f, (a, 5)),
+        (a, 7): (f, (a, 6)),
+        (a, 8): (f, (a, 6)),
+        (b, 3): 1,
+        (c, 4): (f, (b, 3)),
+        (c, 5): (f, (c, 4)),
+        (d, 4): (f, (b, 3)),
+        (d, 5): (f, (d, 4)),
+        (e, 4): (f, (b, 3)),
+        (e, 5): (f, (e, 4)),
+        (e, 6): (f, (e, 5)),
+    }
+    o = order(dsk)
+
+    assert o[(a, 0)] == 0  # probably
+    assert o[(a, 5)] > o[(c, 5)]
+    assert o[(a, 5)] > o[(d, 5)]
+    assert o[(a, 5)] > o[(e, 6)]
+
+
+def test_order_with_equal_dependents(abcde):
+    """ From https://github.com/dask/dask/issues/5859#issuecomment-608422198
+
+    See the visualization of `(maxima, argmax)` example from the above comment.
+
+    This DAG has enough structure to exercise more parts of `order`
+
+    """
+    a, b, c, d, e = abcde
+    dsk = {}
+    abc = [a, b, c, d]
+    for x in abc:
+        dsk.update(
+            {
+                (x, 0): 0,
+                (x, 1): (f, (x, 0)),
+                (x, 2, 0): (f, (x, 0)),
+                (x, 2, 1): (f, (x, 1)),
+            }
+        )
+        for i, y in enumerate(abc):
+            dsk.update(
+                {
+                    (x, 3, i): (f, (x, 2, 0), (y, 2, 1)),  # cross x and y
+                    (x, 4, i): (f, (x, 3, i)),
+                    (x, 5, i, 0): (f, (x, 4, i)),
+                    (x, 5, i, 1): (f, (x, 4, i)),
+                    (x, 6, i, 0): (f, (x, 5, i, 0)),
+                    (x, 6, i, 1): (f, (x, 5, i, 1)),
+                }
+            )
+    o = order(dsk)
+    total = 0
+    for x in abc:
+        for i in range(len(abc)):
+            val = o[(x, 6, i, 1)] - o[(x, 6, i, 0)]
+            assert val > 0  # ideally, val == 2
+            total += val
+    assert total <= 32  # ideally, this should be 2 * 16 = 32
+
+    # Add one to the end of the nine bundles
+    dsk2 = dict(dsk)
+    for x in abc:
+        for i in range(len(abc)):
+            dsk2[(x, 7, i, 0)] = (f, (x, 6, i, 0))
+    o = order(dsk2)
+    total = 0
+    for x in abc:
+        for i in range(len(abc)):
+            val = o[(x, 7, i, 0)] - o[(x, 6, i, 1)]
+            assert val > 0  # ideally, val == 3
+            total += val
+    assert total <= 165  # ideally, this should be 3 * 16 == 48
+
+    # Remove one from each of the nine bundles
+    dsk3 = dict(dsk)
+    for x in abc:
+        for i in range(len(abc)):
+            del dsk3[(x, 6, i, 1)]
+    o = order(dsk3)
+    total = 0
+    for x in abc:
+        for i in range(len(abc)):
+            val = o[(x, 6, i, 0)] - o[(x, 5, i, 1)]
+            assert val > 0  # ideally, val == 2
+            total += val
+    assert total <= 119  # ideally, this should be 2 * 16 == 32
+
+    # Remove another one from each of the nine bundles
+    dsk4 = dict(dsk3)
+    for x in abc:
+        for i in range(len(abc)):
+            del dsk4[(x, 6, i, 0)]
+    o = order(dsk4)
+    total = 0
+    for x in abc:
+        for i in range(len(abc)):
+            assert o[(x, 5, i, 1)] - o[(x, 5, i, 0)] == 1