Use new tokenize module (#8858)

distributed/client.py

@@ -44,11 +44,12 @@
 from tlz import first, groupby, merge, partition_all, valmap
 
 import dask
-from dask.base import collections_to_dsk, tokenize
+from dask.base import collections_to_dsk
 from dask.core import flatten, validate_key
 from dask.highlevelgraph import HighLevelGraph
 from dask.layers import Layer
 from dask.optimization import SubgraphCallable
+from dask.tokenize import tokenize
 from dask.typing import Key, NoDefault, no_default
 from dask.utils import (
     apply,

distributed/diagnostics/progress.py

@@ -9,7 +9,7 @@
 
 from tlz import groupby, valmap
 
-from dask.base import tokenize
+from dask.tokenize import tokenize
 from dask.utils import key_split
 
 from distributed.diagnostics.plugin import SchedulerPlugin

distributed/protocol/serialize.py

@@ -14,8 +14,8 @@
 import msgpack
 
 import dask
-from dask.base import normalize_token
 from dask.sizeof import sizeof
+from dask.tokenize import normalize_token
 from dask.utils import typename
 
 from distributed.metrics import context_meter
@@ -776,7 +776,7 @@ def register_serialization_lazy(toplevel, func):
 
 @partial(normalize_token.register, Serialized)
 def normalize_Serialized(o):
-    return [o.header] + o.frames  # for dask.base.tokenize
+    return [o.header] + o.frames  # for dask.tokenize.tokenize
 
 
 # Teach serialize how to handle bytes

distributed/scheduler.py

@@ -54,8 +54,8 @@
 
 import dask
 import dask.utils
-from dask.base import TokenizationError, normalize_token, tokenize
 from dask.core import get_deps, iskey, validate_key
+from dask.tokenize import TokenizationError, normalize_token, tokenize
 from dask.typing import Key, no_default
 from dask.utils import (
     _deprecated,

distributed/shuffle/_merge.py

@@ -5,9 +5,10 @@
 from typing import TYPE_CHECKING, Any
 
 import dask
-from dask.base import is_dask_collection, tokenize
+from dask.base import is_dask_collection
 from dask.highlevelgraph import HighLevelGraph
 from dask.layers import Layer
+from dask.tokenize import tokenize
 
 from distributed.shuffle._arrow import check_minimal_arrow_version
 from distributed.shuffle._core import ShuffleId, barrier_key, get_worker_plugin

distributed/shuffle/_rechunk.py

@@ -121,9 +121,9 @@
 
 import dask
 import dask.config
-from dask.base import tokenize
 from dask.highlevelgraph import HighLevelGraph
 from dask.layers import Layer
+from dask.tokenize import tokenize
 from dask.typing import Key
 from dask.utils import parse_bytes
 

distributed/shuffle/_shuffle.py

@@ -21,9 +21,9 @@
 from tornado.ioloop import IOLoop
 
 import dask
-from dask.base import tokenize
 from dask.highlevelgraph import HighLevelGraph
 from dask.layers import Layer
+from dask.tokenize import tokenize
 from dask.typing import Key
 
 from distributed.core import PooledRPCCall

distributed/tests/test_client.py

@@ -43,6 +43,7 @@
 import dask.bag as db
 from dask import delayed
 from dask.optimization import SubgraphCallable
+from dask.tokenize import tokenize
 from dask.utils import get_default_shuffle_method, parse_timedelta, tmpfile
 
 from distributed import (
@@ -73,7 +74,6 @@
     futures_of,
     get_task_metadata,
     temp_default_client,
-    tokenize,
     wait,
 )
 from distributed.cluster_dump import load_cluster_dump
@@ -1127,7 +1127,7 @@ async def test_scatter_non_list(c, s, a, b):
 
 @gen_cluster(client=True)
 async def test_scatter_tokenize_local(c, s, a, b):
-    from dask.base import normalize_token
+    from dask.tokenize import normalize_token
 
     class MyObj:
         pass