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

#!/usr/bin/env python3
"""Submit the self-service multi-seed unlearning grid (seeds 42,43,44) using bin-level
influence-targeted forget-sets (expA/expC) + the random sampler (exp1/exp3).

Each recipe x seed -> one seed_unlearn.sbatch job (NGDiff r8 LoRA, PPL early-stop).
Idempotent: skips a (recipe, seed) whose output dir already has a checkpoint.

Usage:
  python submit_multiseed.py --headline [--dry]   # curated subset first
  python submit_multiseed.py --full     [--dry]   # all 24 topics x 4 benches + expC + exp1 + exp3
"""

import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
HOME = Path.home()
FG = HOME / "scratch/n16_selectivity/forgetsets_binlevel"
MS = HOME / "scratch/n16_selectivity/multiseed"
SBATCH = ROOT / "scripts/slurm/analysis/multiseed/seed_unlearn.sbatch"
BASE = (
    HOME
    / "scratch/hf_cache/hub/models--allenai--OLMo-3-1025-7B/snapshots/a81bae42db3975be1671e27b9c9a56da1a9f980f"
)
LOG = "/storage/ice-shared/cs7634/staff/TDA/logs/tom_unlearn"
SEEDS = [42, 43, 44]
BENCHES = ["socialiqa", "mmlu_social_science", "mmlu_stem", "arc_challenge"]
TOPICS = [
    "adult_content",
    "art_and_design",
    "crime_and_law",
    "education_and_jobs",
    "electronics_and_hardware",
    "entertainment",
    "fashion_and_beauty",
    "finance_and_business",
    "food_and_dining",
    "games",
    "health",
    "history_and_geography",
    "home_and_hobbies",
    "industrial",
    "literature",
    "politics",
    "religion",
    "science_math_and_technology",
    "social_life",
    "software",
    "software_development",
    "sports_and_fitness",
    "transportation",
    "travel_and_tourism",
]
# Headline topics that drive the published claims (head-to-head top-1, pilot picks, hindsight best, gate).
HEADLINE_TOPICS = [
    "industrial",
    "science_math_and_technology",
    "social_life",
    "finance_and_business",
    "software_development",
    "sports_and_fitness",
    "health",
    "electronics_and_hardware",
    "literature",
]


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


def recipes(headline: bool):
    topics = HEADLINE_TOPICS if headline else TOPICS
    out = []  # (label, topic_bin, forget_ids_or_None, outdir_stem)
    for bench in BENCHES:
        for topic in topics:
            fg = FG / f"expA__{topic}__{bench}.txt"
            if fg.exists():
                out.append(
                    (f"expA_{topic}_{bench}", topic, str(fg), f"expA_{topic}_{bench}")
                )
        out.append(
            (f"expC_{bench}", "null", str(FG / f"expC__{bench}.txt"), f"expC_{bench}")
        )
    for topic in topics:
        out.append((f"exp1_{topic}", topic, None, f"exp1_{topic}"))  # random-in-topic
    out.append(("exp3_null", "null", None, "exp3_null"))  # random-global
    return out


def main():
    headline = "--full" not in sys.argv
    dry = "--dry" in sys.argv
    recs = recipes(headline)
    n_sub = n_skip = 0
    for label, topic_bin, forget_ids, stem in recs:
        for seed in SEEDS:
            outdir = MS / f"{stem}_seed{seed}"
            if has_ckpt(outdir):
                n_skip += 1
                continue
            env = (
                f"ALL,TOPIC_BIN={topic_bin},OUTPUT_DIR={outdir},SEED={seed},MAXWALL=340"
            )
            if forget_ids:
                env += f",FORGET_IDS={forget_ids}"
            cmd = [
                "sbatch",
                "--parsable",
                "--partition=coe-gpu,ice-gpu",
                "--qos=coe-ice",
                "--account=ic",
                "--gres=gpu:1",
                "--constraint=HX00",
                "--cpus-per-task=8",
                "--mem=128G",
                "--time=06:00:00",
                f"--job-name=ms_{label}_s{seed}",
                f"--output={LOG}/ms_{label}_s{seed}_%j.out",
                f"--error={LOG}/ms_{label}_s{seed}_%j.err",
                f"--export={env}",
                str(SBATCH),
            ]
            if dry:
                print(
                    f"WOULD SUBMIT {label} seed{seed} -> {outdir.name}"
                    + (" [forget_ids]" if forget_ids else " [random]")
                )
                n_sub += 1
                continue
            r = subprocess.run(cmd, capture_output=True, text=True)
            if r.returncode == 0:
                n_sub += 1
            else:
                print(f"FAIL {label} s{seed}: {r.stderr.strip()[:150]}")
    scope = "HEADLINE" if headline else "FULL"
    print(
        f"\n[{scope}] submit={n_sub} skip(existing)={n_skip} recipes={len(recs)} seeds={SEEDS}"
    )


if __name__ == "__main__":
    main()