<s> import zmq <EOL> import time <EOL> from collections import defaultdict <EOL> from cloudasr import Poller <EOL> from cloudasr . messages import HeartbeatMessage <EOL> from cloudasr . messages . helpers import * <EOL> def create_master ( worker_address , frontend_address , monitor_address ) : <EOL> poller = create_poller ( worker_address , frontend_address ) <EOL> context = zmq . Context ( ) <EOL> monitor = context . socket ( zmq . PUSH ) <EOL> monitor . connect ( monitor_address ) <EOL> run_forever = lambda : True <EOL> return Master ( poller , monitor , run_forever ) <EOL> def create_poller ( worker_address , frontend_address ) : <EOL> context = zmq . Context ( ) <EOL> worker_socket = context . socket ( zmq . PULL ) <EOL> worker_socket . bind ( worker_address ) <EOL> frontend_socket = context . socket ( zmq . REP ) <EOL> frontend_socket . bind ( frontend_address ) <EOL> sockets = { <EOL> "<STR_LIT>" : { "<STR_LIT>" : worker_socket , "<STR_LIT>" : worker_socket . recv , "<STR_LIT>" : worker_socket . send_json } , <EOL> "<STR_LIT>" : { "<STR_LIT>" : frontend_socket , "<STR_LIT>" : frontend_socket . recv , "<STR_LIT>" : frontend_socket . send } , <EOL> } <EOL> time_func = time . time <EOL> return Poller ( sockets , time_func ) <EOL> class Master : <EOL> def __init__ ( self , poller , monitor , should_continue ) : <EOL> self . poller = poller <EOL> self . should_continue = should_continue <EOL> self . workers = WorkerPool ( monitor ) <EOL> self . time = <NUM_LIT:0> <EOL> def run ( self ) : <EOL> while self . should_continue ( ) : <EOL> messages , self . time = self . poller . poll ( ) <EOL> if "<STR_LIT>" in messages : <EOL> self . handle_worker_request ( messages [ "<STR_LIT>" ] ) <EOL> if "<STR_LIT>" in messages : <EOL> self . handle_fronted_request ( messages [ "<STR_LIT>" ] ) <EOL> def handle_fronted_request ( self , message ) : <EOL> try : <EOL> request = parseWorkerRequestMessage ( message ) <EOL> model = request . model <EOL> worker = self . workers . get_worker ( model , self . time ) <EOL> message = createMasterResponseMessage ( "<STR_LIT>" , worker ) <EOL> self . poller . send ( "<STR_LIT>" , message . SerializeToString ( ) ) <EOL> except NoWorkerAvailableException : <EOL> message = createMasterResponseMessage ( "<STR_LIT>" ) <EOL> self . poller . send ( "<STR_LIT>" , message . SerializeToString ( ) ) <EOL> def handle_worker_request ( self , message ) : <EOL> statuses = { <EOL> HeartbeatMessage . STARTED : "<STR_LIT>" , <EOL> HeartbeatMessage . WAITING : "<STR_LIT>" , <EOL> HeartbeatMessage . WORKING : "<STR_LIT>" , <EOL> HeartbeatMessage . FINISHED : "<STR_LIT>" <EOL> } <EOL> heartbeat = parseHeartbeatMessage ( message ) <EOL> address = heartbeat . address <EOL> model = heartbeat . model <EOL> status = statuses [ heartbeat . status ] <EOL> self . workers . add_worker ( model , address , status , self . time ) <EOL> class WorkerPool : <EOL> def __init__ ( self , monitor ) : <EOL> self . workers_status = defaultdict ( lambda : { "<STR_LIT:status>" : "<STR_LIT>" , "<STR_LIT>" : <NUM_LIT:0> , "<STR_LIT>" : <NUM_LIT:0> } ) <EOL> self . available_workers = defaultdict ( list ) <EOL> self . monitor = monitor <EOL> def get_worker ( self , model , time ) : <EOL> worker = self . find_available_worker ( model , time ) <EOL> if worker is None : <EOL> raise NoWorkerAvailableException ( ) <EOL> self . update_worker_status ( model , worker , "<STR_LIT>" , time ) <EOL> return worker <EOL> def find_available_worker ( self , model , time ) : <EOL> while len ( self . available_workers [ model ] ) > <NUM_LIT:0> : <EOL> worker = self . available_workers [ model ] . pop ( <NUM_LIT:0> ) <EOL> if self . is_worker_available ( worker , time ) : <EOL> return worker <EOL> return None <EOL> def is_worker_available ( self , worker , time ) : <EOL> status = self . workers_status [ worker ] <EOL> return status [ "<STR_LIT:status>" ] and status [ "<STR_LIT>" ] > time - <NUM_LIT:10> <EOL> def add_worker ( self , model , address , status , time ) : <EOL> worker_status = self . workers_status [ address ] [ "<STR_LIT:status>" ] <EOL> if worker_status == "<STR_LIT>" : <EOL> if status == "<STR_LIT>" or status == "<STR_LIT>" : <EOL> self . available_workers [ model ] . append ( address ) <EOL> self . update_worker_status ( model , address , "<STR_LIT>" , time ) <EOL> if status == "<STR_LIT>" : <EOL> self . update_worker_status ( model , address , "<STR_LIT>" , time ) <EOL> if status == "<STR_LIT>" : <EOL> self . workers_status [ address ] [ "<STR_LIT>" ] += <NUM_LIT:1> <EOL> if self . workers_status [ address ] [ "<STR_LIT>" ] == <NUM_LIT:10> : <EOL> self . available_workers [ model ] . append ( address ) <EOL> self . update_worker_status ( model , address , "<STR_LIT>" , time ) <EOL> elif worker_status == "<STR_LIT>" : <EOL> self . available_workers [ model ] . append ( address ) <EOL> self . update_worker_status ( model , address , "<STR_LIT>" , time ) <EOL> elif worker_status == "<STR_LIT>" : <EOL> self . update_worker_status ( model , address , "<STR_LIT>" , time ) <EOL> def update_worker_status ( self , model , worker , status , time ) : <EOL> self . workers_status [ worker ] = { <EOL> "<STR_LIT:status>" : "<STR_LIT>" if status == "<STR_LIT>" else status , <EOL> "<STR_LIT>" : time , <EOL> "<STR_LIT>" : <NUM_LIT:0> <EOL> } <EOL> worker_status = createWorkerStatusMessage ( worker , model , status , int ( time ) ) <EOL> self . monitor . send ( worker_status . SerializeToString ( ) ) <EOL> class NoWorkerAvailableException ( Exception ) : <EOL> pass </s>
