runtime_records.py

1"""Bounded open mappings over immutable contiguous record shards."""
2from collections import OrderedDict
3import json,re
4from pathlib import Path
5import numpy as np
6from records import DTYPE,sha,digest
7from codec import INPUT,TARGET_POSITIONS
8
9class Records:
10    def __init__(self,root,source,split,verify=True,max_open=32):
11        if type(max_open) is not int or not 1<=max_open<=128:raise ValueError('mapping bound')
12        self.root=Path(root).resolve();self.max_open=max_open;self.cache=OrderedDict()
13        m=json.loads((self.root/'manifest.json').read_text())
14        if m['format']!='chess-state-records-v2' or digest(m['identity'])!=m['identity_sha256'] or m['record_bytes']!=DTYPE.itemsize:raise ValueError('record manifest')
15        if sha(self.root/'shared.json')!=m['shared_sha256']:raise ValueError('shared masks changed')
16        if json.loads((self.root/'shared.json').read_text())!={'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')
17        self.parts=[];ends=[];n=0;seen=set()
18        for f in m['files']:
19            if not re.fullmatch(r'records-[0-9]{6}\.bin',f['name']) or f['name'] in seen or type(f['rows']) is not int or f['rows']<=0:raise ValueError('shard descriptor')
20            seen.add(f['name'])
21            if f['source']!=source or f['split']!=split:continue
22            p=self.root/f['name']
23            if p.stat().st_size!=f['rows']*DTYPE.itemsize or (verify and sha(p)!=f['sha256']):raise ValueError('record integrity')
24            self.parts.append(p);n+=f['rows'];ends.append(n)
25        self.ends=np.asarray(ends,dtype=np.int64);self.n=n
26    def __len__(self):return self.n
27    def take(self,indices):
28        ids=np.asarray(indices,dtype=np.int64)
29        if ids.ndim!=1 or np.any(ids<0) or np.any(ids>=self.n):raise IndexError('record index')
30        which=np.searchsorted(self.ends,ids,side='right');out=np.empty(len(ids),dtype=DTYPE)
31        for part in np.unique(which):
32            if part not in self.cache:
33                if len(self.cache)>=self.max_open:self.cache.popitem(last=False)[1]._mmap.close()
34                self.cache[part]=np.memmap(self.parts[part],dtype=DTYPE,mode='r')
35            self.cache.move_to_end(part);mask=which==part;start=0 if part==0 else self.ends[part-1]
36            out[mask]=self.cache[part][ids[mask]-start]
37        return out
38    def close(self):
39        for value in self.cache.values():value._mmap.close()
40        self.cache.clear()
41    def __del__(self):
42        if hasattr(self,'cache'):self.close()
43