Blame for code/tests/test_audit_g3_trainer_control_stop.py

"""G3 audit regressions: U01-1 (a stop must be able to finish) and U01-2 (an
operator's stop must not be relaunched). CPU only.

Trainer half: single process, REAL signals / files. Supervisor half: REAL child
process groups, as in tests/test_supervisor.py.
"""

from __future__ import annotations

import json
import os
import signal
import subprocess
import sys
import textwrap
import time

import pytest
from test_control_plane import _cfg, _signal_at
from test_model_parity import tiny_config
from test_trainer import _tiny_trainer

from align.train.trainer import LossOutput
from align.utils import supervisor as sv
from align.utils.supervisor import Supervisor


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


def _marker(run_dir, name):
    path = os.path.join(run_dir, name)
    return json.load(open(path)) if os.path.exists(path) else None


def _at_micro(micro, fn):
    """wrap_adapter running `fn()` on the `micro`-th adapter call."""

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

        def adapter(m, batch):
            n["i"] += 1
            if n["i"] == micro:
                fn()
            return inner(m, batch)

        return adapter

    return wrap


# ---------------------------------------------------------------------------
# trainer side
# ---------------------------------------------------------------------------


def test_u01_2_signal_stop_leaves_a_stop_marker_the_supervisor_can_read(tiny, tmp_workdir):
    """torchrun exits 1 whatever its worker exited with, so the checkpointed 143
    never reached the supervisor and `kill <torchrun>` was relaunched. The
    trainer now leaves STOP.json (128+signum). Mutant: drop
    `_write_stop_marker` -> no marker."""
    t0 = time.time()
    t, _ = _tiny_trainer(
        tiny, tmp_workdir, save_every_micro_steps=8, wrap_adapter=_signal_at(3, signal.SIGTERM)
    )
    summary = t.train()
    assert summary["stop_reason"] == "signal"
    rec = sv.read_stop_marker(tmp_workdir, since=t0)
    assert rec is not None and rec["code"] == 128 + int(signal.SIGTERM), rec
    assert t.stop_exit_code == rec["code"]


def test_u01_2_a_run_that_ends_on_its_own_leaves_no_stop_marker(tiny, tmp_workdir):
    t, _ = _tiny_trainer(tiny, tmp_workdir, save_every_micro_steps=8)
    assert t.train()["stop_reason"] == "data_exhausted"
    assert _marker(tmp_workdir, sv.STOP_MARKER_NAME) is None


def test_u01_1_a_stop_request_file_stops_at_the_boundary_and_saves(tiny, tmp_workdir):
    """The signal-free drain channel. The request lands mid-window (micro 3);
    the stop is the window's boundary (micro 4) and writes the checkpoint
    (save_every is 8, so micro 4 exists only because the stop wrote it).
    Mutant: do not poll the file -> the run trains all 8 micros."""
    t, _ = _tiny_trainer(
        tiny,
        tmp_workdir,
        save_every_micro_steps=8,
        wrap_adapter=_at_micro(3, lambda: sv.write_run_marker(tmp_workdir, sv.STOP_REQUEST_NAME)),
    )
    summary = t.train()
    assert summary["stop_reason"] == "signal", summary
    assert summary["micro_step"] == 4
    assert os.path.isdir(os.path.join(tmp_workdir, "checkpoints", "micro_00000004"))
    assert t.stop_exit_code == 128 + int(signal.SIGTERM)
    assert sv.read_stop_marker(tmp_workdir) is not None


def test_u01_1_a_stale_stop_request_from_an_earlier_run_is_ignored(tiny, tmp_workdir):
    """A request left on disk by a previous attempt must not stop the next one."""
    with open(os.path.join(tmp_workdir, sv.STOP_REQUEST_NAME), "w") as f:
        json.dump({"written_at": time.time() - 3600}, f)
    t, _ = _tiny_trainer(tiny, tmp_workdir, save_every_micro_steps=8)
    assert t.train()["stop_reason"] == "data_exhausted"


def test_u01_1_the_trainer_announces_drain_capability_only_while_training(tiny, tmp_workdir):
    seen = {}

    def peek():
        seen["during"] = sv.read_run_marker(tmp_workdir, sv.DRAIN_READY_NAME) is not None

    t, _ = _tiny_trainer(tiny, tmp_workdir, save_every_micro_steps=8, wrap_adapter=_at_micro(1, peek))
    assert sv.read_run_marker(tmp_workdir, sv.DRAIN_READY_NAME) is None
    t.train()
    assert seen["during"] is True
    assert sv.read_run_marker(tmp_workdir, sv.DRAIN_READY_NAME) is None


# ---------------------------------------------------------------------------
# supervisor side
# ---------------------------------------------------------------------------

STOPPED_CHILD = textwrap.dedent(
    """
    import json, os, sys, time
    run = sys.argv[1]
    cnt = os.path.join(run, "attempt.txt")
    n = int(open(cnt).read()) if os.path.exists(cnt) else 0
    open(cnt, "w").write(str(n + 1))
    if n == 0:
        time.sleep(0.05)
        with open(os.path.join(run, "history.jsonl"), "a") as f:
            f.write("{}\\n")
        os.makedirs(os.path.join(run, "checkpoints", "micro_%08d" % 100))
        if sys.argv[2] == "marker":
            with open(os.path.join(run, "STOP.json"), "w") as f:
                json.dump({"code": 143, "reason": "t", "written_at": time.time()}, f)
        sys.exit(1)  # torchrun: exit 1 whatever the worker did
    sys.exit(0)
    """
)


def _sup(tmp, mode, **kw):
    script = os.path.join(tmp, "child.py")
    with open(script, "w") as f:
        f.write(STOPPED_CHILD)
    run = os.path.join(tmp, "run")
    os.makedirs(run, exist_ok=True)
    return Supervisor(
        [sys.executable, script, run, mode], run, backoff_s=0, preflight=lambda: None, **kw
    ), run


def test_u01_2_supervisor_does_not_relaunch_a_stop_the_trainer_marked(tmp_workdir):
    """rc == 1 with progress made would normally be retried; STOP.json says an
    operator asked for the stop. Mutant: ignore the marker -> 2 attempts."""
    sup, run = _sup(tmp_workdir, "marker")
    res = sup.run(install_signal_handlers=False)
    assert (res.reason, res.attempts, res.exit_code) == ("operator_stop", 1, 143)


def test_u01_2_retry_on_stop_still_relaunches(tmp_workdir):
    sup, _ = _sup(tmp_workdir, "marker", retry_on_stop=True)
    res = sup.run(install_signal_handlers=False)
    assert (res.reason, res.attempts) == ("success", 2)


def test_u01_2_without_a_marker_exit_1_with_progress_is_still_a_retryable_crash(tmp_workdir):
    sup, _ = _sup(tmp_workdir, "nomarker")
    res = sup.run(install_signal_handlers=False)
    assert (res.reason, res.attempts) == ("success", 2)


def test_u01_2_a_stale_stop_marker_is_ignored(tmp_workdir):
    sup, run = _sup(tmp_workdir, "nomarker")
    with open(os.path.join(run, "STOP.json"), "w") as f:
        json.dump({"code": 143, "reason": "old", "written_at": time.time() - 3600}, f)
    res = sup.run(install_signal_handlers=False)
    assert (res.reason, res.attempts) == ("success", 2)


# A drain-capable child: announces DRAIN_READY, IGNORES SIGTERM, and on seeing
# STOP_REQUEST takes `take_s` seconds to "checkpoint" before exiting 143. If the
# supervisor signalled or killed it first, `done` is never written.
DRAINING_CHILD = textwrap.dedent(
    """
    import json, os, signal, sys, time
    run, take_s, mode = sys.argv[1], float(sys.argv[2]), sys.argv[3]
    signal.signal(signal.SIGTERM, signal.SIG_IGN)
    if mode == "capable":
        with open(os.path.join(run, "DRAIN_READY.json"), "w") as f:
            json.dump({"written_at": time.time()}, f)
    open(os.path.join(run, "started"), "w").close()
    while True:
        if os.path.exists(os.path.join(run, "STOP_REQUEST")):
            time.sleep(take_s)
            open(os.path.join(run, "done"), "w").close()
            sys.exit(143)
        time.sleep(0.05)
    """
)


def _drain_run(tmp, take_s, mode, drain_s, grace_s=1):
    script = os.path.join(tmp, "drainer.py")
    with open(script, "w") as f:
        f.write(DRAINING_CHILD)
    run = os.path.join(tmp, "run")
    os.makedirs(run)
    runner = (
        "import sys;from align.utils.supervisor import Supervisor;"
        f"r=Supervisor([sys.executable,{script!r},{run!r},{str(take_s)!r},{mode!r}],{run!r},"
        f"grace_s={grace_s},drain_s={drain_s},max_retries=0,"
        "preflight=lambda:None).run();sys.exit(r.exit_code)"
    )
    sup = subprocess.Popen([sys.executable, "-c", runner], env=dict(os.environ, PYTHONPATH=os.getcwd()))
    for _ in range(200):
        if os.path.exists(os.path.join(run, "started")):
            break
        time.sleep(0.05)
    assert os.path.exists(os.path.join(run, "started"))
    time.sleep(0.3)
    t0 = time.time()
    os.kill(sup.pid, signal.SIGTERM)
    return sup, run, t0


def _events(run):
    return [json.loads(l) for l in open(os.path.join(run, "supervisor.jsonl"))]


def test_u01_1_supervisor_drains_a_capable_trainer_instead_of_killing_it(tmp_workdir):
    """The child needs 3s to finish its stop; grace_s is 1. A SIGTERM-then-
    SIGKILL supervisor kills it mid-"save" (no `done`). Mutant: skip the
    drain -> `done` is missing."""
    sup, run, _ = _drain_run(tmp_workdir, take_s=3, mode="capable", drain_s=30)
    assert sup.wait(timeout=30) == 128 + int(signal.SIGTERM)
    assert os.path.exists(os.path.join(run, "done")), "the stop was killed before it finished"
    assert any(e["event"] == "drain_requested" for e in _events(run))
    assert not os.path.exists(os.path.join(run, sv.STOP_REQUEST_NAME)), "request left behind"


def test_u01_1_a_drain_that_never_finishes_is_escalated_after_drain_s(tmp_workdir):
    sup, run, t0 = _drain_run(tmp_workdir, take_s=600, mode="capable", drain_s=1.5)
    assert sup.wait(timeout=30) == 128 + int(signal.SIGTERM)
    assert not os.path.exists(os.path.join(run, "done"))
    assert any(e["event"] == "drain_timeout" for e in _events(run))
    assert time.time() - t0 < 25


def test_u01_1_a_child_that_cannot_drain_is_signalled_immediately(tmp_workdir):
    """No DRAIN_READY -> the legacy path (SIGTERM, grace, SIGKILL): the
    supervisor must not wait drain_s on a child that will never answer."""
    sup, run, t0 = _drain_run(tmp_workdir, take_s=600, mode="plain", drain_s=300)
    assert sup.wait(timeout=30) == 128 + int(signal.SIGTERM)
    assert time.time() - t0 < 20
    assert not any(e["event"] == "drain_requested" for e in _events(run))