diff --git a/python/ray/tests/test_basic.py b/python/ray/tests/test_basic.py index 1fd920dac..e19927c3f 100644 --- a/python/ray/tests/test_basic.py +++ b/python/ray/tests/test_basic.py @@ -507,7 +507,7 @@ def test_put_get(shutdown_only): assert value_before == value_after -def test_custom_serializers(ray_start_regular): +def custom_serializers(): class Foo(object): def __init__(self): self.x = 3 @@ -537,6 +537,30 @@ def test_custom_serializers(ray_start_regular): assert ray.get(f.remote()) == ((3, "string1", Bar.__name__), "string2") +def test_custom_serializers(ray_start_regular): + custom_serializers() + + +def test_custom_serializers_with_pickle(shutdown_only): + ray.init(use_pickle=True) + custom_serializers() + + class Foo(object): + def __init__(self): + self.x = 4 + + # Test the pickle serialization backend without serializer. + # NOTE: 'use_pickle' here is different from 'use_pickle' in + # ray.init + ray.register_custom_serializer(Foo, use_pickle=True) + + @ray.remote + def f(): + return Foo() + + assert type(ray.get(f.remote())) == Foo + + def test_serialization_final_fallback(ray_start_regular): pytest.importorskip("catboost") # This test will only run when "catboost" is installed. diff --git a/python/ray/worker.py b/python/ray/worker.py index 083bd9aa1..ef35798aa 100644 --- a/python/ray/worker.py +++ b/python/ray/worker.py @@ -2090,6 +2090,10 @@ def _register_custom_serializer(cls, "Exactly one of use_pickle, use_dict, or serializer/deserializer must " "be specified.") + if worker.use_pickle and serializer is None: + # In this case it should do nothing. + return + if use_dict: # Raise an exception if cls cannot be serialized efficiently by Ray. serialization.check_serializable(cls)