From 47ee88bcf6538dd28b38a69bbe60ca243b5fb7d6 Mon Sep 17 00:00:00 2001
From: Charan Rathore <180254320+charan-rathore@users.noreply.github.com>
Date: Tue, 6 Oct 2026 21:19:38 +0530
Subject: [PATCH 1/2] Use dask_serialize reducers when falling back to
 cloudpickle

Signed-off-by: Charan Rathore <180254320+charan-rathore@users.noreply.github.com>
---
 distributed/protocol/pickle.py          | 29 +++++++++++++++++++++++--
 distributed/protocol/tests/test_h5py.py | 22 +++++++++++++++++++
 2 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/distributed/protocol/pickle.py b/distributed/protocol/pickle.py
index 561de2ac32d..34652e5598e 100644
--- a/distributed/protocol/pickle.py
+++ b/distributed/protocol/pickle.py
@@ -28,6 +28,31 @@ def reducer_override(self, obj):
             return NotImplemented
 
 
+class _DaskCloudPickler(cloudpickle.Pickler):
+    """cloudpickle with the same ``dask_serialize`` reducers as ``_DaskPickler``
+
+    Objects that plain pickle can't handle, such as h5py datasets, would otherwise
+    fail as soon as the same graph also holds something that needs cloudpickle, e.g.
+    a function defined in ``__main__``.
+    """
+
+    def reducer_override(self, obj):
+        if not _always_use_pickle_for(obj):
+            try:
+                serialize = dask_serialize.dispatch(type(obj))
+                deserialize = dask_deserialize.dispatch(type(obj))
+                return deserialize, serialize(obj)
+            except TypeError:
+                pass
+        return super().reducer_override(obj)
+
+
+def _cloudpickle_dumps(x, **dump_kwargs):
+    f = io.BytesIO()
+    _DaskCloudPickler(f, **dump_kwargs).dump(x)
+    return f.getvalue()
+
+
 def _always_use_pickle_for(x):
     try:
         mod, _, _ = x.__class__.__module__.partition(".")
@@ -73,11 +98,11 @@ def dumps(x, *, buffer_callback=None, protocol=HIGHEST_PROTOCOL):
         ):
             if len(result) < 1000 or not _always_use_pickle_for(x):
                 buffers.clear()
-                result = cloudpickle.dumps(x, **dump_kwargs)
+                result = _cloudpickle_dumps(x, **dump_kwargs)
     except Exception:
         try:
             buffers.clear()
-            result = cloudpickle.dumps(x, **dump_kwargs)
+            result = _cloudpickle_dumps(x, **dump_kwargs)
         except Exception:
             logger.exception("Failed to serialize %s.", x)
             raise
diff --git a/distributed/protocol/tests/test_h5py.py b/distributed/protocol/tests/test_h5py.py
index 35a9ce99baa..09144f32525 100644
--- a/distributed/protocol/tests/test_h5py.py
+++ b/distributed/protocol/tests/test_h5py.py
@@ -119,3 +119,25 @@ async def test_h5py_serialize_2(c, s, a, b):
             y = c.compute(x.sum())
             y = await y
             assert y == (1 + 2 + 3 + 4) * 3
+
+
+@silence_h5py_issue775
+def test_pickle_dataset_next_to_object_needing_cloudpickle():
+    # https://github.com/dask/distributed/issues/9013
+    # A local function can only be pickled by cloudpickle, which used to ignore the
+    # dask_serialize reducers for h5py objects.
+    from distributed.protocol import pickle
+
+    def func(x):
+        return x**2
+
+    with tmpfile() as fn:
+        with h5py.File(fn, mode="a") as f:
+            f.create_dataset("/group/x", data=[1, 2, 3, 4])
+        with h5py.File(fn, mode="r") as f:
+            dset = f["/group/x"]
+            func2, dset2 = pickle.loads(pickle.dumps((func, dset)))
+            assert func2(3) == 9
+            assert isinstance(dset2, h5py.Dataset)
+            assert dset2.name == dset.name
+            assert (dset2[:] == dset[:]).all()

From 66ae10928aec828c892cbd330cf580a97ff93230 Mon Sep 17 00:00:00 2001
From: Charan Rathore <180254320+charan-rathore@users.noreply.github.com>
Date: Tue, 6 Oct 2026 22:12:30 +0530
Subject: [PATCH 2/2] Avoid infinite recursion on self-framed serializers in
 cloudpickle fallback

Serializers for buffer-like objects such as memoryview and bytearray
return the object itself in frames; applying their dask_serialize
reducer inside the cloudpickle pickler recursed until the interpreter
hit the recursion limit (test_warn_when_submitting_large_values_memoryview).
Fall back to native pickling when a serializer frames the object itself.

Signed-off-by: Charan Rathore <180254320+charan-rathore@users.noreply.github.com>
---
 distributed/protocol/pickle.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/distributed/protocol/pickle.py b/distributed/protocol/pickle.py
index 34652e5598e..ae6efc9fcd6 100644
--- a/distributed/protocol/pickle.py
+++ b/distributed/protocol/pickle.py
@@ -41,9 +41,15 @@ def reducer_override(self, obj):
             try:
                 serialize = dask_serialize.dispatch(type(obj))
                 deserialize = dask_deserialize.dispatch(type(obj))
-                return deserialize, serialize(obj)
+                header, frames = serialize(obj)
             except TypeError:
                 pass
+            else:
+                # Serializers for buffer-like objects (e.g. ``memoryview``)
+                # return the object itself in ``frames``, which would recurse
+                # forever here. Native pickling handles those as buffers.
+                if not any(obj is frame for frame in frames):
+                    return deserialize, (header, frames)
         return super().reducer_override(obj)