added code
This commit is contained in:
@@ -0,0 +1,197 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user