| 155 | |
| 156 | |
| 157 | class _TorchDistributedEnvironment: |
| 158 | def __init__(self): |
| 159 | self.master_addr = "127.0.0.1" |
| 160 | self.master_port = 0 |
| 161 | self.rank = -1 |
| 162 | self.world_size = -1 |
| 163 | self.local_rank = -1 |
| 164 | self.local_world_size = -1 |
| 165 | |
| 166 | if _is_slurm_job_process(): |
| 167 | return self._set_from_slurm_env() |
| 168 | |
| 169 | env_vars = _collect_env_vars() |
| 170 | if not env_vars: |
| 171 | # Environment is not set |
| 172 | pass |
| 173 | elif len(env_vars) == len(_TORCH_DISTRIBUTED_ENV_VARS): |
| 174 | # Environment is fully set |
| 175 | return self._set_from_preset_env() |
| 176 | else: |
| 177 | # Environment is partially set |
| 178 | collected_env_vars = ", ".join(env_vars.keys()) |
| 179 | raise RuntimeError(f"Partially set environment: {collected_env_vars}") |
| 180 | |
| 181 | if torch.cuda.device_count() > 0: |
| 182 | return self._set_from_local() |
| 183 | |
| 184 | raise RuntimeError("Can't initialize PyTorch distributed environment") |
| 185 | |
| 186 | # Slurm job created with sbatch, submitit, etc... |
| 187 | def _set_from_slurm_env(self): |
| 188 | # logger.info("Initialization from Slurm environment") |
| 189 | job_id = int(os.environ["SLURM_JOB_ID"]) |
| 190 | node_count = int(os.environ["SLURM_JOB_NUM_NODES"]) |
| 191 | nodes = _parse_slurm_node_list(os.environ["SLURM_JOB_NODELIST"]) |
| 192 | assert len(nodes) == node_count, f"Expected {node_count} nodes, got {nodes}" |
| 193 | |
| 194 | self.master_addr = nodes[0] |
| 195 | self.master_port = _get_master_port(seed=job_id) |
| 196 | self.rank = int(os.environ["SLURM_PROCID"]) |
| 197 | self.world_size = int(os.environ["SLURM_NTASKS"]) |
| 198 | assert self.rank < self.world_size |
| 199 | self.local_rank = int(os.environ["SLURM_LOCALID"]) |
| 200 | self.local_world_size = self.world_size // node_count |
| 201 | assert self.local_rank < self.local_world_size |
| 202 | |
| 203 | # Single node job with preset environment (i.e. torchrun) |
| 204 | def _set_from_preset_env(self): |
| 205 | # logger.info("Initialization from preset environment") |
| 206 | self.master_addr = os.environ["MASTER_ADDR"] |
| 207 | self.master_port = os.environ["MASTER_PORT"] |
| 208 | self.rank = int(os.environ["RANK"]) |
| 209 | self.world_size = int(os.environ["WORLD_SIZE"]) |
| 210 | assert self.rank < self.world_size |
| 211 | self.local_rank = int(os.environ["LOCAL_RANK"]) |
| 212 | self.local_world_size = int(os.environ["LOCAL_WORLD_SIZE"]) |
| 213 | assert self.local_rank < self.local_world_size |
| 214 | |