[unsupervised AI] Use dask_serialize reducers when falling back to cloudpickle

distributed/protocol/pickle.py

@@ -28,6 +28,37 @@ 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))
+                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)
+
+
+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 +104,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

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()