#!/usr/bin/env python3 """ EGFR-VHL ablation driver. Deduplicated layout: 3 baseline runs (weights = 1,1,1,1,1, seeds 42-44) 5 scorers x 2 non-baseline weights {0,2} x 3 seeds = 30 runs 3 acceptance=on runs (baseline weights, seeds 42-44) total = 36 runs The baseline runs cover every (scorer, w=1) cell at once. The acceptance=off case is the baseline (same runs reused), so only acceptance=on is run here. Each run = 100 samples on EGFR-VHL. Outputs go to generation_logs/ablation_egfr//. Resumable (skips done .smi files). """ import argparse import itertools import os import subprocess import time from pathlib import Path # ---------------------------------------------------------------------------- # Config # ---------------------------------------------------------------------------- SCRIPT_DIR = Path(__file__).resolve().parent.parent REPO_ROOT = SCRIPT_DIR.parent CHECKPOINT = REPO_ROOT / "checkpoints" / "protac_linker_finetune" / "epoch=99-v1.ckpt" CLF_PATH = REPO_ROOT / "src" / "classifier" / "protac_linker_clf.pkl" GENERATE = SCRIPT_DIR / "generate.py" OUT_BASE = REPO_ROOT / "generation_logs" / "ablation_egfr" PROTAC_MODULE = "ErlotinibVHL" N_SAMPLES = 100 N_STEPS = 1500 WARMUP_STEPS = 100 SCORER_NAMES = ["classifier", "fast", "sa", "chem", "conn"] # obj index 0..4 WEIGHTS = [0.0, 1.0, 2.0, 4.0] # 1.0 is baseline (shared) NON_BASELINE = [w for w in WEIGHTS if w != 1.0] # [0.0, 2.0, 4.0] SEEDS = [42, 43, 44] BASELINE_W = [1.0, 1.0, 1.0, 1.0, 1.0] def weight_vec(scorer_idx: int, w: float) -> list[float]: """Baseline (1,1,1,1,1) with one slot replaced.""" v = list(BASELINE_W) v[scorer_idx] = w return v def build_runs() -> list[dict]: """Return list of run configs. Each config has a unique tag.""" runs = [] # Shared baseline (weights = 1,1,1,1,1, acceptance off). Reused for all # (scorer, w=1) cells in the main ablation AND for the acceptance=off cell. for seed in SEEDS: runs.append({ "tag": f"egfr_baseline_seed{seed}", "weights": list(BASELINE_W), "seed": seed, "accept": False, }) # Main ablation: per-scorer weight sweep, non-baseline weights only for idx, name in enumerate(SCORER_NAMES): for w in NON_BASELINE: for seed in SEEDS: runs.append({ "tag": f"egfr_{name}_w{w:g}_seed{seed}", "weights": weight_vec(idx, w), "seed": seed, "accept": False, }) # Acceptance=on (the "off" cell is the baseline above) for seed in SEEDS: runs.append({ "tag": f"egfr_accept_on_seed{seed}", "weights": list(BASELINE_W), "seed": seed, "accept": True, }) return runs def run_command(run: dict, gpu: int) -> list[str]: out_dir = OUT_BASE / run["tag"] out_dir.mkdir(parents=True, exist_ok=True) cmd = [ "python", str(GENERATE), "--checkpoint", str(CHECKPOINT), "--protac_module", PROTAC_MODULE, "--fix_termini", "cooh", "--fast_scoring", "--clf_path", str(CLF_PATH), "--num_samples", str(N_SAMPLES), "--num_batches", "1", "--num_steps", str(N_STEPS), "--warmup_steps", str(WARMUP_STEPS), "--weights", *[str(w) for w in run["weights"]], "--seed", str(run["seed"]), "--output_file", str(out_dir / f"{run['tag']}.smi"), "--output_dir", str(out_dir), ] if run["accept"]: cmd.append("--use_acceptance") return cmd def is_done(run: dict) -> bool: out_dir = OUT_BASE / run["tag"] smi = out_dir / f"{run['tag']}.smi" if not smi.exists(): return False # crude check: non-empty file means generation reached the "save SMILES" step return smi.stat().st_size > 0 def launch(cmd: list[str], gpu: int, log_path: Path): env = os.environ.copy() env["CUDA_VISIBLE_DEVICES"] = str(gpu) with open(log_path, "w") as logf: return subprocess.Popen(cmd, env=env, stdout=logf, stderr=subprocess.STDOUT) def main(): ap = argparse.ArgumentParser() ap.add_argument("--gpus", default="2,3,4,5,6", help="comma-separated GPU IDs to use (default: 2,3,4,5,6)") ap.add_argument("--stagger", type=float, default=2.0, help="seconds between successive run starts (avoid CSV collisions)") ap.add_argument("--dry-run", action="store_true", help="print commands without launching") args = ap.parse_args() gpu_list = [int(g) for g in args.gpus.split(",")] runs = build_runs() todo = [r for r in runs if not is_done(r)] done = [r for r in runs if is_done(r)] print(f"Total runs: {len(runs)}") print(f"Already done: {len(done)}") print(f"To launch: {len(todo)}") print(f"GPUs available: {gpu_list}") if args.dry_run: for r in todo[:5]: cmd = run_command(r, gpu_list[0]) print("[DRY]", " ".join(cmd)) if len(todo) > 5: print(f"[DRY] ... and {len(todo)-5} more") return OUT_BASE.mkdir(parents=True, exist_ok=True) active: dict[int, dict] = {} # gpu -> (proc, run, log_path, start_ts) queue = list(todo) def dispatch(gpu: int): if not queue: return False run = queue.pop(0) cmd = run_command(run, gpu) log_path = OUT_BASE / run["tag"] / "run.log" proc = launch(cmd, gpu, log_path) active[gpu] = {"proc": proc, "run": run, "log": log_path, "start": time.time()} print(f"[GPU {gpu}] start {run['tag']} (pid {proc.pid})") time.sleep(args.stagger) return True # Initial fill for gpu in gpu_list: dispatch(gpu) # Poll while active or queue: for gpu, info in list(active.items()): ret = info["proc"].poll() if ret is not None: elapsed = time.time() - info["start"] status = "OK" if ret == 0 else f"FAIL({ret})" print(f"[GPU {gpu}] done {info['run']['tag']} " f"({status}, {elapsed/60:.1f} min)") del active[gpu] if queue: dispatch(gpu) time.sleep(10) print(f"\nAll runs finished. Outputs under {OUT_BASE}/") if __name__ == "__main__": main()