bomb

#!/usr/bin/env python

import argparse, sys
import subprocess
from argparse import RawTextHelpFormatter

__author__ = "Colby Chiang (cc2qe@virginia.edu)"
__version__ = "$Revision: 0.0.1 $"
__date__ = "$Date: 2015-01-13 15:15 $"

# --------------------------------------
# define functions

def get_args():
    parser = argparse.ArgumentParser(formatter_class=RawTextHelpFormatter, description="\
bomb\n\
author: " + __author__ + "\n\
version: " + __version__ + "\n\
description: Constructor for bsub jobs")
    parser.add_argument('-J', '--job_name', metavar='STR', type=str, required=False, help='job name [cmd]')
    parser.add_argument('-g', '--job_group', metavar='STR', type=str, required=False, default=None, help='job group')
    parser.add_argument('-q', '--queue', metavar='STR', type=str, required=False, default=None, help='queue')
    parser.add_argument('-t', '--threads', metavar='INT', type=int, required=False, default=1, help='number of threads [1]')
    parser.add_argument('-m', '--mem', metavar='INT', type=int, required=False, default=8, help='memory in gigabytes [8]')
    parser.add_argument('-n', '--num_hosts', metavar='INT', type=int, required=False, default=1, help='number of hosts [1]')
    parser.add_argument('-x', '--exclude_host', metavar='STR', type=str, required=False, default=None, help='comma delimited list of hosts to exclude')
    parser.add_argument('-i', '--include_host', metavar='STR', type=str, required=False, default=None, help='comma delimited list of hosts to include')
    parser.add_argument('-w', '--dependency_expression', metavar='STR', type=str, required=False, default=None, help='dependency expression')
    parser.add_argument('-o', '--out', metavar='STR', type=str, required=False, default='%J.log', help='output log [%%J.log]')
    parser.add_argument('-e', '--err', metavar='STR', type=str, required=False, default='%J.log', help='error log [%%J.log]')
    parser.add_argument('-E', '--email', metavar='STR', type=str, required=False, default=None, help='email address to send logs')
    parser.add_argument('-a', '--additional_esub_information', metavar='STR', type=str, required=False, default=None, help='additional arguments (e.g.: \'docker(user/image:latest)\')')
    parser.add_argument('-d', '--dry_run', action='store_true', required=False, help='print the LSF command and exit')
    parser.add_argument('cmd', metavar='STR', type=str, help='command to submit to LSF')

    # parse the arguments
    args = parser.parse_args()

    # send back the user input
    return args

# primary function
def bsub(args):
    bcmd = map(str, ['bsub',
                     '-n', args.threads,
                     '-M', '%s000000' % args.mem,
                     '-R', '\"select[mem>%s000] rusage[mem=%s000] span[hosts=%s]\"' % (args.mem,
                                                                                args.mem,
                                                                                args.num_hosts),
                     '-J', args.job_name,
                     '-o', args.out,
                     '-e', args.err])

    if args.queue is not None:
        bcmd += ['-q', args.queue]

    if args.job_group is not None:
        bcmd += ['-g', args.job_group]

    if args.email is not None:
        bcmd += ['-N', '-u', args.email]

    if args.include_host is not None:
        for host in args.include_host.split(','):
            bcmd += ['-R', '\"select[hname==\'%s\']\"' % host]

    if args.exclude_host is not None:
        for host in args.exclude_host.split(','):
            bcmd += ['-R', '\"select[hname!=\'%s\']\"' % host]
    
    if args.dependency_expression is not None:
        bcmd += ['-w', '%s' % args.dependency_expression]

    if args.additional_esub_information is not None:
        bcmd += ['-a', '%s' % args.additional_esub_information]

    bcmd += [args.cmd]

    print ' '.join(bcmd);
    if not args.dry_run:
        proc = subprocess.call(bcmd)

    return

# --------------------------------------
# main function

def main():
    # parse the command line args
    args = get_args()

    if args.job_name is None:
        args.job_name = args.cmd

    # call primary function
    bsub(args)

# initialize the script
if __name__ == '__main__':
    try:
        sys.exit(main())
    except IOError, e:
        if e.errno != 32:  # ignore SIGPIPE
            raise