Initialize distributed process group. Args: rank: Process rank (ignored for OMPI/torchelastic). world_size: Total number of processes (ignored for OMPI/torchelastic). port: Initial port to try. If None, a free port will be selected. shared_port: Optional mp.Value
(
rank: int = 0,
world_size: int = 1,
port: Optional[int] = None,
shared_port: Optional["mp.Value"] = None,
port_ready_barrier: Optional["mp.Barrier"] = None,
max_retries: int = 5,
)
| 150 | |
| 151 | |
| 152 | def initialize( |
| 153 | rank: int = 0, |
| 154 | world_size: int = 1, |
| 155 | port: Optional[int] = None, |
| 156 | shared_port: Optional["mp.Value"] = None, |
| 157 | port_ready_barrier: Optional["mp.Barrier"] = None, |
| 158 | max_retries: int = 5, |
| 159 | ) -> Tuple[int, int]: |
| 160 | """Initialize distributed process group. |
| 161 | |
| 162 | Args: |
| 163 | rank: Process rank (ignored for OMPI/torchelastic). |
| 164 | world_size: Total number of processes (ignored for OMPI/torchelastic). |
| 165 | port: Initial port to try. If None, a free port will be selected. |
| 166 | shared_port: Optional mp.Value for rank 0 to share the final port with other ranks. |
| 167 | port_ready_barrier: Optional mp.Barrier to synchronize port selection. |
| 168 | max_retries: Maximum number of port retry attempts for rank 0. |
| 169 | """ |
| 170 | if is_ompi(): |
| 171 | lib = "OMPI" |
| 172 | local_rank = int(os.environ["OMPI_COMM_WORLD_LOCAL_RANK"]) |
| 173 | world_size = int(os.environ["OMPI_COMM_WORLD_SIZE"]) |
| 174 | elif is_torchelastic(): |
| 175 | lib = "TORCHELASTIC" |
| 176 | local_rank = int(os.environ["LOCAL_RANK"]) |
| 177 | world_size = int(os.environ["WORLD_SIZE"]) |
| 178 | port = int(os.environ["MASTER_PORT"]) |
| 179 | else: |
| 180 | lib = "MP" |
| 181 | local_rank = rank |
| 182 | |
| 183 | if port is None: |
| 184 | assert world_size == 1, "Port is required for world_size > 1." |
| 185 | port = get_free_port() |
| 186 | |
| 187 | ad_logger.set_rank(local_rank) |
| 188 | |
| 189 | # Necessary to assign a device to each rank. |
| 190 | torch.cuda.set_device(local_rank) |
| 191 | |
| 192 | # If we have shared port synchronization (multiprocess spawn mode) |
| 193 | if shared_port is not None and port_ready_barrier is not None: |
| 194 | if local_rank == 0: |
| 195 | # Rank 0: try ports until one works, then share with other ranks |
| 196 | init_success = False |
| 197 | init_error = None |
| 198 | try: |
| 199 | for attempt in range(max_retries): |
| 200 | ad_logger.info( |
| 201 | f"Initializing for: {lib=}, {local_rank=}, {world_size=}, {port=} (attempt {attempt + 1})" |
| 202 | ) |
| 203 | if _try_init_process_group(local_rank, world_size, port): |
| 204 | # Success! Share the working port with other ranks |
| 205 | shared_port.value = port |
| 206 | init_success = True |
| 207 | break |
| 208 | else: |
| 209 | # Port was taken, try a new one |
no test coverage detected