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

sdks/python/apache_beam/options/pipeline_options.py

@@ -1455,6 +1455,14 @@ def _add_argparse_args(cls, parser):
             'responsible for executing the user code and communicating with '
             'the runner. Depending on the runner, there may be more than one '
             'SDK Harness process running on the same worker node.'))
+    parser.add_argument(
+        '--element_processing_timeout_minutes',
+        type=int,
+        default=None,
+        help=(
+            'The time limit (minute) 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/options/pipeline_options_test.py

@@ -405,10 +405,18 @@ def test_experiments(self):
     self.assertEqual(options.get_all_options()['experiments'], None)
 
   def test_worker_options(self):
-    options = PipelineOptions(['--machine_type', 'abc', '--disk_type', 'def'])
+    options = PipelineOptions([
+        '--machine_type',
+        'abc',
+        '--disk_type',
+        'def',
+        '--element_processing_timeout_minutes',
+        '10',
+    ])
     worker_options = options.view_as(WorkerOptions)
     self.assertEqual(worker_options.machine_type, 'abc')
     self.assertEqual(worker_options.disk_type, 'def')
+    self.assertEqual(worker_options.element_processing_timeout_minutes, 10)
 
     options = PipelineOptions(
         ['--worker_machine_type', 'abc', '--worker_disk_type', 'def'])

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

@@ -176,6 +176,7 @@ def __init__(
       # that should be reported to the runner when proocessing the first bundle.
       deferred_exception=None,  # type: Optional[Exception]
       runner_capabilities=frozenset(),  # type: FrozenSet[str]
+      element_processing_timeout_minutes=None,  # type: Optional[int]
   ):
     # type: (...) -> None
     self._alive = True
@@ -207,6 +208,7 @@ def __init__(
     self._profiler_factory = profiler_factory
     self.data_sampler = data_sampler
     self.runner_capabilities = runner_capabilities
+    self.element_processing_timeout_minutes = element_processing_timeout_minutes
 
     def default_factory(id):
       # type: (str) -> beam_fn_api_pb2.ProcessBundleDescriptor
@@ -230,7 +232,12 @@ def default_factory(id):
             status_address,
             self._bundle_processor_cache,
             self._state_cache,
-            enable_heap_dump)  # type: Optional[FnApiWorkerStatusHandler]
+            enable_heap_dump,
+            element_processing_timeout_minutes=self.
+            element_processing_timeout_minutes
+        )  # type: Optional[FnApiWorkerStatusHandler]
+      except TimeoutError as e:
+        self._shutdown_due_to_element_processing_timeout(e)
       except Exception:
         traceback_string = traceback.format_exc()
         _LOGGER.warning(
@@ -409,6 +416,11 @@ def create_worker(self):
     return SdkWorker(
         self._bundle_processor_cache, profiler_factory=self._profiler_factory)
 
+  def _shutdown_due_to_element_processing_timeout(
+      self, err: TimeoutError) -> None:
+    _LOGGER.error('%sThe SDK harness will be terminated.', str(err))
+    sys.exit(1)
+
 
 class BundleProcessorCache(object):
   """A cache for ``BundleProcessor``s.

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

@@ -167,7 +167,9 @@ def create_harness(environment, dry_run=False):
       enable_heap_dump=enable_heap_dump,
       data_sampler=data_sampler,
       deferred_exception=deferred_exception,
-      runner_capabilities=runner_capabilities)
+      runner_capabilities=runner_capabilities,
+      element_processing_timeout_minutes=sdk_pipeline_options.view_as(
+          WorkerOptions).element_processing_timeout_minutes)
   return fn_log_handler, sdk_harness, sdk_pipeline_options
 
 

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

@@ -165,7 +165,8 @@ def __init__(
       state_cache=None,
       enable_heap_dump=False,
       worker_id=None,
-      log_lull_timeout_ns=DEFAULT_LOG_LULL_TIMEOUT_NS):
+      log_lull_timeout_ns=DEFAULT_LOG_LULL_TIMEOUT_NS,
+      element_processing_timeout_minutes=None):
     """Initialize FnApiWorkerStatusHandler.
 
     Args:
@@ -184,6 +185,10 @@ def __init__(
         self._status_channel)
     self._responses = queue.Queue()
     self.log_lull_timeout_ns = log_lull_timeout_ns
+    if element_processing_timeout_minutes:
+      self._element_processing_timeout_ns = element_processing_timeout_minutes * 60 * 1e9
+    else:
+      self._element_processing_timeout_ns = None
     self._last_full_thread_dump_secs = 0.0
     self._last_lull_logged_secs = 0.0
     self._server = threading.Thread(
@@ -250,14 +255,36 @@ def _log_lull_in_bundle_processor(self, bundle_process_cache):
           if processor:
             info = processor.state_sampler.get_info()
             self._log_lull_sampler_info(info, instruction)
+            if self._element_processing_timeout_ns:
+              self._terminate_sdk_worker_lull(info, instruction)
+
+  def _terminate_sdk_worker_lull(self, sampler_info, instruction):
+    if (sampler_info and sampler_info.time_since_transition and
+        sampler_info.time_since_transition
+        > self._element_processing_timeout_ns):
+      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)
+      log_lull_msg = (
+          '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))
+      raise TimeoutError(log_lull_msg + 'The SDK harness will be terminated.')
 
   def _log_lull_sampler_info(self, sampler_info, instruction):
     if not self._passed_lull_timeout_since_last_log():
       return
     if (sampler_info and sampler_info.time_since_transition and
         sampler_info.time_since_transition > self.log_lull_timeout_ns):
       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:

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

@@ -59,7 +59,8 @@ def setUp(self):
     self.test_port = self.server.add_insecure_port('[::]:0')
     self.server.start()
     self.url = 'localhost:%s' % self.test_port
-    self.fn_status_handler = FnApiWorkerStatusHandler(self.url)
+    self.fn_status_handler = FnApiWorkerStatusHandler(
+        self.url, element_processing_timeout_minutes=10)
 
   def tearDown(self):
     self.server.stop(5)
@@ -86,19 +87,19 @@ def test_generate_error(self, mock_method):
       self.assertIsNotNone(response.error)
     self.fn_status_handler.close()
 
-  def test_log_lull_in_bundle_processor(self):
-    def get_state_sampler_info_for_lull(lull_duration_s):
-      return "bundle-id", statesampler.StateSamplerInfo(
-          CounterName('progress-msecs', 'stage_name', 'step_name'),
-          1,
-          lull_duration_s * 1e9,
-          threading.current_thread())
+  def get_state_sampler_info_for_lull(lull_duration_s):
+    return "bundle-id", statesampler.StateSamplerInfo(
+        CounterName('progress-msecs', 'stage_name', 'step_name'),
+        1,
+        lull_duration_s * 1e9,
+        threading.current_thread())
 
+  def test_log_lull_in_bundle_processor(self):
     now = time.time()
     with mock.patch('logging.Logger.warning') as warn_mock:
       with mock.patch('time.time') as time_mock:
         time_mock.return_value = now
-        bundle_id, sampler_info = get_state_sampler_info_for_lull(21 * 60)
+        bundle_id, sampler_info = self.get_state_sampler_info_for_lull(21 * 60)
         self.fn_status_handler._log_lull_sampler_info(sampler_info, bundle_id)
 
         bundle_id_template = warn_mock.call_args[0][1]
@@ -113,19 +114,33 @@ def get_state_sampler_info_for_lull(lull_duration_s):
 
       with mock.patch('time.time') as time_mock:
         time_mock.return_value = now + 6 * 60  # 6 minutes
-        bundle_id, sampler_info = get_state_sampler_info_for_lull(21 * 60)
+        bundle_id, sampler_info = self.get_state_sampler_info_for_lull(21 * 60)
         self.fn_status_handler._log_lull_sampler_info(sampler_info, bundle_id)
 
       with mock.patch('time.time') as time_mock:
         time_mock.return_value = now + 21 * 60  # 21 minutes
-        bundle_id, sampler_info = get_state_sampler_info_for_lull(10 * 60)
+        bundle_id, sampler_info = self.get_state_sampler_info_for_lull(10 * 60)
         self.fn_status_handler._log_lull_sampler_info(sampler_info, bundle_id)
 
       with mock.patch('time.time') as time_mock:
         time_mock.return_value = now + 42 * 60  # 21 minutes after previous one
-        bundle_id, sampler_info = get_state_sampler_info_for_lull(21 * 60)
+        bundle_id, sampler_info = self.get_state_sampler_info_for_lull(21 * 60)
         self.fn_status_handler._log_lull_sampler_info(sampler_info, bundle_id)
 
+  def test__terminate_sdk_worker_lull_in_bundle_processor(self):
+    now = time.time()
+    with mock.patch('time.time') as time_mock:
+      time_mock.return_value = now
+      bundle_id, sampler_info = self.get_state_sampler_info_for_lull(5 * 60)
+      self.fn_status_handler._terminate_sdk_worker_lull(sampler_info, bundle_id)
+
+    with mock.patch('time.time') as time_mock:
+      time_mock.return_value = now + 6 * 60  # 6 minutes
+      bundle_id, sampler_info = self.get_state_sampler_info_for_lull(11 * 60)
+      with self.assertRaises(TimeoutError):
+        self.fn_status_handler._terminate_sdk_worker_lull(
+            sampler_info, bundle_id)
+
 
 class HeapDumpTest(unittest.TestCase):
   @mock.patch('apache_beam.runners.worker.worker_status.hpy', None)