Submit jobs to qsub, the easy way.

qsub.py

#! /usr/bin/env python3
from __future__ import print_function

import argparse
from collections import OrderedDict
from datetime import datetime
from getpass import getuser
import os
import re
import readline
import shlex
import subprocess
import sys

HIDDEN_DUMMY_FILE = f"/home/{getuser()}/.email"

parser = argparse.ArgumentParser(description='Run a script via qsub')
parser.add_argument('--gpu', action='store_true', help='use gpu?')

parser.add_argument('--from_home', dest='here', action='store_false', help="Run from home folder (instead of this folder)?")
parser.add_argument('--host', type=str, help='Force a host other than default gpu/cpu host')
parser.add_argument('--mem', default=10, type=int, help='GB of mem')
parser.add_argument('--threads', default=1, type=int, help="How many threads to use?")
parser.add_argument('--name', type=str, help="Name that will appear when calling `qstat`")

verbosity = parser.add_mutually_exclusive_group()
verbosity.add_argument('-v', '--verbose', action='store_const',
                       const=2, default=1)
verbosity.add_argument('-q', '--quiet', dest='verbose',
                       action='store_const', const=0)

parser.add_argument('script', type=str, help='script to run')
parser.add_argument('script_args', nargs=argparse.REMAINDER)

args = parser.parse_args()

_resources = None  # HACKY AS HECK
def determine_resources(args: argparse.Namespace) -> str:
    resources = OrderedDict()
    # Memory
    resources['mem_free'] = f"{args.mem}G"
    resources['ram_free'] = resources['mem_free']  # Grid requires this twice.
    # Host
    if args.host:
        resources['hostname'] = args.host
    # GPU
    if args.gpu:
        resources['gpu'] = 1
    #if args.gpu and not args.host:
    #    resources['hostname'] = 'c*'  # Because c nodes have fastest GPUs.
    global _resources
    _resources = resources
    return ",".join(f"{k}={v}" for k, v in resources.items())

resources = determine_resources(args)

opts = OrderedDict()

opts['-l'] = resources
# opts['-j'] = 'y'  # Merge output and error files.
if args.gpu:
    opts['-q'] = ",".join(["gpu.q", "all.q", "g.q"])
if args.threads > 1:
    opts['-pe'] = f'smp {args.threads}'
if args.here:
    opts['-cwd'] = ""
if args.name:
    opts['-N'] = args.name

try:
    with open(HIDDEN_DUMMY_FILE) as f:
        email = next(f)
except FileNotFoundError:
    email_regex = re.compile("(^[a-zA-Z0-9_.+-]+@[a-zA-Z0-9-]+\.[a-zA-Z0-9-.]+$)")
    # Prompt until you get a valid email.
    while True:
        data = input("**config: Enter your email address: ")
        res = email_regex.search(data)
        if res:
            with open(HIDDEN_DUMMY_FILE, 'w') as f:
                email = res.group()
                f.write(email)
            break

opts['-m'] = 'e'  # End. (Didn't include 'a' for abort or 'b' for begin.)
opts['-M'] = email

options = " ".join(f"{k} {v}" for k, v in opts.items())

cmd = f"qsub {options} -terse {args.script} {' '.join(args.script_args)}"

if args.verbose > 0:
    gpu_str = "." if not args.gpu else ", AND A GPU!"
    host_string = "anywhere" if 'hostname' not in _resources else f"on {_resources['hostname']}"
    to_print = f"\tGoing to run {args.script} {host_string} with {args.threads} threads, {args.mem} GB memory{gpu_str}"
    print("\x1b[1mPreview:\x1b[0m", cmd)
    with open(HIDDEN_DUMMY_FILE) as f:
        email = next(f)
        to_print += f"\n\tWill email {email} about progress."
        print(to_print)
    if input("Continue? [Y/n] ").lower()[:1] == 'n':
        sys.stderr.write("Aborted.\n")
        sys.exit(1)
else:
    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.now())
    jobnum = success.search(output).group(1)
    with open(f"/home/{getuser()}/hist.txt", 'a+') as outf:
        outf.write(f"[{now}] {jobnum} {os.getcwd()}/{args.script}\n")
    if args.verbose > 1:
        try:
            os.system(f"watch qstat -j {jobnum}")
        except KeyboardInterrupt:
            pass