def test_inactive_bundle_processor_returns_empty_split_response(self):
bundle_processor = mock.MagicMock()
bundle_processor_cache = BundleProcessorCache(None, None, {})
bundle_processor_cache.activate('instruction_id')
worker = SdkWorker(bundle_processor_cache)
split_request = beam_fn_api_pb2.InstructionRequest(
instruction_id='split_instruction_id',
process_bundle_split=beam_fn_api_pb2.ProcessBundleSplitRequest(
instruction_id='instruction_id'))
self.assertEqual(
worker.do_instruction(split_request),
beam_fn_api_pb2.InstructionResponse(
instruction_id='split_instruction_id',
process_bundle_split=beam_fn_api_pb2.
ProcessBundleSplitResponse()))
# Add a mock bundle processor as if it was running before it's released
bundle_processor_cache.active_bundle_processors['instruction_id'] = (
'descriptor_id', bundle_processor)
bundle_processor_cache.release('instruction_id')
self.assertEqual(
worker.do_instruction(split_request),
beam_fn_api_pb2.InstructionResponse(
instruction_id='split_instruction_id',
process_bundle_split=beam_fn_api_pb2.
ProcessBundleSplitResponse()))
def test_log_lull_in_bundle_processor(self):
bundle_processor_cache = mock.MagicMock()
worker = SdkWorker(bundle_processor_cache)
now = time.time()
log_full_thread_dump_fn_name = \
'apache_beam.runners.worker.sdk_worker.SdkWorker._log_full_thread_dump'
with mock.patch('logging.Logger.warning') as warn_mock:
with mock.patch(
log_full_thread_dump_fn_name) as log_full_thread_dump:
with mock.patch('time.time') as time_mock:
time_mock.return_value = now
sampler_info = self._get_state_sampler_info_for_lull(21 *
60)
worker._log_lull_sampler_info(sampler_info)
processing_template = warn_mock.call_args[0][1]
step_name_template = warn_mock.call_args[0][2]
traceback = warn_mock.call_args = warn_mock.call_args[0][3]
self.assertIn('progress-msecs', processing_template)
self.assertIn('step_name', step_name_template)
self.assertIn('test_log_lull_in_bundle_processor',
traceback)
log_full_thread_dump.assert_called_once_with()
with mock.patch(log_full_thread_dump_fn_name) as log_full_thread_dump:
with mock.patch('time.time') as time_mock:
time_mock.return_value = now + 6 * 60 # 6 minutes
sampler_info = self._get_state_sampler_info_for_lull(21 * 60)
worker._log_lull_sampler_info(sampler_info)
self.assertFalse(
log_full_thread_dump.called,
'log_full_thread_dump should not be called because only 6 minutes '
'have passed since the last dump.')
with mock.patch(log_full_thread_dump_fn_name) as log_full_thread_dump:
with mock.patch('time.time') as time_mock:
time_mock.return_value = now + 21 * 60 # 21 minutes
sampler_info = self._get_state_sampler_info_for_lull(10 * 60)
worker._log_lull_sampler_info(sampler_info)
self.assertFalse(
log_full_thread_dump.called,
'log_full_thread_dump should not be called because lull is only '
'for 10 minutes.')
with mock.patch(log_full_thread_dump_fn_name) as log_full_thread_dump:
with mock.patch('time.time') as time_mock:
time_mock.return_value = now + 21 * 60 # 21 minutes
sampler_info = self._get_state_sampler_info_for_lull(21 * 60)
worker._log_lull_sampler_info(sampler_info)
log_full_thread_dump.assert_called_once_with()
def test_failed_bundle_processor_returns_failed_split_response(self):
bundle_processor = mock.MagicMock()
bundle_processor_cache = BundleProcessorCache(None, None, {})
bundle_processor_cache.activate('instruction_id')
worker = SdkWorker(bundle_processor_cache)
# Add a mock bundle processor as if it was running before it's discarded
bundle_processor_cache.active_bundle_processors['instruction_id'] = (
'descriptor_id', bundle_processor)
bundle_processor_cache.discard('instruction_id')
split_request = beam_fn_api_pb2.InstructionRequest(
instruction_id='split_instruction_id',
process_bundle_split=beam_fn_api_pb2.ProcessBundleSplitRequest(
instruction_id='instruction_id'))
hc.assert_that(
worker.do_instruction(split_request).error,
hc.contains_string(
'Bundle processing associated with instruction_id has failed'))
def test_log_lull_in_bundle_processor(self):
bundle_processor_cache = mock.MagicMock()
worker = SdkWorker(bundle_processor_cache)
sampler_info = statesampler.StateSamplerInfo(
CounterName('progress-msecs', 'stage_name', 'step_name'), 1,
400000000000, threading.current_thread())
now = time.time()
log_full_thread_dump_fn_name = \
'apache_beam.runners.worker.sdk_worker.SdkWorker._log_full_thread_dump'
with mock.patch('logging.Logger.warning') as warn_mock:
with mock.patch(
log_full_thread_dump_fn_name) as log_full_thread_dump:
with mock.patch('time.time') as time_mock:
time_mock.return_value = now
worker._log_lull_sampler_info(sampler_info)
processing_template = warn_mock.call_args[0][1]
step_name_template = warn_mock.call_args[0][2]
traceback = warn_mock.call_args = warn_mock.call_args[0][3]
self.assertIn('progress-msecs', processing_template)
self.assertIn('step_name', step_name_template)
self.assertIn('test_log_lull_in_bundle_processor',
traceback)
log_full_thread_dump.assert_called_once_with()
with mock.patch(log_full_thread_dump_fn_name) as log_full_thread_dump:
with mock.patch('time.time') as time_mock:
time_mock.return_value = now + 6 * 60 # 6 minutes
worker._log_lull_sampler_info(sampler_info)
self.assertFalse(log_full_thread_dump.called,
'log_full_thread_dump should not be called.')
with mock.patch(log_full_thread_dump_fn_name) as log_full_thread_dump:
with mock.patch('time.time') as time_mock:
time_mock.return_value = now + 21 * 60 # 21 minutes
worker._log_lull_sampler_info(sampler_info)
log_full_thread_dump.assert_called_once_with()