chore: import upstream snapshot with attribution

release/microbenchmark/experimental/compiled_graph_gpu_microbenchmark.py

@@ -0,0 +1,461 @@
+# coding: utf-8
+import io
+import json
+import logging
+import os
+import socket
+import time
+
+import cupy
+import numpy as np
+import torch
+
+import ray
+import ray.cloudpickle as pickle
+import ray.cluster_utils
+from ray._private.ray_microbenchmark_helpers import timeit
+from ray.air._internal import torch_utils
+from ray.dag import DAGContext, InputNode
+from ray.util.collective.collective_group import nccl_util
+
+logger = logging.getLogger(__name__)
+
+
+SHAPE = None
+DTYPE = torch.float16
+
+NUM_ITERS = 10
+
+
+@ray.remote
+class TorchIpcWorker:
+    def __init__(self):
+        self.device = torch_utils.get_devices()[0]
+
+    def send(self, shape, dtype, value: int):
+        t = torch.ones(shape, dtype=dtype, device=self.device) * value
+        if self.device.type == "cuda":
+            # NOTE(swang): This is needed because the IPC can get sent before
+            # the value has been written to memory. But somehow the read value
+            # is still the wrong one?
+            torch.cuda.synchronize()
+        h = cupy.cuda.runtime.ipcGetMemHandle(t.data_ptr())
+        return h
+
+    def recv(self, device_ptr, num_bytes, shape, dtype):
+        h = cupy.cuda.runtime.ipcOpenMemHandle(device_ptr)
+        m = cupy.cuda.UnownedMemory(h, num_bytes, None)
+        m_ptr = cupy.cuda.MemoryPointer(m, 0)
+        tensor = torch.tensor(cupy.ndarray(shape, dtype, m_ptr), device=self.device)
+        assert tensor.device == self.device
+        return (tensor[0].item(), tensor.shape, tensor.dtype)
+
+
+@ray.remote
+class TorchTensorWorker:
+    def __init__(self):
+        self.device = torch_utils.get_devices()[0]
+
+    def send(self, shape, dtype, _):
+        t = torch.ones(shape, dtype=dtype, device=self.device) * 1
+        return t
+
+    def recv(self, tensor):
+        # This benchmark tests the overhead of sending a tensor between
+        # actors. To minimize the overhead of shared memory transfer,
+        # we return only a byte string.
+        assert tensor.device == self.device
+        return b"x"
+
+
+@ray.remote(num_gpus=1)
+class NcclWorker:
+    def __init__(self, rank):
+        self.rank = rank
+
+    def get_node_id(self):
+        return ray.get_runtime_context().get_node_id()
+
+    def init(self, world_size):
+        from ray.air._internal import torch_utils
+
+        self.device = torch_utils.get_devices()[0]
+        self.world_size = world_size
+
+        torch.distributed.init_process_group(
+            backend="nccl",
+            world_size=world_size,
+            rank=self.rank,
+        )
+
+    def _send(self, buf, num_el, rank):
+        torch.distributed.send(buf, rank)
+
+    def _recv(self, buf, num_el, rank):
+        torch.distributed.recv(buf, rank)
+
+    def do_send_recv(self, shape, dtype):
+        other_rank = (self.rank + 1) % self.world_size
+
+        def _run():
+
+            if self.rank == 0:
+                i = np.random.randint(100)
+                input_buffer = torch.ones(shape, dtype=dtype, device=self.device) * i
+                self._send(input_buffer, input_buffer.numel(), other_rank)
+            else:
+                input_buffer = torch.empty(shape, dtype=dtype, device=self.device)
+                self._recv(input_buffer, input_buffer.numel(), other_rank)
+
+            torch.cuda.synchronize()
+
+        return timeit("exec_nccl_gpu", _run)
+
+
+def exec_ray_dag(
+    label,
+    sender,
+    receiver,
+    use_nccl=False,
+    use_cgraph=True,
+    static_shape=False,
+    direct_return=False,
+):
+    # Test torch.Tensor sent between actors.
+    with InputNode() as inp:
+        dag = sender.send.bind(SHAPE, DTYPE, inp)
+
+        if use_cgraph:
+            dag = dag.with_tensor_transport(
+                transport="nccl" if use_nccl else "auto",
+                _static_shape=static_shape,
+                _direct_return=direct_return,
+            )
+
+        dag = receiver.recv.bind(dag)
+
+    if use_cgraph:
+        dag = dag.experimental_compile()
+
+        def _run():
+            ref = dag.execute(b"x")
+            result = ray.get(ref)
+            assert result == b"x"
+
+    else:
+
+        def _run():
+            result = ray.get(dag.execute(b"x"))
+            assert result == b"x"
+
+    results = timeit(label, _run)
+
+    if use_cgraph:
+        dag.teardown()
+
+    # Workaround for Ray bug in reusing GPUs too quickly.
+    # See https://github.com/ray-project/ray/issues/44821.
+    ray.kill(sender)
+    ray.kill(receiver)
+    time.sleep(1)
+
+    return results
+
+
+def exec_ray_dag_ipc(label, sender, receiver, use_nccl=False):
+    # Test torch.Tensor sent between actors.
+    with InputNode() as inp:
+        dag = sender.send.bind(SHAPE, DTYPE, inp)
+        dag = receiver.recv.bind(
+            dag,
+            # torch.float16 has item size of 2 bytes.
+            SHAPE[0] * 2,
+            SHAPE,
+            nccl_util.TORCH_NUMPY_DTYPE_MAP[DTYPE],
+        )
+
+    compiled_dag = dag.experimental_compile(_buffer_size_bytes=int(SHAPE[0] * 3))
+    # Flag that each run can set if it sees incorrect results.
+    ok = [True]
+
+    def _run():
+        i = np.random.randint(100)
+        ref = compiled_dag.execute(i)
+        result = ray.get(ref)
+        if result != (i, SHAPE, DTYPE):
+            ok[0] = False
+
+    results = timeit(label, _run)
+
+    if not ok[0]:
+        logger.warning("IPC DAG returned incorrect result")
+    compiled_dag.teardown()
+
+    return results
+
+
+def _exec_torch_cpu_cpu():
+    i = np.random.randint(100)
+    t = torch.ones(SHAPE, dtype=DTYPE) * i
+    t2 = t.to(copy=True)
+    assert (t2[0].item(), t2.shape, t2.dtype) == (i, SHAPE, DTYPE)
+
+
+def _exec_torch_gpu():
+    i = np.random.randint(100)
+    device_from = torch.device("cuda:1")
+    device_to = torch.device("cuda:0")
+
+    t = torch.ones(SHAPE, dtype=DTYPE, device=device_from) * i
+    t2 = t.to(device_to)
+    torch.cuda.synchronize(device_to)
+    assert (t2[0].item(), t2.shape, t2.dtype) == (i, SHAPE, DTYPE)
+
+
+def exec_nccl_gpu(sender_hint, receiver_hint):
+    workers = [
+        NcclWorker.options(**sender_hint).remote(0),
+        NcclWorker.options(**receiver_hint).remote(1),
+    ]
+
+    # node_id = ray.get(workers[0].get_node_id.remote())
+    # head_node = [node for node in ray.nodes() if node["NodeID"] == node_id]
+    # assert len(head_node) == 1
+    # head_node = head_node[0]
+    # rank_0_addr = f"{head_node['NodeManagerAddress']}:8888"
+
+    ray.get([worker.init.remote(2) for worker in workers])
+
+    tasks = [worker.do_send_recv.remote(SHAPE, DTYPE) for worker in workers]
+    done_refs, _ = ray.wait(tasks, num_returns=1)
+
+    results = ray.get(done_refs[0])
+
+    # Workaround for Ray bug in reusing GPUs too quickly.
+    # See https://github.com/ray-project/ray/issues/44821.
+    for worker in workers:
+        ray.kill(worker)
+    time.sleep(1)
+
+    return results
+
+
+def _exec_torch_gpu_cpu_gpu():
+    i = np.random.randint(100)
+    device_from = torch.device("cuda:0")
+    device_to = torch.device("cuda:1")
+    t = torch.ones(SHAPE, dtype=DTYPE, device=device_from) * i
+    t = t.to("cpu")
+    t2 = t.to(device_to)
+    torch.cuda.synchronize(device_to)
+    assert (t2[0].item(), t2.shape, t2.dtype) == (i, SHAPE, DTYPE)
+
+
+def _exec_pickle_cpu():
+    i = np.random.randint(100)
+    t = torch.ones(SHAPE, dtype=DTYPE) * i
+    byte_stream = io.BytesIO()
+    pickle.dump(t, byte_stream)
+    byte_stream.seek(0)
+    pickle.load(byte_stream)
+
+
+def _exec_pickle_gpu():
+    i = np.random.randint(100)
+    t = torch.ones(SHAPE, dtype=DTYPE, device="cuda") * i
+    byte_stream = io.BytesIO()
+    pickle.dump(t, byte_stream)
+    byte_stream.seek(0)
+    pickle.load(byte_stream)
+
+
+def _exec_ray_put_cpu():
+    i = np.random.randint(100)
+    t = torch.ones(SHAPE, dtype=DTYPE) * i
+    ray.get(ray.put(t))
+
+
+def _exec_ray_put_np_zero_copy():
+    i = np.random.randint(100)
+    t = torch.ones(SHAPE, dtype=DTYPE) * i
+    torch.as_tensor(ray.get(ray.put(t.numpy())))
+
+
+def _exec_ray_put_gpu():
+    i = np.random.randint(100)
+    t = torch.ones(SHAPE, dtype=DTYPE, device="cuda") * i
+    ray.get(ray.put(t))
+
+
+def exec_ray_dag_cpu(sender_hint, receiver_hint):
+    sender = TorchTensorWorker.options(**sender_hint).remote()
+    receiver = TorchTensorWorker.options(**receiver_hint).remote()
+    return exec_ray_dag("exec_ray_dag_cpu", sender, receiver)
+
+
+def exec_ray_core_cpu(sender_hint, receiver_hint):
+    time.sleep(1)
+    sender = TorchTensorWorker.options(**sender_hint).remote()
+    receiver = TorchTensorWorker.options(**receiver_hint).remote()
+    return exec_ray_dag("exec_ray_core_cpu", sender, receiver, use_cgraph=False)
+
+
+def exec_ray_dag_gpu_ipc_gpu():
+    time.sleep(1)
+    sender = TorchIpcWorker.options(num_gpus=1).remote()
+    receiver = TorchIpcWorker.options(num_gpus=1).remote()
+    return exec_ray_dag_ipc("exec_ray_dag_gpu_ipc_gpu", sender, receiver)
+
+
+def exec_ray_dag_gpu_cpu_gpu(sender_hint, receiver_hint):
+    time.sleep(1)
+    sender = TorchTensorWorker.options(num_gpus=1, **sender_hint).remote()
+    receiver = TorchTensorWorker.options(num_gpus=1, **receiver_hint).remote()
+    return exec_ray_dag("exec_ray_dag_gpu_cpu_gpu", sender, receiver)
+
+
+def exec_ray_dag_gpu_nccl(
+    sender_hint,
+    receiver_hint,
+    static_shape: bool = False,
+    direct_return: bool = False,
+):
+    time.sleep(1)
+    sender = TorchTensorWorker.options(num_gpus=1, **sender_hint).remote()
+    receiver = TorchTensorWorker.options(num_gpus=1, **receiver_hint).remote()
+    return exec_ray_dag(
+        "exec_ray_dag_gpu_nccl"
+        + ("_static_shape" if static_shape else "")
+        + ("_direct_return" if direct_return else ""),
+        sender,
+        receiver,
+        use_nccl=True,
+        static_shape=static_shape,
+        direct_return=direct_return,
+    )
+
+
+def exec_ray_core_gpu(sender_hint, receiver_hint):
+    time.sleep(1)
+    sender = TorchTensorWorker.options(num_gpus=1, **sender_hint).remote()
+    receiver = TorchTensorWorker.options(num_gpus=1, **receiver_hint).remote()
+    return exec_ray_dag("exec_ray_core_gpu", sender, receiver, use_cgraph=False)
+
+
+def main(distributed):
+    results = []
+
+    ray.init(
+        runtime_env={
+            "env_vars": {
+                "CUDA_VISIBLE_DEVICES": "0,1",
+                # Needed for torch distributed.
+                "MASTER_ADDR": socket.gethostbyname(socket.gethostname()),
+                "MASTER_PORT": "8888",
+            }
+        }
+    )
+
+    # NCCL takes a while to warm up on multi node so increase the default
+    # timeout.
+    ctx = DAGContext.get_current()
+    ctx.get_timeout = 120
+
+    sender_hint, receiver_hint = {}, {}
+    if distributed:
+        local_node_id = ray.get_runtime_context().get_node_id()
+        node_ids = [node["NodeID"] for node in ray.nodes()]
+        remote_node_ids = [node_id for node_id in node_ids if node_id != local_node_id]
+        assert remote_node_ids
+        remote_node_id = remote_node_ids[0]
+
+        # Pin sender on local node and receiver on the other node for consistent
+        # results.
+        sender_hint = {"label_selector": {ray._raylet.RAY_NODE_ID_KEY: local_node_id}}
+        receiver_hint = {
+            "label_selector": {ray._raylet.RAY_NODE_ID_KEY: remote_node_id}
+        }
+
+    if not distributed:
+        results += timeit("exec_torch_cpu_cpu", _exec_torch_cpu_cpu)
+        results += timeit("exec_torch_gpu", _exec_torch_gpu)
+        results += timeit("exec_torch_gpu_cpu_gpu", _exec_torch_gpu_cpu_gpu)
+
+    results += exec_nccl_gpu(sender_hint, receiver_hint)
+
+    if not distributed:
+        results += timeit("exec_ray_put_cpu", _exec_ray_put_cpu)
+        results += timeit("exec_ray_put_np_zero_copy", _exec_ray_put_np_zero_copy)
+        results += timeit("exec_ray_put_gpu", _exec_ray_put_gpu)
+
+    results += exec_ray_core_cpu(sender_hint, receiver_hint)
+    results += exec_ray_dag_cpu(sender_hint, receiver_hint)
+    results += exec_ray_core_gpu(sender_hint, receiver_hint)
+    results += exec_ray_dag_gpu_cpu_gpu(sender_hint, receiver_hint)
+    results += exec_ray_dag_gpu_nccl(
+        sender_hint, receiver_hint, static_shape=True, direct_return=True
+    )
+    results += exec_ray_dag_gpu_nccl(
+        sender_hint, receiver_hint, static_shape=False, direct_return=True
+    )
+    results += exec_ray_dag_gpu_nccl(
+        sender_hint, receiver_hint, static_shape=True, direct_return=False
+    )
+    results += exec_ray_dag_gpu_nccl(
+        sender_hint, receiver_hint, static_shape=False, direct_return=False
+    )
+
+    return results
+
+
+def to_dict_key(key: str):
+    for r in [" ", ":", "-"]:
+        key = key.replace(r, "_")
+    for r in ["(", ")"]:
+        key = key.replace(r, "")
+    return key
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--tensor-size-bytes",
+        type=int,
+        # 100KB
+        default=100_000,
+    )
+    parser.add_argument(
+        "--distributed",
+        action="store_true",
+        help="Whether this is running on more than one node",
+    )
+    args = parser.parse_args()
+
+    # Divide by 2 because we're using torch.float16.
+    SHAPE = (args.tensor_size_bytes // 2,)
+
+    results = main(args.distributed)
+
+    result_dict = {
+        f"{to_dict_key(v[0])}": (v[1], v[2]) for v in results if v is not None
+    }
+
+    perf_metrics = [
+        {
+            "perf_metric_name": to_dict_key(v[0]),
+            "perf_metric_value": v[1],
+            "perf_metric_type": "THROUGHPUT",
+        }
+        for v in results
+        if v is not None
+    ]
+    result_dict["perf_metrics"] = perf_metrics
+
+    test_output_json = os.environ.get(
+        "TEST_OUTPUT_JSON", "/tmp/microbenchmark_gpu.json"
+    )
+
+    with open(test_output_json, "wt") as f:
+        json.dump(result_dict, f)

release/microbenchmark/experimental/compute_a100_gpu.yaml

@@ -0,0 +1,6 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: p4d.24xlarge
+
+worker_nodes: []

release/microbenchmark/experimental/compute_gpu_2x1_aws.yaml

@@ -0,0 +1,13 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: g4dn.4xlarge
+    resources:
+        CPU: 16
+        GPU: 1
+
+worker_nodes:
+    - instance_type: g4dn.4xlarge
+      min_nodes: 1
+      max_nodes: 1
+      market_type: ON_DEMAND

release/microbenchmark/experimental/compute_l4_gpu.yaml

@@ -0,0 +1,6 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: g6.12xlarge
+
+worker_nodes: []

release/microbenchmark/experimental/compute_l4_gpu_2x1_aws.yaml

@@ -0,0 +1,13 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: g6.4xlarge
+    resources:
+        CPU: 16
+        GPU: 1
+
+worker_nodes:
+    - instance_type: g6.4xlarge
+      min_nodes: 1
+      max_nodes: 1
+      market_type: ON_DEMAND

release/microbenchmark/experimental/compute_t4_gpu.yaml

@@ -0,0 +1,6 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: g4dn.12xlarge
+
+worker_nodes: []

release/microbenchmark/experimental/rdt_single_node_microbenchmark.py

@@ -0,0 +1,281 @@
+import argparse
+import json
+import os
+import socket
+import time
+
+import torch
+
+import ray
+from ray.experimental.collective import (
+    create_collective_group,
+    destroy_all_collective_groups,
+)
+
+ray.init(
+    runtime_env={
+        "env_vars": {
+            # Needed for torch distributed.
+            "MASTER_ADDR": socket.gethostbyname(socket.gethostname()),
+            "MASTER_PORT": "8888",
+        }
+    }
+)
+
+
+@ray.remote(num_gpus=1, enable_tensor_transport=True)
+class GPUActor:
+    def send(self, size_in_bytes, device):
+        return torch.ones(size_in_bytes, dtype=torch.int8, device=device)
+
+    def recv(self, rdt_tensor: torch.Tensor):
+        return rdt_tensor[0].item()
+
+    def init_torch(self, rank):
+        self.rank = rank
+        torch.distributed.init_process_group(
+            backend="nccl",
+            world_size=2,
+            rank=rank,
+        )
+
+    def send_with_torch(self, size_in_bytes, device, other_rank):
+        buf = torch.ones(size_in_bytes, dtype=torch.int8, device=device)
+        torch.distributed.send(buf, other_rank)
+
+    def recv_with_torch(self, size_in_bytes, device, other_rank):
+        buf = torch.empty(size_in_bytes, dtype=torch.int8, device=device)
+        torch.distributed.recv(buf, other_rank)
+        return buf[0].item()
+
+    def send_many_with_torch(self, size_in_bytes, device, other_rank, num_transfers):
+        for _ in range(num_transfers):
+            buf = torch.ones(size_in_bytes, dtype=torch.int8, device=device)
+            torch.distributed.send(buf, other_rank)
+
+    def recv_many_with_torch(self, size_in_bytes, device, other_rank, num_transfers):
+        results = []
+        for _ in range(num_transfers):
+            buf = torch.empty(size_in_bytes, dtype=torch.int8, device=device)
+            torch.distributed.recv(buf, other_rank)
+            results.append(buf[0].item())
+        return results
+
+
+"""
+THROUGHPUT
+- NEW SEND OBJECT PER RECV
+- SAME SEND OBJECT PER RECV
+LATENCY
+- JUST 1 TRANSFER
+TORCH_LATENCY
+- JUST 1 TRANSFER
+TORCH THROUGHPUT
+- NEW SEND PER RECV (all transfers done inside just 2 ray tasks)
+"""
+
+
+def throughput_new_send_per_recv(
+    num_transfers, transport, size, device, sender, receiver
+):
+    refs = []
+    ########### optional warmup
+    send_ref = sender.send.options(tensor_transport=transport).remote(size, device)
+    ray.get(receiver.recv.remote(send_ref))
+    ############
+    start = time.perf_counter()
+    for _ in range(num_transfers):
+        send_ref = sender.send.options(tensor_transport=transport).remote(size, device)
+        refs.append(receiver.recv.remote(send_ref))
+    ray.get(refs)
+    return time.perf_counter() - start
+
+
+def throughput_same_send_per_recv(
+    num_transfers, transport, size, device, sender, receiver
+):
+    refs = []
+    ########### optional warmup
+    send_ref = sender.send.options(tensor_transport=transport).remote(size, device)
+    ray.get(receiver.recv.remote(send_ref))
+    ############
+    start = time.perf_counter()
+    send_ref = sender.send.options(tensor_transport=transport).remote(size, device)
+    for _ in range(num_transfers):
+        refs.append(receiver.recv.remote(send_ref))
+    ray.get(refs)
+    return time.perf_counter() - start
+
+
+def latency_test(_num_transfers, transport, size, device, sender, receiver):
+    times = []
+    for _ in range(10):
+        start = time.perf_counter()
+        ray.get(
+            receiver.recv.remote(
+                sender.send.options(tensor_transport=transport).remote(size, device)
+            )
+        )
+        times.append(time.perf_counter() - start)
+    return sum(times) / len(times)
+
+
+def torch_latency(_num_transfers, _transport, size, device, sender, receiver):
+    times = []
+    for _ in range(10):
+        start = time.perf_counter()
+        send_ref = sender.send_with_torch.remote(size, device, 1)
+        recv_ref = receiver.recv_with_torch.remote(size, device, 0)
+        ray.get([send_ref, recv_ref])
+        times.append(time.perf_counter() - start)
+    return sum(times) / len(times)
+
+
+def torch_throughput(num_transfers, _transport, size, device, sender, receiver):
+    start_time = time.perf_counter()
+    send_ref = sender.send_many_with_torch.remote(size, device, 1, num_transfers)
+    recv_ref = receiver.recv_many_with_torch.remote(size, device, 0, num_transfers)
+    ray.get([send_ref, recv_ref])
+    return time.perf_counter() - start_time
+
+
+# torch funcs only for when directly testing torch distributed
+TEST_FUNCS = [
+    throughput_new_send_per_recv,
+    throughput_same_send_per_recv,
+    latency_test,
+    # torch_latency, added based on cli arg
+    # torch_throughput, added based on cli arg
+]
+
+# (transport, device)
+TRANSPORTS_AND_DEVICE = [
+    ("nccl", "cuda"),
+    # ("nixl", "cuda"), # nixl enabled based on cli arg
+    # ("nixl", "cpu"),
+    ("gloo", "cpu"),
+    # ("object_store", "cpu"),
+    # ("object_store", "cuda"),
+    # ("torch", "cuda") # only works with torch TEST_FUNCS, added based on cli arg
+]
+
+# (size_str, size, num_transfers)
+SIZES_AND_NUM_TRANSFERS = [
+    ("4B", 4, 50),
+    # ("1KB",   (1024),               50),
+    # ("50KB",  (50 * 1024),          50),
+    ("150KB", (150 * 1024), 50),
+    # ("500KB", (500 * 1024),         50),
+    ("1MB", (1024 * 1024), 50),
+    # ("10MB",  (10 * 1024 * 1024),   50),
+    # ("50MB",  (50 * 1024 * 1024),   50),
+    ("100MB", (100 * 1024 * 1024), 50),
+    # ("512MB", (512 * 1024 * 1024),  20),
+    ("1GB", (1024 * 1024 * 1024), 10),
+    # ("10GB", (10 * 1024 * 1024 * 1024), 1) - added based on cli arg
+]
+
+
+def do_benchmark(transport, device, test_func):
+    # Create actors + collective group
+    sender = GPUActor.remote()
+    receiver = GPUActor.remote()
+    if transport == "nccl" or transport == "gloo":
+        create_collective_group([sender, receiver], transport)
+
+    # Initialize
+    if transport == "torch":
+        ray.get([sender.init_torch.remote(0), receiver.init_torch.remote(1)])
+    else:
+        ray.get(
+            receiver.recv.remote(
+                sender.send.options(tensor_transport=transport).remote(4, device)
+            )
+        )
+
+    # Bench per size
+    bench_times = []
+    print(f"Benchmark times for transport {transport}, test_func {test_func.__name__}")
+    for size_str, size, num_transfers in SIZES_AND_NUM_TRANSFERS:
+        bench_time = test_func(num_transfers, transport, size, device, sender, receiver)
+        bench_times.append(
+            {
+                "transport": transport,
+                "test_func": test_func.__name__,
+                "num_transfers": num_transfers,
+                "size_str": size_str,
+                "bench_time": bench_time,
+            }
+        )
+
+        extra_pad = (10 - len(size_str)) * " "
+        if test_func == latency_test or test_func == torch_latency:
+            print(f"Size {size_str}{extra_pad}: {bench_time}")
+        else:
+            print(
+                f"{num_transfers} Transfers, Size {size_str}{extra_pad}: {bench_time}"
+            )
+
+        # Cool off, GC time
+        time.sleep(2)
+
+    destroy_all_collective_groups()
+    print()
+    return bench_times
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--enable_10gb",
+    action="store_true",
+)
+parser.add_argument(
+    "--enable_nixl",
+    action="store_true",
+)
+parser.add_argument(
+    "--enable_torch_bench",
+    action="store_true",
+)
+args = parser.parse_args()
+if args.enable_10gb:
+    SIZES_AND_NUM_TRANSFERS.append(("10GB", (10 * 1024 * 1024 * 1024), 1))
+
+if args.enable_nixl:
+    TRANSPORTS_AND_DEVICE.append(("nixl", "cuda"))
+
+if args.enable_torch_bench:
+    TEST_FUNCS.append(torch_latency)
+    TEST_FUNCS.append(torch_throughput)
+    TRANSPORTS_AND_DEVICE.append(("torch", "cuda"))
+
+
+bench_results = []
+for test_func in TEST_FUNCS:
+    for transport, device in TRANSPORTS_AND_DEVICE:
+        if (
+            test_func == torch_latency or test_func == torch_throughput
+        ) and transport != "torch":
+            continue
+        if transport == "torch" and (
+            test_func != torch_latency and test_func != torch_throughput
+        ):
+            continue
+        bench_results.extend(do_benchmark(transport, device, test_func))
+
+
+if "TEST_OUTPUT_JSON" in os.environ:
+    with open(os.environ["TEST_OUTPUT_JSON"], "w") as out_file:
+        # NOTE that throughput results are also returned as a time because we have to fix the amount of memory
+        # being moved to avoid GPU memory OOM's.
+        results = {}
+        results["perf_metrics"] = [
+            {
+                "perf_metric_name": f"{res['transport']}-{res['size_str']}-{res['test_func']}",
+                "perf_metric_value": res["bench_time"],
+                "perf_metric_type": "LATENCY",
+            }
+            for res in bench_results
+        ]
+
+        json.dump(results, out_file)

release/microbenchmark/run_microbenchmark.py

@@ -0,0 +1,49 @@
+import argparse
+import json
+import os
+
+
+def to_dict_key(key: str):
+    for r in [" ", ":", "-"]:
+        key = key.replace(r, "_")
+    for r in ["(", ")"]:
+        key = key.replace(r, "")
+    return key
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--experimental",
+        action="store_true",
+        default=False,
+        help="If passed, run ray.experimental microbenchmarks.",
+    )
+    args = parser.parse_args()
+
+    if args.experimental:
+        from ray._private.ray_experimental_perf import main
+    else:
+        from ray._private.ray_perf import main
+
+    results = main() or []
+
+    result_dict = {
+        f"{to_dict_key(v[0])}": (v[1], v[2]) for v in results if v is not None
+    }
+
+    perf_metrics = [
+        {
+            "perf_metric_name": to_dict_key(v[0]),
+            "perf_metric_value": v[1],
+            "perf_metric_type": "THROUGHPUT",
+        }
+        for v in results
+        if v is not None
+    ]
+    result_dict["perf_metrics"] = perf_metrics
+
+    test_output_json = os.environ.get("TEST_OUTPUT_JSON", "/tmp/microbenchmark.json")
+
+    with open(test_output_json, "wt") as f:
+        json.dump(result_dict, f)

release/microbenchmark/tpl_64.yaml

@@ -0,0 +1,6 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m5.16xlarge
+
+worker_nodes: []

release/microbenchmark/tpl_64_gce.yaml

@@ -0,0 +1,7 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+zones: [us-west1-c]
+
+head_node:
+    instance_type: n2-standard-64 # aws m5.16xlarge
+
+worker_nodes: []