From 30e2b6b91bb63a65902f3efdaff44c17529f1092 Mon Sep 17 00:00:00 2001
From: Danyang Zhuo <danyangzhuo@gmail.com>
Date: Fri, 15 Nov 2019 21:39:06 -0800
Subject: [PATCH] Microbenchmark for inter-node object transfer (#6098)

---
 python/ray/ray_cluster_perf.py | 47 ++++++++++++++++++++++++++++++++++
 python/ray/scripts/scripts.py  |  6 +++++
 2 files changed, 53 insertions(+)
 create mode 100644 python/ray/ray_cluster_perf.py

diff --git a/python/ray/ray_cluster_perf.py b/python/ray/ray_cluster_perf.py
new file mode 100644
index 000000000..28869f023
--- /dev/null
+++ b/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()
diff --git a/python/ray/scripts/scripts.py b/python/ray/scripts/scripts.py
index 35b123d8a..1d5c2fb36 100644
--- a/python/ray/scripts/scripts.py
+++ b/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",