chore: import upstream snapshot with attribution

release/benchmarks/object_store/test_object_store.py

@@ -0,0 +1,75 @@
+import json
+import os
+from time import perf_counter
+
+import numpy as np
+from tqdm import tqdm
+
+import ray
+import ray.autoscaler.sdk
+
+NUM_NODES = 50
+OBJECT_SIZE = 2**30
+
+
+def num_alive_nodes():
+    n = 0
+    for node in ray.nodes():
+        if node["Alive"]:
+            n += 1
+    return n
+
+
+def test_object_broadcast():
+    assert num_alive_nodes() == NUM_NODES
+
+    @ray.remote(num_cpus=1, resources={"node": 1})
+    class Actor:
+        def foo(self):
+            pass
+
+        def data_len(self, arr):
+            return len(arr)
+
+    actors = [Actor.remote() for _ in range(NUM_NODES)]
+
+    arr = np.ones(OBJECT_SIZE, dtype=np.uint8)
+    ref = ray.put(arr)
+
+    for actor in tqdm(actors, desc="Ensure all actors have started."):
+        ray.get(actor.foo.remote())
+
+    start = perf_counter()
+    result_refs = []
+    for actor in tqdm(actors, desc="Broadcasting objects"):
+        result_refs.append(actor.data_len.remote(ref))
+
+    results = ray.get(result_refs)
+    end = perf_counter()
+
+    for result in results:
+        assert result == OBJECT_SIZE
+
+    return end - start
+
+
+ray.init(address="auto")
+duration = test_object_broadcast()
+print(f"Broadcast time: {duration} ({OBJECT_SIZE} B x {NUM_NODES} nodes)")
+
+if "TEST_OUTPUT_JSON" in os.environ:
+    with open(os.environ["TEST_OUTPUT_JSON"], "w") as out_file:
+        results = {
+            "broadcast_time": duration,
+            "object_size": OBJECT_SIZE,
+            "num_nodes": NUM_NODES,
+        }
+        perf_metric_name = f"time_to_broadcast_{OBJECT_SIZE}_bytes_to_{NUM_NODES}_nodes"
+        results["perf_metrics"] = [
+            {
+                "perf_metric_name": perf_metric_name,
+                "perf_metric_value": duration,
+                "perf_metric_type": "LATENCY",
+            }
+        ]
+        json.dump(results, out_file)