test_records.py

1import hashlib,json,tempfile,unittest
2from pathlib import Path
3import numpy as np
4from unittest.mock import patch
5import codec as c
6from records import Store,Records,DTYPE
7from prepare import replay,process
8from batches import Balanced,tensors,Prefetch
9from runtime_records import Records as BoundedRecords
10
11def fixture(moves='e2e4 e7e5 g1f3'):
12    b=c.chess.Board();tokens=[69]
13    for u in moves.split():m=c.chess.Move.from_uci(u);tokens.extend(c.triple(m));b.push(m)
14    tokens.append(70)
15    return bytes(tokens),{'plies':len(moves.split()),'final_fen':b.fen(),'sequence_sha256':hashlib.sha256(moves.encode()).hexdigest(),'game_id':'lichess:fixture'}
16
17class Tests(unittest.TestCase):
18    def test_protected_transposition_and_resume(self):
19        class Pool:
20            def map(self,fn,items):return map(fn,items)
21        config={'seed':'fixture','chunk_games':2}
22        protected=fixture('g1f3 g8f6 d2d4 d7d5')
23        # Same final state via a different move order must be excluded, even
24        # though the full-game sequence hash differs. Later position is novel.
25        train=fixture('d2d4 d7d5 g1f3 g8f6 c2c4 e7e6')
26        train[1]['game_id']='lichess:other'
27        with tempfile.TemporaryDirectory() as tmp:
28            s=Store(tmp,{'fixture':True},1_000_000_000)
29            def feed(root,manifest,split,mode,seed,source,done,shuffle):
30                seq,meta=protected if mode=='protect' else train
31                if done<1:yield 1,(seq,meta,mode,seed,source)
32            ref={'root':'unused','manifest_sha256':'fixture'}
33            with patch('prepare.tasks',feed):
34                process(s,Pool(),ref,{},0,'protect',0,config)
35                process(s,Pool(),ref,{},0,'train',4,config)
36                before=s.db.execute('SELECT COUNT(*) FROM seen').fetchone()[0]
37                process(s,Pool(),ref,{},0,'train',4,config)
38            self.assertEqual(before,4);self.assertEqual(s.db.execute('SELECT COUNT(*) FROM seen').fetchone()[0],4)
39            common=replay((*protected,'protect','fixture',0))[1][-1]
40            self.assertIsNone(s.db.execute('SELECT 1 FROM seen WHERE k=?',(b't'+common,)).fetchone())
41            self.assertGreaterEqual(dict(s.db.execute('SELECT k,n FROM counters'))['train:protected_position'],2)
42            s.db.close()
43    def test_whole_game_replay(self):
44        tok,meta=fixture();m,keys,rows=replay((tok,meta,'train','seed',0))
45        self.assertEqual(len(keys),4);self.assertEqual(len(rows),3);self.assertEqual(rows.dtype.itemsize,416)
46        self.assertEqual(rows[0]['y'].tolist(),[12,28,64]);self.assertEqual(rows[0]['x'][-2:].tolist(),[12,28])
47        with self.assertRaises(ValueError):replay((tok[:-4]+bytes([12,12,64,70]),meta,'train','seed',0))
48        with self.assertRaises(ValueError):replay((tok,{**meta,'final_fen':c.chess.Board().fen()},'train','seed',0))
49    def test_history_dropout_deterministic(self):
50        tok,m=fixture();a=replay((tok,m,'train','x',0))[2];b=replay((tok,m,'train','x',0))[2]
51        self.assertEqual(a.tobytes(),b.tobytes())
52    def test_crash_boundaries(self):
53        for fail in ('after_file','before_commit','after_commit'):
54            with self.subTest(fail=fail),tempfile.TemporaryDirectory() as tmp:
55                s=Store(tmp,{'test':1},1_000_000_000)
56                s.db.execute('INSERT INTO seen VALUES(?)',(b'key',))
57                with self.assertRaises(RuntimeError):s.commit(np.zeros(2,dtype=DTYPE),0,0,'cursor',2,fail)
58                s.db.close();s=Store(tmp,{'test':1},1_000_000_000)
59                committed=fail=='after_commit'
60                self.assertEqual(s.get('cursor'),2 if committed else None)
61                self.assertEqual(s.db.execute('SELECT COUNT(*) FROM seen').fetchone()[0],int(committed))
62                self.assertEqual(len(list(Path(tmp).glob('records-*.bin'))),int(committed));s.db.close()
63    def test_manifest_mmap_corruption(self):
64        with tempfile.TemporaryDirectory() as tmp:
65            s=Store(tmp,{'test':2},1_000_000_000);rows=np.zeros(5,dtype=DTYPE);rows['ply']=np.arange(5)
66            s.commit(rows[:2],0,0,'a',2);s.commit(rows[2:],0,0,'a',5);s.seal();s.db.close()
67            r=Records(tmp,0,0);self.assertEqual(r.take([4,1,3])['ply'].tolist(),[4,1,3])
68            bounded=BoundedRecords(tmp,0,0,max_open=1)
69            self.assertEqual(bounded.take([4,1,3,0])['ply'].tolist(),[4,1,3,0]);self.assertEqual(len(bounded.cache),1);bounded.close()
70            with self.assertRaises(IndexError):r.take([5])
71            p=next(Path(tmp).glob('records-*.bin'));p.write_bytes(b'x'+p.read_bytes()[1:])
72            with self.assertRaises(ValueError):Records(tmp,0,0)
73    def test_identity_and_missing_journal(self):
74        with tempfile.TemporaryDirectory() as tmp:
75            s=Store(tmp,{'x':1},1_000_000_000);s.commit(np.zeros(1,dtype=DTYPE),0,0,'x',1);s.db.close()
76            with self.assertRaises(ValueError):Store(tmp,{'x':2},1_000_000_000)
77            (Path(tmp)/'journal.sqlite').unlink()
78            with self.assertRaises(ValueError):Store(tmp,{'x':1},1_000_000_000)
79    def test_balanced_disjoint_final(self):
80        class Source:
81            def __init__(self,s):self.s=s
82            def __len__(self):return 19
83            def take(self,ids):
84                a=np.zeros(len(ids),dtype=DTYPE);a['source']=self.s;a['game']=ids;return a
85        plans=[Balanced(Source(0),Source(1),7,r,2,8) for r in range(2)]
86        allrows=[]
87        for p in plans:
88            pre=Prefetch(p)
89            for i in range(p.steps):
90                a=p.rows(i);self.assertEqual(sum(a['source']==0),sum(a['source']==1));allrows.extend(zip(a['source'].tolist(),a['game'].tolist()))
91                ts,_=pre.take(i);self.assertEqual(ts[0].shape,(len(a),122))
92            pre.close()
93        self.assertEqual(len(allrows),38);self.assertEqual(len(set(allrows)),38)
94
95if __name__=='__main__':unittest.main()
96