예제 소스

mariago_ws_bridge.py

#!/usr/bin/env python3
"""MariaGO 웹 착수 브리지.

지금은 모든 착수를 로컬 휴리스틱으로 둔다. Gemma·R5 포트는 학습 대국이 쓴다.

  .venv/bin/python scripts/mariago_ws_bridge.py --vlm-port 9319 --ws-port 9320

프로토콜 (JSON):
  → {"type":"move","board":[361 ints],"ai_side":1|2,"game_epoch":int,"rejects":[[r,c],...]}
  ← {"type":"move","row":int,"col":int,"pass":false}
  ← {"type":"move","pass":true}
  ← {"type":"error","error":"..."}
  → {"type":"ping"} ← {"type":"pong"}
"""

from __future__ import annotations

import argparse
import asyncio
import importlib.util
import io
import json
import os
import random
import socket
import struct
import sys
from pathlib import Path
from typing import List, Optional, Tuple

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from vlm_server.dataset_baduk import BOARD_SIZE, render_board_png  # noqa: E402
from vlm_server.probe_move import MAGIC_GO19, MAGIC_RPLY, _recv_exact, _u32be  # noqa: E402

try:
    import websockets
except ImportError:
    print("websockets 패키지 필요: .venv/bin/pip install websockets", file=sys.stderr)
    raise SystemExit(1)

# websockets 12+ / 14+ API 모두 대응
try:
    from websockets.asyncio.server import serve as ws_serve
except ImportError:
    from websockets import serve as ws_serve  # type: ignore


def _empty_png() -> bytes:
    img = render_board_png(bytes(BOARD_SIZE * BOARD_SIZE), BOARD_SIZE, px=256)
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return buf.getvalue()


def _send_go19_ext(
    sock: socket.socket,
    png: bytes,
    board: bytes,
    ai_side: int,
    *,
    game_epoch: int,
    rejects: List[Tuple[int, int]],
) -> None:
    body = bytearray()
    body.append(BOARD_SIZE)
    body.append(ai_side)
    body += _u32be(len(png))
    body += png
    body += board
    body += _u32be(int(game_epoch) & 0xFFFFFFFF)
    if rejects:
        n = min(64, len(rejects))
        body.append(n & 0xFF)
        for i in range(n):
            r, c = rejects[i]
            body.append(int(r) & 0xFF)
            body.append(int(c) & 0xFF)
    sock.sendall(MAGIC_GO19 + _u32be(len(body)) + body)


def _recv_reply(sock: socket.socket) -> Tuple[bool, Optional[int], Optional[int], Optional[str]]:
    head = _recv_exact(sock, 4)
    if head != MAGIC_RPLY:
        raise ValueError(f"bad reply header: {head!r}")
    flag = _recv_exact(sock, 1)[0]
    if flag == 1:
        row = _recv_exact(sock, 1)[0]
        col = _recv_exact(sock, 1)[0]
        if row == 0xFF and col == 0xFF:
            return True, None, None, None
        return True, int(row), int(col), None
    if flag == 0:
        (n,) = struct.unpack(">H", _recv_exact(sock, 2))
        msg = _recv_exact(sock, n).decode("utf-8", errors="replace")
        return False, None, None, msg
    raise ValueError(f"unknown flag: {flag}")


def request_move_sync(
    host: str,
    port: int,
    board: List[int],
    ai_side: int,
    game_epoch: int,
    rejects: List[Tuple[int, int]],
    timeout_s: float,
) -> dict:
    if len(board) != BOARD_SIZE * BOARD_SIZE:
        return {"type": "error", "error": f"board length must be {BOARD_SIZE*BOARD_SIZE}"}
    packed = bytes(int(x) & 0xFF for x in board)
    sock = socket.create_connection((host, port), timeout=15)
    sock.settimeout(timeout_s)
    try:
        _send_go19_ext(
            sock,
            _empty_png(),
            packed,
            ai_side,
            game_epoch=game_epoch,
            rejects=rejects,
        )
        ok, row, col, err = _recv_reply(sock)
        if not ok:
            return {"type": "error", "error": err or "vlm error"}
        if row is None or col is None:
            return {"type": "move", "pass": True}
        return {"type": "move", "pass": False, "row": row, "col": col}
    finally:
        try:
            sock.close()
        except Exception:
            pass


_monitor = None


def _monitor_mod():
    global _monitor
    if _monitor is None:
        path = ROOT / "scripts" / "monitor_mariago_web_play.py"
        spec = importlib.util.spec_from_file_location("monitor_mariago_web_play", path)
        if spec is None or spec.loader is None:
            raise RuntimeError("heuristic import failed")
        mod = importlib.util.module_from_spec(spec)
        sys.modules["monitor_mariago_web_play"] = mod
        spec.loader.exec_module(mod)
        _monitor = mod
    return _monitor


def local_ai_move(board: List[int], ai_side: int, rejects: List[Tuple[int, int]], ko: Optional[Tuple[int, int]]) -> dict:
    """이 요청만 로컬 휴리스틱으로 둔다. 서버 엔진은 바꾸지 않는다."""
    mon = _monitor_mod()
    banned = {(int(r), int(c)) for r, c in rejects}
    row, col, is_pass = mon.heuristic_move(list(board), int(ai_side), ko, random.Random(), banned)
    if is_pass or row is None or col is None:
        return {"type": "move", "pass": True}
    ok, reason, _nxt, _cap, _ko = mon.try_play(list(board), int(row), int(col), int(ai_side), ko)
    if not ok:
        return {"type": "error", "error": reason or "illegal"}
    return {"type": "move", "pass": False, "row": int(row), "col": int(col)}


async def handler(ws, *, vlm_host: str, vlm_port: int, timeout_s: float):
    async for raw in ws:
        try:
            msg = json.loads(raw)
        except json.JSONDecodeError:
            await ws.send(json.dumps({"type": "error", "error": "invalid json"}))
            continue
        mtype = msg.get("type")
        if mtype == "ping":
            await ws.send(json.dumps({"type": "pong"}))
            continue
        if mtype != "move":
            await ws.send(json.dumps({"type": "error", "error": f"unknown type: {mtype}"}))
            continue
        board = msg.get("board") or []
        ai_side = int(msg.get("ai_side") or 2)
        rejects_raw = msg.get("rejects") or []
        rejects: List[Tuple[int, int]] = []
        for item in rejects_raw:
            if isinstance(item, (list, tuple)) and len(item) >= 2:
                rejects.append((int(item[0]), int(item[1])))
        ko_raw = msg.get("ko")
        ko = None
        if isinstance(ko_raw, (list, tuple)) and len(ko_raw) >= 2:
            ko = (int(ko_raw[0]), int(ko_raw[1]))
        try:
            result = await asyncio.to_thread(local_ai_move, list(board), ai_side, rejects, ko)
        except Exception as e:
            result = {"type": "error", "error": str(e)}
        await ws.send(json.dumps(result, ensure_ascii=False))


async def main_async(args: argparse.Namespace) -> None:
    print(
        f"[mariago-bridge] ws://********:{args.ws_port} → heuristic only",
        flush=True,
    )

    async def _handler(ws):
        await handler(
            ws,
            vlm_host=args.vlm_host,
            vlm_port=args.vlm_port,
            timeout_s=args.timeout_s,
        )

    async with ws_serve(_handler, "********", args.ws_port):
        await asyncio.Future()


MOVE_PORT_FILE = ROOT / "config/mariago_move_port"


def default_move_port() -> int:
    """착수 서버 포트: MARIAGO_MOVE_PORT > config/mariago_move_port > 9319.

    앙상블 프록시(:9318)는 scripts/install_mariago_ensemble.sh on 이 이 파일을 쓴다.
    """
    env = os.environ.get("MARIAGO_MOVE_PORT", "").strip()
    if env:
        return int(env)
    try:
        return int(MOVE_PORT_FILE.read_text().strip())
    except (OSError, ValueError):
        return 9319


def main() -> int:
    p = argparse.ArgumentParser(description="MariaGO WebSocket → GO19 bridge")
    p.add_argument("--vlm-host", default="127.0.0.1")
    p.add_argument("--vlm-port", type=int, default=default_move_port())
    p.add_argument("--ws-port", type=int, default=9320)
    p.add_argument("--timeout-s", type=float, default=240.0)
    args = p.parse_args()
    asyncio.run(main_async(args))
    return 0


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