(args)
| 469 | |
| 470 | |
| 471 | def mpi_define_env(args): |
| 472 | from mpi4py import MPI |
| 473 | comm = MPI.COMM_WORLD |
| 474 | rank = comm.Get_rank() |
| 475 | world_size = comm.Get_size() |
| 476 | |
| 477 | master_addr = None |
| 478 | if rank == 0: |
| 479 | master_addr = get_hostname() |
| 480 | master_addr = comm.bcast(master_addr, root=0) |
| 481 | |
| 482 | # Determine local rank by assuming hostnames are unique |
| 483 | proc_name = MPI.Get_processor_name() |
| 484 | all_procs = comm.allgather(proc_name) |
| 485 | local_rank = sum([i == proc_name for i in all_procs[:rank]]) |
| 486 | |
| 487 | os.environ['RANK'] = str(rank) |
| 488 | os.environ['WORLD_SIZE'] = str(world_size) |
| 489 | args.local_rank = local_rank |
| 490 | args.world_size = world_size |
| 491 | args.rank = rank |
| 492 | os.environ['MASTER_ADDR'] = master_addr |
| 493 | os.environ['MASTER_PORT'] = "29500" # TORCH_DISTRIBUTED_DEFAULT_PORT = 29500 |
| 494 | |
| 495 | print( |
| 496 | "Discovered MPI settings of world_rank={}, local_rank={}, world_size={}, master_addr={}, master_port={}" |
| 497 | .format(os.environ['RANK'], |
| 498 | args.local_rank, |
| 499 | os.environ['WORLD_SIZE'], |
| 500 | os.environ['MASTER_ADDR'], |
| 501 | os.environ['MASTER_PORT'])) |
no test coverage detected