[BEAM-36736] Add state sampling for timer processing in the Python SDK

sdks/python/apache_beam/runners/worker/operations.pxd

@@ -117,6 +117,7 @@ cdef class DoOperation(Operation):
   cdef dict timer_specs
   cdef public object input_info
   cdef object fn
+  cdef readonly object scoped_timer_processing_state
 
 
 cdef class SdfProcessSizedElements(DoOperation):

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

@@ -49,6 +49,7 @@
 from apache_beam.runners.worker import opcounters
 from apache_beam.runners.worker import operation_specs
 from apache_beam.runners.worker import sideinputs
+from apache_beam.runners.worker import statesampler
 from apache_beam.runners.worker.data_sampler import DataSampler
 from apache_beam.transforms import sideinputs as apache_sideinputs
 from apache_beam.transforms import combiners
@@ -444,12 +445,19 @@ def __init__(
     self.metrics_container = MetricsContainer(self.name_context.metrics_name())
 
     self.state_sampler = state_sampler
-    self.scoped_start_state = self.state_sampler.scoped_state(
-        self.name_context, 'start', metrics_container=self.metrics_container)
-    self.scoped_process_state = self.state_sampler.scoped_state(
-        self.name_context, 'process', metrics_container=self.metrics_container)
-    self.scoped_finish_state = self.state_sampler.scoped_state(
-        self.name_context, 'finish', metrics_container=self.metrics_container)
+    if self.state_sampler:
+      self.scoped_start_state = self.state_sampler.scoped_state(
+          self.name_context, 'start', metrics_container=self.metrics_container)
+      self.scoped_process_state = self.state_sampler.scoped_state(
+          self.name_context,
+          'process',
+          metrics_container=self.metrics_container)
+      self.scoped_finish_state = self.state_sampler.scoped_state(
+          self.name_context, 'finish', metrics_container=self.metrics_container)
+    else:
+      self.scoped_start_state = statesampler.NOOP_SCOPED_STATE
+      self.scoped_process_state = statesampler.NOOP_SCOPED_STATE
+      self.scoped_finish_state = statesampler.NOOP_SCOPED_STATE
     # TODO(ccy): the '-abort' state can be added when the abort is supported in
     # Operations.
     self.receivers = []  # type: List[ConsumerSet]
@@ -808,8 +816,15 @@ def __init__(
     self.user_state_context = user_state_context
     self.tagged_receivers = None  # type: Optional[_TaggedReceivers]
     # A mapping of timer tags to the input "PCollections" they come in on.
+    # Force clean rebuild
     self.input_info = None  # type: Optional[OpInputInfo]
-
+    self.scoped_timer_processing_state = statesampler.NOOP_SCOPED_STATE
+    if self.state_sampler:
+      self.scoped_timer_processing_state = self.state_sampler.scoped_state(
+          self.name_context,
+          'process-timers',
+          metrics_container=self.metrics_container,
+          suffix="-millis")
     # See fn_data in dataflow_runner.py
     # TODO: Store all the items from spec?
     self.fn, _, _, _, _ = (pickler.loads(self.spec.serialized_fn))
@@ -971,14 +986,21 @@ def add_timer_info(self, timer_family_id, timer_info):
     self.user_state_context.add_timer_info(timer_family_id, timer_info)
 
   def process_timer(self, tag, timer_data):
-    timer_spec = self.timer_specs[tag]
-    self.dofn_runner.process_user_timer(
-        timer_spec,
-        timer_data.user_key,
-        timer_data.windows[0],
-        timer_data.fire_timestamp,
-        timer_data.paneinfo,
-        timer_data.dynamic_timer_tag)
+    def process_timer_logic():
+      timer_spec = self.timer_specs[tag]
+      self.dofn_runner.process_user_timer(
+          timer_spec,
+          timer_data.user_key,
+          timer_data.windows[0],
+          timer_data.fire_timestamp,
+          timer_data.paneinfo,
+          timer_data.dynamic_timer_tag)
+
+    if self.scoped_timer_processing_state:
+      with self.scoped_timer_processing_state:
+        process_timer_logic()
+    else:
+      process_timer_logic()
 
   def finish(self):
     # type: () -> None

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

@@ -134,8 +134,8 @@ def scoped_state(
       name_context: Union[str, 'common.NameContext'],
       state_name: str,
       io_target=None,
-      metrics_container: Optional['MetricsContainer'] = None
-  ) -> statesampler_impl.ScopedState:
+      metrics_container: Optional['MetricsContainer'] = None,
+      suffix: str = '-msecs') -> statesampler_impl.ScopedState:
     """Returns a ScopedState object associated to a Step and a State.
 
     Args:
@@ -152,7 +152,7 @@ def scoped_state(
       name_context = common.NameContext(name_context)
 
     counter_name = CounterName(
-        state_name + '-msecs',
+        state_name + suffix,
         stage_name=self._prefix,
         step_name=name_context.metrics_name(),
         io_target=io_target)
@@ -170,3 +170,17 @@ def commit_counters(self) -> None:
     for state in self._states_by_name.values():
       state_msecs = int(1e-6 * state.nsecs)
       state.counter.update(state_msecs - state.counter.value())
+
+
+class NoOpScopedState:
+  def __enter__(self):
+    pass
+
+  def __exit__(self, exc_type, exc_val, exc_tb):
+    pass
+
+  def sampled_msecs_int(self):
+    return 0
+
+
+NOOP_SCOPED_STATE = NoOpScopedState()

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

@@ -127,6 +127,41 @@ def test_sampler_transition_overhead(self):
     # debug mode).
     self.assertLess(overhead_us, 20.0)
 
+  @retry(reraise=True, stop=stop_after_attempt(3))
+  def test_timer_sampler(self):
+    # Set up state sampler.
+    counter_factory = CounterFactory()
+    sampler = statesampler.StateSampler(
+        'timer', counter_factory, sampling_period_ms=1)
+
+    # Duration of the timer processing.
+    state_duration_ms = 100
+    margin_of_error = 0.25
+
+    sampler.start()
+    with sampler.scoped_state('step1', 'process-timers', suffix='-millis'):
+      time.sleep(state_duration_ms / 1000)
+    sampler.stop()
+    sampler.commit_counters()
+
+    if not statesampler.FAST_SAMPLER:
+      # The slow sampler does not implement sampling, so we won't test it.
+      return
+
+    # Test that sampled state timings are close to their expected values.
+    expected_counter_values = {
+        CounterName(
+            'process-timers-millis', step_name='step1', stage_name='timer'): state_duration_ms,
+    }
+    for counter in counter_factory.get_counters():
+      self.assertIn(counter.name, expected_counter_values)
+      expected_value = expected_counter_values[counter.name]
+      actual_value = counter.value()
+      deviation = float(abs(actual_value - expected_value)) / expected_value
+      _LOGGER.info('Sampling deviation from expectation: %f', deviation)
+      self.assertGreater(actual_value, expected_value * (1.0 - margin_of_error))
+      self.assertLess(actual_value, expected_value * (1.0 + margin_of_error))
+
 
 if __name__ == '__main__':
   logging.getLogger().setLevel(logging.INFO)