resume_state.py

1"""Exact committed-state restoration; no new scheduler, sampler or budget."""
2import json,math
3from pathlib import Path
4import torch
5from reader import sha
6from records import digest
7from checkpoint import restore_rng
8def validate_source(saved,cfg,clock,ledger,plan):
9    original=json.loads(Path(cfg['original_config']).read_text())
10    if sha(cfg['original_config'])!=cfg['original_config_sha256'] or saved['identity']!=original or saved['identity_sha256']!=digest(original):raise ValueError('source identity')
11    if saved['format']!='cm7-theory-recovery-v1' or not saved['terminal'] or saved['retention_failures']:raise ValueError('clean loss-only terminal source required')
12    if saved['deadline']!=json.loads(Path(cfg['source_deadline']).read_text()):raise ValueError('source clock changed')
13    consumed=ledger.verify(saved['exposures'])
14    if consumed!=cfg['source_cursor'] or saved['updates']!=cfg['source_updates'] or saved['microstep']!=cfg['source_microstep'] or plan.cursor(saved['microstep'])!=consumed:raise ValueError('source sampler/counter mismatch')
15    if consumed>=len(ledger.data.schedule):raise ValueError('finite schedule already complete')
16    if len(saved['rank_rng'])!=2 or len(saved['rank_group_totals'])!=2 or sum(float(t[:,0].sum()) for t in saved['rank_group_totals'])!=consumed:raise ValueError('source rank state')
17    if saved['scheduler_state_dict']['last_epoch']!=saved['updates']:raise ValueError('source scheduler step')
18    for state in saved['optimizer_state_dict']['state'].values():
19        if float(state['step'])!=19000+saved['updates'] or not torch.isfinite(state['exp_avg']).all() or not torch.isfinite(state['exp_avg_sq']).all():raise ValueError('source Adam state')
20    if saved['best_checkpoint']!=cfg['prior_selected']:raise ValueError('prior selected reference changed')
21    return consumed
22
23def restore_optimizer_scheduler(opt,sched,saved):
24    # LambdaLR construction can reset lr. Load optimizer AFTER construction,
25    # then scheduler state, preserving the saved group's actual next-step LR.
26    opt.load_state_dict(saved['optimizer_state_dict']);sched.load_state_dict(saved['scheduler_state_dict'])
27    if [g['lr'] for g in opt.param_groups]!=list(sched.get_last_lr()):raise ValueError('restored optimizer/scheduler LR disagreement')
28    return {'matching_state_retained':True,'scheduler_last_epoch':sched.last_epoch,'lr':sched.get_last_lr(),'optimizer_steps':sorted({float(x['step']) for x in opt.state.values()})}
29def select_loss(latest,best):
30    if not math.isfinite(latest['nll']) or not latest['retention']['finite']:raise FloatingPointError('nonfinite validation')
31    return latest['nll']<best
32