<s> '''<STR_LIT>''' <EOL> import os <EOL> import time <EOL> import threading <EOL> try : <EOL> import cPickle as pickle <EOL> except ImportError : <EOL> import pickle <EOL> import tempfile <EOL> import shutil <EOL> from cola . functions . counter import CounterServer <EOL> from cola . functions . budget import BudgetApplyServer , ALLFINISHED <EOL> from cola . functions . speed import SpeedControlServer <EOL> from cola . cluster . tracker import WorkerTracker , JobTracker <EOL> from cola . cluster . stage import Stage <EOL> from cola . core . rpc import FileTransportServer , FileTransportClient , client_call <EOL> from cola . core . zip import ZipHandler <EOL> from cola . core . utils import import_job_desc <EOL> from cola . core . logs import get_logger , LogRecordSocketReceiver <EOL> RUNNING , HANGUP , STOPPED = range ( <NUM_LIT:3> ) <EOL> STATUSES = [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] <EOL> CONTINOUS_HEARTBEAT = <NUM_LIT> <EOL> HEARTBEAT_INTERVAL = <NUM_LIT:20> <EOL> HEARTBEAT_CHECK_INTERVAL = <NUM_LIT:3> * HEARTBEAT_INTERVAL <EOL> JOB_CHECK_INTERVAL = <NUM_LIT:5> <EOL> JOB_META_STATUS_FILENAME = '<STR_LIT>' <EOL> class JobMaster ( object ) : <EOL> def __init__ ( self , ctx , job_name , job_desc , workers ) : <EOL> self . working_dir = os . path . join ( ctx . working_dir , '<STR_LIT>' , <EOL> '<STR_LIT>' , job_name ) <EOL> if not os . path . exists ( self . working_dir ) : <EOL> os . makedirs ( self . working_dir ) <EOL> self . job_name = job_name <EOL> self . job_desc = job_desc <EOL> self . settings = job_desc . settings <EOL> self . rpc_server = ctx . master_rpc_server <EOL> self . inited = False <EOL> self . lock = threading . Lock ( ) <EOL> self . init ( ) <EOL> self . workers = workers <EOL> def _init_counter_server ( self ) : <EOL> counter_dir = os . path . join ( self . working_dir , '<STR_LIT>' ) <EOL> self . counter_server = CounterServer ( counter_dir , self . settings , <EOL> rpc_server = self . rpc_server , <EOL> app_name = self . job_name ) <EOL> def _init_budget_server ( self ) : <EOL> budget_dir = os . path . join ( self . working_dir , '<STR_LIT>' ) <EOL> self . budget_server = BudgetApplyServer ( budget_dir , self . settings , <EOL> rpc_server = self . rpc_server , <EOL> app_name = self . job_name ) <EOL> def _init_speed_server ( self ) : <EOL> speed_dir = os . path . join ( self . working_dir , '<STR_LIT>' ) <EOL> self . speed_server = SpeedControlServer ( speed_dir , self . settings , <EOL> rpc_server = self . rpc_server , <EOL> app_name = self . job_name ) <EOL> def init ( self ) : <EOL> with self . lock : <EOL> if self . inited : <EOL> return <EOL> self . _init_counter_server ( ) <EOL> self . _init_budget_server ( ) <EOL> self . _init_speed_server ( ) <EOL> self . inited = True <EOL> def remove_worker ( self , worker ) : <EOL> if worker not in self . workers : <EOL> return <EOL> self . workers . remove ( worker ) <EOL> for node in self . workers : <EOL> client_call ( node , '<STR_LIT>' , worker ) <EOL> def add_worker ( self , worker ) : <EOL> if worker in self . workers : <EOL> return <EOL> for node in self . workers : <EOL> client_call ( node , '<STR_LIT>' , worker ) <EOL> self . workers . append ( worker ) <EOL> def has_worker ( self , worker ) : <EOL> return worker in self . workers <EOL> def shutdown ( self ) : <EOL> with self . lock : <EOL> if not self . inited : <EOL> return <EOL> self . counter_server . shutdown ( ) <EOL> self . budget_server . shutdown ( ) <EOL> self . speed_server . shutdown ( ) <EOL> self . inited = False <EOL> class Master ( object ) : <EOL> def __init__ ( self , ctx ) : <EOL> self . ctx = ctx <EOL> self . rpc_server = self . ctx . master_rpc_server <EOL> assert self . rpc_server is not None <EOL> self . working_dir = os . path . join ( self . ctx . working_dir , '<STR_LIT>' ) <EOL> self . zip_dir = os . path . join ( self . working_dir , '<STR_LIT>' ) <EOL> self . job_dir = os . path . join ( self . working_dir , '<STR_LIT>' ) <EOL> if not os . path . exists ( self . zip_dir ) : <EOL> os . makedirs ( self . zip_dir ) <EOL> if not os . path . exists ( self . job_dir ) : <EOL> os . makedirs ( self . job_dir ) <EOL> self . worker_tracker = WorkerTracker ( ) <EOL> self . job_tracker = JobTracker ( ) <EOL> self . black_list = [ ] <EOL> self . stopped = threading . Event ( ) <EOL> self . logger = get_logger ( "<STR_LIT>" ) <EOL> self . _init_log_server ( self . logger ) <EOL> self . _register_rpc ( ) <EOL> self . load ( ) <EOL> FileTransportServer ( self . rpc_server , self . zip_dir ) <EOL> def load ( self ) : <EOL> self . runned_job_metas = { } <EOL> job_meta_file = os . path . join ( self . working_dir , JOB_META_STATUS_FILENAME ) <EOL> if os . path . exists ( job_meta_file ) and os . path . getsize ( job_meta_file ) > <NUM_LIT:0> : <EOL> try : <EOL> with open ( job_meta_file ) as f : <EOL> self . runned_job_metas = pickle . load ( f ) <EOL> except : <EOL> pass <EOL> def save ( self ) : <EOL> job_meta_file = os . path . join ( self . working_dir , JOB_META_STATUS_FILENAME ) <EOL> with open ( job_meta_file , '<STR_LIT:w>' ) as f : <EOL> pickle . dump ( self . runned_job_metas , f ) <EOL> def _register_rpc ( self ) : <EOL> self . rpc_server . register_function ( self . run_job , '<STR_LIT>' ) <EOL> self . rpc_server . register_function ( self . stop_job , '<STR_LIT>' ) <EOL> self . rpc_server . register_function ( self . pack_job_error , '<STR_LIT>' ) <EOL> self . rpc_server . register_function ( self . list_runnable_jobs , <EOL> '<STR_LIT>' ) <EOL> self . rpc_server . register_function ( lambda : self . job_tracker . running_jobs . keys ( ) , <EOL> '<STR_LIT>' ) <EOL> self . rpc_server . register_function ( self . list_workers , <EOL> '<STR_LIT>' ) <EOL> self . rpc_server . register_function ( self . shutdown , '<STR_LIT>' ) <EOL> self . rpc_server . register_function ( self . register_heartbeat , <EOL> '<STR_LIT>' ) <EOL> def register_heartbeat ( self , worker ) : <EOL> need_to_add_worker_nodes = self . worker_tracker . register_worker ( worker ) <EOL> if need_to_add_worker_nodes is not None : <EOL> for node in need_to_add_worker_nodes : <EOL> client_call ( node , '<STR_LIT>' , worker ) <EOL> return self . worker_tracker . workers . keys ( ) <EOL> def _init_log_server ( self , logger ) : <EOL> self . log_server = LogRecordSocketReceiver ( host = self . ctx . ip , <EOL> logger = self . logger ) <EOL> self . log_t = threading . Thread ( target = self . log_server . serve_forever ) <EOL> self . log_t . start ( ) <EOL> def _shutdown_log_server ( self ) : <EOL> if hasattr ( self , '<STR_LIT>' ) : <EOL> self . log_server . shutdown ( ) <EOL> self . log_t . join ( ) <EOL> def _check_workers ( self ) : <EOL> while not self . stopped . is_set ( ) : <EOL> for worker , info in self . worker_tracker . workers . iteritems ( ) : <EOL> if int ( time . time ( ) ) - info . last_update > HEARTBEAT_CHECK_INTERVAL : <EOL> info . continous_register = <NUM_LIT:0> <EOL> if info . status == RUNNING : <EOL> info . status = HANGUP <EOL> elif info . status == HANGUP : <EOL> info . status = STOPPED <EOL> self . black_list . append ( worker ) <EOL> for job in self . job_tracker . running_jobs : <EOL> self . job_tracker . remove_worker ( job , worker ) <EOL> elif info . continous_register >= CONTINOUS_HEARTBEAT : <EOL> if info . status != RUNNING : <EOL> info . status = RUNNING <EOL> if worker in self . black_list : <EOL> self . black_list . remove ( worker ) <EOL> for job in self . job_tracker . running_jobs : <EOL> if not client_call ( worker , '<STR_LIT>' ) : <EOL> client_call ( worker , '<STR_LIT>' , job ) <EOL> client_call ( worker , '<STR_LIT>' , job ) <EOL> self . job_tracker . add_worker ( job , worker ) <EOL> self . stopped . wait ( HEARTBEAT_CHECK_INTERVAL ) <EOL> def _check_jobs ( self ) : <EOL> while not self . stopped . is_set ( ) : <EOL> for job_master in self . job_tracker . running_jobs . values ( ) : <EOL> if job_master . budget_server . get_status ( ) == ALLFINISHED : <EOL> self . stop_job ( job_master . job_name ) <EOL> self . job_tracker . remove_job ( job_master . job_name ) <EOL> self . stopped . wait ( JOB_CHECK_INTERVAL ) <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 _register_runned_job ( self , job_name , job_desc ) : <EOL> self . runned_job_metas [ job_name ] = { '<STR_LIT>' : job_desc . name , <EOL> '<STR_LIT>' : time . time ( ) } <EOL> def run ( self ) : <EOL> self . _worker_t = threading . Thread ( target = self . _check_workers ) <EOL> self . _worker_t . start ( ) <EOL> self . _job_t = threading . Thread ( target = self . _check_jobs ) <EOL> self . _job_t . start ( ) <EOL> def run_job ( self , job_name , unzip = False , <EOL> wait_for_workers = False ) : <EOL> if wait_for_workers : <EOL> while not self . stopped . is_set ( ) : <EOL> if len ( self . worker_tracker . workers ) > <NUM_LIT:0> : <EOL> break <EOL> stopped = self . stopped . wait ( <NUM_LIT:3> ) <EOL> if stopped : <EOL> return <EOL> if unzip : <EOL> self . _unzip ( job_name ) <EOL> job_path = os . path . join ( self . job_dir , job_name ) <EOL> job_desc = import_job_desc ( job_path ) <EOL> job_master = JobMaster ( self . ctx , job_name , job_desc , <EOL> self . worker_tracker . workers . keys ( ) ) <EOL> job_master . init ( ) <EOL> self . job_tracker . register_job ( job_name , job_master ) <EOL> self . _register_runned_job ( job_name , job_desc ) <EOL> zip_file = os . path . join ( self . zip_dir , job_name + '<STR_LIT>' ) <EOL> for worker in job_master . workers : <EOL> FileTransportClient ( worker , zip_file ) . send_file ( ) <EOL> self . logger . debug ( <EOL> '<STR_LIT>' % job_name ) <EOL> self . logger . debug ( <EOL> '<STR_LIT>' % job_master . workers ) <EOL> stage = Stage ( job_master . workers , '<STR_LIT>' , logger = self . logger ) <EOL> prepared_ok = stage . barrier ( True , job_name ) <EOL> if not prepared_ok : <EOL> self . logger . error ( "<STR_LIT>" ) <EOL> return <EOL> self . logger . debug ( <EOL> '<STR_LIT>' % job_name ) <EOL> stage = Stage ( job_master . workers , '<STR_LIT>' , logger = self . logger ) <EOL> run_ok = stage . barrier ( True , job_name ) <EOL> if not run_ok : <EOL> self . logger . error ( "<STR_LIT>" % job_name ) <EOL> def stop_job ( self , job_name ) : <EOL> job_master = self . job_tracker . get_job_master ( job_name ) <EOL> stage = Stage ( job_master . workers , '<STR_LIT>' ) <EOL> stage . barrier ( True , job_name ) <EOL> stage = Stage ( job_master . workers , '<STR_LIT>' ) <EOL> stage . barrier ( True , job_name ) <EOL> self . job_tracker . remove_job ( job_name ) <EOL> self . logger . debug ( '<STR_LIT>' % job_name ) <EOL> def pack_job_error ( self , job_name ) : <EOL> job_master = self . job_tracker . get_job_master ( job_name ) <EOL> stage = Stage ( job_master . workers , '<STR_LIT>' ) <EOL> stage . barrier ( True , job_name ) <EOL> error_dir = os . path . join ( self . working_dir , '<STR_LIT>' ) <EOL> if not os . path . exists ( error_dir ) : <EOL> os . makedirs ( error_dir ) <EOL> error_filename = os . path . join ( error_dir , '<STR_LIT>' % job_name ) <EOL> suffix = '<STR_LIT>' % job_name <EOL> temp_dir = tempfile . mkdtemp ( ) <EOL> try : <EOL> for name in os . listdir ( self . zip_dir ) : <EOL> if name . endswith ( suffix ) : <EOL> shutil . move ( os . path . join ( self . zip_dir , name ) , temp_dir ) <EOL> ZipHandler . compress ( error_filename , temp_dir ) <EOL> finally : <EOL> shutil . rmtree ( temp_dir ) <EOL> return error_filename <EOL> def list_runnable_jobs ( self ) : <EOL> job_dirs = filter ( lambda s : os . path . isdir ( os . path . join ( self . job_dir , s ) ) , <EOL> os . listdir ( self . job_dir ) ) <EOL> jobs = { } <EOL> for job_dir in job_dirs : <EOL> desc = import_job_desc ( os . path . join ( self . job_dir , job_dir ) ) <EOL> jobs [ job_dir ] = desc . name <EOL> return jobs <EOL> def has_running_jobs ( self ) : <EOL> return len ( self . job_tracker . running_jobs ) > <NUM_LIT:0> <EOL> def list_workers ( self ) : <EOL> return [ ( worker , STATUSES [ worker_info . status ] ) for worker , worker_info in self . worker_tracker . workers . iteritems ( ) ] <EOL> def _stop_all_jobs ( self ) : <EOL> for job_name in self . job_tracker . running_jobs . keys ( ) : <EOL> self . stop_job ( job_name ) <EOL> def _shutdown_all_workers ( self ) : <EOL> stage = Stage ( self . worker_tracker . workers . keys ( ) , '<STR_LIT>' ) <EOL> stage . barrier ( True ) <EOL> def shutdown ( self ) : <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> return <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> return <EOL> self . logger . debug ( '<STR_LIT>' ) <EOL> self . stopped . set ( ) <EOL> self . _stop_all_jobs ( ) <EOL> self . _shutdown_all_workers ( ) <EOL> self . _worker_t . join ( ) <EOL> self . _job_t . join ( ) <EOL> self . save ( ) <EOL> self . rpc_server . shutdown ( ) <EOL> self . logger . debug ( '<STR_LIT>' ) <EOL> self . _shutdown_log_server ( ) </s>
