diff --git a/doc/source/serialization.rst b/doc/source/serialization.rst index 5a6426787..302a59664 100644 --- a/doc/source/serialization.rst +++ b/doc/source/serialization.rst @@ -117,7 +117,7 @@ message, and you should fall back to pickle. .. code-block:: python # This call tells Ray to fall back to using pickle when it encounters objects - # of type function. + # of type function (we actually already do this under the hood). f = lambda x: x + 1 ray.register_class(type(f), pickle=True) diff --git a/python/ray/worker.py b/python/ray/worker.py index e9f02b383..a54fc4277 100644 --- a/python/ray/worker.py +++ b/python/ray/worker.py @@ -805,6 +805,12 @@ def initialize_numbuf(worker=global_worker): register_class(RayTaskError) register_class(RayGetError) register_class(RayGetArgumentError) + # Tell Ray to serialize lambdas with pickle. + register_class(type(lambda: 0), pickle=True) + # Tell Ray to serialize sets with pickle. + register_class(type(set()), pickle=True) + # Tell Ray to serialize types with pickle. + register_class(type(int), pickle=True) def get_address_info_from_redis_helper(redis_address, node_ip_address): diff --git a/test/runtest.py b/test/runtest.py index c1f79ae8f..f9f7c7990 100644 --- a/test/runtest.py +++ b/test/runtest.py @@ -10,7 +10,7 @@ import time import shutil import string import sys -from collections import namedtuple +from collections import defaultdict, namedtuple import ray.test.test_functions as test_functions @@ -209,6 +209,41 @@ class SerializationTest(unittest.TestCase): ray.worker.cleanup() + def testPassingArgumentsByValueOutOfTheBox(self): + ray.init(num_workers=1) + + @ray.remote + def f(x): + return x + + # Test passing lambdas. + + def temp(): + return 1 + + self.assertEqual(ray.get(f.remote(temp))(), 1) + self.assertEqual(ray.get(f.remote(lambda x: x + 1))(3), 4) + + # Test sets. + self.assertEqual(ray.get(f.remote(set())), set()) + s = set([1, (1, 2, "hi")]) + self.assertEqual(ray.get(f.remote(s)), s) + + # Test types. + self.assertEqual(ray.get(f.remote(int)), int) + self.assertEqual(ray.get(f.remote(float)), float) + self.assertEqual(ray.get(f.remote(str)), str) + + class Foo(object): + def __init__(self): + pass + + # Make sure that we can put and get a custom type. Note that the result + # won't be "equal" to Foo. + ray.get(ray.put(Foo)) + + ray.worker.cleanup() + class WorkerTest(unittest.TestCase): @@ -267,13 +302,13 @@ class APITest(unittest.TestCase): # throws an exception. class TempClass(object): pass - self.assertRaises(Exception, lambda: ray.put(Foo)) + self.assertRaises(Exception, lambda: ray.put(TempClass())) # Check that registering a class that Ray cannot serialize efficiently # raises an exception. - self.assertRaises(Exception, lambda: ray.register_class(type(True))) + self.assertRaises(Exception, lambda: ray.register_class(defaultdict)) # Check that registering the same class with pickle works. - ray.register_class(type(float), pickle=True) - self.assertEqual(ray.get(ray.put(float)), float) + ray.register_class(defaultdict, pickle=True) + ray.get(ray.put(defaultdict(lambda: 0))) ray.worker.cleanup()