prepare.py

1"""Bounded CPU preparation from immutable, already-tokenized game corpora.
2
3No source filters are added. Entire games are replay-validated before positions
4are committed. All historical validation/test positions are protected, including
5terminal positions. Training dedup ignores clocks/history as does protection.
6"""
7import argparse,concurrent.futures,fcntl,hashlib,json,os,sys,time
8from pathlib import Path
9import numpy as np
10import codec as c
11from records import Store,DTYPE,canonical,digest,sha,atomic_json
12sys.path.insert(0,str(Path(__file__).resolve().parents[1]/'elite-continuation-v1'))
13from chessdata.reader import _manifest,_read_shard
14
15def aliases(meta):
16    ids=set(meta.get('game_ids',[]))
17    if meta.get('game_id'):ids.add(meta['game_id'])
18    return ['seq:'+meta['sequence_sha256']]+['id:'+v for v in sorted(ids)]
19
20def replay(task):
21    tokens,meta,mode,seed,source=task
22    if tokens[0]!=c.BOS or tokens[-1]!=c.EOS or (len(tokens)-2)%3:raise ValueError('game boundaries')
23    b=c.chess.Board();moves=[];keys=[];rows=[];terminal_at=None
24    for ply,start in enumerate(range(1,len(tokens)-1,3)):
25        move=c.decode(list(tokens[start:start+3]))
26        if move not in b.legal_moves:raise ValueError('illegal source game')
27        if terminal_at is None and b.is_game_over(claim_draw=False):terminal_at=ply
28        keys.append(c.position_key(b));moves.append(move.uci())
29        if mode!='protect' and terminal_at is None:
30            h=hashlib.sha256(f'{seed}:{meta["sequence_sha256"]}:{ply}'.encode()).digest()
31            history=0 if h[0]<64 else 16 # fixed 25% Bernoulli dropout, not exactly 25%.
32            e=c.example(b,move,history)
33            row=np.zeros((),dtype=DTYPE);row['x']=e['input_ids'];row['y']=e['labels'];row['attention']=e['attention_mask']
34            row['source']=source;row['split']=0 if mode=='train' else 1;row['ply']=ply;row['fullmove']=b.fullmove_number;row['key']=np.frombuffer(keys[-1],dtype=np.uint8)
35            rows.append(row)
36        b.push(move)
37    keys.append(c.position_key(b))
38    if meta['plies']!=len(moves) or hashlib.sha256(' '.join(moves).encode()).hexdigest()!=meta['sequence_sha256']:raise ValueError('move sequence identity')
39    # Existing converter uses ordinary FEN (legal en-passant normalization).
40    if b.fen()!=meta['final_fen']:raise ValueError('final board mismatch')
41    if terminal_at is not None and mode!='protect':
42        # Reject the WHOLE recorded game, not merely its post-terminal suffix.
43        # Protection still includes every supplied legal position conservatively.
44        meta={**meta,'state_v2_rejection':'continued_after_automatic_terminal','terminal_before_ply':terminal_at};rows=[]
45    return meta,keys,np.asarray(rows,dtype=DTYPE)
46
47def tasks(root,manifest,split,mode,seed,source,done,shuffle):
48    shards=[s for s in manifest['shards'] if s['split'] in split]
49    if shuffle:shards.sort(key=lambda s:digest([seed,sorted(s['files'])]))
50    ordinal=0
51    for shard in shards:
52        if ordinal+shard['games']<=done:ordinal+=shard['games'];continue
53        tokens,offsets,metas=_read_shard(Path(root),shard)
54        for a,b,meta in zip(offsets,offsets[1:],metas):
55            ordinal+=1
56            if ordinal<=done:continue
57            yield ordinal,(tokens[a:b],meta,mode,seed,source)
58
59def chunks(iterator,size):
60    chunk=[]
61    for item in iterator:
62        chunk.append(item)
63        if len(chunk)==size:yield chunk;chunk=[]
64    if chunk:yield chunk
65
66def process(store,pool,ref,manifest,source,mode,quota,config):
67    key=f'{mode}:{ref["manifest_sha256"]}:{source}'
68    cursor=store.get(key,{'ordinal':0,'rows':0,'finished':False})
69    if cursor['finished']:return
70    selected=('validation','test') if mode=='protect' else ('validation',) if mode=='validation' else ('train',)
71    iterator=tasks(ref['root'],manifest,selected,mode,config['seed'],source,cursor['ordinal'],mode=='train')
72    for chunk in chunks(iterator,config['chunk_games']):
73        # At most chunk_games tasks/results in flight, never unbounded executor.map.
74        outputs=list(pool.map(replay,[item[1] for item in chunk]))
75        accepted=[]
76        try:
77            store.db.execute('BEGIN')
78            for ordinal,(meta,keys,rows) in zip([item[0] for item in chunk],outputs):
79                store.count(mode+':games')
80                ids=aliases(meta)
81                if mode=='protect':
82                    store.db.executemany('INSERT OR IGNORE INTO protected VALUES(?)',((k,) for k in keys))
83                    store.db.executemany('INSERT OR IGNORE INTO protected_games VALUES(?)',((k,) for k in ids))
84                else:
85                    if meta.get('state_v2_rejection'):
86                        store.count(mode+':rejected:'+meta['state_v2_rejection']);cursor['ordinal']=ordinal;continue
87                    if mode=='train' and any(store.db.execute('SELECT 1 FROM protected_games WHERE k=?',(k,)).fetchone() for k in ids):
88                        store.count('train:heldout_game');cursor['ordinal']=ordinal;continue
89                    identity=meta['sequence_sha256']
90                    store.db.execute('INSERT OR IGNORE INTO games(identity,metadata) VALUES(?,?)',(identity,canonical(meta).decode()))
91                    game=store.db.execute('SELECT id FROM games WHERE identity=?',(identity,)).fetchone()[0]
92                    if game>=2**32:raise ValueError('game index overflow')
93                    for row in rows:
94                        if cursor['rows']>=quota:break
95                        k=row['key'].tobytes()
96                        if mode=='train' and store.db.execute('SELECT 1 FROM protected WHERE k=?',(k,)).fetchone():store.count('train:protected_position');continue
97                        # Validation dedup is separate from training dedup, but is
98                        # globally unique across both validation sources.
99                        sk=(b'v' if mode=='validation' else b't')+k
100                        if not store.db.execute('INSERT OR IGNORE INTO seen VALUES(?)',(sk,)).rowcount:store.count(mode+':duplicate_position');continue
101                        row['game']=game;accepted.append(row);cursor['rows']+=1
102                        store.count(f'{mode}:source{source}:rows');store.count(mode+':fen_only',int(row['x'][1]==c.HIST_NONE))
103                cursor['ordinal']=ordinal
104                if mode!='protect' and cursor['rows']>=quota:break
105            cursor['finished']=mode!='protect' and cursor['rows']>=quota
106            store.commit(accepted,source,0 if mode=='train' else 1,key,cursor)
107        except BaseException:store.db.rollback();raise
108        store.capacity()
109        print(json.dumps({'time':time.time(),'phase':mode,'source':source,**cursor,'allocated_bytes':store.capacity()}),flush=True)
110        if cursor['finished']:return
111    if mode!='protect' and cursor['rows']<quota:raise RuntimeError('source exhausted below configured position quota')
112    with store.db:cursor['finished']=True;store.set(key,cursor)
113
114def run(config):
115    root=Path(config['output']);root.mkdir(parents=True,exist_ok=True)
116    with (root/'owner.lock').open('a+b') as lock:
117        fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB)
118        code={p.name:sha(p) for p in [Path(__file__),Path(c.__file__),Path(__file__).with_name('records.py')]}
119        identity={'config':config,'code':code,'codec':c.CONTRACT}
120        store=Store(root,identity,config['storage_cap_bytes'])
121        refs={r['manifest_sha256']:r for r in config['references']+config['sources']}
122        manifests={}
123        for fp,ref in refs.items():
124            if sha(Path(ref['root'])/'manifest.json')!=fp:raise ValueError('source manifest mutated')
125            manifests[fp]=_manifest(Path(ref['root']))
126        with concurrent.futures.ProcessPoolExecutor(max_workers=config['workers']) as pool:
127            for ref in refs.values():process(store,pool,ref,manifests[ref['manifest_sha256']],0,'protect',0,config)
128            with store.db:store.set('protection_complete',True)
129            for source,ref in enumerate(config['sources']):process(store,pool,ref,manifests[ref['manifest_sha256']],source,'validation',config['validation_positions_per_source'],config)
130            for source,ref in enumerate(config['sources']):process(store,pool,ref,manifests[ref['manifest_sha256']],source,'train',config['train_positions_per_source'],config)
131        manifest=store.seal()
132        atomic_json(root/'completion.json',{'manifest_sha256':sha(root/'manifest.json'),'rows':manifest['rows'],'finished_unix':time.time(),'allocated_bytes':store.capacity(),
133            'protected_positions':store.db.execute('SELECT COUNT(*) FROM protected').fetchone()[0],
134            'protected_game_aliases':store.db.execute('SELECT COUNT(*) FROM protected_games').fetchone()[0]})
135        store.db.close()
136
137if __name__=='__main__':
138    p=argparse.ArgumentParser();p.add_argument('config');args=p.parse_args();run(json.loads(Path(args.config).read_text()))
139