#!/usr/bin/python2
# wait.py
# fork from lcg-info-dynamic-scheduler
# $Id: lcg-info-dynamic-scheduler.cin 2296 2011-05-23 10:12:43Z templon $
# Source: $URL: svn+ssh://svn@ndpfsvn.nikhef.nl/repos/pdpsoft/nl.nikhef.pdp.dynsched/trunk/lcg-info-dynamic-scheduler.cin $
# J. A. Templon, NIKHEF/PDP 2011

# purpose: determine longest-waiting job for each of groups present, and append this data to a file
#   also appends data on core rollover time to a file
#   input is standard lrmsinfo-generic format
#   see https://ndpfsvn.nikhef.nl/cgi-bin/viewvc.cgi/pdpsoft/nl.nikhef.pdp.dynsched/trunk/lrmsinfo-generic.txt

# define mapping of attributes in lrmsinfo-generic output to local

jobatts = {
    'id' : { 'source' : 'jobid', 'fmt' : 'varchar(25)' },
    'ncpu' : { 'source' : 'cpucount', 'fmt' : 'int'},
    'maxwall' : { 'source' : 'maxwalltime', 'fmt' : 'int'},
    'name' : { 'source' : 'name', 'fmt' : 'varchar(32)' },
    'qtime' : { 'source' : 'qtime', 'fmt' : 'int' },
    'queue' : { 'source' : 'queue', 'fmt' : 'varchar(16)' },
    'start' : { 'source' : 'start', 'fmt' : 'int' },
    'anchor' : { 'source' : 'startAnchor', 'fmt' : 'varchar(32)' },
    'state' : { 'source' : 'state', 'fmt' : 'varchar(8)' },
    'usr' : { 'source' : 'user', 'fmt' : 'varchar(8)' },
    'grp' : { 'source' : 'group', 'fmt' : 'varchar(8)' },
    'wall' : { 'source' : 'walltime', 'fmt': 'int'}
    }

# create reverse mapping source->local
revmap = dict()
for k in jobatts.keys():
    revmap[jobatts[k]['source']] = k

def usage():
    print "Usage: wait.py -g grplist -<o output_dir> infile"
    print "       grplist is like atlas,alice,lhcb"
    print "       default for output_dir is $HOME/tmp"
    
def abort_without_output(msg):
    print msg
    print "Exiting without output"
    sys.exit(2)

import sys

import getopt
import string

try:
    opts, args = getopt.getopt(sys.argv[1:], "o:g:",
                               ["output_dir=","group_list="])
except getopt.GetoptError:
    # print help information and exit:
    emsg = "Error parsing command line: " + string.join(sys.argv)
    usage()
    abort_without_output(emsg)

import os
odir = os.environ['HOME'] + '/tmp/'
grplist = None

for o, a in opts:
    if o in ("-o", "--output_dir"):
        odir = a
    elif o in ("-g", "--group_list"):
        grplist = a.split(",")

if not grplist:
    usage()
    abort_without_output('Please provide a valid group list')
    
try:
    retval = os.stat(odir)
except OSError:
    abort_without_output("Specified output directory does not exist: " + odir)

def tconv(timeval):  # convert time to desired units
    timeval = min(timeval, 2146060842)
    return repr( int(timeval) ) # seconds->printable

if len(args) < 1:
    abort_without_output("Please specify an input file")
else:
    infh = open(args[0],mode='r')

lrmlines=infh.readlines()

import sqlite
cx = sqlite.connect(":memory:")
cu =  cx.cursor()

# build table creation string

crstr = 'create table jobs ('
ctr = 0
for k in jobatts.keys():
    if ctr > 0 : crstr = crstr + ','
    crstr = crstr + "%s %s" % (k, jobatts[k]['fmt'])  ; ctr += 1
crstr += ')'
    
cu.execute(crstr)

for line in lrmlines:
    s = line.strip()
    f = s.split()
    if s[0] == '{' and s[-1] == '}':  # looks like a dict
        jd = eval(s, {"__builtins__" : {}})

        # build insert string
        kl = jd.keys()
        insstr = 'insert into jobs ('
        ctr = 0
        for k in kl:
            if ctr > 0 : insstr += ','
            insstr += revmap[k]
            ctr += 1
        insstr += ') values ('
        ctr = 0
        for k in kl:
            if ctr > 0 : insstr += ','
            insstr += repr(jd[k])
            ctr += 1
        insstr += ')'

        cu.execute(insstr)

    elif len(f) == 2:
        if f[0]   == 'nactive'    : slotsUp   = int(f[1])
        elif f[0] == 'nfree'      : slotsFree = int(f[1])
        elif f[0] == 'now'        : now       = int(f[1])
        elif f[0] == 'schedCycle' : schedCycle = int(f[1])
        else:
            logging.warn("invalid input in result of lrms plugin command " + cmd)
	    logging.warn("invalid input: " + s)
    else:
	logging.warn("invalid input in result of lrms plugin command " + cmd)
	logging.warn("invalid input: " + s)

import time
nowstr = time.strftime("%b-%d-%Y-%H:%M:%S",time.localtime(now))
for g in grplist:
    selstr = 'select qtime from jobs where state="queued" and grp="' + g + '"'
    cu.execute(selstr)
    rl = cu.fetchall()
    if len(rl) > 0:
        qtl = list()
        for r in rl:
            qtl.append(r[0])
        qtl.sort()
        waitt = now - qtl[0]
    else:
        waitt = 2
        
    of = open(os.environ['HOME']+'/tmp/'+g+'.wdata',mode='a')
    of.write("%s   %s   %6d\n" % (nowstr, g, waitt))
    of.close()

    of = open(os.environ['HOME']+'/tmp/'+g+'.qdata',mode='a')
    of.write("%s   %s   %6d\n" % (nowstr, g, len(rl)))
    of.close()
    
selstr = 'select start from jobs where state="running" and start<>"None"'
cu.execute(selstr)
rl = cu.fetchall()
if len(rl) > 0:
    stl = list()
    for r in rl:
        stl.append(r[0])

    stl.sort()
    lastroll = now - stl[-1]
    if lastroll < 2:
        lastroll = 2

    of = open(os.environ['HOME']+'/tmp/core.wdata',mode='a')
    of.write("%s   %s   %6d\n" % (nowstr, 'core', lastroll))
    of.close()

selstr = 'select start from jobs where state="queued"'
cu.execute(selstr)
rl = cu.fetchall()
of = open(os.environ['HOME']+'/tmp/total.qdata',mode='a')
of.write("%s   %s   %6d\n" % (nowstr, 'total', len(rl)))
of.close()

selstr = 'select grp,count(*) from jobs where state="running" group by grp'
cu.execute(selstr)
rl = cu.fetchall()
of = open(os.environ['HOME']+'/tmp/groups.rdata',mode='a')
of.write("%s\n" % (nowstr))
for r in rl:
    of.write(repr(r) + '\n')
of.write('\n')
of.close()

selstr = 'select grp,count(*) from jobs where state="queued" group by grp'
cu.execute(selstr)
rl = cu.fetchall()
of = open(os.environ['HOME']+'/tmp/groups.qdata',mode='a')
of.write("%s\n" % (nowstr))
for r in rl:
    of.write(repr(r) + '\n')
of.write('\n')
of.close()

### Local Variables: ***
### mode: python ***
### End: ***