1919810
This commit is contained in:
@@ -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
|
||||
|
||||
+251
-15
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user