1919810

2026-09-11 23:07:42 +08:00
parent 5c96699b70
commit 36f354736c
8 changed files with 522 additions and 23 deletions
@@ -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 <checkpoint.zip>` 直接
+续训（`--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
@@ -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 <checkpoint.zip>`` 直接可用。
+        """
+
+        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)


@@ -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,
@@ -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,
@@ -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", `训练异常结束（退出码 <b>${job.returncode === null ? "?" : job.returncode}</b>）：${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: "{}",
  });
@@ -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
@@ -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)