Repository: dask/dask
Pull request #11970 (2025-05-28 15:29:52)

Account for ``__main__`` in ``pickle`` normalization
Closes https://github.com/dask/dask/issues/11964 

In https://github.com/dask/dask/pull/11320 we dropped logic that checks for if an object is defined in `__main__` (where `cloudpickle` should be used instead of `pickle`). Specifically these lines https://github.com/dask/dask/pull/11320/files#diff-10422b02c591d63ee295724faa14f7698b4a742c98ba20771c5f70d1a6926d06L1240-L1243

This PR proposes we just always use `cloudpickle` instead of `pickle` first and then, if some conditions are met, going to `cloudpickle` instead (though I'll admit I'm not totally sure why the previous "try pickle first" logic existed). Another option would be to restore the `if b"__main__" in pik:` logic. 

cc @hendrikmakait 

jrbourbeau: > IIRC, the reason for trying pickle first is performance

Thanks, that's helpful. Just pushed a commit to make sure we try `pickle` first (restoring older behavior) 

Patch:
diff --git a/continuous_integration/scripts/install.sh b/continuous_integration/scripts/install.sh
index b3ebdbf33..5baf53779 100644
--- a/continuous_integration/scripts/install.sh
+++ b/continuous_integration/scripts/install.sh
@@ -19,22 +19,23 @@ if [[ ${UPSTREAM_DEV} ]]; then
     python -m pip install \
         --upgrade \
         locket \
-        git+https://github.com/dask/s3fs \
-        git+https://github.com/intake/filesystem_spec \
+        git+https://github.com/fsspec/s3fs \
         git+https://github.com/dask/partd \
         git+https://github.com/dask/zict \
         git+https://github.com/dask/distributed \
         git+https://github.com/zarr-developers/zarr-python
-    # TODO: Add nightly `h5py` back once https://github.com/h5py/h5py/issues/2563 is resolved
-    # mamba uninstall --force numpy pandas scipy numexpr numba sparse scikit-image h5py numbagg
+    # NOTE: Dev version of `fsspec` needs to be installed after the dev version of
+    # `s3fs` to avoid dependency conflicts
+    python -m pip install --upgrade git+https://github.com/fsspec/filesystem_spec
+    # TODO: Add nightly `scikit-image` back once it's available
     mamba uninstall --force numpy pandas scipy numexpr numba sparse scikit-image numbagg
     python -m pip install --no-deps --pre --retries 10 \
         -i https://pypi.anaconda.org/scientific-python-nightly-wheels/simple \
         numpy \
         pandas \
         scipy \
-        scikit-image
-        # h5py
+        h5py
+        # scikit-image
 
     # Used when automatically opening an issue when the `upstream` CI build fails
     mamba install pytest-reportlog
diff --git a/dask/tokenize.py b/dask/tokenize.py
index 5f780adc2..f2bf86a83 100644
--- a/dask/tokenize.py
+++ b/dask/tokenize.py
@@ -233,26 +233,30 @@ def _normalize_pure_object(o: object) -> tuple[str, int]:
 def _normalize_pickle(o: object) -> tuple:
     buffers: list[pickle.PickleBuffer] = []
     pik: int | None = None
-    pik2: int
-    success = False
-    for mod in [pickle, cloudpickle]:
-        if success:
-            break
-        for _ in range(3):
+    pik2: int | None = None
+    for _ in range(3):
+        buffers.clear()
+        try:
+            out = pickle.dumps(o, protocol=5, buffer_callback=buffers.append)
+            if b"__main__" in out:
+                # Use `cloudpickle` for objects defined in `__main__`
+                buffers.clear()
+                out = cloudpickle.dumps(o, protocol=5, buffer_callback=buffers.append)
+            pickle.loads(out, buffers=buffers)
+            pik2 = hash_buffer_hex(out)
+        except Exception:
             buffers.clear()
             try:
-                out = mod.dumps(o, protocol=5, buffer_callback=buffers.append)
-                mod.loads(out, buffers=buffers)
+                out = cloudpickle.dumps(o, protocol=5, buffer_callback=buffers.append)
+                pickle.loads(out, buffers=buffers)
                 pik2 = hash_buffer_hex(out)
             except Exception:
                 break
-            if pik == pik2:
-                success = True
-                break
-            pik = pik2
-        else:
-            _maybe_raise_nondeterministic("Failed to tokenize deterministically")
+        if pik and pik2 and pik == pik2:
             break
+        pik = pik2
+    else:
+        _maybe_raise_nondeterministic("Failed to tokenize deterministically")
     if pik is None:
         _maybe_raise_nondeterministic("Failed to tokenize deterministically")
         pik = int(uuid.uuid4())
diff --git a/dask/utils.py b/dask/utils.py
index ee9226ecf..6edfda8f8 100644
--- a/dask/utils.py
+++ b/dask/utils.py
@@ -747,7 +747,10 @@ class Dispatch:
             return lk[cls]
         for cls2 in cls.__mro__:
             # Is a lazy registration function present?
-            toplevel, _, _ = cls2.__module__.partition(".")
+            try:
+                toplevel, _, _ = cls2.__module__.partition(".")
+            except Exception:
+                continue
             try:
                 register = self._lazy[toplevel]
             except KeyError:


Test patch:
diff --git a/dask/tests/test_tokenize.py b/dask/tests/test_tokenize.py
index 08da401c5..e15d1d090 100644
--- a/dask/tests/test_tokenize.py
+++ b/dask/tests/test_tokenize.py
@@ -1205,6 +1205,18 @@ def test_tokenize_functions_main():
         return y + 1
 
     assert tokenize(inc2) != tokenize(inc)
+
+    # Test that redefining a function changes the token
+    def func(x):
+        return x + 1
+
+    result = tokenize(func)
+
+    def func(x):
+        return x + 2
+
+    result2 = tokenize(func)
+    assert result != result2
     """
     proc = subprocess.run([sys.executable, "-c", textwrap.dedent(script)])
     proc.check_returncode()