chore: add helper scripts from Andy Grant (#170)

DeveloperPaul123 · DeveloperPaul123 · commit 8a468de4dbba · 2025-11-20T11:24:45.000-05:00
bench: 1049220
diff --git a/scripts/batched_execution_pool.py b/scripts/batched_execution_pool.py
@@ -0,0 +1,66 @@
+import multiprocessing
+import typing
+
+
+class BatchedExecutionPool:
+    def __init__(
+        self,
+        input_generator: typing.Generator,
+        process_function: typing.Callable[..., typing.Any],
+        process_function_args: typing.Optional[typing.Iterable[typing.Any]],
+    ) -> None:
+        self.input_generator = input_generator
+        self.process_function = process_function
+        self.process_function_args = process_function_args
+
+    def execute(self, threads: int, batchsize: int) -> int:
+        in_queue = multiprocessing.Queue()
+        out_queue = multiprocessing.Queue()
+
+        workers = [
+            multiprocessing.Process(
+                target=BatchedExecutionPool._process_function_wrapper,
+                args=(
+                    in_queue,
+                    out_queue,
+                    self.process_function,
+                    self.process_function_args,
+                ),
+                daemon=True,
+            )
+            for f in range(threads)
+        ]
+
+        for worker in workers:
+            worker.start()
+
+        while True:
+            n = self._enqueue_elements(batchsize, in_queue)
+            for f in range(n):
+                yield out_queue.get()
+            if n != batchsize:
+                break
+
+        for f in range(threads):
+            in_queue.put(None)
+
+        for worker in workers:
+            worker.join()
+
+    def _enqueue_elements(self, batchsize: int, in_queue: multiprocessing.Queue) -> int:
+        for f in range(batchsize):
+            try:
+                in_queue.put(next(self.input_generator))
+            except StopIteration:
+                return f
+        return batchsize
+
+    @staticmethod
+    def _process_function_wrapper(
+        in_queue: multiprocessing.Queue,
+        out_queue: multiprocessing.Queue,
+        process_function: typing.Callable[..., typing.Any],
+        process_function_args: typing.Optional[typing.Iterable[typing.Any]],
+    ) -> None:
+        while (data := in_queue.get()) != None:
+            out_queue.put(process_function(data, process_function_args))
diff --git a/scripts/replay.py b/scripts/replay.py
@@ -0,0 +1,133 @@
+#!/usr/bin/env python3
+
+import argparse
+import chess
+import chess.pgn
+import sys
+import traceback
+import io
+
+from subprocess import Popen, PIPE
+from batched_execution_pool import BatchedExecutionPool
+
+class Engine():
+
+    def __init__(self, binary):
+        self.engine = Popen([binary], stdin=PIPE, stdout=PIPE, universal_newlines=True, shell=True)
+        self.uci_ready()
+
+    def write_line(self, line):
+        self.engine.stdin.write(line)
+        self.engine.stdin.flush()
+
+    def read_line(self):
+        return self.engine.stdout.readline().rstrip()
+
+    def uci_ready(self):
+        self.write_line('isready\n')
+        while self.read_line() != 'readyok': pass
+
+    def quit(self):
+        self.write_line('quit\n')
+
+def parse_args():
+
+    p = argparse.ArgumentParser(description='Add UCI options with: --option.Name=Value')
+    p.add_argument('--engine', type=str, required=True, help='Path to the engine or engine name')
+    p.add_argument('--pgn',    type=str, required=True, help='Path to the PGN file')
+    p.add_argument('--player', type=str, required=True, help='Name of the player')
+    args, unknown = p.parse_known_args()
+
+    uci_options = []
+    for value in unknown:
+        if '=' in value and value.startswith('--option.'):
+            uci_options.append(value[len('--option.'):].split('='))
+
+    return args, uci_options
+
+def convert_uci_to_score(score_type, score_value):
+
+    if score_type == 'cp' and int(score_value) == 0:
+        return '0.00'
+
+    if score_type == 'cp':
+        return '%+.2f' % (float(score_value) / 100.0)
+
+    if score_type == 'mate' and int(score_value) < 0:
+        return '-M%d' % (abs(2 * int(score_value)))
+
+    if score_type == 'mate' and int(score_value) > 0:
+        return '+M%d' % (abs(2 * int(score_value) - 1))
+
+    raise Exception('Unable to process Score (%s, %s)' % (score_type, score_value))
+
+
+def game_generator(args):
+    with open(args.pgn) as pgn_file:
+        while game := chess.pgn.read_game(pgn_file):
+            yield (str(game))
+
+def replay_game(game_str, args):
+
+    game = chess.pgn.read_game(io.StringIO(game_str))
+    process_args, uci_options = args
+
+    is_white = game.headers.get('White') == process_args.player
+    is_black = game.headers.get('Black') == process_args.player
+    assert is_white or is_black
+
+    fen = game.headers.get('FEN', None)
+    pos = 'position fen %s moves' % (fen) if fen else 'position startpos moves'
+
+    engine = Engine(process_args.engine)
+    for opt, value in uci_options:
+        engine.write_line('setoption name %s value %s\n' % (opt, value))
+
+    node = game
+    while node.variations:
+
+        next_node = node.variation(0)
+
+        if (node.turn() and is_white) or (not node.turn() and is_black):
+
+            try:
+                pgn_score, pgn_depths, pgn_timems, pgn_nodes = next_node.comment.split()
+                pgn_nodes = int(pgn_nodes)
+
+            except Exception:
+                break
+
+            engine.uci_ready()
+            engine.write_line('%s\ngo nodes %d\n' % (pos, pgn_nodes))
+
+            score = None
+            while 'bestmove' not in (line := engine.read_line()):
+                if ' score ' in line:
+                    score = line.split(' score ')[1].split('nodes ')[0].split()[:2]
+            best_move = line.split()[1]
+
+            try:
+                assert best_move == next_node.move.uci()
+                assert convert_uci_to_score(*score) == pgn_score
+            except AssertionError:
+                engine.quit()
+                return 'Failed: ' + game_str
+
+        pos += ' ' + next_node.move.uci()
+        node = next_node
+
+    engine.quit()
+
+if __name__ == '__main__':
+
+    args, uci_options = parse_args()
+
+    pool = BatchedExecutionPool(
+        input_generator       = game_generator(args),
+        process_function      = replay_game,
+        process_function_args = [args, uci_options],
+    )
+
+    for result in pool.execute(threads=15, batchsize=256):
+        if result:
+            print (result)
diff --git a/scripts/reproduce.py b/scripts/reproduce.py
@@ -0,0 +1,95 @@
+#!/usr/bin/env python3
+
+# Given a .pgn file containing a single game, attempt to create a series of UCI commands
+# for the specified player, which aims to perfectly recreate the game, down to the node.
+#
+# Bugs, oddities, or just interesting situations, are often seen in PGNs from games played.
+# Those games are almost always played with standard Fischer time controls. As a result,
+# reproducing those games can be very challenging.
+#
+# FastChess, and OpenBench's fork of cutechess, supply the node counters most recently
+# reported by the playing engine in the PGN. Engines like Torch, which always produce
+# a final UCI report at the end of the search, and which have perfect determinism for
+# "go nodes <x>", are able to replay games if the exact node counters are known.
+#
+# This script makes use of a UCI extension, a command called "wait", which will block
+# until the current search is done. That allows piping the output of the script, directly
+# into an engine as stdin, to set up the engine state.
+#
+# Suppose that we have a pgn where Torch played an illegal move as white, during an SPRT
+# test. We could get a series of commands to reproduce the engine state, up to the final
+# search where the illegal move occurred, with the following:
+# ./reproduce.py --pgn bug.pgn --white --nodes --option.Hash=16
+
+import argparse
+import chess
+import chess.pgn
+
+def iterate_uci_options(unknown):
+    for arg in unknown:
+        if '=' in arg and arg.startswith('--option.'):
+            yield arg[len('--option.'):].split('=')
+
+def parse_args():
+
+    p = argparse.ArgumentParser(description='Add UCI options with: --option.Name=Value')
+    p.add_argument('--pgn', type=str, required=True, help='.pgn with only a single game')
+
+    # Must pick between replicating game with fixed nodes or fixed depth
+    p.add_argument('--nodes',  action='store_true', help='Generate commands using "go nodes"')
+    p.add_argument('--depth',  action='store_true', help='Generate commands using "go depth"')
+
+    # Provide commands only for the desired colour
+    p.add_argument('--white',  action='store_true', help='Generate commands for White')
+    p.add_argument('--black',  action='store_true', help='Generate commands for Black')
+
+    args, unknown = p.parse_known_args()
+
+    if args.nodes == args.depth:
+        raise Exception('Must use either --nodes or --depth')
+
+    if args.white == args.black:
+        raise Exception('Must use either --white or --black')
+
+    return args, iterate_uci_options(unknown)
+
+def main():
+
+    args, uci_options = parse_args()
+
+    with open(args.pgn) as pgn_file:
+        game = chess.pgn.read_game(pgn_file)
+
+    if not game:
+        raise Exception('Empty PGN file')
+
+    print ('uci')
+    for opt_name, opt_value in uci_options:
+        print ('setoption name %s value %s' % (opt_name, opt_value))
+    print ('ucinewgame')
+    print ('isready')
+
+    fen = game.headers.get('FEN', None)
+    pos = 'position fen %s moves' % (fen) if fen else 'position startpos moves'
+
+    node = game
+    while node.variations:
+
+        next_node = node.variation(0)
+
+        if (node.turn() and args.white) or (not node.turn() and args.black):
+
+            print (pos)
+
+            try: score, depths, timems, nodes = next_node.comment.split()
+            except: exit()
+            depth, seldepth = depths.split('/')
+
+            print ('go depth %s' % (depth) if args.depth else 'go nodes %s' % (nodes))
+            print ('wait')
+
+        pos += ' ' + next_node.move.uci()
+        node = next_node
+
+if __name__ == '__main__':
+    main()
