This commit is contained in:
Horldsence
2026-09-11 23:07:42 +08:00
parent 5c96699b70
commit 36f354736c
8 changed files with 522 additions and 23 deletions
Vendored
BIN
View File
Binary file not shown.
+18
View File
@@ -80,6 +80,24 @@ uv run python -m mcm26b.training_ui --print-preview # 只看离线示例快照
- 产物路径限定在仓库根下解析,模型输出沿用训练脚本自身的 `--force`/`--resume` - 产物路径限定在仓库根下解析,模型输出沿用训练脚本自身的 `--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 ```bash
+251 -15
View File
@@ -7,6 +7,15 @@
- 低层默认使用几何选点;``--local-selector rl`` 仅作为接口预留并记录, - 低层默认使用几何选点;``--local-selector rl`` 仅作为接口预留并记录,
当前 HTTP 执行路径仍由高层宏动作终点决定,不把八方向推演当正式收益。 当前 HTTP 执行路径仍由高层宏动作终点决定,不把八方向推演当正式收益。
落盘保证(长训必需):
- ``--checkpoint-every`` 每 N 个环境步在 ``<模型目录>/checkpoints/`` 落一个可续训
的 checkpoint(zip + 同名 json 元数据),只保留最近 ``--checkpoint-keep`` 个;
- 训练前端的「停止」是对**进程组**发 SIGTERM,因此 subproc worker 在 :func:`_make_env`
里显式忽略 SIGTERM/SIGINT,主进程收到信号后在 step 边界收尾并保存模型,
不会被 worker 先行退出打断;
- 任何异常(含被强杀之外的失败)都会先落一份 ``*_emergency.zip`` 再抛出,
不再出现"跑了很久什么都没有"。
P1 smoke 示例: P1 smoke 示例:
python scripts/train_hrl_high.py \ python scripts/train_hrl_high.py \
--problem 3 --stage 1 --timesteps 2048 --seed 320001 \ --problem 3 --stage 1 --timesteps 2048 --seed 320001 \
@@ -23,6 +32,7 @@ import json
import os import os
import platform import platform
import random import random
import signal
import sys import sys
import time import time
from pathlib import Path from pathlib import Path
@@ -58,6 +68,23 @@ def _reward_from_profile(profile: str):
raise ValueError(f"未知 reward-profile: {profile}") 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, def _make_env(problem: int, stage_index: int, seed: int,
monitor_path: str | None, profile: str, monitor_path: str | None, profile: str,
reward_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 from mcm26b.hrl.hierarchical_env import HierarchicalEnv
_ignore_stop_signals_in_worker()
stage = _stage_from_index(stage_index) stage = _stage_from_index(stage_index)
env = HierarchicalEnv( env = HierarchicalEnv(
problem=problem, problem=problem,
@@ -98,6 +126,41 @@ def _hash_file(path: Path) -> str:
return hashlib.sha256(path.read_bytes()).hexdigest() 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, def _model_metadata(*, args, wall_s: float, n_envs: int,
actual_timesteps: int) -> dict[str, Any]: actual_timesteps: int) -> dict[str, Any]:
import numpy import numpy
@@ -130,6 +193,9 @@ def _model_metadata(*, args, wall_s: float, n_envs: int,
"reward_version": REWARD_VERSION, "reward_version": REWARD_VERSION,
"local_selector": args.local_selector, "local_selector": args.local_selector,
"local_model": args.local_model, "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, "device": args.device,
"wall_seconds": round(wall_s, 3), "wall_seconds": round(wall_s, 3),
"max_wall_s": float(args.max_wall_s), "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, def _evaluate_env(model, *, problem: int, stage_index: int,
episodes: int, seed: int, profile: str, episodes: int, seed: int, profile: str,
reward_profile: str) -> dict[str, Any]: reward_profile: str) -> dict[str, Any]:
@@ -288,6 +388,52 @@ def run_training(args) -> int:
if args.init: if args.init:
pass # 已在上面处理;此分支仅保持语义清晰 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): class _WallClockStop(BaseCallback):
def __init__(self, max_wall_s: float): def __init__(self, max_wall_s: float):
super().__init__() super().__init__()
@@ -312,7 +458,9 @@ def run_training(args) -> int:
def __init__(self, every_steps: int): def __init__(self, every_steps: int):
super().__init__() super().__init__()
self.every_steps = max(1, int(every_steps)) 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 self.emitted = 0
def _next_gap(self) -> int: def _next_gap(self) -> int:
@@ -350,6 +498,7 @@ def run_training(args) -> int:
payload = { payload = {
"event": "progress", "event": "progress",
"timesteps": current, "timesteps": current,
"start_timesteps": self.start,
"n_updates": int(getattr(self.model, "_n_updates", 0) or 0), "n_updates": int(getattr(self.model, "_n_updates", 0) or 0),
"ep_rew_mean": (sum(ep_rewards) / len(ep_rewards) "ep_rew_mean": (sum(ep_rewards) / len(ep_rewards)
if ep_rewards else None), if ep_rewards else None),
@@ -366,43 +515,119 @@ def run_training(args) -> int:
flush=True) flush=True)
return 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 progress_every = (args.progress_every if args.progress_every > 0
else args.rollout_steps * args.envs) 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: if args.progress_json:
callbacks.insert(0, _ProgressJSON(progress_every)) callbacks.insert(0, _ProgressJSON(progress_every))
started = time.monotonic()
actual_timesteps = int(model.num_timesteps) actual_timesteps = int(model.num_timesteps)
try: try:
model.learn(total_timesteps=args.timesteps, model.learn(total_timesteps=args.timesteps,
callback=callbacks if len(callbacks) > 1 else callbacks[0], callback=callbacks if len(callbacks) > 1 else callbacks[0],
progress_bar=False, progress_bar=False,
reset_num_timesteps=not (args.resume or args.init)) reset_num_timesteps=not (args.resume or args.init))
except Exception as exc: # 训练失败也必须保存失败上下文,不静默 except BaseException as exc: # 训练失败也必须留下可续训产物,不静默
failure = { failure = {
"event": "training_failed", "event": "training_failed",
"error": f"{type(exc).__name__}: {exc}", "error": f"{type(exc).__name__}: {exc}",
"model": str(model_path), "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( (run_dir / "training_failure.json").write_text(
json.dumps(failure, ensure_ascii=False, indent=2), encoding="utf-8") json.dumps(failure, ensure_ascii=False, indent=2), encoding="utf-8")
print(json.dumps(failure, ensure_ascii=False), file=sys.stderr,
flush=True)
raise raise
wall_s = time.monotonic() - started wall_s = time.monotonic() - started
actual_timesteps = int(model.num_timesteps) actual_timesteps = int(model.num_timesteps)
stopped_by_signal = bool(stop_state["requested"])
model.save(str(model_path)) model.save(str(model_path))
metadata = _model_metadata(args=args, wall_s=wall_s, metadata = _write_metadata(
n_envs=args.envs, model_path=model_path, args=args, wall_s=wall_s, n_envs=args.envs,
actual_timesteps=actual_timesteps) actual_timesteps=actual_timesteps, monitor_paths=monitor_paths,
metadata["model_sha256"] = _hash_file(model_path) model=model,
metadata["monitor_files"] = monitor_paths extra={"stopped_by_signal": stopped_by_signal,
metadata["reward_note"] = ( "stop_signal": stop_state["signal"],
"reward_profile=time 使用 HRLReward.time_only(P0.4 纯时间口径);" "checkpoint_files": (checkpoint_cb.saved
"reward_profile=prompt 使用原始 §5.2 塑形。模型元数据已固化口径。") if checkpoint_cb is not None else [])})
metadata_path = model_path.with_suffix(".json") 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() 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( evaluation = _evaluate_env(
model, problem=args.problem, stage_index=args.stage, model, problem=args.problem, stage_index=args.stage,
@@ -439,7 +664,8 @@ def main(argv=None) -> int:
choices=("production", "coarse-pretrain"), choices=("production", "coarse-pretrain"),
default="production") default="production")
parser.add_argument("--model", required=True) 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, parser.add_argument("--init", default=None,
help="可选:Q3→Q4 或继续预训练权重路径") help="可选:Q3→Q4 或继续预训练权重路径")
parser.add_argument("--force", action="store_true") parser.add_argument("--force", action="store_true")
@@ -461,6 +687,12 @@ def main(argv=None) -> int:
help="向训练前端输出单行 JSON 进度(UI 自动启用)") help="向训练前端输出单行 JSON 进度(UI 自动启用)")
parser.add_argument("--progress-every", type=int, default=0, parser.add_argument("--progress-every", type=int, default=0,
help="JSON 进度间隔环境步数;0=每个 rollout") 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) parser.add_argument("--verbose", type=int, default=1)
args = parser.parse_args(argv) args = parser.parse_args(argv)
if args.timesteps <= 0 or args.envs <= 0: if args.timesteps <= 0 or args.envs <= 0:
@@ -469,6 +701,10 @@ def main(argv=None) -> int:
parser.error("PPO rollout 至少需要 2 个样本") parser.error("PPO rollout 至少需要 2 个样本")
if args.eval_episodes < 0: if args.eval_episodes < 0:
parser.error("eval-episodes 不能为负") 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) return run_training(args)
+2 -1
View File
@@ -144,7 +144,8 @@ CONFIG_FIELDS: list[dict[str, Any]] = [
{ {
"key": "init", "kind": "text", "group": "paths", "optional": True, "key": "init", "kind": "text", "group": "paths", "optional": True,
"label": "初始权重 --init", "default": None, "label": "初始权重 --init", "default": None,
"help": "留空=从零;填上一阶段 .zip 即课程迁移(步数计数重置)。", "help": "留空=从零;填上一阶段 .zip 即课程迁移。PPO 超参与步数计数沿用"
"该模型(timesteps 按增量理解,进度条按增量显示)。",
}, },
{ {
"key": "local_model", "kind": "text", "group": "paths", "optional": True, "key": "local_model", "kind": "text", "group": "paths", "optional": True,
+15 -4
View File
@@ -739,7 +739,7 @@ _STATUS_VIEW = {
"created": ("已创建", "gray", False, "命令已构造,等待 spawn"), "created": ("已创建", "gray", False, "命令已构造,等待 spawn"),
"running": ("运行中", "ok", False, "子进程持有控制台输出与独立进程组"), "running": ("运行中", "ok", False, "子进程持有控制台输出与独立进程组"),
"stopping": ("停止中", "warn", False, "已发送 SIGTERM,等待子进程退出"), "stopping": ("停止中", "warn", False, "已发送 SIGTERM,等待子进程退出"),
"stopped": ("已手动停止", "warn", True, "收到停止信号后退出,模型可能未保存"), "stopped": ("已手动停止", "warn", True, "收到停止信号后退出;脚本会在 step 边界保存模型与元数据"),
"exited": ("正常结束", "ok", True, "退出码 0:训练与评估流程走完"), "exited": ("正常结束", "ok", True, "退出码 0:训练与评估流程走完"),
"failed": ("异常结束", "bad", True, "非零退出码或启动失败,见日志尾部"), "failed": ("异常结束", "bad", True, "非零退出码或启动失败,见日志尾部"),
} }
@@ -812,15 +812,24 @@ def describe_process(snapshot: dict[str, Any]) -> dict[str, Any]:
done_steps = value done_steps = value
if done_steps == 0.0: if done_steps == 0.0:
done_steps = _finite_float(convergence.get("latest_timesteps")) or 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: if percent is not None:
percent = max(0.0, min(100.0, percent)) percent = max(0.0, min(100.0, percent))
elapsed = float(snapshot.get("elapsed_s") or 0.0) elapsed = float(snapshot.get("elapsed_s") or 0.0)
rate = _latest_rate(metrics) rate = _latest_rate(metrics)
remaining: float | None = None remaining: float | None = None
eta_ts: float | None = None eta_ts: float | None = None
if rate and rate > 0 and total_steps > 0 and done_steps < total_steps: if rate and rate > 0 and total_steps > 0 and trained_steps < total_steps:
remaining = (total_steps - done_steps) / rate remaining = (total_steps - trained_steps) / rate
eta_ts = float(snapshot.get("finished_at") or time.time()) + remaining eta_ts = float(snapshot.get("finished_at") or time.time()) + remaining
if not terminal and remaining is None: if not terminal and remaining is None:
cap = _finite_float(config.get("max_wall_s")) or 0.0 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, "timeline": timeline,
"percent": percent, "percent": percent,
"done_steps": done_steps, "done_steps": done_steps,
"start_steps": start_steps,
"trained_steps": trained_steps,
"total_steps": total_steps, "total_steps": total_steps,
"elapsed_s": round(elapsed, 3), "elapsed_s": round(elapsed, 3),
"rate_per_s": rate, "rate_per_s": rate,
+3 -3
View File
@@ -347,7 +347,7 @@ function renderProcess(job) {
chip("chip-progress", "进度", chip("chip-progress", "进度",
percent === null || percent === undefined ? "—" : `${percent.toFixed(1)}%`, percent === null || percent === undefined ? "—" : `${percent.toFixed(1)}%`,
percent !== null && percent >= 100 ? "ok" : "cyan"); 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("steps-total", fmtInt(process.total_steps));
setText("progress-percent", percent === null || percent === undefined ? "—" : `${percent.toFixed(2)}%`); setText("progress-percent", percent === null || percent === undefined ? "—" : `${percent.toFixed(2)}%`);
const fill = $("bar-fill"); const fill = $("bar-fill");
@@ -428,7 +428,7 @@ function renderBanner(job) {
if (job.status === "failed") { if (job.status === "failed") {
showBanner("bad", `训练异常结束(退出码 <b>${job.returncode === null ? "?" : job.returncode}</b>):${escapeHtml(job.error || "见日志尾部")}。`); showBanner("bad", `训练异常结束(退出码 <b>${job.returncode === null ? "?" : job.returncode}</b>):${escapeHtml(job.error || "见日志尾部")}。`);
} else if (job.status === "stopped") { } else if (job.status === "stopped") {
showBanner("warn", "已按请求停止训练(SIGTERM 独立进程组),本轮模型可能未保存。"); showBanner("warn", "已按请求停止训练(SIGTERM 独立进程组),脚本会在 step 边界保存当前模型与元数据。");
} else if (job.status === "exited") { } else if (job.status === "exited") {
showBanner("ok", "训练与评估流程正常结束,模型与元数据已写入输出目录。"); showBanner("ok", "训练与评估流程正常结束,模型与元数据已写入输出目录。");
} else if (process.remaining_s !== null && process.remaining_s !== undefined } 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 () => { $("btn-stop").addEventListener("click", async () => {
if (!confirm("确认停止当前训练?\n将向独立进程组发送 SIGTERM,本轮模型可能未保存。")) return; if (!confirm("确认停止当前训练?\n将向独立进程组发送 SIGTERM,脚本会先保存当前模型与元数据再退出。")) return;
const res = await fetch("/api/train/stop", { const res = await fetch("/api/train/stop", {
method: "POST", headers: { "Content-Type": "application/json" }, body: "{}", method: "POST", headers: { "Content-Type": "application/json" }, body: "{}",
}); });
+216
View File
@@ -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
+17
View File
@@ -424,6 +424,23 @@ def test_describe_process_progress_and_eta():
assert ended["remaining_s"] is None 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(): def test_demo_snapshot_contract():
for scenario in demo.DEMO_SCENARIOS: for scenario in demo.DEMO_SCENARIOS:
snapshot = demo.demo_snapshot(scenario=scenario) snapshot = demo.demo_snapshot(scenario=scenario)