chore: import upstream snapshot with attribution

python/ray/experimental/collective/collective.py

@@ -0,0 +1,209 @@
+import threading
+import uuid
+from typing import Dict, List, Optional, Union
+
+import ray
+import ray.experimental.internal_kv as internal_kv
+from ray.experimental.collective.communicator import CommunicatorHandle
+from ray.util.annotations import PublicAPI
+from ray.util.collective.collective_group.torch_gloo_collective_group import (
+    get_master_address_metadata_key,
+)
+from ray.util.collective.types import Backend
+
+_remote_communicator_manager: "Optional[RemoteCommunicatorManager]" = None
+_remote_communicator_manager_lock = threading.Lock()
+
+
+class RemoteCommunicatorManager:
+    """Singleton class to store the mapping between actors and communicators
+    that the actors are a part of.
+    """
+
+    def __init__(self):
+        # Handles to communicators that we created. Key is a user-provided
+        # name or UUID.
+        self._remote_communicators: Dict[str, CommunicatorHandle] = {}
+
+    @staticmethod
+    def get() -> "RemoteCommunicatorManager":
+        global _remote_communicator_manager
+        with _remote_communicator_manager_lock:
+            if _remote_communicator_manager is None:
+                _remote_communicator_manager = RemoteCommunicatorManager()
+            return _remote_communicator_manager
+
+    def add_remote_communicator(self, comm_handle: CommunicatorHandle):
+        self._remote_communicators[comm_handle.name] = comm_handle
+
+    def remove_remote_communicator(self, name: str):
+        return self._remote_communicators.pop(name, None)
+
+    def get_collective_groups(
+        self,
+        actors: Optional[List[ray.actor.ActorHandle]] = None,
+        backend: Optional[Backend] = None,
+    ):
+        """
+        Get the collective groups that the given actors are a subset of. Filter by
+        backend if provided.
+        """
+        actors = actors or []
+        actors = set(actors)
+
+        collectives = []
+        # Find all collective groups that the given actors are a subset
+        # of, with the matching backend if provided.
+        for collective in self._remote_communicators.values():
+            if actors.issubset(set(collective.actors)):
+                if backend is None or collective.backend == backend:
+                    collectives.append(collective)
+        return collectives
+
+
+@PublicAPI(stability="alpha")
+def get_collective_groups(
+    actors: List[ray.actor.ActorHandle], backend: Optional[str] = None
+) -> List[CommunicatorHandle]:
+    """
+    Get the collective groups that the given actors are a subset of. Filter by
+    backend if provided.
+
+    Args:
+        actors: List of actors. Return handles to all collective groups that
+            these actors are a subset of.
+        backend: An optional backend to filter by. See
+            ray.util.collective.types.Backend for valid backends.
+
+    Returns:
+        A list of communicator handles that the actors are a subset of.
+    """
+    manager = RemoteCommunicatorManager.get()
+    backend = Backend(backend) if backend is not None else None
+    return manager.get_collective_groups(actors, backend)
+
+
+@PublicAPI(stability="alpha")
+def create_collective_group(
+    actors: List[ray.actor.ActorHandle],
+    backend: str,
+    name: Optional[str] = None,
+) -> CommunicatorHandle:
+    """Create a collective group on the given list of actors. If this function
+    returns successfully, then the collective group has been initialized on all
+    actors, using the given order of actors as the ranks.
+
+    Currently, an actor can only participate in one collective group per
+    backend at a time. To reuse an actor, destroy its collective group and
+    create a new one.
+
+    Args:
+        actors: The actors to participate in the collective group.
+        backend: The backend to use. See ray.util.collective.types.Backend for
+            valid backends.
+        name: A name to use for the collective group. If None is provided, a
+            random name will be generated.
+
+    Returns:
+        Handle to the communicator.
+    """
+    manager = RemoteCommunicatorManager.get()
+
+    if name is None:
+        name = str(uuid.uuid4())
+
+    # Validate the backend.
+    backend = Backend(backend)
+
+    world_size = len(actors)
+
+    for actor in actors:
+        if manager.get_collective_groups([actor], backend):
+            raise RuntimeError(
+                f"Actor {actor} already in group for backend {backend}. Actors can currently only participate in at most one group per backend."
+            )
+
+    actor_ids = [actor._ray_actor_id for actor in actors]
+    if len(set(actor_ids)) != len(actor_ids):
+        raise ValueError(f"All actors must be unique, got: {actors}")
+
+    metadata_key = None
+    if backend == Backend.GLOO:
+        metadata_key = get_master_address_metadata_key(name)
+
+    def _do_init_collective_group(self, rank: int):
+        ray.util.collective.init_collective_group(
+            world_size, rank, backend, group_name=name
+        )
+
+    try:
+        init_tasks = [
+            actor.__ray_call__.remote(
+                _do_init_collective_group,
+                rank,
+            )
+            for rank, actor in enumerate(actors)
+        ]
+        ray.get(init_tasks)
+    finally:
+        # Clean up the metadata once collective group is initialized
+        # (or failed to initialize).
+        if metadata_key is not None:
+            internal_kv._internal_kv_del(metadata_key)
+
+    # Group was successfully created.
+    comm = CommunicatorHandle(actors, name, backend)
+    manager.add_remote_communicator(comm)
+    return comm
+
+
+@PublicAPI(stability="alpha")
+def destroy_collective_group(group_or_name: Union[CommunicatorHandle, str]):
+    """
+    Destroy a collective group. If this functions returns successfully, then
+    the actors that were in the collective can be reused to create a new
+    collective group.
+
+    Args:
+        group_or_name: Either a communicator handle or the name of the group to
+            destroy.
+    """
+    if isinstance(group_or_name, CommunicatorHandle):
+        name = group_or_name.name
+    elif isinstance(group_or_name, str):
+        name = group_or_name
+    else:
+        raise ValueError("Expected CommunicatorHandle or str (group name).")
+
+    manager = RemoteCommunicatorManager.get()
+    group = manager.remove_remote_communicator(name)
+    if group is not None:
+
+        def _do_destroy_collective_group(self):
+            ray.util.collective.destroy_collective_group(name)
+
+        destroy_tasks = [
+            actor.__ray_call__.options(concurrency_group="_ray_system").remote(
+                _do_destroy_collective_group
+            )
+            for actor in group.actors
+        ]
+        try:
+            ray.get(destroy_tasks)
+        except ray.exceptions.ActorDiedError:
+            pass
+    else:
+        raise ValueError(f"No group with name {name} found.")
+
+
+@PublicAPI(stability="alpha")
+def destroy_all_collective_groups():
+    """
+    Destroy all collective groups. This will destroy all collective groups that
+    were previously created by this process. After this function returns, the
+    actors participating in those collective groups can be reused to create a
+    new collective group.
+    """
+    manager = RemoteCommunicatorManager.get()
+    for collective in manager.get_collective_groups():
+        destroy_collective_group(collective.name)