<s> from __future__ import absolute_import <EOL> import logging <EOL> import os <EOL> from pipes import quote <EOL> import subprocess <EOL> import time <EOL> import math <EOL> from Queue import Queue , Empty <EOL> from threading import Thread <EOL> from toil . batchSystems import MemoryString <EOL> from toil . batchSystems . abstractBatchSystem import BatchSystemSupport <EOL> logger = logging . getLogger ( __name__ ) <EOL> sleepSeconds = <NUM_LIT:1> <EOL> class Worker ( Thread ) : <EOL> def __init__ ( self , newJobsQueue , updatedJobsQueue , killQueue , killedJobsQueue , boss ) : <EOL> Thread . __init__ ( self ) <EOL> self . newJobsQueue = newJobsQueue <EOL> self . updatedJobsQueue = updatedJobsQueue <EOL> self . killQueue = killQueue <EOL> self . killedJobsQueue = killedJobsQueue <EOL> self . waitingJobs = list ( ) <EOL> self . runningJobs = set ( ) <EOL> self . boss = boss <EOL> self . allocatedCpus = dict ( ) <EOL> self . sgeJobIDs = dict ( ) <EOL> def getRunningJobIDs ( self ) : <EOL> times = { } <EOL> currentjobs = dict ( ( str ( self . sgeJobIDs [ x ] [ <NUM_LIT:0> ] ) , x ) for x in self . runningJobs ) <EOL> process = subprocess . Popen ( [ "<STR_LIT>" ] , stdout = subprocess . PIPE ) <EOL> stdout , stderr = process . communicate ( ) <EOL> for currline in stdout . split ( '<STR_LIT:\n>' ) : <EOL> items = currline . strip ( ) . split ( ) <EOL> if items : <EOL> if items [ <NUM_LIT:0> ] in currentjobs and items [ <NUM_LIT:4> ] == '<STR_LIT:r>' : <EOL> jobstart = "<STR_LIT:U+0020>" . join ( items [ <NUM_LIT:5> : <NUM_LIT:7> ] ) <EOL> jobstart = time . mktime ( time . strptime ( jobstart , "<STR_LIT>" ) ) <EOL> times [ currentjobs [ items [ <NUM_LIT:0> ] ] ] = time . time ( ) - jobstart <EOL> return times <EOL> def getSgeID ( self , jobID ) : <EOL> if not jobID in self . sgeJobIDs : <EOL> RuntimeError ( "<STR_LIT>" ) <EOL> ( job , task ) = self . sgeJobIDs [ jobID ] <EOL> if task is None : <EOL> return str ( job ) <EOL> else : <EOL> return str ( job ) + "<STR_LIT:.>" + str ( task ) <EOL> def forgetJob ( self , jobID ) : <EOL> self . runningJobs . remove ( jobID ) <EOL> del self . allocatedCpus [ jobID ] <EOL> del self . sgeJobIDs [ jobID ] <EOL> def killJobs ( self ) : <EOL> killList = list ( ) <EOL> while True : <EOL> try : <EOL> jobId = self . killQueue . get ( block = False ) <EOL> except Empty : <EOL> break <EOL> else : <EOL> killList . append ( jobId ) <EOL> if not killList : <EOL> return False <EOL> for jobID in list ( killList ) : <EOL> if jobID in self . runningJobs : <EOL> logger . debug ( '<STR_LIT>' , jobID ) <EOL> subprocess . check_call ( [ '<STR_LIT>' , self . getSgeID ( jobID ) ] ) <EOL> else : <EOL> if jobID in self . waitingJobs : <EOL> self . waitingJobs . remove ( jobID ) <EOL> self . killedJobsQueue . put ( jobID ) <EOL> killList . remove ( jobID ) <EOL> while killList : <EOL> for jobID in list ( killList ) : <EOL> if self . getJobExitCode ( self . sgeJobIDs [ jobID ] ) is not None : <EOL> logger . debug ( '<STR_LIT>' , jobID ) <EOL> self . killedJobsQueue . put ( jobID ) <EOL> killList . remove ( jobID ) <EOL> self . forgetJob ( jobID ) <EOL> if len ( killList ) > <NUM_LIT:0> : <EOL> logger . warn ( "<STR_LIT>" , sleepSeconds ) <EOL> time . sleep ( sleepSeconds ) <EOL> return True <EOL> def createJobs ( self , newJob ) : <EOL> activity = False <EOL> if newJob is not None : <EOL> self . waitingJobs . append ( newJob ) <EOL> while ( len ( self . waitingJobs ) > <NUM_LIT:0> <EOL> and sum ( self . allocatedCpus . values ( ) ) < int ( self . boss . maxCores ) ) : <EOL> activity = True <EOL> jobID , cpu , memory , command = self . waitingJobs . pop ( <NUM_LIT:0> ) <EOL> qsubline = self . prepareQsub ( cpu , memory , jobID ) + [ command ] <EOL> sgeJobID = self . qsub ( qsubline ) <EOL> self . sgeJobIDs [ jobID ] = ( sgeJobID , None ) <EOL> self . runningJobs . add ( jobID ) <EOL> self . allocatedCpus [ jobID ] = cpu <EOL> return activity <EOL> def checkOnJobs ( self ) : <EOL> activity = False <EOL> logger . debug ( '<STR_LIT>' , self . runningJobs ) <EOL> for jobID in list ( self . runningJobs ) : <EOL> status = self . getJobExitCode ( self . sgeJobIDs [ jobID ] ) <EOL> if status is not None : <EOL> activity = True <EOL> self . updatedJobsQueue . put ( ( jobID , status ) ) <EOL> self . forgetJob ( jobID ) <EOL> return activity <EOL> def run ( self ) : <EOL> while True : <EOL> activity = False <EOL> newJob = None <EOL> if not self . newJobsQueue . empty ( ) : <EOL> activity = True <EOL> newJob = self . newJobsQueue . get ( ) <EOL> if newJob is None : <EOL> logger . debug ( '<STR_LIT>' ) <EOL> break <EOL> activity |= self . killJobs ( ) <EOL> activity |= self . createJobs ( newJob ) <EOL> activity |= self . checkOnJobs ( ) <EOL> if not activity : <EOL> logger . debug ( '<STR_LIT>' , sleepSeconds ) <EOL> time . sleep ( sleepSeconds ) <EOL> def prepareQsub ( self , cpu , mem , jobID ) : <EOL> qsubline = [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT:y>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT:y>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' + str ( jobID ) ] <EOL> if self . boss . environment : <EOL> qsubline . append ( '<STR_LIT>' ) <EOL> qsubline . append ( '<STR_LIT:U+002C>' . join ( k + '<STR_LIT:=>' + quote ( os . environ [ k ] if v is None else v ) <EOL> for k , v in self . boss . environment . iteritems ( ) ) ) <EOL> reqline = list ( ) <EOL> if mem is not None : <EOL> memStr = str ( mem / <NUM_LIT> ) + '<STR_LIT>' <EOL> reqline += [ '<STR_LIT>' + memStr , '<STR_LIT>' + memStr ] <EOL> if len ( reqline ) > <NUM_LIT:0> : <EOL> qsubline . extend ( [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT:U+002C>' . join ( reqline ) ] ) <EOL> if cpu is not None and math . ceil ( cpu ) > <NUM_LIT:1> : <EOL> peConfig = os . getenv ( '<STR_LIT>' ) or '<STR_LIT>' <EOL> qsubline . extend ( [ '<STR_LIT>' , peConfig , str ( int ( math . ceil ( cpu ) ) ) ] ) <EOL> return qsubline <EOL> def qsub ( self , qsubline ) : <EOL> logger . debug ( "<STR_LIT>" , qsubline ) <EOL> process = subprocess . Popen ( qsubline , stdout = subprocess . PIPE ) <EOL> result = int ( process . stdout . readline ( ) . strip ( ) . split ( '<STR_LIT:.>' ) [ <NUM_LIT:0> ] ) <EOL> return result <EOL> def getJobExitCode ( self , sgeJobID ) : <EOL> job , task = sgeJobID <EOL> args = [ "<STR_LIT>" , "<STR_LIT>" , str ( job ) ] <EOL> if task is not None : <EOL> args . extend ( [ "<STR_LIT>" , str ( task ) ] ) <EOL> process = subprocess . Popen ( args , stdout = subprocess . PIPE , stderr = subprocess . STDOUT ) <EOL> for line in process . stdout : <EOL> if line . startswith ( "<STR_LIT>" ) and int ( line . split ( ) [ <NUM_LIT:1> ] ) == <NUM_LIT:1> : <EOL> return <NUM_LIT:1> <EOL> elif line . startswith ( "<STR_LIT>" ) : <EOL> logger . debug ( '<STR_LIT>' , line . split ( ) [ <NUM_LIT:1> ] ) <EOL> return int ( line . split ( ) [ <NUM_LIT:1> ] ) <EOL> return None <EOL> class GridengineBatchSystem ( BatchSystemSupport ) : <EOL> """<STR_LIT>""" <EOL> @ classmethod <EOL> def supportsWorkerCleanup ( cls ) : <EOL> return False <EOL> @ classmethod <EOL> def supportsHotDeployment ( cls ) : <EOL> return False <EOL> def __init__ ( self , config , maxCores , maxMemory , maxDisk ) : <EOL> super ( GridengineBatchSystem , self ) . __init__ ( config , maxCores , maxMemory , maxDisk ) <EOL> self . gridengineResultsFile = self . _getResultsFileName ( config . jobStore ) <EOL> self . gridengineResultsFileHandle = open ( self . gridengineResultsFile , '<STR_LIT:w>' ) <EOL> self . gridengineResultsFileHandle . close ( ) <EOL> self . currentJobs = set ( ) <EOL> self . maxCPU , self . maxMEM = self . obtainSystemConstants ( ) <EOL> self . nextJobID = <NUM_LIT:0> <EOL> self . newJobsQueue = Queue ( ) <EOL> self . updatedJobsQueue = Queue ( ) <EOL> self . killQueue = Queue ( ) <EOL> self . killedJobsQueue = Queue ( ) <EOL> self . worker = Worker ( self . newJobsQueue , self . updatedJobsQueue , self . killQueue , <EOL> self . killedJobsQueue , self ) <EOL> self . worker . start ( ) <EOL> def __des__ ( self ) : <EOL> self . gridengineResultsFileHandle . close ( ) <EOL> def issueBatchJob ( self , command , memory , cores , disk , preemptable ) : <EOL> self . checkResourceRequest ( memory , cores , disk ) <EOL> jobID = self . nextJobID <EOL> self . nextJobID += <NUM_LIT:1> <EOL> self . currentJobs . add ( jobID ) <EOL> self . newJobsQueue . put ( ( jobID , cores , memory , command ) ) <EOL> logger . debug ( "<STR_LIT>" , command , str ( jobID ) ) <EOL> return jobID <EOL> def killBatchJobs ( self , jobIDs ) : <EOL> """<STR_LIT>""" <EOL> jobIDs = set ( jobIDs ) <EOL> logger . debug ( '<STR_LIT>' , jobIDs ) <EOL> for jobID in jobIDs : <EOL> self . killQueue . put ( jobID ) <EOL> while jobIDs : <EOL> killedJobId = self . killedJobsQueue . get ( ) <EOL> if killedJobId is None : <EOL> break <EOL> jobIDs . remove ( killedJobId ) <EOL> if killedJobId in self . currentJobs : <EOL> self . currentJobs . remove ( killedJobId ) <EOL> if jobIDs : <EOL> logger . debug ( '<STR_LIT>' , len ( jobIDs ) , <EOL> sleepSeconds ) <EOL> time . sleep ( sleepSeconds ) <EOL> def getIssuedBatchJobIDs ( self ) : <EOL> """<STR_LIT>""" <EOL> return list ( self . currentJobs ) <EOL> def getRunningBatchJobIDs ( self ) : <EOL> return self . worker . getRunningJobIDs ( ) <EOL> def getUpdatedBatchJob ( self , maxWait ) : <EOL> try : <EOL> item = self . updatedJobsQueue . get ( timeout = maxWait ) <EOL> except Empty : <EOL> return None <EOL> logger . debug ( '<STR_LIT>' , item ) <EOL> jobID , retcode = item <EOL> self . currentJobs . remove ( jobID ) <EOL> return jobID , retcode , None <EOL> def shutdown ( self ) : <EOL> """<STR_LIT>""" <EOL> newJobsQueue = self . newJobsQueue <EOL> self . newJobsQueue = None <EOL> newJobsQueue . put ( None ) <EOL> self . worker . join ( ) <EOL> def getWaitDuration ( self ) : <EOL> return <NUM_LIT:0.0> <EOL> @ classmethod <EOL> def getRescueBatchJobFrequency ( cls ) : <EOL> return <NUM_LIT:30> * <NUM_LIT> <EOL> @ staticmethod <EOL> def obtainSystemConstants ( ) : <EOL> lines = filter ( None , map ( str . strip , subprocess . check_output ( [ "<STR_LIT>" ] ) . split ( '<STR_LIT:\n>' ) ) ) <EOL> line = lines [ <NUM_LIT:0> ] <EOL> items = line . strip ( ) . split ( ) <EOL> num_columns = len ( items ) <EOL> cpu_index = None <EOL> mem_index = None <EOL> for i in range ( num_columns ) : <EOL> if items [ i ] == '<STR_LIT>' : <EOL> cpu_index = i <EOL> elif items [ i ] == '<STR_LIT>' : <EOL> mem_index = i <EOL> if cpu_index is None or mem_index is None : <EOL> RuntimeError ( '<STR_LIT>' ) <EOL> maxCPU = <NUM_LIT:0> <EOL> maxMEM = MemoryString ( "<STR_LIT:0>" ) <EOL> for line in lines [ <NUM_LIT:2> : ] : <EOL> items = line . strip ( ) . split ( ) <EOL> if len ( items ) < num_columns : <EOL> RuntimeError ( '<STR_LIT>' ) <EOL> if items [ cpu_index ] != '<STR_LIT:->' and items [ cpu_index ] > maxCPU : <EOL> maxCPU = items [ cpu_index ] <EOL> if items [ mem_index ] != '<STR_LIT:->' and MemoryString ( items [ mem_index ] ) > maxMEM : <EOL> maxMEM = MemoryString ( items [ mem_index ] ) <EOL> if maxCPU is <NUM_LIT:0> or maxMEM is <NUM_LIT:0> : <EOL> RuntimeError ( '<STR_LIT>' ) <EOL> return maxCPU , maxMEM <EOL> def setEnv ( self , name , value = None ) : <EOL> if value and '<STR_LIT:U+002C>' in value : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> return super ( GridengineBatchSystem , self ) . setEnv ( name , value ) </s>
