#!/usr/bin/env python3 """End-to-end NNUE training pipeline driver: feed it a glob of PGN files that already carry per-move score comments (match_play.py's "{+-}NN.NN/DEPTH PV..." format -- e.g. pgn/stockfish_self_play_training_data.pgn, pgn/self_play_games.pgn) and (after a long time -- extract is cheap, but a from-scratch training run over millions of positions is not) it drives every remaining stage through to a .nnue weights file the C engine can load. This does NOT run the engine's own sd-N search (that is annotate.py's job, for PGNs that do NOT already have scores -- see that script and extract_positions.py/build_games.py instead if that is what's needed). Every input PGN here already has its own score/depth/PV comment on each move, so the only extraction step is parsing those comments -- extract_selfplay_positions.py, despite its name, is a general reader for this comment format regardless of which engine(s) produced it (see that script's own MATE_CLAMP_THRESHOLD comment, added specifically to handle a foreign-engine/Stockfish-sourced corpus). Chains the existing, already-individually-tested pipeline scripts in this order: 1. extract_selfplay_positions.py each input PGN -> games.jsonl (one run per file, since that script takes a single PGN; outputs concatenated) 2. split_holdout.py games.jsonl -> train/holdout split (by game, not position) 3. fit_k_and_label.py fits the WDL-sigmoid K on train, labels both splits (sharded, numpy-vectorized) 4. quiet_filter_games.py drops in-check / SEE-positive- capture positions from both labeled splits 5. build_training_tensors.py quiet-filtered games -> cached CSR feature tensors (sharded) 6. nnue_train.py trains a checkpoint per epoch 7. export_nnue_weights.py latest checkpoint -> flat .nnue binary Each stage is skipped if its expected output already exists (pass --force to re-run everything regardless) -- stages 5-6 here can run for hours, so a crash/interrupt partway through should not mean starting over from stage 1. This is a thin orchestrator: it shells out to the existing, individually tested scripts rather than reimplementing any of their logic, and assumes they are run from this directory (annotate_pgns/). Usage: python3 run_nnue_pipeline.py \\ --pgn-glob '../pgn/stockfish*.pgn' \\ --rundir runs/stockfish_labeled \\ --outfile runs/stockfish_labeled/nnue_weights.nnue """ import argparse import glob import json import subprocess import sys import time from pathlib import Path HERE = Path(__file__).resolve().parent def split_pgn(path, n, outdir): """Splits one PGN file into n roughly-equal-game-count chunk files under outdir (named .chunk000.pgn etc.), returning their paths. Game boundaries are found by scanning for lines starting "[Event " (every game in this corpus starts with one -- see extract_selfplay_positions.py's own game headers, build_pgn_game() in match_play.py), tracked as raw byte offsets via f.buffer.tell() rather than parsed/re-rendered by python-chess -- avoids any risk of a PGN round-trip through its writer subtly reformatting a comment or header in a way extract_selfplay_positions.py's own parse wouldn't expect, and avoids the text-mode tell()-vs-byte-offset mismatch a chess.pgn.read_game-based scan would hit. A single-PGN-file run (98k games in one file, 2026-10-19's first real pipeline attempt) was otherwise stuck running extract_selfplay_positions.py single-threaded for well over an hour -- that parse is purely per-game, embarrassingly parallel, just never sharded because the pipeline only ever sharded across *separate input files*, not within one big one.""" offsets = [] with open(path, "rb") as f: pos = 0 for raw_line in f: if raw_line.startswith(b"[Event "): offsets.append(pos) pos += len(raw_line) end = pos n_games = len(offsets) if n_games == 0: return [path] n = min(n, n_games) boundaries = [offsets[round(i * n_games / n)] for i in range(n)] + [end] outdir = Path(outdir) outdir.mkdir(parents=True, exist_ok=True) stem = Path(path).stem chunk_paths = [] with open(path, "rb") as f: for i in range(n): start, stop = boundaries[i], boundaries[i + 1] f.seek(start) data = f.read(stop - start) chunk_path = outdir / f"{stem}.chunk{i:03d}.pgn" with open(chunk_path, "wb") as out: out.write(data) chunk_paths.append(chunk_path) return chunk_paths def run_stage(name, outputs, cmd, force): """Runs cmd (a list of str/Path args) unless every path in outputs already exists and force is False. Exits the whole pipeline on a non-zero return code -- no point starting a later stage against a broken earlier one.""" outputs = [Path(p) for p in outputs] if not force and all(p.exists() for p in outputs): print(f"[skip] {name}: all outputs already exist " f"({', '.join(str(p) for p in outputs)})") return print(f"[run] {name}") print(f" {' '.join(str(c) for c in cmd)}") t0 = time.time() subprocess.run([str(c) for c in cmd], check=True) print(f"[done] {name} ({time.time() - t0:.0f}s)") def best_checkpoint(ckpt_dir): """Returns the best-holdout checkpoint (nnue_train.py's best.json), not just the highest epoch number -- a run that overfits past its best epoch would otherwise silently export the worst-of-the-run checkpoint just because it happened to be trained last.""" ckpt_dir = Path(ckpt_dir) best_path = ckpt_dir / "best.json" if best_path.exists(): best = json.loads(best_path.read_text()) ckpt_path = ckpt_dir / best["ckpt"] print(f"[info] best checkpoint: epoch {best['epoch']}, " f"holdout_mse={best['holdout_mse']:.5f} ({ckpt_path})") return ckpt_path # Fallback for a ckpt_dir produced before best.json existed. ckpts = sorted(ckpt_dir.glob("ckpt_epoch*.pt")) if not ckpts: sys.exit(f"no checkpoints found in {ckpt_dir} -- nnue_train.py " f"stage did not produce any?") print(f"[warn] no best.json in {ckpt_dir} -- falling back to the " f"last-trained epoch, not necessarily the best holdout one") return ckpts[-1] def main(): ap = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--pgn-glob", action="append", required=True, help="glob pattern (quoted, so this script expands " "it rather than the shell) for input PGN " "files, each already carrying per-move " "score/depth/PV comments; repeatable") ap.add_argument("--rundir", required=True, help="directory for all intermediate/output " "artifacts of this run") ap.add_argument("--outfile", required=True, help="final .nnue weights file path") ap.add_argument("--event", default=None, help="extract_selfplay_positions.py: only extract " "games with this Event header (default: any)") ap.add_argument("--dates", nargs="*", default=None, help="extract_selfplay_positions.py: only extract " "games with Date header in this list " "(default: any)") ap.add_argument("--exclude-movers", nargs="*", default=[], help="extract_selfplay_positions.py: skip plies " "whose mover's PGN name is in this list") ap.add_argument("--holdout-frac", type=float, default=0.1, help="split_holdout.py's game-level holdout " "fraction") ap.add_argument("--workers", type=int, default=16, help="worker processes for every sharded Python " "stage (fit_k_and_label.py, " "build_training_tensors.py). Default 16, " "kept below this box's 24 cores to leave " "other jobs some breathing room") ap.add_argument("--threads", type=int, default=16, help="torch CPU thread count for nnue_train.py, " "same reasoning as --workers (ignored if " "--device is not cpu)") ap.add_argument("--device", default="cpu", help="nnue_train.py's --device, e.g. 'cuda' on a " "rented GPU box") ap.add_argument("--lambda", dest="lam", type=float, default=0.9, help="build_training_tensors.py's " "target = lambda*p_search + (1-lambda)*r_outcome " "for non-exact rows") ap.add_argument("--k", type=float, default=None, help="nnue_train.py's --k (the WDL-sigmoid scale, " "pred=sigmoid(raw/k)) -- previously NOT passed " "through at all, silently falling back to " "nnue_train.py's own DEFAULT_K (1.3123, a Phase-A " "placeholder), regardless of what K " "fit_k_and_label.py actually fit for this corpus's " "own p_search labels (stage 3 below) -- a real " "label/loss K mismatch, fixed 2026-10-20. Required " "now (no default) specifically so this can't go " "silently unmatched again: read fit_k_and_label.py's " "own printed 'Fitted K = ...' line for this corpus " "and pass that (times whatever saturation-" "correction factor empirical testing supports -- " "see docs/NNUE.md's K-recalibration status).") ap.add_argument("--epochs", type=int, default=20) ap.add_argument("--patience", type=int, default=0, help="nnue_train.py's --patience (0 disables early " "stopping, always runs the full --epochs)") ap.add_argument("--batch-size", type=int, default=8192) ap.add_argument("--lr", type=float, default=1e-3) ap.add_argument("--loss-exponent", type=float, default=2.0) ap.add_argument("--score-loss-weight", type=float, default=0.0) ap.add_argument("--force", action="store_true", help="re-run every stage even if its outputs " "already exist") args = ap.parse_args() if args.k is None: sys.exit("--k is required (no default, deliberately -- see its " "help text): pass the K you actually want nnue_train.py's " "loss to use. Do not guess -- run fit_k_and_label.py " "first (or check its log from stage 3 below) to see what " "it fits for THIS corpus, and apply whatever saturation-" "correction factor has been validated empirically before " "trusting its raw fit.") pgn_paths = sorted({p for pattern in args.pgn_glob for p in glob.glob(pattern)}) if not pgn_paths: sys.exit(f"--pgn-glob matched no files: {args.pgn_glob!r}") print(f"Input PGNs ({len(pgn_paths)}): {pgn_paths}") rundir = Path(args.rundir) rundir.mkdir(parents=True, exist_ok=True) games_path = rundir / "games.jsonl" train_path = rundir / "games.train.jsonl" holdout_path = rundir / "games.holdout.jsonl" train_labeled_path = rundir / "games.train.labeled.jsonl" holdout_labeled_path = rundir / "games.holdout.labeled.jsonl" train_quiet_path = rundir / "games.train.labeled.quiet.jsonl" holdout_quiet_path = rundir / "games.holdout.labeled.quiet.jsonl" tensors_train_dir = rundir / "tensors_train" tensors_holdout_dir = rundir / "tensors_holdout" ckpt_dir = rundir / "nnue_ckpt" event_flag = ["--event", args.event] if args.event is not None else [] dates_flag = ["--dates", *args.dates] if args.dates else [] exclude_flag = (["--exclude-movers", *args.exclude_movers] if args.exclude_movers else []) # 1. Extract already-scored positions. extract_selfplay_positions.py # is purely per-game, embarrassingly parallel -- shard each input # PGN into args.workers game-aligned chunks (split_pgn, raw byte # copy, no PGN re-render) and run one extraction process per chunk # concurrently, rather than one single-threaded process per whole # input file (which is what left a 98k-game single-file run stuck # for over an hour on 2026-10-19). split_dir = rundir / "pgn_splits" all_chunks = [] # (chunk_path, part_path) pairs, in output order for pgn_path in pgn_paths: chunks = split_pgn(pgn_path, args.workers, split_dir) for chunk_path in chunks: part_path = rundir / f"{chunk_path.stem}.jsonl" all_chunks.append((chunk_path, part_path)) pending = [(c, p) for c, p in all_chunks if args.force or not p.exists()] if pending: print(f"[run] extract_selfplay_positions.py: {len(pending)} chunk(s) " f"across {len(pgn_paths)} input file(s), up to {args.workers} " f"concurrent") t0 = time.time() procs = [] for chunk_path, part_path in pending: cmd = [sys.executable, HERE / "extract_selfplay_positions.py", chunk_path, "--out", part_path, *event_flag, *dates_flag, *exclude_flag] while len(procs) >= args.workers: procs = [p for p in procs if p.poll() is None] if len(procs) >= args.workers: time.sleep(0.5) procs.append(subprocess.Popen([str(c) for c in cmd])) for p in procs: p.wait() failed = [p for p in procs if p.returncode != 0] if failed: sys.exit(f"extract_selfplay_positions.py: {len(failed)} chunk(s) " f"failed -- aborting pipeline") print(f"[done] extract_selfplay_positions.py ({time.time() - t0:.0f}s)") else: print(f"[skip] extract_selfplay_positions.py: all " f"{len(all_chunks)} chunk outputs already exist") if args.force or not games_path.exists(): print(f"[run] concatenating {len(all_chunks)} parts -> {games_path}") with open(games_path, "w") as out: for _, part_path in all_chunks: with open(part_path) as f: out.write(f.read()) else: print(f"[skip] concatenate: {games_path} already exists") # 2. Holdout split (by game). run_stage( "split_holdout.py", [train_path, holdout_path], [sys.executable, HERE / "split_holdout.py", "--in", games_path, "--train-out", train_path, "--holdout-out", holdout_path, "--holdout-frac", args.holdout_frac], args.force) # 3. Fit K + label both splits (sharded, numpy-vectorized). run_stage( "fit_k_and_label.py", [train_labeled_path, holdout_labeled_path], [sys.executable, HERE / "fit_k_and_label.py", "--train", train_path, "--holdout", holdout_path, "--train-out", train_labeled_path, "--holdout-out", holdout_labeled_path, "--workers", args.workers], args.force) # 4. Quiet-filter both labeled splits. run_stage( "quiet_filter_games.py (train)", [train_quiet_path], [sys.executable, HERE / "quiet_filter_games.py", "--in", train_labeled_path, "--out", train_quiet_path, "--workers", args.workers], args.force) run_stage( "quiet_filter_games.py (holdout)", [holdout_quiet_path], [sys.executable, HERE / "quiet_filter_games.py", "--in", holdout_labeled_path, "--out", holdout_quiet_path, "--workers", args.workers], args.force) # 5. Build cached CSR feature tensors (sharded). run_stage( "build_training_tensors.py (train)", [tensors_train_dir / "target.npy"], [sys.executable, HERE / "build_training_tensors.py", "--in", train_quiet_path, "--outdir", tensors_train_dir, "--lambda", args.lam, "--workers", args.workers], args.force) run_stage( "build_training_tensors.py (holdout)", [tensors_holdout_dir / "target.npy"], [sys.executable, HERE / "build_training_tensors.py", "--in", holdout_quiet_path, "--outdir", tensors_holdout_dir, "--lambda", args.lam, "--workers", args.workers], args.force) # 6. Train. Marker: train_complete.marker, written by nnue_train.py # itself right at the end of main() -- NOT the final epoch's # checkpoint file, which --patience early stopping may legitimately # mean never gets written, and which would otherwise make a clean # early-stopped run look unfinished forever to this existence # check. If a previous attempt left partial epoch checkpoints # without that marker (killed partway, or crashed), resume from the # latest one instead of silently restarting at epoch 0 and # discarding completed epochs. train_complete_marker = ckpt_dir / "train_complete.marker" resume_flag = [] if not args.force and not train_complete_marker.exists(): existing_ckpts = sorted(ckpt_dir.glob("ckpt_epoch*.pt")) if existing_ckpts: resume_flag = ["--resume", existing_ckpts[-1]] print(f"[info] nnue_train.py: resuming from " f"{existing_ckpts[-1]} (no train_complete.marker found)") run_stage( "nnue_train.py", [train_complete_marker], [sys.executable, HERE / "nnue_train.py", "--train-dir", tensors_train_dir, "--holdout-dir", tensors_holdout_dir, "--outdir", ckpt_dir, *resume_flag, "--epochs", args.epochs, "--batch-size", args.batch_size, "--lr", args.lr, "--k", args.k, "--loss-exponent", args.loss_exponent, "--score-loss-weight", args.score_loss_weight, "--threads", args.threads, "--device", args.device, "--patience", args.patience], args.force) # 7. Export the latest checkpoint to the flat .nnue binary format. ckpt = best_checkpoint(ckpt_dir) run_stage( "export_nnue_weights.py", [args.outfile], [sys.executable, HERE / "export_nnue_weights.py", ckpt, args.outfile, "--tensor-dir", tensors_train_dir, "--holdout-tensor-dir", tensors_holdout_dir], args.force) # 8. Refit the lazy-eval probe (nnue_vs_head.sh's NNUE_PROBE_WEIGHTS / # command.c's "nnue probeload") against THIS checkpoint -- it's a # separate offline ridge regression, not part of the exported .nnue # file, and must be re-fit every generation (same rule as the # calibration inside export_nnue_weights.py: not portable across # checkpoints). Written alongside outfile, not into it. See # fit_lazy_probe.py and the 2026-10-20 "keeping this stable across # retrains" / lazy-eval conversation -- this is the regression fit # only; the real search-node-traffic margin-safety calibration # (CALIBRATE_MARGIN_SAFETY) is a separate, not-yet-automated step, # deliberately not run here every generation (see that conversation # for why it may not need to be). probe_outfile = Path(args.outfile).with_name( Path(args.outfile).stem + "_probe.bin") run_stage( "fit_lazy_probe.py", [probe_outfile], [sys.executable, HERE / "fit_lazy_probe.py", ckpt, probe_outfile, "--train-tensor-dir", tensors_train_dir, "--holdout-tensor-dir", tensors_holdout_dir], args.force) print(f"\nPipeline complete: {args.outfile} (lazy probe: {probe_outfile})") if __name__ == "__main__": main()