def _register_ray_serializer(op):
# register a custom serializer for Mars operand
try:
ray.register_custom_serializer(type(op),
serializer=operand_serializer,
deserializer=operand_deserializer)
except AttributeError: # ray >= 1.0
from ray.worker import global_worker
global_worker.check_connected()
context = global_worker.get_serialization_context()
context.register_custom_serializer(type(op),
serializer=operand_serializer,
deserializer=operand_deserializer)
def register_ray_serializer(obj_type, serializer=None, deserializer=None):
ray = lazy_import("ray")
if ray:
try:
ray.register_custom_serializer(
obj_type, serializer=serializer, deserializer=deserializer)
except AttributeError: # ray >= 1.0
try:
from ray.worker import global_worker
global_worker.check_connected()
context = global_worker.get_serialization_context()
context.register_custom_serializer(
obj_type, serializer=serializer, deserializer=deserializer)
except AttributeError: # ray >= 1.2.0
ray.util.register_serializer(
obj_type, serializer=serializer, deserializer=deserializer)
def _register_sparse_matrix_serializer():
# register a custom serializer for Mars SparseMatrix
try:
ray.register_custom_serializer(
sparse.matrix.SparseMatrix,
serializer=_mars_sparse_matrix_serializer,
deserializer=_mars_sparse_matrix_deserializer)
except AttributeError: # ray >= 1.0
try:
from ray.worker import global_worker
global_worker.check_connected()
context = global_worker.get_serialization_context()
context.register_custom_serializer(
sparse.matrix.SparseMatrix,
serializer=_mars_sparse_matrix_serializer,
deserializer=_mars_sparse_matrix_deserializer)
except AttributeError: # ray >= 1.2.0
ray.util.register_serializer(
sparse.matrix.SparseMatrix,
serializer=_mars_sparse_matrix_serializer,
deserializer=_mars_sparse_matrix_deserializer)
