"""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 from typing import TYPE_CHECKING, Any from app.harness.validate import Probation from core.evolution.types import ( CaseSample, RejectedEdit, SystemCasePack, ToolCasePack, ) if TYPE_CHECKING: from pathlib import Path CHECKPOINT_SCHEMA_VERSION = 1 # --------------------------------------------------------------------------- # 结构性 / 决策性指纹键 # --------------------------------------------------------------------------- _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())