sbmy_control/slurm_submission.py

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