holy fuking shit

2026-09-12 01:56:34 +08:00
parent 36f354736c
commit d975309a81
83 changed files with 2748 additions and 896 deletions
@@ -1,10 +1,13 @@
-"""高层训练落盘保证：checkpoint、SIGTERM 优雅保存、worker 信号屏蔽。
+"""高层训练落盘保证：checkpoint、控制文件优雅停止、信号兜底、worker 信号屏蔽。

-背景：训练前端的「停止」向整个进程组发 SIGTERM，而旧脚本只在 ``learn()``
-正常返回后才保存，导致"停一下就全丢"。这里把三条保证固定成回归测试：
+背景：训练前端的「停止」以前直接向进程组发 SIGTERM，而旧脚本只在 ``learn()``
+正常返回后才保存，导致"停一下就全丢"；在 Windows 上 ``os.killpg`` 更是直接
+``AttributeError``。现在停止/暂停走**控制文件**（跨平台），信号只作兜底。这里把
+保证固定成回归测试：
 1. 定期 checkpoint 可续训（zip + 同名元数据）；
-2. 进程组收到 SIGTERM 后仍能保存最终模型与元数据；
-3. subproc worker 忽略停止信号，不抢在主进程收尾之前退出。
+2. 控制文件发出 stop 后仍能保存最终模型与元数据（Windows 同样成立）；
+3. POSIX 下进程组 SIGTERM 也能优雅收尾；
+4. subproc worker 忽略停止信号，不抢在主进程收尾之前退出。
 """
 from __future__ import annotations

@@ -15,11 +18,14 @@ import os
 import signal
 import subprocess
 import sys
+import threading
 import time
 from pathlib import Path

 import pytest

+from mcm26b.process_control import spawn_kwargs, terminate_process_tree, write_control
+

 def _train_module():
    """按路径加载训练脚本（不触发 torch 导入，torch 只在函数内部惰性导入）。"""
@@ -39,13 +45,20 @@ def _require_rl():


 class _Run:
-    """一次子进程训练：stdout 走文件，避免管道写满把子进程卡死。"""
+    """一次子进程训练：stdout 走文件，避免管道写满把子进程卡死。
+
+    与训练前端一致：用控制文件请求暂停/停止，进程组隔离走
+    :func:`mcm26b.process_control.spawn_kwargs`（POSIX ``setsid`` /
+    Windows ``CREATE_NEW_PROCESS_GROUP``）。
+    """

    def __init__(self, run_dir: Path, *extra: str):
        script = Path(__file__).parents[1] / "scripts" / "train_hrl_high.py"
        self.run_dir = run_dir
        self.log_path = run_dir / "train_stdout.log"
        self.log_file = self.log_path.open("w", encoding="utf-8")
+        self.control_path = run_dir / "control.json"
+        write_control(self.control_path)
        cmd = [sys.executable, str(script),
               "--problem", "3", "--stage", "1", "--timesteps", "200000",
               "--seed", "7", "--envs", "2", "--vec-env", "subproc",
@@ -53,13 +66,28 @@ class _Run:
               "--reward-profile", "time", "--eval-episodes", "1",
               "--rollout-steps", "32", "--batch-size", "32", "--n-epochs", "1",
               "--max-wall-s", "600", "--force", "--verbose", "0",
+               "--control-file", str(self.control_path),
               "--model", str(run_dir / "q3.zip"), *extra]
-        self.proc = subprocess.Popen(
+        self.proc = subprocess.Popen(   # pylint: disable=consider-using-with
            cmd, stdout=self.log_file, stderr=subprocess.STDOUT, text=True,
-            start_new_session=True)
+            **spawn_kwargs())

    def checkpoints(self) -> list[Path]:
-        return sorted((self.run_dir / "checkpoints").glob("q3_ckpt_*.zip"))
+        """历史 checkpoint（``<stem>_<步数>.zip``）。
+
+        刻意排除滚动恢复点 ``<stem>_latest.zip``：它随写随覆盖，不是"第几个
+        checkpoint"，混进来会让序号断言失去意义。
+        """
+        found = (self.run_dir / "checkpoints").glob("q3_ckpt_[0-9]*.zip")
+        return sorted(path for path in found
+                      if path.stem.rsplit("_", 1)[-1].isdigit())
+
+    def recovery(self) -> Path:
+        return self.run_dir / "checkpoints" / "q3_ckpt_latest.zip"
+
+    def resume_hint(self) -> dict:
+        return json.loads((self.run_dir / "resume.json").read_text(
+            encoding="utf-8"))

    def wait_for(self, predicate, timeout_s: float) -> bool:
        deadline = time.monotonic() + timeout_s
@@ -71,8 +99,35 @@ class _Run:
            time.sleep(0.4)
        return predicate()

+    def logs(self) -> str:
+        try:
+            return self.log_path.read_text(encoding="utf-8", errors="replace")
+        except OSError:  # pragma: no cover - 文件被并发写时的极端情况
+            return ""
+
+    def pause(self) -> None:
+        write_control(self.control_path, paused=True)
+
+    def resume(self) -> None:
+        write_control(self.control_path, paused=False)
+
+    def _finish(self) -> str:
+        self.log_file.close()
+        return self.logs()
+
    def stop(self) -> str:
-        """复现训练前端的停止语义：对整个进程组发 SIGTERM。"""
+        """复现训练前端的停止语义：写控制文件请求优雅收尾（跨平台）。"""
+        write_control(self.control_path, stop=True)
+        forced = terminate_process_tree(self.proc, grace_s=240.0)
+        if forced:
+            self.log_file.close()
+            pytest.fail("控制文件停止请求后训练进程未在 240s 内退出（优雅停止失败）")
+        return self._finish()
+
+    def stop_by_signal(self) -> str:
+        """POSIX 兜底路径：对整个进程组发 SIGTERM（Windows 不支持 killpg）。"""
+        if os.name == "nt":  # pragma: no cover - 仅在 Windows 上走到
+            pytest.skip("Windows 无 os.killpg，信号路径不适用")
        try:
            os.killpg(os.getpgid(self.proc.pid), signal.SIGTERM)
        except (ProcessLookupError, PermissionError):
@@ -84,8 +139,12 @@ class _Run:
            self.proc.wait(timeout=30)
            self.log_file.close()
            pytest.fail("SIGTERM 后训练进程未在 240s 内退出（优雅停止失败）")
+        return self._finish()
+
+    def hard_kill(self) -> None:
+        """模拟"急停"：立即强杀整棵进程树，不给训练脚本任何收尾机会。"""
+        terminate_process_tree(self.proc, grace_s=0.0)
        self.log_file.close()
-        return self.log_path.read_text(encoding="utf-8", errors="replace")


 def _run_dir(tmp_path: Path) -> Path:
@@ -119,7 +178,7 @@ def test_worker_ignores_stop_signals_but_main_does_not(monkeypatch):
    class _Worker:
        name = "SpawnProcess-3"

-    monkeypatch.setattr(multiprocessing, "current_process", lambda: _Worker())
+    monkeypatch.setattr(multiprocessing, "current_process", _Worker)
    try:
        module._ignore_stop_signals_in_worker()
        assert signal.getsignal(signal.SIGTERM) == signal.SIG_IGN
@@ -131,7 +190,7 @@ def test_worker_ignores_stop_signals_but_main_does_not(monkeypatch):
    class _Main:
        name = "MainProcess"

-    monkeypatch.setattr(multiprocessing, "current_process", lambda: _Main())
+    monkeypatch.setattr(multiprocessing, "current_process", _Main)
    module._ignore_stop_signals_in_worker()
    assert signal.getsignal(signal.SIGTERM) == original

@@ -168,14 +227,139 @@ def test_policy_hyperparams_records_effective_values():
    assert module._policy_hyperparams(object())["n_steps"] == 0


+def test_control_poll_state_machine(tmp_path):
+    """控制文件轮询状态机：抽样读盘、暂停阻塞、stop 打断、暂停时长累计。"""
+    module = _train_module()
+    path = tmp_path / "control.json"
+    write_control(path, paused=True)
+    poll = module._ControlPoll(path, poll_every=1, poll_s=0.05)
+    assert poll.poll().paused is True
+
+    events = []
+
+    def _resume_later():
+        time.sleep(0.15)
+        write_control(path, paused=False)
+
+    threading.Thread(target=_resume_later, daemon=True).start()
+    stopped = poll.wait_while_paused(lambda: events.append("pause"),
+                                     lambda: events.append("resume"))
+    assert stopped is False
+    assert events == ["pause", "resume"]
+    assert poll.paused_total_s >= 0.1
+
+    # 暂停期间收到 stop：wait 必须返回 True，让调用方转入收尾
+    write_control(path, paused=True)
+
+    def _stop_later():
+        time.sleep(0.15)
+        write_control(path, stop=True)
+
+    threading.Thread(target=_stop_later, daemon=True).start()
+    assert poll.wait_while_paused(lambda: None, lambda: None) is True
+    assert poll.poll().stop is True
+
+
+def test_atomic_model_save_keeps_previous_snapshot_on_failure(tmp_path):
+    """保存途中失败/被杀：目标路径仍是上一份完整快照，不会出现半截 zip。"""
+    module = _train_module()
+    target = tmp_path / "q3_ckpt_latest.zip"
+    target.write_bytes(b"OLD-COMPLETE")
+
+    class _FailingModel:
+        def save(self, path):
+            Path(path).write_bytes(b"PARTIAL")
+            raise RuntimeError("磁盘写满")
+
+    with pytest.raises(RuntimeError):
+        module._save_model_atomic(_FailingModel(), target)
+    assert target.read_bytes() == b"OLD-COMPLETE"
+    leftovers = [p.name for p in tmp_path.iterdir() if p.name != target.name]
+    assert leftovers and all(name.startswith(".tmp-") for name in leftovers)
+
+
+def test_write_json_atomic_replaces_in_one_step(tmp_path):
+    module = _train_module()
+    path = tmp_path / "resume.json"
+    module._write_json_atomic(path, {"model": "a.zip", "timesteps": 1})
+    assert json.loads(path.read_text(encoding="utf-8"))["timesteps"] == 1
+    module._write_json_atomic(path, {"model": "b.zip", "timesteps": 2})
+    assert json.loads(path.read_text(encoding="utf-8"))["model"] == "b.zip"
+    assert not list(tmp_path.glob(".tmp-*"))
+
+
+def test_prune_checkpoints_never_deletes_recovery_point(tmp_path):
+    """滚动恢复点与临时文件不参与历史清理。"""
+    module = _train_module()
+    directory = tmp_path / "checkpoints"
+    directory.mkdir()
+    latest = directory / "q3_ckpt_latest.zip"
+    latest.write_bytes(b"RECOVERY")
+    (directory / ".tmp-q3_ckpt_000000009.zip").write_bytes(b"PARTIAL")
+    for index in range(5):
+        (directory / f"q3_ckpt_{index:09d}.zip").write_bytes(b"x")
+    module._prune_checkpoints(directory, "q3_ckpt", 1)
+    assert latest.exists() and latest.read_bytes() == b"RECOVERY"
+    assert (directory / ".tmp-q3_ckpt_000000009.zip").exists()
+
+
 # --------------------------------------------------------------- 集成回归
+def test_recovery_point_exists_before_training_and_is_loadable(tmp_path):
+    """起步基线：训练刚起就有可加载、可续训的产物（--checkpoint-every 0 也成立）。"""
+    _require_rl()
+    from sb3_contrib import MaskablePPO
+
+    run = _Run(_run_dir(tmp_path), "--checkpoint-every", "0")
+    try:
+        assert run.wait_for(lambda: run.recovery().exists(), 120.0), \
+            "起步恢复点未落盘"
+        loaded = MaskablePPO.load(str(run.recovery()))
+        assert loaded.num_timesteps >= 0
+        hint = run.resume_hint()
+        assert Path(hint["model"]) == run.recovery()
+        assert hint["timesteps_semantics"] == "additional"
+    finally:
+        run.stop()
+
+
+def test_hard_kill_still_leaves_resumable_recovery_point(tmp_path):
+    """急停（强杀、无收尾）后仍能续训：滚动恢复点 + resume.json 都在且可加载。"""
+    _require_rl()
+    from sb3_contrib import MaskablePPO
+
+    run = _Run(_run_dir(tmp_path), "--checkpoint-every", "0",
+               "--recovery-every", "32")
+    try:
+        assert run.wait_for(lambda: run.recovery().exists(), 120.0), \
+            "起步恢复点未落盘"
+        first = run.recovery().stat().st_mtime_ns
+        assert run.wait_for(
+            lambda: run.recovery().stat().st_mtime_ns != first, 300.0), \
+            "滚动恢复点未按 --recovery-every 刷新"
+        assert run.proc.poll() is None, "训练已自行结束，无法验证急停"
+    finally:
+        run.hard_kill()          # 急停：强杀整棵进程树，没有任何 Python 收尾
+
+    hint = run.resume_hint()
+    assert Path(hint["model"]).is_file(), hint
+    assert Path(hint["model"]) == run.recovery(), hint
+    loaded = MaskablePPO.load(str(run.recovery()))
+    assert loaded.num_timesteps >= 32, loaded.num_timesteps
+    metadata = json.loads(run.recovery().with_suffix(".json").read_text(
+        encoding="utf-8"))
+    assert metadata["recovery"] is True
+    assert metadata["checkpoint_reason"] in ("baseline", "interval")
+    # 急停没有 graceful 收尾，所以不会有"stopped_by_signal"的最终模型
+    assert not (run.run_dir / "q3.zip").exists()
+
+
 def test_checkpoint_is_written_and_loadable(tmp_path):
    _require_rl()
    from sb3_contrib import MaskablePPO

    run = _Run(_run_dir(tmp_path), "--checkpoint-every", "64")
    try:
-        assert run.wait_for(lambda: bool(run.checkpoints()), 300.0), \
+        assert run.wait_for(run.checkpoints, 300.0), \
            "训练未在 300s 内产出 checkpoint"
    finally:
        run.stop()
@@ -193,16 +377,60 @@ def test_checkpoint_is_written_and_loadable(tmp_path):
    assert (hyper["n_steps"], hyper["batch_size"], hyper["n_epochs"]) == (32, 32, 1)


+def test_control_file_stop_saves_model(tmp_path):
+    """控制文件 stop 是训练前端的默认停止路径，Windows 上同样必须成立。"""
+    _require_rl()
+    run = _Run(_run_dir(tmp_path), "--checkpoint-every", "64")
+    model_path = run.run_dir / "q3.zip"
+    try:
+        assert run.wait_for(run.checkpoints, 300.0), \
+            "训练未在 300s 内进入 rollout 阶段"
+        assert run.proc.poll() is None, "训练进程提前退出，无法验证优雅停止"
+    finally:
+        output = run.stop()
+
+    assert run.proc.returncode == 0, \
+        f"控制文件停止应正常退出，实际 {run.proc.returncode}\n{output}"
+    assert model_path.exists(), "控制文件停止后最终模型未落盘"
+    metadata = json.loads(model_path.with_suffix(".json").read_text(
+        encoding="utf-8"))
+    assert metadata["stopped_by_signal"] is True
+    assert metadata["stop_signal"] == "control_file"
+    assert metadata["checkpoint_files"], "元数据应记录已产出的 checkpoint"
+    assert not (run.run_dir / "training_failure.json").exists()
+    assert "saved_after_stop" in output
+
+
+def test_control_file_pause_then_resume_keeps_training_alive(tmp_path):
+    """暂停是协作式的：进程存活、停在 step 边界，continue 后照常收尾。"""
+    _require_rl()
+    run = _Run(_run_dir(tmp_path), "--checkpoint-every", "64")
+    try:
+        assert run.wait_for(run.checkpoints, 300.0), \
+            "训练未在 300s 内进入 rollout 阶段"
+        run.pause()
+        assert run.wait_for(lambda: '"event": "paused"' in run.logs(), 120.0), \
+            "训练脚本未在 120s 内进入暂停"
+        assert run.proc.poll() is None, "暂停不应结束进程"
+        run.resume()
+        assert run.wait_for(lambda: '"event": "resumed"' in run.logs(), 120.0), \
+            "训练脚本未在 120s 内恢复"
+    finally:
+        run.stop()
+    assert run.proc.returncode == 0
+
+
+@pytest.mark.skipif(os.name == "nt", reason="Windows 没有 os.killpg，信号兜底路径不适用")
 def test_sigterm_saves_model_instead_of_losing_the_run(tmp_path):
    _require_rl()
    run = _Run(_run_dir(tmp_path), "--checkpoint-every", "64")
    model_path = run.run_dir / "q3.zip"
    try:
-        assert run.wait_for(lambda: bool(run.checkpoints()), 300.0), \
+        assert run.wait_for(run.checkpoints, 300.0), \
            "训练未在 300s 内进入 rollout 阶段"
        assert run.proc.poll() is None, "训练进程提前退出，无法验证优雅停止"
    finally:
-        output = run.stop()
+        output = run.stop_by_signal()

    assert run.proc.returncode == 0, \
        f"优雅停止应正常退出，实际 {run.proc.returncode}\n{output}"