예제 소스

katago_analyze.py

#!/usr/bin/env python3
"""MariaGO board → KataGo analysis (score lead, winrate, ownership)."""

from __future__ import annotations

import argparse
import json
import os
import shutil
import subprocess
import sys
import threading
import time
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parents[1]
LETTERS = "ABCDEFGHJKLMNOPQRST"


def _brew_katago(prefix: Path) -> tuple[Path, Path, Path] | None:
    model = prefix / "share/katago/kata1-b18c384nbt-s9996604416-d4316597426.bin.gz"
    cfg = prefix / "share/katago/configs/analysis_example.cfg"
    exe = prefix / "bin/katago"
    if model.exists() and cfg.exists() and exe.exists():
        return exe, model, cfg
    return None


def katago_paths() -> tuple[Path, Path, Path]:
    # `brew --prefix` 는 API 데이터 갱신으로 수 초 걸릴 수 있어 흔한 경로를 먼저 본다.
    for prefix in (Path("/opt/homebrew/opt/katago"), Path("/usr/local/opt/katago")):
        found = _brew_katago(prefix)
        if found:
            return found
    brew = shutil.which("brew")
    if brew:
        try:
            prefix = Path(
                subprocess.check_output([brew, "--prefix", "katago"], text=True).strip()
            )
            found = _brew_katago(prefix)
            if found:
                return found
        except subprocess.CalledProcessError:
            pass
    local = ROOT / "tools/katago"
    return local / "katago", local / "model.bin.gz", local / "analysis.cfg"


def move_to_katago(m: dict[str, Any]) -> list[str]:
    color = "B" if int(m.get("color", 1)) == 1 else "W"
    if m.get("pass"):
        return [color, "pass"]
    r, c = int(m["r"]), int(m["c"])
    return [color, f"{LETTERS[c]}{19 - r}"]


def analyze(
    moves: list[dict[str, Any]],
    *,
    komi: float = 6.5,
    max_visits: int = 80,
    timeout_s: float = 45.0,
) -> dict[str, Any]:
    exe, model, cfg = katago_paths()
    if not exe.exists():
        raise FileNotFoundError(f"KataGo not found: {exe}")
    if not model.exists():
        raise FileNotFoundError(f"KataGo model not found: {model}")
    if not cfg.exists():
        raise FileNotFoundError(f"KataGo config not found: {cfg}")

    katago_moves = [move_to_katago(m) for m in moves]
    turn = len(katago_moves)
    req = {
        "id": "mariago",
        "initialStones": [],
        "moves": katago_moves,
        "rules": "korean",
        "komi": komi,
        "boardXSize": 19,
        "boardYSize": 19,
        "analyzeTurns": [turn],
        "maxVisits": max_visits,
        "includeOwnership": True,
        "includePolicy": False,
    }

    log_dir = ROOT / "logs/katago_analysis"
    log_dir.mkdir(parents=True, exist_ok=True)
    env = os.environ.copy()
    env["HOME"] = str(log_dir)

    proc = subprocess.Popen(
        [str(exe), "analysis", "-config", str(cfg), "-model", str(model)],
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        bufsize=1,
        env=env,
    )
    assert proc.stdin and proc.stdout

    def feed() -> None:
        time.sleep(0.5)
        proc.stdin.write(json.dumps(req) + "\n")
        proc.stdin.flush()

    threading.Thread(target=feed, daemon=True).start()

    deadline = time.monotonic() + timeout_s
    last: dict[str, Any] | None = None
    while time.monotonic() < deadline:
        line = proc.stdout.readline()
        if not line:
            if proc.poll() is not None:
                break
            continue
        line = line.strip()
        if not line.startswith("{"):
            continue
        try:
            msg = json.loads(line)
        except json.JSONDecodeError:
            continue
        if msg.get("id") != "mariago":
            continue
        if msg.get("error"):
            proc.terminate()
            raise RuntimeError(msg["error"])
        last = msg
        if msg.get("isDuringSearch") is False:
            break

    proc.terminate()
    try:
        proc.wait(timeout=3)
    except subprocess.TimeoutExpired:
        proc.kill()

    if not last or "rootInfo" not in last:
        err = proc.stderr.read() if proc.stderr else ""
        raise RuntimeError(f"KataGo analysis failed. stderr tail: {err[-500:]}")

    root = last["rootInfo"]
    ownership = last.get("ownership") or []
    to_play = "B" if turn % 2 == 0 else "W"
    score_lead_black = float(root.get("scoreLead", 0))
    winrate_black = float(root.get("winrate", 0.5))

    black_area = sum(1 for x in ownership if x < -0.35)
    white_area = sum(1 for x in ownership if x > 0.35)
    contested = sum(1 for x in ownership if -0.35 <= x <= 0.35)

    if to_play == "W":
        score_lead_side = -score_lead_black
    else:
        score_lead_side = score_lead_black

    return {
        "turn": turn,
        "toPlay": to_play,
        "scoreLeadBlack": round(score_lead_black, 2),
        "scoreLeadSide": round(score_lead_side, 2),
        "winrateBlack": round(winrate_black, 4),
        "winrateWhite": round(1 - winrate_black, 4),
        "ownership": ownership,
        "summary": {
            "blackLikely": black_area,
            "whiteLikely": white_area,
            "contested": contested,
        },
        "rootInfo": root,
    }


def main() -> int:
    p = argparse.ArgumentParser()
    p.add_argument("--komi", type=float, default=6.5)
    p.add_argument("--max-visits", type=int, default=80)
    p.add_argument("--timeout-s", type=float, default=45.0)
    args = p.parse_args()
    raw = sys.stdin.read()
    body = json.loads(raw) if raw.strip() else {"moves": []}
    moves = body.get("moves") or []
    out = analyze(
        moves,
        komi=float(body.get("komi", args.komi)),
        max_visits=int(body.get("maxVisits", args.max_visits)),
        timeout_s=args.timeout_s,
    )
    print(json.dumps(out, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())