social-data-attribution-w2/scripts/unlearning/submit_eval_faithful.py

#!/usr/bin/env python3
"""Submit the throttled OLMES eval array over the faithful multi-seed checkpoints.

Rebuilds the manifest first (build_faithful_eval_manifest.py: excludes already-evaluated
dirs AND dirs whose training job is still live), then submits --array=1-N%THROTTLE.
Idempotent across rounds — re-run after each training wave to catch up new checkpoints.

Usage: python submit_eval_faithful.py [--dry]
"""

import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
MULTISEED = ROOT / "scripts/analysis/multiseed"
HOME = Path.home()
MANIFEST = HOME / "scratch/n16_selectivity/ms_eval_faithful_all.tsv"
BASE = (
    HOME
    / "scratch/hf_cache/hub/models--allenai--OLMo-3-1025-7B/snapshots/a81bae42db3975be1671e27b9c9a56da1a9f980f"
)
LOG = "/storage/ice-shared/cs7634/staff/TDA/logs/tom_unlearn"
SBATCH = ROOT / "scripts/slurm/analysis/multiseed/eval_array.sbatch"
# max concurrent eval tasks (HX00 GPUs). Default 4 when training is sharing GPUs;
# pass --throttle N for the post-training catch-up eval when GPUs are free.
THROTTLE = 4
if "--throttle" in sys.argv:
    THROTTLE = int(sys.argv[sys.argv.index("--throttle") + 1])

r = subprocess.run(
    [sys.executable, str(MULTISEED / "build_faithful_eval_manifest.py")],
    capture_output=True,
    text=True,
)
print(r.stdout.strip())
if r.returncode != 0:
    print("manifest build FAILED:", r.stderr.strip()[:300])
    sys.exit(1)

n = sum(1 for line in MANIFEST.read_text().splitlines() if line.strip())
if n == 0:
    print("nothing to eval (manifest empty)")
    sys.exit(0)

cmd = [
    "sbatch",
    "--partition=coe-gpu,ice-gpu",
    "--qos=coe-ice",
    "--account=ic",
    "--gres=gpu:1",
    "--constraint=HX00",
    "--cpus-per-task=8",
    "--mem=128G",
    "--time=02:00:00",
    f"--array=1-{n}%{THROTTLE}",
    f"--output={LOG}/oevarr_%A_%a.out",
    f"--error={LOG}/oevarr_%A_%a.err",
    f"--export=ALL,MANIFEST={MANIFEST},BASE_LOCAL={BASE},HFOFFLINE=0",
    str(SBATCH),
]

if "--dry" in sys.argv:
    print("WOULD SUBMIT:\n  " + " ".join(cmd))
    sys.exit(0)

out = subprocess.run(cmd, capture_output=True, text=True)
print(
    f"submitted array of {n} (throttle {THROTTLE}):",
    out.stdout.strip() or out.stderr.strip(),
)