src/match.py - master - keorlov/aichess

src/match.py
471 строка17 KB
Konstantin Orlov
wip
06 июн 2026, 13:19
06 июн 2026, 13:19cb645b7
"""Chess match logic — play a match between two engines.

A match consists of two games with colours swapped.
Each move is subject to a timeout enforced via a long-lived subprocess worker
(``EngineWorker``) so the cost of process spawn + engine import is paid once
per match instead of once per move.
"""

import importlib.util
import inspect
import multiprocessing as mp
import os
import queue as _queue
import time
from dataclasses import dataclass, field
from typing import Optional, Tuple

import chess

from .logging_setup import get_error_logger

_WORKER_TIMEOUT = 5


@dataclass
class GameResult:
    white: str
    black: str
    result: str          # "1-0" / "0-1" / "1/2-1/2" / "0-0"
    termination: str     # "checkmate" / "stalemate" / "timeout" / ...
    num_moves: int       # plies (half-moves) played
    pgn: str = ""


@dataclass
class MatchResult:
    engine_a: str
    engine_b: str
    score_a: float = 0.0
    score_b: float = 0.0
    games: list = field(default_factory=list)


class EngineWorker:
    """Long-lived subprocess worker that runs ``choose_move`` on demand.

    A single subprocess per engine is kept alive for the duration of a match
    and reused across all moves. The host pays the cost of process spawn +
    engine import once instead of once per move (a ~10-100x speedup on slow
    engines). On timeout, the subprocess is terminated and a fresh one is
    started for the next request.
    """

    def __init__(self, engine_path: str, engine_name: str,
                 engine_timeout_fraction: float = 1.0):
        self.engine_path = engine_path
        self.engine_name = engine_name
        self.engine_timeout_fraction = engine_timeout_fraction
        self._ctx = mp.get_context("spawn")
        self._input_q: mp.Queue = self._ctx.Queue()
        self._output_q: mp.Queue = self._ctx.Queue()
        self._proc: Optional[mp.Process] = None
        self._start()

    def _start(self):
        self._proc = self._ctx.Process(
            target=_engine_worker_loop,
            args=(self.engine_path, self._input_q, self._output_q),
            name=self.engine_name,
        )
        self._proc.start()

    def get_move(self, fen: str, timeout: float) -> Tuple[Optional[str], Optional[str]]:
        """Get a move from the engine. Returns ``(move_uci, error_kind)``.

        ``error_kind`` is ``None`` on success, or one of ``"timeout"``,
        ``"import_error"``, ``"no_choose_move"``, ``"choose_move_not_callable"``,
        ``"exception"``, ``"bad_return_type"``, ``"import_failed"`` on failure.
        """
        if self._proc is None or not self._proc.is_alive():
            self._start()

        # Drain any stale messages left over from a previous timeout/restart.
        _drain_queue(self._output_q)

        try:
            engine_timeout = timeout * self.engine_timeout_fraction
            self._input_q.put((fen, engine_timeout))
        except Exception:
            return None, "worker_dead"

        try:
            msg = self._output_q.get(timeout=timeout)
        except _queue.Empty:
            self._terminate_and_restart()
            return None, "timeout"

        if isinstance(msg, str):
            return msg, None
        if isinstance(msg, dict):
            reason = msg.get("error", "error")
            # The worker may have died after raising — restart lazily next call.
            if not self._proc.is_alive():
                self._start()
            return None, reason
        return None, "error"

    def _terminate_and_restart(self):
        if self._proc is not None and self._proc.is_alive():
            self._proc.terminate()
            self._proc.join(_WORKER_TIMEOUT)
            if self._proc.is_alive():
                self._proc.kill()
                self._proc.join(_WORKER_TIMEOUT)
        _drain_queue(self._input_q)
        _drain_queue(self._output_q)
        self._start()

    def close(self):
        """Shut the worker down cleanly. Safe to call multiple times."""
        if self._proc is None:
            return
        try:
            if self._proc.is_alive():
                self._input_q.put(None)  # poison pill
                self._proc.join(_WORKER_TIMEOUT)
        except Exception:
            pass
        if self._proc.is_alive():
            self._proc.terminate()
            self._proc.join(_WORKER_TIMEOUT)
            if self._proc.is_alive():
                self._proc.kill()
                self._proc.join(_WORKER_TIMEOUT)
        self._proc = None


def _drain_queue(q):
    try:
        while True:
            q.get_nowait()
    except Exception:
        pass


def _engine_worker_loop(engine_path: str, input_q: mp.Queue, output_q: mp.Queue):
    """Long-lived worker: import the engine once, then answer FEN→UCI requests.

    Protocol:
        - Receive FEN string from ``input_q``. ``None`` is a poison pill to exit.
        - Send back a UCI string on success, or a ``{"error": ..., "detail": ...}``
          dict on failure.
    """
    try:
        spec = importlib.util.spec_from_file_location("_engine", engine_path)
        if spec is None or spec.loader is None:
            output_q.put({"error": "import_failed", "detail": "spec/loader is None"})
            return
        try:
            mod = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(mod)
        except Exception as e:
            output_q.put({"error": "import_error", "detail": str(e)[:200]})
            return

        fn = getattr(mod, "choose_move", None)
        if fn is None:
            output_q.put({"error": "no_choose_move",
                          "detail": "module has no choose_move attribute"})
            return
        if not callable(fn):
            output_q.put({"error": "choose_move_not_callable",
                          "detail": "choose_move is not callable"})
            return

        try:
            sig = inspect.signature(fn)
            accepts_timeout = "move_timeout" in sig.parameters or any(
                p.kind == inspect.Parameter.VAR_KEYWORD
                for p in sig.parameters.values()
            )
            if any(
                p.kind == inspect.Parameter.VAR_POSITIONAL
                for p in sig.parameters.values()
            ) and not accepts_timeout:
                accepts_timeout = True
        except (TypeError, ValueError):
            accepts_timeout = False

        while True:
            try:
                msg = input_q.get()
            except (EOFError, OSError):
                break
            if msg is None:
                break
            if isinstance(msg, tuple) and len(msg) == 2:
                fen, move_timeout = msg
            else:
                fen, move_timeout = msg, 1.0
            try:
                if accepts_timeout:
                    result = fn(fen, move_timeout)
                else:
                    result = fn(fen)
            except Exception as e:
                output_q.put({"error": "exception", "detail": str(e)[:200]})
                continue
            if isinstance(result, str):
                output_q.put(result)
            else:
                output_q.put({"error": "bad_return_type",
                              "detail": f"returned {type(result).__name__} instead of str"})
    except Exception as e:
        try:
            output_q.put({"error": "worker_fatal", "detail": str(e)[:200]})
        except Exception:
            pass


def play_match(
    engine_a_name: str,
    engine_a_path: str,
    engine_b_name: str,
    engine_b_path: str,
    move_timeout: float,
    max_moves: int = 200,
    verbose: int = 0,
    colors: Tuple[str, str] = ("a", "b"),
    engine_timeout_fraction: float = 0.8,
) -> MatchResult:
    """Play a two-game match with colours configurable per game.

    Args:
        engine_a_name: Display name.
        engine_a_path: Absolute path to engine_a's engine.py.
        engine_b_name: Display name.
        engine_b_path: Absolute path to engine_b's engine.py.
        move_timeout: Seconds per move (the hard-kill budget enforced by the host).
        max_moves: Maximum plies (half-moves) per game before adjudicated as draw.
            ``200`` plies = 100 full moves (200 plies is the default).
        verbose: Log level (0 = normal, 1 = detailed, 2 = per-move SAN).
        colors: ``(game1_white, game2_white)`` where each value is ``"a"`` or
            ``"b"`` indicating which engine plays white in that game. Defaults
            to ``("a", "b")`` (the historical behaviour). Use ``("a", "b")``
            for a standard colour-swap; pass an explicit pair for fair
            pairings.
        engine_timeout_fraction: Fraction of ``move_timeout`` actually passed
            to each engine's ``choose_move``. The remaining ``1 - fraction``
            is the safety margin the host keeps to avoid hard-killing the
            engine. Defaults to ``0.8``.
    """

    match = MatchResult(engine_a=engine_a_name, engine_b=engine_b_name)

    worker_a = EngineWorker(engine_a_path, engine_a_name,
                            engine_timeout_fraction=engine_timeout_fraction)
    worker_b = EngineWorker(engine_b_path, engine_b_name,
                            engine_timeout_fraction=engine_timeout_fraction)
    try:
        # Game 1
        w1_name, b1_name, w1_worker, b1_worker = _resolve_colors(
            colors[0], engine_a_name, engine_b_name, worker_a, worker_b
        )
        if verbose >= 1:
            print(f"    Game 1: {w1_name} (W) vs {b1_name} (B)")
        g1 = _play_game(w1_name, w1_worker, b1_name, b1_worker,
                        move_timeout, max_moves, verbose)
        sa, sb = _game_points(engine_a_name, g1)
        match.score_a += sa
        match.score_b += sb
        match.games.append(g1)
        if verbose >= 1:
            print(f"      Result: {g1.result} ({g1.termination}), {g1.num_moves} plies")

        # Game 2
        w2_name, b2_name, w2_worker, b2_worker = _resolve_colors(
            colors[1], engine_a_name, engine_b_name, worker_a, worker_b
        )
        if verbose >= 1:
            print(f"    Game 2: {w2_name} (W) vs {b2_name} (B)")
        g2 = _play_game(w2_name, w2_worker, b2_name, b2_worker,
                        move_timeout, max_moves, verbose)
        sa2, sb2 = _game_points(engine_a_name, g2)
        match.score_a += sa2
        match.score_b += sb2
        match.games.append(g2)
        if verbose >= 1:
            print(f"      Result: {g2.result} ({g2.termination}), {g2.num_moves} plies")
    finally:
        worker_a.close()
        worker_b.close()

    return match


def _resolve_colors(white_is, engine_a_name, engine_b_name, worker_a, worker_b):
    if white_is == "a":
        return engine_a_name, engine_b_name, worker_a, worker_b
    return engine_b_name, engine_a_name, worker_b, worker_a


def _game_points(engine_a_name: str, game: GameResult) -> tuple[float, float]:
    if game.result == "0-0":
        return 0.0, 0.0
    if game.result == "1/2-1/2":
        return 0.5, 0.5

    winner_color = "white" if game.result == "1-0" else "black"
    if (winner_color == "white" and game.white == engine_a_name) or \
       (winner_color == "black" and game.black == engine_a_name):
        return 1.0, 0.0
    return 0.0, 1.0


def _play_game(
    white_name: str,
    white_worker: EngineWorker,
    black_name: str,
    black_worker: EngineWorker,
    move_timeout: float,
    max_moves: int = 200,
    verbose: int = 0,
) -> GameResult:
    """Play a single game using the supplied long-lived workers."""

    board = chess.Board()
    logger = get_error_logger()
    pgn_moves: list[str] = []
    move_count = 0

    for _ in range(max_moves):
        outcome = board.outcome()
        if outcome is not None:
            return _make_result(
                board, outcome, white_name, black_name, move_count, pgn_moves
            )

        move_count += 1
        current_name = white_name if board.turn == chess.WHITE else black_name
        current_worker = white_worker if board.turn == chess.WHITE else black_worker
        current_color = "white" if board.turn == chess.WHITE else "black"

        t0 = time.perf_counter() if verbose >= 1 else None
        move_uci, error_kind = current_worker.get_move(board.fen(), move_timeout)
        elapsed = time.perf_counter() - t0 if t0 is not None else None

        if move_uci is None:
            winner_name = black_name if current_color == "white" else white_name
            loser_name = current_name
            is_timeout = error_kind == "timeout"
            if verbose >= 1:
                kind = "TIMEOUT" if is_timeout else f"FAIL ({error_kind})"
                dur = f" after {elapsed:.3f}s" if elapsed is not None else ""
                print(f"        Move {move_count}: {current_name}: {kind}{dur}")
            if move_count == 1:
                result_str = "1/2-1/2"
                termination = "timeout (round draw)" if is_timeout else "engine error (round draw)"
            else:
                result_str = "0-1" if current_color == "white" else "1-0"
                termination = "timeout" if is_timeout else "engine error"
            logger.info(
                "%s: %s (%s) failed vs %s on move %d (FEN: %s)",
                "TIMEOUT" if is_timeout else "ENGINE_FAIL",
                loser_name, current_color, winner_name, move_count, board.fen()
            )
            return GameResult(
                white=white_name, black=black_name,
                result=result_str, termination=termination,
                num_moves=move_count, pgn=" ".join(pgn_moves),
            )

        try:
            move = chess.Move.from_uci(move_uci)
        except ValueError:
            winner_name = black_name if current_color == "white" else white_name
            result_str = "0-1" if current_color == "white" else "1-0"
            logger.info(
                "ILLEGAL_MOVE: %s played invalid UCI '%s' vs %s (FEN: %s)",
                current_name, move_uci, winner_name, board.fen()
            )
            if verbose >= 1:
                dur = f" after {elapsed:.3f}s" if elapsed is not None else ""
                print(f"        Move {move_count}: {current_name}: invalid UCI '{move_uci}'{dur}")
            return GameResult(
                white=white_name, black=black_name,
                result=result_str, termination="illegal move",
                num_moves=move_count, pgn=" ".join(pgn_moves),
            )

        if move not in board.legal_moves:
            winner_name = black_name if current_color == "white" else white_name
            result_str = "0-1" if current_color == "white" else "1-0"
            logger.info(
                "ILLEGAL_MOVE: %s played illegal move %s vs %s (FEN: %s)",
                current_name, move_uci, winner_name, board.fen()
            )
            if verbose >= 1:
                dur = f" after {elapsed:.3f}s" if elapsed is not None else ""
                print(f"        Move {move_count}: {current_name}: illegal move {move_uci}{dur}")
            return GameResult(
                white=white_name, black=black_name,
                result=result_str, termination="illegal move",
                num_moves=move_count, pgn=" ".join(pgn_moves),
            )

        if verbose >= 1:
            dur = f" in {elapsed:.3f}s" if elapsed is not None else ""
            print(f"        Move {move_count}: {current_name}: returned '{move_uci}'{dur}")

        san_move = board.san(move)
        board.push(move)
        pgn_moves.append(san_move)

        if verbose >= 2:
            print(f"      Move {move_count:3d}: {current_name} plays {san_move}")

    return GameResult(
        white=white_name, black=black_name,
        result="1/2-1/2", termination="max moves",
        num_moves=move_count, pgn=" ".join(pgn_moves),
    )


def _make_result(
    board: chess.Board,
    outcome: chess.Outcome,
    white_name: str,
    black_name: str,
    move_count: int,
    pgn_moves: list[str],
) -> GameResult:
    if outcome.winner is None:
        result = "1/2-1/2"
        termination_map = {
            chess.Termination.STALEMATE: "stalemate",
            chess.Termination.INSUFFICIENT_MATERIAL: "insufficient material",
            chess.Termination.FIFTY_MOVES: "50-move rule",
            chess.Termination.THREEFOLD_REPETITION: "threefold repetition",
        }
        term = termination_map.get(outcome.termination, "draw")
    elif outcome.winner == chess.WHITE:
        result = "1-0"
        term = "checkmate"
    else:
        result = "0-1"
        term = "checkmate"

    return GameResult(
        white=white_name, black=black_name,
        result=result, termination=term,
        num_moves=move_count, pgn=" ".join(pgn_moves),
    )


def _call_engine(
    engine_path: str,
    fen: str,
    timeout: float,
    engine_name: str,
    verbose: int = 0,
) -> tuple[Optional[str], Optional[str]]:
    """One-shot wrapper: create an ``EngineWorker``, request a single move,
    tear it down. Kept for the test suite and any external callers that
    expect the per-call API.
    """
    worker = EngineWorker(engine_path, engine_name)
    try:
        return worker.get_move(fen, timeout)
    finally:
        worker.close()