Repository: dask/dask
Pull request #10904 (2024-02-06 19:12:08)

Tweak sequence tokenization
Allow calling normalize_token() outside of tokenize(). This is useful for debugging purposes.
Additionally, polish some tests in test_tokenize.


Patch:
diff --git a/dask/base.py b/dask/base.py
index ae4745e29..a82ee50dd 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -9,10 +9,11 @@ import pathlib
 import pickle
 import random
 import threading
+import types
 import uuid
 import warnings
 from collections import OrderedDict
-from collections.abc import Callable, Hashable, Iterator, Mapping
+from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping
 from concurrent.futures import Executor
 from contextlib import contextmanager, suppress
 from contextvars import ContextVar
@@ -20,7 +21,6 @@ from enum import Enum
 from functools import partial
 from numbers import Integral, Number
 from operator import getitem
-from types import MappingProxyType
 from typing import Any, Literal, TypeVar
 
 from tlz import curry, groupby, identity, merge
@@ -1060,8 +1060,7 @@ def normalize_type(typ):
     return typ.__name__, typ.__module__
 
 
-@normalize_token.register(MappingProxyType)
-@normalize_token.register(dict)
+@normalize_token.register((dict, types.MappingProxyType))
 def normalize_dict(d):
     return "dict", _normalize_seq_func(sorted(d.items(), key=str))
 
@@ -1073,6 +1072,9 @@ def normalize_ordered_dict(d):
 
 @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))
 
 
@@ -1088,24 +1090,36 @@ def normalize_set(s):
 _seen: ContextVar[dict[int, tuple[int, object]]] = ContextVar("_seen")
 
 
-def _normalize_seq_func(seq):
-    seen = _seen.get()
-    out = []
+def _normalize_seq_func(seq: Iterable[object]) -> list[object]:
+    try:
+        seen = _seen.get()
+        tok = None
+    except LookupError:
+        # This is for debug only, for when normalize_token is called outside of
+        # tokenize(). It is important to reset the token on tokenize to avoid artifacts
+        # when tokenize() is called recursively.
+        seen = {}
+        tok = _seen.set(seen)
 
-    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.
-            norm_item = item
-        elif id(item) in seen:
-            # May or may not be a circular recursion. Maybe just a double reference.
-            seen_when, _ = seen[id(item)]
-            norm_item = "__seen", seen_when
-        else:
-            seen[id(item)] = len(seen), item
-            norm_item = normalize_token(item)
-        out.append(norm_item)
-    return out
+    try:
+        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
+    finally:
+        if tok is not None:
+            _seen.reset(tok)
 
 
 @normalize_token.register((tuple, list))
@@ -1183,9 +1197,6 @@ def normalize_callable(func: Callable) -> Callable | tuple | str | bytes:
         return _normalize_callable(func)
 
 
-normalize_function = normalize_callable
-
-
 def _normalize_callable(func: Callable) -> tuple | str | bytes:
     if isinstance(func, Compose):
         first = getattr(func, "first", None)
@@ -1366,7 +1377,7 @@ def register_numpy():
                 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)
+        return data, x.dtype, x.shape
 
     @normalize_token.register(np.matrix)
     def normalize_matrix(x):
diff --git a/docs/source/deploying.rst b/docs/source/deploying.rst
index 05f7a02c2..4d65e48cd 100644
--- a/docs/source/deploying.rst
+++ b/docs/source/deploying.rst
@@ -170,7 +170,7 @@ deployments this is often handled by other Kubernetes services.
        from dask_kubernetes.operator import KubeCluster
        cluster = KubeCluster(
            name="my-dask-cluster",
-           image='ghcr.io/dask/dask:latest',
+           image="ghcr.io/dask/dask:latest",
            resources={"requests": {"memory": "2Gi"}, "limits": {"memory": "64Gi"}},
        )
        cluster.scale(10)
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 6056b1ce2..7b611cf2e 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -206,7 +206,7 @@ Dask provides several APIs.  Choose one that works best for you:
 How to Install Dask
 -------------------
 
-Installing Dask is easy with ``pip`` or ``conda``
+Installing Dask is easy with ``pip`` or ``conda``.
 
 Learn more at :bdg-link-primary:`Install Documentation <install.html>`
 
@@ -227,7 +227,7 @@ Learn more at :bdg-link-primary:`Install Documentation <install.html>`
 How to Deploy Dask
 ------------------
 
-You can then use Dask on a single machine, or deploy it on distributed hardware
+You can use Dask on a single machine, or deploy it on distributed hardware.
 
 Learn more at :bdg-link-primary:`Deploy Documentation <deploying.html>`
 
@@ -249,97 +249,64 @@ Learn more at :bdg-link-primary:`Deploy Documentation <deploying.html>`
        Alternatively, you can skip this part, and Dask will operate within a
        thread pool contained entirely with your local process.
 
-    .. tab-item:: Kubernetes
-
-       The `dask-kubernetes project <https://kubernetes.dask.org>`_ provides
-       a Dask Kubernetes Operator.
-
-       .. code-block:: python
-
-          from dask_kubernetes.operator import KubeCluster
-          cluster = KubeCluster(
-             name="my-dask-cluster",
-             image='ghcr.io/dask/dask:latest'
-          )
-          cluster.scale(10)
-
-       Learn more at :bdg-link-primary:`Dask Kubernetes Documentation <https://kubernetes.dask.org>`
-
-    .. tab-item:: HPC
-
-        The `dask-jobqueue project <https://jobqueue.dask.org>`_ interfaces
-        with popular job submission projects, like SLURM, PBS, SGE, LSF,
-        Torque, Condor, and others.
-
-
-        .. code-block:: python
-
-           from dask_jobqueue import SLURMCluster
-
-           cluster = SLURMCluster()
-           cluster.scale(jobs=10)
-
+    .. tab-item:: Cloud
 
-        You can also deploy Dask with MPI
+        `Coiled <https://docs.coiled.io/user_guide/index.html?utm_source=dask-docs&utm_medium=homepage>`_
+        is a commercial SaaS product that deploys Dask clusters on cloud platforms like AWS, GCP, and Azure.
 
         .. code-block:: python
 
-           # myscript.py
-           from dask_mpi import initialize
-           initialize()
-
-           from dask.distributed import Client
-           client = Client()  # Connect this local process to remote workers
-
-        .. code-block::
+            import coiled
+            cluster = coiled.Cluster(
+               n_workers=100,
+               region="us-east-2",
+               worker_memory="16 GiB",
+               spot_policy="spot_with_fallback",
+            )
+            client = cluster.get_client()
 
-           $ mpirun -np 4 python myscript.py
+        Learn more at :bdg-link-primary:`Coiled Documentation <https://docs.coiled.io/user_guide/index.html?utm_source=dask-docs&utm_medium=homepage>`
 
-        Learn more at :bdg-link-primary:`Dask Jobqueue Documentation <https://jobqueue.dask.org>` and the :bdg-link-primary:`Dask MPI Documentation <https://mpi.dask.org>`.
 
-    .. tab-item:: Cloud
+    .. tab-item:: HPC
 
-        The `dask-cloudprovider project <https://cloudprovider.dask.org>`_ interfaces
-        with popular cloud platforms like AWS, GCP, Azure, and Digital Ocean.
+        The `Dask-Jobqueue project <https://jobqueue.dask.org>`_ deploys
+        Dask clusters on popular HPC job submission systems like SLURM, PBS, SGE, LSF,
+        Torque, Condor, and others.
 
         .. code-block:: python
 
-           from dask_cloudprovider.aws import FargateCluster
-           cluster = FargateCluster(
-               # Cluster manager specific config kwargs
-           )
-
-        Learn more at :bdg-link-primary:`Dask CloudProvider Documentation <https://cloudprovider.dask.org>`
-
-    .. tab-item:: Managed/Commercial
+            from dask_jobqueue import PBSCluster
+            cluster = PBSCluster(
+               cores=24,
+               memory="100GB",
+               queue="regular",
+               account="my-account",
+            )
+            cluster.scale(jobs=100)
+            client = cluster.get_client()
 
-        Several companies offer commercial Dask products.  These are not open
-        source, but tend to be easier, safer, cheaper, more fully featured,
-        etc..  All options here include solid free offerings for individuals.
+        Learn more at :bdg-link-primary:`Dask-Jobqueue Documentation <https://jobqueue.dask.org>`
 
-        .. _Coiled: https://coiled.io/?utm_source=dask-docs&utm_medium=homepage
-        .. |coiled| replace:: **Coiled**
-        .. _saturn: https://saturncloud.io
-        .. |saturn| replace:: **Saturn Cloud**
-        .. _nebari: https://nebari.dev
-        .. |nebari| replace:: **Nebari**
+    .. tab-item:: Kubernetes
 
-        -  |Coiled|_ provides a standalone Dask deployment product that works
-           in AWS and GCP.
+       The `Dask Kubernetes project <https://kubernetes.dask.org>`_ provides
+       a Dask Kubernetes Operator for deploying Dask on Kubernetes clusters.
 
-           Coiled notably employs many of the active Dask maintainers today.
+       .. code-block:: python
 
-           Learn more at :bdg-link-primary:`Coiled <https://coiled.io/?utm_source=dask-docs&utm_medium=homepage>`
+            from dask_kubernetes.operator import KubeCluster
+            cluster = KubeCluster(
+               name="my-dask-cluster",
+               image="ghcr.io/dask/dask:latest",
+               resources={"requests": {"memory": "2Gi"}, "limits": {"memory": "64Gi"}},
+            )
+            cluster.scale(10)
+            client = cluster.get_client()
 
-        -  |saturn|_ provides Dask as part of their hosted platform
-           including Jupyter and other products.
 
-           Learn more at :bdg-link-primary:`Saturn Cloud <https://saturncloud.io>`
-        -  |nebari|_ from Quansight provides Dask as part of a Kubernetes-based
-           git-ops managed platform along with Jupyter and other products
-           suitable for on-prem deployments.
+       Learn more at :bdg-link-primary:`Dask Kubernetes Documentation <https://kubernetes.dask.org>`
 
-           Learn more at :bdg-link-primary:`Nebari <https://nebari.dev>`
 
 Learn with Examples
 -------------------


Test patch:
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
index 25a1c5b43..06e3abdcf 100644
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -1828,7 +1828,9 @@ def test_assign():
 
     # divisions unknown won't work with pandas
     with pytest.raises(ValueError):
-        ddf_unknown.assign(c=df.a + 1)
+        q = ddf_unknown.assign(c=df.a + 1)
+        if DASK_EXPR_ENABLED:
+            q.optimize()
 
     # unsupported type
     with pytest.raises(TypeError):
@@ -1836,10 +1838,14 @@ def test_assign():
 
     # Fails when assigning known divisions to unknown divisions
     with pytest.raises(ValueError):
-        ddf_unknown.assign(foo=ddf.a)
+        q = ddf_unknown.assign(foo=ddf.a)
+        if DASK_EXPR_ENABLED:
+            q.optimize()
     # Fails when assigning unknown divisions to known divisions
     with pytest.raises(ValueError):
-        ddf.assign(foo=ddf_unknown.a)
+        q = ddf.assign(foo=ddf_unknown.a)
+        if DASK_EXPR_ENABLED:
+            q.optimize()
 
     df = pd.DataFrame({"A": [1, 2]})
     df.assign(B=lambda df: df["A"], C=lambda df: df.A + df.B)
@@ -5148,8 +5154,9 @@ def test_to_datetime(gpu):
             )
         else:
             ctx = contextlib.nullcontext()
+        ctx_expected = contextlib.nullcontext() if gpu else ctx
 
-        with ctx:
+        with ctx_expected:
             expected = xd.to_datetime(s, infer_datetime_format=True)
         with ctx:
             result = dd.to_datetime(ds, infer_datetime_format=True)
@@ -5158,7 +5165,7 @@ def test_to_datetime(gpu):
             result = dd.to_datetime(s, infer_datetime_format=True)
         assert_eq(expected, result, check_dtype=check_dtype)
 
-        with ctx:
+        with ctx_expected:
             expected = xd.to_datetime(s.index, infer_datetime_format=True)
         with ctx:
             result = dd.to_datetime(ds.index, infer_datetime_format=True)
diff --git a/dask/tests/test_tokenize.py b/dask/tests/test_tokenize.py
index 59d9a2c6a..aff284103 100644
--- a/dask/tests/test_tokenize.py
+++ b/dask/tests/test_tokenize.py
@@ -18,7 +18,7 @@ import pytest
 from tlz import compose, curry, partial
 
 import dask
-from dask.base import function_cache, normalize_function, normalize_token, tokenize
+from dask.base import function_cache, normalize_callable, normalize_token, tokenize
 from dask.core import literal
 from dask.utils import tmpfile
 from dask.utils_test import import_or_none
@@ -31,39 +31,6 @@ pa = import_or_none("pyarrow")
 pd = import_or_none("pandas")
 
 
-def f1(a, b, c=1):
-    pass
-
-
-def f2(a, b=1, c=2):
-    pass
-
-
-def f3(a):
-    pass
-
-
-def test_normalize_function():
-    assert normalize_function(f2)
-
-    assert normalize_function(lambda a: a)
-
-    assert normalize_function(partial(f2, b=2)) == normalize_function(partial(f2, b=2))
-
-    assert normalize_function(partial(f2, b=2)) != normalize_function(partial(f2, b=3))
-
-    assert normalize_function(partial(f1, b=2)) != normalize_function(partial(f2, b=2))
-
-    assert normalize_function(compose(f2, f3)) == normalize_function(compose(f2, f3))
-
-    assert normalize_function(compose(f2, f3)) != normalize_function(compose(f2, f1))
-
-    assert normalize_function(curry(f2)) == normalize_function(curry(f2))
-    assert normalize_function(curry(f2)) != normalize_function(curry(f1))
-    assert normalize_function(curry(f2, b=1)) == normalize_function(curry(f2, b=1))
-    assert normalize_function(curry(f2, b=1)) != normalize_function(curry(f2, b=2))
-
-
 def _clear_function_cache():
     from dask.base import function_cache, function_cache_lock
 
@@ -129,7 +96,7 @@ def test_check_tokenize():
         def __init__(self):
             self.tok = random.random()
 
-        def __reduce_ex__(self, protocol):
+        def __reduce__(self):
             return A, ()
 
         def __dask_tokenize__(self):
@@ -381,12 +348,28 @@ def test_tokenize_local_functions():
     assert len(set(tokens)) == len(all_funcs)
 
 
-def my_func(a, b, c=1):
-    return a + b + c
+def f1(a, b, c=1):
+    pass
+
+
+def f2(a, b=1, c=2):
+    pass
+
+
+def f3(a):
+    pass
 
 
 def test_tokenize_callable():
-    check_tokenize(my_func)
+    assert check_tokenize(f1) != check_tokenize(f2)
+
+
+def test_tokenize_composite_functions():
+    assert check_tokenize(partial(f2, b=2)) != check_tokenize(partial(f2, b=3))
+    assert check_tokenize(partial(f1, b=2)) != check_tokenize(partial(f2, b=2))
+    assert check_tokenize(compose(f2, f3)) != check_tokenize(compose(f2, f1))
+    assert check_tokenize(curry(f2)) != check_tokenize(curry(f1))
+    assert check_tokenize(curry(f2, b=1)) != check_tokenize(curry(f2, b=2))
 
 
 @pytest.mark.skipif("not pd")
@@ -572,25 +555,35 @@ def test_staticmethods():
     assert check_tokenize(a.class_method) != check_tokenize(c.class_method)
 
 
-@pytest.mark.skipif("not np")
 def test_tokenize_sequences():
     assert check_tokenize([1]) != check_tokenize([2])
     assert check_tokenize([1]) != check_tokenize((1,))
     assert check_tokenize([1]) == check_tokenize([1])
 
-    x = np.arange(2000)  # long enough to drop information in repr
-    y = np.arange(2000)
-    y[1000] = 0  # middle isn't printed in repr
-    assert check_tokenize([x]) != check_tokenize([y])
+    # 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])]),
+        ],
+    )
 
 
 def test_tokenize_dict():
-    assert check_tokenize({"x": 1, 1: "x"}) == check_tokenize({"x": 1, 1: "x"})
+    # Insertion order is ignored. Keys can be an unsortable mix of types.
+    assert check_tokenize({"x": 1, 1: "x"}) == check_tokenize({1: "x", "x": 1})
+    assert check_tokenize({"x": 1, 1: "x"}) != check_tokenize({"x": 1, 2: "x"})
+    assert check_tokenize({"x": 1, 1: "x"}) != check_tokenize({"x": 2, 1: "x"})
 
 
 def test_tokenize_set():
     assert check_tokenize({1, 2, "x", (1, "x")}) == check_tokenize(
-        {1, 2, "x", (1, "x")}
+        {2, "x", (1, "x"), 1}
     )
 
 
@@ -685,6 +678,14 @@ def test_tokenize_range(other):
     assert check_tokenize(range(5, 10, 2)) != check_tokenize(range(*other))
 
 
+@pytest.mark.skipif("not np")
+def test_tokenize_numpy_array():
+    x = np.arange(2000)  # long enough to drop information in repr
+    y = np.arange(2000)
+    y[1000] = 0  # middle isn't printed in repr
+    assert check_tokenize([x]) != check_tokenize([y])
+
+
 @pytest.mark.skipif("not np")
 def test_tokenize_object_array_with_nans():
     a = np.array(["foo", "Jos\xe9", np.nan], dtype="O")
@@ -749,10 +750,7 @@ def test_tokenize_circular_recursion():
     b = [1, 3]
     b[0] = b
 
-    def copy(x):
-        return pickle.loads(pickle.dumps(x))
-
-    assert check_tokenize(a, copy=copy) != check_tokenize(b, copy=copy)
+    assert check_tokenize(a) != check_tokenize(b)
 
     # Different circular recursions tokenize differently
     c = [[], []]
@@ -762,12 +760,12 @@ def test_tokenize_circular_recursion():
     d = [[], []]
     d[0].append(d[1])
     d[1].append(d[0])
-    assert check_tokenize(c, copy=copy) != check_tokenize(d, copy=copy)
+    assert check_tokenize(c) != check_tokenize(d)
 
     # For dicts, the dict itself is not passed to _normalize_seq_func
     e = {}
     e[0] = e
-    check_tokenize(e, copy=copy)
+    check_tokenize(e)
 
 
 @pytest.mark.parametrize(
@@ -898,7 +896,7 @@ def test_tokenize_functions_main():
 def test_normalize_function_limited_size():
     _clear_function_cache()
     for _ in range(1000):
-        normalize_function(lambda x: x)
+        normalize_callable(lambda x: x)
     assert 50 < len(function_cache) < 600
 
 
@@ -984,13 +982,11 @@ def test_tokenize_random_functions_with_state_numpy():
 def test_tokenize_pyarrow_datatypes_simple():
     a = pa.int64()
     b = pa.float64()
-    assert tokenize(a) == tokenize(a)
-    assert tokenize(a) != tokenize(b)
+    assert check_tokenize(a) != check_tokenize(b)
 
 
 @pytest.mark.skipif("not pa")
 def test_tokenize_pyarrow_datatypes_complex():
     a = pa.struct({"x": pa.int32(), "y": pa.string()})
     b = pa.struct({"x": pa.float64(), "y": pa.int16()})
-    assert tokenize(a) == tokenize(a)
-    assert tokenize(a) != tokenize(b)
+    assert check_tokenize(a) != check_tokenize(b)