Remove register_class from API. (#550)

* Perform ray.register_class under the hood.

* Fix bug.

* Release worker lock when waiting for imports to arrive in get.

* Remove calls to register_class from examples and tests.

* Clear serialization state between tests.

* Fix bug and add test for multiple custom classes with same name.

* Fix failure test.

* Fix linting and cleanups to python code.

* Fixes to documentation.

* Implement recursion depth for recursively registering classes.

* Fix linting.

* Push warning to user if waiting for class for too long.

* Fix typos.

* Don't export FunctionToRun if pickling the function fails.

* Don't broadcast class definition when pickling class.

test/actor_test.py

@@ -141,7 +141,6 @@ class ActorAPI(unittest.TestCase):
     class Foo(object):
       def __init__(self, x):
         self.x = x
-    ray.register_class(Foo)
 
     @ray.remote
     class Actor(object):

test/failure_test.py

@@ -28,42 +28,6 @@ def wait_for_errors(error_type, num_errors, timeout=10):
   print("Timing out of wait.")
 
 
-class FailureTest(unittest.TestCase):
-  def testUnknownSerialization(self):
-    reload(test_functions)
-    ray.init(num_workers=1, driver_mode=ray.SILENT_MODE)
-
-    test_functions.test_unknown_type.remote()
-    wait_for_errors(b"task", 1)
-    self.assertEqual(len(relevant_errors(b"task")), 1)
-
-    ray.worker.cleanup()
-
-
-class TaskSerializationTest(unittest.TestCase):
-  def testReturnAndPassUnknownType(self):
-    ray.init(num_workers=1, driver_mode=ray.SILENT_MODE)
-
-    class Foo(object):
-      pass
-
-    # Check that returning an unknown type from a remote function raises an
-    # exception.
-    @ray.remote
-    def f():
-      return Foo()
-    self.assertRaises(Exception, lambda: ray.get(f.remote()))
-
-    # Check that passing an unknown type into a remote function raises an
-    # exception.
-    @ray.remote
-    def g(x):
-      return 1
-    self.assertRaises(Exception, lambda: g.remote(Foo()))
-
-    ray.worker.cleanup()
-
-
 class TaskStatusTest(unittest.TestCase):
   def testFailedTask(self):
     reload(test_functions)

test/runtest.py

@@ -2,15 +2,16 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import os
-import unittest
-import ray
+from collections import defaultdict, namedtuple
 import numpy as np
-import time
+import os
+import ray
+import re
 import shutil
 import string
 import sys
-from collections import defaultdict, namedtuple
+import time
+import unittest
 
 import ray.test.test_functions as test_functions
 
@@ -169,8 +170,6 @@ class SerializationTest(unittest.TestCase):
     class ClassA(object):
       pass
 
-    ray.register_class(ClassA)
-
     # Make a list that contains itself.
     l = []
     l.append(l)
@@ -201,14 +200,6 @@ class SerializationTest(unittest.TestCase):
     def f(x):
       return x
 
-    ray.register_class(Exception)
-    ray.register_class(CustomError)
-    ray.register_class(Point)
-    ray.register_class(Foo)
-    ray.register_class(Bar)
-    ray.register_class(Baz)
-    ray.register_class(NamedTupleExample)
-
     # Check that we can pass arguments by value to remote functions and that
     # they are uncorrupted.
     for obj in RAY_TEST_OBJECTS:
@@ -303,20 +294,129 @@ class WorkerTest(unittest.TestCase):
 class APITest(unittest.TestCase):
 
   def testRegisterClass(self):
-    ray.init(num_workers=0)
+    ray.init(num_workers=2)
 
     # Check that putting an object of a class that has not been registered
     # throws an exception.
     class TempClass(object):
       pass
-    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(defaultdict))
-    # Check that registering the same class with pickle works.
-    ray.register_class(defaultdict, pickle=True)
+    ray.get(ray.put(TempClass()))
+
+    # Note that the below actually returns a dictionary and not a defaultdict.
+    # This is a bug (https://github.com/ray-project/ray/issues/512).
     ray.get(ray.put(defaultdict(lambda: 0)))
 
+    # Test passing custom classes into remote functions from the driver.
+    @ray.remote
+    def f(x):
+      return x
+
+    foo = ray.get(f.remote(Foo(7)))
+    self.assertEqual(foo, Foo(7))
+
+    regex = re.compile(r"\d+\.\d*")
+    new_regex = ray.get(f.remote(regex))
+    self.assertEqual(regex, new_regex)
+
+    # Test returning custom classes created on workers.
+    @ray.remote
+    def g():
+      return SubQux(), Qux()
+
+    subqux, qux = ray.get(g.remote())
+    self.assertEqual(subqux.objs[2].foo.value, 0)
+
+    # Test exporting custom class definitions from one worker to another when
+    # the worker is blocked in a get.
+    class NewTempClass(object):
+      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))
+
+    self.assertEqual(ray.get(h2.remote(10)).value, 10)
+
+    # Test registering multiple classes with the same name.
+    @ray.remote(num_return_vals=3)
+    def j():
+      class Class0(object):
+        def method0(self):
+          pass
+
+      c0 = Class0()
+
+      class Class0(object):
+        def method1(self):
+          pass
+
+      c1 = Class0()
+
+      class Class0(object):
+        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()
+
+      self.assertFalse(hasattr(c0, "method1"))
+      self.assertFalse(hasattr(c0, "method2"))
+      self.assertFalse(hasattr(c1, "method0"))
+      self.assertFalse(hasattr(c1, "method2"))
+      self.assertFalse(hasattr(c2, "method0"))
+      self.assertFalse(hasattr(c2, "method1"))
+
+    @ray.remote
+    def k():
+      class Class0(object):
+        def method0(self):
+          pass
+
+      c0 = Class0()
+
+      class Class0(object):
+        def method1(self):
+          pass
+
+      c1 = Class0()
+
+      class Class0(object):
+        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()
+
+      self.assertFalse(hasattr(c0, "method1"))
+      self.assertFalse(hasattr(c0, "method2"))
+      self.assertFalse(hasattr(c1, "method0"))
+      self.assertFalse(hasattr(c1, "method2"))
+      self.assertFalse(hasattr(c2, "method0"))
+      self.assertFalse(hasattr(c2, "method1"))
+
     ray.worker.cleanup()
 
   def testKeywordArgs(self):
@@ -666,35 +766,6 @@ class APITest(unittest.TestCase):
 
     ray.worker.cleanup()
 
-  def testPassingInfoToAllWorkers(self):
-    ray.init(num_workers=10, num_cpus=10)
-
-    def f(worker_info):
-      sys.path.append(worker_info)
-    ray.worker.global_worker.run_function_on_all_workers(f)
-
-    @ray.remote
-    def get_path():
-      time.sleep(1)
-      return sys.path
-    # Retrieve the values that we stored in the worker paths.
-    paths = ray.get([get_path.remote() for _ in range(10)])
-    # Add the driver's path to the list.
-    paths.append(sys.path)
-    worker_infos = [path[-1] for path in paths]
-    for worker_info in worker_infos:
-      self.assertEqual(list(worker_info.keys()), ["counter"])
-    counters = [worker_info["counter"] for worker_info in worker_infos]
-    # We use range(11) because the driver also runs the function.
-    self.assertEqual(set(counters), set(range(11)))
-
-    # Clean up the worker paths.
-    def f(worker_info):
-      sys.path.pop(-1)
-    ray.worker.global_worker.run_function_on_all_workers(f)
-
-    ray.worker.cleanup()
-
   def testLoggingAPI(self):
     ray.init(num_workers=1, driver_mode=ray.SILENT_MODE)