Timeout execution tree creation for SDK worker ops.

sdks/python/apache_beam/runners/worker/bundle_processor.py

@@ -24,6 +24,7 @@
 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 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 @@ def __init__(
         '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()

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):

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: