Repository: dask/dask
Pull request #10919 (2024-02-12 15:40:19)

tokenize: Don't call str() on dict values
- XREF https://github.com/dask-contrib/dask-awkward/issues/468



lgray: I think we can drop handling of `key=str` in `set`, at least in my case it will not happen, but it does leave a possibility for future unintended problems like this.
crusaderky: > I think we can drop handling of `key=str` in `set`, at least in my case it will not happen, but it does leave a possibility for future unintended problems like this.

dropping `key=str` would cause a failure for `tokenize({1, (2, 3)})`.
Ultimately I would much rather work on the assumption that `str()` is cheap for all objects.
lgray: Sure - I think that's fine and reasonable for set. It may still be useful to stick in the documentation a louder warning/callout/suggestion that `__dask_tokenize__` is very much the preferred option.
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   3h 26m 4s :stopwatch: + 9m 56s
 12 991 tests + 3   12 062 :white_check_mark: + 3     929 :zzz: ±0  0 :x: ±0 
160 552 runs  +45  144 080 :white_check_mark: +48  16 472 :zzz:  - 3  0 :x: ±0 

Results for commit 65550e09. ± Comparison against base commit 7726d319.

[test-results]:data:application/gzip;base64,H4sIAFFRymUC/03MTQ7CIBCG4as0rF0MU6AdL2PoSBNifwyFlfHuIlba5fvM5HuJ0U9uE9dG6ksjtuTjUfcUbPTr8m1sjcqUr7HckUj++7Yl5oJg8MCHf2YkpEqj9VMmqOBCWMMuIS1l2YDWuHddVgp6qPhblkZ19fE0Xfq8zOs8+5hDGK01OCCre+7k4FgPRtGI2KFi4h4cgyHbivcHNduB4xkBAAA=


Patch:
diff --git a/dask/base.py b/dask/base.py
index 795ebd536..7e3e33436 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -1062,7 +1062,7 @@ def normalize_type(typ):
 
 @normalize_token.register((dict, types.MappingProxyType))
 def normalize_dict(d):
-    return "dict", _normalize_seq_func(sorted(d.items(), key=str))
+    return "dict", _normalize_seq_func(sorted(d.items(), key=lambda kv: str(kv[0])))
 
 
 @normalize_token.register(OrderedDict)


Test patch:
diff --git a/dask/tests/test_tokenize.py b/dask/tests/test_tokenize.py
index f418e4543..91b010e51 100644
--- a/dask/tests/test_tokenize.py
+++ b/dask/tests/test_tokenize.py
@@ -29,6 +29,7 @@ np = import_or_none("numpy")
 sp = import_or_none("scipy.sparse")
 pa = import_or_none("pyarrow")
 pd = import_or_none("pandas")
+numba = import_or_none("numba")
 
 
 def check_tokenize(*args, **kwargs):
@@ -790,6 +791,42 @@ def test_tokenize_ordered_dict():
     assert check_tokenize(a) != check_tokenize(c)
 
 
+def test_tokenize_dict_doesnt_call_str_on_values():
+    class C:
+        def __dask_tokenize__(self):
+            return "C"
+
+        def __repr__(self):
+            assert False
+
+    check_tokenize({1: C(), "2": C()})
+
+
+def test_tokenize_sorts_dict_before_seen_map():
+    """When sequence values are repeated, the 2nd+ entry is tokenized as (__seen, 0).
+    This makes it important to ensure that dicts are sorted *before* you call
+    normalize_token() on their elements.
+    """
+    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)
+
+
+def test_tokenize_sorts_set_before_seen_map():
+    """Same as test_tokenize_sorts_dict_before_seen_map, but for sets.
+
+    Note that this test is only meaningful if set insertion order impacts iteration
+    order, which is an implementation detail of the Python interpreter.
+    """
+    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)
+
+
 def test_tokenize_timedelta():
     assert check_tokenize(datetime.timedelta(days=1)) == check_tokenize(
         datetime.timedelta(days=1)
@@ -1201,3 +1238,100 @@ def test_tokenize_opaque_object_with_buffers():
             self.x = np.array(x)
 
     assert check_tokenize(C([1, 2])) != check_tokenize(C([1, 3]))
+
+
+if not numba:
+
+    class NumbaDummy:
+        def __bool__(self):
+            return False
+
+        def _dummy_decorator(self, *args, **kwargs):
+            def wrapper(func):
+                return func
+
+            return wrapper
+
+        jit = vectorize = guvectorize = _dummy_decorator
+
+    numba = NumbaDummy()
+
+
+@numba.jit(nopython=True)
+def numba_jit(x, y):
+    return x + y
+
+
+@numba.jit("f8(f8, f8)", nopython=True)
+def numba_jit_with_signature(x, y):
+    return x + y
+
+
+@numba.vectorize(nopython=True)
+def numba_vectorize(x, y):
+    return x + y
+
+
+@numba.vectorize("f8(f8, f8)", nopython=True)
+def numba_vectorize_with_signature(x, y):
+    return x + y
+
+
+@numba.guvectorize(["f8,f8,f8[:]"], "(),()->()")
+def numba_guvectorize(x, y, out):
+    out[0] = x + y
+
+
+all_numba_funcs = [
+    numba_jit,
+    numba_jit_with_signature,
+    numba_vectorize,
+    numba_vectorize_with_signature,
+    numba_guvectorize,
+]
+
+
+@pytest.mark.skipif("not numba")
+@pytest.mark.parametrize("func", all_numba_funcs)
+def test_tokenize_numba(func):
+    assert func(1, 2) == 3
+    check_tokenize(func)
+
+
+@pytest.mark.skipif("not numba")
+def test_tokenize_numba_unique_token():
+    tokens = [check_tokenize(func) for func in all_numba_funcs]
+    assert len(tokens) == len(set(tokens))
+
+
+@pytest.mark.skipif("not numba")
+def test_numba_local():
+    @numba.jit(nopython=True)
+    def local_jit(x, y):
+        return x + y
+
+    @numba.jit("f8(f8, f8)", nopython=True)
+    def local_jit_with_signature(x, y):
+        return x + y
+
+    @numba.vectorize(nopython=True)
+    def local_vectorize(x, y):
+        return x + y
+
+    @numba.vectorize("f8(f8, f8)", nopython=True)
+    def local_vectorize_with_signature(x, y):
+        return x + y
+
+    @numba.guvectorize(["f8,f8,f8[:]"], "(),()->()")
+    def local_guvectorize(x, y, out):
+        out[0] = x + y
+
+    all_funcs = [
+        local_jit,
+        local_jit_with_signature,
+        local_vectorize,
+        local_vectorize_with_signature,
+        local_guvectorize,
+    ]
+    tokens = [check_tokenize(func) for func in all_funcs]
+    assert len(tokens) == len(set(tokens))