Repository: dask/dask
Pull request #11903 (2025-04-24 09:44:56)

Ensure culling changes layer names
This is pretty much an internal thing but the way HLGs are performing culling is dangerous and makes working with the API more difficult than it has to.

Culling does mutate layers but the internal tracking does not change the layer name to reflect this. Merging culled layers with `HighlevelGraph.merge` can therefore produce corrupted graphs.

I stumbled over this while fixing https://github.com/dask/dask/issues/11901

github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/dask/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

      9 files  ±0        9 suites  ±0   3h 24m 2s ⏱️ + 6m 4s
 18 012 tests +1   16 798 ✅ + 1   1 214 💤 ±0  0 ❌ ±0 
161 182 runs  +9  149 077 ✅ +11  12 105 💤  - 2  0 ❌ ±0 

Results for commit 6c84cac8. ± Comparison against base commit 3a810721.

[test-results]:data:application/gzip;base64,H4sIANEUCmgC/02MSwrDMAwFrxK87sIy/ii9THGUBEzzKY69Kr17lZCq2QhmHpq3GtM0bOretLdGbTUVgb7mWNK6MIIx1rDisewzoAbhx1aJdulDi3/5TK/jE6y4MaaJnRYx5Lzm0+S6HGkPgOZkSdtWhyBS0tr93CV98LVM6zynwqA8oaXIR3cRutAb5wYESzb6SM4hr4SjQfX5AlUNWXwYAQAA


Patch:
diff --git a/dask/highlevelgraph.py b/dask/highlevelgraph.py
index 4a8603537..a2b11d06e 100644
--- a/dask/highlevelgraph.py
+++ b/dask/highlevelgraph.py
@@ -23,7 +23,7 @@ import dask
 from dask import config
 from dask.base import clone_key, flatten, is_dask_collection
 from dask.core import keys_in_tasks, reverse_dict
-from dask.tokenize import normalize_token
+from dask.tokenize import normalize_token, tokenize
 from dask.typing import DaskCollection, Graph, Key
 from dask.utils import ensure_dict, import_required, key_split
 from dask.widgets import get_template
@@ -729,8 +729,10 @@ class HighLevelGraph(Graph):
         all_ext_keys = self.get_all_external_keys()
         ret_layers: dict = {}
         ret_key_deps: dict = {}
+        layer_dependencies = {}
         for layer_name in reversed(self._toposort_layers()):
             layer = self.layers[layer_name]
+            new_layer_name = layer_name
             # Let's cull the layer to produce its part of `keys`.
             # Note: use .intersection rather than & because the RHS is
             # a collections.abc.Set rather than a real set, and using &
@@ -738,6 +740,7 @@ class HighLevelGraph(Graph):
             # if there is no culling can be much bigger than the RHS.
             output_keys = keys_set.intersection(layer.get_output_keys())
             if output_keys:
+                new_layer_name = f"{layer_name}-{tokenize(output_keys)}"
                 culled_layer, culled_deps = layer.cull(output_keys, all_ext_keys)
                 # Update `keys` with all layer's external key dependencies, which
                 # are all the layer's dependencies (`culled_deps`) excluding
@@ -749,7 +752,8 @@ class HighLevelGraph(Graph):
                 keys_set |= external_deps
 
                 # Save the culled layer and its key dependencies
-                ret_layers[layer_name] = culled_layer
+                ret_layers[new_layer_name] = culled_layer
+                layer_dependencies[new_layer_name] = self.dependencies[layer_name]
                 if (
                     isinstance(layer, Blockwise)
                     or isinstance(layer, MaterializedLayer)
@@ -769,7 +773,7 @@ class HighLevelGraph(Graph):
         # intersection to iterate over the smaller of the two sets.
         ret_layers_keys = set(ret_layers.keys())
         ret_dependencies = {
-            layer_name: self.dependencies[layer_name] & ret_layers_keys
+            layer_name: layer_dependencies[layer_name] & ret_layers_keys
             for layer_name in ret_layers
         }
 


Test patch:
diff --git a/dask/tests/test_highgraph.py b/dask/tests/test_highgraph.py
index 409ef7a6a..96066cee9 100644
--- a/dask/tests/test_highgraph.py
+++ b/dask/tests/test_highgraph.py
@@ -9,10 +9,11 @@ from concurrent.futures import ThreadPoolExecutor
 import pytest
 
 import dask
-from dask.base import tokenize
+from dask.base import collections_to_expr, tokenize
 from dask.blockwise import Blockwise
+from dask.delayed import Delayed
 from dask.highlevelgraph import HighLevelGraph, Layer, MaterializedLayer, to_graphviz
-from dask.utils_test import inc
+from dask.utils_test import dec, inc
 
 
 def test_visualize(tmpdir):
@@ -252,11 +253,12 @@ def test_blockwise_cull(flat):
     keys = {(x._name, *select)}
     dsk_cull = dsk.cull(keys)
     for name, layer in dsk_cull.layers.items():
+        old_name = name.rsplit("-", 1)[0]
         if not isinstance(layer, dask.blockwise.Blockwise):
             # The original layer shouldn't be Blockwise if the new one isn't
-            assert not isinstance(dsk.layers[name], dask.blockwise.Blockwise)
+            assert not isinstance(dsk.layers[old_name], dask.blockwise.Blockwise)
             continue
-        assert isinstance(dsk.layers[name], dask.blockwise.Blockwise)
+        assert isinstance(dsk.layers[old_name], dask.blockwise.Blockwise)
         assert not layer.is_materialized()
         out_keys = layer.get_output_keys()
         assert out_keys == {(layer.output, *select)}
@@ -318,3 +320,45 @@ def test_tokenize_hlg():
     c = db.from_sequence(list(range(10)), npartitions=3).max()
     assert tokenize(a.dask) == tokenize(b.dask)
     assert tokenize(a.dask) != tokenize(c.dask)
+
+
+def test_culling_changes_layer_names():
+    """Test that culling changes the layer names in the graph."""
+    dsk = MaterializedLayer({"x": 1, "y": (inc, "x"), "z": (dec, "x")})
+    hg = HighLevelGraph({"a": dsk}, {"a": set()})
+
+    # This construction of Delayed objects is similar to array.to_delayed
+    y = Delayed("y", hg, layer="a")
+    z = Delayed("z", hg, layer="a")
+
+    from dask.delayed import optimize
+
+    # This is roughly what dask.optimize is doing
+    # The exception is that collection_to_dsk takes the Delayed object and
+    # extracts the layer as a low level graph in which case we're loosing the
+    # layer name, i.e. all internal manipulations of the graphs with and without
+    # optimization could be affected but triggering this with user APIs is
+    # almost impossible
+    yopt_hlg = optimize(y._dask, ["y"])
+    assert isinstance(yopt_hlg, HighLevelGraph)
+    layers = yopt_hlg.layers
+    # This test doesn't care about there being a single layer but after
+    # optimization we no longer want to hardcode the delayed layer key
+    assert len(layers) == 1
+    layer_key = next(iter(layers))
+    yopt = Delayed("y", yopt_hlg, layer=layer_key)
+
+    zopt_hlg = optimize(z._dask, ["z"])
+    assert isinstance(zopt_hlg, HighLevelGraph)
+    layers = zopt_hlg.layers
+    # This test doesn't care about there being a single layer but after
+    # optimization we no longer want to hardcode the delayed layer key
+    assert len(layers) == 1
+    layer_key = next(iter(layers))
+    zopt = Delayed("z", zopt_hlg, layer=layer_key)
+
+    # This internally merged the HLG graphs i.e. if layer names are not
+    # deduplicated we will loose keys.
+    expr_opt = collections_to_expr([yopt, zopt]).optimize()
+    dsk_out = expr_opt.__dask_graph__()
+    assert set(dsk_out) == set("xyz")