<s> '''<STR_LIT>''' <EOL> import os <EOL> import threading <EOL> import socket <EOL> import shutil <EOL> from cola . core . utils import import_job_desc , Clock , pack_local_job_error <EOL> from cola . core . rpc import FileTransportServer , client_call , FileTransportClient <EOL> from cola . core . zip import ZipHandler <EOL> from cola . core . logs import get_logger <EOL> from cola . job import Job <EOL> HEARTBEAT_INTERVAL = <NUM_LIT:20> <EOL> class WorkerJobInfo ( object ) : <EOL> def __init__ ( self , job_name , working_dir ) : <EOL> self . job_name = job_name <EOL> self . working_dir = working_dir <EOL> self . job = None <EOL> self . thread = None <EOL> self . clock = None <EOL> class Worker ( object ) : <EOL> def __init__ ( self , ctx ) : <EOL> self . ctx = ctx <EOL> self . master = self . ctx . master_addr <EOL> addr_dirname = self . ctx . addr . replace ( '<STR_LIT:.>' , '<STR_LIT:_>' ) . replace ( '<STR_LIT::>' , '<STR_LIT:_>' ) <EOL> self . working_dir = os . path . join ( self . ctx . working_dir , '<STR_LIT>' , addr_dirname ) <EOL> self . job_dir = os . path . join ( self . working_dir , '<STR_LIT>' ) <EOL> self . zip_dir = os . path . join ( self . working_dir , '<STR_LIT>' ) <EOL> self . running_jobs = { } <EOL> self . rpc_server = self . ctx . worker_rpc_server <EOL> assert self . rpc_server is not None <EOL> self . _register_rpc ( ) <EOL> self . stopped = threading . Event ( ) <EOL> self . logger = get_logger ( '<STR_LIT>' , server = self . ctx . master_ip ) <EOL> self . _ensure_exists ( self . job_dir ) <EOL> self . _ensure_exists ( self . zip_dir ) <EOL> FileTransportServer ( self . rpc_server , self . zip_dir ) <EOL> def _ensure_exists ( self , dir_ ) : <EOL> if not os . path . exists ( dir_ ) : <EOL> os . makedirs ( dir_ ) <EOL> def _register_rpc ( self ) : <EOL> if self . rpc_server : <EOL> self . rpc_server . register_function ( self . prepare , '<STR_LIT>' ) <EOL> self . rpc_server . register_function ( self . run_job , '<STR_LIT>' ) <EOL> self . rpc_server . register_function ( self . has_job , '<STR_LIT>' ) <EOL> self . rpc_server . register_function ( self . stop_running_job , <EOL> '<STR_LIT>' ) <EOL> self . rpc_server . register_function ( self . clear_running_job , <EOL> '<STR_LIT>' ) <EOL> self . rpc_server . register_function ( self . pack_job_error , <EOL> '<STR_LIT>' ) <EOL> self . rpc_server . register_function ( self . add_node , '<STR_LIT>' ) <EOL> self . rpc_server . register_function ( self . remove_node , '<STR_LIT>' ) <EOL> self . rpc_server . register_function ( self . shutdown , '<STR_LIT>' ) <EOL> def run ( self ) : <EOL> def _report ( ) : <EOL> while not self . stopped . is_set ( ) : <EOL> workers = client_call ( self . master , '<STR_LIT>' , <EOL> self . ctx . worker_addr ) <EOL> self . ctx . addrs = [ self . ctx . fix_addr ( worker ) for worker in workers ] <EOL> self . ctx . ips = [ self . ctx . fix_ip ( worker ) for worker in workers ] <EOL> self . stopped . wait ( HEARTBEAT_INTERVAL ) <EOL> self . _t = threading . Thread ( target = _report ) <EOL> self . _t . start ( ) <EOL> def _unzip ( self , job_name ) : <EOL> zip_file = os . path . join ( self . zip_dir , job_name + '<STR_LIT>' ) <EOL> job_path = os . path . join ( self . job_dir , job_name ) <EOL> if os . path . exists ( job_path ) : <EOL> shutil . rmtree ( job_path ) <EOL> if os . path . exists ( zip_file ) : <EOL> ZipHandler . uncompress ( zip_file , self . job_dir ) <EOL> def prepare ( self , job_name , unzip = True , overwrite = False , <EOL> settings = None ) : <EOL> self . logger . debug ( '<STR_LIT>' % job_name ) <EOL> if unzip : <EOL> self . _unzip ( job_name ) <EOL> src_job_name = job_name <EOL> job_path = os . path . join ( self . job_dir , job_name ) <EOL> if not os . path . exists ( job_path ) : <EOL> return False <EOL> job_desc = import_job_desc ( job_path ) <EOL> if settings is not None : <EOL> job_desc . update_settings ( settings ) <EOL> job_id = self . ctx . ips . index ( self . ctx . ip ) <EOL> clear = job_desc . settings . job . clear if self . ctx . is_local_mode else False <EOL> job_name , working_dir = self . ctx . _get_name_and_dir ( <EOL> self . working_dir , job_name , overwrite = overwrite , clear = clear ) <EOL> job = Job ( self . ctx , job_path , job_name , job_desc = job_desc , <EOL> working_dir = working_dir , rpc_server = self . rpc_server , <EOL> manager = self . ctx . manager , job_offset = job_id ) <EOL> t = threading . Thread ( target = job . run , args = ( True , ) ) <EOL> job_info = WorkerJobInfo ( job_name , working_dir ) <EOL> job_info . job = job <EOL> job_info . thread = t <EOL> self . running_jobs [ src_job_name ] = job_info <EOL> self . logger . debug ( '<STR_LIT>' % job_name ) <EOL> return True <EOL> def run_job ( self , job_name ) : <EOL> self . logger . debug ( '<STR_LIT>' % job_name ) <EOL> if job_name not in self . running_jobs : <EOL> self . logger . debug ( <EOL> '<STR_LIT>' % job_name ) <EOL> return False <EOL> job_info = self . running_jobs [ job_name ] <EOL> clock = Clock ( ) <EOL> job_info . clock = clock <EOL> job_info . thread . start ( ) <EOL> self . logger . debug ( '<STR_LIT>' % job_name ) <EOL> return True <EOL> def stop_running_job ( self , job_name ) : <EOL> job_info = self . running_jobs . get ( job_name ) <EOL> if job_info : <EOL> job_info . job . stop_running ( ) <EOL> def clear_running_job ( self , job_name ) : <EOL> job_info = self . running_jobs . get ( job_name ) <EOL> if job_info : <EOL> job_info . job . clear_running ( ) <EOL> job_info . thread . join ( ) <EOL> del self . running_jobs [ job_name ] <EOL> return job_info . clock . clock ( ) <EOL> def has_job ( self , job_name ) : <EOL> return job_name in self . running_jobs <EOL> def pack_job_error ( self , job_name ) : <EOL> working_dir = os . path . join ( self . working_dir , job_name ) <EOL> pack_dir = pack_local_job_error ( job_name , working_dir = working_dir , <EOL> logger = self . logger ) <EOL> zip_filename = os . path . join ( self . zip_dir , <EOL> '<STR_LIT>' % ( self . ctx . ip . replace ( '<STR_LIT:.>' , '<STR_LIT:_>' ) , job_name ) ) <EOL> if os . path . exists ( zip_filename ) : <EOL> os . remove ( zip_filename ) <EOL> ZipHandler . compress ( zip_filename , pack_dir ) <EOL> FileTransportClient ( self . master , zip_filename ) . send_file ( ) <EOL> def add_node ( self , worker ) : <EOL> self . ctx . add_node ( worker ) <EOL> for job_info in self . running_jobs . values ( ) : <EOL> job_info . job . add_node ( worker ) <EOL> def remove_node ( self , worker ) : <EOL> self . ctx . remove_node ( worker ) <EOL> for job_info in self . running_jobs . values ( ) : <EOL> job_info . job . remove_node ( worker ) <EOL> def shutdown ( self ) : <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> return <EOL> self . logger . debug ( '<STR_LIT>' ) <EOL> for job_info in self . running_jobs . values ( ) : <EOL> job_info . job . shutdown ( ) <EOL> job_info . thread . join ( ) <EOL> try : <EOL> self . ctx . manager . shutdown ( ) <EOL> except socket . error : <EOL> pass <EOL> self . stopped . set ( ) <EOL> self . _t . join ( ) <EOL> self . rpc_server . shutdown ( ) <EOL> self . logger . debug ( '<STR_LIT>' ) </s>
