loadl

#!/usr/bin/env python3

import sys
import argparse
import subprocess
from loadleveller import jobfile
import os

parser = argparse.ArgumentParser(description='Helper script for running and managing loadleveller Monte Carlo jobs.', usage='''loadl <command> <jobscript> [<args>]

<jobscript> is an executable that prints the job parameter YAML-file to stdout. It is convenient to use the taskmaker python module for this purpose.

Possible commands and their shorthands are
    delete, d  delete all data related to a job
    merge, m   merges results of an unfinished job into an output file
    run, r     runs the job
    status, s  print job completion information''')
    
parser.add_argument('command')
parser.add_argument('jobfile')

args, leftover_args = parser.parse_known_args()

# all paths are relative to the jobscript
jobdir = os.path.dirname(args.jobfile)
if jobdir != '':
    print('$ cd '+jobdir)
    os.chdir(jobdir)

try:
    job = jobfile.JobFile(os.path.basename(args.jobfile))
except jobfile.JobFileGenError as e:
    print('Error: {}'.format(e))
    sys.exit(1)

def run():
    import glob
    from loadleveller import clusterutils
    
    parser = argparse.ArgumentParser(description='run a loadleveller job on a cluster or locally')

    parser.add_argument('-s','--single', action='store_true', help='run in the single core scheduler mode')
    parser.add_argument('-f', '--force', action='store_true', help='ignore warnings about possible job corruption')
    parser.add_argument('-r', '--restart', action='store_true', help='delete all existing job data before starting.')
    args_run = parser.parse_args(leftover_args)

    if args_run.restart:
        delete()
    else:
        # check age of the different files
        binary_modtime = os.stat(job.jobconfig['mc_binary']).st_mtime
        try:
            f = next(glob.iglob('{}.data/*/*.h5'.format(job.jobname))) # only check one of the output files for speed
            data_modtime = os.stat(f).st_mtime

            label = 'Warning' if args_run.force else 'Error'
            if binary_modtime > data_modtime:
                print('{}: binary \'{}\' is newer than the checkpoint files.'.format(label, job.jobconfig['mc_binary']))
                if not args_run.force:
                    print('Use \'--restart\' to start from a blank run or use \'--force\' to proceed if you are sure\nthe changes you made are compatible.')
                    sys.exit(1)
        except StopIteration:
            pass

    job_input_filename = job.write_job_input_file()
    if args_run.single:
        cmd = '{} single "{}"'.format(job.jobconfig['mc_binary'], job_input_filename)
        print('$ '+cmd)
        subprocess.run(cmd, shell=True)
    else:
        clusterutils.run(job.jobname, job.jobconfig, [job.jobconfig['mc_binary'], job_input_filename])

def delete():
    import shutil
    datadir = '{}.data'.format(job.jobname)
    results_file = '{}.results.yml'.format(job.jobname)

    if os.path.exists(datadir):
        print('$ rm -r {}'.format(datadir))
        shutil.rmtree(datadir)
    if os.path.exists(results_file):
        print('$ rm {}'.format(results_file))
        os.unlink(results_file)

def merge():
    job_input_filename = job.write_job_input_file()
    cmd = '{} merge "{}"'.format(job.jobconfig['mc_binary'], job_input_filename)
    print('$ '+cmd)
    subprocess.run(cmd, shell=True)

def status():
    from loadleveller import jobstatus
    rc = jobstatus.print_status(job, leftover_args)
    sys.exit(rc)


if args.command == 'delete' or args.command == 'd':
    delete()
elif args.command == 'merge' or args.command == 'm':
    merge()
elif args.command == 'run' or args.command == 'r':
    run()
elif args.command == 'status' or args.command == 's':
    status()
else:
    print('Unknown command \'{}\'.'.format(args.command))
    parser.print_help()
    sys.exit(1)