Skip to content

Fix a memory leak in the failed instruction id cache. #39405

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tvalentyn merged 2 commits into from
Jul 21, 2026
+48 −2
Merged

Fix a memory leak in the failed instruction id cache. #39405

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
90516bc
Sanitize the exception message stored for failed instructions to avoi…
tvalentyn Jul 21, 2026
6e61efd
Use Py3.10 syntax
tvalentyn Jul 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 10 additions & 2 deletions sdks/python/apache_beam/runners/worker/sdk_worker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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({
Expand Down Expand Up @@ -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))

@gemini-code-assist gemini-code-assist Bot Jul 21, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

In Python versions prior to 3.10, traceback.format_exception does not accept a single exception object as a positional argument and will raise a TypeError (missing value and tb arguments). Since Apache Beam supports Python 3.8 and 3.9, you should pass type(exception), exception, exception.__traceback__ to maintain backward compatibility.

      tb_str = "".join(
          traceback.format_exception(
              type(exception), exception, exception.__traceback__))

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:
Expand Down
38 changes: 38 additions & 0 deletions sdks/python/apache_beam/runners/worker/sdk_worker_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Loading