A script to run scripts via qsub

run_qsub.py

from __future__ import print_function

import os
import re
import sys
import random
import datetime
import argparse
import subprocess

parser = argparse.ArgumentParser(description='Run a script via qsub')
parser.add_argument('script', type=str, help='script to run')
parser.add_argument('--gpu', default=False, type=bool, help='use gpu?')
parser.add_argument('--mem', default=10, type=int, help='GB of mem')
parser.add_argument('--ram', default=10, type=int, help='GB of ram')
parser.add_argument('--host', default="any", type=str, help='Force a host other than default gpu/cpu host')
parser.add_argument('--v', default=1, type=int, help="How verbose to be? 0: No prompt 1: Prompt before submitting 2: Prompt and immediately run 'watch qstat -j <jobnum>'")
parser.add_argument('--threads', default=1, type=int, help="How many threads to use?")
parser.add_argument('--cwd', default=False, type=bool, help="Run with cwd specified?")

args = parser.parse_args()

io = "-o $HOME/jobs -e $HOME/jobs"

resources = " -l mem_free=%dG,ram_free=%dG" % (args.mem, args.ram)
if args.host != 'any':
    resources += ',hostname=%s' % args.host
if args.gpu:
    resources += ",gpu=1"
    if args.host == 'any':
        resources += ",hostname=b1[12345678]*"

    resources += " -q gpu.q,all.q,g.q"

if args.threads > 1:
    resources += " -pe smp {}".format(args.threads) 

if args.v > 0:
    gpu_str = "" if not args.gpu else ", AND A GPU!"
    print("Going to run %s on %s host with %d threads, %sG RAM, %sG memory%s." % (
            args.script, args.host, args.threads, args.ram, args.mem, gpu_str
        ))
    if str(input("Continue? [y/n] ")).lower()[:1] == "n":
        sys.exit(1)

cmd = "qsub %s %s %s" % (resources, io, args.script)
print(cmd)
output = subprocess.check_output(cmd.split()).decode('utf8')
print(output)

success = re.compile(r"Your job (\d+) \([^\)]*\) has been submitted")
if success.search(output):
    now = str(datetime.datetime.now())
    jobnum = success.search(output).group(1)
    with open("$HOME/jobs/hist.txt", 'a') as outf:
        outf.write("%s %s %s/%s\n" % (now, jobnum, os.getcwd(), args.script))
    if args.v > 1:
        try:
            os.system("watch qstat -j %s" % jobnum)
        except KeyboardInterrupt:
            pass