Files
ray-project--ray/python/ray/tests/test_serialization.py
T
2026-07-13 13:17:40 +08:00

842 lines
22 KiB
Python

# coding: utf-8
import collections
import io
import logging
import re
import string
import sys
import weakref
from dataclasses import make_dataclass
import numpy as np
import pytest
from numpy import log
import ray
import ray.cluster_utils
import ray.exceptions
from ray import cloudpickle
from ray._common.test_utils import is_named_tuple
logger = logging.getLogger(__name__)
def test_simple_serialization(ray_start_regular):
primitive_objects = [
# Various primitive types.
0,
0.0,
0.9,
1 << 62,
1 << 999,
b"",
b"a",
"a",
string.printable,
"\u262f",
"hello world",
"\xff\xfe\x9c\x001\x000\x00",
None,
True,
False,
[],
(),
{},
type,
int,
set(),
# Collections types.
collections.Counter([np.random.randint(0, 10) for _ in range(100)]),
collections.OrderedDict([("hello", 1), ("world", 2)]),
collections.defaultdict(lambda: 0, [("hello", 1), ("world", 2)]),
collections.defaultdict(lambda: [], [("hello", 1), ("world", 2)]),
collections.deque([1, 2, 3, "a", "b", "c", 3.5]),
# Numpy dtypes.
np.int8(3),
np.int32(4),
np.int64(5),
np.uint8(3),
np.uint32(4),
np.uint64(5),
np.float32(1.9),
np.float64(1.9),
]
composite_objects = (
[[obj] for obj in primitive_objects]
+ [(obj,) for obj in primitive_objects]
+ [{(): obj} for obj in primitive_objects]
)
@ray.remote
def f(x):
return x
# Check that we can pass arguments by value to remote functions and
# that they are uncorrupted.
for obj in primitive_objects + composite_objects:
new_obj_1 = ray.get(f.remote(obj))
new_obj_2 = ray.get(ray.put(obj))
assert obj == new_obj_1
assert obj == new_obj_2
# TODO(rkn): The numpy dtypes currently come back as regular integers
# or floats.
if type(obj).__module__ != "numpy":
assert type(obj) is type(new_obj_1)
assert type(obj) is type(new_obj_2)
def test_complex_serialization(ray_start_regular):
def assert_equal(obj1, obj2):
module_numpy = (
type(obj1).__module__ == np.__name__ or type(obj2).__module__ == np.__name__
)
if module_numpy:
empty_shape = (hasattr(obj1, "shape") and obj1.shape == ()) or (
hasattr(obj2, "shape") and obj2.shape == ()
)
if empty_shape:
# This is a special case because currently
# np.testing.assert_equal fails because we do not properly
# handle different numerical types.
assert obj1 == obj2, "Objects {} and {} are different.".format(
obj1, obj2
)
else:
np.testing.assert_equal(obj1, obj2)
elif hasattr(obj1, "__dict__") and hasattr(obj2, "__dict__"):
special_keys = ["_pytype_"]
assert set(list(obj1.__dict__.keys()) + special_keys) == set(
list(obj2.__dict__.keys()) + special_keys
), "Objects {} and {} are different.".format(obj1, obj2)
for key in obj1.__dict__.keys():
if key not in special_keys:
assert_equal(obj1.__dict__[key], obj2.__dict__[key])
elif type(obj1) is dict or type(obj2) is dict:
assert_equal(obj1.keys(), obj2.keys())
for key in obj1.keys():
assert_equal(obj1[key], obj2[key])
elif type(obj1) is list or type(obj2) is list:
assert len(obj1) == len(
obj2
), "Objects {} and {} are lists with different lengths.".format(obj1, obj2)
for i in range(len(obj1)):
assert_equal(obj1[i], obj2[i])
elif type(obj1) is tuple or type(obj2) is tuple:
assert len(obj1) == len(
obj2
), "Objects {} and {} are tuples with different lengths.".format(obj1, obj2)
for i in range(len(obj1)):
assert_equal(obj1[i], obj2[i])
elif is_named_tuple(type(obj1)) or is_named_tuple(type(obj2)):
assert len(obj1) == len(
obj2
), "Objects {} and {} are named tuples with different lengths.".format(
obj1, obj2
)
for i in range(len(obj1)):
assert_equal(obj1[i], obj2[i])
else:
assert obj1 == obj2, "Objects {} and {} are different.".format(obj1, obj2)
long_extras = [0, np.array([["hi", "hi"], [1.3, 1]])]
PRIMITIVE_OBJECTS = [
0,
0.0,
0.9,
1 << 62,
1 << 100,
1 << 999,
[1 << 100, [1 << 100]],
"a",
string.printable,
"\u262f",
"hello world",
"\xff\xfe\x9c\x001\x000\x00",
None,
True,
False,
[],
(),
{},
np.int8(3),
np.int32(4),
np.int64(5),
np.uint8(3),
np.uint32(4),
np.uint64(5),
np.float32(1.9),
np.float64(1.9),
np.zeros([100, 100]),
np.random.normal(size=[100, 100]),
np.array(["hi", 3]),
np.array(["hi", 3], dtype=object),
] + long_extras
COMPLEX_OBJECTS = [
[[[[[[[[[[[[]]]]]]]]]]]],
{"obj{}".format(i): np.random.normal(size=[100, 100]) for i in range(10)},
# {(): {(): {(): {(): {(): {(): {(): {(): {(): {(): {
# (): {(): {}}}}}}}}}}}}},
((((((((((),),),),),),),),),),
{"a": {"b": {"c": {"d": {}}}}},
]
class Foo:
def __init__(self, value=0):
self.value = value
def __hash__(self):
return hash(self.value)
def __eq__(self, other):
return other.value == self.value
class Bar:
def __init__(self):
for i, val in enumerate(PRIMITIVE_OBJECTS + COMPLEX_OBJECTS):
setattr(self, "field{}".format(i), val)
class Baz:
def __init__(self):
self.foo = Foo()
self.bar = Bar()
def method(self, arg):
pass
class Qux:
def __init__(self):
self.objs = [Foo(), Bar(), Baz()]
class SubQux(Qux):
def __init__(self):
Qux.__init__(self)
class CustomError(Exception):
pass
Point = collections.namedtuple("Point", ["x", "y"])
NamedTupleExample = collections.namedtuple(
"Example", "field1, field2, field3, field4, field5"
)
CUSTOM_OBJECTS = [
Exception("Test object."),
CustomError(),
Point(11, y=22),
Foo(),
Bar(),
Baz(), # Qux(), SubQux(),
NamedTupleExample(1, 1.0, "hi", np.zeros([3, 5]), [1, 2, 3]),
]
DataClass0 = make_dataclass("DataClass0", [("number", int)])
CUSTOM_OBJECTS.append(DataClass0(number=3))
class CustomClass:
def __init__(self, value):
self.value = value
DataClass1 = make_dataclass("DataClass1", [("custom", CustomClass)])
class DataClass2(DataClass1):
@classmethod
def from_custom(cls, data):
custom = CustomClass(data)
return cls(custom)
def __reduce__(self):
return (self.from_custom, (self.custom.value,))
CUSTOM_OBJECTS.append(DataClass2(custom=CustomClass(43)))
BASE_OBJECTS = PRIMITIVE_OBJECTS + COMPLEX_OBJECTS + CUSTOM_OBJECTS
LIST_OBJECTS = [[obj] for obj in BASE_OBJECTS]
TUPLE_OBJECTS = [(obj,) for obj in BASE_OBJECTS]
# The check that type(obj).__module__ != "numpy" should be unnecessary, but
# otherwise this seems to fail on Mac OS X on Travis.
DICT_OBJECTS = (
[
{obj: obj}
for obj in PRIMITIVE_OBJECTS
if (obj.__hash__ is not None and type(obj).__module__ != "numpy")
]
+ [{0: obj} for obj in BASE_OBJECTS]
+ [{Foo(123): Foo(456)}]
)
RAY_TEST_OBJECTS = BASE_OBJECTS + LIST_OBJECTS + TUPLE_OBJECTS + DICT_OBJECTS
@ray.remote
def f(x):
return x
# Check that we can pass arguments by value to remote functions and
# that they are uncorrupted.
for obj in RAY_TEST_OBJECTS:
assert_equal(obj, ray.get(f.remote(obj)))
assert_equal(obj, ray.get(ray.put(obj)))
# Test StringIO serialization
s = io.StringIO("Hello, world!\n")
s.seek(0)
line = s.readline()
s.seek(0)
assert ray.get(ray.put(s)).readline() == line
def test_numpy_serialization(ray_start_regular):
array = np.zeros(314)
from ray.cloudpickle import dumps
buffers = []
inband = dumps(array, protocol=5, buffer_callback=buffers.append)
assert len(inband) < array.nbytes
assert len(buffers) == 1
def test_numpy_subclass_serialization_pickle(ray_start_regular):
class MyNumpyConstant(np.ndarray):
def __init__(self, value):
super().__init__()
self.constant = value
def __str__(self):
print(self.constant)
constant = MyNumpyConstant(123)
repr_orig = repr(constant)
repr_ser = repr(ray.get(ray.put(constant)))
assert repr_orig == repr_ser
def test_inspect_serialization(enable_pickle_debug):
import threading
from ray.cloudpickle import dumps_debug
lock = threading.Lock()
with pytest.raises(TypeError):
dumps_debug(lock)
def test_func():
print(lock)
with pytest.raises(TypeError):
dumps_debug(test_func)
class test_class:
def test(self):
self.lock = lock
from ray.util.check_serialize import inspect_serializability
results = inspect_serializability(lock)
assert list(results[1])[0].obj == lock, results
results = inspect_serializability(test_func)
assert list(results[1])[0].obj == lock, results
results = inspect_serializability(test_class)
assert list(results[1])[0].obj == lock, results
# Test path tracking
results = inspect_serializability(test_func, name="my_func")
failures = list(results[1])
assert len(failures) == 1
path = failures[0].path
assert "my_func" in path
assert "lock" in path
results = inspect_serializability(test_class, name="my_class")
failures = list(results[1])
assert len(failures) == 1
path = failures[0].path
assert "my_class" in path
assert "lock" in path
def test_serialization_final_fallback(ray_start_regular):
pytest.importorskip("catboost")
# This test will only run when "catboost" is installed.
from catboost import CatBoostClassifier
model = CatBoostClassifier(
iterations=2,
depth=2,
learning_rate=1,
loss_function="Logloss",
logging_level="Verbose",
)
reconstructed_model = ray.get(ray.put(model))
assert set(model.get_params().items()) == set(
reconstructed_model.get_params().items()
)
def test_register_class(ray_start_2_cpus):
# Check that putting an object of a class that has not been registered
# throws an exception.
class TempClass:
pass
ray.get(ray.put(TempClass()))
# Test passing custom classes into remote functions from the driver.
@ray.remote
def f(x):
return x
class Foo:
def __init__(self, value=0):
self.value = value
def __hash__(self):
return hash(self.value)
def __eq__(self, other):
return other.value == self.value
foo = ray.get(f.remote(Foo(7)))
assert foo == Foo(7)
regex = re.compile(r"\d+\.\d*")
new_regex = ray.get(f.remote(regex))
# This seems to fail on the system Python 3 that comes with
# Ubuntu, so it is commented out for now:
# assert regex == new_regex
# Instead, we do this:
assert regex.pattern == new_regex.pattern
class TempClass1:
def __init__(self):
self.value = 1
# Test returning custom classes created on workers.
@ray.remote
def g():
class TempClass2:
def __init__(self):
self.value = 2
return TempClass1(), TempClass2()
object_1, object_2 = ray.get(g.remote())
assert object_1.value == 1
assert object_2.value == 2
# Test exporting custom class definitions from one worker to another
# when the worker is blocked in a get.
class NewTempClass:
def __init__(self, value):
self.value = value
@ray.remote
def h1(x):
return NewTempClass(x)
@ray.remote
def h2(x):
return ray.get(h1.remote(x))
assert ray.get(h2.remote(10)).value == 10
# Test registering multiple classes with the same name.
@ray.remote(num_returns=3)
def j():
class Class0:
def method0(self):
pass
c0 = Class0()
class Class0:
def method1(self):
pass
c1 = Class0()
class Class0:
def method2(self):
pass
c2 = Class0()
return c0, c1, c2
results = []
for _ in range(5):
results += j.remote()
for i in range(len(results) // 3):
c0, c1, c2 = ray.get(results[(3 * i) : (3 * (i + 1))])
c0.method0()
c1.method1()
c2.method2()
assert not hasattr(c0, "method1")
assert not hasattr(c0, "method2")
assert not hasattr(c1, "method0")
assert not hasattr(c1, "method2")
assert not hasattr(c2, "method0")
assert not hasattr(c2, "method1")
@ray.remote
def k():
class Class0:
def method0(self):
pass
c0 = Class0()
class Class0:
def method1(self):
pass
c1 = Class0()
class Class0:
def method2(self):
pass
c2 = Class0()
return c0, c1, c2
results = ray.get([k.remote() for _ in range(5)])
for c0, c1, c2 in results:
c0.method0()
c1.method1()
c2.method2()
assert not hasattr(c0, "method1")
assert not hasattr(c0, "method2")
assert not hasattr(c1, "method0")
assert not hasattr(c1, "method2")
assert not hasattr(c2, "method0")
assert not hasattr(c2, "method1")
def test_deserialized_from_buffer_immutable(ray_start_regular_shared):
x = np.full((2, 2), 1.0)
o = ray.put(x)
y = ray.get(o)
with pytest.raises(ValueError, match="assignment destination is read-only"):
y[0, 0] = 9.0
def test_reducer_override_no_reference_cycle(ray_start_regular_shared):
# bpo-39492: reducer_override used to induce a spurious reference cycle
# inside the Pickler object, that could prevent all serialized objects
# from being garbage-collected without explicity invoking gc.collect.
# test a dynamic function
def f():
return 4669201609102990671853203821578
wr = weakref.ref(f)
bio = io.BytesIO()
from ray.cloudpickle import CloudPickler, dumps, loads
p = CloudPickler(bio, protocol=5)
p.dump(f)
new_f = loads(bio.getvalue())
assert new_f() == 4669201609102990671853203821578
del p
del f
assert wr() is None
# test a dynamic class
class ShortlivedObject:
def __del__(self):
print("Went out of scope!")
obj = ShortlivedObject()
new_obj = weakref.ref(obj)
dumps(obj)
del obj
assert new_obj() is None
def test_buffer_alignment(ray_start_regular_shared):
# Deserialized large numpy arrays should be 64-byte aligned.
x = np.random.normal(size=(10, 20, 30))
y = ray.get(ray.put(x))
assert y.ctypes.data % 64 == 0
# Unlike PyArrow, Ray aligns small numpy arrays to 8
# bytes to be memory efficient.
xs = [np.random.normal(size=i) for i in range(100)]
ys = ray.get(ray.put(xs))
for y in ys:
assert y.ctypes.data % 8 == 0
xs = [np.random.normal(size=i * (1,)) for i in range(20)]
ys = ray.get(ray.put(xs))
for y in ys:
assert y.ctypes.data % 8 == 0
xs = [np.random.normal(size=i * (5,)) for i in range(1, 8)]
xs = [xs[i][(i + 1) * (slice(1, 3),)] for i in range(len(xs))]
ys = ray.get(ray.put(xs))
for y in ys:
assert y.ctypes.data % 8 == 0
def test_custom_serializer(ray_start_regular_shared):
import threading
class A:
def __init__(self, x):
self.x = x
self.lock = threading.Lock()
def custom_serializer(a):
return a.x
def custom_deserializer(x):
return A(x)
ray.util.register_serializer(
A, serializer=custom_serializer, deserializer=custom_deserializer
)
ray.get(ray.put(A(1)))
ray.util.deregister_serializer(A)
with pytest.raises(Exception):
ray.get(ray.put(A(1)))
# deregister again takes no effects
ray.util.deregister_serializer(A)
def test_numpy_ufunc(ray_start_regular_shared):
@ray.remote
def f():
# add reference to the numpy ufunc
_ = log
ray.get(f.remote())
class _SelfDereferenceObject:
"""A object that dereferences itself during deserialization"""
def __init__(self, ref: ray.ObjectRef):
self.ref = ref
def __reduce__(self):
return ray.get, (self.ref,)
def test_recursive_resolve(ray_start_regular_shared):
ref = ray.put(42)
for _ in range(10):
ref = ray.put(_SelfDereferenceObject(ref))
assert ray.get(ref) == 42
def test_serialization_before_init(shutdown_only):
"""This test checks if serializers registered before initializing Ray
works after initialization."""
# make sure ray is shutdown
ray.shutdown()
assert ray._private.worker.global_worker.current_job_id.is_nil()
import threading
class A:
def __init__(self, x):
self.x = x
self.lock = threading.Lock() # could not be serialized!
def custom_serializer(a):
return a.x
def custom_deserializer(b):
return A(b)
# Register serializer and deserializer for class A:
ray.util.register_serializer(
A, serializer=custom_serializer, deserializer=custom_deserializer
)
# Initialize Ray later.
ray.init()
ray.get(ray.put(A(1))) # success!
def test_serialization_pydantic_runtime_env(ray_start_regular):
@ray.remote
def test(pydantic_model):
return pydantic_model.x
@ray.remote(runtime_env={"pip": ["pydantic>=2"]})
def py():
from pydantic import BaseModel
class Foo(BaseModel):
x: int
return ray.get(test.remote(Foo(x=1)))
assert ray.get(py.remote()) == 1
def test_usage_with_dataclass(ray_start_regular):
import dataclasses
@dataclasses.dataclass
class Test:
v: str
@ray.remote
def test(x, expect):
assert dataclasses.asdict(x) == expect, dataclasses.asdict(x)
return x
expect_dict = {"v": "x"}
x = Test(v="x")
new_x = ray.get(test.remote(x, expect=expect_dict))
assert new_x == x
assert dataclasses.asdict(new_x) == dataclasses.asdict(x)
assert dataclasses.asdict(new_x) == expect_dict
y = Test(v="y")
expect_dict = {"v": "y"}
new_y = ray.get(test.remote(y, expect=expect_dict))
assert new_y == y
assert dataclasses.asdict(new_y) == dataclasses.asdict(y)
assert dataclasses.asdict(new_y) == expect_dict
def test_cannot_out_of_band_serialize_object_ref(shutdown_only, monkeypatch):
monkeypatch.setenv("RAY_allow_out_of_band_object_ref_serialization", "0")
ray.init()
# Use ray.remote as a workaround because
# RAY_allow_out_of_band_object_ref_serialization cannot be set dynamically.
@ray.remote
def test():
ref = ray.put(1)
@ray.remote
def f():
_ = ref
with pytest.raises(ray.exceptions.OufOfBandObjectRefSerializationException):
ray.get(f.remote())
@ray.remote
def f():
cloudpickle.dumps(ray.put(1))
with pytest.raises(ray.exceptions.OufOfBandObjectRefSerializationException):
ray.get(f.remote())
return ray.get(test.remote())
def test_can_out_of_band_serialize_object_ref_with_env_var(shutdown_only, monkeypatch):
monkeypatch.setenv("RAY_allow_out_of_band_object_ref_serialization", "1")
ray.init()
# Use ray.remote as a workaround because
# RAY_allow_out_of_band_object_ref_serialization cannot be set dynamically.
@ray.remote
def test():
ref = ray.put(1)
@ray.remote
def f():
_ = ref
ray.get(f.remote())
@ray.remote
def f():
ref = ray.put(1)
cloudpickle.dumps(ref)
ray.get(f.remote())
# It should pass.
ray.get(test.remote())
def test_inspect_serializability_warning_message_is_actionable():
"""Regression test: WARNING message should include actionable guidance,
not just say 'this may be an oversight'."""
from ray.util.check_serialize import inspect_serializability
# A custom __reduce__ that lies to cloudpickle but trips
# the traversal — produces the WARNING branch.
class Tricky:
def __reduce__(self):
raise TypeError("cannot pickle 'Tricky' object")
def method(self):
pass
output = io.StringIO()
inspect_serializability(Tricky(), print_file=output)
result = output.getvalue()
# The warning must exist (Tricky trips the warning branch)
assert "WARNING" in result
# It must now contain actionable guidance, not the old dead-end message
assert "inspect_serializability" in result
assert "This may be an oversight" not in result
def test_inspect_func_serialization_prints_qualname():
"""Regression test for https://github.com/ray-project/ray/issues/48759.
The qualified function name should appear in traversal output when
inspecting a closure that captures a non-serializable object.
Uses a two-level nested function to showcase the context printed
at each closure boundary.
"""
import io
import threading
from ray.util.check_serialize import inspect_serializability
def make_task():
lock = threading.Lock()
def inner():
return lock
def middle():
return inner()
return middle
out = io.StringIO()
serializable, _ = inspect_serializability(make_task(), print_file=out)
output = out.getvalue()
assert not serializable
assert (
"make_task.<locals>.middle':" in output
), f"Expected middle closure qualname in output, got:\n{output}"
assert (
"make_task.<locals>.inner':" in output
), f"Expected inner closure qualname in output, got:\n{output}"
if __name__ == "__main__":
sys.exit(pytest.main(["-sv", __file__]))