<s> """<STR_LIT>""" <EOL> import multiprocessing <EOL> import os <EOL> import subprocess <EOL> import signal <EOL> import sys <EOL> import logging <EOL> import jip . cluster <EOL> from jip . logger import getLogger <EOL> class LocalCluster ( jip . cluster . Cluster ) : <EOL> def __init__ ( self , _start = True , <EOL> _remote_ids = False , <EOL> cores = None , log_level = logging . INFO ) : <EOL> self . _current_id = <NUM_LIT:0> <EOL> self . log = getLogger ( "<STR_LIT>" ) <EOL> self . master_requests = None <EOL> self . master_response = None <EOL> self . master_process = None <EOL> self . _current_id <EOL> self . _remote_ids = _remote_ids <EOL> self . cores = cores <EOL> self . loglevel = log_level <EOL> self . log . setLevel ( log_level ) <EOL> if _start : <EOL> self . start ( ) <EOL> def _next_id ( self ) : <EOL> """<STR_LIT>""" <EOL> self . _current_id += <NUM_LIT:1> <EOL> return self . _current_id <EOL> def start ( self ) : <EOL> if self . master_process is None : <EOL> self . log . info ( "<STR_LIT>" ) <EOL> self . master_requests = multiprocessing . Queue ( ) <EOL> self . master_response = multiprocessing . Queue ( ) <EOL> self . master_process = multiprocessing . Process ( <EOL> target = _GridMaster . create_master , <EOL> args = [ self . master_requests , self . master_response ] , <EOL> kwargs = { "<STR_LIT>" : self . cores , <EOL> "<STR_LIT>" : self . loglevel } , <EOL> name = "<STR_LIT>" <EOL> ) <EOL> self . master_process . start ( ) <EOL> def shutdown ( self ) : <EOL> if self . master_process is None : <EOL> return <EOL> self . log . info ( "<STR_LIT>" ) <EOL> self . master_requests . put ( [ '<STR_LIT>' ] ) <EOL> self . master_process . join ( ) <EOL> self . master_requests . close ( ) <EOL> self . master_response . close ( ) <EOL> self . master_process = None <EOL> def wait ( self ) : <EOL> if self . master_process is None : <EOL> return <EOL> self . log . info ( "<STR_LIT>" ) <EOL> self . master_requests . put ( [ "<STR_LIT>" ] ) <EOL> self . master_process . join ( ) <EOL> self . master_requests . close ( ) <EOL> self . master_response . close ( ) <EOL> def list ( self ) : <EOL> if self . master_process is None : <EOL> return [ ] <EOL> self . master_requests . put ( [ "<STR_LIT>" ] ) <EOL> return self . master_response . get ( ) <EOL> def submit ( self , job ) : <EOL> if self . master_process is None : <EOL> raise jip . cluster . SubmissionError ( "<STR_LIT>" ) <EOL> local_job = _Job . from_job ( job ) if not isinstance ( job , _Job ) else job <EOL> if not self . _remote_ids : <EOL> local_job . job_id = self . _next_id ( ) <EOL> self . master_requests . put ( [ <EOL> "<STR_LIT>" , local_job , <EOL> ] ) <EOL> if self . _remote_ids : <EOL> job . job_id = self . master_response . get ( ) <EOL> else : <EOL> job . job_id = local_job . job_id <EOL> self . log . info ( "<STR_LIT>" , job . job_id ) <EOL> return job <EOL> def resolve_log ( self , job , path ) : <EOL> if path is None : <EOL> return None <EOL> return path . replace ( "<STR_LIT>" , str ( job . id ) ) <EOL> def cancel ( self , job ) : <EOL> if self . master_process is None : <EOL> return <EOL> self . log . info ( "<STR_LIT>" , job . job_id ) <EOL> self . master_requests . put ( [ "<STR_LIT>" , job . job_id ] ) <EOL> class JIP ( jip . cluster . Cluster ) : <EOL> def __init__ ( self ) : <EOL> self . log = getLogger ( "<STR_LIT>" ) <EOL> cfg = jip . config . get ( '<STR_LIT>' , { } ) <EOL> self . host = cfg . get ( '<STR_LIT:host>' , '<STR_LIT:127.0.0.1>' ) <EOL> self . port = cfg . get ( '<STR_LIT:port>' , '<STR_LIT>' ) <EOL> self . log = getLogger ( "<STR_LIT>" ) <EOL> def _connect ( self ) : <EOL> import zmq <EOL> context = zmq . Context ( ) <EOL> socket = context . socket ( zmq . REQ ) <EOL> socket . setsockopt ( zmq . LINGER , <NUM_LIT:0> ) <EOL> socket . RCVTIMEO = <NUM_LIT:1000> <EOL> self . log . info ( "<STR_LIT>" , self . host , self . port ) <EOL> socket . connect ( "<STR_LIT>" % ( self . host , self . port ) ) <EOL> return context , socket <EOL> def list ( self ) : <EOL> try : <EOL> context , socket = self . _connect ( ) <EOL> socket . send_json ( { "<STR_LIT>" : "<STR_LIT:list>" } ) <EOL> jobs = socket . recv_json ( ) <EOL> except : <EOL> self . log . error ( "<STR_LIT>" ) <EOL> raise <EOL> finally : <EOL> socket . close ( ) <EOL> context . term ( ) <EOL> return jobs <EOL> def submit ( self , job ) : <EOL> try : <EOL> context , socket = self . _connect ( ) <EOL> local = _Job . from_job ( job ) <EOL> socket . send_json ( { "<STR_LIT>" : "<STR_LIT>" , "<STR_LIT>" : { <EOL> "<STR_LIT>" : local . job_id , <EOL> "<STR_LIT:id>" : local . id , <EOL> "<STR_LIT>" : local . cmd , <EOL> "<STR_LIT>" : local . stdout , <EOL> "<STR_LIT>" : local . stderr , <EOL> "<STR_LIT>" : local . working_directory , <EOL> "<STR_LIT>" : [ f for f in local . dependencies ] , <EOL> "<STR_LIT>" : [ f for f in local . children ] <EOL> } } ) <EOL> submitted = socket . recv_json ( ) <EOL> job . job_id = submitted [ '<STR_LIT>' ] <EOL> job . working_directory = submitted [ '<STR_LIT>' ] <EOL> job . stdout = submitted [ '<STR_LIT>' ] <EOL> job . stderr = submitted [ '<STR_LIT>' ] <EOL> return job <EOL> except : <EOL> self . log . error ( "<STR_LIT>" , exc_info = True ) <EOL> raise jip . cluster . SubmissionError ( "<STR_LIT>" ) <EOL> finally : <EOL> socket . close ( ) <EOL> context . term ( ) <EOL> return job <EOL> def resolve_log ( self , job , path ) : <EOL> if path is None : <EOL> return None <EOL> return path . replace ( "<STR_LIT>" , str ( job . id ) ) <EOL> def cancel ( self , job ) : <EOL> try : <EOL> context , socket = self . _connect ( ) <EOL> self . log . info ( "<STR_LIT>" , job . job_id ) <EOL> socket . send_json ( { "<STR_LIT>" : "<STR_LIT>" , '<STR_LIT:id>' : job . job_id } ) <EOL> socket . recv ( ) <EOL> except : <EOL> self . log . error ( "<STR_LIT>" ) <EOL> raise <EOL> finally : <EOL> socket . close ( ) <EOL> context . term ( ) <EOL> class _Job ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , cmd = None , cwd = None , stdout = None , stderr = None , <EOL> dependencies = None , threads = <NUM_LIT:1> , job_id = None ) : <EOL> self . id = job_id <EOL> self . job_id = job_id <EOL> self . dependencies = [ ] if dependencies is None else dependencies <EOL> self . threads = threads <EOL> self . cmd = cmd <EOL> self . working_directory = cwd <EOL> self . stdout = stdout <EOL> self . stderr = stderr <EOL> self . children = [ ] <EOL> self . process = None <EOL> @ classmethod <EOL> def from_job ( cls , job ) : <EOL> cwd = job . working_directory if job . working_directory is not None else os . getcwd ( ) <EOL> if job . stderr is None : <EOL> job . stderr = os . path . join ( cwd , "<STR_LIT>" ) <EOL> if job . stdout is None : <EOL> job . stdout = os . path . join ( cwd , "<STR_LIT>" ) <EOL> deps = set ( [ ] ) <EOL> for dep in [ d for d in job . dependencies if d . job_id ] : <EOL> deps . add ( int ( dep . job_id ) ) <EOL> local_job = _Job ( <EOL> job . get_cluster_command ( ) , <EOL> cwd , <EOL> job . stdout , <EOL> job . stderr , <EOL> dependencies = deps , <EOL> threads = job . threads <EOL> ) <EOL> local_job . id = job . id <EOL> return local_job <EOL> def __lt__ ( self , other ) : <EOL> num_deps = len ( self . dependencies ) <EOL> o_deps = len ( other . dependencies ) <EOL> if num_deps == o_deps : <EOL> return self . job_id < other . job_id <EOL> return num_deps < o_deps <EOL> def __repr__ ( self ) : <EOL> return "<STR_LIT>" % ( <EOL> str ( self . job_id ) , self . dependencies , self . children <EOL> ) <EOL> class _GridMaster ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , requests , response , cores = None , loglevel = logging . INFO ) : <EOL> """<STR_LIT>""" <EOL> self . requests = requests <EOL> self . response = response <EOL> self . queued = { } <EOL> self . running = { } <EOL> self . slots_available = multiprocessing . cpu_count ( ) if not cores else cores <EOL> self . slots_total = self . slots_available <EOL> self . log = getLogger ( "<STR_LIT>" ) <EOL> self . wait_mode = False <EOL> self . _current_id = <NUM_LIT:0> <EOL> self . log . setLevel ( loglevel ) <EOL> self . log . info ( "<STR_LIT>" , <EOL> self . slots_available ) <EOL> def _next_id ( self ) : <EOL> """<STR_LIT>""" <EOL> self . _current_id += <NUM_LIT:1> <EOL> return self . _current_id <EOL> def _remove_children ( self , job ) : <EOL> """<STR_LIT>""" <EOL> self . log . debug ( "<STR_LIT>" , <EOL> job . job_id ) <EOL> for c in job . children : <EOL> if c in self . queued : <EOL> child = self . queued [ c ] <EOL> self . _remove_children ( child ) <EOL> self . _update_dependencies ( child ) <EOL> del self . queued [ c ] <EOL> def _update_dependencies ( self , job ) : <EOL> """<STR_LIT>""" <EOL> for c in job . children : <EOL> if c in self . queued : <EOL> self . log . debug ( "<STR_LIT>" <EOL> "<STR_LIT>" , job . job_id , c ) <EOL> try : <EOL> self . queued [ c ] . dependencies . remove ( job . job_id ) <EOL> except KeyError : <EOL> self . log . error ( "<STR_LIT>" <EOL> "<STR_LIT>" , job . job_id , c ) <EOL> def _run_job ( self , job ) : <EOL> """<STR_LIT>""" <EOL> if not job . job_id in self . queued : <EOL> self . log . error ( "<STR_LIT>" , job . job_id ) <EOL> return <EOL> job_id = job . job_id <EOL> self . log . info ( "<STR_LIT>" , job_id ) <EOL> process = multiprocessing . Process ( <EOL> target = _execute_job , <EOL> args = [ self . requests , job_id , job ] , <EOL> name = "<STR_LIT>" % str ( job_id ) <EOL> ) <EOL> job . process = process <EOL> self . running [ job_id ] = job <EOL> del self . queued [ job_id ] <EOL> process . start ( ) <EOL> def schedule ( self ) : <EOL> self . log . info ( "<STR_LIT>" , <EOL> self . slots_available , self . slots_total ) <EOL> sorted_jobs = sorted ( self . queued . values ( ) ) <EOL> for job in sorted_jobs : <EOL> if len ( job . dependencies ) > <NUM_LIT:0> : <EOL> self . log . info ( "<STR_LIT>" ) <EOL> break <EOL> if job . threads <= self . slots_available : <EOL> self . log . info ( "<STR_LIT>" , <EOL> job . job_id ) <EOL> self . _run_job ( job ) <EOL> self . slots_available -= job . threads <EOL> if self . slots_available >= <NUM_LIT:1> : <EOL> self . schedule ( ) <EOL> return <EOL> def _resolve_log ( self , job_id , path ) : <EOL> if path is None : <EOL> return None <EOL> return path . replace ( "<STR_LIT>" , str ( job_id ) ) <EOL> def _handle_exit ( self , * args ) : <EOL> """<STR_LIT>""" <EOL> self . log . info ( "<STR_LIT>" ) <EOL> for job_id , p in self . running . iteritems ( ) : <EOL> self . log . warn ( "<STR_LIT>" , job_id ) <EOL> p . process . terminate ( ) <EOL> p . process . join ( ) <EOL> return False <EOL> def _handle_wait ( self , * args ) : <EOL> """<STR_LIT>""" <EOL> if not self . wait_mode : <EOL> self . log . info ( "<STR_LIT>" ) <EOL> self . wait_mode = True <EOL> return True <EOL> def _handle_jobs ( self , * args ) : <EOL> """<STR_LIT>""" <EOL> self . log . info ( "<STR_LIT>" ) <EOL> self . response . put ( list ( self . queued . keys ( ) ) + <EOL> list ( self . running . keys ( ) ) ) <EOL> return True <EOL> def _handle_submit ( self , * args ) : <EOL> """<STR_LIT>""" <EOL> job = args [ <NUM_LIT:1> ] <EOL> create_id = job . job_id is None <EOL> job_id = self . _next_id ( ) if create_id else job . job_id <EOL> job . job_id = job_id <EOL> job . stdout = self . _resolve_log ( job . id , job . stdout ) <EOL> job . stderr = self . _resolve_log ( job . id , job . stderr ) <EOL> self . queued [ job_id ] = job <EOL> for d in job . dependencies : <EOL> if d in self . queued : <EOL> self . queued [ d ] . children . append ( job_id ) <EOL> if d in self . running : <EOL> self . running [ d ] . children . append ( job_id ) <EOL> self . log . info ( "<STR_LIT>" , job_id ) <EOL> if create_id : <EOL> self . response . put ( job_id ) <EOL> if self . slots_available >= <NUM_LIT:1> : <EOL> self . schedule ( ) <EOL> return True <EOL> def _handle_done ( self , * args ) : <EOL> """<STR_LIT>""" <EOL> job_id = int ( args [ <NUM_LIT:1> ] ) <EOL> state = args [ <NUM_LIT:2> ] <EOL> if job_id in self . running : <EOL> job = self . running [ job_id ] <EOL> job . process . join ( ) <EOL> job . process = None <EOL> self . slots_available += job . threads <EOL> self . _update_dependencies ( job ) <EOL> if state != <NUM_LIT:0> : <EOL> self . log . error ( "<STR_LIT>" , <EOL> job_id , state ) <EOL> self . _remove_children ( job ) <EOL> else : <EOL> self . log . info ( "<STR_LIT>" , <EOL> job_id , state ) <EOL> job . process = None <EOL> del self . running [ job_id ] <EOL> self . schedule ( ) <EOL> else : <EOL> self . log . warn ( "<STR_LIT>" <EOL> "<STR_LIT>" , job_id ) <EOL> return True <EOL> def _handle_failed ( self , * args ) : <EOL> """<STR_LIT>""" <EOL> job_id = int ( args [ <NUM_LIT:1> ] ) <EOL> error = args [ <NUM_LIT:2> ] <EOL> self . log . error ( "<STR_LIT>" , job_id , error ) <EOL> if job_id in self . unning : <EOL> job = self . running [ job_id ] <EOL> job . process = None <EOL> self . _update_dependencies ( job ) <EOL> self . _remove_children ( job ) <EOL> self . slots_available += job . threads <EOL> del self . running [ job_id ] <EOL> self . schedule ( ) <EOL> else : <EOL> self . log . error ( "<STR_LIT>" <EOL> "<STR_LIT>" % job_id ) <EOL> return True <EOL> def _handle_cancel ( self , * args ) : <EOL> """<STR_LIT>""" <EOL> job_id = int ( args [ <NUM_LIT:1> ] ) <EOL> self . log . warn ( "<STR_LIT>" , job_id ) <EOL> job = None <EOL> if job_id in self . running : <EOL> self . log . warn ( "<STR_LIT>" , job_id ) <EOL> job = self . running [ job_id ] <EOL> job . process . terminate ( ) <EOL> job . process . join ( ) <EOL> job . process = None <EOL> self . slots_available += job . threads <EOL> del self . running [ job_id ] <EOL> elif job_id in self . queued : <EOL> job = self . queued [ job_id ] <EOL> del self . queued [ job_id ] <EOL> else : <EOL> self . log . warn ( "<STR_LIT>" <EOL> "<STR_LIT>" , job_id ) <EOL> if job : <EOL> self . _update_dependencies ( job ) <EOL> self . _remove_children ( job ) <EOL> self . schedule ( ) <EOL> return True <EOL> def start ( self ) : <EOL> """<STR_LIT>""" <EOL> handlers = { <EOL> "<STR_LIT>" : self . _handle_exit , <EOL> "<STR_LIT>" : self . _handle_jobs , <EOL> "<STR_LIT>" : self . _handle_submit , <EOL> "<STR_LIT>" : self . _handle_wait , <EOL> "<STR_LIT>" : self . _handle_done , <EOL> "<STR_LIT>" : self . _handle_failed , <EOL> "<STR_LIT>" : self . _handle_cancel , <EOL> } <EOL> while True : <EOL> if self . wait_mode and self . _num_jobs ( ) == <NUM_LIT:0> : <EOL> break <EOL> try : <EOL> msg = self . requests . get ( ) <EOL> except KeyboardInterrupt as err : <EOL> self . log . warn ( "<STR_LIT>" ) <EOL> continue <EOL> self . log . debug ( "<STR_LIT>" , msg ) <EOL> try : <EOL> command = msg [ <NUM_LIT:0> ] <EOL> handler = handlers [ command ] <EOL> if not handler ( * msg ) : <EOL> break <EOL> except Exception as err : <EOL> self . log . error ( "<STR_LIT>" , <EOL> msg , str ( err ) , exc_info = True ) <EOL> break <EOL> self . log . info ( "<STR_LIT>" ) <EOL> def _num_jobs ( self ) : <EOL> return len ( self . queued ) + len ( self . running ) <EOL> @ staticmethod <EOL> def create_master ( request , response , cores = None , loglevel = logging . INFO ) : <EOL> master = _GridMaster ( request , response , cores = cores , <EOL> loglevel = loglevel ) <EOL> master . start ( ) <EOL> def _execute_job ( requests , job_id , job ) : <EOL> """<STR_LIT>""" <EOL> log = getLogger ( "<STR_LIT>" ) <EOL> process = None <EOL> def handle_term ( sig , frame ) : <EOL> log . error ( "<STR_LIT>" ) <EOL> _terminate_process ( process , log ) <EOL> requests . put ( [ '<STR_LIT>' , job_id , <NUM_LIT:1> ] ) <EOL> sys . exit ( <NUM_LIT:1> ) <EOL> signal . signal ( signal . SIGTERM , handle_term ) <EOL> log . info ( "<STR_LIT>" , job_id ) <EOL> result = - <NUM_LIT:1> <EOL> try : <EOL> cwd = job . working_directory <EOL> stdout = open ( job . stdout , '<STR_LIT:w>' ) <EOL> stderr = open ( job . stderr , '<STR_LIT:w>' ) <EOL> cmd = job . cmd <EOL> process = subprocess . Popen ( <EOL> "<STR_LIT>" + cmd , <EOL> stdout = stdout , <EOL> stderr = stderr , <EOL> shell = True , <EOL> cwd = cwd <EOL> ) <EOL> result = process . wait ( ) <EOL> except Exception as err : <EOL> log . error ( "<STR_LIT>" , err , exc_info = True ) <EOL> requests . put ( [ '<STR_LIT>' , job_id , str ( err ) ] ) <EOL> log . info ( "<STR_LIT>" , job_id , result ) <EOL> requests . put ( [ '<STR_LIT>' , job_id , result ] ) <EOL> def _terminate_process ( process , log ) : <EOL> """<STR_LIT>""" <EOL> if process is not None : <EOL> log . info ( "<STR_LIT>" ) <EOL> process . send_signal ( signal . SIGTERM ) <EOL> if process . poll ( ) is None : <EOL> import time <EOL> for t in [ <NUM_LIT> , <NUM_LIT> , <NUM_LIT> , <EOL> <NUM_LIT> , <NUM_LIT> , <NUM_LIT> , <EOL> <NUM_LIT> , <NUM_LIT> , <NUM_LIT> , <EOL> <NUM_LIT> , <NUM_LIT:1> , <NUM_LIT:2> , <NUM_LIT:5> ] : <EOL> time . sleep ( t ) <EOL> if process . poll ( ) is not None : <EOL> log . info ( "<STR_LIT>" ) <EOL> break <EOL> else : <EOL> log . info ( "<STR_LIT>" ) <EOL> os . kill ( process . _popen . pid , signal . SIGKILL ) </s>
