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)
