"""Two-rank CPU DDP training under torchrun.

Every rank calls lht.init; rank 0 creates the run and each rank joins as
its own participant named ``rank-<r>``. Launch with:

    LAUNCHHELM_URL=... LAUNCHHELM_API_KEY=... \
    LAUNCHHELM_ORGANIZATION=... LAUNCHHELM_PROJECT=... \
    torchrun --master_addr=127.0.0.1 --master_port=29500 \
    --nproc_per_node=2 torch_ddp.py
"""

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel

import launchhelm.torch as lht


def main() -> None:
    dist.init_process_group("gloo")
    rank = dist.get_rank()
    torch.manual_seed(0)

    model = DistributedDataParallel(
        torch.nn.Sequential(
            torch.nn.Linear(8, 16), torch.nn.ReLU(), torch.nn.Linear(16, 1)
        )
    )
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    x = torch.randn(64, 8)
    y = x.sum(dim=1, keepdim=True)

    tr = lht.init("quickstart", name="ddp-run", system_metrics=True)
    tr.watch(model, optimizer)

    for step in range(20):
        optimizer.zero_grad()
        loss = torch.nn.functional.mse_loss(model(x), y)
        loss.backward()
        optimizer.step()
        # step() is the collective control point: call it on every rank at
        # the same cadence.
        tr.step(step=step, loss=loss.item())

    tr.finish()
    dist.destroy_process_group()
    if rank == 0:
        print("run finished:", tr.run.run_id)


if __name__ == "__main__":
    main()