curriculum_loader.py

1"""CPU gather for composite curriculum. Schedule entries already partition DDP ranks."""
2import json
3from pathlib import Path
4import numpy as np
5from build_curriculum import DTYPE,sha
6from records import Records
7from runtime_records import Records as BoundedRecords
8SCHEDULE_DTYPE=np.dtype([('kind','u1'),('source','u1'),('row','<u4'),('exposure','u1')])
9
10def apply_tie(row,metadata,exposure):
11    choices=metadata['tied_target_triples'];target=choices[int(exposure)%len(choices)]
12    row['y']=target;row['x'][-2:]=target[:2]
13    return row
14
15class Curriculum:
16    def __init__(self,root,verify=True):
17        self.root=Path(root);self.manifest=json.loads((self.root/'manifest.json').read_text())
18        m=self.manifest
19        if m['readiness']!='independently_reviewed_ready':raise ValueError('curriculum not approved')
20        for f in m['dependencies']:
21            if sha(Path(f['path']))!=f['sha256']:raise ValueError('dependency changed')
22        if verify and sha(self.root/'schedule.bin')!=m['schedule_sha256']:raise ValueError('schedule changed')
23        self.schedule=np.memmap(self.root/'schedule.bin',dtype=SCHEDULE_DTYPE,mode='r')
24        if len(self.schedule)!=m['schedule_entries']:raise ValueError('schedule length')
25        self.sources={(kind,source):(BoundedRecords if kind=='ordinary' else Records)(m['roots'][kind],source,0,verify=verify) for kind in ('ordinary','theory') for source in (0,1)}
26        theory_root=Path(m['roots']['theory']);theory_manifest=json.loads((theory_root/'manifest.json').read_text())
27        if sha(theory_root/'metadata.jsonl')!=theory_manifest['metadata_sha256']:raise ValueError('tie metadata changed')
28        self.metadata=[json.loads(x) for x in (theory_root/'metadata.jsonl').read_text().splitlines()]
29    def take(self,indices):
30        ids=np.asarray(indices,dtype=np.int64)
31        if ids.ndim!=1 or np.any(ids<0) or np.any(ids>=len(self.schedule)):raise IndexError('schedule index')
32        entries=self.schedule[ids];out=np.empty(len(entries),dtype=DTYPE)
33        for kind,name in ((0,'ordinary'),(1,'theory')):
34            for src in (0,1):
35                mask=(entries['kind']==kind)&(entries['source']==src)
36                if not mask.any():continue
37                out[mask]=self.sources[name,src].take(entries['row'][mask])
38                if kind:
39                    for i in np.flatnonzero(mask):out[i]=apply_tie(out[i],self.metadata[int(out[i]['game'])],entries[i]['exposure'])
40        return out
41    def rank_indices(self,start,global_batch,rank,world):
42        if world<1 or rank<0 or rank>=world or start<0 or global_batch<1 or global_batch%world:raise ValueError('invalid rank/batch')
43        if start+global_batch>len(self.schedule):return np.array([],dtype=np.int64)
44        stop=start+global_batch
45        # Drop final partial global batch, never pad with repeated ordinary rows.
46        return np.arange(start+rank,stop,world,dtype=np.int64)
47