<s> from __future__ import print_function <EOL> import time <EOL> import sys <EOL> import os <EOL> from configargparse import ArgParser <EOL> from datetime import datetime <EOL> from multiprocessing import Process , Pool <EOL> import subprocess <EOL> import shlex <EOL> import pydpiper . queueing as q <EOL> import atoms_and_modules . registration_functions as rf <EOL> import logging <EOL> import socket <EOL> import signal <EOL> import threading <EOL> import Pyro4 <EOL> Pyro4 . config . SERVERTYPE = "<STR_LIT>" <EOL> WAIT_TIMEOUT = <NUM_LIT> <EOL> HEARTBEAT_INTERVAL = <NUM_LIT> <EOL> logger = logging . getLogger ( __name__ ) <EOL> sys . excepthook = Pyro4 . util . excepthook <EOL> def addExecutorArgumentGroup ( parser ) : <EOL> group = parser . add_argument_group ( "<STR_LIT>" , <EOL> "<STR_LIT>" ) <EOL> group . add_argument ( "<STR_LIT>" , dest = "<STR_LIT>" , <EOL> type = str , default = None , <EOL> help = "<STR_LIT>" ) <EOL> group . add_argument ( "<STR_LIT>" , dest = "<STR_LIT>" , <EOL> action = "<STR_LIT:store_true>" , default = False , <EOL> help = "<STR_LIT>" ) <EOL> group . add_argument ( "<STR_LIT>" , dest = "<STR_LIT>" , <EOL> type = float , default = <NUM_LIT> , <EOL> help = "<STR_LIT>" ) <EOL> group . add_argument ( "<STR_LIT>" , dest = "<STR_LIT>" , <EOL> type = int , default = - <NUM_LIT:1> , <EOL> help = "<STR_LIT>" ) <EOL> group . add_argument ( "<STR_LIT>" , dest = "<STR_LIT>" , <EOL> type = int , default = <NUM_LIT:2> , <EOL> help = "<STR_LIT>" ) <EOL> group . add_argument ( "<STR_LIT>" , dest = "<STR_LIT>" , <EOL> action = "<STR_LIT>" , default = True , <EOL> help = "<STR_LIT>" ) <EOL> group . add_argument ( "<STR_LIT>" , dest = "<STR_LIT:time>" , <EOL> type = str , default = None , <EOL> help = "<STR_LIT>" ) <EOL> group . add_argument ( "<STR_LIT>" , dest = "<STR_LIT>" , <EOL> type = int , default = <NUM_LIT:1> , <EOL> help = "<STR_LIT>" ) <EOL> group . add_argument ( "<STR_LIT>" , dest = "<STR_LIT>" , <EOL> type = float , default = <NUM_LIT:6> , <EOL> help = "<STR_LIT>" ) <EOL> group . add_argument ( "<STR_LIT>" , dest = "<STR_LIT>" , <EOL> type = str , default = None , <EOL> help = "<STR_LIT>" ) <EOL> group . add_argument ( "<STR_LIT>" , dest = "<STR_LIT>" , <EOL> type = str , default = "<STR_LIT>" , <EOL> help = "<STR_LIT>" ) <EOL> group . add_argument ( "<STR_LIT>" , dest = "<STR_LIT>" , <EOL> action = "<STR_LIT:store_true>" , default = False , <EOL> help = "<STR_LIT>" ) <EOL> group . add_argument ( "<STR_LIT>" , dest = "<STR_LIT>" , <EOL> type = int , default = <NUM_LIT:8> , <EOL> help = "<STR_LIT>" ) <EOL> group . add_argument ( "<STR_LIT>" , dest = "<STR_LIT>" , type = str , default = None , <EOL> help = "<STR_LIT>" ) <EOL> group . add_argument ( "<STR_LIT>" , dest = "<STR_LIT>" , type = str , default = None , <EOL> help = """<STR_LIT>""" ) <EOL> group . add_argument ( "<STR_LIT>" , dest = "<STR_LIT>" , <EOL> type = str , default = "<STR_LIT>" , <EOL> help = "<STR_LIT>" ) <EOL> group . add_argument ( "<STR_LIT>" , dest = "<STR_LIT>" , type = int , default = <NUM_LIT> , <EOL> help = "<STR_LIT>" ) <EOL> group . add_argument ( "<STR_LIT>" , dest = "<STR_LIT>" , <EOL> type = int , default = <NUM_LIT:1> , <EOL> help = "<STR_LIT>" ) <EOL> group . add_argument ( "<STR_LIT>" , dest = "<STR_LIT>" , <EOL> type = int , <EOL> help = "<STR_LIT>" ) <EOL> group . add_argument ( '<STR_LIT>' , dest = "<STR_LIT>" , action = '<STR_LIT:store_true>' , default = False , help = "<STR_LIT>" ) <EOL> group . add_argument ( "<STR_LIT>" , type = str , metavar = '<STR_LIT>' , is_config_file = True , <EOL> required = False , help = '<STR_LIT>' ) <EOL> group . add_argument ( "<STR_LIT>" , type = str , metavar = '<STR_LIT:file>' , <EOL> help = "<STR_LIT>" ) <EOL> group . add_argument ( "<STR_LIT>" , dest = "<STR_LIT>" , type = int , default = <NUM_LIT:0> , <EOL> help = "<STR_LIT>" ) <EOL> group . add_argument ( "<STR_LIT>" , dest = "<STR_LIT>" , type = int , default = None , <EOL> help = "<STR_LIT>" ) <EOL> group . add_argument ( "<STR_LIT>" , dest = "<STR_LIT>" , <EOL> type = float , default = <NUM_LIT> , <EOL> help = "<STR_LIT>" ) <EOL> def noExecSpecified ( numExec ) : <EOL> if numExec < <NUM_LIT:0> : <EOL> logger . info ( "<STR_LIT>" ) <EOL> print ( "<STR_LIT>" ) <EOL> sys . exit ( ) <EOL> def launchExecutor ( executor ) : <EOL> network_address = Pyro4 . socketutil . getIpAddress ( socket . gethostname ( ) , <EOL> workaround127 = True , ipVersion = <NUM_LIT:4> ) <EOL> daemon = Pyro4 . core . Daemon ( host = network_address ) <EOL> clientURI = daemon . register ( executor ) <EOL> if executor . ns : <EOL> ns = Pyro4 . locateNS ( ) <EOL> serverURI = ns . lookup ( "<STR_LIT>" ) <EOL> else : <EOL> try : <EOL> uf = open ( executor . uri_file ) <EOL> serverURI = Pyro4 . URI ( uf . readline ( ) ) <EOL> uf . close ( ) <EOL> except : <EOL> logger . exception ( "<STR_LIT>" ) <EOL> raise <EOL> p = Pyro4 . Proxy ( serverURI ) <EOL> p . registerClient ( clientURI . asString ( ) , executor . mem ) <EOL> executor . registeredWithServer ( ) <EOL> executor . setClientURI ( clientURI . asString ( ) ) <EOL> executor . setServerURI ( serverURI . asString ( ) ) <EOL> executor . setProxyForServer ( p ) <EOL> logger . info ( "<STR_LIT>" , serverURI ) <EOL> logger . info ( "<STR_LIT>" , clientURI ) <EOL> executor . connection_time_with_server = time . time ( ) <EOL> logger . info ( "<STR_LIT>" , datetime . isoformat ( datetime . now ( ) , "<STR_LIT:U+0020>" ) ) <EOL> executor . initializePool ( ) <EOL> logger . debug ( "<STR_LIT>" , daemon . locationStr ) <EOL> try : <EOL> t = threading . Thread ( target = daemon . requestLoop ) <EOL> t . daemon = True <EOL> t . start ( ) <EOL> h = threading . Thread ( target = executor . heartbeat ) <EOL> h . daemon = True <EOL> h . start ( ) <EOL> executor . mainLoop ( ) <EOL> except KeyboardInterrupt : <EOL> logger . exception ( "<STR_LIT>" ) <EOL> executor . generalShutdownCall ( ) <EOL> sys . exit ( <NUM_LIT:0> ) <EOL> except Exception : <EOL> logger . exception ( "<STR_LIT>" ) <EOL> executor . generalShutdownCall ( ) <EOL> sys . exit ( <NUM_LIT:0> ) <EOL> else : <EOL> executor . completeAndExitChildren ( ) <EOL> logger . info ( "<STR_LIT>" ) <EOL> daemon . shutdown ( ) <EOL> t . join ( ) <EOL> def runStage ( serverURI , clientURI , i ) : <EOL> p = Pyro4 . core . Proxy ( serverURI ) <EOL> client = Pyro4 . core . Proxy ( clientURI ) <EOL> try : <EOL> logger . info ( "<STR_LIT>" , i , clientURI ) <EOL> logger . debug ( "<STR_LIT>" , p . getStageMem ( i ) ) <EOL> p . setStageStarted ( i , clientURI ) <EOL> try : <EOL> command_to_run = str ( p . getStageCommand ( i ) ) <EOL> logger . info ( command_to_run ) <EOL> command_logfile = p . getStageLogfile ( i ) <EOL> of = open ( command_logfile , '<STR_LIT:a>' ) <EOL> of . write ( "<STR_LIT>" + str ( i ) + "<STR_LIT>" + socket . gethostname ( ) + "<STR_LIT>" + datetime . isoformat ( datetime . now ( ) , "<STR_LIT:U+0020>" ) + "<STR_LIT>" ) <EOL> of . write ( command_to_run + "<STR_LIT:\n>" ) <EOL> of . flush ( ) <EOL> args = shlex . split ( command_to_run ) <EOL> process = subprocess . Popen ( args , stdout = of , stderr = of , shell = False ) <EOL> client . addPIDtoRunningList ( process . pid ) <EOL> process . communicate ( ) <EOL> client . removePIDfromRunningList ( process . pid ) <EOL> ret = process . returncode <EOL> of . close ( ) <EOL> except : <EOL> logger . exception ( "<STR_LIT>" , i , clientURI ) <EOL> client . notifyStageTerminated ( i ) <EOL> else : <EOL> logger . info ( "<STR_LIT>" , i , ret , clientURI ) <EOL> client . notifyStageTerminated ( i , ret ) <EOL> return ( p . getStageMem ( i ) , p . getStageProcs ( i ) ) <EOL> except : <EOL> logger . exception ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> raise <EOL> """<STR_LIT>""" <EOL> class ChildProcess ( object ) : <EOL> def __init__ ( self , stage , result , mem , procs ) : <EOL> self . stage = stage <EOL> self . result = result <EOL> self . mem = mem <EOL> self . procs = procs <EOL> class InsufficientResources ( Exception ) : <EOL> pass <EOL> class pipelineExecutor ( object ) : <EOL> def __init__ ( self , options , memNeeded = None ) : <EOL> logger . debug ( "<STR_LIT>" , memNeeded ) <EOL> self . mem = memNeeded or options . mem <EOL> logger . debug ( "<STR_LIT>" , self . mem ) <EOL> if self . mem > options . mem : <EOL> raise InsufficientResources ( "<STR_LIT>" % ( options . mem , self . mem ) ) <EOL> self . procs = options . proc <EOL> self . ppn = options . ppn <EOL> self . pe = options . pe <EOL> self . mem_request_variable = options . mem_request_variable <EOL> self . queue_type = options . queue_type <EOL> self . queue_name = options . queue_name <EOL> self . queue_opts = options . queue_opts <EOL> self . ns = options . use_ns <EOL> self . uri_file = options . urifile <EOL> if self . uri_file is None : <EOL> self . uri_file = os . path . abspath ( os . path . join ( os . curdir , options . pipeline_name + "<STR_LIT>" ) ) <EOL> self . idle_time = <NUM_LIT:0> <EOL> self . prev_time = None <EOL> self . current_time = None <EOL> self . time_to_seppuku = options . time_to_seppuku <EOL> self . time_to_accept_jobs = options . time_to_accept_jobs <EOL> self . connection_time_with_server = None <EOL> self . runningMem = <NUM_LIT:0.0> <EOL> self . runningProcs = <NUM_LIT:0> <EOL> self . runningChildren = [ ] <EOL> self . pool = None <EOL> self . pyro_proxy_for_server = None <EOL> self . clientURI = None <EOL> self . serverURI = None <EOL> self . current_running_job_pids = [ ] <EOL> self . registered_with_server = False <EOL> self . heartbeat_thread_crashed = False <EOL> self . e = threading . Event ( ) <EOL> def registeredWithServer ( self ) : <EOL> self . registered_with_server = True <EOL> def addPIDtoRunningList ( self , pid ) : <EOL> self . current_running_job_pids . append ( pid ) <EOL> def removePIDfromRunningList ( self , pid ) : <EOL> self . current_running_job_pids . remove ( pid ) <EOL> def initializePool ( self ) : <EOL> self . pool = Pool ( processes = self . procs ) <EOL> def setClientURI ( self , cURI ) : <EOL> self . clientURI = cURI <EOL> def setServerURI ( self , sURI ) : <EOL> self . serverURI = sURI <EOL> def setProxyForServer ( self , proxy ) : <EOL> self . pyro_proxy_for_server = proxy <EOL> def generalShutdownCall ( self ) : <EOL> logger . debug ( "<STR_LIT>" ) <EOL> for subprocID in self . current_running_job_pids : <EOL> os . kill ( subprocID , signal . SIGTERM ) <EOL> self . unregister_with_server ( ) <EOL> def completeAndExitChildren ( self ) : <EOL> if len ( self . current_running_job_pids ) > <NUM_LIT:0> : <EOL> self . pool . close ( ) <EOL> self . pool . join ( ) <EOL> self . unregister_with_server ( ) <EOL> def unregister_with_server ( self ) : <EOL> if self . registered_with_server : <EOL> logger . info ( "<STR_LIT>" , self . clientURI ) <EOL> self . registered_with_server = False <EOL> logger . info ( "<STR_LIT>" , self . clientURI ) <EOL> self . pyro_proxy_for_server . unregisterClient ( self . clientURI ) <EOL> def submitToQueue ( self , programName = None ) : <EOL> """<STR_LIT>""" <EOL> if self . queue_type == "<STR_LIT>" : <EOL> strprocs = str ( self . procs ) <EOL> strmem = "<STR_LIT>" % ( self . mem_request_variable , float ( self . mem ) ) <EOL> jobname = "<STR_LIT>" <EOL> if programName is not None : <EOL> executablePath = os . path . abspath ( programName ) <EOL> jobname = os . path . basename ( executablePath ) + "<STR_LIT:->" <EOL> now = datetime . now ( ) . strftime ( "<STR_LIT>" ) <EOL> ident = "<STR_LIT>" + now <EOL> jobname += ident <EOL> queue_opts = [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT:yes>' , <EOL> '<STR_LIT>' , jobname , <EOL> '<STR_LIT>' , strmem , <EOL> '<STR_LIT>' , os . path . join ( os . getcwd ( ) , <EOL> ident + '<STR_LIT>' ) ] + ( [ '<STR_LIT>' , self . queue_name ] <EOL> if self . queue_name else [ ] ) + ( [ '<STR_LIT>' , self . pe , strprocs ] <EOL> if self . pe else [ ] ) + shlex . split ( self . queue_opts ) <EOL> qsub_cmd = [ '<STR_LIT>' ] + queue_opts <EOL> cmd = [ "<STR_LIT>" , "<STR_LIT>" ] <EOL> cmd += [ '<STR_LIT>' , self . uri_file ] <EOL> cmd += [ "<STR_LIT>" , str ( <NUM_LIT:1> ) ] <EOL> cmd += q . remove_flags ( [ '<STR_LIT>' , '<STR_LIT>' ] , sys . argv [ <NUM_LIT:1> : ] ) <EOL> cmd += [ '<STR_LIT>' , str ( self . mem ) ] <EOL> script = "<STR_LIT>" % '<STR_LIT:U+0020>' . join ( cmd ) <EOL> env = os . environ . copy ( ) <EOL> env [ '<STR_LIT>' ] = os . path . join ( os . getcwd ( ) , ident + "<STR_LIT>" ) <EOL> p = subprocess . Popen ( qsub_cmd , stdin = subprocess . PIPE , shell = False , env = env ) <EOL> p . communicate ( script ) <EOL> else : <EOL> logger . info ( "<STR_LIT>" % ( self . queue_type ) ) <EOL> logger . info ( "<STR_LIT>" ) <EOL> logger . info ( "<STR_LIT>" ) <EOL> sys . exit ( ) <EOL> def canRun ( self , stageMem , stageProcs , runningMem , runningProcs ) : <EOL> """<STR_LIT>""" <EOL> return stageMem <= self . mem - runningMem and stageProcs <= self . procs - runningProcs <EOL> def is_seppuku_time ( self ) : <EOL> if self . time_to_seppuku != None : <EOL> if ( self . time_to_seppuku * <NUM_LIT> ) < self . idle_time : <EOL> return True <EOL> return False <EOL> def is_time_to_drain ( self ) : <EOL> if ( self . time_to_accept_jobs != None ) : <EOL> current_time = time . time ( ) <EOL> time_take_so_far = current_time - self . connection_time_with_server <EOL> minutes_so_far , seconds_so_far = divmod ( time_take_so_far , <NUM_LIT> ) <EOL> if self . time_to_accept_jobs < minutes_so_far : <EOL> return True <EOL> return False <EOL> def free_resources ( self ) : <EOL> for child in self . runningChildren : <EOL> if child . result . ready ( ) : <EOL> logger . debug ( "<STR_LIT>" , child . stage ) <EOL> self . runningMem -= child . mem <EOL> self . runningProcs -= child . procs <EOL> self . runningChildren . remove ( child ) <EOL> def notifyStageTerminated ( self , i , returncode = None ) : <EOL> if returncode == <NUM_LIT:0> : <EOL> self . pyro_proxy_for_server . setStageFinished ( i , self . clientURI ) <EOL> else : <EOL> self . pyro_proxy_for_server . setStageFailed ( i , self . clientURI ) <EOL> self . e . set ( ) <EOL> def idle ( self ) : <EOL> return self . runningMem == <NUM_LIT:0> and self . runningProcs == <NUM_LIT:0> and self . prev_time <EOL> def heartbeat ( self ) : <EOL> try : <EOL> tick = <NUM_LIT:0> <EOL> while self . registered_with_server : <EOL> logger . debug ( "<STR_LIT>" , tick ) <EOL> tick += <NUM_LIT:1> <EOL> self . pyro_proxy_for_server . updateClientTimestamp ( self . clientURI , tick ) <EOL> time . sleep ( HEARTBEAT_INTERVAL ) <EOL> except : <EOL> logger . exception ( "<STR_LIT>" ) <EOL> self . heartbeat_thread_crashed = True <EOL> def mainLoop ( self ) : <EOL> while self . mainFn ( ) : <EOL> self . e . wait ( WAIT_TIMEOUT ) <EOL> self . e . clear ( ) <EOL> logger . debug ( "<STR_LIT>" ) <EOL> def mainFn ( self ) : <EOL> """<STR_LIT>""" <EOL> self . prev_time = self . current_time <EOL> self . current_time = time . time ( ) <EOL> self . free_resources ( ) <EOL> if self . heartbeat_thread_crashed : <EOL> logger . debug ( "<STR_LIT>" ) <EOL> return False <EOL> if self . idle ( ) : <EOL> self . idle_time += self . current_time - self . prev_time <EOL> logger . debug ( "<STR_LIT>" , self . idle_time , self . time_to_seppuku * <NUM_LIT> ) <EOL> if self . is_seppuku_time ( ) : <EOL> logger . warn ( "<STR_LIT>" ) <EOL> return False <EOL> if self . is_time_to_drain ( ) : <EOL> logger . info ( "<STR_LIT>" ) <EOL> return False <EOL> cmd , i = self . pyro_proxy_for_server . getCommand ( clientURIstr = self . clientURI , <EOL> clientMemFree = self . mem - self . runningMem , <EOL> clientProcsFree = self . procs - self . runningProcs ) <EOL> if cmd == "<STR_LIT>" : <EOL> logger . debug ( '<STR_LIT>' ) <EOL> return False <EOL> elif cmd == "<STR_LIT>" : <EOL> return True <EOL> elif cmd == "<STR_LIT>" : <EOL> stageMem , stageProcs = self . pyro_proxy_for_server . getStageMem ( i ) , self . pyro_proxy_for_server . getStageProcs ( i ) <EOL> self . idle_time = <NUM_LIT:0> <EOL> self . runningMem += stageMem <EOL> self . runningProcs += stageProcs <EOL> result = self . pool . apply_async ( runStage , ( self . serverURI , self . clientURI , i ) ) <EOL> self . runningChildren . append ( ChildProcess ( i , result , stageMem , stageProcs ) ) <EOL> logger . debug ( "<STR_LIT>" , i ) <EOL> return True <EOL> else : <EOL> raise Exception ( "<STR_LIT>" % cmd ) <EOL> if __name__ == "<STR_LIT:__main__>" : <EOL> default_config_file = os . getenv ( "<STR_LIT>" ) <EOL> if default_config_file is not None : <EOL> files = [ default_config_file ] <EOL> else : <EOL> files = [ ] <EOL> parser = ArgParser ( default_config_files = files ) <EOL> rf . addGenRegArgumentGroup ( parser ) <EOL> addExecutorArgumentGroup ( parser ) <EOL> options = parser . parse_known_args ( ) [ <NUM_LIT:0> ] <EOL> noExecSpecified ( options . num_exec ) <EOL> def local_launch ( options ) : <EOL> pe = pipelineExecutor ( options ) <EOL> ps = [ Process ( target = launchExecutor , args = ( pe , ) ) <EOL> for _ in range ( options . num_exec ) ] <EOL> for p in ps : <EOL> p . start ( ) <EOL> for p in ps : <EOL> p . join ( ) <EOL> if options . local : <EOL> local_launch ( options ) <EOL> elif options . queue_type == "<STR_LIT>" : <EOL> roq = q . runOnQueueingSystem ( options , sysArgs = sys . argv ) <EOL> for i in range ( options . num_exec ) : <EOL> roq . createAndSubmitExecutorJobFile ( i , after = None , <EOL> time = q . timestr_to_secs ( options . time ) ) <EOL> elif options . queue_type == "<STR_LIT>" : <EOL> for i in range ( options . num_exec ) : <EOL> pe = pipelineExecutor ( options ) <EOL> pe . submitToQueue ( ) <EOL> else : <EOL> local_launch ( options ) </s>
