Barrier for all tasks

Is there any recommended practice to wait for preparation complete for all tasks? That is, to insert a barrier.

I've implemented one (for server-client scenario) with PyTorch, though I believe there might be better options. For example, something that has been natively supported by pai runtime.

#!/usr/bin/env python3
import functools
import os

import torch

print = functools.partial(print, flush=True)

master_addr = os.environ['PAI_HOST_IP_server_0']
master_port = os.environ['PAI_PORT_LIST_server_0_torch']
task_count = int(os.environ['PAI_TASK_ROLE_TASK_COUNT_client']) + 1
if os.environ['PAI_CURRENT_TASK_ROLE_NAME'] == 'server':
    task_rank = 0
else:
    task_rank = int(os.environ['PAI_CURRENT_TASK_ROLE_CURRENT_TASK_INDEX']) + 1
torch.distributed.init_process_group('gloo', init_method=f'tcp://{master_addr}:{master_port}',
                                     world_size=task_count, rank=task_rank)
print('PAI barrier: waiting for other workers...')
torch.distributed.barrier()