278 lines
9.7 KiB
Python
278 lines
9.7 KiB
Python
"""step 级续训 checkpoint：_TrainState 可持久化字段的序列化 / 反序列化。
|
||
|
||
_TrainState 的累加包均为扁平纯数据 dataclass，经 dataclasses.asdict 序列化为
|
||
纯 JSON dict；反序列化时用 Cls(**d) 还原，其中 SystemCasePack 含嵌套 CaseSample
|
||
列表、Probation 含嵌套 RejectedEdit 列表，需逐个重建。
|
||
|
||
不持久化的字段：gate_pools / baseline_cache（各自文件级自持久化，resume 时按
|
||
指纹重载）、best_*（从 manifest best 指针读）、global_step（存 progress 块，
|
||
由 train 单独赋值）。gate_epoch_observed 持久化：warm p-hat 在 gate_pools.json
|
||
幸存，观测开关须随行，否则 resume 后阶梯排序回退冷启动序。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
from dataclasses import asdict, dataclass, field
|
||
from typing import TYPE_CHECKING, Any
|
||
|
||
from core.evolution.types import (
|
||
CaseSample,
|
||
RejectedEdit,
|
||
SystemCasePack,
|
||
ToolCasePack,
|
||
)
|
||
|
||
if TYPE_CHECKING:
|
||
from pathlib import Path
|
||
|
||
CHECKPOINT_SCHEMA_VERSION = 1
|
||
|
||
|
||
from app.harness.validate import Probation # noqa: E402
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 结构性 / 决策性指纹键
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_STRUCTURAL_KEYS = (
|
||
"batch_size",
|
||
"min_class_per_batch",
|
||
"epochs",
|
||
"diag_size",
|
||
"val_size",
|
||
"batch_correct_ratio",
|
||
)
|
||
|
||
_DECISION_KEYS = (
|
||
"edit_budget_start",
|
||
"edit_budget_end",
|
||
"early_stop_patience",
|
||
"use_slow_momentum",
|
||
"skill_update_mode",
|
||
"appendix_consolidate_threshold",
|
||
"momentum_samples",
|
||
"gate_e_confirm",
|
||
"gate_e_provisional",
|
||
"gate_w_net_min",
|
||
"gate_delta_min",
|
||
"gate_lambda_dir",
|
||
"gate_e_rollback",
|
||
"gate_block",
|
||
"gate_n_max",
|
||
"gate_p_low",
|
||
"gate_p_high",
|
||
"gate_probe_quota",
|
||
"gate_gamma_decay",
|
||
"gate_cooldown_steps",
|
||
"gate_guard_err",
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 序列化 / 反序列化
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def serialize_state(state: Any) -> dict[str, Any]:
|
||
"""把 _TrainState 的可持久化字段转为纯 JSON dict。
|
||
|
||
参数:
|
||
state: _TrainState 实例（duck-typed，仅需含可持久化字段）。
|
||
|
||
返回:
|
||
纯 JSON 可序列化的 dict，不含 gate_pools / baseline_cache /
|
||
best_* / global_step。
|
||
|
||
关键实现细节:
|
||
- changed_task_types_this_epoch 是 set，JSON 无 set，故 sorted 成有序列表。
|
||
- dataclass 均经 asdict 递归转 dict（含 SystemCasePack 嵌套 CaseSample、
|
||
Probation 嵌套 RejectedEdit）。
|
||
"""
|
||
return {
|
||
"correctness": state.correctness,
|
||
"eval_prev_acc": state.eval_prev_acc,
|
||
"eval_prev_run_id": state.eval_prev_run_id,
|
||
"baseline_skills_version": state.baseline_skills_version,
|
||
"baseline_prompts_version": state.baseline_prompts_version,
|
||
"steps_since_best_improved": state.steps_since_best_improved,
|
||
"epoch_start_skills": state.epoch_start_skills,
|
||
"changed_task_types_this_epoch": sorted(state.changed_task_types_this_epoch),
|
||
"rejected_buffer": {k: [asdict(x) for x in v] for k, v in state.rejected_buffer.items()},
|
||
"system_packs": [asdict(x) for x in state.system_packs],
|
||
"tool_packs": [asdict(x) for x in state.tool_packs],
|
||
"probations": {t: asdict(p) for t, p in state.probations.items()},
|
||
"gate_cooldown": state.gate_cooldown,
|
||
"gate_epoch_observed": state.gate_epoch_observed,
|
||
}
|
||
|
||
|
||
def _restore_system_pack(d: dict[str, Any]) -> SystemCasePack:
|
||
"""还原 SystemCasePack，含嵌套 CaseSample 列表。
|
||
|
||
参数:
|
||
d: asdict(SystemCasePack) 产出的 dict。
|
||
|
||
返回:
|
||
复活的 SystemCasePack；failure_cases / success_cases 重建为 CaseSample 实例。
|
||
"""
|
||
return SystemCasePack(
|
||
stats=d["stats"],
|
||
failure_cases=[CaseSample(**c) for c in d["failure_cases"]],
|
||
success_cases=[CaseSample(**c) for c in d["success_cases"]],
|
||
)
|
||
|
||
|
||
def deserialize_state_fields(d: dict[str, Any]) -> dict[str, Any]:
|
||
"""把序列化 dict 还原为可填入 _TrainState 的字段字典（dataclass 复活）。
|
||
|
||
参数:
|
||
d: serialize_state 产出并经 JSON 往返的 dict。
|
||
|
||
返回:
|
||
字段名 -> 值的 dict，可直接铺到 _TrainState；其中各 dataclass 已复活、
|
||
changed_task_types_this_epoch 还原为 set。
|
||
|
||
关键实现细节:
|
||
- RejectedEdit / ToolCasePack 字段均为标量/dict/list[dict]，Cls(**d) 直接构造。
|
||
- SystemCasePack 含嵌套 CaseSample，交由 _restore_system_pack 重建。
|
||
- Probation 含嵌套 RejectedEdit 列表（pending_edits），先重建内层再构造外层。
|
||
- 直接取 d[...] 不用 .get 兜底：serialize 后的 checkpoint 必带全部键，
|
||
缺键即 checkpoint 损坏，应硬失败（P5 不掩盖）。
|
||
"""
|
||
return {
|
||
"correctness": d["correctness"],
|
||
"eval_prev_acc": d["eval_prev_acc"],
|
||
"eval_prev_run_id": d["eval_prev_run_id"],
|
||
"baseline_skills_version": d["baseline_skills_version"],
|
||
"baseline_prompts_version": d["baseline_prompts_version"],
|
||
"steps_since_best_improved": d["steps_since_best_improved"],
|
||
"epoch_start_skills": d["epoch_start_skills"],
|
||
"changed_task_types_this_epoch": set(d["changed_task_types_this_epoch"]),
|
||
"rejected_buffer": {
|
||
k: [RejectedEdit(**x) for x in v] for k, v in d["rejected_buffer"].items()
|
||
},
|
||
"system_packs": [_restore_system_pack(x) for x in d["system_packs"]],
|
||
"tool_packs": [ToolCasePack(**x) for x in d["tool_packs"]],
|
||
"probations": {
|
||
t: Probation(
|
||
**{
|
||
**d_p,
|
||
"pending_edits": [RejectedEdit(**x) for x in d_p["pending_edits"]],
|
||
}
|
||
)
|
||
for t, d_p in d["probations"].items()
|
||
},
|
||
"gate_cooldown": d["gate_cooldown"],
|
||
"gate_epoch_observed": d["gate_epoch_observed"],
|
||
}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 配置指纹
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def compute_fingerprint(config: Any) -> dict[str, Any]:
|
||
"""采集影响训练轨迹的配置项（结构性 + 决策性）。
|
||
|
||
参数:
|
||
config: 训练配置对象（duck-typed，需含 _STRUCTURAL_KEYS + _DECISION_KEYS 属性）。
|
||
|
||
返回:
|
||
指纹 dict，键为配置项名，值为对应配置值。
|
||
"""
|
||
return {k: getattr(config, k) for k in _STRUCTURAL_KEYS + _DECISION_KEYS}
|
||
|
||
|
||
def check_fingerprint(saved: dict[str, Any], config: Any) -> tuple[list[str], list[str]]:
|
||
"""比对保存的指纹与当前配置。返回 (结构性不一致项, 决策性不一致项)。
|
||
|
||
参数:
|
||
saved: checkpoint 中保存的 config_fingerprint。
|
||
config: 当前训练配置对象。
|
||
|
||
返回:
|
||
(structural, decision) 两个不一致项名列表。
|
||
|
||
关键实现细节:
|
||
结构性不一致（batch_size/min_class_per_batch/epochs/diag_size/val_size/
|
||
batch_correct_ratio）→ 调用方应拒绝 resume；决策性不一致 → 仅告警放行。
|
||
"""
|
||
cur = compute_fingerprint(config)
|
||
structural = [k for k in _STRUCTURAL_KEYS if saved.get(k) != cur[k]]
|
||
decision = [k for k in _DECISION_KEYS if saved.get(k) != cur[k]]
|
||
return structural, decision
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 读写 checkpoint
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def write_checkpoint(
|
||
workspace_dir: Path,
|
||
*,
|
||
state: Any,
|
||
epoch: int,
|
||
step_completed: int,
|
||
phase: str,
|
||
global_step: int,
|
||
total_steps: int,
|
||
version_snapshot: dict[str, str],
|
||
epoch_batches: list[list[str]],
|
||
config: Any,
|
||
) -> None:
|
||
"""原子写 checkpoint.json（.tmp 再 os.replace）。
|
||
|
||
参数:
|
||
workspace_dir: workspace 目录，checkpoint.json 写入其下。
|
||
state: _TrainState 实例，交由 serialize_state 序列化。
|
||
epoch: 当前 epoch 序号。
|
||
step_completed: 本 epoch 内已完成的 step 数。
|
||
phase: 续训阶段标识（如 "in_epoch"）。
|
||
global_step: 全局 step 序号。
|
||
total_steps: 全局总 step 数。
|
||
version_snapshot: skills/prompts 版本快照。
|
||
epoch_batches: 本 epoch 的 batch 划分（question_id 列表的列表）。
|
||
config: 训练配置对象，用于计算 config_fingerprint。
|
||
|
||
关键实现细节:
|
||
先写 checkpoint.json.tmp 再 os.replace，保证 checkpoint 不被写一半的中断破坏。
|
||
"""
|
||
payload = {
|
||
"schema_version": CHECKPOINT_SCHEMA_VERSION,
|
||
"progress": {
|
||
"epoch": epoch,
|
||
"step_completed": step_completed,
|
||
"phase": phase,
|
||
"global_step": global_step,
|
||
"total_steps": total_steps,
|
||
},
|
||
"version_snapshot": version_snapshot,
|
||
"epoch_batches": epoch_batches,
|
||
"config_fingerprint": compute_fingerprint(config),
|
||
"state": serialize_state(state),
|
||
}
|
||
path = workspace_dir / "checkpoint.json"
|
||
tmp = path.with_name("checkpoint.json.tmp")
|
||
tmp.write_text(json.dumps(payload, ensure_ascii=False, indent=2))
|
||
os.replace(tmp, path)
|
||
|
||
|
||
def load_checkpoint(workspace_dir: Path) -> dict[str, Any] | None:
|
||
"""读 checkpoint.json，不存在返回 None。
|
||
|
||
参数:
|
||
workspace_dir: workspace 目录。
|
||
|
||
返回:
|
||
checkpoint payload dict；checkpoint.json 不存在时返回 None。
|
||
"""
|
||
path = workspace_dir / "checkpoint.json"
|
||
if not path.exists():
|
||
return None
|
||
return json.loads(path.read_text())
|