Extended/Enhanced Sun Grid Engine qstat

#! /usr/bin/env python
 
###############################################################################
##  Copyright 2009 Jeet Sukumaran.
##
##  This program is free software; you can redistribute it and/or modify
##  it under the terms of the GNU General Public License as published by
##  the Free Software Foundation; either version 3 of the License, or
##  (at your option) any later version.
##
##  This program is distributed in the hope that it will be useful,
##  but WITHOUT ANY WARRANTY; without even the implied warranty of
##  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
##  GNU General Public License for more details.
##
##  You should have received a copy of the GNU General Public License along
##  with this program. If not, see <http://www.gnu.org/licenses/>.
##
###############################################################################
 
"""
Wrapper around qstat for more compact representation and flexible sorting.
"""
 
import sys
import os
import subprocess
import datetime
import re
from xml.etree import ElementTree
from optparse import OptionGroup
from optparse import OptionParser
 
_prog_usage = '%prog [options] [qstat arguments]'
_prog_version = 'QSTAT-X Version 1.0'
_prog_description = 'Extended/enhanced qstat.'
_prog_author = 'Jeet Sukumaran'
_prog_copyright = 'Copyright (C) 2009 Jeet Sukumaran.'
 
### from phylocratics
def format_table(rows, column_names=None, max_column_width=None, border_style=2):
    """
    'Pretty-prints' a tuple of dictionaries in a table format. This method
    can read the column names directly off the dictionary keys, but if a tuple of
    these keys is provided in the 'column_names' variable, then the order of column_names will follow
    the order of the fields/keys in that variable.
    """
    if column_names or len(rows) > 0:
        lengths = {}
        rules = {}
        if column_names:
            column_list = column_names
        else:
            try:
                column_list = rows[0].keys()
            except:
                column_list = None
        if column_list:
            # characters that make up the table rules
            border_style = int(border_style)
            #border_style = 0
            if border_style >= 1:
                vertical_rule = '  '
                horizontal_rule = '-'
                rule_junction = '---'                
            elif border_style >= 2:
                vertical_rule = ' | '
                horizontal_rule = '-'
                rule_junction = '-+-'
            else:
                vertical_rule = '  '
                horizontal_rule = ''
                rule_junction = ''                
            if border_style >= 3:
                left_table_edge_rule = '| '
                right_table_edge_rule = ' |'
                left_table_edge_rule_junction = '+-'
                right_table_edge_rule_junction = '-+'
            else:
                left_table_edge_rule = ''
                right_table_edge_rule = ''
                left_table_edge_rule_junction = ''
                right_table_edge_rule_junction = ''
 
            if max_column_width:
                column_list = [c[:max_column_width] for c in column_list]
                trunc_rows = []
                for row in rows:
                    new_row = {}
                    for k in row.keys():
                        new_row[k[:max_column_width]] = str(row[k])[:max_column_width]
                    trunc_rows.append(new_row)
                rows = trunc_rows
 
            for col in column_list:
                rls = [len(str(row[col])) for row in rows]
                lengths[col] = max(rls+[len(col)])
                rules[col] = horizontal_rule*lengths[col]
 
            template_elements = ["%%(%s)-%ss" % (col, lengths[col]) for col in column_list]
            row_template = vertical_rule.join(template_elements)
            border_template = rule_junction.join(template_elements)
            full_line = left_table_edge_rule_junction + (border_template % rules) + right_table_edge_rule_junction 
            display = []
            if border_style > 0:
                display.append(full_line)
            display.append(left_table_edge_rule + (row_template % dict(zip(column_list, column_list))) + right_table_edge_rule)
            if border_style > 0:
                display.append(full_line)            
            for row in rows:       
                display.append(left_table_edge_rule + (row_template % row) + right_table_edge_rule)      
            if border_style > 0:
                display.append(full_line)
            return "\n".join(display)
        else:
            return ''
    else:
        return ''
 
def parse_error(qstat_cmd, stdout, stderr):
    sys.stderr.write("No information or failed to parse output.\n")
    sys.stderr.write("Command executed was: \"%s\"\n" % qstat_cmd)
    sys.stderr.write("Raw result was:\n")
    sys.stderr.write(stdout)
    sys.stderr.write("\n")
    sys.exit(1)
 
def main():
    """
    Main CLI handler.
    """
 
    parser = OptionParser(usage=_prog_usage, 
        add_help_option=True, 
        version=_prog_version, 
        description=_prog_description)    
 
    parser.add_option('-i', '--id',
        action='store_const',
        dest='sort_by',
        const='id',
        default='id',
        help='sort by job id [default]')
 
    parser.add_option('-n', '--name',
        action='store_const',
        dest='sort_by',
        const='name',
        help='sort by job name')
 
    parser.add_option('-o', '--owner', '-u', '--user', 
        action='store_const',
        dest='sort_by',
        const='owner',
        help='sort by owner/user')
 
    parser.add_option('-t', '--submitted', 
        action='store_const',
        dest='sort_by',
        const='submitted',
        help='sort by submission time') 
 
    parser.add_option('-s', '--state', 
        action='store_const',
        dest='sort_by',
        const='state',
        help='sort by state')
 
    parser.add_option('-q', '--queue', 
        action='store_const',
        dest='sort_by',
        const='queue',
        help='sort by queue')
 
    parser.add_option('-c', '--node', 
        action='store_const',
        dest='sort_by',
        const='node_index',
        help='sort by node')        
 
    parser.add_option('-l', '--slots', 
        action='store_const',
        dest='sort_by',
        const='slots',
        help='sort by slots')
 
    parser.add_option('-r', '--reverse', 
        action='store_true',
        dest='reverse',
        default=False,
        help='reverse sort')          
 
    (opts, args) = parser.parse_args()
 
    qstat_cmd = "qstat -xml -u '*' " + " ".join(args)
 
    qstat_proc = subprocess.Popen(qstat_cmd,
            shell=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            env=os.environ)
    stdout, stderr = qstat_proc.communicate()
 
    try:
        root = ElementTree.fromstring(stdout)
    except:
        parse_error(qstat_cmd, stdout, stderr)
    queue_info = root.find('queue_info')
    if queue_info is None:
        parse_error(qstat_cmd, stdout, stderr)    
 
    fields = ["id",
              "name",
              "owner",
              "submitted",
              "state",
#               "cpu",
#               "mem",
#               "io",
              "queue",
              "node",
              "slots"]
    jobs = []
 
    for job_list in queue_info:
        job = {}
        job['id'] = job_list.find("JB_job_number").text
        job['name'] = job_list.find("JB_name").text
        job['owner'] = job_list.find("JB_owner").text
        job['state'] = job_list.find("state").text
#         job['cpu'] = job_list.find("cpu_usage").text
#         job['mem'] = job_list.find("mem_usage").text
#         job['io'] = job_list.find("io_usage").text
        qn = job_list.find("queue_name").text
        job['queue'], node = qn.split("@")
        job['node'] = node[:-6]
        job['node_index'] = int(node[10:-6])
        job['slots'] = job_list.find("slots").text
        stime = datetime.datetime.strptime(job_list.find("JAT_start_time").text, "%Y-%m-%dT%H:%M:%S")
        job["submitted"] = stime.strftime("%Y-%m-%d %H:%M:%S")
 
        jobs.append(job)        
 
    jobs.sort(key=lambda x : x[opts.sort_by], reverse=opts.reverse)
    sys.stdout.write(format_table(jobs, fields, border_style=0))
    sys.stdout.write("\n")
 
if __name__ == "__main__":
    main()