Repository: dask/dask
Pull request #10660 (2023-11-30 13:11:25)

Dask.order rewrite using a critical path approach
This is another ~attempt to~ rewrite of dask.order

Supersedes https://github.com/dask/dask/pull/10557
Closes https://github.com/dask/distributed/issues/8255


It does not incorporate https://github.com/dask/dask/pull/10619 yet but I will likely want to follow up adding this

This takes a quite different approach to our ordering that the PRs before this and actually rewrites it from scratch.

I still have a couple of unit tests to take care of but the failing tests right now are mostly artificial test cases. There are one or two mildly concerning things I still have to look at but generally speaking I'm optimistic about this approach

High level, this change was inspired by what the root task queuing is doing in that it tries to finish off whatever runnable computation branches exist before it starts loading new stuff. On such a high level, the algorithm can be described in a couple of simple steps

1. Compute a critical path by reverse DFS
2. Use the critical path as a stack and walk the stack in forward direction. Whenever a node is encountered, that is not runnable, add the missing dependencies
3. Whenever that happens, look at all currently runnable tasks and walk them as far as possible. Those paths are remembered but not immediately added to the result
4. If the runnable paths meet a reducer (or final node) and that reducer would be reachable given all runnable paths, execute those paths
5. If not, proceed with the critical path until all nodes are processed


The nuances of the implementation are a little harder to describe but I'll add some more in-code comments.


fjetter: coiled benchmark results are in

https://github.com/coiled/benchmarks/actions/runs/7046212948

Wall clock time is mostly unaffected. The little signal there is likely just noise

Looking at peak memory, we see a reduction for double_diff and dataframe_align

![image](https://github.com/dask/dask/assets/8629629/74fc3c55-db47-4dd6-9cb6-29bd08996f79)

and a very solid reduction for the zarr filtering tests

![image](https://github.com/dask/dask/assets/8629629/1358a145-5931-4988-b8b0-8bd928b47b07)

Similar patterns can be seen in average memory use


fjetter: `doublediff` does not drop as strongly as in https://github.com/dask/dask/pull/10557 so there is likely still room for improvement. Given that the cost function / sorting key I am using is almost random, I assume this is the most interesting knob I have to tweak
fjetter: This critical path approach has one flaw I haven't quite figured out, yet. 

The following example is a smaller version of what was reported in https://github.com/dask/dask/issues/3055

![image](https://github.com/dask/dask/assets/8629629/9ce73fdc-3f6d-4620-a435-dd092a549ef3)

The visualization shows clearly that we're very eagerly processing the critical path and are working off unlocked branches as soon as they are available. However, with this particular graph, a root task that is required to reduce everything early on is only required by the _critical_ path very late on such that those branches never get a chance to run
fjetter: So, I changed how the critical path is determined such that we have two different strategies based on a heuristic about graph symmetry. That is working pretty well for the test suite and the benchmark suite.

https://github.com/coiled/benchmarks/actions/runs/7088167915

Wall time looks great for all H20 tests and the zarr filtering

<img width="587" alt="Screenshot 2023-12-05 at 09 25 38" src="https://github.com/dask/dask/assets/8629629/48a90e3b-cce5-430d-ad4c-037b47e543e9">
<img width="619" alt="Screenshot 2023-12-05 at 09 25 59" src="https://github.com/dask/dask/assets/8629629/37ab208a-55e1-418c-b2c9-e411fad6e433">

There appears to be a mild regression in the shuffle tests I still want to investigate since I'm surprised those are affected at all (maybe this is runtime of the algo, idk)

<img width="411" alt="Screenshot 2023-12-05 at 09 26 17" src="https://github.com/dask/dask/assets/8629629/ff739694-77a9-450b-916d-02350c786e29">


Memory looks great overall (below peak memory, avg looks similar)

<img width="617" alt="Screenshot 2023-12-05 at 09 27 03" src="https://github.com/dask/dask/assets/8629629/7a6e9452-c0df-4f22-8f00-17ababf1bdc7">
<img width="567" alt="Screenshot 2023-12-05 at 09 27 27" src="https://github.com/dask/dask/assets/8629629/c3632f4b-3e12-4cee-8409-e51f50abff6a">

mrocklin: I'm mildly curious when we'll be able to merge and release this.  There's probably some mild benefit to it going out before AGU (next week) if that's doable.
mrocklin: cc @dcherian 
fjetter: I reran the benchmarks and the shuffle related regressions that showed up earlier did vanish. The improvements in H20 unfortunately also vanished but that's somewhat expected. 

https://github.com/coiled/benchmarks/actions/runs/7130181112

Note, that this rewrite does not include the "skip data tasks" as proposed inhttps://github.com/dask/dask/pull/10619 yet but it does fix the test case I added there to close https://github.com/dask/dask/issues/10618 so I added it also to this PR
fjetter: I also checked again https://github.com/pydata/xarray/issues/8414 which now produces this ordering

![image](https://github.com/dask/dask/assets/8629629/39d59ebd-da4b-4180-b7c9-f76bb7469d7a)

which I believe is just right. I'll add this as a test case as well.
fjetter: A note about the implementation:

The complexity went through the roof but I think it's worth it. I tried factorizing it but it doesn't break apart easily. However, there are very clear stages and I split off those stages and modularized it as good as possible. I didn't manage to factor this stuff out to functions but I still ended up using local functions to manage scope and namespace pollution.
If anybody is brave enough to review I suggest to start towards the bottom of the `order` function where the core algorithm starts and is documented. I'll eventually follow up with more docs about the other sections.
mrocklin: It looks like the bus factor went from zero to one.
fjetter: Apparently, the reordering made us hit https://github.com/dask/dask/issues/10034 in the `test_groupby_reduction_split`. What's even stranger is that it also causes other failures that are more meta related, e.g. a groupby.shift is losing it's `freq` argument in `test_groupby_shift_with_freq`.
fjetter: Re algorithmic performance, this implementation is now slower than before but it look to be reasonable and in real world testing it doesn't seem to have an impact (the other stuff we're doing on the scheduler is just much slower).

Particularly, the dask-benchmark tests I've been running locally to test performance are often hitting the worst cases of this new algorithm.

Overall, I'm not concerned by a (~linear complexity) increase in runtime
fjetter: So, I saw again regressions in set_index benchmarks that somehow pop up on the ordering. I just cannot understand how ordering would have an impact on P2P shuffling. When inspecting the dashboard, two things pop up that probably cause the benchmark failure

![image](https://github.com/dask/dask/assets/8629629/531865c2-1602-4fc9-8ec9-be2afa530e02)


For some reason the unpack tasks are flagged as root tasks and are queued which causes a heavy skew in assigned tasks which causes work stealing and a lot of irrelevant data transfer to compute the `size`. @hendrikmakait the p2p tasks previously carried worker annotations which effectively blocked them from every be flagged rootish. By removing those annotations we might have opened ourselves up to this scheduling anomaly. I don't understand how this connects to ordering and I suspect this is unrelated. Maybe the ordering change changed some subtle timing somewhere that makes it more likely to hit this.

Patch:
diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml
index 5e2a69e24..57defd884 100644
--- a/.github/workflows/pre-commit.yml
+++ b/.github/workflows/pre-commit.yml
@@ -12,7 +12,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4.1.1
-      - uses: actions/setup-python@v4
+      - uses: actions/setup-python@v5
         with:
           python-version: '3.9'
       - uses: pre-commit/action@v3.0.0
diff --git a/dask/base.py b/dask/base.py
index c015c4b70..eddbf9b1a 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -722,7 +722,30 @@ def visualize(
 
     dsk = dict(collections_to_dsk(args, optimize_graph=optimize_graph))
 
+    return visualize_dsk(
+        dsk=dsk,
+        filename=filename,
+        traverse=traverse,
+        optimize_graph=optimize_graph,
+        maxval=maxval,
+        engine=engine,
+        **kwargs,
+    )
+
+
+def visualize_dsk(
+    dsk,
+    filename="mydask",
+    traverse=True,
+    optimize_graph=False,
+    maxval=None,
+    o=None,
+    engine: Literal["cytoscape", "ipycytoscape", "graphviz"] | None = None,
+    limit=None,
+    **kwargs,
+):
     color = kwargs.get("color")
+    from dask.order import diagnostics, order
 
     if color in {
         "order",
@@ -736,12 +759,20 @@ def visualize(
         "memoryincreases",
         "memorydecreases",
         "memorypressure",
+        "critical",
+        "cpath",
     }:
         import matplotlib.pyplot as plt
 
-        from dask.order import diagnostics, order
+        if o is None:
+            o_stats = order(dsk, return_stats=True)
+            o = {k: v.priority for k, v in o_stats.items()}
+        elif isinstance(next(iter(o.values())), int):
+            o_stats = order(dsk, return_stats=True)
+        else:
+            o_stats = o
+            o = {k: v.priority for k, v in o.items()}
 
-        o = order(dsk)
         try:
             cmap = kwargs.pop("cmap")
         except KeyError:
@@ -771,11 +802,15 @@ def visualize(
                     key: max(0, val.num_data_when_released - val.num_data_when_run)
                     for key, val in info.items()
                 }
-            else:  # memorydecreases
+            elif color.endswith("memorydecreases"):
                 values = {
                     key: max(0, val.num_data_when_run - val.num_data_when_released)
                     for key, val in info.items()
                 }
+            elif color.split("-")[-1] in {"critical", "cpath"}:
+                values = {key: val.critical_path for key, val in o_stats.items()}
+            else:
+                raise NotImplementedError(color)
 
             if color.startswith("order-"):
 
@@ -819,7 +854,6 @@ def visualize(
                 engine = "cytoscape"
             except ImportError:
                 pass
-
     if engine == "graphviz":
         from dask.dot import dot_graph
 
diff --git a/dask/optimization.py b/dask/optimization.py
index 91a91d827..4ec7957d9 100644
--- a/dask/optimization.py
+++ b/dask/optimization.py
@@ -571,7 +571,10 @@ def fuse(
     if dependencies is None:
         deps = {k: get_dependencies(dsk, k, as_list=True) for k in dsk}
     else:
-        deps = dict(dependencies)
+        deps = {
+            k: v if isinstance(v, list) else get_dependencies(dsk, k, as_list=True)
+            for k, v in dependencies.items()
+        }
 
     rdeps = {}
     for k, vals in deps.items():
diff --git a/dask/order.py b/dask/order.py
index 9a8678660..c29196c3f 100644
--- a/dask/order.py
+++ b/dask/order.py
@@ -45,51 +45,46 @@ the docstrings of tests in dask/tests/test_order.py.  These tests usually have
 graph types laid out very carefully to show the kinds of situations that often
 arise, and the order we would like to be determined.
 
+"""
+from collections import defaultdict, deque, namedtuple
+from collections.abc import Iterable, Mapping, MutableMapping
+from typing import Any, Callable, Literal, NamedTuple, overload
 
-Policy
-------
-
-Work towards *small goals* with *big steps*.
-
-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.
+from dask.core import get_dependencies, get_deps, getcycle, istask, reverse_dict
+from dask.typing import Key
 
-    With more detail, we compute the total number of dependencies that each
-    task depends on (both its own dependencies, and the dependencies of its
-    dependencies, and so on), and then we choose those tasks that drive towards
-    results with a low number of total dependencies.  We choose to prioritize
-    tasks that work towards finishing shorter computations first.
 
-2.  **Big steps**: prefer tasks with many dependents
+class Order(NamedTuple):
+    priority: int
+    critical_path: float | int
 
-    However, many tasks work towards the same final dependents.  Among those,
-    we choose those tasks with the most work left to do.  We want to finish
-    the larger portions of a sub-computation before we start on the smaller
-    ones.
 
-3.  **Name comparison**: break ties with key name
+@overload
+def order(
+    dsk: Mapping[Key, Any],
+    dependencies: Mapping[Key, set[Key]] | None = None,
+    *,
+    return_stats: Literal[True],
+) -> dict[Key, Order]:
+    ...
 
-    Often graphs are made with regular keynames.  When no other structural
-    difference exists between two keys, use the key name to break ties.
-    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, namedtuple
-from collections.abc import Mapping, MutableMapping
-from math import log
-from typing import Any, cast
 
-from dask.core import get_dependencies, get_deps, getcycle, istask, reverse_dict
-from dask.typing import Key
+@overload
+def order(
+    dsk: Mapping[Key, Any],
+    dependencies: Mapping[Key, set[Key]] | None = None,
+    *,
+    return_stats: Literal[False],
+) -> dict[Key, int]:
+    ...
 
 
 def order(
-    dsk: MutableMapping[Key, Any],
-    dependencies: MutableMapping[Key, set[Key]] | None = None,
-) -> dict[Key, int]:
+    dsk: Mapping[Key, Any],
+    dependencies: Mapping[Key, set[Key]] | None = None,
+    *,
+    return_stats: bool = False,
+) -> dict[Key, Order] | dict[Key, int]:
     """Order nodes in dask graph
 
     This produces an ordering over our tasks that we use to break ties when
@@ -97,14 +92,7 @@ def order(
     scheduler and also to assist in static analysis.
 
     This currently traverses the graph as a single-threaded scheduler would
-    traverse it.  It breaks ties in the following ways:
-
-    1.  Begin at a leaf node that is a dependency of a root node that has the
-        largest subgraph (start hard things first)
-    2.  Prefer tall branches with few dependents (start hard things first and
-        try to avoid memory usage)
-    3.  Prefer dependents that are dependencies of root nodes that have
-        the smallest subgraph (do small goals that can terminate quickly)
+    traverse it.
 
     Examples
     --------
@@ -115,787 +103,518 @@ def order(
     {'a': 0, 'c': 1, 'b': 2, 'd': 3}
     """
     if not dsk:
-        return {}
+        return {}  # type: ignore
 
     dsk = dict(dsk)
 
     if dependencies is None:
         dependencies = {k: get_dependencies(dsk, k) for k in dsk}
+
     dependents = reverse_dict(dependencies)
     num_needed, total_dependencies = ndependencies(dependencies, dependents)
-    metrics = graph_metrics(dependencies, dependents, total_dependencies)
-
-    if len(metrics) != len(dsk):
+    if len(total_dependencies) != len(dsk):
         cycle = getcycle(dsk, None)
         raise RuntimeError(
             "Cycle detected between the following keys:\n  -> %s"
             % "\n  -> ".join(str(x) for x in cycle)
         )
 
-    # Single root nodes that depend on everything. These cause issues for
-    # the current ordering algorithm, since we often hit the root node
-    # and fell back to the key tie-breaker to choose which immediate dependency
-    # to finish next, rather than finishing off subtrees.
-    # So under the special case of a single root node that depends on the entire
-    # 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}
-    if len(root_nodes) > 1:
-        # This is also nice because it makes us robust to difference when
-        # computing vs persisting collections
-        root = cast(Key, object())
-
-        def _f(*args: Any, **kwargs: Any) -> None:
-            pass
-
-        dsk[root] = (_f, *root_nodes)
-        dependencies[root] = root_nodes
-        o = order(dsk, dependencies)
-        del o[root]
-        return o
-
-    init_stack: dict[Key, tuple] | set[Key] | list[Key]
-    # 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.
-    init_stack = {
-        # 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)
-            num_dependents - max_heights,
-            # tactically, finish small connected jobs first
-            num_dependents - min_heights,  # prefer tall and narrow
-            -total_dependents,  # take a big step
-            # try to be memory efficient
-            num_dependents,
-            # tie-breaker
-            StrComparable(key),
-        )
-        for key, num_dependents, (
-            total_dependents,
-            _,
-            _,
-            min_heights,
-            max_heights,
-        ) in (
-            (key, len(dependents[key]), metrics[key])
-            for key, val in dependencies.items()
-            if not val
-        )
-    }
-    # `initial_stack_key` chooses which task to run at the very beginning.
-    # This value is static, so we pre-compute as the value of this dict.
-    initial_stack_key = init_stack.__getitem__
+    leaf_nodes = {k for k, v in dependents.items() if not v}
+    root_nodes = {k for k, v in dependencies.items() if not v}
+    assert dependencies is not None
+    roots_connected, max_dependents = _connecting_to_roots(dependencies, dependents)
+    leafs_connected, _ = _connecting_to_roots(dependents, dependencies)
+    i = 0
+    result: dict[Key, Order | int] = {}
 
-    def dependents_key(x: Key) -> tuple:
-        """Choose a path from our starting task to our tactical goal
+    runnable_hull = set()
+    reachable_hull = set()
 
-        This path is connected to a large goal, but focuses on completing
-        a small goal and being memory efficient.
-        """
-        assert dependencies is not 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),
-        )
+    runnable = []
 
-    def dependencies_key(x: Key) -> tuple:
-        """Choose which dependency to run as part of a reverse DFS
+    known_runnable_paths: dict[Key, list[list[Key]]] = {}
+    known_runnable_paths_pop = known_runnable_paths.pop
 
-        This is very similar to both ``initial_stack_key``.
-        """
-        assert dependencies is not None
-        num_dependents = len(dependents[x])
-        (
-            total_dependents,
-            _,
-            _,
-            min_heights,
-            max_heights,
-        ) = metrics[x]
-        # Prefer short and narrow instead of tall in narrow, because we're going in
-        # reverse along dependencies.
-        return (
-            num_dependents + max_heights,
-            num_dependents + min_heights,  # prefer short and narrow
-            -total_dependencies[x],  # go where the work is
-            # try to be memory efficient
-            num_dependents - len(dependencies[x]) + num_needed[x],
-            num_dependents,
-            total_dependents,  # already found work, so don't add more
-            # tie-breaker
-            StrComparable(x),
-        )
-
-    root_total_dependencies = total_dependencies[list(root_nodes)[0]]
-    # Computing this for all keys can sometimes be relatively expensive :(
-    partition_keys = {
-        key: (
-            (root_total_dependencies - total_dependencies[key] + 1)
-            * (total_dependents - min_heights)
-        )
-        for key, (
-            total_dependents,
-            _,
-            _,
-            min_heights,
-            _,
-        ) in metrics.items()
-    }
+    crit_path_counter = 0
+    scrit_path: set[Key] = set()
+    _crit_path_counter_offset: int | float = 0
 
-    result: dict[Key, int] = {}
-    i = 0
+    _sort_keys_cache: dict[Key, tuple[int, int, int, int, str]] = {}
 
-    # `inner_stack` 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_stack = [min(init_stack, key=initial_stack_key)]
-    inner_stack_pop = inner_stack.pop
-    inner_stacks: list[list[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,
-    # 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[int, list[list[Key] | set[Key]]] = 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: list[list[Key]] = []
-    outer_stack_extend = outer_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(root_nodes)
-    seen_update = seen.update
-    seen_add = seen.add
-
-    # "singles" are tasks that are available to run, and when run may free a dependency.
-    # Although running a task to free a dependency may seem like a wash (net zero), it
-    # can be beneficial by providing more opportunities for a later task to free even
-    # more data.  So, it costs us little in the short term to more eagerly compute
-    # chains of tasks that keep the same number of data in memory, and the longer term
-    # rewards are potentially high.  I would expect a dynamic scheduler to have similar
-    # behavior, so I think it makes sense to do the same thing here in `dask.order`.
-    #
-    # When we gather tasks in `singles`, we do so optimistically: running the task *may*
-    # free the parent, but it also may not, because other dependents of the parent may
-    # be in the inner stacks.  When we process `singles`, we run tasks that *will* free
-    # the parent, otherwise we move the task to `later_singles`.  `later_singles` is run
-    # when there are no inner stacks, so it is safe to run all of them (because no other
-    # dependents will be hiding in the inner stacks to keep hold of the parent).
-    # `singles` is processed when the current item on the stack needs to compute
-    # dependencies before it can be run.
-    #
-    # Processing singles is meant to be a detour.  Doing so should not change our
-    # tactical goal in most cases.  Hence, we set `add_to_inner_stack = False`.
-    #
-    # In reality, this is a pretty limited strategy for running a task to free a
-    # dependency.  A thorough strategy would be to check whether running a dependent
-    # with `num_needed[dep] == 0` would free *any* of its dependencies.  This isn't
-    # what we do.  This data isn't readily or cheaply available.  We only check whether
-    # it will free its last dependency that was computed (the current `item`).  This is
-    # probably okay.  In general, our tactics and strategies for ordering try to be
-    # memory efficient, so we shouldn't try too hard to work around what we already do.
-    # However, sometimes the DFS nature of it leaves "easy-to-compute" stragglers behind.
-    # The current approach is very fast to compute, can be beneficial, and is generally
-    # low-risk.  There could be more we could do here though.  Our static scheduling
-    # here primarily looks at "what dependent should we run next?" instead of "what
-    # dependency should we try to free?"  Two sides to the same question, but a dynamic
-    # scheduler is much better able to answer the latter one, because it knows the size
-    # of data and can react to current state.  Does adding a little more dynamic-like
-    # behavior to `dask.order` add any tension to running with an actual dynamic
-    # scheduler?  Should we defer to dynamic schedulers and let them behave like this
-    # if they so choose?  Maybe.  However, I'm sensitive to the multithreaded scheduler,
-    # which is heavily dependent on the ordering obtained here.
-    singles: dict[Key, Key] = {}
-    singles_clear = singles.clear
-    later_singles: list[Key] = []
-    later_singles_append = later_singles.append
-    later_singles_clear = later_singles.clear
-
-    # Priority of being processed
-    #   1. inner_stack
-    #   2. singles (may be moved to later_singles)
-    #   3. inner_stacks
-    #   4. later_singles
-    #   5. next_nodes
-    #   6. outer_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
-            deps = set()
-            add_to_inner_stack = True
-            if inner_stack:
-                item = inner_stack_pop()
-                if item in result:
-                    continue
-                if num_needed[item]:
-                    if item not in root_nodes:
-                        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)
-                        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
-                    i += 1
-                    deps = dependents[item]
-                    add_to_inner_stack = True
-
-                    if deps:
-                        for dep in deps:
-                            num_needed[dep] -= 1
-                        process_singles = False
-                    else:
-                        continue
-            elif inner_stacks:
-                inner_stack = inner_stacks_pop()
-                inner_stack_pop = inner_stack.pop
+    def sort_key(x: Key) -> tuple[int, int, int, int, str]:
+        try:
+            return _sort_keys_cache[x]
+        except KeyError:
+            assert dependencies is not None
+            _sort_keys_cache[x] = rv = (
+                len(dependencies[x]),
+                total_dependencies[x],
+                len(roots_connected[x]),
+                -max_dependents[x],
+                # Converting to str is actually faster than creating some
+                # wrapper class and comparisons that come this far are
+                # relatively rare so we prefer fast init over fast comparison
+                str(x),
+            )
+            return rv
+
+    def add_to_result(item: Key) -> None:
+        nonlocal crit_path_counter
+        # Earlier versions recursed into this method but this could cause
+        # recursion depth errors. This is the only reason for the while loop
+        next_items = [item]
+        nonlocal i
+        while next_items:
+            item = next_items.pop()
+            assert not num_needed[item]
+            runnable_hull.discard(item)
+            reachable_hull.discard(item)
+            leaf_nodes.discard(item)
+            if item in result:
                 continue
-            elif singles:
-                process_singles = True
-            elif later_singles:
-                # No need to be optimistic: all nodes in `later_singles` will free a dependency
-                # when run, so no need to check whether dependents are in `seen`.
-                for single in later_singles:
-                    if single in result:
-                        continue
-                    while True:
-                        deps_singles = dependents[single]
-                        result[single] = i
-                        i += 1
-                        if deps_singles:
-                            for dep in deps_singles:
-                                num_needed[dep] -= 1
-                            if len(deps_singles) == 1:
-                                # Fast path!  We trim down `dep2` above hoping to reach here.
-                                (single,) = deps_singles
-                                if not num_needed[single]:
-                                    # Keep it going!
-                                    deps_singles = dependents[single]
-                                    continue
-                            deps |= deps_singles
-                        del deps_singles
-                        break
-                later_singles_clear()
-                deps = set_difference(deps, result)
-                if not deps:
-                    continue
-                add_to_inner_stack = False
-                process_singles = True
+            if return_stats:
+                result[item] = Order(i, crit_path_counter - _crit_path_counter_offset)
             else:
-                break
-
-            if process_singles and singles:
-                # We gather all dependents of all singles into `deps`, which we then process below.
-
-                add_to_inner_stack = True if inner_stack or inner_stacks else False
-                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 (
-                        len(
-                            set_difference(
-                                set_difference(dependents[parent], result),
-                                seen,
-                            )
-                        )
-                        > 1
-                    ):
-                        later_singles_append(single)
-                        continue
-                    while True:
-                        deps_singles = dependents[single]
-                        result[single] = i
-                        i += 1
-                        if deps_singles:
-                            for dep in deps_singles:
-                                num_needed[dep] -= 1
-                            if add_to_inner_stack:
-                                already_seen = deps_singles & seen
-                                if already_seen:
-                                    # This means that the singles path also
-                                    # leads to the current or previous strategic
-                                    # path
-                                    if len(deps_singles) == len(already_seen):
-                                        if len(already_seen) == 1:
-                                            (single,) = already_seen
-                                            if not num_needed[single]:
-                                                deps_singles = dependents[single]
-                                                continue
-                                        break
-                                    deps_singles = deps_singles - already_seen
-                            else:
-                                already_seen = set()
-                            if len(deps_singles) == 1:
-                                # Fast path!  We trim down `dep2` above hoping to reach here.
-                                (single,) = deps_singles
-                                if not num_needed[single]:
-                                    if not already_seen:
-                                        # Keep it going!
-                                        deps_singles = dependents[single]
-                                        continue
-                                    later_singles_append(single)
-                                    break
-                            deps |= deps_singles
-                        del deps_singles
-                        break
-                del singles_keys
-                deps = set_difference(deps, result)
-                singles_clear()
-                if not deps:
-                    continue
-                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 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.
-            already_seen = deps & seen
-            if already_seen:
-                if len(deps) == len(already_seen):
-                    if len(already_seen) == 1:
-                        (dep,) = already_seen
-                        if not num_needed[dep]:
-                            singles[dep] = item
-                        del dep
-                    continue
-                add_to_inner_stack = False
-                deps = deps - already_seen
-
-            if len(deps) == 1:
-                # Fast path!  We trim down `deps` above hoping to reach here.
-                (dep,) = deps
-                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]
+                result[item] = i
+            i += 1
+            # Note: This is a `set` and therefore this introduces a certain
+            # randomness. However, this randomness should not have any impact on
+            # the final result since the `process_runnable` should produce
+            # equivalent results regardless of the order in which runnable is
+            # populated (not identical but equivalent)
+            for dep in dependents[item]:
+                num_needed[dep] -= 1
+                reachable_hull.add(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
-                else:
-                    next_nodes[key].append(deps)
-                del dep, key
-            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.
-                #
-                # Let me explain why this is a special case.  If we put the better dependent
-                # onto the inner stack, then it's guaranteed to run next.  After it's run,
-                # then running the other dependent *may* allow their parent to be freed.
-                dep, dep2 = deps
-                key = partition_keys[dep]
-                key2 = partition_keys[dep2]
-                if (
-                    key2 < key
-                    or key == key2
-                    and dependents_key(dep2) < dependents_key(dep)
-                ):
-                    dep, dep2 = dep2, dep
-                    key, key2 = key2, key
-                if inner_stack:
-                    prev_key = partition_keys[inner_stack[0]]
-                    if key2 < prev_key:
-                        inner_stacks_append(inner_stack)
-                        inner_stacks_append([dep2])
-                        inner_stack = [dep]
-                        inner_stack_pop = inner_stack.pop
-                        seen_update(deps)
-                        if not num_needed[dep2]:
-                            if process_singles:
-                                later_singles_append(dep2)
-                            else:
-                                singles[dep2] = item
-                    elif key < prev_key:
-                        inner_stacks_append(inner_stack)
-                        inner_stack = [dep]
-                        inner_stack_pop = inner_stack.pop
-                        seen_add(dep)
-                        if not num_needed[dep2]:
-                            if process_singles:
-                                later_singles_append(dep2)
-                            else:
-                                singles[dep2] = item
-                        else:
-                            next_nodes[key2].append([dep2])
+                    if len(dependents[item]) == 1:
+                        next_items.append(dep)
                     else:
-                        item_key = partition_keys[item]
-                        for k, d in [(key, dep), (key2, dep2)]:
-                            if not num_needed[d]:
-                                if process_singles:
-                                    later_singles_append(d)
+                        runnable.append(dep)
+
+    def _with_offset(func: Callable[..., None]) -> Callable[..., None]:
+        # This decorator is only used to reduce indentation levels. The offset
+        # is purely cosmetical and used for some visualizations and I haven't
+        # settled on how to implement this best so I didn't want to have large
+        # indentations that make things harder to read
+        nonlocal _crit_path_counter_offset
+
+        def wrapper(*args: Any, **kwargs: Any) -> None:
+            nonlocal _crit_path_counter_offset
+            _crit_path_counter_offset = 0.5
+            try:
+                func(*args, **kwargs)
+            finally:
+                _crit_path_counter_offset = 0
+
+        return wrapper
+
+    @_with_offset
+    def process_runnables() -> None:
+        """Compute all currently runnable paths and either cache or execute them
+
+        This is designed to ensure we are running tasks that are free to execute
+        (e.g. the result of a splitter task) not too eagerly. If we executed
+        such free tasks too early we'd be walking the graph in a too wide /
+        breadth first fashion that is not optimal. If instead we were to only
+        execute them once they are needed for a final result, this can cause
+        very high memory pressure since valueable reducers are executed too
+        late.
+
+        The strategy here is to take all runnable tasks and walk forwards until
+        we hit a reducer node (i.e. a node with more than one dependency). We
+        will remember/cache the path to this reducer node.
+        If this path leads to a leaf or if we find enough runnable paths for a
+        reducer to be runnable, we will execute the path.
+
+        If instead of a reducer a splitter is encountered that is runnable, we
+        will follow it's splitter paths individually and apply the same logic to
+        each branch.
+        """
+        while runnable:
+            candidates = runnable.copy()
+            runnable.clear()
+            while candidates:
+                key = candidates.pop()
+                if key in runnable_hull or key in result:
+                    continue
+                if key in leaf_nodes:
+                    add_to_result(key)
+                    continue
+                path = [key]
+                branches = deque([path])
+                while branches:
+                    path = branches.popleft()
+                    while True:
+                        current = path[-1]
+                        runnable_hull.add(current)
+                        deps_downstream = dependents[current]
+                        deps_upstream = dependencies[current]  # type: ignore
+                        if not deps_downstream:
+                            # FIXME: The fact that it is possible for
+                            # num_needed[current] == 0 means we're doing some
+                            # work twice
+                            if num_needed[current] <= 1 or (
+                                not branches
+                                # FIXME: This is a very magical number
+                                and len(path) > 2
+                            ):
+                                for k in path[:-1]:
+                                    add_to_result(k)
+                                if not num_needed[current]:
+                                    add_to_result(current)
+                        elif len(path) == 1 or len(deps_upstream) == 1:
+                            if len(deps_downstream) > 1:
+                                for d in sorted(deps_downstream, key=sort_key):
+                                    # This ensures we're only considering splitters
+                                    # that are genuinely splitting and not
+                                    # interleaving
+                                    if len(dependencies[d]) == 1:  # type: ignore
+                                        branch = path.copy()
+                                        branch.append(d)
+                                        branches.append(branch)
+                                break
+                            runnable_hull.update(deps_downstream)
+                            path.extend(sorted(deps_downstream, key=sort_key))
+                            continue
+                        elif current in known_runnable_paths:
+                            known_runnable_paths[current].append(path)
+                            if (
+                                len(known_runnable_paths[current])
+                                >= num_needed[current]
+                            ):
+                                pruned_branches: deque[list[Key]] = deque()
+                                for path in known_runnable_paths_pop(current):
+                                    if path[-2] not in result:
+                                        pruned_branches.append(path)
+                                if len(pruned_branches) < num_needed[current]:
+                                    known_runnable_paths[current] = list(
+                                        pruned_branches
+                                    )
                                 else:
-                                    singles[d] = item
-                            else:
-                                next_nodes[k].append([d])
-                        del item_key
-                    del prev_key
-                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]:
-                            singles[dep2] = item
-                        elif key == key2 and 5 * partition_keys[item] > 22 * key:
-                            inner_stacks_append([dep2])
-                            seen_add(dep2)
+                                    while pruned_branches:
+                                        path = pruned_branches.popleft()
+                                        for k in path:
+                                            if num_needed[k]:
+                                                pruned_branches.append(path)
+                                                break
+                                            add_to_result(k)
                         else:
-                            next_nodes[key2].append([dep2])
-                    else:
-                        for k, d in [(key, dep), (key2, dep2)]:
-                            next_nodes[k].append([d])
-                del dep, dep2, key, key2
-            else:
-                # Slow path :(.  This requires grouping by partition_key.
-                dep_pools = defaultdict(set)
-                possible_singles = defaultdict(set)
-                for dep in deps:
-                    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
-                    prev_key = partition_keys[inner_stack[0]]
-                    now_keys = []  # < inner_stack[0]
-                    psingles = set()
-                    for key, vals in dep_pools.items():
-                        if key < prev_key:
-                            now_keys.append(key)
-                        else:
-                            psingles = possible_singles[key]
-                            for s in psingles:
-                                singles[s] = item
-                            vals -= psingles
-                            next_nodes[key].append(vals)
-                        del vals, key
-                    del psingles
-                    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: set[Key] | list[Key]
-                            pool = dep_pools[key]
-                            if 1 < len(pool) < 100:
-                                pool = sorted(pool, key=dependents_key, reverse=True)
-                            inner_stacks_extend([dep] for dep in pool)
-                            seen_update(pool)
-                            del pool
-                        inner_stack = inner_stacks_pop()
-                        inner_stack_pop = inner_stack.pop
-                    del now_keys, prev_key
-                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_pool: list[Key] | set[Key]
-                        min_key = min(dep_pools)
-                        min_pool = dep_pools.pop(min_key)
-                        if len(min_pool) == 1:
-                            inner_stack = list(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 = 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)
-                        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)]
+                            if (
+                                len(dependencies[current]) > 1  # type: ignore
+                                and num_needed[current] <= 1
+                            ):
+                                for k in path:
+                                    add_to_result(k)
                             else:
-                                inner_stack = [min_pool.pop()]
-                            next_nodes[min_key].append(min_pool)
-                            seen_update(inner_stack)
-                        del min_pool, min_key
-                        inner_stack_pop = inner_stack.pop
-                    for key, vals in dep_pools.items():
-                        psingles = possible_singles[key]
-                        for s in psingles:
-                            singles[s] = item
-                        vals -= psingles
-                        next_nodes[key].append(vals)
-                        del key, 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(list(el) for el in reversed(next_nodes[key]))
-            next_nodes.clear()
-
-        outer_deps = []
-        while outer_stack:
-            # Try to add a few items to `inner_stacks`
-            outer_deps = [x for x in outer_stack_pop() if x not in result]
-            if outer_deps:
-                if 1 < len(outer_deps) < 100:
-                    outer_deps.sort(key=dependents_key, reverse=True)
-                inner_stacks_extend([dep] for dep in outer_deps)
-                seen_update(outer_deps)
-                break
-        del outer_deps
-
-        if inner_stacks:
-            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.
-        # If we have many small groups left, then it's best to sort.
-        # If we have many tiny groups left, then it's best to simply iterate.
-        if not is_init_sorted:
-            prev_len = len(init_stack)
-            init_stack = set(init_stack)
-            init_stack = set_difference(init_stack, result)
-            N = len(init_stack)
-            m = prev_len - N
-            # is `min` likely better than `sort`?
-            if m >= N or N + (N - m) * log(N - m) < N * log(N):
-                item = min(init_stack, key=initial_stack_key)
-                continue
+                                known_runnable_paths[current] = [path]
+                        break
 
-            if len(init_stack) < 10000:
-                init_stack = sorted(init_stack, key=initial_stack_key, reverse=True)
-            else:
-                init_stack = list(init_stack)
-            init_stack_pop = init_stack.pop
-            is_init_sorted = True
+    # Pick strategy
+    # Note: We're trying to be smart here by picking a strategy on how to
+    # determine the critical path. This is not always clear and we may want to
+    # consider just calculating both orderings and picking the one with less
+    # pressure. The only concern to this would be performance but at time of
+    # writing, the most expensive part of ordering is the prep work (mostly
+    # connected roots + sort_key) which can be reused for multiple orderings.
+
+    def use_longest_path() -> bool:
+        size = 0
+        # Heavy reducer / splitter topologies often benefit from a very
+        # traditional critical path that expresses the longest chain of
+        # tasks.
+        if abs(len(root_nodes) - len(leaf_nodes)) / len(root_nodes) < 0.8:
+            # If the graph stays about the same, we are checking for symmetry
+            # and choose a "quickest path first" approach if the graph appears
+            # to be asymmetrical
+            for r in root_nodes:
+                if not size:
+                    size = len(leafs_connected[r])
+                elif size != len(leafs_connected[r]):
+                    return False
+
+        return True
+
+    longest_path = use_longest_path()
+
+    def get_target() -> Key:
+        raise NotImplementedError()
+
+    if not longest_path:
+
+        def _build_get_target() -> Callable[[], Key]:
+            occurences: defaultdict[Key, int] = defaultdict(int)
+            for t in leaf_nodes:
+                for r in roots_connected[t]:
+                    occurences[r] += 1
+            occurences_grouped = defaultdict(set)
+            for root, occ in occurences.items():
+                occurences_grouped[occ].add(root)
+            occurences_grouped_sorted = {}
+            for k, v in occurences_grouped.items():
+                occurences_grouped_sorted[k] = sorted(v, key=sort_key, reverse=True)
+            del occurences_grouped, occurences
+
+            def pick_seed() -> Key | None:
+                while occurences_grouped_sorted:
+                    key = max(occurences_grouped_sorted)
+                    picked_root = occurences_grouped_sorted[key][-1]
+                    if picked_root in result:
+                        occurences_grouped_sorted[key].pop()
+                        if not occurences_grouped_sorted[key]:
+                            del occurences_grouped_sorted[key]
+                        continue
+                    return picked_root
+                return None
+
+            def get_target() -> Key:
+                candidates = leaf_nodes
+                skey: Callable = sort_key
+
+                if runnable_hull:
+                    skey = lambda k: (num_needed[k], sort_key(k))
+                    candidates = runnable_hull & candidates
+                elif reachable_hull:
+                    skey = lambda k: (num_needed[k], sort_key(k))
+                    candidates = reachable_hull & candidates
+
+                if not candidates:
+                    if seed := pick_seed():
+                        candidates = leafs_connected[seed]
+                    else:
+                        candidates = runnable_hull or reachable_hull
+                # FIXME: This can be very expensive
+                return min(candidates, key=skey)
 
-        if item in root_nodes:
-            item = init_stack_pop()
+            return get_target
 
-        while item in result:
-            item = init_stack_pop()
-        inner_stack.append(item)
+        get_target = _build_get_target()
+    else:
+        leaf_nodes_sorted = sorted(leaf_nodes, key=sort_key, reverse=False)
+        get_target = leaf_nodes_sorted.pop
+        del leaf_nodes_sorted
 
-    return result
+    # *************************************************************************
+    # CORE ALGORITHM STARTS HERE
+    #
+    # 0. Nomenclature
+    #
+    #   - roots: Nodes that have no dependencies (i.e. typically data producers)
+    #   - leafs: Nodes that have no dependents (i.e. user requested keys)
+    #   - critical_path: The strategic path through the graph.
+    #   - walking forwards: Starting from a root node we walk the graph as if we
+    #     were to compute the individual nodes, i.e. along dependents
+    #   - walking backwards: Starting from a leaf node we walk the graph in
+    #     reverse direction, i.e. along dependencies
+    #   - runnable: Nodes that are ready to be computed because all their
+    #     dependencies are in result
+    #   - runnable_hull: Nodes that could be reached and executed without
+    #     "walking back". This typically means that these are tasks than can be
+    #     executed without loading additional data/executing additional root
+    #     nodes
+    #  -  reachable_hull: Nodes that are touching the result, i.e. all nodes in
+    #     reachable_hull have at least one dependency in result
+    #
+    # A. Build the critical path
+    #
+    #   To build the critical path we will use a provided `get_target` function
+    #   that returns a node that is anywhere in the graph, typically a leaf
+    #   node. This node is not required to be runnable. We will walk the graph
+    #   backwards, i.e. from leafs to roots and append nodes to the graph as we
+    #   go. The critical path is a
+    #   linear path in the graph. While this is a viable strategy, it is not
+    #   required for the critical path to be a classical "longest path" but it
+    #   can define any route through the graph that should be considered as top
+    #   priority.
+    #
+    #   1. Determine the target node by calling `get_target`` and append the
+    #      target to the critical path stack
+    #   2. Take the _most valuable_ (max given a `sort_key`) of it's dependends
+    #      and append it to the critical path stack. This key is the new target.
+    #   3. Repeat step 2 until we reach a node that has no dependencies and is
+    #      therefore runnable
+    #
+    # B. Walk the critical path
+    #
+    #   Only the first element of the critical path is an actually runnable node
+    #   and this is where we're starting the sort. Strategically, this is the
+    #   most important goal to achieve but since not all of the nodes are
+    #   immediately runnable we have to walk back and compute other nodes first
+    #   before we can unlock the critical path. This typically requires us also
+    #   to load more data / run more root tasks.
+    #   While walking the critical path we will also unlock non-critical tasks
+    #   that could be run but are not contributing to our strategic goal. Under
+    #   certain circumstances, those runnable tasks are allowed to be run right
+    #   away to reduce memory pressure. This is described in more detail in
+    #   `process_runnable`.
+    #   Given this, the algorithm is as follows:
+    #
+    #   1. Pop the first element of the critical path
+    #   2a. If the node is already in the result, continue
+    #   2b. If the node is not runnable, we will put it back on the stack and
+    #       put all its dependencies on the stack and continue with step 1. This
+    #       is what we refer to as "walking back"
+    #   2c. Else, we add the node to the result
+    #   3.  If we previously had to walk back we will consider running
+    #       non-critical tasks (by calling process_runnables)
+    #   4a. If critical path is not empty, repeat step 1
+    #   4b. Go back to A.) and build a new critical path given a new target that
+    #       accounts for the already computed nodes.
+    #
+    # *************************************************************************
 
+    critical_path: list[Key] = []
+    cpath_append = critical_path.append
+    scpath_add = scrit_path.add
 
-def graph_metrics(
-    dependencies: Mapping[Key, set[Key]],
-    dependents: Mapping[Key, set[Key]],
-    total_dependencies: Mapping[Key, int],
-) -> dict[Key, tuple[int, int, int, int, int]]:
-    r"""Useful measures of a graph used by ``dask.order.order``
+    def path_append(item: Key) -> None:
+        cpath_append(item)
+        scpath_add(item)
 
-    Example DAG (a1 has no dependencies; b2 and c1 are root nodes):
+    scpath_update = scrit_path.update
+    cpath_extend = critical_path.extend
 
-    c1
-    |
-    b1  b2
-     \  /
-      a1
+    def path_extend(items: Iterable[Key]) -> None:
+        cpath_extend(items)
+        scpath_update(items)
 
-    For each key we return:
+    cpath_pop = critical_path.pop
+    scpath_discard = scrit_path.discard
 
-    1.  **total_dependents**: The number of keys that can only be run
-        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.
+    def path_pop() -> Key:
+        item = cpath_pop()
+        scpath_discard(item)
+        return item
 
-        1
-        |
-        2   1
-         \ /
-          4
+    while len(result) < len(dsk):
+        crit_path_counter += 1
+        assert not critical_path
+        assert not scrit_path
 
-    2.  **min_dependencies**: The minimum value of the total number of
-        dependencies of all final dependents (see module-level comment for more).
-        In other words, the minimum of ``ndependencies`` of root
-        nodes connected to the current node.
+        # A. Build the critical path
+        target = get_target()
+        next_deps = dependencies[target]
+        path_append(target)
 
-        3
-        |
-        3   2
-         \ /
-          2
+        while next_deps:
+            item = max(next_deps, key=sort_key)
+            path_append(item)
+            next_deps = dependencies[item]
+            path_extend(next_deps)
 
-    3.  **max_dependencies**: The maximum value of the total number of
-        dependencies of all final dependents (see module-level comment for more).
-        In other words, the maximum of ``ndependencies`` of root
-        nodes connected to the current node.
+        # B. Walk the critical path
 
-        3
-        |
-        3   2
-         \ /
-          3
+        walked_back = False
+        while critical_path:
+            item = path_pop()
+            if item in result:
+                continue
+            if num_needed[item]:
+                if item in known_runnable_paths:
+                    for path in known_runnable_paths_pop(item):
+                        path_extend(reversed(path))
+                    continue
+                path_append(item)
+                deps = dependencies[item].difference(result)
+                unknown: list[Key] = []
+                known: list[Key] = []
+                k_append = known.append
+                uk_append = unknown.append
+                for d in sorted(deps, key=sort_key):
+                    if d in known_runnable_paths:
+                        k_append(d)
+                    else:
+                        uk_append(d)
+                if len(unknown) > 1:
+                    walked_back = True
 
-    4.  **min_height**: The minimum height from a root node
+                for d in unknown:
+                    path_append(d)
+                for d in known:
+                    for path in known_runnable_paths_pop(d):
+                        path_extend(reversed(path))
 
-        0
-        |
-        1   0
-         \ /
-          1
+                del deps
+                continue
+            else:
+                if walked_back and len(runnable) < len(critical_path):
+                    process_runnables()
+                add_to_result(item)
+        process_runnables()
 
-    5.  **max_height**: The maximum height from a root node
+    return result  # type: ignore
 
-        0
-        |
-        1   0
-         \ /
-          2
 
-    Examples
-    --------
-    >>> inc = lambda x: x + 1
-    >>> dsk = {'a1': 1, 'b1': (inc, 'a1'), 'b2': (inc, 'a1'), 'c1': (inc, 'b1')}
-    >>> dependencies, dependents = get_deps(dsk)
-    >>> _, total_dependencies = ndependencies(dependencies, dependents)
-    >>> metrics = graph_metrics(dependencies, dependents, total_dependencies)
-    >>> sorted(metrics.items())
-    [('a1', (4, 2, 3, 1, 2)), ('b1', (2, 3, 3, 1, 1)), ('b2', (1, 2, 2, 0, 0)), ('c1', (1, 3, 3, 0, 0))]
+def _connecting_to_roots(
+    dependencies: Mapping[Key, set[Key]], dependents: Mapping[Key, set[Key]]
+) -> tuple[dict[Key, set[Key]], dict[Key, int]]:
+    """Determine for every node which root nodes are connected to it (i.e.
+    ancestors). If arguments of dependencies and dependents are switched, this
+    can also be used to determine which leaf nodes are connected to which node
+    (i.e. descendants).
+
+    Also computes a weight that is defined as (cheaper to compute here)
+
+            `max(len(dependents[k]) for k in connected_roots[key])`
 
-    Returns
-    -------
-    metrics: Dict[key, Tuple[int, int, int, int, int]]
     """
     result = {}
-    num_needed = {k: len(v) for k, v in dependents.items() if v}
-    current: list[Key] = []
-    current_pop = current.pop
-    current_append = current.append
-    for key, deps in dependents.items():
-        if not deps:
-            val = total_dependencies[key]
-            result[key] = (1, val, val, 0, 0)
-            for child in dependencies[key]:
+    current = []
+    num_needed = {k: len(v) for k, v in dependencies.items() if v}
+    max_dependents = {}
+    roots = set()
+    for k, v in dependencies.items():
+        if not v:
+            # Note: Hashing the full keys is relatively expensive. Hashing
+            # integers would be much faster so this could be sped up by just
+            # introducing a counter here. However, the order algorithm is also
+            # sometimes interested in the actual keys and the only way to
+            # benefit from the speedup of using integers would be to convert
+            # this back on demand which makes the code very hard to read.
+            roots.add(k)
+            result[k] = {k}
+            deps = dependents[k]
+            max_dependents[k] = len(deps)
+            for child in deps:
                 num_needed[child] -= 1
                 if not num_needed[child]:
-                    current_append(child)
-
+                    current.append(child)
     while current:
-        key = current_pop()
-        parents = dependents[key]
-        if len(parents) == 1:
-            (parent,) = parents
-            (
-                total_dependents,
-                min_dependencies,
-                max_dependencies,
-                min_heights,
-                max_heights,
-            ) = result[parent]
-            result[key] = (
-                1 + total_dependents,
-                min_dependencies,
-                max_dependencies,
-                1 + min_heights,
-                1 + max_heights,
-            )
-        else:
-            (
-                total_dependents_,
-                min_dependencies_,
-                max_dependencies_,
-                min_heights_,
-                max_heights_,
-            ) = zip(*(result[parent] for parent in dependents[key]))
-            result[key] = (
-                1 + sum(total_dependents_),
-                min(min_dependencies_),
-                max(max_dependencies_),
-                1 + min(min_heights_),
-                1 + max(max_heights_),
-            )
+        key = current.pop()
+        for parent in dependents[key]:
+            num_needed[parent] -= 1
+            if not num_needed[parent]:
+                current.append(parent)
+        # At some point, all the roots are the same, particularly for dense
+        # graphs. We don't want to create new sets over and over again
+        new_set = None
+        identical_sets = True
+        result_first = None
+
         for child in dependencies[key]:
-            num_needed[child] -= 1
-            if not num_needed[child]:
-                current_append(child)
-    return result
+            r_child = result[child]
+            if not result_first:
+                result_first = r_child
+                max_dependents[key] = max_dependents[child]
+            # This clause is written such that it can circuit break early
+            elif not (  # type: ignore[unreachable]
+                identical_sets
+                and (result_first is r_child or r_child.issubset(result_first))
+            ):
+                identical_sets = False
+                if not new_set:
+                    new_set = result_first.copy()
+                max_dependents[key] = max(max_dependents[child], max_dependents[key])
+                new_set.update(r_child)
+
+        assert new_set is not None or result_first is not None
+        result[key] = new_set or result_first
+
+    # The order algo doesn't care about this but this makes it easier to
+    # understand and shouldn't take that much time
+    empty_set: set[Key] = set()
+    for r in roots:
+        result[r] = empty_set
+    return result, max_dependents
 
 
 def ndependencies(
@@ -947,37 +666,6 @@ def ndependencies(
     return num_dependencies, result
 
 
-class StrComparable:
-    """Wrap object so that it defaults to string comparison
-
-    When comparing two objects of different types Python fails
-
-    >>> 'a' < 1
-    Traceback (most recent call last):
-        ...
-    TypeError: '<' not supported between instances of 'str' and 'int'
-
-    This class wraps the object so that, when this would occur it instead
-    compares the string representation
-
-    >>> StrComparable('a') < StrComparable(1)
-    False
-    """
-
-    __slots__ = ("obj",)
-
-    obj: Any
-
-    def __init__(self, obj: Any):
-        self.obj = obj
-
-    def __lt__(self, other: Any) -> bool:
-        try:
-            return self.obj < other.obj
-        except Exception:
-            return str(self.obj) < str(other.obj)
-
-
 OrderInfo = namedtuple(
     "OrderInfo",
     (
@@ -1012,8 +700,9 @@ def diagnostics(
         dependencies, dependents = get_deps(dsk)
     else:
         dependents = reverse_dict(dependencies)
+    assert dependencies is not None
     if o is None:
-        o = order(dsk, dependencies=dependencies)
+        o = order(dsk, dependencies=dependencies, return_stats=False)
 
     pressure = []
     num_in_memory = 0
@@ -1068,8 +757,10 @@ def _convert_task(task: Any) -> Any:
             elif isinstance(el, list):
                 new_spec.append([_convert_task(e) for e in el])
         return (_f, *new_spec)
+    elif isinstance(task, tuple):
+        return (_f, task)
     else:
-        return task
+        return (_f, *task)
 
 
 def sanitize_dsk(dsk: MutableMapping[Key, Any]) -> dict:


Test patch:
diff --git a/dask/dataframe/tests/test_groupby.py b/dask/dataframe/tests/test_groupby.py
index f5a5b49bc..4d604b92d 100644
--- a/dask/dataframe/tests/test_groupby.py
+++ b/dask/dataframe/tests/test_groupby.py
@@ -838,7 +838,9 @@ def test_split_apply_combine_on_series(empty):
 
 
 @pytest.mark.parametrize("keyword", ["split_every", "split_out"])
-def test_groupby_reduction_split(keyword, agg_func):
+def test_groupby_reduction_split(keyword, agg_func, shuffle_method):
+    if agg_func in {"first", "last"} and shuffle_method == "disk":
+        pytest.skip(reason="https://github.com/dask/dask/issues/10034")
     pdf = pd.DataFrame(
         {"a": [1, 2, 6, 4, 4, 6, 4, 3, 7] * 100, "b": [4, 2, 7, 3, 3, 1, 1, 1, 2] * 100}
     )
@@ -2667,7 +2669,7 @@ def test_groupby_shift_within_partition_sorting():
         )
 
 
-def test_groupby_shift_with_freq():
+def test_groupby_shift_with_freq(shuffle_method):
     pdf = pd.DataFrame(
         dict(a=[1, 2, 3, 4, 5, 6], b=[0, 0, 0, 1, 1, 1]),
         index=pd.date_range(start="20100101", periods=6),
@@ -2685,7 +2687,13 @@ def test_groupby_shift_with_freq():
         check_freq=False,
     )
     df_result = pdf.groupby("b").shift(periods=-2, freq="D")
-    assert_eq(df_result, ddf.groupby("b").shift(periods=-2, freq="D", meta=df_result))
+    assert_eq(
+        df_result,
+        ddf.groupby("b").shift(periods=-2, freq="D", meta=df_result),
+        # Somehow https://github.com/dask/dask/issues/10034 causes us to
+        # drop the `freq`
+        check_freq=shuffle_method != "disk",
+    )
 
 
 @pytest.mark.parametrize(
diff --git a/dask/tests/test_local.py b/dask/tests/test_local.py
index 31a2370d6..148f38c74 100644
--- a/dask/tests/test_local.py
+++ b/dask/tests/test_local.py
@@ -169,7 +169,7 @@ def test_ordering():
 
     get_sync(dsk, "y")
 
-    assert L == sorted(L)
+    assert L == sorted(L, reverse=True)
 
 
 def test_complex_ordering():
diff --git a/dask/tests/test_optimization.py b/dask/tests/test_optimization.py
index b0df94163..397f6fc23 100644
--- a/dask/tests/test_optimization.py
+++ b/dask/tests/test_optimization.py
@@ -227,6 +227,48 @@ def test_fuse_keys():
     )
 
 
+def test_donot_substitute_same_key_multiple_times():
+    already_called = False
+
+    def inc_only_once(x):
+        nonlocal already_called
+        if already_called:
+            raise RuntimeError
+        already_called = True
+        return x + 1
+
+    # This is the graph topology of a Z.T @ Z after blockwise fusion as given in
+    # https://github.com/dask/dask/issues/10645
+    dsk = {
+        # Note that there is some logic in there that actually checks the type
+        # of this and an integer is inlined while an array is not. However, data
+        # keys are also flagged as forbidden to be fused so this should not be a
+        # problem
+        "A": 42,  # Array
+        "B": (inc_only_once, "A"),  # Add
+        "C": (add, "B", "B"),  # matmul
+        "D": (inc, "C"),  # getitem
+    }
+
+    # fuse accepts dependencies to avoid having to recompute it. However,
+    # internally it actually requires it to be in the format of dict[Key,
+    # list[Key]] which is rarely used elsewhere since most applications are
+    # using lists.
+    # It uses these lists to infer duplicates in dependencies and avoids fusing
+    # those sicne substituting the same key multiple times also causes them to
+    # be computed multiple times. A better logic would replace those with a
+    # SubGraphCallable
+    dependencies = {"A": set(), "B": {"A"}, "C": {"B"}, "D": {"C"}}
+    fused_dsk = fuse(
+        dsk,
+        keys=["A", "D"],
+        dependencies=dependencies,
+    )[0]
+    from dask.core import get
+
+    assert get(fused_dsk, "D") > 1
+
+
 def test_inline():
     d = {"a": 1, "b": (inc, "a"), "c": (inc, "b"), "d": (add, "a", "c")}
     assert inline(d) == {"a": 1, "b": (inc, 1), "c": (inc, "b"), "d": (add, 1, "c")}
diff --git a/dask/tests/test_order.py b/dask/tests/test_order.py
index de43cc14a..16e387acc 100644
--- a/dask/tests/test_order.py
+++ b/dask/tests/test_order.py
@@ -1,27 +1,84 @@
 from __future__ import annotations
 
+import inspect
+
 import pytest
 
 import dask
-from dask.base import collections_to_dsk
+from dask import delayed
+from dask.base import collections_to_dsk, key_split, visualize_dsk
 from dask.core import get_deps
-from dask.order import diagnostics, ndependencies, order
+from dask.order import _connecting_to_roots, diagnostics, ndependencies, order
 from dask.utils_test import add, inc
 
 
-@pytest.fixture(params=["abcde", "edcba"])
+@pytest.fixture(
+    params=[
+        "abcde",
+        "edcba",
+    ]
+)
 def abcde(request):
     return request.param
 
 
-def issorted(L, reverse=False):
-    return sorted(L, reverse=reverse) == L
-
-
 def f(*args):
     pass
 
 
+def visualize(dsk, **kwargs):
+    """Utility to visualize the raw low level graphs in this tests suite. This
+    automatically generates a set of visualizations with different metrics and
+    writes them out to a file prefixed by the test name and suffixed by the
+    measure used."""
+    funcname = inspect.stack()[1][3]
+    if hasattr(dsk, "__dask_graph__"):
+        dsk = collections_to_dsk([dsk], optimize_graph=True)
+
+    node_attrs = {"penwidth": "6"}
+    visualize_dsk(dsk, filename=f"{funcname}.pdf", node_attr=node_attrs, **kwargs)
+    o = order(dsk, return_stats=True)
+    visualize_dsk(
+        dsk,
+        o=o,
+        filename=f"{funcname}-order.pdf",
+        color="order",
+        node_attr=node_attrs,
+        cmap="viridis",
+        collapse_outputs=False,
+        **kwargs,
+    )
+    visualize_dsk(
+        dsk,
+        o=o,
+        filename=f"{funcname}-age.pdf",
+        color="age",
+        node_attr=node_attrs,
+        cmap="Reds",
+        collapse_outputs=False,
+        **kwargs,
+    )
+    visualize_dsk(
+        dsk,
+        o=o,
+        filename=f"{funcname}-pressure.pdf",
+        color="memorypressure",
+        node_attr=node_attrs,
+        cmap="Reds",
+        collapse_outputs=False,
+        **kwargs,
+    )
+    visualize_dsk(
+        dsk,
+        o=o,
+        filename=f"{funcname}-cpath.pdf",
+        color="cpath",
+        node_attr=node_attrs,
+        collapse_outputs=False,
+        **kwargs,
+    )
+
+
 def test_ordering_keeps_groups_together(abcde):
     a, b, c, d, e = abcde
     d = {(a, i): (f,) for i in range(4)}
@@ -46,7 +103,9 @@ def test_avoid_broker_nodes(abcde):
     |      \  /
     a0      a1
 
-    a0 should be run before a1
+    There are good arguments for both a0 or a1 to run first. Regardless of what
+    we choose to run first, we should finish the computation branch before
+    moving to the other one
     """
     a, b, c, d, e = abcde
     dsk = {
@@ -57,29 +116,7 @@ def test_avoid_broker_nodes(abcde):
         (b, 2): (f, (a, 1)),
     }
     o = order(dsk)
-    assert o[(a, 0)] < o[(a, 1)]
-
-    # Switch name of 0, 1 to ensure that this isn't due to string comparison
-    dsk = {
-        (a, 1): (f,),
-        (a, 0): (f,),
-        (b, 0): (f, (a, 1)),
-        (b, 1): (f, (a, 0)),
-        (b, 2): (f, (a, 0)),
-    }
-    o = order(dsk)
-    assert o[(a, 0)] > o[(a, 1)]
-
-    # Switch name of 0, 1 for "b"s too
-    dsk = {
-        (a, 0): (f,),
-        (a, 1): (f,),
-        (b, 1): (f, (a, 0)),
-        (b, 0): (f, (a, 1)),
-        (b, 2): (f, (a, 1)),
-    }
-    o = order(dsk)
-    assert o[(a, 0)] < o[(a, 1)]
+    assert o[(a, 1)] < o[(b, 0)] or (o[(b, 1)] < o[(a, 0)] and o[(b, 2)] < o[(a, 0)])
 
 
 def test_base_of_reduce_preferred(abcde):
@@ -201,7 +238,6 @@ def test_deep_bases_win_over_dependents(abcde):
     dsk = {a: (f, b, c, d), b: (f, d, e), c: (f, d), d: 1, e: 2}
 
     o = order(dsk)
-    assert o[e] < o[d]  # ambiguous, but this is what we currently expect
     assert o[b] < o[c]
 
 
@@ -237,7 +273,7 @@ def test_break_ties_by_str(abcde):
 
     o = order(dsk)
     expected = {"y": 10}
-    expected.update({k: i for i, k in enumerate(x_keys)})
+    expected.update({k: i for i, k in enumerate(x_keys[::-1])})
 
     assert o == expected
 
@@ -257,6 +293,7 @@ def test_gh_3055():
 
     dsk = dict(w.__dask_graph__())
     o = order(dsk)
+    assert max(diagnostics(dsk, o=o)[1]) <= 8
     L = [o[k] for k in w.__dask_keys__()]
     assert sum(x < len(o) / 2 for x in L) > len(L) / 3  # some complete quickly
 
@@ -402,21 +439,35 @@ def test_nearest_neighbor(abcde):
 
     assert 3 < sum(o[a + i] < len(o) / 2 for i in "123456789") < 7
     assert 1 < sum(o[b + i] < len(o) / 2 for i in "1234") < 4
-    assert o[min([b1, b2, b3, b4])] == 0
 
 
 def test_string_ordering():
     """Prefer ordering tasks by name first"""
     dsk = {("a", 1): (f,), ("a", 2): (f,), ("a", 3): (f,)}
     o = order(dsk)
-    assert o == {("a", 1): 0, ("a", 2): 1, ("a", 3): 2}
+    assert o == {("a", 1): 0, ("a", 2): 1, ("a", 3): 2} or o == {
+        ("a", 1): 2,
+        ("a", 2): 1,
+        ("a", 3): 0,
+    }
 
 
+@pytest.mark.xfail(reason="see comment", strict=False)
 def test_string_ordering_dependents():
     """Prefer ordering tasks by name first even when in dependencies"""
+    # XFAIL This is a little too artificial. While we can construct the
+    # algorithm in a way that respects the key ordering strictly, there is not
+    # necessarily a point to it and we can save ourselves one sorting step with
+    # this.
+    # See comment in add_to_result
     dsk = {("a", 1): (f, "b"), ("a", 2): (f, "b"), ("a", 3): (f, "b"), "b": (f,)}
     o = order(dsk)
-    assert o == {"b": 0, ("a", 1): 1, ("a", 2): 2, ("a", 3): 3}
+    assert o == {"b": 0, ("a", 1): 1, ("a", 2): 2, ("a", 3): 3} or o == {
+        "b": 0,
+        ("a", 1): 3,
+        ("a", 2): 2,
+        ("a", 3): 1,
+    }
 
 
 def test_prefer_short_narrow(abcde):
@@ -563,6 +614,8 @@ def test_use_structure_not_keys(abcde):
         (b, 0): (f, (a, 3), (a, 8), (a, 1)),
     }
     o = order(dsk)
+
+    assert max(diagnostics(dsk, o=o)[1]) == 3
     As = sorted(val for (letter, _), val in o.items() if letter == a)
     Bs = sorted(val for (letter, _), val in o.items() if letter == b)
     assert Bs[0] in {1, 3}
@@ -653,61 +706,6 @@ 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
 
@@ -719,7 +717,7 @@ def test_order_with_equal_dependents(abcde):
     # 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
+    max_pressure = 10
     a, b, c, d, e = abcde
     dsk = {}
     abc = [a, b, c, d]
@@ -747,13 +745,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, 1)] - o[(x, 6, i, 0)]
-            assert val > 0  # ideally, val == 2
+            val = abs(o[(x, 6, i, 1)] - o[(x, 6, i, 0)])
             total += val
-    assert total <= 56  # ideally, this should be 2 * 16 == 32
+
+    assert total <= 32  # 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)
     for x in abc:
@@ -763,10 +760,9 @@ def test_order_with_equal_dependents(abcde):
     total = 0
     for x in abc:
         for i in range(len(abc)):
-            val = o[(x, 6, i, 1)] - o[(x, 7, i, 0)]
-            assert val > 0  # ideally, val == 3
+            val = abs(o[(x, 6, i, 1)] - o[(x, 7, i, 0)])
             total += val
-    assert total <= 75  # ideally, this should be 3 * 16 == 48
+    assert total <= 48  # ideally, this should be 3 * 16 == 48
     pressure = diagnostics(dsk2, o=o)[1]
     assert max(pressure) <= max_pressure
 
@@ -779,14 +775,13 @@ def test_order_with_equal_dependents(abcde):
     total = 0
     for x in abc:
         for i in range(len(abc)):
-            val = o[(x, 5, i, 1)] - o[(x, 6, i, 0)]
-            assert val > 0
+            val = abs(o[(x, 5, i, 1)] - o[(x, 6, i, 0)])
             total += val
-    assert total <= 45  # ideally, this should be 2 * 16 == 32
+    assert total <= 32  # 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
+    # # Remove another one from each of the nine bundles
     dsk4 = dict(dsk3)
     for x in abc:
         for i in range(len(abc)):
@@ -796,7 +791,7 @@ def test_order_with_equal_dependents(abcde):
     assert max(pressure) <= max_pressure
     for x in abc:
         for i in range(len(abc)):
-            assert abs(o[(x, 5, i, 1)] - o[(x, 5, i, 0)]) <= 10
+            assert abs(o[(x, 5, i, 1)] - o[(x, 5, i, 0)]) <= 2
 
 
 def test_terminal_node_backtrack():
@@ -872,7 +867,6 @@ def test_array_store_final_order(tmpdir):
     dest = root.empty_like(name="dest", data=x, chunks=x.chunksize, overwrite=True)
     d = x.store(dest, lock=False, compute=False)
     o = order(d.dask)
-
     # Find the lowest store. Dask starts here.
     stores = [k for k in o if isinstance(k, tuple) and k[0].startswith("store-map-")]
     first_store = min(stores, key=lambda k: o[k])
@@ -976,16 +970,27 @@ def test_eager_to_compute_dependent_to_free_parent():
     }
     dependencies, dependents = get_deps(dsk)
     o = order(dsk)
-    parents = {deps.pop() for key, deps in dependents.items() if not dependencies[key]}
 
-    def cost(deps):
-        a, b = deps
-        return abs(o[a] - o[b])
-
-    cost_of_pairs = {key: cost(dependents[key]) for key in parents}
-    # Allow one to be bad, b/c this is hard!
-    costs = sorted(cost_of_pairs.values())
-    assert sum(costs[:-1]) <= 25 or sum(costs) <= 31
+    _, pressure = diagnostics(dsk, o=o)
+    assert max(pressure) <= 8
+    # Visualizing the graph shows that there are two deep, thick branches and
+    # two shallow ones. We prefer to run the deep ones first and reduce them as
+    # far as possible before starting another computation branch.
+    # We don't care about the ordering so deep_roots1 and deep_roots2 are
+    # ambiguous. At time of writing, deep_roots1 was executed first.
+    shallow_roots = ["a64", "a66"]
+    deep_roots1 = ["a68", "a63", "a65", "a67"]
+    # These are the final reducers of the first, thick branch
+    reducers_1 = ["a06", "a41"]
+    deep_roots2 = ["a69", "a70", "a71", "a62"]
+    reducers_2 = ["a39", "a00"]
+    for deep_roots in [deep_roots1, deep_roots2]:
+        assert max(o[r] for r in deep_roots) < min(o[r] for r in shallow_roots)
+    # These two are the reduction nodes furthest along before additional roots
+    # have to be loaded
+    for reducer in [reducers_1, reducers_2]:
+        for red in reducer:
+            assert o[red] < min(o[r] for r in shallow_roots)
 
 
 def test_diagnostics(abcde):
@@ -1010,71 +1015,115 @@ def test_diagnostics(abcde):
         (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()} == {
-        (a, 0): 0,
-        (b, 0): 1,
-        (c, 0): 2,
-        (d, 0): 4,
-        (e, 0): 6,
-        (a, 1): 3,
-        (b, 1): 5,
-        (c, 1): 7,
-        (d, 1): 8,
-        (e, 1): 9,
-    }
-    assert {key: val.age for key, val in info.items()} == {
-        (a, 0): 3,
-        (b, 0): 4,
-        (c, 0): 5,
-        (d, 0): 4,
-        (e, 0): 3,
-        (a, 1): 0,
-        (b, 1): 0,
-        (c, 1): 0,
-        (d, 1): 0,
-        (e, 1): 0,
-    }
-    assert {key: val.num_dependencies_freed for key, val in info.items()} == {
-        (a, 0): 0,
-        (b, 0): 0,
-        (c, 0): 0,
-        (d, 0): 0,
-        (e, 0): 0,
-        (a, 1): 1,
-        (b, 1): 1,
-        (c, 1): 1,
-        (d, 1): 1,
-        (e, 1): 1,
-    }
-    assert {key: val.num_data_when_run for key, val in info.items()} == {
-        (a, 0): 0,
-        (b, 0): 1,
-        (c, 0): 2,
-        (d, 0): 2,
-        (e, 0): 2,
-        (a, 1): 3,
-        (b, 1): 3,
-        (c, 1): 3,
-        (d, 1): 2,
-        (e, 1): 1,
-    }
-    assert {key: val.num_data_when_released for key, val in info.items()} == {
-        (a, 0): 3,
-        (b, 0): 3,
-        (c, 0): 3,
-        (d, 0): 2,
-        (e, 0): 1,
-        (a, 1): 3,
-        (b, 1): 3,
-        (c, 1): 3,
-        (d, 1): 2,
-        (e, 1): 1,
-    }
+    # this is ambiguous, depending on whether we start from left or right
+    assert all(o[key] == val.order for key, val in info.items())
+    if o[(e, 1)] == 1:
+        assert o[(e, 1)] == 1
+        assert o[(d, 1)] == 3
+        assert o[(c, 1)] == 5
+        assert memory_over_time == [0, 1, 1, 2, 2, 3, 2, 3, 2, 3]
+        assert {key: val.age for key, val in info.items()} == {
+            (a, 0): 1,
+            (b, 0): 3,
+            (c, 0): 5,
+            (d, 0): 5,
+            (e, 0): 5,
+            (a, 1): 0,
+            (b, 1): 0,
+            (c, 1): 0,
+            (d, 1): 0,
+            (e, 1): 0,
+        }
+        assert {key: val.num_dependencies_freed for key, val in info.items()} == {
+            (a, 0): 0,
+            (b, 0): 0,
+            (c, 0): 0,
+            (d, 0): 0,
+            (e, 0): 0,
+            (a, 1): 3,
+            (b, 1): 1,
+            (c, 1): 1,
+            (d, 1): 0,
+            (e, 1): 0,
+        }
+        assert {key: val.num_data_when_run for key, val in info.items()} == {
+            (a, 0): 2,
+            (b, 0): 2,
+            (c, 0): 2,
+            (d, 0): 1,
+            (e, 0): 0,
+            (a, 1): 3,
+            (b, 1): 3,
+            (c, 1): 3,
+            (d, 1): 2,
+            (e, 1): 1,
+        }
+        assert {key: val.num_data_when_released for key, val in info.items()} == {
+            (a, 0): 3,
+            (b, 0): 3,
+            (c, 0): 3,
+            (d, 0): 3,
+            (e, 0): 3,
+            (a, 1): 3,
+            (b, 1): 3,
+            (c, 1): 3,
+            (d, 1): 2,
+            (e, 1): 1,
+        }
+    else:
+        assert o[(e, 1)] == len(dsk) - 1
+        assert o[(d, 1)] == len(dsk) - 2
+        assert o[(c, 1)] == len(dsk) - 3
+        assert memory_over_time == [0, 1, 2, 3, 2, 3, 2, 3, 2, 1]
+        assert {key: val.age for key, val in info.items()} == {
+            (a, 0): 3,
+            (b, 0): 4,
+            (c, 0): 5,
+            (d, 0): 4,
+            (e, 0): 3,
+            (a, 1): 0,
+            (b, 1): 0,
+            (c, 1): 0,
+            (d, 1): 0,
+            (e, 1): 0,
+        }
+        assert {key: val.num_dependencies_freed for key, val in info.items()} == {
+            (a, 0): 0,
+            (b, 0): 0,
+            (c, 0): 0,
+            (d, 0): 0,
+            (e, 0): 0,
+            (a, 1): 1,
+            (b, 1): 1,
+            (c, 1): 1,
+            (d, 1): 1,
+            (e, 1): 1,
+        }
+        assert {key: val.num_data_when_run for key, val in info.items()} == {
+            (a, 0): 0,
+            (b, 0): 1,
+            (c, 0): 2,
+            (d, 0): 2,
+            (e, 0): 2,
+            (a, 1): 3,
+            (b, 1): 3,
+            (c, 1): 3,
+            (d, 1): 2,
+            (e, 1): 1,
+        }
+        assert {key: val.num_data_when_released for key, val in info.items()} == {
+            (a, 0): 3,
+            (b, 0): 3,
+            (c, 0): 3,
+            (d, 0): 2,
+            (e, 0): 1,
+            (a, 1): 3,
+            (b, 1): 3,
+            (c, 1): 3,
+            (d, 1): 2,
+            (e, 1): 1,
+        }
 
 
 def test_xarray_like_reduction():
@@ -1107,7 +1156,10 @@ def test_xarray_like_reduction():
 
 @pytest.mark.parametrize(
     "optimize",
-    [True, False],
+    [
+        True,
+        False,
+    ],
 )
 def test_array_vs_dataframe(optimize):
     xr = pytest.importorskip("xarray")
@@ -1144,88 +1196,120 @@ def test_anom_mean():
     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),
-    )
+    data = da.random.random((200, 1), chunks=(1, -1))
 
-    ngroups = data.shape[0] // 4
+    ngroups = 5
     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():
+    graph = collections_to_dsk([anom_mean])
+    dependencies, dependents = get_deps(graph)
+    diags, pressure = diagnostics(graph)
+    # Encoding the "best" ordering for this graph is tricky. When inspecting the
+    # visualization, one sees that there are small, connected tree-like steps at
+    # the beginning (tested well below in anom_mean_raw) followed by a
+    # concat+transpose per group (see ngroups above). This transpose task fans
+    # out into many (20-30) getitem tasks that are tiny and feed into a
+    # `mean_chunk` which is the primary reducer in this graph. Therefore we want
+    # to run those as quickly as possible.
+    # This is difficult to assert on but the pressure is an ok-ish proxy
+    # visualize(graph)
+    assert max(pressure) <= 177
+    from collections import defaultdict
+
+    count_dependents = defaultdict(set)
+    for k in dict(graph).keys():
+        count_dependents[len(dependents[k])].add(k)
+    n_splits = max(count_dependents)
+    # There is a transpose/stack group that is splitting into many tasks
+    # see https://github.com/dask/dask/pull/10660#discussion_r1420571664
+    # the name depends on the version of xarray
+    assert n_splits > 30  # at time of writing 40
+    transpose_tasks = count_dependents[n_splits]
+    transpose_metrics = {k: diags[k] for k in transpose_tasks}
+    assert len(transpose_metrics) == ngroups, {key_split(k) for k in diags}
+    # This is a pretty tightly connected graph overall and we'll have to hold
+    # many tasks in memory until this can complete. However, we should ensure
+    # that we get to the mean_chunks asap while the transposes are released
+    # quickly.
+    # If this breaks, I suggest to visually inspect the graph and run the above
+    # on a single threaded LocalCluster and verify that the progress is indeed
+    # in five steps (i.e. five groups)
+    ages_mean_chunks = {k: v.age for k, v in diags.items() if "mean_chunk" in k[0]}
+    avg_age_mean_chunks = sum(ages_mean_chunks.values()) / len(ages_mean_chunks)
+    max_age_mean_chunks = max(ages_mean_chunks.values())
+    ages_tranpose = {k: v.age for k, v in transpose_metrics.items()}
+    assert max_age_mean_chunks > 900
+    assert avg_age_mean_chunks > 100
+    avg_age_transpose = sum(ages_tranpose.values()) / len(ages_tranpose)
+    max_age_transpose = max(ages_tranpose.values())
+    assert max_age_transpose < 150
+    assert avg_age_transpose < 100
+    assert sum(pressure) / len(pressure) < 100
+
+
+def test_anom_mean_raw(abcde):
+    a, b, c, d, e = abcde
+    g, h = "gh"
     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)]),
+        (d, 0, 0): (f, (a, 0, 0), (b, 1, 0, 0)),
+        (d, 1, 0): (f, (a, 1, 0), (b, 1, 1, 0)),
+        (d, 2, 0): (f, (a, 2, 0), (b, 1, 2, 0)),
+        (d, 3, 0): (f, (a, 3, 0), (b, 1, 3, 0)),
+        (d, 4, 0): (f, (a, 4, 0), (b, 1, 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, (g, 1, 0)),
+        (e, 1, 0): (f, (g, 3, 0)),
+        (b, 0, 0, 0): (f, (a, 0, 0)),
+        (b, 0, 1, 0): (f, (a, 2, 0)),
+        (b, 0, 2, 0): (f, (a, 4, 0)),
+        (c, 0, 0, 0): (f, (b, 0, 0, 0)),
+        (c, 0, 1, 0): (f, (b, 0, 1, 0)),
+        (c, 0, 2, 0): (f, (b, 0, 2, 0)),
+        (g, 1, 0): (f, [(c, 0, 0, 0), (c, 0, 1, 0), (c, 0, 2, 0)]),
+        (b, 2, 0, 0): (f, (a, 1, 0)),
+        (b, 2, 1, 0): (f, (a, 3, 0)),
+        (c, 1, 0, 0): (f, (b, 2, 0, 0)),
+        (c, 1, 1, 0): (f, (b, 2, 1, 0)),
+        (g, 3, 0): (f, [(c, 1, 0, 0), (c, 1, 1, 0)]),
+        (b, 1, 0, 0): (f, (e, 0, 0)),
+        (b, 1, 1, 0): (f, (e, 1, 0)),
+        (b, 1, 2, 0): (f, (e, 0, 0)),
+        (b, 1, 3, 0): (f, (e, 1, 0)),
+        (b, 1, 4, 0): (f, (e, 0, 0)),
+        (c, 2, 0, 0): (f, (d, 0, 0)),
+        (c, 2, 1, 0): (f, (d, 1, 0)),
+        (c, 2, 2, 0): (f, (d, 2, 0)),
+        (c, 2, 3, 0): (f, (d, 3, 0)),
+        (c, 2, 4, 0): (f, (d, 4, 0)),
+        (h, 0, 0): (f, [(c, 2, 0, 0), (c, 2, 1, 0), (c, 2, 2, 0), (c, 2, 3, 0)]),
+        (h, 1, 0): (f, [(c, 2, 4, 0)]),
+        (g, 2, 0): (f, [(h, 0, 0), (h, 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),
+        (h, 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)]
+        assert o[n] < o[(a, 1, 0)]
+        assert o[n] < o[(a, 3, 0)]
 
 
 def test_flaky_array_reduction():
@@ -1635,36 +1719,402 @@ def test_flaky_array_reduction():
     assert first_pressure == second_pressure
 
 
-def test_flox_reduction():
-    # TODO: It would be nice to scramble keys to ensure we're not comparing keys
+def test_flox_reduction(abcde):
+    a, b, c, d, e = abcde
+    g = "g"
+    dsk = {
+        (a, 0): (f, 1),
+        (a, 1, 0): (f, 1),
+        (a, 1, 1): (f, 1),
+        (a, 2, 0): (f, 1),
+        (a, 2, 1): (f, 1),
+        (b, 1, 0): (f, [(f, (a, 2, 0))]),
+        (b, 1, 1): (f, [(f, (a, 2, 1))]),
+        (b, 2, 1): (f, [(f, (a, 1, 1))]),
+        (b, 2, 0): (f, [(f, (a, 1, 0))]),
+        (b, 1, 1, 0): (b, 1, 0),
+        (b, 1, 1, 1): (b, 1, 1),
+        (b, 2, 2, 0): (b, 2, 0),
+        (b, 2, 2, 1): (b, 2, 1),
+        (c, 1, 0): (f, (b, 2, 2, 0)),
+        (c, 1, 1): (f, (b, 2, 2, 1)),
+        (c, 2, 0): (f, (b, 1, 1, 0)),
+        (c, 2, 1): (f, (b, 1, 1, 1)),
+        (d, 1): (f, [(f, (a, 1, 1), (a, 2, 1), (c, 1, 1), (c, 2, 1))]),
+        (d, 0): (f, [(f, (a, 1, 0), (a, 2, 0), (c, 1, 0), (c, 2, 0))]),
+        (e, 0): (d, 0),
+        (e, 1): (d, 1),
+        (g, 1, 0): (f, (a, 0), (b, 1, 1, 0)),
+        (g, 1, 1): (f, (a, 0), (b, 2, 2, 0)),
+        (g, 1, 2): (f, (a, 0), (e, 0)),
+        (g, 2, 0): (f, (a, 0), (b, 1, 1, 1)),
+        (g, 2, 1): (f, (a, 0), (b, 2, 2, 1)),
+        (g, 2, 2): (f, (a, 0), (e, 1)),
+    }
+    o = order(dsk)
+    of1 = list(o[(g, 1, ix)] for ix in range(3))
+    of2 = list(o[(g, 2, ix)] for ix in range(3))
+    assert max(of1) < min(of2) or max(of2) < min(of1)
+
+
+@pytest.mark.parametrize("ndeps", [2, 5])
+@pytest.mark.parametrize("n_reducers", [4, 7])
+def test_reduce_with_many_common_dependents(ndeps, n_reducers):
+    da = pytest.importorskip("dask.array")
+    import numpy as np
+
+    def random(**kwargs):
+        assert len(kwargs) == ndeps
+        return np.random.random((10, 10))
+
+    trivial_deps = {
+        f"k{i}": delayed(object(), name=f"object-{i}") for i in range(ndeps)
+    }
+    x = da.blockwise(
+        random,
+        "yx",
+        new_axes={"y": (10,) * n_reducers, "x": (10,) * n_reducers},
+        dtype=float,
+        **trivial_deps,
+    )
+    graph = x.sum(axis=1, split_every=20)
+    from dask.order import order
+
+    dsk = collections_to_dsk([graph])
+    dependencies, dependents = get_deps(dsk)
+    # Verify assumptions
+    o = order(dsk)
+    # Verify assumptions (specifically that the reducers are sum-aggregate)
+    assert {key_split(k) for k in o} == {"object", "sum", "sum-aggregate"}
+
+    reducers = {k for k in o if key_split(k) == "sum-aggregate"}
+    drift = dict()
+    for r in reducers:
+        prios_deps = []
+        for dep in dependencies[r]:
+            prios_deps.append(o[dep])
+        drift[r] = (min(prios_deps), max(prios_deps))
+        assert max(prios_deps) - min(prios_deps) == len(dependencies[r]) - 1
+
+
+def test_doublediff(abcde):
+    a, b, c, d, e = abcde
+    dsk = {
+        (a, 3, 0, 1): (f, 1),
+        (a, 3, 1, 1): (f, 1),
+        (a, 3, 1, 0): (f, 1),
+        (a, 2, 2, 2): (f, 1),
+        (a, 4, 0, 1): (f, 1),
+        (a, 4, 1, 0): (f, 1),
+        (a, 4, 0, 0): (f, 1),
+        (a, 5, 0, 0): (f, 1),
+        (c, 0, 0, 0): (b, 1, 0, 0),
+        (c, 0, 0, 1): (b, 1, 0, 1),
+        (c, 0, 0, 2): (b, 1, 0, 2),
+        (c, 0, 1, 0): (b, 1, 1, 0),
+        (c, 0, 1, 1): (b, 1, 1, 1),
+        (c, 0, 1, 2): (b, 1, 1, 2),
+        (c, 0, 2, 0): (b, 1, 2, 0),
+        (c, 0, 2, 1): (b, 1, 2, 1),
+        (c, 1, 0, 0): (a, 5, 0, 0),
+        (c, 1, 0, 1): (b, 0, 0, 1),
+        (c, 1, 0, 2): (b, 0, 0, 2),
+        (c, 1, 1, 0): (b, 0, 1, 0),
+        (c, 1, 1, 1): (b, 0, 1, 1),
+        (c, 1, 1, 2): (b, 0, 1, 2),
+        (c, 1, 2, 0): (b, 0, 2, 0),
+        (b, 1, 0, 0): (f, (a, 4, 0, 0)),
+        (b, 1, 0, 1): (f, (a, 4, 0, 0)),
+        (b, 1, 0, 2): (f, (a, 4, 0, 1)),
+        (b, 1, 1, 0): (f, (a, 4, 0, 0)),
+        (b, 1, 1, 1): (f, (a, 4, 0, 0)),
+        (b, 1, 1, 2): (f, (a, 4, 0, 1)),
+        (b, 1, 2, 0): (f, (a, 4, 1, 0)),
+        (b, 1, 2, 1): (f, (a, 4, 1, 0)),
+        (c, 1, 2, 1): (b, 0, 2, 1),
+        (c, 1, 2, 2): (b, 0, 2, 2),
+        (d, 0, 0, 0): (f, (c, 1, 0, 0), (c, 0, 0, 0)),
+        (d, 0, 0, 1): (f, (c, 1, 0, 1), (c, 0, 0, 1)),
+        (d, 0, 0, 2): (f, (c, 1, 0, 2), (c, 0, 0, 2)),
+        (d, 0, 1, 0): (f, (c, 1, 1, 0), (c, 0, 1, 0)),
+        (d, 0, 1, 1): (f, (c, 1, 1, 1), (c, 0, 1, 1)),
+        (d, 0, 1, 2): (f, (c, 1, 1, 2), (c, 0, 1, 2)),
+        (d, 0, 2, 0): (f, (c, 1, 2, 0), (c, 0, 2, 0)),
+        (d, 0, 2, 1): (f, (c, 1, 2, 1), (c, 0, 2, 1)),
+        (d, 0, 2, 2): (f, (c, 1, 2, 2), (a, 2, 2, 2)),
+        (b, 0, 0, 1): (f, (a, 3, 0, 1)),
+        (b, 0, 0, 2): (f, (a, 3, 0, 1)),
+        (b, 0, 1, 0): (f, (a, 3, 1, 0)),
+        (b, 0, 1, 1): (f, (a, 3, 1, 1)),
+        (b, 0, 1, 2): (f, (a, 3, 1, 1)),
+        (b, 0, 2, 0): (f, (a, 3, 1, 0)),
+        (b, 0, 2, 1): (f, (a, 3, 1, 1)),
+        (b, 0, 2, 2): (f, (a, 3, 1, 1)),
+        (e, 0, 0, 0): (f, (d, 0, 0, 0)),
+        (e, 0, 0, 1): (f, (d, 0, 0, 1)),
+        (e, 0, 0, 2): (f, (d, 0, 0, 2)),
+        (e, 0, 1, 0): (f, (d, 0, 1, 0)),
+        (e, 0, 1, 1): (f, (d, 0, 1, 1)),
+        (e, 0, 1, 2): (f, (d, 0, 1, 2)),
+        (e, 0, 2, 0): (f, (d, 0, 2, 0)),
+        (e, 0, 2, 1): (f, (d, 0, 2, 1)),
+        (e, 0, 2, 2): (f, (d, 0, 2, 2)),
+        "END": [
+            (e, 0, 0, 0),
+            (e, 0, 0, 1),
+            (e, 0, 0, 2),
+            (e, 0, 1, 0),
+            (e, 0, 1, 1),
+            (e, 0, 1, 2),
+            (e, 0, 2, 0),
+            (e, 0, 2, 1),
+            (e, 0, 2, 2),
+        ],
+    }
+    _, pressure = diagnostics(dsk)
+    assert max(pressure) <= 11, max(pressure)
+
+
+def test_recursion_depth_long_linear_chains():
+    dsk = {"-1": (f, 1)}
+    for ix in range(10000):
+        dsk[str(ix)] = (f, str(ix - 1))
+    assert len(order(dsk)) == len(dsk)
+
+
+def test_gh_3055_explicit(abcde):
+    # This is a subgraph extracted from gh_3055
+    # From a critical path perspective, the root a, 2 only has to be loaded
+    # towards the very end. However, loading it so late means that we are
+    # blocking many possible reductions which is bad for memory pressure
+    a, b, c, d, e = abcde
+    g = "g"
+    dsk = {
+        ("root", 0): (f, 1),
+        (a, 0): (f, ("root", 0)),
+        (a, 1): (f, 1),
+        (a, 2): (f, 1),
+        (a, 3): (f, 1),
+        (a, 4): (f, 1),
+        (b, 0, 0): (f, (a, 0)),
+        (b, 0, 1): (f, (a, 0)),
+        (c, 0, 1): (f, (a, 0)),
+        (b, 1, 0): (f, (a, 1)),
+        (b, 1, 2): (f, (a, 1)),
+        (c, 0, 0): (f, (b, 0), (a, 2), (a, 1)),
+        (d, 0, 0): (f, (c, 0, 1), (c, 0, 0)),
+        (d, 0, 1): (f, (c, 0, 1), (c, 0, 0)),
+        (f, 1, 1): (f, (d, 0, 1)),
+        (c, 1, 0): (f, (b, 1, 0), (b, 1, 2)),
+        (c, 0, 2): (f, (b, 0, 0), (b, 0, 1)),
+        (e, 0): (f, (c, 1, 0), (c, 0, 2)),
+        (g, 1): (f, (e, 0), (a, 3)),
+        (g, 2): (f, (g, 1), (a, 4), (d, 0, 0)),
+    }
+    dependencies, dependents = get_deps(dsk)
+    con_r, _ = _connecting_to_roots(dependencies, dependents)
+    assert len(con_r) == len(dsk)
+    assert con_r[(e, 0)] == {("root", 0), (a, 1)}
+    o = order(dsk)
+    assert max(diagnostics(dsk, o=o)[1]) <= 5
+    assert o[(e, 0)] < o[(a, 3)] < o[(a, 4)]
+    assert o[(a, 2)] < o[(a, 3)] < o[(a, 4)]
+
+
+def test_order_flux_reduction_2(abcde):
+    # https://github.com/dask/dask/issues/10618
+    a, b, c, d, e = abcde
     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)),
+        (a, 0): 0,
+        (a, 1): 0,
+        (a, 2): 0,
+        (b, 0, 0, 0): (f, (a, 0)),
+        (b, 0, 0, 1): (f, (a, 1)),
+        (b, 0, 0, 2): (f, (a, 2)),
+        (b, 0, 1, 0): (f, (a, 0)),
+        (b, 0, 1, 1): (f, (a, 1)),
+        (b, 0, 1, 2): (f, (a, 2)),
+        (b, 1, 0, 0): (f, (a, 0)),
+        (b, 1, 0, 1): (f, (a, 1)),
+        (b, 1, 0, 2): (f, (a, 2)),
+        (b, 1, 1, 0): (f, (a, 0)),
+        (b, 1, 1, 1): (f, (a, 1)),
+        (b, 1, 1, 2): (f, (a, 2)),
+        (c, 0, 0): (f, [(b, 0, 0, 0), (b, 0, 0, 1), (b, 0, 0, 2)]),
+        (c, 0, 1): (f, [(b, 0, 1, 0), (b, 0, 1, 1), (b, 0, 1, 2)]),
+        (c, 1, 0): (f, [(b, 1, 0, 0), (b, 1, 0, 1), (b, 1, 0, 2)]),
+        (c, 1, 1): (f, [(b, 1, 1, 0), (b, 1, 1, 1), (b, 1, 1, 2)]),
+        (d, 0, 0): (c, 0, 0),
+        (d, 0, 1): (c, 0, 1),
+        (d, 1, 0): (c, 1, 0),
+        (d, 1, 1): (c, 1, 1),
     }
+
+    o = order(dsk)
+    final_nodes = sorted(
+        [(d, ix, jx) for ix in range(2) for jx in range(2)], key=o.__getitem__
+    )
+    for ix in range(1, len(final_nodes)):
+        # This assumes that all the data tasks are scheduled first.
+        # Then, there are exactly four dependencies to load for every final
+        # task.
+        assert o[final_nodes[ix]] - o[final_nodes[ix - 1]] == 5
+
+
+def test_xarray_8414():
+    # https://github.com/pydata/xarray/issues/8414#issuecomment-1793860552
+    np = pytest.importorskip("numpy")
+    xr = pytest.importorskip("xarray")
+
+    ds = xr.Dataset(
+        data_vars=dict(
+            a=(("x", "y"), np.arange(80).reshape(8, 10)),
+            b=(("y"), np.arange(10)),
+        )
+    ).chunk(x=1)
+
+    def f(ds):
+        d = ds.a * ds.b
+        return (d * ds.b).sum("y")
+
+    graph = ds.map_blocks(f).__dask_graph__()
+
+    # Two tasks always reduce to one leaf
+    # plus a commonly shared root node makes three
+    assert (p := max(diagnostics(graph)[1])) == 3, p
+
+    # If the pressure is actually at three, there is not much that could go
+    # wrong. Still, this asserts that the dependents of the shared dependency
+    # are only loaded only after the reducers have been computed which given a
+    # symmetrical graph, means that the steps between the dependents are the
+    # same. If the xarray implementation changes or fusing enters, the step size
+    # could grow or shrink but should still be the same everywhere
+    dsk = dict(graph)
+    dependencies, dependents = get_deps(dsk)
+    roots = {k for k, v in dependencies.items() if not v}
+    shared_roots = {r for r in roots if len(dependents[r]) > 1}
+    assert len(shared_roots) == 1
+    shared_root = shared_roots.pop()
     o = order(dsk)
-    assert max(o[("F1", ix)] for ix in range(3)) < min(o[("F2", ix)] for ix in range(3))
+
+    previous = None
+    step = 0
+    assert len(dependents[shared_root]) > 2
+    for dep in sorted(dependents[shared_root], key=o.__getitem__):
+        if previous is not None:
+            if step == 0:
+                step = o[dep] - o[shared_root] - 1
+                assert step > 1
+            else:
+                assert o[dep] == previous + step
+        previous = o[dep]
+
+
+def test_connecting_to_roots_single_root():
+    dsk = {
+        "a": (f, 1),
+        "b1": (f, "a"),
+        "b2": (f, "a"),
+        "c": (f, "b1", "b2"),
+    }
+    dependencies, dependents = get_deps(dsk)
+    connected_roots, max_dependents = _connecting_to_roots(dependencies, dependents)
+    assert connected_roots == {k: {"a"} if k != "a" else set() for k in dsk}
+    assert len({id(v) for v in connected_roots.values()}) == 2
+    assert max_dependents == {
+        "a": 2,
+        "b1": 2,
+        "b2": 2,
+        "c": 2,
+    }, max_dependents
+    connected_roots, max_dependents = _connecting_to_roots(dependents, dependencies)
+    assert connected_roots == {k: {"c"} if k != "c" else set() for k in dsk}
+    assert len({id(v) for v in connected_roots.values()}) == 2
+    assert max_dependents == {
+        "a": 2,
+        "b1": 2,
+        "b2": 2,
+        "c": 2,
+    }, max_dependents
+
+
+def test_connecting_to_roots_tree_reduction():
+    dsk = {
+        "a0": (f, 1),
+        "a1": (f, 1),
+        "a2": (f, 1),
+        "a3": (f, 1),
+        "b1": (f, "a0"),
+        "b2": (f, "a1"),
+        "b3": (f, "a2"),
+        "b4": (f, "a3"),
+        "c1": (f, "b1", "b2"),
+        "c2": (f, "b3", "b4"),
+        "d": (f, "c1", "c2"),
+    }
+    dependencies, dependents = get_deps(dsk)
+    connected_roots, max_dependents = _connecting_to_roots(dependencies, dependents)
+    assert connected_roots == {
+        "a0": set(),
+        "a1": set(),
+        "a2": set(),
+        "a3": set(),
+        "b1": {"a0"},
+        "b2": {"a1"},
+        "b3": {"a2"},
+        "b4": {"a3"},
+        "c1": {"a1", "a0"},
+        "c2": {"a3", "a2"},
+        "d": {"a1", "a2", "a3", "a0"},
+    }
+    assert all(v == 1 for v in max_dependents.values())
+
+    connected_roots, max_dependents = _connecting_to_roots(dependents, dependencies)
+    assert connected_roots.pop("d") == set()
+    assert all(v == {"d"} for v in connected_roots.values()), set(
+        connected_roots.values()
+    )
+    assert all(v == 2 for v in max_dependents.values())
+
+
+def test_connecting_to_roots_asym():
+    dsk = {
+        "a0": (f, 1),
+        "a1": (f, 1),
+        "a2": (f, 1),
+        "a3": (f, 1),
+        "a4": (f, 1),
+        # Diamond
+        "b1": (f, "a0", "a1"),
+        "c1": (f, "b1"),
+        "c2": (f, "b1"),
+        "d1": (f, "c1", "c2"),
+        # Multi stage reducers
+        "b2": (f, "a2"),
+        "b3": (f, "a3"),
+        "c3": (f, "b3", "b2"),
+        "d2": (f, "b3", "c3", "a4"),
+    }
+    dependencies, dependents = get_deps(dsk)
+    connected_roots, max_dependents = _connecting_to_roots(dependencies, dependents)
+    assert connected_roots == {
+        "a0": set(),
+        "a1": set(),
+        "a2": set(),
+        "a3": set(),
+        "a4": set(),
+        "b1": {"a0", "a1"},
+        "c1": {"a0", "a1"},
+        "c2": {"a0", "a1"},
+        "d1": {"a0", "a1"},
+        "b2": {"a2"},
+        "b3": {"a3"},
+        "c3": {"a2", "a3"},
+        "d2": {"a2", "a3", "a4"},
+    }
+    # Max dependents is just pre-computed for performance but it is itself just
+    # a derived property of connected_roots
+    assert max_dependents == {
+        k: max((len(dependents[r]) for r in v), default=len(dependents[k]))
+        for k, v in connected_roots.items()
+    }