Export remote functions when first used and also fix bug in which rem… (#4844)

* Export remote functions when first used and also fix bug in which remote functions and actor classes are not exported from workers during subsequent ray sessions.

* Documentation update

* Fix tests.

* Fix grammar

python/ray/tests/test_basic.py

@@ -303,6 +303,23 @@ def test_complex_serialization(ray_start_regular):
         assert_equal(obj, ray.get(ray.put(obj)))
 
 
+def test_nested_functions(ray_start_regular):
+    # Make sure that remote functions can use other values that are defined
+    # after the remote function but before the first function invocation.
+    @ray.remote
+    def f():
+        return g(), ray.get(h.remote())
+
+    def g():
+        return 1
+
+    @ray.remote
+    def h():
+        return 2
+
+    assert ray.get(f.remote()) == (1, 2)
+
+
 def test_ray_recursive_objects(ray_start_regular):
     class ClassA(object):
         pass
@@ -2968,3 +2985,17 @@ def test_export_after_shutdown(ray_start_regular):
     ray.get(f.remote())
     a = Actor.remote()
     ray.get(a.method.remote())
+
+    ray.shutdown()
+
+    # Start Ray again and make sure that these definitions can be exported from
+    # workers.
+    ray.init(num_cpus=2)
+
+    @ray.remote
+    def export_definitions_from_worker(remote_function, actor_class):
+        ray.get(remote_function.remote())
+        actor_handle = actor_class.remote()
+        ray.get(actor_handle.method.remote())
+
+    ray.get(export_definitions_from_worker.remote(f, Actor))

python/ray/tests/test_failure.py

@@ -95,7 +95,15 @@ def temporary_helper_function():
     # fail when it is unpickled.
     @ray.remote
     def g():
-        return module.temporary_python_file()
+        try:
+            module.temporary_python_file()
+        except Exception:
+            # This test is not concerned with the error from running this
+            # function. Only from unpickling the remote function.
+            pass
+
+    # Invoke the function so that the definition is exported.
+    g.remote()
 
     wait_for_errors(ray_constants.REGISTER_REMOTE_FUNCTION_PUSH_ERROR, 2)
     errors = relevant_errors(ray_constants.REGISTER_REMOTE_FUNCTION_PUSH_ERROR)
@@ -499,6 +507,9 @@ def test_export_large_objects(ray_start_regular):
     def f():
         large_object
 
+    # Invoke the function so that the definition is exported.
+    f.remote()
+
     # Make sure that a warning is generated.
     wait_for_errors(ray_constants.PICKLING_LARGE_OBJECT_PUSH_ERROR, 1)
 

python/ray/tests/test_monitors.py

@@ -46,13 +46,6 @@ def _test_cleanup_on_driver_exit(num_redis_shards):
         # Two new objects.
         ray.get(ray.put(1111))
         ray.get(ray.put(1111))
-        attempts = 0
-        while (2, 1, summary_start[2]) != StateSummary():
-            time.sleep(0.1)
-            attempts += 1
-            if attempts == max_attempts_before_failing:
-                success.value = False
-                break
 
         @ray.remote
         def f():
@@ -61,7 +54,7 @@ def _test_cleanup_on_driver_exit(num_redis_shards):
 
         # 1 new function.
         attempts = 0
-        while (2, 1, summary_start[2] + 1) != StateSummary():
+        while (2, 1, summary_start[2]) != StateSummary():
             time.sleep(0.1)
             attempts += 1
             if attempts == max_attempts_before_failing: