code/tests/test_audit_g3_trainer_control_stop.py

align / code /tests /test_audit_g3_trainer_control_stop.py

Add files using upload-large-folder tool

bcc8d35 verified 16 days ago

10.1 kB

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

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