#!/usr/bin/env python3 """Shard a PGN file across N single-threaded typhoon engine instances, score every position in every game at a fixed search depth, and reassemble an annotated PGN with { +score/depth PV } comments matching the self-play annotation format already used in ../pgn/self_play_games.pgn. Reusable across depths: --depth is a plain CLI argument, not baked into any file name or code path, so re-running at sd12/sd16/etc. is just a different --depth and --outdir. Usage: python3 annotate.py --input ../pgn/twic_filtered.pgn \ --depth 10 --shards 16 --hash 32m --outdir runs/sd10_twic See the module-level functions below for the four stages: shard, build engine scripts, run engines in parallel, parse + reassemble. """ import argparse import os import pty import re import select import signal import subprocess import sys import threading import time from pathlib import Path import chess import chess.pgn from quiet_common import position_is_quiet TELLOTHERS_RE = re.compile( r"^tellothers d(\d+), ([+-]?[\w.]+), [\d.]+s,\s*\d+knps, PV=(.*)$" ) FLOAT_RE = re.compile(r"^[+-]?\d+\.\d+$") MATE_RE = re.compile(r"^([+-])MATE(\d*)$") # digit(s) optional -- the # "FORCED MATE --> stop searching now" immediate-short-circuit path # (search.c) prints a bare "+MATE"/"-MATE" with no distance number at # all, not just the ordinary "+MATE" case. # --engine-protocol uci support: drives a real UCI engine (Stockfish) # directly over stdin/stdout with the native UCI protocol -- no xboard, # no PolyGlot adapter needed. Matches the standard "info depth N ... # score cp/mate X ... pv ..." thinking-output line Stockfish prints once # per completed depth. UCI_ITER_RE = re.compile( r"^info depth (\d+) .*?\bscore (cp|mate) (-?\d+).*?\bpv (.*)$" ) # A naive first attempt at this driver tried to interrupt each search # early (xboard's "force", or UCI's "stop") the instant the target depth # was seen, rather than waiting for "bestmove" -- motivated by "go depth # N" appearing to hang forever with no bestmove ever printed. That hang # turned out to be an artifact of this tool's own sandboxed subprocess # environment, NOT a Stockfish or protocol issue (confirmed 2026-10-20: # the identical "uci; isready; position fen ...; go depth 10" sequence # got a normal, prompt "bestmove" reply when run directly in an # interactive terminal outside that sandbox). So this driver just sends # "go depth " and waits for "bestmove" like any normal UCI # client -- no interrupt-and-abort workaround needed. UCI_HANDSHAKE_TIMEOUT_SEC = 15 UCI_GO_TIMEOUT_SEC = 120 # per position, generous but bounded def load_games(input_paths, max_games=None): """Load games, skipping ones with an unknown/unterminated result ('*') -- can't derive a WDL training label without a real outcome. Short-agreed-draw games are not separately filtered here: the >40-ply / >=2400-Elo filter already applied when twic_filtered.pgn was built (eval_tune/filter_pgn.py) rules those out by construction (confirmed empirically: no games below 40 plies in a 5000-game sample).""" games = [] for path in input_paths: with open(path, encoding="utf-8", errors="replace") as f: while max_games is None or len(games) < max_games: game = chess.pgn.read_game(f) if game is None: break if game.headers.get("Result", "*") == "*": continue games.append(game) if max_games is not None and len(games) >= max_games: break return games def shard_games(games, n_shards): shards = [[] for _ in range(n_shards)] for i, g in enumerate(games): shards[i % n_shards].append(g) return shards def positions_for_game(game, prefilter_quiet=False): """Yield the FEN *before* each move is played, skipping any ply at or after the game is already over (no legal moves to search). If prefilter_quiet, also skip any ply whose position fails the quiet-position test (quiet_common.position_is_quiet) -- checked against the move actually played in this game, since that's the only move available before a search has run. This is the same check/exclude criteria as quiet_filter.py's post-hoc filter, but SEE-checked against the *played* move rather than the engine's own PV first move -- see quiet_common.py's module docstring for why those are different training distributions, not interchangeable. ply numbering (pi, used by write_shard_script/annotate_shard/ build_games.py to re-align results) still counts only the plies actually yielded, so this must be called with the same prefilter_quiet value everywhere a run's ply indices are recomputed (build_games.py in particular).""" board = game.board() for move in game.mainline_moves(): if board.is_game_over(): break if not prefilter_quiet or position_is_quiet(board, move): yield board.fen() board.push(move) def build_shard_positions(shard_games_list, prefilter_quiet=False): """Return (fens, index) for one shard: the flat list of FENs in script order, and the parallel index recording which (game_idx_in_shard, ply_idx) each FEN corresponds to -- shared by both write_shard_script (typhoon protocol, writes an engine command script to disk) and run_shard_uci (uci protocol, drives the engine interactively with no script file at all).""" fens = [] index = [] for gi, game in enumerate(shard_games_list): for pi, fen in enumerate(positions_for_game(game, prefilter_quiet)): fens.append(fen) index.append((gi, pi)) return fens, index def write_shard_script(shard_games_list, script_path, depth, prefilter_quiet=False): """Write the engine command script for one shard. The engine processes 'go' commands strictly in script order with no interleaving, so the index returned by build_shard_positions zips back onto the flat sequence of tellothers output lines.""" fens, index = build_shard_positions(shard_games_list, prefilter_quiet) with open(script_path, "w") as out: for fen in fens: out.write(f"setboard {fen}\n") out.write(f"sd {depth}\n") out.write("go\n") return index def run_shard(engine_path, script_path, out_path, hash_size, cpus): cmd = [ str(engine_path), "--cpus", str(cpus), "--hash", hash_size, "--logfile", "-", "--batch", "--command", f"force; book name /nonexistent.book.bin; script {script_path}", ] with open(out_path, "w") as out: return subprocess.Popen(cmd, stdout=out, stderr=subprocess.STDOUT) def _uci_format_tellothers(depth, score_kind, score_val, pv): """Render one uci-protocol result as a synthetic tellothers line, so the existing parse_shard_output/_format_comment/annotate_shard pipeline (written for typhoon's own tellothers output) can be reused unchanged for uci-mode results too. `pv` must already be SAN (see _uci_pv_to_san) -- Stockfish's own PV is UCI long algebraic ("e2e4 d7d5"), not the SAN typhoon's tellothers line uses ("e4 d5"), and downstream readers (extract_positions.py, extract_selfplay_positions.py) key off SAN.""" if score_kind == "mate": score_tok = "+MATE" if score_val > 0 else "-MATE" else: score_tok = f"{score_val / 100:+.2f}" return f"tellothers d{depth}, {score_tok}, 0.0s, 0knps, PV={pv}\n" def _uci_pv_to_san(fen, uci_pv): """Convert a UCI-long-algebraic PV string ("e2e4 d7d5 ...") rooted at `fen` into typhoon's own tellothers SAN convention ("e4 d5 ..."). Stops at the first move that doesn't parse/isn't legal (shouldn't happen for a PV the engine itself just reported, but an illegal/ambiguous tail is far less useful than no PV, not more).""" board = chess.Board(fen) san_moves = [] for uci_move in uci_pv.split(): try: move = chess.Move.from_uci(uci_move) if move not in board.legal_moves: break san_moves.append(board.san(move)) board.push(move) except (ValueError, AssertionError): break return " ".join(san_moves) def run_shard_uci(engine_path, fens, depth, out_path): """Drive one persistent native-UCI engine subprocess (e.g. a Stockfish binary) through every FEN in `fens`, in order, writing one synthetic tellothers line per FEN to out_path. Spawned under a pty (via pty.fork), not a plain subprocess.Popen pipe -- verified live (2026-10-20) that this particular Stockfish build hangs after printing just its banner + "id name" line, with no more output ever arriving, when its stdout is a plain pipe (the "id author"/options/"uciok" lines, and everything after, simply never get produced -- not merely buffered: a pty.fork-based read of the exact same binary gets the whole handshake instantly). The earlier guess that this was xboard/PolyGlot-specific, and later that it was this tool's own sandboxed subprocess environment, were both wrong -- it reproduces under dangerouslyDisableSandbox too. Root cause not fully isolated (likely an isatty()-gated code path in Stockfish's own startup or NUMA/processor-count detection), but the pty workaround is confirmed sufficient.""" pid, fd = pty.fork() if pid == 0: try: os.execvp(str(engine_path), [str(engine_path)]) finally: os._exit(127) # os.execvp only returns on failure buf = b"" def send(cmd): os.write(fd, (cmd + "\n").encode()) def readline(timeout): nonlocal buf while b"\n" not in buf: r, _, _ = select.select([fd], [], [], timeout) if not r: return None try: chunk = os.read(fd, 65536) except OSError: return None if not chunk: return None buf += chunk line, buf = buf.split(b"\n", 1) # The pty's line-discipline translates outgoing "\n" to "\r\n" # (termios ONLCR), so every line read back here has a trailing # "\r" that a plain subprocess.Popen pipe would never produce -- # strip it, or sentinel comparisons ("uciok", "readyok") and the # captured PV text silently never match/contain a stray "\r". return line.rstrip(b"\r").decode(errors="replace") + "\n" def wait_for(sentinel, timeout_sec): deadline = time.time() + timeout_sec while time.time() < deadline: line = readline(2) if line is None: continue if line.rstrip("\n") == sentinel: return True return False send("uci") wait_for("uciok", UCI_HANDSHAKE_TIMEOUT_SEC) # Match ../src/stockfish_xboard's own default: Stockfish's tbprobe.cpp # is Fathom-derived, so it reads the exact same .rtbw/.rtbz files # typhoon does -- no separate Stockfish-format tablebase needed. # SYZYGY_PATH="" explicitly disables TB probing rather than falling # back to the default, same as that wrapper. syzygy_path = os.environ.get("SYZYGY_PATH", "/zscratch/egtb") if syzygy_path: send(f"setoption name SyzygyPath value {syzygy_path}") send("isready") wait_for("readyok", UCI_HANDSHAKE_TIMEOUT_SEC) with open(out_path, "w") as out: for fen in fens: send(f"position fen {fen}") send(f"go depth {depth}") best = None deadline = time.time() + UCI_GO_TIMEOUT_SEC while time.time() < deadline: line = readline(2) if line is None: continue line = line.rstrip("\n") if line.startswith("bestmove"): break m = UCI_ITER_RE.match(line) if m: d = int(m.group(1)) if d >= depth: best = (d, m.group(2), int(m.group(3)), m.group(4)) # Keep reading up to "bestmove" even once the # target depth is seen: with "go depth N", # Stockfish always reports depth N as its final # iteration before bestmove, so there's nothing # more useful to capture, but exiting early here # (like the earlier interrupt-based design did) # would again require aborting the search instead # of letting it finish on its own. # Always write exactly one line per "go", even if the # deadline was hit with no captured depth -- annotate_shard # zips index and parsed results positionally, so skipping a # write here would desync every later position in this # shard, not just leave this one unannotated (the same class # of bug the module docstring describes for a dropped # tellothers line). "NA" matches neither FLOAT_RE nor # MATE_RE, so _format_comment's fallback passes it through # as-is instead of mis-rendering it as a real score. if best is not None: d, kind, val, uci_pv = best out.write(_uci_format_tellothers( d, kind, val, _uci_pv_to_san(fen, uci_pv))) else: out.write(f"tellothers d{depth}, NA, 0.0s, 0knps, PV=\n") send("quit") deadline = time.time() + 5 while time.time() < deadline: wpid, _ = os.waitpid(pid, os.WNOHANG) if wpid == pid: break time.sleep(0.1) else: try: os.kill(pid, signal.SIGKILL) os.waitpid(pid, 0) except ProcessLookupError: pass os.close(fd) def parse_shard_output(out_path): """Return a list of (depth, score_token, pv_string) in script order. score_token is kept as the engine's raw tellothers token (an ordinary "+1.23"-style float string, or "+MATE"/"-MATE") rather than eagerly parsed to float here -- a MATE token has no pawns magnitude to parse in the first place, and silently dropping any line that doesn't match a plain float is exactly the bug that used to desync this function's output from `index` (every position after a dropped mate/EGTB line in a shard got the wrong score) -- see docs/NNUE.md's "Bug found and fixed in annotate.py's output parser" status entry, previously patched only in the separate extract_positions.py reader, not here.""" results = [] with open(out_path, encoding="utf-8", errors="replace") as f: for line in f: m = TELLOTHERS_RE.match(line.rstrip("\n")) if m: depth, score_tok, pv = m.groups() results.append((int(depth), score_tok, pv.strip())) return results def _format_comment(depth, score_tok, pv): """Render one (depth, score_token, pv) result as a PGN comment in the same "{+-}SCORE/DEPTH PV" convention match_play.py's _format_eval_comment uses for self-play games (and extract_selfplay_positions.py's COMMENT_RE expects on read) -- translates the engine's own raw "+MATE"/"-MATE" tellothers token into that convention's "+M"/"-M" (no distance digit, matching match_play.py exactly -- mate distance isn't independently verified here either). The PV's own sentinel token (//<=>/<#>), when present, is passed through verbatim since that is exactly what both extract_positions.py and extract_selfplay_positions.py key their exactness detection off of.""" m = MATE_RE.match(score_tok) if FLOAT_RE.match(score_tok): score_str = score_tok elif m: score_str = "+M" if m.group(1) == "+" else "-M" else: score_str = score_tok # unrecognized token: pass through as-is comment = f"{score_str}/{depth}" if pv: comment += f" {pv}" return comment def annotate_shard(shard_games_list, index, results): """Attach comments to each game's mainline nodes in place. index and results are both in script order and the same length; index maps each result back to (game_idx_in_shard, ply_idx).""" by_game = {} for (gi, pi), res in zip(index, results): by_game.setdefault(gi, {})[pi] = res for gi, game in enumerate(shard_games_list): scores = by_game.get(gi, {}) node = game pi = 0 while node.variations: node = node.variations[0] if pi in scores: depth, score_tok, pv = scores[pi] node.comment = _format_comment(depth, score_tok, pv) pi += 1 def main(): ap = argparse.ArgumentParser(description=__doc__) ap.add_argument("--input", nargs="+", required=True, help="one or more input PGN files") ap.add_argument("--depth", type=int, required=True, help="fixed search depth (sd N) for every position") ap.add_argument("--shards", type=int, default=16) ap.add_argument("--hash", default="32m", help="hash table size per shard engine, e.g. 32m") ap.add_argument("--cpus", type=int, default=1, help="--cpus passed to each shard's engine") ap.add_argument("--engine", default="../src/typhoon") ap.add_argument("--engine-protocol", choices=["typhoon", "uci"], default="typhoon", help="'typhoon' drives --engine with its own " "'--command \"script \"' batch mode and " "parses its tellothers output (default). 'uci' " "drives --engine (a UCI engine binary, e.g. a " "Stockfish build) directly over stdin/stdout " "with the native UCI protocol instead ('uci', " "'position fen ...', 'go depth N', wait for " "'bestmove'). --hash/--cpus are ignored in this " "mode; set hash/threads via UCI setoption " "yourself if needed.") ap.add_argument("--outdir", required=True) ap.add_argument("--max-games", type=int, default=None, help="limit to the first N games total, for smoke testing") ap.add_argument("--prefilter-quiet", action="store_true", help="skip in-check positions and SEE-positive-capture " "played moves before ever sending them to the " "engine, instead of searching everything and " "filtering after the fact (see quiet_filter.py / " "quiet_filter_games.py). Changes the quiet " "criterion's move source from the engine's own PV " "to the move actually played -- see quiet_common.py.") args = ap.parse_args() outdir = Path(args.outdir) outdir.mkdir(parents=True, exist_ok=True) engine_path = Path(args.engine).resolve() if not engine_path.exists(): sys.exit(f"engine not found: {engine_path}") print(f"Loading games from {args.input} ...") games = load_games(args.input, args.max_games) print(f" {len(games)} games loaded") shards = shard_games(games, args.shards) if args.prefilter_quiet: print("Pre-filtering: skipping in-check / SEE-positive-played-move " "positions before search.") indices = [] shard_fens = [] # only populated/used in uci mode if args.engine_protocol == "typhoon": print(f"Writing {args.shards} shard scripts at sd {args.depth} ...") for i, shard in enumerate(shards): script_path = outdir / f"shard_{i:02d}.script" idx = write_shard_script(shard, script_path, args.depth, args.prefilter_quiet) indices.append(idx) print(f" shard {i:02d}: {len(shard)} games, {len(idx)} positions") else: print(f"Building {args.shards} shard position lists at sd " f"{args.depth} ...") for i, shard in enumerate(shards): fens, idx = build_shard_positions(shard, args.prefilter_quiet) shard_fens.append(fens) indices.append(idx) print(f" shard {i:02d}: {len(shard)} games, {len(idx)} positions") print("Launching shard engines in parallel ...") t0 = time.time() if args.engine_protocol == "typhoon": procs = [] for i in range(args.shards): script_path = outdir / f"shard_{i:02d}.script" out_path = outdir / f"shard_{i:02d}.out" if indices[i]: procs.append((i, run_shard(engine_path, script_path, out_path, args.hash, args.cpus))) # Wait for the slowest shard. for i, p in procs: rc = p.wait() if rc != 0: print(f" WARNING: shard {i:02d} exited with code {rc}", file=sys.stderr) else: # No per-engine OS process launched up front here the way # run_shard does -- run_shard_uci owns its own subprocess # internally and blocks for the whole shard, so parallelism # across shards comes from running each one in its own thread # instead (these are I/O-bound waits on each engine subprocess's # pipe, so the GIL isn't a bottleneck here). threads = [] for i in range(args.shards): if not indices[i]: continue out_path = outdir / f"shard_{i:02d}.out" t = threading.Thread( target=run_shard_uci, args=(engine_path, shard_fens[i], args.depth, out_path), ) t.start() threads.append(t) for t in threads: t.join() elapsed = time.time() - t0 total_positions = sum(len(idx) for idx in indices) print(f"All shards done in {elapsed:.1f}s " f"({total_positions / elapsed:.2f} positions/sec aggregate)") print("Parsing output and reassembling annotated PGN ...") final_path = outdir / "annotated.pgn" total_annotated = 0 with open(final_path, "w") as out: for i, shard in enumerate(shards): out_path = outdir / f"shard_{i:02d}.out" results = parse_shard_output(out_path) if indices[i] else [] if len(results) != len(indices[i]): print(f" WARNING: shard {i:02d} expected {len(indices[i])} " f"results, got {len(results)} -- some positions may be " f"unannotated", file=sys.stderr) annotate_shard(shard, indices[i], results) total_annotated += len(results) for game in shard: print(game, file=out, end="\n\n") print(f"Wrote {final_path} ({len(games)} games, " f"{total_annotated} annotated positions)") if __name__ == "__main__": main()