composite.py

1"""Adapter for builder-owned immutable schedule, including the final partial batch."""
2import sys,math,time,hashlib
3from pathlib import Path
4from concurrent.futures import ThreadPoolExecutor
5import numpy as np
6import torch
7STATE=Path(__file__).resolve().parent.parent/'state-v2'
8sys.path.insert(0,str(STATE))
9from batches import tensors
10def open_curriculum(root,adapter_directory):
11    sys.path.insert(0,str(Path(adapter_directory).resolve()))
12    from curriculum_loader import Curriculum
13    return Curriculum(root,verify=True)
14def phase(rows):
15    pieces=np.count_nonzero(rows['x'][:,51:115]!=128,axis=1)
16    return np.where(pieces<=10,2,np.where(rows['ply']<16,0,1)).astype(np.int64)
17class Plan:
18    def __init__(self,data,rank=0,world=2,micro=64):
19        if world!=2 or rank not in (0,1) or micro<1:raise ValueError('two ranks required')
20        self.data=data;self.rank=rank;self.world=world;self.micro=micro;self.steps=math.ceil(len(data.schedule)/(world*micro))
21    def indices(self,step):
22        if not 0<=step<self.steps:raise IndexError('microstep')
23        start=step*self.world*self.micro
24        return np.arange(start+self.rank,min(start+self.world*self.micro,len(self.data.schedule)),self.world,dtype=np.int64)
25    def global_decisions(self,step):return max(0,min(self.world*self.micro,len(self.data.schedule)-step*self.world*self.micro))
26    def cursor(self,step):return min(step*self.world*self.micro,len(self.data.schedule))
27    def load(self,step,pin):
28        began=time.perf_counter();ids=self.indices(step);rows=self.data.take(ids);e=self.data.schedule[ids]
29        tags=rows['source'].astype(np.int64)*6+phase(rows)*2+e['kind'].astype(np.int64)
30        return tensors(rows,pin),torch.from_numpy(np.ascontiguousarray(tags)),time.perf_counter()-began
31class Prefetch:
32    def __init__(self,plan,start=0,pin=False):
33        self.plan=plan;self.pin=pin;self.next=start;self.pool=ThreadPoolExecutor(max_workers=1)
34        self.futures={i:self.pool.submit(plan.load,i,pin) for i in range(start,min(start+2,plan.steps))}
35    def take(self,step):
36        if step!=self.next:raise ValueError('ordered cursor')
37        result=self.futures.pop(step).result();self.next+=1;n=step+2
38        if n<self.plan.steps:self.futures[n]=self.pool.submit(self.plan.load,n,self.pin)
39        return result
40    def close(self):self.pool.shutdown(wait=True,cancel_futures=True)
41
42def start_prefetch(data,plan,start=0,pin=False):
43    # Readers own an unsynchronized mmap cache. Finish main-thread reads before
44    # the sole prefetch thread can touch it; close/join before reader cleanup.
45    dummy=tensors(data.take([0]),pin)
46    return dummy,Prefetch(plan,start,pin)
47
48class ExposureLedger:
49    def __init__(self,data):
50        self.data=data;self.offsets={};self.sizes={};n=0
51        for kind,name in ((0,'ordinary'),(1,'theory')):
52            for source in (0,1):
53                self.offsets[kind,source]=n;self.sizes[kind,source]=len(data.sources[name,source]);n+=self.sizes[kind,source]
54        self.n=n
55    def counts(self,cursor):
56        if not 0<=cursor<=len(self.data.schedule):raise ValueError('schedule cursor')
57        counts=np.zeros(self.n,dtype=np.uint32)
58        for start in range(0,cursor,65536):
59            e=self.data.schedule[start:min(cursor,start+65536)];ids=e['row'].astype(np.int64)
60            if np.any(e['kind']>1) or np.any(e['source']>1):raise ValueError('unknown schedule kind/source')
61            for (kind,source),offset in self.offsets.items():
62                mask=(e['kind']==kind)&(e['source']==source)
63                if np.any(ids[mask]>=self.sizes[kind,source]):raise ValueError('schedule row outside source')
64                ids[mask]+=offset
65            np.add.at(counts,ids,1)
66        for (kind,source),offset in self.offsets.items():
67            if np.any(counts[offset:offset+self.sizes[kind,source]]>(1 if kind==0 else 3)):raise ValueError('overexposure')
68        return counts.astype(np.uint8)
69    def snapshot(self,cursor):
70        counts=self.counts(cursor)
71        return {'cursor':cursor,'counts':torch.from_numpy(counts),'sha256':hashlib.sha256(counts.tobytes()).hexdigest(),'unique_seen':int(np.count_nonzero(counts)),'exposures':int(counts.astype(np.uint64).sum())}
72    def verify(self,value):
73        actual=self.snapshot(value['cursor'])
74        if any(value[k]!=actual[k] for k in ('sha256','unique_seen','exposures')) or not torch.equal(value['counts'],actual['counts']):raise ValueError('exposure checkpoint mismatch')
75        return actual['cursor']
76