GridEngine Array Job Support Tool

arrayrunner.py

#!/usr/bin/env python

import argparse
import os
import subprocess
import shlex
import sys

def _main():
    """
    """

    parser = argparse.ArgumentParser(description='Array Job Runner. Try all argument combinations')
    parser.add_argument('-c', '--command', default='./array_runner.py -h {1} {2} {3} {4}', help='Job command')
    parser.add_argument('-r', '--run', action='store_true', help='Run command')
    parser.add_argument('-1', dest='a1', nargs='*', help='Argument 1')
    parser.add_argument('-2', dest='a2', nargs='*', help='Argument 2')
    parser.add_argument('-3', dest='a3', nargs='*', help='Argument 3')
    parser.add_argument('-4', dest='a4', nargs='*', help='Argument 4')
    parser.add_argument('-5', dest='a5', nargs='*', help='Argument 5')
    parser.add_argument('-6', dest='a6', nargs='*', help='Argument 6')
    parser.add_argument('-7', dest='a7', nargs='*', help='Argument 7')
    parser.add_argument('-8', dest='a8', nargs='*', help='Argument 8')
    parser.add_argument('-n', '--no-combination', action='store_true')
    parser.add_argument('-s', '--submit', action='store_true', help='Submit to GridEngine')
    parser.add_argument('-m', '--memory', help='Memory requirements (For GridEngine)', default="2G")
    parser.add_argument('-p', '--parallel', type=int, help="# of threads used in one job (For GridEngine)")
    options = parser.parse_args()

    count = 0
    exit_status = 0

    if options.no_combination:
        if options.a8:
            combinations = zip(options.a1, options.a2, options.a3, options.a4, options.a5, options.a6, options.a7, options.a8)
        elif options.a7:
            combinations = zip(options.a1, options.a2, options.a3, options.a4, options.a5, options.a6, options.a7)
        elif options.a6:
            combinations = zip(options.a1, options.a2, options.a3, options.a4, options.a5, options.a6)
        elif options.a5:
            combinations = zip(options.a1, options.a2, options.a3, options.a4, options.a5)
        elif options.a4:
            combinations = zip(options.a1, options.a2, options.a3, options.a4)
        elif options.a3:
            combinations = zip(options.a1, options.a2, options.a3)
        elif options.a2:
            combinations = zip(options.a1, options.a2)
        else:
            combinations = zip(options.a1)
    else:
        combinations = list(make_combination(options.a1, options.a2, options.a3, options.a4, options.a5, options.a6, options.a7, options.a8))
    #print combinations

    for one in slice_by_jobnumber(combinations):
        #print one
        print options.command.format(*([None]+list(one)))
        if options.run and not options.submit:
            this_exit = subprocess.call(shlex.split(options.command.format(*([None]+list(one)))))
            #this_exit = os.system(options.command.format(*([None]+one)))
            print this_exit
            exit_status = max(exit_status, this_exit)
        count += 1
    if 'SGE_TASK_ID' not in os.environ:
        print 'Count:', count
    if options.submit:
        additional_options = ['-cwd', '-S', '/bin/sh']
        if options.parallel and options.parallel > 1:
            additional_options += ["-pe", "def_slot", str(options.parallel)]
        if options.memory:
            additional_options += ['-l', 's_vmem={0},mem_req={0}'.format(options.memory)]
        
        script = '#!/bin/sh\n\n' + \
                 "PATH="+os.environ['PATH']+"\n"+ \
                 'export PATH\n' + \
                 "LD_LIBRARY_PATH="+os.environ['LD_LIBRARY_PATH']+"\n"+ \
                 'export LD_LIBRARY_PATH\n' + \
                 ' '.join(["python"] + ['"{}"'.format(x.replace('"', '\"')) for x in sys.argv if x != '-s'] + ['-r'])
        print script
        
        sp = subprocess.Popen(['qsub', '-t', '1-{}'.format(count)] + additional_options, stdin=subprocess.PIPE)
        sp.communicate(script)
    
    #exit(100)
    exit(exit_status)

def make_combination(*args):
    """
    
    Arguments:
    - `*args`:
    """

    #print args

    if not len(args) or args[0] is None:
        yield []
        return

    for one in args[0]:
        for two in make_combination(*args[1:]):
            yield [one] + two
    

def get_job_number():
    """
    """

    job_number = 0
    number_of_jobs = 1
    if 'SGE_TASK_ID' in os.environ:
        array_task_id = os.environ['SGE_TASK_ID']
        array_task_last = os.environ['SGE_TASK_LAST']
        #array_task_first = os.environ['SGE_TASK_FIRST']
        try:
            job_number = int(array_task_id)-1
            number_of_jobs = int(array_task_last)
        except:
            job_number = 0
            number_of_jobs = 1
    return (job_number,number_of_jobs)

def slice_by_jobnumber(sequence):
    """
    
    Arguments:
    - `sequence`:
    """

    num, jobs = get_job_number()
    size = float(len(sequence))/jobs
    if num != jobs-1:
        return sequence[int(size*num):int(size*(num+1))]
    return sequence[int(size*num):]


if __name__ == '__main__':
    _main()