"""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 ` 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))