예제 소스

nightly_sqlite_train.py

#!/usr/bin/env python3
"""매일 00:00(KST)에 전날까지 쌓인 matches.sqlite 대국을 학습한다.

대상: logs/mariago_web_monitor/matches.sqlite
  monitor_mariago_web_play.py 가 이어서 쓰는 게임 로그.
  finished_at 을 KST 날짜로 묶어, 아직 학습하지 않은 완료일 중 가장 오래된 하루만 돌린다.
  오늘(자정 이후 진행 중인 날)은 제외한다.

흐름:
  1) 그날 백 차례 국면 → KataGo 라벨 (sft_data/onpolicy_v1)
  2) train_ladder_lora --shape onpolicy --stop-vlm
     베이스는 직전 야간 어댑터, 없으면 lora_onpolicy_qwen_v1, 없으면 lora_mix_qwen_r5
     결과: vlm_models/lora_onpolicy_nightly/YYYYMMDD
     운영 서버 어댑터는 바꾸지 않는다 (기존 학습 스크립트와 동일).

상태: logs/onpolicy/nightly_state.json
로그: logs/onpolicy/nightly/YYYYMMDD/
"""
from __future__ import annotations

import argparse
import fcntl
import json
import os
import sqlite3
import subprocess
import sys
from datetime import datetime, timedelta, timezone
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
PY = str(ROOT / ".venv" / "bin" / "python")
DB = ROOT / "logs" / "mariago_web_monitor" / "matches.sqlite"
STATE = ROOT / "logs" / "onpolicy" / "nightly_state.json"
KST = timezone(timedelta(hours=9))

DAY_SQL = "date(replace(substr(finished_at,1,19),'T',' '), '+9 hours')"


def kst_today():
    return datetime.now(KST).date()


def log(msg: str, fh) -> None:
    line = f"[{datetime.now(KST).strftime('%Y-%m-%d %H:%M:%S')}] {msg}"
    print(line, flush=True)
    if fh is not None:
        fh.write(line + "\n")
        fh.flush()


def connect() -> sqlite3.Connection:
    uri = f"file:{DB}?mode=ro"
    return sqlite3.connect(uri, uri=True)


def completed_days(today: str) -> list[str]:
    con = connect()
    rows = con.execute(
        f"""select distinct {DAY_SQL} d from games
            where finished_at is not null and finished_at != ''
              and {DAY_SQL} < ?
            order by d""",
        (today,),
    ).fetchall()
    con.close()
    return [r[0] for r in rows if r[0]]


def count_games(day: str) -> tuple[int, int]:
    con = connect()
    games = con.execute(
        f"select count(*) from games where {DAY_SQL} = ?",
        (day,),
    ).fetchone()[0]
    moves = con.execute(
        f"""select count(*) from moves m
            where m.color='W' and m.attempt='first' and m.kind in ('move','pass')
              and exists (
                select 1 from games g
                where g.run_id = m.run_id and g.game_id = m.game_id
                  and {DAY_SQL.replace('finished_at', 'g.finished_at')} = ?
              )""",
        (day,),
    ).fetchone()[0]
    con.close()
    return int(games), int(moves)


def load_state() -> dict:
    if not STATE.exists():
        return {"done": []}
    return json.loads(STATE.read_text(encoding="utf-8"))


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


def mark(state: dict, day: str, status: str, extra: dict | None = None) -> None:
    done = [d for d in state.get("done", []) if d.get("day") != day]
    row = {"day": day, "status": status, "at": datetime.now(KST).strftime("%Y-%m-%d %H:%M:%S")}
    if extra:
        row.update(extra)
    done.append(row)
    state["done"] = done
    save_state(state)


def done_days(state: dict) -> set[str]:
    # seed: 자동화 등록 전에 이미 쌓인 날. ok/empty: 야간 작업이 처리한 날.
    # 실패는 기록하지 않아서 다음 밤에 다시 시도한다.
    return {d["day"] for d in state.get("done", []) if d.get("status") in ("ok", "empty", "seed")}


def pick_base() -> str:
    nightly = ROOT / "vlm_models" / "lora_onpolicy_nightly"
    if nightly.is_dir():
        cands = sorted(
            p for p in nightly.iterdir() if p.is_dir() and (p / "adapters.safetensors").exists()
        )
        if cands:
            return str(cands[-1].relative_to(ROOT))
    if (ROOT / "vlm_models" / "lora_onpolicy_qwen_v1" / "adapters.safetensors").exists():
        return "********"
    return "vlm_models/lora_mix_qwen_r5"


def busy() -> str:
    for pat in ("vlm_server.train_lora", "vlm_server.dataset_onpolicy", "train_ladder_lora.py"):
        r = subprocess.run(["pgrep", "-f", pat], capture_output=True)
        if r.returncode == 0:
            return pat
    return ""


def run(cmd: list[str], log_path: Path, fh) -> int:
    log("RUN " + " ".join(cmd), fh)
    env = os.environ.copy()
    env["PYTHONUNBUFFERED"] = "1"
    env["HF_HOME"] = str(ROOT / "vlm_models")
    env["HF_HUB_CACHE"] = str(ROOT / "vlm_models")
    env["TOKENIZERS_PARALLELISM"] = "********"
    log_path.parent.mkdir(parents=True, exist_ok=True)
    with log_path.open("a", encoding="utf-8") as f:
        f.write(f"\n===== {datetime.now(KST).isoformat()} =====\n")
        f.flush()
        return subprocess.run(cmd, cwd=ROOT, env=env, stdout=f, stderr=subprocess.STDOUT).returncode


PENDING = ROOT / "logs" / "onpolicy" / "pending_sources.json"


def ready_pending() -> dict | None:
    if not PENDING.exists():
        return None
    data = json.loads(PENDING.read_text(encoding="utf-8"))
    for src in data.get("sources", []):
        if src.get("status") == "ready" and Path(src.get("db", "")).exists():
            return src
    return None


def mark_pending(label: str, status: str, extra: dict | None = None) -> None:
    data = json.loads(PENDING.read_text(encoding="utf-8"))
    for src in data.get("sources", []):
        if src.get("label") == label:
            src["status"] = status
            src["at"] = datetime.now(KST).strftime("%Y-%m-%d %H:%M:%S")
            if extra:
                src.update(extra)
    PENDING.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def train_pending(src: dict, dry_run: bool) -> int:
    """등록된 SQLite 한 건을 KataGo 라벨 뒤 LoRA 학습한다."""
    label = str(src["label"])
    db = str(src["db"])
    out_name = "".join(ch if ch.isalnum() else "_" for ch in label)[:40]
    print(f"pending label={label} db={db}")
    if dry_run:
        print(f"base={pick_base()}")
        print(f"out=vlm_models/lora_onpolicy_nightly/{out_name}")
        return 0
    why = busy()
    if why:
        print(f"다른 학습이 돌아 {why} — 등록은 유지하고 다음 밤에 시도")
        return 0
    day_dir = ROOT / "logs" / "onpolicy" / "nightly" / out_name
    day_dir.mkdir(parents=True, exist_ok=True)
    data_out = ROOT / "sft_data" / "onpolicy_v1"
    with (day_dir / "run.log").open("a", encoding="utf-8") as fh:
        log(f"pending start {label} {db}", fh)
        rc = run(
            [PY, "-m", "vlm_server.dataset_onpolicy", "--db", db, "--out", str(data_out)],
            day_dir / "dataset.log",
            fh,
        )
        stats_path = data_out / "stats.json"
        stats = json.loads(stats_path.read_text(encoding="utf-8")) if stats_path.exists() else {}
        rows = int(stats.get("rows") or 0)
        log(f"dataset rc={rc} rows={rows}", fh)
        if rc != 0 or rows == 0:
            if rc == 0:
                mark_pending(label, "empty", {"rows": 0})
            return rc
        base = pick_base()
        out = f"vlm_models/lora_onpolicy_nightly/{out_name}"
        cmd = [
            PY, "scripts/train_ladder_lora.py",
            "--shape", "onpolicy", "--stop-vlm", "--skip-base-eval",
            "--shape-n", "1200", "--base", base, "--out", out,
        ]
        if (ROOT / "sft_data" / "baduk_sgf_rest" / "DONE").exists():
            cmd += ["--retain", "baduk_sgf_rest:400"]
        rc = run(cmd, day_dir / "train.log", fh)
        log(f"train rc={rc} base={base} out={out}", fh)
        if rc != 0:
            return rc
        mark_pending(label, "ok", {"rows": rows, "base": base, "out": out})
    return 0



def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--day", default="", help="KST YYYY-MM-DD. 비우면 미학습 완료일 중 가장 오래된 날")
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()

    if not DB.exists():
        print(f"sqlite 없음: {DB}", file=sys.stderr)
        return 1

    lock_path = ROOT / "logs" / "onpolicy" / "nightly.lock"
    lock_path.parent.mkdir(parents=True, exist_ok=True)
    lockf = lock_path.open("a", encoding="utf-8")
    try:
        fcntl.flock(lockf, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except BlockingIOError:
        print("이미 야간 학습이 실행 중이라 건너뜁니다.")
        return 0

    if not args.day:
        pend = ready_pending()
        if pend:
            return train_pending(pend, args.dry_run)

    today = kst_today().isoformat()
    state = load_state()
    days = completed_days(today)
    finished = done_days(state)
    if args.day:
        target = args.day
    else:
        pending = [d for d in days if d not in finished]
        target = pending[0] if pending else ""

    print(f"today={today} completed={days} done={sorted(finished)} target={target or '-'}")
    if not target:
        print("학습할 완료일이 없습니다.")
        return 0
    games, moves = count_games(target)
    print(f"day={target} games={games} white_positions={moves}")
    if args.dry_run:
        print(f"base={pick_base()}")
        print(f"out=vlm_models/lora_onpolicy_nightly/{target.replace('-', '')}")
        return 0

    day_dir = ROOT / "logs" / "onpolicy" / "nightly" / target.replace("-", "")
    day_dir.mkdir(parents=True, exist_ok=True)
    summary = day_dir / "run.log"
    with summary.open("a", encoding="utf-8") as fh:
        log(f"start day={target} games={games} positions={moves}", fh)
        if games == 0 or moves == 0:
            log("국면이 없어 학습을 건너뜁니다.", fh)
            mark(state, target, "empty", {"games": games, "positions": moves})
            return 0
        why = busy()
        if why:
            log(f"다른 학습이 돌아 {why} — 이번 밤은 건너뛰고 다음날 다시 시도", fh)
            return 0

        out_name = target.replace("-", "")
        data_out = ROOT / "sft_data" / "onpolicy_v1"
        rc = run(
            [
                PY, "-m", "vlm_server.dataset_onpolicy",
                "--db", str(DB),
                "--day", target,
                "--out", str(data_out),
            ],
            day_dir / "dataset.log",
            fh,
        )
        stats_path = data_out / "stats.json"
        stats = json.loads(stats_path.read_text(encoding="utf-8")) if stats_path.exists() else {}
        (day_dir / "stats.json").write_text(
            json.dumps(stats, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
        )
        rows = int(stats.get("rows") or 0)
        log(f"dataset rc={rc} rows={rows}", fh)
        if rc != 0:
            log("데이터 생성 실패. 상태를 완료로 표시하지 않습니다.", fh)
            return rc
        if rows == 0:
            log("라벨이 남지 않아 학습을 건너뜁니다.", fh)
            mark(state, target, "empty", {"games": games, "positions": moves, "rows": 0})
            return 0

        base = pick_base()
        out = f"vlm_models/lora_onpolicy_nightly/{out_name}"
        rc = run(
            [
                PY, "scripts/train_ladder_lora.py",
                "--shape", "onpolicy",
                "--stop-vlm",
                "--skip-base-eval",
                "--shape-n", "1200",
                "--base", base,
                "--out", out,
            ],
            day_dir / "train.log",
            fh,
        )
        status_src = ROOT / "logs" / "onpolicy" / "TRAIN_STATUS.md"
        if status_src.exists():
            (day_dir / "TRAIN_STATUS.md").write_text(status_src.read_text(encoding="utf-8"), encoding="utf-8")
        log(f"train rc={rc} base={base} out={out}", fh)
        if rc != 0:
            log("학습 실패. 다음날 같은 날을 다시 시도합니다.", fh)
            return rc
        mark(state, target, "ok", {"games": games, "positions": moves, "rows": rows, "base": base, "out": out})
        log("done", fh)
    return 0


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