Support pass by ref args in for direct actor calls (#6040)

python/ray/ray_perf.py

@@ -1,27 +1,46 @@
 """This is the script for `ray microbenchmark`."""
 
+import os
 import time
 import numpy as np
 import multiprocessing
 import ray
 
+# Only run tests matching this filter pattern.
+filter_pattern = os.environ.get("TESTS_TO_RUN", "")
+
 
 @ray.remote
 class Actor(object):
     def small_value(self):
         return 0
 
+    def small_value_arg(self, x):
+        return 0
+
     def small_value_batch(self, n):
         ray.get([small_value.remote() for _ in range(n)])
 
 
 @ray.remote
 class Client(object):
-    def __init__(self, server):
-        self.server = server
+    def __init__(self, servers):
+        if not isinstance(servers, list):
+            servers = [servers]
+        self.servers = servers
 
     def small_value_batch(self, n):
-        ray.get([self.server.small_value.remote() for _ in range(n)])
+        results = []
+        for s in self.servers:
+            results.extend([s.small_value.remote() for _ in range(n)])
+        ray.get(results)
+
+    def small_value_batch_arg(self, n):
+        x = ray.put(0)
+        results = []
+        for s in self.servers:
+            results.extend([s.small_value_arg.remote(x) for _ in range(n)])
+        ray.get(results)
 
 
 @ray.remote
@@ -37,6 +56,8 @@ def small_value_batch(n):
 
 
 def timeit(name, fn, multiplier=1):
+    if filter_pattern not in name:
+        return
     # warmup
     start = time.time()
     while time.time() - start < 1:
@@ -56,6 +77,7 @@ def timeit(name, fn, multiplier=1):
 
 
 def main():
+    print("Tip: set TESTS_TO_RUN='pattern' to run a subset of benchmarks")
     ray.init()
     value = ray.put(0)
     arr = np.zeros(100 * 1024 * 1024, dtype=np.int64)
@@ -141,23 +163,17 @@ def main():
 
     timeit("multi client actor calls async", actor_multi2, m * n)
 
-    a = Actor._remote(is_direct_call=True)
-
-    def actor_sync_direct():
-        ray.get(a.small_value.remote())
-
-    timeit("single client direct actor calls sync", actor_sync_direct)
-
-    a = Actor._remote(is_direct_call=True)
-
-    def actor_async_direct():
-        ray.get([a.small_value.remote() for _ in range(1000)])
-
-    timeit("single client direct actor calls async", actor_async_direct, 1000)
-
     n = 5000
     n_cpu = multiprocessing.cpu_count() // 2
     actors = [Actor._remote(is_direct_call=True) for _ in range(n_cpu)]
+    client = Client.remote(actors)
+
+    def actor_async_direct():
+        ray.get(client.small_value_batch.remote(n))
+
+    timeit("single client direct actor calls async", actor_async_direct,
+           n * len(actors))
+
     clients = [Client.remote(a) for a in actors]
 
     def actor_multi2_direct():
@@ -166,6 +182,16 @@ def main():
     timeit("multi client direct actor calls async", actor_multi2_direct,
            n * len(clients))
 
+    n = 1000
+    actors = [Actor._remote(is_direct_call=True) for _ in range(n_cpu)]
+    clients = [Client.remote(a) for a in actors]
+
+    def actor_multi2_direct_arg():
+        ray.get([c.small_value_batch_arg.remote(n) for c in clients])
+
+    timeit("multi client direct actor calls with arg async",
+           actor_multi2_direct_arg, n * len(clients))
+
 
 if __name__ == "__main__":
     main()

python/ray/tests/test_basic.py

@@ -1223,9 +1223,70 @@ def test_direct_actor_errors(ray_start_regular):
     with pytest.raises(Exception):
         ray.get(f.remote([a.f.remote(2)]))
 
-    # by ref args not implemented
-    with pytest.raises(ray.exceptions.RayletError):
-        a.f.remote(f.remote(2))
+
+def test_direct_actor_pass_by_ref(ray_start_regular):
+    @ray.remote
+    class Actor(object):
+        def __init__(self):
+            pass
+
+        def f(self, x):
+            return x * 2
+
+    @ray.remote
+    def f(x):
+        return x
+
+    @ray.remote
+    def error():
+        sys.exit(0)
+
+    a = Actor._remote(is_direct_call=True)
+    assert ray.get(a.f.remote(f.remote(1))) == 2
+
+    fut = [a.f.remote(f.remote(i)) for i in range(100)]
+    assert ray.get(fut) == [i * 2 for i in range(100)]
+
+    # propagates errors for pass by ref
+    with pytest.raises(Exception):
+        ray.get(a.f.remote(error.remote()))
+
+
+def test_direct_actor_pass_by_ref_order_optimization(shutdown_only):
+    ray.init(num_cpus=4)
+
+    @ray.remote
+    class Actor(object):
+        def __init__(self):
+            pass
+
+        def f(self, x):
+            pass
+
+    a = Actor._remote(is_direct_call=True)
+
+    @ray.remote
+    def fast_value():
+        print("fast value")
+        pass
+
+    @ray.remote
+    def slow_value():
+        print("start sleep")
+        time.sleep(30)
+
+    @ray.remote
+    def runner(f):
+        print("runner", a, f)
+        return ray.get(a.f.remote(f.remote()))
+
+    runner.remote(slow_value)
+    time.sleep(1)
+    x2 = runner.remote(fast_value)
+    start = time.time()
+    ray.get(x2)
+    delta = time.time() - start
+    assert delta < 10, "did not skip slow value"
 
 
 def test_direct_actor_recursive(ray_start_regular):