ray/python/ray/serialization.py

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import cloudpickle as pickle

import ray.numbuf


class RaySerializationException(Exception):
  def __init__(self, message, example_object):
    Exception.__init__(self, message)
    self.example_object = example_object


class RayDeserializationException(Exception):
  def __init__(self, message, class_id):
    Exception.__init__(self, message)
    self.class_id = class_id


class RayNotDictionarySerializable(Exception):
  pass


def check_serializable(cls):
  """Throws an exception if Ray cannot serialize this class efficiently.

  Args:
    cls (type): The class to be serialized.

  Raises:
    Exception: An exception is raised if Ray cannot serialize this class
      efficiently.
  """
  if is_named_tuple(cls):
    # This case works.
    return
  if not hasattr(cls, "__new__"):
    print("The class {} does not have a '__new__' attribute and is probably "
          "an old-stye class. Please make it a new-style class by inheriting "
          "from 'object'.")
    raise RayNotDictionarySerializable("The class {} does not have a "
                                       "'__new__' attribute and is probably "
                                       "an old-style class. We do not support "
                                       "this. Please make it a new-style "
                                       "class by inheriting from 'object'."
                                       .format(cls))
  try:
    obj = cls.__new__(cls)
  except:
    raise RayNotDictionarySerializable("The class {} has overridden '__new__'"
                                       ", so Ray may not be able to serialize "
                                       "it efficiently.".format(cls))
  if not hasattr(obj, "__dict__"):
    raise RayNotDictionarySerializable("Objects of the class {} do not have a "
                                       "'__dict__' attribute, so Ray cannot "
                                       "serialize it efficiently.".format(cls))
  if hasattr(obj, "__slots__"):
    raise RayNotDictionarySerializable("The class {} uses '__slots__', so Ray "
                                       "may not be able to serialize it "
                                       "efficiently.".format(cls))


# This field keeps track of a whitelisted set of classes that Ray will
# serialize.
type_to_class_id = dict()
whitelisted_classes = dict()
classes_to_pickle = set()
custom_serializers = dict()
custom_deserializers = dict()


def is_named_tuple(cls):
  """Return True if cls is a namedtuple and False otherwise."""
  b = cls.__bases__
  if len(b) != 1 or b[0] != tuple:
    return False
  f = getattr(cls, "_fields", None)
  if not isinstance(f, tuple):
    return False
  return all(type(n) == str for n in f)


def add_class_to_whitelist(cls, class_id, pickle=False, custom_serializer=None,
                           custom_deserializer=None):
  """Add cls to the list of classes that we can serialize.

  Args:
    cls (type): The class that we can serialize.
    class_id: A string of bytes used to identify the class.
    pickle (bool): True if the serialization should be done with pickle. False
      if it should be done efficiently with Ray.
    custom_serializer: This argument is optional, but can be provided to
      serialize objects of the class in a particular way.
    custom_deserializer: This argument is optional, but can be provided to
      deserialize objects of the class in a particular way.
  """
  type_to_class_id[cls] = class_id
  whitelisted_classes[class_id] = cls
  if pickle:
    classes_to_pickle.add(class_id)
  if custom_serializer is not None:
    custom_serializers[class_id] = custom_serializer
    custom_deserializers[class_id] = custom_deserializer


def serialize(obj):
  """This is the callback that will be used by numbuf.

  If numbuf does not know how to serialize an object, it will call this method.

  Args:
    obj (object): A Python object.

  Returns:
    A dictionary that has the key "_pyttype_" to identify the class, and
      contains all information needed to reconstruct the object.
  """
  if type(obj) not in type_to_class_id:
    raise RaySerializationException("Ray does not know how to serialize "
                                    "objects of type {}.".format(type(obj)),
                                    obj)
  class_id = type_to_class_id[type(obj)]

  if class_id in classes_to_pickle:
    serialized_obj = {"data": pickle.dumps(obj),
                      "pickle": True}
  elif class_id in custom_serializers:
    serialized_obj = {"data": custom_serializers[class_id](obj)}
  else:
    # Handle the namedtuple case.
    if is_named_tuple(type(obj)):
      serialized_obj = {}
      serialized_obj["_ray_getnewargs_"] = obj.__getnewargs__()
    elif hasattr(obj, "__dict__"):
      serialized_obj = obj.__dict__
    else:
      raise RaySerializationException("We do not know how to serialize the "
                                      "object '{}'".format(obj), obj)
  result = dict(serialized_obj, **{"_pytype_": class_id})
  return result


def deserialize(serialized_obj):
  """This is the callback that will be used by numbuf.

  If numbuf encounters a dictionary that contains the key "_pytype_" during
    deserialization, it will ask this callback to deserialize the object.

  Args:
    serialized_obj (object): A dictionary that contains the key "_pytype_".

  Returns:
    A Python object.

  Raises:
    An exception is raised if we do not know how to deserialize the object.
  """
  class_id = serialized_obj["_pytype_"]

  if "pickle" in serialized_obj:
    # The object was pickled, so unpickle it.
    obj = pickle.loads(serialized_obj["data"])
  else:
    assert class_id not in classes_to_pickle
    if class_id not in whitelisted_classes:
      # If this happens, that means that the call to _register_class, which
      # should have added the class to the list of whitelisted classes, has not
      # yet propagated to this worker. It should happen if we wait a little
      # longer.
      raise RayDeserializationException("The class {} is not one of the "
                                        "whitelisted classes."
                                        .format(class_id), class_id)
    cls = whitelisted_classes[class_id]
    if class_id in custom_deserializers:
      obj = custom_deserializers[class_id](serialized_obj["data"])
    else:
      # In this case, serialized_obj should just be the __dict__ field.
      if "_ray_getnewargs_" in serialized_obj:
        obj = cls.__new__(cls, *serialized_obj["_ray_getnewargs_"])
      else:
        obj = cls.__new__(cls)
        serialized_obj.pop("_pytype_")
        obj.__dict__.update(serialized_obj)
  return obj


def set_callbacks():
  """Register the custom callbacks with numbuf.

  The serialize callback is used to serialize objects that numbuf does not know
  how to serialize (for example custom Python classes). The deserialize
  callback is used to serialize objects that were serialized by the serialize
  callback.
  """
  ray.numbuf.register_callbacks(serialize, deserialize)


def clear_state():
  type_to_class_id.clear()
  whitelisted_classes.clear()
  classes_to_pickle.clear()
  custom_serializers.clear()
  custom_deserializers.clear()