"""Full benchmarks of the REAP-pruned checkpoints with apex_moe.lmeval, the dense full-run recipe:
MATH-500, GPQA-Diamond, HumanEval+, MBPP+ with thinking and a 64k budget, MMLU-Pro without thinking
and a 16k budget; seeds 0, 31323, 67210 (the dense full runs' seeds); tag full-reap-official; no routes saved. Qwen3.6 (NVFP4) jobs take
one idle GPU, Qwen3.8 (FP8, TP 2) jobs any two idle GPUs (NVSwitch: any pair works), DeepSeek-V4.1-Flash
(TP 4) any four, with its full-run recipe (deep_gemm, batch 36, MMLU-Pro batch 128,
VLLM_USE_BREAKABLE_CUDAGRAPH=0); --gpus limits the queue to some GPUs; a model's jobs start once its config exists
(make_config.py). Jobs with an ok ledger row or a live process (a restarted queue) are skipped; a
failed job is retried once. Seed 0 of every model first, in order: while an earlier two-GPU job
waits for a second idle GPU, later jobs do not take the idle one. Prints every launch and exit.

    cd apex-moe && setsid nohup .venv/bin/python external/reap/experiments/apex/eval_queue.py \
        > data/logs/reap-official-eval-queue.log 2>&1 &
"""

import json
import os
import subprocess
import time
from pathlib import Path

ROOT = Path("/localhome/mkleineg/dgx-moe/apex-moe")
TAG = "full-reap-official"
MODELS = {  # config name -> (GPUs per job, batch)
    "qwen3.6-35b-a3b-nvfp4-reapofficial-keep50": (1, 76),
    "qwen3.6-35b-a3b-nvfp4-reapofficial-keep25": (1, 76),
    "qwen3.8-flash-next-fp8-reapofficial-keep50": (2, 80),
    "qwen3.8-flash-next-fp8-reapofficial-keep25": (2, 80),
    "deepseek-v4.1-flash-reapofficial-keep50": (4, 36),
    "deepseek-v4.1-flash-reapofficial-keep25": (4, 36),
}
DEEPSEEK = {  # the DeepSeek full runs' extra settings (configs/models/deepseek-v4.1-flash.yaml)
    "sets": ["model.model_args.moe_backend=deep_gemm"],
    "mmlu_batch": 128,
    "env": {"VLLM_USE_BREAKABLE_CUDAGRAPH": "0"},
}
TASKS = {
    "minerva": "minerva_math500",
    "gpqa": "gpqa_diamond_cot_zeroshot",
    "humaneval": "humaneval_plus_think",
    "mbpp": "mbpp_plus_think",
    "mmlu": "mmlu_pro",
}
SEEDS = (0, 31323, 67210)


def gpu_memory() -> dict[int, int]:
    out = subprocess.run(
        ["nvidia-smi", "--query-gpu=index,memory.used", "--format=csv,noheader,nounits"],
        capture_output=True, text=True, check=True,
    ).stdout
    return {int(i): int(m) for i, m in (line.split(",") for line in out.strip().splitlines())}


def done() -> set[tuple[str, int]]:
    ok = set()
    for line in (ROOT / "experiments/ledger.jsonl").open():
        row = json.loads(line)
        if row.get("tag") == TAG and row.get("status") == "ok":
            ok.add((row["run_id"].rsplit("-", 2)[0], row["seed"]))
    return ok


def live() -> set[tuple[str, str, int]]:
    """(model, task, seed) of every apex_moe.lmeval process of this tag already running."""
    out = set()
    for proc in Path("/proc").iterdir():
        try:
            argv = (proc / "cmdline").read_bytes().decode().split("\0")
        except (OSError, ValueError):
            continue
        if "apex_moe.lmeval" not in argv or TAG not in argv:
            continue
        model = Path(argv[argv.index("--model") + 1]).stem
        task = next(k for k, v in TASKS.items() if v == argv[argv.index("--tasks") + 1])
        out.add((model, task, int(argv[argv.index("--seed") + 1])))
    return out


def name_of(model: str, task: str) -> str:
    return f"{model}-lmeval-{task}-docs20000-{'nothink-' if task == 'mmlu' else ''}dense"


def command(model: str, task: str, seed: int) -> list[str]:
    batch = MODELS[model][1]
    deepseek = model.startswith("deepseek")
    mmlu_batch = DEEPSEEK["mmlu_batch"] if deepseek else 140
    cmd = [
        ".venv/bin/python", "-m", "apex_moe.lmeval", "--model", f"configs/models/{model}.yaml",
        "--docs", "20000", "--seed", str(seed), "--tasks", TASKS[task], "--tag", TAG,
    ]
    if task == "mmlu":
        cmd += [
            "--set", "model.model_args.enable_thinking=false",
            "--set", "model.gen_kwargs.chat_template_kwargs.enable_thinking=false",
            "--set", "model.gen_kwargs.max_gen_toks=16384",
            "--set", "model.model_args.max_model_len=18432",
            "--set", f"model.model_args.max_num_seqs={mmlu_batch}",
            "--set", f"name={name_of(model, task)}",
            "--note", f"official REAP pruned, full MMLU-Pro without thinking, batched {mmlu_batch}",
        ]
    else:
        cmd += [
            "--set", "model.gen_kwargs.max_gen_toks=65536",
            "--set", "model.model_args.max_model_len=69632",
            "--set", f"model.model_args.max_num_seqs={batch}",
            "--note", f"official REAP pruned, full test set, 64k budget, batched {batch}",
        ]
        if task in ("humaneval", "mbpp"):
            cmd.append("--allow-code-execution")
    if deepseek:
        for setting in DEEPSEEK["sets"]:
            cmd += ["--set", setting]
    return cmd


def main() -> None:
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--gpus", default="0,1,2,3,4,5,6,7", help="GPUs the queue may use")
    allowed = {int(g) for g in parser.parse_args().gpus.split(",")}
    os.chdir(ROOT)
    jobs = [(m, t, s) for s in SEEDS for m in MODELS for t in TASKS]
    running: dict[tuple, tuple[subprocess.Popen, tuple[int, ...]]] = {}
    failures: dict[tuple, int] = {}
    print(f"{time.strftime('%T')} queue: {len(jobs)} jobs", flush=True)
    while True:
        for job, (proc, gpus) in list(running.items()):
            if proc.poll() is not None:
                print(f"{time.strftime('%T')} exit {proc.returncode}: {job} on GPUs {gpus}", flush=True)
                if proc.returncode != 0:
                    failures[job] = failures.get(job, 0) + 1
                del running[job]
        finished = done()
        elsewhere = live() - set(running)
        pending = [
            j for j in jobs
            if (name_of(j[0], j[1]), j[2]) not in finished and j not in running and failures.get(j, 0) < 2
            and j not in elsewhere
            and (ROOT / f"configs/models/{j[0]}.yaml").exists()
        ]
        waiting = [j for j in jobs if not (ROOT / f"configs/models/{j[0]}.yaml").exists()]
        if not pending and not running and not waiting and not elsewhere:
            break
        busy = {g for _, gpus in running.values() for g in gpus}
        memory = gpu_memory()
        idle = [g for g, used in sorted(memory.items()) if g in allowed and g not in busy and used < 2000]
        for job in pending:
            width = MODELS[job[0]][0]
            gpus = tuple(idle[:width]) if len(idle) >= width else None
            if gpus is None:
                break
            idle = [g for g in idle if g not in gpus]
            env = os.environ | {
                "PATH": f"{ROOT}/.venv/bin:{os.environ['PATH']}",
                "OMP_NUM_THREADS": "8",
                "CUDA_VISIBLE_DEVICES": ",".join(map(str, gpus)),
                "APEX_SAVE_ROUTES": "0",
                "HF_ALLOW_CODE_EVAL": "1",
            } | (DEEPSEEK["env"] if job[0].startswith("deepseek") else {})
            log = ROOT / f"data/logs/reapoff-{job[0].split('-reapofficial-')[1]}-{job[0][:8]}-{job[1]}-s{job[2]}.log"
            proc = subprocess.Popen(
                command(*job), env=env, stdout=log.open("w"), stderr=subprocess.STDOUT,
                stdin=subprocess.DEVNULL, start_new_session=True,
            )
            running[job] = (proc, gpus)
            where = ",".join(map(str, gpus))
            print(f"{time.strftime('%T')} GPUs {where}: {job} (PID {proc.pid}) -> {log.name}", flush=True)
            time.sleep(20)
        time.sleep(30)
    print(f"{time.strftime('%T')} queue done; failed twice: {[j for j, n in failures.items() if n >= 2]}", flush=True)


if __name__ == "__main__":
    main()