import copy import json import os import platform import queue import shlex import subprocess import threading import time import traceback import sys from typing import Callable, Dict, List, Optional from katrain.core.constants import ( OUTPUT_DEBUG, OUTPUT_ERROR, OUTPUT_EXTRA_DEBUG, OUTPUT_KATAGO_STDERR, DATA_FOLDER, KATAGO_EXCEPTION, PONDERING_REPORT_DT, ) from katrain.core.events import event_sink from katrain.core.game_node import GameNode from katrain.core.lang import i18n from katrain.core.sgf_parser import Move from katrain.core.utils import find_package_resource, json_truncate_arrays ENGINE_BACKENDS = ("local", "remote", "custom") def resolve_engine_backend(config) -> str: """Which engine backend to use: 'local', 'remote' or 'custom'. Uses the explicit `engine/backend` setting when present (set by the engine settings tabs). Falls back to deriving it from the legacy fields so older configs keep working: a remote_url means remote, an altcommand means custom. """ backend = (config.get("backend") or "").strip().lower() if backend in ENGINE_BACKENDS: return backend if (config.get("remote_url") or "").strip(): return "remote" if (config.get("altcommand") or "").strip(): return "custom" return "local" class BaseEngine: # some common elements between analysis and contribute engine RULESETS_ABBR = [ ("jp", "japanese"), ("cn", "chinese"), ("ko", "korean"), ("aga", "aga"), ("tt", "tromp-taylor"), ("nz", "new zealand"), ("stone_scoring", "stone_scoring"), ] RULESETS = {fromkey: name for abbr, name in RULESETS_ABBR for fromkey in [abbr, name]} def __init__(self, katrain, config): self.katrain = katrain self.config = config @staticmethod def get_rules(ruleset): if ruleset.strip().startswith("{"): try: ruleset = json.loads(ruleset) except json.JSONDecodeError: pass if isinstance(ruleset, dict): return ruleset return KataGoEngine.RULESETS.get(str(ruleset).lower(), "japanese") def advance_showing_game(self): pass # avoid transitional error def status(self): return "" # avoid transitional error def get_engine_path(self, exe): if not exe: if os.name == "nt": exe = "katrain/KataGo/katago.exe" elif sys.platform.startswith("linux"): exe = "katrain/KataGo/katago" else: exe = find_package_resource("katrain/KataGo/katago-osx") # github actions built if not os.path.isfile(exe) or "arm64" in platform.version().lower(): exe = "katago" # e.g. MacOS after brewing if exe.startswith("katrain"): exe = find_package_resource(exe) exepath, exename = os.path.split(exe) if exepath and not os.path.isfile(exe): self.on_error(i18n._("Kata exe not found").format(exe=exe), "KATAGO-EXE") return None elif not exepath: paths = os.getenv("PATH", ".").split(os.pathsep) + ["/opt/homebrew/bin/"] exe_with_paths = [os.path.join(path, exe) for path in paths if os.path.isfile(os.path.join(path, exe))] if not exe_with_paths: self.on_error(i18n._("Kata exe not found in path").format(exe=exe), "KATAGO-EXE") return None exe = exe_with_paths[0] return exe def on_error(self, message, code, allow_popup): print("ERROR", message, code) class KataGoEngine(BaseEngine): """Starts and communicates with the KataGO analysis engine""" PONDER_KEY = "_kt_continuous" ENGINE_TYPE = "local" # passed to the recovery popup so it can tailor its advice def __init__(self, katrain, config): super().__init__(katrain, config) self.allow_recovery = self.config.get("allow_recovery", True) # if false, don't give popups self.queries = {} # outstanding query id -> start time and callback self.ponder_query = None self.query_counter = 0 self.katago_process = None self.base_priority = 0 self.override_settings = {"reportAnalysisWinratesAs": "BLACK"} # force these settings self.analysis_thread = None self.stderr_thread = None self.write_stdin_thread = None self.shell = False self.write_queue = queue.Queue() self.thread_lock = threading.Lock() if resolve_engine_backend(config) == "custom": self.command = config["altcommand"] self.shell = True else: model = find_package_resource(config["model"]) cfg = find_package_resource(config["config"]) exe = self.get_engine_path(config.get("katago", "").strip()) if not exe: return # Add human model to command if provided if config.get("humanlike_model", ""): human_model_path = find_package_resource(config.get("humanlike_model","")) if os.path.isfile(human_model_path): self.command = shlex.split( f'"{exe}" analysis -model "{model}" -human-model "{human_model_path}" -config "{cfg}" -override-config "homeDataDir={os.path.expanduser(DATA_FOLDER)}"' ) else: self.katrain.log(f"Human model not found at {human_model_path}", -1) # Fall back to regular command without human model self.command = shlex.split( f'"{exe}" analysis -model "{model}" -config "{cfg}" -override-config "homeDataDir={os.path.expanduser(DATA_FOLDER)}"' ) else: # Regular command without human model self.command = shlex.split( f'"{exe}" analysis -model "{model}" -config "{cfg}" -override-config "homeDataDir={os.path.expanduser(DATA_FOLDER)}"' ) self.start() def on_error(self, message, code=None, allow_popup=True): self.katrain.log(message, OUTPUT_ERROR) if self.allow_recovery and allow_popup: event_sink(self.katrain).engine_failure(message, code, self.ENGINE_TYPE) def start(self): with self.thread_lock: self.write_queue = queue.Queue() try: self.katrain.log(f"Starting KataGo with {self.command}", OUTPUT_DEBUG) 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, startupinfo=startupinfo, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=self.shell, ) except (FileNotFoundError, PermissionError, OSError) as e: self.on_error(i18n._("Starting Kata failed").format(command=self.command, error=e), code="c") return # don't start self.analysis_thread = threading.Thread(target=self._analysis_read_thread, daemon=True) self.stderr_thread = threading.Thread(target=self._read_stderr_thread, daemon=True) self.write_stdin_thread = threading.Thread(target=self._write_stdin_thread, daemon=True) self.analysis_thread.start() self.stderr_thread.start() self.write_stdin_thread.start() def on_new_game(self): self.base_priority += 1 if not self.is_idle(): with self.thread_lock: self.write_queue = queue.Queue() self.terminate_queries(only_for_node=None, lock=False) self.ponder_query = None self.queries = {} def terminate_queries(self, only_for_node=None, lock=True): if lock: with self.thread_lock: return self.terminate_queries(only_for_node=only_for_node, lock=False) for query_id, (_, _, _, _, node) in list(self.queries.items()): if only_for_node is None or only_for_node is node: self.terminate_query(query_id) def stop_pondering(self): pq = self.ponder_query if pq: self.terminate_query(pq["id"], ignore_further_results=False) self.ponder_query = None def terminate_query(self, query_id, ignore_further_results=True): self.katrain.log(f"Terminating query {query_id}", OUTPUT_DEBUG) if query_id is not None: self.send_query({"action": "terminate", "terminateId": query_id}, None, None) if ignore_further_results: self.queries.pop(query_id, None) def restart(self): self.queries = {} self.shutdown(finish=False) self.start() def check_alive(self, os_error="", exception_if_dead=False, maybe_open_recovery=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: died_msg = i18n._("Engine died unexpectedly").format(error=f"{os_error} status {code}") if code != 1: # deliberate exit self.on_error(died_msg, code, allow_popup=maybe_open_recovery) self.katago_process = None # return from threads else: self.katrain.log(i18n._("Engine died unexpectedly").format(error=os_error), OUTPUT_DEBUG) return ok def wait_to_finish(self): while self.queries and self.katago_process and self.katago_process.poll() is None: time.sleep(0.1) def shutdown(self, finish=False): process = self.katago_process if finish and process: self.wait_to_finish() if process: self.katago_process = None self.katrain.log("Terminating KataGo process", OUTPUT_DEBUG) process.terminate() self.katrain.log("Terminated KataGo process", OUTPUT_DEBUG) if finish is not None: # don't care if exiting app for t in [self.write_stdin_thread, self.analysis_thread, self.stderr_thread]: if t: t.join() def is_idle(self): return not self.queries and self.write_queue.empty() def queries_remaining(self): return len(self.queries) + int(not self.write_queue.empty()) def _read_stderr_thread(self): while self.katago_process is not None: try: line = self.katago_process.stderr.readline() if line: if b"Uncaught exception" in line or b"what()" in line: # linux=what msg = f"KataGo Engine Failed: {line.decode(errors='ignore')[9:].strip()}" self.on_error(msg, KATAGO_EXCEPTION) return try: self.katrain.log(line.decode(errors="ignore").strip(), OUTPUT_KATAGO_STDERR) except Exception as e: print("ERROR in processing KataGo stderr:", line, "Exception", e) elif not self.check_alive(exception_if_dead=True): return except Exception as e: self.katrain.log(f"Exception in reading stderr: {e}", OUTPUT_DEBUG) return def _analysis_read_thread(self): while self.katago_process is not None: try: line = self.katago_process.stdout.readline().strip() if self.katago_process and not line: if not self.check_alive(exception_if_dead=True, maybe_open_recovery=True): return except OSError as e: self.check_alive(os_error=str(e), exception_if_dead=True, maybe_open_recovery=True) return if b"Uncaught exception" in line: msg = f"KataGo Engine Failed: {line.decode(errors='ignore')}" self.on_error(msg, KATAGO_EXCEPTION) return if not line: continue try: analysis = json.loads(line) if "id" not in analysis: self.katrain.log(f"Error without ID {analysis} received from KataGo", OUTPUT_ERROR) continue query_id = analysis["id"] if query_id not in self.queries: if analysis.get("action") != "terminate": self.katrain.log( f"Query result {query_id} discarded -- recent new game or node reset?", OUTPUT_DEBUG ) continue callback, error_callback, start_time, next_move, _ = self.queries[query_id] if "error" in analysis: del self.queries[query_id] if error_callback: error_callback(analysis) elif not (next_move and "Illegal move" in analysis["error"]): # sweep self.katrain.log(f"{analysis} received from KataGo", OUTPUT_ERROR) elif "warning" in analysis: self.katrain.log(f"{analysis} received from KataGo", OUTPUT_DEBUG) elif "terminateId" in analysis: self.katrain.log(f"{analysis} received from KataGo", OUTPUT_DEBUG) else: partial_result = analysis.get("isDuringSearch", False) if not partial_result: del self.queries[query_id] time_taken = time.time() - start_time results_exist = not analysis.get("noResults", False) self.katrain.log( f"[{time_taken:.1f}][{query_id}][{'....' if partial_result else 'done'}] KataGo analysis received: {len(analysis.get('moveInfos',[]))} candidate moves, {analysis['rootInfo']['visits'] if results_exist else 'n/a'} visits", OUTPUT_DEBUG, ) self.katrain.log(json_truncate_arrays(analysis), OUTPUT_EXTRA_DEBUG) try: if callback and results_exist: callback(analysis, partial_result) except Exception as e: self.katrain.log(f"Error in engine callback for query {query_id}: {e}", OUTPUT_ERROR) traceback.print_exc() event_sink(self.katrain).state_changed() except Exception as e: self.katrain.log(f"Unexpected exception {e} while processing KataGo output {line}", OUTPUT_ERROR) traceback.print_exc() def _write_stdin_thread(self): # flush only in a thread since it returns only when the other program reads while self.katago_process is not None: try: query, callback, error_callback, next_move, node = self.write_queue.get(block=True, timeout=0.1) except queue.Empty: continue with self.thread_lock: if "id" not in query: self.query_counter += 1 query["id"] = f"QUERY:{str(self.query_counter)}" ponder = query.pop(self.PONDER_KEY, False) if ponder: # handle pondering in here to be in lock and such pq = self.ponder_query or {} # basically we handle pondering by just asking for these queries a lot and ignoring duplicates # when a different ponder query comes in, e.g. due to selecting a roi or different node, switch differences = { k: (pq.get(k), query.get(k)) for k in (query.keys() | pq.keys()) - {"id", "maxVisits", "reportDuringSearchEvery"} if pq.get(k) != query.get(k) } if differences: self.stop_pondering() query["maxVisits"] = 10_000_000 query["reportDuringSearchEvery"] = PONDERING_REPORT_DT self.ponder_query = query else: continue terminate = query.get("action") == "terminate" if not terminate: self.queries[query["id"]] = (callback, error_callback, time.time(), next_move, node) tag = "ponder " if ponder else ("terminate " if terminate else "") self.katrain.log(f"Sending {tag}query {query['id']}: {json.dumps(query)}", OUTPUT_DEBUG) try: self.katago_process.stdin.write((json.dumps(query) + "\n").encode()) self.katago_process.stdin.flush() except OSError as e: self.katrain.log(f"Exception in writing to katago: {e}", OUTPUT_DEBUG) return # some other thread will take care of this def send_query(self, query, callback, error_callback, next_move=None, node=None): self.write_queue.put((query, callback, error_callback, next_move, node)) def request_analysis( self, analysis_node: GameNode, callback: Callable, error_callback: Optional[Callable] = None, visits: int = None, analyze_fast: bool = False, time_limit=True, find_alternatives: bool = False, region_of_interest: Optional[List] = None, priority: int = 0, ponder=False, # infinite visits, cancellable ownership: Optional[bool] = None, next_move: Optional[GameNode] = None, extra_settings: Optional[Dict] = None, include_policy=True, report_every: Optional[float] = None, ): nodes = analysis_node.nodes_from_root moves = [m for node in nodes for m in node.moves] initial_stones = [m for node in nodes for m in node.placements] clear_placements = [m for node in nodes for m in node.clear_placements] if clear_placements: # TODO: support these self.katrain.log(f"Not analyzing node {analysis_node} as there are AE commands in the path", OUTPUT_DEBUG) return if next_move: moves.append(next_move) if ownership is None: ownership = self.config["_enable_ownership"] and not next_move if visits is None: visits = self.config["max_visits"] if analyze_fast and self.config.get("fast_visits"): visits = self.config["fast_visits"] size_x, size_y = analysis_node.board_size if find_alternatives: avoid = [ { "moves": list(analysis_node.analysis["moves"].keys()), "player": analysis_node.next_player, "untilDepth": 1, } ] elif region_of_interest: xmin, xmax, ymin, ymax = region_of_interest avoid = [ { "moves": [ Move((x, y)).gtp() for x in range(0, size_x) for y in range(0, size_y) if x < xmin or x > xmax or y < ymin or y > ymax ], "player": player, "untilDepth": 1, # tried a large number here, or 2, but this seems more natural } for player in "BW" ] else: avoid = [] settings = copy.copy(self.override_settings) settings["wideRootNoise"] = self.config["wide_root_noise"] if time_limit: settings["maxTime"] = self.config["max_time"] query = { "rules": self.get_rules(analysis_node.ruleset), "priority": self.base_priority + priority, "analyzeTurns": [len(moves)], "maxVisits": visits, "komi": analysis_node.komi, "boardXSize": size_x, "boardYSize": size_y, "includeOwnership": ownership and not next_move, "includeMovesOwnership": ownership and not next_move, "includePolicy": include_policy, "initialStones": [[m.player, m.gtp()] for m in initial_stones], "initialPlayer": analysis_node.initial_player, "moves": [[m.player, m.gtp()] for m in moves], "overrideSettings": {**settings, **(extra_settings or {})}, self.PONDER_KEY: ponder, } if report_every is not None: query["reportDuringSearchEvery"] = report_every if avoid: query["avoidMoves"] = avoid self.send_query(query, callback, error_callback, next_move, analysis_node) analysis_node.analysis_visits_requested = max(analysis_node.analysis_visits_requested, visits)