pytorch_distributed_point_to_point_send_recv_quickstart.py

1import os
2import torch
3import torch.distributed as dist
4import torch.multiprocessing as mp
5
6def run(rank, size):
7    """ Distributed function to be implemented later. """
8    tensor = torch.zeros(1)
9    if rank == 0:
10        tensor += 1
11        # Send the tensor to process 1
12        dist.send(tensor=tensor, dst=1)
13    else:
14        # Receive tensor from process 0
15        dist.recv(tensor=tensor, src=0)
16    print('Rank ', rank, ' has data ', tensor[0])
17
18def init_process(rank, size, fn, backend='gloo'):
19    """ Initialize the distributed environment. """
20    os.environ['MASTER_ADDR'] = '127.0.0.1'
21    os.environ['MASTER_PORT'] = '29500'
22    dist.init_process_group(backend, rank=rank, world_size=size)
23    fn(rank, size)
24
25if __name__ == "__main__":
26    size = 2
27    processes = []
28    mp.set_start_method("spawn")
29    for rank in range(size):
30        p = mp.Process(target=init_process, args=(rank, size, run))
31        p.start()
32        processes.append(p)
33
34    for p in processes:
35        p.join()