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_kill_task_terminate(self): task_id = get_random_task_id() stdout_name = ensure_directory('build/stdout.' + str(task_id)) stderr_name = ensure_directory('build/stderr.' + str(task_id)) stdout = open(stdout_name, 'w+') stderr = open(stderr_name, 'w+') try: command = 'sleep 100' process = subprocess.Popen(command, shell=True, stdout=stdout, stderr=stderr) process_info = process, stdout, stderr shutdown_grace_period_ms = 2000 ce.kill_task(process_info, shutdown_grace_period_ms) # await process termination while process.poll() is None: time.sleep(0.01) self.assertEqual(-1 * signal.SIGTERM, process.poll()) self.assertTrue(stdout.closed) self.assertTrue(stderr.closed) finally: cleanup_output(stdout_name, stderr_name)
def test_await_process_completion_normal(self): task_id = get_random_task_id() stdout_name = ensure_directory('build/stdout.' + str(task_id)) stderr_name = ensure_directory('build/stderr.' + str(task_id)) stdout = open(stdout_name, 'w+') stderr = open(stderr_name, 'w+') try: command = 'sleep 2' process = subprocess.Popen(command, shell=True, stdout=stdout, stderr=stderr) process_info = process, stdout, stderr shutdown_grace_period_ms = 1000 stop_signal = Event() ce.await_process_completion(stop_signal, process_info, shutdown_grace_period_ms) self.assertFalse(stop_signal.isSet()) self.assertEqual(0, process.returncode) self.assertTrue(stdout.closed) self.assertTrue(stderr.closed) finally: cleanup_output(stdout_name, stderr_name)
def test_await_process_completion_killed(self): task_id = get_random_task_id() stdout_name = 'build/stdout.' + str(task_id) stderr_name = 'build/stderr.' + str(task_id) stdout = open(stdout_name, 'w+') stderr = open(stderr_name, 'w+') try: command = 'sleep 100' process = subprocess.Popen(command, shell=True, stdout=stdout, stderr=stderr) process_info = process, stdout, stderr shutdown_grace_period_ms = 2000 stop_signal = Event() def sleep_and_set_stop_signal(): time.sleep(2 * cook.RUNNING_POLL_INTERVAL_SECS) stop_signal.set() thread = Thread(target=sleep_and_set_stop_signal, args=()) thread.start() ce.await_process_completion(stop_signal, process_info, shutdown_grace_period_ms) self.assertTrue(process.returncode < 0) self.assertTrue(stdout.closed) self.assertTrue(stderr.closed) finally: cleanup_output(stdout_name, stderr_name)
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_launch_task_handle_exception(self): task_id = tu.get_random_task_id() task = {'task_id': {'value': task_id}} process = ce.launch_task(task, os.environ) self.assertIsNone(process)
def test_send_progress_does_not_trim_unknown_field(self): driver = tu.FakeMesosExecutorDriver() task_id = tu.get_random_task_id() max_message_length = 30 poll_interval_ms = 10 send_progress_message = self.send_progress_message_helper( driver, max_message_length) progress_updater = cp.ProgressUpdater(task_id, max_message_length, poll_interval_ms, send_progress_message) progress_data_0 = { 'progress-message': b' pm', 'progress-sequence': 1, 'unknown': 'Unknown field has a really long lorem ipsum dolor sit amet exceed limit text' } progress_updater.send_progress_update(progress_data_0) self.assertEqual(1, len(driver.messages)) actual_encoded_message_0 = driver.messages[0] expected_message_0 = { 'progress-message': 'pm', 'progress-sequence': 1, 'task-id': task_id, 'unknown': 'Unknown field has a really long lorem ipsum dolor sit amet exceed limit text' } tu.assert_message(self, expected_message_0, actual_encoded_message_0)
def test_await_process_completion_killed(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)) tu.redirect_stdout_to_file(stdout_name) tu.redirect_stderr_to_file(stderr_name) try: command = 'sleep 100' process = subprocess.Popen(command, preexec_fn=os.setpgrp, shell=True) shutdown_grace_period_ms = 2000 stop_signal = Event() sleep_and_set_stop_signal_task(stop_signal, 2) ce.await_process_completion(process, stop_signal, shutdown_grace_period_ms) self.assertTrue(process.returncode < 0) finally: tu.cleanup_output(stdout_name, stderr_name)
def test_send_progress_update_trims_progress_message(self): driver = tu.FakeMesosExecutorDriver() task_id = tu.get_random_task_id() max_message_length = 30 poll_interval_ms = 10 send_progress_message = self.send_progress_message_helper( driver, max_message_length) progress_updater = cp.ProgressUpdater(task_id, max_message_length, poll_interval_ms, send_progress_message) progress_data_0 = { 'progress-message': b' Progress message-0 is really long lorem ipsum dolor sit amet text', 'progress-sequence': 1 } progress_updater.send_progress_update(progress_data_0) self.assertEqual(1, len(driver.messages)) actual_encoded_message_0 = driver.messages[0] expected_message_0 = { 'progress-message': 'Progress message-0 is reall...', 'progress-sequence': 1, 'task-id': task_id } tu.assert_message(self, expected_message_0, actual_encoded_message_0)
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_kill_task_terminate_with_sigkill(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)) tu.redirect_stdout_to_file(stdout_name) tu.redirect_stderr_to_file(stderr_name) try: command = "trap '' TERM SIGTERM; sleep 200" process = cs.launch_process(command, {}) shutdown_grace_period_ms = 1000 group_id = cs.find_process_group(process.pid) self.assertGreater(len(find_process_ids_in_group(group_id)), 0) cs.kill_process(process, shutdown_grace_period_ms) # await process termination for i in range(1, 10 * shutdown_grace_period_ms): if process.poll() is None: time.sleep(0.01) if process.poll() is None: process.kill() self.assertTrue( ((-1 * signal.SIGKILL) == process.poll()) or ((128 + signal.SIGKILL) == process.poll()), 'Process exited with code {}'.format(process.poll())) self.assertEqual(len(find_process_ids_in_group(group_id)), 0) finally: tu.cleanup_output(stdout_name, stderr_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 test_send_message_max_length_exceeded(self): driver = object() task_id = get_random_task_id() message = json.dumps({'task-id': task_id, 'message': 'test-message'}) max_message_length = 1 result = ce.send_message(driver, message, max_message_length) self.assertFalse(result)
def test_launch_task_handle_exception(self): task_id = get_random_task_id() task = {'task_id': {'value': task_id}} stdout_name = '' stderr_name = '' result = ce.launch_task(task, stdout_name, stderr_name) self.assertIsNone(result)
def test_manage_task_successful_exit_despite_faulty_driver(self): def assertions(driver, task_id, sandbox_directory): 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_FINISHED }] tu.assert_statuses(self, expected_statuses, driver.statuses) expected_core_messages = [{ 'sandbox-directory': sandbox_directory, 'task-id': task_id, 'type': 'directory' }, { 'exit-code': 0, 'task-id': task_id }] expected_progress_messages = [ { 'progress-message': 'ninety percent', 'progress-percent': 90, 'progress-sequence': 1, 'task-id': task_id }, # retried because the previous send fails { 'progress-message': 'ninety percent', 'progress-percent': 90, 'progress-sequence': 1, 'task-id': task_id } ] tu.assert_messages(self, expected_core_messages, expected_progress_messages, driver.messages) test_file_name = tu.ensure_directory('build/file.' + tu.get_random_task_id()) command = ('echo "^^^^JOB-PROGRESS: 90 ninety percent"'.format( test_file_name)) socket_error = OSError( 'socket.error') # socket.error is an alias of OSError self.manage_task_runner( command, assertions, driver=tu.ErrorMesosExecutorDriver(socket_error))
def test_launch_task_no_command(self): task_id = get_random_task_id() task = {'task_id': {'value': task_id}, 'data': encode_data(json.dumps({'command': ''}).encode('utf8'))} stdout_name = '' stderr_name = '' result = ce.launch_task(task, stdout_name, stderr_name) self.assertIsNone(result)
def test_create_status_running(self): task_id = get_random_task_id() actual_status = ce.create_status(task_id, cook.TASK_RUNNING) expected_status = { 'task_id': { 'value': task_id }, 'state': cook.TASK_RUNNING } assert_status(self, expected_status, actual_status)
def test_update_status(self): driver = FakeMesosExecutorDriver() task_id = get_random_task_id() task_state = "TEST_TASK_STATE" ce.update_status(driver, task_id, task_state) self.assertEqual(1, len(driver.statuses)) actual_status = driver.statuses[0] expected_status = {'task_id': {'value': task_id}, 'state': task_state} assert_status(self, expected_status, actual_status)
def test_collect_progress_updates_dev_null(self): file_name = tu.ensure_directory('build/collect_progress_test.' + tu.get_random_task_id()) progress_regex = '\^\^\^\^JOB-PROGRESS:\s+([0-9]*\.?[0-9]+)($|\s+.*)' location = '/dev/null' stop = Event() completed = Event() termination = Event() file = open(file_name, 'w+') file.flush() counter = cp.ProgressSequenceCounter() dn_watcher = cp.ProgressWatcher(location, 'dn', counter, 1024, progress_regex, stop, completed, termination) out_watcher = cp.ProgressWatcher(file_name, 'so', counter, 1024, progress_regex, stop, completed, termination) try: def print_to_file(): file.write('Stage One complete\n') file.write('^^^^JOB-PROGRESS: 100 100-percent\n') file.flush() file.close() completed.set() print_thread = Thread(target=print_to_file, args=()) print_thread.start() progress_states = [{ 'progress-message': b' 100-percent', 'progress-percent': 100, 'progress-sequence': 1 }] for actual_progress_state in out_watcher.retrieve_progress_states( ): expected_progress_state = progress_states.pop(0) self.assertEqual(expected_progress_state, actual_progress_state) self.assertEqual(expected_progress_state, out_watcher.current_progress()) self.assertFalse(progress_states) iterable = dn_watcher.retrieve_progress_states() exhausted = object() self.assertEqual(exhausted, next(iterable, exhausted)) self.assertIsNone(dn_watcher.current_progress()) print_thread.join() finally: completed.set() tu.cleanup_file(file_name)
def test_send_message(self): driver = tu.FakeMesosExecutorDriver() task_id = tu.get_random_task_id() expected_message = {'task-id': task_id, 'message': 'test-message'} result = ce.send_message(driver, tu.fake_os_error_handler, expected_message) self.assertTrue(result) self.assertEqual(1, len(driver.messages)) actual_encoded_message = driver.messages[0] tu.assert_message(self, expected_message, actual_encoded_message)
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_create_status_running(self): driver = tu.FakeMesosExecutorDriver() task_id = tu.get_random_task_id() status_updater = ce.StatusUpdater(driver, task_id) actual_status = status_updater.create_status(cook.TASK_RUNNING) expected_status = { 'task_id': { 'value': task_id }, 'state': cook.TASK_RUNNING } tu.assert_status(self, expected_status, actual_status)
def test_send_message(self): driver = FakeMesosExecutorDriver() task_id = get_random_task_id() expected_message = {'task-id': task_id, 'message': 'test-message'} message = json.dumps(expected_message) max_message_length = 512 result = ce.send_message(driver, message, max_message_length) self.assertTrue(result) self.assertEqual(1, len(driver.messages)) actual_encoded_message = driver.messages[0] assert_message(self, expected_message, actual_encoded_message)
def test_launch_task_no_command(self): task_id = tu.get_random_task_id() task = { 'task_id': { 'value': task_id }, 'data': pm.encode_data(json.dumps({ 'command': '' }).encode('utf8')) } process = ce.launch_task(task, os.environ) self.assertIsNone(process)
def test_send_message_handles_generic_exception(self): exception_handler_calls = [] exception_handler = functools.partial(tu.store_exception_handler, exception_handler_calls) exception = Exception('Generic Exception') driver = tu.ErrorMesosExecutorDriver(exception) task_id = tu.get_random_task_id() message = {'task-id': task_id, 'message': 'test-message'} result = ce.send_message(driver, exception_handler, message) self.assertEqual([], exception_handler_calls) self.assertFalse(result) self.assertEqual(1, len(driver.messages))
def test_send_message_handles_os_error_non_memory(self): exception_handler_calls = [] exception_handler = functools.partial(tu.store_exception_handler, exception_handler_calls) exception = OSError(errno.EACCES, 'Permission denied') driver = tu.ErrorMesosExecutorDriver(exception) task_id = tu.get_random_task_id() message = {'task-id': task_id, 'message': 'test-message'} result = ce.send_message(driver, exception_handler, message) self.assertEqual([], exception_handler_calls) self.assertFalse(result) self.assertEqual(1, len(driver.messages))
def test_collect_progress_updates_faulty_regex(self): file_name = tu.ensure_directory( 'build/collect_progress_updates_skip_faulty.' + tu.get_random_task_id()) progress_regex = '\^\^\^\^JOB-PROGRESS: (\S+)(?: )?(.*)' stop = Event() completed = Event() termination = Event() file = open(file_name, 'w+') file.flush() counter = cp.ProgressSequenceCounter() watcher = cp.ProgressWatcher(file_name, 'test', counter, 1024, progress_regex, stop, completed, termination) try: def print_to_file(): file.write('^^^^JOB-PROGRESS: ABCDEF string percent\n') file.write('^^^^JOB-PROGRESS: F50 Fifty percent\n') file.write( '^^^^JOB-PROGRESS: 1019101010101010101010101018101101010101010110171010110 Sixty percent\n' ) file.write('^^^^JOB-PROGRESS: 75 75% percent\n') file.flush() file.close() completed.set() print_thread = Thread(target=print_to_file, args=()) print_thread.start() progress_states = [{ 'progress-message': b'75% percent', 'progress-percent': 75, 'progress-sequence': 1 }] for actual_progress_state in watcher.retrieve_progress_states(): expected_progress_state = progress_states.pop(0) self.assertEqual(expected_progress_state, actual_progress_state) self.assertEqual(expected_progress_state, watcher.current_progress()) self.assertFalse(progress_states) print_thread.join() finally: completed.set() tu.cleanup_file(file_name)
def test_send_progress_update(self): driver = tu.FakeMesosExecutorDriver() task_id = tu.get_random_task_id() max_message_length = 30 poll_interval_ms = 100 send_progress_message = self.send_progress_message_helper( driver, max_message_length) progress_updater = cp.ProgressUpdater(task_id, max_message_length, poll_interval_ms, send_progress_message) progress_data_0 = { 'progress-message': b' Progress message-0', 'progress-sequence': 1 } progress_updater.send_progress_update(progress_data_0) self.assertEqual(1, len(driver.messages)) actual_encoded_message_0 = driver.messages[0] expected_message_0 = { 'progress-message': 'Progress message-0', 'progress-sequence': 1, 'task-id': task_id } tu.assert_message(self, expected_message_0, actual_encoded_message_0) progress_data_1 = { 'progress-message': b' Progress message-1', 'progress-sequence': 2 } progress_updater.send_progress_update(progress_data_1) self.assertEqual(1, len(driver.messages)) time.sleep(poll_interval_ms / 1000.0) progress_data_2 = { 'progress-message': b' Progress message-2', 'progress-sequence': 3 } progress_updater.send_progress_update(progress_data_2) self.assertEqual(2, len(driver.messages)) actual_encoded_message_2 = driver.messages[1] expected_message_2 = { 'progress-message': 'Progress message-2', 'progress-sequence': 3, 'task-id': task_id } tu.assert_message(self, expected_message_2, actual_encoded_message_2)
def test_launch_task_interactive_output(self): task_id = tu.get_random_task_id() command = 'echo "Start"; echo "Hello"; sleep 100; echo "World"; echo "Done"; ' 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) try: process = ce.launch_task(task, os.environ) self.assertIsNotNone(process) # let the process run for up to 50 seconds for _ in range(5000): if cs.is_process_running(process): time.sleep(0.01) with open(stdout_name) as f: stdout_content = f.read() if 'Start' in stdout_content and 'Hello' in stdout_content: break try: with open(stdout_name) as f: stdout_content = f.read() logging.info( 'Contents of stdout: {}'.format(stdout_content)) self.assertTrue("Start" in stdout_content) self.assertTrue("Hello" in stdout_content) self.assertFalse("World" in stdout_content) self.assertFalse("Done" in stdout_content) finally: if process.poll() is None: logging.info('Killing launched process') process.kill() finally: tu.cleanup_output(stdout_name, stderr_name)
def test_cleanup_process(self): task_id = get_random_task_id() stdout_name = ensure_directory('build/stdout.' + str(task_id)) stderr_name = ensure_directory('build/stderr.' + str(task_id)) stdout = open(stdout_name, 'w+') stderr = open(stderr_name, 'w+') try: process_info = None, stdout, stderr ce.cleanup_process(process_info) self.assertTrue(stdout.closed) self.assertTrue(stderr.closed) finally: cleanup_output(stdout_name, stderr_name)