social-data-attribution-w2/scripts/slurm/analysis/multiseed/submit_binlevel_arc_sharedvol.py

#!/usr/bin/env python3
"""Bin-level ARC 3-variation REDO (challenge/easy/combined) on the shared volume.

Replaces the broken 3-variation runs whose +forget_ids_file key was silently
ignored (they trained on random forget docs). Forget-sets are injected via the
supported +forget_texts_file mechanism from forget_texts_binlevel_<variation>/
parquets. Passes MIN_TOKENS=0, DOCS_PER_RETAIN_BIN=null, RESAMPLE_INTERVAL=0
so the data pipeline matches the manuscript-era runs (retain = 9000
random-capped, no min-token filter, no retain resampling). Seed 42 only,
matching the original 3-variation design. Outputs go to binlevel_redo/ so the
broken outputs are left untouched. queue/checkpoint-aware, idempotent.
--smoke = expC challenge seed 42 only. --dry = print, don't submit.
"""
import os
import subprocess
import sys
from pathlib import Path

HOME = Path.home()
SHARED = Path("/storage/ice-shared/cs7634/staff/TDA/arc_rescore")
OUT = SHARED / "binlevel_redo"
DOLMA_CACHE = SHARED / "dolma_cache/dolma3_6t_filtered"
SBATCH = HOME / "dev/data-attribution/scripts/slurm/analysis/multiseed/seed_unlearn_sharedvol.sbatch"
LOG = "/storage/ice-shared/cs7634/staff/TDA/logs/tom_unlearn"
SEED = 42
VARIATIONS = ["challenge", "easy", "combined"]
EXCLUDE = os.environ.get(
    "EXCLUDE",
    "atl1-1-03-010-25-0,atl1-1-03-014-16-0,atl1-1-03-011-13-0",
)


def has_ckpt(o: Path) -> bool:
    return bool(list(o.glob("checkpoint-*/adapter_model.safetensors")))


def queued() -> set[str]:
    r = subprocess.run(
        ["squeue", "-u", os.environ.get("USER", "gmatlin3"), "-h", "-o", "%200j"],
        capture_output=True,
        text=True,
    )
    return {ln.strip() for ln in r.stdout.splitlines() if ln.strip().startswith("blr_")}


def recipes():
    out = []  # (variation, recipe, topic_bin, parquet)
    for var in VARIATIONS:
        fgdir = SHARED / f"forget_texts_binlevel_{var}"
        for fg in sorted(fgdir.glob("*.parquet")):
            recipe = fg.stem
            tb = recipe.split("__")[1] if recipe.startswith("expA__") else "null"
            out.append((var, recipe, tb, fg))
    return out


def main():
    dry = "--dry" in sys.argv
    limit = 0
    if "--limit" in sys.argv:
        limit = int(sys.argv[sys.argv.index("--limit") + 1])
    recs = recipes()
    if "--smoke" in sys.argv:
        recs = [r for r in recs if r[0] == "challenge" and r[1].startswith("expC")]
    inq = queued()
    n = skipped = 0
    for var, recipe, tb, fg in recs:
        if limit and n >= limit:
            break
        stem = f"{recipe}_seed{SEED}"
        jobname = f"blr_{var}_{stem}"
        outdir = OUT / var / stem
        if has_ckpt(outdir) or jobname in inq:
            skipped += 1
            continue
        env = (
            f"ALL,TOPIC_BIN={tb},OUTPUT_DIR={outdir},SEED={SEED},MAXWALL=340,"
            f"FORGET_TEXTS={fg},DOLMA_CACHE={DOLMA_CACHE},"
            "MIN_TOKENS=0,DOCS_PER_RETAIN_BIN=null,RESAMPLE_INTERVAL=0"
        )
        cmd = [
            "sbatch", "--parsable", f"--exclude={EXCLUDE}",
            f"--job-name={jobname}",
            f"--output={LOG}/{jobname}_%j.out",
            f"--error={LOG}/{jobname}_%j.err",
            f"--export={env}", str(SBATCH),
        ]
        if dry:
            print(f"WOULD SUBMIT {jobname} -> {outdir}")
            n += 1
            continue
        r = subprocess.run(cmd, capture_output=True, text=True)
        if r.returncode == 0:
            print(f"SUBMITTED {jobname} job={r.stdout.strip()}")
            n += 1
        else:
            print(f"FAIL {jobname}: {r.stderr.strip()[:140]}")
    print(f"submit={n} skipped={skipped} recipes={len(recs)}")


if __name__ == "__main__":
    main()