[tune] Support object refs in with_parameters

Author: richardliaw

python/ray/tune/registry.py

@@ -168,7 +168,10 @@ def get(self, k):
 
     def flush(self):
         for k, v in self.to_flush.items():
-            self.references[k] = ray.put(v)
+            if isinstance(v, ray.ObjectRef):
+                self.references[k] = v
+            else:
+                self.references[k] = ray.put(v)
         self.to_flush.clear()
 
 

python/ray/tune/tests/test_api.py

@@ -1256,6 +1256,27 @@ def step(self):
         dumped = cp.dumps(trainable)
         assert sys.getsizeof(dumped) < 100 * 1024
 
+    def testWithParameters3(self):
+        class Data:
+            def __init__(self):
+                import numpy as np
+                self.data = np.random.rand((2 * 1024 * 1024))
+
+        class TestTrainable(Trainable):
+            def setup(self, config, data):
+                self.data = data.data
+
+            def step(self):
+                return dict(metric=len(self.data), done=True)
+
+        new_data = Data()
+        ref = ray.put(new_data)
+        trainable = tune.with_parameters(TestTrainable, data=ref)
+        # ray.cloudpickle will crash for some reason
+        import cloudpickle as cp
+        dumped = cp.dumps(trainable)
+        assert sys.getsizeof(dumped) < 100 * 1024
+
 
 class SerializabilityTest(unittest.TestCase):
     @classmethod