#!/usr/bin/python
# run slurm jobs

import re
import os
import string

def run(jobname,nodelist):
    nodelisttext = string.join(nodelist,',')
    cmd = "/usr/local/bin/scontrol update JobId=%s Priority=1 ReqNodeList=%s" % (jobname,nodelisttext)
    #returntext = os.popen(cmd).read()
    #return returntext
    return cmd

def cancel(jobname):
    cmd = "/usr/local/bin/scancel %s" % (jobname,)
    #returntext = os.popen(cmd).read()
    #return returntext
    return cmd

def getnodes():
    #NodeName=popocatzin5 Arch=x86_64 CoresPerSocket=1 CPUAlloc=8 CPUErr=0 CPUTot=8 Features=(null) OS=Linux RealMemory=1 Sockets=8 State=ALLOCATED ThreadsPerCore=1 TmpDisk=0 Weight=1 Reason=(null)
    #NodeName=popocatzin6 Arch=x86_64 CoresPerSocket=1 CPUAlloc=0 CPUErr=0 CPUTot=8 Features=(null) OS=Linux RealMemory=1 Sockets=8 State=IDLE ThreadsPerCore=1 TmpDisk=0 Weight=1 Reason=(null)
    nodepat = r"NodeName=(?P<nodename>\S+)\s+.*State=(?P<state>\S+)\s+.*"
    node_reo = re.compile(nodepat)
    cmd = "/usr/local/bin/scontrol --oneliner show node"
    returnlines = os.popen(cmd).readlines()
    nodedict = {}
    for returnline in returnlines:
        node_mo = node_reo.match(returnline)
        if node_mo != None:
            nodedict[node_mo.group('nodename')] = {'state' : node_mo.group('state'),
                                                  }
    return nodedict

def getjobs():
    #JobId=64 Name=r.script UserId=lyn(1113) GroupId=lyn(1113) Priority=0 Account=(null) QOS=(null) JobState=RUNNING Reason=None Dependency=(null) TimeLimit=01:00:00 Requeue=1 Restarts=0 BatchFlag=1 ExitCode=0:0 SubmitTime=2010-12-02T17:57:36 EligibleTime=2010-12-02T17:57:36 StartTime=2010-12-02T18:19:02 EndTime=2010-12-02T19:19:02 SuspendTime=None SecsPreSuspend=0 Partition=trabajoslyn AllocNode:Sid=popocatzin:23871 ReqNodeList=popocatzin3,popocatzin4 ExcNodeList=(null) NodeList=popocatzin[3-4] NumNodes=2 NumCPUs=16 CPUs/Task=1 ReqS:C:T=65534:65534:65534 MinCPUsNode=8 MinMemoryNode=0 MinTmpDiskNode=0 Features=(null) Reservation=(null) Shared=0 Contiguous=0 Licenses=(null) Network=(null) Command=/home/lyn/jobscripts/r.script -o=/home/lyn/job*/output/r.script.%t%j.out WorkDir=/home/lyn/jobscripts
    #JobId=65 Name=r.script UserId=lyn(1113) GroupId=lyn(1113) Priority=0 Account=(null) QOS=(null) JobState=PENDING Reason=JobHeld Dependency=(null) TimeLimit=01:00:00 Requeue=1 Restarts=0 BatchFlag=1 ExitCode=0:0 SubmitTime=2010-12-02T17:57:38 EligibleTime=2010-12-02T17:57:38 StartTime=Unknown EndTime=Unknown SuspendTime=None SecsPreSuspend=0 Partition=trabajoslyn AllocNode:Sid=popocatzin:23871 ReqNodeList=(null) ExcNodeList=(null) NodeList=(null) NumNodes=2-2 NumCPUs=16 CPUs/Task=1 ReqS:C:T=65534:65534:65534 MinCPUsNode=8 MinMemoryNode=0 MinTmpDiskNode=0 Features=(null) Reservation=(null) Shared=OK Contiguous=0 Licenses=(null) Network=(null) Command=/home/lyn/jobscripts/r.script -o=/home/lyn/job*/output/r.script.%t%j.out WorkDir=/home/lyn/jobscripts
    jobpat = r"JobId=(?P<jobid>\S+)\s+.*JobState=(?P<jobstate>\S+)\s+.*NumNodes=(?P<numnodes>\S+)\s+.*"
    job_reo = re.compile(jobpat)
    cmd = "/usr/local/bin/scontrol --oneliner show job"
    returnlines = os.popen(cmd).readlines()
    jobdict = {}
    for returnline in returnlines:
        job_mo = job_reo.match(returnline)
        if job_mo != None:
            jobdict[job_mo.group('jobid')] = { 'jobstate' : job_mo.group('jobstate'),
                                               'numnodes' : job_mo.group('numnodes')
                                             }
    return jobdict

nodedict = getnodes()
jobdict = getjobs()

idlenodeslist = []
for nodename in nodedict.keys():
    if nodedict[nodename]['state'] == 'IDLE':
        idlenodeslist.append(nodename)

# additional code to filter and sort nodes
# can be placed here

idlejobslist = []
for jobid in jobdict.keys():
    if jobdict[jobid]['jobstate'] == 'PENDING':
        idlejobslist.append(jobid)

# additional code to filter and sort jobs
# can be placed here

for idlejob in idlejobslist:
    minnodes_string,maxnodes_string = string.split(jobdict[idlejob]['numnodes'],'-')
    minnodes = int(minnodes_string)
    if minnodes <= len(idlenodeslist):
        nodelist = idlenodeslist[:minnodes]
        newidlenodeslist = idlenodeslist[minnodes:]
        return_text = run(idlejob,nodelist)
        print return_text
        idlenodeslist = newidlenodeslist