| """Process-group context and rank-local device selection.""" |
|
|
| from __future__ import annotations |
|
|
| import os |
| from contextlib import contextmanager |
| from dataclasses import dataclass |
| from datetime import timedelta |
| from typing import Iterator |
|
|
| import torch |
| import torch.distributed as dist |
|
|
|
|
| @dataclass(frozen=True) |
| class DistributedContext: |
| enabled: bool = False |
| rank: int = 0 |
| local_rank: int = 0 |
| world_size: int = 1 |
| device: torch.device = torch.device("cpu") |
|
|
| @property |
| def is_main_process(self) -> bool: |
| return self.rank == 0 |
|
|
| @classmethod |
| def single_process(cls, device: torch.device | str = "cpu") -> "DistributedContext": |
| return cls(device=torch.device(device)) |
|
|
|
|
| def _launched() -> bool: |
| return "RANK" in os.environ or "WORLD_SIZE" in os.environ |
|
|
|
|
| def initialize( |
| *, |
| enabled: bool | None = None, |
| backend: str | None = None, |
| device: torch.device | str | None = None, |
| timeout_seconds: int = 1800, |
| ) -> DistributedContext: |
| """Initialize ``torchrun``'s ``env://`` process group when requested.""" |
|
|
| launched = _launched() |
| if enabled is False and launched and int(os.environ.get("WORLD_SIZE", "1")) > 1: |
| raise ValueError( |
| "distributed launcher variables are present but distributed.enabled=false" |
| ) |
| active = launched if enabled is None else enabled |
| if not active: |
| requested = torch.device( |
| device or ("cuda" if torch.cuda.is_available() else "cpu") |
| ) |
| if requested.type == "cuda" and not torch.cuda.is_available(): |
| requested = torch.device("cpu") |
| return DistributedContext.single_process(requested) |
|
|
| try: |
| rank = int(os.environ["RANK"]) |
| local_rank = int(os.environ.get("LOCAL_RANK", rank)) |
| world_size = int(os.environ["WORLD_SIZE"]) |
| except KeyError as error: |
| raise ValueError( |
| "distributed execution requires torchrun RANK and WORLD_SIZE" |
| ) from error |
| if world_size < 2: |
| return DistributedContext.single_process(device or "cpu") |
| requested = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu")) |
| if requested.type == "cuda": |
| if not torch.cuda.is_available(): |
| raise RuntimeError( |
| "CUDA was requested for distributed execution but is unavailable" |
| ) |
| torch.cuda.set_device(local_rank) |
| assigned = torch.device("cuda", local_rank) |
| selected_backend = backend or "nccl" |
| else: |
| assigned = torch.device("cpu") |
| selected_backend = backend or "gloo" |
| if not dist.is_initialized(): |
| dist.init_process_group( |
| backend=selected_backend, |
| init_method="env://", |
| timeout=timedelta(seconds=timeout_seconds), |
| ) |
| return DistributedContext(True, rank, local_rank, world_size, assigned) |
|
|
|
|
| def finalize(context: DistributedContext) -> None: |
| """Destroy a process group created for ``context``.""" |
|
|
| if context.enabled and dist.is_available() and dist.is_initialized(): |
| dist.destroy_process_group() |
|
|
|
|
| @contextmanager |
| def process_group(**kwargs: object) -> Iterator[DistributedContext]: |
| context = initialize(**kwargs) |
| try: |
| yield context |
| finally: |
| finalize(context) |
|
|
