freealloc

#!/usr/bin/python
#
# freealloc - display resources requested by jobs in an allocation
#

INSTALL_ROOT = '/usr/local/flux-utils';

fluxod_maxproc = 4385


import site
site.addsitedir( "%s/lib/python2.7/site-packages/" % INSTALL_ROOT )


import sys
import os
import re
import subprocess
from collections import defaultdict
import argparse

from torque import *

# Use libc's atoi() function per
# http://stackoverflow.com/questions/1665511/python-equivalent-to-atoi-atof
# ...because using regular expressions followed by int() seems overkill,
# and there are several points where we need to get integers that are at
# the beginning of strings.
import ctypes.util
libc = ctypes.cdll.LoadLibrary(ctypes.util.find_library('c'))


def run_moab(argv, type='xml'):
    success = False
    attempt = 1
    while attempt <= 60:
        try:
            pipe = subprocess.Popen( argv, stdout=subprocess.PIPE )
            if type == 'xml':
		pass # so we don't have to import lxml.etree for this script
                #output = etree.parse( pipe.stdout )
            else:
                output = pipe.stdout.readlines()
            pipe.stdout.close()
            pipe.wait()
            success = True
            break
        except:
            pass
        sys.stderr.write( 'Trying %s.' % argv[0].split('/')[-1] if attempt == 1 else '.' )
        sys.stderr.flush()
        attempt += 1
        time.sleep( 1 )
    if not success:
        sys.stderr.write( "\nERROR: unable to run " + argv[0].split('/')[-1] )
        sys.exit( 1 )
    if attempt > 1:
        sys.stderr.write( "\n" )
        sys.stderr.flush()
    return output


def main():

    global fluxod_maxproc

    parser = argparse.ArgumentParser(description='Displays unused cores and memory for an allocation')
    parser.add_argument("allocation_name", help="Name of the allocation")
    parser.add_argument("--jobs", help="display core and memory usage for each job", action="store_true")
    args = parser.parse_args()
    if not re.search( r'^[a-zA-Z][a-zA-Z0-9_-]+$', args.allocation_name ):
        print "ERROR: bad allocation name: allocation names must begin with a letter followed by letters, digits, and underscores."
        sys.exit( 1 );

    maxproc = -1
    maxmem = -1 # memory in kb

    output = run_moab( [ '/opt/moab/bin/mdiag', '-a', args.allocation_name ], type='text' )
    
    for line in output:
        m = re.search( r'\sMAXPROC=(\d+,)?(\d+)\s', line )
        if m and maxproc == -1: maxproc = int( m.group(2) )
        m = re.search( r'ALLOC_SUM:(\d+)', line )
        if m: maxproc = int( m.group(1) )
        m = re.search( r'\sMAXMEM=(\d+,)?(\d+)\s', line )
        if m: maxmem = int( m.group(2) ) * 1024  # convert mb to kb
    if args.allocation_name.endswith('_fluxod') and maxproc == -1:
        maxproc = fluxod_maxproc
    if maxproc == -1:
        print "ERROR: could not determine MAXPROC for allocation %s" % args.allocation_name
	print "It is possible that %s may not currently be active." % args.allocation_name
	print "To determine if this is the case, try running:\n    mdiag -a %s" % args.allocation_name
        sys.exit( 2 )
    if maxmem == -1:
        maxmem = maxproc * 4 * 1024 * 1024
        if not args.allocation_name.endswith('_fluxod'):
            print "WARNING: could not determine MAXMEM for allocation %s, assuming 4 GB per core" % args.allocation_name
        #print "ERROR: could not determine MAXMEM for allocation %s" % args.allocation_name
        #sys.exit( 2 )

    total_cores = 0
    total_mem = 0  # memory in kb
    uncertain = 0

    if args.jobs:
        print "Job ID           User     Cores    Memory"
        print "---------------- -------- ----- ---------"

    pbs = PBS()
    pbs.connect( pbs.default() )
    sel_list = ATTROPL( None, "job_state", None, "R", BATCH_OP.EQ )
    jobs = pbs.selstat( sel_list, None )

    for job in jobs:

        attrib = job['attrib']

        if not 'Account_Name' in attrib or attrib['Account_Name'] != args.allocation_name:
            continue

        cores = -1
        r = attrib['Resource_List']
        if 'procs_bitmap' in r:
            # see http://docs.adaptivecomputing.com/torque/4-2-6/help.htm#topics/2-jobs/requestingRes.htm#procs_bitmap
            cores = len( r['procs_bitmap'] )
        if 'procs' in r:
            c = libc.atoi( r['procs'] )
            if c > cores and c > 0: cores = c
        if 'nodes' in r:
            c = 0
            for nodespec in r['nodes'].split('+'):
                p = 1
                n = libc.atoi( nodespec )
                if n < 1: n = 1
                m = re.search( r'ppn=(\d+)', nodespec )
                if m: p = int( m.group(1) )
                if p < 1: p = 1
                c += n*p
            if c > cores and c > 0: cores = c
        if 'exec_host' in attrib:
            c = 0
            for host in attrib['exec_host'].split('+'):
                # we know how to handle host specs of the form nyx5678/3
                if not re.search( r'\w+/\d+', host ):
                    # something we don't know how to interpret, bail
                    c = 0
                    break
                c += 1
            if c > cores and c > 0: cores = c

        if cores > 0: total_cores += cores

        mem = -1
        if 'pmem' in r and cores > 0:
            mem = get_memory( r['pmem'] ) * cores;
        if 'mem' in r:
            m = get_memory( r['mem'] );
            if m > mem and m > 0: mem = m

        if mem > 0: total_mem += mem

        if cores <= 0 or mem <= 0: uncertain = 1

        id = job['name'].split('.')[0]
        user = attrib['Job_Owner'].split('@')[0]

        if args.jobs:
            print "%-16s %-8s %5s %9s" % ( id, user, cores if cores > 0 else '???', show_memory( mem ) )

    cores_avail = maxproc - total_cores
    if cores_avail < 0: cores_avail = 0
    mem_avail = maxmem - total_mem
    if mem_avail < 0: mem_avail = 0

    if args.jobs: print ""
    print "%s%d of %d core%s in use, %d %score%s available " % ( 'AT LEAST ' if uncertain else '', total_cores, maxproc, '' if maxproc == 1 else 's', cores_avail, 'OR FEWER ' if uncertain else '', '' if cores_avail == 1 else 's' )
    print "%s%s of %s memory in use, %s %smemory available " % ( 'AT LEAST ' if uncertain else '', show_memory( total_mem ), show_memory( maxmem ), show_memory( mem_avail ), 'OR LESS ' if uncertain else '' )


if __name__ == "__main__":
    main()