Function enable

astroclip/astrodino/distributed.py:245–281

Enable distributed mode Args: set_cuda_current_device: If True, call torch.cuda.set_device() to set the current PyTorch CUDA device to the one matching the local rank. overwrite: If True, overwrites already set variables. Else fails.

(
    *,
    set_cuda_current_device: bool = True,
    overwrite: bool = False,
    allow_nccl_timeout: bool = False,
)

def enable(
    *,
    set_cuda_current_device: bool = True,
    overwrite: bool = False,
    allow_nccl_timeout: bool = False,
):
    """Enable distributed mode

    Args:
        set_cuda_current_device: If True, call torch.cuda.set_device() to set the
            current PyTorch CUDA device to the one matching the local rank.
        overwrite: If True, overwrites already set variables. Else fails.
    """

    global _LOCAL_RANK, _LOCAL_WORLD_SIZE
    if _LOCAL_RANK >= 0 or _LOCAL_WORLD_SIZE >= 0:
        raise RuntimeError("Distributed mode has already been enabled")
    torch_env = _TorchDistributedEnvironment()
    torch_env.export(overwrite=overwrite)

    if set_cuda_current_device:
        torch.cuda.set_device(torch_env.local_rank)

    if allow_nccl_timeout:
        # This allows to use torch distributed timeout in a NCCL backend
        key, value = "NCCL_ASYNC_ERROR_HANDLING", "1"
        if not overwrite:
            _check_env_variable(key, value)
        os.environ[key] = value

    dist.init_process_group(backend="nccl")
    dist.barrier()

    # Finalize setup
    _LOCAL_RANK = torch_env.local_rank
    _LOCAL_WORLD_SIZE = torch_env.local_world_size
    _restrict_print_to_main_process()