ops/run_eval.py

#!/usr/bin/env python3
"""One-command evaluation entry — works anywhere, no manual steps.

Quick start (from the repo root):

    python ops/run_eval.py                       # v2 checkpoint from HF + mock game
    python ops/run_eval.py --games all25         # the 25 public games (needs arc-agi)
    python ops/run_eval.py --backend tiny --no-checkpoint --max-actions 60   # CPU harness test
    python ops/run_eval.py --checkpoint /path/to/dir-with-qwenjev_s1.pt     # local dir

What it does:
  1. Resolves the checkpoint: `hf:<repo>@<subpath>` (downloads qwenjev_s1.pt +
     hf_lora parts, reassembles adapter_model.safetensors, verifies sha256)
     or a local dir laid out as  <dir>/qwenjev_s1.pt + <dir>/hf_lora/...
  2. Picks the backend: `auto` -> hf (GPU, fp16 base + LoRA merged) when CUDA
     is available, else tiny (CPU). The hf path never quantizes for eval:
     merge=True is illegal on a kbit base (see qwenjev/hf_compat.py).
  3. Runs the hybrid agent per game (heuristic fallback if System-1 fails to
     load) and writes a JSON report with provenance.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import shutil
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
sys.path.insert(0, str(ROOT / "train"))

DEFAULT_CKPT = "hf:tchbcb/qwenjev@t4_assets_v2/checkpoints"


def log(msg: str) -> None:
    print(f"[run_eval {time.strftime('%H:%M:%S')}] {msg}", flush=True)


# --------------------------------------------------------------- checkpoint --
def resolve_checkpoint(spec: str, cache_root: Path) -> Path | None:
    """Return a dir containing qwenjev_s1.pt (+ hf_lora/), or None if 'none'."""
    if spec == "none":
        return None
    if not spec.startswith("hf:"):
        d = Path(spec)
        if not (d / "qwenjev_s1.pt").exists():
            raise SystemExit(f"[error] {d} has no qwenjev_s1.pt")
        ensure_adapter_safetensors(d / "hf_lora")
        return d

    repo, _, sub = spec[3:].partition("@")
    sub = sub.strip("/")
    # hf_hub_download(local_dir=base) preserves the repo-relative path, so
    # passing the REPO-level dir makes files land exactly at <out>/...
    # (= base/<sub>). Passing out.parent instead would double the first sub
    # component (t4_assets_v2/t4_assets_v2/...) — caught in live testing.
    base = cache_root / repo.replace("/", "__")
    out = base / sub
    out.mkdir(parents=True, exist_ok=True)
    from huggingface_hub import hf_hub_download, list_repo_files

    wanted = [f for f in list_repo_files(repo, repo_type="model")
              if f.startswith(sub + "/")
              and (f == f"{sub}/qwenjev_s1.pt" or "/hf_lora/" in f)]
    for f in wanted:
        hf_hub_download(repo_id=repo, repo_type="model", filename=f,
                        local_dir=str(base))
        log(f"downloaded {f}")
    if not (out / "qwenjev_s1.pt").exists():
        raise SystemExit(f"[error] download failed: {out / 'qwenjev_s1.pt'} missing")
    ensure_adapter_safetensors(out / "hf_lora")
    return out


def ensure_adapter_safetensors(lora_dir: Path) -> None:
    """Reassemble parts/ -> adapter_model.safetensors when only shards exist.
    v0.4.8: also accept adapter_model.pt (single-file form) — HF strips
    .safetensors from commits on our repo, so weights are stored as .pt."""
    if (lora_dir / "adapter_model.safetensors").exists():
        return
    pt = lora_dir / "adapter_model.pt"
    if pt.exists() and pt.stat().st_size > 1_000_000:
        shutil.copyfile(pt, lora_dir / "adapter_model.safetensors")
        log(f"adapter restored from adapter_model.pt ({pt.stat().st_size} B)")
        return
    manifest = lora_dir / "parts" / "manifest.json"
    if not manifest.exists():
        return
    m = json.loads(manifest.read_text())
    nonce = int(m.get("nonce_bytes", 16))
    buf = bytearray()
    for p in m["parts"]:
        piece = (lora_dir / "parts" / p["name"]).read_bytes()
        buf += piece[nonce:]
    h = hashlib.sha256(buf).hexdigest()
    if h != m["sha256"]:
        raise SystemExit(f"[error] adapter sha256 mismatch: {h} != {m['sha256']}")
    (lora_dir / "adapter_model.safetensors").write_bytes(bytes(buf))
    log(f"adapter reassembled: {len(buf)} B, sha256 {h[:12]}... OK")


# ------------------------------------------------------------------ backend --
def pick_backend(backend: str) -> str:
    if backend != "auto":
        return backend
    try:
        import torch
        return "hf" if torch.cuda.is_available() else "tiny"
    except ImportError:
        raise SystemExit("[error] torch not installed — pip install torch "
                         "(GPU: default PyPI wheel; CPU: --index-url "
                         "https://download.pytorch.org/whl/cpu)")


# --------------------------------------------------------------------- main --
def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--checkpoint", default=DEFAULT_CKPT,
                    help=f"hf:<repo>@<subpath> | local dir | none (default {DEFAULT_CKPT})")
    ap.add_argument("--backend", default="auto", choices=["auto", "hf", "tiny"])
    ap.add_argument("--games", nargs="*", default=["mock"],
                    help="mock | all25 | explicit game ids (space separated)")
    ap.add_argument("--list-games", action="store_true",
                    help="list the public games (live from the arcade API) and exit")
    ap.add_argument("--download-only", action="store_true",
                    help="cache all 25 game sources locally, then exit")
    ap.add_argument("--env-dir", default=None,
                    help="game cache dir (default: ./environment_files, "
                         "overridden by ARC_ENVIRONMENTS_DIR)")
    ap.add_argument("--max-actions", type=int, default=300)
    ap.add_argument("--out", default="eval_report.json")
    ap.add_argument("--cache", default=str(Path.home() / ".cache" / "qwenjev_eval"))
    ap.add_argument("--no-checkpoint", action="store_true",
                    help="skip checkpoint entirely (untrained heads / harness test)")
    args = ap.parse_args()

    if args.env_dir:
        import os
        os.environ["ARC_ENVIRONMENTS_DIR"] = args.env_dir

    from data.collect_selfplay import PUBLIC_GAMES_FALLBACK
    if args.list_games:
        try:
            import arc_agi
            from arc_agi import OperationMode
            arc = arc_agi.Arcade(operation_mode=OperationMode.NORMAL)
            live = [e.game_id.split("-")[0] for e in arc.get_environments()]
        except Exception:
            live = []
        print("\n".join(live or PUBLIC_GAMES_FALLBACK))
        return
    if args.download_only:
        import subprocess
        r = subprocess.run([sys.executable, str(ROOT / "ops" / "fetch_games.py")]
                           + (["--dir", args.env_dir] if args.env_dir else []))
        sys.exit(r.returncode)

    games: list[str] = []
    for g in args.games:
        if g == "mock":
            games.append("mock-pushbox")
        elif g == "all25":
            games.extend(PUBLIC_GAMES_FALLBACK)
        else:
            games.append(g)
    if any(not g.startswith("mock") for g in games):
        try:
            import arc_agi  # noqa: F401
        except ImportError:
            raise SystemExit(
                "[error] real games need the official offline engine:\n"
                "  pip install arc-agi   (PyPI, pure-python engine)\n"
                "Game sources are then fetched automatically on first use "
                "(anonymous arcade key, no Kaggle account needed) and cached "
                "in ./environment_files/ — or pre-cache them with:\n"
                "  python ops/fetch_games.py")

    backend = pick_backend(args.backend)
    ckpt_dir = None if args.no_checkpoint else resolve_checkpoint(args.checkpoint, Path(args.cache))

    cfg = {
        "system1": {
            "backend": backend,
            "hf_model": "Qwen/Qwen3-VL-2B-Instruct",
            "max_leaf_tokens": 512,
            "score_levels": 10,
            "leaf_batch_size": 2,
            "eval_leaf_batch": 8,
        },
        "eval_max_actions": args.max_actions,
        "out_dir": str(ckpt_dir) if ckpt_dir else "checkpoints/none",
    }
    if ckpt_dir is None:
        cfg["out_dir"] = "checkpoints/definitely-missing"  # build_s1 skips silently

    from evaluate_local import build_s1, run_game
    s1 = None
    try:
        s1 = build_s1(cfg)
    except Exception as exc:
        log(f"[warn] System-1 unavailable ({exc}); heuristic policy will run")
    if backend == "hf" and s1 is None and not args.no_checkpoint:
        log("[warn] hf backend failed to load — continuing with heuristics only")

    from qwenjev.hybrid import HybridAgent, HybridConfig
    brain = HybridAgent(s1=s1, s2=None,
                        config=HybridConfig(game_time_limit_s=600))

    def write_report(report: list, t0: float) -> dict:
        """Atomically dump the report; called after EVERY game so a Colab
        disconnect mid-run keeps all completed results."""
        out = {
            "meta": {
                "backend": backend,
                "checkpoint": args.checkpoint if not args.no_checkpoint else "none",
                "games": games,
                "max_actions": args.max_actions,
                "seconds_total": round(time.time() - t0, 1),
                "timestamp": time.strftime("%F %T"),
            },
            "games": report,
            "summary": {
                "played": len(report),
                "wins": sum(1 for r in report if r["win"]),
                "best_level_sum": sum(r["best_level"] for r in report),
                "actions_sum": sum(r["actions"] for r in report),
            },
        }
        tmp = Path(args.out + ".tmp")
        tmp.write_text(json.dumps(out, indent=2))
        tmp.replace(args.out)
        return out

    from data.collect_selfplay import make_env
    report, t0 = [], time.time()
    write_report(report, t0)  # header-only report, proves out path works
    for game_id in games:
        try:
            env = make_env(game_id, None)
            res = run_game(env, game_id, brain, max_actions=args.max_actions)
            res["steps_per_s"] = round(res["actions"] / max(res["seconds"], 0.1), 2)
            report.append(res)
            log(f"{game_id}: {res}")
        except Exception as exc:
            log(f"[warn] {game_id}: {type(exc).__name__}: {exc}")
        write_report(report, t0)

    out = write_report(report, t0)
    log(f"done: {out['summary']} -> {args.out}")


if __name__ == "__main__":
    main()