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