def main(args): context = Context.singleton() executor = ExampleExecutor() driver = PesosExecutorDriver(executor, context=context) print('Starting driver') driver.start() print('Joining driver') driver.join() context.stop()
def test_mesos_executor_abort_on_disconnection(self): os.environ['MESOS_SLAVE_PID'] = str(self.slave.pid) shutdown_event, shutdown_call, shutdown_side_effect = self.mock_method( ) executor = mock.create_autospec(Executor, spec_set=True) executor.shutdown = mock.Mock(side_effect=shutdown_side_effect) driver_context = Context() driver_context.start() self.close_contexts.append(driver_context) driver = PesosExecutorDriver(executor, context=driver_context) assert driver.start() == mesos_pb2.DRIVER_RUNNING driver.executor_process.connected.wait(timeout=MAX_TIMEOUT) assert driver.executor_process.connected.is_set() self.slave.register_event.wait(timeout=MAX_TIMEOUT) assert self.slave.register_event.is_set() # kill connection conn = driver_context._connections[self.slave.pid] conn.close() # abort event shutdown_event.wait(timeout=MAX_TIMEOUT) assert shutdown_event.is_set() assert shutdown_call.mock_calls == [mock.call(driver)]
def test_mesos_executor_driver_status_update(self): os.environ['MESOS_SLAVE_PID'] = str(self.slave.pid) executor = Executor() driver = PesosExecutorDriver(executor, context=self.context) assert driver.start() == mesos_pb2.DRIVER_RUNNING # wait until registered driver.executor_process.connected.wait(timeout=MAX_TIMEOUT) assert driver.executor_process.connected.is_set() task_status = mesos_pb2.TaskStatus( task_id=mesos_pb2.TaskID(value='task-id'), state=mesos_pb2.TASK_FAILED, ) assert driver.send_status_update( task_status) == mesos_pb2.DRIVER_RUNNING self.slave.status_update_event.wait(timeout=MAX_TIMEOUT) assert self.slave.status_update_event.is_set() assert len(self.slave.status_updates) == 1 assert self.slave.status_updates[0][0] == driver.executor_process.pid status_update_message = self.slave.status_updates[0][1] status_update = status_update_message.update assert status_update.slave_id == self.slave.slave_id assert status_update.framework_id == mesos_pb2.FrameworkID( value=self.FRAMEWORK_ID) assert status_update.executor_id == mesos_pb2.ExecutorID( value=self.EXECUTOR_ID) assert status_update.status.task_id == task_status.task_id assert status_update.status.state == task_status.state assert status_update.HasField('timestamp') assert status_update.HasField('uuid')
def test_mesos_executor_framework_message_delivery(self): os.environ['MESOS_SLAVE_PID'] = str(self.slave.pid) framework_message_event, framework_message_call, framework_message_side_effect = ( self.mock_method()) executor = mock.create_autospec(Executor, spec_set=True) executor.frameworkMessage = mock.Mock( side_effect=framework_message_side_effect) driver = PesosExecutorDriver(executor, context=self.context) assert driver.start() == mesos_pb2.DRIVER_RUNNING # wait until registered driver.executor_process.connected.wait(timeout=MAX_TIMEOUT) assert driver.executor_process.connected.is_set() self.slave.send_framework_message( driver.executor_process.pid, mesos_pb2.FrameworkID(value=self.FRAMEWORK_ID), mesos_pb2.ExecutorID(value=self.EXECUTOR_ID), data=b'beep boop beep', ) framework_message_event.wait(timeout=MAX_TIMEOUT) assert framework_message_event.is_set() assert framework_message_call.mock_calls == [ mock.call(driver, b'beep boop beep') ]
def test_mesos_executor_kill_task(self): os.environ['MESOS_SLAVE_PID'] = str(self.slave.pid) kill_task_event, kill_task_call, kill_task_side_effect = self.mock_method() executor = mock.create_autospec(Executor, spec_set=True) executor.killTask = mock.Mock(side_effect=kill_task_side_effect) driver = PesosExecutorDriver(executor, context=self.context) assert driver.start() == mesos_pb2.DRIVER_RUNNING # wait until registered driver.executor_process.connected.wait(timeout=MAX_TIMEOUT) assert driver.executor_process.connected.is_set() # now launch task task_id = mesos_pb2.TaskID(value='task-id') self.slave.send_kill_task( driver.executor_process.pid, mesos_pb2.FrameworkID(value=self.FRAMEWORK_ID), task_id, ) kill_task_event.wait(timeout=MAX_TIMEOUT) assert kill_task_event.is_set() assert kill_task_call.mock_calls == [mock.call(driver, task_id)]
def test_mesos_executor_abort_on_disconnection(self): os.environ['MESOS_SLAVE_PID'] = str(self.slave.pid) shutdown_event, shutdown_call, shutdown_side_effect = self.mock_method() executor = mock.create_autospec(Executor, spec_set=True) executor.shutdown = mock.Mock(side_effect=shutdown_side_effect) driver_context = Context() driver_context.start() self.close_contexts.append(driver_context) driver = PesosExecutorDriver(executor, context=driver_context) assert driver.start() == mesos_pb2.DRIVER_RUNNING driver.executor_process.connected.wait(timeout=MAX_TIMEOUT) assert driver.executor_process.connected.is_set() self.slave.register_event.wait(timeout=MAX_TIMEOUT) assert self.slave.register_event.is_set() # kill connection conn = driver_context._connections[self.slave.pid] conn.close() # abort event shutdown_event.wait(timeout=MAX_TIMEOUT) assert shutdown_event.is_set() assert shutdown_call.mock_calls == [mock.call(driver)]
def test_mesos_executor_driver_framework_message(self): os.environ['MESOS_SLAVE_PID'] = str(self.slave.pid) executor = Executor() driver = PesosExecutorDriver(executor, context=self.context) assert driver.start() == mesos_pb2.DRIVER_RUNNING # wait until registered driver.executor_process.connected.wait(timeout=MAX_TIMEOUT) assert driver.executor_process.connected.is_set() # send and wait for framework message assert driver.send_framework_message( b'beep boop beep') == mesos_pb2.DRIVER_RUNNING self.slave.framework_message_event.wait(timeout=MAX_TIMEOUT) assert self.slave.framework_message_event.is_set() framework_message = internal.ExecutorToFrameworkMessage( slave_id=self.slave.slave_id, framework_id=mesos_pb2.FrameworkID(value=self.FRAMEWORK_ID), executor_id=mesos_pb2.ExecutorID(value=self.EXECUTOR_ID), data=b'beep boop beep', ) assert len(self.slave.framework_messages) == 1 assert self.slave.framework_messages[0][ 0] == driver.executor_process.pid assert self.slave.framework_messages[0][1] == framework_message, ( '%s != %s' % (self.slave.framework_messages[0][1], framework_message))
def test_mesos_executor_driver_status_update(self): os.environ['MESOS_SLAVE_PID'] = str(self.slave.pid) executor = Executor() driver = PesosExecutorDriver(executor, context=self.context) assert driver.start() == mesos_pb2.DRIVER_RUNNING # wait until registered driver.executor_process.connected.wait(timeout=MAX_TIMEOUT) assert driver.executor_process.connected.is_set() task_status = mesos_pb2.TaskStatus( task_id=mesos_pb2.TaskID(value='task-id'), state=mesos_pb2.TASK_FAILED, ) assert driver.send_status_update(task_status) == mesos_pb2.DRIVER_RUNNING self.slave.status_update_event.wait(timeout=MAX_TIMEOUT) assert self.slave.status_update_event.is_set() assert len(self.slave.status_updates) == 1 assert self.slave.status_updates[0][0] == driver.executor_process.pid status_update_message = self.slave.status_updates[0][1] status_update = status_update_message.update assert status_update.slave_id == self.slave.slave_id assert status_update.framework_id == mesos_pb2.FrameworkID(value=self.FRAMEWORK_ID) assert status_update.executor_id == mesos_pb2.ExecutorID(value=self.EXECUTOR_ID) assert status_update.status.task_id == task_status.task_id assert status_update.status.state == task_status.state assert status_update.HasField('timestamp') assert status_update.HasField('uuid')
def test_mesos_executor_driver_framework_message(self): os.environ['MESOS_SLAVE_PID'] = str(self.slave.pid) executor = Executor() driver = PesosExecutorDriver(executor, context=self.context) assert driver.start() == mesos_pb2.DRIVER_RUNNING # wait until registered driver.executor_process.connected.wait(timeout=MAX_TIMEOUT) assert driver.executor_process.connected.is_set() # send and wait for framework message assert driver.send_framework_message(b'beep boop beep') == mesos_pb2.DRIVER_RUNNING self.slave.framework_message_event.wait(timeout=MAX_TIMEOUT) assert self.slave.framework_message_event.is_set() framework_message = internal.ExecutorToFrameworkMessage( slave_id=self.slave.slave_id, framework_id=mesos_pb2.FrameworkID(value=self.FRAMEWORK_ID), executor_id=mesos_pb2.ExecutorID(value=self.EXECUTOR_ID), data=b'beep boop beep', ) assert len(self.slave.framework_messages) == 1 assert self.slave.framework_messages[0][0] == driver.executor_process.pid assert self.slave.framework_messages[0][1] == framework_message, ( '%s != %s' % (self.slave.framework_messages[0][1], framework_message))
def test_mesos_executor_kill_task(self): os.environ['MESOS_SLAVE_PID'] = str(self.slave.pid) kill_task_event, kill_task_call, kill_task_side_effect = self.mock_method( ) executor = mock.create_autospec(Executor, spec_set=True) executor.killTask = mock.Mock(side_effect=kill_task_side_effect) driver = PesosExecutorDriver(executor, context=self.context) assert driver.start() == mesos_pb2.DRIVER_RUNNING # wait until registered driver.executor_process.connected.wait(timeout=MAX_TIMEOUT) assert driver.executor_process.connected.is_set() # now launch task task_id = mesos_pb2.TaskID(value='task-id') self.slave.send_kill_task( driver.executor_process.pid, mesos_pb2.FrameworkID(value=self.FRAMEWORK_ID), task_id, ) kill_task_event.wait(timeout=MAX_TIMEOUT) assert kill_task_event.is_set() assert kill_task_call.mock_calls == [mock.call(driver, task_id)]
def test_mesos_executor_framework_message_delivery(self): os.environ['MESOS_SLAVE_PID'] = str(self.slave.pid) framework_message_event, framework_message_call, framework_message_side_effect = ( self.mock_method()) executor = mock.create_autospec(Executor, spec_set=True) executor.frameworkMessage = mock.Mock(side_effect=framework_message_side_effect) driver = PesosExecutorDriver(executor, context=self.context) assert driver.start() == mesos_pb2.DRIVER_RUNNING # wait until registered driver.executor_process.connected.wait(timeout=MAX_TIMEOUT) assert driver.executor_process.connected.is_set() self.slave.send_framework_message( driver.executor_process.pid, mesos_pb2.FrameworkID(value=self.FRAMEWORK_ID), mesos_pb2.ExecutorID(value=self.EXECUTOR_ID), data=b'beep boop beep', ) framework_message_event.wait(timeout=MAX_TIMEOUT) assert framework_message_event.is_set() assert framework_message_call.mock_calls == [mock.call(driver, b'beep boop beep')]
def test_mesos_executor_driver_init(self): os.environ['MESOS_SLAVE_PID'] = str(self.slave.pid) executor = Executor() driver = PesosExecutorDriver(executor, context=self.context) assert driver.start() == mesos_pb2.DRIVER_RUNNING assert driver.stop() == mesos_pb2.DRIVER_STOPPED assert driver.join() == mesos_pb2.DRIVER_STOPPED
def test_mesos_executor_driver_abort(self): os.environ['MESOS_SLAVE_PID'] = str(self.slave.pid) executor = Executor() driver = PesosExecutorDriver(executor, context=self.context) assert driver.start() == mesos_pb2.DRIVER_RUNNING assert driver.abort() == mesos_pb2.DRIVER_ABORTED driver.executor_process.aborted.wait(timeout=MAX_TIMEOUT) assert driver.executor_process.aborted.is_set() assert driver.join() == mesos_pb2.DRIVER_ABORTED
def test_mesos_executor_driver_init(): context = Context() context.start() # spawn slave slave = MockSlave() context.spawn(slave) os.environ['MESOS_SLAVE_PID'] = str(slave.pid) executor = Executor() driver = PesosExecutorDriver(executor, context=context) assert driver.start() == mesos.DRIVER_RUNNING assert driver.stop() == mesos.DRIVER_STOPPED assert driver.join() == mesos.DRIVER_STOPPED context.stop()
def test_mesos_executor_register(self): os.environ['MESOS_SLAVE_PID'] = str(self.slave.pid) registered_event, registered_call, registered_side_effect = self.mock_method() executor = mock.create_autospec(Executor, spec_set=True) executor.registered = mock.Mock(side_effect=registered_side_effect) driver = PesosExecutorDriver(executor, context=self.context) assert driver.start() == mesos_pb2.DRIVER_RUNNING driver.executor_process.connected.wait(timeout=MAX_TIMEOUT) assert driver.executor_process.connected.is_set() registered_event.wait(timeout=MAX_TIMEOUT) assert registered_event.is_set() assert registered_call.mock_calls == [ mock.call(driver, self.executor_info, self.framework_info, self.slave.slave_info)] assert driver.stop() == mesos_pb2.DRIVER_STOPPED
def test_mesos_executor_shutdown(self): os.environ['MESOS_SLAVE_PID'] = str(self.slave.pid) shutdown_event, shutdown_call, shutdown_side_effect = self.mock_method() executor = mock.create_autospec(Executor, spec_set=True) executor.shutdown = mock.Mock(side_effect=shutdown_side_effect) driver = PesosExecutorDriver(executor, context=self.context) assert driver.start() == mesos_pb2.DRIVER_RUNNING # wait until registered driver.executor_process.connected.wait(timeout=MAX_TIMEOUT) assert driver.executor_process.connected.is_set() self.slave.send_shutdown(driver.executor_process.pid) shutdown_event.wait(timeout=MAX_TIMEOUT) assert shutdown_event.is_set() assert shutdown_call.mock_calls == [mock.call(driver)]
def test_mesos_executor_shutdown(self): os.environ['MESOS_SLAVE_PID'] = str(self.slave.pid) shutdown_event, shutdown_call, shutdown_side_effect = self.mock_method( ) executor = mock.create_autospec(Executor, spec_set=True) executor.shutdown = mock.Mock(side_effect=shutdown_side_effect) driver = PesosExecutorDriver(executor, context=self.context) assert driver.start() == mesos_pb2.DRIVER_RUNNING # wait until registered driver.executor_process.connected.wait(timeout=MAX_TIMEOUT) assert driver.executor_process.connected.is_set() self.slave.send_shutdown(driver.executor_process.pid) shutdown_event.wait(timeout=MAX_TIMEOUT) assert shutdown_event.is_set() assert shutdown_call.mock_calls == [mock.call(driver)]
def test_mesos_executor_register(): context = Context() context.start() # spawn slave slave = MockSlave() context.spawn(slave) os.environ['MESOS_SLAVE_PID'] = str(slave.pid) executor = mock.MagicMock() executor.registered = mock.MagicMock() driver = PesosExecutorDriver(executor, context=context) assert driver.start() == mesos.DRIVER_RUNNING command_info = mesos.CommandInfo(value='wat') framework_id = mesos.FrameworkID(value='fake_framework_id') executor_id = mesos.ExecutorID(value='fake_executor_id') executor_info = mesos.ExecutorInfo( executor_id=executor_id, framework_id=framework_id, command=command_info ) framework_info = mesos.FrameworkInfo(user='******', name='fake_framework_name') slave.send_registered( driver.executor_process.pid, executor_info, framework_id, framework_info ) driver.executor_process.connected.wait(timeout=1) assert driver.executor_process.connected.is_set() executor.registered.assert_called_with(driver, executor_info, framework_info, slave.slave_info) assert driver.stop() == mesos.DRIVER_STOPPED context.stop()
def test_mesos_executor_reregister(self): os.environ['MESOS_SLAVE_PID'] = str(self.slave.pid) os.environ['MESOS_CHECKPOINT'] = '1' os.environ['MESOS_RECOVERY_TIMEOUT'] = '15mins' reregistered_event, reregistered_call, reregistered_side_effect = self.mock_method( ) executor = mock.create_autospec(Executor, spec_set=True) executor.reregistered = mock.Mock(side_effect=reregistered_side_effect) driver_context = Context() driver_context.start() self.close_contexts.append(driver_context) driver = PesosExecutorDriver(executor, context=driver_context) assert driver.start() == mesos_pb2.DRIVER_RUNNING self.slave.register_event.wait(timeout=MAX_TIMEOUT) assert self.slave.register_event.is_set() driver.executor_process.connected.wait(timeout=MAX_TIMEOUT) assert driver.executor_process.connected.is_set() # kill connection, but expect reconnect conn = driver_context._connections[self.slave.pid] conn.close() # simulate slave reconnecting self.slave.send_reconnect(driver.executor_process.pid) # wait for reregister event self.slave.reregister_event.wait(timeout=MAX_TIMEOUT) assert self.slave.reregister_event.is_set() driver.executor_process.connected.wait(timeout=MAX_TIMEOUT) assert driver.executor_process.connected.is_set() # reregistered event reregistered_event.wait(timeout=MAX_TIMEOUT) assert reregistered_event.is_set() assert reregistered_call.mock_calls == [ mock.call(driver, self.slave.slave_info) ]
def test_mesos_executor_register(self): os.environ['MESOS_SLAVE_PID'] = str(self.slave.pid) registered_event, registered_call, registered_side_effect = self.mock_method( ) executor = mock.create_autospec(Executor, spec_set=True) executor.registered = mock.Mock(side_effect=registered_side_effect) driver = PesosExecutorDriver(executor, context=self.context) assert driver.start() == mesos_pb2.DRIVER_RUNNING driver.executor_process.connected.wait(timeout=MAX_TIMEOUT) assert driver.executor_process.connected.is_set() registered_event.wait(timeout=MAX_TIMEOUT) assert registered_event.is_set() assert registered_call.mock_calls == [ mock.call(driver, self.executor_info, self.framework_info, self.slave.slave_info) ] assert driver.stop() == mesos_pb2.DRIVER_STOPPED
def test_mesos_executor_reregister(self): os.environ['MESOS_SLAVE_PID'] = str(self.slave.pid) os.environ['MESOS_CHECKPOINT'] = '1' os.environ['MESOS_RECOVERY_TIMEOUT'] = '15mins' reregistered_event, reregistered_call, reregistered_side_effect = self.mock_method() executor = mock.create_autospec(Executor, spec_set=True) executor.reregistered = mock.Mock(side_effect=reregistered_side_effect) driver_context = Context() driver_context.start() self.close_contexts.append(driver_context) driver = PesosExecutorDriver(executor, context=driver_context) assert driver.start() == mesos_pb2.DRIVER_RUNNING self.slave.register_event.wait(timeout=MAX_TIMEOUT) assert self.slave.register_event.is_set() driver.executor_process.connected.wait(timeout=MAX_TIMEOUT) assert driver.executor_process.connected.is_set() # kill connection, but expect reconnect conn = driver_context._connections[self.slave.pid] conn.close() # simulate slave reconnecting self.slave.send_reconnect(driver.executor_process.pid) # wait for reregister event self.slave.reregister_event.wait(timeout=MAX_TIMEOUT) assert self.slave.reregister_event.is_set() driver.executor_process.connected.wait(timeout=MAX_TIMEOUT) assert driver.executor_process.connected.is_set() # reregistered event reregistered_event.wait(timeout=MAX_TIMEOUT) assert reregistered_event.is_set() assert reregistered_call.mock_calls == [mock.call(driver, self.slave.slave_info)]