#!/usr/bin/env python3
"""torchrun/NCCL 멀티 노드 bootstrap만 확인하는 최소 DDP probe."""

from __future__ import annotations

import argparse
import datetime as dt
import os
import socket
import sys
import time

import torch
import torch.distributed as dist


def _print(message: str) -> None:
    print(f"[ddp_probe][pid={os.getpid()}] {message}", flush=True)


def _env_line() -> str:
    keys = [
        "RANK",
        "LOCAL_RANK",
        "WORLD_SIZE",
        "MASTER_ADDR",
        "MASTER_PORT",
        "CUDA_VISIBLE_DEVICES",
        "NCCL_SOCKET_IFNAME",
    ]
    return " ".join(f"{key}={os.environ.get(key, '<unset>')}" for key in keys)


def main() -> None:
    parser = argparse.ArgumentParser(description="Minimal torch.distributed bootstrap probe.")
    parser.add_argument("--backend", default="", help="Distributed backend. Default: nccl when CUDA exists, else gloo.")
    parser.add_argument("--timeout-sec", type=int, default=300, help="init_process_group timeout in seconds.")
    parser.add_argument("--sleep-sec", type=int, default=0, help="Sleep after barrier for nvidia-smi inspection.")
    args = parser.parse_args()

    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    backend = args.backend or ("nccl" if torch.cuda.is_available() else "gloo")

    _print(f"host={socket.gethostname()} fqdn={socket.getfqdn()} python={sys.executable}")
    _print(f"env {_env_line()}")
    _print(
        f"torch={torch.__version__} cuda_available={torch.cuda.is_available()} "
        f"cuda_device_count={torch.cuda.device_count()} backend={backend}"
    )

    if torch.cuda.is_available():
        if local_rank >= torch.cuda.device_count():
            raise RuntimeError(
                f"LOCAL_RANK={local_rank} but torch.cuda.device_count()={torch.cuda.device_count()}"
            )
        torch.cuda.set_device(local_rank)
        _print(
            f"after set_device local_rank={local_rank} "
            f"current_device={torch.cuda.current_device()} name={torch.cuda.get_device_name(local_rank)}"
        )
    elif backend == "nccl":
        raise RuntimeError("NCCL backend requires CUDA, but torch.cuda.is_available() is false.")

    _print("before init_process_group")
    dist.init_process_group(
        backend=backend,
        init_method="env://",
        timeout=dt.timedelta(seconds=args.timeout_sec),
    )
    rank = dist.get_rank()
    world = dist.get_world_size()
    _print(f"after init_process_group rank={rank} world={world} backend={dist.get_backend()}")

    device = torch.device("cuda", local_rank) if backend == "nccl" else torch.device("cpu")
    value = torch.tensor([rank + 1], dtype=torch.float32, device=device)
    _print(f"before all_reduce value={value.item()} device={value.device}")
    dist.all_reduce(value, op=dist.ReduceOp.SUM)
    expected = world * (world + 1) / 2
    _print(f"after all_reduce value={value.item()} expected={expected}")

    _print("before barrier")
    dist.barrier()
    _print("after barrier")

    if args.sleep_sec > 0:
        _print(f"sleeping {args.sleep_sec}s before destroy_process_group")
        time.sleep(args.sleep_sec)

    dist.destroy_process_group()
    _print("destroyed process group; OK")


if __name__ == "__main__":
    main()