1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
|
#!/usr/bin/env python3
"""Parallelize mine_king_mates.py across N shards by game-boundary byte
offset, since a full 642k-game scan is CPU-bound (python-chess replay +
SAN parsing per game) and this box has 24 cores. Finds "[Event " line
starts via a fast raw byte scan (not a full PGN parse) to get exact game
boundaries, splits into N even-sized shards by game count, launches one
mine_king_mates.py subprocess per shard with --byte-start/--byte-end, then
concatenates the resulting CSVs (one header, rows from all shards).
"""
import argparse
import csv
import os
import subprocess
import sys
import time
def find_game_offsets(pgn_file):
offsets = []
pos = 0
with open(pgn_file, "rb") as f:
for line in f:
if line.startswith(b"[Event "):
offsets.append(pos)
pos += len(line)
return offsets, pos # pos ends up as total file size
def main():
ap = argparse.ArgumentParser()
ap.add_argument("pgn_file")
ap.add_argument("out_dir")
ap.add_argument("--num-shards", type=int, default=8)
ap.add_argument("--extra-args", default="",
help="Extra args passed through to mine_king_mates.py, "
"e.g. '--control-ratio 1.0'")
args = ap.parse_args()
os.makedirs(args.out_dir, exist_ok=True)
script_dir = os.path.dirname(os.path.abspath(__file__))
mine_script = os.path.join(script_dir, "mine_king_mates.py")
print("Scanning for game boundaries...", file=sys.stderr)
t0 = time.time()
offsets, file_size = find_game_offsets(args.pgn_file)
print(f"Found {len(offsets)} games, {file_size} bytes, "
f"in {time.time() - t0:.1f}s", file=sys.stderr)
n = args.num_shards
games_per_shard = len(offsets) // n
shard_bounds = []
for i in range(n):
start_idx = i * games_per_shard
end_idx = (i + 1) * games_per_shard if i < n - 1 else len(offsets)
byte_start = offsets[start_idx]
byte_end = offsets[end_idx] if end_idx < len(offsets) else file_size
shard_bounds.append((byte_start, byte_end))
procs = []
for i, (byte_start, byte_end) in enumerate(shard_bounds):
out_csv = os.path.join(args.out_dir, f"shard_{i}.csv")
log_file = os.path.join(args.out_dir, f"shard_{i}.log")
cmd = [
sys.executable, mine_script, args.pgn_file, out_csv,
"--byte-start", str(byte_start),
"--byte-end", str(byte_end),
] + args.extra_args.split()
print(f"Launching shard {i}: bytes [{byte_start}, {byte_end}) -> {out_csv}",
file=sys.stderr)
with open(log_file, "w") as logf:
p = subprocess.Popen(cmd, stdout=logf, stderr=subprocess.STDOUT)
procs.append((i, p, out_csv))
failed = []
for i, p, out_csv in procs:
rc = p.wait()
print(f"Shard {i} exited with code {rc}", file=sys.stderr)
if rc != 0:
failed.append(i)
if failed:
print(f"FAILED shards: {failed} -- not merging, check logs in {args.out_dir}",
file=sys.stderr)
sys.exit(1)
merged_path = os.path.join(args.out_dir, "merged.csv")
header = None
total_rows = 0
with open(merged_path, "w", newline="") as out_f:
writer = csv.writer(out_f)
for i, _, out_csv in procs:
with open(out_csv, newline="") as in_f:
reader = csv.reader(in_f)
this_header = next(reader)
if header is None:
header = this_header
writer.writerow(header)
elif this_header != header:
print(f"WARNING: shard {i} header mismatch: {this_header}",
file=sys.stderr)
for row in reader:
writer.writerow(row)
total_rows += 1
print(f"Merged {total_rows} rows from {n} shards into {merged_path}",
file=sys.stderr)
if __name__ == "__main__":
main()
|
