"""G0 control plane (single process): coordinated stop, rank-reduced HALT,
skip caps, process-group timeout + flight recorder env, checkpoint durability
and completeness, RNG + pin persistence.

The collective halves of these properties -- "a signal / HALT / NaN loss on ONE
rank acts on BOTH" -- need two ranks and live in tests/test_trainer_dist.py.
What is pinned here is everything a single process can see: the decisions
themselves, the on-disk protocol, and the helpers' contracts.

Run: .venv/bin/python -m pytest tests/test_control_plane.py -q -m "not gpu and not oracle" -p no:cacheprovider
"""

from __future__ import annotations

import datetime
import logging
import os
import random
import signal
import threading

import numpy as np
import pytest
import torch
from test_model_parity import tiny_config
from test_trainer import _tiny_trainer

import align.train.trainer as trainer_mod
import align.utils.dist as udist
from align.train.probes import TripwireResult
from align.train.trainer import HaltError, LossOutput, checkpoint_problem


@pytest.fixture
def tiny(tiny_model_constants):
    return tiny_config(tiny_model_constants)


def _cfg(**train_attrs):
    """cfg_hook that sets TrainConfig attributes the Trainer reads in __init__
    (they are read with getattr, so they need not exist on the dataclass)."""

    def hook(cfg):
        for k, v in train_attrs.items():
            setattr(cfg.train, k, v)

    return hook


def _signal_at(micro: int, signum=signal.SIGTERM, also_at=()):
    """wrap_adapter that delivers a REAL signal to this process on the
    `micro`-th adapter call (one adapter call == one micro-step)."""

    def wrap(inner):
        n = {"i": 0}

        def adapter(m, batch):
            n["i"] += 1
            if n["i"] == micro or n["i"] in also_at:
                os.kill(os.getpid(), signum)
            return inner(m, batch)

        return adapter

    return wrap


def _loss_transform(fn):
    """wrap_adapter that rewrites each micro's loss: fn(micro_index, loss)."""

    def wrap(inner):
        n = {"i": 0}

        def adapter(m, batch):
            n["i"] += 1
            out = inner(m, batch)
            return LossOutput(loss=fn(n["i"], out.loss), n_tokens=out.n_tokens, metrics=out.metrics)

        return adapter

    return wrap


def _weights(trainer) -> dict:
    return {k: v.detach().clone() for k, v in trainer.model.state_dict().items()}


def _same(a: dict, b: dict) -> bool:
    return a.keys() == b.keys() and all(torch.equal(a[k], b[k]) for k in a)


# ---------------------------------------------------------------------------
# 1. coordinated stop: signal + wall clock
# ---------------------------------------------------------------------------


def test_sigterm_mid_window_stops_at_the_boundary_and_resumes_bit_exact(tiny, tmp_workdir):
    """SIGTERM-MID-SAVE / MID-WINDOW (REFERENCE_GAP_ANALYSIS gap 1).

    The signal arrives at micro 3: the first micro of the second accumulation
    window (accum 2). The run must NOT stop there -- a checkpoint holds no
    gradients, so a mid-window save resumes into a different run. It stops at
    micro 4, the boundary, with stop_reason "signal", and writes a checkpoint
    that a fresh process resumes from to weights BIT-EQUAL to the run nobody
    interrupted. save_every is 8, so the micro-4 checkpoint can only exist
    because the stop wrote it.
    """
    ref, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "ref"), save_every_micro_steps=8)
    ref.train()
    want = _weights(ref)

    run = os.path.join(tmp_workdir, "run")
    t, _ = _tiny_trainer(
        tiny, run, save_every_micro_steps=8, wrap_adapter=_signal_at(3)
    )
    summary = t.train()
    assert summary["stop_reason"] == "signal", summary
    assert summary["micro_step"] == 4 and summary["optimizer_step"] == 2, summary
    assert summary["checkpoints_written"] == [4], summary
    assert os.listdir(os.path.join(run, "checkpoints")) == ["micro_00000004"]
    assert checkpoint_problem(os.path.join(run, "checkpoints", "micro_00000004")) is None

    t2, _ = _tiny_trainer(tiny, run, save_every_micro_steps=8)
    assert t2.resume_from_latest().endswith("micro_00000004")
    assert t2.micro_step == 4
    # vacuity guard: the stopped run really was behind the finished one
    assert not _same(_weights(t2), want)
    t2.train()
    assert _same(_weights(t2), want), "resume from a signal stop diverged from the uninterrupted run"


def test_signal_handlers_are_restored_after_train_and_after_an_exception(tiny, tmp_workdir):
    """Handlers live only for the duration of `train()`: a Trainer that leaked
    its SIGTERM handler would swallow the NEXT stage's signals (stage scripts
    run several trainers in one process)."""
    before = {s: signal.getsignal(s) for s in (signal.SIGTERM, signal.SIGINT)}
    t, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "a"))
    t.train()
    assert {s: signal.getsignal(s) for s in before} == before

    def boom(inner):
        def adapter(m, batch):
            raise ValueError("adapter failed")

        return adapter

    t2, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "b"), wrap_adapter=boom)
    with pytest.raises(ValueError, match="adapter failed"):
        t2.train()
    assert {s: signal.getsignal(s) for s in before} == before


def test_handler_is_installed_during_train_and_only_sets_a_flag(tiny, tmp_workdir):
    """While train() runs, SIGTERM is OURS (flag only) -- not the default
    action that kills the process wherever it is."""
    seen = {}

    def wrap(inner):
        def adapter(m, batch):
            seen.setdefault("handler", signal.getsignal(signal.SIGTERM))
            return inner(m, batch)

        return adapter

    t, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "a"), wrap_adapter=wrap)
    t.train()
    assert seen["handler"] not in (signal.SIG_DFL, signal.SIG_IGN, None)
    assert getattr(seen["handler"], "__self__", None) is t._stop_signals


def test_second_sigint_escapes_and_leaves_no_partial_checkpoint(tiny, tmp_workdir):
    """The first SIGINT is honoured at the boundary; a SECOND one raises
    KeyboardInterrupt so an operator is never locked out by a slow stop. The
    unwind must not leave a `.partial` that looks resumable."""
    run = os.path.join(tmp_workdir, "run")
    t, _ = _tiny_trainer(
        tiny, run, wrap_adapter=_signal_at(1, signal.SIGINT, also_at=(2,)), save_every_micro_steps=8
    )
    with pytest.raises(KeyboardInterrupt):
        t.train()
    assert not [d for d in os.listdir(os.path.join(run, "checkpoints")) if d.endswith(".partial")]
    assert signal.getsignal(signal.SIGINT) is signal.default_int_handler


def test_a_stop_request_wins_over_max_micro_steps(tiny, tmp_workdir):
    """A capped run normally writes NO end checkpoint (probes must not spend
    minutes gathering 14GB). A SIGTERM during one must still save: otherwise
    the window's work is discarded, which is the defect."""
    run = os.path.join(tmp_workdir, "run")
    t, _ = _tiny_trainer(
        tiny, run, save_every_micro_steps=8, max_micro_steps=4, wrap_adapter=_signal_at(3)
    )
    summary = t.train()
    assert summary["stop_reason"] == "signal" and summary["micro_step"] == 4
    assert os.listdir(os.path.join(run, "checkpoints")) == ["micro_00000004"]


def test_a_capped_run_without_a_signal_still_writes_no_end_checkpoint(tiny, tmp_workdir):
    """The other direction: the stop-save must not leak into ordinary capped
    runs (the probe-artifact rule in train())."""
    run = os.path.join(tmp_workdir, "run")
    t, _ = _tiny_trainer(tiny, run, save_every_micro_steps=8, max_micro_steps=4)
    assert t.train()["stop_reason"] == "max_micro_steps"
    assert os.listdir(os.path.join(run, "checkpoints")) == []


def test_wall_clock_limit_stops_at_the_first_boundary_and_saves(tiny, tmp_workdir):
    """`train.max_wall_seconds`: a 1ns limit has expired by the first boundary
    (micro 2); the stop is at that boundary, never at micro 1."""
    run = os.path.join(tmp_workdir, "run")
    t, _ = _tiny_trainer(
        tiny, run, save_every_micro_steps=8, cfg_hook=_cfg(max_wall_seconds=1e-9)
    )
    summary = t.train()
    assert summary["stop_reason"] == "wall_clock", summary
    assert summary["micro_step"] == 2 and summary["optimizer_step"] == 1, summary
    assert os.listdir(os.path.join(run, "checkpoints")) == ["micro_00000002"]


def test_wall_clock_generous_limit_and_default_do_not_stop_the_run(tiny, tmp_workdir):
    for name, hook in (("none", None), ("big", _cfg(max_wall_seconds=3600.0))):
        t, _ = _tiny_trainer(
            tiny, os.path.join(tmp_workdir, name), save_every_micro_steps=8, cfg_hook=hook
        )
        summary = t.train()
        assert summary["stop_reason"] == "data_exhausted" and summary["micro_step"] == 8, summary


@pytest.mark.parametrize("bad", [0, 0.0, -5])
def test_wall_clock_limit_must_be_positive(tiny, tmp_workdir, bad):
    with pytest.raises(ValueError, match="max_wall_seconds"):
        _tiny_trainer(tiny, os.path.join(tmp_workdir, "a"), cfg_hook=_cfg(max_wall_seconds=bad))


def test_train_off_the_main_thread_runs_without_handlers(tiny, tmp_workdir, caplog):
    """signal.signal() raises ValueError off the main thread; a harness that
    runs the trainer in a worker thread must degrade (loudly), not crash."""
    t, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "a"))
    box = {}
    with caplog.at_level(logging.WARNING, logger="align.train.trainer"):
        th = threading.Thread(target=lambda: box.update(summary=t.train()))
        th.start()
        th.join(120)
    assert box["summary"]["micro_step"] == 8
    assert "cannot be installed" in caplog.text
    assert t._stop_signals.installed is False


# ---------------------------------------------------------------------------
# 2. rank-reduced HALT (the multi-rank half is in test_trainer_dist.py)
# ---------------------------------------------------------------------------


class _Halting:
    HALT = True
    INPUT_KEYS: tuple = ()

    def __init__(self, at):
        self.at, self.n = at, 0

    def update(self, _x):
        self.n += 1
        return TripwireResult(self.n >= self.at, "stop now", value=3.5)


def test_halt_tripwire_raises_haltError_with_the_original_message(tiny, tmp_workdir):
    """Single process: the HALT is raised exactly as before (RuntimeError, same
    text) -- the rank-reduction must not change what an operator reads."""
    t, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "a"))
    t.tripwires = {"grad_norm": _Halting(at=3)}
    with pytest.raises(RuntimeError, match=r"HALT tripwire 'grad_norm' fired at micro_step 3: stop now \(value=3.5\)") as ei:
        t.train()
    assert isinstance(ei.value, HaltError)
    assert t.tripwire_alarms[-1]["alarm"] == "grad_norm"


def test_agree_or_raise_peer_failure_makes_a_healthy_rank_raise(tiny, tmp_workdir, monkeypatch):
    """The reduction itself, simulated: THIS rank is healthy, rank 1 reports a
    failure. Without the reduction this rank would carry on into the next
    collective alone."""
    t, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "a"))
    monkeypatch.setattr(trainer_mod, "_dist_world_size", lambda: 2)
    monkeypatch.setattr(udist, "any_rank", lambda flag: True)  # a peer said yes
    monkeypatch.setattr(
        trainer_mod, "all_gather_objects", lambda obj: [None, ("HaltError: boom on rank 1", 80)]
    )
    with pytest.raises(HaltError, match=r"\[rank 1\] HaltError: boom on rank 1.*rank-reduced during x") as ei:
        t._agree_or_raise(None, "x")
    assert type(ei.value) is HaltError and ei.value.exit_code == 80


def test_agree_or_raise_keeps_the_original_exception_on_the_failing_rank(tiny, tmp_workdir, monkeypatch):
    """A non-HALT failure (a KeyError from a missing metric) is re-raised AS
    ITSELF where it happened -- the traceback an engineer needs -- while peers
    get the HaltError summary."""
    t, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "a"))
    monkeypatch.setattr(trainer_mod, "_dist_world_size", lambda: 2)
    monkeypatch.setattr(udist, "any_rank", lambda flag: True)
    err = KeyError("rewards/margins")
    monkeypatch.setattr(trainer_mod, "all_gather_objects", lambda obj: [obj, None])
    with pytest.raises(KeyError):
        t._agree_or_raise(err, "tripwire check")


def test_agree_or_raise_is_silent_when_no_rank_failed(tiny, tmp_workdir, monkeypatch):
    t, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "a"))
    monkeypatch.setattr(trainer_mod, "_dist_world_size", lambda: 2)
    monkeypatch.setattr(udist, "any_rank", lambda flag: False)
    t._agree_or_raise(None, "x")  # must not raise


def test_a_missing_tripwire_input_still_fails_loudly(tiny, tmp_workdir):
    """The 'KeyError surfaces' contract of _check_tripwires survives the
    reduction wrapper (single process re-raises the original)."""
    t, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "a"), armed=("logps_chosen",))
    t.tripwires["logps_chosen"].INPUT_KEYS = ()  # let the arm-time test pass
    with pytest.raises(KeyError):
        t.train()


# ---------------------------------------------------------------------------
# 3. skip caps
# ---------------------------------------------------------------------------

NAN = float("nan")


def _nan_grads_for(micros):
    """Loss * NaN on the listed micros: NaN loss AND NaN grads (the grad gate
    skips the step)."""
    return _loss_transform(lambda i, loss: loss * NAN if i in micros else loss)


def test_consecutive_skipped_steps_beyond_the_cap_halt_the_run(tiny, tmp_workdir):
    """THE SILENT-SKIP RUN. Micros 1-8 are four consecutive NaN windows. The
    default cap is 3 CONSECUTIVE, so the fourth is the one that halts -- at its
    boundary (micro 8), with a message naming the knob."""
    t, _ = _tiny_trainer(
        tiny, os.path.join(tmp_workdir, "a"), n_batches=12, wrap_adapter=_nan_grads_for(range(1, 9))
    )
    with pytest.raises(HaltError, match=r"4 consecutive bad optimizer steps exceed train\.max_consecutive_skipped_steps=3"):
        t.train()
    assert t.micro_step == 8
    assert t.train_module.skipped_steps == 4, "vacuous: the grad gate did not skip"
    assert t.tripwire_alarms[-1]["alarm"] == "skip_cap"


def test_exactly_the_cap_is_allowed_and_a_good_step_resets_the_streak(tiny, tmp_workdir):
    """3 bad windows (== cap) then a good one then 3 more: never more than 3 in
    a row, so no halt. Guards the off-by-one in both directions."""
    bad = set(range(1, 7)) | set(range(9, 15))  # windows 1-3 and 5-7 bad, 4 good
    t, _ = _tiny_trainer(
        tiny, os.path.join(tmp_workdir, "a"), n_batches=10, epochs=2, wrap_adapter=_nan_grads_for(bad),
        save_every_micro_steps=100, max_micro_steps=16,
    )
    t.train()
    assert t.skip_stats["bad_total"] == 6 and t.skip_stats["consecutive"] == 0
    assert t.skip_stats["boundaries"] == 8


def test_non_finite_loss_counts_even_when_the_gradients_are_finite(tiny, tmp_workdir):
    """'Non-finite loss must also count, not only grads.' `loss + nan` has a
    NaN value and an UNCHANGED, finite gradient: the optimizer applies every
    step (skipped_steps stays 0), so only the loss check can see it."""
    t, _ = _tiny_trainer(
        tiny,
        os.path.join(tmp_workdir, "a"),
        n_batches=12,
        wrap_adapter=_loss_transform(lambda i, loss: loss + NAN),
    )
    with pytest.raises(HaltError, match="non-finite loss"):
        t.train()
    assert t.train_module.skipped_steps == 0, "vacuous: the grad gate fired, not the loss check"
    assert t.optimizer_step == 4 and t.micro_step == 8


def test_one_non_finite_micro_in_a_window_marks_the_whole_step_bad(tiny, tmp_workdir):
    t, _ = _tiny_trainer(
        tiny,
        os.path.join(tmp_workdir, "a"),
        wrap_adapter=_loss_transform(lambda i, loss: loss + NAN if i == 3 else loss),
        save_every_micro_steps=100,
        max_micro_steps=4,
    )
    t.train()
    assert t.skip_stats == {"consecutive": 1, "bad_total": 1, "boundaries": 2}
    assert t.train_module.skipped_steps == 0


def _health(tiny, tmp_workdir, name, **cfg_attrs):
    t, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, name), cfg_hook=_cfg(**cfg_attrs))
    return t


def test_fractional_cap_waits_for_100_steps_then_halts_above_5_percent(tiny, tmp_workdir):
    """>5% bad, evaluated only once 100 boundaries exist. Driven through
    `_record_step_health` (the decision) with a pattern that is never more than
    one bad step in a row, so the consecutive cap cannot be what fires."""
    t = _health(tiny, tmp_workdir, "frac")
    # 3 bad of the first 10 = 30%: far above 5%, but under the minimum sample
    for i in range(10):
        t._record_step_health(False, i in (1, 5, 8))
    assert t.skip_stats["bad_total"] == 3  # no halt yet
    # 90 more good steps: boundary 100, 3/100 = 3% -> still fine
    for _ in range(89):
        t._record_step_health(False, False)
    t._record_step_health(False, False)
    assert t.skip_stats["boundaries"] == 100
    # push to 6/106 > 5%
    with pytest.raises(HaltError, match=r"train\.max_skipped_fraction=0\.05"):
        for i in range(10):
            t._record_step_health(False, i % 2 == 0)


def test_fractional_cap_exactly_at_the_threshold_does_not_halt(tiny, tmp_workdir):
    t = _health(tiny, tmp_workdir, "edge")
    for i in range(100):  # 5 of 100 = 5.0%: not ABOVE the cap
        t._record_step_health(False, i % 20 == 3)
    assert t.skip_stats == {"consecutive": 0, "bad_total": 5, "boundaries": 100}


def test_fractional_cap_halts_at_the_hundredth_step_when_already_above(tiny, tmp_workdir):
    t = _health(tiny, tmp_workdir, "fire")
    with pytest.raises(HaltError, match="10 of 100 optimizer steps were bad"):
        for i in range(100):
            t._record_step_health(False, i % 10 == 0)
    assert t.skip_stats["boundaries"] == 100


def test_skip_caps_can_be_overridden_or_disabled_by_config(tiny, tmp_workdir):
    t = _health(tiny, tmp_workdir, "off", max_consecutive_skipped_steps=None, max_skipped_fraction=None)
    for _ in range(300):
        t._record_step_health(True, True)  # never halts
    t2 = _health(tiny, tmp_workdir, "tight", max_consecutive_skipped_steps=0)
    with pytest.raises(HaltError, match=r"1 consecutive bad optimizer steps exceed .*=0"):
        t2._record_step_health(False, True)
    t3 = _health(tiny, tmp_workdir, "default")
    assert (t3.max_consecutive_skipped, t3.max_skipped_fraction) == (3, 0.05)


@pytest.mark.parametrize(
    "attrs",
    [{"max_consecutive_skipped_steps": -1}, {"max_skipped_fraction": 1.5}, {"max_skipped_fraction": -0.1}],
)
def test_skip_cap_config_is_validated(tiny, tmp_workdir, attrs):
    with pytest.raises(ValueError):
        _health(tiny, tmp_workdir, "bad", **attrs)


# ---------------------------------------------------------------------------
# 4. process-group timeout + NCCL flight recorder
# ---------------------------------------------------------------------------


def test_pg_timeout_default_and_env_override():
    assert udist.pg_timeout_s({}) == 1800.0
    assert udist.pg_timeout_s({"ALIGN_PG_TIMEOUT_S": "90"}) == 90.0
    assert udist.pg_timeout_s({"ALIGN_PG_TIMEOUT_S": " "}) == 1800.0


@pytest.mark.parametrize("raw", ["30m", "0", "-1", "inf", "nan", "abc"])
def test_pg_timeout_rejects_values_that_would_silently_disarm_the_watchdog(raw):
    with pytest.raises(ValueError, match="ALIGN_PG_TIMEOUT_S"):
        udist.pg_timeout_s({"ALIGN_PG_TIMEOUT_S": raw})


def _isolate_env(monkeypatch, *names):
    """Make monkeypatch restore these env vars (to absent) at teardown, so what
    `init_distributed` exports does not leak into later tests."""
    for n in names:
        monkeypatch.setenv(n, "placeholder")
        monkeypatch.delenv(n)


FR_VARS = (
    "TORCH_FR_BUFFER_SIZE",
    "TORCH_NCCL_TRACE_BUFFER_SIZE",
    "TORCH_NCCL_DUMP_ON_TIMEOUT",
    "TORCH_FR_DUMP_TEMP_FILE",
    "TORCH_NCCL_DEBUG_INFO_TEMP_FILE",
)


def test_init_distributed_passes_an_explicit_timeout(monkeypatch):
    import torch.distributed as dist

    calls = []
    monkeypatch.setenv("WORLD_SIZE", "2")
    monkeypatch.setenv("ALIGN_PG_TIMEOUT_S", "77")
    monkeypatch.setattr(dist, "is_initialized", lambda: False)
    monkeypatch.setattr(dist, "init_process_group", lambda **kw: calls.append(kw))
    udist.init_distributed(backend="gloo")
    assert calls == [{"backend": "gloo", "timeout": datetime.timedelta(seconds=77)}]
    # explicit argument beats the environment
    calls.clear()
    udist.init_distributed(backend="gloo", timeout_s=5)
    assert calls[0]["timeout"] == datetime.timedelta(seconds=5)
    # the default (no env) is 30 minutes, not torch's backend-dependent one
    calls.clear()
    monkeypatch.delenv("ALIGN_PG_TIMEOUT_S")
    udist.init_distributed(backend="gloo")
    assert calls[0]["timeout"] == datetime.timedelta(minutes=30)


def test_init_distributed_nccl_arms_the_flight_recorder_before_init(monkeypatch, tmp_workdir):
    import torch.distributed as dist

    _isolate_env(monkeypatch, *FR_VARS)
    order = []
    monkeypatch.setenv("WORLD_SIZE", "2")
    monkeypatch.setattr(dist, "is_initialized", lambda: False)
    monkeypatch.setattr(
        dist,
        "init_process_group",
        lambda **kw: order.append(("init", os.environ.get("TORCH_NCCL_DUMP_ON_TIMEOUT"))),
    )
    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
    udist.init_distributed(backend="nccl", run_dir=tmp_workdir)
    # the env was already set when init_process_group ran
    assert order == [("init", "1")]


def test_gloo_init_does_not_touch_the_nccl_environment(monkeypatch):
    import torch.distributed as dist

    _isolate_env(monkeypatch, *FR_VARS)
    monkeypatch.setenv("WORLD_SIZE", "2")
    monkeypatch.setattr(dist, "is_initialized", lambda: False)
    monkeypatch.setattr(dist, "init_process_group", lambda **kw: None)
    udist.init_distributed(backend="gloo")
    assert not [v for v in FR_VARS if v in os.environ]


def test_flight_recorder_env_defaults_live_under_the_run_dir_and_never_override(tmp_workdir):
    env: dict = {}
    applied = udist.configure_nccl_diagnostics(tmp_workdir, environ=env, torch_version="2.11.0+cu128")
    assert env["TORCH_FR_BUFFER_SIZE"] == "20000"
    assert "TORCH_NCCL_TRACE_BUFFER_SIZE" not in env
    assert env["TORCH_NCCL_DUMP_ON_TIMEOUT"] == "1"
    prefix = env["TORCH_FR_DUMP_TEMP_FILE"]
    assert prefix.startswith(os.path.join(tmp_workdir, "nccl_trace") + os.sep), prefix
    assert env["TORCH_NCCL_DEBUG_INFO_TEMP_FILE"] == prefix
    assert os.path.isdir(os.path.dirname(prefix)), "the dump directory must exist before a timeout"
    assert not prefix.startswith("/tmp"), "the torch default (/tmp/nccl_trace_rank_) is forbidden here"
    assert set(applied) == set(FR_VARS) - {"TORCH_NCCL_TRACE_BUFFER_SIZE"}

    keep = {"TORCH_FR_BUFFER_SIZE": "5", "TORCH_NCCL_DUMP_ON_TIMEOUT": "0"}
    env2 = dict(keep)
    udist.configure_nccl_diagnostics(tmp_workdir, environ=env2, torch_version="2.11.0")
    assert {k: env2[k] for k in keep} == keep, "an operator's explicit export must win"


def test_flight_recorder_variable_names_exist_in_the_installed_torch():
    """The names are not guesses (CLAUDE.md: verify against the installed
    torch). If a torch upgrade renames one, this fails instead of the recorder
    going silently dark."""
    import mmap

    root = os.path.dirname(torch.__file__)
    hpp = os.path.join(root, "include/torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp")
    lib = os.path.join(root, "lib/libtorch_cpu.so")
    if not (os.path.isfile(hpp) and os.path.isfile(lib)):
        pytest.skip("torch headers / libtorch_cpu.so not shipped in this install")
    header = open(hpp).read()
    for name in ("TORCH_NCCL_TRACE_BUFFER_SIZE", "TORCH_NCCL_DUMP_ON_TIMEOUT"):
        assert name in header, name
    with open(lib, "rb") as f, mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) as mm:
        for name in ("TORCH_FR_DUMP_TEMP_FILE", "TORCH_NCCL_DEBUG_INFO_TEMP_FILE"):
            assert mm.find(name.encode()) != -1, name


# ---------------------------------------------------------------------------
# 5. checkpoint durability + completeness
# ---------------------------------------------------------------------------


def _two_checkpoints(tiny, run_dir, **kw):
    t, _ = _tiny_trainer(tiny, run_dir, save_every_micro_steps=4, **kw)
    t.train()
    root = os.path.join(run_dir, "checkpoints")
    assert sorted(os.listdir(root)) == ["micro_00000004", "micro_00000008"]
    return t, root


def test_every_file_and_directory_is_fsynced_before_the_rename(tiny, tmp_workdir, monkeypatch):
    """THE ATOMIC-RENAME-IS-NOT-DURABLE DEFECT. Recorded as a timeline: every
    file in the checkpoint, and the staging directory, are fsynced BEFORE
    `os.rename(.partial -> final)`; the checkpoint root is fsynced AFTER it so
    the new directory entry itself survives a power loss."""
    events = []
    real_fsync, real_rename = os.fsync, os.rename

    def spy_fsync(fd):
        events.append(("fsync", os.readlink(f"/proc/self/fd/{fd}")))
        return real_fsync(fd)

    def spy_rename(src, dst, *a, **k):
        events.append(("rename", src, dst))
        return real_rename(src, dst, *a, **k)

    run = os.path.join(tmp_workdir, "run")
    t, _ = _tiny_trainer(tiny, run, save_every_micro_steps=8, max_micro_steps=8)
    monkeypatch.setattr(os, "fsync", spy_fsync)
    monkeypatch.setattr(os, "rename", spy_rename)
    t.save_checkpoint()
    monkeypatch.undo()

    root = os.path.realpath(os.path.join(run, "checkpoints"))
    staging = os.path.join(root, "micro_00000000.partial")
    renames = [i for i, e in enumerate(events) if e[0] == "rename" and e[1].endswith(".partial")]
    assert len(renames) == 1, events
    r = renames[0]
    before = {os.path.relpath(e[1], staging) for e in events[:r] if e[0] == "fsync" and e[1].startswith(staging)}
    final_files = set(os.listdir(os.path.join(root, "micro_00000000")))
    assert {"model.safetensors", "resume.pt", "config.json"} <= final_files
    assert final_files <= before, f"not fsynced before the rename: {sorted(final_files - before)}"
    assert "." in before, "the staging DIRECTORY entry must be fsynced before the rename"
    assert ("fsync", root) in events[r + 1 :], "the checkpoint root must be fsynced after the rename"


def test_a_complete_checkpoint_passes_and_the_newest_is_returned(tiny, tmp_workdir):
    t, root = _two_checkpoints(tiny, os.path.join(tmp_workdir, "run"))
    for d in os.listdir(root):
        assert checkpoint_problem(os.path.join(root, d)) is None
    assert t.latest_checkpoint() == os.path.join(root, "micro_00000008")


def _truncate_half(path):
    size = os.path.getsize(path)
    with open(path, "r+b") as f:
        f.truncate(size // 2)


MUTATIONS = {
    "weights_truncated": lambda d: _truncate_half(os.path.join(d, "model.safetensors")),
    "weights_zero_length": lambda d: open(os.path.join(d, "model.safetensors"), "wb").close(),
    "weights_missing": lambda d: os.remove(os.path.join(d, "model.safetensors")),
    "resume_missing": lambda d: os.remove(os.path.join(d, "resume.pt")),
    "resume_truncated": lambda d: _truncate_half(os.path.join(d, "resume.pt")),
    "resume_empty": lambda d: open(os.path.join(d, "resume.pt"), "wb").close(),
}


@pytest.mark.parametrize("mutation", sorted(MUTATIONS))
def test_incomplete_newest_checkpoint_is_skipped_with_a_warning_and_kept(
    tiny, tmp_workdir, mutation, caplog
):
    """A host seizure can leave a `micro_N` directory whose rename survived and
    whose data did not. `latest_checkpoint` must hand back the previous good one,
    warn, and NOT delete the evidence."""
    run = os.path.join(tmp_workdir, "run")
    t, root = _two_checkpoints(tiny, run)
    bad = os.path.join(root, "micro_00000008")
    MUTATIONS[mutation](bad)
    assert checkpoint_problem(bad) is not None

    with caplog.at_level(logging.WARNING, logger="align.train.trainer"):
        got = t.latest_checkpoint()
    assert got == os.path.join(root, "micro_00000004"), got
    assert "INCOMPLETE checkpoint micro_00000008" in caplog.text
    assert os.path.isdir(bad), "an incomplete checkpoint must be skipped, never deleted"

    # and a resume actually takes the good one
    t2, _ = _tiny_trainer(tiny, run, save_every_micro_steps=4)
    assert t2.resume_from_latest() == got and t2.micro_step == 4


def test_all_checkpoints_incomplete_means_no_checkpoint(tiny, tmp_workdir):
    t, root = _two_checkpoints(tiny, os.path.join(tmp_workdir, "run"))
    for d in os.listdir(root):
        os.remove(os.path.join(root, d, "resume.pt"))
    assert t.latest_checkpoint() is None


def test_sharded_weights_are_checked_against_the_index(tiny, tmp_workdir):
    """A multi-shard layout (`model.safetensors.index.json`): every shard it
    names must exist and be complete."""
    import json

    _t, root = _two_checkpoints(tiny, os.path.join(tmp_workdir, "run"))
    d = os.path.join(root, "micro_00000008")
    index = os.path.join(d, "model.safetensors.index.json")
    with open(index, "w") as f:
        json.dump({"weight_map": {"a": "model.safetensors"}}, f)
    assert checkpoint_problem(d) is None
    with open(index, "w") as f:
        json.dump({"weight_map": {"a": "model.safetensors", "b": "model-00002-of-00002.safetensors"}}, f)
    assert "missing weights file model-00002-of-00002.safetensors" in checkpoint_problem(d)
    with open(index, "w") as f:
        f.write("{not json")
    assert "index.json" in checkpoint_problem(d)


def test_per_rank_optimizer_files_must_all_be_present(tiny, tmp_workdir):
    """When the manifest says per_rank_files, the newest checkpoint is complete
    only if EVERY saved rank's optimizer file is: a missing one makes a resume
    die on that rank alone, mid-load."""
    run = os.path.join(tmp_workdir, "run")
    t, root = _two_checkpoints(tiny, run, cfg_hook=_cfg(per_rank_optimizer_state=True))
    d = os.path.join(root, "micro_00000008")
    assert checkpoint_problem(d) is None
    assert os.path.isfile(os.path.join(d, "optimizer_rank0.pt"))

    # the manifest names two ranks but only rank 0's file exists
    state = torch.load(os.path.join(d, "resume.pt"), weights_only=True, map_location="cpu")
    state["optimizer_state"] = {"layout": "per_rank_files", "world_size": 2}
    torch.save(state, os.path.join(d, "resume.pt"))
    assert "optimizer_rank1.pt" in checkpoint_problem(d)

    state["optimizer_state"] = {"layout": "per_rank_files", "world_size": 1}
    torch.save(state, os.path.join(d, "resume.pt"))
    os.remove(os.path.join(d, "optimizer_rank0.pt"))
    assert "optimizer_rank0.pt" in checkpoint_problem(d)
    assert t.latest_checkpoint() == os.path.join(root, "micro_00000004")

    _truncate_half(os.path.join(root, "micro_00000004", "optimizer_rank0.pt"))
    assert t.latest_checkpoint() is None, "a cut optimizer file must disqualify its checkpoint"


def test_partial_directories_are_still_discarded_at_startup_and_never_listed(tiny, tmp_workdir):
    run = os.path.join(tmp_workdir, "run")
    os.makedirs(os.path.join(run, "checkpoints", "micro_00000100.partial"))
    t, _ = _tiny_trainer(tiny, run)
    assert os.listdir(os.path.join(run, "checkpoints")) == []
    assert t.latest_checkpoint() is None


# ---------------------------------------------------------------------------
# 6. RNG state + pinned checkpoints + skip stats persisted in resume.pt
# ---------------------------------------------------------------------------


def _draw():
    return [float(torch.rand(1)), random.random(), float(np.random.rand()), int(np.random.randint(1 << 30))]


def test_rng_state_is_restored_for_torch_python_and_numpy(tiny, tmp_workdir):
    """A resumed process re-seeds from scratch and replays the original run's
    first draws; with the state in resume.pt it CONTINUES the stream the saving
    process was on."""
    run = os.path.join(tmp_workdir, "run")
    t, _ = _tiny_trainer(tiny, run, save_every_micro_steps=4, max_micro_steps=4)
    t.train()  # the cadence writes micro 4
    want = _draw()

    udist.seed_everything(999)  # a different stream, as a fresh process would have
    assert _draw() != want, "vacuous: the reseed reproduced the saved stream"
    t2, _ = _tiny_trainer(tiny, run, save_every_micro_steps=4)
    t2.resume_from_latest()
    assert _draw() == want


def test_resume_payload_holds_rng_pins_and_skip_stats_and_is_weights_only_loadable(tiny, tmp_workdir):
    run = os.path.join(tmp_workdir, "run")
    t, _ = _tiny_trainer(tiny, run, save_every_micro_steps=4, max_micro_steps=4)
    t._pin_next_save = True
    t.train()
    state = torch.load(
        os.path.join(run, "checkpoints", "micro_00000004", "resume.pt"),
        weights_only=True,
        map_location="cpu",
    )
    assert state["pinned_checkpoints"] == [4]
    assert state["skip_stats"] == {"consecutive": 0, "bad_total": 0, "boundaries": 2}
    (rng,) = state["rng_per_rank"]
    assert {"torch", "python", "numpy"} <= rng.keys()
    assert rng["torch"].dtype == torch.uint8


def test_pinned_checkpoint_survives_a_resume_and_the_following_prune(tiny, tmp_workdir):
    """R5 PERSISTENCE. Before: the pin set lived in process memory. A resumed
    run registered the resumed checkpoint as ITS OWN (so keep-N could prune it)
    and had forgotten it was pinned -- the first save after the resume deleted
    the one checkpoint it had been told to keep. keep_checkpoints=1 makes the
    prune immediate."""
    run = os.path.join(tmp_workdir, "run")
    a, _ = _tiny_trainer(
        tiny, run, n_batches=10, epochs=2, save_every_micro_steps=4, keep_checkpoints=1, max_micro_steps=4
    )
    a._pin_next_save = True
    a.train()
    root = os.path.join(run, "checkpoints")
    assert os.listdir(root) == ["micro_00000004"]

    b, _ = _tiny_trainer(
        tiny, run, n_batches=10, epochs=2, save_every_micro_steps=4, keep_checkpoints=1, max_micro_steps=16
    )
    b.resume_from_latest()
    assert b.pinned_checkpoints == {4}
    b.train()
    assert sorted(os.listdir(root)) == ["micro_00000004", "micro_00000016"], os.listdir(root)


def test_resume_adopts_older_checkpoints_so_keep_n_does_not_leak_one_per_resume(tiny, tmp_workdir):
    """TRACK M (2026-10-12, OLMoE 4-GPU SIGTERM/resume): keep_checkpoints=2 left micro_40, _44 and _80 on disk
    (83GB, disk 95% full). The interrupted attempt wrote 40 and 44; the new process registered only 44, so after
    saving 80 its list was [44, 80] and 40 was unprunable forever -- one leaked checkpoint per resume. Directories
    numbered BELOW the resumed one are this lineage's older saves and are adopted. A HIGHER-numbered directory is a
    stale-run zombie (the Arm A rerun incident that made pruning per-process) and must survive untouched.
    MUTANT (resume registers only the resumed step): micro_00000004 is still on disk at the end."""
    run = os.path.join(tmp_workdir, "run")
    a, _ = _tiny_trainer(
        tiny, run, n_batches=10, epochs=2, save_every_micro_steps=4, keep_checkpoints=3, max_micro_steps=8
    )
    a.train()
    root = os.path.join(run, "checkpoints")
    assert sorted(os.listdir(root)) == ["micro_00000004", "micro_00000008"]
    os.makedirs(os.path.join(root, "micro_00000999"))  # a zombie from an unrelated older run: incomplete, higher-numbered

    b, _ = _tiny_trainer(
        tiny, run, n_batches=10, epochs=2, save_every_micro_steps=4, keep_checkpoints=2, max_micro_steps=12
    )
    assert b.resume_from_latest().endswith("micro_00000008")
    assert b.checkpoints_written == [4, 8], "older save adopted, zombie not, and the resumed one stays LAST (= newest)"
    b.train()
    assert sorted(os.listdir(root)) == ["micro_00000008", "micro_00000012", "micro_00000999"], os.listdir(root)


def test_skip_stats_survive_a_resume(tiny, tmp_workdir):
    """A resume must not forgive a run that was already skipping."""
    run = os.path.join(tmp_workdir, "run")
    a, _ = _tiny_trainer(
        tiny, run, save_every_micro_steps=4, max_micro_steps=4, wrap_adapter=_nan_grads_for({1, 2, 3, 4})
    )
    a.train()
    assert a.skip_stats == {"consecutive": 2, "bad_total": 2, "boundaries": 2}
    b, _ = _tiny_trainer(tiny, run, save_every_micro_steps=4)
    b.resume_from_latest()
    assert b.skip_stats == a.skip_stats


def test_a_checkpoint_from_before_g0_still_resumes(tiny, tmp_workdir, caplog):
    """Back-compat: strip the new keys from resume.pt (what every existing
    checkpoint looks like). Resume must work, warn about the RNG, and leave the
    defaults in place -- not KeyError."""
    run = os.path.join(tmp_workdir, "run")
    a, _ = _tiny_trainer(tiny, run, save_every_micro_steps=4, max_micro_steps=4)
    a.train()
    path = os.path.join(run, "checkpoints", "micro_00000004", "resume.pt")
    state = torch.load(path, weights_only=True, map_location="cpu")
    for k in ("rng_per_rank", "pinned_checkpoints", "skip_stats"):
        del state[k]
    torch.save(state, path)
    # a pre-G0 checkpoint also predates the integrity sidecar (audit U01-3)
    os.remove(os.path.join(os.path.dirname(path), "integrity.json"))
    b, _ = _tiny_trainer(tiny, run, save_every_micro_steps=4)
    with caplog.at_level(logging.WARNING, logger="align.train.trainer"):
        b.resume_from_latest()
    assert "predates RNG persistence" in caplog.text
    assert b.micro_step == 4 and b.pinned_checkpoints == set()
    assert b.skip_stats == {"consecutive": 0, "bad_total": 0, "boundaries": 0}


def test_rng_capture_roundtrips_exactly_and_is_weights_only_clean(tmp_workdir):
    udist.seed_everything(3)
    sd = udist.capture_rng_state()
    want = _draw()
    path = os.path.join(tmp_workdir, "rng.pt")
    torch.save(sd, path)
    loaded = torch.load(path, weights_only=True, map_location="cpu")
    udist.seed_everything(4)
    udist.restore_rng_state(loaded)
    assert _draw() == want


def test_cuda_rng_is_captured_for_the_current_device_only(monkeypatch):
    """Faked CUDA (CPU-only host rules): only the CURRENT device's state is
    captured -- `get_rng_state_all` would create a context on every visible GPU
    from every rank -- and only when CUDA is already initialised."""
    calls = []
    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
    monkeypatch.setattr(torch.cuda, "is_initialized", lambda: True)
    monkeypatch.setattr(torch.cuda, "current_device", lambda: 2)
    monkeypatch.setattr(torch.cuda, "device_count", lambda: 4)
    monkeypatch.setattr(torch.cuda, "get_rng_state", lambda d: torch.tensor([7, 8, 9], dtype=torch.uint8))
    monkeypatch.setattr(torch.cuda, "set_rng_state", lambda s, device: calls.append((s.tolist(), device)))
    monkeypatch.setattr(
        torch.cuda, "get_rng_state_all", lambda: (_ for _ in ()).throw(AssertionError("all devices"))
    )
    sd = udist.capture_rng_state()
    assert list(sd["cuda"]) == ["2"] and sd["cuda"]["2"].tolist() == [7, 8, 9]
    udist.restore_rng_state(sd)
    assert calls == [([7, 8, 9], 2)]
    monkeypatch.setattr(torch.cuda, "is_initialized", lambda: False)
    assert "cuda" not in udist.capture_rng_state()


def test_all_reduce_helpers_are_identity_on_one_process():
    assert udist.all_reduce_max([0, 3, 1]) == [0, 3, 1]
    assert udist.any_rank(False) is False and udist.any_rank(True) is True
    assert udist.all_gather_objects({"a": 1}) == [{"a": 1}]