register_collective_backend

ray.util.collective.backend_registry.register_collective_backend(name: str, group_cls: Type[BaseGroup])

Register a custom collective backend with Ray.

This function registers a custom backend class that can be used for collective operations. The backend must be a subclass of BaseGroup and implement all required collective operations.

Important: The backend must be registered on both the driver and all actors before creating collective groups. This is because each process (driver and each actor) needs to know about your backend class to instantiate it.

Parameters:

• name (str) – The name of the backend (e.g., “MY_BACKEND”). This will be automatically added to the Backend enum as Backend.MY_BACKEND.

• group_cls (Type[BaseGroup]) – The backend class, which must be a subclass of BaseGroup.

Example

>>> import ray
>>> from ray.util.collective import create_collective_group, init_collective_group
>>> from ray.util.collective.backend_registry import register_collective_backend
>>> from ray.util.collective.collective_group.base_collective_group import BaseGroup
>>>
>>> class MyCustomBackend(BaseGroup):
...     def __init__(self, world_size, rank, group_name):
...         super().__init__(world_size, rank, group_name)
...     @classmethod
...     def backend(cls):
...         return "MY_BACKEND"
...     @classmethod
...     def check_backend_availability(cls) -> bool:
...         return True
...     def allreduce(self, tensor, allreduce_options=None):
...         pass
...     def broadcast(self, tensor, broadcast_options=None):
...         pass
...     def barrier(self, barrier_options=None):
...         pass
>>>
>>> # Register on the driver
>>> register_collective_backend("MY_BACKEND", MyCustomBackend)
>>>
>>> ray.init()
>>>
>>> @ray.remote
... class Worker:
...     def __init__(self, rank):
...         self.rank = rank
...     def setup(self, world_size):
...         # IMPORTANT: Register on each worker too
...         register_collective_backend("MY_BACKEND", MyCustomBackend)
...         init_collective_group(
...             world_size=world_size,
...             rank=self.rank,
...             backend="MY_BACKEND",
...             group_name="default",
...         )
>>>
>>> actors = [Worker.remote(rank=i) for i in range(2)]
>>> create_collective_group(
...     actors=actors,
...     world_size=2,
...     ranks=[0, 1],
...     backend="MY_BACKEND",
...     group_name="default",
... )
>>> ray.get([a.setup.remote(2) for a in actors])

PublicAPI (alpha): This API is in alpha and may change before becoming stable.