r"""Receive a tensor from another process. Args: src_rank: Rank of source process. device: The specific device to execute this operator. None default device means the device of inp will be used. Specify "gpu0:1" to execute this operator on diffrent cuda s
(src_rank: int, device: Optional[str] = None, inp=None)
| 853 | |
| 854 | |
| 855 | def remote_recv(src_rank: int, device: Optional[str] = None, inp=None) -> Tensor: |
| 856 | r"""Receive a tensor from another process. |
| 857 | |
| 858 | Args: |
| 859 | src_rank: Rank of source process. |
| 860 | device: The specific device to execute this operator. |
| 861 | None default device means the device of inp will be used. |
| 862 | Specify "gpu0:1" to execute this operator on diffrent cuda stream, |
| 863 | 1 is stream id, and default stream id is 0. |
| 864 | inp: Dummy input to determine received tensor type. |
| 865 | |
| 866 | Returns: |
| 867 | Received tensor. |
| 868 | |
| 869 | Examples: |
| 870 | |
| 871 | .. code-block:: |
| 872 | |
| 873 | if rank == 0: |
| 874 | data = mge.tensor(1) |
| 875 | # Tensor(1) |
| 876 | F.distributed.remote_send(data, 1) # return None |
| 877 | else: |
| 878 | data = F.distributed.remote_recv(0) |
| 879 | # Tensor(1) |
| 880 | """ |
| 881 | group = _SendRecvGroup(src_rank, get_rank()) |
| 882 | shape, dtype = _bcast_shape_dtype(group, None) |
| 883 | |
| 884 | if device is None: |
| 885 | device = get_default_device() |
| 886 | # dummy input |
| 887 | if inp is None: |
| 888 | inp = Tensor(0, device=device) |
| 889 | _bcast_tracer_state(group, inp) |
| 890 | |
| 891 | op = RemoteRecv() |
| 892 | op.key = group.key |
| 893 | op.cn = device |
| 894 | op.shape = shape |
| 895 | op.dtype = dtype |
| 896 | op.addr, op.port = get_mm_server_addr() |
| 897 | op.rank_from = src_rank |
| 898 | op.backend = _backend() |
| 899 | ret = _RemoteRecv(op)(inp) |
| 900 | return ret |
| 901 | |
| 902 | |
| 903 | def _remote_send_nobackward(inp: Tensor, dest_rank: int): |