예제 소스

run_gemma_x10_nolimit_loop.py

#!/usr/bin/env python3
"""10에폭 학습과 수 무제한 10국을 번갈아 반복한다.

이미 끝난 lora_gemma4_coord_x10 을 첫 10에폭으로 본다.
그 다음부터 수 무제한 10국 → 좌표 10에폭 → 수 무제한 10국 → 좌표 10에폭 순서로 계속한다.
수 무제한은 play_heuristic_vs_model 의 안전 상한 1000수다.
세트·에폭 사이에 logs/gemma_x10_nolimit/STOP 이 있으면 멈춘다.

  python scripts/run_gemma_x10_nolimit_loop.py
  상태: logs/gemma_x10_nolimit/REPORT.md
"""

from __future__ import annotations

import argparse
import json
import os
import shutil
import signal
import socket
import sqlite3
import subprocess
import time
from datetime import datetime
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
PY = str(ROOT / ".venv" / "bin" / "python")
OUT = ROOT / "logs" / "gemma_x10_nolimit"
HOLD = ROOT / "logs" / "mariago_web_watch" / "GEMMA_HOLD"
STOP = OUT / "STOP"
MODEL = "********"
START_LORA = "********"
RESUME_STAGE = ROOT / "vlm_models" / "lora_gemma4_coord_cycle_resume"
MIX = ROOT / "sft_data" / "mix_gemma_coord_r2" / "train.jsonl"
PORT = 9341
TRAIN_EPOCHS = 10


def log(msg: str) -> None:
    line = f"[{datetime.now():%m-%d %H:%M:%S}] {msg}"
    print(line, flush=True)
    OUT.mkdir(parents=True, exist_ok=True)
    with (OUT / "runner.log").open("a", encoding="utf-8") as f:
        f.write(line + "\n")


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


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


def gemma_pids() -> list[int]:
    try:
        out = subprocess.check_output(["pgrep", "-f", rf"vlm_server serve.*--port {PORT}"], text=True)
    except subprocess.CalledProcessError:
        return []
    return [int(x) for x in out.split()]


def stop_gemma() -> None:
    pids = gemma_pids()
    if not pids:
        return
    log(f"stop gemma {pids}")
    for pid in pids:
        try:
            os.kill(pid, signal.SIGTERM)
        except ProcessLookupError:
            pass
    deadline = time.time() + 40
    while time.time() < deadline and (port_open(PORT) or gemma_pids()):
        time.sleep(1)
    for pid in gemma_pids():
        try:
            os.kill(pid, signal.SIGKILL)
        except ProcessLookupError:
            pass
    time.sleep(2)


def gemma_is_ours(lora: str, jsonl: Path) -> bool:
    pids = gemma_pids()
    if not pids or not port_open(PORT):
        return False
    try:
        cmd = subprocess.check_output(["ps", "-p", str(pids[0]), "-o", "args="], text=True)
    except subprocess.CalledProcessError:
        return False
    return lora in cmd and str(jsonl) in cmd


def start_gemma(lora: str, jsonl: Path) -> None:
    stop_gemma()
    jsonl.parent.mkdir(parents=True, exist_ok=True)
    cmd = [
        PY, "-m", "vlm_server", "serve", "--host", "127.0.0.1", "--port", str(PORT),
        "--model", MODEL, "--lora-dir", lora,
        "--text-only-move", "--board-repr", "aligned", "--no-score-guided",
        "--illegal-retry-max", "16", "--illegal-retry-budget-s", "240",
        "--log-jsonl", str(jsonl),
    ]
    logf = open(OUT / "gemma_serve.stdout", "ab")
    log(f"start gemma lora={lora}")
    subprocess.Popen(cmd, cwd=ROOT, stdout=logf, stderr=subprocess.STDOUT, start_new_session=True)
    deadline = time.time() + 600
    while time.time() < deadline:
        if port_open(PORT):
            log("gemma listening")
            return
        time.sleep(5)
    raise SystemExit("gemma did not listen")


def state() -> dict:
    p = OUT / "state.json"
    return json.loads(p.read_text(encoding="utf-8")) if p.exists() else {}


def save_state(st: dict) -> None:
    OUT.mkdir(parents=True, exist_ok=True)
    (OUT / "state.json").write_text(json.dumps(st, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def completed_games(log_path: Path) -> int:
    if not log_path.exists():
        return 0
    return sum(1 for ln in log_path.read_text(encoding="utf-8", errors="replace").splitlines()
               if ln.startswith("[match] game ") and " plies=" in ln)


def edge_line(row: int, col: int) -> int:
    return min(row, col, 18 - row, 18 - col) + 1


def match_stats(set_dir: Path) -> dict:
    moves = []
    jl = set_dir / "gemma4.jsonl"
    if jl.exists():
        for ln in jl.open(encoding="utf-8", errors="replace"):
            try:
                rec = json.loads(ln)
            except json.JSONDecodeError:
                continue
            if rec.get("event") == "move":
                moves.append(rec)
    n = len(moves) or 1
    st = {
        "gemma_moves": len(moves),
        "first_try_pct": round(100 * sum(1 for r in moves if not r.get("retry_count") and not r.get("fallback_used")) / n, 1),
        "fallback_pct": round(100 * sum(1 for r in moves if r.get("fallback_used")) / n, 1),
        "attempts_per_move": round(sum(len(r.get("vlm_seeds") or []) for r in moves) / n, 2),
        "occupied_pct": round(100 * sum(1 for r in moves if r.get("illegal_reason") == "occupied") / n, 1),
    }
    db = set_dir / "matches.sqlite"
    if db.exists():
        con = sqlite3.connect(db)
        games = con.execute(
            "select game_id, end_reason, plies, winner, katago_lead_b from games where finished_at is not null"
        ).fetchall()
        rows = con.execute(
            "select ply, row, col from moves where color='W' and kind='move' and row is not null"
        ).fetchall()
        con.close()
        leads = [g[4] for g in games if g[4] is not None]
        lines: dict[str, int] = {}
        for _ply, row, col in rows:
            key = str(edge_line(int(row), int(col)))
            lines[key] = lines.get(key, 0) + 1
        white_n = sum(lines.values()) or 1
        st.update(
            games=len(games),
            gemma_wins=sum(1 for g in games if g[3] == "white"),
            heuristic_wins=sum(1 for g in games if g[3] == "black"),
            avg_lead_b=round(sum(leads) / len(leads), 1) if leads else None,
            leads=[round(x, 2) for x in leads],
            plies=[g[2] for g in games],
            end_reasons=[g[1] for g in games],
            line1=lines.get("1", 0),
            line1_rate=round(lines.get("1", 0) / white_n, 3),
            lines=lines,
        )
    return st


def prune_checkpoints(directory: Path) -> None:
    for path in directory.glob("*_adapters.safetensors"):
        if path.name[0].isdigit():
            path.unlink()
            log(f"pruned {path.name}")


def publish_resume(src: Path) -> None:
    weights = src / "adapters.safetensors"
    cfg = src / "adapter_config.json"
    if not weights.exists() or not cfg.exists():
        raise SystemExit(f"adapter incomplete: {src}")
    RESUME_STAGE.mkdir(parents=True, exist_ok=True)
    shutil.copy2(weights, RESUME_STAGE / "adapters.safetensors")
    shutil.copy2(cfg, RESUME_STAGE / "adapter_config.json")


def train_epoch(out_lora: str, resume_from: Path, log_path: Path) -> int:
    cmd = [
        PY, "-m", "vlm_server.train_lora",
        "--dataset-jsonl", str(MIX),
        "--model", MODEL,
        "--output-dir", out_lora,
        "--epochs", "1",
        "--batch-size", "1",
        "--grad-accum", "8",
        "--lr", "1e-5",
        "--lora-rank", "16",
        "--lora-alpha", "32",
        "--max-seq-length", "4096",
        # 앞 2600개: 좌표 읽기 + 휴리스틱 수 + R5 수. dataset_r5_moves.py 가 이 길이에 맞춘다.
        "--max-samples", "2600",
        "--grad-checkpoint",
        "--resume-from", str(resume_from),
    ]
    return run(cmd, log_path)


def write_report(st: dict) -> None:
    lines = [
        "# Gemma 4 좌표 10에폭과 수 무제한 10국 반복",
        "",
        f"- 갱신: {datetime.now():%Y-%m-%d %H:%M}",
        f"- 시작 어댑터: `{START_LORA}`",
        f"- 현재 어댑터: `{st.get('lora', START_LORA)}`",
        "- 순서: 10에폭(이미 완료) → 수 무제한 10국 → 좌표 10에폭 → 수 무제한 10국 → …",
        "- 안전 상한 1000수, 휴리스틱 흑 vs Gemma 백, 점수 유도 없음",
        f"- 끝난 대국 세트: {st.get('matches_done', 0)}",
        f"- 끝난 추가 학습: {st.get('trains_done', 0)}회 × {TRAIN_EPOCHS}에폭",
        "",
        "## 대국",
        "",
        "| 세트 | 어댑터 | Gemma 승 | 휴리스틱 승 | 평균 흑 우세 | 평균 수 | 첫 시도 합법 | 서버 대체 | 1선 |",
        "|---|---|---|---|---|---|---|---|---|",
    ]
    for key, s in sorted(st.get("sets", {}).items(), key=lambda kv: int(kv[0])):
        plies = s.get("plies") or []
        avg_plies = round(sum(plies) / len(plies), 1) if plies else "-"
        lines.append(
            f"| {key} | `{Path(s.get('lora', '')).name}` | {s.get('gemma_wins', '-')} | {s.get('heuristic_wins', '-')} | "
            f"{s.get('avg_lead_b', '-')} | {avg_plies} | {s.get('first_try_pct', '-')}% | "
            f"{s.get('fallback_pct', '-')}% | {s.get('line1_rate', '-')} |"
        )
    lines += ["", "## 학습", ""]
    for key, t in sorted(st.get("trains", {}).items(), key=lambda kv: int(kv[0])):
        lines.append(f"- {key}회차 `{t.get('lora')}` {t.get('epochs_done', 0)}/{TRAIN_EPOCHS}에폭, 기준 `{t.get('base')}`")
    lines.append("")
    (OUT / "REPORT.md").write_text("\n".join(lines), encoding="utf-8")


def play_set(st: dict, games: int) -> bool:
    index = st["matches_done"] + 1
    set_dir = OUT / f"set{index:03d}"
    done = completed_games(set_dir / "play.log")
    left = games - done
    lora = st["lora"]
    if left <= 0:
        stats = match_stats(set_dir)
        stats.update(done=True, lora=lora)
        st["sets"][str(index)] = stats
        st["matches_done"] = index
        st["sets_done"] = index
        save_state(st)
        write_report(st)
        log(f"set {index} already complete {stats}")
        return True
    if not gemma_is_ours(lora, set_dir / "gemma4.jsonl"):
        start_gemma(lora, set_dir / "gemma4.jsonl")
    rc = run([
        PY, "scripts/play_heuristic_vs_model.py",
        "--port", str(PORT),
        "--model-name", Path(lora).name,
        "--games", str(left),
        "--start-id", str(done + 1),
        "--max-moves", "0",
        "--seed", str(19 + index),
        "--out-dir", str(set_dir.relative_to(ROOT)),
    ], set_dir / "play.log")
    finished = completed_games(set_dir / "play.log")
    if finished < games:
        log(f"set {index} stopped rc={rc} games={finished}/{games}")
        stop_gemma()
        return False
    stats = match_stats(set_dir)
    stats.update(done=True, lora=lora)
    st["sets"][str(index)] = stats
    st["matches_done"] = index
    st["sets_done"] = index
    save_state(st)
    write_report(st)
    log(f"set {index} {stats}")
    return True


def train_round(st: dict) -> bool:
    rnd = st["trains_done"] + 1
    info = st.setdefault("trains", {}).setdefault(str(rnd), {})
    out_lora = info.get("lora") or f"vlm_models/lora_gemma4_coord_cycle{rnd:02d}"
    info.setdefault("base", st["lora"])
    info.setdefault("epochs_done", 0)
    info["lora"] = out_lora
    save_state(st)
    stop_gemma()
    while info["epochs_done"] < TRAIN_EPOCHS:
        if STOP.exists():
            log("STOP during training")
            return False
        epoch = info["epochs_done"] + 1
        resume = Path(info["base"]) if epoch == 1 else RESUME_STAGE
        if epoch > 1 and not (RESUME_STAGE / "adapters.safetensors").exists():
            raise SystemExit(f"resume adapter missing for cycle {rnd} epoch {epoch}")
        tdir = OUT / f"train{rnd:02d}" / f"epoch{epoch:02d}"
        t0 = time.time()
        log(f"train {rnd} epoch {epoch}/{TRAIN_EPOCHS} resume={resume}")
        rc = train_epoch(out_lora, resume, tdir / "train.log")
        if rc != 0:
            log(f"train {rnd} epoch {epoch} failed rc={rc}")
            return False
        out_dir = ROOT / out_lora
        prune_checkpoints(out_dir)
        publish_resume(out_dir)
        info["epochs_done"] = epoch
        info[f"epoch{epoch}_hours"] = round((time.time() - t0) / 3600, 2)
        save_state(st)
        write_report(st)
        log(f"train {rnd} epoch {epoch} done {info[f'epoch{epoch}_hours']}h")
    st["lora"] = out_lora
    st["trains_done"] = rnd
    save_state(st)
    write_report(st)
    log(f"train {rnd} adopted {out_lora}")
    return True


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--games", type=int, default=10)
    args = ap.parse_args()
    if not MIX.exists():
        raise SystemExit(f"missing mix {MIX}")
    OUT.mkdir(parents=True, exist_ok=True)
    (OUT / "runner.pid").write_text(str(os.getpid()), encoding="utf-8")
    st = state()
    if "matches_done" not in st:
        st["matches_done"] = int(st.get("sets_done") or 0)
    st.setdefault("trains_done", 0)
    st.setdefault("sets", {})
    st.setdefault("trains", {})
    st.setdefault("lora", START_LORA)
    save_state(st)
    HOLD.write_text("gemma x10 cycle\n", encoding="utf-8")

    def _term(_signum, _frame):
        raise SystemExit(0)

    signal.signal(signal.SIGTERM, _term)
    try:
        while not STOP.exists():
            if st["matches_done"] == st["trains_done"]:
                if not play_set(st, args.games):
                    time.sleep(10)
                    continue
            elif st["matches_done"] == st["trains_done"] + 1:
                if not train_round(st):
                    return 1
            else:
                raise SystemExit(f"cycle state mismatch matches={st['matches_done']} trains={st['trains_done']}")
            st = state()
        log("STOP file seen")
        return 0
    finally:
        stop_gemma()
        if HOLD.exists():
            HOLD.unlink()
        write_report(state())
        log("cycle stopped")


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