checkpoint.py

1"""Committed-cursor resume contract, including per-state exposure reconciliation."""
2import random,hashlib
3import numpy as np
4import torch
5from schedule import counters
6def capture_rng(device=None):
7    n=np.random.get_state()
8    return {'python':random.getstate(),'numpy':(n[0],n[1].tolist(),n[2],n[3],n[4]),'torch':torch.get_rng_state(),
9            'cuda':None if device is None else torch.cuda.get_rng_state(device)}
10def restore_rng(value,device=None):
11    random.setstate(value['python']);n=value['numpy'];np.random.set_state((n[0],np.asarray(n[1],dtype=np.uint32),n[2],n[3],n[4]))
12    torch.set_rng_state(value['torch'])
13    if device is not None:
14        if value['cuda'] is None:raise ValueError('missing CUDA RNG')
15        torch.cuda.set_rng_state(value['cuda'],device)
16def exposure_state(schedule,cursor,n):
17    values=counters(schedule,cursor,n)
18    return {'cursor':cursor,'counts':torch.from_numpy(values.copy()),'counts_sha256':hashlib.sha256(values.tobytes()).hexdigest(),
19            'unique_states_seen':int(np.count_nonzero(values)),'decision_exposures':int(values.astype(np.uint64).sum())}
20def validate_resume(saved,identity,schedule,n,deadline):
21    if saved['format']!='cm6-curriculum-checkpoint-v1' or saved['identity']!=identity:raise ValueError('checkpoint identity mismatch')
22    if saved['terminal']:raise ValueError('terminal checkpoint cannot resume')
23    if saved['deadline']!=deadline:raise ValueError('original deadline mismatch')
24    if len(saved['rank_rng'])!=2:raise ValueError('both-rank RNG required')
25    expected=exposure_state(schedule,saved['exposures']['cursor'],n)
26    if saved['exposures']['counts_sha256']!=expected['counts_sha256'] or not torch.equal(saved['exposures']['counts'],expected['counts']):raise ValueError('exposure counters/cursor mismatch')
27    for key in ('decision_exposures','unique_states_seen'):
28        if saved['exposures'][key]!=expected[key]:raise ValueError('exposure aggregate mismatch')
29    return expected['cursor']
30def breakdown(meta,schedule,cursor):
31    """Exposure and unique-state totals; repeated theory is never called unique."""
32    counts=counters(schedule,cursor,len(meta));out={}
33    for source in (0,1):
34        for phase in (0,1,2):
35            for kind in (0,1,2):
36                mask=(meta['source']==source)&(meta['phase']==phase)&(meta['kind']==kind)
37                out[f'{source}/{phase}/{kind}']={'unique_states':int(np.count_nonzero(counts[mask])),
38                    'decision_exposures':int(counts[mask].astype(np.uint64).sum()),'prediction_targets':3*int(counts[mask].astype(np.uint64).sum())}
39    return out
40