"""Audit 2026-10-14 (G4-rl-batch) regression for `align/ppo/trainer.py` (U05-2)."""

from __future__ import annotations

import torch
from test_ppo import _build_module, _make_rollout, _micro, _perturb, _ppo_cfg, tiny  # noqa: F401


def test_u05_2_a_skipped_micro_with_non_finite_logprobs_still_reaches_every_parameter(tiny, monkeypatch):
    """FSDP2 reduce-scatter fires from the parameters' grad hooks. A rank whose `new_lp` is
    non-finite used to skip on a DETACHED leaf: backward touched no parameter, the rank never
    joined the collective its finite-`new_lp` peers sat in, and the job hung until the watchdog.
    The skip must back-propagate through the model whatever `new_lp` holds."""
    import align.ppo.trainer as T

    cfg = _ppo_cfg()
    mod = _build_module(tiny, cfg)
    exp = mod.prepare_experience(_make_rollout(tiny, cfg))
    _perturb(mod.model, "policy", seed=5, scale=0.1)
    mod.set_train()
    mod.train_batch(_micro(mod, exp, iteration=1))  # a healthy micro first (arms the tripwire)
    for p in mod.model.parameters():
        p.grad = None

    real = T.policy_forward

    def nan_forward(*a, **k):
        new_lp, values = real(*a, **k)
        return new_lp * float("nan"), values  # graph-connected, non-finite

    monkeypatch.setattr(T, "policy_forward", nan_forward)
    out = mod.train_batch(_micro(mod, exp, iteration=1, rows=(2, 3)))
    assert out.metrics["ppo/skipped_micro"] == 1.0
    grads = [p.grad for p in mod.model.parameters() if p.requires_grad]
    assert grads and all(g is not None for g in grads), (
        "a non-finite-new_lp skip must still back-propagate through every adapter (FSDP2 lockstep)"
    )