diff --git a/distributed/protocol/pickle.py b/distributed/protocol/pickle.py index 561de2ac32..ae6efc9fcd 100644 --- a/distributed/protocol/pickle.py +++ b/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 diff --git a/distributed/protocol/tests/test_h5py.py b/distributed/protocol/tests/test_h5py.py index 35a9ce99ba..09144f3252 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()