eval_runner.py

import subprocess
import itertools
import os
import signal
import sys

if __name__ == "__main__":
    eval_types = [
        "core",
        "sparse_probing",
        # "absorption",
        # "unlearning",
    ]
    model_names = [
        "smollm2-135m"
        
        # "pythia-160m",

        # "pythia-410m",
        # "gemma-2-2b",
        # "llama-1b"
    ]
    # CUDA_VISIBLE_DEVICES=i python -m sae_bench.custom_saes.run_all_evals_custom_sae <model> <eval>
    def run_eval(model, eval_type, device):
        cmd = f"CUDA_VISIBLE_DEVICES={device} python -m sae_bench.custom_saes.run_all_evals_custom_sae {model} {eval_type}"
        subprocess.run(cmd, shell=True)

    def get_free_device(available_devices):
        # Return a free device from the set of available devices
        return available_devices.pop()

    available_devices = {5, 7}
    processes = []

    def handle_interrupt(signal, frame):
        print("Interrupt received, terminating processes...")
        for proc, device in processes:
            proc.terminate()
        sys.exit(0)

    signal.signal(signal.SIGINT, handle_interrupt)

    for model, eval_type in itertools.product(model_names, eval_types):
        if not available_devices:
            # Wait for any process to finish before starting a new one
            while not available_devices:
                for proc, device in processes:
                    if proc.poll() is not None:
                        processes.remove((proc, device))
                        available_devices.add(device)
                        break

        device = get_free_device(available_devices)
        p = subprocess.Popen(
            f"CUDA_VISIBLE_DEVICES={device} .venv/bin/python -m sae_bench.custom_saes.run_all_evals_custom_saes {model} {eval_type}",
            shell=True)
        processes.append((p, device))

    # Wait for all processes to finish
    for proc, device in processes:
        proc.wait()