예제 소스

text_client.py

"""텍스트 VLM 클라이언트 — GO19 로 VLM vs VLM (또는 self-play 단일 서버).

예:
  python -m vlm_server.text_client --host 127.0.0.1 --port 9319 --gamess 3

파이프라인:
  python -m scripts.daemon_text_client_harden --foreground
"""

from __future__ import annotations

import argparse
import io
import json
import socket
import struct
import time
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from rl_env.board import BLACK, BOARD_SIZE, WHITE

from .dataset_baduk import render_board_png
from .probe_move import MAGIC_GO19, MAGIC_RPLY, _recv_exact, _u32be
from .text_client_core import (
    ClientCore,
    CoreConfig,
    GameEndReason,
    RejectReason,
)


def _send_go19_ext(
    sock: socket.socket,
    png: bytes,
    board: bytes,
    ai_side: int,
    *,
    game_epoch: int = 1,
    rejects: Optional[List[Tuple[int, int]]] = None,
) -> 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}")


@dataclass
class Endpoint:
    host: str
    port: int
    label: str = ""


@dataclass
class MoveEvent:
    ply: int
    side: str
    row: Optional[int]
    col: Optional[int]
    is_pass: bool
    accepted: bool
    reject: Optional[str] = None
    ms: float = 0.0
    retries: int = 0


@dataclass
class GameResult:
    game_id: int
    end_reason: str
    error: Optional[str]
    moves: List[MoveEvent] = field(default_factory=list)
    plies: int = 0
    started_at: str = ""
    finished_at: str = ""
    black_endpoint: str = ""
    white_endpoint: str = ""


class TextVlmClient:
    def __init__(self, ep: Endpoint, timeout_s: float = 180.0):
        self.ep = ep
        self.timeout_s = timeout_s
        self._sock: Optional[socket.socket] = None
        self._png = self._empty_png()

    @staticmethod
    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 close(self) -> None:
        if self._sock is not None:
            try:
                self._sock.close()
            except Exception:
                pass
            self._sock = None

    def _ensure(self) -> socket.socket:
        if self._sock is not None:
            return self._sock
        s = socket.create_connection((self.ep.host, self.ep.port), timeout=15)
        s.settimeout(self.timeout_s)
        self._sock = s
        return s

    def request_move(
        self,
        board: bytes,
        ai_side: int,
        *,
        game_epoch: int,
        rejects: List[Tuple[int, int]],
    ) -> Tuple[bool, Optional[int], Optional[int], Optional[str], float]:
        t0 = time.time()
        try:
            sock = self._ensure()
            _send_go19_ext(
                sock,
                self._png,
                board,
                ai_side,
                game_epoch=game_epoch,
                rejects=rejects,
            )
            ok, row, col, err = _recv_reply(sock)
            return ok, row, col, err, (time.time() - t0) * 1000.0
        except Exception as e:
            self.close()
            return False, None, None, str(e), (time.time() - t0) * 1000.0


def play_one_game(
    *,
    game_id: int,
    black: TextVlmClient,
    white: TextVlmClient,
    max_moves: int = 400,
    core_cfg: Optional[CoreConfig] = None,
) -> GameResult:
    cfg = core_cfg or CoreConfig()
    core = ClientCore(cfg=cfg, game_epoch=game_id)
    started = datetime.now(timezone.utc).isoformat()
    events: List[MoveEvent] = []
    end = GameEndReason.MAX_MOVES
    err: Optional[str] = None

    while core.ply < max_moves:
        who = core.to_play
        side = "B" if who == BLACK else "W"
        client = black if who == BLACK else white
        turn_done = False
        comm_fails = 0
        while not turn_done:
            ok, row, col, cerr, ms = client.request_move(
                core.board.packed(),
                who,
                game_epoch=core.game_epoch,
                rejects=list(core.forbid),
            )
            if not ok:
                comm_fails += 1
                events.append(
                    MoveEvent(
                        core.ply,
                        side,
                        None,
                        None,
                        True,
                        False,
                        cerr or "comm_error",
                        ms,
                        core._turn_retries,
                    )
                )
                if comm_fails <= cfg.max_comm_retries:
                    client.close()
                    time.sleep(cfg.comm_retry_sleep_s)
                    continue
                end = GameEndReason.COMM_ERROR
                err = cerr or "comm_error"
                turn_done = True
                break

            comm_fails = 0
            is_pass = row is None or col is None
            if not is_pass:
                sticky = core.on_server_offer(int(row), int(col))
                if sticky is not None:
                    end = sticky
                    err = core.last_error
                    events.append(
                        MoveEvent(
                            core.ply, side, row, col, False, False, "sticky", ms, core._turn_retries
                        )
                    )
                    turn_done = True
                    break

            decision = core.evaluate_move(row, col, is_pass=is_pass)
            if decision.accept:
                if is_pass:
                    events.append(
                        MoveEvent(core.ply, side, None, None, True, True, None, ms, core._turn_retries)
                    )
                    maybe_end = core.apply_pass()
                    if maybe_end is not None:
                        end = maybe_end
                else:
                    events.append(
                        MoveEvent(
                            core.ply, side, int(row), int(col), False, True, None, ms, core._turn_retries
                        )
                    )
                    core.apply_move(int(row), int(col))
                turn_done = True
                break

            assert row is not None and col is not None
            loop = core.on_reject(int(row), int(col), decision.reason)
            events.append(
                MoveEvent(
                    core.ply,
                    side,
                    int(row),
                    int(col),
                    False,
                    False,
                    decision.reason.value,
                    ms,
                    core._turn_retries,
                )
            )
            if loop is not None:
                end = loop
                err = core.last_error
                turn_done = True
                break

        if end != GameEndReason.MAX_MOVES:
            break

    finished = datetime.now(timezone.utc).isoformat()
    return GameResult(
        game_id=game_id,
        end_reason=end.value,
        error=err,
        moves=events,
        plies=core.ply,
        started_at=started,
        finished_at=finished,
        black_endpoint=f"{black.ep.host}:{black.ep.port}",
        white_endpoint=f"{white.ep.host}:{white.ep.port}",
    )


def main(argv: Optional[List[str]] = None) -> int:
    p = argparse.ArgumentParser(description="Text VLM client (VLM vs VLM / self-play)")
    p.add_argument("--host", default="127.0.0.1")
    p.add_argument("--port", type=int, default=9319, help="흑(및 백 미지정 시) 포트")
    p.add_argument("--white-host", default=None)
    p.add_argument("--white-port", type=int, default=None, help="없으면 흑과 동일(self-play)")
    p.add_argument("--games", type=int, default=1)
    p.add_argument("--max-moves", type=int, default=400)
    p.add_argument("--out-dir", default="logs/text_client")
    p.add_argument("--timeout-s", type=float, default=180.0)
    args = p.parse_args(argv)

    out = Path(args.out_dir)
    out.mkdir(parents=True, exist_ok=True)
    games_path = out / "games.jsonl"

    bep = Endpoint(args.host, args.port, "black")
    wep = Endpoint(
        args.white_host or args.host,
        args.white_port if args.white_port is not None else args.port,
        "white",
    )
    black = TextVlmClient(bep, timeout_s=args.timeout_s)
    white = TextVlmClient(wep, timeout_s=args.timeout_s)

    try:
        for gid in range(1, args.games + 1):
            print(f"[text_client] game {gid} start ({bep.host}:{bep.port} vs {wep.host}:{wep.port})")
            res = play_one_game(
                game_id=gid, black=black, white=white, max_moves=args.max_moves
            )
            rec = asdict(res)
            with open(games_path, "a", encoding="utf-8") as f:
                f.write(json.dumps(rec, ensure_ascii=False) + "\n")
            tag = "ERROR" if res.error or res.end_reason in (
                "reject_loop",
                "sticky_offer",
                "retry_exhausted",
                "comm_error",
            ) else "ok"
            print(
                f"[text_client] game {gid} end={res.end_reason} plies={res.plies} "
                f"{tag} {res.error or ''}"
            )
    finally:
        black.close()
        white.close()
    return 0


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