예제 소스

board_viz_server.py

"""보드 시각화 TCP 서버 (BRD1 / RTXT).

클라이언트가 PNG + 361바이트 packed 보드를 보내면 VLM 으로 격자 JSON 을 생성하고,
가능하면 ASCII 격자까지 덧붙여 UTF-8 텍스트로 돌려줍니다.

프로토콜:
  요청: "BRD1" | u32_be body_len | u8 size | u8 reserved |
        u32_be png_len | png | size*size board (0/1/2)
  응답: "RTXT" | u8 ok (1=성공) | ok=1: u32_be utf8_len | text
                              | ok=0: u16_be utf8_len | err
  PING → PONG (4바이트, GO19 서버와 동일)
"""

from __future__ import annotations

import argparse
import io
import logging
import os
import socket
import socketserver
import struct
import sys
import time
from dataclasses import dataclass
from typing import Optional

from PIL import Image

from .server import (
    BOARD_SIZE,
    LOGGER,
    MAGIC_PING,
    MAGIC_PONG,
    _MlxVlmModel,
    _recv_exact,
    packed_board_ascii_for_log,
)
from .verify_board import (
    VERIFY_SYSTEM,
    build_verify_user_prompt,
    describe_model_grid_issues,
    parse_model_grid,
    repair_model_grid_with_packed_board,
)
from .verify_board import expected_rows

MAGIC_BRD1 = b"BRD1"
MAGIC_RTXT = b"RTXT"


class _Handled(Exception):
    """PING 처리 후 다음 프레임으로."""


def _read_brd1_frame(sock: socket.socket) -> tuple[Image.Image, bytes]:
    head = _recv_exact(sock, 4)
    if head == MAGIC_PING:
        sock.sendall(MAGIC_PONG)
        raise _Handled()
    if head != MAGIC_BRD1:
        raise ValueError(f"알 수 없는 헤더: {head!r} (BRD1 기대)")
    (body_len,) = struct.unpack(">I", _recv_exact(sock, 4))
    if body_len > 32 * 1024 * 1024:
        raise ValueError("body 너무 큼")
    body = _recv_exact(sock, body_len)
    if len(body) < 6:
        raise ValueError("body 너무 짧음")
    size = body[0]
    # body[1] reserved — 무시
    o = 2
    (png_len,) = struct.unpack(">I", body[o : o + 4])
    o += 4
    png = body[o : o + png_len]
    o += png_len
    board_len = size * size
    board = body[o : o + board_len]
    o += board_len
    if o != len(body):
        raise ValueError("body 길이 불일치")
    if size != BOARD_SIZE:
        raise ValueError(f"이 서버는 {BOARD_SIZE}×{BOARD_SIZE} 만 지원 (받음 {size})")
    if len(board) != board_len:
        raise ValueError("board 길이 불일치")
    img = Image.open(io.BytesIO(png)).convert("RGB")
    return img, bytes(board)


def _send_rtxt_ok(sock: socket.socket, text: str) -> None:
    b = text.encode("utf-8")
    if len(b) > 0xFFF_FFFF:
        b = b[: 0xFFF_FFFF]
    sock.sendall(MAGIC_RTXT + bytes([1]) + struct.pack(">I", len(b)) + b)


def _send_rtxt_err(sock: socket.socket, msg: str) -> None:
    b = msg.encode("utf-8")[:65535]
    sock.sendall(MAGIC_RTXT + bytes([0]) + struct.pack(">H", len(b)) + b)


def _ascii_from_rows(rows: list[str]) -> str:
    col_header = "     " + " ".join(f"{c:2d}" for c in range(BOARD_SIZE))
    lines = [col_header]
    for r, row in enumerate(rows):
        cells = [row[j : j + 1] for j in range(BOARD_SIZE)]
        lines.append(f"{r:3d}  " + " ".join(cells))
    return "\n".join(lines)


def _postprocess_model_text(raw: str, board: bytes) -> str:
    if not raw or not raw.strip():
        return raw
    parsed = parse_model_grid(raw)
    repair_note: Optional[str] = None
    if parsed is None:
        repaired, note = repair_model_grid_with_packed_board(raw, board)
        if repaired is not None:
            parsed = repaired
            repair_note = note
            LOGGER.info("[board-viz] %s", note)
    if parsed is None:
        diag = describe_model_grid_issues(raw)
        return (
            raw.strip()
            + "\n\n--- rows 검증 실패 (모델이 19×19 규격을 지키지 않음) ---\n"
            + diag
            + "\n(18줄만 온 경우 등은 서버가 packed 보드로 한 줄 보정을 시도합니다. "
            "그 외는 칸 수·공백·문자 규칙 위반입니다.)"
        )
    out_parts = [raw.strip()]
    if repair_note:
        out_parts.append(f"\n\n--- 서버 보정 ---\n{repair_note}")
    out_parts.append("\n\n--- ASCII board (from parsed rows) ---\n")
    out_parts.append(_ascii_from_rows(parsed))
    return "".join(out_parts)


@dataclass
class BoardVizConfig:
    host: str
    port: int
    model_id: str
    max_new_tokens: int
    temperature: float
    load_model: bool
    lora_dir: Optional[str]
    infer_timeout_s: float
    with_image: bool = True
    packed_only: bool = False


class _BoardVizTcpHandler(socketserver.BaseRequestHandler):
    server: "_BoardVizTcpServer"  # type: ignore[assignment]

    def handle(self) -> None:
        peer = self.client_address
        LOGGER.info("[board-viz] client connected: %s", peer)
        sock: socket.socket = self.request
        sock.settimeout(None)
        try:
            while True:
                try:
                    img, board = _read_brd1_frame(sock)
                except _Handled:
                    continue
                except EOFError:
                    break
                except Exception as e:
                    LOGGER.warning("[board-viz] 프레임 파싱 실패: %s", e)
                    try:
                        _send_rtxt_err(sock, f"프레임 파싱 실패: {e}")
                    except Exception:
                        pass
                    break

                n_b = sum(1 for b in board if b == 1)
                n_w = sum(1 for b in board if b == 2)
                LOGGER.info(
                    "[board-viz] BRD1 stones B=%d W=%d png=%s — packed board (test):\n%s",
                    n_b,
                    n_w,
                    img.size,
                    packed_board_ascii_for_log(board, BOARD_SIZE),
                )

                t0 = time.time()
                cfg = self.server.cfg
                model = self.server.model
                try:
                    if cfg.packed_only:
                        rows = expected_rows(board)
                        raw = (
                            '{"rows": '
                            + __import__("json").dumps(rows, ensure_ascii=False)
                            + "}"
                        )
                        out = raw + "\n\n--- ASCII board (packed, deterministic) ---\n" + _ascii_from_rows(rows)
                        _send_rtxt_ok(sock, out)
                        continue
                    if not cfg.load_model or model is None:
                        user = build_verify_user_prompt(board)
                        stub = (
                            "[no-model 모드] 요청 수신. user 프롬프트 길이="
                            f"{len(user)} bytes, 이미지={img.size}\n"
                            "모델 없이는 빈 응답입니다. --no-load 를 끄고 실행하세요."
                        )
                        _send_rtxt_ok(sock, stub)
                    else:
                        user = build_verify_user_prompt(board)
                        if cfg.with_image:
                            user = (
                                user
                                + "\n\n위 국면의 보드 이미지도 함께 참고하세요. "
                                "이미지와 텍스트가 다르면 텍스트를 우선하세요."
                            )
                        raw = model.predict_raw(
                            VERIFY_SYSTEM,
                            user,
                            image=img if cfg.with_image else None,
                            num_images=1 if cfg.with_image else 0,
                            timeout_s=cfg.infer_timeout_s,
                            max_tokens=cfg.max_new_tokens,
                            temperature=cfg.temperature,
                        )
                        out = _postprocess_model_text(raw, board)
                        _send_rtxt_ok(sock, out)
                except Exception as e:
                    LOGGER.exception("[board-viz] 추론/응답 실패")
                    try:
                        _send_rtxt_err(sock, str(e))
                    except Exception:
                        pass
                    break
                LOGGER.info("[board-viz] served in %.2fs", time.time() - t0)
        finally:
            LOGGER.info("[board-viz] client disconnected: %s", peer)


class _BoardVizTcpServer(socketserver.ThreadingTCPServer):
    allow_reuse_address = True

    def __init__(
        self,
        address,
        handler_cls,
        cfg: BoardVizConfig,
        model: Optional[_MlxVlmModel],
    ):
        super().__init__(address, handler_cls)
        self.cfg = cfg
        self.model = model


def run_board_viz_server(
    *,
    host: str,
    port: int,
    model_id: str,
    max_new_tokens: int,
    temperature: float,
    load_model: bool,
    lora_dir: Optional[str],
    infer_timeout_s: float = 300.0,
    with_image: bool = True,
    packed_only: bool = False,
) -> None:
    logging.basicConfig(
        level=os.environ.get("VLM_LOGLEVEL", "INFO"),
        format="%(asctime)s %(levelname)s %(message)s",
    )

    cfg = BoardVizConfig(
        host=host,
        port=port,
        model_id=model_id,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        load_model=load_model,
        lora_dir=lora_dir,
        infer_timeout_s=float(infer_timeout_s),
        with_image=bool(with_image),
        packed_only=bool(packed_only),
    )

    model: Optional[_MlxVlmModel] = None
    if load_model:
        model = _MlxVlmModel(
            model_id,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            lora_dir=lora_dir,
            text_only_moves=False,
            board_repr="both",
        )
        try:
            model.start(wait_load=True)
        except FileNotFoundError:
            LOGGER.error("LoRA 어댑터 부재 등으로 모델을 로드하지 못했습니다.")
            raise SystemExit(2)
        except BaseException as e:
            if lora_dir:
                LOGGER.exception("모델 로드 실패 → 종료: %s", e)
                raise SystemExit(2)
            LOGGER.exception("모델 로드 실패 → no-model 스텁 모드: %s", e)
            model = None
            cfg.load_model = False

    LOGGER.info(
        "board-viz listening on %s:%d (model=%s, load=%s, with_image=%s, packed_only=%s)",
        host,
        port,
        model_id,
        bool(model),
        bool(cfg.with_image),
        bool(cfg.packed_only),
    )
    with _BoardVizTcpServer((host, port), _BoardVizTcpHandler, cfg, model) as srv:
        try:
            srv.serve_forever()
        except KeyboardInterrupt:
            LOGGER.info("board-viz shutting down")
            srv.shutdown()
        finally:
            if model is not None:
                model.stop()


def main(argv: Optional[list[str]] = None) -> int:
    p = argparse.ArgumentParser(description="보드 시각화 BRD1/RTXT TCP 서버")
    p.add_argument("--host", default="********")
    p.add_argument("--port", type=int, default=9321)
    p.add_argument("--model", default="********")
    p.add_argument("--max-new-tokens", type=int, default=1024)
    p.add_argument("--temperature", type=float, default=0.0)
    p.add_argument("--lora-dir", default=None)
    p.add_argument("--infer-timeout-s", type=float, default=300.0)
    p.add_argument(
        "--no-image",
        action="store_true",
        help="이미지 입력 없이 텍스트 격자만으로 rows 를 복원(형식 안정성 향상용)",
    )
    p.add_argument(
        "--packed-only",
        action="store_true",
        help="모델을 거치지 않고 요청의 packed 보드로 rows 를 결정론적으로 생성(항상 정확)",
    )
    p.add_argument(
        "--no-load",
        action="store_true",
        help="모델 없이 프로토콜만 확인 (스텁 텍스트 응답)",
    )
    args = p.parse_args(argv)
    run_board_viz_server(
        host=args.host,
        port=args.port,
        model_id=args.model,
        max_new_tokens=args.max_new_tokens,
        temperature=args.temperature,
        load_model=not args.no_load,
        lora_dir=args.lora_dir,
        infer_timeout_s=args.infer_timeout_s,
        with_image=not args.no_image,
        packed_only=args.packed_only,
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())