from argparse import ArgumentParser
from args_main import parse_arguments_main
path_to_monofonic_binary = "/home/aubin/monofonic/build/monofonIC"
def register_arguments_slurm(parser:ArgumentParser):
"""
Register the arguments for the SLURM parameters used in the slurm submissions scripts for the different binary calls.
Binary calls:
- monofonic
- simbelmyne
- scola
"""
parser.add_argument("-smf","--slurm_monofonic", type=str, default=None, help="Path to the monofonic SLURM submission script template.")
parser.add_argument("-ssbmy","--slurm_simbelmyne", type=str, default=None, help="Path to the simbelmyne SLURM submission script template.")
parser.add_argument("-sscola","--slurm_scola", type=str, default=None, help="Path to the scola SLURM submission script template.")
parser.add_argument("--slurm_logs", type=str, default=None, help="Path to the directory where the SLURM logs will be saved.")
parser.add_argument("--slurm_scripts", type=str, default=None, help="Path to the directory where the SLURM scripts will be saved.")
def parse_arguments_slurm(parsed_args):
"""
Parse the arguments for the SLURM parameters used in the slurm submissions scripts for the different binary calls.
"""
from pathlib import Path
main_dict = parse_arguments_main(parsed_args)
slurm_dict = dict(
monofonic_template=parsed_args.slurm_monofonic,
simbelmyne_template=parsed_args.slurm_simbelmyne,
scola_template=parsed_args.slurm_scola,
logs=parsed_args.slurm_logs,
scripts=parsed_args.slurm_scripts
)
if slurm_dict["monofonic_template"] is None:
slurm_dict["monofonic_template"]=main_dict["paramdir"]+"slurm_monofonic.template"
if slurm_dict["simbelmyne_template"] is None:
slurm_dict["simbelmyne_template"]=main_dict["paramdir"]+"slurm_simbelmyne.template"
if slurm_dict["scola_template"] is None:
slurm_dict["scola_template"]=main_dict["paramdir"]+"slurm_scola.template"
if slurm_dict["logs"] is None:
slurm_dict["logs"]=main_dict["directory"]+"slurm_logs/"
if slurm_dict["scripts"] is None:
slurm_dict["scripts"]=main_dict["directory"]+"slurm_scripts/"
Path(slurm_dict["logs"]).mkdir(parents=True, exist_ok=True)
Path(slurm_dict["scripts"]).mkdir(parents=True, exist_ok=True)
return slurm_dict
def create_slurm_template(
slurm_template:str,
job_name:str,
ntasks:int,
nthreads:int,
partition:str,
time:str,
mem:int,
log_out:str,
log_err:str,
array:tuple|None=None,
):
"""
Creates a SLURM submission script template.
"""
with open(slurm_template, "w") as f:
f.write("#!/bin/bash\n")
f.write(f"#SBATCH --job-name={job_name}\n")
f.write(f"#SBATCH --ntasks={ntasks}\n")
f.write(f"#SBATCH --cpus-per-task={nthreads}\n")
f.write(f"#SBATCH --partition={partition}\n")
f.write(f"#SBATCH --time={time}\n")
f.write(f"#SBATCH --mem={mem}G\n")
if array is not None:
f.write(f"#SBATCH --array={array[0]}-{array[1]}\n")
f.write(f"#SBATCH --output={log_out}%x_%a_%A.out\n")
f.write(f"#SBATCH --error={log_err}%x_%a_%A.err\n")
else:
f.write(f"#SBATCH --output={log_out}%x_%j.out\n")
f.write(f"#SBATCH --error={log_err}%x_%j.err\n")
f.write("\n")
f.write("echo '################## SLURM VARIABLES ##################'\n")
f.write("echo SLURM_JOB_ID: $SLURM_JOB_ID\n")
f.write("echo SLURM_JOB_NAME: $SLURM_JOB_NAME\n")
f.write("echo SLURM_JOB_NODELIST: $SLURM_JOB_NODELIST\n")
f.write("echo SLURM_NNODES: $SLURM_NNODES\n")
f.write("echo SLURM_NTASKS: $SLURM_NTASKS\n")
f.write("echo SLURM_CPUS_PER_TASK: $SLURM_CPUS_PER_TASK\n")
f.write("echo SLURM_JOB_CPUS_PER_NODE: $SLURM_JOB_CPUS_PER_NODE\n")
f.write("echo SLURM_MEM_PER_CPU: $SLURM_MEM_PER_CPU\n")
f.write("echo SLURM_MEM_PER_NODE: $SLURM_MEM_PER_NODE\n")
f.write("echo '#####################################################'\n")
f.write("\n\n")
f.write(f"export OMP_NUM_THREADS={nthreads}\n\n")
f.write("start=`date +%s`\n")
f.write("%COMMAND%\n")
f.write("end=`date +%s`\n")
f.write("runtime=$((end-start))\n")
f.write("\n")
f.write("echo ''\n")
f.write("echo ''\n")
f.write("echo \"################## RUNTIME ##################\"\n")
f.write("echo \"Runtime was $runtime seconds\"\n")
#TODO: fix the two lines below
# f.write("echo \"Runtime per task was $((${runtime}.0/${SLURM_NTASKS}.0)) seconds\"\n") # not working yet...
# f.write("echo \"Runtime per total cpus was $((${runtime}.0/${($SLURM_NTASKS*$SLURM_CPUS_PER_TASK)}.0)) seconds\"\n")
f.write("echo \"#############################################\"\n")
f.write("\n")
f.write("exit 0\n")
def create_slurm_script(slurm_template:str,
slurm_script:str,
job:str,
job_config_file:str,
job_log:str,
array:tuple|None=None,
job_name:str|None=None,
):
"""
Creates a SLURM submission script based on the provided template.
For three different kind of jobs:
- monofonic
- simbelmyne
- scola
"""
if array is not None and job != "scola":
raise ValueError(f"Array job range provided for job type {job}.")
if array is None and job == "scola":
raise ValueError(f"Array job range not provided for job type {job}.")
from os.path import isfile
if not isfile(slurm_template):
raise FileNotFoundError(f"SLURM template {slurm_template} does not exist.")
# Copy template content
with open(slurm_template, "r") as f:
template = f.readlines()
command_line = ""
# Add the job command
match job:
case "monofonic":
command_line = f"{path_to_monofonic_binary} {job_config_file} > {job_log}"
case "simbelmyne":
command_line = f"{job} {job_config_file} {job_log}"
case "scola":
command_line = f"{job} {job_config_file} {job_log} "+"-b ${SLURM_ARRAY_TASK_ID}"
case _:
raise ValueError(f"Job type {job} not recognized.")
# Create the script file
with open(slurm_script, "w") as f:
for line in template:
if job_name is not None and "--job-name" in line:
line = f"#SBATCH --job-name={job_name}\n"
if array is not None and "--array" in line:
line = f"#SBATCH --array={array[0]}-{array[1]}\n"
if array is not None and ("--output" in line or "--error" in line):
line = line.replace("%j","%a_%A")
if "%COMMAND%" in line:
line = command_line+"\n"
f.write(line)
if __name__ == "__main__":
from argparse import ArgumentParser
parser = ArgumentParser(description="Generate slurm submission templates.")
parser.add_argument("-j","--job", type=str, default="monofonic", help="Job type: monofonic, simbelmyne, scola.")
parser.add_argument("-N","--ntasks", type=int, default=1, help="Number of tasks.")
parser.add_argument("-n","--nthreads", type=int, default=32, help="Number of threads per task.")
parser.add_argument("-p","--partition", type=str, default="comp,pscomp,compl", help="Partition to use.")
parser.add_argument("-t","--time", type=str, default="0-00:10:00", help="Time limit.")
parser.add_argument("-m","--mem", type=int, default=64, help="Memory limit.")
parser.add_argument("-d", "--directory", type=str, default="./", help="Main directory where the output will be saved (if other dir and filenames are not specified).")
parser.add_argument("-o","--log_out", type=str, default=None, help="File root for the output logs.")
parser.add_argument("-e","--log_err", type=str, default=None, help="File root for the error logs.")
parser.add_argument("-a","--array", type=int, nargs=2, default=None, help="Array job range.")
parser.add_argument("-s","--slurm_template", type=str, default=None, help="Path to the SLURM template.")
parser.add_argument("-jn","--job_name", type=str, default=None, help="Job name.")
parsed_args = parser.parse_args()
job_name = parsed_args.job if parsed_args.job_name is None else parsed_args.job_name
slurm_template = parsed_args.slurm_template if parsed_args.slurm_template is not None else f"{parsed_args.directory}params/slurm_{job_name}.template"
log_out = parsed_args.log_out if parsed_args.log_out is not None else f"{parsed_args.directory}slurm_logs/{job_name}_"
log_err = parsed_args.log_err if parsed_args.log_err is not None else f"{parsed_args.directory}slurm_logs/{job_name}_"
create_slurm_template(
slurm_template=slurm_template,
job_name=job_name,
ntasks=parsed_args.ntasks,
nthreads=parsed_args.nthreads,
partition=parsed_args.partition,
time=parsed_args.time,
mem=parsed_args.mem,
log_out=log_out,
log_err=log_err,
array=parsed_args.array,
)