convert.py

1"""Disk-backed ingestion journal and atomic, replayable dataset publication."""
2import argparse
3import hashlib
4import json
5import os
6from pathlib import Path
7import re
8import sqlite3
9import struct
10import sys
11from itertools import islice
12from concurrent.futures import ProcessPoolExecutor
13import chess
14from .pgn import open_pgn, records, parse, Rejected, filter_reason
15from .tokenizer import VERSION
16from .limits import MAX_SHARDS, MAX_MANIFEST_BYTES, MAX_SOURCES, MAX_CONFIG_BYTES
17
18
19def canonical(value):
20    return json.dumps(value, sort_keys=True, separators=(',', ':'), ensure_ascii=False).encode('utf-8')
21
22
23def digest(value):
24    return hashlib.sha256(value).hexdigest()
25
26
27def file_hash(path):
28    h = hashlib.sha256()
29    with open(path, 'rb') as stream:
30        while chunk := stream.read(1024 * 1024):
31            h.update(chunk)
32    return h.hexdigest()
33
34
35def fault(point):
36    # Explicit test-only failpoint. Abrupt process death exercises actual recovery.
37    if os.environ.get('CHESSDATA_TEST_CRASH') == point:
38        os._exit(86)
39
40
41def sync_dir(path):
42    fd = os.open(path, os.O_RDONLY | os.O_DIRECTORY)
43    try:
44        os.fsync(fd)
45    finally:
46        os.close(fd)
47
48
49def atomic(path, data):
50    temporary = path.with_name(path.name + '.tmp')
51    with open(temporary, 'wb') as stream:
52        stream.write(data)
53        stream.flush()
54        os.fsync(stream.fileno())
55    os.replace(temporary, path)
56    sync_dir(path.parent)
57
58
59def game_id(headers):
60    # Only explicit IDs or canonical Lichess game URLs, never generic Site strings.
61    explicit = headers.get('GameId')
62    if explicit:
63        return 'explicit:' + explicit
64    match = re.fullmatch(r'https?://lichess.org/([A-Za-z0-9]{8})(?:/(?:white|black))?/?', headers.get('Site', ''))
65    return 'lichess:' + match[1] if match else None
66
67
68def split_for(sequence_hash, config):
69    value = int(digest((config['seed'] + ':' + sequence_hash).encode())[:16], 16) % 10000
70    return 'test' if value < config['test_bps'] else ('validation' if value < config['test_bps'] + config['validation_bps'] else 'train')
71
72
73def _parse_result(record):
74    try:
75        return None, parse(record)
76    except Rejected as error:
77        return str(error), None
78
79
80def _worker_parent_guard(parent_pid):
81    # Linux-only writer already requires fcntl/fsync. Prevent orphan parser workers
82    # after an abrupt parent crash; this affects only this owned child process.
83    import ctypes
84    import signal
85    if ctypes.CDLL(None).prctl(1,signal.SIGTERM,0,0,0)!=0:
86        raise OSError('could not set parser-worker parent-death cleanup')
87    if os.getppid()!=parent_pid:os._exit(1)
88
89
90def _parsed_records(stream, config, done):
91    pending=((i,r) for i,r in enumerate(records(stream,config['max_record_bytes'])) if i>done)
92    workers=config.get('workers',1)
93    if workers==1:
94        for i,record in pending:yield i,_parse_result(record)
95    else:
96        # Python 3.12 Executor.map is eager: bound each submission explicitly.
97        with ProcessPoolExecutor(max_workers=workers,initializer=_worker_parent_guard,initargs=(os.getpid(),)) as executor:
98            while batch:=list(islice(pending,64)):
99                for (i,_),parsed in zip(batch,executor.map(_parse_result,(r for i,r in batch))):
100                    yield i,parsed
101
102
103class _Transactions:
104    """Dedup, tokens, counts and checkpoint commit together in bounded groups."""
105    def __init__(self,db,size):self.db,self.size,self.pending,self.committed=db,size,0,False
106    def __enter__(self):
107        self.committed=False
108        if not self.db.in_transaction:self.db.execute('BEGIN')
109        return self
110    def __exit__(self,kind,value,tb):
111        if kind:
112            self.db.rollback();self.pending=0
113        else:
114            self.pending+=1
115            if self.pending>=self.size:self.flush()
116    def flush(self):
117        if self.pending:
118            self.db.commit();self.pending=0;self.committed=True
119
120
121def convert(inputs, output, config):
122    import fcntl  # Linux publication durability/locking contract.
123    config = dict(config)
124    config.setdefault('max_shards', MAX_SHARDS)
125    if type(config.get('workers',1)) is not int or not 1<=config.get('workers',1)<=8:
126        raise ValueError('workers must be 1..8')
127    if type(config.get('commit_records',1)) is not int or not 1<=config.get('commit_records',1)<=4096:
128        raise ValueError('commit_records must be 1..4096')
129    if config.get('max_game_tokens') is not None and (type(config['max_game_tokens']) is not int or config['max_game_tokens']<2):
130        raise ValueError('max_game_tokens must be an integer >=2')
131    if type(config['max_shards']) is not int or not 1 <= config['max_shards'] <= MAX_SHARDS:
132        raise ValueError(f'max_shards must be between 1 and {MAX_SHARDS}')
133    if len(canonical(config)) > MAX_CONFIG_BYTES:
134        raise ValueError('configuration exceeds v1 size limit')
135    inputs = list(islice(iter(inputs), MAX_SOURCES + 1))
136    if not 1 <= len(inputs) <= MAX_SOURCES:
137        raise ValueError(f'v1 requires 1..{MAX_SOURCES} source files')
138    output = Path(output)
139    output.mkdir(parents=True, exist_ok=True)
140    with open(output / '.writer.lock', 'a') as lock:
141        try:
142            fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
143        except BlockingIOError as error:
144            raise ValueError('another writer holds the output lock') from error
145        return _convert(inputs, output, config)
146
147
148def _convert(inputs, output, config):
149    inputs = [Path(p).resolve() for p in inputs]
150    source = [{'path': str(p), 'bytes': p.stat().st_size, 'sha256': file_hash(p)} for p in inputs]
151    code = {p.name: file_hash(p) for p in sorted(Path(__file__).parent.glob('*.py'))}
152    provenance = {'contract': VERSION, 'config': config, 'sources': source,
153                  'code_sha256': code, 'chess_version': chess.__version__, 'python': sys.version.split()[0]}
154    fingerprint = digest(canonical(provenance))
155    database = output / 'journal.sqlite3'
156    if not database.exists() and any(p.name != '.writer.lock' for p in output.iterdir()):
157        raise ValueError('output is not an empty or resumable dataset directory')
158    db = sqlite3.connect(database)
159    try:
160        db.execute('PRAGMA journal_mode=WAL')
161        db.execute('PRAGMA synchronous=FULL')
162        db.execute('PRAGMA cache_size=-4096')
163        db.execute('PRAGMA temp_store=FILE')
164        db.executescript('''
165            CREATE TABLE IF NOT EXISTS meta(key TEXT PRIMARY KEY, value TEXT NOT NULL);
166            CREATE TABLE IF NOT EXISTS progress(source INTEGER PRIMARY KEY, ordinal INTEGER NOT NULL);
167            CREATE TABLE IF NOT EXISTS counts(reason TEXT PRIMARY KEY, n INTEGER NOT NULL);
168            CREATE TABLE IF NOT EXISTS games(seq TEXT PRIMARY KEY, split TEXT NOT NULL,
169                tokens BLOB, ntokens INTEGER NOT NULL, metadata TEXT NOT NULL, source INTEGER, ordinal INTEGER);
170            CREATE TABLE IF NOT EXISTS ids(id TEXT PRIMARY KEY, seq TEXT NOT NULL);
171            CREATE TABLE IF NOT EXISTS excluded(seq TEXT PRIMARY KEY, ntokens INTEGER NOT NULL, metadata TEXT NOT NULL);
172            CREATE INDEX IF NOT EXISTS game_order ON games(split,source,ordinal);
173        ''')
174        row = db.execute("SELECT value FROM meta WHERE key='fingerprint'").fetchone()
175        if db.execute("SELECT 1 FROM meta WHERE key='invalid_source'").fetchone():
176            raise ValueError('source mutated during an earlier run; use a fresh output directory')
177        if row and row[0] != fingerprint:
178            raise ValueError('resume refused: source/config/implementation/runtime fingerprint changed')
179        with db:
180            db.execute("INSERT OR IGNORE INTO meta VALUES('fingerprint', ?)", (fingerprint,))
181        for source_index, path in enumerate(inputs):
182            row = db.execute('SELECT ordinal FROM progress WHERE source=?', (source_index,)).fetchone()
183            done = row[0] if row else -1
184            transactions=_Transactions(db,config.get('commit_records',1))
185            with open_pgn(path) as stream:
186                for ordinal, (parse_error, parsed) in _parsed_records(stream,config,done):
187                    with transactions:
188                        reason = 'accepted'
189                        if parse_error is not None:
190                            reason = 'rejected:' + parse_error
191                        else:
192                            headers, tokens, uci, final_fen = parsed
193                            filtered = filter_reason(headers, config['filter'])
194                            if filtered:
195                                reason = 'filtered:' + filtered
196                            else:
197                                seq = digest(uci.encode('ascii'))
198                                gid = game_id(headers)
199                                previous = db.execute('SELECT seq FROM ids WHERE id=?', (gid,)).fetchone() if gid else None
200                                exists = db.execute('SELECT 1 FROM games WHERE seq=? UNION ALL SELECT 1 FROM excluded WHERE seq=?', (seq,seq)).fetchone()
201                                if previous and previous[0] != seq:
202                                    reason = 'rejected:game_id_conflict'
203                                elif exists or previous:
204                                    reason = 'duplicate'
205                                else:
206                                    metadata = {'headers': headers, 'sequence_sha256': seq, 'game_id': gid,
207                                                'plies': (len(tokens)-2)//3, 'final_fen': final_fen,
208                                                'source_index': source_index, 'record_ordinal': ordinal}
209                                    if config.get('max_game_tokens') is not None and len(tokens)>config['max_game_tokens']:
210                                        db.execute('INSERT INTO excluded VALUES(?,?,?)',(seq,len(tokens),canonical(metadata).decode()))
211                                        reason='filtered:overlength_unique'
212                                    else:
213                                        db.execute('INSERT INTO games VALUES(?,?,?,?,?,?,?)',
214                                                   (seq, split_for(seq, config), None if config['count_only'] else bytes(tokens),
215                                                    len(tokens), canonical(metadata).decode(), source_index, ordinal))
216                                if gid and reason in ('accepted', 'duplicate','filtered:overlength_unique'):
217                                    db.execute('INSERT OR IGNORE INTO ids VALUES(?,?)', (gid, seq))
218                        db.execute('INSERT INTO counts VALUES(?,1) ON CONFLICT(reason) DO UPDATE SET n=n+1', (reason,))
219                        db.execute('INSERT INTO progress VALUES(?,?) ON CONFLICT(source) DO UPDATE SET ordinal=excluded.ordinal', (source_index, ordinal))
220                        fault('before_record_commit')
221                    if transactions.committed:fault('after_record_commit')
222                    if transactions.committed and ordinal%5000<config.get('commit_records',1):
223                        atomic(output/'ingest-progress.json',canonical({'source':source_index,'committed_ordinal':ordinal,
224                               'counts':dict(db.execute('SELECT reason,n FROM counts')),'published':False}))
225                transactions.flush()
226                if transactions.committed:fault('after_record_commit')
227            # A concurrent source edit is fatal, even when all old records replayed.
228            if file_hash(path) != source[source_index]['sha256']:
229                with db:
230                    db.execute("INSERT OR REPLACE INTO meta VALUES('invalid_source','true')")
231                raise ValueError('source changed during conversion; use a fresh output directory')
232        fault('after_ingest')
233        shards = []
234        if not config['count_only']:
235            for split in ('train', 'validation', 'test'):
236                number, data, offsets, metadata, metadata_bytes = 0, bytearray(), [0], [], 0
237                def flush():
238                    nonlocal number, data, offsets, metadata, metadata_bytes
239                    if not metadata:
240                        return
241                    if len(shards) >= config['max_shards']:
242                        raise ValueError('shard limit exceeded; no new manifest published. Use a fresh output with larger shard_tokens or a higher max_shards within the v1 hard limit')
243                    stem = f'{split}-{number:06d}'
244                    files = {}
245                    for suffix, content in [('tokens.u8', bytes(data)), ('offsets.u64le', struct.pack('<' + 'Q'*len(offsets), *offsets)),
246                                            ('metadata.jsonl', b''.join(canonical(m)+b'\n' for m in metadata))]:
247                        name = stem + '.' + suffix
248                        atomic(output / name, content)
249                        files[name] = {'sha256': digest(content), 'bytes': len(content)}
250                        fault('after_shard_file')
251                    shards.append({'split': split, 'games': len(metadata), 'tokens': len(data), 'files': files})
252                    number += 1
253                    data, offsets, metadata, metadata_bytes = bytearray(), [0], [], 0
254                cursor = db.execute('SELECT tokens,metadata FROM games WHERE split=? ORDER BY source,ordinal', (split,))
255                for blob, meta in cursor:
256                    if data and (len(data) + len(blob) > config['shard_tokens'] or metadata_bytes + len(meta) > config['shard_tokens']):
257                        flush()
258                    data.extend(blob)
259                    offsets.append(len(data))
260                    metadata.append(json.loads(meta))
261                    metadata_bytes += len(meta)
262                flush()
263        stats = dict(db.execute('SELECT reason,n FROM counts ORDER BY reason'))
264        split_counts = {s: {'games': n, 'tokens': t} for s,n,t in db.execute('SELECT split,count(*),sum(ntokens) FROM games GROUP BY split')}
265        manifest = {'format': 'chess-dataset-v1', 'fingerprint': fingerprint, 'provenance': provenance,
266                    'counts': stats, 'splits': split_counts, 'shards': shards,
267                    'accepted_tokens': sum(s['tokens'] for s in split_counts.values()),
268                    'overlength_unique':dict(zip(('games','tokens'),db.execute('SELECT count(*),coalesce(sum(ntokens),0) FROM excluded').fetchone())),
269                    'attention_contract': 'Each offset interval is one game. Never attend or train EOS->BOS across intervals. Use single-game causal windows; pad loss=False. EOS only at actual recorded endings.'}
270        fault('before_manifest')
271        serialized = canonical(manifest) + b'\n'
272        if len(serialized) > MAX_MANIFEST_BYTES:
273            raise ValueError('manifest exceeds v1 byte limit; no new manifest published')
274        atomic(output / 'manifest.json', serialized)
275        fault('after_manifest')
276        return manifest
277    finally:
278        db.close()
279
280
281def main(argv=None):
282    parser = argparse.ArgumentParser(description=__doc__)
283    parser.add_argument('inputs', nargs='+')
284    parser.add_argument('--output', required=True)
285    parser.add_argument('--count-only', action='store_true')
286    parser.add_argument('--strong', action='store_true', help='Provisional: both Elo fields >=2000, base+40*inc >=480, completed, no BOT titles')
287    parser.add_argument('--min-rating', type=int)
288    parser.add_argument('--min-time', type=int)
289    parser.add_argument('--completed', action='store_true')
290    parser.add_argument('--exclude-bots', action='store_true')
291    parser.add_argument('--shard-tokens', type=int, default=4*1024*1024)
292    parser.add_argument('--max-shards', type=int, default=MAX_SHARDS, help=f'manifest index bound, 1..{MAX_SHARDS}')
293    parser.add_argument('--max-record-bytes', type=int, default=1024*1024)
294    parser.add_argument('--workers',type=int,default=1)
295    parser.add_argument('--commit-records',type=int,default=1)
296    parser.add_argument('--max-game-tokens',type=int)
297    parser.add_argument('--seed', default='chess-v1')
298    parser.add_argument('--validation-bps', type=int, default=100)
299    parser.add_argument('--test-bps', type=int, default=100)
300    args = parser.parse_args(argv)
301    if args.shard_tokens < 1 or args.max_record_bytes < 64 or min(args.validation_bps,args.test_bps) < 0 or args.validation_bps + args.test_bps > 10000:
302        parser.error('invalid size or split configuration')
303    config = {k: getattr(args, k) for k in ('count_only','shard_tokens','max_shards','max_record_bytes','seed','validation_bps','test_bps','workers','commit_records','max_game_tokens')}
304    config['filter'] = {'min_rating': args.min_rating if args.min_rating is not None else (2000 if args.strong else None),
305                        'min_time': args.min_time if args.min_time is not None else (480 if args.strong else None),
306                        'completed': args.completed or args.strong, 'exclude_bots': args.exclude_bots or args.strong}
307    try:
308        result = convert(args.inputs, args.output, config)
309    except (ValueError, OSError, UnicodeError, sqlite3.Error) as error:
310        parser.exit(2, f'conversion failed: {error}\n')
311    print(json.dumps({k: result[k] for k in ('fingerprint','counts','splits','accepted_tokens')}, indent=2))
312
313
314if __name__ == '__main__':
315    main()
316