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()