Repository: dask/dask
Pull request #10794 (2024-01-12 23:28:39)

Generalize ``CreationDispatch`` for dask-expr
Minor change to simplify migration of creation dispatching to dask-expr (see https://github.com/dask-contrib/dask-expr/pull/728).


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._

     15 files  ±0       15 suites  ±0   3h 29m 47s :stopwatch: - 15m 17s
 12 894 tests ±0   11 469 :white_check_mark: +1     929 :zzz: ±0  496 :x:  - 1 
159 765 runs  ±0  142 361 :white_check_mark: +1  16 908 :zzz: ±0  496 :x:  - 1 

For more details on these failures, see [this check](https://github.com/dask/dask/runs/20444782822).

Results for commit ca305a35. ± Comparison against base commit 91dd4252.

[test-results]:data:application/gzip;base64,H4sIAFjToWUC/1XNTQ6DIBCG4asY1l0AMgPTyzQEISH1p0FYmd69qFXb5ftM8s3CQuz9zO6NgFvD5hLzVV1JNsdpXFuC0ZXqNW93aUgd/ZiLcysKhXThM74qkrwo2NhXUoQn+ZSmVI1XSWXcf5NG+Pa5rWSL4sR9WyBxc9jf+Ca/224ahphrMGdbDrYFr7X0WnSOQ8BAhgLvCNA7UlqhEOz9AfFomBwdAQAA


Patch:
diff --git a/dask/backends.py b/dask/backends.py
index ffa382c37..4dc11fcfd 100644
--- a/dask/backends.py
+++ b/dask/backends.py
@@ -59,6 +59,7 @@ class CreationDispatch(Generic[BackendEntrypointType]):
     _config_field: str
     _default: str
     _entrypoint_class: type[BackendEntrypointType]
+    _entrypoint_root: str
 
     def __init__(
         self,
@@ -66,12 +67,14 @@ class CreationDispatch(Generic[BackendEntrypointType]):
         default: str,
         entrypoint_class: type[BackendEntrypointType],
         name: str | None = None,
+        entrypoint_root: str = "dask",
     ):
         self._lookup = {}
         self._module_name = module_name
         self._config_field = f"{module_name}.backend"
         self._default = default
         self._entrypoint_class = entrypoint_class
+        self._entrypoint_root = entrypoint_root
         if name:
             self.__name__ = name
 
@@ -94,7 +97,9 @@ class CreationDispatch(Generic[BackendEntrypointType]):
             impl = self._lookup[backend]
         except KeyError:
             # Check entrypoints for the specified backend
-            entrypoints = detect_entrypoints(f"dask.{self._module_name}.backends")
+            entrypoints = detect_entrypoints(
+                f"{self._entrypoint_root}.{self._module_name}.backends"
+            )
             if backend in entrypoints:
                 return self.register_backend(backend, entrypoints[backend].load()())
         else:
diff --git a/dask/highlevelgraph.py b/dask/highlevelgraph.py
index 0f13212ce..5448019ee 100644
--- a/dask/highlevelgraph.py
+++ b/dask/highlevelgraph.py
@@ -21,7 +21,7 @@ import tlz as toolz
 
 import dask
 from dask import config
-from dask.base import clone_key, flatten, is_dask_collection
+from dask.base import clone_key, flatten, is_dask_collection, normalize_token
 from dask.core import keys_in_tasks, reverse_dict
 from dask.typing import DaskCollection, Graph, Key
 from dask.utils import ensure_dict, import_required, key_split
@@ -993,3 +993,8 @@ def _get_some_layer_name(collection) -> str:
         # collection does not define the optional __dask_layers__ method
         # or it spuriously returns more than one layer
         return str(id(collection))
+
+
+@normalize_token.register(HighLevelGraph)
+def register_highlevelgraph(hlg):
+    return normalize_token(list(hlg.layers.keys()))


Test patch:
diff --git a/dask/tests/test_highgraph.py b/dask/tests/test_highgraph.py
index 87283c233..03a21c763 100644
--- a/dask/tests/test_highgraph.py
+++ b/dask/tests/test_highgraph.py
@@ -9,6 +9,7 @@ from concurrent.futures import ThreadPoolExecutor
 import pytest
 
 import dask
+from dask.base import normalize_token
 from dask.blockwise import Blockwise, blockwise_token
 from dask.highlevelgraph import HighLevelGraph, Layer, MaterializedLayer, to_graphviz
 from dask.utils_test import inc
@@ -300,3 +301,13 @@ def test_node_tooltips_exist():
             end = layer.find('"', start)
             tooltip = layer[start:end]
             assert len(tooltip) > 0
+
+
+def test_tokenize_hlg():
+    import dask.bag as db
+
+    a = db.from_sequence(list(range(10)), npartitions=2).max()
+    b = db.from_sequence(list(range(10)), npartitions=2).max()
+    c = db.from_sequence(list(range(10)), npartitions=3).max()
+    assert normalize_token(a.dask) == normalize_token(b.dask)
+    assert normalize_token(a.dask) != normalize_token(c.dask)