1919810

tests/test_hrl_checkpoint.py

@@ -0,0 +1,216 @@
+"""高层训练落盘保证：checkpoint、SIGTERM 优雅保存、worker 信号屏蔽。
+
+背景：训练前端的「停止」向整个进程组发 SIGTERM，而旧脚本只在 ``learn()``
+正常返回后才保存，导致"停一下就全丢"。这里把三条保证固定成回归测试：
+1. 定期 checkpoint 可续训（zip + 同名元数据）；
+2. 进程组收到 SIGTERM 后仍能保存最终模型与元数据；
+3. subproc worker 忽略停止信号，不抢在主进程收尾之前退出。
+"""
+from __future__ import annotations
+
+import importlib.util
+import json
+import multiprocessing
+import os
+import signal
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+import pytest
+
+
+def _train_module():
+    """按路径加载训练脚本（不触发 torch 导入，torch 只在函数内部惰性导入）。"""
+    path = Path(__file__).parents[1] / "scripts" / "train_hrl_high.py"
+    spec = importlib.util.spec_from_file_location("train_hrl_high_ckpt", path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _require_rl():
+    pytest.importorskip("stable_baselines3")
+    pytest.importorskip("sb3_contrib")
+    pytest.importorskip("torch")
+
+
+class _Run:
+    """一次子进程训练：stdout 走文件，避免管道写满把子进程卡死。"""
+
+    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")
+        cmd = [sys.executable, str(script),
+               "--problem", "3", "--stage", "1", "--timesteps", "200000",
+               "--seed", "7", "--envs", "2", "--vec-env", "subproc",
+               "--device", "cpu", "--config-profile", "coarse-pretrain",
+               "--reward-profile", "time", "--eval-episodes", "1",
+               "--rollout-steps", "32", "--batch-size", "32", "--n-epochs", "1",
+               "--max-wall-s", "600", "--force", "--verbose", "0",
+               "--model", str(run_dir / "q3.zip"), *extra]
+        self.proc = subprocess.Popen(
+            cmd, stdout=self.log_file, stderr=subprocess.STDOUT, text=True,
+            start_new_session=True)
+
+    def checkpoints(self) -> list[Path]:
+        return sorted((self.run_dir / "checkpoints").glob("q3_ckpt_*.zip"))
+
+    def wait_for(self, predicate, timeout_s: float) -> bool:
+        deadline = time.monotonic() + timeout_s
+        while time.monotonic() < deadline:
+            if predicate():
+                return True
+            if self.proc.poll() is not None:
+                return predicate()
+            time.sleep(0.4)
+        return predicate()
+
+    def stop(self) -> str:
+        """复现训练前端的停止语义：对整个进程组发 SIGTERM。"""
+        try:
+            os.killpg(os.getpgid(self.proc.pid), signal.SIGTERM)
+        except (ProcessLookupError, PermissionError):
+            self.proc.terminate()
+        try:
+            self.proc.wait(timeout=240)
+        except subprocess.TimeoutExpired:
+            self.proc.kill()
+            self.proc.wait(timeout=30)
+            self.log_file.close()
+            pytest.fail("SIGTERM 后训练进程未在 240s 内退出（优雅停止失败）")
+        self.log_file.close()
+        return self.log_path.read_text(encoding="utf-8", errors="replace")
+
+
+def _run_dir(tmp_path: Path) -> Path:
+    run_dir = tmp_path / "run"
+    run_dir.mkdir(parents=True, exist_ok=True)
+    return run_dir
+
+
+# --------------------------------------------------------------- 纯函数回归
+def test_prune_checkpoints_keeps_newest(tmp_path):
+    module = _train_module()
+    directory = tmp_path / "checkpoints"
+    directory.mkdir()
+    for index in range(5):
+        for suffix in (".zip", ".json"):
+            target = directory / f"q3_ckpt_{index:09d}{suffix}"
+            target.write_bytes(b"x")
+            stamp = 1_700_000_000 + index
+            os.utime(target, (stamp, stamp))
+    removed = module._prune_checkpoints(directory, "q3_ckpt", 2)
+    assert len(removed) == 3
+    remaining = sorted(path.name for path in directory.iterdir())
+    assert remaining == ["q3_ckpt_000000003.json", "q3_ckpt_000000003.zip",
+                         "q3_ckpt_000000004.json", "q3_ckpt_000000004.zip"]
+
+
+def test_worker_ignores_stop_signals_but_main_does_not(monkeypatch):
+    module = _train_module()
+    original = signal.getsignal(signal.SIGTERM)
+
+    class _Worker:
+        name = "SpawnProcess-3"
+
+    monkeypatch.setattr(multiprocessing, "current_process", lambda: _Worker())
+    try:
+        module._ignore_stop_signals_in_worker()
+        assert signal.getsignal(signal.SIGTERM) == signal.SIG_IGN
+        assert signal.getsignal(signal.SIGINT) == signal.SIG_IGN
+    finally:
+        signal.signal(signal.SIGTERM, original)
+        signal.signal(signal.SIGINT, signal.default_int_handler)
+
+    class _Main:
+        name = "MainProcess"
+
+    monkeypatch.setattr(multiprocessing, "current_process", lambda: _Main())
+    module._ignore_stop_signals_in_worker()
+    assert signal.getsignal(signal.SIGTERM) == original
+
+
+def test_policy_hyperparams_records_effective_values():
+    """元数据必须记实际生效的超参（resume 时来自 checkpoint，而非命令行）。"""
+    module = _train_module()
+
+    class _Schedule:
+        def __init__(self, value):
+            self._value = value
+
+        def current_value(self):
+            return self._value
+
+    class _Model:
+        n_steps = 256
+        batch_size = 256
+        n_epochs = 10
+        gamma = 1.0
+        gae_lambda = 0.95
+        learning_rate = _Schedule(3e-4)
+        clip_range = _Schedule(0.2)
+        ent_coef = 0.005
+        vf_coef = 0.5
+        max_grad_norm = 0.5
+
+    hyper = module._policy_hyperparams(_Model())
+    assert hyper["n_steps"] == 256 and hyper["batch_size"] == 256
+    assert hyper["n_epochs"] == 10
+    assert hyper["learning_rate"] == 3e-4
+    assert hyper["clip_range"] == 0.2
+    assert hyper["gamma"] == 1.0 and hyper["ent_coef"] == 0.005
+    assert module._policy_hyperparams(object())["n_steps"] == 0
+
+
+# --------------------------------------------------------------- 集成回归
+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), \
+            "训练未在 300s 内产出 checkpoint"
+    finally:
+        run.stop()
+
+    archives = run.checkpoints()
+    assert archives, "checkpoint 目录为空"
+    loaded = MaskablePPO.load(str(archives[-1]))
+    assert loaded.num_timesteps >= 64
+    metadata = json.loads(archives[-1].with_suffix(".json").read_text(
+        encoding="utf-8"))
+    assert metadata["checkpoint"] is True
+    assert metadata["checkpoint_timesteps"] == loaded.num_timesteps
+    assert metadata["model_sha256"]
+    hyper = metadata["ppo_hyperparams"]
+    assert (hyper["n_steps"], hyper["batch_size"], hyper["n_epochs"]) == (32, 32, 1)
+
+
+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), \
+            "训练未在 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(), "SIGTERM 后最终模型未落盘"
+    metadata = json.loads(model_path.with_suffix(".json").read_text(
+        encoding="utf-8"))
+    assert metadata["stopped_by_signal"] is True
+    assert metadata["stop_signal"] == "SIGTERM"
+    assert metadata["checkpoint_files"], "元数据应记录已产出的 checkpoint"
+    assert not (run.run_dir / "training_failure.json").exists()
+    assert "saved_after_stop" in output

tests/test_training_ui.py

@@ -424,6 +424,23 @@ def test_describe_process_progress_and_eta():
     assert ended["remaining_s"] is None
 
 
+def test_describe_process_curriculum_transfer_uses_delta_steps():
+    """--init 课程迁移：计数器从上一阶段继续，进度/ETA 必须按本次增量算。"""
+    snapshot = {
+        "status": "running", "elapsed_s": 10.0,
+        "config": {"timesteps": 8000, "max_wall_s": 3600.0},
+        "metrics": [{"timesteps": 110000, "start_timesteps": 109776, "fps": 50.0},
+                    {"timesteps": 111776, "start_timesteps": 109776, "fps": 50.0}],
+        "convergence": {}, "log_total": 2, "metric_total": 2,
+    }
+    process = td.describe_process(snapshot)
+    assert process["done_steps"] == 111776      # 绝对计数照实显示
+    assert process["start_steps"] == 109776
+    assert process["trained_steps"] == 2000     # 本次增量
+    assert process["percent"] == 25.0           # 而不是被夹到 100%
+    assert process["remaining_s"] == pytest.approx(120.0, rel=1e-6)
+
+
 def test_demo_snapshot_contract():
     for scenario in demo.DEMO_SCENARIOS:
         snapshot = demo.demo_snapshot(scenario=scenario)