<s> import collections <EOL> import logging <EOL> import os <EOL> from pulsar . mesos import ( <EOL> Scheduler , <EOL> MesosSchedulerDriver , <EOL> mesos_pb2 , <EOL> ) <EOL> from pulsar . client . util import to_base64_json <EOL> from pulsar . main import ( <EOL> PULSAR_ROOT_DIR , <EOL> ) <EOL> from pulsar . messaging import bind_amqp <EOL> log = logging . getLogger ( __name__ ) <EOL> DEFAULT_FRAMEWORK_USER = "<STR_LIT>" <EOL> DEFAULT_FRAMEWORK_NAME = "<STR_LIT>" <EOL> DEFAULT_FRAMEWORK_PRINCIPAL = "<STR_LIT>" <EOL> DEFAULT_EXECUTOR_ID = "<STR_LIT:default>" <EOL> DEFAULT_EXECUTOR_NAME = "<STR_LIT>" <EOL> DEFAULT_EXECUTOR_SOURCE = "<STR_LIT>" <EOL> class PulsarScheduler ( Scheduler ) : <EOL> def __init__ ( self , executor , manager_options , mesos_url ) : <EOL> self . executor = executor <EOL> self . manager_options = manager_options <EOL> self . mesos_url = mesos_url <EOL> self . taskData = { } <EOL> self . tasksLaunched = <NUM_LIT:0> <EOL> self . tasksFinished = <NUM_LIT:0> <EOL> self . messagesSent = <NUM_LIT:0> <EOL> self . messagesReceived = <NUM_LIT:0> <EOL> self . in_memory_queue = collections . deque ( ) <EOL> def registered ( self , driver , frameworkId , masterInfo ) : <EOL> log . info ( "<STR_LIT>" % frameworkId . value ) <EOL> def resourceOffers ( self , driver , offers ) : <EOL> log . info ( "<STR_LIT>" % len ( offers ) ) <EOL> for offer in offers : <EOL> tasks = self . _tasks_for_offer ( offer ) <EOL> if tasks : <EOL> log . info ( "<STR_LIT>" % tasks ) <EOL> try : <EOL> driver . launchTasks ( offer . id , tasks ) <EOL> except Exception : <EOL> log . exception ( "<STR_LIT>" ) <EOL> raise <EOL> def _tasks_for_offer ( self , offer ) : <EOL> tasks = [ ] <EOL> log . info ( "<STR_LIT>" % offer . id . value ) <EOL> try : <EOL> next_job = self . in_memory_queue . pop ( ) <EOL> except IndexError : <EOL> log . info ( "<STR_LIT>" ) <EOL> return tasks <EOL> tid = self . tasksLaunched <EOL> self . tasksLaunched += <NUM_LIT:1> <EOL> log . info ( "<STR_LIT>" % ( offer . hostname , tid ) ) <EOL> task = mesos_pb2 . TaskInfo ( ) <EOL> task . task_id . value = str ( tid ) <EOL> task . slave_id . value = offer . slave_id . value <EOL> task . name = "<STR_LIT>" % tid <EOL> task . executor . MergeFrom ( self . executor ) <EOL> self . _populate_task_data_for_job ( task , next_job ) <EOL> cpus = task . resources . add ( ) <EOL> cpus . name = "<STR_LIT>" <EOL> cpus . type = mesos_pb2 . Value . SCALAR <EOL> cpus . scalar . value = <NUM_LIT:1> <EOL> mem = task . resources . add ( ) <EOL> mem . name = "<STR_LIT>" <EOL> mem . type = mesos_pb2 . Value . SCALAR <EOL> mem . scalar . value = <NUM_LIT:32> <EOL> tasks . append ( task ) <EOL> self . taskData [ task . task_id . value ] = ( <EOL> offer . slave_id , <EOL> task . executor . executor_id <EOL> ) <EOL> return tasks <EOL> def statusUpdate ( self , driver , update ) : <EOL> log . info ( "<STR_LIT:%s>" % update . SerializeToString ( ) ) <EOL> if update . state == mesos_pb2 . TASK_FINISHED : <EOL> self . tasksFinished += <NUM_LIT:1> <EOL> slave_id , executor_id = self . taskData [ update . task_id . value ] <EOL> self . messagesSent += <NUM_LIT:1> <EOL> driver . sendFrameworkMessage ( <EOL> executor_id , <EOL> slave_id , <EOL> '<STR_LIT>' <EOL> ) <EOL> def frameworkMessage ( self , driver , executorId , slaveId , message ) : <EOL> self . messagesReceived += <NUM_LIT:1> <EOL> log . info ( "<STR_LIT>" , repr ( str ( message ) ) ) <EOL> def handle_setup_message ( self , body , message ) : <EOL> try : <EOL> self . __queue_setup_message ( body ) <EOL> finally : <EOL> message . ack ( ) <EOL> def _populate_task_data_for_job ( self , task , job ) : <EOL> if "<STR_LIT>" not in job : <EOL> job [ "<STR_LIT>" ] = [ ] <EOL> job [ "<STR_LIT>" ] . extend ( self . _mesos_env_vars ( ) ) <EOL> task_data = dict ( <EOL> job = job , <EOL> manager = self . manager_options <EOL> ) <EOL> task . data = to_base64_json ( <EOL> task_data <EOL> ) <EOL> def __queue_setup_message ( self , body ) : <EOL> self . in_memory_queue . appendleft ( body ) <EOL> def _mesos_env_vars ( self ) : <EOL> return [ <EOL> dict ( name = "<STR_LIT>" , value = self . mesos_url ) , <EOL> ] <EOL> def run ( master , manager_options , config ) : <EOL> executor = mesos_pb2 . ExecutorInfo ( ) <EOL> executor . executor_id . value = DEFAULT_EXECUTOR_ID <EOL> executor . command . value = os . path . join ( <EOL> PULSAR_ROOT_DIR , <EOL> "<STR_LIT>" , <EOL> "<STR_LIT>" <EOL> ) <EOL> executor . name = DEFAULT_EXECUTOR_NAME <EOL> executor . source = DEFAULT_EXECUTOR_SOURCE <EOL> framework = mesos_pb2 . FrameworkInfo ( ) <EOL> framework . user = DEFAULT_FRAMEWORK_USER <EOL> framework . name = DEFAULT_FRAMEWORK_NAME <EOL> framework . principal = DEFAULT_FRAMEWORK_PRINCIPAL <EOL> scheduler = PulsarScheduler ( <EOL> executor , <EOL> manager_options = manager_options , <EOL> mesos_url = master , <EOL> ) <EOL> driver = MesosSchedulerDriver ( <EOL> scheduler , <EOL> framework , <EOL> master <EOL> ) <EOL> message_queue_url = config . get ( "<STR_LIT>" , None ) <EOL> exchange = bind_amqp . get_exchange ( <EOL> message_queue_url , <EOL> manager_name = manager_options [ "<STR_LIT>" ] , <EOL> conf = config , <EOL> ) <EOL> def drain ( ) : <EOL> exchange . consume ( "<STR_LIT>" , callback = scheduler . handle_setup_message , check = True ) <EOL> log . info ( "<STR_LIT>" ) <EOL> bind_amqp . start_setup_consumer ( <EOL> exchange , <EOL> drain , <EOL> ) <EOL> try : <EOL> log . info ( "<STR_LIT>" ) <EOL> if driver . run ( ) != mesos_pb2 . DRIVER_STOPPED : <EOL> raise Exception ( "<STR_LIT>" ) <EOL> except Exception : <EOL> log . exception ( "<STR_LIT>" ) <EOL> raise <EOL> finally : <EOL> driver . stop ( ) </s>
