"""Gradient averaging across ranks (gloo, CPU) for the multi-GPU trainer."""

from __future__ import annotations

import os

import pytest

torch = pytest.importorskip("torch")


def _worker(rank: int, world: int, port: int) -> None:
    import torch.distributed as dist

    from deg.train import _allreduce_grads

    os.environ.update(MASTER_ADDR="127.0.0.1", MASTER_PORT=str(port))
    dist.init_process_group("gloo", rank=rank, world_size=world)
    params = [torch.nn.Parameter(torch.zeros(n)) for n in (3, 5, 7)]
    params[0].grad = torch.full((3,), float(rank + 1))
    params[1].grad = torch.full((5,), float(10 * (rank + 1)))
    # params[2] has no grad on any rank -> zeros
    _allreduce_grads(params, world, bucket=4)  # several buckets
    assert torch.equal(params[0].grad, torch.full((3,), 1.5))
    assert torch.equal(params[1].grad, torch.full((5,), 15.0))
    assert torch.equal(params[2].grad, torch.zeros(7))
    dist.destroy_process_group()


def test_allreduce_grads_averages_across_ranks():
    import socket

    import torch.multiprocessing as mp

    with socket.socket() as s:
        s.bind(("127.0.0.1", 0))
        port = s.getsockname()[1]
    mp.spawn(_worker, args=(2, port), nprocs=2, join=True)