|
|
|
|
|
from typing import Dict, List, Optional, Union |
|
|
|
|
|
import torch |
|
|
from torch._C._distributed_rpc import _TensorPipeRpcBackendOptionsBase |
|
|
from . import constants as rpc_contants |
|
|
|
|
|
|
|
|
DeviceType = Union[int, str, torch.device] |
|
|
|
|
|
__all__ = ["TensorPipeRpcBackendOptions"] |
|
|
|
|
|
def _to_device(device: DeviceType) -> torch.device: |
|
|
device = torch.device(device) |
|
|
if device.type != "cuda": |
|
|
raise ValueError( |
|
|
"`set_devices` expect a list of CUDA devices, but got " |
|
|
f"device type {device.type}." |
|
|
) |
|
|
return device |
|
|
|
|
|
|
|
|
def _to_device_map( |
|
|
device_map: Dict[DeviceType, DeviceType] |
|
|
) -> Dict[torch.device, torch.device]: |
|
|
full_device_map: Dict[torch.device, torch.device] = {} |
|
|
reverse_map: Dict[torch.device, torch.device] = {} |
|
|
for k, v in device_map.items(): |
|
|
k, v = torch.device(k), torch.device(v) |
|
|
if v in reverse_map: |
|
|
raise ValueError( |
|
|
"`device_map` only supports 1-to-1 mapping, " |
|
|
f"trying to map {k} and {reverse_map[v]} to {v}" |
|
|
) |
|
|
full_device_map[k] = v |
|
|
reverse_map[v] = k |
|
|
return full_device_map |
|
|
|
|
|
|
|
|
def _to_device_list(devices: List[DeviceType]) -> List[torch.device]: |
|
|
return list(map(_to_device, devices)) |
|
|
|
|
|
|
|
|
class TensorPipeRpcBackendOptions(_TensorPipeRpcBackendOptionsBase): |
|
|
r""" |
|
|
The backend options for |
|
|
:class:`~torch.distributed.rpc.TensorPipeAgent`, derived from |
|
|
:class:`~torch.distributed.rpc.RpcBackendOptions`. |
|
|
|
|
|
Args: |
|
|
num_worker_threads (int, optional): The number of threads in the |
|
|
thread-pool used by |
|
|
:class:`~torch.distributed.rpc.TensorPipeAgent` to execute |
|
|
requests (default: 16). |
|
|
rpc_timeout (float, optional): The default timeout, in seconds, |
|
|
for RPC requests (default: 60 seconds). If the RPC has not |
|
|
completed in this timeframe, an exception indicating so will |
|
|
be raised. Callers can override this timeout for individual |
|
|
RPCs in :meth:`~torch.distributed.rpc.rpc_sync` and |
|
|
:meth:`~torch.distributed.rpc.rpc_async` if necessary. |
|
|
init_method (str, optional): The URL to initialize the distributed |
|
|
store used for rendezvous. It takes any value accepted for the |
|
|
same argument of :meth:`~torch.distributed.init_process_group` |
|
|
(default: ``env://``). |
|
|
device_maps (Dict[str, Dict], optional): Device placement mappings from |
|
|
this worker to the callee. Key is the callee worker name and value |
|
|
the dictionary (``Dict`` of ``int``, ``str``, or ``torch.device``) |
|
|
that maps this worker's devices to the callee worker's devices. |
|
|
(default: ``None``) |
|
|
devices (List[int, str, or ``torch.device``], optional): all local |
|
|
CUDA devices used by RPC agent. By Default, it will be initialized |
|
|
to all local devices from its own ``device_maps`` and corresponding |
|
|
devices from its peers' ``device_maps``. When processing CUDA RPC |
|
|
requests, the agent will properly synchronize CUDA streams for |
|
|
all devices in this ``List``. |
|
|
""" |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
*, |
|
|
num_worker_threads: int = rpc_contants.DEFAULT_NUM_WORKER_THREADS, |
|
|
rpc_timeout: float = rpc_contants.DEFAULT_RPC_TIMEOUT_SEC, |
|
|
init_method: str = rpc_contants.DEFAULT_INIT_METHOD, |
|
|
device_maps: Optional[Dict[str, Dict[DeviceType, DeviceType]]] = None, |
|
|
devices: Optional[List[DeviceType]] = None, |
|
|
_transports: Optional[List] = None, |
|
|
_channels: Optional[List] = None, |
|
|
): |
|
|
full_device_maps = ( |
|
|
{} |
|
|
if device_maps is None |
|
|
else {k: _to_device_map(v) for k, v in device_maps.items()} |
|
|
) |
|
|
full_device_list = [] if devices is None else _to_device_list(devices) |
|
|
super().__init__( |
|
|
num_worker_threads, |
|
|
_transports, |
|
|
_channels, |
|
|
rpc_timeout, |
|
|
init_method, |
|
|
full_device_maps, |
|
|
full_device_list, |
|
|
) |
|
|
|
|
|
def set_device_map(self, to: str, device_map: Dict[DeviceType, DeviceType]): |
|
|
r""" |
|
|
Set device mapping between each RPC caller and callee pair. This |
|
|
function can be called multiple times to incrementally add |
|
|
device placement configurations. |
|
|
|
|
|
Args: |
|
|
to (str): Callee name. |
|
|
device_map (Dict of int, str, or torch.device): Device placement |
|
|
mappings from this worker to the callee. This map must be |
|
|
invertible. |
|
|
|
|
|
Example: |
|
|
>>> # xdoctest: +SKIP("distributed") |
|
|
>>> # both workers |
|
|
>>> def add(x, y): |
|
|
>>> print(x) # tensor([1., 1.], device='cuda:1') |
|
|
>>> return x + y, (x + y).to(2) |
|
|
>>> |
|
|
>>> # on worker 0 |
|
|
>>> options = TensorPipeRpcBackendOptions( |
|
|
>>> num_worker_threads=8, |
|
|
>>> device_maps={"worker1": {0: 1}} |
|
|
>>> # maps worker0's cuda:0 to worker1's cuda:1 |
|
|
>>> ) |
|
|
>>> options.set_device_map("worker1", {1: 2}) |
|
|
>>> # maps worker0's cuda:1 to worker1's cuda:2 |
|
|
>>> |
|
|
>>> rpc.init_rpc( |
|
|
>>> "worker0", |
|
|
>>> rank=0, |
|
|
>>> world_size=2, |
|
|
>>> backend=rpc.BackendType.TENSORPIPE, |
|
|
>>> rpc_backend_options=options |
|
|
>>> ) |
|
|
>>> |
|
|
>>> x = torch.ones(2) |
|
|
>>> rets = rpc.rpc_sync("worker1", add, args=(x.to(0), 1)) |
|
|
>>> # The first argument will be moved to cuda:1 on worker1. When |
|
|
>>> # sending the return value back, it will follow the invert of |
|
|
>>> # the device map, and hence will be moved back to cuda:0 and |
|
|
>>> # cuda:1 on worker0 |
|
|
>>> print(rets[0]) # tensor([2., 2.], device='cuda:0') |
|
|
>>> print(rets[1]) # tensor([2., 2.], device='cuda:1') |
|
|
""" |
|
|
full_device_map = _to_device_map(device_map) |
|
|
curr_device_maps = super().device_maps |
|
|
|
|
|
if to in curr_device_maps: |
|
|
for k, v in full_device_map.items(): |
|
|
if k in curr_device_maps[to] and v != curr_device_maps[to][k]: |
|
|
raise ValueError( |
|
|
"`set_device_map` only supports 1-to-1 mapping, trying" |
|
|
f" to map {k} to {v} and {curr_device_maps[to][k]}" |
|
|
) |
|
|
|
|
|
super()._set_device_map(to, full_device_map) |
|
|
|
|
|
def set_devices(self, devices: List[DeviceType]): |
|
|
r""" |
|
|
Set local devices used by the TensorPipe RPC agent. When processing |
|
|
CUDA RPC requests, the TensorPipe RPC agent will properly synchronize |
|
|
CUDA streams for all devices in this ``List``. |
|
|
|
|
|
Args: |
|
|
devices (List of int, str, or torch.device): local devices used by |
|
|
the TensorPipe RPC agent. |
|
|
""" |
|
|
self.devices = _to_device_list(devices) |
|
|
|
