wehub-resource-sync
86 lines
2.1 KiB
Python
86 lines
2.1 KiB
Python
"""Some utility class for Collectives."""
|
|
import asyncio
|
|
import logging
|
|
|
|
import ray
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@ray.remote
|
|
class NCCLUniqueIDStore:
|
|
"""NCCLUniqueID Store as a named actor class.
|
|
|
|
Args:
|
|
name: the unique name for this named actor.
|
|
|
|
Attributes:
|
|
name: the unique name for this named actor.
|
|
nccl_id: the NCCLUniqueID held in this store.
|
|
"""
|
|
|
|
def __init__(self, name: str):
|
|
self.name = name
|
|
self.nccl_id = None
|
|
self.event = asyncio.Event()
|
|
|
|
async def set_id(self, uid: bytes):
|
|
"""
|
|
Initialize the NCCL unique ID for this store.
|
|
|
|
Args:
|
|
uid: the unique ID generated via the NCCL generate_communicator_id API.
|
|
|
|
Returns:
|
|
The NCCL unique ID set.
|
|
"""
|
|
self.nccl_id = uid
|
|
self.event.set()
|
|
return uid
|
|
|
|
async def wait_and_get_id(self):
|
|
"""Wait for the NCCL unique ID to be set and return it."""
|
|
await self.event.wait()
|
|
return self.nccl_id
|
|
|
|
def get_id(self):
|
|
"""Get the NCCL unique ID held in this store."""
|
|
if not self.nccl_id:
|
|
logger.warning(
|
|
"The NCCL ID has not been set yet for store {}.".format(self.name)
|
|
)
|
|
return self.nccl_id
|
|
|
|
|
|
@ray.remote
|
|
class Info:
|
|
"""Store the group information created via `create_collective_group`.
|
|
|
|
Note: Should be used as a NamedActor.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.ids = None
|
|
self.world_size = -1
|
|
self.rank = -1
|
|
self.backend = None
|
|
self.gloo_timeout = 30000
|
|
|
|
def set_info(self, ids, world_size, rank, backend, gloo_timeout):
|
|
"""Store collective information."""
|
|
self.ids = ids
|
|
self.world_size = world_size
|
|
self.rank = rank
|
|
self.backend = backend
|
|
self.gloo_timeout = gloo_timeout
|
|
|
|
def get_info(self):
|
|
"""Get previously stored collective information."""
|
|
return (
|
|
self.ids,
|
|
self.world_size,
|
|
self.rank,
|
|
self.backend,
|
|
self.gloo_timeout,
|
|
)
|