code/local_eval/cli.py

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path


def _ensure_cuda_home() -> None:
    from .cuda_env import apply

    apply()

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT / "src") not in sys.path:
    sys.path.insert(0, str(ROOT / "src"))
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from .constants import (
    DEFAULT_DATA_ROOT,
    DEFAULT_GPU_MEMORY_UTILIZATION,
    DEFAULT_MODEL_DIR,
    DEFAULT_RUNS_DIR,
    GENESIS_REPO,
    GENESIS_REVISION,
)
from .judge import DEFAULT_JUDGE_RUN
from .download import download_genesis, download_lite_data, verify_model_dir


def _require_model_dir(path: Path, label: str) -> Path:
    resolved = path.expanduser().resolve()
    if not resolved.is_dir() or not (resolved / "config.json").is_file():
        raise SystemExit(
            f"{label} is not a local model directory: {path}\n"
            "Need a folder with config.json and *.safetensors.\n"
            "You only have genesis downloaded — smoke-test it with:\n"
            "  python -m local_eval run --self-check --samples 1 --turns 1"
        )
    return resolved


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Local Albedo SN97 eval harness")
    sub = parser.add_subparsers(dest="cmd", required=True)

    p = sub.add_parser("download-genesis", help="pull pinned genesis weights")
    p.add_argument("--dir", type=Path, default=DEFAULT_MODEL_DIR)

    p = sub.add_parser("download-lite-data", help="pull 2 mini-coder shards for local sampling")
    p.add_argument("--root", type=Path, default=DEFAULT_DATA_ROOT)

    p = sub.add_parser("verify", help="check a local checkpoint against the file allowlist")
    p.add_argument("--path", type=Path, default=DEFAULT_MODEL_DIR)

    p = sub.add_parser("run", help="duel challenger vs genesis with official gates + offline env")
    p.add_argument(
        "--challenger",
        type=Path,
        default=None,
        help="local checkpoint directory. Omit this and pass --self-check to run genesis alone.",
    )
    p.add_argument("--king", type=Path, default=DEFAULT_MODEL_DIR)
    p.add_argument("--dataset-root", type=Path, default=DEFAULT_DATA_ROOT)
    p.add_argument("--samples", type=int, default=8)
    p.add_argument("--seed", default="local-eval")
    p.add_argument("--turns", type=int, default=4, help="cap assistant turns (official is 12/16)")
    p.add_argument("--max-model-len", type=int, default=65536)
    p.add_argument("--king-gpus", default="0,1,2,3")
    p.add_argument("--chal-gpus", default="4,5,6,7")
    p.add_argument(
        "--gpu-memory-utilization",
        type=float,
        default=DEFAULT_GPU_MEMORY_UTILIZATION,
        help="vLLM per-GPU reservation (default 0.80; 0.90 fails when ~22 GiB is already held)",
    )
    p.add_argument("--runs-dir", type=Path, default=DEFAULT_RUNS_DIR)
    p.add_argument("--skip-king", action="store_true", help="score challenger gates only")
    p.add_argument(
        "--self-check",
        action="store_true",
        help="run genesis as the challenger (no second model). Implies --skip-king.",
    )
    p.add_argument(
        "--compile",
        action="store_true",
        help="allow vLLM CUDA compile (needs nvcc). Default is eager — this box has no /usr/local/cuda.",
    )

    p = sub.add_parser(
        "judge",
        help="score a finished local duel with the official GLM-5.2 checklist (OpenRouter)",
    )
    p.add_argument(
        "--run",
        type=Path,
        default=DEFAULT_JUDGE_RUN,
        help="eval-run directory with generated-samples.jsonl (default: v11 8x12)",
    )
    p.add_argument("--dataset-root", type=Path, default=DEFAULT_DATA_ROOT)
    p.add_argument("--seed", default=None, help="override verdict seed used for submit-protocol salt")
    p.add_argument(
        "--limit",
        type=int,
        default=2,
        help="score only the first N pairs (0 = all). Default 2 for a cheap smoke.",
    )
    p.add_argument(
        "--dry-run",
        action="store_true",
        help="rebuild JudgeSamples and print the plan; do not call OpenRouter",
    )

    args = parser.parse_args(argv)
    if args.cmd == "download-genesis":
        dest = download_genesis(args.dir)
        print(json.dumps(verify_model_dir(dest), indent=2))
        return 0
    if args.cmd == "download-lite-data":
        download_lite_data(args.root)
        return 0
    if args.cmd == "verify":
        report = verify_model_dir(args.path)
        print(json.dumps(report, indent=2))
        return 0 if report.get("ok") else 1
    if args.cmd == "run":
        _ensure_cuda_home()
        from .run import run_duel

        king = _require_model_dir(args.king, "--king")
        if args.self_check:
            challenger = king
            skip_king = True
            chal_gpus = args.king_gpus
        elif args.challenger is None:
            raise SystemExit(
                "pass --challenger /path/to/local/checkpoint, or --self-check to run genesis alone"
            )
        else:
            challenger = _require_model_dir(args.challenger, "--challenger")
            skip_king = args.skip_king
            chal_gpus = args.chal_gpus
        run_duel(
            challenger=challenger,
            king=king,
            dataset_root=args.dataset_root,
            sample_count=args.samples,
            seed=args.seed,
            max_turns=args.turns,
            king_gpus=args.king_gpus.split(","),
            chal_gpus=chal_gpus.split(","),
            max_model_len=args.max_model_len,
            runs_dir=args.runs_dir,
            skip_king=skip_king,
            enforce_eager=not args.compile,
            gpu_memory_utilization=args.gpu_memory_utilization,
        )
        return 0
    if args.cmd == "judge":
        from .judge import run_judge

        result = run_judge(
            run_dir=args.run,
            dataset_root=args.dataset_root,
            seed=args.seed,
            limit=args.limit,
            dry_run=args.dry_run,
        )
        if args.dry_run:
            return 0
        return 0 if result.get("state") == "succeeded" else 2
    return 2


def banner() -> str:
    return f"genesis {GENESIS_REPO}@{GENESIS_REVISION}"