diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker.py b/sdks/python/apache_beam/runners/worker/sdk_worker.py index a79cb0e8de6e..db25a40a405c 100644 --- a/sdks/python/apache_beam/runners/worker/sdk_worker.py +++ b/sdks/python/apache_beam/runners/worker/sdk_worker.py @@ -80,7 +80,7 @@ MAX_KNOWN_NOT_RUNNING_INSTRUCTIONS = 1000 # The number of ProcessBundleRequest instruction ids that BundleProcessorCache # will remember for failed instructions. -MAX_FAILED_INSTRUCTIONS = 10000 +MAX_FAILED_INSTRUCTIONS = 1000 # retry on transient UNAVAILABLE grpc error from state channels. _GRPC_SERVICE_CONFIG = json.dumps({ @@ -559,7 +559,15 @@ def discard(self, instruction_id, exception): """ processor = None with self._lock: - self.failed_instruction_ids[instruction_id] = exception + tb_str = "".join(traceback.format_exception(exception)) + if len(tb_str) > 10240: + tb_str = ( + tb_str[:5000] + "\n... [traceback truncated] ...\n" + + tb_str[-5000:]) + clean_exception = RuntimeError( + f"Original Exception: {type(exception).__name__}: {str(exception)[:2000]}\n{tb_str}" + ) + self.failed_instruction_ids[instruction_id] = clean_exception while len(self.failed_instruction_ids) > MAX_FAILED_INSTRUCTIONS: self.failed_instruction_ids.popitem(last=False) if instruction_id in self.active_bundle_processors: diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker_test.py b/sdks/python/apache_beam/runners/worker/sdk_worker_test.py index 76e428f06464..bea313a4d2fd 100644 --- a/sdks/python/apache_beam/runners/worker/sdk_worker_test.py +++ b/sdks/python/apache_beam/runners/worker/sdk_worker_test.py @@ -296,6 +296,44 @@ def test_failed_bundle_processor_returns_failed_split_response(self): worker.do_instruction(split_request).error, hc.contains_string('test message')) + def test_failed_instruction_id_cache_size_is_capped(self): + data_channel_factory = mock.create_autospec( + data_plane.GrpcClientDataChannelFactory) + bundle_processor_cache = BundleProcessorCache( + None, None, data_channel_factory, {}) + if bundle_processor_cache.periodic_shutdown: + bundle_processor_cache.periodic_shutdown.cancel() + + with mock.patch( + 'apache_beam.runners.worker.sdk_worker.MAX_FAILED_INSTRUCTIONS', 10): + for i in range(15): + bundle_processor_cache.discard(f'inst_{i}', RuntimeError(f'error {i}')) + + for i in range(5): + self.assertNotIn( + f'inst_{i}', bundle_processor_cache.failed_instruction_ids) + for i in range(5, 15): + self.assertIn( + f'inst_{i}', bundle_processor_cache.failed_instruction_ids) + + def test_failed_instruction_tracebacks_are_truncated_when_too_long(self): + data_channel_factory = mock.create_autospec( + data_plane.GrpcClientDataChannelFactory) + bundle_processor_cache = BundleProcessorCache( + None, None, data_channel_factory, {}) + if bundle_processor_cache.periodic_shutdown: + bundle_processor_cache.periodic_shutdown.cancel() + + long_message = "x" * 15000 + bundle_processor_cache.discard('instruction_id', RuntimeError(long_message)) + + stored_exception = bundle_processor_cache.failed_instruction_ids[ + 'instruction_id'] + tb_str = str(stored_exception) + + self.assertLessEqual(len(tb_str), 13000) + self.assertIn('[traceback truncated]', tb_str) + def test_data_sampling_response(self): # Create a data sampler with some fake sampled data. This data will be seen # in the sample response.