File size: 2,009 Bytes
cf6a8b4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 | import ray
import ray.cloudpickle as pickle
from ray.util.annotations import DeveloperAPI, PublicAPI
@PublicAPI
def register_serializer(cls: type, *, serializer: callable, deserializer: callable):
"""Use the given serializer to serialize instances of type ``cls``,
and use the deserializer to deserialize the serialized object.
Args:
cls: A Python class/type.
serializer: A function that converts an instances of
type ``cls`` into a serializable object (e.g. python dict
of basic objects).
deserializer: A function that constructs the
instance of type ``cls`` from the serialized object.
This function itself must be serializable.
"""
context = ray._private.worker.global_worker.get_serialization_context()
context._register_cloudpickle_serializer(cls, serializer, deserializer)
@PublicAPI
def deregister_serializer(cls: type):
"""Deregister the serializer associated with the type ``cls``.
There is no effect if the serializer is unavailable.
Args:
cls: A Python class/type.
"""
context = ray._private.worker.global_worker.get_serialization_context()
context._unregister_cloudpickle_reducer(cls)
@DeveloperAPI
class StandaloneSerializationContext:
# NOTE(simon): Used for registering custom serializers. We cannot directly
# use the SerializationContext because it requires Ray workers. Please
# make sure to keep the API consistent.
def _register_cloudpickle_reducer(self, cls, reducer):
pickle.CloudPickler.dispatch[cls] = reducer
def _unregister_cloudpickle_reducer(self, cls):
pickle.CloudPickler.dispatch.pop(cls, None)
def _register_cloudpickle_serializer(
self, cls, custom_serializer, custom_deserializer
):
def _CloudPicklerReducer(obj):
return custom_deserializer, (custom_serializer(obj),)
# construct a reducer
pickle.CloudPickler.dispatch[cls] = _CloudPicklerReducer
|