| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| |
| from contextlib import contextmanager |
|
|
| import torch |
|
|
|
|
| def is_dist_avail_and_initialized() -> bool: |
| return torch.distributed.is_available() and torch.distributed.is_initialized() |
|
|
|
|
| def get_rank() -> int: |
| if is_dist_avail_and_initialized(): |
| return torch.distributed.get_rank() |
| return 0 |
|
|
|
|
| def barrier(): |
| if is_dist_avail_and_initialized(): |
| torch.distributed.barrier() |
|
|
|
|
| def _collective_device() -> torch.device: |
| """Pick the device used for status-broadcast collectives. |
| |
| NCCL backend requires CUDA tensors; gloo / mpi accept CPU. We rely on the |
| caller (typically `experiment.run` after `torch.cuda.set_device(local_rank)`) |
| having already pinned the current CUDA device. |
| """ |
| if torch.distributed.get_backend() == "nccl": |
| return torch.device(f"cuda:{torch.cuda.current_device()}") |
| return torch.device("cpu") |
|
|
|
|
| @contextmanager |
| def run_or_wait_on_rank0(label: str | None = None): |
| """Run the body on rank-0 only, broadcasting any rank-0 error to all ranks. |
| |
| Yields ``True`` on rank-0, ``False`` elsewhere; callers gate the body with |
| ``if is_rank0:``. If rank-0 raises, every rank raises consistently |
| (rank-0 re-raises the original; others get ``RuntimeError(label: ...)``). |
| Without this contract a rank-0 raise inside a plain ``if get_rank() == 0:`` |
| block leaves other ranks hanging at the next NCCL collective until the |
| 30-min timeout fires. |
| |
| Degenerates to a plain ``with`` block when distributed is not initialized. |
| |
| with run_or_wait_on_rank0(label="generate_stats") as is_rank0: |
| if is_rank0: |
| ... |
| """ |
| if not is_dist_avail_and_initialized(): |
| yield True |
| return |
|
|
| rank = torch.distributed.get_rank() |
| is_rank0 = rank == 0 |
| device = _collective_device() |
| status = torch.zeros(1, dtype=torch.int, device=device) |
| rank0_err: BaseException | None = None |
| error_summary: list[str | None] = [None] |
| try: |
| yield is_rank0 |
| except BaseException as exc: |
| if is_rank0: |
| status += 1 |
| rank0_err = exc |
| error_summary[0] = f"{type(exc).__name__}: {exc}" |
| else: |
| |
| |
| raise |
| finally: |
| torch.distributed.all_reduce(status, op=torch.distributed.ReduceOp.MAX) |
| if status.item() != 0: |
| torch.distributed.broadcast_object_list(error_summary, src=0) |
| torch.distributed.barrier() |
|
|
| if status.item() == 0: |
| return |
| if rank0_err is not None: |
| raise rank0_err |
| label_prefix = f"{label}: " if label else "" |
| raise RuntimeError( |
| f"{label_prefix}rank-0 failed inside run_or_wait_on_rank0() with " |
| f"{error_summary[0] or '<no error info broadcast>'}; " |
| "see rank-0's traceback for the full stack." |
| ) from None |
|
|
|
|
| def run_on_rank0(fn, *args, label: str | None = None, **kwargs): |
| """Call ``fn(*args, **kwargs)`` on rank-0 only; return its result there, ``None`` elsewhere. |
| |
| A rank-0 failure is broadcast to every rank (see ``run_or_wait_on_rank0``). |
| ``args`` / ``kwargs`` are evaluated on every rank, so keep rank-0-only work inside ``fn``. |
| """ |
| result = None |
| with run_or_wait_on_rank0(label=label or getattr(fn, "__qualname__", None)) as is_rank0: |
| if is_rank0: |
| result = fn(*args, **kwargs) |
| return result |
|
|
