Repository: dask/dask
Pull request #10913 (2024-02-09 11:47:27)

Override tokenize.ensure-deterministic config flag
- Add optional parameter to `tokenize(ensure_deterministic=True)` to force raising exception
- Can now run `test_tokenize.py` with either value for the config variable
- Subclass RuntimeError to restrict exception catching (we're calling pickle internally, which is very complex)
- Better tests for the `_seen` variable


crusaderky: @hendrikmakait @fjetter all comments have been addressed

Patch:
diff --git a/dask/base.py b/dask/base.py
index f32095d55..aaeec20e4 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -1005,9 +1005,13 @@ def persist(*args, traverse=True, optimize_graph=True, scheduler=None, **kwargs)
 ############
 # Tokenize #
 ############
+class TokenizationError(RuntimeError):
+    pass
 
 
-def tokenize(*args, **kwargs):
+def tokenize(
+    *args: object, ensure_deterministic: bool | None = None, **kwargs: object
+) -> str:
     """Deterministic token
 
     >>> tokenize([1, 2, '3'])
@@ -1015,19 +1019,89 @@ def tokenize(*args, **kwargs):
 
     >>> 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.
     """
-    tok = _seen.set({})
-    try:
-        token = _normalize_seq_func(args)
+    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()))
-    finally:
-        _seen.reset(tok)
 
     # 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(
     (
@@ -1068,30 +1142,8 @@ def normalize_set(s):
     return "set", _normalize_seq_func(sorted(s, key=str))
 
 
-# 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")
-
-
 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)
-
-    try:
+    with _seen_ctx(reset=False) as seen:
         out = []
         for item in seq:
             if isinstance(item, (str, bytes, int, float, bool, type(None))):
@@ -1107,9 +1159,6 @@ def _normalize_seq_func(seq: Iterable[object]) -> list[object]:
                 item = normalize_token(item)
             out.append(item)
         return out
-    finally:
-        if tok is not None:
-            _seen.reset(tok)
 
 
 @normalize_token.register((tuple, list))
@@ -1147,9 +1196,6 @@ def normalize_builtin_function_or_method(func):
         return normalize_object(func)
 
 
-_seen_objects = set()
-
-
 @normalize_token.register(object)
 def normalize_object(o):
     method = getattr(o, "__dask_tokenize__", None)
@@ -1157,15 +1203,7 @@ def normalize_object(o):
         return method()
 
     if type(o) is object:
-        if config.get("tokenize.ensure-deterministic"):
-            raise RuntimeError(
-                "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)
+        return _normalize_pure_object(o)
 
     if dataclasses.is_dataclass(o) and not isinstance(o, type):
         return _normalize_dataclass(o)
@@ -1173,15 +1211,28 @@ def normalize_object(o):
     try:
         return _normalize_pickle(o)
     except Exception:
-        if config.get("tokenize.ensure-deterministic"):
-            raise RuntimeError(
-                f"Object {o!r} cannot be deterministically hashed. See "
-                "https://docs.dask.org/en/latest/custom-collections.html#implementing-deterministic-hashing "
-                "for more information."
-            )
+        _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: bytes | None
@@ -1303,12 +1354,11 @@ def register_numpy():
         try:
             return _normalize_pickle(func)
         except Exception:
-            if config.get("tokenize.ensure-deterministic"):
-                raise RuntimeError(
-                    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"
-                )
+            _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
 
 
diff --git a/dask/dataframe/groupby.py b/dask/dataframe/groupby.py
index b85102336..6dc737e0c 100644
--- a/dask/dataframe/groupby.py
+++ b/dask/dataframe/groupby.py
@@ -745,6 +745,12 @@ def _cov_agg(_t, levels, ddof, std=False, sort=False):
 
 def _nunique_df_chunk(df, *by, **kwargs):
     name = kwargs.pop("name")
+    try:
+        # This is a lot faster but kind of a pain to implement when by
+        # has a boolean series in it.
+        return df.drop_duplicates(subset=list(by) + [name]).set_index(list(by))
+    except Exception:
+        pass
     group_keys = {}
     if PANDAS_GE_150:
         group_keys["group_keys"] = True


Test patch:
diff --git a/dask/tests/test_tokenize.py b/dask/tests/test_tokenize.py
index adce0573a..a7f6060bd 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 normalize_token, tokenize
+from dask.base import TokenizationError, normalize_token, tokenize
 from dask.core import literal
 from dask.utils import tmpfile
 from dask.utils_test import import_or_none
@@ -32,6 +32,24 @@ pd = import_or_none("pandas")
 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.base import _ensure_deterministic, _seen
+
+    with pytest.raises(LookupError):
+        _ensure_deterministic.get()
+    with pytest.raises(LookupError):
+        _seen.get()
+    yield
+    with pytest.raises(LookupError):
+        _ensure_deterministic.get()
+    with pytest.raises(LookupError):
+        _seen.get()
+
+
 def check_tokenize(*args, **kwargs):
     with dask.config.set({"tokenize.ensure-deterministic": True}):
         before = tokenize(*args, **kwargs)
@@ -264,13 +282,24 @@ def test_normalize_numpy_ufunc_unserializable():
     # Make a ufunc that isn't in the numpy namespace and can't be serialized
     inc = np.frompyfunc(lambda x: x + 1, 1, 1)
     with dask.config.set({"tokenize.ensure-deterministic": False}):
+        # Not idempotent
         assert tokenize(inc) != tokenize(inc)
+        # You can call normalize_token directly
+        assert normalize_token(inc) != normalize_token(inc)
+
     with dask.config.set({"tokenize.ensure-deterministic": True}):
         with pytest.raises(
-            RuntimeError, match=r"Cannot tokenize.*dask\.array\.ufunc.*instead"
+            TokenizationError, match=r"Cannot tokenize.*dask\.array\.ufunc.*instead"
         ):
             tokenize(inc)
 
+    # Test env override
+    assert tokenize(inc, ensure_deterministic=False) != tokenize(
+        inc, ensure_deterministic=False
+    )
+    with pytest.raises(TokenizationError, match="Cannot tokenize"):
+        tokenize(inc, ensure_deterministic=True)
+
 
 def test_normalize_object_unserializable():
     class C:
@@ -280,11 +309,24 @@ def test_normalize_object_unserializable():
     c = C()
 
     with dask.config.set({"tokenize.ensure-deterministic": False}):
+        # Not idempotent
         assert tokenize(c) != tokenize(c)
+        # You can call normalize_token directly
+        assert normalize_token(c) != normalize_token(c)
+
     with dask.config.set({"tokenize.ensure-deterministic": True}):
-        with pytest.raises(RuntimeError, match="cannot be deterministically hashed"):
+        with pytest.raises(
+            TokenizationError, match="cannot be deterministically hashed"
+        ):
             tokenize(c)
 
+    # Test env override
+    assert tokenize(c, ensure_deterministic=False) != tokenize(
+        c, ensure_deterministic=False
+    )
+    with pytest.raises(TokenizationError, match="cannot be deterministically hashed"):
+        tokenize(c, ensure_deterministic=True)
+
 
 def test_tokenize_partial_func_args_kwargs_consistent():
     f = partial(f3, f2, c=f1)
@@ -308,18 +350,50 @@ def test_normalize_base():
 
 
 def test_tokenize_object():
-    # object() tokenization is idempotent
-    o = object()
-    assert tokenize(o) == tokenize(o)
+    with dask.config.set({"tokenize.ensure-deterministic": False}):
+        # object() tokenization is idempotent...
+        o = object()
+        assert tokenize(o) == tokenize(o)
+        # ...but not deterministic
+        assert tokenize(object()) != tokenize(object())
+
+        # Two objects don't tokenize to the same token even if their pickle output is
+        # identical. Stress id reuse by creating and dereferencing many objects in quick
+        # succession.
+        assert len({tokenize(object()) for _ in range(100)}) == 100
+
+        # You can call normalize_token even if the _ensure_deterministic context
+        # variable hasn't been set
+        assert normalize_token(o) == normalize_token(o)
 
     with dask.config.set({"tokenize.ensure-deterministic": True}):
-        with pytest.raises(RuntimeError, match="deterministic"):
+        with pytest.raises(TokenizationError, match="deterministic"):
             tokenize(o)
+        with pytest.raises(TokenizationError, match="deterministic"):
+            normalize_token(o)
+
+    # Test env override
+    assert tokenize(o, ensure_deterministic=False) == tokenize(
+        o, ensure_deterministic=False
+    )
+    with pytest.raises(TokenizationError, match="deterministic"):
+        tokenize(o, ensure_deterministic=True)
 
-    # Two objects don't tokenize to the same token even if their pickle output is
-    # identical. Stress id reuse by creating and dereferencing many objects in quick
-    # succession.
-    assert len({tokenize(object()) for _ in range(100)}) == 100
+
+def nested_tokenize_ensure_deterministic():
+    """Test that the ensure_deterministic override is not lost if tokenize() is
+    called recursively
+    """
+
+    class C:
+        def __dask_tokenize__(self):
+            return tokenize(object())
+
+    assert tokenize(C(), ensure_deterministic=False) != tokenize(
+        C(), ensure_deterministic=False
+    )
+    with pytest.raises(TokenizationError):
+        tokenize(C())
 
 
 _GLOBAL = 1
@@ -757,6 +831,28 @@ def test_tokenize_sequences():
     )
 
 
+def test_nested_tokenize_seen():
+    """Test that calling tokenize() recursively doesn't alter the output due to
+    memoization of already-seen objects
+    """
+    o = [1, 2, 3]
+
+    class C:
+        def __init__(self, x):
+            self.x = x
+            self.tok = None
+
+        def __dask_tokenize__(self):
+            if not self.tok:
+                self.tok = tokenize(self.x)
+            return self.tok
+
+    c1, c2 = C(o), C(o)
+    check_tokenize(o, c1, o)
+    assert c1.tok
+    assert check_tokenize(c1) == check_tokenize(c2)
+
+
 def test_tokenize_dict():
     # 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})