Repository: dask/dask
Pull request #10535 (2023-09-26 16:49:44)

[dask.order] Reduce memory pressure for multi array reductions by releasing splitter tasks more eagerly
Related https://github.com/dask/dask/pull/10505
Closes https://github.com/dask/dask/issues/10384
Closes https://github.com/dask/dask/issues/9995


- It removes the ambiguity between dataframe and array computes by normalizing topologies by adding a final reducing node to all graphs (this requires a little more work)
- It adds branched out dependents to the singles list. this is not exactly what it was designed for but this allows us to process those tasks before opening a new branch of computation
- It removes a "change strategic goal" code path. Mostly because it is necessary to make the code work but also because I currently distrust the partition_keys target function


Edit: The above description is a little outdated. See comment farthest below and inline comments for an accurate description

fjetter: At least on smaller graphs, this is fixing the root task problem mentioned in https://github.com/dask/dask/issues/9995 by ensuring to finish a computation branch before opening a new one
![anom-mean-order-raw-colored](https://github.com/dask/dask/assets/8629629/909a24e7-aeba-46db-8011-0019a0db188b)

The xarray style reduction reported in https://github.com/dask/dask/issues/10384 now looks like 
![xarray-reduction-array-colored](https://github.com/dask/dask/assets/8629629/af8153a7-f38d-4889-ad51-d6d368f5ef6e)
the simulated pressure for dataframes and arrays are not identical due to the adjustment to the graph topologies. What's even better, the pressure with this PR is now smaller than the DF pressure before.

I will move to real world testing/benchmarking tomorrow. If that looks good, I will have to do a little refactoring since there is now plenty of dead code I'm effectively skipping due to the root task addition. 🤞 


fjetter: Ok, I got through the unit tests and think everything is in order. I got sidetracked a little when looking at the cost functions because the graph_metrics confused me a little. Turns out (obvious in hindsight) they are only accurate for trees and are double counting quite a bit. Oh well, now there are a couple of tests. I also changed the definition slightly in a way that makes more sense to me. I may break this out into a dedicated PR, this is quite unrelated to the important stuff in here
mrocklin: Cool.  How excited should I be about this?  "Solves all Xarray pain?"  "Solves a couple of problems we've seen recently, but still lots more to do?"
fjetter: > Cool. How excited should I be about this? "Solves all Xarray pain?" "Solves a couple of problems we've seen recently, but still lots more to do?"

Let's wait for the first round of benchmarks before I answer :)
fjetter: OK, so benchmark results will be a little delayed. There is some kind of issue with our CI.

From manual testing I can say that the xarray example that's been reported in https://github.com/pangeo-data/distributed-array-examples/issues/2 now runs smoothly with constant memory.
However, if I dial up the size, we are running again into a bad ordering problem. I suspect something with the graph metrics is off.

(Legend below: scaling the time dimension / number of tasks by that factor; The green line is only constant because we're spilling)

![image](https://github.com/dask/dask/assets/8629629/1f83ea38-c24a-42be-b628-3199c02c6812)

The up side is that on main, the 1.0 measure was already saturating memory fully... Looks like I have some more work to do
fjetter: Well, at the very least simulated memory pressure, i.e. the number of tasks a single threaded scheduler would have to hold in memory, is behaving reasonably well. Not constant but no [longer linear](https://github.com/dask/dask/issues/10384#issuecomment-1715841811) and even better than the DF version before.

(This goes up to 100_000, i.e. 20x more than the reported example)

![image](https://github.com/dask/dask/assets/8629629/d1ee6010-9f38-499f-bc37-11c9d8014e99)

fjetter: I'm pretty sure that what I'm seeing right now is not an ordering problem anymore. The behavior appears to be quite flaky. I assume this there is some kind distributed scheduler state that is impacting the worker placement decision.

For instance, I now managed to get another round of runs through where the _small_ dataset actually misbehaved but I didn't have an issue scaling up to 5 times the data (i.e. total input data is 10x more than the cluster memory)

![image](https://github.com/dask/dask/assets/8629629/e23cea35-71ac-4702-97f4-9a317d6c1d6a)



fjetter: OK, turns out it is an ordering problem after all. Order appears to not be stable.

I patched the scheduler to compute a hash of the order result and compute the simulated pressure and log this. The hash is always different. The pressure is sometimes bad and when it is, the run is horrible.

Just running the same thing in a for loop for a couple of times

![image](https://github.com/dask/dask/assets/8629629/9bb25fef-073f-4d6d-8456-8e58df5f3f3a)

At least I know what to look for now... kind of

fjetter: Well... good and bad news. Order seems to be fine (ish?). If I disable optimization, everything seems to be very, very stable. However, once optimization is on, this becomes flaky!
going through the code, once I remove/comment out `optimize_blockwise` everything works fine. With `optimize_blockwise` this is flaky

fjetter: By best guess is that this is connected to keys somehow. This issue is only flaky if I regenerate the graph, i.e. regenerate the keys. We probably no longer have sufficient structural information to break ties w/out str comp after the fusion.
fjetter: TLDR I shouldn't have touch the graph_metrics. However, everything is sooo fragile that there is more work to do still

Who was I to think touching the graph_metrics definition was a good idea? Stupid me! The two diverging orderings were indeed triggered by a str comparison but the fundamental bug was introduced by me changing the graph_metrics. I don't know if this will ever be relevant but for my own sanity, I want to document what happened. Interested readers may skip this safely.

I was looking at why these two graphs diverged

raw graph
![first](https://github.com/dask/dask/assets/8629629/23001bf2-6031-4b94-8ae8-7ffabaf6aebe)

snapshot good ordering
![image](https://github.com/dask/dask/assets/8629629/a3b01ed3-de96-4291-82bb-5ad9a8d3e3fe)


snapshot bad ordering

![image](https://github.com/dask/dask/assets/8629629/d1d22a65-c853-41e1-9e62-e5e5a0ccb511)


The good ordering properly targeted the first reducer and (not shown here) would clean up after itself. That's perfect. The bad ordering regressed into a breadth first search that was releasing roots relatively quickly but wouldn't go deep, such that we cannot truly release anything until rather late.

The divergence happens already at the very fist decision point. The initial root node is random (str compared) because the graph is pretty much symmetrical. It would then decide which one of the dependents to pick by using `dependents_key` which is identical for both after my graph_metrics changes (because with my changes, _coincidentally_ min_heights and total_dependents is equal causing this to zero out https://github.com/dask/dask/blob/2b45b215e13e6d7893853ef33cb4b1796c477a9e/dask/order.py#L224C47-L224C47)
This decision here https://github.com/dask/dask/blob/2b45b215e13e6d7893853ef33cb4b1796c477a9e/dask/order.py#L593-L599

which leads us to these if/elif/else clauses

https://github.com/dask/dask/blob/2b45b215e13e6d7893853ef33cb4b1796c477a9e/dask/order.py#L640-L652

The thing is, these dependents are **not** identical since one of them can already be run while the other has a dependency to load first. Depending on which one of the two is `dep2` the `if not num_needed[dep2]` clause will evaluate differently.
The bad case actually runs in to the `num_needed[dep2] == 0` zero case and puts a "single" relation on the other dependency
https://github.com/dask/dask/blob/2b45b215e13e6d7893853ef33cb4b1796c477a9e/dask/order.py#L645 which forces `order` to work on this "single" _before_ loading any more dependents. Effectively these _singles_ is a deviation from the depth first search (which is sometimes necessary to free memory).
The good path actually picks the `key == key2 and 5 * partition_keys[item] > 22 * key` path (where do these magical numbers come from) https://github.com/fjetter/dask/blob/e71033821ecedab254750bd6f158714ba589acf8/dask/order.py#L650
which does almost the same thing but doesn't put a "single" relationship on the other dep.

Alas, reverting the graph_metrics change ensure that both graphs are now producing identical orderings. The chosen path is actually the bad one... but once we move to larger graphs it appears to work again. This just shows that there is still some work to do...
fjetter: alright, the above example now works nicely with constant memory. The below shows runs with increasingly more data/tasks. Runtime is increasing linearily and memory is constant. Just the way it's supposed to.

![image](https://github.com/dask/dask/assets/8629629/dcb358b2-b5d2-435f-a117-bf0b0915ae7a)

dcherian: ❤️ 

Patch:
diff --git a/dask/array/core.py b/dask/array/core.py
index 1a1122d96..283faefec 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -5803,10 +5803,10 @@ class BlockView:
 
         keys = product(*(range(len(c)) for c in chunks))
 
-        layer = {(name,) + key: tuple(new_keys[key].tolist()) for key in keys}
+        graph: Graph = {(name,) + key: tuple(new_keys[key].tolist()) for key in keys}
 
-        graph = HighLevelGraph.from_collections(name, layer, dependencies=[self._array])
-        return Array(graph, name, chunks, meta=self._array)
+        hlg = HighLevelGraph.from_collections(name, graph, dependencies=[self._array])
+        return Array(hlg, name, chunks, meta=self._array)
 
     def __eq__(self, other: Any) -> bool:
         if isinstance(other, BlockView):
diff --git a/dask/bag/core.py b/dask/bag/core.py
index 4603ae57c..e3a6f675d 100644
--- a/dask/bag/core.py
+++ b/dask/bag/core.py
@@ -7,7 +7,7 @@ import operator
 import uuid
 import warnings
 from collections import defaultdict
-from collections.abc import Iterable, Iterator, Mapping, Sequence
+from collections.abc import Iterable, Iterator, Sequence
 from functools import partial, reduce, wraps
 from random import Random
 from urllib.request import urlopen
@@ -469,7 +469,7 @@ class Bag(DaskMethodsMixin):
     30
     """
 
-    def __init__(self, dsk: Mapping, name: str, npartitions: int):
+    def __init__(self, dsk: Graph, name: str, npartitions: int):
         if not isinstance(dsk, HighLevelGraph):
             dsk = HighLevelGraph.from_collections(name, dsk, dependencies=[])
         self.dask = dsk
diff --git a/dask/dataframe/io/csv.py b/dask/dataframe/io/csv.py
index 1d5189e8b..5a847789c 100644
--- a/dask/dataframe/io/csv.py
+++ b/dask/dataframe/io/csv.py
@@ -486,10 +486,12 @@ def read_pandas(
         lineterminator = "\n"
     if include_path_column and isinstance(include_path_column, bool):
         include_path_column = "path"
-    if "index" in kwargs or "index_col" in kwargs:
+    if "index" in kwargs or (
+        "index_col" in kwargs and kwargs.get("index_col") is not False
+    ):
         raise ValueError(
-            "Keywords 'index' and 'index_col' not supported. "
-            f"Use dd.{reader_name}(...).set_index('my-index') instead"
+            "Keywords 'index' and 'index_col' not supported, except for "
+            "'index_col=False'. Use dd.{reader_name}(...).set_index('my-index') instead"
         )
     for kw in ["iterator", "chunksize"]:
         if kw in kwargs:
diff --git a/dask/graph_manipulation.py b/dask/graph_manipulation.py
index a80fd8112..3c4cbec8d 100644
--- a/dask/graph_manipulation.py
+++ b/dask/graph_manipulation.py
@@ -20,7 +20,7 @@ from dask.blockwise import blockwise
 from dask.core import flatten
 from dask.delayed import Delayed, delayed
 from dask.highlevelgraph import HighLevelGraph, Layer, MaterializedLayer
-from dask.typing import Key
+from dask.typing import Graph, Key
 
 __all__ = ("bind", "checkpoint", "clone", "wait_on")
 
@@ -78,7 +78,7 @@ def _checkpoint_one(collection, split_every) -> Delayed:
         next(keys_iter)
     except StopIteration:
         # Collection has 0 or 1 keys; no need for a map step
-        layer = {name: (chunks.checkpoint, collection.__dask_keys__())}
+        layer: Graph = {name: (chunks.checkpoint, collection.__dask_keys__())}
         dsk = HighLevelGraph.from_collections(name, layer, dependencies=(collection,))
         return Delayed(name, dsk)
 
diff --git a/dask/highlevelgraph.py b/dask/highlevelgraph.py
index 19aa3befe..eb13ff013 100644
--- a/dask/highlevelgraph.py
+++ b/dask/highlevelgraph.py
@@ -3,7 +3,18 @@ from __future__ import annotations
 import abc
 import copy
 import html
-from collections.abc import Collection, Hashable, Iterable, KeysView, Mapping, Set
+from collections.abc import (
+    Collection,
+    Hashable,
+    ItemsView,
+    Iterable,
+    Iterator,
+    KeysView,
+    Mapping,
+    Sequence,
+    Set,
+    ValuesView,
+)
 from typing import Any
 
 import tlz as toolz
@@ -12,7 +23,7 @@ import dask
 from dask import config
 from dask.base import clone_key, flatten, is_dask_collection
 from dask.core import keys_in_tasks, reverse_dict
-from dask.typing import Graph, Key
+from dask.typing import DaskCollection, Graph, Key
 from dask.utils import ensure_dict, import_required, key_split
 from dask.widgets import get_template
 
@@ -34,7 +45,7 @@ def compute_layer_dependencies(layers):
     return ret
 
 
-class Layer(Mapping):
+class Layer(Graph):
     """High level graph layer
 
     This abstract class establish a protocol for high level graph layers.
@@ -327,7 +338,7 @@ class MaterializedLayer(Layer):
         return self.keys()
 
 
-class HighLevelGraph(Mapping):
+class HighLevelGraph(Graph):
     """Task graph composed of layers of dependent subgraphs
 
     This object encodes a Dask task graph that is composed of layers of
@@ -433,7 +444,12 @@ class HighLevelGraph(Mapping):
         return cls(layers, deps)
 
     @classmethod
-    def from_collections(cls, name, layer, dependencies=()):
+    def from_collections(
+        cls,
+        name: str,
+        layer: Graph,
+        dependencies: Sequence[DaskCollection] = (),
+    ) -> HighLevelGraph:
         """Construct a HighLevelGraph from a new layer and a set of collections
 
         This constructs a HighLevelGraph in the common case where we have a single
@@ -470,34 +486,35 @@ class HighLevelGraph(Mapping):
         if len(dependencies) == 1:
             return cls._from_collection(name, layer, dependencies[0])
         layers = {name: layer}
-        deps = {name: set()}
+        name_dep: set[str] = set()
+        deps: dict[str, Set[str]] = {name: name_dep}
         for collection in toolz.unique(dependencies, key=id):
             if is_dask_collection(collection):
                 graph = collection.__dask_graph__()
                 if isinstance(graph, HighLevelGraph):
                     layers.update(graph.layers)
                     deps.update(graph.dependencies)
-                    deps[name] |= set(collection.__dask_layers__())
+                    name_dep |= set(collection.__dask_layers__())
                 else:
                     key = _get_some_layer_name(collection)
                     layers[key] = graph
-                    deps[name].add(key)
+                    name_dep.add(key)
                     deps[key] = set()
             else:
                 raise TypeError(type(collection))
 
         return cls(layers, deps)
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: Key) -> Any:
         # Attempt O(1) direct access first, under the assumption that layer names match
         # either the keys (Scalar, Item, Delayed) or the first element of the key tuples
         # (Array, Bag, DataFrame, Series). This assumption is not always true.
         try:
-            return self.layers[key][key]
+            return self.layers[key][key]  # type: ignore
         except KeyError:
             pass
         try:
-            return self.layers[key[0]][key]
+            return self.layers[key[0]][key]  # type: ignore
         except (KeyError, IndexError, TypeError):
             pass
 
@@ -518,10 +535,10 @@ class HighLevelGraph(Mapping):
         # https://github.com/dask/dask/issues/7271
         return sum(len(layer) for layer in self.layers.values())
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[Key]:
         return iter(self.to_dict())
 
-    def to_dict(self) -> dict:
+    def to_dict(self) -> dict[Key, Any]:
         """Efficiently convert to plain dict. This method is faster than dict(self)."""
         try:
             return self._to_dict
@@ -537,7 +554,7 @@ class HighLevelGraph(Mapping):
         """
         return self.to_dict().keys()
 
-    def get_all_external_keys(self) -> set:
+    def get_all_external_keys(self) -> set[Key]:
         """Get all output keys of all layers
 
         This will in most cases _not_ materialize any layers, which makes
@@ -560,10 +577,10 @@ class HighLevelGraph(Mapping):
             self._all_external_keys = keys
             return keys
 
-    def items(self):
+    def items(self) -> ItemsView[Key, Any]:
         return self.to_dict().items()
 
-    def values(self):
+    def values(self) -> ValuesView[Any]:
         return self.to_dict().values()
 
     def get_all_dependencies(self) -> dict[Key, Set[Key]]:
@@ -586,10 +603,10 @@ class HighLevelGraph(Mapping):
         return self.key_dependencies
 
     @property
-    def dependents(self):
+    def dependents(self) -> dict[str, set[str]]:
         return reverse_dict(self.dependencies)
 
-    def copy(self):
+    def copy(self) -> HighLevelGraph:
         return HighLevelGraph(
             ensure_dict(self.layers, copy=True),
             ensure_dict(self.dependencies, copy=True),
@@ -597,16 +614,16 @@ class HighLevelGraph(Mapping):
         )
 
     @classmethod
-    def merge(cls, *graphs):
-        layers = {}
-        dependencies = {}
+    def merge(cls, *graphs: Graph) -> HighLevelGraph:
+        layers: dict[str, Graph] = {}
+        dependencies: dict[str, Set[str]] = {}
         for g in graphs:
             if isinstance(g, HighLevelGraph):
                 layers.update(g.layers)
                 dependencies.update(g.dependencies)
             elif isinstance(g, Mapping):
-                layers[id(g)] = g
-                dependencies[id(g)] = set()
+                layers[str(id(g))] = g
+                dependencies[str(id(g))] = set()
             else:
                 raise TypeError(g)
         return cls(layers, dependencies)
@@ -655,7 +672,7 @@ class HighLevelGraph(Mapping):
         graphviz_to_file(g, filename, format)
         return g
 
-    def _toposort_layers(self):
+    def _toposort_layers(self) -> list[str]:
         """Sort the layers in a high level graph topologically
 
         Parameters
@@ -669,7 +686,7 @@ class HighLevelGraph(Mapping):
             List of layer names sorted topologically
         """
         degree = {k: len(v) for k, v in self.dependencies.items()}
-        reverse_deps = {k: [] for k in self.dependencies}
+        reverse_deps: dict[str, list[str]] = {k: [] for k in self.dependencies}
         ready = []
         for k, v in self.dependencies.items():
             for dep in v:
@@ -686,7 +703,7 @@ class HighLevelGraph(Mapping):
                     ready.append(rdep)
         return ret
 
-    def cull(self, keys: Iterable) -> HighLevelGraph:
+    def cull(self, keys: Iterable[Key]) -> HighLevelGraph:
         """Return new HighLevelGraph with only the tasks required to calculate keys.
 
         In other words, remove unnecessary tasks from dask.
@@ -779,7 +796,7 @@ class HighLevelGraph(Mapping):
 
         return HighLevelGraph(ret_layers, ret_dependencies)
 
-    def validate(self):
+    def validate(self) -> None:
         # Check dependencies
         for layer_name, deps in self.dependencies.items():
             if layer_name not in self.layers:
@@ -820,7 +837,7 @@ class HighLevelGraph(Mapping):
             representation += f" {i}. {layerkey}\n"
         return representation
 
-    def _repr_html_(self):
+    def _repr_html_(self) -> str:
         return get_template("highlevelgraph.html.j2").render(
             type=type(self).__name__,
             layers=self.layers,
diff --git a/dask/order.py b/dask/order.py
index 46e8c3145..457841e49 100644
--- a/dask/order.py
+++ b/dask/order.py
@@ -80,7 +80,7 @@ Work towards *small goals* with *big steps*.
 from collections import defaultdict, namedtuple
 from math import log
 
-from dask.core import get_dependencies, get_deps, getcycle, reverse_dict
+from dask.core import get_dependencies, get_deps, getcycle, istask, reverse_dict
 
 
 def order(dsk, dependencies=None):
@@ -110,6 +110,7 @@ def order(dsk, dependencies=None):
     """
     if not dsk:
         return {}
+    dsk = dict(dsk)
 
     if dependencies is None:
         dependencies = {k: get_dependencies(dsk, k) for k in dsk}
@@ -133,7 +134,19 @@ def order(dsk, dependencies=None):
     # tree, we skip processing it normally.
     # See https://github.com/dask/dask/issues/6745
     root_nodes = {k for k, v in dependents.items() if not v}
-    skip_root_node = len(root_nodes) == 1 and len(dsk) > 1
+    if len(root_nodes) > 1:
+        # This is also nice because it makes us robust to difference when
+        # computing vs persisting collections
+        root = object()
+
+        def _f(*args, **kwargs):
+            pass
+
+        dsk[root] = (_f, *root_nodes)
+        dependencies[root] = root_nodes
+        o = order(dsk, dependencies)
+        del o[root]
+        return o
 
     # Leaf nodes.  We choose one--the initial node--for each weakly connected subgraph.
     # Let's calculate the `initial_stack_key` as we determine `init_stack` set.
@@ -141,10 +154,8 @@ def order(dsk, dependencies=None):
         # First prioritize large, tall groups, then prioritize the same as ``dependents_key``.
         key: (
             # at a high-level, work towards a large goal (and prefer tall and narrow)
-            -max_dependencies,
             num_dependents - max_heights,
             # 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
@@ -154,8 +165,8 @@ def order(dsk, dependencies=None):
         )
         for key, num_dependents, (
             total_dependents,
-            min_dependencies,
-            max_dependencies,
+            _,
+            _,
             min_heights,
             max_heights,
         ) in (
@@ -177,6 +188,9 @@ def order(dsk, dependencies=None):
         return (
             # Focus on being memory-efficient
             len(dependents[x]) - len(dependencies[x]) + num_needed[x],
+            # Do we favor deep or shallow branches?
+            #  -1: deep
+            #  +1: shallow
             -metrics[x][3],  # min_heights
             # tie-breaker
             StrComparable(x),
@@ -190,19 +204,15 @@ def order(dsk, dependencies=None):
         num_dependents = len(dependents[x])
         (
             total_dependents,
-            min_dependencies,
-            max_dependencies,
+            _,
+            _,
             min_heights,
             max_heights,
         ) = metrics[x]
         # Prefer short and narrow instead of tall in narrow, because we're going in
         # reverse along dependencies.
         return (
-            # at a high-level, work towards a large goal (and prefer short and narrow)
-            -max_dependencies,
             num_dependents + max_heights,
-            # tactically, finish small connected jobs first
-            min_dependencies,
             num_dependents + min_heights,  # prefer short and narrow
             -total_dependencies[x],  # go where the work is
             # try to be memory efficient
@@ -213,14 +223,11 @@ def order(dsk, dependencies=None):
             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))
-
+    root_total_dependencies = total_dependencies[list(root_nodes)[0]]
     # Computing this for all keys can sometimes be relatively expensive :(
     partition_keys = {
         key: (
-            (min_dependencies - total_dependencies[key] + 1)
+            (root_total_dependencies - total_dependencies[key] + 1)
             * (total_dependents - min_heights)
         )
         for key, (
@@ -258,13 +265,10 @@ def order(dsk, dependencies=None):
     # 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.
+    # When the inner stacks are depleted, we process `next_nodes`.
     # 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.
@@ -279,10 +283,7 @@ def order(dsk, dependencies=None):
 
     # Keep track of nodes that are in `inner_stack` or `inner_stacks` so we don't
     # process them again.
-    if skip_root_node:
-        seen = set(root_nodes)
-    else:
-        seen = set()  # seen in an inner_stack (and has dependencies)
+    seen = set(root_nodes)
     seen_update = seen.update
     seen_add = seen.add
 
@@ -325,7 +326,6 @@ def order(dsk, dependencies=None):
     # if they so choose?  Maybe.  However, I'm sensitive to the multithreaded scheduler,
     # which is heavily dependent on the ordering obtained here.
     singles = {}
-    singles_items = singles.items()
     singles_clear = singles.clear
     later_singles = []
     later_singles_append = later_singles.append
@@ -338,13 +338,13 @@ def order(dsk, dependencies=None):
     #   4. later_singles
     #   5. next_nodes
     #   6. outer_stack
-    #   7. later_nodes
-    #   8. init_stack
+    #   7. init_stack
 
     # alias for speed
     set_difference = set.difference
 
     is_init_sorted = False
+
     while True:
         while True:
             # Perform a DFS along dependencies until we complete our tactical goal
@@ -353,7 +353,7 @@ def order(dsk, dependencies=None):
                 if item in result:
                     continue
                 if num_needed[item]:
-                    if not skip_root_node or item not in root_nodes:
+                    if item not in root_nodes:
                         inner_stack.append(item)
                         deps = set_difference(dependencies[item], result)
                         if 1 < len(deps) < 1000:
@@ -365,6 +365,13 @@ def order(dsk, dependencies=None):
                         seen_update(deps)
                     if not singles:
                         continue
+                    # Only process singles once the inner_stack is fully
+                    # resolved. This is important because the singles path later
+                    # on verifies that running the single indeed opens an
+                    # opportunity to release soon by comparing the singles
+                    # parent's dependents with the inner_stack(s)
+                    if inner_stack and num_needed[inner_stack[-1]]:
+                        continue
                     process_singles = True
                 else:
                     result[item] = i
@@ -372,28 +379,6 @@ def order(dsk, dependencies=None):
                     deps = dependents[item]
                     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
-                        elif skip_root_node:
-                            for dep in root_nodes:
-                                num_needed[dep] -= 1
-                                # Use remove here to complain loudly if our assumptions change
-                                deps.remove(dep)  # Safe to mutate
-
                     if deps:
                         for dep in deps:
                             num_needed[dep] -= 1
@@ -417,26 +402,6 @@ def order(dsk, dependencies=None):
                         dep2 = dependents[single]
                         result[single] = i
                         i += 1
-                        if metrics[single][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 dep2
-                                if not dependents[dep] and num_needed[dep] == 1
-                            }
-                            if finish_now:
-                                dep2 -= 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
-                            elif skip_root_node:
-                                for dep in root_nodes:
-                                    num_needed[dep] -= 1
-                                    # Use remove here to complain loudly if our assumptions change
-                                    dep2.remove(dep)  # Safe to mutate
                         if dep2:
                             for dep in dep2:
                                 num_needed[dep] -= 1
@@ -460,51 +425,48 @@ def order(dsk, dependencies=None):
 
             if process_singles and singles:
                 # We gather all dependents of all singles into `deps`, which we then process below.
-                # A lingering question is: what should we use for `item`?  `item_key` is used to
-                # determine whether each dependent goes to `next_nodes` or `later_nodes`.  Currently,
-                # we use the last value of `item` (i.e., we don't do anything).
+
                 deps = set()
                 add_to_inner_stack = True if inner_stack or inner_stacks else False
-                for single, parent in singles_items:
-                    if single in result:
-                        continue
+                singles_keys = set_difference(set(singles), result)
+
+                # NOTE: If this was too slow, LIFO would be a decent
+                # approximation
+                for single in sorted(singles_keys, key=lambda x: partition_keys[x]):
+                    # We want to run the singles if they are either releasing a
+                    # dependency directly or that they may be releasing a
+                    # dependency once the current critical path / inner_stack is
+                    # walked.
+                    # By using `seen` here this is more permissive since it also
+                    # includes tasks in a future critical path / inner_stacks
+                    # but it would require additional state to make this
+                    # distinction and we don't have enough data to dermine if
+                    # this is worth it.
+                    parent = singles[single]
                     if (
-                        add_to_inner_stack
-                        and len(set_difference(dependents[parent], result)) > 1
+                        len(
+                            set_difference(
+                                set_difference(dependents[parent], result),
+                                seen,
+                            )
+                        )
+                        > 1
                     ):
                         later_singles_append(single)
                         continue
-
                     while True:
                         dep2 = dependents[single]
                         result[single] = i
                         i += 1
-                        if metrics[single][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 dep2
-                                if not dependents[dep] and num_needed[dep] == 1
-                            }
-                            if finish_now:
-                                dep2 -= 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
-                            elif skip_root_node:
-                                for dep in root_nodes:
-                                    num_needed[dep] -= 1
-                                    # Use remove here to complain loudly if our assumptions change
-                                    dep2.remove(dep)  # Safe to mutate
                         if dep2:
                             for dep in dep2:
                                 num_needed[dep] -= 1
                             if add_to_inner_stack:
                                 already_seen = dep2 & seen
                                 if already_seen:
+                                    # This means that the singles path also
+                                    # leads to the current or previous strategic
+                                    # path
                                     if len(dep2) == len(already_seen):
                                         if len(already_seen) == 1:
                                             (single,) = already_seen
@@ -535,8 +497,7 @@ def order(dsk, dependencies=None):
                 add_to_inner_stack = False
 
             # 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" above
-            # or if a dependent is already on an inner_stack.  In this case, we add the
+            # However, we don't add to it 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,
@@ -556,30 +517,18 @@ def order(dsk, dependencies=None):
             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 add_to_inner_stack and not inner_stack:
+                    inner_stack = [dep]
+                    inner_stack_pop = inner_stack.pop
+                    seen_add(dep)
+                    continue
+                key = partition_keys[dep]
                 if not num_needed[dep]:
                     # We didn't put the single dependency on the stack, but we should still
                     # run it soon, because doing so may free its parent.
                     singles[dep] = item
-                elif key < partition_keys[item]:
-                    next_nodes[key].append(deps)
                 else:
-                    later_nodes[key].append(deps)
+                    next_nodes[key].append(deps)
             elif len(deps) == 2:
                 # We special-case when len(deps) == 2 so that we may place a dep on singles.
                 # Otherwise, the logic here is the same as when `len(deps) > 2` below.
@@ -620,60 +569,43 @@ def order(dsk, dependencies=None):
                                 later_singles_append(dep2)
                             else:
                                 singles[dep2] = item
-                        elif key2 < partition_keys[item]:
-                            next_nodes[key2].append([dep2])
                         else:
-                            later_nodes[key2].append([dep2])
+                            next_nodes[key2].append([dep2])
                     else:
                         item_key = partition_keys[item]
-                        if key2 < item_key:
-                            next_nodes[key].append([dep])
-                            next_nodes[key2].append([dep2])
-                        elif key < item_key:
-                            next_nodes[key].append([dep])
-                            later_nodes[key2].append([dep2])
-                        else:
-                            later_nodes[key].append([dep])
-                            later_nodes[key2].append([dep2])
+                        for k, d in [(key, dep), (key2, dep2)]:
+                            if not num_needed[d]:
+                                if process_singles:
+                                    later_singles_append(d)
+                                else:
+                                    singles[d] = item
+                            else:
+                                next_nodes[k].append([d])
                 else:
+                    assert not inner_stack
                     if add_to_inner_stack:
+                        inner_stack = [dep]
+                        inner_stack_pop = inner_stack.pop
+                        seen_add(dep)
                         if not num_needed[dep2]:
-                            inner_stacks_append(inner_stack)
-                            inner_stack = [dep]
-                            inner_stack_pop = inner_stack.pop
-                            seen_add(dep)
                             singles[dep2] = item
                         elif key == key2 and 5 * partition_keys[item] > 22 * key:
-                            inner_stacks_append(inner_stack)
                             inner_stacks_append([dep2])
-                            inner_stack = [dep]
-                            inner_stack_pop = inner_stack.pop
-                            seen_update(deps)
+                            seen_add(dep2)
                         else:
-                            inner_stacks_append(inner_stack)
-                            inner_stack = [dep]
-                            inner_stack_pop = inner_stack.pop
-                            seen_add(dep)
-                            if key2 < partition_keys[item]:
-                                next_nodes[key2].append([dep2])
-                            else:
-                                later_nodes[key2].append([dep2])
-                    else:
-                        item_key = partition_keys[item]
-                        if key2 < item_key:
-                            next_nodes[key].append([dep])
                             next_nodes[key2].append([dep2])
-                        elif key < item_key:
-                            next_nodes[key].append([dep])
-                            later_nodes[key2].append([dep2])
-                        else:
-                            later_nodes[key].append([dep])
-                            later_nodes[key2].append([dep2])
+                    else:
+                        for k, d in [(key, dep), (key2, dep2)]:
+                            next_nodes[k].append([d])
             else:
                 # Slow path :(.  This requires grouping by partition_key.
-                dep_pools = defaultdict(list)
+                dep_pools = defaultdict(set)
+                possible_singles = defaultdict(set)
                 for dep in deps:
-                    dep_pools[partition_keys[dep]].append(dep)
+                    pkey = partition_keys[dep]
+                    if not num_needed[dep] and not process_singles:
+                        possible_singles[pkey].add(dep)
+                    dep_pools[pkey].add(dep)
                 item_key = partition_keys[item]
                 if inner_stack:
                     # If we have an inner_stack, we need to look for a "better" path
@@ -682,10 +614,12 @@ def order(dsk, dependencies=None):
                     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)
+                            psingles = possible_singles[key]
+                            for s in psingles:
+                                singles[s] = item
+                            vals -= psingles
+                            next_nodes[key].append(vals)
                     if now_keys:
                         # Run before `inner_stack` (change tactical goal!)
                         inner_stacks_append(inner_stack)
@@ -694,7 +628,7 @@ def order(dsk, dependencies=None):
                         for key in now_keys:
                             pool = dep_pools[key]
                             if 1 < len(pool) < 100:
-                                pool.sort(key=dependents_key, reverse=True)
+                                pool = sorted(pool, key=dependents_key, reverse=True)
                             inner_stacks_extend([dep] for dep in pool)
                             seen_update(pool)
                         inner_stack = inner_stacks_pop()
@@ -706,7 +640,7 @@ def order(dsk, dependencies=None):
                         min_key = min(dep_pools)
                         min_pool = dep_pools.pop(min_key)
                         if len(min_pool) == 1:
-                            inner_stack = min_pool
+                            inner_stack = list(min_pool)
                             seen_update(inner_stack)
                         elif (
                             10 * item_key > 11 * len(min_pool) * len(min_pool) * min_key
@@ -721,7 +655,9 @@ def order(dsk, dependencies=None):
                             # 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)
+                                min_pool = sorted(
+                                    min_pool, key=dependents_key, reverse=True
+                                )
                             inner_stacks_extend([dep] for dep in min_pool)
                             inner_stack = inner_stacks_pop()
                             seen_update(min_pool)
@@ -736,10 +672,11 @@ def order(dsk, dependencies=None):
 
                         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)
+                        psingles = possible_singles[key]
+                        for s in psingles:
+                            singles[s] = item
+                        vals -= psingles
+                        next_nodes[key].append(vals)
 
         if len(dependencies) == len(result):
             break  # all done!
@@ -765,12 +702,6 @@ def order(dsk, dependencies=None):
         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.
@@ -795,7 +726,7 @@ def order(dsk, dependencies=None):
             init_stack_pop = init_stack.pop
             is_init_sorted = True
 
-        if skip_root_node and item in root_nodes:
+        if item in root_nodes:
             item = init_stack_pop()
 
         while item in result:
@@ -819,8 +750,10 @@ def graph_metrics(dependencies, dependents, total_dependencies):
     For each key we return:
 
     1.  **total_dependents**: The number of keys that can only be run
-        after this key is run.  The root nodes have value 1 while deep child
-        nodes will have larger values.
+        after this key is run.
+        Note that this is only exact for trees. (undirected) cycles will cause
+        double counting of nodes. Therefore, this metric is an upper bound
+        approximation.
 
         1
         |
@@ -1076,3 +1009,40 @@ def diagnostics(dsk, o=None, dependencies=None):
         for key, val in o.items()
     }
     return rv, pressure
+
+
+def _f():
+    ...
+
+
+def _convert_task(task):
+    if istask(task):
+        assert callable(task[0])
+        new_spec = []
+        for el in task[1:]:
+            if isinstance(el, (str, int)):
+                new_spec.append(el)
+            elif isinstance(el, tuple):
+                if istask(el):
+                    new_spec.append(_convert_task(el))
+                else:
+                    new_spec.append(el)
+            elif isinstance(el, list):
+                new_spec.append([_convert_task(e) for e in el])
+        return (_f, *new_spec)
+    else:
+        return task
+
+
+def sanitize_dsk(dsk):
+    """Take a dask graph and replace callables with a dummy function and remove
+    payload data like numpy arrays, dataframes, etc.
+    """
+    new = {}
+    for key, values in dsk.items():
+        new_key = key
+        new[new_key] = _convert_task(values)
+    if get_deps(new) != get_deps(dsk):
+        # The switch statement in _convert likely dropped some keys
+        raise RuntimeError("Sanitization failed to preserve topology.")
+    return new
diff --git a/dask/utils.py b/dask/utils.py
index 6713fb3b8..e2ebb6a58 100644
--- a/dask/utils.py
+++ b/dask/utils.py
@@ -1851,10 +1851,11 @@ def key_split(s):
     >>> key_split('_(x)')  # strips unpleasant characters
     'x'
     """
+    # If we convert the key, recurse to utilize LRU cache better
     if type(s) is bytes:
-        s = s.decode()
+        return key_split(s.decode())
     if type(s) is tuple:
-        s = s[0]
+        return key_split(s[0])
     try:
         words = s.split("-")
         if not words[0][0].isalpha():
@@ -1873,7 +1874,7 @@ def key_split(s):
         else:
             if result[0] == "<":
                 result = result.strip("<>").split()[0].split(".")[-1]
-            return result
+            return sys.intern(result)
     except Exception:
         return "Other"
 
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
index c15083aaa..23c057ac2 100644
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
@@ -1,6 +1,39 @@
 Changelog
 =========
 
+.. _v2023.9.3:
+
+2023.9.3
+--------
+
+Released on September 29, 2023
+
+Highlights
+^^^^^^^^^^
+
+Restore previous configuration override behavior
+""""""""""""""""""""""""""""""""""""""""""""""""
+The 2023.9.2 release introduced an unintentional breaking change in
+how configuration options are overriden in ``dask.config.get`` with
+the ``override_with=`` keyword (see :issue:`10519`).
+This release restores the previous behavior.
+
+See :pr:`10521` from `crusaderky`_ for details.
+    
+Complex dtypes in Dask Array reductions
+"""""""""""""""""""""""""""""""""""""""
+This release includes improved support for using common reductions
+in Dask Array (e.g. ``var``, ``std``, ``moment``) with complex dtypes.
+
+See :pr:`10009` from `wkrasnicki`_ for details.
+
+.. dropdown:: Additional changes
+
+    - Bump ``actions/checkout`` from 4.0.0 to 4.1.0 (:pr:`10532`)
+    - Match ``pandas`` reverting ``apply`` deprecation (:pr:`10531`) `James Bourbeau`_
+    - Update gpuCI ``RAPIDS_VER`` to ``23.12`` (:pr:`10526`)
+    - Temporarily skip failing tests with ``fsspec==2023.9.1`` (:pr:`10520`) `James Bourbeau`_
+
 .. _v2023.9.2:
 
 2023.9.2
@@ -7024,3 +7057,4 @@ Other
 .. _`Alexander Clausen`: https://github.com/sk1p
 .. _`Swayam Patil`: https://github.com/Swish78
 .. _`Johan Olsson`: https://github.com/johanols
+.. _`wkrasnicki`: https://github.com/wkrasnicki
diff --git a/pyproject.toml b/pyproject.toml
index 6b861e5e1..c29035253 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -55,7 +55,7 @@ dataframe = [
     "dask[array]",
     "pandas >= 1.3",
 ]
-distributed = ["distributed == 2023.9.2"]
+distributed = ["distributed == 2023.9.3"]
 diagnostics = [
     "bokeh >= 2.4.2",
     "jinja2 >= 2.10.3",


Test patch:
diff --git a/dask/bag/tests/test_bag.py b/dask/bag/tests/test_bag.py
index ab46224b5..2926122b4 100644
--- a/dask/bag/tests/test_bag.py
+++ b/dask/bag/tests/test_bag.py
@@ -34,10 +34,11 @@ from dask.bag.core import (
 from dask.bag.utils import assert_eq
 from dask.blockwise import Blockwise
 from dask.delayed import Delayed
+from dask.typing import Graph
 from dask.utils import filetexts, tmpdir, tmpfile
 from dask.utils_test import add, hlg_layer, hlg_layer_topological, inc
 
-dsk = {("x", 0): (range, 5), ("x", 1): (range, 5), ("x", 2): (range, 5)}
+dsk: Graph = {("x", 0): (range, 5), ("x", 1): (range, 5), ("x", 2): (range, 5)}
 
 L = list(range(5)) * 3
 
diff --git a/dask/dataframe/io/tests/test_csv.py b/dask/dataframe/io/tests/test_csv.py
index 67df4c9df..1df7202f2 100644
--- a/dask/dataframe/io/tests/test_csv.py
+++ b/dask/dataframe/io/tests/test_csv.py
@@ -1117,6 +1117,10 @@ def test_index_col():
         except ValueError as e:
             assert "set_index" in str(e)
 
+        df = pd.read_csv(fn, index_col=False)
+        ddf = dd.read_csv(fn, blocksize=30, index_col=False)
+        assert_eq(df, ddf, check_index=False)
+
 
 def test_read_csv_with_datetime_index_partitions_one():
     with filetext(timeseries) as fn:
diff --git a/dask/tests/test_order.py b/dask/tests/test_order.py
index cd8bca788..de43cc14a 100644
--- a/dask/tests/test_order.py
+++ b/dask/tests/test_order.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 import pytest
 
 import dask
+from dask.base import collections_to_dsk
 from dask.core import get_deps
 from dask.order import diagnostics, ndependencies, order
 from dask.utils_test import add, inc
@@ -109,22 +110,19 @@ def test_base_of_reduce_preferred(abcde):
     assert o[(b, 1)] <= 3
 
 
-@pytest.mark.xfail(reason="Can't please 'em all")
 def test_avoid_upwards_branching(abcde):
     r"""
-         a1
-         |
-         a2
-         |
-         a3    d1
-        /  \  /
-      b1    c1
-      |     |
-      b2    c2
-            |
-            c3
-
-    Prefer b1 over c1 because it won't stick around waiting for d1 to complete
+       a1
+       |
+       a2
+       |
+       a3    d1
+      /  \  /
+    b1    c1
+    |     |
+    b2    c2
+          |
+          c3
     """
     a, b, c, d, e = abcde
     dsk = {
@@ -135,11 +133,13 @@ def test_avoid_upwards_branching(abcde):
         (c, 1): (f, (c, 2)),
         (c, 2): (f, (c, 3)),
         (d, 1): (f, (c, 1)),
+        (c, 3): 1,
+        (b, 2): 1,
     }
 
     o = order(dsk)
 
-    assert o[(b, 1)] < o[(c, 1)]
+    assert o[(d, 1)] < o[(b, 1)]
 
 
 def test_avoid_upwards_branching_complex(abcde):
@@ -272,27 +272,24 @@ def test_type_comparisions_ok(abcde):
     order(dsk)  # this doesn't err
 
 
-def test_prefer_short_dependents(abcde):
+def test_favor_longest_critical_path(abcde):
     r"""
 
-         a
-         |
-      d  b  e
-       \ | /
-         c
+       a
+       |
+    d  b  e
+     \ | /
+       c
 
-    Prefer to finish d and e before starting b.  That way c can be released
-    during the long computations.
     """
     a, b, c, d, e = abcde
     dsk = {c: (f,), d: (f, c), e: (f, c), b: (f, c), a: (f, b)}
 
     o = order(dsk)
-    assert o[d] < o[b]
-    assert o[e] < o[b]
+    assert o[d] > o[b]
+    assert o[e] > o[b]
 
 
-@pytest.mark.xfail(reason="This is challenging to do precisely")
 def test_run_smaller_sections(abcde):
     r"""
             aa
@@ -301,36 +298,26 @@ def test_run_smaller_sections(abcde):
      / \ /|  | /
     a   c e  cc
 
-    Prefer to run acb first because then we can get that out of the way
     """
     a, b, c, d, e = abcde
     aa, bb, cc, dd = (x * 2 for x in [a, b, c, d])
 
-    expected = [a, c, b, e, d, cc, bb, aa, dd]
-
-    log = []
-
-    def f(x):
-        def _(*args):
-            log.append(x)
-
-        return _
-
     dsk = {
-        a: (f(a),),
-        c: (f(c),),
-        e: (f(e),),
-        cc: (f(cc),),
-        b: (f(b), a, c),
-        d: (f(d), c, e),
-        bb: (f(bb), cc),
-        aa: (f(aa), d, bb),
-        dd: (f(dd), cc),
+        a: (f,),
+        c: (f,),
+        e: (f,),
+        cc: (f,),
+        b: (f, a, c),
+        d: (f, c, e),
+        bb: (f, cc),
+        aa: (f, d, bb),
+        dd: (f, cc),
     }
-
-    dask.get(dsk, [aa, b, dd])  # trigger computation
-
-    assert log == expected
+    o = order(dsk)
+    assert max(diagnostics(dsk)[1]) <= 4  # optimum is 3
+    # This is a mildly ambiguous example
+    # https://github.com/dask/dask/pull/10535/files#r1337528255
+    assert (o[aa] < o[a] and o[dd] < o[a]) or (o[b] < o[e] and o[b] < o[cc])
 
 
 def test_local_parents_of_reduction(abcde):
@@ -591,12 +578,18 @@ def test_dont_run_all_dependents_too_early(abcde):
     """From https://github.com/dask/dask-ml/issues/206#issuecomment-395873372"""
     a, b, c, d, e = abcde
     depth = 10
-    dsk = {(a, 0): 0, (b, 0): 1, (c, 0): 2, (d, 0): (f, (a, 0), (b, 0), (c, 0))}
+    dsk = {
+        (a, 0): (f, 0),
+        (b, 0): (f, 1),
+        (c, 0): (f, 2),
+        (d, 0): (f, (a, 0), (b, 0), (c, 0)),
+    }
     for i in range(1, depth):
         dsk[(b, i)] = (f, (b, 0))
         dsk[(c, i)] = (f, (c, 0))
         dsk[(d, i)] = (f, (d, i - 1), (b, i), (c, i))
     o = order(dsk)
+
     expected = [3, 6, 9, 12, 15, 18, 21, 24, 27, 30]
     actual = sorted(v for (letter, num), v in o.items() if letter == d)
     assert expected == actual
@@ -723,13 +716,17 @@ def test_order_with_equal_dependents(abcde):
     This DAG has enough structure to exercise more parts of `order`
 
     """
+    # Lower pressure is better but this is where we are right now. Important is
+    # that no variation below should be worse since all variations below should
+    # reduce to the same graph when optimized/fused.
+    max_pressure = 11
     a, b, c, d, e = abcde
     dsk = {}
     abc = [a, b, c, d]
     for x in abc:
         dsk.update(
             {
-                (x, 0): 0,
+                (x, 0): (f, 0),
                 (x, 1): (f, (x, 0)),
                 (x, 2, 0): (f, (x, 0)),
                 (x, 2, 1): (f, (x, 1)),
@@ -753,7 +750,9 @@ def test_order_with_equal_dependents(abcde):
             val = o[(x, 6, i, 1)] - o[(x, 6, i, 0)]
             assert val > 0  # ideally, val == 2
             total += val
-    assert total <= 110  # ideally, this should be 2 * 16 = 32
+    assert total <= 56  # ideally, this should be 2 * 16 == 32
+    pressure = diagnostics(dsk, o=o)[1]
+    assert max(pressure) <= max_pressure
 
     # Add one to the end of the nine bundles
     dsk2 = dict(dsk)
@@ -764,10 +763,12 @@ def test_order_with_equal_dependents(abcde):
     total = 0
     for x in abc:
         for i in range(len(abc)):
-            val = o[(x, 7, i, 0)] - o[(x, 6, i, 1)]
+            val = o[(x, 6, i, 1)] - o[(x, 7, i, 0)]
             assert val > 0  # ideally, val == 3
             total += val
-    assert total <= 138  # ideally, this should be 3 * 16 == 48
+    assert total <= 75  # ideally, this should be 3 * 16 == 48
+    pressure = diagnostics(dsk2, o=o)[1]
+    assert max(pressure) <= max_pressure
 
     # Remove one from each of the nine bundles
     dsk3 = dict(dsk)
@@ -778,10 +779,12 @@ def test_order_with_equal_dependents(abcde):
     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
+            val = o[(x, 5, i, 1)] - o[(x, 6, i, 0)]
+            assert val > 0
             total += val
-    assert total <= 98  # ideally, this should be 2 * 16 == 32
+    assert total <= 45  # ideally, this should be 2 * 16 == 32
+    pressure = diagnostics(dsk3, o=o)[1]
+    assert max(pressure) <= max_pressure
 
     # Remove another one from each of the nine bundles
     dsk4 = dict(dsk3)
@@ -789,10 +792,11 @@ def test_order_with_equal_dependents(abcde):
         for i in range(len(abc)):
             del dsk4[(x, 6, i, 0)]
     o = order(dsk4)
-    total = 0
+    pressure = diagnostics(dsk4, o=o)[1]
+    assert max(pressure) <= max_pressure
     for x in abc:
         for i in range(len(abc)):
-            assert o[(x, 5, i, 1)] - o[(x, 5, i, 0)] == 1
+            assert abs(o[(x, 5, i, 1)] - o[(x, 5, i, 0)]) <= 10
 
 
 def test_terminal_node_backtrack():
@@ -986,7 +990,7 @@ def test_eager_to_compute_dependent_to_free_parent():
 
 def test_diagnostics(abcde):
     r"""
-        a1  b1  c2  d1  e1
+        a1  b1  c1  d1  e1
         /|\ /|\ /|\ /|  /
        / | X | X | X | /
       /  |/ \|/ \|/ \|/
@@ -1005,6 +1009,10 @@ def test_diagnostics(abcde):
         (d, 1): (f, (d, 0), (e, 0)),
         (e, 1): (f, (e, 0)),
     }
+    o = order(dsk)
+    assert o[(e, 1)] == len(dsk) - 1
+    assert o[(d, 1)] == len(dsk) - 2
+    assert o[(c, 1)] == len(dsk) - 3
     info, memory_over_time = diagnostics(dsk)
     assert memory_over_time == [0, 1, 2, 3, 2, 3, 2, 3, 2, 1]
     assert {key: val.order for key, val in info.items()} == {
@@ -1067,3 +1075,596 @@ def test_diagnostics(abcde):
         (d, 1): 2,
         (e, 1): 1,
     }
+
+
+def test_xarray_like_reduction():
+    a, b, c, d, e = list("abcde")
+
+    dsk = {}
+    for ix in range(3):
+        part = {
+            # Part1
+            (a, 0, ix): (f,),
+            (a, 1, ix): (f,),
+            (b, 0, ix): (f, (a, 0, ix)),
+            (b, 1, ix): (f, (a, 0, ix), (a, 1, ix)),
+            (b, 2, ix): (f, (a, 1, ix)),
+            (c, 0, ix): (f, (b, 0, ix)),
+            (c, 1, ix): (f, (b, 1, ix)),
+            (c, 2, ix): (f, (b, 2, ix)),
+        }
+        dsk.update(part)
+    for ix in range(3):
+        dsk.update(
+            {
+                (d, ix): (f, (c, ix, 0), (c, ix, 1), (c, ix, 2)),
+            }
+        )
+    o = order(dsk)
+    _, pressure = diagnostics(dsk, o=o)
+    assert max(pressure) <= 9
+
+
+@pytest.mark.parametrize(
+    "optimize",
+    [True, False],
+)
+def test_array_vs_dataframe(optimize):
+    xr = pytest.importorskip("xarray")
+
+    import dask.array as da
+
+    size = 5000
+    ds = xr.Dataset(
+        dict(
+            anom_u=(
+                ["time", "face", "j", "i"],
+                da.random.random((size, 1, 987, 1920), chunks=(10, 1, -1, -1)),
+            ),
+            anom_v=(
+                ["time", "face", "j", "i"],
+                da.random.random((size, 1, 987, 1920), chunks=(10, 1, -1, -1)),
+            ),
+        )
+    )
+
+    quad = ds**2
+    quad["uv"] = ds.anom_u * ds.anom_v
+    mean = quad.mean("time")
+    diag_array = diagnostics(collections_to_dsk([mean], optimize_graph=optimize))
+    diag_df = diagnostics(
+        collections_to_dsk([mean.to_dask_dataframe()], optimize_graph=optimize)
+    )
+    assert max(diag_df[1]) == max(diag_array[1])
+    assert max(diag_array[1]) < 50
+
+
+def test_anom_mean():
+    np = pytest.importorskip("numpy")
+    xr = pytest.importorskip("xarray")
+
+    import dask.array as da
+    from dask.utils import parse_bytes
+
+    data = da.random.random(
+        (260, 1310720),
+        chunks=(1, parse_bytes("10MiB") // 8),
+    )
+
+    ngroups = data.shape[0] // 4
+    arr = xr.DataArray(
+        data,
+        dims=["time", "x"],
+        coords={"day": ("time", np.arange(data.shape[0]) % ngroups)},
+    )
+    data = da.random.random((5, 1), chunks=(1, 1))
+
+    arr = xr.DataArray(
+        data,
+        dims=["time", "x"],
+        coords={"day": ("time", np.arange(5) % 2)},
+    )
+
+    clim = arr.groupby("day").mean(dim="time")
+    anom = arr.groupby("day") - clim
+    anom_mean = anom.mean(dim="time")
+    _, pressure = diagnostics(anom_mean.__dask_graph__())
+    assert max(pressure) <= 9
+
+
+def test_anom_mean_raw():
+    dsk = {
+        ("d", 0, 0): (f, ("a", 0, 0), ("b1", 0, 0)),
+        ("d", 1, 0): (f, ("a", 1, 0), ("b1", 1, 0)),
+        ("d", 2, 0): (f, ("a", 2, 0), ("b1", 2, 0)),
+        ("d", 3, 0): (f, ("a", 3, 0), ("b1", 3, 0)),
+        ("d", 4, 0): (f, ("a", 4, 0), ("b1", 4, 0)),
+        ("a", 0, 0): (f, f, "random_sample", None, (1, 1), [], {}),
+        ("a", 1, 0): (f, f, "random_sample", None, (1, 1), [], {}),
+        ("a", 2, 0): (f, f, "random_sample", None, (1, 1), [], {}),
+        ("a", 3, 0): (f, f, "random_sample", None, (1, 1), [], {}),
+        ("a", 4, 0): (f, f, "random_sample", None, (1, 1), [], {}),
+        ("e", 0, 0): (f, ("g1", 0)),
+        ("e", 1, 0): (f, ("g3", 0)),
+        ("b0", 0, 0): (f, ("a", 0, 0)),
+        ("b0", 1, 0): (f, ("a", 2, 0)),
+        ("b0", 2, 0): (f, ("a", 4, 0)),
+        ("c0", 0, 0): (f, ("b0", 0, 0)),
+        ("c0", 1, 0): (f, ("b0", 1, 0)),
+        ("c0", 2, 0): (f, ("b0", 2, 0)),
+        ("g1", 0): (f, [("c0", 0, 0), ("c0", 1, 0), ("c0", 2, 0)]),
+        ("b2", 0, 0): (f, ("a", 1, 0)),
+        ("b2", 1, 0): (f, ("a", 3, 0)),
+        ("c1", 0, 0): (f, ("b2", 0, 0)),
+        ("c1", 1, 0): (f, ("b2", 1, 0)),
+        ("g3", 0): (f, [("c1", 0, 0), ("c1", 1, 0)]),
+        ("b1", 0, 0): (f, ("e", 0, 0)),
+        ("b1", 1, 0): (f, ("e", 1, 0)),
+        ("b1", 2, 0): (f, ("e", 0, 0)),
+        ("b1", 3, 0): (f, ("e", 1, 0)),
+        ("b1", 4, 0): (f, ("e", 0, 0)),
+        ("c2", 0, 0): (f, ("d", 0, 0)),
+        ("c2", 1, 0): (f, ("d", 1, 0)),
+        ("c2", 2, 0): (f, ("d", 2, 0)),
+        ("c2", 3, 0): (f, ("d", 3, 0)),
+        ("c2", 4, 0): (f, ("d", 4, 0)),
+        ("f", 0, 0): (f, [("c2", 0, 0), ("c2", 1, 0), ("c2", 2, 0), ("c2", 3, 0)]),
+        ("f", 1, 0): (f, [("c2", 4, 0)]),
+        ("g2", 0): (f, [("f", 0, 0), ("f", 1, 0)]),
+    }
+
+    o = order(dsk)
+
+    # The left hand computation branch should complete before we start loading
+    # more data
+    nodes_to_finish_before_loading_more_data = [
+        ("f", 1, 0),
+        ("d", 0, 0),
+        ("d", 2, 0),
+        ("d", 4, 0),
+    ]
+    for n in nodes_to_finish_before_loading_more_data:
+        assert o[n] < o[("a", 1, 0)]
+        assert o[n] < o[("a", 3, 0)]
+
+
+def test_flaky_array_reduction():
+    first = {
+        ("mean_agg-aggregate-10d721567ef5a0d6a0e1afae8a87c066", 0, 0, 0): (
+            f,
+            [
+                ("mean_combine-partial-17c7b5c6eed42e203858b3f6dde16003", 0, 0, 0, 0),
+                ("mean_combine-partial-17c7b5c6eed42e203858b3f6dde16003", 1, 0, 0, 0),
+            ],
+        ),
+        ("mean_combine-partial-17c7b5c6eed42e203858b3f6dde16003", 0, 0, 0, 0): (
+            f,
+            [
+                ("mean_chunk-98a32cd9f4fadbed908fffb32e0c9679", 0, 0, 0, 0),
+                ("mean_chunk-98a32cd9f4fadbed908fffb32e0c9679", 1, 0, 0, 0),
+                ("mean_chunk-98a32cd9f4fadbed908fffb32e0c9679", 2, 0, 0, 0),
+                ("mean_chunk-98a32cd9f4fadbed908fffb32e0c9679", 3, 0, 0, 0),
+            ],
+        ),
+        ("mean_combine-partial-17c7b5c6eed42e203858b3f6dde16003", 1, 0, 0, 0): (
+            "mean_chunk-mean_combine-partial-17c7b5c6eed42e203858b3f6dde16003",
+            1,
+            0,
+            0,
+            0,
+        ),
+        ("mean_chunk-98a32cd9f4fadbed908fffb32e0c9679", 0, 0, 0, 0): (
+            f,
+            ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 0, 0, 0, 0),
+            ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 0, 0, 0, 0),
+        ),
+        ("mean_chunk-98a32cd9f4fadbed908fffb32e0c9679", 1, 0, 0, 0): (
+            f,
+            ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 1, 0, 0, 0),
+            ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 1, 0, 0, 0),
+        ),
+        ("mean_chunk-98a32cd9f4fadbed908fffb32e0c9679", 2, 0, 0, 0): (
+            f,
+            ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 2, 0, 0, 0),
+            ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 2, 0, 0, 0),
+        ),
+        ("mean_chunk-98a32cd9f4fadbed908fffb32e0c9679", 3, 0, 0, 0): (
+            f,
+            ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 3, 0, 0, 0),
+            ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 3, 0, 0, 0),
+        ),
+        ("mean_agg-aggregate-fdb340546b01334890192fcfa55fa0d9", 0, 0, 0): (
+            f,
+            [
+                ("mean_combine-partial-23adb4747560e6e33afd63c5bb179709", 0, 0, 0, 0),
+                ("mean_combine-partial-23adb4747560e6e33afd63c5bb179709", 1, 0, 0, 0),
+            ],
+        ),
+        ("mean_combine-partial-23adb4747560e6e33afd63c5bb179709", 0, 0, 0, 0): (
+            f,
+            [
+                ("mean_chunk-7edba1c5a284fcec88b9efdda6c2135f", 0, 0, 0, 0),
+                ("mean_chunk-7edba1c5a284fcec88b9efdda6c2135f", 1, 0, 0, 0),
+                ("mean_chunk-7edba1c5a284fcec88b9efdda6c2135f", 2, 0, 0, 0),
+                ("mean_chunk-7edba1c5a284fcec88b9efdda6c2135f", 3, 0, 0, 0),
+            ],
+        ),
+        ("mean_combine-partial-23adb4747560e6e33afd63c5bb179709", 1, 0, 0, 0): (
+            "mean_chunk-mean_combine-partial-23adb4747560e6e33afd63c5bb179709",
+            1,
+            0,
+            0,
+            0,
+        ),
+        ("mean_chunk-7edba1c5a284fcec88b9efdda6c2135f", 0, 0, 0, 0): (
+            f,
+            ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 0, 0, 0, 0),
+            2,
+        ),
+        ("mean_chunk-7edba1c5a284fcec88b9efdda6c2135f", 1, 0, 0, 0): (
+            f,
+            ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 1, 0, 0, 0),
+            2,
+        ),
+        ("mean_chunk-7edba1c5a284fcec88b9efdda6c2135f", 2, 0, 0, 0): (
+            f,
+            ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 2, 0, 0, 0),
+            2,
+        ),
+        ("mean_chunk-7edba1c5a284fcec88b9efdda6c2135f", 3, 0, 0, 0): (
+            f,
+            ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 3, 0, 0, 0),
+            2,
+        ),
+        ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 0, 0, 0, 0): (f, 1),
+        ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 1, 0, 0, 0): (f, 1),
+        ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 2, 0, 0, 0): (f, 1),
+        ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 3, 0, 0, 0): (f, 1),
+        ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 4, 0, 0, 0): (f, 1),
+        ("mean_agg-aggregate-cc19342c8116d616fc6573f5d20b5762", 0, 0, 0): (
+            f,
+            [
+                ("mean_combine-partial-0c98c5a4517f58f8268985e7464daace", 0, 0, 0, 0),
+                ("mean_combine-partial-0c98c5a4517f58f8268985e7464daace", 1, 0, 0, 0),
+            ],
+        ),
+        ("mean_combine-partial-0c98c5a4517f58f8268985e7464daace", 0, 0, 0, 0): (
+            f,
+            [
+                ("mean_chunk-540e88b7d9289f6b5461b95a0787af3e", 0, 0, 0, 0),
+                ("mean_chunk-540e88b7d9289f6b5461b95a0787af3e", 1, 0, 0, 0),
+                ("mean_chunk-540e88b7d9289f6b5461b95a0787af3e", 2, 0, 0, 0),
+                ("mean_chunk-540e88b7d9289f6b5461b95a0787af3e", 3, 0, 0, 0),
+            ],
+        ),
+        ("mean_combine-partial-0c98c5a4517f58f8268985e7464daace", 1, 0, 0, 0): (
+            "mean_chunk-mean_combine-partial-0c98c5a4517f58f8268985e7464daace",
+            1,
+            0,
+            0,
+            0,
+        ),
+        ("mean_chunk-540e88b7d9289f6b5461b95a0787af3e", 0, 0, 0, 0): (
+            f,
+            ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 0, 0, 0, 0),
+        ),
+        ("mean_chunk-540e88b7d9289f6b5461b95a0787af3e", 1, 0, 0, 0): (
+            f,
+            ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 1, 0, 0, 0),
+        ),
+        ("mean_chunk-540e88b7d9289f6b5461b95a0787af3e", 2, 0, 0, 0): (
+            f,
+            ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 2, 0, 0, 0),
+        ),
+        ("mean_chunk-540e88b7d9289f6b5461b95a0787af3e", 3, 0, 0, 0): (
+            f,
+            ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 3, 0, 0, 0),
+        ),
+        ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 0, 0, 0, 0): (f, 1),
+        ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 1, 0, 0, 0): (f, 1),
+        ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 2, 0, 0, 0): (f, 1),
+        ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 3, 0, 0, 0): (f, 1),
+        ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 4, 0, 0, 0): (f, 1),
+        (
+            "mean_chunk-mean_combine-partial-17c7b5c6eed42e203858b3f6dde16003",
+            1,
+            0,
+            0,
+            0,
+        ): (
+            f,
+            [
+                (
+                    f,
+                    ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 4, 0, 0, 0),
+                    ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 4, 0, 0, 0),
+                )
+            ],
+        ),
+        (
+            "mean_chunk-mean_combine-partial-0c98c5a4517f58f8268985e7464daace",
+            1,
+            0,
+            0,
+            0,
+        ): (
+            f,
+            [(f, ("random_sample-e16bcfb15a013023c98a21e2f03d66a9", 4, 0, 0, 0), 2)],
+        ),
+        (
+            "mean_chunk-mean_combine-partial-23adb4747560e6e33afd63c5bb179709",
+            1,
+            0,
+            0,
+            0,
+        ): (
+            f,
+            [(f, ("random_sample-02eaa4a8dbb23fac4db22ad034c401b3", 4, 0, 0, 0), 2)],
+        ),
+    }
+
+    other = {
+        ("mean_agg-aggregate-e79dd3b9757c9fb2ad7ade96f3f6c814", 0, 0, 0): (
+            f,
+            [
+                ("mean_combine-partial-e7d9fd7c132e12007a4b4f62ce443a75", 0, 0, 0, 0),
+                ("mean_combine-partial-e7d9fd7c132e12007a4b4f62ce443a75", 1, 0, 0, 0),
+            ],
+        ),
+        ("mean_combine-partial-e7d9fd7c132e12007a4b4f62ce443a75", 0, 0, 0, 0): (
+            f,
+            [
+                ("mean_chunk-0df65d9a6e168673f32082f59f19576a", 0, 0, 0, 0),
+                ("mean_chunk-0df65d9a6e168673f32082f59f19576a", 1, 0, 0, 0),
+                ("mean_chunk-0df65d9a6e168673f32082f59f19576a", 2, 0, 0, 0),
+                ("mean_chunk-0df65d9a6e168673f32082f59f19576a", 3, 0, 0, 0),
+            ],
+        ),
+        ("mean_combine-partial-e7d9fd7c132e12007a4b4f62ce443a75", 1, 0, 0, 0): (
+            "mean_chunk-mean_combine-partial-e7d9fd7c132e12007a4b4f62ce443a75",
+            1,
+            0,
+            0,
+            0,
+        ),
+        ("mean_chunk-0df65d9a6e168673f32082f59f19576a", 0, 0, 0, 0): (
+            f,
+            ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 0, 0, 0, 0),
+            ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 0, 0, 0, 0),
+        ),
+        ("mean_chunk-0df65d9a6e168673f32082f59f19576a", 1, 0, 0, 0): (
+            f,
+            ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 1, 0, 0, 0),
+            ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 1, 0, 0, 0),
+        ),
+        ("mean_chunk-0df65d9a6e168673f32082f59f19576a", 2, 0, 0, 0): (
+            f,
+            ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 2, 0, 0, 0),
+            ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 2, 0, 0, 0),
+        ),
+        ("mean_chunk-0df65d9a6e168673f32082f59f19576a", 3, 0, 0, 0): (
+            f,
+            ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 3, 0, 0, 0),
+            ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 3, 0, 0, 0),
+        ),
+        ("mean_agg-aggregate-c7647920facf0e557f947b7a6626b7be", 0, 0, 0): (
+            f,
+            [
+                ("mean_combine-partial-57413f0bb18da78db0f689a096c7fbbf", 0, 0, 0, 0),
+                ("mean_combine-partial-57413f0bb18da78db0f689a096c7fbbf", 1, 0, 0, 0),
+            ],
+        ),
+        ("mean_combine-partial-57413f0bb18da78db0f689a096c7fbbf", 0, 0, 0, 0): (
+            f,
+            [
+                ("mean_chunk-d6bd425ea61739f1eaa71762fe3bbbd7", 0, 0, 0, 0),
+                ("mean_chunk-d6bd425ea61739f1eaa71762fe3bbbd7", 1, 0, 0, 0),
+                ("mean_chunk-d6bd425ea61739f1eaa71762fe3bbbd7", 2, 0, 0, 0),
+                ("mean_chunk-d6bd425ea61739f1eaa71762fe3bbbd7", 3, 0, 0, 0),
+            ],
+        ),
+        ("mean_combine-partial-57413f0bb18da78db0f689a096c7fbbf", 1, 0, 0, 0): (
+            "mean_chunk-mean_combine-partial-57413f0bb18da78db0f689a096c7fbbf",
+            1,
+            0,
+            0,
+            0,
+        ),
+        ("mean_chunk-d6bd425ea61739f1eaa71762fe3bbbd7", 0, 0, 0, 0): (
+            f,
+            ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 0, 0, 0, 0),
+            2,
+        ),
+        ("mean_chunk-d6bd425ea61739f1eaa71762fe3bbbd7", 1, 0, 0, 0): (
+            f,
+            ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 1, 0, 0, 0),
+            2,
+        ),
+        ("mean_chunk-d6bd425ea61739f1eaa71762fe3bbbd7", 2, 0, 0, 0): (
+            f,
+            ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 2, 0, 0, 0),
+            2,
+        ),
+        ("mean_chunk-d6bd425ea61739f1eaa71762fe3bbbd7", 3, 0, 0, 0): (
+            f,
+            ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 3, 0, 0, 0),
+            2,
+        ),
+        ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 0, 0, 0, 0): (
+            f,
+            "random_sample",
+            (10, 1, 987, 1920),
+            [],
+        ),
+        ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 1, 0, 0, 0): (
+            f,
+            "random_sample",
+            (10, 1, 987, 1920),
+            [],
+        ),
+        ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 2, 0, 0, 0): (
+            f,
+            "random_sample",
+            (10, 1, 987, 1920),
+            [],
+        ),
+        ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 3, 0, 0, 0): (
+            f,
+            "random_sample",
+            (10, 1, 987, 1920),
+            [],
+        ),
+        ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 4, 0, 0, 0): (
+            f,
+            "random_sample",
+            (10, 1, 987, 1920),
+            [],
+        ),
+        ("mean_agg-aggregate-05071ebaabb68a64c180f6f443c5c8f4", 0, 0, 0): (
+            f,
+            [
+                ("mean_combine-partial-a7c475f79a46af4265b189ffdc000bb3", 0, 0, 0, 0),
+                ("mean_combine-partial-a7c475f79a46af4265b189ffdc000bb3", 1, 0, 0, 0),
+            ],
+        ),
+        ("mean_combine-partial-a7c475f79a46af4265b189ffdc000bb3", 0, 0, 0, 0): (
+            f,
+            [
+                ("mean_chunk-fd17feaf0728ea7a89d119d3fd172c75", 0, 0, 0, 0),
+                ("mean_chunk-fd17feaf0728ea7a89d119d3fd172c75", 1, 0, 0, 0),
+                ("mean_chunk-fd17feaf0728ea7a89d119d3fd172c75", 2, 0, 0, 0),
+                ("mean_chunk-fd17feaf0728ea7a89d119d3fd172c75", 3, 0, 0, 0),
+            ],
+        ),
+        ("mean_combine-partial-a7c475f79a46af4265b189ffdc000bb3", 1, 0, 0, 0): (
+            "mean_chunk-mean_combine-partial-a7c475f79a46af4265b189ffdc000bb3",
+            1,
+            0,
+            0,
+            0,
+        ),
+        ("mean_chunk-fd17feaf0728ea7a89d119d3fd172c75", 0, 0, 0, 0): (
+            f,
+            ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 0, 0, 0, 0),
+            2,
+        ),
+        ("mean_chunk-fd17feaf0728ea7a89d119d3fd172c75", 1, 0, 0, 0): (
+            f,
+            ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 1, 0, 0, 0),
+            2,
+        ),
+        ("mean_chunk-fd17feaf0728ea7a89d119d3fd172c75", 2, 0, 0, 0): (
+            f,
+            ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 2, 0, 0, 0),
+            2,
+        ),
+        ("mean_chunk-fd17feaf0728ea7a89d119d3fd172c75", 3, 0, 0, 0): (
+            f,
+            ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 3, 0, 0, 0),
+            2,
+        ),
+        ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 0, 0, 0, 0): (
+            f,
+            "random_sample",
+            (10, 1, 987, 1920),
+            [],
+        ),
+        ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 1, 0, 0, 0): (
+            f,
+            "random_sample",
+            (10, 1, 987, 1920),
+            [],
+        ),
+        ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 2, 0, 0, 0): (
+            f,
+            "random_sample",
+            (10, 1, 987, 1920),
+            [],
+        ),
+        ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 3, 0, 0, 0): (
+            f,
+            "random_sample",
+            (10, 1, 987, 1920),
+            [],
+        ),
+        ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 4, 0, 0, 0): (
+            f,
+            "random_sample",
+            (10, 1, 987, 1920),
+            [],
+        ),
+        (
+            "mean_chunk-mean_combine-partial-a7c475f79a46af4265b189ffdc000bb3",
+            1,
+            0,
+            0,
+            0,
+        ): (
+            f,
+            [(f, ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 4, 0, 0, 0), 2)],
+        ),
+        (
+            "mean_chunk-mean_combine-partial-e7d9fd7c132e12007a4b4f62ce443a75",
+            1,
+            0,
+            0,
+            0,
+        ): (
+            f,
+            [
+                (
+                    f,
+                    ("random_sample-a155d5a37ac5e09ede89c98a3bfcadff", 4, 0, 0, 0),
+                    ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 4, 0, 0, 0),
+                )
+            ],
+        ),
+        (
+            "mean_chunk-mean_combine-partial-57413f0bb18da78db0f689a096c7fbbf",
+            1,
+            0,
+            0,
+            0,
+        ): (
+            f,
+            [(f, ("random_sample-241fdbadc062900adc59d1a79c4c41e1", 4, 0, 0, 0), 2)],
+        ),
+    }
+    first_pressure = max(diagnostics(first)[1])
+    second_pressure = max(diagnostics(other)[1])
+    assert first_pressure == second_pressure
+
+
+def test_flox_reduction():
+    # TODO: It would be nice to scramble keys to ensure we're not comparing keys
+    dsk = {
+        "A0": (f, 1),
+        ("A1", 0): (f, 1),
+        ("A1", 1): (f, 1),
+        ("A2", 0): (f, 1),
+        ("A2", 1): (f, 1),
+        ("B1", 0): (f, [(f, ("A2", 0))]),
+        ("B1", 1): (f, [(f, ("A2", 1))]),
+        ("B2", 1): (f, [(f, ("A1", 1))]),
+        ("B2", 0): (f, [(f, ("A1", 0))]),
+        ("B11", 0): ("B1", 0),
+        ("B11", 1): ("B1", 1),
+        ("B22", 0): ("B2", 0),
+        ("B22", 1): ("B2", 1),
+        ("C1", 0): (f, ("B22", 0)),
+        ("C1", 1): (f, ("B22", 1)),
+        ("C2", 0): (f, ("B11", 0)),
+        ("C2", 1): (f, ("B11", 1)),
+        ("E", 1): (f, [(f, ("A1", 1), ("A2", 1), ("C1", 1), ("C2", 1))]),
+        ("E", 0): (f, [(f, ("A1", 0), ("A2", 0), ("C1", 0), ("C2", 0))]),
+        ("EE", 0): ("E", 0),
+        ("EE", 1): ("E", 1),
+        ("F1", 0): (f, "A0", ("B11", 0)),
+        ("F1", 1): (f, "A0", ("B22", 0)),
+        ("F1", 2): (f, "A0", ("EE", 0)),
+        ("F2", 0): (f, "A0", ("B11", 1)),
+        ("F2", 1): (f, "A0", ("B22", 1)),
+        ("F2", 2): (f, "A0", ("EE", 1)),
+    }
+    o = order(dsk)
+    assert max(o[("F1", ix)] for ix in range(3)) < min(o[("F2", ix)] for ix in range(3))