예제 소스

run_gemma_r5_match.py

#!/usr/bin/env python3
"""Gemma 4(흑)와 R5(백) 무제한 5판 후, 각자 국면을 KataGo 정답으로 순서대로 학습한다.

R5가 9319에서 이미 떠 있으면 그대로 쓴다. Gemma 는 9341에만 띄운다.
학습이 끝나면 서빙 중인 R5 어댑터는 바꾸지 않고, 새 어댑터는 models.json 에 이름만 추가한다.
"""

from __future__ import annotations

import json
import os
import socket
import subprocess
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
PY = ROOT / ".venv" / "bin" / "python"
OUT = ROOT / "logs" / "gemma4_r5_20261009"
REGISTRY = ROOT / "scripts" / "models.json"


def log(msg: str) -> None:
    print(msg, flush=True)


def port_open(port: int) -> bool:
    sock = socket.socket()
    sock.settimeout(1)
    try:
        sock.connect(("127.0.0.1", port))
        return True
    except OSError:
        return False
    finally:
        sock.close()


def wait_port(port: int, timeout_s: float) -> bool:
    deadline = time.time() + timeout_s
    while time.time() < deadline:
        if port_open(port):
            return True
        time.sleep(5)
    return False


def start_serve(name: str, spec: dict, jsonl: Path) -> subprocess.Popen | None:
    if port_open(int(spec["port"])):
        log(f"[runner] {name} port {spec['port']} already up")
        return None
    reg = json.loads(REGISTRY.read_text(encoding="utf-8"))
    cmd = [
        str(PY),
        "-m",
        reg["serve"]["module"],
        "serve",
        "--host",
        reg["serve"]["host"],
        "--port",
        str(spec["port"]),
        "--model",
        spec["model"],
        "--log-jsonl",
        str(jsonl),
        *reg["serve"]["flags"],
    ]
    if spec.get("lora"):
        cmd.extend(["--lora-dir", spec["lora"]])
    jsonl.parent.mkdir(parents=True, exist_ok=True)
    logf = open(OUT / f"{name}_serve.stdout", "ab")
    log(f"[runner] start {name} port {spec['port']}")
    proc = subprocess.Popen(
        cmd,
        cwd=ROOT,
        stdout=logf,
        stderr=subprocess.STDOUT,
        start_new_session=True,
    )
    if not wait_port(int(spec["port"]), 600):
        raise SystemExit(f"{name} did not listen on {spec['port']}")
    log(f"[runner] {name} listening pid={proc.pid}")
    return proc


def stop_proc(proc: subprocess.Popen | None) -> None:
    if proc is None or proc.poll() is not None:
        return
    log(f"[runner] stop pid {proc.pid}")
    proc.terminate()
    try:
        proc.wait(timeout=30)
    except subprocess.TimeoutExpired:
        proc.kill()


def run(cmd: list[str], log_path: Path) -> int:
    log("[runner] " + " ".join(cmd))
    log_path.parent.mkdir(parents=True, exist_ok=True)
    with log_path.open("ab") as fh:
        rc = subprocess.call(cmd, cwd=ROOT, stdout=fh, stderr=subprocess.STDOUT)
    log(f"[runner] rc={rc} log={log_path}")
    return rc


def finished_games(db: Path) -> int:
    import sqlite3

    con = sqlite3.connect(db)
    n = con.execute("select count(*) from games where finished_at is not null and end_reason!='error'").fetchone()[0]
    con.close()
    return int(n)


def save_registry(reg: dict) -> None:
    REGISTRY.write_text(json.dumps(reg, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def main() -> int:
    OUT.mkdir(parents=True, exist_ok=True)
    reg = json.loads(REGISTRY.read_text(encoding="utf-8"))
    models = reg["models"]
    gemma_proc = start_serve("gemma4", models["gemma4"], OUT / "gemma4.jsonl")
    if not port_open(int(models["r5"]["port"])):
        raise SystemExit("R5 port 9319 is down; not starting a second Qwen")

    rc = run(
        [
            str(PY),
            "scripts/play_two_models.py",
            "--black",
            "gemma4",
            "--white",
            "r5",
            "--games",
            "5",
            "--max-moves",
            "0",
            "--out-dir",
            str(OUT.relative_to(ROOT)),
        ],
        OUT / "play.log",
    )
    stop_proc(gemma_proc)
    if rc != 0:
        return rc
    db = OUT / "matches.sqlite"
    n = finished_games(db)
    log(f"[runner] finished games={n}")
    if n < 3:
        log("[runner] too few games, skip learning")
        return 0

    label_jobs = [
        ("W", "sft_data/onpolicy_r5_vs_gemma4", OUT / "label_r5.log"),
        ("B", "sft_data/onpolicy_gemma4_vs_r5", OUT / "label_gemma4.log"),
    ]
    for color, dest, log_path in label_jobs:
        rc = run(
            [
                str(PY),
                "-m",
                "vlm_server.dataset_onpolicy",
                "--db",
                str(db),
                "--color",
                color,
                "--out",
                dest,
            ],
            log_path,
        )
        if rc != 0:
            return rc

    qwen = models["r5"]["model"]
    gemma = models["gemma4"]["model"]
    r5_out = "********"
    gemma_out = "********"
    rc = run(
        [
            str(PY),
            "scripts/train_ladder_lora.py",
            "--shape",
            "onpolicy",
            "--data-name",
            "onpolicy_r5_vs_gemma4",
            "--log-name",
            "r5_from_gemma4",
            "--model",
            qwen,
            "--base",
            models["r5"]["lora"],
            "--out",
            r5_out,
            "--shape-n",
            "1200",
            "--skip-base-eval",
        ],
        OUT / "train_r5.log",
    )
    if rc != 0:
        return rc
    rc = run(
        [
            str(PY),
            "scripts/train_ladder_lora.py",
            "--shape",
            "onpolicy",
            "--data-name",
            "onpolicy_gemma4_vs_r5",
            "--log-name",
            "gemma4_from_r5",
            "--model",
            gemma,
            "--base",
            "-",
            "--out",
            gemma_out,
            "--shape-n",
            "1200",
            "--skip-base-eval",
        ],
        OUT / "train_gemma4.log",
    )
    if rc != 0:
        return rc

    reg = json.loads(REGISTRY.read_text(encoding="utf-8"))
    reg["models"]["gemma4"]["lora"] = gemma_out
    reg["models"]["gemma4"]["note"] = "R5 대국 학습 어댑터. 9341이 비어 있을 때만 이 가중치로 띄운다."
    reg["models"]["r5_from_gemma4"] = {
        "port": 9343,
        "model": qwen,
        "lora": r5_out,
        "note": "Gemma4 대국으로 배운 R5. 9319 서빙 어댑터는 바꾸지 않았다.",
    }
    save_registry(reg)
    log("[runner] learning done; serving R5 adapter unchanged")
    return 0


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