repro_bundle/code/jobsrc/job_eval.py

# /// script
# requires-python = ">=3.11"
# dependencies = ["torch", "numpy", "scipy", "huggingface_hub"]
# ///
"""Fallback: evaluate already-trained L-SR1 models sitting in the output bucket
(runs only eval_analytic + efficiency, no training). Use if the main job times
out before its in-job evaluation runs."""
import os
import subprocess
import sys

CODE = os.environ.get("CODE", "/code")
OUT = os.environ.get("OUT", "/out")
DEVICE = os.environ.get("DEVICE", "cuda")
EVAL_QUICK = os.environ.get("EVAL_QUICK", "0")


def run(cmd):
    print(">>>", " ".join(map(str, cmd)), flush=True)
    print("<<< exit", subprocess.run(cmd, cwd=CODE).returncode, flush=True)


run([sys.executable, f"{CODE}/eval_analytic.py", "--device", DEVICE,
     "--quick", EVAL_QUICK,
     "--quad2", f"{OUT}/quad2tr_proj1.pt",
     "--quad2-noproj", f"{OUT}/quad2tr_proj0.pt",
     "--quad100", f"{OUT}/quad100_proj1.pt",
     "--rosen100", f"{OUT}/rosen100_proj1.pt",
     "--rastr100", f"{OUT}/rastr100_proj1.pt",
     "--out", f"{OUT}/eval.json"])
run([sys.executable, f"{CODE}/efficiency.py", "--device", DEVICE,
     "--batch", "256", "--dim", "256", "--buffer", "4",
     "--out", f"{OUT}/efficiency.json"])
print("eval-only job done; outputs:", sorted(os.listdir(OUT)), flush=True)