"""高层训练落盘保证：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