(beam) branch master updated: Timeout execution tree creation for SDK worker ops. (#36200)

[tv...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

tvalentyn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/master by this push:
     new d91fb6d6987 Timeout execution tree creation for SDK worker ops. (#36200)
d91fb6d6987 is described below

commit d91fb6d6987310e8ce6e50e3f3af7afc628f038d
Author: tvalentyn <tv...@users.noreply.github.com>
AuthorDate: Thu Oct 16 12:10:55 2025 -0700

    Timeout execution tree creation for SDK worker ops. (#36200)
    
    * Timeout execution tree creation for SDK worker ops.
    
    * Run tests
    
    * Update the error message
    
    * Minor change to termination logic
    
    * Grammar
---
 .../apache_beam/runners/worker/bundle_processor.py | 27 +++++++++++++++++++++-
 .../apache_beam/runners/worker/sdk_worker_main.py  |  4 ++++
 .../apache_beam/runners/worker/worker_status.py    | 16 ++++++++++---
 3 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/sdks/python/apache_beam/runners/worker/bundle_processor.py b/sdks/python/apache_beam/runners/worker/bundle_processor.py
index ad48358d588..85f1e43d603 100644
--- a/sdks/python/apache_beam/runners/worker/bundle_processor.py
+++ b/sdks/python/apache_beam/runners/worker/bundle_processor.py
@@ -24,6 +24,7 @@ from __future__ import annotations
 import base64
 import bisect
 import collections
+import concurrent.futures
 import copy
 import heapq
 import itertools
@@ -76,6 +77,7 @@ from apache_beam.runners.worker import data_sampler
 from apache_beam.runners.worker import operation_specs
 from apache_beam.runners.worker import operations
 from apache_beam.runners.worker import statesampler
+from apache_beam.runners.worker.worker_status import thread_dump
 from apache_beam.transforms import TimeDomain
 from apache_beam.transforms import core
 from apache_beam.transforms import environments
@@ -1130,7 +1132,30 @@ class BundleProcessor(object):
         'fnapi-step-%s' % self.process_bundle_descriptor.id,
         self.counter_factory)
 
-    self.ops = self.create_execution_tree(self.process_bundle_descriptor)
+    with concurrent.futures.ThreadPoolExecutor(
+        max_workers=1, thread_name_prefix='ExecutionTreeCreator') as executor:
+      future = executor.submit(
+          self.create_execution_tree, self.process_bundle_descriptor)
+      try:
+        self.ops = future.result(timeout=3600)
+      except concurrent.futures.TimeoutError:
+        # In rare cases, unpickling a DoFn might get permanently stuck,
+        # for example when unpickling involves importing a module and
+        # a subprocess is launched during the import operation.
+        _LOGGER.error(
+            'Timed out while reconstructing a pipeline fragment for: %s.\n'
+            'This is likely a transient error. The SDK harness '
+            'will self-terminate, and the runner can retry the operation. '
+            'If the error is frequent, check whether the stuckness happens '
+            'while deserializing (unpickling) a dependency of your pipeline '
+            'in the stacktrace below: \n%s\n',
+            self.process_bundle_descriptor.id,
+            thread_dump('ExecutionTreeCreator'))
+        # Raising an exception here doesn't interrupt the left-over thread.
+        # Out of caution, terminate the SDK harness process.
+        from apache_beam.runners.worker.sdk_worker_main import terminate_sdk_harness
+        terminate_sdk_harness()
+
     for op in reversed(self.ops.values()):
       op.setup(self.data_sampler)
     self.splitting_lock = threading.Lock()
diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker_main.py b/sdks/python/apache_beam/runners/worker/sdk_worker_main.py
index 7ea0e0eb109..cdb807e8dbc 100644
--- a/sdks/python/apache_beam/runners/worker/sdk_worker_main.py
+++ b/sdks/python/apache_beam/runners/worker/sdk_worker_main.py
@@ -233,6 +233,10 @@ def terminate_sdk_harness():
   if _FN_LOG_HANDLER:
     _FN_LOG_HANDLER.close()
   os.kill(os.getpid(), signal.SIGINT)
+  # Delay further control flow in the caller until process is terminated.
+  time.sleep(60)
+  # Try to force-terminate if still running.
+  os.kill(os.getpid(), signal.SIGKILL)
 
 
 def _load_pipeline_options(options_json):
diff --git a/sdks/python/apache_beam/runners/worker/worker_status.py b/sdks/python/apache_beam/runners/worker/worker_status.py
index f4102b19389..1d54a3ee176 100644
--- a/sdks/python/apache_beam/runners/worker/worker_status.py
+++ b/sdks/python/apache_beam/runners/worker/worker_status.py
@@ -66,13 +66,23 @@ def _current_frames():
     return sys._current_frames()  # pylint: disable=protected-access
 
 
-def thread_dump():
-  """Get a thread dump for the current SDK worker harness. """
+def thread_dump(thread_prefix=None):
+  """Get a thread dump for the current SDK harness.
+
+  Args:
+    thread_prefix: (str) An optional prefix to filter threads by.
+  """
   # deduplicate threads with same stack trace
   stack_traces = defaultdict(list)
   frames = _current_frames()
 
-  for t in threading.enumerate():
+  threads_to_dump = threading.enumerate()
+  if thread_prefix:
+    threads_to_dump = [
+        t for t in threads_to_dump if t.name.startswith(thread_prefix)
+    ]
+
+  for t in threads_to_dump:
     try:
       stack_trace = ''.join(traceback.format_stack(frames[t.ident]))
     except KeyError: