445 lines
18 KiB
Python
445 lines
18 KiB
Python
"""高层训练落盘保证：checkpoint、控制文件优雅停止、信号兜底、worker 信号屏蔽。
|
||
|
||
背景：训练前端的「停止」以前直接向进程组发 SIGTERM，而旧脚本只在 ``learn()``
|
||
正常返回后才保存，导致"停一下就全丢"；在 Windows 上 ``os.killpg`` 更是直接
|
||
``AttributeError``。现在停止/暂停走**控制文件**（跨平台），信号只作兜底。这里把
|
||
保证固定成回归测试：
|
||
1. 定期 checkpoint 可续训（zip + 同名元数据）；
|
||
2. 控制文件发出 stop 后仍能保存最终模型与元数据（Windows 同样成立）；
|
||
3. POSIX 下进程组 SIGTERM 也能优雅收尾；
|
||
4. subproc worker 忽略停止信号，不抢在主进程收尾之前退出。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import importlib.util
|
||
import json
|
||
import multiprocessing
|
||
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 只在函数内部惰性导入）。"""
|
||
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 走文件，避免管道写满把子进程卡死。
|
||
|
||
与训练前端一致：用控制文件请求暂停/停止，进程组隔离走
|
||
: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",
|
||
"--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",
|
||
"--control-file", str(self.control_path),
|
||
"--model", str(run_dir / "q3.zip"), *extra]
|
||
self.proc = subprocess.Popen( # pylint: disable=consider-using-with
|
||
cmd, stdout=self.log_file, stderr=subprocess.STDOUT, text=True,
|
||
**spawn_kwargs())
|
||
|
||
def checkpoints(self) -> list[Path]:
|
||
"""历史 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
|
||
while time.monotonic() < deadline:
|
||
if predicate():
|
||
return True
|
||
if self.proc.poll() is not None:
|
||
return predicate()
|
||
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:
|
||
"""复现训练前端的停止语义：写控制文件请求优雅收尾（跨平台）。"""
|
||
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):
|
||
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 内退出（优雅停止失败）")
|
||
return self._finish()
|
||
|
||
def hard_kill(self) -> None:
|
||
"""模拟"急停"：立即强杀整棵进程树，不给训练脚本任何收尾机会。"""
|
||
terminate_process_tree(self.proc, grace_s=0.0)
|
||
self.log_file.close()
|
||
|
||
|
||
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", _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", _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_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(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_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(run.checkpoints, 300.0), \
|
||
"训练未在 300s 内进入 rollout 阶段"
|
||
assert run.proc.poll() is None, "训练进程提前退出，无法验证优雅停止"
|
||
finally:
|
||
output = run.stop_by_signal()
|
||
|
||
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
|