<s> from __future__ import absolute_import <EOL> from abc import ABCMeta , abstractmethod <EOL> from contextlib import contextmanager <EOL> from fractions import Fraction <EOL> from inspect import getsource <EOL> import logging <EOL> import os <EOL> import tempfile <EOL> from textwrap import dedent <EOL> import time <EOL> import multiprocessing <EOL> import sys <EOL> import subprocess <EOL> from unittest import skipIf <EOL> from toil . common import Config <EOL> from toil . batchSystems . mesos . test import MesosTestSupport <EOL> from toil . batchSystems . parasolTestSupport import ParasolTestSupport <EOL> from toil . batchSystems . parasol import ParasolBatchSystem <EOL> from toil . batchSystems . singleMachine import SingleMachineBatchSystem <EOL> from toil . batchSystems . abstractBatchSystem import ( InsufficientSystemResources , <EOL> AbstractBatchSystem , <EOL> BatchSystemSupport ) <EOL> from toil . job import Job <EOL> from toil . test import ToilTest , needs_mesos , needs_parasol , needs_gridengine , needs_slurm <EOL> log = logging . getLogger ( __name__ ) <EOL> numCores = <NUM_LIT:2> <EOL> preemptable = True <EOL> defaultRequirements = dict ( memory = <NUM_LIT> , cores = <NUM_LIT:1> , disk = <NUM_LIT:1000> , preemptable = preemptable ) <EOL> class hidden : <EOL> """<STR_LIT>""" <EOL> class AbstractBatchSystemTest ( ToilTest ) : <EOL> """<STR_LIT>""" <EOL> __metaclass__ = ABCMeta <EOL> @ abstractmethod <EOL> def createBatchSystem ( self ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError <EOL> def supportsWallTime ( self ) : <EOL> return False <EOL> def _createDummyConfig ( self ) : <EOL> return Config ( ) <EOL> @ classmethod <EOL> def setUpClass ( cls ) : <EOL> super ( hidden . AbstractBatchSystemTest , cls ) . setUpClass ( ) <EOL> logging . basicConfig ( level = logging . DEBUG ) <EOL> def setUp ( self ) : <EOL> super ( hidden . AbstractBatchSystemTest , self ) . setUp ( ) <EOL> self . config = self . _createDummyConfig ( ) <EOL> self . batchSystem = self . createBatchSystem ( ) <EOL> self . tempDir = self . _createTempDir ( '<STR_LIT>' ) <EOL> def tearDown ( self ) : <EOL> self . batchSystem . shutdown ( ) <EOL> super ( hidden . AbstractBatchSystemTest , self ) . tearDown ( ) <EOL> def testAvailableCores ( self ) : <EOL> self . assertTrue ( multiprocessing . cpu_count ( ) >= numCores ) <EOL> def testRunJobs ( self ) : <EOL> testPath = os . path . join ( self . tempDir , "<STR_LIT>" ) <EOL> job1 = self . batchSystem . issueBatchJob ( "<STR_LIT>" , ** defaultRequirements ) <EOL> job2 = self . batchSystem . issueBatchJob ( "<STR_LIT>" , ** defaultRequirements ) <EOL> issuedIDs = self . _waitForJobsToIssue ( <NUM_LIT:2> ) <EOL> self . assertEqual ( set ( issuedIDs ) , { job1 , job2 } ) <EOL> runningJobIDs = self . _waitForJobsToStart ( <NUM_LIT:2> ) <EOL> self . assertEqual ( set ( runningJobIDs ) , { job1 , job2 } ) <EOL> self . batchSystem . killBatchJobs ( [ job1 , job2 ] ) <EOL> self . assertEqual ( { } , self . batchSystem . getRunningBatchJobIDs ( ) ) <EOL> self . assertFalse ( os . path . exists ( testPath ) ) <EOL> job3 = self . batchSystem . issueBatchJob ( "<STR_LIT>" % testPath , ** defaultRequirements ) <EOL> updatedID , exitStatus , wallTime = self . batchSystem . getUpdatedBatchJob ( maxWait = <NUM_LIT:1000> ) <EOL> self . assertEqual ( exitStatus , <NUM_LIT:0> ) <EOL> self . assertEqual ( updatedID , job3 ) <EOL> if self . supportsWallTime ( ) : <EOL> self . assertTrue ( wallTime > <NUM_LIT:0> ) <EOL> else : <EOL> self . assertIsNone ( wallTime ) <EOL> self . assertTrue ( os . path . exists ( testPath ) ) <EOL> self . assertFalse ( self . batchSystem . getUpdatedBatchJob ( <NUM_LIT:0> ) ) <EOL> self . batchSystem . killBatchJobs ( [ <NUM_LIT:10> ] ) <EOL> def testSetEnv ( self ) : <EOL> def assertEnv ( ) : <EOL> import os , sys <EOL> sys . exit ( <NUM_LIT:0> if os . getenv ( '<STR_LIT>' ) == '<STR_LIT:bar>' else <NUM_LIT> ) <EOL> script_body = dedent ( '<STR_LIT:\n>' . join ( getsource ( assertEnv ) . split ( '<STR_LIT:\n>' ) [ <NUM_LIT:1> : ] ) ) <EOL> with tempFileContaining ( script_body ) as script_path : <EOL> command = sys . executable + '<STR_LIT:U+0020>' + script_path <EOL> job4 = self . batchSystem . issueBatchJob ( command , ** defaultRequirements ) <EOL> updatedID , exitStatus , wallTime = self . batchSystem . getUpdatedBatchJob ( maxWait = <NUM_LIT:1000> ) <EOL> self . assertEqual ( exitStatus , <NUM_LIT> ) <EOL> self . assertEqual ( updatedID , job4 ) <EOL> self . batchSystem . setEnv ( '<STR_LIT>' , '<STR_LIT:bar>' ) <EOL> job5 = self . batchSystem . issueBatchJob ( command , ** defaultRequirements ) <EOL> updatedID , exitStatus , wallTime = self . batchSystem . getUpdatedBatchJob ( maxWait = <NUM_LIT:1000> ) <EOL> self . assertEqual ( exitStatus , <NUM_LIT:0> ) <EOL> self . assertEqual ( updatedID , job5 ) <EOL> def testCheckResourceRequest ( self ) : <EOL> if isinstance ( self . batchSystem , BatchSystemSupport ) : <EOL> checkResourceRequest = self . batchSystem . checkResourceRequest <EOL> self . assertRaises ( InsufficientSystemResources , checkResourceRequest , <EOL> memory = <NUM_LIT:1000> , cores = <NUM_LIT:200> , disk = <NUM_LIT> ) <EOL> self . assertRaises ( InsufficientSystemResources , checkResourceRequest , <EOL> memory = <NUM_LIT:5> , cores = <NUM_LIT:200> , disk = <NUM_LIT> ) <EOL> self . assertRaises ( InsufficientSystemResources , checkResourceRequest , <EOL> memory = <NUM_LIT> , cores = <NUM_LIT:1> , disk = <NUM_LIT> ) <EOL> self . assertRaises ( InsufficientSystemResources , checkResourceRequest , <EOL> memory = <NUM_LIT:5> , cores = <NUM_LIT:1> , disk = <NUM_LIT> ) <EOL> self . assertRaises ( AssertionError , checkResourceRequest , <EOL> memory = None , cores = <NUM_LIT:1> , disk = <NUM_LIT:1000> ) <EOL> self . assertRaises ( AssertionError , checkResourceRequest , <EOL> memory = <NUM_LIT:10> , cores = None , disk = <NUM_LIT:1000> ) <EOL> checkResourceRequest ( memory = <NUM_LIT:10> , cores = <NUM_LIT:1> , disk = <NUM_LIT:100> ) <EOL> def testGetRescueJobFrequency ( self ) : <EOL> self . assertTrue ( self . batchSystem . getRescueBatchJobFrequency ( ) > <NUM_LIT:0> ) <EOL> def testScalableBatchSystem ( self ) : <EOL> pass <EOL> def _waitForJobsToIssue ( self , numJobs ) : <EOL> issuedIDs = [ ] <EOL> while not len ( issuedIDs ) == numJobs : <EOL> issuedIDs = self . batchSystem . getIssuedBatchJobIDs ( ) <EOL> return issuedIDs <EOL> def _waitForJobsToStart ( self , numJobs ) : <EOL> runningIDs = [ ] <EOL> while not len ( runningIDs ) == numJobs : <EOL> time . sleep ( <NUM_LIT:0.1> ) <EOL> runningIDs = self . batchSystem . getRunningBatchJobIDs ( ) . keys ( ) <EOL> return runningIDs <EOL> @ needs_mesos <EOL> class MesosBatchSystemTest ( hidden . AbstractBatchSystemTest , MesosTestSupport ) : <EOL> """<STR_LIT>""" <EOL> def supportsWallTime ( self ) : <EOL> return True <EOL> def createBatchSystem ( self ) : <EOL> from toil . batchSystems . mesos . batchSystem import MesosBatchSystem <EOL> self . _startMesos ( numCores ) <EOL> return MesosBatchSystem ( config = self . config , <EOL> maxCores = numCores , maxMemory = <NUM_LIT> , maxDisk = <NUM_LIT> , <EOL> masterAddress = '<STR_LIT>' ) <EOL> def tearDown ( self ) : <EOL> self . _stopMesos ( ) <EOL> super ( MesosBatchSystemTest , self ) . tearDown ( ) <EOL> class SingleMachineBatchSystemTest ( hidden . AbstractBatchSystemTest ) : <EOL> """<STR_LIT>""" <EOL> def supportsWallTime ( self ) : <EOL> return True <EOL> def createBatchSystem ( self ) : <EOL> return SingleMachineBatchSystem ( config = self . config , <EOL> maxCores = numCores , maxMemory = <NUM_LIT> , maxDisk = <NUM_LIT> ) <EOL> class MaxCoresSingleMachineBatchSystemTest ( ToilTest ) : <EOL> """<STR_LIT>""" <EOL> @ classmethod <EOL> def setUpClass ( cls ) : <EOL> super ( MaxCoresSingleMachineBatchSystemTest , cls ) . setUpClass ( ) <EOL> logging . basicConfig ( level = logging . DEBUG ) <EOL> def setUp ( self ) : <EOL> super ( MaxCoresSingleMachineBatchSystemTest , self ) . setUp ( ) <EOL> def writeTempFile ( s ) : <EOL> fd , path = tempfile . mkstemp ( ) <EOL> try : <EOL> assert os . write ( fd , s ) == len ( s ) <EOL> except : <EOL> os . unlink ( path ) <EOL> raise <EOL> else : <EOL> return path <EOL> finally : <EOL> os . close ( fd ) <EOL> self . counterPath = writeTempFile ( '<STR_LIT>' ) <EOL> def script ( ) : <EOL> import os , sys , fcntl , time <EOL> def count ( delta ) : <EOL> """<STR_LIT>""" <EOL> fd = os . open ( sys . argv [ <NUM_LIT:1> ] , os . O_RDWR ) <EOL> try : <EOL> fcntl . flock ( fd , fcntl . LOCK_EX ) <EOL> try : <EOL> s = os . read ( fd , <NUM_LIT:10> ) <EOL> value , maxValue = map ( int , s . split ( '<STR_LIT:U+002C>' ) ) <EOL> value += delta <EOL> if value > maxValue : maxValue = value <EOL> os . lseek ( fd , <NUM_LIT:0> , <NUM_LIT:0> ) <EOL> os . ftruncate ( fd , <NUM_LIT:0> ) <EOL> os . write ( fd , '<STR_LIT:U+002C>' . join ( map ( str , ( value , maxValue ) ) ) ) <EOL> finally : <EOL> fcntl . flock ( fd , fcntl . LOCK_UN ) <EOL> finally : <EOL> os . close ( fd ) <EOL> if len ( sys . argv ) < <NUM_LIT:3> : <EOL> count ( <NUM_LIT:1> ) <EOL> try : <EOL> time . sleep ( <NUM_LIT:1> ) <EOL> finally : <EOL> count ( - <NUM_LIT:1> ) <EOL> else : <EOL> count ( int ( sys . argv [ <NUM_LIT:2> ] ) ) <EOL> self . scriptPath = writeTempFile ( dedent ( '<STR_LIT:\n>' . join ( getsource ( script ) . split ( '<STR_LIT:\n>' ) [ <NUM_LIT:1> : ] ) ) ) <EOL> def tearDown ( self ) : <EOL> os . unlink ( self . scriptPath ) <EOL> os . unlink ( self . counterPath ) <EOL> def scriptCommand ( self ) : <EOL> return '<STR_LIT:U+0020>' . join ( [ sys . executable , self . scriptPath , self . counterPath ] ) <EOL> def test ( self ) : <EOL> F = Fraction <EOL> minCores = F ( <NUM_LIT:1> , <NUM_LIT:10> ) <EOL> self . assertEquals ( float ( minCores ) , SingleMachineBatchSystem . minCores ) <EOL> for maxCores in { F ( minCores ) , minCores * <NUM_LIT:10> , F ( <NUM_LIT:1> ) , F ( numCores , <NUM_LIT:2> ) , F ( numCores ) } : <EOL> for coresPerJob in { F ( minCores ) , F ( minCores * <NUM_LIT:10> ) , F ( <NUM_LIT:1> ) , F ( maxCores , <NUM_LIT:2> ) , F ( maxCores ) } : <EOL> for load in ( F ( <NUM_LIT:1> , <NUM_LIT:10> ) , F ( <NUM_LIT:1> ) , F ( <NUM_LIT:10> ) ) : <EOL> jobs = int ( maxCores / coresPerJob * load ) <EOL> if jobs >= <NUM_LIT:1> and minCores <= coresPerJob < maxCores : <EOL> self . assertEquals ( maxCores , float ( maxCores ) ) <EOL> bs = SingleMachineBatchSystem ( config = Config ( ) , <EOL> maxCores = float ( maxCores ) , <EOL> maxMemory = jobs * <NUM_LIT:10> , <EOL> maxDisk = jobs * <NUM_LIT:10> ) <EOL> try : <EOL> jobIds = set ( ) <EOL> for i in range ( <NUM_LIT:0> , int ( jobs ) ) : <EOL> jobIds . add ( bs . issueBatchJob ( command = self . scriptCommand ( ) , <EOL> cores = float ( coresPerJob ) , <EOL> memory = <NUM_LIT:1> , <EOL> disk = <NUM_LIT:1> , <EOL> preemptable = preemptable ) ) <EOL> self . assertEquals ( len ( jobIds ) , jobs ) <EOL> while jobIds : <EOL> job = bs . getUpdatedBatchJob ( maxWait = <NUM_LIT:10> ) <EOL> self . assertIsNotNone ( job ) <EOL> jobId , status , wallTime = job <EOL> self . assertEquals ( status , <NUM_LIT:0> ) <EOL> jobIds . remove ( jobId ) <EOL> finally : <EOL> bs . shutdown ( ) <EOL> concurrentTasks , maxConcurrentTasks = self . getCounters ( ) <EOL> self . assertEquals ( concurrentTasks , <NUM_LIT:0> ) <EOL> log . info ( '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' . format ( ** locals ( ) ) ) <EOL> expectedMaxConcurrentTasks = min ( maxCores / coresPerJob , jobs ) <EOL> self . assertEquals ( maxConcurrentTasks , expectedMaxConcurrentTasks ) <EOL> self . resetCounters ( ) <EOL> def getCounters ( self ) : <EOL> with open ( self . counterPath , '<STR_LIT>' ) as f : <EOL> s = f . read ( ) <EOL> log . info ( '<STR_LIT>' , s ) <EOL> concurrentTasks , maxConcurrentTasks = map ( int , s . split ( '<STR_LIT:U+002C>' ) ) <EOL> return concurrentTasks , maxConcurrentTasks <EOL> def resetCounters ( self ) : <EOL> with open ( self . counterPath , '<STR_LIT:w>' ) as f : <EOL> f . write ( '<STR_LIT>' ) <EOL> @ skipIf ( SingleMachineBatchSystem . numCores < <NUM_LIT:3> , '<STR_LIT>' ) <EOL> def testServices ( self ) : <EOL> options = Job . Runner . getDefaultOptions ( self . _getTestJobStorePath ( ) ) <EOL> options . logDebug = True <EOL> options . maxCores = <NUM_LIT:3> <EOL> self . assertTrue ( options . maxCores <= SingleMachineBatchSystem . numCores ) <EOL> Job . Runner . startToil ( Job . wrapJobFn ( parentJob , self . scriptCommand ( ) ) , options ) <EOL> with open ( self . counterPath , '<STR_LIT>' ) as f : <EOL> s = f . read ( ) <EOL> log . info ( '<STR_LIT>' , s ) <EOL> self . assertEqual ( self . getCounters ( ) , ( <NUM_LIT:0> , <NUM_LIT:3> ) ) <EOL> def parentJob ( job , cmd ) : <EOL> job . addChildJobFn ( childJob , cmd ) <EOL> def childJob ( job , cmd ) : <EOL> job . addService ( Service ( cmd ) ) <EOL> job . addChildJobFn ( grandChildJob , cmd ) <EOL> subprocess . check_call ( cmd , shell = True ) <EOL> def grandChildJob ( job , cmd ) : <EOL> job . addService ( Service ( cmd ) ) <EOL> job . addChildFn ( greatGrandChild , cmd ) <EOL> subprocess . check_call ( cmd , shell = True ) <EOL> def greatGrandChild ( cmd ) : <EOL> subprocess . check_call ( cmd , shell = True ) <EOL> class Service ( Job . Service ) : <EOL> def __init__ ( self , cmd ) : <EOL> super ( Service , self ) . __init__ ( ) <EOL> self . cmd = cmd <EOL> def start ( self , fileStore ) : <EOL> subprocess . check_call ( self . cmd + '<STR_LIT>' , shell = True ) <EOL> def check ( self ) : <EOL> return True <EOL> def stop ( self , fileStore ) : <EOL> subprocess . check_call ( self . cmd + '<STR_LIT>' , shell = True ) <EOL> @ needs_parasol <EOL> class ParasolBatchSystemTest ( hidden . AbstractBatchSystemTest , ParasolTestSupport ) : <EOL> """<STR_LIT>""" <EOL> def supportsWallTime ( self ) : <EOL> return True <EOL> def _createDummyConfig ( self ) : <EOL> config = super ( ParasolBatchSystemTest , self ) . _createDummyConfig ( ) <EOL> config . jobStore = self . _createTempDir ( '<STR_LIT>' ) <EOL> return config <EOL> def createBatchSystem ( self ) : <EOL> memory = int ( <NUM_LIT> ) <EOL> self . _startParasol ( numCores = numCores , memory = memory ) <EOL> return ParasolBatchSystem ( config = self . config , <EOL> maxCores = numCores , <EOL> maxMemory = memory , <EOL> maxDisk = <NUM_LIT> ) <EOL> def tearDown ( self ) : <EOL> super ( ParasolBatchSystemTest , self ) . tearDown ( ) <EOL> self . _stopParasol ( ) <EOL> def testBatchResourceLimits ( self ) : <EOL> job1 = self . batchSystem . issueBatchJob ( "<STR_LIT>" , memory = <NUM_LIT:1> << <NUM_LIT:30> , cores = <NUM_LIT:1> , disk = <NUM_LIT:1000> , <EOL> preemptable = preemptable ) <EOL> self . assertIsNotNone ( job1 ) <EOL> job2 = self . batchSystem . issueBatchJob ( "<STR_LIT>" , memory = <NUM_LIT:2> << <NUM_LIT:30> , cores = <NUM_LIT:1> , disk = <NUM_LIT:1000> , <EOL> preemptable = preemptable ) <EOL> self . assertIsNotNone ( job2 ) <EOL> batches = self . _getBatchList ( ) <EOL> self . assertEqual ( len ( batches ) , <NUM_LIT:2> ) <EOL> self . assertNotEqual ( batches [ <NUM_LIT:0> ] [ '<STR_LIT>' ] , batches [ <NUM_LIT:1> ] [ '<STR_LIT>' ] ) <EOL> self . batchSystem . killBatchJobs ( [ job2 ] ) <EOL> job3 = self . batchSystem . issueBatchJob ( "<STR_LIT>" , memory = <NUM_LIT:1> << <NUM_LIT:30> , cores = <NUM_LIT:1> , disk = <NUM_LIT:1000> , <EOL> preemptable = preemptable ) <EOL> self . assertIsNotNone ( job3 ) <EOL> batches = self . _getBatchList ( ) <EOL> self . assertEqual ( len ( batches ) , <NUM_LIT:1> ) <EOL> def _parseBatchString ( self , batchString ) : <EOL> import re <EOL> batchInfo = dict ( ) <EOL> memPattern = re . compile ( "<STR_LIT>" ) <EOL> items = batchString . split ( ) <EOL> batchInfo [ "<STR_LIT>" ] = int ( items [ <NUM_LIT:7> ] ) <EOL> name = str ( items [ <NUM_LIT:11> ] ) <EOL> memMatch = memPattern . match ( items [ <NUM_LIT:8> ] ) <EOL> ramValue = float ( memMatch . group ( <NUM_LIT:1> ) ) <EOL> ramUnits = memMatch . group ( <NUM_LIT:2> ) <EOL> ramConversion = { '<STR_LIT:b>' : <NUM_LIT> , '<STR_LIT:k>' : <NUM_LIT> , '<STR_LIT:m>' : <NUM_LIT> , '<STR_LIT:g>' : <NUM_LIT> , '<STR_LIT:t>' : <NUM_LIT> } <EOL> batchInfo [ "<STR_LIT>" ] = ramValue * ramConversion [ ramUnits ] <EOL> return batchInfo <EOL> def _getBatchList ( self ) : <EOL> exitStatus , batchLines = self . batchSystem . _runParasol ( [ '<STR_LIT:list>' , '<STR_LIT>' ] ) <EOL> self . assertEqual ( exitStatus , <NUM_LIT:0> ) <EOL> return [ self . _parseBatchString ( line ) for line in batchLines [ <NUM_LIT:1> : ] if line ] <EOL> @ needs_gridengine <EOL> class GridEngineBatchSystemTest ( hidden . AbstractBatchSystemTest ) : <EOL> """<STR_LIT>""" <EOL> def _createDummyConfig ( self ) : <EOL> config = super ( GridEngineBatchSystemTest , self ) . _createDummyConfig ( ) <EOL> config . jobStore = self . _createTempDir ( '<STR_LIT>' ) <EOL> return config <EOL> def createBatchSystem ( self ) : <EOL> from toil . batchSystems . gridengine import GridengineBatchSystem <EOL> return GridengineBatchSystem ( config = self . config , maxCores = numCores , maxMemory = <NUM_LIT> , <EOL> maxDisk = <NUM_LIT> ) <EOL> @ needs_slurm <EOL> class SlurmBatchSystemTest ( hidden . AbstractBatchSystemTest ) : <EOL> """<STR_LIT>""" <EOL> def _createDummyConfig ( self ) : <EOL> config = super ( SlurmBatchSystemTest , self ) . _createDummyConfig ( ) <EOL> config . jobStore = self . _createTempDir ( '<STR_LIT>' ) <EOL> return config <EOL> def createBatchSystem ( self ) : <EOL> from toil . batchSystems . slurm import SlurmBatchSystem <EOL> return SlurmBatchSystem ( config = self . config , maxCores = numCores , maxMemory = <NUM_LIT> , <EOL> maxDisk = <NUM_LIT> ) <EOL> def tearDown ( self ) : <EOL> super ( SlurmBatchSystemTest , self ) . tearDown ( ) <EOL> from glob import glob <EOL> for f in glob ( '<STR_LIT>' ) : <EOL> os . unlink ( f ) <EOL> @ contextmanager <EOL> def tempFileContaining ( content ) : <EOL> fd , path = tempfile . mkstemp ( suffix = '<STR_LIT>' ) <EOL> try : <EOL> os . write ( fd , content ) <EOL> except : <EOL> os . close ( fd ) <EOL> raise <EOL> else : <EOL> os . close ( fd ) <EOL> yield path <EOL> finally : <EOL> os . unlink ( path ) </s>
