Blame for code/tests/test_audit_g4_rl_batch_source.py

"""Audit 2026-10-14 (G4-rl-batch) regressions for `align/rl/batch_source.py`.

U04-1  fewer live rows than world x accum slots killed the run (RuntimeError after the
       rollout was paid for);
U04-2  row quantisation threw away up to ~50% of live rows and toggled the step size 2x;
U04-3  peers waited for rank 0's rollout on a collective bounded by the process-group
       timeout, and a rank-0 failure left them blocked.
"""

from __future__ import annotations

import os
import tempfile
import time

import pytest
import torch
from test_grpo_module import _batch, _module
from test_grpo_driver import FakeClient, _source

from align.rl.batch_source import GrpoBatchSource, RolloutMicroHandle


def _payload(B: int, T: int = 5) -> dict[str, torch.Tensor]:
    """B group-contiguous rows; row i is identifiable by input_ids[i, 0] == i + 1."""
    ids = torch.zeros((B, T), dtype=torch.int64)
    ids[:, 0] = torch.arange(1, B + 1)
    rm = torch.zeros((B, T), dtype=torch.int64)
    rm[:, 2:] = 1
    return {
        "input_ids": ids,
        "attention_mask": torch.ones((B, T), dtype=torch.int64),
        "response_mask": rm,
        "advantages": torch.ones(B),
        "vllm_logprobs": torch.zeros((B, T - 1)),
        "group_ids": torch.arange(B) // 16,
        "reward::reward/total": torch.ones(B),
    }


def _bare(world: int, accum: int, rank: int = 0, prompts: int = 32, g: int = 16) -> GrpoBatchSource:
    src = GrpoBatchSource.__new__(GrpoBatchSource)
    src.world_size, src.grad_accum, src.rank = world, accum, rank
    src.prompts_per_step, src.n_samples_per_prompt = prompts, g
    src.dropped_remainder_rows = 0
    return src


def _real_rows(slices) -> list[int]:
    """Row ids carried by real (non-padding) rows of one rank's slices."""
    out = []
    for s in slices:
        for i in range(s["input_ids"].shape[0]):
            if int(s["response_mask"][i].sum()) > 0:
                out.append(int(s["input_ids"][i, 0]))
    return out


def _all_ranks(world, accum, B):
    per_rank = []
    for r in range(world):
        src = _bare(world, accum, rank=r)
        per_rank.append(src._slice_rank(_payload(B)))
    return per_rank


def test_u04_1_fewer_live_rows_than_slots_still_trains_every_row():
    """World 3 x accum 85 = 255 slots; 240 live rows (17 of 32 groups dead) used to raise."""
    per_rank = _all_ranks(3, 85, 240)
    for slices in per_rank:
        assert len(slices) == 85
        assert all(s["input_ids"].shape[0] >= 1 for s in slices), "a rank would run an empty micro"
    used = sorted(i for slices in per_rank for i in _real_rows(slices))
    assert used == list(range(1, 241)), "every live row must reach exactly one micro"


def test_u04_1_padding_rows_contribute_nothing_and_are_flagged():
    per_rank = _all_ranks(3, 85, 240)
    pad_rows = [
        s
        for slices in per_rank
        for s in slices
        if int(s["response_mask"].sum()) == 0
    ]
    assert len(pad_rows) == 255 - 240
    for s in pad_rows:
        assert float(s["advantages"].abs().sum()) == 0.0
        assert s["input_ids"].shape[0] == 1


def test_u04_2_one_dead_group_no_longer_discards_half_the_rows():
    """B=496 (one dead group of 32x16, 3 ranks, accum 85) used 255 rows and dropped 241."""
    per_rank = _all_ranks(3, 85, 496)
    used = sorted(i for slices in per_rank for i in _real_rows(slices))
    assert used == list(range(1, 497))
    src = _bare(3, 85)
    src._slice_rank(_payload(496))
    assert src.dropped_remainder_rows == 0


def test_u04_2_full_step_shape_is_unchanged():
    """The nominal full-live shape keeps its per-micro row cap (memory budget): 512 rows over 255
    slots is 2 rows per micro, 2 rows dropped and counted -- exactly the previous behaviour."""
    src = _bare(3, 85)
    slices = src._slice_rank(_payload(512))
    assert all(s["input_ids"].shape[0] == 2 for s in slices)
    assert src.dropped_remainder_rows == 2


def test_u04_2_micro_sizes_never_exceed_the_nominal_cap_and_ranks_stay_balanced():
    for B in (255, 300, 400, 496, 510, 512):
        per_rank = _all_ranks(3, 85, B)
        for m in range(85):
            sizes = [per_rank[r][m]["input_ids"].shape[0] for r in range(3)]
            assert max(sizes) <= 2 and max(sizes) - min(sizes) <= 1, (B, m, sizes)


def test_u04_1_pad_micro_runs_through_the_real_module_with_zero_loss_and_live_graph():
    mod, model = _module(grad_accum=1)
    src = _bare(1, 4, prompts=1, g=8)
    src.device = "cpu"
    src.old_logprob_fn = lambda ids, attn: mod.compute_logprobs_nograd(ids, attn)
    b = _batch(model, B=2, T=10)
    payload = {
        "input_ids": b.input_ids,
        "attention_mask": b.attention_mask,
        "response_mask": b.response_mask,
        "advantages": b.advantages,
        "vllm_logprobs": b.vllm_logprobs,
        "group_ids": b.group_ids,
    }
    slices = src._slice_rank(payload)  # 2 rows over 4 slots: two micros are padding
    pads = [s for s in slices if int(s["response_mask"].sum()) == 0]
    assert len(pads) == 2
    batch = src._finalize(model, pads[0])
    out = mod._grpo_loss_adapter(model, batch)
    assert float(out.loss) == 0.0
    out.loss.backward()
    assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in model.parameters())


def test_u04_1_tripwire_is_not_fed_from_a_padding_micro():
    """A padding micro reports entropy 0 / length 0; observing it would fake a collapse."""
    seen = []

    class _Wire:
        fired = False
        fire_reason = None

        def observe(self, step, entropy, length):
            seen.append((step, entropy, length))
            return False

    src = _source(g=8, accum=2, steps=1)
    src.collapse_tripwire = _Wire()
    src._on_step_boundary(0, {"entropy": 0.0, "lengths/response_mean": 0.0}, padded=True)
    assert seen == []
    src._on_step_boundary(0, {"entropy": 1.0, "lengths/response_mean": 3.0})
    assert seen == [(0, 1.0, 3.0)]


# ---------------------------------------------------------------------------- U04-3 (2 real ranks)


def _wait_worker(rank, world, init_file, mode, q):
    import datetime

    import torch.distributed as dist

    os.environ["ALIGN_PG_TIMEOUT_S"] = "3"
    os.environ["WORLD_SIZE"] = str(world)
    os.environ["RANK"] = str(rank)
    os.environ["ALIGN_ROLLOUT_WAIT_TIMEOUT_S"] = "60"
    dist.init_process_group(
        "gloo", init_method=f"file://{init_file}", rank=rank, world_size=world,
        timeout=datetime.timedelta(seconds=3),
    )

    class _Src(GrpoBatchSource):
        def _produce_step(self, optimizer_step):
            if mode == "slow":
                time.sleep(7)  # longer than the 3 s PG timeout, far shorter than the rollout budget
                return _payload(4)
            raise RuntimeError("rollout engine died (test)")

    src = _Src(
        [{"problem": "p", "answer": "1"}], prompts_per_step=1, n_samples_per_prompt=4,
        max_response_len=4, temperature=1.0, optimizer_steps=1, grad_accum=1, seed=0,
        prompt_builder=lambda p: p,
    )
    src.bind(client=None, reward=None, verifier=None, old_logprob_fn=lambda i, a: None,
             weights_dir_fn=lambda: "w", rank=rank, world_size=world)
    src.old_logprob_fn = lambda ids, attn: torch.zeros(ids.shape[0], ids.shape[1] - 1)
    t0 = time.time()
    try:
        batch = src._materialize_micro(None, RolloutMicroHandle(0, 0))
        q.put((rank, "ok", int(batch.input_ids.shape[0]), time.time() - t0))
    except Exception as e:  # noqa: BLE001
        q.put((rank, "err", str(e), time.time() - t0))


def _run_two_ranks(mode):
    import torch.multiprocessing as mp

    ctx = mp.get_context("spawn")
    q = ctx.Queue()
    with tempfile.TemporaryDirectory(dir="/apps/development/rlhf/align/.tmp/audit_2026-10-14/scratch_g4_rl_batch") as d:
        f = os.path.join(d, "store")
        ps = [ctx.Process(target=_wait_worker, args=(r, 2, f, mode, q)) for r in range(2)]
        for p in ps:
            p.start()
        res = {}
        for _ in range(2):
            r = q.get(timeout=90)
            res[r[0]] = r
        for p in ps:
            p.join(30)
    return res


def test_u04_3_peers_survive_a_rank0_rollout_longer_than_the_pg_timeout():
    res = _run_two_ranks("slow")
    assert res[0][1] == "ok" and res[1][1] == "ok", res
    assert res[1][2] == 2  # 4 rows over 2 ranks


def test_u04_3_a_rank0_failure_reaches_the_peers_as_a_named_error_not_a_timeout():
    res = _run_two_ranks("fail")
    assert res[0][1] == "err" and "rollout engine died (test)" in res[0][2], res
    assert res[1][1] == "err" and "rollout engine died (test)" in res[1][2] and "rank 0" in res[1][2], res
    assert res[1][3] < 30, "the peer must be told promptly, not wait out a timeout"