chore: import upstream snapshot with attribution

python/ray/util/collective/tests/conftest.py

@@ -0,0 +1,106 @@
+"""Some fixtures for collective tests."""
+import logging
+
+import pytest
+
+import ray
+
+try:
+    from ray.util.collective.collective_group.nccl_collective_group import (
+        _get_comm_key_from_devices,
+        _get_comm_key_send_recv,
+    )
+except Exception:  # Cupy/NCCL may be unavailable on CPU-only setups
+    _get_comm_key_from_devices = None
+    _get_comm_key_send_recv = None
+from ray.util.collective.const import get_store_name
+
+logger = logging.getLogger(__name__)
+logger.setLevel("INFO")
+
+
+# TODO (Hao): remove this clean_up function as it sometimes crashes Ray.
+def clean_up():
+    # If NCCL helpers are unavailable (e.g., no cupy), skip cleanup.
+    if _get_comm_key_from_devices is None or _get_comm_key_send_recv is None:
+        return
+    group_names = ["default", "test", "123?34!", "default2", "random"]
+    group_names.extend([str(i) for i in range(10)])
+    max_world_size = 4
+    all_keys = []
+    for name in group_names:
+        devices = [[0], [0, 1], [1, 0]]
+        for d in devices:
+            collective_communicator_key = _get_comm_key_from_devices(d)
+            all_keys.append(collective_communicator_key + "@" + name)
+        for i in range(max_world_size):
+            for j in range(max_world_size):
+                if i < j:
+                    p2p_communicator_key = _get_comm_key_send_recv(i, 0, j, 0)
+                    all_keys.append(p2p_communicator_key + "@" + name)
+    for group_key in all_keys:
+        store_name = get_store_name(group_key)
+        try:
+            actor = ray.get_actor(store_name)
+        except ValueError:
+            actor = None
+        if actor:
+            logger.debug(
+                "Killing actor with group_key: '{}' and store: '{}'.".format(
+                    group_key, store_name
+                )
+            )
+            ray.kill(actor)
+
+
+@pytest.fixture
+def ray_start_single_node_2_gpus():
+    # Please start this fixture in a cluster with 2 GPUs.
+    address_info = ray.init(num_gpus=2)
+    yield address_info
+    ray.shutdown()
+
+
+# Hao: this fixture is a bit tricky.
+# I use a bash script to start a ray cluster on
+# my own on-premise cluster before run this fixture.
+@pytest.fixture
+def ray_start_distributed_2_nodes_4_gpus():
+    # The cluster has a setup of 2 nodes, each node with 2
+    # GPUs. Each actor will be allocated 1 GPU.
+    ray.init("auto")
+    yield
+    clean_up()
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_distributed_multigpu_2_nodes_4_gpus():
+    # The cluster has a setup of 2 nodes, each node with 2
+    # GPUs. Each actor will be allocated 2 GPUs.
+    ray.init("auto")
+    yield
+    clean_up()
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_single_node():
+    address_info = ray.init(num_cpus=8)
+    yield address_info
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_distributed_2_nodes():
+    # The cluster has a setup of 2 nodes.
+    # no GPUs!
+    ray.init("auto")
+    yield
+    ray.shutdown()
+
+
+@pytest.fixture
+def shutdown_only():
+    yield None
+    ray.shutdown()

python/ray/util/collective/tests/cpu_util.py

@@ -0,0 +1,137 @@
+import logging
+
+import numpy as np
+import torch
+
+import ray
+import ray.util.collective as col
+from ray.util.collective.types import Backend, ReduceOp
+
+logger = logging.getLogger(__name__)
+
+
+@ray.remote(num_cpus=1)
+class Worker:
+    def __init__(self):
+        self.buffer = None
+        self.list_buffer = None
+
+    def init_tensors(self):
+        self.buffer = np.ones((10,), dtype=np.float32)
+        self.list_buffer = [np.ones((10,), dtype=np.float32) for _ in range(2)]
+        return True
+
+    def init_group(self, world_size, rank, backend=Backend.NCCL, group_name="default"):
+        col.init_collective_group(world_size, rank, backend, group_name)
+        return True
+
+    def set_buffer(self, data):
+        self.buffer = data
+        return self.buffer
+
+    def get_buffer(self):
+        return self.buffer
+
+    def set_list_buffer(self, list_of_arrays, copy=False):
+        if copy:
+            copy_list = []
+            for tensor in list_of_arrays:
+                if isinstance(tensor, np.ndarray):
+                    copy_list.append(tensor.copy())
+                elif isinstance(tensor, torch.Tensor):
+                    copy_list.append(tensor.clone().detach())
+            self.list_buffer = copy_list
+        else:
+            self.list_buffer = list_of_arrays
+        return self.list_buffer
+
+    def do_allreduce(self, group_name="default", op=ReduceOp.SUM):
+        col.allreduce(self.buffer, group_name, op)
+        return self.buffer
+
+    def do_reduce(self, group_name="default", dst_rank=0, op=ReduceOp.SUM):
+        col.reduce(self.buffer, dst_rank, group_name, op)
+        return self.buffer
+
+    def do_broadcast(self, group_name="default", src_rank=0):
+        col.broadcast(self.buffer, src_rank, group_name)
+        return self.buffer
+
+    def do_allgather(self, group_name="default"):
+        col.allgather(self.list_buffer, self.buffer, group_name)
+        return self.list_buffer
+
+    def do_reducescatter(self, group_name="default", op=ReduceOp.SUM):
+        col.reducescatter(self.buffer, self.list_buffer, group_name, op)
+        return self.buffer
+
+    def do_send(self, group_name="default", dst_rank=0):
+        col.send(self.buffer, dst_rank, group_name)
+        return self.buffer
+
+    def do_recv(self, group_name="default", src_rank=0):
+        col.recv(self.buffer, src_rank, group_name)
+        return self.buffer
+
+    def destroy_group(self, group_name="default"):
+        col.destroy_collective_group(group_name)
+        return True
+
+    def report_rank(self, group_name="default"):
+        rank = col.get_rank(group_name)
+        return rank
+
+    def report_world_size(self, group_name="default"):
+        ws = col.get_collective_group_size(group_name)
+        return ws
+
+    def report_nccl_availability(self):
+        avail = col.nccl_available()
+        return avail
+
+    def report_gloo_availability(self):
+        avail = col.gloo_available()
+        return avail
+
+    def report_is_group_initialized(self, group_name="default"):
+        is_init = col.is_group_initialized(group_name)
+        return is_init
+
+
+def create_collective_workers(num_workers=2, group_name="default", backend="nccl"):
+    actors = [None] * num_workers
+    for i in range(num_workers):
+        actor = Worker.remote()
+        ray.get([actor.init_tensors.remote()])
+        actors[i] = actor
+    world_size = num_workers
+    init_results = ray.get(
+        [
+            actor.init_group.remote(world_size, i, backend, group_name)
+            for i, actor in enumerate(actors)
+        ]
+    )
+    return actors, init_results
+
+
+def init_tensors_for_gather_scatter(
+    actors, array_size=10, dtype=np.float32, tensor_backend="numpy"
+):
+    world_size = len(actors)
+    for i, a in enumerate(actors):
+        if tensor_backend == "numpy":
+            t = np.ones(array_size, dtype=dtype) * (i + 1)
+        elif tensor_backend == "torch":
+            t = torch.ones(array_size, dtype=torch.float32) * (i + 1)
+        else:
+            raise RuntimeError("Unsupported tensor backend.")
+        ray.get([a.set_buffer.remote(t)])
+    if tensor_backend == "numpy":
+        list_buffer = [np.ones(array_size, dtype=dtype) for _ in range(world_size)]
+    elif tensor_backend == "torch":
+        list_buffer = [
+            torch.ones(array_size, dtype=torch.float32) for _ in range(world_size)
+        ]
+    else:
+        raise RuntimeError("Unsupported tensor backend.")
+    ray.get([a.set_list_buffer.remote(list_buffer, copy=True) for a in actors])

python/ray/util/collective/tests/distributed_cpu_tests/test_distributed_allgather.py

@@ -0,0 +1,143 @@
+"""Test the allgather API on a distributed Ray cluster."""
+import numpy as np
+import pytest
+import torch
+
+import ray
+from ray.util.collective.tests.cpu_util import (
+    create_collective_workers,
+    init_tensors_for_gather_scatter,
+)
+from ray.util.collective.types import Backend
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("tensor_backend", ["numpy", "torch"])
+@pytest.mark.parametrize(
+    "array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
+)
+def test_allgather_different_array_size(
+    ray_start_distributed_2_nodes, array_size, tensor_backend, backend
+):
+    world_size = 8
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    init_tensors_for_gather_scatter(
+        actors, array_size=array_size, tensor_backend=tensor_backend
+    )
+    results = ray.get([a.do_allgather.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(world_size):
+            if tensor_backend == "numpy":
+                assert (
+                    results[i][j] == np.ones(array_size, dtype=np.float32) * (j + 1)
+                ).all()
+            else:
+                assert (
+                    results[i][j]
+                    == torch.ones(array_size, dtype=torch.float32) * (j + 1)
+                ).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("dtype", [np.uint8, np.float16, np.float32, np.float64])
+def test_allgather_different_dtype(ray_start_distributed_2_nodes, dtype, backend):
+    world_size = 8
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    init_tensors_for_gather_scatter(actors, dtype=dtype)
+    results = ray.get([a.do_allgather.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(world_size):
+            assert (results[i][j] == np.ones(10, dtype=dtype) * (j + 1)).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("length", [0, 1, 3, 4, 7, 8])
+def test_unmatched_tensor_list_length(ray_start_distributed_2_nodes, length, backend):
+    world_size = 8
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    list_buffer = [np.ones(10, dtype=np.float32) for _ in range(length)]
+    ray.wait([a.set_list_buffer.remote(list_buffer, copy=True) for a in actors])
+    if length != world_size:
+        with pytest.raises(RuntimeError):
+            ray.get([a.do_allgather.remote() for a in actors])
+    else:
+        ray.get([a.do_allgather.remote() for a in actors])
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("shape", [10, 20, [4, 5], [1, 3, 5, 7]])
+def test_unmatched_tensor_shape(ray_start_distributed_2_nodes, shape, backend):
+    world_size = 8
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    init_tensors_for_gather_scatter(actors, array_size=10)
+    list_buffer = [np.ones(shape, dtype=np.float32) for _ in range(world_size)]
+    ray.get([a.set_list_buffer.remote(list_buffer, copy=True) for a in actors])
+    if shape != 10:
+        with pytest.raises(RuntimeError):
+            ray.get([a.do_allgather.remote() for a in actors])
+    else:
+        ray.get([a.do_allgather.remote() for a in actors])
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+def test_allgather_torch_numpy(ray_start_distributed_2_nodes, backend):
+    world_size = 8
+    shape = [10, 10]
+    actors, _ = create_collective_workers(world_size, backend=backend)
+
+    # tensor is pytorch, list is numpy
+    for i, a in enumerate(actors):
+        t = torch.ones(shape, dtype=torch.float32) * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = [np.ones(shape, dtype=np.float32) for _ in range(world_size)]
+        ray.wait([a.set_list_buffer.remote(list_buffer, copy=True)])
+    results = ray.get([a.do_allgather.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(world_size):
+            assert (results[i][j] == np.ones(shape, dtype=np.float32) * (j + 1)).all()
+
+    # tensor is numpy, list is pytorch
+    for i, a in enumerate(actors):
+        t = np.ones(shape, dtype=np.float32) * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = [
+            torch.ones(shape, dtype=torch.float32) for _ in range(world_size)
+        ]
+        ray.wait([a.set_list_buffer.remote(list_buffer)])
+    results = ray.get([a.do_allgather.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(world_size):
+            assert (
+                results[i][j] == torch.ones(shape, dtype=torch.float32) * (j + 1)
+            ).all()
+
+    # some tensors in the list are pytorch, some are numpy
+    for i, a in enumerate(actors):
+        t = np.ones(shape, dtype=np.float32) * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = []
+        for j in range(world_size):
+            if j % 2 == 0:
+                list_buffer.append(torch.ones(shape, dtype=torch.float32))
+            else:
+                list_buffer.append(np.ones(shape, dtype=np.float32))
+        ray.wait([a.set_list_buffer.remote(list_buffer, copy=True)])
+    results = ray.get([a.do_allgather.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(world_size):
+            if j % 2 == 0:
+                assert (
+                    results[i][j] == torch.ones(shape, dtype=torch.float32) * (j + 1)
+                ).all()
+            else:
+                assert (
+                    results[i][j] == np.ones(shape, dtype=np.float32) * (j + 1)
+                ).all()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/util/collective/tests/distributed_cpu_tests/test_distributed_allreduce.py

@@ -0,0 +1,179 @@
+"""Test the collective allreduice API on a distributed Ray cluster."""
+import numpy as np
+import pytest
+import torch
+
+import ray
+from ray.util.collective.tests.cpu_util import create_collective_workers
+from ray.util.collective.types import Backend, ReduceOp
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
+@pytest.mark.parametrize("world_size", [5, 6, 7, 8])
+def test_allreduce_different_name(
+    ray_start_distributed_2_nodes, group_name, world_size, backend
+):
+    actors, _ = create_collective_workers(
+        num_workers=world_size, group_name=group_name, backend=backend
+    )
+    results = ray.get([a.do_allreduce.remote(group_name) for a in actors])
+    assert (results[0] == np.ones((10,), dtype=np.float32) * world_size).all()
+    assert (results[1] == np.ones((10,), dtype=np.float32) * world_size).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
+def test_allreduce_different_array_size(
+    ray_start_distributed_2_nodes, array_size, backend
+):
+    world_size = 8
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    ray.wait(
+        [a.set_buffer.remote(np.ones(array_size, dtype=np.float32)) for a in actors]
+    )
+    results = ray.get([a.do_allreduce.remote() for a in actors])
+    assert (results[0] == np.ones((array_size,), dtype=np.float32) * world_size).all()
+    assert (results[1] == np.ones((array_size,), dtype=np.float32) * world_size).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+def test_allreduce_destroy(
+    ray_start_distributed_2_nodes, backend, group_name="default"
+):
+    world_size = 8
+    actors, _ = create_collective_workers(world_size, backend=backend)
+
+    results = ray.get([a.do_allreduce.remote() for a in actors])
+    assert (results[0] == np.ones((10,), dtype=np.float32) * world_size).all()
+    assert (results[1] == np.ones((10,), dtype=np.float32) * world_size).all()
+
+    # destroy the group and try do work, should fail
+    ray.get([a.destroy_group.remote() for a in actors])
+    with pytest.raises(RuntimeError):
+        results = ray.get([a.do_allreduce.remote() for a in actors])
+
+    # reinit the same group and all reduce
+    ray.get(
+        [
+            actor.init_group.remote(world_size, i, backend, group_name)
+            for i, actor in enumerate(actors)
+        ]
+    )
+    results = ray.get([a.do_allreduce.remote() for a in actors])
+    assert (
+        results[0] == np.ones((10,), dtype=np.float32) * world_size * world_size
+    ).all()
+    assert (
+        results[1] == np.ones((10,), dtype=np.float32) * world_size * world_size
+    ).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+def test_allreduce_multiple_group(ray_start_distributed_2_nodes, backend, num_groups=5):
+    world_size = 8
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    for group_name in range(1, num_groups):
+        ray.get(
+            [
+                actor.init_group.remote(world_size, i, backend, str(group_name))
+                for i, actor in enumerate(actors)
+            ]
+        )
+    for i in range(num_groups):
+        group_name = "default" if i == 0 else str(i)
+        results = ray.get([a.do_allreduce.remote(group_name) for a in actors])
+        assert (
+            results[0] == np.ones((10,), dtype=np.float32) * (world_size ** (i + 1))
+        ).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+def test_allreduce_different_op(ray_start_distributed_2_nodes, backend):
+    world_size = 8
+    actors, _ = create_collective_workers(world_size, backend=backend)
+
+    # check product
+    ray.wait(
+        [
+            a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get([a.do_allreduce.remote(op=ReduceOp.PRODUCT) for a in actors])
+    product = 1
+    for i in range(world_size):
+        product = product * (i + 2)
+    assert (results[0] == np.ones((10,), dtype=np.float32) * product).all()
+    assert (results[1] == np.ones((10,), dtype=np.float32) * product).all()
+
+    # check min
+    ray.wait(
+        [
+            a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get([a.do_allreduce.remote(op=ReduceOp.MIN) for a in actors])
+    assert (results[0] == np.ones((10,), dtype=np.float32) * 2).all()
+    assert (results[1] == np.ones((10,), dtype=np.float32) * 2).all()
+
+    # check max
+    ray.wait(
+        [
+            a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get([a.do_allreduce.remote(op=ReduceOp.MAX) for a in actors])
+    assert (results[0] == np.ones((10,), dtype=np.float32) * 9).all()
+    assert (results[1] == np.ones((10,), dtype=np.float32) * 9).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("dtype", [np.uint8, np.float16, np.float32, np.float64])
+def test_allreduce_different_dtype(ray_start_distributed_2_nodes, dtype, backend):
+    world_size = 8
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    ray.wait([a.set_buffer.remote(np.ones(10, dtype=dtype)) for a in actors])
+    results = ray.get([a.do_allreduce.remote() for a in actors])
+    assert (results[0] == np.ones((10,), dtype=dtype) * world_size).all()
+    assert (results[1] == np.ones((10,), dtype=dtype) * world_size).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+def test_allreduce_torch_numpy(ray_start_distributed_2_nodes, backend):
+    # import torch
+    world_size = 8
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    ray.wait(
+        [
+            actors[1].set_buffer.remote(
+                torch.ones(
+                    10,
+                )
+            )
+        ]
+    )
+    results = ray.get([a.do_allreduce.remote() for a in actors])
+    assert (results[0] == np.ones((10,)) * world_size).all()
+
+    ray.wait(
+        [
+            actors[0].set_buffer.remote(
+                torch.ones(
+                    10,
+                )
+            )
+        ]
+    )
+    ray.wait([actors[1].set_buffer.remote(np.ones(10, dtype=np.float32))])
+    results = ray.get([a.do_allreduce.remote() for a in actors])
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/util/collective/tests/distributed_cpu_tests/test_distributed_basic_apis.py

@@ -0,0 +1,138 @@
+"""Test the collective group APIs."""
+from random import shuffle
+
+import pytest
+
+import ray
+from ray.util.collective.tests.cpu_util import Worker, create_collective_workers
+from ray.util.collective.types import Backend
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("world_size", [2, 3, 4])
+@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
+def test_init_two_actors(
+    ray_start_distributed_2_nodes, world_size, group_name, backend
+):
+    actors, results = create_collective_workers(world_size, group_name, backend=backend)
+    for i in range(world_size):
+        assert results[i]
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("world_size", [2, 3, 4])
+def test_init_multiple_groups(ray_start_distributed_2_nodes, world_size, backend):
+    num_groups = 5
+    actors = [Worker.remote() for _ in range(world_size)]
+    for i in range(num_groups):
+        group_name = str(i)
+        init_results = ray.get(
+            [
+                actor.init_group.remote(
+                    world_size, i, group_name=group_name, backend=backend
+                )
+                for i, actor in enumerate(actors)
+            ]
+        )
+        for j in range(world_size):
+            assert init_results[j]
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("world_size", [5, 6, 7, 8])
+def test_get_rank(ray_start_distributed_2_nodes, world_size, backend):
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    actor0_rank = ray.get(actors[0].report_rank.remote())
+    assert actor0_rank == 0
+    actor1_rank = ray.get(actors[1].report_rank.remote())
+    assert actor1_rank == 1
+
+    # create a second group with a different name, and different
+    # orders of ranks.
+    new_group_name = "default2"
+    ranks = list(range(world_size))
+    shuffle(ranks)
+    _ = ray.get(
+        [
+            actor.init_group.remote(
+                world_size, ranks[i], group_name=new_group_name, backend=backend
+            )
+            for i, actor in enumerate(actors)
+        ]
+    )
+    actor0_rank = ray.get(actors[0].report_rank.remote(new_group_name))
+    assert actor0_rank == ranks[0]
+    actor1_rank = ray.get(actors[1].report_rank.remote(new_group_name))
+    assert actor1_rank == ranks[1]
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("world_size", [5, 6, 7, 8])
+def test_get_world_size(ray_start_distributed_2_nodes, world_size, backend):
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    actor0_world_size = ray.get(actors[0].report_world_size.remote())
+    actor1_world_size = ray.get(actors[1].report_world_size.remote())
+    assert actor0_world_size == actor1_world_size == world_size
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+def test_is_group_initialized(ray_start_distributed_2_nodes, backend):
+    world_size = 8
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    # check group is_init
+    actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
+    assert actor0_is_init
+    actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("random"))
+    assert not actor0_is_init
+    actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("123"))
+    assert not actor0_is_init
+    actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
+    assert actor1_is_init
+    actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote("456"))
+    assert not actor1_is_init
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+def test_destroy_group(ray_start_distributed_2_nodes, backend):
+    world_size = 8
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    # Now destroy the group at actor0
+    ray.wait([actors[0].destroy_group.remote()])
+    actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
+    assert not actor0_is_init
+
+    # should go well as the group `random` does not exist at all
+    ray.wait([actors[0].destroy_group.remote("random")])
+
+    actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
+    assert actor1_is_init
+    ray.wait([actors[1].destroy_group.remote("random")])
+    actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
+    assert actor1_is_init
+    ray.wait([actors[1].destroy_group.remote("default")])
+    actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
+    assert not actor1_is_init
+    for i in range(2, world_size):
+        ray.wait([actors[i].destroy_group.remote("default")])
+
+    # Now reconstruct the group using the same name
+    init_results = ray.get(
+        [
+            actor.init_group.remote(world_size, i, backend=backend)
+            for i, actor in enumerate(actors)
+        ]
+    )
+    for i in range(world_size):
+        assert init_results[i]
+    actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
+    assert actor0_is_init
+    actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
+    assert actor1_is_init
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/util/collective/tests/distributed_cpu_tests/test_distributed_broadcast.py

@@ -0,0 +1,96 @@
+"""Test the broadcast API."""
+import numpy as np
+import pytest
+
+import ray
+from ray.util.collective.tests.cpu_util import create_collective_workers
+from ray.util.collective.types import Backend
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
+@pytest.mark.parametrize("src_rank", [0, 2, 5, 6, 7])
+def test_broadcast_different_name(
+    ray_start_distributed_2_nodes, group_name, src_rank, backend
+):
+    world_size = 8
+    actors, _ = create_collective_workers(
+        num_workers=world_size, group_name=group_name, backend=backend
+    )
+    ray.wait(
+        [
+            a.set_buffer.remote(np.ones((10,), dtype=np.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get(
+        [
+            a.do_broadcast.remote(group_name=group_name, src_rank=src_rank)
+            for a in actors
+        ]
+    )
+    for i in range(world_size):
+        assert (results[i] == np.ones((10,), dtype=np.float32) * (src_rank + 2)).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
+@pytest.mark.parametrize("src_rank", [0, 2, 5, 6, 7])
+def test_broadcast_different_array_size(
+    ray_start_distributed_2_nodes, array_size, src_rank, backend
+):
+    world_size = 8
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    ray.wait(
+        [
+            a.set_buffer.remote(np.ones(array_size, dtype=np.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
+    for i in range(world_size):
+        assert (
+            results[i] == np.ones((array_size,), dtype=np.float32) * (src_rank + 2)
+        ).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("src_rank", [0, 2, 5, 6, 7])
+def test_broadcast_torch_numpy(ray_start_distributed_2_nodes, src_rank, backend):
+    import torch
+
+    world_size = 8
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    ray.wait(
+        [
+            actors[1].set_buffer.remote(
+                torch.ones(
+                    10,
+                )
+                * world_size
+            )
+        ]
+    )
+    results = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
+    if src_rank == 0:
+        assert (results[0] == np.ones((10,))).all()
+        assert (results[1] == torch.ones((10,))).all()
+    else:
+        assert (results[0] == np.ones((10,)) * world_size).all()
+        assert (results[1] == torch.ones((10,)) * world_size).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+def test_broadcast_invalid_rank(ray_start_distributed_2_nodes, backend, src_rank=9):
+    world_size = 8
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    with pytest.raises(ValueError):
+        _ = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/util/collective/tests/distributed_cpu_tests/test_distributed_reduce.py

@@ -0,0 +1,147 @@
+"""Test the reduce API."""
+import numpy as np
+import pytest
+
+import ray
+from ray.util.collective.tests.cpu_util import create_collective_workers
+from ray.util.collective.types import Backend, ReduceOp
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
+@pytest.mark.parametrize("dst_rank", [0, 2, 5, 6, 7])
+def test_reduce_different_name(
+    ray_start_distributed_2_nodes, group_name, backend, dst_rank
+):
+    world_size = 8
+    actors, _ = create_collective_workers(
+        num_workers=world_size, group_name=group_name, backend=backend
+    )
+    results = ray.get([a.do_reduce.remote(group_name, dst_rank) for a in actors])
+    for i in range(world_size):
+        if i == dst_rank:
+            assert (results[i] == np.ones((10,), dtype=np.float32) * world_size).all()
+        else:
+            assert (results[i] == np.ones((10,), dtype=np.float32)).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
+@pytest.mark.parametrize("dst_rank", [0, 2, 5, 6, 7])
+def test_reduce_different_array_size(
+    ray_start_distributed_2_nodes, backend, array_size, dst_rank
+):
+    world_size = 8
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    ray.wait(
+        [a.set_buffer.remote(np.ones(array_size, dtype=np.float32)) for a in actors]
+    )
+    results = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
+    for i in range(world_size):
+        if i == dst_rank:
+            assert (
+                results[i] == np.ones((array_size,), dtype=np.float32) * world_size
+            ).all()
+        else:
+            assert (results[i] == np.ones((array_size,), dtype=np.float32)).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("dst_rank", [0, 2, 5, 6, 7])
+def test_reduce_different_op(ray_start_distributed_2_nodes, backend, dst_rank):
+    world_size = 8
+    actors, _ = create_collective_workers(world_size, backend=backend)
+
+    # check product
+    ray.wait(
+        [
+            a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get(
+        [a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.PRODUCT) for a in actors]
+    )
+
+    product = 1
+    for i in range(world_size):
+        product = product * (i + 2)
+    for i in range(world_size):
+        if i == dst_rank:
+            assert (results[i] == np.ones((10,), dtype=np.float32) * product).all()
+        else:
+            assert (results[i] == np.ones((10,), dtype=np.float32) * (i + 2)).all()
+    # check min
+    ray.wait(
+        [
+            a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get(
+        [a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.MIN) for a in actors]
+    )
+    for i in range(world_size):
+        if i == dst_rank:
+            assert (results[i] == np.ones((10,), dtype=np.float32) * 2).all()
+        else:
+            assert (results[i] == np.ones((10,), dtype=np.float32) * (i + 2)).all()
+
+    # check max
+    ray.wait(
+        [
+            a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get(
+        [a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.MAX) for a in actors]
+    )
+    for i in range(world_size):
+        if i == dst_rank:
+            assert (
+                results[i] == np.ones((10,), dtype=np.float32) * (world_size + 1)
+            ).all()
+        else:
+            assert (results[i] == np.ones((10,), dtype=np.float32) * (i + 2)).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("dst_rank", [0, 2, 5, 6, 7])
+def test_reduce_torch_numpy(ray_start_distributed_2_nodes, backend, dst_rank):
+    import torch
+
+    world_size = 8
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    ray.get(
+        [
+            actors[1].set_buffer.remote(
+                torch.ones(
+                    10,
+                )
+            )
+        ]
+    )
+    results = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
+    if dst_rank == 0:
+        assert (results[0] == np.ones((10,)) * world_size).all()
+        assert (results[1] == torch.ones((10,))).all()
+    else:
+        assert (results[0] == np.ones((10,))).all()
+        assert (results[1] == torch.ones((10,))).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+def test_reduce_invalid_rank(ray_start_distributed_2_nodes, backend, dst_rank=9):
+    world_size = 8
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    with pytest.raises(ValueError):
+        _ = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/util/collective/tests/distributed_cpu_tests/test_distributed_reducescatter.py

@@ -0,0 +1,131 @@
+"""Test the collective reducescatter API on a distributed Ray cluster."""
+import numpy as np
+import pytest
+import torch
+
+import ray
+from ray.util.collective.tests.cpu_util import (
+    create_collective_workers,
+    init_tensors_for_gather_scatter,
+)
+from ray.util.collective.types import Backend
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("tensor_backend", ["numpy", "torch"])
+@pytest.mark.parametrize(
+    "array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
+)
+def test_reducescatter_different_array_size(
+    ray_start_distributed_2_nodes, array_size, tensor_backend, backend
+):
+    world_size = 8
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    init_tensors_for_gather_scatter(
+        actors, array_size=array_size, tensor_backend=tensor_backend
+    )
+    results = ray.get([a.do_reducescatter.remote() for a in actors])
+    for i in range(world_size):
+        if tensor_backend == "numpy":
+            assert (
+                results[i] == np.ones(array_size, dtype=np.float32) * world_size
+            ).all()
+        else:
+            assert (
+                results[i] == torch.ones(array_size, dtype=torch.float32) * world_size
+            ).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("dtype", [np.uint8, np.float16, np.float32, np.float64])
+def test_reducescatter_different_dtype(ray_start_distributed_2_nodes, dtype, backend):
+    world_size = 8
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    init_tensors_for_gather_scatter(actors, dtype=dtype)
+    results = ray.get([a.do_reducescatter.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(world_size):
+            assert (results[i] == np.ones(10, dtype=dtype) * world_size).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+def test_reducescatter_torch_numpy(ray_start_distributed_2_nodes, backend):
+    world_size = 8
+    shape = [10, 10]
+    actors, _ = create_collective_workers(world_size, backend=backend)
+
+    # tensor is pytorch, list is numpy
+    for i, a in enumerate(actors):
+        t = torch.ones(shape, dtype=torch.float32) * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = [np.ones(shape, dtype=np.float32) for _ in range(world_size)]
+        ray.wait([a.set_list_buffer.remote(list_buffer)])
+    results = ray.get([a.do_reducescatter.remote() for a in actors])
+    for i in range(world_size):
+        assert (results[i] == torch.ones(shape, dtype=torch.float32) * world_size).all()
+
+    # tensor is numpy, list is pytorch
+    for i, a in enumerate(actors):
+        t = np.ones(shape, dtype=np.float32) * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = [
+            torch.ones(shape, dtype=torch.float32) for _ in range(world_size)
+        ]
+        ray.wait([a.set_list_buffer.remote(list_buffer)])
+    results = ray.get([a.do_reducescatter.remote() for a in actors])
+    for i in range(world_size):
+        assert (results[i] == np.ones(shape, dtype=np.float32) * world_size).all()
+
+    # some tensors in the list are pytorch, some are numpy
+    for i, a in enumerate(actors):
+        if i % 2 == 0:
+            t = torch.ones(shape, dtype=torch.float32) * (i + 1)
+        else:
+            t = np.ones(shape, dtype=np.float32) * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = []
+        for j in range(world_size):
+            if j % 2 == 0:
+                list_buffer.append(torch.ones(shape, dtype=torch.float32))
+            else:
+                list_buffer.append(np.ones(shape, dtype=np.float32))
+        ray.wait([a.set_list_buffer.remote(list_buffer)])
+    results = ray.get([a.do_reducescatter.remote() for a in actors])
+    for i in range(world_size):
+        if i % 2 == 0:
+            assert (
+                results[i] == torch.ones(shape, dtype=torch.float32) * world_size
+            ).all()
+        else:
+            assert (results[i] == np.ones(shape, dtype=np.float32) * world_size).all()
+
+    # mixed case
+    for i, a in enumerate(actors):
+        if i % 2 == 0:
+            t = torch.ones(shape, dtype=torch.float32) * (i + 1)
+        else:
+            t = np.ones(shape, dtype=np.float32) * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = []
+        for j in range(world_size):
+            if j % 2 == 0:
+                list_buffer.append(np.ones(shape, dtype=np.float32))
+            else:
+                list_buffer.append(torch.ones(shape, dtype=torch.float32))
+        ray.wait([a.set_list_buffer.remote(list_buffer)])
+    results = ray.get([a.do_reducescatter.remote() for a in actors])
+    for i in range(world_size):
+        if i % 2 == 0:
+            assert (
+                results[i] == torch.ones(shape, dtype=torch.float32) * world_size
+            ).all()
+        else:
+            assert (results[i] == np.ones(shape, dtype=np.float32) * world_size).all()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/util/collective/tests/distributed_cpu_tests/test_distributed_sendrecv.py

@@ -0,0 +1,52 @@
+"""Test the send/recv API."""
+import numpy as np
+import pytest
+
+import ray
+from ray.util.collective.tests.cpu_util import create_collective_workers
+from ray.util.collective.types import Backend
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
+@pytest.mark.parametrize("dst_rank", [0, 1, 3, 6])
+@pytest.mark.parametrize("src_rank", [0, 2, 4, 7])
+@pytest.mark.parametrize(
+    "array_size", [2**10, 2**15, 2**20, [2, 2], [5, 9, 10, 85]]
+)
+def test_sendrecv(
+    ray_start_distributed_2_nodes, group_name, array_size, src_rank, dst_rank, backend
+):
+    if src_rank == dst_rank:
+        return
+    world_size = 8
+    actors, _ = create_collective_workers(
+        num_workers=world_size, group_name=group_name, backend=backend
+    )
+    ray.get(
+        [
+            a.set_buffer.remote(np.ones(array_size, dtype=np.float32) * (i + 1))
+            for i, a in enumerate(actors)
+        ]
+    )
+    refs = []
+    for i in range(world_size):
+        refs.append(actors[i].get_buffer.remote())
+    refs[src_rank] = actors[src_rank].do_send.remote(group_name, dst_rank)
+    refs[dst_rank] = actors[dst_rank].do_recv.remote(group_name, src_rank)
+    results = ray.get(refs)
+    assert (
+        results[src_rank] == np.ones(array_size, dtype=np.float32) * (src_rank + 1)
+    ).all()
+    assert (
+        results[dst_rank] == np.ones(array_size, dtype=np.float32) * (src_rank + 1)
+    ).all()
+    ray.get([a.destroy_group.remote(group_name) for a in actors])
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/util/collective/tests/distributed_gpu_tests/test_distributed_allgather.py

@@ -0,0 +1,138 @@
+"""Test the allgather API on a distributed Ray cluster."""
+import cupy as cp
+import pytest
+import torch
+
+import ray
+from ray.util.collective.tests.util import (
+    create_collective_workers,
+    init_tensors_for_gather_scatter,
+)
+
+
+@pytest.mark.parametrize("tensor_backend", ["cupy", "torch"])
+@pytest.mark.parametrize(
+    "array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
+)
+def test_allgather_different_array_size(
+    ray_start_distributed_2_nodes_4_gpus, array_size, tensor_backend
+):
+    world_size = 4
+    actors, _ = create_collective_workers(world_size)
+    init_tensors_for_gather_scatter(
+        actors, array_size=array_size, tensor_backend=tensor_backend
+    )
+    results = ray.get([a.do_allgather.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(world_size):
+            if tensor_backend == "cupy":
+                assert (
+                    results[i][j] == cp.ones(array_size, dtype=cp.float32) * (j + 1)
+                ).all()
+            else:
+                assert (
+                    results[i][j]
+                    == torch.ones(array_size, dtype=torch.float32).cuda() * (j + 1)
+                ).all()
+
+
+@pytest.mark.parametrize("dtype", [cp.uint8, cp.float16, cp.float32, cp.float64])
+def test_allgather_different_dtype(ray_start_distributed_2_nodes_4_gpus, dtype):
+    world_size = 4
+    actors, _ = create_collective_workers(world_size)
+    init_tensors_for_gather_scatter(actors, dtype=dtype)
+    results = ray.get([a.do_allgather.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(world_size):
+            assert (results[i][j] == cp.ones(10, dtype=dtype) * (j + 1)).all()
+
+
+@pytest.mark.parametrize("length", [0, 1, 3, 4, 7, 8])
+def test_unmatched_tensor_list_length(ray_start_distributed_2_nodes_4_gpus, length):
+    world_size = 4
+    actors, _ = create_collective_workers(world_size)
+    list_buffer = [cp.ones(10, dtype=cp.float32) for _ in range(length)]
+    ray.wait([a.set_list_buffer.remote(list_buffer) for a in actors])
+    if length != world_size:
+        with pytest.raises(RuntimeError):
+            ray.get([a.do_allgather.remote() for a in actors])
+    else:
+        ray.get([a.do_allgather.remote() for a in actors])
+
+
+@pytest.mark.parametrize("shape", [10, 20, [4, 5], [1, 3, 5, 7]])
+def test_unmatched_tensor_shape(ray_start_distributed_2_nodes_4_gpus, shape):
+    world_size = 4
+    actors, _ = create_collective_workers(world_size)
+    init_tensors_for_gather_scatter(actors, array_size=10)
+    list_buffer = [cp.ones(shape, dtype=cp.float32) for _ in range(world_size)]
+    ray.get([a.set_list_buffer.remote(list_buffer) for a in actors])
+    if shape != 10:
+        with pytest.raises(RuntimeError):
+            ray.get([a.do_allgather.remote() for a in actors])
+    else:
+        ray.get([a.do_allgather.remote() for a in actors])
+
+
+def test_allgather_torch_cupy(ray_start_distributed_2_nodes_4_gpus):
+    world_size = 4
+    shape = [10, 10]
+    actors, _ = create_collective_workers(world_size)
+
+    # tensor is pytorch, list is cupy
+    for i, a in enumerate(actors):
+        t = torch.ones(shape, dtype=torch.float32).cuda() * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = [cp.ones(shape, dtype=cp.float32) for _ in range(world_size)]
+        ray.wait([a.set_list_buffer.remote(list_buffer)])
+    results = ray.get([a.do_allgather.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(world_size):
+            assert (results[i][j] == cp.ones(shape, dtype=cp.float32) * (j + 1)).all()
+
+    # tensor is cupy, list is pytorch
+    for i, a in enumerate(actors):
+        t = cp.ones(shape, dtype=cp.float32) * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = [
+            torch.ones(shape, dtype=torch.float32).cuda() for _ in range(world_size)
+        ]
+        ray.wait([a.set_list_buffer.remote(list_buffer)])
+    results = ray.get([a.do_allgather.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(world_size):
+            assert (
+                results[i][j] == torch.ones(shape, dtype=torch.float32).cuda() * (j + 1)
+            ).all()
+
+    # some tensors in the list are pytorch, some are cupy
+    for i, a in enumerate(actors):
+        t = cp.ones(shape, dtype=cp.float32) * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = []
+        for j in range(world_size):
+            if j % 2 == 0:
+                list_buffer.append(torch.ones(shape, dtype=torch.float32).cuda())
+            else:
+                list_buffer.append(cp.ones(shape, dtype=cp.float32))
+        ray.wait([a.set_list_buffer.remote(list_buffer)])
+    results = ray.get([a.do_allgather.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(world_size):
+            if j % 2 == 0:
+                assert (
+                    results[i][j]
+                    == torch.ones(shape, dtype=torch.float32).cuda() * (j + 1)
+                ).all()
+            else:
+                assert (
+                    results[i][j] == cp.ones(shape, dtype=cp.float32) * (j + 1)
+                ).all()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/util/collective/tests/distributed_gpu_tests/test_distributed_allreduce.py

@@ -0,0 +1,170 @@
+"""Test the collective allreduice API on a distributed Ray cluster."""
+import cupy as cp
+import pytest
+import torch
+
+import ray
+from ray.util.collective.tests.util import create_collective_workers
+from ray.util.collective.types import ReduceOp
+
+
+@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
+@pytest.mark.parametrize("world_size", [2, 3, 4])
+def test_allreduce_different_name(
+    ray_start_distributed_2_nodes_4_gpus, group_name, world_size
+):
+    actors, _ = create_collective_workers(num_workers=world_size, group_name=group_name)
+    results = ray.get([a.do_allreduce.remote(group_name) for a in actors])
+    assert (results[0] == cp.ones((10,), dtype=cp.float32) * world_size).all()
+    assert (results[1] == cp.ones((10,), dtype=cp.float32) * world_size).all()
+
+
+@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
+def test_allreduce_different_array_size(
+    ray_start_distributed_2_nodes_4_gpus, array_size
+):
+    world_size = 4
+    actors, _ = create_collective_workers(world_size)
+    ray.wait(
+        [a.set_buffer.remote(cp.ones(array_size, dtype=cp.float32)) for a in actors]
+    )
+    results = ray.get([a.do_allreduce.remote() for a in actors])
+    assert (results[0] == cp.ones((array_size,), dtype=cp.float32) * world_size).all()
+    assert (results[1] == cp.ones((array_size,), dtype=cp.float32) * world_size).all()
+
+
+def test_allreduce_destroy(
+    ray_start_distributed_2_nodes_4_gpus, backend="nccl", group_name="default"
+):
+    world_size = 4
+    actors, _ = create_collective_workers(world_size)
+
+    results = ray.get([a.do_allreduce.remote() for a in actors])
+    assert (results[0] == cp.ones((10,), dtype=cp.float32) * world_size).all()
+    assert (results[1] == cp.ones((10,), dtype=cp.float32) * world_size).all()
+
+    # destroy the group and try do work, should fail
+    ray.get([a.destroy_group.remote() for a in actors])
+    with pytest.raises(RuntimeError):
+        results = ray.get([a.do_allreduce.remote() for a in actors])
+
+    # reinit the same group and all reduce
+    ray.get(
+        [
+            actor.init_group.remote(world_size, i, backend, group_name)
+            for i, actor in enumerate(actors)
+        ]
+    )
+    results = ray.get([a.do_allreduce.remote() for a in actors])
+    assert (
+        results[0] == cp.ones((10,), dtype=cp.float32) * world_size * world_size
+    ).all()
+    assert (
+        results[1] == cp.ones((10,), dtype=cp.float32) * world_size * world_size
+    ).all()
+
+
+def test_allreduce_multiple_group(
+    ray_start_distributed_2_nodes_4_gpus, backend="nccl", num_groups=5
+):
+    world_size = 4
+    actors, _ = create_collective_workers(world_size)
+    for group_name in range(1, num_groups):
+        ray.get(
+            [
+                actor.init_group.remote(world_size, i, backend, str(group_name))
+                for i, actor in enumerate(actors)
+            ]
+        )
+    for i in range(num_groups):
+        group_name = "default" if i == 0 else str(i)
+        results = ray.get([a.do_allreduce.remote(group_name) for a in actors])
+        assert (
+            results[0] == cp.ones((10,), dtype=cp.float32) * (world_size ** (i + 1))
+        ).all()
+
+
+def test_allreduce_different_op(ray_start_distributed_2_nodes_4_gpus):
+    world_size = 4
+    actors, _ = create_collective_workers(world_size)
+
+    # check product
+    ray.wait(
+        [
+            a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get([a.do_allreduce.remote(op=ReduceOp.PRODUCT) for a in actors])
+    assert (results[0] == cp.ones((10,), dtype=cp.float32) * 120).all()
+    assert (results[1] == cp.ones((10,), dtype=cp.float32) * 120).all()
+
+    # check min
+    ray.wait(
+        [
+            a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get([a.do_allreduce.remote(op=ReduceOp.MIN) for a in actors])
+    assert (results[0] == cp.ones((10,), dtype=cp.float32) * 2).all()
+    assert (results[1] == cp.ones((10,), dtype=cp.float32) * 2).all()
+
+    # check max
+    ray.wait(
+        [
+            a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get([a.do_allreduce.remote(op=ReduceOp.MAX) for a in actors])
+    assert (results[0] == cp.ones((10,), dtype=cp.float32) * 5).all()
+    assert (results[1] == cp.ones((10,), dtype=cp.float32) * 5).all()
+
+
+@pytest.mark.parametrize("dtype", [cp.uint8, cp.float16, cp.float32, cp.float64])
+def test_allreduce_different_dtype(ray_start_distributed_2_nodes_4_gpus, dtype):
+    world_size = 4
+    actors, _ = create_collective_workers(world_size)
+    ray.wait([a.set_buffer.remote(cp.ones(10, dtype=dtype)) for a in actors])
+    results = ray.get([a.do_allreduce.remote() for a in actors])
+    assert (results[0] == cp.ones((10,), dtype=dtype) * world_size).all()
+    assert (results[1] == cp.ones((10,), dtype=dtype) * world_size).all()
+
+
+def test_allreduce_torch_cupy(ray_start_distributed_2_nodes_4_gpus):
+    # import torch
+    world_size = 4
+    actors, _ = create_collective_workers(world_size)
+    ray.wait(
+        [
+            actors[1].set_buffer.remote(
+                torch.ones(
+                    10,
+                ).cuda()
+            )
+        ]
+    )
+    results = ray.get([a.do_allreduce.remote() for a in actors])
+    assert (results[0] == cp.ones((10,)) * world_size).all()
+
+    ray.wait(
+        [
+            actors[0].set_buffer.remote(
+                torch.ones(
+                    10,
+                )
+            )
+        ]
+    )
+    ray.wait(
+        [
+            actors[1].set_buffer.remote(
+                cp.ones(
+                    10,
+                )
+            )
+        ]
+    )
+    with pytest.raises(RuntimeError):
+        results = ray.get([a.do_allreduce.remote() for a in actors])

python/ray/util/collective/tests/distributed_gpu_tests/test_distributed_basic_apis.py

@@ -0,0 +1,122 @@
+"""Test the collective group APIs."""
+from random import shuffle
+
+import pytest
+
+import ray
+from ray.util.collective.tests.util import Worker, create_collective_workers
+
+
+@pytest.mark.parametrize("world_size", [2, 3, 4])
+@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
+def test_init_two_actors(ray_start_distributed_2_nodes_4_gpus, world_size, group_name):
+    actors, results = create_collective_workers(world_size, group_name)
+    for i in range(world_size):
+        assert results[i]
+
+
+@pytest.mark.parametrize("world_size", [2, 3, 4])
+def test_init_multiple_groups(ray_start_distributed_2_nodes_4_gpus, world_size):
+    num_groups = 1
+    actors = [Worker.remote() for _ in range(world_size)]
+    for i in range(num_groups):
+        group_name = str(i)
+        init_results = ray.get(
+            [
+                actor.init_group.remote(world_size, k, group_name=group_name)
+                for k, actor in enumerate(actors)
+            ]
+        )
+        for j in range(world_size):
+            assert init_results[j]
+
+
+@pytest.mark.parametrize("world_size", [2, 3, 4])
+def test_get_rank(ray_start_distributed_2_nodes_4_gpus, world_size):
+    actors, _ = create_collective_workers(world_size)
+    actor0_rank = ray.get(actors[0].report_rank.remote())
+    assert actor0_rank == 0
+    actor1_rank = ray.get(actors[1].report_rank.remote())
+    assert actor1_rank == 1
+
+    # create a second group with a different name, and different
+    # orders of ranks.
+    new_group_name = "default2"
+    ranks = list(range(world_size))
+    shuffle(ranks)
+    ray.get(
+        [
+            actor.init_group.remote(world_size, ranks[i], group_name=new_group_name)
+            for i, actor in enumerate(actors)
+        ]
+    )
+    actor0_rank = ray.get(actors[0].report_rank.remote(new_group_name))
+    assert actor0_rank == ranks[0]
+    actor1_rank = ray.get(actors[1].report_rank.remote(new_group_name))
+    assert actor1_rank == ranks[1]
+
+
+@pytest.mark.parametrize("world_size", [2, 3, 4])
+def test_get_collective_group_size(ray_start_distributed_2_nodes_4_gpus, world_size):
+    actors, _ = create_collective_workers(world_size)
+    actor0_world_size = ray.get(actors[0].report_world_size.remote())
+    actor1_world_size = ray.get(actors[1].report_world_size.remote())
+    assert actor0_world_size == actor1_world_size == world_size
+
+
+def test_is_group_initialized(ray_start_distributed_2_nodes_4_gpus):
+    world_size = 4
+    actors, _ = create_collective_workers(world_size)
+    # check group is_init
+    actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
+    assert actor0_is_init
+    actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("random"))
+    assert not actor0_is_init
+    actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("123"))
+    assert not actor0_is_init
+    actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
+    assert actor1_is_init
+    actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote("456"))
+    assert not actor1_is_init
+
+
+def test_destroy_group(ray_start_distributed_2_nodes_4_gpus):
+    world_size = 4
+    actors, _ = create_collective_workers(world_size)
+    # Now destroy the group at actor0
+    ray.wait([actors[0].destroy_group.remote()])
+    actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
+    assert not actor0_is_init
+
+    # should go well as the group `random` does not exist at all
+    ray.wait([actors[0].destroy_group.remote("random")])
+
+    actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
+    assert actor1_is_init
+    ray.wait([actors[1].destroy_group.remote("random")])
+    actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
+    assert actor1_is_init
+    ray.wait([actors[1].destroy_group.remote("default")])
+    actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
+    assert not actor1_is_init
+    for i in [2, 3]:
+        ray.wait([actors[i].destroy_group.remote("default")])
+
+    # Now reconstruct the group using the same name
+    init_results = ray.get(
+        [actor.init_group.remote(world_size, i) for i, actor in enumerate(actors)]
+    )
+    for i in range(world_size):
+        assert init_results[i]
+    actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
+    assert actor0_is_init
+    actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
+    assert actor1_is_init
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/util/collective/tests/distributed_gpu_tests/test_distributed_broadcast.py

@@ -0,0 +1,81 @@
+"""Test the broadcast API."""
+import cupy as cp
+import pytest
+
+import ray
+from ray.util.collective.tests.util import create_collective_workers
+
+
+@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
+@pytest.mark.parametrize("src_rank", [0, 1, 2, 3])
+def test_broadcast_different_name(
+    ray_start_distributed_2_nodes_4_gpus, group_name, src_rank
+):
+    world_size = 4
+    actors, _ = create_collective_workers(num_workers=world_size, group_name=group_name)
+    ray.wait(
+        [
+            a.set_buffer.remote(cp.ones((10,), dtype=cp.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get(
+        [
+            a.do_broadcast.remote(group_name=group_name, src_rank=src_rank)
+            for a in actors
+        ]
+    )
+    for i in range(world_size):
+        assert (results[i] == cp.ones((10,), dtype=cp.float32) * (src_rank + 2)).all()
+
+
+@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
+@pytest.mark.parametrize("src_rank", [0, 1, 2, 3])
+def test_broadcast_different_array_size(
+    ray_start_distributed_2_nodes_4_gpus, array_size, src_rank
+):
+    world_size = 4
+    actors, _ = create_collective_workers(world_size)
+    ray.wait(
+        [
+            a.set_buffer.remote(cp.ones(array_size, dtype=cp.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
+    for i in range(world_size):
+        assert (
+            results[i] == cp.ones((array_size,), dtype=cp.float32) * (src_rank + 2)
+        ).all()
+
+
+@pytest.mark.parametrize("src_rank", [0, 1])
+def test_broadcast_torch_cupy(ray_start_distributed_2_nodes_4_gpus, src_rank):
+    import torch
+
+    world_size = 4
+    actors, _ = create_collective_workers(world_size)
+    ray.wait(
+        [
+            actors[1].set_buffer.remote(
+                torch.ones(
+                    10,
+                ).cuda()
+                * world_size
+            )
+        ]
+    )
+    results = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
+    if src_rank == 0:
+        assert (results[0] == cp.ones((10,))).all()
+        assert (results[1] == torch.ones((10,)).cuda()).all()
+    else:
+        assert (results[0] == cp.ones((10,)) * world_size).all()
+        assert (results[1] == torch.ones((10,)).cuda() * world_size).all()
+
+
+def test_broadcast_invalid_rank(ray_start_distributed_2_nodes_4_gpus, src_rank=3):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+    with pytest.raises(ValueError):
+        ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])

python/ray/util/collective/tests/distributed_gpu_tests/test_distributed_reduce.py

@@ -0,0 +1,127 @@
+"""Test the reduce API."""
+import cupy as cp
+import pytest
+
+import ray
+from ray.util.collective.tests.util import create_collective_workers
+from ray.util.collective.types import ReduceOp
+
+
+@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
+@pytest.mark.parametrize("dst_rank", [0, 1, 2, 3])
+def test_reduce_different_name(
+    ray_start_distributed_2_nodes_4_gpus, group_name, dst_rank
+):
+    world_size = 4
+    actors, _ = create_collective_workers(num_workers=world_size, group_name=group_name)
+    results = ray.get([a.do_reduce.remote(group_name, dst_rank) for a in actors])
+    for i in range(world_size):
+        if i == dst_rank:
+            assert (results[i] == cp.ones((10,), dtype=cp.float32) * world_size).all()
+        else:
+            assert (results[i] == cp.ones((10,), dtype=cp.float32)).all()
+
+
+@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
+@pytest.mark.parametrize("dst_rank", [0, 1, 2, 3])
+def test_reduce_different_array_size(
+    ray_start_distributed_2_nodes_4_gpus, array_size, dst_rank
+):
+    world_size = 4
+    actors, _ = create_collective_workers(world_size)
+    ray.wait(
+        [a.set_buffer.remote(cp.ones(array_size, dtype=cp.float32)) for a in actors]
+    )
+    results = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
+    for i in range(world_size):
+        if i == dst_rank:
+            assert (
+                results[i] == cp.ones((array_size,), dtype=cp.float32) * world_size
+            ).all()
+        else:
+            assert (results[i] == cp.ones((array_size,), dtype=cp.float32)).all()
+
+
+@pytest.mark.parametrize("dst_rank", [0, 1, 2, 3])
+def test_reduce_different_op(ray_start_distributed_2_nodes_4_gpus, dst_rank):
+    world_size = 4
+    actors, _ = create_collective_workers(world_size)
+
+    # check product
+    ray.wait(
+        [
+            a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get(
+        [a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.PRODUCT) for a in actors]
+    )
+    for i in range(world_size):
+        if i == dst_rank:
+            assert (results[i] == cp.ones((10,), dtype=cp.float32) * 120).all()
+        else:
+            assert (results[i] == cp.ones((10,), dtype=cp.float32) * (i + 2)).all()
+
+    # check min
+    ray.wait(
+        [
+            a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get(
+        [a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.MIN) for a in actors]
+    )
+    for i in range(world_size):
+        if i == dst_rank:
+            assert (results[i] == cp.ones((10,), dtype=cp.float32) * 2).all()
+        else:
+            assert (results[i] == cp.ones((10,), dtype=cp.float32) * (i + 2)).all()
+
+    # check max
+    ray.wait(
+        [
+            a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get(
+        [a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.MAX) for a in actors]
+    )
+    for i in range(world_size):
+        if i == dst_rank:
+            assert (results[i] == cp.ones((10,), dtype=cp.float32) * 5).all()
+        else:
+            assert (results[i] == cp.ones((10,), dtype=cp.float32) * (i + 2)).all()
+
+
+@pytest.mark.parametrize("dst_rank", [0, 1])
+def test_reduce_torch_cupy(ray_start_distributed_2_nodes_4_gpus, dst_rank):
+    import torch
+
+    world_size = 4
+    actors, _ = create_collective_workers(world_size)
+    ray.wait(
+        [
+            actors[1].set_buffer.remote(
+                torch.ones(
+                    10,
+                ).cuda()
+            )
+        ]
+    )
+    results = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
+    if dst_rank == 0:
+        assert (results[0] == cp.ones((10,)) * world_size).all()
+        assert (results[1] == torch.ones((10,)).cuda()).all()
+    else:
+        assert (results[0] == cp.ones((10,))).all()
+        assert (results[1] == torch.ones((10,)).cuda() * world_size).all()
+
+
+def test_reduce_invalid_rank(ray_start_distributed_2_nodes_4_gpus, dst_rank=7):
+    world_size = 4
+    actors, _ = create_collective_workers(world_size)
+    with pytest.raises(ValueError):
+        ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])

python/ray/util/collective/tests/distributed_gpu_tests/test_distributed_reducescatter.py

@@ -0,0 +1,130 @@
+"""Test the collective reducescatter API on a distributed Ray cluster."""
+import cupy as cp
+import pytest
+import torch
+
+import ray
+from ray.util.collective.tests.util import (
+    create_collective_workers,
+    init_tensors_for_gather_scatter,
+)
+
+
+@pytest.mark.parametrize("tensor_backend", ["cupy", "torch"])
+@pytest.mark.parametrize(
+    "array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
+)
+def test_reducescatter_different_array_size(
+    ray_start_distributed_2_nodes_4_gpus, array_size, tensor_backend
+):
+    world_size = 4
+    actors, _ = create_collective_workers(world_size)
+    init_tensors_for_gather_scatter(
+        actors, array_size=array_size, tensor_backend=tensor_backend
+    )
+    results = ray.get([a.do_reducescatter.remote() for a in actors])
+    for i in range(world_size):
+        if tensor_backend == "cupy":
+            assert (
+                results[i] == cp.ones(array_size, dtype=cp.float32) * world_size
+            ).all()
+        else:
+            assert (
+                results[i]
+                == torch.ones(array_size, dtype=torch.float32).cuda() * world_size
+            ).all()
+
+
+@pytest.mark.parametrize("dtype", [cp.uint8, cp.float16, cp.float32, cp.float64])
+def test_reducescatter_different_dtype(ray_start_distributed_2_nodes_4_gpus, dtype):
+    world_size = 4
+    actors, _ = create_collective_workers(world_size)
+    init_tensors_for_gather_scatter(actors, dtype=dtype)
+    results = ray.get([a.do_reducescatter.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(world_size):
+            assert (results[i] == cp.ones(10, dtype=dtype) * world_size).all()
+
+
+def test_reducescatter_torch_cupy(ray_start_distributed_2_nodes_4_gpus):
+    world_size = 4
+    shape = [10, 10]
+    actors, _ = create_collective_workers(world_size)
+
+    # tensor is pytorch, list is cupy
+    for i, a in enumerate(actors):
+        t = torch.ones(shape, dtype=torch.float32).cuda() * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = [cp.ones(shape, dtype=cp.float32) for _ in range(world_size)]
+        ray.wait([a.set_list_buffer.remote(list_buffer)])
+    results = ray.get([a.do_reducescatter.remote() for a in actors])
+    for i in range(world_size):
+        assert (
+            results[i] == torch.ones(shape, dtype=torch.float32).cuda() * world_size
+        ).all()
+
+    # tensor is cupy, list is pytorch
+    for i, a in enumerate(actors):
+        t = cp.ones(shape, dtype=cp.float32) * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = [
+            torch.ones(shape, dtype=torch.float32).cuda() for _ in range(world_size)
+        ]
+        ray.wait([a.set_list_buffer.remote(list_buffer)])
+    results = ray.get([a.do_reducescatter.remote() for a in actors])
+    for i in range(world_size):
+        assert (results[i] == cp.ones(shape, dtype=cp.float32) * world_size).all()
+
+    # some tensors in the list are pytorch, some are cupy
+    for i, a in enumerate(actors):
+        if i % 2 == 0:
+            t = torch.ones(shape, dtype=torch.float32).cuda() * (i + 1)
+        else:
+            t = cp.ones(shape, dtype=cp.float32) * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = []
+        for j in range(world_size):
+            if j % 2 == 0:
+                list_buffer.append(torch.ones(shape, dtype=torch.float32).cuda())
+            else:
+                list_buffer.append(cp.ones(shape, dtype=cp.float32))
+        ray.wait([a.set_list_buffer.remote(list_buffer)])
+    results = ray.get([a.do_reducescatter.remote() for a in actors])
+    for i in range(world_size):
+        if i % 2 == 0:
+            assert (
+                results[i] == torch.ones(shape, dtype=torch.float32).cuda() * world_size
+            ).all()
+        else:
+            assert (results[i] == cp.ones(shape, dtype=cp.float32) * world_size).all()
+
+    # mixed case
+    for i, a in enumerate(actors):
+        if i % 2 == 0:
+            t = torch.ones(shape, dtype=torch.float32).cuda() * (i + 1)
+        else:
+            t = cp.ones(shape, dtype=cp.float32) * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = []
+        for j in range(world_size):
+            if j % 2 == 0:
+                list_buffer.append(cp.ones(shape, dtype=cp.float32))
+            else:
+                list_buffer.append(torch.ones(shape, dtype=torch.float32).cuda())
+        ray.wait([a.set_list_buffer.remote(list_buffer)])
+    results = ray.get([a.do_reducescatter.remote() for a in actors])
+    for i in range(world_size):
+        if i % 2 == 0:
+            assert (
+                results[i] == torch.ones(shape, dtype=torch.float32).cuda() * world_size
+            ).all()
+        else:
+            assert (results[i] == cp.ones(shape, dtype=cp.float32) * world_size).all()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/util/collective/tests/distributed_gpu_tests/test_distributed_sendrecv.py

@@ -0,0 +1,40 @@
+"""Test the send/recv API."""
+import cupy as cp
+import pytest
+
+import ray
+from ray.util.collective.tests.util import create_collective_workers
+
+
+@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
+@pytest.mark.parametrize("dst_rank", [0, 1, 2, 3])
+@pytest.mark.parametrize("src_rank", [0, 1, 2, 3])
+@pytest.mark.parametrize(
+    "array_size", [2**10, 2**15, 2**20, [2, 2], [5, 9, 10, 85]]
+)
+def test_sendrecv(
+    ray_start_distributed_2_nodes_4_gpus, group_name, array_size, src_rank, dst_rank
+):
+    if src_rank == dst_rank:
+        return
+    world_size = 4
+    actors, _ = create_collective_workers(num_workers=world_size, group_name=group_name)
+    ray.get(
+        [
+            a.set_buffer.remote(cp.ones(array_size, dtype=cp.float32) * (i + 1))
+            for i, a in enumerate(actors)
+        ]
+    )
+    refs = []
+    for i in range(world_size):
+        refs.append(actors[i].get_buffer.remote())
+    refs[src_rank] = actors[src_rank].do_send.remote(group_name, dst_rank)
+    refs[dst_rank] = actors[dst_rank].do_recv.remote(group_name, src_rank)
+    results = ray.get(refs)
+    assert (
+        results[src_rank] == cp.ones(array_size, dtype=cp.float32) * (src_rank + 1)
+    ).all()
+    assert (
+        results[dst_rank] == cp.ones(array_size, dtype=cp.float32) * (src_rank + 1)
+    ).all()
+    ray.get([a.destroy_group.remote(group_name) for a in actors])

python/ray/util/collective/tests/distributed_multigpu_tests/test_distributed_multigpu_allgather.py

@@ -0,0 +1,87 @@
+"""Test the allgather API on a distributed Ray cluster."""
+import cupy as cp
+import pytest
+import torch
+
+import ray
+from ray.util.collective.tests.util import (
+    create_collective_multigpu_workers,
+    init_tensors_for_gather_scatter_multigpu,
+)
+
+
+@pytest.mark.parametrize("tensor_backend", ["cupy", "torch"])
+@pytest.mark.parametrize(
+    "array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
+)
+def test_allgather_different_array_size(
+    ray_start_distributed_multigpu_2_nodes_4_gpus, array_size, tensor_backend
+):
+    world_size = 2
+    num_gpu_per_worker = 2
+    actual_world_size = world_size * num_gpu_per_worker
+    actors, _ = create_collective_multigpu_workers(world_size)
+    init_tensors_for_gather_scatter_multigpu(
+        actors, array_size=array_size, tensor_backend=tensor_backend
+    )
+    results = ray.get([a.do_allgather_multigpu.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(num_gpu_per_worker):
+            for k in range(actual_world_size):
+                if tensor_backend == "cupy":
+                    assert (
+                        results[i][j][k] == cp.ones(array_size, dtype=cp.float32)
+                    ).all()
+                else:
+                    assert (
+                        results[i][j][k]
+                        == torch.ones(array_size, dtype=torch.float32).cuda(j)
+                    ).all()
+
+
+def test_allgather_torch_cupy(ray_start_distributed_multigpu_2_nodes_4_gpus):
+    world_size = 2
+    num_gpu_per_worker = 2
+    actual_world_size = world_size * num_gpu_per_worker
+    shape = [10, 10]
+    actors, _ = create_collective_multigpu_workers(world_size)
+
+    # tensor is pytorch, list is cupy
+    for i, a in enumerate(actors):
+        ray.get(
+            [a.set_buffer.remote(shape, tensor_type0="torch", tensor_type1="torch")]
+        )
+        ray.get(
+            [a.set_list_buffer.remote(shape, tensor_type0="cupy", tensor_type1="cupy")]
+        )
+    results = ray.get([a.do_allgather_multigpu.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(num_gpu_per_worker):
+            for k in range(actual_world_size):
+                assert (results[i][j][k] == cp.ones(shape, dtype=cp.float32)).all()
+
+    # tensor is cupy, list is pytorch
+    for i, a in enumerate(actors):
+        ray.get([a.set_buffer.remote(shape, tensor_type0="cupy", tensor_type1="cupy")])
+        ray.get(
+            [
+                a.set_list_buffer.remote(
+                    shape, tensor_type0="torch", tensor_type1="torch"
+                )
+            ]
+        )
+    results = ray.get([a.do_allgather_multigpu.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(num_gpu_per_worker):
+            for k in range(actual_world_size):
+                assert (
+                    results[i][j][k] == torch.ones(shape, dtype=torch.float32).cuda(j)
+                ).all()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/util/collective/tests/distributed_multigpu_tests/test_distributed_multigpu_allreduce.py

@@ -0,0 +1,166 @@
+"""Test the collective allreduice API on a distributed Ray cluster."""
+import logging
+
+import cupy as cp
+import pytest
+
+import ray
+from ray.util.collective.tests.util import create_collective_multigpu_workers
+from ray.util.collective.types import ReduceOp
+
+logger = logging.getLogger(__name__)
+logger.setLevel("DEBUG")
+
+
+@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
+def test_allreduce_multigpu_different_name(
+    ray_start_distributed_multigpu_2_nodes_4_gpus, group_name
+):
+    world_size = 2
+    num_gpu_per_worker = 2
+    actual_world_size = world_size * num_gpu_per_worker
+    actors, _ = create_collective_multigpu_workers(
+        num_workers=world_size, group_name=group_name
+    )
+    results = ray.get([a.do_allreduce_multigpu.remote(group_name) for a in actors])
+    assert (results[0] == cp.ones((10,), dtype=cp.float32) * actual_world_size).all()
+    assert (results[1] == cp.ones((10,), dtype=cp.float32) * actual_world_size).all()
+
+
+@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
+def test_allreduce_multigpu_different_array_size(
+    ray_start_distributed_multigpu_2_nodes_4_gpus, array_size
+):
+    world_size = 2
+    num_gpu_per_worker = 2
+    actual_world_size = world_size * num_gpu_per_worker
+    actors, _ = create_collective_multigpu_workers(world_size)
+    ray.get([a.set_buffer.remote(array_size) for a in actors])
+    results = ray.get([a.do_allreduce_multigpu.remote() for a in actors])
+    assert (
+        results[0] == cp.ones((array_size,), dtype=cp.float32) * actual_world_size
+    ).all()
+    assert (
+        results[1] == cp.ones((array_size,), dtype=cp.float32) * actual_world_size
+    ).all()
+
+
+def test_allreduce_multigpu_destroy(
+    ray_start_distributed_multigpu_2_nodes_4_gpus, backend="nccl", group_name="default"
+):
+    world_size = 2
+    num_gpu_per_worker = 2
+    actual_world_size = world_size * num_gpu_per_worker
+    actors, _ = create_collective_multigpu_workers(world_size)
+
+    results = ray.get([a.do_allreduce_multigpu.remote() for a in actors])
+    assert (results[0] == cp.ones((10,), dtype=cp.float32) * actual_world_size).all()
+    assert (results[1] == cp.ones((10,), dtype=cp.float32) * actual_world_size).all()
+
+    # destroy the group and try do work, should fail
+    ray.get([a.destroy_group.remote() for a in actors])
+    with pytest.raises(RuntimeError):
+        results = ray.get([a.do_allreduce_multigpu.remote() for a in actors])
+
+    # reinit the same group and all reduce
+    ray.get(
+        [
+            actor.init_group.remote(world_size, i, backend, group_name)
+            for i, actor in enumerate(actors)
+        ]
+    )
+    results = ray.get([a.do_allreduce_multigpu.remote() for a in actors])
+    assert (
+        results[0]
+        == cp.ones((10,), dtype=cp.float32) * actual_world_size * actual_world_size
+    ).all()
+    assert (
+        results[1]
+        == cp.ones((10,), dtype=cp.float32) * actual_world_size * actual_world_size
+    ).all()
+
+
+def test_allreduce_multigpu_multiple_group(
+    ray_start_distributed_multigpu_2_nodes_4_gpus, backend="nccl", num_groups=5
+):
+    world_size = 2
+    num_gpu_per_worker = 2
+    actual_world_size = world_size * num_gpu_per_worker
+    actors, _ = create_collective_multigpu_workers(world_size)
+    for group_name in range(1, num_groups):
+        ray.get(
+            [
+                actor.init_group.remote(world_size, i, backend, str(group_name))
+                for i, actor in enumerate(actors)
+            ]
+        )
+    for i in range(num_groups):
+        group_name = "default" if i == 0 else str(i)
+        results = ray.get([a.do_allreduce_multigpu.remote(group_name) for a in actors])
+        assert (
+            results[0]
+            == cp.ones((10,), dtype=cp.float32) * (actual_world_size ** (i + 1))
+        ).all()
+
+
+def test_allreduce_multigpu_different_op(ray_start_distributed_multigpu_2_nodes_4_gpus):
+    world_size = 2
+    actors, _ = create_collective_multigpu_workers(world_size)
+    # check product
+    ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
+    ray.get(actors[1].set_buffer.remote([10], value0=4, value1=5))
+    results = ray.get(
+        [a.do_allreduce_multigpu.remote(op=ReduceOp.PRODUCT) for a in actors]
+    )
+    assert (results[0] == cp.ones((10,), dtype=cp.float32) * 120).all()
+    assert (results[1] == cp.ones((10,), dtype=cp.float32) * 120).all()
+
+    # check min
+    ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
+    ray.get(actors[1].set_buffer.remote([10], value0=4, value1=5))
+    results = ray.get([a.do_allreduce_multigpu.remote(op=ReduceOp.MIN) for a in actors])
+    assert (results[0] == cp.ones((10,), dtype=cp.float32) * 2).all()
+    assert (results[1] == cp.ones((10,), dtype=cp.float32) * 2).all()
+
+    # check max
+    ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
+    ray.get(actors[1].set_buffer.remote([10], value0=4, value1=5))
+    results = ray.get([a.do_allreduce_multigpu.remote(op=ReduceOp.MAX) for a in actors])
+    assert (results[0] == cp.ones((10,), dtype=cp.float32) * 5).all()
+    assert (results[1] == cp.ones((10,), dtype=cp.float32) * 5).all()
+
+
+@pytest.mark.parametrize("dtype", [cp.uint8, cp.float16, cp.float32, cp.float64])
+def test_allreduce_multigpu_different_dtype(
+    ray_start_distributed_multigpu_2_nodes_4_gpus, dtype
+):
+    world_size = 2
+    num_gpu_per_worker = 2
+    actual_world_size = world_size * num_gpu_per_worker
+    actors, _ = create_collective_multigpu_workers(world_size)
+    ray.get([a.set_buffer.remote([10], dtype=dtype) for a in actors])
+    results = ray.get([a.do_allreduce_multigpu.remote() for a in actors])
+    assert (results[0] == cp.ones((10,), dtype=dtype) * actual_world_size).all()
+    assert (results[1] == cp.ones((10,), dtype=dtype) * actual_world_size).all()
+
+
+def test_allreduce_torch_cupy(ray_start_distributed_multigpu_2_nodes_4_gpus):
+    # import torch
+    world_size = 2
+    actual_world_size = 4
+    actors, _ = create_collective_multigpu_workers(world_size)
+    ray.get(actors[0].set_buffer.remote([10]))
+    ray.get(
+        actors[1].set_buffer.remote([10], tensor_type0="torch", tensor_type1="torch")
+    )
+    results = ray.get([a.do_allreduce_multigpu.remote() for a in actors])
+    assert (results[0] == cp.ones((10,)) * actual_world_size).all()
+
+    ray.get(
+        actors[0].set_buffer.remote([10], tensor_type0="cupy", tensor_type1="torch")
+    )
+    ray.get(
+        actors[1].set_buffer.remote([10], tensor_type0="torch", tensor_type1="cupy")
+    )
+    results = ray.get([a.do_allreduce_multigpu.remote() for a in actors])
+    assert (results[0] == cp.ones((10,)) * actual_world_size).all()

python/ray/util/collective/tests/distributed_multigpu_tests/test_distributed_multigpu_basic_apis.py

@@ -0,0 +1,103 @@
+"""Test the collective group APIs."""
+from random import shuffle
+
+import pytest
+
+import ray
+from ray.util.collective.tests.util import create_collective_multigpu_workers
+
+
+@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
+def test_init_two_actors(ray_start_distributed_multigpu_2_nodes_4_gpus, group_name):
+    world_size = 2
+    actors, results = create_collective_multigpu_workers(world_size, group_name)
+    for i in range(world_size):
+        assert results[i]
+
+
+def test_report_num_gpus(ray_start_distributed_multigpu_2_nodes_4_gpus):
+    world_size = 2
+    actors, results = create_collective_multigpu_workers(world_size)
+    num_gpus = ray.get([actor.report_num_gpus.remote() for actor in actors])
+    assert num_gpus == [2, 2]
+
+
+def test_get_rank(ray_start_distributed_multigpu_2_nodes_4_gpus):
+    world_size = 2
+    actors, _ = create_collective_multigpu_workers(world_size)
+    actor0_rank = ray.get(actors[0].report_rank.remote())
+    assert actor0_rank == 0
+    actor1_rank = ray.get(actors[1].report_rank.remote())
+    assert actor1_rank == 1
+
+    # create a second group with a different name, and different
+    # orders of ranks.
+    new_group_name = "default2"
+    ranks = list(range(world_size))
+    shuffle(ranks)
+    ray.get(
+        [
+            actor.init_group.remote(world_size, ranks[i], group_name=new_group_name)
+            for i, actor in enumerate(actors)
+        ]
+    )
+    actor0_rank = ray.get(actors[0].report_rank.remote(new_group_name))
+    assert actor0_rank == ranks[0]
+    actor1_rank = ray.get(actors[1].report_rank.remote(new_group_name))
+    assert actor1_rank == ranks[1]
+
+
+def test_is_group_initialized(ray_start_distributed_multigpu_2_nodes_4_gpus):
+    world_size = 2
+    actors, _ = create_collective_multigpu_workers(world_size)
+    # check group is_init
+    actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
+    assert actor0_is_init
+    actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("random"))
+    assert not actor0_is_init
+    actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("123"))
+    assert not actor0_is_init
+    actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
+    assert actor1_is_init
+    actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote("456"))
+    assert not actor1_is_init
+
+
+def test_destroy_group(ray_start_distributed_multigpu_2_nodes_4_gpus):
+    world_size = 2
+    actors, _ = create_collective_multigpu_workers(world_size)
+    # Now destroy the group at actor0
+    ray.wait([actors[0].destroy_group.remote()])
+    actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
+    assert not actor0_is_init
+
+    # should go well as the group `random` does not exist at all
+    ray.wait([actors[0].destroy_group.remote("random")])
+
+    actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
+    assert actor1_is_init
+    ray.wait([actors[1].destroy_group.remote("random")])
+    actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
+    assert actor1_is_init
+    ray.wait([actors[1].destroy_group.remote("default")])
+    actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
+    assert not actor1_is_init
+
+    # Now reconstruct the group using the same name
+    init_results = ray.get(
+        [actor.init_group.remote(world_size, i) for i, actor in enumerate(actors)]
+    )
+    for i in range(world_size):
+        assert init_results[i]
+    actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
+    assert actor0_is_init
+    actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
+    assert actor1_is_init
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/util/collective/tests/distributed_multigpu_tests/test_distributed_multigpu_broadcast.py

@@ -0,0 +1,112 @@
+"""Test the broadcast API."""
+import cupy as cp
+import pytest
+
+import ray
+from ray.util.collective.tests.util import create_collective_multigpu_workers
+
+
+@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
+@pytest.mark.parametrize("src_rank", [0, 1])
+@pytest.mark.parametrize("src_gpu_index", [0, 1])
+def test_broadcast_different_name(
+    ray_start_distributed_multigpu_2_nodes_4_gpus, group_name, src_rank, src_gpu_index
+):
+    world_size = 2
+    num_gpu_per_worker = 2
+    actors, _ = create_collective_multigpu_workers(
+        num_workers=world_size, group_name=group_name
+    )
+    ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
+    ray.get(actors[1].set_buffer.remote([10], value0=4, value1=5))
+
+    results = ray.get(
+        [
+            a.do_broadcast_multigpu.remote(
+                group_name=group_name, src_rank=src_rank, src_gpu_index=src_gpu_index
+            )
+            for a in actors
+        ]
+    )
+    for i in range(world_size):
+        for j in range(num_gpu_per_worker):
+            val = (src_rank + 1) * 2 + src_gpu_index
+            assert (results[i][j] == cp.ones([10], dtype=cp.float32) * val).all()
+
+
+@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
+@pytest.mark.parametrize("src_rank", [0, 1])
+@pytest.mark.parametrize("src_gpu_index", [0, 1])
+def test_broadcast_different_array_size(
+    ray_start_distributed_multigpu_2_nodes_4_gpus, array_size, src_rank, src_gpu_index
+):
+    world_size = 2
+    num_gpu_per_worker = 2
+    actors, _ = create_collective_multigpu_workers(world_size)
+    ray.get(actors[0].set_buffer.remote([array_size], value0=2, value1=3))
+    ray.get(actors[1].set_buffer.remote([array_size], value0=4, value1=5))
+    results = ray.get(
+        [
+            a.do_broadcast_multigpu.remote(
+                src_rank=src_rank, src_gpu_index=src_gpu_index
+            )
+            for a in actors
+        ]
+    )
+    for i in range(world_size):
+        for j in range(num_gpu_per_worker):
+            val = (src_rank + 1) * 2 + src_gpu_index
+            assert (
+                results[i][j] == cp.ones((array_size,), dtype=cp.float32) * val
+            ).all()
+
+
+@pytest.mark.parametrize("src_rank", [0, 1])
+@pytest.mark.parametrize("src_gpu_index", [0, 1])
+def test_broadcast_torch_cupy(
+    ray_start_distributed_multigpu_2_nodes_4_gpus, src_rank, src_gpu_index
+):
+    import torch
+
+    world_size = 2
+    num_gpu_per_worker = 2
+    actors, _ = create_collective_multigpu_workers(world_size)
+    ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
+    ray.get(
+        actors[1].set_buffer.remote(
+            [10], value0=4, value1=5, tensor_type0="torch", tensor_type1="torch"
+        )
+    )
+    results = ray.get(
+        [
+            a.do_broadcast_multigpu.remote(
+                src_rank=src_rank, src_gpu_index=src_gpu_index
+            )
+            for a in actors
+        ]
+    )
+    for i in range(world_size):
+        for j in range(num_gpu_per_worker):
+            val = (src_rank + 1) * 2 + src_gpu_index
+            if i == 0:
+                assert (results[i][j] == cp.ones([10], dtype=cp.float32) * val).all()
+            else:
+                assert (results[i][j] == torch.ones([10]).cuda(j) * val).all()
+
+
+@pytest.mark.parametrize("src_rank", [3, 4])
+@pytest.mark.parametrize("src_gpu_index", [2, 3])
+def test_broadcast_invalid_rank(
+    ray_start_distributed_multigpu_2_nodes_4_gpus, src_rank, src_gpu_index
+):
+    world_size = 2
+    actors, _ = create_collective_multigpu_workers(world_size)
+    with pytest.raises(ValueError):
+        _ = ray.get(
+            [
+                a.do_broadcast_multigpu.remote(
+                    src_rank=src_rank, src_gpu_index=src_gpu_index
+                )
+                for a in actors
+            ]
+        )

python/ray/util/collective/tests/distributed_multigpu_tests/test_distributed_multigpu_reduce.py

@@ -0,0 +1,194 @@
+"""Test the reduce API."""
+import cupy as cp
+import pytest
+
+import ray
+from ray.util.collective.tests.util import create_collective_multigpu_workers
+from ray.util.collective.types import ReduceOp
+
+
+@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
+@pytest.mark.parametrize("dst_rank", [0, 1])
+@pytest.mark.parametrize("dst_gpu_index", [0, 1])
+def test_reduce_different_name(
+    ray_start_distributed_multigpu_2_nodes_4_gpus, group_name, dst_rank, dst_gpu_index
+):
+    world_size = 2
+    num_gpu_per_worker = 2
+    actual_world_size = world_size * num_gpu_per_worker
+    actors, _ = create_collective_multigpu_workers(
+        num_workers=world_size, group_name=group_name
+    )
+    results = ray.get(
+        [
+            a.do_reduce_multigpu.remote(
+                group_name, dst_rank=dst_rank, dst_gpu_index=dst_gpu_index
+            )
+            for a in actors
+        ]
+    )
+    for i in range(world_size):
+        for j in range(num_gpu_per_worker):
+            if i == dst_rank and j == dst_gpu_index:
+                assert (
+                    results[i][j]
+                    == cp.ones((10,), dtype=cp.float32) * actual_world_size
+                ).all()
+            else:
+                assert (results[i][j] == cp.ones((10,), dtype=cp.float32)).all()
+
+
+@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
+@pytest.mark.parametrize("dst_rank", [0, 1])
+@pytest.mark.parametrize("dst_gpu_index", [0, 1])
+def test_reduce_different_array_size(
+    ray_start_distributed_multigpu_2_nodes_4_gpus, array_size, dst_rank, dst_gpu_index
+):
+    world_size = 2
+    num_gpu_per_worker = 2
+    actual_world_size = world_size * num_gpu_per_worker
+    actors, _ = create_collective_multigpu_workers(num_workers=world_size)
+
+    ray.get(actors[0].set_buffer.remote(array_size))
+    ray.get(actors[1].set_buffer.remote(array_size))
+    results = ray.get(
+        [
+            a.do_reduce_multigpu.remote(dst_rank=dst_rank, dst_gpu_index=dst_gpu_index)
+            for a in actors
+        ]
+    )
+    for i in range(world_size):
+        for j in range(num_gpu_per_worker):
+            if i == dst_rank and j == dst_gpu_index:
+                assert (
+                    results[i][j]
+                    == cp.ones((array_size,), dtype=cp.float32) * actual_world_size
+                ).all()
+            else:
+                assert (results[i][j] == cp.ones((array_size,), dtype=cp.float32)).all()
+
+
+@pytest.mark.parametrize("dst_rank", [0, 1])
+@pytest.mark.parametrize("dst_gpu_index", [0, 1])
+def test_reduce_different_op(
+    ray_start_distributed_multigpu_2_nodes_4_gpus, dst_rank, dst_gpu_index
+):
+    world_size = 2
+    num_gpu_per_worker = 2
+    actors, _ = create_collective_multigpu_workers(world_size)
+
+    # check product
+    ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
+    ray.get(actors[1].set_buffer.remote([10], value0=4, value1=5))
+    results = ray.get(
+        [
+            a.do_reduce_multigpu.remote(
+                dst_rank=dst_rank, dst_gpu_index=dst_gpu_index, op=ReduceOp.PRODUCT
+            )
+            for a in actors
+        ]
+    )
+    for i in range(world_size):
+        for j in range(num_gpu_per_worker):
+            if i == dst_rank and j == dst_gpu_index:
+                assert (results[i][j] == cp.ones((10,), dtype=cp.float32) * 120).all()
+            else:
+                val = (i + 1) * 2 + j
+                assert (results[i][j] == cp.ones((10,), dtype=cp.float32) * val).all()
+
+    # check min
+    ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
+    ray.get(actors[1].set_buffer.remote([10], value0=4, value1=5))
+    results = ray.get(
+        [
+            a.do_reduce_multigpu.remote(
+                dst_rank=dst_rank, dst_gpu_index=dst_gpu_index, op=ReduceOp.MIN
+            )
+            for a in actors
+        ]
+    )
+    for i in range(world_size):
+        for j in range(num_gpu_per_worker):
+            if i == dst_rank and j == dst_gpu_index:
+                assert (results[i][j] == cp.ones((10,), dtype=cp.float32) * 2).all()
+            else:
+                val = (i + 1) * 2 + j
+                assert (results[i][j] == cp.ones((10,), dtype=cp.float32) * val).all()
+
+    # check max
+    ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
+    ray.get(actors[1].set_buffer.remote([10], value0=4, value1=5))
+    results = ray.get(
+        [
+            a.do_reduce_multigpu.remote(
+                dst_rank=dst_rank, dst_gpu_index=dst_gpu_index, op=ReduceOp.MAX
+            )
+            for a in actors
+        ]
+    )
+    for i in range(world_size):
+        for j in range(num_gpu_per_worker):
+            if i == dst_rank and j == dst_gpu_index:
+                assert (results[i][j] == cp.ones((10,), dtype=cp.float32) * 5).all()
+            else:
+                val = (i + 1) * 2 + j
+                assert (results[i][j] == cp.ones((10,), dtype=cp.float32) * val).all()
+
+
+@pytest.mark.parametrize("dst_rank", [0, 1])
+@pytest.mark.parametrize("dst_gpu_index", [0, 1])
+def test_reduce_torch_cupy(
+    ray_start_distributed_multigpu_2_nodes_4_gpus, dst_rank, dst_gpu_index
+):
+    import torch
+
+    world_size = 2
+    num_gpu_per_worker = 2
+    actors, _ = create_collective_multigpu_workers(world_size)
+    ray.get(actors[0].set_buffer.remote([10], value0=2, value1=3))
+    ray.get(
+        actors[1].set_buffer.remote(
+            [10], value0=4, value1=5, tensor_type0="torch", tensor_type1="torch"
+        )
+    )
+
+    results = ray.get(
+        [
+            a.do_reduce_multigpu.remote(dst_rank=dst_rank, dst_gpu_index=dst_gpu_index)
+            for a in actors
+        ]
+    )
+
+    for i in range(world_size):
+        for j in range(num_gpu_per_worker):
+            val = (i + 1) * 2 + j
+            if dst_rank == i and dst_gpu_index == j:
+                if i == 0:
+                    assert (results[i][j] == cp.ones([10], dtype=cp.float32) * 14).all()
+                else:
+                    assert (results[i][j] == torch.ones([10]).cuda(j) * 14).all()
+            else:
+                if i == 0:
+                    assert (
+                        results[i][j] == cp.ones([10], dtype=cp.float32) * val
+                    ).all()
+                else:
+                    assert (results[i][j] == torch.ones([10]).cuda(j) * val).all()
+
+
+@pytest.mark.parametrize("dst_rank", [3, 4])
+@pytest.mark.parametrize("dst_gpu_index", [2, 3])
+def test_reduce_invalid_rank(
+    ray_start_distributed_multigpu_2_nodes_4_gpus, dst_rank, dst_gpu_index
+):
+    world_size = 2
+    actors, _ = create_collective_multigpu_workers(world_size)
+    with pytest.raises(ValueError):
+        _ = ray.get(
+            [
+                a.do_reduce_multigpu.remote(
+                    dst_rank=dst_rank, dst_gpu_index=dst_gpu_index
+                )
+                for a in actors
+            ]
+        )

python/ray/util/collective/tests/distributed_multigpu_tests/test_distributed_multigpu_reducescatter.py

@@ -0,0 +1,90 @@
+"""Test the collective reducescatter API on a distributed Ray cluster."""
+import cupy as cp
+import pytest
+import torch
+
+import ray
+from ray.util.collective.tests.util import (
+    create_collective_multigpu_workers,
+    init_tensors_for_gather_scatter_multigpu,
+)
+
+
+@pytest.mark.parametrize("tensor_backend", ["cupy", "torch"])
+@pytest.mark.parametrize(
+    "array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
+)
+def test_reducescatter_different_array_size(
+    ray_start_distributed_multigpu_2_nodes_4_gpus, array_size, tensor_backend
+):
+    world_size = 2
+    num_gpu_per_worker = 2
+    actual_world_size = world_size * num_gpu_per_worker
+    actors, _ = create_collective_multigpu_workers(world_size)
+
+    init_tensors_for_gather_scatter_multigpu(
+        actors, array_size=array_size, tensor_backend=tensor_backend
+    )
+    results = ray.get([a.do_reducescatter_multigpu.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(num_gpu_per_worker):
+            if tensor_backend == "cupy":
+                assert (
+                    results[i][j]
+                    == cp.ones(array_size, dtype=cp.float32) * actual_world_size
+                ).all()
+            else:
+                assert (
+                    results[i][j]
+                    == torch.ones(array_size, dtype=torch.float32).cuda(j)
+                    * actual_world_size
+                ).all()
+
+
+def test_reducescatter_torch_cupy(ray_start_distributed_multigpu_2_nodes_4_gpus):
+    world_size = 2
+    num_gpu_per_worker = 2
+    actual_world_size = world_size * num_gpu_per_worker
+    shape = [10, 10]
+    actors, _ = create_collective_multigpu_workers(world_size)
+
+    # tensor is pytorch, list is cupy
+    for i, a in enumerate(actors):
+        ray.get(
+            [a.set_buffer.remote(shape, tensor_type0="torch", tensor_type1="torch")]
+        )
+        ray.get(
+            [a.set_list_buffer.remote(shape, tensor_type0="cupy", tensor_type1="cupy")]
+        )
+    results = ray.get([a.do_reducescatter_multigpu.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(num_gpu_per_worker):
+            assert (
+                results[i][j]
+                == torch.ones(shape, dtype=torch.float32).cuda(j) * actual_world_size
+            ).all()
+
+    # tensor is cupy, list is pytorch
+    for i, a in enumerate(actors):
+        ray.get([a.set_buffer.remote(shape, tensor_type0="cupy", tensor_type1="cupy")])
+        ray.get(
+            [
+                a.set_list_buffer.remote(
+                    shape, tensor_type0="torch", tensor_type1="torch"
+                )
+            ]
+        )
+    results = ray.get([a.do_reducescatter_multigpu.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(num_gpu_per_worker):
+            assert (
+                results[i][j] == cp.ones(shape, dtype=cp.float32) * actual_world_size
+            ).all()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/util/collective/tests/distributed_multigpu_tests/test_distributed_multigpu_sendrecv.py

@@ -0,0 +1,54 @@
+"""Test the send/recv API."""
+import cupy as cp
+import pytest
+
+import ray
+from ray.util.collective.tests.util import create_collective_multigpu_workers
+
+
+# @pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
+@pytest.mark.parametrize("dst_rank", [0, 1])
+@pytest.mark.parametrize("src_rank", [0, 1])
+@pytest.mark.parametrize("dst_gpu_index", [0, 1])
+@pytest.mark.parametrize("src_gpu_index", [0, 1])
+@pytest.mark.parametrize(
+    "array_size", [2**10, 2**15, 2**20, [2, 2], [5, 9, 10, 85]]
+)
+def test_sendrecv(
+    ray_start_distributed_multigpu_2_nodes_4_gpus,
+    array_size,
+    src_rank,
+    dst_rank,
+    src_gpu_index,
+    dst_gpu_index,
+):
+    if src_rank == dst_rank:
+        return
+    world_size = 2
+    actors, _ = create_collective_multigpu_workers(num_workers=world_size)
+
+    ray.get(actors[0].set_buffer.remote(array_size, value0=2, value1=3))
+    ray.get(actors[1].set_buffer.remote(array_size, value0=4, value1=5))
+
+    refs = []
+    for i in range(world_size):
+        refs.append(actors[i].get_buffer.remote())
+    refs[src_rank][src_gpu_index] = actors[src_rank].do_send_multigpu.remote(
+        dst_rank=dst_rank, dst_gpu_index=dst_gpu_index, src_gpu_index=src_gpu_index
+    )
+    refs[dst_rank][dst_gpu_index] = actors[dst_rank].do_recv_multigpu.remote(
+        src_rank=src_rank, src_gpu_index=src_gpu_index, dst_gpu_index=dst_gpu_index
+    )
+    results = []
+    results_flattend = ray.get(refs[0] + refs[1])
+    results.append([results_flattend[0], results_flattend[1]])
+    results.append([results_flattend[2], results_flattend[3]])
+    assert (
+        results[src_rank][src_gpu_index]
+        == cp.ones(array_size, dtype=cp.float32) * ((src_rank + 1) * 2 + src_gpu_index)
+    ).all()
+    assert (
+        results[dst_rank][dst_gpu_index]
+        == cp.ones(array_size, dtype=cp.float32) * ((src_rank + 1) * 2 + src_gpu_index)
+    ).all()
+    ray.get([a.destroy_group.remote() for a in actors])

python/ray/util/collective/tests/single_node_cpu_tests/test_allgather.py

@@ -0,0 +1,143 @@
+"""Test the collective allgather API."""
+import numpy as np
+import pytest
+import torch
+
+import ray
+from ray.util.collective.tests.cpu_util import (
+    create_collective_workers,
+    init_tensors_for_gather_scatter,
+)
+from ray.util.collective.types import Backend
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("tensor_backend", ["numpy", "torch"])
+@pytest.mark.parametrize(
+    "array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
+)
+def test_allgather_different_array_size(
+    ray_start_single_node, array_size, tensor_backend, backend
+):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    init_tensors_for_gather_scatter(
+        actors, array_size=array_size, tensor_backend=tensor_backend
+    )
+    results = ray.get([a.do_allgather.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(world_size):
+            if tensor_backend == "numpy":
+                assert (
+                    results[i][j] == np.ones(array_size, dtype=np.float32) * (j + 1)
+                ).all()
+            else:
+                assert (
+                    results[i][j]
+                    == torch.ones(array_size, dtype=torch.float32) * (j + 1)
+                ).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("dtype", [np.uint8, np.float16, np.float32, np.float64])
+def test_allgather_different_dtype(ray_start_single_node, dtype, backend):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    init_tensors_for_gather_scatter(actors, dtype=dtype)
+    results = ray.get([a.do_allgather.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(world_size):
+            assert (results[i][j] == np.ones(10, dtype=dtype) * (j + 1)).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("length", [0, 1, 2, 3])
+def test_unmatched_tensor_list_length(ray_start_single_node, length, backend):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    list_buffer = [np.ones(10, dtype=np.float32) for _ in range(length)]
+    ray.wait([a.set_list_buffer.remote(list_buffer, copy=True) for a in actors])
+    if length != world_size:
+        with pytest.raises(RuntimeError):
+            ray.get([a.do_allgather.remote() for a in actors])
+    else:
+        ray.get([a.do_allgather.remote() for a in actors])
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("shape", [10, 20, [4, 5], [1, 3, 5, 7]])
+def test_unmatched_tensor_shape(ray_start_single_node, shape, backend):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    init_tensors_for_gather_scatter(actors, array_size=10)
+    list_buffer = [np.ones(shape, dtype=np.float32) for _ in range(world_size)]
+    ray.get([a.set_list_buffer.remote(list_buffer, copy=True) for a in actors])
+    if shape != 10:
+        with pytest.raises(RuntimeError):
+            ray.get([a.do_allgather.remote() for a in actors])
+    else:
+        ray.get([a.do_allgather.remote() for a in actors])
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+def test_allgather_torch_numpy(ray_start_single_node, backend):
+    world_size = 2
+    shape = [10, 10]
+    actors, _ = create_collective_workers(world_size, backend=backend)
+
+    # tensor is pytorch, list is numpy
+    for i, a in enumerate(actors):
+        t = torch.ones(shape, dtype=torch.float32) * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = [np.ones(shape, dtype=np.float32) for _ in range(world_size)]
+        ray.wait([a.set_list_buffer.remote(list_buffer, copy=True)])
+    results = ray.get([a.do_allgather.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(world_size):
+            assert (results[i][j] == np.ones(shape, dtype=np.float32) * (j + 1)).all()
+
+    # tensor is numpy, list is pytorch
+    for i, a in enumerate(actors):
+        t = np.ones(shape, dtype=np.float32) * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = [
+            torch.ones(shape, dtype=torch.float32) for _ in range(world_size)
+        ]
+        ray.wait([a.set_list_buffer.remote(list_buffer, copy=True)])
+    results = ray.get([a.do_allgather.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(world_size):
+            assert (
+                results[i][j] == torch.ones(shape, dtype=torch.float32) * (j + 1)
+            ).all()
+
+    # some tensors in the list are pytorch, some are numpy
+    for i, a in enumerate(actors):
+        t = np.ones(shape, dtype=np.float32) * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = []
+        for j in range(world_size):
+            if j % 2 == 0:
+                list_buffer.append(torch.ones(shape, dtype=torch.float32))
+            else:
+                list_buffer.append(np.ones(shape, dtype=np.float32))
+        ray.wait([a.set_list_buffer.remote(list_buffer, copy=True)])
+    results = ray.get([a.do_allgather.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(world_size):
+            if j % 2 == 0:
+                assert (
+                    results[i][j] == torch.ones(shape, dtype=torch.float32) * (j + 1)
+                ).all()
+            else:
+                assert (
+                    results[i][j] == np.ones(shape, dtype=np.float32) * (j + 1)
+                ).all()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/util/collective/tests/single_node_cpu_tests/test_allreduce.py

@@ -0,0 +1,166 @@
+"""Test the collective allreduice API."""
+import numpy as np
+import pytest
+import torch
+
+import ray
+from ray.util.collective.tests.cpu_util import create_collective_workers
+from ray.util.collective.types import Backend, ReduceOp
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
+def test_allreduce_different_name(ray_start_single_node, group_name, backend):
+    world_size = 2
+    actors, _ = create_collective_workers(
+        num_workers=world_size, group_name=group_name, backend=backend
+    )
+    results = ray.get([a.do_allreduce.remote(group_name) for a in actors])
+    assert (results[0] == np.ones((10,), dtype=np.float32) * world_size).all()
+    assert (results[1] == np.ones((10,), dtype=np.float32) * world_size).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
+def test_allreduce_different_array_size(ray_start_single_node, array_size, backend):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    ray.wait(
+        [a.set_buffer.remote(np.ones(array_size, dtype=np.float32)) for a in actors]
+    )
+    results = ray.get([a.do_allreduce.remote() for a in actors])
+    assert (results[0] == np.ones((array_size,), dtype=np.float32) * world_size).all()
+    assert (results[1] == np.ones((array_size,), dtype=np.float32) * world_size).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+def test_allreduce_destroy(ray_start_single_node, backend, group_name="default"):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+
+    results = ray.get([a.do_allreduce.remote() for a in actors])
+    assert (results[0] == np.ones((10,), dtype=np.float32) * world_size).all()
+    assert (results[1] == np.ones((10,), dtype=np.float32) * world_size).all()
+
+    # destroy the group and try do work, should fail
+    ray.get([a.destroy_group.remote() for a in actors])
+    with pytest.raises(RuntimeError):
+        results = ray.get([a.do_allreduce.remote() for a in actors])
+
+    # reinit the same group and all reduce
+    ray.get(
+        [
+            actor.init_group.remote(world_size, i, backend, group_name)
+            for i, actor in enumerate(actors)
+        ]
+    )
+    results = ray.get([a.do_allreduce.remote() for a in actors])
+    assert (results[0] == np.ones((10,), dtype=np.float32) * world_size * 2).all()
+    assert (results[1] == np.ones((10,), dtype=np.float32) * world_size * 2).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+def test_allreduce_multiple_group(ray_start_single_node, backend, num_groups=5):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    for group_name in range(1, num_groups):
+        ray.get(
+            [
+                actor.init_group.remote(world_size, i, backend, str(group_name))
+                for i, actor in enumerate(actors)
+            ]
+        )
+    for i in range(num_groups):
+        group_name = "default" if i == 0 else str(i)
+        results = ray.get([a.do_allreduce.remote(group_name) for a in actors])
+        assert (
+            results[0] == np.ones((10,), dtype=np.float32) * (world_size ** (i + 1))
+        ).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+def test_allreduce_different_op(ray_start_single_node, backend):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+
+    # check product
+    ray.wait(
+        [
+            a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get([a.do_allreduce.remote(op=ReduceOp.PRODUCT) for a in actors])
+    assert (results[0] == np.ones((10,), dtype=np.float32) * 6).all()
+    assert (results[1] == np.ones((10,), dtype=np.float32) * 6).all()
+
+    # check min
+    ray.wait(
+        [
+            a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get([a.do_allreduce.remote(op=ReduceOp.MIN) for a in actors])
+    assert (results[0] == np.ones((10,), dtype=np.float32) * 2).all()
+    assert (results[1] == np.ones((10,), dtype=np.float32) * 2).all()
+
+    # check max
+    ray.wait(
+        [
+            a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get([a.do_allreduce.remote(op=ReduceOp.MAX) for a in actors])
+    assert (results[0] == np.ones((10,), dtype=np.float32) * 3).all()
+    assert (results[1] == np.ones((10,), dtype=np.float32) * 3).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("dtype", [np.uint8, np.float16, np.float32, np.float64])
+def test_allreduce_different_dtype(ray_start_single_node, dtype, backend):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    ray.wait([a.set_buffer.remote(np.ones(10, dtype=dtype)) for a in actors])
+    results = ray.get([a.do_allreduce.remote() for a in actors])
+    assert (results[0] == np.ones((10,), dtype=dtype) * world_size).all()
+    assert (results[1] == np.ones((10,), dtype=dtype) * world_size).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+def test_allreduce_torch_numpy(ray_start_single_node, backend):
+    # import torch
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    ray.wait(
+        [
+            actors[1].set_buffer.remote(
+                torch.ones(
+                    10,
+                )
+            )
+        ]
+    )
+    results = ray.get([a.do_allreduce.remote() for a in actors])
+    assert (results[0] == np.ones((10,)) * world_size).all()
+
+    ray.wait(
+        [
+            actors[0].set_buffer.remote(
+                torch.ones(
+                    10,
+                )
+            )
+        ]
+    )
+    ray.wait([actors[1].set_buffer.remote(np.ones(10, dtype=np.float32))])
+    ray.get([a.do_allreduce.remote() for a in actors])
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/util/collective/tests/single_node_cpu_tests/test_basic_apis.py

@@ -0,0 +1,130 @@
+"""Test the collective group APIs."""
+import pytest
+
+import ray
+from ray.util.collective.tests.cpu_util import Worker, create_collective_workers
+from ray.util.collective.types import Backend
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
+def test_init_two_actors(ray_start_single_node, group_name, backend):
+    world_size = 2
+    actors, results = create_collective_workers(world_size, group_name, backend=backend)
+    for i in range(world_size):
+        assert results[i]
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+def test_init_multiple_groups(ray_start_single_node, backend):
+    world_size = 2
+    num_groups = 10
+    actors = [Worker.remote() for i in range(world_size)]
+    for i in range(num_groups):
+        group_name = str(i)
+        init_results = ray.get(
+            [
+                actor.init_group.remote(
+                    world_size, k, group_name=group_name, backend=backend
+                )
+                for k, actor in enumerate(actors)
+            ]
+        )
+        for j in range(world_size):
+            assert init_results[j]
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+def test_get_rank(ray_start_single_node, backend):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    actor0_rank = ray.get(actors[0].report_rank.remote())
+    assert actor0_rank == 0
+    actor1_rank = ray.get(actors[1].report_rank.remote())
+    assert actor1_rank == 1
+
+    # create a second group with a different name,
+    # and different order of ranks.
+    new_group_name = "default2"
+    ray.get(
+        [
+            actor.init_group.remote(
+                world_size,
+                world_size - 1 - i,
+                group_name=new_group_name,
+                backend=backend,
+            )
+            for i, actor in enumerate(actors)
+        ]
+    )
+    actor0_rank = ray.get(actors[0].report_rank.remote(new_group_name))
+    assert actor0_rank == 1
+    actor1_rank = ray.get(actors[1].report_rank.remote(new_group_name))
+    assert actor1_rank == 0
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+def test_get_world_size(ray_start_single_node, backend):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    actor0_world_size = ray.get(actors[0].report_world_size.remote())
+    actor1_world_size = ray.get(actors[1].report_world_size.remote())
+    assert actor0_world_size == actor1_world_size == world_size
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+def test_is_group_initialized(ray_start_single_node, backend):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    # check group is_init
+    actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
+    assert actor0_is_init
+    actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("random"))
+    assert not actor0_is_init
+    actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("123"))
+    assert not actor0_is_init
+    actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
+    assert actor1_is_init
+    actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote("456"))
+    assert not actor1_is_init
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+def test_destroy_group(ray_start_single_node, backend):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    # Now destroy the group at actor0
+    ray.wait([actors[0].destroy_group.remote()])
+    actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
+    assert not actor0_is_init
+
+    # should go well as the group `random` does not exist at all
+    ray.wait([actors[0].destroy_group.remote("random")])
+
+    actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
+    assert actor1_is_init
+    ray.wait([actors[1].destroy_group.remote("random")])
+    actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
+    assert actor1_is_init
+    ray.wait([actors[1].destroy_group.remote("default")])
+    actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
+    assert not actor1_is_init
+
+    # Now reconstruct the group using the same name
+    init_results = ray.get(
+        [actor.init_group.remote(world_size, i) for i, actor in enumerate(actors)]
+    )
+    for i in range(world_size):
+        assert init_results[i]
+    actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
+    assert actor0_is_init
+    actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
+    assert actor1_is_init
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/util/collective/tests/single_node_cpu_tests/test_broadcast.py

@@ -0,0 +1,94 @@
+"""Test the broadcast API."""
+import numpy as np
+import pytest
+
+import ray
+from ray.util.collective.tests.cpu_util import create_collective_workers
+from ray.util.collective.types import Backend
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
+@pytest.mark.parametrize("src_rank", [0, 1])
+def test_broadcast_different_name(ray_start_single_node, group_name, src_rank, backend):
+    world_size = 2
+    actors, _ = create_collective_workers(
+        num_workers=world_size, group_name=group_name, backend=backend
+    )
+    ray.get(
+        [
+            a.set_buffer.remote(np.ones((10,), dtype=np.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get(
+        [
+            a.do_broadcast.remote(group_name=group_name, src_rank=src_rank)
+            for a in actors
+        ]
+    )
+    for i in range(world_size):
+        assert (results[i] == np.ones((10,), dtype=np.float32) * (src_rank + 2)).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
+@pytest.mark.parametrize("src_rank", [0, 1])
+def test_broadcast_different_array_size(
+    ray_start_single_node, array_size, src_rank, backend
+):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    ray.get(
+        [
+            a.set_buffer.remote(np.ones(array_size, dtype=np.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
+    for i in range(world_size):
+        assert (
+            results[i] == np.ones((array_size,), dtype=np.float32) * (src_rank + 2)
+        ).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("src_rank", [0, 1])
+def test_broadcast_torch_numpy(ray_start_single_node, src_rank, backend):
+    import torch
+
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    ray.wait(
+        [
+            actors[1].set_buffer.remote(
+                torch.ones(
+                    10,
+                )
+                * world_size
+            )
+        ]
+    )
+    results = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
+    if src_rank == 0:
+        assert (results[0] == np.ones((10,))).all()
+        assert (results[1] == torch.ones((10,))).all()
+    else:
+        assert (results[0] == np.ones((10,)) * world_size).all()
+        assert (results[1] == torch.ones((10,)) * world_size).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+def test_broadcast_invalid_rank(ray_start_single_node, backend, src_rank=3):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    with pytest.raises(ValueError):
+        ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/util/collective/tests/single_node_cpu_tests/test_gloo_group_isolation.py

@@ -0,0 +1,67 @@
+import time
+
+import ray
+import ray.util.collective as col
+from ray.util.collective.collective_group.torch_gloo_collective_group import (
+    TorchGLOOGroup as GLOOGroup,
+)
+from ray.util.collective.types import Backend
+
+
+@ray.remote
+class Worker:
+    def __init__(self):
+        pass
+
+    def init_gloo_group(
+        self, world_size: int, rank: int, group_name: str, gloo_timeout: int = 30000
+    ):
+        col.init_collective_group(
+            world_size, rank, Backend.GLOO, group_name, gloo_timeout
+        )
+        return True
+
+    def get_gloo_timeout(self, group_name: str) -> int:
+        g = col.get_group_handle(group_name)
+        # Check if the group is initialized correctly
+        assert isinstance(g, GLOOGroup)
+        return g._gloo_context.getTimeout()
+
+
+def test_two_groups_in_one_cluster(ray_start_single_node):
+    name1 = "name_1"
+    name2 = "name_2"
+    time1 = 40000
+    time2 = 60000
+    w1 = Worker.remote()
+    ret1 = w1.init_gloo_group.remote(1, 0, name1, time1)
+    w2 = Worker.remote()
+    ret2 = w2.init_gloo_group.remote(1, 0, name2, time2)
+    assert ray.get(ret1)
+    assert ray.get(ret2)
+    assert ray.get(w1.get_gloo_timeout.remote(name1)) == time1
+    assert ray.get(w2.get_gloo_timeout.remote(name2)) == time2
+
+
+def test_failure_when_initializing(shutdown_only):
+    # job1
+    ray.init()
+    w1 = Worker.remote()
+    ret1 = w1.init_gloo_group.remote(2, 0, "name_1")
+    ray.wait([ret1], timeout=1)
+    time.sleep(5)
+    ray.shutdown()
+
+    # job2
+    ray.init()
+    w2 = Worker.remote()
+    ret2 = w2.init_gloo_group.remote(1, 0, "name_1")
+    assert ray.get(ret2)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/util/collective/tests/single_node_cpu_tests/test_reduce.py

@@ -0,0 +1,167 @@
+"""Test the reduce API."""
+import numpy as np
+import pytest
+
+import ray
+from ray.util.collective.tests.cpu_util import create_collective_workers
+from ray.util.collective.types import Backend, ReduceOp
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
+@pytest.mark.parametrize("dst_rank", [0, 1])
+def test_reduce_different_name(ray_start_single_node, group_name, dst_rank, backend):
+    world_size = 2
+    actors, _ = create_collective_workers(
+        num_workers=world_size, group_name=group_name, backend=backend
+    )
+    results = ray.get([a.do_reduce.remote(group_name, dst_rank) for a in actors])
+    for i in range(world_size):
+        if i == dst_rank:
+            assert (results[i] == np.ones((10,), dtype=np.float32) * world_size).all()
+        else:
+            assert (results[i] == np.ones((10,), dtype=np.float32)).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
+@pytest.mark.parametrize("dst_rank", [0, 1])
+def test_reduce_different_array_size(
+    ray_start_single_node, array_size, dst_rank, backend
+):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    ray.wait(
+        [a.set_buffer.remote(np.ones(array_size, dtype=np.float32)) for a in actors]
+    )
+    results = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
+    for i in range(world_size):
+        if i == dst_rank:
+            assert (
+                results[i] == np.ones((array_size,), dtype=np.float32) * world_size
+            ).all()
+        else:
+            assert (results[i] == np.ones((array_size,), dtype=np.float32)).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("dst_rank", [0, 1])
+def test_reduce_multiple_group(ray_start_single_node, dst_rank, backend, num_groups=5):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    for group_name in range(1, num_groups):
+        ray.get(
+            [
+                actor.init_group.remote(world_size, i, backend, str(group_name))
+                for i, actor in enumerate(actors)
+            ]
+        )
+    for i in range(num_groups):
+        group_name = "default" if i == 0 else str(i)
+        results = ray.get(
+            [
+                a.do_reduce.remote(dst_rank=dst_rank, group_name=group_name)
+                for a in actors
+            ]
+        )
+        for j in range(world_size):
+            if j == dst_rank:
+                assert (results[j] == np.ones((10,), dtype=np.float32) * (i + 2)).all()
+            else:
+                assert (results[j] == np.ones((10,), dtype=np.float32)).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("dst_rank", [0, 1])
+def test_reduce_different_op(ray_start_single_node, dst_rank, backend):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+
+    # check product
+    ray.wait(
+        [
+            a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get(
+        [a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.PRODUCT) for a in actors]
+    )
+    for i in range(world_size):
+        if i == dst_rank:
+            assert (results[i] == np.ones((10,), dtype=np.float32) * 6).all()
+        else:
+            assert (results[i] == np.ones((10,), dtype=np.float32) * (i + 2)).all()
+
+    # check min
+    ray.wait(
+        [
+            a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get(
+        [a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.MIN) for a in actors]
+    )
+    for i in range(world_size):
+        if i == dst_rank:
+            assert (results[i] == np.ones((10,), dtype=np.float32) * 2).all()
+        else:
+            assert (results[i] == np.ones((10,), dtype=np.float32) * (i + 2)).all()
+
+    # check max
+    ray.wait(
+        [
+            a.set_buffer.remote(np.ones(10, dtype=np.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get(
+        [a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.MAX) for a in actors]
+    )
+    for i in range(world_size):
+        if i == dst_rank:
+            assert (results[i] == np.ones((10,), dtype=np.float32) * 3).all()
+        else:
+            assert (results[i] == np.ones((10,), dtype=np.float32) * (i + 2)).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("dst_rank", [0, 1])
+def test_reduce_torch_numpy(ray_start_single_node, dst_rank, backend):
+    import torch
+
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    ray.wait(
+        [
+            actors[1].set_buffer.remote(
+                torch.ones(
+                    10,
+                )
+            )
+        ]
+    )
+    results = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
+    if dst_rank == 0:
+        assert (results[0] == np.ones((10,)) * world_size).all()
+        assert (results[1] == torch.ones((10,))).all()
+    else:
+        assert (results[0] == np.ones((10,))).all()
+        assert (results[1] == torch.ones((10,)) * world_size).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+def test_reduce_invalid_rank(ray_start_single_node, backend, dst_rank=3):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    with pytest.raises(ValueError):
+        ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/util/collective/tests/single_node_cpu_tests/test_reducescatter.py

@@ -0,0 +1,131 @@
+"""Test the collective reducescatter API."""
+import numpy as np
+import pytest
+import torch
+
+import ray
+from ray.util.collective.tests.cpu_util import (
+    create_collective_workers,
+    init_tensors_for_gather_scatter,
+)
+from ray.util.collective.types import Backend
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("tensor_backend", ["numpy", "torch"])
+@pytest.mark.parametrize(
+    "array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
+)
+def test_reducescatter_different_array_size(
+    ray_start_single_node, array_size, tensor_backend, backend
+):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    init_tensors_for_gather_scatter(
+        actors, array_size=array_size, tensor_backend=tensor_backend
+    )
+    results = ray.get([a.do_reducescatter.remote() for a in actors])
+    for i in range(world_size):
+        if tensor_backend == "numpy":
+            assert (
+                results[i] == np.ones(array_size, dtype=np.float32) * world_size
+            ).all()
+        else:
+            assert (
+                results[i] == torch.ones(array_size, dtype=torch.float32) * world_size
+            ).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("dtype", [np.uint8, np.float16, np.float32, np.float64])
+def test_reducescatter_different_dtype(ray_start_single_node, dtype, backend):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    init_tensors_for_gather_scatter(actors, dtype=dtype)
+    results = ray.get([a.do_reducescatter.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(world_size):
+            assert (results[i] == np.ones(10, dtype=dtype) * world_size).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+def test_reducescatter_torch_numpy(ray_start_single_node, backend):
+    world_size = 2
+    shape = [10, 10]
+    actors, _ = create_collective_workers(world_size, backend=backend)
+
+    # tensor is pytorch, list is numpy
+    for i, a in enumerate(actors):
+        t = torch.ones(shape, dtype=torch.float32) * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = [np.ones(shape, dtype=np.float32) for _ in range(world_size)]
+        ray.wait([a.set_list_buffer.remote(list_buffer)])
+    results = ray.get([a.do_reducescatter.remote() for a in actors])
+    for i in range(world_size):
+        assert (results[i] == torch.ones(shape, dtype=torch.float32) * world_size).all()
+
+    # tensor is numpy, list is pytorch
+    for i, a in enumerate(actors):
+        t = np.ones(shape, dtype=np.float32) * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = [
+            torch.ones(shape, dtype=torch.float32) for _ in range(world_size)
+        ]
+        ray.wait([a.set_list_buffer.remote(list_buffer)])
+    results = ray.get([a.do_reducescatter.remote() for a in actors])
+    for i in range(world_size):
+        assert (results[i] == np.ones(shape, dtype=np.float32) * world_size).all()
+
+    # some tensors in the list are pytorch, some are numpy
+    for i, a in enumerate(actors):
+        if i % 2 == 0:
+            t = torch.ones(shape, dtype=torch.float32) * (i + 1)
+        else:
+            t = np.ones(shape, dtype=np.float32) * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = []
+        for j in range(world_size):
+            if j % 2 == 0:
+                list_buffer.append(torch.ones(shape, dtype=torch.float32))
+            else:
+                list_buffer.append(np.ones(shape, dtype=np.float32))
+        ray.wait([a.set_list_buffer.remote(list_buffer)])
+    results = ray.get([a.do_reducescatter.remote() for a in actors])
+    for i in range(world_size):
+        if i % 2 == 0:
+            assert (
+                results[i] == torch.ones(shape, dtype=torch.float32) * world_size
+            ).all()
+        else:
+            assert (results[i] == np.ones(shape, dtype=np.float32) * world_size).all()
+
+    # mixed case
+    for i, a in enumerate(actors):
+        if i % 2 == 0:
+            t = torch.ones(shape, dtype=torch.float32) * (i + 1)
+        else:
+            t = np.ones(shape, dtype=np.float32) * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = []
+        for j in range(world_size):
+            if j % 2 == 0:
+                list_buffer.append(np.ones(shape, dtype=np.float32))
+            else:
+                list_buffer.append(torch.ones(shape, dtype=torch.float32))
+        ray.wait([a.set_list_buffer.remote(list_buffer)])
+    results = ray.get([a.do_reducescatter.remote() for a in actors])
+    for i in range(world_size):
+        if i % 2 == 0:
+            assert (
+                results[i] == torch.ones(shape, dtype=torch.float32) * world_size
+            ).all()
+        else:
+            assert (results[i] == np.ones(shape, dtype=np.float32) * world_size).all()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/util/collective/tests/single_node_cpu_tests/test_sendrecv.py

@@ -0,0 +1,92 @@
+"""Test the send/recv API."""
+import numpy as np
+import pytest
+
+import ray
+from ray.util.collective.tests.cpu_util import create_collective_workers
+from ray.util.collective.types import Backend
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
+@pytest.mark.parametrize("dst_rank", [0, 1])
+@pytest.mark.parametrize(
+    "array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 9, 10, 85]]
+)
+def test_reduce_different_name(
+    ray_start_single_node, group_name, array_size, dst_rank, backend
+):
+    world_size = 2
+    actors, _ = create_collective_workers(
+        num_workers=world_size, group_name=group_name, backend=backend
+    )
+    ray.wait(
+        [
+            a.set_buffer.remote(np.ones(array_size, dtype=np.float32) * (i + 1))
+            for i, a in enumerate(actors)
+        ]
+    )
+    src_rank = 1 - dst_rank
+    refs = []
+    for i, actor in enumerate(actors):
+        if i != dst_rank:
+            ref = actor.do_send.remote(group_name, dst_rank)
+        else:
+            ref = actor.do_recv.remote(group_name, src_rank)
+        refs.append(ref)
+    results = ray.get(refs)
+    for i in range(world_size):
+        assert (
+            results[i] == np.ones(array_size, dtype=np.float32) * (src_rank + 1)
+        ).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+@pytest.mark.parametrize("dst_rank", [0, 1])
+def test_sendrecv_torch_numpy(ray_start_single_node, dst_rank, backend):
+    import torch
+
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    ray.wait(
+        [
+            actors[1].set_buffer.remote(
+                torch.ones(
+                    10,
+                )
+                * 2
+            )
+        ]
+    )
+    src_rank = 1 - dst_rank
+
+    refs = []
+    for i, actor in enumerate(actors):
+        if i != dst_rank:
+            ref = actor.do_send.remote(dst_rank=dst_rank)
+        else:
+            ref = actor.do_recv.remote(src_rank=src_rank)
+        refs.append(ref)
+    results = ray.get(refs)
+    if dst_rank == 0:
+        assert (results[0] == np.ones((10,)) * 2).all()
+        assert (results[1] == torch.ones((10,)) * 2).all()
+    else:
+        assert (results[0] == np.ones((10,))).all()
+        assert (results[1] == torch.ones((10,))).all()
+
+
+@pytest.mark.parametrize("backend", [Backend.GLOO])
+def test_sendrecv_invalid_rank(ray_start_single_node, backend, dst_rank=3):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size, backend=backend)
+    with pytest.raises(ValueError):
+        ray.get([a.do_send.remote(dst_rank=dst_rank) for a in actors])
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/util/collective/tests/single_node_gpu_tests/test_allgather.py

@@ -0,0 +1,138 @@
+"""Test the collective allgather API."""
+import cupy as cp
+import pytest
+import torch
+
+import ray
+from ray.util.collective.tests.util import (
+    create_collective_workers,
+    init_tensors_for_gather_scatter,
+)
+
+
+@pytest.mark.parametrize("tensor_backend", ["cupy", "torch"])
+@pytest.mark.parametrize(
+    "array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
+)
+def test_allgather_different_array_size(
+    ray_start_single_node_2_gpus, array_size, tensor_backend
+):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+    init_tensors_for_gather_scatter(
+        actors, array_size=array_size, tensor_backend=tensor_backend
+    )
+    results = ray.get([a.do_allgather.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(world_size):
+            if tensor_backend == "cupy":
+                assert (
+                    results[i][j] == cp.ones(array_size, dtype=cp.float32) * (j + 1)
+                ).all()
+            else:
+                assert (
+                    results[i][j]
+                    == torch.ones(array_size, dtype=torch.float32).cuda() * (j + 1)
+                ).all()
+
+
+@pytest.mark.parametrize("dtype", [cp.uint8, cp.float16, cp.float32, cp.float64])
+def test_allgather_different_dtype(ray_start_single_node_2_gpus, dtype):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+    init_tensors_for_gather_scatter(actors, dtype=dtype)
+    results = ray.get([a.do_allgather.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(world_size):
+            assert (results[i][j] == cp.ones(10, dtype=dtype) * (j + 1)).all()
+
+
+@pytest.mark.parametrize("length", [0, 1, 2, 3])
+def test_unmatched_tensor_list_length(ray_start_single_node_2_gpus, length):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+    list_buffer = [cp.ones(10, dtype=cp.float32) for _ in range(length)]
+    ray.wait([a.set_list_buffer.remote(list_buffer) for a in actors])
+    if length != world_size:
+        with pytest.raises(RuntimeError):
+            ray.get([a.do_allgather.remote() for a in actors])
+    else:
+        ray.get([a.do_allgather.remote() for a in actors])
+
+
+@pytest.mark.parametrize("shape", [10, 20, [4, 5], [1, 3, 5, 7]])
+def test_unmatched_tensor_shape(ray_start_single_node_2_gpus, shape):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+    init_tensors_for_gather_scatter(actors, array_size=10)
+    list_buffer = [cp.ones(shape, dtype=cp.float32) for _ in range(world_size)]
+    ray.get([a.set_list_buffer.remote(list_buffer) for a in actors])
+    if shape != 10:
+        with pytest.raises(RuntimeError):
+            ray.get([a.do_allgather.remote() for a in actors])
+    else:
+        ray.get([a.do_allgather.remote() for a in actors])
+
+
+def test_allgather_torch_cupy(ray_start_single_node_2_gpus):
+    world_size = 2
+    shape = [10, 10]
+    actors, _ = create_collective_workers(world_size)
+
+    # tensor is pytorch, list is cupy
+    for i, a in enumerate(actors):
+        t = torch.ones(shape, dtype=torch.float32).cuda() * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = [cp.ones(shape, dtype=cp.float32) for _ in range(world_size)]
+        ray.wait([a.set_list_buffer.remote(list_buffer)])
+    results = ray.get([a.do_allgather.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(world_size):
+            assert (results[i][j] == cp.ones(shape, dtype=cp.float32) * (j + 1)).all()
+
+    # tensor is cupy, list is pytorch
+    for i, a in enumerate(actors):
+        t = cp.ones(shape, dtype=cp.float32) * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = [
+            torch.ones(shape, dtype=torch.float32).cuda() for _ in range(world_size)
+        ]
+        ray.wait([a.set_list_buffer.remote(list_buffer)])
+    results = ray.get([a.do_allgather.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(world_size):
+            assert (
+                results[i][j] == torch.ones(shape, dtype=torch.float32).cuda() * (j + 1)
+            ).all()
+
+    # some tensors in the list are pytorch, some are cupy
+    for i, a in enumerate(actors):
+        t = cp.ones(shape, dtype=cp.float32) * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = []
+        for j in range(world_size):
+            if j % 2 == 0:
+                list_buffer.append(torch.ones(shape, dtype=torch.float32).cuda())
+            else:
+                list_buffer.append(cp.ones(shape, dtype=cp.float32))
+        ray.wait([a.set_list_buffer.remote(list_buffer)])
+    results = ray.get([a.do_allgather.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(world_size):
+            if j % 2 == 0:
+                assert (
+                    results[i][j]
+                    == torch.ones(shape, dtype=torch.float32).cuda() * (j + 1)
+                ).all()
+            else:
+                assert (
+                    results[i][j] == cp.ones(shape, dtype=cp.float32) * (j + 1)
+                ).all()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/util/collective/tests/single_node_gpu_tests/test_allreduce.py

@@ -0,0 +1,170 @@
+"""Test the collective allreduice API."""
+import cupy as cp
+import pytest
+import torch
+
+import ray
+from ray.util.collective.tests.util import create_collective_workers
+from ray.util.collective.types import ReduceOp
+
+
+@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
+def test_allreduce_different_name(ray_start_single_node_2_gpus, group_name):
+    world_size = 2
+    actors, _ = create_collective_workers(num_workers=world_size, group_name=group_name)
+    results = ray.get([a.do_allreduce.remote(group_name) for a in actors])
+    assert (results[0] == cp.ones((10,), dtype=cp.float32) * world_size).all()
+    assert (results[1] == cp.ones((10,), dtype=cp.float32) * world_size).all()
+
+
+@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
+def test_allreduce_different_array_size(ray_start_single_node_2_gpus, array_size):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+    ray.wait(
+        [a.set_buffer.remote(cp.ones(array_size, dtype=cp.float32)) for a in actors]
+    )
+    results = ray.get([a.do_allreduce.remote() for a in actors])
+    assert (results[0] == cp.ones((array_size,), dtype=cp.float32) * world_size).all()
+    assert (results[1] == cp.ones((array_size,), dtype=cp.float32) * world_size).all()
+
+
+def test_allreduce_destroy(
+    ray_start_single_node_2_gpus, backend="nccl", group_name="default"
+):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+
+    results = ray.get([a.do_allreduce.remote() for a in actors])
+    assert (results[0] == cp.ones((10,), dtype=cp.float32) * world_size).all()
+    assert (results[1] == cp.ones((10,), dtype=cp.float32) * world_size).all()
+
+    # destroy the group and try do work, should fail
+    ray.get([a.destroy_group.remote() for a in actors])
+    with pytest.raises(RuntimeError):
+        results = ray.get([a.do_allreduce.remote() for a in actors])
+
+    # reinit the same group and all reduce
+    ray.get(
+        [
+            actor.init_group.remote(world_size, i, backend, group_name)
+            for i, actor in enumerate(actors)
+        ]
+    )
+    results = ray.get([a.do_allreduce.remote() for a in actors])
+    assert (results[0] == cp.ones((10,), dtype=cp.float32) * world_size * 2).all()
+    assert (results[1] == cp.ones((10,), dtype=cp.float32) * world_size * 2).all()
+
+
+def test_allreduce_multiple_group(
+    ray_start_single_node_2_gpus, backend="nccl", num_groups=5
+):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+    for group_name in range(1, num_groups):
+        ray.get(
+            [
+                actor.init_group.remote(world_size, i, backend, str(group_name))
+                for i, actor in enumerate(actors)
+            ]
+        )
+    for i in range(num_groups):
+        group_name = "default" if i == 0 else str(i)
+        results = ray.get([a.do_allreduce.remote(group_name) for a in actors])
+        assert (
+            results[0] == cp.ones((10,), dtype=cp.float32) * (world_size ** (i + 1))
+        ).all()
+
+
+def test_allreduce_different_op(ray_start_single_node_2_gpus):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+
+    # check product
+    ray.wait(
+        [
+            a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get([a.do_allreduce.remote(op=ReduceOp.PRODUCT) for a in actors])
+    assert (results[0] == cp.ones((10,), dtype=cp.float32) * 6).all()
+    assert (results[1] == cp.ones((10,), dtype=cp.float32) * 6).all()
+
+    # check min
+    ray.wait(
+        [
+            a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get([a.do_allreduce.remote(op=ReduceOp.MIN) for a in actors])
+    assert (results[0] == cp.ones((10,), dtype=cp.float32) * 2).all()
+    assert (results[1] == cp.ones((10,), dtype=cp.float32) * 2).all()
+
+    # check max
+    ray.wait(
+        [
+            a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get([a.do_allreduce.remote(op=ReduceOp.MAX) for a in actors])
+    assert (results[0] == cp.ones((10,), dtype=cp.float32) * 3).all()
+    assert (results[1] == cp.ones((10,), dtype=cp.float32) * 3).all()
+
+
+@pytest.mark.parametrize("dtype", [cp.uint8, cp.float16, cp.float32, cp.float64])
+def test_allreduce_different_dtype(ray_start_single_node_2_gpus, dtype):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+    ray.wait([a.set_buffer.remote(cp.ones(10, dtype=dtype)) for a in actors])
+    results = ray.get([a.do_allreduce.remote() for a in actors])
+    assert (results[0] == cp.ones((10,), dtype=dtype) * world_size).all()
+    assert (results[1] == cp.ones((10,), dtype=dtype) * world_size).all()
+
+
+def test_allreduce_torch_cupy(ray_start_single_node_2_gpus):
+    # import torch
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+    ray.wait(
+        [
+            actors[1].set_buffer.remote(
+                torch.ones(
+                    10,
+                ).cuda()
+            )
+        ]
+    )
+    results = ray.get([a.do_allreduce.remote() for a in actors])
+    assert (results[0] == cp.ones((10,)) * world_size).all()
+
+    ray.wait(
+        [
+            actors[0].set_buffer.remote(
+                torch.ones(
+                    10,
+                )
+            )
+        ]
+    )
+    ray.wait(
+        [
+            actors[1].set_buffer.remote(
+                cp.ones(
+                    10,
+                )
+            )
+        ]
+    )
+    with pytest.raises(RuntimeError):
+        results = ray.get([a.do_allreduce.remote() for a in actors])
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/util/collective/tests/single_node_gpu_tests/test_basic_apis.py

@@ -0,0 +1,118 @@
+"""Test the collective group APIs."""
+import pytest
+
+import ray
+from ray.util.collective.tests.util import Worker, create_collective_workers
+
+
+@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
+def test_init_two_actors(ray_start_single_node_2_gpus, group_name):
+    world_size = 2
+    actors, results = create_collective_workers(world_size, group_name)
+    for i in range(world_size):
+        assert results[i]
+
+
+def test_init_multiple_groups(ray_start_single_node_2_gpus):
+    world_size = 2
+    num_groups = 10
+    actors = [Worker.remote() for i in range(world_size)]
+    for i in range(num_groups):
+        group_name = str(i)
+        init_results = ray.get(
+            [
+                actor.init_group.remote(world_size, k, group_name=group_name)
+                for k, actor in enumerate(actors)
+            ]
+        )
+        for j in range(world_size):
+            assert init_results[j]
+
+
+def test_get_rank(ray_start_single_node_2_gpus):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+    actor0_rank = ray.get(actors[0].report_rank.remote())
+    assert actor0_rank == 0
+    actor1_rank = ray.get(actors[1].report_rank.remote())
+    assert actor1_rank == 1
+
+    # create a second group with a different name,
+    # and different order of ranks.
+    new_group_name = "default2"
+    ray.get(
+        [
+            actor.init_group.remote(
+                world_size, world_size - 1 - i, group_name=new_group_name
+            )
+            for i, actor in enumerate(actors)
+        ]
+    )
+    actor0_rank = ray.get(actors[0].report_rank.remote(new_group_name))
+    assert actor0_rank == 1
+    actor1_rank = ray.get(actors[1].report_rank.remote(new_group_name))
+    assert actor1_rank == 0
+
+
+def test_get_collective_group_size(ray_start_single_node_2_gpus):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+    actor0_world_size = ray.get(actors[0].report_world_size.remote())
+    actor1_world_size = ray.get(actors[1].report_world_size.remote())
+    assert actor0_world_size == actor1_world_size == world_size
+
+
+def test_is_group_initialized(ray_start_single_node_2_gpus):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+    # check group is_init
+    actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
+    assert actor0_is_init
+    actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("random"))
+    assert not actor0_is_init
+    actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote("123"))
+    assert not actor0_is_init
+    actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
+    assert actor1_is_init
+    actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote("456"))
+    assert not actor1_is_init
+
+
+def test_destroy_group(ray_start_single_node_2_gpus):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+    # Now destroy the group at actor0
+    ray.wait([actors[0].destroy_group.remote()])
+    actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
+    assert not actor0_is_init
+
+    # should go well as the group `random` does not exist at all
+    ray.wait([actors[0].destroy_group.remote("random")])
+
+    actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
+    assert actor1_is_init
+    ray.wait([actors[1].destroy_group.remote("random")])
+    actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
+    assert actor1_is_init
+    ray.wait([actors[1].destroy_group.remote("default")])
+    actor1_is_init = ray.get(actors[1].report_is_group_initialized.remote())
+    assert not actor1_is_init
+
+    # Now reconstruct the group using the same name
+    init_results = ray.get(
+        [actor.init_group.remote(world_size, i) for i, actor in enumerate(actors)]
+    )
+    for i in range(world_size):
+        assert init_results[i]
+    actor0_is_init = ray.get(actors[0].report_is_group_initialized.remote())
+    assert actor0_is_init
+    actor1_is_init = ray.get(actors[0].report_is_group_initialized.remote())
+    assert actor1_is_init
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/util/collective/tests/single_node_gpu_tests/test_broadcast.py

@@ -0,0 +1,79 @@
+"""Test the broadcast API."""
+import cupy as cp
+import pytest
+
+import ray
+from ray.util.collective.tests.util import create_collective_workers
+
+
+@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
+@pytest.mark.parametrize("src_rank", [0, 1])
+def test_broadcast_different_name(ray_start_single_node_2_gpus, group_name, src_rank):
+    world_size = 2
+    actors, _ = create_collective_workers(num_workers=world_size, group_name=group_name)
+    ray.wait(
+        [
+            a.set_buffer.remote(cp.ones((10,), dtype=cp.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get(
+        [
+            a.do_broadcast.remote(group_name=group_name, src_rank=src_rank)
+            for a in actors
+        ]
+    )
+    for i in range(world_size):
+        assert (results[i] == cp.ones((10,), dtype=cp.float32) * (src_rank + 2)).all()
+
+
+@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
+@pytest.mark.parametrize("src_rank", [0, 1])
+def test_broadcast_different_array_size(
+    ray_start_single_node_2_gpus, array_size, src_rank
+):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+    ray.wait(
+        [
+            a.set_buffer.remote(cp.ones(array_size, dtype=cp.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
+    for i in range(world_size):
+        assert (
+            results[i] == cp.ones((array_size,), dtype=cp.float32) * (src_rank + 2)
+        ).all()
+
+
+@pytest.mark.parametrize("src_rank", [0, 1])
+def test_broadcast_torch_cupy(ray_start_single_node_2_gpus, src_rank):
+    import torch
+
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+    ray.wait(
+        [
+            actors[1].set_buffer.remote(
+                torch.ones(
+                    10,
+                ).cuda()
+                * world_size
+            )
+        ]
+    )
+    results = ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])
+    if src_rank == 0:
+        assert (results[0] == cp.ones((10,))).all()
+        assert (results[1] == torch.ones((10,)).cuda()).all()
+    else:
+        assert (results[0] == cp.ones((10,)) * world_size).all()
+        assert (results[1] == torch.ones((10,)).cuda() * world_size).all()
+
+
+def test_broadcast_invalid_rank(ray_start_single_node_2_gpus, src_rank=3):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+    with pytest.raises(ValueError):
+        ray.get([a.do_broadcast.remote(src_rank=src_rank) for a in actors])

python/ray/util/collective/tests/single_node_gpu_tests/test_reduce.py

@@ -0,0 +1,151 @@
+"""Test the reduce API."""
+import cupy as cp
+import pytest
+
+import ray
+from ray.util.collective.tests.util import create_collective_workers
+from ray.util.collective.types import ReduceOp
+
+
+@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
+@pytest.mark.parametrize("dst_rank", [0, 1])
+def test_reduce_different_name(ray_start_single_node_2_gpus, group_name, dst_rank):
+    world_size = 2
+    actors, _ = create_collective_workers(num_workers=world_size, group_name=group_name)
+    results = ray.get([a.do_reduce.remote(group_name, dst_rank) for a in actors])
+    for i in range(world_size):
+        if i == dst_rank:
+            assert (results[i] == cp.ones((10,), dtype=cp.float32) * world_size).all()
+        else:
+            assert (results[i] == cp.ones((10,), dtype=cp.float32)).all()
+
+
+@pytest.mark.parametrize("array_size", [2, 2**5, 2**10, 2**15, 2**20])
+@pytest.mark.parametrize("dst_rank", [0, 1])
+def test_reduce_different_array_size(
+    ray_start_single_node_2_gpus, array_size, dst_rank
+):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+    ray.wait(
+        [a.set_buffer.remote(cp.ones(array_size, dtype=cp.float32)) for a in actors]
+    )
+    results = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
+    for i in range(world_size):
+        if i == dst_rank:
+            assert (
+                results[i] == cp.ones((array_size,), dtype=cp.float32) * world_size
+            ).all()
+        else:
+            assert (results[i] == cp.ones((array_size,), dtype=cp.float32)).all()
+
+
+@pytest.mark.parametrize("dst_rank", [0, 1])
+def test_reduce_multiple_group(ray_start_single_node_2_gpus, dst_rank, num_groups=5):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+    for group_name in range(1, num_groups):
+        ray.get(
+            [
+                actor.init_group.remote(world_size, i, "nccl", str(group_name))
+                for i, actor in enumerate(actors)
+            ]
+        )
+    for i in range(num_groups):
+        group_name = "default" if i == 0 else str(i)
+        results = ray.get(
+            [
+                a.do_reduce.remote(dst_rank=dst_rank, group_name=group_name)
+                for a in actors
+            ]
+        )
+        for j in range(world_size):
+            if j == dst_rank:
+                assert (results[j] == cp.ones((10,), dtype=cp.float32) * (i + 2)).all()
+            else:
+                assert (results[j] == cp.ones((10,), dtype=cp.float32)).all()
+
+
+@pytest.mark.parametrize("dst_rank", [0, 1])
+def test_reduce_different_op(ray_start_single_node_2_gpus, dst_rank):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+
+    # check product
+    ray.wait(
+        [
+            a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get(
+        [a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.PRODUCT) for a in actors]
+    )
+    for i in range(world_size):
+        if i == dst_rank:
+            assert (results[i] == cp.ones((10,), dtype=cp.float32) * 6).all()
+        else:
+            assert (results[i] == cp.ones((10,), dtype=cp.float32) * (i + 2)).all()
+
+    # check min
+    ray.wait(
+        [
+            a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get(
+        [a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.MIN) for a in actors]
+    )
+    for i in range(world_size):
+        if i == dst_rank:
+            assert (results[i] == cp.ones((10,), dtype=cp.float32) * 2).all()
+        else:
+            assert (results[i] == cp.ones((10,), dtype=cp.float32) * (i + 2)).all()
+
+    # check max
+    ray.wait(
+        [
+            a.set_buffer.remote(cp.ones(10, dtype=cp.float32) * (i + 2))
+            for i, a in enumerate(actors)
+        ]
+    )
+    results = ray.get(
+        [a.do_reduce.remote(dst_rank=dst_rank, op=ReduceOp.MAX) for a in actors]
+    )
+    for i in range(world_size):
+        if i == dst_rank:
+            assert (results[i] == cp.ones((10,), dtype=cp.float32) * 3).all()
+        else:
+            assert (results[i] == cp.ones((10,), dtype=cp.float32) * (i + 2)).all()
+
+
+@pytest.mark.parametrize("dst_rank", [0, 1])
+def test_reduce_torch_cupy(ray_start_single_node_2_gpus, dst_rank):
+    import torch
+
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+    ray.wait(
+        [
+            actors[1].set_buffer.remote(
+                torch.ones(
+                    10,
+                ).cuda()
+            )
+        ]
+    )
+    results = ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])
+    if dst_rank == 0:
+        assert (results[0] == cp.ones((10,)) * world_size).all()
+        assert (results[1] == torch.ones((10,)).cuda()).all()
+    else:
+        assert (results[0] == cp.ones((10,))).all()
+        assert (results[1] == torch.ones((10,)).cuda() * world_size).all()
+
+
+def test_reduce_invalid_rank(ray_start_single_node_2_gpus, dst_rank=3):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+    with pytest.raises(ValueError):
+        ray.get([a.do_reduce.remote(dst_rank=dst_rank) for a in actors])

python/ray/util/collective/tests/single_node_gpu_tests/test_reducescatter.py

@@ -0,0 +1,130 @@
+"""Test the collective reducescatter API."""
+import cupy as cp
+import pytest
+import torch
+
+import ray
+from ray.util.collective.tests.util import (
+    create_collective_workers,
+    init_tensors_for_gather_scatter,
+)
+
+
+@pytest.mark.parametrize("tensor_backend", ["cupy", "torch"])
+@pytest.mark.parametrize(
+    "array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
+)
+def test_reducescatter_different_array_size(
+    ray_start_single_node_2_gpus, array_size, tensor_backend
+):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+    init_tensors_for_gather_scatter(
+        actors, array_size=array_size, tensor_backend=tensor_backend
+    )
+    results = ray.get([a.do_reducescatter.remote() for a in actors])
+    for i in range(world_size):
+        if tensor_backend == "cupy":
+            assert (
+                results[i] == cp.ones(array_size, dtype=cp.float32) * world_size
+            ).all()
+        else:
+            assert (
+                results[i]
+                == torch.ones(array_size, dtype=torch.float32).cuda() * world_size
+            ).all()
+
+
+@pytest.mark.parametrize("dtype", [cp.uint8, cp.float16, cp.float32, cp.float64])
+def test_reducescatter_different_dtype(ray_start_single_node_2_gpus, dtype):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+    init_tensors_for_gather_scatter(actors, dtype=dtype)
+    results = ray.get([a.do_reducescatter.remote() for a in actors])
+    for i in range(world_size):
+        for j in range(world_size):
+            assert (results[i] == cp.ones(10, dtype=dtype) * world_size).all()
+
+
+def test_reducescatter_torch_cupy(ray_start_single_node_2_gpus):
+    world_size = 2
+    shape = [10, 10]
+    actors, _ = create_collective_workers(world_size)
+
+    # tensor is pytorch, list is cupy
+    for i, a in enumerate(actors):
+        t = torch.ones(shape, dtype=torch.float32).cuda() * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = [cp.ones(shape, dtype=cp.float32) for _ in range(world_size)]
+        ray.wait([a.set_list_buffer.remote(list_buffer)])
+    results = ray.get([a.do_reducescatter.remote() for a in actors])
+    for i in range(world_size):
+        assert (
+            results[i] == torch.ones(shape, dtype=torch.float32).cuda() * world_size
+        ).all()
+
+    # tensor is cupy, list is pytorch
+    for i, a in enumerate(actors):
+        t = cp.ones(shape, dtype=cp.float32) * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = [
+            torch.ones(shape, dtype=torch.float32).cuda() for _ in range(world_size)
+        ]
+        ray.wait([a.set_list_buffer.remote(list_buffer)])
+    results = ray.get([a.do_reducescatter.remote() for a in actors])
+    for i in range(world_size):
+        assert (results[i] == cp.ones(shape, dtype=cp.float32) * world_size).all()
+
+    # some tensors in the list are pytorch, some are cupy
+    for i, a in enumerate(actors):
+        if i % 2 == 0:
+            t = torch.ones(shape, dtype=torch.float32).cuda() * (i + 1)
+        else:
+            t = cp.ones(shape, dtype=cp.float32) * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = []
+        for j in range(world_size):
+            if j % 2 == 0:
+                list_buffer.append(torch.ones(shape, dtype=torch.float32).cuda())
+            else:
+                list_buffer.append(cp.ones(shape, dtype=cp.float32))
+        ray.wait([a.set_list_buffer.remote(list_buffer)])
+    results = ray.get([a.do_reducescatter.remote() for a in actors])
+    for i in range(world_size):
+        if i % 2 == 0:
+            assert (
+                results[i] == torch.ones(shape, dtype=torch.float32).cuda() * world_size
+            ).all()
+        else:
+            assert (results[i] == cp.ones(shape, dtype=cp.float32) * world_size).all()
+
+    # mixed case
+    for i, a in enumerate(actors):
+        if i % 2 == 0:
+            t = torch.ones(shape, dtype=torch.float32).cuda() * (i + 1)
+        else:
+            t = cp.ones(shape, dtype=cp.float32) * (i + 1)
+        ray.wait([a.set_buffer.remote(t)])
+        list_buffer = []
+        for j in range(world_size):
+            if j % 2 == 0:
+                list_buffer.append(cp.ones(shape, dtype=cp.float32))
+            else:
+                list_buffer.append(torch.ones(shape, dtype=torch.float32).cuda())
+        ray.wait([a.set_list_buffer.remote(list_buffer)])
+    results = ray.get([a.do_reducescatter.remote() for a in actors])
+    for i in range(world_size):
+        if i % 2 == 0:
+            assert (
+                results[i] == torch.ones(shape, dtype=torch.float32).cuda() * world_size
+            ).all()
+        else:
+            assert (results[i] == cp.ones(shape, dtype=cp.float32) * world_size).all()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/util/collective/tests/single_node_gpu_tests/test_sendrecv.py

@@ -0,0 +1,78 @@
+"""Test the send/recv API."""
+import cupy as cp
+import pytest
+
+import ray
+from ray.util.collective.tests.util import create_collective_workers
+
+
+@pytest.mark.parametrize("group_name", ["default", "test", "123?34!"])
+@pytest.mark.parametrize("dst_rank", [0, 1])
+@pytest.mark.parametrize(
+    "array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 9, 10, 85]]
+)
+def test_reduce_different_name(
+    ray_start_single_node_2_gpus, group_name, array_size, dst_rank
+):
+    world_size = 2
+    actors, _ = create_collective_workers(num_workers=world_size, group_name=group_name)
+    ray.wait(
+        [
+            a.set_buffer.remote(cp.ones(array_size, dtype=cp.float32) * (i + 1))
+            for i, a in enumerate(actors)
+        ]
+    )
+    src_rank = 1 - dst_rank
+    refs = []
+    for i, actor in enumerate(actors):
+        if i != dst_rank:
+            ref = actor.do_send.remote(group_name, dst_rank)
+        else:
+            ref = actor.do_recv.remote(group_name, src_rank)
+        refs.append(ref)
+    results = ray.get(refs)
+    for i in range(world_size):
+        assert (
+            results[i] == cp.ones(array_size, dtype=cp.float32) * (src_rank + 1)
+        ).all()
+
+
+@pytest.mark.parametrize("dst_rank", [0, 1])
+def test_sendrecv_torch_cupy(ray_start_single_node_2_gpus, dst_rank):
+    import torch
+
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+    ray.wait(
+        [
+            actors[1].set_buffer.remote(
+                torch.ones(
+                    10,
+                ).cuda()
+                * 2
+            )
+        ]
+    )
+    src_rank = 1 - dst_rank
+
+    refs = []
+    for i, actor in enumerate(actors):
+        if i != dst_rank:
+            ref = actor.do_send.remote(dst_rank=dst_rank)
+        else:
+            ref = actor.do_recv.remote(src_rank=src_rank)
+        refs.append(ref)
+    results = ray.get(refs)
+    if dst_rank == 0:
+        assert (results[0] == cp.ones((10,)) * 2).all()
+        assert (results[1] == torch.ones((10,)).cuda() * 2).all()
+    else:
+        assert (results[0] == cp.ones((10,))).all()
+        assert (results[1] == torch.ones((10,)).cuda()).all()
+
+
+def test_sendrecv_invalid_rank(ray_start_single_node_2_gpus, dst_rank=3):
+    world_size = 2
+    actors, _ = create_collective_workers(world_size)
+    with pytest.raises(ValueError):
+        ray.get([a.do_send.remote(dst_rank=dst_rank) for a in actors])

python/ray/util/collective/tests/util.py

@@ -0,0 +1,373 @@
+import logging
+
+import cupy as cp
+import torch
+
+import ray
+import ray.util.collective as col
+from ray.util.collective.collective_group.nccl_util import get_num_gpus
+from ray.util.collective.types import Backend, ReduceOp
+
+logger = logging.getLogger(__name__)
+
+
+@ray.remote(num_gpus=1)
+class Worker:
+    def __init__(self):
+        self.buffer = None
+        self.list_buffer = None
+
+    def init_tensors(self):
+        self.buffer = cp.ones((10,), dtype=cp.float32)
+        self.list_buffer = [cp.ones((10,), dtype=cp.float32) for _ in range(2)]
+        cp.cuda.Stream.null.synchronize()
+        return True
+
+    def init_group(self, world_size, rank, backend=Backend.NCCL, group_name="default"):
+        col.init_collective_group(world_size, rank, backend, group_name)
+        return True
+
+    def set_buffer(self, data):
+        self.buffer = data
+        return self.buffer
+
+    def get_buffer(self):
+        return self.buffer
+
+    def set_list_buffer(self, list_of_arrays):
+        self.list_buffer = list_of_arrays
+        return self.list_buffer
+
+    def do_allreduce(self, group_name="default", op=ReduceOp.SUM):
+        col.allreduce(self.buffer, group_name, op)
+        return self.buffer
+
+    def do_reduce(self, group_name="default", dst_rank=0, op=ReduceOp.SUM):
+        col.reduce(self.buffer, dst_rank, group_name, op)
+        return self.buffer
+
+    def do_broadcast(self, group_name="default", src_rank=0):
+        col.broadcast(self.buffer, src_rank, group_name)
+        return self.buffer
+
+    def do_allgather(self, group_name="default"):
+        col.allgather(self.list_buffer, self.buffer, group_name)
+        return self.list_buffer
+
+    def do_reducescatter(self, group_name="default", op=ReduceOp.SUM):
+        col.reducescatter(self.buffer, self.list_buffer, group_name, op)
+        return self.buffer
+
+    def do_send(self, group_name="default", dst_rank=0):
+        col.send(self.buffer, dst_rank, group_name)
+        return self.buffer
+
+    def do_recv(self, group_name="default", src_rank=0):
+        col.recv(self.buffer, src_rank, group_name)
+        return self.buffer
+
+    def destroy_group(self, group_name="default"):
+        col.destroy_collective_group(group_name)
+        return True
+
+    def report_rank(self, group_name="default"):
+        rank = col.get_rank(group_name)
+        return rank
+
+    def report_world_size(self, group_name="default"):
+        ws = col.get_collective_group_size(group_name)
+        return ws
+
+    def report_nccl_availability(self):
+        avail = col.nccl_available()
+        return avail
+
+    def report_gloo_availability(self):
+        avail = col.gloo_available()
+        return avail
+
+    def report_is_group_initialized(self, group_name="default"):
+        is_init = col.is_group_initialized(group_name)
+        return is_init
+
+
+def create_collective_workers(num_workers=2, group_name="default", backend="nccl"):
+    actors = [None] * num_workers
+    for i in range(num_workers):
+        actor = Worker.remote()
+        ray.get([actor.init_tensors.remote()])
+        actors[i] = actor
+    world_size = num_workers
+    init_results = ray.get(
+        [
+            actor.init_group.remote(world_size, i, backend, group_name)
+            for i, actor in enumerate(actors)
+        ]
+    )
+    return actors, init_results
+
+
+def init_tensors_for_gather_scatter(
+    actors, array_size=10, dtype=cp.float32, tensor_backend="cupy"
+):
+    world_size = len(actors)
+    for i, a in enumerate(actors):
+        if tensor_backend == "cupy":
+            t = cp.ones(array_size, dtype=dtype) * (i + 1)
+        elif tensor_backend == "torch":
+            t = torch.ones(array_size, dtype=torch.float32).cuda() * (i + 1)
+        else:
+            raise RuntimeError("Unsupported tensor backend.")
+        ray.get([a.set_buffer.remote(t)])
+    if tensor_backend == "cupy":
+        list_buffer = [cp.ones(array_size, dtype=dtype) for _ in range(world_size)]
+    elif tensor_backend == "torch":
+        list_buffer = [
+            torch.ones(array_size, dtype=torch.float32).cuda()
+            for _ in range(world_size)
+        ]
+    else:
+        raise RuntimeError("Unsupported tensor backend.")
+    ray.get([a.set_list_buffer.remote(list_buffer) for a in actors])
+
+
+@ray.remote(num_gpus=2)
+class MultiGPUWorker:
+    def __init__(self):
+        self.buffer0 = None
+        self.buffer1 = None
+        self.list_buffer0 = None
+        self.list_buffer1 = None
+
+    def __del__(self):
+        self.buffer0 = None
+        self.buffer1 = None
+        self.list_buffer0 = None
+        self.list_buffer1 = None
+
+    def init_tensors(self):
+        with cp.cuda.Device(0):
+            self.buffer0 = cp.ones((10,), dtype=cp.float32)
+            self.list_buffer0 = [cp.ones((10,), dtype=cp.float32) for _ in range(4)]
+        with cp.cuda.Device(1):
+            self.buffer1 = cp.ones((10,), dtype=cp.float32)
+            self.list_buffer1 = [cp.ones((10,), dtype=cp.float32) for _ in range(4)]
+        cp.cuda.Stream.null.synchronize()
+        return True
+
+    def init_group(self, world_size, rank, backend=Backend.NCCL, group_name="default"):
+        col.init_collective_group(world_size, rank, backend, group_name)
+        return True
+
+    def set_buffer(
+        self,
+        size,
+        value0=1.0,
+        value1=1.0,
+        dtype=cp.float32,
+        tensor_type0="cupy",
+        tensor_type1="cupy",
+    ):
+        if tensor_type0 == "cupy":
+            with cp.cuda.Device(0):
+                self.buffer0 = cp.ones(size, dtype=dtype) * value0
+        elif tensor_type0 == "torch":
+            self.buffer0 = torch.ones(size, dtype=torch.float32).cuda(0) * value0
+        else:
+            raise RuntimeError()
+
+        if tensor_type1 == "cupy":
+            with cp.cuda.Device(1):
+                self.buffer1 = cp.ones(size, dtype=dtype) * value1
+        elif tensor_type1 == "torch":
+            self.buffer1 = torch.ones(size, dtype=torch.float32).cuda(1) * value1
+        else:
+            raise RuntimeError()
+        cp.cuda.Device(0).synchronize()
+        cp.cuda.Device(1).synchronize()
+        # cp.cuda.Stream.null.synchronize()
+        return True
+
+    def set_list_buffer(
+        self,
+        size,
+        value0=1.0,
+        value1=1.0,
+        dtype=cp.float32,
+        tensor_type0="cupy",
+        tensor_type1="cupy",
+    ):
+        if tensor_type0 == "cupy":
+            with cp.cuda.Device(0):
+                self.list_buffer0 = [
+                    cp.ones(size, dtype=dtype) * value0 for _ in range(4)
+                ]
+        elif tensor_type0 == "torch":
+            self.list_buffer0 = [
+                torch.ones(size, dtype=torch.float32).cuda(0) * value0 for _ in range(4)
+            ]
+        else:
+            raise RuntimeError()
+
+        if tensor_type1 == "cupy":
+            with cp.cuda.Device(1):
+                self.list_buffer1 = [
+                    cp.ones(size, dtype=dtype) * value1 for _ in range(4)
+                ]
+        elif tensor_type1 == "torch":
+            self.list_buffer1 = [
+                torch.ones(size, dtype=torch.float32).cuda(1) * value1 for _ in range(4)
+            ]
+        else:
+            raise RuntimeError()
+        cp.cuda.Device(0).synchronize()
+        cp.cuda.Device(1).synchronize()
+        return True
+
+    @ray.method(num_returns=2)
+    def get_buffer(self):
+        return self.buffer0, self.buffer1
+
+    def do_allreduce_multigpu(self, group_name="default", op=ReduceOp.SUM):
+        col.allreduce_multigpu([self.buffer0, self.buffer1], group_name, op)
+        cp.cuda.Device(0).synchronize()
+        cp.cuda.Device(1).synchronize()
+        return self.buffer0
+
+    def do_reduce_multigpu(
+        self, group_name="default", dst_rank=0, dst_gpu_index=0, op=ReduceOp.SUM
+    ):
+        col.reduce_multigpu(
+            [self.buffer0, self.buffer1], dst_rank, dst_gpu_index, group_name, op
+        )
+        cp.cuda.Device(0).synchronize()
+        cp.cuda.Device(1).synchronize()
+        return self.buffer0, self.buffer1
+
+    def do_broadcast_multigpu(self, group_name="default", src_rank=0, src_gpu_index=0):
+        col.broadcast_multigpu(
+            [self.buffer0, self.buffer1], src_rank, src_gpu_index, group_name
+        )
+        return self.buffer0, self.buffer1
+
+    def do_allgather_multigpu(self, group_name="default"):
+        col.allgather_multigpu(
+            [self.list_buffer0, self.list_buffer1],
+            [self.buffer0, self.buffer1],
+            group_name,
+        )
+        cp.cuda.Device(0).synchronize()
+        cp.cuda.Device(1).synchronize()
+        return self.list_buffer0, self.list_buffer1
+
+    def do_reducescatter_multigpu(self, group_name="default", op=ReduceOp.SUM):
+        col.reducescatter_multigpu(
+            [self.buffer0, self.buffer1],
+            [self.list_buffer0, self.list_buffer1],
+            group_name,
+            op,
+        )
+        cp.cuda.Device(0).synchronize()
+        cp.cuda.Device(1).synchronize()
+        return self.buffer0, self.buffer1
+
+    def do_send_multigpu(
+        self, group_name="default", dst_rank=0, dst_gpu_index=0, src_gpu_index=0
+    ):
+        if src_gpu_index == 0:
+            col.send_multigpu(self.buffer0, dst_rank, dst_gpu_index, group_name)
+            cp.cuda.Device(0).synchronize()
+            return self.buffer0
+        elif src_gpu_index == 1:
+            col.send_multigpu(self.buffer1, dst_rank, dst_gpu_index, group_name)
+            cp.cuda.Device(1).synchronize()
+            return self.buffer1
+        else:
+            raise RuntimeError()
+
+    def do_recv_multigpu(
+        self, group_name="default", src_rank=0, src_gpu_index=0, dst_gpu_index=0
+    ):
+        if dst_gpu_index == 0:
+            col.recv_multigpu(self.buffer0, src_rank, src_gpu_index, group_name)
+            cp.cuda.Device(0).synchronize()
+            return self.buffer0
+        elif dst_gpu_index == 1:
+            col.recv_multigpu(self.buffer1, src_rank, src_gpu_index, group_name)
+            cp.cuda.Device(1).synchronize()
+            return self.buffer1
+        else:
+            raise RuntimeError()
+
+    def destroy_group(self, group_name="default"):
+        col.destroy_collective_group(group_name)
+        return True
+
+    def report_rank(self, group_name="default"):
+        rank = col.get_rank(group_name)
+        return rank
+
+    def report_world_size(self, group_name="default"):
+        ws = col.get_collective_group_size(group_name)
+        return ws
+
+    def report_nccl_availability(self):
+        avail = col.nccl_available()
+        return avail
+
+    def report_gloo_availability(self):
+        avail = col.gloo_available()
+        return avail
+
+    def report_is_group_initialized(self, group_name="default"):
+        is_init = col.is_group_initialized(group_name)
+        return is_init
+
+    def report_num_gpus(self):
+        n_gpus = get_num_gpus()
+        return n_gpus
+
+
+def create_collective_multigpu_workers(
+    num_workers=2, group_name="default", backend="nccl"
+):
+    actors = [None] * num_workers
+    for i in range(num_workers):
+        actor = MultiGPUWorker.remote()
+        ray.get([actor.set_buffer.remote([10])], timeout=10)
+        ray.get([actor.set_list_buffer.remote([10])], timeout=10)
+        actors[i] = actor
+    world_size = num_workers
+    init_results = ray.get(
+        [
+            actor.init_group.remote(world_size, i, backend, group_name)
+            for i, actor in enumerate(actors)
+        ]
+    )
+    return actors, init_results
+
+
+def init_tensors_for_gather_scatter_multigpu(
+    actors, array_size=10, tensor_backend="cupy"
+):
+    for i, a in enumerate(actors):
+        if tensor_backend == "cupy":
+            ray.get([a.set_buffer.remote(array_size)])
+            ray.get([a.set_list_buffer.remote(array_size)])
+        elif tensor_backend == "torch":
+            ray.get(
+                [
+                    a.set_buffer.remote(
+                        array_size, tensor_type0="torch", tensor_type1="torch"
+                    )
+                ]
+            )
+            ray.get(
+                [
+                    a.set_list_buffer.remote(
+                        array_size, tensor_type0="torch", tensor_type1="torch"
+                    )
+                ]
+            )
+        else:
+            raise RuntimeError("Unsupported tensor backend.")