Custom serialization does not seem to work

I need a custom serialization for my trainable. However, tune seems to ignore my own serialization/deserialization functions.

Using the first two approaches mentioned in here: Serialization — Ray v2.0.0.dev0,
I’ve done the following brief example showing neither option works.

What am I missing? Thank you in advance.

EDIT: Forgot to mention that the commented part creating dump directly from a MyTrainable instance does work for both approaches mentioned in docs. The problem is, the serialization functions are not used in tune.run().

import sys

import ray
from ray import tune

class MyTrainable(tune.Trainable):
    def step(self):
        return { "score" : 5 }
    def __reduce__(self):
        print("Reduce invoked")
        return MyTrainable, (self.config,)

def custom_serializer(a):
    print("Serializer invoked")
    return a.config

def custom_deserializer(b):
    return MyTrainable(b)

#ray.init()
#ray.util.register_serializer(
#  MyTrainable, serializer=custom_serializer, deserializer=custom_deserializer)

#ray.cloudpickle.dumps(MyTrainable({"a" : 10}))
#sys.exit()

analysis = tune.run(
    MyTrainable,
    stop={"training_iteration": 5},
    config={
        "a": tune.choice([1,2])
    })
1 Like

Maybe try something like:

analysis = tune.run(MyTrainable()...)

__reduce__ is an instance method, and tune serializes the class (not the instance).

1 Like