dask / dask/distributed

Registration of custom (de)serializer is recognized by dask_loads, dask_dumps, but not when computing graph

Open
#2,469 12 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
1.7k
Forks
778
Avg merge
2h 50m
Merged PRs (30d)
3

Description

I've created an object wrapper type to carry along metadata with otherwise typical data objects such as NumPy arrays. It seems that for whatever reason, after registering the associated (de)serializers, while the Distributed functions for serialization recognize this registration when called directly, during execution of the computational graph, these registrations are ignored, and Pickle is used instead.

## Self-contained demonstration of issue:

Must be in two separate module files for whatever reason.

main.py

```python
from imageobject import ImageObject
from distributed.protocol.serialize import dask_loads, dask_dumps

import numpy as np

from dask.distributed import Client, LocalCluster

import dask.bag as db

def get_imageobject(*args, **kwargs):
return ImageObject(np.zeros(6))

if __name__ == '__main__':

# #######################
# serialize manually

im = get_imageobject()

im_ser = dask_loads(*dask_dumps(im))

print(np.all(im == im_ser))

# #########################################
# serialize through Dask graph execution

cluster = LocalCluster()

with Client(cluster) as client:
im_ser = client.compute(
db.from_sequence([None, None, None]).map(get_imageobject),
sync=True
)

print(np.all(im == im_ser))
```

imageobject.py

```python
import wrapt
import dill

from typing import Any, Dict, Tuple, List, Union
from distributed.protocol.serialize import dask_serialize, dask_deserialize

class ImageObject(wrapt.ObjectProxy):
def __init__(self, object_to_wrap: Any):
super().__init__(object_to_wrap)

def serialize(self):
print('ser')

obj_dict = dict(
dataobject_type=type(self),
object_to_wrap=self.__wrapped__,
)

return dill.dumps(obj_dict)

@classmethod
def deserialize(cls, serialized_dict: Union[bytes, str]):
print('deser')

obj_dict = dill.loads(serialized_dict)

dataobject_type = obj_dict['dataobject_type']
object_to_wrap = obj_dict['object_to_wrap']

return dataobject_type(object_to_wrap)

def dask_serialize_imageobject(dataobject: ImageObject) -> Tuple[Dict, List[bytes]]:
header = {}
frames = [dataobject.serialize()]

return header, frames

def dask_deserialize_imageobject(header: Dict, frames: List[bytes]) -> ImageObject:
return ImageObject.deserialize(frames[0])

dask_serialize.register(ImageObject)(dask_serialize_imageobject)
dask_deserialize.register(ImageObject)(dask_deserialize_imageobject)
```

## Execution output:

```
/scratch/anaconda3/envs/beads/bin/python /data/dtk-pipeline/scripts/test_dask_serialization.py
ser
deser
True
distributed.protocol.core - CRITICAL - Failed to deserialize
Traceback (most recent call last):
File "/scratch/anaconda3/envs/beads/lib/python3.7/site-packages/distributed/protocol/core.py", line 132, in loads
value = _deserialize(head, fs, deserializers=deserializers)
File "/scratch/anaconda3/envs/beads/lib/python3.7/site-packages/distributed/protocol/serialize.py", line 184, in deserialize
return loads(header, frames)
File "/scratch/anaconda3/envs/beads/lib/python3.7/site-packages/distributed/protocol/serialize.py", line 81, in serialization_error_loads
raise TypeError(msg)
TypeError: Could not serialize object of type list.
Traceback (most recent call last):
File "/scratch/anaconda3/envs/beads/lib/python3.7/site-packages/distributed/protocol/pickle.py", line 38, in dumps
result = pickle.dumps(x, protocol=pickle.HIGHEST_PROTOCOL)
TypeError: can't pickle ImageObject objects
```

Which seems to just keep popping up in an endless loop.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.