oargen.py

#!/usr/bin/env python
# coding: utf-8

from __future__ import print_function
from __future__ import unicode_literals

import os
import argparse
import subprocess


def argparser():
    parser = argparse.ArgumentParser()
    # parser.add_argument('command', nargs='?',
    #                     help="Command to use for the job (in passive mode)")
    parser.add_argument('command', nargs=argparse.REMAINDER,
                        help="Command to use for the job (in passive mode)")
    parser.add_argument('-n', '--name',
                        help="Name to give to the job")
    parser.add_argument('-d', '--directory',
                        help="Directory in which will be stored oarsub outputs")
    parser.add_argument('-b', '--besteffort', action="store_true",
                        help="Launch job in besteffort mode")
    parser.add_argument('-t', '--time', default="10",
                        help="Estimated maximum duration of the job (format: h[:m[:s]]) (default: %(default)s)")
    parser.add_argument('-g', '--gpu', action="store_true",
                        help="If True, reserves only cores with GPUs")
    parser.add_argument('-c', '--core', default=1, type=int,
                        help="Number of cores to reserve. (default: %(default)s)")
    parser.add_argument('-H', '--host', nargs="+", default=[],
                        help="Name of the hosts (SQL 'LIKE' syntax accepted)")
    parser.add_argument('-I', '--ignore-host', nargs="+", default=[],
                        help="Name of the hosts to ignore (SQL 'NOT LIKE' syntax accepted)")
    parser.add_argument('-i', '--interactive', action="store_true",
                        help="Launch job in interactive mode")
    parser.add_argument('-C', '--checkpoint', type=int, metavar="SECONDS",
                        help="Enable checkpoint signals with the given delay (in seconds)")
    parser.add_argument('-r', '--run', action="store_true",
                        help="Run the command")
    parser.add_argument('-a', '--anterior',
                        help="Anterior job id that must be terminated to start this new one")
    args = parser.parse_args()

    return args


def prepare_oarsub(gpu, hosts, core, time,
                   ignore_hosts=[],
                   command=None,
                   interactive=False,
                   name=None, output_directory=None,
                   besteffort=False,
                   checkpoint=None, anterior=None):
    oar_cmd = ["oarsub"]
    oar_cmd.append("-p")
    properties = ""
    if gpu:
        properties += "(gpu IS NOT NULL)"
    else:
        properties += "(gpu IS NULL)"
    if hosts:
        properties += " AND ("
        for idx, host in enumerate(hosts):
            if idx != 0:
                properties += " OR "
            properties += "host LIKE '{}'".format(host)
        properties += ")"
    if ignore_hosts:
        for host in ignore_hosts:
            properties += " AND host NOT LIKE '{}'".format(host)
    properties += ""
    oar_cmd.append(properties)

    oar_cmd.append("-l")
    time = time.split(':')
    hour = time[0]
    minutes = time[1] if len(time) >= 2 else "00"
    seconds = time[2] if len(time) >= 3 else "00"
    ressources = "core={},walltime={}:{}:{}".format(core, hour, minutes, seconds)
    oar_cmd.append(ressources)

    if name is not None:
        oar_cmd.append("-n")
        oar_cmd.append(name)
        directory = output_directory + "/" if output_directory is not None else ""
        oar_cmd.append("-O")
        oar_cmd.append("{}{}.%jobid%.stdout".format(directory, name))
        oar_cmd.append("-E")
        oar_cmd.append("{}{}.%jobid%.stderr".format(directory, name))

    if besteffort:
        oar_cmd.extend(["-t", "besteffort", "-t", "idempotent"])

    if checkpoint is not None:
        oar_cmd.extend(["--checkpoint", checkpoint])

    if anterior is not None:
        oar_cmd.extend(["-a", anterior])

    if interactive:
        oar_cmd.append('-I')
    else:
        job_command = command if isinstance(command, str) else " ".join(command)
        oar_cmd.append(job_command)
    return oar_cmd


def run_oarsub(command, print_cmd=False, fake_run=False, return_output=False):
    if print_cmd:
        print(subprocess.list2cmdline(command))
    if fake_run:
        return None
    if not return_output:
        subprocess.call(command)
        return None
    return subprocess.check_output(command).decode("utf8")


def main():
    args = argparser()

    if args.directory is not None and not os.path.isdir(args.directory) and args.run:
            raise RuntimeError("'{}' is not a directory!".format(args.directory))

    oar_command = prepare_oarsub(args.gpu, args.host, args.core, args.time,
                                 ignore_hosts=args.ignore_host,
                                 command=args.command,
                                 interactive=args.interactive,
                                 name=args.name, output_directory=args.directory,
                                 besteffort=args.besteffort, checkpoint=args.checkpoint,
                                 anterior=args.anterior)

    run_oarsub(oar_command, print_cmd=True, fake_run=not args.run)


if __name__ == '__main__':
    main()