replication-run_model_on_csv.py

import argparse
import os
import os.path
import bz2
import csv
import multiprocessing
import humanize
import time
import queue

import chess

import maia_chess_backend

@maia_chess_backend.logged_main
def main():
    parser = argparse.ArgumentParser(description='Run model on all the lines of the csv', formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument('nets', help='nets dir')
    parser.add_argument('input', help='input CSV')
    parser.add_argument('output', help='output dir')
    parser.add_argument('--nrows', type=int, help='number of rows to read in', default = None)
    parser.add_argument('--lc0_depth', type=int, help='Number of rollouts/nodes for lc0 engines', default = None)
    parser.add_argument('--pool_size', type=int, help='Number of models to run in parallel', default = 64)
    parser.add_argument('--queueSize', type=int, help='Max number of games to cache', default = 1000)
    parser.add_argument('--multipv', type=int, help='Number of potetial boards to consider', default = None)
    parser.add_argument('--no_hist', help='Disable history', action = 'store_true')
    args = parser.parse_args()

    maia_chess_backend.printWithDate(f"Starting model {args.nets} analysis of {args.input} writing to {args.output}")

    models = []
    for name, _, files in os.walk(args.nets):
        if 'config.yaml' in files:
            models.append(name)

    maia_chess_backend.printWithDate(f"Found {len(models)} models")
    os.makedirs(args.output, exist_ok = True)
    with multiprocessing.Pool(args.pool_size + 1) as pool, multiprocessing.Manager() as manager:
        for model_path in models:
            runModel(model_path, manager, pool, args.pool_size,args.queueSize, args.input, args.output, args.nrows, args.lc0_depth, args.multipv, args.no_hist)
    maia_chess_backend.printWithDate(f"Done all models")

def runModel(model_path, manager, pool, pool_size, queueSize, data_path, output_dir, nrows, lc0_depth, multipv, no_hist):
    unproccessedQueue = manager.Queue(queueSize)
    resultsQueue = manager.Queue(queueSize)

    try:
        model, config = maia_chess_backend.load_model_config(model_path, lc0_depth = lc0_depth)
        name = config['name']
        display_name = config['display_name']
    except NotImplementedError:
        maia_chess_backend.printWithDate(f"{model_path} model not implemented", colour = 'red')
        return
    output_file = os.path.join(output_dir, f"{name}_results.csv.bz2")
    runners = []
    for _ in range(pool_size - 1):
        runner = pool.apply_async(sequenceRunner, (unproccessedQueue, resultsQueue, model_path, lc0_depth, multipv))
        runners.append(runner)
    maia_chess_backend.printWithDate(f"{display_name} Started {len(runners)} runners", flush = True)
    writer = pool.apply_async(writerWorker, (output_file, resultsQueue, pool_size - 1, name, display_name, lc0_depth, multipv))
    tstart = time.time()
    with bz2.open(data_path, 'rt') as fin:
        reader = csv.DictReader(fin)
        board = chess.Board()
        current_game = None
        for i, row in enumerate(reader):
            if nrows is not None and i >= nrows:
                break
            if no_hist or row['game_id'] != current_game:
                current_game = row['game_id']
                board = chess.Board(fen = row['board'])
            unproccessedQueue.put((board, {
                            'game_id' : row['game_id'],
                            'move_ply' : row['move_ply'],
                            'move' : row['move'],
                            }
                        ))
            try:
                board.push_uci(row['move'])
            except ValueError:
                current_game = row['game_id']
            if i % 1000 == 0:
                maia_chess_backend.printWithDate(f"{name} row {i} in {humanize.naturaldelta(time.time() - tstart)} {unproccessedQueue.qsize()} in {resultsQueue.qsize()} out".ljust(50), end = '\r')
    for _ in range(pool_size - 1):
        unproccessedQueue.put('kill')
    writer_done = False
    if not writer_done and writer.ready():
        writer.get()
        writer_done = True
    while len(runners) > 0:
        for i in range(len(runners)):
            try:
                runners[i].get(1)
            except multiprocessing.TimeoutError:
                pass
            else:
                del runners[i]
                break
        if not writer_done and writer.ready():
            writer.get()
            writer_done = True
    if not writer_done:
        writer.get()

    maia_chess_backend.printWithDate(f"{name} done {i} rows in {humanize.naturaldelta(time.time() - tstart)}".ljust(70))

def sequenceRunner(inputQueue, outputQueue, model_path, lc0_depth, multipv):
    model, config = maia_chess_backend.load_model_config(model_path, lc0_depth = lc0_depth)
    while True:
        try:
            dat = inputQueue.get()
        except queue.Empty:
            break
        if dat == 'kill':
            outputQueue.put('kill', True)
            break
        else:
            board, row = dat
            if multipv is None:
                m_move_obj, m_cp = model.getMoveWithCP(board)
                m_move = m_move_obj.uci()
                outputQueue.put((m_move, m_cp, row))
            else:
                moves = model.getTopMovesCP(board, multipv)
                outputQueue.put(([m for m, c in moves], [c for m, c in moves], row))

def writerWorker(outputFile, inputQueue, num_readers, name, display_name, rl_depth, multipv):
    num_kill_remaining = num_readers
    with bz2.open(outputFile, 'wt') as f:
        if multipv is None:
            writer = csv.DictWriter(f, ['game_id', 'move_ply', 'player_move', 'model_move', 'model_cp', 'model_correct', 'model_name', 'model_display_name', 'rl_depth'])
        else:
            moveheader = []
            for i in range(multipv):
                moveheader+= [f'model_move_{i}', f'model_cp_{i}']
            writer = csv.DictWriter(f, ['game_id', 'move_ply', 'player_move', 'model_correct', 'model_name', 'model_display_name', 'rl_depth'] + moveheader)
        writer.writeheader()
        while True:
            dat = inputQueue.get()
            if dat == 'kill':
                num_kill_remaining -= 1
                if num_kill_remaining <= 0:
                    break
            else:
                m_move, m_cp, row = dat
                write_dict = {
                    'game_id' : row['game_id'],
                    'move_ply' : row['move_ply'],
                    'player_move' : row['move'],
                    'model_correct' : row['move'] == m_move,
                    'model_name' : name,
                    'model_display_name' : display_name,
                    'rl_depth' : rl_depth,
                }
                if multipv is None:
                    write_dict['model_move'] = m_move
                    write_dict['model_cp'] = m_cp
                else:
                    for i, m in enumerate(m_move):
                        write_dict[f'model_move_{i}'] = m
                    for i, c in enumerate(m_cp):
                        write_dict[f'model_cp_{i}'] = c
                writer.writerow(write_dict)

if __name__ == '__main__':
    main()