Repository: dask/dask
Pull request #11371 (2024-09-04 08:41:18)

Move ``tokenize`` to dedicated submodule
By now, tokenize is a rather nontrivial feature. It is still known to have problems so we'll likely need to iterate on this a bit. Moving it to a dedicated module allows us to work on it a little more easily.


I am not raising any deprecation warnings, yet. I first would like to move distributed as well (not top prio right now)


Patch:
diff --git a/dask/array/ma.py b/dask/array/ma.py
index b1cf1a4c4..cc00aaecc 100644
--- a/dask/array/ma.py
+++ b/dask/array/ma.py
@@ -9,7 +9,7 @@ from dask.array.core import asanyarray, blockwise, elemwise, map_blocks
 from dask.array.reductions import reduction
 from dask.array.routines import _average
 from dask.array.routines import nonzero as _nonzero
-from dask.base import normalize_token
+from dask.tokenize import normalize_token
 from dask.utils import derived_from
 
 
diff --git a/dask/array/ufunc.py b/dask/array/ufunc.py
index 1a8d790b7..2af747a40 100644
--- a/dask/array/ufunc.py
+++ b/dask/array/ufunc.py
@@ -7,8 +7,9 @@ import numpy as np
 
 from dask import core
 from dask.array.core import Array, apply_infer_dtype, asarray, blockwise, elemwise
-from dask.base import is_dask_collection, normalize_token
+from dask.base import is_dask_collection
 from dask.highlevelgraph import HighLevelGraph
+from dask.tokenize import normalize_token
 from dask.utils import derived_from, funcname
 
 
diff --git a/dask/base.py b/dask/base.py
index b7a4cd00e..6394ac28f 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -1,18 +1,11 @@
 from __future__ import annotations
 
-import copyreg
 import dataclasses
-import datetime
-import decimal
-import hashlib
 import inspect
-import pathlib
-import pickle
-import types
 import uuid
 import warnings
 from collections import OrderedDict
-from collections.abc import Hashable, Iterable, Iterator, Mapping
+from collections.abc import Hashable, Iterator, Mapping
 from concurrent.futures import Executor
 from contextlib import contextmanager, suppress
 from contextvars import ContextVar
@@ -21,20 +14,16 @@ from numbers import Integral, Number
 from operator import getitem
 from typing import Any, Literal, TypeVar
 
-import cloudpickle
-from tlz import curry, groupby, identity, merge
-from tlz.functoolz import Compose
+from tlz import groupby, merge
 
 from dask import config, local
 from dask._compatibility import EMSCRIPTEN
 from dask.core import flatten
 from dask.core import get as simple_get
-from dask.core import literal, quote
-from dask.hashing import hash_buffer_hex
+from dask.core import quote
 from dask.system import CPU_COUNT
 from dask.typing import Key, SchedulerGetCallable
 from dask.utils import (
-    Dispatch,
     apply,
     ensure_dict,
     is_namedtuple_instance,
@@ -79,6 +68,9 @@ __all__ = (
     "clone_key",
 )
 
+# Backwards compat
+from dask.tokenize import TokenizationError, normalize_token, tokenize  # noqa: F401
+
 _annotations: ContextVar[dict[str, Any] | None] = ContextVar(
     "annotations", default=None
 )
@@ -1009,442 +1001,6 @@ def persist(*args, traverse=True, optimize_graph=True, scheduler=None, **kwargs)
     return repack(results2)
 
 
-############
-# Tokenize #
-############
-class TokenizationError(RuntimeError):
-    pass
-
-
-def tokenize(
-    *args: object, ensure_deterministic: bool | None = None, **kwargs: object
-) -> str:
-    """Deterministic token
-
-    >>> tokenize([1, 2, '3'])
-    '06961e8de572e73c2e74b51348177918'
-
-    >>> tokenize('Hello') == tokenize('Hello')
-    True
-
-    Parameters
-    ----------
-    args, kwargs:
-        objects to tokenize
-    ensure_deterministic: bool, optional
-        If True, raise TokenizationError if the objects cannot be deterministically
-        tokenized, e.g. two identical objects will return different tokens.
-        Defaults to the `tokenize.ensure-deterministic` configuration parameter.
-    """
-    with _seen_ctx(reset=True), _ensure_deterministic_ctx(ensure_deterministic):
-        token: object = _normalize_seq_func(args)
-        if kwargs:
-            token = token, _normalize_seq_func(sorted(kwargs.items()))
-
-    # Pass `usedforsecurity=False` to support FIPS builds of Python
-    return hashlib.md5(str(token).encode(), usedforsecurity=False).hexdigest()
-
-
-# tokenize.ensure-deterministic flag, potentially overridden by tokenize()
-_ensure_deterministic: ContextVar[bool] = ContextVar("_ensure_deterministic")
-
-# Circular reference breaker used by _normalize_seq_func.
-# This variable is recreated anew every time you call tokenize(). Note that this means
-# that you could call tokenize() from inside tokenize() and they would be fully
-# independent.
-#
-# It is a map of {id(obj): (<first seen incremental int>, obj)} which causes an object
-# to be tokenized as ("__seen", <incremental>) the second time it's encountered while
-# traversing collections. A strong reference to the object is stored in the context to
-# prevent ids from being reused by different objects.
-_seen: ContextVar[dict[int, tuple[int, object]]] = ContextVar("_seen")
-
-
-@contextmanager
-def _ensure_deterministic_ctx(ensure_deterministic: bool | None) -> Iterator[bool]:
-    try:
-        ensure_deterministic = _ensure_deterministic.get()
-        # There's a call of tokenize() higher up in the stack
-        tok = None
-    except LookupError:
-        # Outermost tokenize(), or normalize_token() was called directly
-        if ensure_deterministic is None:
-            ensure_deterministic = config.get("tokenize.ensure-deterministic")
-        tok = _ensure_deterministic.set(ensure_deterministic)
-
-    try:
-        yield ensure_deterministic
-    finally:
-        if tok:
-            tok.var.reset(tok)
-
-
-def _maybe_raise_nondeterministic(msg: str) -> None:
-    with _ensure_deterministic_ctx(None) as ensure_deterministic:
-        if ensure_deterministic:
-            raise TokenizationError(msg)
-
-
-@contextmanager
-def _seen_ctx(reset: bool) -> Iterator[dict[int, tuple[int, object]]]:
-    if reset:
-        # It is important to reset the token on tokenize() to avoid artifacts when
-        # it is called recursively
-        seen: dict[int, tuple[int, object]] = {}
-        tok = _seen.set(seen)
-    else:
-        try:
-            seen = _seen.get()
-            tok = None
-        except LookupError:
-            # This is for debug only, for when normalize_token is called outside of
-            # tokenize()
-            seen = {}
-            tok = _seen.set(seen)
-
-    try:
-        yield seen
-    finally:
-        if tok:
-            tok.var.reset(tok)
-
-
-normalize_token = Dispatch()
-normalize_token.register(
-    (
-        int,
-        float,
-        str,
-        bytes,
-        type(None),
-        slice,
-        complex,
-        type(Ellipsis),
-        decimal.Decimal,
-        datetime.date,
-        datetime.time,
-        datetime.datetime,
-        datetime.timedelta,
-        pathlib.PurePath,
-    ),
-    identity,
-)
-
-
-@normalize_token.register((types.MappingProxyType, dict))
-def normalize_dict(d):
-    return "dict", _normalize_seq_func(sorted(d.items(), key=lambda kv: str(kv[0])))
-
-
-@normalize_token.register(OrderedDict)
-def normalize_ordered_dict(d):
-    return _normalize_seq_func((type(d), list(d.items())))
-
-
-@normalize_token.register(set)
-def normalize_set(s):
-    # Note: in some Python version / OS combinations, set order changes every
-    # time you recreate the set (even within the same interpreter).
-    # In most other cases, set ordering is consistent within the same interpreter.
-    return "set", _normalize_seq_func(sorted(s, key=str))
-
-
-def _normalize_seq_func(seq: Iterable[object]) -> list[object]:
-    with _seen_ctx(reset=False) as seen:
-        out = []
-        for item in seq:
-            if isinstance(item, (str, bytes, int, float, bool, type(None))):
-                # Basic data type. This is just for performance and compactness of the
-                # output. It doesn't need to be a comprehensive list.
-                pass
-            elif id(item) in seen:
-                # May or may not be a circular recursion. Maybe just a double reference.
-                seen_when, _ = seen[id(item)]
-                item = "__seen", seen_when
-            else:
-                seen[id(item)] = len(seen), item
-                item = normalize_token(item)
-            out.append(item)
-        return out
-
-
-@normalize_token.register((tuple, list))
-def normalize_seq(seq):
-    return type(seq).__name__, _normalize_seq_func(seq)
-
-
-@normalize_token.register(literal)
-def normalize_literal(lit):
-    return "literal", normalize_token(lit())
-
-
-@normalize_token.register(Compose)
-def normalize_compose(func):
-    return _normalize_seq_func((func.first,) + func.funcs)
-
-
-@normalize_token.register((partial, curry))
-def normalize_partial(func):
-    return _normalize_seq_func((func.func, func.args, func.keywords))
-
-
-@normalize_token.register((types.MethodType, types.MethodWrapperType))
-def normalize_bound_method(meth):
-    return normalize_token(meth.__self__), meth.__name__
-
-
-@normalize_token.register(types.BuiltinFunctionType)
-def normalize_builtin_function_or_method(func):
-    # Note: BuiltinMethodType is BuiltinFunctionType
-    self = getattr(func, "__self__", None)
-    if self is not None and not inspect.ismodule(self):
-        return normalize_bound_method(func)
-    else:
-        return normalize_object(func)
-
-
-@normalize_token.register(object)
-def normalize_object(o):
-    method = getattr(o, "__dask_tokenize__", None)
-    if method is not None and not isinstance(o, type):
-        return method()
-
-    if type(o) is object:
-        return _normalize_pure_object(o)
-
-    if isinstance(o, type):
-        copyreg._slotnames(o)
-
-    if dataclasses.is_dataclass(o) and not isinstance(o, type):
-        return _normalize_dataclass(o)
-
-    try:
-        return _normalize_pickle(o)
-    except Exception:
-        _maybe_raise_nondeterministic(
-            f"Object {o!r} cannot be deterministically hashed. See "
-            "https://docs.dask.org/en/latest/custom-collections.html#implementing-deterministic-hashing "
-            "for more information."
-        )
-        return uuid.uuid4().hex
-
-
-_seen_objects = set()
-
-
-def _normalize_pure_object(o: object) -> tuple[str, int]:
-    _maybe_raise_nondeterministic(
-        "object() cannot be deterministically hashed. See "
-        "https://docs.dask.org/en/latest/custom-collections.html#implementing-deterministic-hashing "
-        "for more information."
-    )
-    # Idempotent, but not deterministic. Make sure that the id is not reused.
-    _seen_objects.add(o)
-    return "object", id(o)
-
-
-def _normalize_pickle(o: object) -> tuple:
-    buffers: list[pickle.PickleBuffer] = []
-    pik: int | None = None
-    pik2: int
-
-    for mod in [pickle, cloudpickle]:
-        for _ in range(3):
-            buffers.clear()
-            try:
-                out = mod.dumps(o, protocol=5, buffer_callback=buffers.append)
-                mod.loads(out, buffers=buffers)
-                pik2 = hash_buffer_hex(out)
-            except Exception:
-                break
-            if pik == pik2:
-                break
-            pik = pik2
-        else:
-            _maybe_raise_nondeterministic("Failed to tokenize deterministically")
-            break
-    if pik is None:
-        _maybe_raise_nondeterministic("Failed to tokenize deterministically")
-        pik = int(uuid.uuid4())
-    return pik, [hash_buffer_hex(buf) for buf in buffers]
-
-
-def _normalize_dataclass(obj):
-    fields = [
-        (field.name, normalize_token(getattr(obj, field.name, None)))
-        for field in dataclasses.fields(obj)
-    ]
-    params = obj.__dataclass_params__
-    params = [(attr, getattr(params, attr)) for attr in params.__slots__]
-
-    return normalize_object(type(obj)), params, fields
-
-
-@normalize_token.register_lazy("pandas")
-def register_pandas():
-    import pandas as pd
-
-    @normalize_token.register(pd.Index)
-    def normalize_index(ind):
-        values = ind.array
-        return [ind.name, normalize_token(values)]
-
-    @normalize_token.register(pd.MultiIndex)
-    def normalize_index(ind):
-        codes = ind.codes
-        return (
-            [ind.name]
-            + [normalize_token(x) for x in ind.levels]
-            + [normalize_token(x) for x in codes]
-        )
-
-    @normalize_token.register(pd.Categorical)
-    def normalize_categorical(cat):
-        return [normalize_token(cat.codes), normalize_token(cat.dtype)]
-
-    @normalize_token.register(pd.arrays.PeriodArray)
-    @normalize_token.register(pd.arrays.DatetimeArray)
-    @normalize_token.register(pd.arrays.TimedeltaArray)
-    def normalize_period_array(arr):
-        return [normalize_token(arr.asi8), normalize_token(arr.dtype)]
-
-    @normalize_token.register(pd.arrays.IntervalArray)
-    def normalize_interval_array(arr):
-        return [
-            normalize_token(arr.left),
-            normalize_token(arr.right),
-            normalize_token(arr.closed),
-        ]
-
-    @normalize_token.register(pd.Series)
-    def normalize_series(s):
-        return [
-            s.name,
-            s.dtype,
-            normalize_token(s._values),
-            normalize_token(s.index),
-        ]
-
-    @normalize_token.register(pd.DataFrame)
-    def normalize_dataframe(df):
-        mgr = df._mgr
-        data = list(mgr.arrays) + [df.columns, df.index]
-        return list(map(normalize_token, data))
-
-    @normalize_token.register(pd.arrays.ArrowExtensionArray)
-    def normalize_extension_array(arr):
-        try:
-            return (type(arr), normalize_token(arr._pa_array))
-        except AttributeError:
-            return (type(arr), normalize_token(arr._data))
-
-    @normalize_token.register(pd.api.extensions.ExtensionArray)
-    def normalize_extension_array(arr):
-        import numpy as np
-
-        return normalize_token(np.asarray(arr))
-
-    # Dtypes
-    @normalize_token.register(pd.api.types.CategoricalDtype)
-    def normalize_categorical_dtype(dtype):
-        return [normalize_token(dtype.categories), normalize_token(dtype.ordered)]
-
-    @normalize_token.register(pd.api.extensions.ExtensionDtype)
-    def normalize_period_dtype(dtype):
-        return normalize_token(dtype.name)
-
-    @normalize_token.register(type(pd.NA))
-    def normalize_na(na):
-        return pd.NA
-
-    @normalize_token.register(pd.offsets.BaseOffset)
-    def normalize_offset(offset):
-        return offset.freqstr
-
-
-@normalize_token.register_lazy("pyarrow")
-def register_pyarrow():
-    import pyarrow as pa
-
-    @normalize_token.register(pa.DataType)
-    def normalize_datatype(dt):
-        return pickle.dumps(dt, protocol=4)
-
-    @normalize_token.register(pa.Table)
-    def normalize_table(dt):
-        return (
-            "pa.Table",
-            normalize_token(dt.schema),
-            normalize_token(dt.columns),
-        )
-
-    @normalize_token.register(pa.ChunkedArray)
-    def normalize_chunked_array(arr):
-        return (
-            "pa.ChunkedArray",
-            normalize_token(arr.type),
-            normalize_token(arr.chunks),
-        )
-
-    @normalize_token.register(pa.Array)
-    def normalize_chunked_array(arr):
-        return (
-            "pa.Array",
-            normalize_token(arr.type),
-            normalize_token(arr.buffers()),
-        )
-
-    @normalize_token.register(pa.Buffer)
-    def normalize_chunked_array(buf):
-        return ("pa.Buffer", hash_buffer_hex(buf))
-
-
-@normalize_token.register_lazy("numpy")
-def register_numpy():
-    import numpy as np
-
-    @normalize_token.register(np.ndarray)
-    def normalize_array(x):
-        if not x.shape:
-            return (x.item(), x.dtype)
-        if x.dtype.hasobject:
-            try:
-                try:
-                    # string fast-path
-                    data = hash_buffer_hex(
-                        "-".join(x.flat).encode(
-                            encoding="utf-8", errors="surrogatepass"
-                        )
-                    )
-                except UnicodeDecodeError:
-                    # bytes fast-path
-                    data = hash_buffer_hex(b"-".join(x.flat))
-            except (TypeError, UnicodeDecodeError):
-                return normalize_object(x)
-        else:
-            try:
-                data = hash_buffer_hex(x.ravel(order="K").view("i1"))
-            except (BufferError, AttributeError, ValueError):
-                data = hash_buffer_hex(x.copy().ravel(order="K").view("i1"))
-        return (data, x.dtype, x.shape)
-
-    @normalize_token.register(np.memmap)
-    def normalize_mmap(mm):
-        return hash_buffer_hex(np.ascontiguousarray(mm))
-
-    @normalize_token.register(np.ufunc)
-    def normalize_ufunc(func):
-        try:
-            return _normalize_pickle(func)
-        except Exception:
-            _maybe_raise_nondeterministic(
-                f"Cannot tokenize numpy ufunc {func!r}. Please use functions "
-                "of the dask.array.ufunc module instead. See also "
-                "https://docs.dask.org/en/latest/array-numpy-compatibility.html"
-            )
-            return uuid.uuid4().hex
-
-
 def _colorize(t):
     """Convert (r, g, b) triple to "#RRGGBB" string
 
diff --git a/dask/dataframe/io/parquet/arrow.py b/dask/dataframe/io/parquet/arrow.py
index c674443e4..aba8fcc1b 100644
--- a/dask/dataframe/io/parquet/arrow.py
+++ b/dask/dataframe/io/parquet/arrow.py
@@ -21,7 +21,6 @@ from pyarrow import dataset as pa_ds
 from pyarrow import fs as pa_fs
 
 import dask
-from dask.base import normalize_token, tokenize
 from dask.core import flatten
 from dask.dataframe._compat import PANDAS_GE_220
 from dask.dataframe.backends import pyarrow_schema_dispatch
@@ -39,6 +38,7 @@ from dask.dataframe.io.parquet.utils import (
 from dask.dataframe.io.utils import _get_pyarrow_dtypes, _is_local_fs, _open_input_files
 from dask.dataframe.utils import clear_known_categories, pyarrow_strings_enabled
 from dask.delayed import Delayed
+from dask.tokenize import normalize_token, tokenize
 from dask.utils import getargspec, natural_sort_key
 
 PYARROW_NULLABLE_DTYPE_MAPPING = {
diff --git a/dask/highlevelgraph.py b/dask/highlevelgraph.py
index 01658122b..33bbb5fc6 100644
--- a/dask/highlevelgraph.py
+++ b/dask/highlevelgraph.py
@@ -21,8 +21,9 @@ import tlz as toolz
 
 import dask
 from dask import config
-from dask.base import clone_key, flatten, is_dask_collection, normalize_token
+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.typing import DaskCollection, Graph, Key
 from dask.utils import ensure_dict, import_required, key_split
 from dask.widgets import get_template
diff --git a/dask/optimization.py b/dask/optimization.py
index 68944594c..49fab6157 100644
--- a/dask/optimization.py
+++ b/dask/optimization.py
@@ -7,7 +7,6 @@ from enum import Enum
 from typing import Any
 
 from dask import config, core, utils
-from dask.base import normalize_token, tokenize
 from dask.core import (
     flatten,
     get_dependencies,
@@ -17,6 +16,7 @@ from dask.core import (
     subs,
     toposort,
 )
+from dask.tokenize import normalize_token, tokenize
 from dask.typing import Graph, Key
 
 
diff --git a/dask/tokenize.py b/dask/tokenize.py
new file mode 100644
index 000000000..96c0b5846
--- /dev/null
+++ b/dask/tokenize.py
@@ -0,0 +1,460 @@
+from __future__ import annotations
+
+import copyreg
+import dataclasses
+import datetime
+import decimal
+import hashlib
+import inspect
+import pathlib
+import pickle
+import types
+import uuid
+from collections import OrderedDict
+from collections.abc import Iterable, Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
+from functools import partial
+
+import cloudpickle
+from tlz import curry, identity
+from tlz.functoolz import Compose
+
+from dask import config
+from dask.core import literal
+from dask.hashing import hash_buffer_hex
+from dask.utils import Dispatch
+
+
+class TokenizationError(RuntimeError):
+    pass
+
+
+def tokenize(
+    *args: object, ensure_deterministic: bool | None = None, **kwargs: object
+) -> str:
+    """Deterministic token
+
+    >>> tokenize([1, 2, '3'])
+    '06961e8de572e73c2e74b51348177918'
+
+    >>> tokenize('Hello') == tokenize('Hello')
+    True
+
+    Parameters
+    ----------
+    args, kwargs:
+        objects to tokenize
+    ensure_deterministic: bool, optional
+        If True, raise TokenizationError if the objects cannot be deterministically
+        tokenized, e.g. two identical objects will return different tokens.
+        Defaults to the `tokenize.ensure-deterministic` configuration parameter.
+    """
+    with _seen_ctx(reset=True), _ensure_deterministic_ctx(ensure_deterministic):
+        token: object = _normalize_seq_func(args)
+        if kwargs:
+            token = token, _normalize_seq_func(sorted(kwargs.items()))
+
+    # Pass `usedforsecurity=False` to support FIPS builds of Python
+    return hashlib.md5(str(token).encode(), usedforsecurity=False).hexdigest()
+
+
+# tokenize.ensure-deterministic flag, potentially overridden by tokenize()
+_ensure_deterministic: ContextVar[bool] = ContextVar("_ensure_deterministic")
+
+# Circular reference breaker used by _normalize_seq_func.
+# This variable is recreated anew every time you call tokenize(). Note that this means
+# that you could call tokenize() from inside tokenize() and they would be fully
+# independent.
+#
+# It is a map of {id(obj): (<first seen incremental int>, obj)} which causes an object
+# to be tokenized as ("__seen", <incremental>) the second time it's encountered while
+# traversing collections. A strong reference to the object is stored in the context to
+# prevent ids from being reused by different objects.
+_seen: ContextVar[dict[int, tuple[int, object]]] = ContextVar("_seen")
+
+
+@contextmanager
+def _ensure_deterministic_ctx(ensure_deterministic: bool | None) -> Iterator[bool]:
+    try:
+        ensure_deterministic = _ensure_deterministic.get()
+        # There's a call of tokenize() higher up in the stack
+        tok = None
+    except LookupError:
+        # Outermost tokenize(), or normalize_token() was called directly
+        if ensure_deterministic is None:
+            ensure_deterministic = config.get("tokenize.ensure-deterministic")
+            if ensure_deterministic is None:
+                ensure_deterministic = False
+        tok = _ensure_deterministic.set(ensure_deterministic)
+
+    try:
+        yield ensure_deterministic
+    finally:
+        if tok:
+            tok.var.reset(tok)
+
+
+def _maybe_raise_nondeterministic(msg: str) -> None:
+    with _ensure_deterministic_ctx(None) as ensure_deterministic:
+        if ensure_deterministic:
+            raise TokenizationError(msg)
+
+
+@contextmanager
+def _seen_ctx(reset: bool) -> Iterator[dict[int, tuple[int, object]]]:
+    if reset:
+        # It is important to reset the token on tokenize() to avoid artifacts when
+        # it is called recursively
+        seen: dict[int, tuple[int, object]] = {}
+        tok = _seen.set(seen)
+    else:
+        try:
+            seen = _seen.get()
+            tok = None
+        except LookupError:
+            # This is for debug only, for when normalize_token is called outside of
+            # tokenize()
+            seen = {}
+            tok = _seen.set(seen)
+    try:
+        yield seen
+    finally:
+        if tok:
+            tok.var.reset(tok)
+
+
+normalize_token = Dispatch()
+normalize_token.register(
+    (
+        int,
+        float,
+        str,
+        bytes,
+        type(None),
+        slice,
+        complex,
+        type(Ellipsis),
+        decimal.Decimal,
+        datetime.date,
+        datetime.time,
+        datetime.datetime,
+        datetime.timedelta,
+        pathlib.PurePath,
+    ),
+    identity,
+)
+
+
+@normalize_token.register((types.MappingProxyType, dict))
+def normalize_dict(d):
+    return "dict", _normalize_seq_func(sorted(d.items(), key=lambda kv: str(kv[0])))
+
+
+@normalize_token.register(OrderedDict)
+def normalize_ordered_dict(d):
+    return _normalize_seq_func((type(d), list(d.items())))
+
+
+@normalize_token.register(set)
+def normalize_set(s):
+    # Note: in some Python version / OS combinations, set order changes every
+    # time you recreate the set (even within the same interpreter).
+    # In most other cases, set ordering is consistent within the same interpreter.
+    return "set", _normalize_seq_func(sorted(s, key=str))
+
+
+def _normalize_seq_func(seq: Iterable[object]) -> list[object]:
+    with _seen_ctx(reset=False) as seen:
+        out = []
+        for item in seq:
+            if isinstance(item, (str, bytes, int, float, bool, type(None))):
+                # Basic data type. This is just for performance and compactness of the
+                # output. It doesn't need to be a comprehensive list.
+                pass
+            elif id(item) in seen:
+                # May or may not be a circular recursion. Maybe just a double reference.
+                seen_when, _ = seen[id(item)]
+                item = "__seen", seen_when
+            else:
+                seen[id(item)] = len(seen), item
+                item = normalize_token(item)
+            out.append(item)
+        return out
+
+
+@normalize_token.register((tuple, list))
+def normalize_seq(seq):
+    return type(seq).__name__, _normalize_seq_func(seq)
+
+
+@normalize_token.register(literal)
+def normalize_literal(lit):
+    return "literal", normalize_token(lit())
+
+
+@normalize_token.register(Compose)
+def normalize_compose(func):
+    return _normalize_seq_func((func.first,) + func.funcs)
+
+
+@normalize_token.register((partial, curry))
+def normalize_partial(func):
+    return _normalize_seq_func((func.func, func.args, func.keywords))
+
+
+@normalize_token.register((types.MethodType, types.MethodWrapperType))
+def normalize_bound_method(meth):
+    return normalize_token(meth.__self__), meth.__name__
+
+
+@normalize_token.register(types.BuiltinFunctionType)
+def normalize_builtin_function_or_method(func):
+    # Note: BuiltinMethodType is BuiltinFunctionType
+    self = getattr(func, "__self__", None)
+    if self is not None and not inspect.ismodule(self):
+        return normalize_bound_method(func)
+    else:
+        return normalize_object(func)
+
+
+@normalize_token.register(object)
+def normalize_object(o):
+    method = getattr(o, "__dask_tokenize__", None)
+    if method is not None and not isinstance(o, type):
+        return method()
+
+    if type(o) is object:
+        return _normalize_pure_object(o)
+
+    if isinstance(o, type):
+        copyreg._slotnames(o)
+
+    if dataclasses.is_dataclass(o) and not isinstance(o, type):
+        return _normalize_dataclass(o)
+
+    try:
+        return _normalize_pickle(o)
+    except Exception:
+        _maybe_raise_nondeterministic(
+            f"Object {o!r} cannot be deterministically hashed. See "
+            "https://docs.dask.org/en/latest/custom-collections.html#implementing-deterministic-hashing "
+            "for more information."
+        )
+        return uuid.uuid4().hex
+
+
+_seen_objects = set()
+
+
+def _normalize_pure_object(o: object) -> tuple[str, int]:
+    _maybe_raise_nondeterministic(
+        "object() cannot be deterministically hashed. See "
+        "https://docs.dask.org/en/latest/custom-collections.html#implementing-deterministic-hashing "
+        "for more information."
+    )
+    # Idempotent, but not deterministic. Make sure that the id is not reused.
+    _seen_objects.add(o)
+    return "object", id(o)
+
+
+def _normalize_pickle(o: object) -> tuple:
+    buffers: list[pickle.PickleBuffer] = []
+    pik: int | None = None
+    pik2: int
+
+    for mod in [pickle, cloudpickle]:
+        for _ in range(3):
+            buffers.clear()
+            try:
+                out = mod.dumps(o, protocol=5, buffer_callback=buffers.append)
+                mod.loads(out, buffers=buffers)
+                pik2 = hash_buffer_hex(out)
+            except Exception:
+                break
+            if pik == pik2:
+                break
+            pik = pik2
+        else:
+            _maybe_raise_nondeterministic("Failed to tokenize deterministically")
+            break
+    if pik is None:
+        _maybe_raise_nondeterministic("Failed to tokenize deterministically")
+        pik = int(uuid.uuid4())
+    return pik, [hash_buffer_hex(buf) for buf in buffers]
+
+
+def _normalize_dataclass(obj):
+    fields = [
+        (field.name, normalize_token(getattr(obj, field.name, None)))
+        for field in dataclasses.fields(obj)
+    ]
+    params = obj.__dataclass_params__
+    params = [(attr, getattr(params, attr)) for attr in params.__slots__]
+
+    return normalize_object(type(obj)), params, fields
+
+
+@normalize_token.register_lazy("pandas")
+def register_pandas():
+    import pandas as pd
+
+    @normalize_token.register(pd.Index)
+    def normalize_index(ind):
+        values = ind.array
+        return [ind.name, normalize_token(values)]
+
+    @normalize_token.register(pd.MultiIndex)
+    def normalize_index(ind):
+        codes = ind.codes
+        return (
+            [ind.name]
+            + [normalize_token(x) for x in ind.levels]
+            + [normalize_token(x) for x in codes]
+        )
+
+    @normalize_token.register(pd.Categorical)
+    def normalize_categorical(cat):
+        return [normalize_token(cat.codes), normalize_token(cat.dtype)]
+
+    @normalize_token.register(pd.arrays.PeriodArray)
+    @normalize_token.register(pd.arrays.DatetimeArray)
+    @normalize_token.register(pd.arrays.TimedeltaArray)
+    def normalize_period_array(arr):
+        return [normalize_token(arr.asi8), normalize_token(arr.dtype)]
+
+    @normalize_token.register(pd.arrays.IntervalArray)
+    def normalize_interval_array(arr):
+        return [
+            normalize_token(arr.left),
+            normalize_token(arr.right),
+            normalize_token(arr.closed),
+        ]
+
+    @normalize_token.register(pd.Series)
+    def normalize_series(s):
+        return [
+            s.name,
+            s.dtype,
+            normalize_token(s._values),
+            normalize_token(s.index),
+        ]
+
+    @normalize_token.register(pd.DataFrame)
+    def normalize_dataframe(df):
+        mgr = df._mgr
+        data = list(mgr.arrays) + [df.columns, df.index]
+        return list(map(normalize_token, data))
+
+    @normalize_token.register(pd.arrays.ArrowExtensionArray)
+    def normalize_extension_array(arr):
+        try:
+            return (type(arr), normalize_token(arr._pa_array))
+        except AttributeError:
+            return (type(arr), normalize_token(arr._data))
+
+    @normalize_token.register(pd.api.extensions.ExtensionArray)
+    def normalize_extension_array(arr):
+        import numpy as np
+
+        return normalize_token(np.asarray(arr))
+
+    # Dtypes
+    @normalize_token.register(pd.api.types.CategoricalDtype)
+    def normalize_categorical_dtype(dtype):
+        return [normalize_token(dtype.categories), normalize_token(dtype.ordered)]
+
+    @normalize_token.register(pd.api.extensions.ExtensionDtype)
+    def normalize_period_dtype(dtype):
+        return normalize_token(dtype.name)
+
+    @normalize_token.register(type(pd.NA))
+    def normalize_na(na):
+        return pd.NA
+
+    @normalize_token.register(pd.offsets.BaseOffset)
+    def normalize_offset(offset):
+        return offset.freqstr
+
+
+@normalize_token.register_lazy("pyarrow")
+def register_pyarrow():
+    import pyarrow as pa
+
+    @normalize_token.register(pa.DataType)
+    def normalize_datatype(dt):
+        return pickle.dumps(dt, protocol=4)
+
+    @normalize_token.register(pa.Table)
+    def normalize_table(dt):
+        return (
+            "pa.Table",
+            normalize_token(dt.schema),
+            normalize_token(dt.columns),
+        )
+
+    @normalize_token.register(pa.ChunkedArray)
+    def normalize_chunked_array(arr):
+        return (
+            "pa.ChunkedArray",
+            normalize_token(arr.type),
+            normalize_token(arr.chunks),
+        )
+
+    @normalize_token.register(pa.Array)
+    def normalize_chunked_array(arr):
+        return (
+            "pa.Array",
+            normalize_token(arr.type),
+            normalize_token(arr.buffers()),
+        )
+
+    @normalize_token.register(pa.Buffer)
+    def normalize_chunked_array(buf):
+        return ("pa.Buffer", hash_buffer_hex(buf))
+
+
+@normalize_token.register_lazy("numpy")
+def register_numpy():
+    import numpy as np
+
+    @normalize_token.register(np.ndarray)
+    def normalize_array(x):
+        if not x.shape:
+            return (x.item(), x.dtype)
+        if x.dtype.hasobject:
+            try:
+                try:
+                    # string fast-path
+                    data = hash_buffer_hex(
+                        "-".join(x.flat).encode(
+                            encoding="utf-8", errors="surrogatepass"
+                        )
+                    )
+                except UnicodeDecodeError:
+                    # bytes fast-path
+                    data = hash_buffer_hex(b"-".join(x.flat))
+            except (TypeError, UnicodeDecodeError):
+                return normalize_object(x)
+        else:
+            try:
+                data = hash_buffer_hex(x.ravel(order="K").view("i1"))
+            except (BufferError, AttributeError, ValueError):
+                data = hash_buffer_hex(x.copy().ravel(order="K").view("i1"))
+        return (data, x.dtype, x.shape)
+
+    @normalize_token.register(np.memmap)
+    def normalize_mmap(mm):
+        return hash_buffer_hex(np.ascontiguousarray(mm))
+
+    @normalize_token.register(np.ufunc)
+    def normalize_ufunc(func):
+        try:
+            return _normalize_pickle(func)
+        except Exception:
+            _maybe_raise_nondeterministic(
+                f"Cannot tokenize numpy ufunc {func!r}. Please use functions "
+                "of the dask.array.ufunc module instead. See also "
+                "https://docs.dask.org/en/latest/array-numpy-compatibility.html"
+            )
+            return uuid.uuid4().hex


Test patch:
diff --git a/dask/tests/test_tokenize.py b/dask/tests/test_tokenize.py
index f93e8626e..87e9fa94a 100644
--- a/dask/tests/test_tokenize.py
+++ b/dask/tests/test_tokenize.py
@@ -18,8 +18,8 @@ import pytest
 from tlz import compose, curry, partial
 
 import dask
-from dask.base import TokenizationError, normalize_token, tokenize
 from dask.core import flatten, literal
+from dask.tokenize import TokenizationError, normalize_token, tokenize
 from dask.utils import tmpfile
 from dask.utils_test import import_or_none
 
@@ -37,7 +37,7 @@ def check_contextvars():
     """Test that tokenize() and normalize_token() properly clean up context
     variables at all times
     """
-    from dask.base import _ensure_deterministic, _seen
+    from dask.tokenize import _ensure_deterministic, _seen
 
     with pytest.raises(LookupError):
         _ensure_deterministic.get()