Repository: dask/dask
Pull request #11431 (2024-10-15 13:47:48)

Ensure that recursively calling tokenize respects ensure_deterministic
In https://github.com/dask/dask/pull/11373 I removed contextvars for the seen set to reduce overhead. The logic I put in place for the `_ENSURE_DETERMINISTIC` variable was erroneous and not safe for recursively calling tokenize. This PR addresses this while maintaining the performance critical global variable for set.

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  ± 0       13 suites  ±0   3h 10m 14s :stopwatch: + 2m 42s
 13 218 tests + 2   12 161 :white_check_mark: + 1   1 056 :zzz: ±0  1 :x: +1 
137 874 runs  +26  118 859 :white_check_mark: +23  19 014 :zzz: +2  1 :x: +1 

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

Results for commit b829b19a. ± Comparison against base commit aa5e392d.

[test-results]:data:application/gzip;base64,H4sIABB7DmcC/03NSw7CIBSF4a00jB1weV83Y9oCCbEPQ2Fk3LtgK+3w/y45vIkPk9vIvQN+68iWQzrL5tinsC61QYAoVK5pvzMw/35seRwrMlBw4jO8KlKpmvk+TNUauBjXWIQWiXnZp7XR4ug2DcZIbHhMI4X28DL96+vyuM5zSCXIYBgOgL1UaBx1XFsrJTrglnrOtC5fAxMI5PMFGiACmBoBAAA=


Patch:
diff --git a/dask/tokenize.py b/dask/tokenize.py
index 0c733db5a..72ba96b0b 100644
--- a/dask/tokenize.py
+++ b/dask/tokenize.py
@@ -13,6 +13,7 @@ import types
 import uuid
 from collections import OrderedDict
 from collections.abc import Iterable
+from contextvars import ContextVar
 from functools import partial
 
 import cloudpickle
@@ -40,7 +41,7 @@ def _tokenize(*args: object, **kwargs: object) -> str:
 
 tokenize_lock = threading.RLock()
 _SEEN: dict[int, tuple[int, object]] = {}
-_ENSURE_DETERMINISTIC = None
+_ENSURE_DETERMINISTIC: ContextVar[bool | None] = ContextVar("_ENSURE_DETERMINISTIC")
 
 
 def tokenize(
@@ -66,20 +67,25 @@ def tokenize(
     global _SEEN, _ENSURE_DETERMINISTIC
     with tokenize_lock:
         seen_before, _SEEN = _SEEN, {}
-        _ENSURE_DETERMINISTIC = ensure_deterministic
+        token = None
+        try:
+            _ENSURE_DETERMINISTIC.get()
+        except LookupError:
+            token = _ENSURE_DETERMINISTIC.set(ensure_deterministic)
         try:
             return _tokenize(*args, **kwargs)
         finally:
+            if token:
+                _ENSURE_DETERMINISTIC.reset(token)
             _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")
-    ):
+    try:
+        val = _ENSURE_DETERMINISTIC.get()
+    except LookupError:
+        val = None
+    if val or val is None and config.get("tokenize.ensure-deterministic"):
         raise TokenizationError(msg)
 
 


Test patch:
diff --git a/dask/tests/test_task_spec.py b/dask/tests/test_task_spec.py
index b5cef98ad..23e8ff57f 100644
--- a/dask/tests/test_task_spec.py
+++ b/dask/tests/test_task_spec.py
@@ -745,3 +745,8 @@ def test_execute_tasks_in_graph():
     res = execute_graph(dsk)
     assert len(res) == 1
     assert res["key-4"] == "foo-bar-a-b=c"
+
+
+def test_deterministic_tokenization_respected():
+    with pytest.raises(RuntimeError, match="deterministic"):
+        tokenize(Task("key", func, object()), ensure_deterministic=True)
diff --git a/dask/tests/test_tokenize.py b/dask/tests/test_tokenize.py
index a3f447027..d361b6bc8 100644
--- a/dask/tests/test_tokenize.py
+++ b/dask/tests/test_tokenize.py
@@ -38,10 +38,12 @@ def check_clean_state():
     from dask.tokenize import _ENSURE_DETERMINISTIC, _SEEN
 
     assert not _SEEN
-    assert _ENSURE_DETERMINISTIC is None
+    with pytest.raises(LookupError):
+        _ENSURE_DETERMINISTIC.get()
     yield
     assert not _SEEN
-    assert _ENSURE_DETERMINISTIC is None
+    with pytest.raises(LookupError):
+        _ENSURE_DETERMINISTIC.get()
 
 
 def check_tokenize(*args, **kwargs):
@@ -1423,3 +1425,12 @@ def test_tokenize_pandas_arrow_strings():
     tokens = normalize_token(ser)
     # Maybe a little brittle but will do for now
     assert any(str(tok) == "string" for tok in flatten(tokens))
+
+
+def test_tokenize_recursive_respects_ensure_deterministic():
+    class Foo:
+        def __dask_tokenize__(self):
+            return tokenize(object())
+
+    with pytest.raises(RuntimeError):
+        tokenize(Foo(), ensure_deterministic=True)