[ray_client] add client microbenchmarks (#13007)

python/ray/_private/ray_client_microbenchmark.py

@@ -0,0 +1,83 @@
+import inspect
+import logging
+import sys
+
+from ray.experimental.client.ray_client_helpers import ray_start_client_server
+
+from ray._private.ray_microbenchmark_helpers import timeit
+from ray._private.ray_microbenchmark_helpers import ray_setup_and_teardown
+
+
+def benchmark_get_calls(ray):
+    value = ray.put(0)
+
+    def get_small():
+        ray.get(value)
+
+    timeit("client: get calls", get_small)
+
+
+def benchmark_put_calls(ray):
+    def put_small():
+        ray.put(0)
+
+    timeit("client: put calls", put_small)
+
+
+def benchmark_remote_put_calls(ray):
+    @ray.remote
+    def do_put_small():
+        for _ in range(100):
+            ray.put(0)
+
+    def put_multi_small():
+        ray.get([do_put_small.remote() for _ in range(10)])
+
+    timeit("client: remote put calls", put_multi_small, 1000)
+
+
+def benchmark_simple_actor(ray):
+    @ray.remote(num_cpus=0)
+    class Actor:
+        def small_value(self):
+            return b"ok"
+
+        def small_value_arg(self, x):
+            return b"ok"
+
+        def small_value_batch(self, n):
+            ray.get([self.small_value.remote() for _ in range(n)])
+
+    a = Actor.remote()
+
+    def actor_sync():
+        ray.get(a.small_value.remote())
+
+    timeit("client: 1:1 actor calls sync", actor_sync)
+
+    def actor_async():
+        ray.get([a.small_value.remote() for _ in range(1000)])
+
+    timeit("client: 1:1 actor calls async", actor_async, 1000)
+
+    a = Actor.options(max_concurrency=16).remote()
+
+    def actor_concurrent():
+        ray.get([a.small_value.remote() for _ in range(1000)])
+
+    timeit("client: 1:1 actor calls concurrent", actor_concurrent, 1000)
+
+
+def main():
+    system_config = {"put_small_object_in_memory_store": True}
+    with ray_setup_and_teardown(
+            logging_level=logging.WARNING, _system_config=system_config):
+        for name, obj in inspect.getmembers(sys.modules[__name__]):
+            if not name.startswith("benchmark_"):
+                continue
+            with ray_start_client_server() as ray:
+                obj(ray)
+
+
+if __name__ == "__main__":
+    main()

python/ray/_private/ray_microbenchmark_helpers.py

@@ -0,0 +1,39 @@
+import time
+import os
+import ray
+import numpy as np
+
+from contextlib import contextmanager
+
+# Only run tests matching this filter pattern.
+filter_pattern = os.environ.get("TESTS_TO_RUN", "")
+
+
+def timeit(name, fn, multiplier=1):
+    if filter_pattern not in name:
+        return
+    # warmup
+    start = time.time()
+    while time.time() - start < 1:
+        fn()
+    # real run
+    stats = []
+    for _ in range(4):
+        start = time.time()
+        count = 0
+        while time.time() - start < 2:
+            fn()
+            count += 1
+        end = time.time()
+        stats.append(multiplier * count / (end - start))
+    print(name, "per second", round(np.mean(stats), 2), "+-",
+          round(np.std(stats), 2))
+
+
+@contextmanager
+def ray_setup_and_teardown(**init_args):
+    ray.init(**init_args)
+    try:
+        yield None
+    finally:
+        ray.shutdown()

python/ray/experimental/client/ray_client_helpers.py

@@ -0,0 +1,16 @@
+from contextlib import contextmanager
+
+import ray.experimental.client.server.server as ray_client_server
+from ray.experimental.client import ray, reset_api
+
+
+@contextmanager
+def ray_start_client_server():
+    server = ray_client_server.serve("localhost:50051", test_mode=True)
+    ray.connect("localhost:50051")
+    try:
+        yield ray
+    finally:
+        ray.disconnect()
+        server.stop(0)
+        reset_api()

python/ray/ray_perf.py

@@ -2,17 +2,15 @@
 
 import asyncio
 import logging
-import os
-import time
+from ray._private.ray_microbenchmark_helpers import timeit
+from ray._private.ray_client_microbenchmark import (main as
+                                                    client_microbenchmark_main)
 import numpy as np
 import multiprocessing
 import ray
 
 logger = logging.getLogger(__name__)
 
-# Only run tests matching this filter pattern.
-filter_pattern = os.environ.get("TESTS_TO_RUN", "")
-
 
 @ray.remote(num_cpus=0)
 class Actor:
@@ -71,27 +69,6 @@ def small_value_batch(n):
     return 0
 
 
-def timeit(name, fn, multiplier=1):
-    if filter_pattern not in name:
-        return
-    # warmup
-    start = time.time()
-    while time.time() - start < 1:
-        fn()
-    # real run
-    stats = []
-    for _ in range(4):
-        start = time.time()
-        count = 0
-        while time.time() - start < 2:
-            fn()
-            count += 1
-        end = time.time()
-        stats.append(multiplier * count / (end - start))
-    print(name, "per second", round(np.mean(stats), 2), "+-",
-          round(np.std(stats), 2))
-
-
 def check_optimized_build():
     if not ray._raylet.OPTIMIZED:
         msg = ("WARNING: Unoptimized build! "
@@ -277,6 +254,9 @@ def main():
         ray.get([async_actor_work.remote(a) for _ in range(m)])
 
     timeit("n:n async-actor calls async", async_actor_multi, m * n)
+    ray.shutdown()
+
+    client_microbenchmark_main()
 
 
 if __name__ == "__main__":

python/ray/tests/test_experimental_client.py

@@ -2,23 +2,9 @@ import pytest
 import time
 import sys
 import logging
-from contextlib import contextmanager
 
-import ray.experimental.client.server.server as ray_client_server
-from ray.experimental.client import ray, reset_api
 from ray.experimental.client.common import ClientObjectRef
-
-
-@contextmanager
-def ray_start_client_server():
-    server = ray_client_server.serve("localhost:50051", test_mode=True)
-    ray.connect("localhost:50051")
-    try:
-        yield ray
-    finally:
-        ray.disconnect()
-        server.stop(0)
-        reset_api()
+from ray.experimental.client.ray_client_helpers import ray_start_client_server
 
 
 def test_real_ray_fallback(ray_start_regular_shared):

python/ray/tests/test_experimental_client_metadata.py

@@ -1,4 +1,4 @@
-from ray.tests.test_experimental_client import ray_start_client_server
+from ray.experimental.client.ray_client_helpers import ray_start_client_server
 
 
 def test_get_ray_metadata(ray_start_regular_shared):

python/ray/tests/test_experimental_client_references.py

@@ -1,4 +1,4 @@
-from ray.tests.test_experimental_client import ray_start_client_server
+from ray.experimental.client.ray_client_helpers import ray_start_client_server
 from ray.test_utils import wait_for_condition
 import ray as real_ray
 from ray.core.generated.gcs_pb2 import ActorTableData

python/ray/tests/test_experimental_client_terminate.py

@@ -1,5 +1,5 @@
 import pytest
-from ray.tests.test_experimental_client import ray_start_client_server
+from ray.experimental.client.ray_client_helpers import ray_start_client_server
 from ray.tests.client_test_utils import create_remote_signal_actor
 from ray.test_utils import wait_for_condition
 from ray.exceptions import TaskCancelledError