diff --git a/.DS_Store b/.DS_Store index 12952a0..f384548 100644 Binary files a/.DS_Store and b/.DS_Store differ diff --git a/docs/training-ui.md b/docs/training-ui.md index cd3717d..f2f646a 100644 --- a/docs/training-ui.md +++ b/docs/training-ui.md @@ -80,6 +80,24 @@ uv run python -m mcm26b.training_ui --print-preview # 只看离线示例快照 - 产物路径限定在仓库根下解析，模型输出沿用训练脚本自身的 `--force`/`--resume` 语义（两者互斥，由后端校验）。 +## 停止与落盘保证 + +训练前端的「停止」是对独立进程组发 SIGTERM。`train_hrl_high.py` 有三层兜底， +点一下「停止」不再等于这一局白跑： + +| 机制 | 行为 | +| --- | --- | +| 定期 checkpoint | `--checkpoint-every`（默认 20000 环境步）在 `<模型目录>/checkpoints/` 落 `*_ckpt_<步数>.zip` + 同名元数据；`--checkpoint-keep` 只保留最近 K 个（默认 3），`--checkpoint-dir` 可改目录 | +| 优雅停止 | subproc worker 忽略 SIGTERM/SIGINT，主进程收到信号后在 step 边界收尾，保存最终模型与元数据（`stopped_by_signal=true`、`stop_signal`、`checkpoint_files`）并正常关闭环境；收尾阶段恢复默认处理器，再按一次可立即强杀 | +| 异常兜底 | `learn()` 抛出任何异常时先落 `*_emergency.zip` + 元数据，再写 `training_failure.json` 并抛出 | + +checkpoint 与最终模型共用同一份 json 契约，因此 `--resume ` 直接 +续训（`--timesteps` 为增量步数）。注意 `--resume`/`--init` 会沿用 checkpoint 内的 +PPO 超参（`n_steps`/`batch_size`/`n_epochs`）：命令行这三项被忽略，脚本会打一条 +`hyperparams_from_checkpoint` 事件说明差异，元数据 `ppo_hyperparams` 记录**实际生效**值。 +回归测试见 `tests/test_hrl_checkpoint.py`， +其中端到端用例真起 subproc 子进程并用 `killpg(SIGTERM)` 复现前端的停止路径。 + ## 测试 ```bash diff --git a/scripts/train_hrl_high.py b/scripts/train_hrl_high.py index 809e1d2..711160a 100644 --- a/scripts/train_hrl_high.py +++ b/scripts/train_hrl_high.py @@ -7,6 +7,15 @@ - 低层默认使用几何选点；``--local-selector rl`` 仅作为接口预留并记录， 当前 HTTP 执行路径仍由高层宏动作终点决定，不把八方向推演当正式收益。 +落盘保证（长训必需）： +- ``--checkpoint-every`` 每 N 个环境步在 ``<模型目录>/checkpoints/`` 落一个可续训 + 的 checkpoint（zip + 同名 json 元数据），只保留最近 ``--checkpoint-keep`` 个； +- 训练前端的「停止」是对**进程组**发 SIGTERM，因此 subproc worker 在 :func:`_make_env` + 里显式忽略 SIGTERM/SIGINT，主进程收到信号后在 step 边界收尾并保存模型， + 不会被 worker 先行退出打断； +- 任何异常（含被强杀之外的失败）都会先落一份 ``*_emergency.zip`` 再抛出， + 不再出现"跑了很久什么都没有"。 + P1 smoke 示例： python scripts/train_hrl_high.py \ --problem 3 --stage 1 --timesteps 2048 --seed 320001 \ @@ -23,6 +32,7 @@ import json import os import platform import random +import signal import sys import time from pathlib import Path @@ -58,6 +68,23 @@ def _reward_from_profile(profile: str): raise ValueError(f"未知 reward-profile: {profile}") +def _ignore_stop_signals_in_worker() -> None: + """subproc worker 忽略 SIGTERM/SIGINT，保证主进程能优雅收尾落盘。 + + 训练前端的「停止」向整个进程组发 SIGTERM；若 worker 同时退出，主进程会在 + 保存模型之前就拿到断管/EOF 错误，整局训练白跑。worker 活在管道后面，主进程 + 收尾时通过 ``vec_env.close()`` 的 ``close`` 指令正常结束它们。 + """ + try: + import multiprocessing + if multiprocessing.current_process().name == "MainProcess": + return + for signum in (signal.SIGTERM, signal.SIGINT): + signal.signal(signum, signal.SIG_IGN) + except (ImportError, ValueError, OSError): # 非主线程/平台不支持：不影响训练 + return + + def _make_env(problem: int, stage_index: int, seed: int, monitor_path: str | None, profile: str, reward_profile: str): @@ -66,6 +93,7 @@ def _make_env(problem: int, stage_index: int, seed: int, from mcm26b.hrl.hierarchical_env import HierarchicalEnv + _ignore_stop_signals_in_worker() stage = _stage_from_index(stage_index) env = HierarchicalEnv( problem=problem, @@ -98,6 +126,41 @@ def _hash_file(path: Path) -> str: return hashlib.sha256(path.read_bytes()).hexdigest() +def _policy_hyperparams(model) -> dict[str, Any]: + """记录**实际生效**的 PPO 超参。 + + ``--resume``/``--init`` 时 ``MaskablePPO.load`` 会一并恢复 checkpoint 内的 + n_steps/batch_size/n_epochs 等，命令行同名参数被忽略；把它们写进元数据， + 事后才能核对一局到底用什么配置训的。 + """ + + def _scalar(value): + # SB3 2.9 的 learning_rate/clip_range 可能是 Schedule 对象 + current = getattr(value, "current_value", None) + if callable(current): + try: + value = current() + except Exception: + return str(value) + try: + return float(value) + except (TypeError, ValueError): + return str(value) + + return { + "n_steps": int(getattr(model, "n_steps", 0) or 0), + "batch_size": int(getattr(model, "batch_size", 0) or 0), + "n_epochs": int(getattr(model, "n_epochs", 0) or 0), + "gamma": _scalar(getattr(model, "gamma", 0.0)), + "gae_lambda": _scalar(getattr(model, "gae_lambda", 0.0)), + "learning_rate": _scalar(getattr(model, "learning_rate", 0.0)), + "clip_range": _scalar(getattr(model, "clip_range", 0.0)), + "ent_coef": _scalar(getattr(model, "ent_coef", 0.0)), + "vf_coef": _scalar(getattr(model, "vf_coef", 0.0)), + "max_grad_norm": _scalar(getattr(model, "max_grad_norm", 0.0)), + } + + def _model_metadata(*, args, wall_s: float, n_envs: int, actual_timesteps: int) -> dict[str, Any]: import numpy @@ -130,6 +193,9 @@ def _model_metadata(*, args, wall_s: float, n_envs: int, "reward_version": REWARD_VERSION, "local_selector": args.local_selector, "local_model": args.local_model, + "checkpoint_every": int(getattr(args, "checkpoint_every", 0) or 0), + "checkpoint_keep": int(getattr(args, "checkpoint_keep", 0) or 0), + "checkpoint_dir": getattr(args, "checkpoint_dir", None), "device": args.device, "wall_seconds": round(wall_s, 3), "max_wall_s": float(args.max_wall_s), @@ -145,6 +211,40 @@ def _model_metadata(*, args, wall_s: float, n_envs: int, } +def _write_metadata(*, model_path: Path, args, wall_s: float, n_envs: int, + actual_timesteps: int, monitor_paths: list[str], + model=None, + extra: dict[str, Any] | None = None) -> dict[str, Any]: + """写模型旁挂元数据（含 sha256），训练前端/审计都读它。""" + metadata = _model_metadata(args=args, wall_s=wall_s, n_envs=n_envs, + actual_timesteps=actual_timesteps) + if model is not None: + metadata["ppo_hyperparams"] = _policy_hyperparams(model) + metadata["model_sha256"] = _hash_file(model_path) + metadata["monitor_files"] = list(monitor_paths) + metadata["reward_note"] = ( + "reward_profile=time 使用 HRLReward.time_only（P0.4 纯时间口径）；" + "reward_profile=prompt 使用原始 §5.2 塑形。模型元数据已固化口径。") + if extra: + metadata.update(extra) + model_path.with_suffix(".json").write_text( + json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8") + return metadata + + +def _prune_checkpoints(directory: Path, stem: str, keep: int) -> list[str]: + """只保留最近 ``keep`` 个 checkpoint（zip+json 同删），返回被删的 zip 路径。""" + files = sorted(directory.glob(f"{stem}_*.zip"), + key=lambda path: path.stat().st_mtime, reverse=True) + removed: list[str] = [] + for stale in files[max(1, int(keep)):]: + for candidate in (stale, stale.with_suffix(".json")): + if candidate.exists(): + candidate.unlink() + removed.append(str(stale)) + return removed + + def _evaluate_env(model, *, problem: int, stage_index: int, episodes: int, seed: int, profile: str, reward_profile: str) -> dict[str, Any]: @@ -288,6 +388,52 @@ def run_training(args) -> int: if args.init: pass # 已在上面处理；此分支仅保持语义清晰 + # resume/init 会恢复 checkpoint 内的 PPO 超参，命令行同名三项目前被忽略； + # 显式说出差异，避免"以为换了 n_steps 其实没换"。 + if args.resume or args.init: + effective = _policy_hyperparams(model) + requested = {"n_steps": int(args.rollout_steps), + "batch_size": int(args.batch_size), + "n_epochs": int(args.n_epochs)} + if {key: effective[key] for key in requested} != requested: + print(json.dumps({ + "event": "hyperparams_from_checkpoint", + "effective": {key: effective[key] for key in requested}, + "requested": requested, + "note": ("resume/init 沿用 checkpoint 内的 PPO 超参，" + "命令行 rollout-steps/batch-size/n-epochs 被忽略")}, + ensure_ascii=False), file=sys.stderr, flush=True) + + # ------------------------------------------------------------ 优雅停止 + # 训练前端「停止」= killpg(SIGTERM)：worker 已在 _make_env 里忽略该信号， + # 主进程在此装处理器，等当前 step 结束后收尾落盘。 + stop_state: dict[str, Any] = {"requested": False, "signal": None} + previous_handlers: dict[int, Any] = {} + + def _on_stop_signal(signum, _frame): + stop_state["requested"] = True + stop_state["signal"] = signal.Signals(signum).name + print(json.dumps({"event": "stop_requested", + "signal": signal.Signals(signum).name}), flush=True) + + for _signum in (signal.SIGTERM, signal.SIGINT): + try: + previous_handlers[_signum] = signal.signal(_signum, _on_stop_signal) + except (ValueError, OSError): # 非主线程/平台不支持：不阻塞训练 + pass + + class _GracefulStop(BaseCallback): + """收到停止信号后在最近一次 step 边界结束 rollout，让模型能落盘。""" + + def _on_step(self) -> bool: + if stop_state["requested"]: + print(json.dumps({"event": "graceful_stop", + "signal": stop_state["signal"], + "timesteps": int(self.model.num_timesteps)}), + flush=True) + return False + return True + class _WallClockStop(BaseCallback): def __init__(self, max_wall_s: float): super().__init__() @@ -312,7 +458,9 @@ def run_training(args) -> int: def __init__(self, every_steps: int): super().__init__() self.every_steps = max(1, int(every_steps)) - self.last_emit = int(model.num_timesteps) + # init/resume 时计数器从上一阶段继续，前端要用增量算进度 + self.start = int(model.num_timesteps) + self.last_emit = self.start self.emitted = 0 def _next_gap(self) -> int: @@ -350,6 +498,7 @@ def run_training(args) -> int: payload = { "event": "progress", "timesteps": current, + "start_timesteps": self.start, "n_updates": int(getattr(self.model, "_n_updates", 0) or 0), "ep_rew_mean": (sum(ep_rewards) / len(ep_rewards) if ep_rewards else None), @@ -366,43 +515,119 @@ def run_training(args) -> int: flush=True) return True + started = time.monotonic() + checkpoints_dir = (Path(args.checkpoint_dir) if args.checkpoint_dir + else run_dir / "checkpoints") + checkpoint_stem = f"{model_path.stem}_ckpt" + + class _CheckpointJSON(BaseCallback): + """每 N 个环境步落一个可续训 checkpoint（zip + 同名元数据）。 + + 只保留最近 ``keep`` 个；checkpoint 与最终模型共用同名 json 契约， + 因此 ``--resume `` 直接可用。 + """ + + def __init__(self, every_steps: int, keep: int): + super().__init__() + self.every_steps = max(1, int(every_steps)) + self.keep = max(1, int(keep)) + self.last_saved = int(model.num_timesteps) + self.saved: list[str] = [] + + def _on_step(self) -> bool: + current = int(self.model.num_timesteps) + if current - self.last_saved < self.every_steps: + return True + self.last_saved = current + self._write(current) + return True + + def _write(self, current: int) -> None: + checkpoints_dir.mkdir(parents=True, exist_ok=True) + target = checkpoints_dir / f"{checkpoint_stem}_{current:09d}.zip" + self.model.save(str(target)) + _write_metadata( + model_path=target, args=args, + wall_s=time.monotonic() - started, n_envs=args.envs, + actual_timesteps=current, monitor_paths=monitor_paths, + model=self.model, + extra={"checkpoint": True, "checkpoint_timesteps": current, + "checkpoint_of": str(model_path)}) + removed = _prune_checkpoints(checkpoints_dir, checkpoint_stem, + self.keep) + self.saved.append(str(target)) + print(json.dumps({"event": "checkpoint", "timesteps": current, + "model": str(target), "pruned": removed}, + ensure_ascii=False), flush=True) + progress_every = (args.progress_every if args.progress_every > 0 else args.rollout_steps * args.envs) - callbacks = [_WallClockStop(args.max_wall_s)] + checkpoint_cb = None + if args.checkpoint_every > 0: + checkpoint_cb = _CheckpointJSON(args.checkpoint_every, + args.checkpoint_keep) + callbacks = [_GracefulStop(), _WallClockStop(args.max_wall_s)] + if checkpoint_cb is not None: + callbacks.append(checkpoint_cb) if args.progress_json: callbacks.insert(0, _ProgressJSON(progress_every)) - started = time.monotonic() actual_timesteps = int(model.num_timesteps) try: model.learn(total_timesteps=args.timesteps, callback=callbacks if len(callbacks) > 1 else callbacks[0], progress_bar=False, reset_num_timesteps=not (args.resume or args.init)) - except Exception as exc: # 训练失败也必须保存失败上下文，不静默 + except BaseException as exc: # 训练失败也必须留下可续训产物，不静默 failure = { "event": "training_failed", "error": f"{type(exc).__name__}: {exc}", "model": str(model_path), + "timesteps": int(model.num_timesteps), } + emergency = run_dir / f"{model_path.stem}_emergency.zip" + try: + model.save(str(emergency)) + _write_metadata(model_path=emergency, args=args, + wall_s=time.monotonic() - started, + n_envs=args.envs, + actual_timesteps=int(model.num_timesteps), + monitor_paths=monitor_paths, model=model, + extra={"emergency": True, + "error": failure["error"]}) + failure["emergency_model"] = str(emergency) + except Exception as save_exc: # 落盘失败不能再掩盖原始异常 + failure["emergency_save_error"] = ( + f"{type(save_exc).__name__}: {save_exc}") (run_dir / "training_failure.json").write_text( json.dumps(failure, ensure_ascii=False, indent=2), encoding="utf-8") + print(json.dumps(failure, ensure_ascii=False), file=sys.stderr, + flush=True) raise wall_s = time.monotonic() - started actual_timesteps = int(model.num_timesteps) + stopped_by_signal = bool(stop_state["requested"]) model.save(str(model_path)) - metadata = _model_metadata(args=args, wall_s=wall_s, - n_envs=args.envs, - actual_timesteps=actual_timesteps) - metadata["model_sha256"] = _hash_file(model_path) - metadata["monitor_files"] = monitor_paths - metadata["reward_note"] = ( - "reward_profile=time 使用 HRLReward.time_only（P0.4 纯时间口径）；" - "reward_profile=prompt 使用原始 §5.2 塑形。模型元数据已固化口径。") + metadata = _write_metadata( + model_path=model_path, args=args, wall_s=wall_s, n_envs=args.envs, + actual_timesteps=actual_timesteps, monitor_paths=monitor_paths, + model=model, + extra={"stopped_by_signal": stopped_by_signal, + "stop_signal": stop_state["signal"], + "checkpoint_files": (checkpoint_cb.saved + if checkpoint_cb is not None else [])}) metadata_path = model_path.with_suffix(".json") - metadata_path.write_text(json.dumps(metadata, ensure_ascii=False, indent=2), - encoding="utf-8") vec_env.close() + # 收尾阶段恢复默认处理器：再次 Ctrl+C / 停止可立即强杀，不再等待评估。 + for _signum, _handler in previous_handlers.items(): + try: + signal.signal(_signum, _handler) + except (ValueError, OSError): + pass + if stopped_by_signal: + print(json.dumps({"event": "saved_after_stop", + "model": str(model_path), + "timesteps": actual_timesteps}), flush=True) evaluation = _evaluate_env( model, problem=args.problem, stage_index=args.stage, @@ -439,7 +664,8 @@ def main(argv=None) -> int: choices=("production", "coarse-pretrain"), default="production") parser.add_argument("--model", required=True) - parser.add_argument("--resume", action="store_true") + parser.add_argument("--resume", action="store_true", + help="从 --model 续训（沿用 checkpoint 内的 PPO 超参）") parser.add_argument("--init", default=None, help="可选：Q3→Q4 或继续预训练权重路径") parser.add_argument("--force", action="store_true") @@ -461,6 +687,12 @@ def main(argv=None) -> int: help="向训练前端输出单行 JSON 进度（UI 自动启用）") parser.add_argument("--progress-every", type=int, default=0, help="JSON 进度间隔环境步数；0=每个 rollout") + parser.add_argument("--checkpoint-every", type=int, default=20000, + help="每 N 个环境步落一个 checkpoint；0=关闭（默认 20000）") + parser.add_argument("--checkpoint-keep", type=int, default=3, + help="checkpoint 保留最近 K 个（默认 3）") + parser.add_argument("--checkpoint-dir", default=None, + help="checkpoint 目录；默认 <模型目录>/checkpoints") parser.add_argument("--verbose", type=int, default=1) args = parser.parse_args(argv) if args.timesteps <= 0 or args.envs <= 0: @@ -469,6 +701,10 @@ def main(argv=None) -> int: parser.error("PPO rollout 至少需要 2 个样本") if args.eval_episodes < 0: parser.error("eval-episodes 不能为负") + if args.checkpoint_every < 0: + parser.error("checkpoint-every 不能为负") + if args.checkpoint_keep < 1: + parser.error("checkpoint-keep 至少为 1") return run_training(args) diff --git a/src/mcm26b/training_ui/schema.py b/src/mcm26b/training_ui/schema.py index 04d9da6..bad3c8b 100644 --- a/src/mcm26b/training_ui/schema.py +++ b/src/mcm26b/training_ui/schema.py @@ -144,7 +144,8 @@ CONFIG_FIELDS: list[dict[str, Any]] = [ { "key": "init", "kind": "text", "group": "paths", "optional": True, "label": "初始权重 --init", "default": None, - "help": "留空=从零；填上一阶段 .zip 即课程迁移（步数计数重置）。", + "help": "留空=从零；填上一阶段 .zip 即课程迁移。PPO 超参与步数计数沿用" + "该模型（timesteps 按增量理解，进度条按增量显示）。", }, { "key": "local_model", "kind": "text", "group": "paths", "optional": True, diff --git a/src/mcm26b/training_ui/server.py b/src/mcm26b/training_ui/server.py index f899d54..ed5e581 100644 --- a/src/mcm26b/training_ui/server.py +++ b/src/mcm26b/training_ui/server.py @@ -739,7 +739,7 @@ _STATUS_VIEW = { "created": ("已创建", "gray", False, "命令已构造，等待 spawn"), "running": ("运行中", "ok", False, "子进程持有控制台输出与独立进程组"), "stopping": ("停止中", "warn", False, "已发送 SIGTERM，等待子进程退出"), - "stopped": ("已手动停止", "warn", True, "收到停止信号后退出，模型可能未保存"), + "stopped": ("已手动停止", "warn", True, "收到停止信号后退出；脚本会在 step 边界保存模型与元数据"), "exited": ("正常结束", "ok", True, "退出码 0：训练与评估流程走完"), "failed": ("异常结束", "bad", True, "非零退出码或启动失败，见日志尾部"), } @@ -812,15 +812,24 @@ def describe_process(snapshot: dict[str, Any]) -> dict[str, Any]: done_steps = value if done_steps == 0.0: done_steps = _finite_float(convergence.get("latest_timesteps")) or 0.0 - percent = (100.0 * done_steps / total_steps) if total_steps > 0 else None + # init/resume（课程迁移）时 num_timesteps 从上一阶段继续，进度必须按本次 + # 增量算；start_timesteps 由训练脚本的进度 JSON 提供，缺省按 0 处理。 + start_steps = 0.0 + for point in metrics: + value = _finite_float(point.get("start_timesteps")) + if value is not None: + start_steps = value + break + trained_steps = max(0.0, done_steps - start_steps) + percent = (100.0 * trained_steps / total_steps) if total_steps > 0 else None if percent is not None: percent = max(0.0, min(100.0, percent)) elapsed = float(snapshot.get("elapsed_s") or 0.0) rate = _latest_rate(metrics) remaining: float | None = None eta_ts: float | None = None - if rate and rate > 0 and total_steps > 0 and done_steps < total_steps: - remaining = (total_steps - done_steps) / rate + if rate and rate > 0 and total_steps > 0 and trained_steps < total_steps: + remaining = (total_steps - trained_steps) / rate eta_ts = float(snapshot.get("finished_at") or time.time()) + remaining if not terminal and remaining is None: cap = _finite_float(config.get("max_wall_s")) or 0.0 @@ -835,6 +844,8 @@ def describe_process(snapshot: dict[str, Any]) -> dict[str, Any]: "timeline": timeline, "percent": percent, "done_steps": done_steps, + "start_steps": start_steps, + "trained_steps": trained_steps, "total_steps": total_steps, "elapsed_s": round(elapsed, 3), "rate_per_s": rate, diff --git a/src/mcm26b/training_ui/static/app.js b/src/mcm26b/training_ui/static/app.js index bc1bf48..ccdd744 100644 --- a/src/mcm26b/training_ui/static/app.js +++ b/src/mcm26b/training_ui/static/app.js @@ -347,7 +347,7 @@ function renderProcess(job) { chip("chip-progress", "进度", percent === null || percent === undefined ? "—" : `${percent.toFixed(1)}%`, percent !== null && percent >= 100 ? "ok" : "cyan"); - setText("steps-done", fmtInt(process.done_steps)); + setText("steps-done", fmtInt(process.trained_steps ?? process.done_steps)); setText("steps-total", fmtInt(process.total_steps)); setText("progress-percent", percent === null || percent === undefined ? "—" : `${percent.toFixed(2)}%`); const fill = $("bar-fill"); @@ -428,7 +428,7 @@ function renderBanner(job) { if (job.status === "failed") { showBanner("bad", `训练异常结束（退出码 ${job.returncode === null ? "?" : job.returncode}）：${escapeHtml(job.error || "见日志尾部")}。`); } else if (job.status === "stopped") { - showBanner("warn", "已按请求停止训练（SIGTERM 独立进程组），本轮模型可能未保存。"); + showBanner("warn", "已按请求停止训练（SIGTERM 独立进程组），脚本会在 step 边界保存当前模型与元数据。"); } else if (job.status === "exited") { showBanner("ok", "训练与评估流程正常结束，模型与元数据已写入输出目录。"); } else if (process.remaining_s !== null && process.remaining_s !== undefined @@ -708,7 +708,7 @@ $("config-form").addEventListener("submit", async (event) => { }); $("btn-stop").addEventListener("click", async () => { - if (!confirm("确认停止当前训练？\n将向独立进程组发送 SIGTERM，本轮模型可能未保存。")) return; + if (!confirm("确认停止当前训练？\n将向独立进程组发送 SIGTERM，脚本会先保存当前模型与元数据再退出。")) return; const res = await fetch("/api/train/stop", { method: "POST", headers: { "Content-Type": "application/json" }, body: "{}", }); diff --git a/tests/test_hrl_checkpoint.py b/tests/test_hrl_checkpoint.py new file mode 100644 index 0000000..0d7ddb9 --- /dev/null +++ b/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 diff --git a/tests/test_training_ui.py b/tests/test_training_ui.py index d07817b..afe76db 100644 --- a/tests/test_training_ui.py +++ b/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)