import os

import torch.distributed as dist


dist.init_process_group("gloo")
rank = dist.get_rank()
world_size = dist.get_world_size()
local_rank = int(os.environ["LOCAL_RANK"])
local_world_size = int(os.environ["LOCAL_WORLD_SIZE"])

message = (
    f"rank={rank} local_rank={local_rank} "
    f"world_size={world_size} local_world_size={local_world_size}"
)

gathered = [None] * world_size if rank == 0 else None
dist.gather_object(message, gathered, dst=0)

if rank == 0:
    for line in gathered:
        print(line)
    print("torchrun_launch_ok=True")

dist.destroy_process_group()