Repository: dask/dask
Pull request #11935 (2025-05-07 15:35:39)

Optimize dask order for worst case of get_target
There is an edge case where we have to reevaluate get_target over and over again with a long list of possible candidates that require sorting. That performs awfully.

This introduces an optimization that eliminates this sorting in most cases and for the cases where it is still required it should shrink the candidates before sorting significantly (I'll follow up with benchmarks shortly)

fjetter: This graph triggers a worst case (maybe not actually _worst_ case but pretty bad) complexity for `get_target`. Since the graph is heavily skewed it is falling into the custom path of `get_target` which is bad. 

![image](https://github.com/user-attachments/assets/668211ff-899a-4b79-bba2-8de571d1f166)


```python
    def get_target_worst_case(depth):
        def func(): ...

        depth -= 1
        root_nodes = 2 ** depth
        all_tasks = []
        for level in range(depth + 1):
            if level == 0:
                level_tasks = [
                    Task(("foo", level, ix), func) for ix in range(root_nodes)
                ]
            else:
                level_tasks = [
                    Task(
                        ("foo", level, ix),
                        func,
                        level_tasks[ix].ref(),
                        level_tasks[ix + 1].ref() if ix + 1 < len(level_tasks) else None,
                    )
                    for ix in range(len(level_tasks) // 2)
                ]
            all_tasks.extend(level_tasks)
        return all_tasks
```

At a depth of 15

```
# depth: duration
# depth -> number of nodes grows exponentially with a factor of 2, i.e. ideally runtime also increases by 2x

# main

14: 3.372095917351544
15: 11.977342499885708  # ~4x

# this PR
14: 0.17605512496083975
15: 0.3922020411118865  # ~2x
```

This is now so fast because the "nontrivial" path of the get_target function actually falls back to a single possible target after the set intersections are applied. Previously we had to walk thousands of nodes every iteration.

Patch:
diff --git a/dask/local.py b/dask/local.py
index 66b72f5ea..0c276aecb 100644
--- a/dask/local.py
+++ b/dask/local.py
@@ -187,6 +187,12 @@ def start_state_from_dask(dsk, cache=None, sortkey=None, keys=None):
         dependencies[key]
         task = dsk.get(key, None)
         if task is None:
+            if dependents[key] and not cache.get(key, None):
+                raise ValueError(
+                    "Missing dependency {} for dependents {}".format(
+                        key, dependents[key]
+                    )
+                )
             continue
         elif isinstance(task, DataNode):
             cache[key] = task()
diff --git a/dask/order.py b/dask/order.py
index b6db4578b..e33489029 100644
--- a/dask/order.py
+++ b/dask/order.py
@@ -195,7 +195,7 @@ def order(
     runnable_hull = set()
     reachable_hull = set()
 
-    runnable = []
+    runnable: list[Key] = []
 
     known_runnable_paths: dict[Key, list[list[Key]]] = {}
     known_runnable_paths_pop = known_runnable_paths.pop
@@ -203,12 +203,8 @@ def order(
     crit_path_counter = 0
     scrit_path: set[Key] = set()
     _crit_path_counter_offset: int | float = 0
-
     _sort_keys_cache: dict[Key, tuple[int, int, int, int, str]] = {}
 
-    leafs_connected_to_loaded_roots: set[Key] = set()
-    processed_roots = set()
-
     def sort_key(x: Key) -> tuple[int, int, int, int, str]:
         try:
             return _sort_keys_cache[x]
@@ -231,6 +227,7 @@ def order(
         # recursion depth errors. This is the only reason for the while loop
         next_items = [item]
         nonlocal i
+        nonlocal min_leaf_degree  # type: ignore[misc]
         while next_items:
             item = next_items.pop()
             runnable_hull.discard(item)
@@ -248,7 +245,25 @@ def order(
             if item not in external_keys:
                 i += 1
             if item in root_nodes:
-                processed_roots.add(item)
+                for leaf in leafs_connected[item]:
+                    if leaf in leaf_nodes:
+                        degree = leafs_degree.pop(leaf)
+                        leafs_per_degree[degree].remove(leaf)
+
+                        new_degree = degree - 1
+                        if new_degree > 0:
+                            if new_degree < min_leaf_degree:
+                                min_leaf_degree = new_degree
+                            leafs_per_degree[new_degree].add(leaf)
+                            leafs_degree[leaf] = new_degree
+                        elif not leafs_per_degree[degree]:
+                            assert degree == min_leaf_degree
+                            while min_leaf_degree != max_leaf_degree and (
+                                min_leaf_degree not in leafs_per_degree
+                                or not leafs_per_degree[min_leaf_degree]
+                            ):
+                                min_leaf_degree += 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
@@ -396,42 +411,38 @@ def order(
     # writing, the most expensive part of ordering is the prep work (mostly
     # connected roots + sort_key) which can be reused for multiple orderings.
 
-    def get_target(longest_path: bool = False) -> Key:
-        # Some topologies benefit if the node with the most dependencies
-        # is used as first choice, others benefit from the opposite.
-        candidates = leaf_nodes
-        skey: Callable = sort_key
-        if reachable_hull:
-            skey = lambda k: (num_needed[k], sort_key(k))
-
-        all_leafs_accessible = len(leafs_connected_to_loaded_roots) >= len(candidates)
-        if reachable_hull and not all_leafs_accessible:
-            # Avoid this branch if we can already access all leafs. Just pick
-            # one of them without the expensive selection process in this case.
-
-            candidates = reachable_hull & candidates
+    # Degree in this context is the number root nodes that have to be loaded for
+    # this leaf to become accessible. Zero means the leaf is already accessible
+    # in which case it _should_ either already be in result or be accessible vie
+    # process_runnables
+    # When picking a new target, we prefer the leafs with the least number of
+    # roots that need loading.
+    leafs_degree = {}
+    leafs_per_degree = defaultdict(set)
+    min_leaf_degree = len(dsk)
+    max_leaf_degree = len(dsk)
+    for leaf in leaf_nodes - root_nodes:
+        degree = len(roots_connected[leaf])
+        min_leaf_degree = min(min_leaf_degree, degree)
+        max_leaf_degree = max(max_leaf_degree, degree)
+        leafs_degree[leaf] = degree
+        leafs_per_degree[degree].add(leaf)
+
+    def get_target() -> Key:
+        # If we're already mid run and there is a runnable_hull we'll attempt to
+        # pick the next target in a way that minimizes the number of additional
+        # root nodes that are needed
+        all_leafs_accessible = min_leaf_degree == max_leaf_degree
+        is_trivial_lookup = not reachable_hull or all_leafs_accessible
+        if not is_trivial_lookup:
+            candidates = reachable_hull & leafs_per_degree[min_leaf_degree]
             if not candidates:
-                # We can't reach a leaf node directly, but we still have nodes
-                # with results in memory, these notes can inform our path towards
-                # a new preferred leaf node.
-                for r in processed_roots:
-                    leafs_connected_to_loaded_roots.update(leafs_connected[r])
-                processed_roots.clear()
-
-                leafs_connected_to_loaded_roots.intersection_update(leaf_nodes)
-                candidates = leafs_connected_to_loaded_roots
-            else:
-                leafs_connected_to_loaded_roots.update(candidates)
-
-        elif not reachable_hull:
-            # We reach a new and independent branch, so throw away previous branch
-            leafs_connected_to_loaded_roots.clear()
-
-        if longest_path and (not reachable_hull or all_leafs_accessible):
-            return leaf_nodes_sorted.pop()
+                candidates = leafs_per_degree[min_leaf_degree]
+            # Even without reachable hull overlap this should be relatively
+            # small so one full pass should be fine
+            return min(candidates, key=sort_key)
         else:
-            # FIXME: This can be very expensive
-            return min(candidates, key=skey)
+            return leaf_nodes_sorted.pop()
 
     def use_longest_path() -> bool:
         size = 0
@@ -450,10 +461,10 @@ def order(
 
         return True
 
+    # Some topologies benefit if the node with the most dependencies
+    # is used as first choice, others benefit from the opposite.
     longest_path = use_longest_path()
-    leaf_nodes_sorted = []
-    if longest_path:
-        leaf_nodes_sorted = sorted(leaf_nodes, key=sort_key, reverse=False)
+    leaf_nodes_sorted = sorted(leaf_nodes, key=sort_key, reverse=not longest_path)
 
     # *************************************************************************
     # CORE ALGORITHM STARTS HERE
@@ -553,7 +564,7 @@ def order(
         assert not scrit_path
 
         # A. Build the critical path
-        target = get_target(longest_path=longest_path)
+        target = get_target()
         next_deps = dependencies[target]
         path_append(target)
 


Test patch:
diff --git a/dask/tests/test_local.py b/dask/tests/test_local.py
index 684e43949..016443820 100644
--- a/dask/tests/test_local.py
+++ b/dask/tests/test_local.py
@@ -3,7 +3,7 @@ from __future__ import annotations
 import pytest
 
 import dask
-from dask._task_spec import Task
+from dask._task_spec import Task, TaskRef
 from dask.local import finish_task, get_sync, sortkey, start_state_from_dask
 from dask.order import order
 from dask.utils_test import GetFunctionTestMixin, add, inc
@@ -213,3 +213,10 @@ def test_ensure_calculate_only_whats_needed():
     ]
     dsk = {t.key: t for t in tasks}
     assert get_sync(dsk, "y") == 2
+
+
+def test_detect_malformed_graph():
+    dsk = {"y": Task("y", lambda x: x, TaskRef("x"))}
+    with pytest.raises(ValueError, match="Missing dependency"):
+        get_sync(dsk, "y")
+    assert get_sync(dsk, "y", cache={"x": 1}) == 1
diff --git a/dask/tests/test_order.py b/dask/tests/test_order.py
index 4074c5b78..b93c95eb0 100644
--- a/dask/tests/test_order.py
+++ b/dask/tests/test_order.py
@@ -560,10 +560,10 @@ def test_run_smaller_sections(abcde):
     }
     o = order(dsk)
     assert_topological_sort(dsk, o)
-    assert max(diagnostics(dsk)[1]) <= 4  # optimum is 3
-    # This is a mildly ambiguous example
+    # The actual order is a bit ambiguous but the optimal pressure is well
+    # defined, see also
     # https://github.com/dask/dask/pull/10535/files#r1337528255
-    assert (o[aa] < o[a] and o[dd] < o[a]) or (o[b] < o[e] and o[b] < o[cc])
+    assert max(diagnostics(dsk)[1]) == 3
 
 
 def test_local_parents_of_reduction(abcde):