<s> from __future__ import absolute_import <EOL> import logging <EOL> import os <EOL> import re <EOL> import sys <EOL> import subprocess <EOL> import tempfile <EOL> import time <EOL> from Queue import Empty <EOL> from Queue import Queue <EOL> from threading import Thread <EOL> from bd2k . util . iterables import concat <EOL> from bd2k . util . processes import which <EOL> from toil . batchSystems . abstractBatchSystem import BatchSystemSupport <EOL> from toil . lib . bioio import getTempFile <EOL> logger = logging . getLogger ( __name__ ) <EOL> class ParasolBatchSystem ( 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 ( ParasolBatchSystem , self ) . __init__ ( config , maxCores , maxMemory , maxDisk ) <EOL> if maxMemory != sys . maxint : <EOL> logger . warn ( '<STR_LIT>' ) <EOL> command = config . parasolCommand <EOL> if os . path . sep not in command : <EOL> try : <EOL> command = next ( which ( command ) ) <EOL> except StopIteration : <EOL> raise RuntimeError ( "<STR_LIT>" % command ) <EOL> logger . info ( '<STR_LIT>' , command ) <EOL> self . parasolCommand = command <EOL> self . parasolResultsDir = tempfile . mkdtemp ( dir = config . jobStore ) <EOL> self . resultsFiles = dict ( ) <EOL> self . maxBatches = config . parasolMaxBatches <EOL> self . cpuUsageQueue = Queue ( ) <EOL> self . updatedJobsQueue = Queue ( ) <EOL> self . running = True <EOL> self . worker = Thread ( target = self . updatedJobWorker , args = ( ) ) <EOL> self . worker . start ( ) <EOL> self . usedCpus = <NUM_LIT:0> <EOL> self . jobIDsToCpu = { } <EOL> self . runningJobs = set ( ) <EOL> def _runParasol ( self , command , autoRetry = True ) : <EOL> """<STR_LIT>""" <EOL> command = list ( concat ( self . parasolCommand , command ) ) <EOL> while True : <EOL> logger . debug ( '<STR_LIT>' , command ) <EOL> process = subprocess . Popen ( command , <EOL> stdout = subprocess . PIPE , <EOL> stderr = subprocess . PIPE , <EOL> bufsize = - <NUM_LIT:1> ) <EOL> stdout , stderr = process . communicate ( ) <EOL> status = process . wait ( ) <EOL> for line in stderr . split ( '<STR_LIT:\n>' ) : <EOL> if line : logger . warn ( line ) <EOL> if status == <NUM_LIT:0> : <EOL> return <NUM_LIT:0> , stdout . split ( '<STR_LIT:\n>' ) <EOL> message = '<STR_LIT>' % ( command , status ) <EOL> if autoRetry : <EOL> logger . warn ( message ) <EOL> else : <EOL> logger . error ( message ) <EOL> return status , None <EOL> logger . warn ( '<STR_LIT>' ) <EOL> time . sleep ( <NUM_LIT:10> ) <EOL> parasolOutputPattern = re . compile ( "<STR_LIT>" ) <EOL> def issueBatchJob ( self , command , memory , cores , disk , preemptable ) : <EOL> """<STR_LIT>""" <EOL> self . checkResourceRequest ( memory , cores , disk ) <EOL> MiB = <NUM_LIT:1> << <NUM_LIT:20> <EOL> truncatedMemory = ( memory / MiB ) * MiB <EOL> if len ( self . resultsFiles ) >= self . maxBatches : <EOL> raise RuntimeError ( '<STR_LIT>' % self . maxBatches ) <EOL> try : <EOL> results = self . resultsFiles [ ( truncatedMemory , cores ) ] <EOL> except KeyError : <EOL> results = getTempFile ( rootDir = self . parasolResultsDir ) <EOL> self . resultsFiles [ ( truncatedMemory , cores ) ] = results <EOL> command = '<STR_LIT:U+0020>' . join ( concat ( '<STR_LIT>' , self . __environment ( ) , command ) ) <EOL> parasolCommand = [ '<STR_LIT>' , <EOL> '<STR_LIT>' % memory , <EOL> '<STR_LIT>' % cores , <EOL> '<STR_LIT>' + results , <EOL> '<STR_LIT>' , '<STR_LIT>' , command ] <EOL> self . usedCpus += cores <EOL> while True : <EOL> try : <EOL> jobID = self . cpuUsageQueue . get_nowait ( ) <EOL> except Empty : <EOL> break <EOL> if jobID in self . jobIDsToCpu . keys ( ) : <EOL> self . usedCpus -= self . jobIDsToCpu . pop ( jobID ) <EOL> assert self . usedCpus >= <NUM_LIT:0> <EOL> while self . usedCpus > self . maxCores : <EOL> jobID = self . cpuUsageQueue . get ( ) <EOL> if jobID in self . jobIDsToCpu . keys ( ) : <EOL> self . usedCpus -= self . jobIDsToCpu . pop ( jobID ) <EOL> assert self . usedCpus >= <NUM_LIT:0> <EOL> while True : <EOL> line = self . _runParasol ( parasolCommand ) [ <NUM_LIT:1> ] [ <NUM_LIT:0> ] <EOL> match = self . parasolOutputPattern . match ( line ) <EOL> if match is None : <EOL> logger . info ( '<STR_LIT>' ) <EOL> time . sleep ( <NUM_LIT:5> ) <EOL> else : <EOL> jobID = int ( match . group ( <NUM_LIT:1> ) ) <EOL> self . jobIDsToCpu [ jobID ] = cores <EOL> self . runningJobs . add ( jobID ) <EOL> logger . debug ( "<STR_LIT>" % ( jobID , line ) ) <EOL> return jobID <EOL> def setEnv ( self , name , value = None ) : <EOL> if value and '<STR_LIT:U+0020>' in value : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> return super ( ParasolBatchSystem , self ) . setEnv ( name , value ) <EOL> def __environment ( self ) : <EOL> return ( k + '<STR_LIT:=>' + ( os . environ [ k ] if v is None else v ) for k , v in self . environment . items ( ) ) <EOL> def killBatchJobs ( self , jobIDs ) : <EOL> """<STR_LIT>""" <EOL> while True : <EOL> for jobID in jobIDs : <EOL> if jobID in self . runningJobs : <EOL> self . runningJobs . remove ( jobID ) <EOL> exitValue = self . _runParasol ( [ '<STR_LIT>' , '<STR_LIT>' , str ( jobID ) ] , <EOL> autoRetry = False ) [ <NUM_LIT:0> ] <EOL> logger . info ( "<STR_LIT>" % ( jobID , exitValue ) ) <EOL> runningJobs = self . getIssuedBatchJobIDs ( ) <EOL> if set ( jobIDs ) . difference ( set ( runningJobs ) ) == set ( jobIDs ) : <EOL> break <EOL> logger . warn ( '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> time . sleep ( <NUM_LIT:5> ) <EOL> for jobID in jobIDs : <EOL> if jobID in self . jobIDsToCpu . keys ( ) : <EOL> self . usedCpus -= self . jobIDsToCpu . pop ( jobID ) <EOL> queuePattern = re . compile ( r'<STR_LIT>' ) <EOL> runningPattern = re . compile ( r'<STR_LIT>' ) <EOL> def getJobIDsForResultsFile ( self , resultsFile ) : <EOL> """<STR_LIT>""" <EOL> jobIDs = [ ] <EOL> for line in self . _runParasol ( [ '<STR_LIT>' + resultsFile , '<STR_LIT>' ] ) [ <NUM_LIT:1> ] : <EOL> runningJobMatch = self . runningPattern . match ( line ) <EOL> queuedJobMatch = self . queuePattern . match ( line ) <EOL> if runningJobMatch : <EOL> jobID = runningJobMatch . group ( <NUM_LIT:1> ) <EOL> elif queuedJobMatch : <EOL> jobID = queuedJobMatch . group ( <NUM_LIT:1> ) <EOL> else : <EOL> continue <EOL> jobIDs . append ( int ( jobID ) ) <EOL> return set ( jobIDs ) <EOL> def getIssuedBatchJobIDs ( self ) : <EOL> """<STR_LIT>""" <EOL> issuedJobs = set ( ) <EOL> for resultsFile in self . resultsFiles . itervalues ( ) : <EOL> issuedJobs . update ( self . getJobIDsForResultsFile ( resultsFile ) ) <EOL> return list ( issuedJobs ) <EOL> def getRunningBatchJobIDs ( self ) : <EOL> """<STR_LIT>""" <EOL> runningJobs = { } <EOL> issuedJobs = self . getIssuedBatchJobIDs ( ) <EOL> for line in self . _runParasol ( [ '<STR_LIT>' ] ) [ <NUM_LIT:1> ] : <EOL> if line != '<STR_LIT>' : <EOL> match = self . runningPattern . match ( line ) <EOL> if match is not None : <EOL> jobID = int ( match . group ( <NUM_LIT:1> ) ) <EOL> startTime = int ( match . group ( <NUM_LIT:2> ) ) <EOL> if jobID in issuedJobs : <EOL> runningJobs [ jobID ] = time . time ( ) - startTime <EOL> return runningJobs <EOL> def getUpdatedBatchJob ( self , maxWait ) : <EOL> while True : <EOL> try : <EOL> jobID , status , wallTime = self . updatedJobsQueue . get ( timeout = maxWait ) <EOL> except Empty : <EOL> return None <EOL> try : <EOL> self . runningJobs . remove ( jobID ) <EOL> except KeyError : <EOL> pass <EOL> else : <EOL> return jobID , status , wallTime <EOL> @ classmethod <EOL> def getRescueBatchJobFrequency ( cls ) : <EOL> """<STR_LIT>""" <EOL> return <NUM_LIT> <EOL> def updatedJobWorker ( self ) : <EOL> """<STR_LIT>""" <EOL> resultsFiles = set ( ) <EOL> resultsFileHandles = [ ] <EOL> try : <EOL> while self . running : <EOL> newResultsFiles = set ( os . listdir ( self . parasolResultsDir ) ) . difference ( resultsFiles ) <EOL> for newFile in newResultsFiles : <EOL> newFilePath = os . path . join ( self . parasolResultsDir , newFile ) <EOL> resultsFileHandles . append ( open ( newFilePath , '<STR_LIT:r>' ) ) <EOL> resultsFiles . add ( newFile ) <EOL> for fileHandle in resultsFileHandles : <EOL> while self . running : <EOL> line = fileHandle . readline ( ) <EOL> if not line : <EOL> break <EOL> assert line [ - <NUM_LIT:1> ] == '<STR_LIT:\n>' <EOL> ( status , host , jobId , exe , usrTicks , sysTicks , submitTime , startTime , <EOL> endTime , user , errFile , command ) = line [ : - <NUM_LIT:1> ] . split ( None , <NUM_LIT:11> ) <EOL> status = int ( status ) <EOL> jobId = int ( jobId ) <EOL> if os . WIFEXITED ( status ) : <EOL> status = os . WEXITSTATUS ( status ) <EOL> else : <EOL> status = - status <EOL> self . cpuUsageQueue . put ( jobId ) <EOL> startTime = int ( startTime ) <EOL> endTime = int ( endTime ) <EOL> if endTime == startTime : <EOL> usrTicks = int ( usrTicks ) <EOL> sysTicks = int ( sysTicks ) <EOL> wallTime = float ( max ( <NUM_LIT:1> , usrTicks + sysTicks ) ) * <NUM_LIT> <EOL> else : <EOL> wallTime = float ( endTime - startTime ) <EOL> self . updatedJobsQueue . put ( ( jobId , status , wallTime ) ) <EOL> time . sleep ( <NUM_LIT:1> ) <EOL> except : <EOL> logger . warn ( "<STR_LIT>" ) <EOL> raise <EOL> finally : <EOL> for fileHandle in resultsFileHandles : <EOL> fileHandle . close ( ) <EOL> def shutdown ( self ) : <EOL> self . killBatchJobs ( self . getIssuedBatchJobIDs ( ) ) <EOL> for results in self . resultsFiles . itervalues ( ) : <EOL> exitValue = self . _runParasol ( [ '<STR_LIT>' + results , '<STR_LIT>' , '<STR_LIT>' ] , <EOL> autoRetry = False ) [ <NUM_LIT:0> ] <EOL> if exitValue is not None : <EOL> logger . warn ( "<STR_LIT>" % results ) <EOL> exitValue = self . _runParasol ( [ '<STR_LIT>' + results , '<STR_LIT>' ] , <EOL> autoRetry = False ) [ <NUM_LIT:0> ] <EOL> if exitValue is not None : <EOL> logger . warn ( "<STR_LIT>" % results ) <EOL> self . running = False <EOL> logger . debug ( '<STR_LIT>' ) <EOL> self . worker . join ( ) <EOL> logger . debug ( '<STR_LIT>' ) <EOL> for results in self . resultsFiles . values ( ) : <EOL> os . remove ( results ) <EOL> os . rmdir ( self . parasolResultsDir ) </s>
