217 lines
8.0 KiB
Python
217 lines
8.0 KiB
Python
"""高层训练落盘保证：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
|