chore: import upstream snapshot with attribution

release/benchmarks/object_store/callback_throughput.yaml

@@ -0,0 +1,14 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m6i.2xlarge
+    resources:
+      node: 1
+
+worker_nodes:
+    - instance_type: m6i.2xlarge
+      min_nodes: 10
+      max_nodes: 10
+      market_type: ON_DEMAND
+      resources:
+        node: 1

release/benchmarks/object_store/large_objects.yaml

@@ -0,0 +1,14 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m6i.16xlarge
+    resources:
+      node: 1
+
+worker_nodes:
+    - instance_type: m6i.4xlarge
+      min_nodes: 9
+      max_nodes: 9
+      market_type: ON_DEMAND
+      resources:
+        node: 1

release/benchmarks/object_store/small_objects.yaml

@@ -0,0 +1,10 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m6i.16xlarge
+
+worker_nodes:
+    - instance_type: m6i.4xlarge
+      min_nodes: 4
+      max_nodes: 4
+      market_type: ON_DEMAND

release/benchmarks/object_store/test_callback_throughput.py

@@ -0,0 +1,108 @@
+import json
+import os
+import threading
+import time
+
+import numpy as np
+
+import ray
+import ray._private.worker
+
+NUM_WORKERS = 10
+OBJECT_SIZE = 1024 * 1024  # 1 MiB, above the 100 KB inlining threshold
+
+
+@ray.remote(num_cpus=1)
+def produce_block():
+    return np.zeros(OBJECT_SIZE, dtype=np.uint8)
+
+
+@ray.remote(num_cpus=1)
+def consume_block(block_ref):
+    return len(block_ref)
+
+
+def test_callback_pipeline(num_blocks, timeout_s=60):
+    core_worker = ray._private.worker.global_worker.core_worker
+
+    latencies = []
+    drop_times = {}
+    lock = threading.Lock()
+    done = threading.Event()
+
+    def on_freed(id_bytes):
+        with lock:
+            latencies.append(time.perf_counter() - drop_times[id_bytes])
+            if len(latencies) == num_blocks:
+                done.set()
+
+    refs = [
+        produce_block.options(scheduling_strategy="SPREAD").remote()
+        for _ in range(num_blocks)
+    ]
+    ray.wait(refs, num_returns=len(refs))
+
+    # live_refs keeps each block ref alive until its consumer completes.
+    live_refs = {}
+    for ref in refs:
+        assert core_worker.add_object_out_of_scope_callback(ref, on_freed)
+        consumer = consume_block.remote(ref)
+        live_refs[consumer] = ref
+    del refs
+
+    # Release each ref as its consumer completes.
+    pending = list(live_refs.keys())
+    while pending:
+        done_list, pending = ray.wait(pending, num_returns=1)
+        for consumer in done_list:
+            ref = live_refs.pop(consumer)
+            drop_times[ref.binary()] = time.perf_counter()
+            del ref
+
+    if not done.wait(timeout=timeout_s):
+        raise TimeoutError(
+            f"Only {len(latencies)}/{num_blocks} callbacks fired within {timeout_s}s"
+        )
+
+    latencies.sort()
+    p95 = latencies[int(len(latencies) * 0.95)]
+    print(f"  {num_blocks} blocks: p95={p95:.4f}s")
+    return p95
+
+
+ray.init(address="auto")
+
+# Warm up gRPC connections and worker pools.
+ray.get(
+    [
+        produce_block.options(scheduling_strategy="SPREAD").remote()
+        for _ in range(NUM_WORKERS)
+    ]
+)
+
+p95_100 = test_callback_pipeline(100)
+p95_1k = test_callback_pipeline(1000)
+
+print("\nSummary:")
+print(f"  100 blocks: p95={p95_100:.4f}s")
+print(f"  1k blocks: p95={p95_1k:.4f}s")
+
+if "TEST_OUTPUT_JSON" in os.environ:
+    with open(os.environ["TEST_OUTPUT_JSON"], "w") as out_file:
+        results = {
+            "p95_100": p95_100,
+            "p95_1k": p95_1k,
+            "perf_metrics": [
+                {
+                    "perf_metric_name": "callback_p95_latency_100_blocks_s",
+                    "perf_metric_value": p95_100,
+                    "perf_metric_type": "LATENCY",
+                },
+                {
+                    "perf_metric_name": "callback_p95_latency_1k_blocks_s",
+                    "perf_metric_value": p95_1k,
+                    "perf_metric_type": "LATENCY",
+                },
+            ],
+        }
+        json.dump(results, out_file)

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)

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)

release/benchmarks/object_store/test_small_objects.py

@@ -0,0 +1,81 @@
+import json
+import os
+import time
+
+import numpy as np
+
+import ray
+
+
+def test_small_objects_many_to_one():
+    @ray.remote(num_cpus=1)
+    class Actor:
+        def send(self, _, actor_idx):
+            # this size is chosen because it's >100kb so big enough to be stored in plasma
+            numpy_arr = np.ones((20, 1024))
+            return (numpy_arr, actor_idx)
+
+    actors = [Actor.remote() for _ in range(64)]
+    not_ready = []
+    for index, actor in enumerate(actors):
+        not_ready.append(actor.send.remote(0, index))
+    num_messages = 0
+    start_time = time.time()
+    while time.time() - start_time < 60:
+        ready, not_ready = ray.wait(not_ready, num_returns=10)
+        for ready_ref in ready:
+            _, actor_idx = ray.get(ready_ref)
+            not_ready.append(actors[actor_idx].send.remote(0, actor_idx))
+        num_messages += 10
+    return num_messages / 60
+
+
+def test_small_objects_one_to_many():
+    @ray.remote(num_cpus=1)
+    class Actor:
+        def receive(self, numpy_arr, actor_idx):
+            return actor_idx
+
+    actors = [Actor.remote() for _ in range(64)]
+    numpy_arr_ref = ray.put(np.ones((20, 1024)))
+    not_ready = []
+
+    num_messages = 0
+    start_time = time.time()
+    for idx, actor in enumerate(actors):
+        not_ready.append(actor.receive.remote(numpy_arr_ref, idx))
+    while time.time() - start_time < 60:
+        ready, not_ready = ray.wait(not_ready, num_returns=10)
+        actor_idxs = ray.get(ready)
+        for actor_idx in actor_idxs:
+            not_ready.append(actors[actor_idx].receive.remote(numpy_arr_ref, actor_idx))
+        num_messages += 10
+    return num_messages / 60
+
+
+ray.init(address="auto")
+many_to_one_throughput = test_small_objects_many_to_one()
+print(f"Number of messages per second many_to_one: {many_to_one_throughput}")
+one_to_many_throughput = test_small_objects_one_to_many()
+print(f"Number of messages per second one_to_many: {one_to_many_throughput}")
+
+
+if "TEST_OUTPUT_JSON" in os.environ:
+    with open(os.environ["TEST_OUTPUT_JSON"], "w") as out_file:
+        results = {
+            "num_messages_many_to_one": many_to_one_throughput,
+            "num_messages_one_to_many": one_to_many_throughput,
+        }
+        results["perf_metrics"] = [
+            {
+                "perf_metric_name": "num_small_objects_many_to_one",
+                "perf_metric_value": many_to_one_throughput,
+                "perf_metric_type": "THROUGHPUT",
+            },
+            {
+                "perf_metric_name": "num_small_objects_one_to_many_per_second",
+                "perf_metric_value": one_to_many_throughput,
+                "perf_metric_type": "THROUGHPUT",
+            },
+        ]
+        json.dump(results, out_file)