chore: import upstream snapshot with attribution

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)