Repository: dask/dask
Pull request #11305 (2024-08-13 11:46:51)

order: ensure runnable tasks are certainly runnable
This fixes a bug in `process_runnables` that caused the `runnable_hull` to be populated with leaf nodes which is not intended to happen. If a leaf node was encountered during `process_runnables` we'd always want it to be executed greedily.

Fixing this bug allows us to eliminate a couple of code branches that are now no longer reachable. It also reveals a bug that was covered up that I think should/will be fixed by https://github.com/dask/dask/pull/11303

While the ordering for `test_run_smaller_sections` is now actually better in terms of max pressure, we are indeed ordering in a not intuitive way and are not considering already loaded root tasks.




Patch:
diff --git a/dask/array/numpy_compat.py b/dask/array/numpy_compat.py
index dcb32ef8c..7137f034f 100644
--- a/dask/array/numpy_compat.py
+++ b/dask/array/numpy_compat.py
@@ -16,6 +16,7 @@ NUMPY_GE_123 = _np_version.release >= (1, 23)
 NUMPY_GE_124 = _np_version.release >= (1, 24)
 NUMPY_GE_125 = _np_version.release >= (1, 25)
 NUMPY_GE_200 = _np_version.release >= (2, 0)
+NUMPY_GE_210 = _np_version.release >= (2, 1)
 
 
 if NUMPY_GE_200:
diff --git a/dask/order.py b/dask/order.py
index 90b7c223b..4de05c2b4 100644
--- a/dask/order.py
+++ b/dask/order.py
@@ -315,6 +315,8 @@ def order(
                 while branches:
                     path = branches.popleft()
                     while True:
+                        # Loop invariant. Too expensive to compute at runtime
+                        # assert not set(known_runnable_paths).intersection(runnable_hull)
                         current = path[-1]
                         runnable_hull.add(current)
                         deps_downstream = dependents[current]
@@ -326,6 +328,8 @@ def order(
                             if num_needed[current] <= 1:
                                 for k in path:
                                     add_to_result(k)
+                            else:
+                                runnable_hull.discard(current)
                         elif len(path) == 1 or len(deps_upstream) == 1:
                             if len(deps_downstream) > 1:
                                 for d in sorted(deps_downstream, key=sort_key):
@@ -337,11 +341,11 @@ def order(
                                         branch.append(d)
                                         branches.append(branch)
                                 break
-                            runnable_hull.update(deps_downstream)
-                            path.extend(sorted(deps_downstream, key=sort_key))
+                            path.extend(deps_downstream)
                             continue
                         elif current in known_runnable_paths:
                             known_runnable_paths[current].append(path)
+                            runnable_hull.discard(current)
                             if (
                                 len(known_runnable_paths[current])
                                 >= num_needed[current]
@@ -371,6 +375,7 @@ def order(
                                     add_to_result(k)
                             else:
                                 known_runnable_paths[current] = [path]
+                                runnable_hull.discard(current)
                         break
 
     # Pick strategy
@@ -434,10 +439,11 @@ def order(
                 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:
+                # We're not considering runnable_hull because if there was a
+                # runnable leaf node, process_runnables should've taken care of
+                # it already, i.e. the intersection of runnable_hull and
+                # candidates is always empty
+                if reachable_hull:
                     skey = lambda k: (num_needed[k], sort_key(k))
                     candidates = reachable_hull & candidates
 
@@ -445,6 +451,7 @@ def order(
                     if seed := pick_seed():
                         candidates = leafs_connected[seed]
                     else:
+                        # FIXME: This seems to be dead code (at least untested)
                         candidates = runnable_hull or reachable_hull
                 # FIXME: This can be very expensive
                 return min(candidates, key=skey)
@@ -573,10 +580,6 @@ def order(
             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] = []


Test patch:
diff --git a/dask/array/tests/test_creation.py b/dask/array/tests/test_creation.py
index db11190dd..bbf9f88cb 100644
--- a/dask/array/tests/test_creation.py
+++ b/dask/array/tests/test_creation.py
@@ -11,7 +11,7 @@ from tlz import concat
 import dask
 import dask.array as da
 from dask.array.core import normalize_chunks
-from dask.array.numpy_compat import AxisError
+from dask.array.numpy_compat import NUMPY_GE_210, AxisError
 from dask.array.utils import assert_eq, same_keys
 
 
@@ -975,6 +975,8 @@ def test_nan_empty_like(shape_chunks, dtype):
 @pytest.mark.parametrize("shape_chunks", [((50, 4), (10, 2)), ((50,), (10,))])
 @pytest.mark.parametrize("dtype", ["u4", np.float32, None, np.int64])
 def test_nan_full_like(val, shape_chunks, dtype):
+    if NUMPY_GE_210 and val == -1 and dtype == "u4":
+        pytest.xfail("can't insert negative numbers into unsigned integer")
     shape, chunks = shape_chunks
     x1 = da.random.standard_normal(size=shape, chunks=chunks)
     y1 = x1[x1 < 0.5]
diff --git a/dask/tests/test_order.py b/dask/tests/test_order.py
index 70a1c6923..0a456cbcc 100644
--- a/dask/tests/test_order.py
+++ b/dask/tests/test_order.py
@@ -339,6 +339,9 @@ def test_favor_longest_critical_path(abcde):
     assert o[e] > o[b]
 
 
+@pytest.mark.xfail(
+    reason="Second target pick does not include already computed nodes properly"
+)
 def test_run_smaller_sections(abcde):
     r"""
             aa