Cluster Cosmology

#!/usr/bin/env python

'''

Created on March 23rd 2016

@author: mkuemmel@usm.lmu.de

Date:     $Date: 2018-05-24 10:25:02 +0200 (Thu, 24 May 2018) $

Revision: $Revision: 340 $

Author:   $Author: Martin.Kuemmel $

'''

import os.path

import subprocess

def theNodeStatus():

    # run the 'sinfo' command

    # to list all nodes (no quotes necessary for format string)

    comSeq = ['sinfo', '-p', 'usm-cl', '-o', '%n  %t %e %m %O %C']

    allReturns = execCommand('sinfo', comSeq, prompt='execCommand>')

    if allReturns['retcode']:

        errMsg = 'Command: "%s" finished with exit code: %i!' % (comSeq, allReturns['retcode'])

        raise Exception(errMsg)

    # read in the 

    allNodes = {}

    lineIndex=0

    for aLine in open(allReturns['stdout']):

        lineIndex += 1

        # skip the first line

        if lineIndex < 2:

            continue

        # isolate the items

        lineItems = aLine.strip().split()

        # get key and state

        key     = lineItems[0]

        state   = lineItems[1]

        # get the other info        

        if state != 'down':

            fracMem = 100.0*float(lineItems[2])/float(lineItems[3])

            fmStr   = '%4.0f%s' % (fracMem, '%')

            load    = float(lineItems[4])

            cpuUsed = '%s/%s'%(lineItems[5].split("/")[0],lineItems[5].split("/")[-1])

        else:

            fracMem = '---'

            fmStr   = '---'

            load    = '---'

            cpuUsed = '---'

        # append the info to the keylist

        allNodes[key] = {'state': state, 'load': load, 'mem': fmStr, 'cpu': cpuUsed}

    # destroy the output files

    os.unlink(allReturns['stdout'])

    os.unlink(allReturns['stderr'])

    # return the keys

    return allNodes

def execCommand(command, commSequence, prompt='execCommand>', verbose=True):

    """

    """

    # make stdout and stderr files ready

    stdoutFile = '%s.out' % command

    stderrFile = '%s.err' % command

    if os.path.isfile(stdoutFile):

        os.unlink(stdoutFile)

    if os.path.isfile(stderrFile):

        os.unlink(stderrFile)

    sout = open(stdoutFile, "w+")

    serr = open(stderrFile, "w+")

    # print the entire command

    if verbose:

        allCommand = ' '.join(commSequence)

        serr.write(allCommand+'\n')

        serr.flush()

    # run the command

    retcode = subprocess.call(commSequence, shell=False, env=None, stdout=sout, stderr=serr )

    # close stdout and stderr

    sout.close()

    serr.close()

    # give back the return

    return {'retcode': retcode, 'stdout': stdoutFile, 'stderr': stderrFile}

def getScontrolList():

    """

    """

    # run the 'scontrol' command

    # to list all jobs

    comSeq = ['scontrol', 'show', 'job']

    allReturns = execCommand('scontrol', comSeq, prompt='execCommand>')

    if allReturns['retcode']:

        errMsg = 'Command: "%s" finished with exit code: %i!' % (comSeq, allReturns['retcode'])

        raise Exception(errMsg)

    # go over all lines

    # of the output files

    nmRunJobs  = []

    nmPendJobs = []

    dbRunJobs  = []

    dbPendJobs = []

    oneJob=None

    for aLine in open(allReturns['stdout']):

        # identify a new job;

        # append the old one, open a new dict

        if aLine.find('JobId')>-1:

            if oneJob:

                if oneJob['JobState']=='RUNNING':

                    if oneJob['Partition'] == 'usm-cl':

                        dbRunJobs.append(oneJob)

                    else:

                        nmRunJobs.append(oneJob)

                elif oneJob['JobState']=='PENDING':

                    if oneJob['Partition'] == 'usm-cl':

                        dbPendJobs.append(oneJob)

                    else:

                        nmPendJobs.append(oneJob)

            oneJob={}

        if aLine.find(',') > -1:

            # split line according to ','

            lineParts = aLine.strip().split(',')

        else:

            # split the line according to blanks

            lineParts = aLine.strip().split()

        # go over the parts

        for aPart in lineParts:

            # identify keywords and append to the dict

            jobParts=aPart.split('=')

            if len(jobParts) < 2:

                continue

            oneJob[jobParts[0]] = jobParts[1]

    # append the last job

    if oneJob:

        if oneJob['JobState']=='RUNNING':

            if oneJob['Partition'] == 'usm-cl':

                dbRunJobs.append(oneJob)

            else:

                nmRunJobs.append(oneJob)

        elif oneJob['JobState']=='PENDING':

            if oneJob['Partition'] == 'usm-cl':

                dbPendJobs.append(oneJob)

            else:

                nmPendJobs.append(oneJob)

    # destroy the output files

    os.unlink(allReturns['stdout'])

    os.unlink(allReturns['stderr'])

    # return the various job lists

    #return {'normalRun': nmRunJobs, 'normalPend' :nmPendJobs, 'lowpriRun': dbRunJobs, 'lowpriPend': dbPendJobs}

    return {'usm-clRun': dbRunJobs, 'usm-clPend': dbPendJobs}

def getUserInfo(jobLists):

    """

    """

    # go over all lists

    allLists = {}

    for aJob in jobLists:

        # skip empty lists

        if len(jobLists[aJob]) < 1:

            continue

        # make a dict for the users    

        allUsers = {}

        for oneJob in jobLists[aJob]:

            # try getting the number of CPU's;

            # get the memory information

            # define the user ID accordingly

            try:

                # get ncpu and user=ID

                nCpu = int(oneJob['NumCPUs'])

                userId = oneJob['UserId']

                # get the memory information, which is different

                # for pending and running jobs

                if 'mem' in oneJob:

                    if oneJob['mem'][-1] == 'G' or oneJob['mem'][-1] == 'M':

                        userMem = float(oneJob['mem'][:-1])

                        unitMem = oneJob['mem'][-1]

                    else:

                        userMem = float(oneJob['mem'])

                        unitMem = 'M'

                elif 'MinMemoryNode' in oneJob:

                    userMem = float(oneJob['MinMemoryNode'][:-1])

                    unitMem = oneJob['MinMemoryNode'][-1]

                elif 'MinMemoryCPU' in oneJob:

                    userMem = float(oneJob['MinMemoryCPU'][:-1])

                    unitMem = oneJob['MinMemoryCPU'][-1]

                # convert to Gb if necessary

                if unitMem == 'M':

                    userMem /= 1000.0

                hasNCpu = True

            except ValueError:

                nCpu = None 

                userId = '%sS' % oneJob['UserId']

                hasNCpu = False

            # append the job to an ID

            # or create a new ID

            if not userId in allUsers:

                if hasNCpu:

                    allUsers[userId] = {'hasNCpu': True, 'njobs': 1, 'ncpu': nCpu, 'memsize': userMem, 'jobids': [oneJob['JobId']]}

                else:

                    allUsers[userId] = {'hasNCpu': False,'njobs': 1, 'ncpu': oneJob['NumCPUs'], 'memsize': oneJob['MinMemoryNode'], 'jobids': [oneJob['JobId']]}

            else:

                if hasNCpu:

                    allUsers[userId]['njobs']   += 1

                    allUsers[userId]['ncpu']    += nCpu

                    allUsers[userId]['memsize'] += userMem

                    allUsers[userId]['jobids'].append(oneJob['JobId'])

                else:

                    allUsers[userId]['njobs']   += 1

                    allUsers[userId]['ncpu']    += '%,'%str(oneJob['NumCPUs'])

                    allUsers[userId]['memsize'] += '%,'%str(oneJob['MinMemoryNode'])

                    allUsers[userId]['jobids'].append(oneJob['JobId'])

        # append the list the combined list

        allLists[aJob] = allUsers

    # return the combined list

    return allLists

def printUserJobList(userJobLists):

    """

    """

    # go over all queues

    for aList in userJobLists:

        # print all users for a queue

        for oneUser in userJobLists[aList]:

            if userJobLists[aList][oneUser]['hasNCpu']:

                print('%10s> %-15s njobs: %3i ncpus: %4i memory: %6.1fG' % (aList, oneUser, userJobLists[aList][oneUser]['njobs'], userJobLists[aList][oneUser]['ncpu'], userJobLists[aList][oneUser]['memsize']))

            else:

                print('%10s> %-15s njobs: %3i ncpus: %s memory: %s' % (aList, oneUser, userJobLists[aList][oneUser]['njobs'], userJobLists[aList][oneUser]['ncpu'], userJobLists[aList][oneUser]['memsize']))

        # print a separator

        print('')

def doScontrol():

    """

    """

    # get the job on all queues via 'scontrol'

    jobLists = getScontrolList()

    # get the user information for all job lists

    userJobLists = getUserInfo(jobLists)

    # print the information

    printUserJobList(userJobLists)

def checkNodes():

    """

    """

    # the list of nodes to check for

    checkNodes=['usm-cl-bt01n1', 'usm-cl-bt01n2', 'usm-cl-bt01n3', 'usm-cl-bt01n4', 'usm-cl-bt02n1', 'usm-cl-bt02n2', 'usm-cl-bt02n3', 'usm-cl-bt02n4']

    allNodes = theNodeStatus()

    print("\n\

+-----------------------------------------------+\n\

+      USM cluster at the physics depart.       +\n\

+-----------------------------------------------+")

    print(str("%15s%8s%8s%8s%8s |" %

              ("Node:     ","Status","Slots"," Mem ","Load")))

    print(str("%15s%8s%8s%8s%8s |" %

              ("----------","------", "-----"," --- ","----")))

    numTot  = 0

    numUsed = 0

    for oneNode in checkNodes:

        if oneNode in allNodes:

            print(str("%15s%8s%8s%8s%8s |" % (oneNode.ljust(10),allNodes[oneNode]['state'],allNodes[oneNode]['cpu'],allNodes[oneNode]['mem'],allNodes[oneNode]['load'])))

            if allNodes[oneNode]['cpu'].find('/') > -1:

                procInfo = allNodes[oneNode]['cpu'].split('/')

                numUsed += int(procInfo[0])

                numTot  += int(procInfo[1])

    print('CPUs used/total:   % 3i/%03i                |' % (numUsed, numTot))

    print("+-----------------------------------------------+\n")

def main():

    checkNodes()

    doScontrol()

if __name__ == '__main__':

    main()