wehub-resource-sync
817 lines
26 KiB
Python
817 lines
26 KiB
Python
"""APIs exposed under the namespace ray.util.collective."""
|
|
|
|
import logging
|
|
import os
|
|
import socket
|
|
import threading
|
|
from typing import Any, List, Tuple
|
|
|
|
import numpy as np
|
|
|
|
import ray
|
|
from . import types
|
|
from ray._common.network_utils import find_free_port, is_ipv6
|
|
from ray.util.collective.backend_registry import (
|
|
_global_registry,
|
|
register_collective_backend,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
try:
|
|
from ray.util.collective.collective_group.nccl_collective_group import NCCLGroup
|
|
|
|
register_collective_backend("NCCL", NCCLGroup)
|
|
except ImportError:
|
|
pass
|
|
|
|
try:
|
|
from ray.util.collective.collective_group.torch_gloo_collective_group import (
|
|
TorchGLOOGroup,
|
|
)
|
|
|
|
register_collective_backend("GLOO", TorchGLOOGroup)
|
|
except ImportError:
|
|
pass
|
|
|
|
|
|
def nccl_available():
|
|
return is_backend_available("NCCL")
|
|
|
|
|
|
def gloo_available():
|
|
return is_backend_available("GLOO")
|
|
|
|
|
|
def is_backend_available(backend: str) -> bool:
|
|
"""Check if a collective backend is available.
|
|
|
|
Args:
|
|
backend: The name of the backend to check (e.g., "NCCL", "GLOO").
|
|
|
|
Returns:
|
|
True if the backend is available, False otherwise.
|
|
"""
|
|
return _global_registry.check(backend)
|
|
|
|
|
|
def get_address_and_port() -> Tuple[str, int]:
|
|
"""Returns the IP address and a free port on this node."""
|
|
addr = ray.util.get_node_ip_address()
|
|
port = find_free_port(socket.AF_INET6 if is_ipv6(addr) else socket.AF_INET)
|
|
return addr, port
|
|
|
|
|
|
class GroupManager(object):
|
|
"""Use this class to manage the collective groups we created so far.
|
|
|
|
Each process will have an instance of `GroupManager`. Each process
|
|
could belong to multiple collective groups. The membership information
|
|
and other metadata are stored in the global `_group_mgr` object.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._name_group_map = {}
|
|
self._registry = _global_registry
|
|
|
|
def create_collective_group(
|
|
self, backend, world_size, rank, group_name, gloo_timeout=None
|
|
):
|
|
"""The entry to create new collective groups in the manager.
|
|
|
|
Put the registration and the group information into the manager
|
|
metadata as well.
|
|
"""
|
|
backend = backend.upper()
|
|
backend_cls = self._registry.get(backend)
|
|
|
|
if not backend_cls.check_backend_availability():
|
|
raise RuntimeError(
|
|
f"Backend {backend} is not available. Please check the installation."
|
|
)
|
|
|
|
if backend == "GLOO":
|
|
logger.debug(
|
|
"Creating torch.distributed GLOO group: '{}'...".format(group_name)
|
|
)
|
|
g = backend_cls(world_size, rank, group_name, gloo_timeout)
|
|
else:
|
|
g = backend_cls(world_size, rank, group_name)
|
|
|
|
self._name_group_map[group_name] = g
|
|
return self._name_group_map[group_name]
|
|
|
|
def is_group_exist(self, group_name):
|
|
return group_name in self._name_group_map
|
|
|
|
def get_group_by_name(self, group_name):
|
|
"""Get the collective group handle by its name."""
|
|
if not self.is_group_exist(group_name):
|
|
logger.warning("The group '{}' is not initialized.".format(group_name))
|
|
return None
|
|
return self._name_group_map[group_name]
|
|
|
|
def destroy_collective_group(self, group_name):
|
|
"""Group destructor."""
|
|
if not self.is_group_exist(group_name):
|
|
logger.warning("The group '{}' does not exist.".format(group_name))
|
|
return
|
|
|
|
# release the collective group resource
|
|
g = self._name_group_map[group_name]
|
|
# clean up the dicts
|
|
del self._name_group_map[group_name]
|
|
# Release the communicator resources
|
|
g.destroy_group()
|
|
|
|
# Release the detached actors spawned by `create_collective_group()`
|
|
name = "info_" + group_name
|
|
try:
|
|
store = ray.get_actor(name)
|
|
ray.kill(store)
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
_group_mgr = GroupManager()
|
|
# This lock is used to make external calls to the _group_mgr thread-safe.
|
|
_group_mgr_lock = threading.Lock()
|
|
|
|
|
|
def is_group_initialized(group_name):
|
|
"""Check if the group is initialized in this process by the group name."""
|
|
global _group_mgr
|
|
global _group_mgr_lock
|
|
with _group_mgr_lock:
|
|
return _group_mgr.is_group_exist(group_name)
|
|
|
|
|
|
def init_collective_group(
|
|
world_size: int,
|
|
rank: int,
|
|
backend: types.Backend = types.Backend.NCCL,
|
|
group_name: str = "default",
|
|
gloo_timeout: int = 30000,
|
|
):
|
|
"""Initialize a collective group inside an actor process.
|
|
|
|
Args:
|
|
world_size: the total number of processes in the group.
|
|
rank: the rank of the current process.
|
|
backend: the CCL backend to use, NCCL or GLOO.
|
|
group_name: the name of the collective group.
|
|
gloo_timeout: timeout in milliseconds for GLOO operations.
|
|
"""
|
|
_check_inside_actor()
|
|
|
|
global _group_mgr
|
|
global _group_mgr_lock
|
|
|
|
# TODO(Hao): implement a group auto-counter.
|
|
if not group_name:
|
|
raise ValueError("group_name '{}' needs to be a string.".format(group_name))
|
|
|
|
with _group_mgr_lock:
|
|
if _group_mgr.is_group_exist(group_name):
|
|
raise RuntimeError("Trying to initialize a group a second time.")
|
|
|
|
assert world_size > 0
|
|
assert rank >= 0
|
|
assert rank < world_size
|
|
_group_mgr.create_collective_group(
|
|
backend, world_size, rank, group_name, gloo_timeout
|
|
)
|
|
|
|
|
|
def create_collective_group(
|
|
actors: List[Any],
|
|
world_size: int,
|
|
ranks: List[int],
|
|
backend: types.Backend = types.Backend.NCCL,
|
|
group_name: str = "default",
|
|
gloo_timeout: int = 30000,
|
|
):
|
|
"""Declare a list of actors as a collective group.
|
|
|
|
Note: This function should be called in a driver process.
|
|
|
|
Args:
|
|
actors: a list of actors to be set in a collective group.
|
|
world_size: the total number of processes in the group.
|
|
ranks: the rank of each actor.
|
|
backend: the CCL backend to use, NCCL or GLOO.
|
|
group_name: the name of the collective group.
|
|
gloo_timeout: timeout in milliseconds for GLOO operations.
|
|
"""
|
|
|
|
name = "info_" + group_name
|
|
try:
|
|
ray.get_actor(name)
|
|
raise RuntimeError("Trying to initialize a group twice.")
|
|
except ValueError:
|
|
pass
|
|
|
|
if len(ranks) != len(actors):
|
|
raise RuntimeError(
|
|
"Each actor should correspond to one rank. Got '{}' "
|
|
"ranks but '{}' actors".format(len(ranks), len(actors))
|
|
)
|
|
|
|
if set(ranks) != set(range(len(ranks))):
|
|
raise RuntimeError(
|
|
"Ranks must be a permutation from 0 to '{}'. Got '{}'.".format(
|
|
len(ranks), "".join([str(r) for r in ranks])
|
|
)
|
|
)
|
|
|
|
if world_size <= 0:
|
|
raise RuntimeError(
|
|
"World size must be greater than zero. Got '{}'.".format(world_size)
|
|
)
|
|
if not all(ranks) >= 0:
|
|
raise RuntimeError("Ranks must be non-negative.")
|
|
if not all(ranks) < world_size:
|
|
raise RuntimeError("Ranks cannot be greater than world_size.")
|
|
|
|
# Check if backend is registered and available
|
|
backend_upper = backend.upper()
|
|
if not _global_registry.is_registered(backend_upper):
|
|
raise RuntimeError(
|
|
f"Backend {backend_upper} is not registered. "
|
|
f"Please register it using register_collective_backend('{backend_upper}', YourBackendClass)."
|
|
)
|
|
if not _global_registry.check(backend_upper):
|
|
raise RuntimeError(
|
|
f"Backend {backend_upper} is registered but not available. "
|
|
f"Please check the installation requirements for this backend."
|
|
)
|
|
|
|
# avoid a circular dependency
|
|
from ray.util.collective.util import Info
|
|
|
|
# store the information into a NamedActor that can be accessed later.
|
|
name = "info_" + group_name
|
|
actors_id = [a._ray_actor_id for a in actors]
|
|
# TODO (Dacheng): how do we recycle this name actor?
|
|
info = Info.options(name=name, lifetime="detached").remote()
|
|
ray.get([info.set_info.remote(actors_id, world_size, ranks, backend, gloo_timeout)])
|
|
|
|
|
|
# TODO (we need a declarative destroy() API here.)
|
|
def destroy_collective_group(group_name: str = "default") -> None:
|
|
"""Destroy a collective group given its group name."""
|
|
_check_inside_actor()
|
|
global _group_mgr
|
|
global _group_mgr_lock
|
|
with _group_mgr_lock:
|
|
_group_mgr.destroy_collective_group(group_name)
|
|
|
|
|
|
def get_rank(group_name: str = "default") -> int:
|
|
"""Return the rank of this process in the given group.
|
|
|
|
Args:
|
|
group_name: the name of the group to query
|
|
|
|
Returns:
|
|
the rank of this process in the named group,
|
|
-1 if the group does not exist or the process does
|
|
not belong to the group.
|
|
"""
|
|
_check_inside_actor()
|
|
|
|
global _group_mgr
|
|
global _group_mgr_lock
|
|
with _group_mgr_lock:
|
|
if not _group_mgr.is_group_exist(group_name):
|
|
return -1
|
|
g = _group_mgr.get_group_by_name(group_name)
|
|
return g.rank
|
|
|
|
|
|
def get_collective_group_size(group_name: str = "default") -> int:
|
|
"""Return the size of the collective group with the given name.
|
|
|
|
Args:
|
|
group_name: the name of the group to query
|
|
|
|
Returns:
|
|
The world size of the collective group, -1 if the group does
|
|
not exist or the process does not belong to the group.
|
|
"""
|
|
_check_inside_actor()
|
|
global _group_mgr
|
|
global _group_mgr_lock
|
|
with _group_mgr_lock:
|
|
if not _group_mgr.is_group_exist(group_name):
|
|
return -1
|
|
g = _group_mgr.get_group_by_name(group_name)
|
|
return g.world_size
|
|
|
|
|
|
def allreduce(
|
|
tensor: Any,
|
|
group_name: str = "default",
|
|
op: types.ReduceOp = types.ReduceOp.SUM,
|
|
):
|
|
"""Collective allreduce the tensor across the group.
|
|
|
|
Args:
|
|
tensor: the tensor to be all-reduced on this process.
|
|
group_name: the collective group name to perform allreduce.
|
|
op: The reduce operation.
|
|
"""
|
|
_check_single_tensor_input(tensor)
|
|
g = get_group_handle(group_name)
|
|
opts = types.AllReduceOptions
|
|
opts.reduceOp = op
|
|
g.allreduce([tensor], opts)
|
|
|
|
|
|
def allreduce_multigpu(
|
|
tensor_list: list,
|
|
group_name: str = "default",
|
|
op: types.ReduceOp = types.ReduceOp.SUM,
|
|
):
|
|
"""Collective allreduce a list of tensors across the group.
|
|
|
|
Args:
|
|
tensor_list: list of tensors to be allreduced, each on a GPU.
|
|
group_name: the collective group name to perform allreduce.
|
|
op: The reduce operation.
|
|
"""
|
|
if not types.cupy_available():
|
|
raise RuntimeError("Multigpu calls requires NCCL and Cupy.")
|
|
_check_tensor_list_input(tensor_list)
|
|
g = get_group_handle(group_name)
|
|
opts = types.AllReduceOptions
|
|
opts.reduceOp = op
|
|
g.allreduce(tensor_list, opts)
|
|
|
|
|
|
def barrier(group_name: str = "default"):
|
|
"""Barrier all processes in the collective group.
|
|
|
|
Args:
|
|
group_name: the name of the group to barrier.
|
|
"""
|
|
g = get_group_handle(group_name)
|
|
g.barrier()
|
|
|
|
|
|
def reduce(
|
|
tensor: Any,
|
|
dst_rank: int = 0,
|
|
group_name: str = "default",
|
|
op: types.ReduceOp = types.ReduceOp.SUM,
|
|
):
|
|
"""Reduce the tensor across the group to the destination rank.
|
|
|
|
Args:
|
|
tensor: the tensor to be reduced on this process.
|
|
dst_rank: the rank of the destination process.
|
|
group_name: the collective group name to perform reduce.
|
|
op: The reduce operation.
|
|
"""
|
|
_check_single_tensor_input(tensor)
|
|
g = get_group_handle(group_name)
|
|
|
|
# check dst rank
|
|
_check_rank_valid(g, dst_rank)
|
|
opts = types.ReduceOptions()
|
|
opts.reduceOp = op
|
|
opts.root_rank = dst_rank
|
|
opts.root_tensor = 0
|
|
g.reduce([tensor], opts)
|
|
|
|
|
|
def reduce_multigpu(
|
|
tensor_list: list,
|
|
dst_rank: int = 0,
|
|
dst_tensor: int = 0,
|
|
group_name: str = "default",
|
|
op: types.ReduceOp = types.ReduceOp.SUM,
|
|
):
|
|
"""Reduce the tensor across the group to the destination rank
|
|
and destination tensor.
|
|
|
|
Args:
|
|
tensor_list: the list of tensors to be reduced on this process;
|
|
each tensor located on a GPU.
|
|
dst_rank: the rank of the destination process.
|
|
dst_tensor: the index of GPU at the destination.
|
|
group_name: the collective group name to perform reduce.
|
|
op: The reduce operation.
|
|
"""
|
|
if not types.cupy_available():
|
|
raise RuntimeError("Multigpu calls requires NCCL and Cupy.")
|
|
_check_tensor_list_input(tensor_list)
|
|
g = get_group_handle(group_name)
|
|
|
|
# check dst rank
|
|
_check_rank_valid(g, dst_rank)
|
|
_check_root_tensor_valid(len(tensor_list), dst_tensor)
|
|
opts = types.ReduceOptions()
|
|
opts.reduceOp = op
|
|
opts.root_rank = dst_rank
|
|
opts.root_tensor = dst_tensor
|
|
g.reduce(tensor_list, opts)
|
|
|
|
|
|
def broadcast(tensor: Any, src_rank: int = 0, group_name: str = "default"):
|
|
"""Broadcast the tensor from a source process to all others.
|
|
|
|
Args:
|
|
tensor: the tensor to be broadcasted (src) or received (destination).
|
|
src_rank: the rank of the source process.
|
|
group_name: the collective group name to perform broadcast.
|
|
"""
|
|
_check_single_tensor_input(tensor)
|
|
g = get_group_handle(group_name)
|
|
|
|
# check src rank
|
|
_check_rank_valid(g, src_rank)
|
|
opts = types.BroadcastOptions()
|
|
opts.root_rank = src_rank
|
|
opts.root_tensor = 0
|
|
g.broadcast([tensor], opts)
|
|
|
|
|
|
def broadcast_multigpu(
|
|
tensor_list: list,
|
|
src_rank: int = 0,
|
|
src_tensor: int = 0,
|
|
group_name: str = "default",
|
|
):
|
|
"""Broadcast the tensor from a source GPU to all other GPUs.
|
|
|
|
Args:
|
|
tensor_list: the tensors to broadcast (src) or receive (dst).
|
|
src_rank: the rank of the source process.
|
|
src_tensor: the index of the source GPU on the source process.
|
|
group_name: the collective group name to perform broadcast.
|
|
"""
|
|
if not types.cupy_available():
|
|
raise RuntimeError("Multigpu calls requires NCCL and Cupy.")
|
|
_check_tensor_list_input(tensor_list)
|
|
g = get_group_handle(group_name)
|
|
|
|
# check src rank
|
|
_check_rank_valid(g, src_rank)
|
|
_check_root_tensor_valid(len(tensor_list), src_tensor)
|
|
opts = types.BroadcastOptions()
|
|
opts.root_rank = src_rank
|
|
opts.root_tensor = src_tensor
|
|
g.broadcast(tensor_list, opts)
|
|
|
|
|
|
def allgather(tensor_list: list, tensor: Any, group_name: str = "default"):
|
|
"""Allgather tensors from each process of the group into a list.
|
|
|
|
Args:
|
|
tensor_list: the results, stored as a list of tensors.
|
|
tensor: the tensor (to be gathered) in the current process
|
|
group_name: the name of the collective group.
|
|
"""
|
|
_check_single_tensor_input(tensor)
|
|
_check_tensor_list_input(tensor_list)
|
|
g = get_group_handle(group_name)
|
|
if len(tensor_list) != g.world_size:
|
|
# Typically CLL lib requires len(tensor_list) >= world_size;
|
|
# Here we make it more strict: len(tensor_list) == world_size.
|
|
raise RuntimeError(
|
|
"The length of the tensor list operands to allgather "
|
|
"must be equal to world_size."
|
|
)
|
|
opts = types.AllGatherOptions()
|
|
g.allgather([tensor_list], [tensor], opts)
|
|
|
|
|
|
def allgather_multigpu(
|
|
output_tensor_lists: list, input_tensor_list: list, group_name: str = "default"
|
|
):
|
|
"""Allgather tensors from each gpus of the group into lists.
|
|
|
|
Args:
|
|
output_tensor_lists: gathered results, with shape
|
|
must be num_gpus * world_size * shape(tensor).
|
|
input_tensor_list: a list of tensors, with shape
|
|
num_gpus * shape(tensor).
|
|
group_name: the name of the collective group.
|
|
"""
|
|
if not types.cupy_available():
|
|
raise RuntimeError("Multigpu calls requires NCCL and Cupy.")
|
|
_check_tensor_lists_input(output_tensor_lists)
|
|
_check_tensor_list_input(input_tensor_list)
|
|
g = get_group_handle(group_name)
|
|
opts = types.AllGatherOptions()
|
|
g.allgather(output_tensor_lists, input_tensor_list, opts)
|
|
|
|
|
|
def reducescatter(
|
|
tensor: Any,
|
|
tensor_list: list,
|
|
group_name: str = "default",
|
|
op: types.ReduceOp = types.ReduceOp.SUM,
|
|
):
|
|
"""Reducescatter a list of tensors across the group.
|
|
|
|
Reduce the list of the tensors across each process in the group, then
|
|
scatter the reduced list of tensors -- one tensor for each process.
|
|
|
|
Args:
|
|
tensor: the resulted tensor on this process.
|
|
tensor_list: The list of tensors to be reduced and scattered.
|
|
group_name: the name of the collective group.
|
|
op: The reduce operation.
|
|
"""
|
|
_check_single_tensor_input(tensor)
|
|
_check_tensor_list_input(tensor_list)
|
|
g = get_group_handle(group_name)
|
|
opts = types.ReduceScatterOptions()
|
|
opts.reduceOp = op
|
|
if len(tensor_list) != g.world_size:
|
|
raise RuntimeError(
|
|
"The length of the tensor list operands to reducescatter "
|
|
"must not be equal to world_size."
|
|
)
|
|
g.reducescatter([tensor], [tensor_list], opts)
|
|
|
|
|
|
def reducescatter_multigpu(
|
|
output_tensor_list: list,
|
|
input_tensor_lists: list,
|
|
group_name: str = "default",
|
|
op: types.ReduceOp = types.ReduceOp.SUM,
|
|
):
|
|
"""Reducescatter a list of tensors across all GPUs.
|
|
|
|
Args:
|
|
output_tensor_list: the resulted list of tensors, with
|
|
shape: num_gpus * shape(tensor).
|
|
input_tensor_lists: the original tensors, with shape:
|
|
num_gpus * world_size * shape(tensor).
|
|
group_name: the name of the collective group.
|
|
op: The reduce operation.
|
|
"""
|
|
if not types.cupy_available():
|
|
raise RuntimeError("Multigpu calls requires NCCL and Cupy.")
|
|
_check_tensor_lists_input(input_tensor_lists)
|
|
_check_tensor_list_input(output_tensor_list)
|
|
g = get_group_handle(group_name)
|
|
opts = types.ReduceScatterOptions()
|
|
opts.reduceOp = op
|
|
g.reducescatter(output_tensor_list, input_tensor_lists, opts)
|
|
|
|
|
|
def send(tensor: Any, dst_rank: int, group_name: str = "default"):
|
|
"""Send a tensor to a remote process synchronously.
|
|
|
|
Args:
|
|
tensor: the tensor to send.
|
|
dst_rank: the rank of the destination process.
|
|
group_name: the name of the collective group.
|
|
"""
|
|
_check_single_tensor_input(tensor)
|
|
g = get_group_handle(group_name)
|
|
_check_rank_valid(g, dst_rank)
|
|
if dst_rank == g.rank:
|
|
raise RuntimeError("The destination rank '{}' is self.".format(dst_rank))
|
|
opts = types.SendOptions()
|
|
opts.dst_rank = dst_rank
|
|
g.send([tensor], opts)
|
|
|
|
|
|
def send_multigpu(
|
|
tensor: Any,
|
|
dst_rank: int,
|
|
dst_gpu_index: int,
|
|
group_name: str = "default",
|
|
n_elements: int = 0,
|
|
):
|
|
"""Send a tensor to a remote GPU synchronously.
|
|
|
|
The function assumes each process owns >1 GPUs, and the sender
|
|
process and receiver process has equal number of GPUs.
|
|
|
|
Args:
|
|
tensor: the tensor to send, located on a GPU.
|
|
dst_rank: the rank of the destination process.
|
|
dst_gpu_index: the destination gpu index.
|
|
group_name: the name of the collective group.
|
|
n_elements: if specified, send the next n elements
|
|
from the starting address of tensor.
|
|
"""
|
|
if not types.cupy_available():
|
|
raise RuntimeError("send_multigpu call requires NCCL.")
|
|
_check_single_tensor_input(tensor)
|
|
g = get_group_handle(group_name)
|
|
_check_rank_valid(g, dst_rank)
|
|
if dst_rank == g.rank:
|
|
raise RuntimeError(
|
|
"The dst_rank '{}' is self. Considering "
|
|
"doing GPU to GPU memcpy instead?".format(dst_rank)
|
|
)
|
|
if n_elements < 0:
|
|
raise RuntimeError("The n_elements '{}' should >= 0.".format(n_elements))
|
|
opts = types.SendOptions()
|
|
opts.dst_rank = dst_rank
|
|
opts.dst_gpu_index = dst_gpu_index
|
|
opts.n_elements = n_elements
|
|
g.send([tensor], opts)
|
|
|
|
|
|
def recv(tensor: Any, src_rank: int, group_name: str = "default"):
|
|
"""Receive a tensor from a remote process synchronously.
|
|
|
|
Args:
|
|
tensor: the received tensor.
|
|
src_rank: the rank of the source process.
|
|
group_name: the name of the collective group.
|
|
"""
|
|
_check_single_tensor_input(tensor)
|
|
g = get_group_handle(group_name)
|
|
_check_rank_valid(g, src_rank)
|
|
if src_rank == g.rank:
|
|
raise RuntimeError("The destination rank '{}' is self.".format(src_rank))
|
|
opts = types.RecvOptions()
|
|
opts.src_rank = src_rank
|
|
g.recv([tensor], opts)
|
|
|
|
|
|
def recv_multigpu(
|
|
tensor: Any,
|
|
src_rank: int,
|
|
src_gpu_index: int,
|
|
group_name: str = "default",
|
|
n_elements: int = 0,
|
|
):
|
|
"""Receive a tensor from a remote GPU synchronously.
|
|
|
|
The function asssume each process owns >1 GPUs, and the sender
|
|
process and receiver process has equal nubmer of GPUs.
|
|
|
|
Args:
|
|
tensor: The received tensor, located on a GPU.
|
|
src_rank: The rank of the source process.
|
|
src_gpu_index: The index of the source GPU on the src process.
|
|
group_name: The name of the collective group.
|
|
n_elements: if specified, receive the next n elements
|
|
from the starting address of tensor.
|
|
"""
|
|
if not types.cupy_available():
|
|
raise RuntimeError("recv_multigpu call requires NCCL.")
|
|
_check_single_tensor_input(tensor)
|
|
g = get_group_handle(group_name)
|
|
_check_rank_valid(g, src_rank)
|
|
if src_rank == g.rank:
|
|
raise RuntimeError(
|
|
"The dst_rank '{}' is self. Considering "
|
|
"doing GPU to GPU memcpy instead?".format(src_rank)
|
|
)
|
|
if n_elements < 0:
|
|
raise RuntimeError("The n_elements '{}' should be >= 0.".format(n_elements))
|
|
opts = types.RecvOptions()
|
|
opts.src_rank = src_rank
|
|
opts.src_gpu_index = src_gpu_index
|
|
opts.n_elements = n_elements
|
|
g.recv([tensor], opts)
|
|
|
|
|
|
def synchronize(gpu_id: int):
|
|
"""Synchronize the current process to a give device.
|
|
|
|
Args:
|
|
gpu_id: the GPU device id to synchronize.
|
|
"""
|
|
if not types.cupy_available():
|
|
raise RuntimeError("synchronize call requires CUDA and NCCL.")
|
|
import cupy as cp
|
|
|
|
cp.cuda.Device(gpu_id).synchronize()
|
|
|
|
|
|
def get_group_handle(group_name: str = "default"):
|
|
"""Check if the group is initialized and return the group handle.
|
|
|
|
Args:
|
|
group_name: the name of the collective group.
|
|
|
|
Returns:
|
|
The collective group handle.
|
|
"""
|
|
_check_inside_actor()
|
|
global _group_mgr
|
|
global _group_mgr_lock
|
|
with _group_mgr_lock:
|
|
if not _group_mgr.is_group_exist(group_name):
|
|
# try loading from remote info store
|
|
try:
|
|
# if the information is stored in an Info object,
|
|
# get and create the group.
|
|
name = "info_" + group_name
|
|
mgr = ray.get_actor(name=name)
|
|
ids, world_size, rank, backend, gloo_timeout = ray.get(
|
|
mgr.get_info.remote()
|
|
)
|
|
worker = ray._private.worker.global_worker
|
|
id_ = worker.core_worker.get_actor_id()
|
|
r = rank[ids.index(id_)]
|
|
_group_mgr.create_collective_group(
|
|
backend, world_size, r, group_name, gloo_timeout
|
|
)
|
|
except ValueError as exc:
|
|
# check if this group is initialized using options()
|
|
if (
|
|
"collective_group_name" in os.environ
|
|
and os.environ["collective_group_name"] == group_name
|
|
):
|
|
rank = int(os.environ["collective_rank"])
|
|
world_size = int(os.environ["collective_world_size"])
|
|
backend = os.environ["collective_backend"]
|
|
gloo_timeout = int(os.getenv("collective_gloo_timeout", 30000))
|
|
_group_mgr.create_collective_group(
|
|
backend, world_size, rank, group_name, gloo_timeout
|
|
)
|
|
else:
|
|
raise RuntimeError(
|
|
"The collective group '{}' is not "
|
|
"initialized in the process.".format(group_name)
|
|
) from exc
|
|
g = _group_mgr.get_group_by_name(group_name)
|
|
return g
|
|
|
|
|
|
def _check_single_tensor_input(tensor):
|
|
"""Check if the tensor is with a supported type."""
|
|
if isinstance(tensor, np.ndarray):
|
|
return
|
|
if types.cupy_available():
|
|
if isinstance(tensor, types.cp.ndarray):
|
|
return
|
|
if types.torch_available():
|
|
if isinstance(tensor, types.th.Tensor):
|
|
return
|
|
raise RuntimeError(
|
|
"Unrecognized tensor type '{}'. Supported types are: "
|
|
"np.ndarray, torch.Tensor, cupy.ndarray.".format(type(tensor))
|
|
)
|
|
|
|
|
|
def _check_inside_actor():
|
|
"""Check if currently it is inside a Ray actor/task."""
|
|
worker = ray._private.worker.global_worker
|
|
if worker.mode == ray.WORKER_MODE:
|
|
return
|
|
else:
|
|
raise RuntimeError(
|
|
"The collective APIs shall be only used inside a Ray actor or task."
|
|
)
|
|
|
|
|
|
def _check_rank_valid(g, rank: int):
|
|
"""Check the rank: 0 <= rank < world_size."""
|
|
if rank < 0:
|
|
raise ValueError("rank '{}' is negative.".format(rank))
|
|
if rank >= g.world_size:
|
|
raise ValueError(
|
|
"rank '{}' must be less than world size '{}'".format(rank, g.world_size)
|
|
)
|
|
|
|
|
|
def _check_tensor_list_input(tensor_list):
|
|
"""Check if the input is a list of supported tensor types."""
|
|
if not isinstance(tensor_list, list):
|
|
raise RuntimeError(
|
|
"The input must be a list of tensors. "
|
|
"Got '{}'.".format(type(tensor_list))
|
|
)
|
|
if not tensor_list:
|
|
raise RuntimeError("Got an empty list of tensors.")
|
|
for t in tensor_list:
|
|
_check_single_tensor_input(t)
|
|
|
|
|
|
def _check_tensor_lists_input(tensor_lists):
|
|
"""Check if the input is a list of lists of supported tensor types."""
|
|
if not isinstance(tensor_lists, list):
|
|
raise RuntimeError(
|
|
"The input must be a list of lists of tensors. "
|
|
"Got '{}'.".format(type(tensor_lists))
|
|
)
|
|
if not tensor_lists:
|
|
raise RuntimeError(f"Did not receive tensors. Got: {tensor_lists}")
|
|
for t in tensor_lists:
|
|
_check_tensor_list_input(t)
|
|
|
|
|
|
def _check_root_tensor_valid(length, root_tensor):
|
|
"""Check the root_tensor device is 0 <= root_tensor < length"""
|
|
if root_tensor < 0:
|
|
raise ValueError("root_tensor '{}' is negative.".format(root_tensor))
|
|
if root_tensor >= length:
|
|
raise ValueError(
|
|
"root_tensor '{}' is greater than the number of GPUs: "
|
|
"'{}'".format(root_tensor, length)
|
|
)
|