ad244ffc7e
Co-authored-by: Artur Tamazian <artur-tamazian> Co-authored-by: Javier Gonzalez <javierggt@yahoo.com>
492 lines
21 KiB
Python
492 lines
21 KiB
Python
import copy
|
|
import json
|
|
import os
|
|
import platform
|
|
import queue
|
|
import shlex
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
import traceback
|
|
from typing import Callable, Dict, List, Optional
|
|
|
|
from kivy.utils import platform as kivy_platform
|
|
|
|
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.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 kivy_platform == "win":
|
|
exe = "katrain/KataGo/katago.exe"
|
|
elif kivy_platform == "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:
|
|
self.katrain("engine_recovery_popup", 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()
|
|
if getattr(self.katrain, "update_state", None): # easier mocking etc
|
|
self.katrain.update_state()
|
|
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)
|