records.py

1"""Disk-ready fixed records and a single-owner recoverable publication journal."""
2import hashlib,json,os,sqlite3
3from pathlib import Path
4import numpy as np
5from codec import INPUT,TARGET_POSITIONS,CONTRACT
6
7DTYPE=np.dtype([('x','<u2',(INPUT,)),('y','<u2',(3,)),('attention','u1',(INPUT,)),
8    ('source','u1'),('split','u1'),('game','<u4'),('ply','<u2'),('fullmove','<u4'),('key','u1',(32,))])
9assert DTYPE.itemsize==416
10def canonical(x):return json.dumps(x,sort_keys=True,separators=(',',':')).encode()
11def digest(x):return hashlib.sha256(canonical(x)).hexdigest()
12def sha(path):
13    h=hashlib.sha256()
14    with Path(path).open('rb') as f:
15        for b in iter(lambda:f.read(8<<20),b''):h.update(b)
16    return h.hexdigest()
17def syncdir(path):
18    fd=os.open(path,os.O_RDONLY|os.O_DIRECTORY)
19    try:os.fsync(fd)
20    finally:os.close(fd)
21def atomic_json(path,obj):
22    path=Path(path);tmp=path.with_suffix('.tmp')
23    with tmp.open('wb') as f:f.write(canonical(obj));f.flush();os.fsync(f.fileno())
24    os.replace(tmp,path);syncdir(path.parent)
25def allocation(root):return sum(p.stat().st_size for p in Path(root).rglob('*') if p.is_file())
26
27class Store:
28    """Caller holds exclusive flock. File fsync precedes the SQLite commit.
29
30    Crash before commit leaves an orphan, removed on reopen; after commit the file
31    is retained and hashed before use. Seen keys, counts, files and source cursor
32    share ONE SQLite transaction. No checkpoint advances independently of dedup.
33    """
34    def __init__(self,root,identity,cap):
35        self.root=Path(root);self.root.mkdir(parents=True,exist_ok=True);self.cap=cap
36        prior=(self.root/'journal.sqlite').exists()
37        if not prior and any(self.root.glob('records-*.bin')):raise ValueError('missing journal with prior records')
38        self.db=sqlite3.connect(self.root/'journal.sqlite')
39        self.db.executescript('''PRAGMA journal_mode=DELETE; PRAGMA synchronous=FULL;
40        CREATE TABLE IF NOT EXISTS settings(k TEXT PRIMARY KEY,v TEXT NOT NULL);
41        CREATE TABLE IF NOT EXISTS protected(k BLOB PRIMARY KEY) WITHOUT ROWID;
42        CREATE TABLE IF NOT EXISTS protected_games(k TEXT PRIMARY KEY) WITHOUT ROWID;
43        CREATE TABLE IF NOT EXISTS seen(k BLOB PRIMARY KEY) WITHOUT ROWID;
44        CREATE TABLE IF NOT EXISTS games(id INTEGER PRIMARY KEY,identity TEXT UNIQUE,metadata TEXT NOT NULL);
45        CREATE TABLE IF NOT EXISTS files(name TEXT PRIMARY KEY,sha TEXT NOT NULL,rows INTEGER,source INTEGER,split INTEGER);
46        CREATE TABLE IF NOT EXISTS counters(k TEXT PRIMARY KEY,n INTEGER NOT NULL);
47        CREATE TABLE IF NOT EXISTS progress(k TEXT PRIMARY KEY,v TEXT NOT NULL);''')
48        fp=digest(identity);row=self.db.execute("SELECT v FROM settings WHERE k='identity'").fetchone()
49        if row and row[0]!=fp:raise ValueError('source/config/code changed; resume refused')
50        with self.db:self.db.execute("INSERT OR IGNORE INTO settings VALUES('identity',?)",(fp,))
51        self.identity=identity
52        names=set()
53        for name,expected,rows,_,_ in self.db.execute('SELECT * FROM files'):
54            p=self.root/name
55            if p.stat().st_size!=rows*DTYPE.itemsize or sha(p)!=expected:raise ValueError('committed records corrupt')
56            names.add(name)
57        for p in self.root.glob('records-*.bin'):
58            if p.name not in names:p.unlink()
59        for p in self.root.glob('records-*.tmp'):p.unlink()
60        self.capacity()
61    def capacity(self,extra=0):
62        # Reserve enough for one rollback journal, one output chunk, and metadata.
63        import shutil
64        n=allocation(self.root)
65        if n+extra+(128<<20)>self.cap:raise RuntimeError('state-v2 namespace storage cap')
66        if shutil.disk_usage(self.root).free<50_000_000_000+extra:raise RuntimeError('native free reserve')
67        return n
68    def get(self,k,default=None):
69        row=self.db.execute('SELECT v FROM progress WHERE k=?',(k,)).fetchone()
70        return json.loads(row[0]) if row else default
71    def set(self,k,v):self.db.execute('INSERT OR REPLACE INTO progress VALUES(?,?)',(k,canonical(v).decode()))
72    def count(self,k,n=1):self.db.execute('INSERT INTO counters VALUES(?,?) ON CONFLICT(k) DO UPDATE SET n=n+excluded.n',(k,n))
73    def commit(self,rows,source,split,cursor_key,cursor,fail=None):
74        rows=np.asarray(rows,dtype=DTYPE)
75        if len(rows):
76            name=f'records-{self.db.execute("SELECT COUNT(*) FROM files").fetchone()[0]:06d}.bin'
77            self.capacity(rows.nbytes)
78            tmp=self.root/name.replace('.bin','.tmp');final=self.root/name
79            with tmp.open('wb') as f:f.write(rows.tobytes());f.flush();os.fsync(f.fileno())
80            os.replace(tmp,final);syncdir(self.root)
81            if fail=='after_file':raise RuntimeError('injected after file')
82            self.db.execute('INSERT INTO files VALUES(?,?,?,?,?)',(name,sha(final),len(rows),source,split))
83        self.set(cursor_key,cursor)
84        if fail=='before_commit':raise RuntimeError('injected before commit')
85        self.db.commit()
86        if fail=='after_commit':raise RuntimeError('injected after commit')
87    def seal(self):
88        files=[dict(zip(('name','sha256','rows','source','split'),r)) for r in self.db.execute('SELECT * FROM files ORDER BY name')]
89        shared={'position_ids':list(range(INPUT)),'loss_mask':[int(i in TARGET_POSITIONS) for i in range(INPUT)],'target_positions':list(TARGET_POSITIONS)}
90        atomic_json(self.root/'shared.json',shared)
91        manifest={'format':'chess-state-records-v2','identity':self.identity,'identity_sha256':digest(self.identity),
92            'contract':CONTRACT,'record_bytes':DTYPE.itemsize,'dtype':DTYPE.descr,'files':files,
93            'shared_sha256':sha(self.root/'shared.json'),'counts':dict(self.db.execute('SELECT k,n FROM counters')),
94            'rows':sum(f['rows'] for f in files),'allocation_bytes_before_manifest':self.capacity()}
95        atomic_json(self.root/'manifest.json',manifest);return manifest
96
97class Records:
98    """Vectorized row gathering from bounded mmap shards; no chess dependency use."""
99    def __init__(self,root,source,split,verify=True):
100        self.root=Path(root);m=json.loads((self.root/'manifest.json').read_text())
101        if m['format']!='chess-state-records-v2' or digest(m['identity'])!=m['identity_sha256'] or m['record_bytes']!=DTYPE.itemsize:raise ValueError('record manifest')
102        if sha(self.root/'shared.json')!=m['shared_sha256']:raise ValueError('shared masks changed')
103        shared=json.loads((self.root/'shared.json').read_text())
104        if shared!={'position_ids':list(range(INPUT)),'loss_mask':[int(i in TARGET_POSITIONS) for i in range(INPUT)],'target_positions':list(TARGET_POSITIONS)}:raise ValueError('shift contract')
105        self.parts=[];ends=[];n=0
106        for f in m['files']:
107            if f['source']!=source or f['split']!=split:continue
108            p=self.root/f['name']
109            if p.parent!=self.root or p.stat().st_size!=f['rows']*DTYPE.itemsize or (verify and sha(p)!=f['sha256']):raise ValueError('record integrity')
110            a=np.memmap(p,dtype=DTYPE,mode='r');self.parts.append(a);n+=len(a);ends.append(n)
111        self.ends=np.asarray(ends,dtype=np.int64);self.n=n
112    def __len__(self):return self.n
113    def take(self,indices):
114        ids=np.asarray(indices,dtype=np.int64)
115        if ids.ndim!=1 or np.any(ids<0) or np.any(ids>=self.n):raise IndexError('record index')
116        which=np.searchsorted(self.ends,ids,side='right');out=np.empty(len(ids),dtype=DTYPE)
117        for part in np.unique(which):
118            mask=which==part;start=0 if part==0 else self.ends[part-1]
119            out[mask]=self.parts[part][ids[mask]-start]
120        return out
121