wehub-resource-sync
210 lines
7.1 KiB
Python
210 lines
7.1 KiB
Python
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)
|