예제 소스
바둑 on-policy 학습: 대국 로그를 KataGo 라벨로 바꾸기 · 파일 169개
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())