chore: import upstream snapshot with attribution

release/benchmarks/object_store/test_large_objects.py

@@ -0,0 +1,101 @@
+import json
+import os
+from time import perf_counter
+
+import numpy as np
+from tqdm import tqdm
+
+import ray
+
+NUM_NODES = 9
+OBJECT_SIZE = 2**32
+
+
+def test_object_many_to_one():
+    @ray.remote(num_cpus=1, resources={"node": 1})
+    class Actor:
+        def foo(self):
+            pass
+
+        def send_objects(self):
+            return np.ones(OBJECT_SIZE, dtype=np.uint8)
+
+    actors = [Actor.remote() for _ in range(NUM_NODES)]
+
+    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="Tasks kickoff"):
+        result_refs.append(actor.send_objects.remote())
+
+    results = ray.get(result_refs)
+    end = perf_counter()
+
+    for result in results:
+        assert len(result) == OBJECT_SIZE
+
+    return end - start
+
+
+def test_object_one_to_many():
+    @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="Tasks kickoff"):
+        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")
+many_to_one_duration = test_object_many_to_one()
+print(f"many_to_one time: {many_to_one_duration} ({OBJECT_SIZE} B x {NUM_NODES} nodes)")
+one_to_many_duration = test_object_one_to_many()
+print(f"one_to_many time: {one_to_many_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 = {
+            "many_to_one_time": many_to_one_duration,
+            "one_to_many_time": one_to_many_duration,
+            "object_size": OBJECT_SIZE,
+            "num_nodes": NUM_NODES,
+        }
+        results["perf_metrics"] = [
+            {
+                "perf_metric_name": f"time_many_to_one_{OBJECT_SIZE}_bytes_from_{NUM_NODES}_nodes",
+                "perf_metric_value": many_to_one_duration,
+                "perf_metric_type": "LATENCY",
+            },
+            {
+                "perf_metric_name": f"time_one_to_many_{OBJECT_SIZE}_bytes_to_{NUM_NODES}_nodes",
+                "perf_metric_value": one_to_many_duration,
+                "perf_metric_type": "LATENCY",
+            },
+        ]
+
+        json.dump(results, out_file)