def test_retrieve_process_environment(self): self.assertEqual({'EXECUTOR_PROGRESS_OUTPUT_FILE': 'stdout'}, ce.retrieve_process_environment( cc.ExecutorConfig(), {})) self.assertEqual( { 'CUSTOM_PROGRESS_OUTPUT_FILE': 'stdout', 'FOO': 'BAR', 'MESOS_SANDBOX': '/path/to/sandbox', 'PROGRESS_OUTPUT_FILE': 'executor.progress' }, ce.retrieve_process_environment( cc.ExecutorConfig( progress_output_env_variable='CUSTOM_PROGRESS_OUTPUT_FILE' ), { 'FOO': 'BAR', 'MESOS_SANDBOX': '/path/to/sandbox', 'PROGRESS_OUTPUT_FILE': 'executor.progress' })) self.assertEqual( { 'CUSTOM_PROGRESS_OUTPUT_FILE': 'custom.progress', 'EXECUTOR_PROGRESS_OUTPUT_FILE_ENV': 'CUSTOM_PROGRESS_OUTPUT_FILE' }, ce.retrieve_process_environment( cc.ExecutorConfig( progress_output_env_variable='CUSTOM_PROGRESS_OUTPUT_FILE', progress_output_name='custom.progress'), { 'CUSTOM_PROGRESS_OUTPUT_FILE': 'executor.progress', 'EXECUTOR_PROGRESS_OUTPUT_FILE_ENV': 'CUSTOM_PROGRESS_OUTPUT_FILE' })) self.assertEqual( { 'CUSTOM_PROGRESS_OUTPUT_FILE': 'custom.progress', 'EXECUTOR_PROGRESS_OUTPUT_FILE_ENV': 'CUSTOM_PROGRESS_OUTPUT_FILE', 'PROGRESS_OUTPUT_FILE': 'stdout' }, ce.retrieve_process_environment( cc.ExecutorConfig( progress_output_env_variable='CUSTOM_PROGRESS_OUTPUT_FILE', progress_output_name='custom.progress'), { 'CUSTOM_PROGRESS_OUTPUT_FILE': 'executor.progress', 'EXECUTOR_PROGRESS_OUTPUT_FILE_ENV': 'CUSTOM_PROGRESS_OUTPUT_FILE', 'PROGRESS_OUTPUT_FILE': 'stdout' }))
def test_executor_config(self): max_bytes_read_per_line = 16 * 1024 max_message_length = 300 progress_output_name = 'stdout_name' progress_regex_string = 'some-regex-string' progress_sample_interval_ms = 100 sandbox_directory = '/location/to/task/sandbox/task_id' shutdown_grace_period_secs = '5secs' config = cc.ExecutorConfig( max_bytes_read_per_line=max_bytes_read_per_line, max_message_length=max_message_length, progress_output_name=progress_output_name, progress_regex_string=progress_regex_string, progress_sample_interval_ms=progress_sample_interval_ms, sandbox_directory=sandbox_directory, shutdown_grace_period=shutdown_grace_period_secs) self.assertEqual(max_bytes_read_per_line, config.max_bytes_read_per_line) self.assertEqual(max_message_length, config.max_message_length) self.assertEqual(progress_output_name, config.progress_output_name) self.assertEqual(progress_regex_string, config.progress_regex_string) self.assertEqual(progress_sample_interval_ms, config.progress_sample_interval_ms) self.assertEqual(sandbox_directory, config.sandbox_directory) self.assertEqual(5000, config.shutdown_grace_period_ms)
def manage_task_runner(self, command, assertions_fn, stop_signal=Event()): driver = FakeMesosExecutorDriver() task_id = get_random_task_id() task = {'task_id': {'value': task_id}, 'data': encode_data(json.dumps({'command': command}).encode('utf8'))} stdout_name = 'build/stdout.' + str(task_id) stderr_name = 'build/stderr.' + str(task_id) completed_signal = Event() max_message_length = 300 progress_sample_interval_ms = 100 sandbox_location = '/location/to/task/sandbox/{}'.format(task_id) progress_output_name = stdout_name progress_regex_string = '\^\^\^\^JOB-PROGRESS: (\d*)(?: )?(.*)' config = cc.ExecutorConfig(max_message_length=max_message_length, progress_output_name=progress_output_name, progress_regex_string=progress_regex_string, progress_sample_interval_ms=progress_sample_interval_ms, sandbox_location=sandbox_location) try: ce.manage_task(driver, task, stop_signal, completed_signal, config, stdout_name, stderr_name) self.assertTrue(completed_signal.isSet()) assertions_fn(driver, task_id, sandbox_location) finally: cleanup_output(stdout_name, stderr_name)
def test_collect_progress_updates_lots_of_writes(self): file_name = 'build/collect_progress_test.' + get_random_task_id() progress_regex_string = 'progress: (\d*), (.*)' config = cc.ExecutorConfig(progress_output_name=file_name, progress_regex_string=progress_regex_string) items_to_write = 250000 stop_signal = Event() completed_signal = Event() def write_to_file(): target_file = open(file_name, 'w+') unit_progress_granularity = int(items_to_write / 100) for item in range(items_to_write): remainder = (item + 1) % unit_progress_granularity if remainder == 0: progress_percent = math.ceil(item / unit_progress_granularity) target_file.write( 'progress: {0}, completed-{0}-percent\n'.format( progress_percent)) target_file.flush() target_file.write("{}\n".format(item)) target_file.flush() target_file.close() time.sleep(0.15) completed_signal.set() write_thread = Thread(target=write_to_file, args=()) write_thread.start() progress_watcher = cp.ProgressWatcher(config, stop_signal, completed_signal) try: collected_data = [] def read_progress_states(): for progress in progress_watcher.retrieve_progress_states(): logging.info('Received: {}'.format(progress)) collected_data.append(progress) read_progress_states_thread = Thread(target=read_progress_states, args=()) read_progress_states_thread.start() read_progress_states_thread.join() expected_data = list( map( lambda x: { 'progress-message': 'completed-{}-percent'.format(x), 'progress-percent': x }, range(1, 101))) self.assertEqual(expected_data, collected_data) finally: completed_signal.set() if os.path.isfile(file_name): os.remove(file_name)
def test_progress_watcher_tail_lot_of_writes(self): file_name = 'build/tail_progress_test.' + get_random_task_id() config = cc.ExecutorConfig(progress_output_name=file_name) items_to_write = 250000 stop_signal = Event() completed_signal = Event() tail_sleep_ms = 25 try: def write_to_file(): file = open(file_name, 'w+') for item in range(items_to_write): file.write("{}\n".format(item)) file.flush() file.close() time.sleep(0.15) completed_signal.set() Thread(target=write_to_file, args=()).start() progress_watcher = cp.ProgressWatcher(config, stop_signal, completed_signal) collected_data = [] for line in progress_watcher.tail(tail_sleep_ms): collected_data.append(line.strip()) logging.info('Items read: {}'.format(len(collected_data))) self.assertEqual(items_to_write, len(collected_data)) expected_data = list(map(lambda x: str(x), range(items_to_write))) self.assertEqual(expected_data, collected_data) finally: if os.path.isfile(file_name): os.remove(file_name)
def test_collect_progress_updates_with_empty_regex(self): file_name = ensure_directory('build/collect_progress_test.' + get_random_task_id()) progress_regex_string = '' config = cc.ExecutorConfig(progress_output_name=file_name, progress_regex_string=progress_regex_string) stop_signal = Event() completed_signal = Event() file = open(file_name, 'w+') file.flush() progress_watcher = cp.ProgressWatcher(config, stop_signal, completed_signal) try: def read_progress_states(): for _ in progress_watcher.retrieve_progress_states(): pass Thread(target=read_progress_states, args=()).start() file.write("Stage One complete\n") file.flush() file.write("^^^^JOB-PROGRESS: 25 Twenty-Fine percent\n") file.flush() file.write("Stage Two complete\n") file.flush() file.write("^^^^JOB-PROGRESS: 50 Fifty percent\n") file.flush() time.sleep(0.10) self.assertIsNone(progress_watcher.current_progress()) file.write("Stage Three complete\n") file.flush() time.sleep(0.10) self.assertIsNone(progress_watcher.current_progress()) file.write("^^^^JOB-PROGRESS: 55 Fifty-five percent\n") file.flush() time.sleep(0.10) self.assertIsNone(progress_watcher.current_progress()) file.write("Stage Four complete\n") file.flush() file.write("^^^^JOB-PROGRESS: 100 Hundred percent\n") file.flush() time.sleep(0.10) self.assertIsNone(progress_watcher.current_progress()) finally: completed_signal.set() file.close() if os.path.isfile(file_name): os.remove(file_name)
def manage_task_runner(self, command, assertions_fn, stop_signal=None, task_id=None, config=None, driver=None): if driver is None: driver = tu.FakeMesosExecutorDriver() if stop_signal is None: stop_signal = Event() if task_id is None: task_id = tu.get_random_task_id() task = { 'task_id': { 'value': task_id }, 'data': pm.encode_data(json.dumps({ 'command': command }).encode('utf8')) } stdout_name = tu.ensure_directory('build/stdout.{}'.format(task_id)) stderr_name = tu.ensure_directory('build/stderr.{}'.format(task_id)) tu.redirect_stdout_to_file(stdout_name) tu.redirect_stderr_to_file(stderr_name) completed_signal = Event() if config is None: sandbox_directory = '/location/to/task/sandbox/{}'.format(task_id) config = cc.ExecutorConfig( max_message_length=300, progress_output_name=stdout_name, progress_regex_string= '\^\^\^\^JOB-PROGRESS:\s+([0-9]*\.?[0-9]+)($|\s+.*)', progress_sample_interval_ms=100, sandbox_directory=sandbox_directory) else: sandbox_directory = config.sandbox_directory try: ce.manage_task(driver, task, stop_signal, completed_signal, config) self.assertTrue(completed_signal.isSet()) assertions_fn(driver, task_id, sandbox_directory) finally: tu.cleanup_output(stdout_name, stderr_name)
def test_executor_exit_env_variable(self, mock_exit): os.environ['EXECUTOR_TEST_EXIT'] = '5' try: config = cc.ExecutorConfig() stop_signal = Event() executor = ce.CookExecutor(stop_signal, config) driver = tu.FakeMesosExecutorDriver() executor_info = {'executor_id': {'value': 'test'}} framework_info = {'id': 'framework'} agent_info = {'id': {'value': 'agent'}} executor.registered(driver, executor_info, framework_info, agent_info) mock_exit.assert_called_with(5) finally: del os.environ['EXECUTOR_TEST_EXIT']
def test_progress_watcher_tail_with_read_limit(self): file_name = ensure_directory('build/tail_progress_test.' + get_random_task_id()) config = cc.ExecutorConfig(max_bytes_read_per_line=10, progress_output_name=file_name) stop_signal = Event() completed_signal = Event() tail_sleep_ms = 25 try: def write_to_file(): file = open(file_name, 'w+') file.write("abcd\n") file.flush() file.write("abcdefghijkl\n") file.flush() file.write("abcdefghijklmnopqrstuvwxyz\n") file.flush() file.close() time.sleep(0.15) completed_signal.set() Thread(target=write_to_file, args=()).start() progress_watcher = cp.ProgressWatcher(config, stop_signal, completed_signal) collected_data = [] for line in progress_watcher.tail(tail_sleep_ms): collected_data.append(line.strip()) logging.debug('collected_data = {}'.format(collected_data)) expected_data = [ 'abcd', 'abcdefghij', 'kl', 'abcdefghij', 'klmnopqrst', 'uvwxyz' ] self.assertEqual(expected_data, collected_data) finally: if os.path.isfile(file_name): os.remove(file_name)
def test_executor_config(self): checkpoint = 1 max_bytes_read_per_line = 16 * 1024 max_message_length = 300 memory_usage_interval_secs = 150 mesos_directory = '/mesos/directory' progress_output_env_variable = 'PROGRESS_OUTPUT_ENV_VARIABLE' progress_output_name = 'stdout_name' progress_regex_string = 'some-regex-string' progress_sample_interval_ms = 100 recovery_timeout = '5mins' reset_vars = ['a', 'b'] sandbox_directory = '/location/to/task/sandbox/task_id' shutdown_grace_period_secs = '5secs' config = cc.ExecutorConfig(checkpoint=checkpoint, max_bytes_read_per_line=max_bytes_read_per_line, max_message_length=max_message_length, memory_usage_interval_secs=memory_usage_interval_secs, mesos_directory=mesos_directory, progress_output_env_variable=progress_output_env_variable, progress_output_name=progress_output_name, progress_regex_string=progress_regex_string, progress_sample_interval_ms=progress_sample_interval_ms, recovery_timeout=recovery_timeout, reset_vars=reset_vars, sandbox_directory=sandbox_directory, shutdown_grace_period=shutdown_grace_period_secs) self.assertEqual(checkpoint, True) self.assertEqual(max_bytes_read_per_line, config.max_bytes_read_per_line) self.assertEqual(max_message_length, config.max_message_length) self.assertEqual(memory_usage_interval_secs, config.memory_usage_interval_secs) self.assertEqual(mesos_directory, config.mesos_directory) self.assertEqual(progress_output_env_variable, config.progress_output_env_variable) self.assertEqual(progress_output_name, config.progress_output_name) self.assertEqual(progress_regex_string, config.progress_regex_string) self.assertEqual(progress_sample_interval_ms, config.progress_sample_interval_ms) self.assertEqual(5 * 60 * 1000, config.recovery_timeout_ms) self.assertEqual(reset_vars, reset_vars) self.assertEqual(sandbox_directory, config.sandbox_directory) self.assertEqual(5000, config.shutdown_grace_period_ms) self.assertEqual(os.path.join(sandbox_directory, 'foo.bar'), config.sandbox_file('foo.bar')) self.assertEqual(os.path.join(sandbox_directory, 'stderr'), config.stderr_file()) self.assertEqual(os.path.join(sandbox_directory, 'stdout'), config.stdout_file())
def test_executor_launch_task_and_disconnect(self): task_id = tu.get_random_task_id() stdout_name = tu.ensure_directory('build/stdout.{}'.format(task_id)) stderr_name = tu.ensure_directory('build/stderr.{}'.format(task_id)) output_name = tu.ensure_directory('build/output.' + str(task_id)) tu.redirect_stdout_to_file(stdout_name) tu.redirect_stderr_to_file(stderr_name) try: config = cc.ExecutorConfig() stop_signal = Event() executor = ce.CookExecutor(stop_signal, config) driver = tu.FakeMesosExecutorDriver() command = 'echo "Start" >> {}; sleep 100; echo "Done." >> {}; '.format( output_name, output_name) task = { 'task_id': { 'value': task_id }, 'data': pm.encode_data( json.dumps({ 'command': command }).encode('utf8')) } executor.launchTask(driver, task) # let the process run for up to 10 seconds for _ in range(1000): time.sleep(0.01) if os.path.isfile(output_name): with open(output_name) as f: content = f.read() if 'Start' in content: break executor.disconnected(driver) self.assertTrue(executor.disconnect_signal.isSet()) self.assertTrue(executor.stop_signal.isSet()) executor.await_completion() logging.info('Task completed') if os.path.isfile(output_name): with open(output_name) as f: file_contents = f.read() self.assertTrue('Start' in file_contents) self.assertTrue('Done' not in file_contents) else: self.fail('{} does not exist.'.format(stderr_name)) expected_statuses = [{ 'task_id': { 'value': task_id }, 'state': cook.TASK_STARTING }, { 'task_id': { 'value': task_id }, 'state': cook.TASK_RUNNING }, { 'task_id': { 'value': task_id }, 'state': cook.TASK_KILLED }] tu.assert_statuses(self, expected_statuses, driver.statuses) expected_message_0 = { 'sandbox-directory': '', 'task-id': task_id, 'type': 'directory' } expected_message_1 = {'exit-code': -15, 'task-id': task_id} tu.assert_messages(self, [expected_message_0, expected_message_1], [], driver.messages) finally: tu.cleanup_output(stdout_name, stderr_name) tu.cleanup_file(output_name)