chore: import upstream snapshot with attribution

python/ray/experimental/collective/conftest.py

@@ -0,0 +1,243 @@
+import uuid
+from typing import Dict, FrozenSet, List, Optional, Set, Tuple, Type
+
+import torch
+
+import ray
+from ray.experimental.channel.common import ChannelContext
+from ray.experimental.channel.communicator import (
+    Communicator,
+    ReduceOp,
+    TorchTensorAllocator,
+)
+
+
+class AbstractNcclGroup(Communicator):
+    """
+    A dummy NCCL group for testing.
+    """
+
+    def __init__(self, actor_handles: List[ray.actor.ActorHandle]):
+        self._actor_handles = actor_handles
+        self._rank = None
+
+    def initialize(self, rank: int) -> None:
+        self._rank = rank
+
+    def get_rank(self, actor: ray.actor.ActorHandle) -> int:
+        return self._actor_handles.index(actor)
+
+    def get_world_size(self) -> int:
+        return len(self._actor_handles)
+
+    def get_self_rank(self) -> Optional[int]:
+        return self._rank
+
+    def get_actor_handles(self) -> List["ray.actor.ActorHandle"]:
+        return self._actor_handles
+
+    def send(self, value: "torch.Tensor", peer_rank: int) -> None:
+        raise NotImplementedError
+
+    def recv(
+        self,
+        shape: Tuple[int],
+        dtype: "torch.dtype",
+        peer_rank: int,
+        allocator: Optional[TorchTensorAllocator] = None,
+    ) -> "torch.Tensor":
+        raise NotImplementedError
+
+    def allgather(
+        self,
+        send_buf: "torch.Tensor",
+        recv_buf: "torch.Tensor",
+    ) -> None:
+        raise NotImplementedError
+
+    def allreduce(
+        self,
+        send_buf: "torch.Tensor",
+        recv_buf: "torch.Tensor",
+        op: ReduceOp = ReduceOp.SUM,
+    ) -> None:
+        raise NotImplementedError
+
+    def reducescatter(
+        self,
+        send_buf: "torch.Tensor",
+        recv_buf: "torch.Tensor",
+        op: ReduceOp = ReduceOp.SUM,
+    ) -> None:
+        raise NotImplementedError
+
+    @property
+    def recv_stream(self):
+        return None
+
+    @property
+    def send_stream(self):
+        return None
+
+    def destroy(self) -> None:
+        pass
+
+    def get_transport_name(self) -> str:
+        return "accelerator"
+
+    @classmethod
+    def generate_communicator_id(cls) -> str:
+        pass
+
+
+class MockNcclGroupSet:
+    def __init__(self):
+        # Represents a mapping from a NCCL group ID to a set of actors and a custom
+        # NCCL group.
+        self.ids_to_actors_and_custom_comms: Dict[
+            str, Tuple[FrozenSet["ray.actor.ActorHandle"], Optional[Communicator]]
+        ] = {}
+
+    def __call__(
+        self,
+        actors: List["ray.actor.ActorHandle"],
+        custom_nccl_group: Optional[Communicator] = None,
+        use_communication_streams: bool = False,
+        accelerator_module_name: Optional[str] = None,
+        accelerator_communicator_cls: Optional[Type[Communicator]] = None,
+    ) -> str:
+        group_id = str(uuid.uuid4())
+        self.ids_to_actors_and_custom_comms[group_id] = (
+            frozenset(actors),
+            custom_nccl_group,
+        )
+
+        if custom_nccl_group is None:
+            ranks = list(range(len(actors)))
+        else:
+            ranks = [custom_nccl_group.get_rank(actor) for actor in actors]
+        init_tasks = [
+            actor.__ray_call__.remote(
+                mock_do_init_nccl_group,
+                group_id,
+                rank,
+                actors,
+                custom_nccl_group,
+            )
+            for rank, actor in zip(ranks, actors)
+        ]
+        ray.get(init_tasks, timeout=30)
+
+        ctx = ChannelContext.get_current()
+        if custom_nccl_group is not None:
+            ctx.communicators[group_id] = custom_nccl_group
+        else:
+            ctx.communicators[group_id] = AbstractNcclGroup(actors)
+
+        return group_id
+
+    def mock_destroy_nccl_group(self, group_id: str) -> None:
+        ctx = ChannelContext.get_current()
+        if group_id not in ctx.communicators:
+            return
+
+        actors, _ = self.ids_to_actors_and_custom_comms[group_id]
+        destroy_tasks = [
+            actor.__ray_call__.remote(
+                mock_do_destroy_nccl_group,
+                group_id,
+            )
+            for actor in actors
+        ]
+        ray.wait(destroy_tasks, timeout=30)
+
+        if group_id in self.ids_to_actors_and_custom_comms:
+            del self.ids_to_actors_and_custom_comms[group_id]
+        ctx.communicators[group_id].destroy()
+        del ctx.communicators[group_id]
+
+    def check_teardown(self, nccl_group_ids: List[str]) -> None:
+        ctx = ChannelContext.get_current()
+        for nccl_group_id in nccl_group_ids:
+            assert nccl_group_id not in self.ids_to_actors_and_custom_comms
+            assert nccl_group_id not in ctx.communicators
+
+
+@ray.remote
+class CPUTorchTensorWorker:
+    def __init__(self):
+        self.device = "cpu"
+
+    def return_tensor(
+        self, size: int, dtype: Optional[torch.dtype] = None
+    ) -> torch.Tensor:
+        return torch.ones(size, dtype=dtype, device=self.device)
+
+    def recv(self, tensor: torch.Tensor) -> Tuple[int, int]:
+        assert tensor.device == self.device
+        return tensor.shape, tensor[0]
+
+    def recv_tensors(self, *tensors) -> Tuple[torch.Tensor, ...]:
+        return tuple(tensors)
+
+
+def mock_do_init_nccl_group(
+    self,
+    group_id: str,
+    rank: int,
+    actors: List[ray.actor.ActorHandle],
+    custom_nccl_group: Optional[Communicator],
+) -> None:
+    ctx = ChannelContext.get_current()
+    if custom_nccl_group is None:
+        nccl_group = AbstractNcclGroup(actors)
+        nccl_group.initialize(rank)
+        ctx.communicators[group_id] = nccl_group
+    else:
+        custom_nccl_group.initialize(rank)
+        ctx.communicators[group_id] = custom_nccl_group
+
+
+def mock_do_destroy_nccl_group(self, group_id: str) -> None:
+    ctx = ChannelContext.get_current()
+    if group_id not in ctx.communicators:
+        return
+    ctx.communicators[group_id].destroy()
+    del ctx.communicators[group_id]
+
+
+def check_nccl_group_init(
+    monkeypatch,
+    dag: "ray.dag.DAGNode",
+    actors_and_custom_comms: Set[
+        Tuple[FrozenSet["ray.actor.ActorHandle"], Optional[Communicator]]
+    ],
+) -> "ray.dag.CompiledDAG":
+    mock_nccl_group_set = MockNcclGroupSet()
+    monkeypatch.setattr(
+        "ray.dag.compiled_dag_node._init_communicator",
+        mock_nccl_group_set,
+    )
+
+    compiled_dag = dag.experimental_compile()
+    assert (
+        set(mock_nccl_group_set.ids_to_actors_and_custom_comms.values())
+        == actors_and_custom_comms
+    )
+
+    return compiled_dag, mock_nccl_group_set
+
+
+def check_nccl_group_teardown(
+    monkeypatch,
+    compiled_dag: "ray.dag.CompiledDAG",
+    mock_nccl_group_set: MockNcclGroupSet,
+):
+    monkeypatch.setattr(
+        "ray.dag.compiled_dag_node._destroy_communicator",
+        mock_nccl_group_set.mock_destroy_nccl_group,
+    )
+
+    created_communicator_ids = compiled_dag._actors_to_created_communicator_id.values()
+    compiled_dag.teardown()
+    mock_nccl_group_set.check_teardown(created_communicator_ids)