"""使用 torchrun 演示 DDP 训练与 torch.distributed.checkpoint 恢复。"""

from pathlib import Path
import argparse
import os
import sys

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))

from learning_labs.distributed import checkpoint_manifest


def main() -> None:
    parser = argparse.ArgumentParser(description="DDP + Distributed Checkpoint 最小实验")
    parser.add_argument("--steps", type=int, default=4)
    parser.add_argument("--checkpoint", type=Path, default=ROOT / "artifacts" / "distributed" / "dcp_checkpoint")
    parser.add_argument("--backend", choices=("gloo", "nccl"), default=None)
    args = parser.parse_args()

    try:
        import torch
        import torch.distributed as dist
        import torch.distributed.checkpoint as dcp
    except ModuleNotFoundError as error:
        raise SystemExit("需要 PyTorch。请使用 Python 3.11/3.12 并先安装 requirements.lock。") from error
    from torch.nn.parallel import DistributedDataParallel

    launched = "RANK" in os.environ
    backend = args.backend or ("nccl" if torch.cuda.is_available() else "gloo")
    if launched:
        dist.init_process_group(backend=backend)
    else:
        dist.init_process_group(backend=backend, init_method="tcp://127.0.0.1:29619", rank=0, world_size=1)
    rank, world_size = dist.get_rank(), dist.get_world_size()
    device = torch.device(f"cuda:{rank}" if backend == "nccl" else "cpu")
    try:
        torch.manual_seed(20260901 + rank)
        model = torch.nn.Linear(4, 1).to(device)
        ddp = DistributedDataParallel(model, device_ids=[rank] if backend == "nccl" else None)
        optimizer = torch.optim.SGD(ddp.parameters(), lr=0.08)
        for _ in range(args.steps):
            features = torch.randn(12, 4, device=device)
            target = features.sum(dim=1, keepdim=True)
            loss = torch.nn.functional.mse_loss(ddp(features), target)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        args.checkpoint.mkdir(parents=True, exist_ok=True)
        state = {"model": ddp.module.state_dict(), "optimizer": optimizer.state_dict()}
        dcp.save(state_dict=state, checkpoint_id=args.checkpoint)
        dist.barrier()
        restored = torch.nn.Linear(4, 1).to(device)
        restore_state = {"model": restored.state_dict(), "optimizer": optimizer.state_dict()}
        dcp.load(state_dict=restore_state, checkpoint_id=args.checkpoint)
        restored.load_state_dict(restore_state["model"])
        if rank == 0:
            manifest = checkpoint_manifest(world_size=world_size, backend=backend, steps=args.steps)
            (args.checkpoint / "manifest.json").write_text(__import__("json").dumps(manifest, indent=2), encoding="utf-8")
            print({"loss": round(float(loss), 6), "checkpoint": str(args.checkpoint), **manifest})
    finally:
        dist.destroy_process_group()


if __name__ == "__main__":
    main()