Adds a pipeline option to terminate processing of an element after a timeout (python)

sdks/python/apache_beam/options/pipeline_options.py

@@ -1460,9 +1460,9 @@ def _add_argparse_args(cls, parser):
         type=int,
         default=None,
         help=(
-            'The time limit (in minutes) that an SDK worker allows for a PTransform'
-            ' operation to process one element before signaling the runner harness'
-            ' to restart the SDK worker.'))
+            'The time limit (in minutes) that an SDK worker allows for a'
+            ' PTransform operation to process one element before signaling'
+            ' the runner harness to restart the SDK worker.'))
 
   def validate(self, validator):
     errors = []

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

@@ -24,7 +24,9 @@
 import logging
 import os
 import re
+from signal import signal
 import sys
+from time import time
 import traceback
 
 from google.protobuf import text_format
@@ -224,12 +226,13 @@ def main(unused_argv):
       fn_log_handler.close()
 
 
-def flush_fn_log_handler():
+def terminate_sdk_harness():
   """Flushes the FnApiLogRecordHandler if it exists."""
-  _LOGGER.error('The Sdk harness will be terminated now.')
+  _LOGGER.error('The SDK harness will be terminated in 5 seconds.')
+  time.sleep(5)
   if _FN_LOG_HANDLER:
     _FN_LOG_HANDLER.close()
-  os._exit(1)
+  os.kill(os.getpid(), signal.SIGINT)
 
 
 def _load_pipeline_options(options_json):

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

@@ -257,45 +257,50 @@ def _log_lull_in_bundle_processor(self, bundle_process_cache):
             self._log_lull_sampler_info(info, instruction)
 
   def _log_lull_sampler_info(self, sampler_info, instruction):
-    if (sampler_info and sampler_info.time_since_transition):
-      lull_seconds = sampler_info.time_since_transition / 1e9
-      step_name = sampler_info.state_name.step_name
-      state_name = sampler_info.state_name.name
-      if step_name and state_name:
-        step_name_log = (
-            ' for PTransform{name=%s, state=%s}' % (step_name, state_name))
-      else:
-        step_name_log = ''
-
-      stack_trace = self._get_stack_trace(sampler_info)
-      if (self._passed_lull_timeout_since_last_log() and
-          sampler_info.time_since_transition > self.log_lull_timeout_ns):
-        _LOGGER.warning(
-            (
-                'Operation ongoing in bundle %s%s for at least %.2f seconds'
-                ' without outputting or completing.\n'
-                'Current Traceback:\n%s'),
-            instruction,
-            step_name_log,
-            lull_seconds,
-            stack_trace,
-        )
-      if (self._element_processing_timeout_ns and
-          sampler_info.time_since_transition
-          > self._element_processing_timeout_ns):
-        _LOGGER.error(
-            (
-                'Operation ongoing in bundle %s%s for at least %.2f seconds'
-                ' without outputting or completing.\n'
-                'Current Traceback:\n%s'),
-            instruction,
-            step_name_log,
-            lull_seconds,
-            stack_trace,
-        )
-        from apache_beam.runners.worker.sdk_worker_main import flush_fn_log_handler
-
-        flush_fn_log_handler()
+    if (not sampler_info or not sampler_info.time_since_transition):
+      return
+
+    log_lull = self._passed_lull_timeout_since_last_log(
+    ) and sampler_info.time_since_transition > self.log_lull_timeout_ns
+    timeout_exceeded = self._element_processing_timeout_ns and sampler_info.time_since_transition > self._element_processing_timeout_ns
+    if (not log_lull or not timeout_exceeded):
+      return
+
+    lull_seconds = sampler_info.time_since_transition / 1e9
+    step_name = sampler_info.state_name.step_name
+    state_name = sampler_info.state_name.name
+    if step_name and state_name:
+      step_name_log = (
+          ' for PTransform{name=%s, state=%s}' % (step_name, state_name))
+    else:
+      step_name_log = ''
+    stack_trace = self._get_stack_trace(sampler_info)
+
+    if (log_lull):
+      _LOGGER.warning(
+          (
+              'Operation ongoing in bundle %s%s for at least %.2f seconds'
+              ' without outputting or completing.\n'
+              'Current Traceback:\n%s'),
+          instruction,
+          step_name_log,
+          lull_seconds,
+          stack_trace,
+      )
+    if (timeout_exceeded):
+      _LOGGER.error(
+          (
+              'Operation ongoing in bundle %s%s for at least %.2f seconds'
+              ' without outputting or completing.\n'
+              'Current Traceback:\n%s'),
+          instruction,
+          step_name_log,
+          self._element_processing_timeout_ns / 60 / 1e9,
+          stack_trace,
+      )
+      from apache_beam.runners.worker.sdk_worker_main import terminate_sdk_harness
+
+      terminate_sdk_harness()
 
   def _get_stack_trace(self, sampler_info):
     exec_thread = getattr(sampler_info, 'tracked_thread', None)

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

@@ -98,7 +98,7 @@ def get_state_sampler_info_for_lull(lull_duration_s):
     now = time.time()
     with mock.patch('logging.Logger.warning') as warn_mock:
       with mock.patch(
-          'apache_beam.runners.worker.sdk_worker_main.flush_fn_log_handler'
+          'apache_beam.runners.worker.sdk_worker_main.terminate_sdk_harness'
       ) as flush_mock:
         with mock.patch('time.time') as time_mock:
           time_mock.return_value = now