qtools.py

#!/usr/bin/env python

"""
Script for managing PBS queue systems. Only dependencies required are 
xmltodict and tabulate, which are easily pip installable.
"""

import os
import sys
import re
import subprocess
import xmltodict
from tabulate import tabulate
from pathlib import Path


def get_jobs_data():
    output = subprocess.check_output([
        "qstat", "-f", "-u", os.environ["USER"], "-x"]).decode("utf-8")
    d = xmltodict.parse(output)
    jobs = d["Data"]["Job"]
    if not isinstance(jobs, list):
        return [jobs]
    return jobs


def print_info(args):
    jobs = get_jobs_data()
    output = []
    for j in jobs:
        jid = j["Job_Id"].split(".")[0]
        time = j.get("resources_used", {}).get("walltime", 0)
        output.append([
            jid, j["Job_Name"], j["queue"], time, j["job_state"], j["Output_Path"].split(":")[-1]
        ])
    print(tabulate(output, headers=["id", "name", "queue", "walltime", "s", "filepath"]))


def find_bad(args):
    jobs = get_jobs_data()
    inq = []
    for j in jobs:
        p = Path(j["Output_Path"].split(":")[-1])
        inq.append(str(p.parent))
    
    for d in os.listdir(args.path):
        p = Path(args.path) / Path(d)
        if p.is_dir():
            p = p.absolute()
            if str(p) not in inq:
                vasprun = p / args.filename
                if not vasprun.exists():
                    print("%s incomplete!" % p)


def tail(args):
    jobs = get_jobs_data()
    inq = []
    n = args.nlines
    for j in jobs:
        if j.get("job_state", "") == "R":
            p = Path(j["Output_Path"].split(":")[-1])
            p = p.parent
            for parent, _, fns in os.walk(p, followlinks=True):
                for f in fns:
                    if f == args.filename:
                        print(p / parent / f)
                        print("----------")
                        with open(p / parent / f) as f:
                            lines = f.readlines()
                            print("".join(lines[-n:]))
                        print()
                        break


def kill_jobs(args):
    jobs = get_jobs_data()
    
    ids = []
    if args.range:
        ids = list(range(args.range[0], args.range[1] + 1))
    
    if args.name:
        p = re.compile(args.name)
        for job in jobs:
            if p.search(job["Job_Name"]):
                ids.append(int(job["Job_Id"].split(".")[0]))
    
    for i in ids:
        if args.simulate:
            print("Simulated deleting job %d..." % i)
        else:
            print("Deleting job %d..." % i)
            subprocess.check_output(["qdel", str(i)])


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Tool for managing PBS.")

    subparsers = parser.add_subparsers()

    parser_info = subparsers.add_parser(
        "info", 
        help="Summary info on all jobs.")
    parser_info.set_defaults(func=print_info)

    parser_find_bad = subparsers.add_parser(
        "find_bad", 
        help="Find bad runs based on existence of file with a certain name.")
    parser_find_bad.add_argument("-p", "--path", dest="path",
        type=str, default=".",
        help="Path in which to look for runs.")
    parser_find_bad.add_argument("-f", "--filename", dest="filename",
        type=str, required=True,
        help="Filename to check for")
    parser_find_bad.set_defaults(func=find_bad)

    parser_tail = subparsers.add_parser(
        "tail", 
        help="Print last few lines of a output file from all running jobs.")
    parser_tail.add_argument("-N", "--nlines", dest="nlines",
        type=int, default=10,
        help="Number of lines to print.")
    parser_tail.add_argument("-f", "--filename", dest="filename",
        type=str, required=True,
        help="Filename to check for")
    parser_tail.set_defaults(func=tail)

    parser_kill = subparsers.add_parser(
        "kill", help="Kill jobs."
    )
    parser_kill.add_argument(
        "-s", "--simulate", dest="simulate", action="store_true",
        help="Simulation mode. No actual qdel is called."
    )
    parser_kill.add_argument(
        "-n", "--name", dest="name", type=str,
        nargs="?",
        help="Specify a substring of a name to delete "
             "jobs. E.g., 'job' will delete all jobs with "
             "'job' in the job name.")
    parser_kill.add_argument(
        "-r", "--range", dest="range", type=int,
        nargs=2,
        help="Delete jobs by range. Specify start and end job ids (inclusive) to delete.")
    parser_kill.set_defaults(func=kill_jobs)
    
    args = parser.parse_args()

    try:
        args.func(args)
    except Exception as ex:
        print(ex)
        parser.print_help()
        sys.exit(0)