qres

#/usr/bin/env python 

import commands
from commands import getoutput
import sys
import os
import argparse
from string import strip
from collections import defaultdict 
import re

__LICENSE__ = "GPL v3"
__AUTHOR__ = "Jaime Huerta-Cepas"

DEFAULT_JOB_PRIO = -512
MEM_STRING_LEN = 15
MAXVMEMMATCH = re.compile("maxvmem=(\d+(\.\d+)?\w?)")
VMEMMATCH = re.compile("vmem=(\d+(\.\d+)?\w?)")
H_VMEMMATCH = re.compile("h_vmem=(\d+(\.\d+)?\w?)")
H_VMEMMATCH_m = re.compile("h_vmem=(\d+(\.\d+)?\w?)", re.MULTILINE)
USAGE_MATCH = re.compile("^usage\s+(\d+):(.+)")
# SOME CONSTANTS
mK = 1024
mM = 1024*1024
mG = 1024*1024*1024
mT = 1024*1024*1024*1024

mk = 1000
mm = 1000*1000
mg = 1000*1000*1000
mt = 1000*1000*1000*1000

DAY = 3600 * 24
HOUR = 3600
MINUTE = 60

def runcheck():
    DEFAULT_RESERVED_MEM, DEFAULT_SLOTS, PER_SLOT_MEM = 0, 0, False
    USERS = set([user.strip() for user in
                 commands.getoutput('qconf -suserl').split("\n")])
    QUEUES = set([queue.strip() for queue in
                 commands.getoutput('qconf -sql').split("\n")])

    HOSTNAMES = set([host.split(".")[0] for host in
                     commands.getoutput('qconf -sel').split("\n")])
    consumables = getoutput("qconf -sc")
    for line in consumables.split("\n"):
        fields = map(strip, line.split())
        if fields[0] == "h_vmem":
            if fields[4].upper() == "YES" and fields[5].upper() in ["YES", "JOB"]:
                DEFAULT_RESERVED_MEM = mem2bytes(fields[6])
                if fields[5].upper() == "YES":
                    PER_SLOT_MEM = True
            else:
                print "ERROR: h_vmem is not defined as a consumable resource:"
                print line
                sys.exit(1)
            
        elif line.startswith("slots"):
            if fields[4].upper() == "YES" and fields[5].upper() == "YES":
                DEFAULT_SLOTS = mem2bytes(fields[6])
            else:
                print "ERROR: slots is not defined as a consumable resource:"
                print line
                sys.exit(1)
                
    ## detect consumables of each host
    host2avail_mem = {}
    host2avail_slots = {}
    for h in HOSTNAMES:
        hinfo = commands.getoutput("qconf -se %s" %h)
        mem_match = re.search(H_VMEMMATCH_m, hinfo)
        slots_match = re.search("[\s,]?slots=(\d+)", hinfo, re.MULTILINE)
        if not slots_match:
            slots_match = re.search("[\s,]?processors\s*(\d+)", hinfo, re.MULTILINE)


        
        if mem_match:
            host2avail_mem[h] = mem2bytes(mem_match.groups()[0])
        else:
            print 'h_vmem should be defined as consumable resource in ' + h
            sys.exit(1)
        if slots_match:
            host2avail_slots[h] = int(slots_match.groups()[0])
        else:
            print 'slot should be defined as consumable resource in ' + h
            sys.exit(1)
        
    return (USERS, QUEUES, HOSTNAMES, DEFAULT_RESERVED_MEM,
            DEFAULT_SLOTS, PER_SLOT_MEM, host2avail_mem, host2avail_slots)
        
def load_general_values():
    USERS = set([user.strip() for user in
                 commands.getoutput('qconf -suserl').split("\n")])
    QUEUES = set([queue.strip() for queue in
                 commands.getoutput('qconf -sql').split("\n")])

    HOSTNAMES = set([host.split("@")[0] for host in
                     commands.getoutput('qconf -sel').split("\n")])
    
    return QUEUES, USERS, HOSTNAMES
        
def color(color, string):
    color2code = {
        "header": '\033[95m',
        "lblue": '\033[94m',
        "lgreen": '\033[92m',
        "yellow": '\033[93m',
        "lred": '\033[91m',
        "magenta": "\033[35m",
        "white": "\033[37m",
        "red": '\033[31m',
        "blue": '\033[34m',
        "purple": '\033[35m',
        "green": '\033[32m',
        "cyan": '\033[36m',
        "brown": '\033[33m',
        
    }
    END = '\033[0m'
    return ''.join([color2code[color], string, END])

def get_hosts_in_queues(queues):
    all_hosts = set()
    for q in queues:
        qinfo = commands.getoutput('qconf -sq %s' %q)
        qinfo = qinfo.replace("\\\n", " ")
        match = re.search("hostlist\s(.+)", qinfo, re.MULTILINE)
        if match:
            for host in [h.strip() for h in match.groups()[0].split()]:
                if host.startswith("@"):
                    hginfo = commands.getoutput('qconf -shgrp %s' %host)
                    hginfo = hginfo.replace("\\\n", " ")
                    match = re.search("hostlist\s(.+)", hginfo, re.MULTILINE)
                    all_hosts.update([h.strip() for h in match.groups()[0].split()])
                else:
                    all_hosts.update([host])

    # update to keep just the hostname without domain
    for host in all_hosts:
        hostname = host.split('.')[0]
        all_hosts.remove(host)
        all_hosts.add(hostname)

    # convert to list to show nodes sorted
    all_hosts_list = []
    for host in all_hosts:
        all_hosts_list.append(host)

    return list(set(all_hosts_list))

def get_hosts_for_user(user):
    hosts = map(strip, commands.getoutput('qstat -s r -u ' + user + ' | grep -v job-ID |grep -v\
    \'\-\-\'|awk {\'print $8\'}|awk -F \"@\" {\'print $2\'} |awk -F \'.\'\
    {\'print $1\'}').split("\n"))
    
    return list(set(hosts))

def print_as_table(rows, header=None, fields=None, print_header=True, stdout=sys.stdout):
    """ Print >>Stdout, a list matrix as a formated table. row must be a list of
    dicts or lists."""
    if header is None:
        header = []
        
    def _str(i):
        if isinstance(i, float):
            return "%0.2f" %i
        else:
            return str(i)

    def _safe_len(i):
        return len(re.sub('\\033\[\d+m', '',  _str(i)))

    def _safe_rjust(s, just):
        return (" " * (just - _safe_len(s))) + s
        
    vtype = None
    for v in rows:
        if vtype != None and type(v)!=vtype:
            raise ValueError("Mixed row types in input")
        else:
            vtype = type(v)

    lengths  = {}
    if vtype == list or vtype == tuple:
        v_len = len(fields) if fields else len(rows[0])
        
        if header and len(header)!=v_len:
            raise Exception("Bad header length")

        # Get max size of each field
        if not fields:
            fields = range(v_len)
        
        for i,iv in enumerate(fields):
            header_length = 0
            if header != []:
                #header_length = len(_str(header[i]))
                header_length = _safe_len(header[i])
            max_field_length = max( [_safe_len(r[iv]) for r in rows] )
            lengths[i] = max( [ header_length, max_field_length ] )

        if header and print_header:
            # Print >>Stdout, header names
            for i in xrange(len(fields)):
                print >>stdout, _str(header[i]).rjust(lengths[i])+" | ",
            print >>stdout, ""
            # Print >>Stdout, underlines
            for i in xrange(len(fields)):
                print >>stdout, "".rjust(lengths[i],"-")+" | ",
            print >>stdout, ""
        # Print >>Stdout, table lines
        for r in rows:
            for i,iv in enumerate(fields):
                #print >>stdout, _str(r[iv]).rjust(lengths[i])+" | ",
                print >>stdout, _safe_rjust(_str(r[iv]), lengths[i])+" | ",
            print >>stdout, ""

    elif vtype == dict:
        if header == []:
            header = rows[0].keys()
        for ppt in header:
            lengths[ppt] = max( [len(_str(ppt))]+[ len(_str(p.get(ppt,""))) for p in rows])
        if header:
            for ppt in header:
                print >>stdout, _str(ppt).rjust(lengths[ppt])+" | ",
            print >>stdout, ""
            for ppt in header:
                print >>stdout, "".rjust(lengths[ppt],"-")+" | ",
            print >>stdout, ""

        for p in rows:
            for ppt in header:
                print >>stdout, _str(p.get(ppt,"")).rjust(lengths[ppt])+" | ",
            print >>stdout, ""
            page_counter +=1

def tm2sec(tm):
    try:
        return time.mktime(time.strptime(tm))
    except Exception: 
        return 0.0

def mem2bytes(mem):
    try: 
        bytes = float(mem)
    except ValueError: 
        mod = mem[-1]
        mem = mem[:-1]
        if mod == "K": 
            bytes = float(mem) * mK
        elif mod == "M":
            bytes = float(mem) * mM
        elif mod == "G":
            bytes = float(mem) * mG
        elif mod == "k": 
            bytes =  float(mem) * mk
        elif mod == "m":
            bytes = float(mem) * mm
        elif mod == "g":
            bytes = float(mem) * mg

    return bytes

def bytes2mem(bytes):
    if bytes > mt:
        v = "%0.1f" %(float(bytes)/mt)
        v = v.rstrip("0").rstrip(".") +"T"
    elif bytes > mg:
        v = "%0.1f" %(float(bytes)/mg)
        v = v.rstrip("0").rstrip(".") +"G"
    elif bytes > mm:
        v = "%0.1f" %(float(bytes)/mm)
        v = v.rstrip("0").rstrip(".") +"M"
    elif bytes > mk:
        v = "%0.1f" %(float(bytes)/mk)
        v = v.rstrip("0").rstrip(".") +"K"
    else:
        v = str(bytes)
    return  v

def get_host_info():
    host2load = {}
    for line in getoutput("qhost").split("\n")[3:]:
        #HOSTNAME                ARCH         NCPU  LOAD  MEMTOT  MEMUSE  SWAPTO  SWAPUS
        hostname, arch, slots, load, mem, mem_used, swap, swap_used = line.split()
        hostname = hostname.split(".")[0]
        try:
            host2load[hostname] = float(load)
        except ValueError:
            host2load[hostname] = -1
            
    return host2load 

def parse_reservations(reservations, valid_hosts):
    valid_hosts = set(valid_hosts)
    queue2ARcpu = defaultdict(int)
    host2ARcpu = defaultdict(int)
    for res in reservations[2:]:
        # ['ar-id', 'name', 'owner', 'state', 'start', 'at', 'end', 'at', 'duration']
        # ['277', 'NADA', 'jhuerta', 'r', '02/25/2013', '18:07:00', '02/26/2013', '19:07:00', '25:00:00']
        try:
            rid, name, user, state, start_day, start_tm, end_day, end_tm, duration = res.split()
        except ValueError:
            rid, user, state, start_day, start_tm, end_day, end_tm, duration = res.split()
        
        if state == "r":
            for line in getoutput("qrstat -ar %s" %rid).split("\n"):
                if line.startswith("granted_slots_list"):
                    for queue, host, slots in  re.findall("[\s,]([^@]+)@([^=]+)=(\d+)", line):
                        slots = int(slots)
                        host = host.strip().split(".")[0]
                        if host in valid_hosts: 
                            queue2ARcpu[queue.strip()] += slots
                            host2ARcpu[host] += slots
    return host2ARcpu

def parse_running_jobs(jobs, valid_hosts):
    valid_hosts = set(valid_hosts)
    ## Load info about running jobs
    job2info = defaultdict(dict)
    tasks = []
    # Load basic info about running jobs and tasks
    for x in jobs[2:]:
        try:
            jid, pri, name, user, status, date, tm,  queue, slots, task = fields = x.split()
        except ValueError:
            task = 1
            jid, pri, name, user, status, date, tm,  queue, slots = fields = x.split()
            
        queue, host = queue.split("@")
        host = host.split(".")[0]
        if host not in valid_hosts:
            continue
           
        jid, task = map(int, [jid, task])
        tasks.append([jid, task, queue, host])
        if jid not in job2info:
            job2info[jid]["raw"] = getoutput("qstat -j %s" %jid)
            job2info[jid]["smp"] = 1 # Assumes 1 cpu if no other SMP info is available
            job2info[jid]["r.mem"] = DEFAULT_RESERVED_MEM
            job2info[jid]["prio"] = DEFAULT_JOB_PRIO
            
            for line in job2info[jid]["raw"].split("\n"):
                usage_match = re.search(USAGE_MATCH, line)
                if usage_match:
                    taskid = int(usage_match.groups()[0])
                    mem_match = re.search(VMEMMATCH, line)
                    maxmem_match = re.search(MAXVMEMMATCH, line)
                    job2info[jid][taskid] = {}
                    if mem_match: 
                        job2info[jid][taskid]["vmem"] = mem2bytes(mem_match.groups()[0])
                    else:
                        job2info[jid][taskid]["vmem"] = 0
                    if maxmem_match:
                        job2info[jid][taskid]["maxvmem"] = mem2bytes(maxmem_match.groups()[0])
                    else:
                        job2info[jid][taskid]["maxvmem"] = 0
                        
                elif line.startswith("hard resource_list:"):
                    rt_match = re.search("h_rt=(\d+)", line)
                    mem_match = re.search(H_VMEMMATCH, line)

                    if rt_match:
                        job2info[jid]["h_rt"] = int(rt_match.groups()[0])
                    if mem_match:
                        job2info[jid]["r.mem"] = mem2bytes(mem_match.groups()[0])
                    else:
                        job2info[jid]["r.mem"] = DEFAULT_RESERVED_MEM

                elif line.startswith("parallel environment:"):
                    smp_match = re.search("smp range:([^\n]+)", line)                    
                    if smp_match:
                        job2info[jid]["smp"] = int(smp_match.groups()[0])
                        
                elif line.startswith("priority:"):
                    job2info[jid]["prio"] = int(re.search("priority:([^\n]+)", line).groups()[0])
                    
                elif line.startswith("owner:"):
                    job2info[jid]["user"] = re.search("owner:([^\n]+)", line).groups()[0].strip()

                    
    return tasks, job2info



def sort_resources_by_host(tasks, job2info):
    host2Umem = defaultdict(int)
    host2Rmem = defaultdict(int)
    host2Ucpu = defaultdict(int)
    
    for jid, taskid, queue, host in tasks:
        tinfo = job2info[jid][taskid]
        jinfo = job2info[jid]
        
        if PER_SLOT_MEM:
            rmem = jinfo["r.mem"] * jinfo["smp"]
        else:
            rmem = jinfo["r.mem"]
                    
        host2Rmem[host] += rmem
        host2Umem[host] += tinfo["vmem"]
        host2Ucpu[host] += jinfo["smp"]

    return host2Umem, host2Rmem, host2Ucpu



if __name__ == "__main__":
    # Let's first check if resources are defined as consumable in sge
    # config
    (USERS, QUEUES, HOSTNAMES, DEFAULT_RESERVED_MEM, DEFAULT_SLOTS,
     PER_SLOT_MEM, host2Amem, host2Acpu) = runcheck()
    host2load = get_host_info()
    
    # Parse options
    parser = argparse.ArgumentParser(description='Show your cluster usage')

    parser.add_argument('-u', dest='user', default="'*'",
                        help='show stats just for this user')

    parser.add_argument('-q', dest='queue',
                        help='show stats just for this queue')

    parser.add_argument('-e', dest='hosts', nargs="+",
                        help='show stats just for list of execution hostnames')

    
    options = parser.parse_args()
        
    # check if user exists
    if options.user != "'*'" and options.user not in USERS: 
        print 'user ' + options.user + ' not found'
        sys.exit()

    # check if queue exists
    if options.queue and options.queue not in QUEUES:
        print 'queue ' + options.queue + ' not found'
        sys.exit()

    if options.hosts:
        user_hosts = set(options.hosts)
        if user_hosts - HOSTNAMES:
            print user_hosts - HOSTNAMES, "are not found"
            sys.exit()
        else:
            hosts = sorted(user_hosts)
    else:
        hosts = sorted(HOSTNAMES)
    
    ## detect host names
    # I had to use shorter host names (without domain), because they were truncated in qstat output
    # hosts = map(strip, commands.getoutput("qconf -sel").split("\n"))

    # is user arg is passed only take the node where that
    # user is running jobs
    #if options.user != '\'*\'':
    #    hosts = get_hosts_for_user(options.user)
    #    # if both user and queue args are passed...
    #    if options.queue:
    #        hostsInQueue = map(strip, get_hosts_in_queues([options.queue]))
    #        toRemove = []
    #        for x in hosts:
    #            if x not in hostsInQueue:
    #                toRemove.append(x)
    #        for x in toRemove:
    #            hosts.remove(x)
    #    hosts.sort()
    #else:
    #    # if just queue arg is passed...
    #    if options.queue:
    #        hosts = map(strip, get_hosts_in_queues([options.queue]))
    #        hosts.sort()
    #    # if no arg is passed take all exec nodes
    #    else:
    #        hosts = map(strip, commands.getoutput('qconf -sel|cut -f1 -d"."').split("\n"))
    #        hosts.sort()

    ## Detect running jobs
    
    if options.queue:
        running_jobs = map(strip, commands.getoutput('qstat -s r -u ' +\
            options.user + ' -q ' + options.queue).split("\n"))


    else:
        running_jobs = map(strip, commands.getoutput('qstat -s r -u ' +\
                options.user).split("\n"))
        
    reservations = map(strip, commands.getoutput('qrstat -u %s' %options.user).split("\n"))
        
    tasks, job2info = parse_running_jobs(running_jobs, hosts)
    host2Umem, host2Rmem, host2Ucpu = sort_resources_by_host(tasks, job2info)
    host2ARcpu = parse_reservations(reservations, hosts)
        
    ## show collected info

    entries = []

    total_mem_used = 0
    total_mem_reserved = 0
    total_mem = 0
    total_cpu = 0
    total_cpu_used = 0

    for x in hosts:
        total_mem_used += host2Umem.get(x, 0)
        total_mem_reserved += host2Rmem.get(x, 0)
        total_mem += host2Amem[x]
        total_cpu += host2Acpu[x]
        total_cpu_used += host2Ucpu.get(x, 0)

        mem_used = host2Umem.get(x, 0)
        mem_reserved = host2Rmem.get(x, 0)
        mem_avail = host2Amem[x]

        mem_factor_used = (mem_used * MEM_STRING_LEN) / mem_avail
        mem_factor_unused = (mem_reserved * MEM_STRING_LEN) / mem_avail
        mem_factor_unused -= mem_factor_used

        mem_char_used = int(round(mem_factor_used))
        mem_char_unused = int(round(mem_factor_unused))
        mem_char_free = MEM_STRING_LEN- mem_char_used - mem_char_unused

        if mem_used > 0.66 * mem_reserved:
            USED_MEM_COL = "lgreen"
        elif mem_used > 0.33 * mem_reserved:
            USED_MEM_COL = "yellow"
        else:
            USED_MEM_COL = "lred"

        mem_eff = (mem_used * 100)  / mem_reserved if mem_reserved else 0.0
        MEM_EFFICIENCY = color(USED_MEM_COL, "%s%%" %( ("%d" %mem_eff).rjust(2)))
        HOST_INFO = "%0.1f/%s (%s)" %(round(host2load.get(x, -1),1), color("brown", str(host2Acpu[x])), color("lblue", bytes2mem(mem_avail)))

        empty_slots = host2Acpu[x] - host2Ucpu.get(x, 0) - host2ARcpu.get(x, 0)
        fields = [x.ljust(8),
                  HOST_INFO,
                  "%s / %s" %(color("cyan", str(host2Ucpu.get(x, 0))), host2Acpu[x]),
                  "%s / %s (%s)" % (color("cyan", bytes2mem(mem_used)), bytes2mem(mem_reserved), MEM_EFFICIENCY),
                  color("red", "#" * mem_char_used) +
                  color("white", "~" * mem_char_unused) +
                  color("white", " " * mem_char_free), 
                  color("red", "#"*host2Ucpu.get(x, 0)) +
                  color("blue", "#"*host2ARcpu.get(x, 0)) +
                  ("." * empty_slots)
              ]
        entries.append(fields)

    # Create totals entry
    HOST_INFO = "%s CPU, %s" %(color("brown", str(total_cpu)), color("lblue", bytes2mem(int(total_mem))))
    if total_mem_used > 0.66 * total_mem_reserved:
        USED_MEM_COL = "lgreen"
    elif total_mem_used > 0.33 * total_mem_reserved:
        USED_MEM_COL = "yellow"
    else:
        USED_MEM_COL = "lred"

    mem_eff = (total_mem_used * 100)  / total_mem_reserved if total_mem_reserved else 0.0
    MEM_EFFICIENCY = color(USED_MEM_COL, "%s%%" %( ("%d" %mem_eff).rjust(2)))
    entries.append([""]*6)
    fields = ["TOTALS".ljust(8),
              HOST_INFO,
                  "%s / %s" %(color("cyan", str(total_cpu_used)), total_cpu),
                  "%s / %s (%s)" % (color("cyan", bytes2mem(total_mem_used)), bytes2mem(total_mem_reserved), MEM_EFFICIENCY),
                  "",
                  ""
              ]
    entries.append(fields)

    # add a description (with colors) before the output with information about
    # what is being printed to stdout
    if options.user != '\'*\'':
        description =  "\n" + color("white", "User " + options.user + "\n")
        if options.queue:
            desk =  "\n" + color("white", "User " + options.user +\
                                     " in queue " + options.queue + "\n")
    else:
        if options.queue:
            description =  "\n" + color("white", "Jobs in queue " + options.queue)
        else:
            description =  "\n" + color("white", "CLUSTER STATS\n")

    print description    
    header = "hostname", "host info", "CPU", "Mem: used/reserved", "Mem usage", "CPU usage" 

    print_as_table(entries, header=header)


    # update cpu usage per array job
    jid2Ucpu = defaultdict(int)
    for jid, taskid, queue, hostname in tasks:
        jid2Ucpu[jid] += job2info[jid]["smp"]
       
    # sort jobs by number of cpus in use
    sorted_jobs = sorted(jid2Ucpu.keys(), lambda x,y: cmp(jid2Ucpu[x], jid2Ucpu[y]), reverse=True)
    
    # Print info by user
    user2jid = defaultdict(list)
    jid2prio = {}
    
    # update per user counter and priorities
    for jid in sorted_jobs:
        info = job2info[jid]
        prio = info["prio"] - DEFAULT_JOB_PRIO
        user = info["user"]
        if prio > 0:
            prio_label = color("lred", "%sx" %prio)
        elif prio < 0:
            prio_label = color("white", "%sx" %prio)
        else:
            prio_label = color("lgreen", "0x")

        jid2prio[jid] = prio_label
        user2jid[user].append(jid)
        
    print
    for user, jobs in user2jid.iteritems():
        print "User", color("green",user.ljust(12)), "has",\
            color("yellow", str(len(jobs))), "running job(s) and is using", \
            color("yellow", str(sum([jid2Ucpu[j] for j in jobs])) ), "CPUs:", \
            ', '.join(["%d=%s" %(jid2Ucpu[jid], jid2prio[jid]) for jid in jobs])
              
        
    print color("white", "\n\nNN=PPx: NNN cpu-tasks running with priority PP" )
    print color("white", "Note that priority is expressed as a fold factor:")
    print color("white", "  i.e. 10x = 10 times more priority than the standard level")
    print
    print color("red", "#"),": used"
    print color("blue", "#"),": reserved"
    print color("white", "~"),": reserved but not used"    
    print ". : not used"
    print
    print "User 'qres -h' for more info"