code/local_train/cli.py

from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
from pathlib import Path

from .constants import (
    DEFAULT_ADAPTER_DIR,
    DEFAULT_DATA_ROOT,
    DEFAULT_DPO_ADAPTER_DIR,
    DEFAULT_DPO_EXPORT_DIR,
    DEFAULT_DPO_PACK,
    DEFAULT_EXPORT_DIR,
    DEFAULT_PACK_DIR,
    DEFAULT_RL_ADAPTER_DIR,
    DEFAULT_RL_EXPORT_DIR,
    GENESIS_DIR,
)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Albedo challenger SFT → export → duel → GRPO")
    sub = parser.add_subparsers(dest="cmd", required=True)

    p = sub.add_parser("pack", help="build the SFT jsonl pack from official parquet")
    p.add_argument("--dataset-root", type=Path, default=DEFAULT_DATA_ROOT)
    p.add_argument("--out-dir", type=Path, default=DEFAULT_PACK_DIR)
    p.add_argument("--max-examples", type=int, default=20_000)
    p.add_argument("--seed", default="sft-pack")
    p.add_argument("--n-salts", type=int, default=2)
    p.add_argument("--submit-frac", type=float, default=0.20)
    p.add_argument("--edit-frac", type=float, default=0.35)
    p.add_argument(
        "--expand-submit-salts",
        action="store_true",
        help="repeat each gold submit under every salt so all markers/tails are seen",
    )
    p.add_argument("--smoke", action="store_true", help="64-example dry pack")

    p = sub.add_parser("sft", help="LoRA SFT from genesis on a packed jsonl")
    p.add_argument("--pack", type=Path, default=None)
    p.add_argument("--out-dir", type=Path, default=DEFAULT_ADAPTER_DIR)
    p.add_argument("--model", type=Path, default=GENESIS_DIR)
    p.add_argument("--max-steps", type=int, default=None)
    p.add_argument("--epochs", type=float, default=1.0)
    p.add_argument("--max-seq-len", type=int, default=16384)
    p.add_argument("--lr", type=float, default=1e-4)
    p.add_argument("--smoke", action="store_true")

    p = sub.add_parser("export", help="merge LoRA and copy genesis metadata byte-for-byte")
    p.add_argument("--adapter", type=Path, default=DEFAULT_ADAPTER_DIR)
    p.add_argument("--genesis", type=Path, default=GENESIS_DIR)
    p.add_argument("--out-dir", type=Path, default=DEFAULT_EXPORT_DIR)
    p.add_argument("--base", type=Path, default=None)

    p = sub.add_parser("check", help="sha256 metadata vs genesis + albedo check-model")
    p.add_argument("--path", type=Path, required=True)
    p.add_argument("--genesis", type=Path, default=GENESIS_DIR)

    p = sub.add_parser("duel", help="local_eval challenger vs genesis + go/no-go")
    p.add_argument("--challenger", type=Path, default=DEFAULT_EXPORT_DIR)
    p.add_argument("--king", type=Path, default=GENESIS_DIR)
    p.add_argument("--samples", type=int, default=8)
    p.add_argument("--turns", type=int, default=12)
    p.add_argument("--seed", default="local-eval")
    p.add_argument("--skip-king", action="store_true")

    p = sub.add_parser("chain-pack", help="SFT rows for live sanity chain (edit→submit→reject→edit→submit)")
    p.add_argument("--dataset-root", type=Path, default=DEFAULT_DATA_ROOT)
    p.add_argument("--out-dir", type=Path, default=DEFAULT_PACK_DIR)
    p.add_argument("--max-groups", type=int, default=200)
    p.add_argument("--seed", default="chain-pack")
    p.add_argument("--n-salts", type=int, default=3)
    p.add_argument("--mix-pack", type=Path, default=None)
    p.add_argument("--mix-n", type=int, default=2000)
    p.add_argument(
        "--mix-edit-n",
        type=int,
        default=0,
        help="official at_edit rows to mix (Stage A). Overrides --mix-n when set with --mix-submit-n",
    )
    p.add_argument(
        "--mix-submit-n",
        type=int,
        default=0,
        help="official submit-after-edit rows to mix (Stage A)",
    )
    p.add_argument("--smoke", action="store_true")

    p = sub.add_parser(
        "chain-gold",
        help="compile 6-turn sessions that pass evaluate_chain, then slice for SFT",
    )
    p.add_argument("--dataset-root", type=Path, default=DEFAULT_DATA_ROOT)
    p.add_argument("--out-dir", type=Path, default=DEFAULT_PACK_DIR)
    p.add_argument("--max-groups", type=int, default=150)
    p.add_argument("--seed", default="chain-gold")
    p.add_argument("--n-salts", type=int, default=3)
    p.add_argument("--mix-pack", type=Path, default=None)
    p.add_argument("--mix-edit-n", type=int, default=800)
    p.add_argument(
        "--mix-submit-n",
        type=int,
        default=0,
        help="official post-edit submit rows to mix (keep small vs edits)",
    )
    p.add_argument("--smoke", action="store_true")

    p = sub.add_parser("chain", help="local live-sanity chain eval (reject first submit)")
    p.add_argument("--challenger", type=Path, default=DEFAULT_RL_EXPORT_DIR)
    p.add_argument(
        "--samples",
        type=int,
        default=6,
        help="3 live-gate files first, then random extras (default 6 = 3+3 transfer)",
    )
    p.add_argument("--turns", type=int, default=32)
    p.add_argument("--seed", default="chain-eval")
    p.add_argument(
        "--no-live-gate",
        action="store_true",
        help="skip the 3 oversampled gate files; draw random samples only",
    )

    p = sub.add_parser(
        "dpo-pack",
        help="chosen/rejected pairs for the live empty-submit policy (not the local 3 files)",
    )
    p.add_argument("--out-dir", type=Path, default=DEFAULT_PACK_DIR)
    p.add_argument("--seed", default="dpo-live-fault-v16")
    p.add_argument(
        "--narrow",
        action="store_true",
        help="only after-submit empty-echo pairs on toml/config paths (v16 random-6 fails)",
    )

    p = sub.add_parser("dpo", help="LoRA DPO on live-fault pairs (prefer this over GRPO)")
    p.add_argument("--pack", type=Path, default=DEFAULT_DPO_PACK)
    p.add_argument("--model", type=Path, default=Path("/workspace/models/albedo-challenger-v15"))
    p.add_argument("--out-dir", type=Path, default=DEFAULT_DPO_ADAPTER_DIR)
    p.add_argument("--max-steps", type=int, default=30)
    p.add_argument("--max-seq-len", type=int, default=4096)
    p.add_argument("--lr", type=float, default=5e-6)
    p.add_argument("--beta", type=float, default=0.1)
    p.add_argument("--smoke", action="store_true")

    p = sub.add_parser(
        "policy",
        help="one-turn live-fault prefixes: next command must be sed -i, never a marker",
    )
    p.add_argument("--challenger", type=Path, default=DEFAULT_DPO_EXPORT_DIR)

    p = sub.add_parser("rl", help="light GRPO on the SFT export (not proxy_score)")
    p.add_argument("--pack", type=Path, default=None)
    p.add_argument("--model", type=Path, default=DEFAULT_EXPORT_DIR)
    p.add_argument("--out-dir", type=Path, default=DEFAULT_RL_ADAPTER_DIR)
    p.add_argument("--max-steps", type=int, default=40)
    p.add_argument("--smoke", action="store_true")

    p = sub.add_parser("inspect", help="print a few packed examples (markers / think / bash)")
    p.add_argument("--pack", type=Path, required=True)
    p.add_argument("--n", type=int, default=8)

    p = sub.add_parser(
        "judge",
        help="official GLM-5.2 checklist on a finished local duel (needs OpenRouter)",
    )
    p.add_argument(
        "--run",
        type=Path,
        default=Path("/workspace/data/eval-runs/20260824T182728Z-cdd1388b"),
        help="eval-run directory (default: v11 8x12)",
    )
    p.add_argument("--limit", type=int, default=2)
    p.add_argument("--dry-run", action="store_true")

    args = parser.parse_args(argv)
    if args.cmd == "pack":
        from .pack import pack

        pack(
            dataset_root=args.dataset_root,
            out_dir=args.out_dir,
            max_examples=64 if args.smoke else args.max_examples,
            seed="smoke" if args.smoke else args.seed,
            n_salts=args.n_salts,
            submit_frac=args.submit_frac,
            edit_frac=args.edit_frac,
            expand_submit_salts=args.expand_submit_salts,
        )
        return 0
    if args.cmd == "sft":
        _maybe_accelerate()
        from .sft import default_pack, train_sft

        train_sft(
            pack_path=args.pack or default_pack(),
            output_dir=args.out_dir,
            model_dir=args.model,
            max_seq_len=args.max_seq_len,
            max_steps=args.max_steps,
            num_epochs=args.epochs,
            lr=args.lr,
            smoke=args.smoke,
        )
        return 0
    if args.cmd == "export":
        from .export import export_merged

        export_merged(
            adapter_dir=args.adapter,
            genesis_dir=args.genesis,
            out_dir=args.out_dir,
            base_dir=args.base,
        )
        return 0
    if args.cmd == "check":
        from .check import check_model

        report = check_model(args.path, args.genesis)
        return 0 if report["ok"] else 1
    if args.cmd == "duel":
        from .duel import duel

        report = duel(
            challenger=args.challenger,
            king=args.king,
            samples=args.samples,
            turns=args.turns,
            seed=args.seed,
            skip_king=args.skip_king,
        )
        return 0 if report["go"] else 2
    if args.cmd == "chain-pack":
        from .chain_pack import pack_chain

        mix = args.mix_pack
        if mix is None:
            candidate = args.out_dir / "sft-4000-sft-pack.jsonl"
            mix = candidate if candidate.is_file() else None
        pack_chain(
            dataset_root=args.dataset_root,
            out_dir=args.out_dir,
            max_groups=8 if args.smoke else args.max_groups,
            seed="smoke" if args.smoke else args.seed,
            n_salts=args.n_salts,
            mix_pack=mix,
            mix_n=0 if args.smoke else args.mix_n,
            mix_edit_n=0 if args.smoke else args.mix_edit_n,
            mix_submit_n=0 if args.smoke else args.mix_submit_n,
        )
        return 0
    if args.cmd == "chain-gold":
        from .chain_gold import pack_chain_gold

        mix = args.mix_pack
        if mix is None:
            candidate = args.out_dir / "sft-4000-sft-pack.jsonl"
            mix = candidate if candidate.is_file() else None
        pack_chain_gold(
            dataset_root=args.dataset_root,
            out_dir=args.out_dir,
            max_groups=4 if args.smoke else args.max_groups,
            seed="smoke" if args.smoke else args.seed,
            n_salts=args.n_salts,
            mix_pack=mix,
            mix_edit_n=0 if args.smoke else args.mix_edit_n,
            mix_submit_n=0 if args.smoke else args.mix_submit_n,
        )
        return 0
    if args.cmd == "chain":
        from .chain_eval import run_chain

        report = run_chain(
            challenger=args.challenger,
            samples=args.samples,
            turns=args.turns,
            seed=args.seed,
            live_gate=not args.no_live_gate,
        )
        return 0 if report["go"] else 2
    if args.cmd == "dpo-pack":
        from .dpo_pack import pack_dpo

        pack_dpo(out_dir=args.out_dir, seed=args.seed, narrow=args.narrow)
        return 0
    if args.cmd == "dpo":
        _maybe_accelerate()
        from .dpo import train_dpo

        train_dpo(
            pack_path=args.pack,
            output_dir=args.out_dir,
            model_dir=args.model,
            max_seq_len=args.max_seq_len,
            max_steps=args.max_steps,
            lr=args.lr,
            beta=args.beta,
            smoke=args.smoke,
        )
        return 0
    if args.cmd == "policy":
        from .policy_eval import run_policy

        report = run_policy(challenger=args.challenger)
        return 0 if report["go"] else 2
    if args.cmd == "rl":
        _maybe_accelerate()
        from .rl import train_grpo

        train_grpo(
            pack_path=args.pack,
            model_dir=args.model,
            output_dir=args.out_dir,
            max_steps=args.max_steps,
            smoke=args.smoke,
        )
        return 0
    if args.cmd == "inspect":
        return _inspect(args.pack, args.n)
    if args.cmd == "judge":
        from local_eval.judge import run_judge

        result = run_judge(run_dir=args.run, 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 _inspect(pack_path: Path, n: int) -> int:
    from albedo_eval_service.shared.observation_format import first_bash_block
    from albedo_eval_service.shared.submit_protocol import ANY_MARKER_RE

    rows = [
        json.loads(line)
        for line in Path(pack_path).read_text().splitlines()
        if line.strip()
    ]
    print(f"n={len(rows)} showing {min(n, len(rows))}", flush=True)
    for row in rows[:n]:
        completion = row.get("completion") or ""
        think_ok = "<think>" in completion and "</think>" in completion
        bash = first_bash_block(completion)
        print("=" * 72, flush=True)
        print(
            f"{row.get('sample_id')} kind={row.get('kind')} phase={row.get('phase')} "
            f"rewrite={row.get('rewrite_mode')} marker={row.get('submit_marker')!r}",
            flush=True,
        )
        print(f"  think_closed={think_ok} bash={bash[:140]!r}", flush=True)
        if row.get("kind") == "submit":
            print(f"  submit={row.get('submit_command')!r} marker_in_bash={bool(ANY_MARKER_RE.search(bash or ''))}", flush=True)
    return 0


def _maybe_accelerate() -> None:
    """Re-exec under accelerate when several GPUs are free and we are not already launched."""
    if os.environ.get("LOCAL_RANK") is not None or os.environ.get("ALBEDO_NO_ACCELERATE"):
        return
    try:
        import torch

        nproc = torch.cuda.device_count()
    except Exception:
        return
    if nproc <= 1:
        return
    cmd = [
        sys.executable,
        "-m",
        "accelerate.commands.launch",
        "--num_processes",
        str(nproc),
        "--mixed_precision",
        "bf16",
        "-m",
        "local_train",
        *sys.argv[1:],
    ]
    print("launch:", " ".join(cmd), flush=True)
    raise SystemExit(subprocess.call(cmd))