#!/usr/bin/env python
#
# A script to submit jobs to the SLURM queuing system
#
# F. Nedelec, 10.2007 --- 29.09.2018
# Adapted from LSF to SLURM on 17.02.2017

"""
    Submit an array of jobs to the SLURM system, to be handled by 'go_sim.py'
    
Syntax:
    
    submit.py ARG [mem=????] [queue=????] [hours=INT] [days=INT] [cpu=INT] file1 [file2] [file3] [...]
    
    ARG is passed as arguments to 'go_sim.py' (please check go_sim.py help),
    and you must use quotes if you have multiple arguments to group them together.
    
    Unless specified otherwise, the queue is 'medium_priority'.
    The amount of requested memory (default=2GB) should be specified in MB:
       mem=1024 (for 1 GB)
       mem=512  (for 512 MB)
       ...
    
Examples:
    
    1. submit.py sim config.cym
       Submit one job to run the config file provided
    
    2. submit.py sim config1.cym config2.cym config3.cym ...
       Submit one job for each config file provided. They will run in parallel.

    3. submit.py 'sim preconfig.py' config.cym.tpl
       Submit one job where the config file will be generated by preconfig.py from config.cym.tpl
       The executable will run sequentially on one node, for each file that is generated.
    
    4. submit.py 'sim 10' config.cym
       Submit one job to run the same config file 10 times sequentially on one node.
    
    5. submit.py 'sim 10' config1.cym config2.cym config3.cym ...
       Submit a different job for each config file provided, with 10 repeats of each
    
    
Last updated 26.09.2018
F. Nedelec
"""


import sys, os, shutil, subprocess

# default parameters for submission:
subcmd  = 'sbatch'
queue   = 'htc'
runtime = '5-00:00:00' # 5 days
memory  = '2096'       # in MB
ncpu    = 1 # nb of threads per job

# parameters of the program:
jdir    = 'job00'
out     = sys.stderr

#-------------------------------------------------------------------------------

def execute(cmd):
    """execute given command with subprocess.call()"""
    try:
        val = subprocess.call(cmd)
        if val:
            out.write("ERROR: command failed with value %i\n" % val)
            print(cmd)
    except OSError:
        out.write("ERROR: command failed:\n")
        print('> '+' '.join(cmd))


def makeNumberedDirectory(root):
    """Create an empty directories that start with 'root'"""
    cnt = 0
    while cnt < 10000:
        name = root + '%02i' % cnt
        try:
            os.mkdir(name)
            return name
        except OSError:
            cnt += 1
    out.write("Error: directory '%s' could not be created \n" % root)
    return ''


def write_script(filename, cmd):
    """create an executable file containing the commands"""
    fid = open(filename, 'w')
    fid.write("#!/bin/bash\n")
    for s in cmd:
        fid.write(s+'\n')
    fid.close()
    os.chmod(filename, 0700)

#------------------------------------------------------------------------

def job(cwd, conf, jarg):
    """return bash script that will run one simulation"""
    cmd  = ['module purge']
    cmd += ['module load foss']
    cmd += ['module load LAPACK OpenBLAS']
    cmd += ['cd %s;' % cwd]
    cmd += ['touch %s;' % conf]
    # the job will call go_sim.py once:
    cmd += ['python go_sim.py %s %s;' % (jarg, conf)]
    cmd += ['mv '+conf+' '+jdir+'/done/.;']
    # cleanup
    cmd += ['mv $0 '+jdir+'/done/.;']
    return cmd


def sub(exe):
    """return command that will submit one job"""
    # specify memory, shell, minimum number of cores and queue
    cmd  = [subcmd, '--nodes=1', '--ntasks=1']
    # specify number of threads if executable is threaded:
    if ncpu > 1:
        cmd += ['--cpus-per-task=%i' % ncpu]
    cmd += ['--job-name='+jdir]
    cmd += ['--partition='+queue] 
    cmd += ['--time='+runtime] 
    cmd += ['--mem='+memory]
    # define signals sent if time is exceeded:
    cmd += ['--signal=15@120']
    cmd += ['--signal=2@60']
    # request special hardware:
    #cmd += ['--constraint=avx2']
    # redirect stderr and sdtout to files:
    cmd += ['--output='+jdir+'/logs/out']
    cmd += ['--error='+jdir+'/logs/err']
    # call executable:
    cmd += [exe]
    return cmd


def array(jobcnt):
    """return command that will submit a job-array"""
    # define parameters directly in the script:
    cmd  = ['#SBATCH --nodes=1']
    cmd += ['#SBATCH --ntasks=1']
    # specify number of threads if executable is threaded:
    if ncpu > 1:
        cmd += ['--cpus-per-task=%i' % ncpu]
    cmd += ['#SBATCH --partition='+queue]
    cmd += ['#SBATCH --time='+runtime]
    cmd += ['#SBATCH --mem='+memory]
    # define signals sent if time is exceeded:
    cmd += ['#SBATCH --signal=INT@60']
    cmd += ['#SBATCH --signal=TERM@120']
    # request special hardware:
    #cmd += ['#SBATCH --constraint=avx2']
    # redirect stderr and sdtout to files:
    cmd += ['#SBATCH --output='+jdir+'/logs/%a.out']
    cmd += ['#SBATCH --error='+jdir+'/logs/%a.err']
    # defines the job name and index range:
    cmd += ['#SBATCH --job-name='+jdir]
    cmd += ['#SBATCH --array=1-%i' % jobcnt]
    cmd += ['']
    cmd += [jdir+'/todo/R$SLURM_ARRAY_TASK_ID']
    return cmd

#-------------------------------------------------------------------------------

def main(args):
    """submit jobs, depending on the arguments provided"""
    global subcmd, memory, runtime, queue, jdir, ncpu
    
    #find subcmd command:
    proc = subprocess.Popen(['which', subcmd], stdout=subprocess.PIPE)
    if proc.wait():
        out.write("Error: submit command `"+subcmd+"' not found!\n")
    else:
        subcmd = proc.stdout.readline().strip()
        #print('|'+subcmd+'|')
    
    # first argument is used for go_sim.py:
    jarg = args.pop(0)
    
    # catch old-style invocation
    if jarg.endswith(".cym"):
        print(__doc__)
        sys.exit()
    
    # make new directories for this job
    jdir = makeNumberedDirectory('job')
    os.mkdir(os.path.join(jdir, 'todo'))
    os.mkdir(os.path.join(jdir, 'done'))
    os.mkdir(os.path.join(jdir, 'save'))
    os.mkdir(os.path.join(jdir, 'logs'))

    print("    go_sim.py will run `%s' in %s" % (jarg, jdir))

    jdup  = 1
    jcnt  = 0
    jname = ''
    cwd   = os.getcwd()
    todo  = os.path.join(jdir,'todo')
    
    for arg in args:
        if os.path.isfile(arg):
            if os.access(arg, os.X_OK):
                out.write("Error: file `%s' should not be executable\n" % arg)
                sys.exit()
            for x in range(jdup):
                jcnt += 1
                conf = todo + '/config%04i.cym' % jcnt
                shutil.copyfile(arg, conf)
                jname = todo + '/R' + str(jcnt)
                cmd = job(cwd, conf, jarg+' name=run%04i'%jcnt+' park='+jdir+'/save')
                write_script(jname, cmd)
        elif arg.isdigit():
            jdup = int(arg)
        elif arg.startswith('mem='):
            memory = arg[4:]
        elif arg.startswith('memory='):
            memory = arg[7:]
        elif arg.startswith('cpu='):
            ncpu = arg[4:]
        elif arg.startswith('ncpu='):
            ncpu = arg[5:]
        elif arg.startswith('day='):
            runtime = arg[4:]+'-00:00:00'
        elif arg.startswith('days='):
            runtime = arg[5:]+'-00:00:00'
        elif arg.startswith('hour='):
            runtime = arg[5:]+':00:00'
        elif arg.startswith('hours='):
            runtime = arg[6:]+':00:00'
        elif arg.startswith('minute='):
            runtime = arg[7:]+':00'
        elif arg.startswith('minutes='):
            runtime = arg[8:]+':00'
        elif arg.startswith('time='):
            runtime = arg[5:]
        elif arg.startswith('queue='):
            queue = arg[6:]
        else:
            out.write("Error: I do not understand argument `%s'\n" % arg)
            sys.exit()

    if memory < 128:
        out.write("Error: requested memory (%s MB) seems too low\n" % memory)
        sys.exit()

    if ncpu < 1:
        out.write("Error: number of cpu/job must be >= 1\n")
        sys.exit()

    if jcnt < 1:
        out.write("Error: you need to specify at least one config file\n" % arg)
        sys.exit()

    if jcnt > 1:
        print("    submit.py created %i scripts in `%s'" % (jcnt, todo))
        cmd = array(jcnt)
        # create script file:
        name = todo + '/job.bash';
        write_script(name, cmd)
        # make command to submit this script:
        cmd = (subcmd, name)
    else:
        cmd = sub(jname);
    execute(cmd)


#-------------------------------------------------------------------------------

if __name__ == "__main__":
    if len(sys.argv) < 2 or sys.argv[1].endswith("help"):
        print(__doc__)
    else:
        main(sys.argv[1:])

