Repository: dask/dask
Pull request #11709 (2025-01-31 09:55:26)

Ensure `normalize_token` is threadsafe
Initially I assumed we're only accessing these functions through proper `tokenize` which is protected with this recursive lock already. However, we're also using `normalize_token` individually in the distributed scheduler.

If the scheduler is running in the same process as the client (which is typically the case for LocalCluster) this can cause race conditions which can manifest in wrong tokens or even in KeyErrors when normalizing a nested sequence.

I wasn't able to reproduce the KeyError with the test but it is sensitive to a wrong result

This was introduced in https://github.com/dask/dask/pull/11373

fjetter: I performed some micro benchmarks and I believe the rlock is not adding any noticeable overhead
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._

     15 files  ± 0       15 suites  ±0   4h 28m 1s ⏱️ + 1m 25s
 17 143 tests + 1   15 963 ✅ + 1   1 180 💤 ±0  0 ❌ ±0 
211 241 runs  +15  194 208 ✅ +19  17 033 💤  - 4  0 ❌ ±0 

Results for commit 3471d0e7. ± Comparison against base commit 7f490b91.

[test-results]:data:application/gzip;base64,H4sIAAKnnGcC/03MSQ7DIAyF4atErLuwzdzLVDQBCTVDRWAV9e4lUabl/1l+Cwux9zN7NigfDZtLzFd1Jbkcp3FtBQYr1Wve7hoFP/o1l7bdnqy64Sd+V0QDpwUX+2oX+JSmtEsq4zpNiCRw73PaCgJz4j6tgfPDbtNb35fbaRhirsG40NiB19qTlMSDVaCkE2/QyhL3QcgAJL1hvz8qDe14GgEAAA==


Patch:
diff --git a/dask/tokenize.py b/dask/tokenize.py
index d6601778c..b0f0a5a77 100644
--- a/dask/tokenize.py
+++ b/dask/tokenize.py
@@ -114,15 +114,16 @@ normalize_token.register(
 
 @normalize_token.register((types.MappingProxyType, dict))
 def normalize_dict(d):
-    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)
+    with tokenize_lock:
+        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)
@@ -145,13 +146,14 @@ def _normalize_seq_func(seq: Iterable[object]) -> tuple[object, ...]:
             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)]
+    with tokenize_lock:
+        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))


Test patch:
diff --git a/dask/tests/test_tokenize.py b/dask/tests/test_tokenize.py
index a2a5f625d..4bf3ca224 100644
--- a/dask/tests/test_tokenize.py
+++ b/dask/tests/test_tokenize.py
@@ -10,6 +10,7 @@ import random
 import subprocess
 import sys
 import textwrap
+from concurrent.futures import ThreadPoolExecutor
 from enum import Enum, Flag, IntEnum, IntFlag
 from typing import Union
 
@@ -1446,3 +1447,20 @@ def test_tokenize_recursive_respects_ensure_deterministic():
 
     with pytest.raises(RuntimeError):
         tokenize(Foo(), ensure_deterministic=True)
+
+
+def test_tokenize_nested_sequence_thread_safe():
+    nested_list = []
+    for ix in range(100):
+        nested_list = [ix, nested_list]
+
+    nested_dict = {}
+    for ix in range(50):
+        nested_dict = {ix: nested_dict}
+
+    normalize_token(nested_list)
+    with ThreadPoolExecutor() as pool:
+        futures = [pool.submit(normalize_token, nested_list) for _ in range(1000)]
+        assert len({f.result() for f in futures}) == 1
+        futures = [pool.submit(normalize_token, nested_dict) for _ in range(1000)]
+        assert len({f.result() for f in futures}) == 1