Repository: dask/dask
Pull request #11373 (2024-09-04 10:32:11)

Reduce overhead in tokenize
Builds on https://github.com/dask/dask/pull/11371

I noticed that tokenize was using some non-trivial overhead. That overhead primarily came from the contextlib contextmanager but also from the ContextVars themselves. I threw both away...

The bottom line is an overhead reduction by about 60-70%. Further speedup would be possible if we ditched md5 instead of stdlib hash (or another hash function). It's not the fastest and I guess we don't need a very good hash function here (but I don't know)

# main (py3.10)

```python
In [1]: from dask.base import tokenize
   ...: trivial_items = [ix for ix in range(10000)]

In [2]: %timeit tokenize(trivial_items)
1.95 ms ± 282 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [3]: %timeit [tokenize(it) for it in trivial_items]
43.5 ms ± 489 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)
```

# This PR

```python
In [1]: from dask.base import tokenize
   ...: trivial_items = [ix for ix in range(10000)]

In [2]: %timeit tokenize(trivial_items)
1.02 ms ± 9.43 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

In [3]: %timeit [tokenize(it) for it in trivial_items]
14.4 ms ± 142 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
```

I personally find this logic now also simpler and easier to reason about but this may be subjective. There is also a change in how memoization works. It now only memoizes collections if they are indeed recursive / self references. This may be slower in some (rare) situations but I believe it also avoids problems I was running into over in https://github.com/dask/dask/pull/11320

if performance for these edge cases are a problem we could choose to cache results instead (memoization remembers that a value was seen and is storing a dummy value, a cache would use the properly tokenized result instead of a dummy token)

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

     13 files  +     2       13 suites  +2   3h 0m 35s :stopwatch: + 40m 1s
 13 174 tests +     1   12 116 :white_check_mark:  -      1   1 056 :zzz: +    1   2 :x: + 1 
137 304 runs  +26 103  118 260 :white_check_mark: +22 794  19 018 :zzz: +3 285  26 :x: +24 

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

Results for commit 062f178e. ± Comparison against base commit 4be4880b.

[test-results]:data:application/gzip;base64,H4sIAN0+2GYC/03MTQ6DIBCG4asY1l3MDIjQyzSIkJD60yCsmt69WEW7fJ/JfG/mw+hWdm+Q3xq25pCuGnI0KSzz1qB4W6hc037HTtR+rNnaDQlRXvgMr99ne5k3YSxGJ7gYl1gEisQ879MdB3H0OY2KJJx4TGtAVa1Oywr/03aZppBKMJDksVNOaAtKk265x14qjZqkFIb63pEYLLDPF1Z/AUYbAQAA

phofl: this needs a rebase now :(

Patch:
diff --git a/continuous_integration/environment-mindeps-optional.yaml b/continuous_integration/environment-mindeps-optional.yaml
index efb422182..6681cc765 100644
--- a/continuous_integration/environment-mindeps-optional.yaml
+++ b/continuous_integration/environment-mindeps-optional.yaml
@@ -16,7 +16,7 @@ dependencies:
   - numpy=1.24
   - pandas=2.0
   # optional dependencies pulled in by pip install dask[diagnostics]
-  - bokeh=2.4.2
+  - bokeh=3.1.0
   - jinja2=2.10.3
   # optional dependencies pulled in by pip install dask[complete]
   - pyarrow=14.0.1
diff --git a/dask/_compatibility.py b/dask/_compatibility.py
index 87b94c1a6..6272469c1 100644
--- a/dask/_compatibility.py
+++ b/dask/_compatibility.py
@@ -34,7 +34,7 @@ def entry_points(group=None):
 VERSIONS = {
     "numpy": "1.21.0",
     "pandas": "2.0.0",
-    "bokeh": "2.4.2",
+    "bokeh": "3.1.0",
     "jinja2": "2.10.3",
     "pyarrow": "7.0.0",
     "lz4": "4.3.2",
diff --git a/dask/array/core.py b/dask/array/core.py
index 76ab569f4..d7e80038c 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -22,7 +22,7 @@ from typing import Any, Literal, TypeVar, Union, cast
 
 import numpy as np
 from numpy.typing import ArrayLike
-from tlz import accumulate, concat, first, frequencies, groupby, partition
+from tlz import accumulate, concat, first, groupby, partition
 from tlz.curried import pluck
 
 from dask import compute, config, core
@@ -3138,17 +3138,7 @@ def normalize_chunks(chunks, shape=None, limit=None, dtype=None, previous_chunks
     if chunks and shape is not None:
         # allints: did we start with chunks as a simple tuple of ints?
         allints = all(isinstance(c, int) for c in chunks)
-        chunks = sum(
-            (
-                (
-                    blockdims_from_blockshape((s,), (c,))
-                    if not isinstance(c, (tuple, list))
-                    else (c,)
-                )
-                for s, c in zip(shape, chunks)
-            ),
-            (),
-        )
+        chunks = _convert_int_chunk_to_tuple(shape, chunks)
     for c in chunks:
         if not c:
             raise ValueError(
@@ -3178,6 +3168,20 @@ def normalize_chunks(chunks, shape=None, limit=None, dtype=None, previous_chunks
     )
 
 
+def _convert_int_chunk_to_tuple(shape, chunks):
+    return sum(
+        (
+            (
+                blockdims_from_blockshape((s,), (c,))
+                if not isinstance(c, (tuple, list))
+                else (c,)
+            )
+            for s, c in zip(shape, chunks)
+        ),
+        (),
+    )
+
+
 def _compute_multiplier(limit: int, dtype, largest_block: int, result):
     """
     Utility function for auto_chunk, to fin how much larger or smaller the ideal
@@ -3187,7 +3191,7 @@ def _compute_multiplier(limit: int, dtype, largest_block: int, result):
         limit
         / dtype.itemsize
         / largest_block
-        / math.prod(r for r in result.values() if r)
+        / math.prod(max(r) if isinstance(r, tuple) else r for r in result.values() if r)
     )
 
 
@@ -3216,9 +3220,7 @@ def auto_chunks(chunks, shape, limit, dtype, previous_chunks=None):
     normalize_chunks: for full docstring and parameters
     """
     if previous_chunks is not None:
-        previous_chunks = tuple(
-            c if isinstance(c, tuple) else (c,) for c in previous_chunks
-        )
+        previous_chunks = _convert_int_chunk_to_tuple(shape, previous_chunks)
     chunks = list(chunks)
 
     autos = {i for i, c in enumerate(chunks) if c == "auto"}
@@ -3252,6 +3254,7 @@ def auto_chunks(chunks, shape, limit, dtype, previous_chunks=None):
             )
 
     limit = max(1, limit)
+    chunksize_tolerance = config.get("array.chunk-size-tolerance")
 
     largest_block = math.prod(
         cs if isinstance(cs, Number) else max(cs) for cs in chunks if cs != "auto"
@@ -3259,47 +3262,96 @@ def auto_chunks(chunks, shape, limit, dtype, previous_chunks=None):
 
     if previous_chunks:
         # Base ideal ratio on the median chunk size of the previous chunks
-        result = {a: np.median(previous_chunks[a]) for a in autos}
-
-        ideal_shape = []
-        for i, s in enumerate(shape):
-            chunk_frequencies = frequencies(previous_chunks[i])
-            mode, count = max(chunk_frequencies.items(), key=lambda kv: kv[1])
-            if mode > 1 and count >= len(previous_chunks[i]) / 2:
-                ideal_shape.append(mode)
-            else:
-                ideal_shape.append(s)
+        median_chunks = {a: np.median(previous_chunks[a]) for a in autos}
+        result = {}
 
         # How much larger or smaller the ideal chunk size is relative to what we have now
-        multiplier = _compute_multiplier(limit, dtype, largest_block, result)
-
-        last_multiplier = 0
-        last_autos = set()
-        while (
-            multiplier != last_multiplier or autos != last_autos
-        ):  # while things change
-            last_multiplier = multiplier  # record previous values
+        multiplier = _compute_multiplier(limit, dtype, largest_block, median_chunks)
+
+        def _trivial_aggregate(a):
+            autos.remove(a)
+            del median_chunks[a]
+            return True
+
+        multiplier_remaining = True
+        while multiplier_remaining:  # while things change
             last_autos = set(autos)  # record previous values
+            multiplier_remaining = False
 
             # Expand or contract each of the dimensions appropriately
             for a in sorted(autos):
-                if ideal_shape[a] == 0:
-                    result[a] = 0
-                    continue
-                proposed = result[a] * multiplier ** (1 / len(autos))
+                this_multiplier = multiplier ** (1 / len(last_autos))
+
+                proposed = median_chunks[a] * this_multiplier
+                this_chunksize_tolerance = chunksize_tolerance ** (1 / len(last_autos))
+                max_chunk_size = proposed * this_chunksize_tolerance
+
                 if proposed > shape[a]:  # we've hit the shape boundary
-                    autos.remove(a)
-                    largest_block *= shape[a]
                     chunks[a] = shape[a]
-                    del result[a]
+                    multiplier_remaining = _trivial_aggregate(a)
+                    largest_block *= shape[a]
+                    continue
+                elif this_multiplier <= 1:
+                    if max(previous_chunks[a]) == 1:
+                        multiplier_remaining = _trivial_aggregate(a)
+                        chunks[a] = tuple(previous_chunks[a])
+                        continue
+
+                    cache, dimension_result = {}, []
+                    for c in previous_chunks[a]:
+                        if c in cache:
+                            # We potentially split the same chunk over and over again
+                            # so cache that
+                            dimension_result.extend(cache[c])
+                            continue
+                        elif c <= max(max_chunk_size, 1):
+                            new_chunks = [c]
+                        else:
+                            new_chunks = _split_up_single_chunk(
+                                c, this_chunksize_tolerance, max_chunk_size, proposed
+                            )
+
+                        cache[c] = new_chunks
+                        dimension_result.extend(new_chunks)
+
+                    if max(dimension_result) == 1:
+                        # See if we can split up other axes further
+                        multiplier_remaining = _trivial_aggregate(a)
+
                 else:
-                    result[a] = round_to(proposed, ideal_shape[a])
+                    dimension_result, new_chunk = [], 0
+                    for c in previous_chunks[a]:
+                        if c + new_chunk <= proposed:
+                            # keep increasing the chunk
+                            new_chunk += c
+                        else:
+                            # We reach the boundary so start a new chunk
+                            dimension_result.append(new_chunk)
+                            new_chunk = c
+                            if new_chunk > max_chunk_size:
+                                # our new chunk is greater than the allowed, so split it up
+                                dimension_result.extend(
+                                    _split_up_single_chunk(
+                                        c,
+                                        this_chunksize_tolerance,
+                                        max_chunk_size,
+                                        proposed,
+                                    )
+                                )
+                                new_chunk = 0
+                    if new_chunk > 0:
+                        dimension_result.append(new_chunk)
+
+                result[a] = tuple(dimension_result)
 
             # recompute how much multiplier we have left, repeat
-            multiplier = _compute_multiplier(limit, dtype, largest_block, result)
+            if multiplier_remaining:
+                multiplier = _compute_multiplier(
+                    limit, dtype, largest_block, median_chunks
+                )
 
         for k, v in result.items():
-            chunks[k] = v
+            chunks[k] = v if v else 0
         return tuple(chunks)
 
     else:
@@ -3321,6 +3373,25 @@ def auto_chunks(chunks, shape, limit, dtype, previous_chunks=None):
         return tuple(chunks)
 
 
+def _split_up_single_chunk(
+    c: int, this_chunksize_tolerance: float, max_chunk_size: int, proposed: int
+) -> list[int]:
+    # Calculate by what factor we have to split this chunk
+    m = c / proposed
+    if math.ceil(m) / m > this_chunksize_tolerance:
+        # We want to smooth things potentially if rounding up would change the result
+        # by a lot
+        m = math.ceil(c / max_chunk_size)
+    else:
+        m = math.ceil(m)
+    # split the chunk
+    new_c, remainder = divmod(c, min(m, c))
+    x = [new_c] * min(m, c)
+    for i in range(remainder):
+        x[i] += 1
+    return x
+
+
 def round_to(c, s):
     """Return a chunk dimension that is close to an even multiple or factor
 
diff --git a/dask/base.py b/dask/base.py
index 6394ac28f..518c47d30 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -1256,11 +1256,11 @@ def clone_key(key: KeyOrStrT, seed: Hashable) -> KeyOrStrT:
 
     Examples
     --------
-    >>> clone_key("x", 123)
+    >>> clone_key("x", 123)  # doctest: +SKIP
     'x-c4fb64ccca807af85082413d7ef01721'
-    >>> clone_key("inc-cbb1eca3bafafbb3e8b2419c4eebb387", 123)
+    >>> clone_key("inc-cbb1eca3bafafbb3e8b2419c4eebb387", 123)  # doctest: +SKIP
     'inc-bc629c23014a4472e18b575fdaf29ee7'
-    >>> clone_key(("sum-cbb1eca3bafafbb3e8b2419c4eebb387", 4, 3), 123)
+    >>> clone_key(("sum-cbb1eca3bafafbb3e8b2419c4eebb387", 4, 3), 123)  # doctest: +SKIP
     ('sum-c053f3774e09bd0f7de6044dbc40e71d', 4, 3)
     """
     if isinstance(key, tuple) and key and isinstance(key[0], str):
diff --git a/dask/diagnostics/profile_visualize.py b/dask/diagnostics/profile_visualize.py
index d4b8ce9f6..32d3f37cc 100644
--- a/dask/diagnostics/profile_visualize.py
+++ b/dask/diagnostics/profile_visualize.py
@@ -211,17 +211,11 @@ def visualize(
             f.x_range = top.x_range
             f.title = None
             f.min_border_top = 20
-            if BOKEH_VERSION().major < 3:
-                f.plot_height -= 30
-            else:
-                f.height -= 30
+            f.height -= 30
         for f in figs[:-1]:
             f.xaxis.axis_label = None
             f.min_border_bottom = 20
-            if BOKEH_VERSION().major < 3:
-                f.plot_height -= 30
-            else:
-                f.height -= 30
+            f.height -= 30
         for f in figs:
             f.min_border_left = 75
             f.min_border_right = 75
diff --git a/dask/tokenize.py b/dask/tokenize.py
index 96c0b5846..3f646c5b8 100644
--- a/dask/tokenize.py
+++ b/dask/tokenize.py
@@ -8,12 +8,11 @@ import hashlib
 import inspect
 import pathlib
 import pickle
+import threading
 import types
 import uuid
 from collections import OrderedDict
-from collections.abc import Iterable, Iterator
-from contextlib import contextmanager
-from contextvars import ContextVar
+from collections.abc import Iterable
 from functools import partial
 
 import cloudpickle
@@ -30,12 +29,26 @@ class TokenizationError(RuntimeError):
     pass
 
 
+def _tokenize(*args: object, **kwargs: object) -> str:
+    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_lock = threading.RLock()
+_SEEN: dict[int, tuple[int, object]] = {}
+_ENSURE_DETERMINISTIC = None
+
+
 def tokenize(
     *args: object, ensure_deterministic: bool | None = None, **kwargs: object
 ) -> str:
     """Deterministic token
 
-    >>> tokenize([1, 2, '3'])
+    >>> tokenize([1, 2, '3'])  # doctest: +SKIP
     '06961e8de572e73c2e74b51348177918'
 
     >>> tokenize('Hello') == tokenize('Hello')
@@ -50,105 +63,60 @@ def tokenize(
         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:
+    global _SEEN, _ENSURE_DETERMINISTIC
+    with tokenize_lock:
+        seen_before, _SEEN = _SEEN, {}
+        _ENSURE_DETERMINISTIC = ensure_deterministic
         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)
+            return _tokenize(*args, **kwargs)
+        finally:
+            _SEEN = seen_before
+            _ENSURE_DETERMINISTIC = None
 
 
+def _maybe_raise_nondeterministic(msg: str) -> None:
+    if (
+        _ENSURE_DETERMINISTIC
+        or _ENSURE_DETERMINISTIC is None
+        and config.get("tokenize.ensure-deterministic")
+    ):
+        raise TokenizationError(msg)
+
+
+_IDENTITY_DISPATCH = (
+    int,
+    float,
+    str,
+    bytes,
+    type(None),
+    slice,
+    complex,
+    type(Ellipsis),
+    decimal.Decimal,
+    datetime.date,
+    datetime.time,
+    datetime.datetime,
+    datetime.timedelta,
+    pathlib.PurePath,
+)
 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_DISPATCH,
     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])))
+    if id(d) in _SEEN:
+        return "__seen", _SEEN[id(d)][0]
+    _SEEN[id(d)] = len(_SEEN), d
+    try:
+        return "dict", _normalize_seq_func(
+            sorted(d.items(), key=lambda kv: hash(kv[0]))
+        )
+    finally:
+        _SEEN.pop(id(d), None)
 
 
 @normalize_token.register(OrderedDict)
@@ -164,23 +132,20 @@ def normalize_set(s):
     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
+def _normalize_seq_func(seq: Iterable[object]) -> tuple[object, ...]:
+    def _inner_normalize_token(item):
+        # Don't go through Dispatch. That's slow
+        if isinstance(item, _IDENTITY_DISPATCH):
+            return item
+        return normalize_token(item)
+
+    if id(seq) in _SEEN:
+        return "__seen", _SEEN[id(seq)][0]
+    _SEEN[id(seq)] = len(_SEEN), seq
+    try:
+        return tuple(map(_inner_normalize_token, seq))
+    finally:
+        del _SEEN[id(seq)]
 
 
 @normalize_token.register((tuple, list))
diff --git a/docs/source/install.rst b/docs/source/install.rst
index 65e2728da..da138b5e6 100644
--- a/docs/source/install.rst
+++ b/docs/source/install.rst
@@ -104,7 +104,7 @@ These optional dependencies and their minimum supported versions are listed belo
 +------------------+-----------------+---------------------------------------------------------------------------------------------------------+
 | Dependency       | Version         | Description                                                                                             |
 +==================+=================+=========================================================================================================+
-| `bokeh`_         | ``>=2.4.2``     | Generate profiles of Dask execution (required for ``dask.diagnostics``)                                 |
+| `bokeh`_         | ``>=3.1.0``     | Generate profiles of Dask execution (required for ``dask.diagnostics``)                                 |
 +------------------+-----------------+---------------------------------------------------------------------------------------------------------+
 | `cachey`_        | ``>=0.1.1``     | Use caching for computation                                                                             |
 +------------------+-----------------+---------------------------------------------------------------------------------------------------------+
diff --git a/pyproject.toml b/pyproject.toml
index 2fa195567..715a9d807 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -56,7 +56,7 @@ dataframe = [
 ]
 distributed = ["distributed == 2024.8.2"]
 diagnostics = [
-    "bokeh >= 2.4.2",
+    "bokeh >= 3.1.0",
     "jinja2 >= 2.10.3",
 ]
 # keeping for backwards compatibility


Test patch:
diff --git a/dask/array/tests/test_rechunk.py b/dask/array/tests/test_rechunk.py
index 02eea953f..26a7e405e 100644
--- a/dask/array/tests/test_rechunk.py
+++ b/dask/array/tests/test_rechunk.py
@@ -825,7 +825,18 @@ def test_rechunk_avoid_needless_chunking():
         (100, 50, 10, (10,) * 10),
         (100, 100, 10, (10,) * 10),
         (20, 7, 10, (7, 7, 6)),
-        (20, (1, 1, 1, 1, 6, 2, 1, 7), 5, (5, 5, 5, 5)),
+        (
+            20,
+            (1, 1, 1, 1, 6, 2, 1, 7),
+            5,
+            (4, 6, 3, 4, 3),
+        ),  # 6 is in tolerance, 7 is not
+        (21, (1,) * 21, 10, (10, 10, 1)),  # ensure that we squash together properly
+        (20, ((2, 2, 2, 9, 1, 2, 2)), 5, (4, 2, 5, 4, 5)),
+        (20, ((10, 10)), 4, (4, 3, 3, 4, 3, 3)),
+        (21, ((10, 11)), 4, (4, 3, 3, 4, 4, 3)),
+        (20, ((1, 18, 1)), 5, (1, 5, 5, 4, 4, 1)),
+        (38, ((10, 18, 10)), 5, (5, 5, 5, 5, 4, 4, 5, 5)),
     ],
 )
 def test_rechunk_auto_1d(shape, chunks, bs, expected):
@@ -834,6 +845,26 @@ def test_rechunk_auto_1d(shape, chunks, bs, expected):
     assert y.chunks == (expected,)
 
 
+@pytest.mark.parametrize(
+    "previous_chunks,bs,expected",
+    [
+        (((1, 1, 1), (10, 10, 10, 10, 10, 10, 10, 10)), 160, ((3,), (50, 30))),
+        (((2, 2), (20,)), 5, ((1, 1, 1, 1), (5, 5, 5, 5))),
+        (((1, 1), (20,)), 5, ((1, 1), (5, 5, 5, 5))),
+    ],
+)
+def test_normalize_chunks_auto_2d(previous_chunks, bs, expected):
+    shape = tuple(map(sum, previous_chunks))
+    result = normalize_chunks(
+        {0: "auto", 1: "auto"},
+        shape,
+        limit=bs,
+        dtype=np.int8,
+        previous_chunks=previous_chunks,
+    )
+    assert result == expected
+
+
 def test_rechunk_auto_2d():
     x = da.ones((20, 20), chunks=(2, 2))
     y = x.rechunk({0: -1, 1: "auto"}, block_size_limit=20 * x.dtype.itemsize)
@@ -880,7 +911,12 @@ def test_rechunk_auto_image_stack(n):
     with dask.config.set({"array.chunk-size": "1MiB"}):
         x = da.ones((n, 1000, 1000), chunks=(1, 1000, 1000), dtype="float64")
         z = x.rechunk("auto")
-        assert z.chunks == ((1,) * n, (362, 362, 276), (362, 362, 276))
+        assert z.chunks == ((1,) * n, (334, 333, 333), (334, 333, 333))
+
+    with dask.config.set({"array.chunk-size": "1MiB"}):
+        x = da.ones((n, 2000, 2000), chunks=(1, 1000, 1000), dtype="float64")
+        z = x.rechunk("auto")
+        assert z.chunks == ((1,) * n, (334, 333, 333) * 2, (334, 333, 333) * 2)
 
 
 def test_rechunk_down():
@@ -891,14 +927,14 @@ def test_rechunk_down():
 
     with dask.config.set({"array.chunk-size": "1MiB"}):
         z = y.rechunk("auto")
-        assert z.chunks == ((4,) * 25, (511, 489), (511, 489))
+        assert z.chunks == ((5,) * 20, (500, 500), (500, 500))
 
     with dask.config.set({"array.chunk-size": "1MiB"}):
         z = y.rechunk({0: "auto"})
         assert z.chunks == ((1,) * 100, (1000,), (1000,))
 
         z = y.rechunk({1: "auto"})
-        assert z.chunks == ((10,) * 10, (104,) * 9 + (64,), (1000,))
+        assert z.chunks == ((10,) * 10, (100,) * 10, (1000,))
 
 
 def test_rechunk_zero():
diff --git a/dask/bag/tests/test_bag.py b/dask/bag/tests/test_bag.py
index 1554ebfd8..fe1b5ffc9 100644
--- a/dask/bag/tests/test_bag.py
+++ b/dask/bag/tests/test_bag.py
@@ -800,11 +800,9 @@ def test_product():
 def test_partition_collect():
     with partd.Pickle() as p:
         partition(identity, range(6), 3, p)
-        assert set(p.get(0)) == {3, 5}
-        assert set(p.get(1)) == {1}
-        assert set(p.get(2)) == {0, 2, 4}
-
-        assert sorted(collect(identity, 2, p, "")) == [(0, [0]), (2, [2]), (4, [4])]
+        for i in range(3):
+            assert p.get(i)
+            assert sorted(collect(identity, i, p, "")) == [(j, [j]) for j in p.get(i)]
 
 
 def test_groupby():
diff --git a/dask/tests/test_base.py b/dask/tests/test_base.py
index 2648b914e..919d9e987 100644
--- a/dask/tests/test_base.py
+++ b/dask/tests/test_base.py
@@ -34,10 +34,11 @@ from dask.base import (
     unpack_collections,
     visualize,
 )
+from dask.core import validate_key
 from dask.delayed import Delayed, delayed
 from dask.diagnostics import Profiler
 from dask.highlevelgraph import HighLevelGraph
-from dask.utils import tmpdir, tmpfile
+from dask.utils import key_split, tmpdir, tmpfile
 from dask.utils_test import dec, import_or_none, inc
 
 da = import_or_none("dask.array")
@@ -1011,15 +1012,13 @@ def test_optimizations_ctd():
 
 
 def test_clone_key():
-    assert clone_key("inc-1-2-3", 123) == "inc-73db79fdf4518507ddc84796726d4844"
-    assert clone_key("x", 123) == "x-c4fb64ccca807af85082413d7ef01721"
-    assert clone_key("x", 456) == "x-d4b538b4d4cf68fca214077609feebae"
-    assert clone_key(("x", 1), 456) == ("x-d4b538b4d4cf68fca214077609feebae", 1)
-    assert clone_key(("sum-1-2-3", h1, 1), 123) == (
-        "sum-822e7622aa1262cef988b3033c32aa37",
-        h1,
-        1,
-    )
+    for key, seed in [("x", 123), (("x", 1), 456), (("sum-1-2-3", h1, 1), 123)]:
+        validate_key(clone_key(key, seed))
+        assert clone_key(key, seed) != key
+        assert clone_key(key, seed) == clone_key(key, seed)
+        assert clone_key(key, seed) != clone_key(key, seed + 1)
+        assert key_split(clone_key(key, seed)) == key_split(key)
+
     with pytest.raises(TypeError):
         clone_key(1, 2)
 
diff --git a/dask/tests/test_tokenize.py b/dask/tests/test_tokenize.py
index 87e9fa94a..a5dde0b58 100644
--- a/dask/tests/test_tokenize.py
+++ b/dask/tests/test_tokenize.py
@@ -33,21 +33,15 @@ numba = import_or_none("numba")
 
 
 @pytest.fixture(autouse=True)
-def check_contextvars():
-    """Test that tokenize() and normalize_token() properly clean up context
-    variables at all times
-    """
-    from dask.tokenize import _ensure_deterministic, _seen
+def check_clean_state():
+    """Test that tokenize() and normalize_token() properly clean up state"""
+    from dask.tokenize import _ENSURE_DETERMINISTIC, _SEEN
 
-    with pytest.raises(LookupError):
-        _ensure_deterministic.get()
-    with pytest.raises(LookupError):
-        _seen.get()
+    assert not _SEEN
+    assert _ENSURE_DETERMINISTIC is None
     yield
-    with pytest.raises(LookupError):
-        _ensure_deterministic.get()
-    with pytest.raises(LookupError):
-        _seen.get()
+    assert not _SEEN
+    assert _ENSURE_DETERMINISTIC is None
 
 
 def check_tokenize(*args, **kwargs):
@@ -820,17 +814,9 @@ def test_tokenize_sequences():
     assert check_tokenize([1]) == check_tokenize([1])
 
     # You can call normalize_token directly.
-    # Repeated objects are memoized.
     x = (1, 2)
     y = [x, x, [x, (2, 3)]]
-    assert normalize_token(y) == (
-        "list",
-        [
-            ("tuple", [1, 2]),
-            ("__seen", 0),
-            ("list", [("__seen", 0), ("tuple", [2, 3])]),
-        ],
-    )
+    assert normalize_token(y)
 
 
 def test_nested_tokenize_seen():
@@ -898,7 +884,6 @@ def test_tokenize_sorts_dict_before_seen_map():
     v = (1, 2, 3)
     d1 = {1: v, 2: v}
     d2 = {2: v, 1: v}
-    assert "__seen" in str(normalize_token(d1))
     assert check_tokenize(d1) == check_tokenize(d2)
 
 
@@ -911,7 +896,6 @@ def test_tokenize_sorts_set_before_seen_map():
     v = (1, 2, 3)
     s1 = {(i, v) for i in range(100)}
     s2 = {(i, v) for i in reversed(range(100))}
-    assert "__seen" in str(normalize_token(s1))
     assert check_tokenize(s1) == check_tokenize(s2)