198 lines
7.5 KiB
Python
198 lines
7.5 KiB
Python
import argparse
|
|
import os
|
|
import sys
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
# Constants
|
|
VALID_HARDWARE_CONFIGS = {
|
|
"clara": {"gpus": ["v100", "rtx2080ti"]},
|
|
"paula": {"gpus": ["a30"]},
|
|
}
|
|
RAM_PER_GPU = 128 # in GB
|
|
|
|
|
|
def parse_unknown_args_to_kwargs(unknown_args):
|
|
kwargs = {}
|
|
key = None
|
|
|
|
for arg in unknown_args:
|
|
if arg.startswith('--'):
|
|
if '=' in arg:
|
|
k, v = arg[2:].split('=', 1)
|
|
kwargs[k.replace('-', '_')] = v
|
|
key = None
|
|
else:
|
|
key = arg[2:].replace('-', '_')
|
|
kwargs[key] = True # Might be a flag if no value follows
|
|
elif key:
|
|
kwargs[key] = arg
|
|
key = None
|
|
else:
|
|
# Handle or log unexpected positional args if desired
|
|
pass
|
|
|
|
return kwargs
|
|
|
|
|
|
def validate_args(args):
|
|
if args.run_type not in ["slurm", "local"]:
|
|
sys.exit("Error: run_type must be 'slurm' or 'local'.")
|
|
|
|
if args.run_type == "slurm":
|
|
# validate partition and hardware
|
|
if args.partition not in VALID_HARDWARE_CONFIGS:
|
|
sys.exit(
|
|
f"Error: Invalid partition '{args.partition}'. Valid partitions are: {', '.join(VALID_HARDWARE_CONFIGS.keys())}.")
|
|
|
|
if args.gpu_type not in VALID_HARDWARE_CONFIGS[args.partition]["gpus"]:
|
|
sys.exit(
|
|
f"Error: Invalid GPU type '{args.gpu_type}' for partition '{args.partition}'. Valid GPU types are: {', '.join(VALID_HARDWARE_CONFIGS[args.partition]['gpus'])}.")
|
|
|
|
if args.num_gpus <= 0 or args.num_cpus_per_gpu <= 0:
|
|
sys.exit("Error: Number of GPUs and CPUs per GPU must be positive integers.")
|
|
|
|
if args.item_limit < -1:
|
|
sys.exit("Error: Item limit must be -1 or a positive integer.")
|
|
|
|
|
|
def compose_sbatch_script(run_config,
|
|
partition,
|
|
gpu_type,
|
|
num_gpus,
|
|
ram_per_gpu,
|
|
num_cpus_per_gpu,
|
|
item_limit,
|
|
time_limit,
|
|
evaluate_only=False,
|
|
**kwargs):
|
|
run_config_name, run_name = run_config.replace(".py", "").split("/")[-2:]
|
|
num_cpus = num_gpus * num_cpus_per_gpu
|
|
total_ram = num_gpus * ram_per_gpu
|
|
|
|
current_dir = Path(__file__).resolve().parent
|
|
# venv is two layers up from this file
|
|
venv_path = current_dir.parents[1] / "venv/bin/activate"
|
|
|
|
base_log_dir = "/work/rr41qemu-MA/logs"
|
|
|
|
print(f"Preparing to submit job for {run_name} from {run_config_name}...")
|
|
print(f"Partition: {partition}\nGPU type: {gpu_type}\nNumber of GPUs: {num_gpus}")
|
|
print(f"CPUs per GPU: {num_cpus_per_gpu} (Total: {num_cpus})\nTotal RAM: {total_ram} GB")
|
|
print(f"Item limit: {item_limit}\n")
|
|
|
|
sbatch_script = f"""#!/bin/bash
|
|
#SBATCH --job-name={run_config_name}_{run_name}
|
|
#SBATCH --output={base_log_dir}/%x_%j.out
|
|
#SBATCH --error={base_log_dir}/%x_%j.err
|
|
#SBATCH --time={time_limit}
|
|
#SBATCH --ntasks=1
|
|
#SBATCH --nodes=1
|
|
#SBATCH --ntasks-per-node=1
|
|
#SBATCH --cpus-per-task={num_cpus}
|
|
#SBATCH --mem={total_ram}G
|
|
#SBATCH --partition={partition}
|
|
#SBATCH --gpus={gpu_type}:{num_gpus}
|
|
|
|
echo "Loading python virtual environment..."
|
|
source {venv_path}
|
|
|
|
echo "Loading python 3.10..."
|
|
module load Python/3.10.4-GCCcore-11.3.0
|
|
|
|
cd {current_dir}
|
|
torchrun --nproc_per_node={num_gpus} --rdzv_backend=c10d --rdzv_endpoint=localhost:0 training_wrapper.py {run_config} --item_limit {item_limit} --num_dataloader_workers {num_cpus_per_gpu} {'--evaluate_only' if evaluate_only else ''}
|
|
"""
|
|
return sbatch_script
|
|
|
|
|
|
def compose_local_script(run_config,
|
|
num_gpus,
|
|
num_cpus_per_gpu,
|
|
item_limit,
|
|
evaluate_only=False,
|
|
**kwargs):
|
|
run_config_name, run_name = run_config.replace(".py", "").split("/")[-2:]
|
|
print(f"Running {run_name} from {run_config_name} with {num_gpus} GPUs and {num_cpus_per_gpu} CPUs per GPU.")
|
|
print(f"Number of GPUs: {num_gpus}\nCPUs per GPU: {num_cpus_per_gpu} (Total: {num_gpus * num_cpus_per_gpu})")
|
|
print(f"Item limit: {item_limit}\n")
|
|
|
|
local_script = (f"torchrun --nproc_per_node={num_gpus} --rdzv_backend=c10d --rdzv_endpoint=localhost:0 "
|
|
f"training_wrapper.py {run_config} --item_limit {item_limit}"
|
|
f" --num_dataloader_workers {num_cpus_per_gpu} {'--evaluate_only' if evaluate_only else ''}")
|
|
return local_script
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Submit a SLURM training job.")
|
|
parser.add_argument("run_type",
|
|
choices=["slurm", "local"],
|
|
help="Type of job to submit: 'slurm' or 'local'")
|
|
parser.add_argument("run_config", help="Path to run configuration module")
|
|
parser.add_argument("--partition",
|
|
required=False,
|
|
help="Partition to submit job to")
|
|
parser.add_argument("--gpu_type",
|
|
required=False,
|
|
help="GPU type to request")
|
|
parser.add_argument("--num_gpus",
|
|
required=False,
|
|
default=1,
|
|
type=int, help="Number of GPUs to request")
|
|
parser.add_argument("--num_cpus_per_gpu",
|
|
required=False,
|
|
default=2,
|
|
type=int, help="Number of CPUs per GPU, also number of dataloaders per gpu, defaults to to 2")
|
|
parser.add_argument("--ram_per_gpu",
|
|
required=False,
|
|
default=True,
|
|
type=int,
|
|
help="Number of RAM per GPU, also number of dataloaders per gpu, defaults to to True")
|
|
parser.add_argument("--item_limit",
|
|
required=False,
|
|
default=-1,
|
|
type=int, help="Item limit for the run (-1 for no limit)")
|
|
parser.add_argument("--time_limit",
|
|
required=False,
|
|
type=str,
|
|
default="2-0", # 1 day, 0 hours
|
|
help="Time limit for the job (e.g. 01:00:00)")
|
|
parser.add_argument("--evaluate_only",
|
|
action="store_true",
|
|
help="If set, only run evaluation without training.")
|
|
|
|
args, unknown_args = parser.parse_known_args()
|
|
# parse unknown args for proper kwargs usage
|
|
unknown_args = parse_unknown_args_to_kwargs(unknown_args)
|
|
validate_args(args)
|
|
|
|
if args.run_type == "local":
|
|
script = compose_local_script(
|
|
run_config=args.run_config,
|
|
num_gpus=args.num_gpus,
|
|
num_cpus_per_gpu=args.num_cpus_per_gpu,
|
|
item_limit=args.item_limit,
|
|
evaluate_only=args.evaluate_only,
|
|
)
|
|
print("Local script generated:\n", script)
|
|
result = subprocess.run(script, shell=True, check=True)
|
|
else:
|
|
script = compose_sbatch_script(
|
|
run_config=args.run_config,
|
|
partition=args.partition,
|
|
gpu_type=args.gpu_type,
|
|
ram_per_gpu=args.ram_per_gpu,
|
|
num_gpus=args.num_gpus,
|
|
num_cpus_per_gpu=args.num_cpus_per_gpu,
|
|
item_limit=args.item_limit,
|
|
time_limit=args.time_limit,
|
|
evaluate_only=args.evaluate_only,
|
|
)
|
|
result = subprocess.run(["sbatch"], input=script.encode(), check=True)
|
|
print("Job submitted.")
|
|
print(script)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|