Custom serializer support for dumps_function()
- Dominant language
- Python
- Stars
- 1.7k
- Forks
- 778
- Avg merge
- 2h 50m
- Merged PRs (30d)
- 3
Description
I'm having a use case which requires me to parallelize nested Cython functions using Dask, and from the documentation, I've learned that the we can set the serializers/deserializers for a [Client object](https://distributed.dask.org/en/latest/api.html#distributed.Client.map). This means I can implement my own serializer, in this case, dill, to parallelize my nested functions. It looks like this.
```python
import dill
from distributed.protocol.serialize import register_serialization_family
def dill_dumps(x):
header = {'serializer': 'dill'}
frames = [dill.dumps(x)]
return header, frames
def dill_loads(header, frames):
if len(frames) > 1: # this may be cut up for network reasons
frame = ''.join(frames)
else:
frame = frames[0]
return dill.loads(frame)
register_serialization_family('dill', dill_dumps, dill_loads)
...
client = Client(address='tcp://0.0.0.0:8786', set_as_default=False, serializers=['dill'])
```
However, when I'm using `map()` from the client, I am still unable to serialize it, and upon inspecting the error message, it turns out it is using pickle (which will not work for nested function) in `dumps_function()` to serialize my function in the class.
Here's the error message:
```
Traceback (most recent call last):
File "training_engine.py", line 2226, in training_engine.train.train_and_update
File "accml.py", line 459, in accml.AccML.fit
File "pipeline_selector.py", line 621, in pipeline_selector.ClassifierPipelineSelector.fit
File "/usr/local/lib/python3.7/site-packages/distributed/client.py", line 1779, in map
actors=actor,
File "/usr/local/lib/python3.7/site-packages/distributed/client.py", line 2590, in _graph_to_futures
"tasks": valmap(dumps_task, dsk3),
File "/usr/local/lib/python3.7/site-packages/toolz/dicttoolz.py", line 83, in valmap
rv.update(zip(iterkeys(d), map(func, itervalues(d))))
File "/usr/local/lib/python3.7/site-packages/distributed/worker.py", line 3354, in dumps_task
return {"function": dumps_function(task[0]), "args": warn_dumps(task[1:])}
File "/usr/local/lib/python3.7/site-packages/distributed/worker.py", line 3318, in dumps_function
result = pickle.dumps(func)
File "/usr/local/lib/python3.7/site-packages/distributed/protocol/pickle.py", line 51, in dumps
result = cloudpickle.dumps(x, **dump_kwargs)
File "/usr/local/lib/python3.7/site-packages/cloudpickle/cloudpickle.py", line 1148, in dumps
cp.dump(obj)
File "/usr/local/lib/python3.7/site-packages/cloudpickle/cloudpickle.py", line 491, in dump
return Pickler.dump(self, obj)
File "/usr/local/lib/python3.7/pickle.py", line 437, in dump
self.save(obj)
File "/usr/local/lib/python3.7/pickle.py", line 535, in save
self.save_global(obj, rv)
File "/usr/local/lib/python3.7/site-packages/cloudpickle/cloudpickle.py", line 897, in save_global
Pickler.save_global(self, obj, name=name)
File "/usr/local/lib/python3.7/pickle.py", line 960, in save_global
(obj, module_name, name)) from None
_pickle.PicklingError: Can't pickle .dask_fit at 0x7fe2b25de870>: it's not found as pipeline_selector.dask_fit
```
I would like to request for the feature to use the custom serializer, as set from the Client object.
Contributor guide
Assessment
This issue has not been assessed yet.