path: root/annotate_pgns/run_nnue_pipeline.py

#!/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 <path.stem>.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("--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("--material-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()

    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, "--loss-exponent", args.loss_exponent,
         "--material-loss-weight", args.material_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],
        args.force)

    print(f"\nPipeline complete: {args.outfile}")


if __name__ == "__main__":
    main()