From 90516bcbd66c4ad5e92fc6ae82c773766d810322 Mon Sep 17 00:00:00 2001
From: Valentyn Tymofieiev <valentyn@google.com>
Date: Tue, 21 Jul 2026 05:37:28 -0700
Subject: [PATCH 1/2] Sanitize the exception message stored for failed
 instructions to avoid inadvertently capturing the content of stack frames in
 the heap and limit RAM growth due to cache size.

---
 .../apache_beam/runners/worker/sdk_worker.py  | 14 ++++++-
 .../runners/worker/sdk_worker_test.py         | 38 +++++++++++++++++++
 2 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker.py b/sdks/python/apache_beam/runners/worker/sdk_worker.py
index a79cb0e8de6e..9dcf31cad290 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,17 @@ def discard(self, instruction_id, exception):
     """
     processor = None
     with self._lock:
-      self.failed_instruction_ids[instruction_id] = exception
+      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:
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.

From 6e61efd883a792b79032e528bf68be9374ff85dd Mon Sep 17 00:00:00 2001
From: Valentyn Tymofieiev <valentyn@google.com>
Date: Tue, 21 Jul 2026 05:48:20 -0700
Subject: [PATCH 2/2] Use Py3.10 syntax

---
 sdks/python/apache_beam/runners/worker/sdk_worker.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker.py b/sdks/python/apache_beam/runners/worker/sdk_worker.py
index 9dcf31cad290..db25a40a405c 100644
--- a/sdks/python/apache_beam/runners/worker/sdk_worker.py
+++ b/sdks/python/apache_beam/runners/worker/sdk_worker.py
@@ -559,9 +559,7 @@ def discard(self, instruction_id, exception):
     """
     processor = None
     with self._lock:
-      tb_str = "".join(
-          traceback.format_exception(
-              type(exception), exception, exception.__traceback__))
+      tb_str = "".join(traceback.format_exception(exception))
       if len(tb_str) > 10240:
         tb_str = (
             tb_str[:5000] + "\n... [traceback truncated] ...\n" +