"""Run with: modal run --profile samuel modal_benchmark.py --gpu L40S"""
import json
from pathlib import Path
import modal

app = modal.App("lfm25-350m-rlcd-benchmark")
image = (modal.Image.debian_slim(python_version="3.11")
         .pip_install_from_requirements("requirements.txt")
         .add_local_python_source("rlcd")
         .add_local_dir("tests", "/root/tests"))
cache = modal.Volume.from_name("lfm25-350m-rlcd-hf-cache", create_if_missing=True)


def execute(suite="diagnostic"):
    import os
    import subprocess
    from rlcd.benchmark import run
    os.environ["RLCD_TEST_DEVICE"] = "cuda"
    test = subprocess.run(["python", "-m", "pytest", "-q", "-s", "/root/tests"], capture_output=True, text=True)
    print(test.stdout, flush=True)
    if test.returncode:
        raise RuntimeError(test.stdout + test.stderr)
    result = run(device="cuda", dtype="float16", warmups=2, repeats=3, suite=suite, max_new_tokens=768 if suite == "stress" else 192)
    result["validation"] = {"pytest_stdout": test.stdout, "pytest_stderr": test.stderr, "returncode": test.returncode}
    return result

@app.function(image=image, gpu="L40S", timeout=3600, volumes={"/cache": cache}, env={"HF_HOME":"/cache"})
def l40s(suite="diagnostic"):
    return execute(suite)

@app.function(image=image, gpu="H100", timeout=3600, volumes={"/cache": cache}, env={"HF_HOME":"/cache"})
def h100(suite="diagnostic"):
    return execute(suite)

@app.local_entrypoint()
def main(gpu: str = "L40S", suite: str = "diagnostic"):
    if suite not in {"diagnostic", "stress"}:
        raise ValueError(suite)
    fn = {"L40S":l40s, "H100":h100}[gpu]
    result = fn.remote(suite)
    path = Path("results") / (gpu.lower() + ("-stress" if suite == "stress" else "") + ".json")
    path.parent.mkdir(exist_ok=True)
    path.write_text(json.dumps(result, indent=2) + "\n")
    print(json.dumps(result["summary"], indent=2))