katrain-qt/katrain/core/contribute_engine.py

import json
import os
import random
import shlex
import shutil
import subprocess
import threading
import time
import traceback
from collections import defaultdict

from katrain.core.constants import OUTPUT_DEBUG, OUTPUT_ERROR, OUTPUT_INFO, OUTPUT_KATAGO_STDERR, DATA_FOLDER
from katrain.core.engine import EngineDiedException, BaseEngine
from katrain.core.game import BaseGame
from katrain.core.lang import i18n
from katrain.core.sgf_parser import Move
from katrain.core.utils import find_package_resource


class KataGoContributeEngine(BaseEngine):
    """Starts and communicates with the KataGo contribute program"""

    DEFAULT_MAX_GAMES = 8

    SHOW_RESULT_TIME = 5
    GIVE_UP_AFTER = 120

    def __init__(self, katrain):
        super().__init__(katrain, katrain.config("contribute"))
        self.katrain = katrain
        base_dir = os.path.expanduser("~/.katrain/katago_contribute")
        self.katago_process = None
        self.stdout_thread = None
        self.stderr_thread = None
        self.shell = False
        self.active_games = {}
        self.finished_games = set()
        self.showing_game = None
        self.last_advance = 0
        self.move_count = 0
        self.uploaded_games_count = 0
        self.last_move_for_game = defaultdict(int)
        self.visits_count = 0
        self.start_time = 0
        self.server_error = None

        self.save_sgf = self.config.get("savesgf", False)
        self.save_path = self.config.get("savepath", "./dist_sgf/")
        self.move_speed = self.config.get("movespeed", 2.0)

        exe = self.get_engine_path(self.config.get("katago"))
        cacert_path = os.path.join(os.path.split(exe)[0], "cacert.pem")
        if not os.path.isfile(cacert_path):
            try:
                shutil.copyfile(find_package_resource("katrain/KataGo/cacert.pem"), cacert_path)
            except Exception as e:
                self.katrain.log(
                    f"Could not copy cacert file ({e}), please add it manually to your katago.exe directory",
                    OUTPUT_ERROR,
                )
        cfg = find_package_resource(self.config.get("config"))

        settings_dict = {
            "username": self.config.get("username"),
            "password": self.config.get("password"),
            "maxSimultaneousGames": self.config.get("maxgames") or self.DEFAULT_MAX_GAMES,
            "includeOwnership": self.config.get("ownership") or False,
            "logGamesAsJson": True,
            "homeDataDir": os.path.expanduser(DATA_FOLDER),
        }
        self.max_buffer_games = 2 * settings_dict["maxSimultaneousGames"]
        settings = {f"{k}={v}" for k, v in settings_dict.items()}
        self.command = shlex.split(
            f'"{exe}" contribute -config "{cfg}" -base-dir "{base_dir}" -override-config "{",".join(settings)}"'
        )
        self.start()

    @staticmethod
    def game_ended(game):
        cn = game.current_node
        if cn.is_pass and cn.analysis_exists:
            moves = cn.candidate_moves
            if moves and moves[0]["move"] == "pass":
                game.play(Move(None, player=game.current_node.next_player))  # play pass
        return game.end_result

    def advance_showing_game(self):
        current_game = self.active_games.get(self.showing_game)
        if current_game:
            end_result = self.game_ended(current_game)
            if end_result is not None:
                self.finished_games.add(self.showing_game)
                if time.time() - self.last_advance > self.SHOW_RESULT_TIME:
                    del self.active_games[self.showing_game]
                    if self.save_sgf:
                        filename = os.path.join(self.save_path, f"{self.showing_game}.sgf")
                        self.katrain.log(current_game.write_sgf(filename, self.katrain.config("trainer")), OUTPUT_INFO)

                    self.katrain.log(f"Game {self.showing_game} finished, finding a new one", OUTPUT_INFO)
                    self.showing_game = None
            elif time.time() - self.last_advance > self.move_speed or len(self.active_games) > self.max_buffer_games:
                if current_game.current_node.children:
                    current_game.redo(1)
                    self.last_advance = time.time()
                    self.katrain("update-state")
                elif time.time() - self.last_advance > self.GIVE_UP_AFTER:
                    self.katrain.log(
                        f"Giving up on game {self.showing_game} which appears stuck, finding a new one", OUTPUT_INFO
                    )
                    self.showing_game = None
        else:
            if self.active_games:
                self.showing_game = None
                best_count = -1
                for game_id, game in self.active_games.items():  # find game with most moves left to show
                    count = 0
                    node = game.current_node
                    while node.children:
                        node = node.children[0]
                        count += 1
                    if count > best_count:
                        best_count = count
                        self.showing_game = game_id
                self.last_advance = time.time()
                self.katrain.log(f"Showing game {self.showing_game}, {best_count} moves left to show.", OUTPUT_INFO)

                self.katrain.game = self.active_games[self.showing_game]
                self.katrain("update-state", redraw_board=True)

    def status(self):
        return f"Contributing to distributed training\nGames: {self.uploaded_games_count} uploaded, {len(self.active_games)} in buffer, {len(self.finished_games)} shown\n{self.move_count} moves played ({60*self.move_count/(time.time()-self.start_time):.1f}/min, {self.visits_count / (time.time() - self.start_time):.1f} visits/s)\n"

    def is_idle(self):
        return False

    def queries_remaining(self):
        return 1

    def start(self):
        try:
            self.katrain.log(f"Starting Distributed KataGo with {self.command}", OUTPUT_INFO)
            startupinfo = None
            if hasattr(subprocess, "STARTUPINFO"):
                startupinfo = subprocess.STARTUPINFO()
                startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW  # stop command box popups on win/pyinstaller
            self.katago_process = subprocess.Popen(
                self.command,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                stdin=subprocess.PIPE,
                startupinfo=startupinfo,
                shell=self.shell,
            )
        except (FileNotFoundError, PermissionError, OSError) as e:
            self.katrain.log(
                i18n._("Starting Kata failed").format(command=self.command, error=e),
                OUTPUT_ERROR,
            )
            return  # don't start
        self.stdout_thread = threading.Thread(target=self._read_stdout_thread, daemon=True)
        self.stderr_thread = threading.Thread(target=self._read_stderr_thread, daemon=True)
        self.stdout_thread.start()
        self.stderr_thread.start()

    def check_alive(self, os_error="", exception_if_dead=False):
        ok = self.katago_process and self.katago_process.poll() is None
        if not ok and exception_if_dead:
            if self.katago_process:
                code = self.katago_process and self.katago_process.poll()
                if code == 3221225781:
                    died_msg = i18n._("Engine missing DLL")
                else:
                    os_error += f"status {code}"
                    died_msg = i18n._("Engine died unexpectedly").format(error=os_error)
                if code != 1 and not self.server_error:  # deliberate exit, already showed message?
                    self.katrain.log(died_msg, OUTPUT_ERROR)
                self.katago_process = None
            else:
                died_msg = i18n._("Engine died unexpectedly").format(error=os_error)
            if not self.server_error:  # dont raise if already know what happened
                raise EngineDiedException(died_msg)
        return ok

    def shutdown(self, finish=False):
        process = self.katago_process
        if process:
            self.katago_process = None
            process.terminate()
        if finish is not None:
            for t in [self.stderr_thread, self.stdout_thread]:
                if t:
                    t.join()

    def _read_stderr_thread(self):
        while self.katago_process is not None:
            try:
                line = self.katago_process.stderr.readline()
                if line:
                    try:
                        message = line.decode(errors="ignore").strip()
                        if any(
                            s in message
                            for s in ["not status code 200 OK", "Server returned error", "Uncaught exception:"]
                        ):
                            message = message.replace("what():", "").replace("Uncaught exception:", "").strip()
                            self.server_error = message  # don't be surprised by engine dying
                            self.katrain.log(message, OUTPUT_ERROR)
                            return
                        else:
                            self.katrain.log(message, OUTPUT_KATAGO_STDERR)
                    except Exception as e:
                        print("ERROR in processing KataGo stderr:", line, "Exception", e)
                elif self.katago_process:
                    self.check_alive(exception_if_dead=True)
            except Exception as e:
                self.katrain.log(f"Exception in reading stdout {e}", OUTPUT_DEBUG)
                return

    def _read_stdout_thread(self):
        while self.katago_process is not None:
            try:
                line = self.katago_process.stdout.readline()
                if line:
                    line = line.decode(errors="ignore").strip()
                    if line.startswith("{"):
                        try:
                            analysis = json.loads(line)
                            if "gameId" in analysis:
                                game_id = analysis["gameId"]
                                if game_id in self.finished_games:
                                    continue
                                current_game = self.active_games.get(game_id)
                                new_game = current_game is None
                                if new_game:
                                    board_size = [analysis["boardXSize"], analysis["boardYSize"]]
                                    placements = {
                                        f"A{bw}": [
                                            Move.from_gtp(move, pl).sgf(board_size)
                                            for pl, move in analysis["initialStones"]
                                            if pl == bw
                                        ]
                                        for bw in "BW"
                                    }
                                    game_properties = {k: v for k, v in placements.items() if v}
                                    game_properties["SZ"] = f"{board_size[0]}:{board_size[1]}"
                                    game_properties["KM"] = analysis["rules"]["komi"]
                                    game_properties["RU"] = json.dumps(analysis["rules"])
                                    game_properties["PB"] = analysis["blackPlayer"]
                                    game_properties["PW"] = analysis["whitePlayer"]
                                    current_game = BaseGame(self.katrain, game_properties=game_properties)
                                    self.active_games[game_id] = current_game
                                last_node = current_game.sync_branch(
                                    [Move.from_gtp(coord, pl) for pl, coord in analysis["moves"]]
                                )
                                last_node.set_analysis(analysis)
                                if new_game:
                                    current_game.set_current_node(last_node)
                                self.start_time = self.start_time or time.time() - 1
                                self.move_count += 1
                                self.visits_count += analysis["rootInfo"]["visits"]
                                last_move = self.last_move_for_game[game_id]
                                self.last_move_for_game[game_id] = time.time()
                                dt = self.last_move_for_game[game_id] - last_move if last_move else 0
                                self.katrain.log(
                                    f"[{time.time()-self.start_time:.1f}] Game {game_id} Move {analysis['turnNumber']}: {' '.join(analysis['move'])} Visits {analysis['rootInfo']['visits']} Time {dt:.1f}s\t Moves/min {60*self.move_count/(time.time()-self.start_time):.1f} Visits/s {self.visits_count/(time.time()-self.start_time):.1f}",
                                    OUTPUT_DEBUG,
                                )
                                self.katrain("update-state")
                        except Exception as e:
                            traceback.print_exc()
                            self.katrain.log(f"Exception {e} in parsing or processing JSON: {line}", OUTPUT_ERROR)
                    elif "uploaded sgf" in line:
                        self.uploaded_games_count += 1
                    else:
                        self.katrain.log(line, OUTPUT_KATAGO_STDERR)
                elif self.katago_process:
                    self.check_alive(exception_if_dead=False)  # stderr will do this
            except Exception as e:
                self.katrain.log(f"Exception in reading stdout {e}", OUTPUT_DEBUG)
                return