iomgaa
6c6fb576ee
核心算法保真 #3（断点续跑）：checkpoint 从逐题 question_id 改为存 unit_id 序列（孪生对折叠为单个 unit_id），恢复时 build_units + 按完整 unit 展开， 续跑后 pair 两成员同进同出、绝不被劈开。 - _batch_unit_ids/_batch_from_ids 对称折叠/展开，保序去重，纯非 AR 下 unit_id==question_id、与旧逐题序列逐字节一致。 - momentum 采样抽取为 _sample_momentum_candidates 纯函数，docstring 显式 记录 Phase 1 设计偏差：仅保证纯非 AR byte-identical，混格 momentum 不保证。 - 新增 test_checkpoint_pair（unit_id 落盘往返、pair 不拆）与 test_non_ar_byte_identical（pools→batching→checkpoint→momentum 端到端黄金）。
255 lines
9.7 KiB
Python
255 lines
9.7 KiB
Python
"""断点续跑 checkpoint 的 unit_id 粒度契约集成测试（Task 10）。
|
||
|
||
覆盖三条铁律：
|
||
- checkpoint 存 **unit_id 序列**而非逐题 question_id（孪生对折叠为单个 unit_id）；
|
||
- 续跑恢复以 **完整 unit** 展开（pair 两成员同进同出，绝不被劈开）；
|
||
- 真实经 write_checkpoint → load_checkpoint 落盘往返后，unit_id 序列无损存活、
|
||
重建的 batch 与原 batch 逐字节一致（core 算法保真 #3 断点续跑）。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from dataclasses import dataclass, field
|
||
from typing import TYPE_CHECKING
|
||
|
||
from app.harness.batching import build_batches
|
||
from app.harness.checkpoint import load_checkpoint, write_checkpoint
|
||
from app.harness.pools import Pools
|
||
from app.harness.question_units import build_units, validate_units
|
||
from app.harness.runner import _batch_from_ids, _batch_unit_ids
|
||
from core.types import GeneratedQuestion
|
||
|
||
if TYPE_CHECKING:
|
||
from pathlib import Path
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 辅助构造
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _make_single(qid: str, task_type: str = "RETRIEVAL", video_id: str = "v1") -> GeneratedQuestion:
|
||
"""构造最小 single 题目。"""
|
||
return GeneratedQuestion(
|
||
question_id=qid,
|
||
video_id=video_id,
|
||
task_type=task_type,
|
||
question=f"q_{qid}",
|
||
options=("A. a", "B. b", "C. c", "D. d"),
|
||
answer="A",
|
||
source_nodes=("n1",),
|
||
difficulty="medium",
|
||
question_role="single",
|
||
)
|
||
|
||
|
||
def _make_pair(
|
||
pair_id: str, task_type: str = "AR", video_id: str = "v1"
|
||
) -> tuple[GeneratedQuestion, GeneratedQuestion]:
|
||
"""构造一个孪生对（original + mirror），共享 pair_id / flip_axis。"""
|
||
common = {
|
||
"video_id": video_id,
|
||
"task_type": task_type,
|
||
"question": "?",
|
||
"options": ("A. a", "B. b", "C. c", "D. d"),
|
||
"answer": "A",
|
||
"source_nodes": ("n1",),
|
||
"difficulty": "medium",
|
||
"pair_id": pair_id,
|
||
"flip_axis": "before_after",
|
||
}
|
||
original = GeneratedQuestion(
|
||
question_id=f"{pair_id}_o", question_role="pair_original", **common
|
||
)
|
||
mirror = GeneratedQuestion(question_id=f"{pair_id}_m", question_role="pair_mirror", **common)
|
||
return original, mirror
|
||
|
||
|
||
def _pools_with(diagnosis: list[GeneratedQuestion]) -> Pools:
|
||
"""用给定诊断池构造最小 Pools（其余池置空，_batch_from_ids 只读 diagnosis）。"""
|
||
return Pools(
|
||
diagnosis=diagnosis,
|
||
validation=[],
|
||
test=[],
|
||
baseline_run_id="baseline",
|
||
baseline_val_accuracy=0.0,
|
||
correctness={q.question_id: False for q in diagnosis},
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# checkpoint 落盘所需的最小 _TrainState / RunConfig 替身
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@dataclass
|
||
class _FakeState:
|
||
"""serialize_state 只读的可持久化字段（本测试用空累加包即可）。"""
|
||
|
||
correctness: dict[str, bool] = field(default_factory=dict)
|
||
eval_prev_acc: float = 0.0
|
||
eval_prev_run_id: str = "run-0"
|
||
baseline_skills_version: str = "v1"
|
||
baseline_prompts_version: str = "v1"
|
||
steps_since_best_improved: int = 0
|
||
epoch_start_skills: str = "v1"
|
||
changed_task_types_this_epoch: set[str] = field(default_factory=set)
|
||
rejected_buffer: dict = field(default_factory=dict)
|
||
system_packs: list = field(default_factory=list)
|
||
tool_packs: list = field(default_factory=list)
|
||
probations: dict = field(default_factory=dict)
|
||
gate_cooldown: dict = field(default_factory=dict)
|
||
gate_epoch_observed: dict = field(default_factory=dict)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class _FakeConfig:
|
||
"""compute_fingerprint 读取的结构性 + 决策性字段。"""
|
||
|
||
batch_size: int = 4
|
||
min_class_per_batch: int = 2
|
||
epochs: int = 3
|
||
diag_size: int = 30
|
||
val_size: int = 50
|
||
batch_correct_ratio: float = 0.0
|
||
edit_budget_start: int = 6
|
||
edit_budget_end: int = 3
|
||
early_stop_patience: int = 3
|
||
use_slow_momentum: bool = True
|
||
skill_update_mode: str = "patch"
|
||
appendix_consolidate_threshold: int = 10
|
||
momentum_samples: int = 20
|
||
gate_e_confirm: float = 20.0
|
||
gate_e_provisional: float = 6.0
|
||
gate_w_net_min: int = 2
|
||
gate_delta_min: float = 0.02
|
||
gate_lambda_dir: float = -3.0
|
||
gate_e_rollback: float = 10.0
|
||
gate_block: int = 4
|
||
gate_n_max: int = 40
|
||
gate_p_low: float = 0.1
|
||
gate_p_high: float = 0.9
|
||
gate_probe_quota: float = 0.2
|
||
gate_gamma_decay: float = 0.9
|
||
gate_cooldown_steps: int = 2
|
||
gate_guard_err: float = 0.3
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# checkpoint 存 unit_id 序列（孪生对折叠为单个 unit_id）
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestCheckpointStoresUnitIds:
|
||
"""一个含孪生对的 batch，checkpoint 序列应折叠为 unit_id（pair 只记一次）。"""
|
||
|
||
def test_pair_folds_to_single_unit_id(self) -> None:
|
||
po, pm = _make_pair("pA")
|
||
s1, s2 = _make_single("s1"), _make_single("s2")
|
||
batch = [s1, po, pm, s2] # pair 两成员相邻
|
||
|
||
unit_ids = _batch_unit_ids(batch)
|
||
|
||
# 4 题折叠为 3 个 unit：s1 / pA(pair) / s2
|
||
assert unit_ids == ["s1", "pA", "s2"]
|
||
# 逐题 question_id 序列（旧行为）会含 4 项且把 pair 拆成两条
|
||
assert [q.question_id for q in batch] == ["s1", "pA_o", "pA_m", "s2"]
|
||
|
||
def test_pure_single_unit_ids_equal_question_ids(self) -> None:
|
||
"""纯 single batch 的 unit_id 序列与逐题 question_id 序列逐字节一致。"""
|
||
batch = [_make_single(f"s{i}") for i in range(5)]
|
||
assert _batch_unit_ids(batch) == [q.question_id for q in batch]
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 恢复以完整 unit 展开：pair 不拆
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestResumeExpandsFullUnit:
|
||
"""按 unit_id 序列恢复时，孪生对整体展开、绝不被劈开。"""
|
||
|
||
def test_pair_not_split_on_resume(self) -> None:
|
||
po, pm = _make_pair("pA")
|
||
s1, s2 = _make_single("s1"), _make_single("s2")
|
||
batch = [s1, po, pm, s2]
|
||
# 诊断池顺序故意与 batch 不同，验证恢复以池折叠的 unit 为准
|
||
pools = _pools_with([s2, pm, s1, po])
|
||
|
||
unit_ids = _batch_unit_ids(batch)
|
||
restored = _batch_from_ids(pools, unit_ids)
|
||
|
||
# 逐字节一致（顺序保留）
|
||
assert [q.question_id for q in restored] == ["s1", "pA_o", "pA_m", "s2"]
|
||
# 恢复后可无损重建为完整孪生对单元
|
||
units = validate_units(build_units(restored))
|
||
pair_units = [u for u in units if u.kind == "pair"]
|
||
assert len(pair_units) == 1
|
||
assert pair_units[0].unit_id == "pA"
|
||
assert pair_units[0].size == 2
|
||
|
||
def test_roundtrip_preserves_order_mixed(self) -> None:
|
||
"""混格 batch 经 unit_id 折叠 → 展开 round-trip 逐字节还原原顺序。"""
|
||
po, pm = _make_pair("pA")
|
||
qo, qm = _make_pair("pB")
|
||
singles = [_make_single(f"s{i}") for i in range(3)]
|
||
batch = [singles[0], po, pm, singles[1], qo, qm, singles[2]]
|
||
pools = _pools_with(list(batch))
|
||
|
||
restored = _batch_from_ids(pools, _batch_unit_ids(batch))
|
||
assert [q.question_id for q in restored] == [q.question_id for q in batch]
|
||
|
||
def test_missing_unit_id_raises(self) -> None:
|
||
"""checkpoint 引用了诊断池不存在的 unit_id 时硬失败（防静默兜底）。"""
|
||
pools = _pools_with([_make_single("s1")])
|
||
try:
|
||
_batch_from_ids(pools, ["s1", "ghost"])
|
||
except KeyError:
|
||
return
|
||
raise AssertionError("引用缺失 unit_id 应触发 KeyError")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 真实落盘往返：write_checkpoint → load_checkpoint → 重建 batch
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestCheckpointPersistenceRoundtrip:
|
||
"""unit_id 序列经真实 checkpoint.json 原子写/读回后无损、重建逐字节一致。"""
|
||
|
||
def test_persisted_unit_ids_rebuild_batches(self, tmp_path: Path) -> None:
|
||
po, pm = _make_pair("pA")
|
||
singles = [_make_single(f"s{i}", task_type="RETRIEVAL") for i in range(6)]
|
||
items = [*singles, po, pm]
|
||
correctness = {q.question_id: False for q in items}
|
||
batches, _ = build_batches(items, correctness, batch_size=4, min_class_per_batch=2, seed=7)
|
||
pools = _pools_with(list(items))
|
||
|
||
# 存 unit_id 序列（train() 落盘路径）
|
||
epoch_batches = [_batch_unit_ids(b) for b in batches]
|
||
write_checkpoint(
|
||
tmp_path,
|
||
state=_FakeState(correctness=correctness),
|
||
epoch=1,
|
||
step_completed=0,
|
||
phase="in_epoch",
|
||
global_step=1,
|
||
total_steps=len(batches),
|
||
version_snapshot={"skills": "skills/v1", "prompts": "prompts/v1"},
|
||
epoch_batches=epoch_batches,
|
||
config=_FakeConfig(),
|
||
)
|
||
|
||
ckpt = load_checkpoint(tmp_path)
|
||
assert ckpt is not None
|
||
rebuilt = [_batch_from_ids(pools, ids) for ids in ckpt["epoch_batches"]]
|
||
|
||
# 落盘往返后逐批逐字节还原
|
||
assert [[q.question_id for q in b] for b in rebuilt] == [
|
||
[q.question_id for q in b] for b in batches
|
||
]
|
||
# 含孪生对的批次恢复后 pair 仍成对
|
||
for b in rebuilt:
|
||
for u in validate_units(build_units(b)):
|
||
if u.kind == "pair":
|
||
assert u.size == 2
|