Microbenchmark for inter-node object transfer (#6098)

python/ray/ray_cluster_perf.py

@@ -0,0 +1,47 @@
+"""This is the script for `ray clusterbenchmark`."""
+
+import time
+import numpy as np
+import ray
+
+from ray.tests.cluster_utils import Cluster
+
+
+def main():
+    cluster = Cluster(
+        initialize_head=True,
+        connect=True,
+        head_node_args={
+            "object_store_memory": 20 * 1024 * 1024 * 1024,
+            "num_cpus": 16
+        })
+    cluster.add_node(
+        object_store_memory=20 * 1024 * 1024 * 1024, num_gpus=1, num_cpus=16)
+
+    object_id_list = []
+    for i in range(0, 10):
+        object_id = ray.put(np.random.rand(1024 * 128, 1024))
+        object_id_list.append(object_id)
+
+    @ray.remote(num_gpus=1)
+    def f(object_id_list):
+        diffs = []
+        for object_id in object_id_list:
+            before = time.time()
+            ray.get(object_id)
+            after = time.time()
+            diffs.append(after - before)
+            time.sleep(1)
+        return np.mean(diffs), np.std(diffs)
+
+    time_diff, time_diff_std = ray.get(f.remote(object_id_list))
+
+    print("latency to get an 1G object over network", round(time_diff, 2),
+          "+-", round(time_diff_std, 2))
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
+if __name__ == "__main__":
+    main()

python/ray/scripts/scripts.py

@@ -788,6 +788,12 @@ def microbenchmark():
     main()
 
 
+@cli.command()
+def clusterbenchmark():
+    from ray.ray_cluster_perf import main
+    main()
+
+
 @cli.command()
 @click.option(
     "--redis-address",