Add different tests, remove no op scoped state, and address formatting/lint issues

Author: ktalluri456

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

@@ -49,7 +49,6 @@
 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
@@ -445,19 +444,12 @@ def __init__(
     self.metrics_container = MetricsContainer(self.name_context.metrics_name())
 
     self.state_sampler = state_sampler
-    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
+    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)
     # TODO(ccy): the '-abort' state can be added when the abort is supported in
     # Operations.
     self.receivers = []  # type: List[ConsumerSet]
@@ -817,12 +809,10 @@ def __init__(
     self.tagged_receivers = None  # type: Optional[_TaggedReceivers]
     # A mapping of timer tags to the input "PCollections" they come in on.
     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)
+    self.scoped_timer_processing_state = self.state_sampler.scoped_state(
+        self.name_context,
+        'process-timers',
+        metrics_container=self.metrics_container)
     # See fn_data in dataflow_runner.py
     # TODO: Store all the items from spec?
     self.fn, _, _, _, _ = (pickler.loads(self.spec.serialized_fn))
@@ -984,7 +974,7 @@ 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):
-    def process_timer_logic():
+    with self.scoped_timer_processing_state:
       timer_spec = self.timer_specs[tag]
       self.dofn_runner.process_user_timer(
           timer_spec,
@@ -994,12 +984,6 @@ def process_timer_logic():
           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
     super(DoOperation, self).finish()

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

@@ -170,17 +170,3 @@ 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

@@ -26,13 +26,14 @@
 from tenacity import retry
 from tenacity import stop_after_attempt
 
-from apache_beam.runners.worker import statesampler
-from apache_beam.utils.counters import CounterFactory
-from apache_beam.utils.counters import CounterName
+from apache_beam.internal import pickler
+from apache_beam.runners import common
 from apache_beam.runners.worker import operation_specs
 from apache_beam.runners.worker import operations
-from apache_beam.internal import pickler
+from apache_beam.runners.worker import statesampler
 from apache_beam.transforms import core
+from apache_beam.utils.counters import CounterFactory
+from apache_beam.utils.counters import CounterName
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -143,7 +144,7 @@ def test_process_timers_metric_is_recorded(self):
     sampler = statesampler.StateSampler(
         'test_stage', counter_factory, sampling_period_ms=1)
 
-    # Keeps range between 50-350 ms, which is fair. 
+    # Keeps range between 50-350 ms, which is fair.
     state_duration_ms = 200
     margin_of_error = 0.75
 
@@ -184,52 +185,149 @@ def test_process_timers_metric_is_recorded(self):
         expected_value * (1.0 + margin_of_error),
         "The timer metric was higher than expected.")
 
-  def test_do_operation_with_sampler(self):
+  def test_do_operation_process_timer_metric(self):
+    """
+    Tests that the 'process-timers-msecs' metric is correctly recorded
+    when a timer is processed within a DoOperation.
     """
-      Tests that a DoOperation with an active state_sampler correctly
-      creates a real ScopedState object for timer processing.
-      """
+    counter_factory = CounterFactory()
+    sampler = statesampler.StateSampler(
+        'test_stage', counter_factory, sampling_period_ms=1)
+
     mock_spec = operation_specs.WorkerDoFn(
         serialized_fn=pickler.dumps((core.DoFn(), None, None, None, None)),
         output_tags=[],
         input=None,
         side_inputs=[],
         output_coders=[])
 
-    sampler = statesampler.StateSampler(
-        'test_stage', CounterFactory(), sampling_period_ms=1)
-
-    # 1. Create the operation WITHOUT the unexpected keyword argument
-    op = operations.create_operation(
-        name_context='test_op',
+    op = operations.DoOperation(
+        name=common.NameContext('test_op'),
         spec=mock_spec,
-        counter_factory=CounterFactory(),
-        state_sampler=sampler)
+        counter_factory=counter_factory,
+        sampler=sampler)
+
+    op.dofn_runner = Mock()
+    op.timer_specs = {'timer_id': Mock()}
+    state_duration_ms = 200
+    margin_of_error = 0.75
+
+    def mock_process_user_timer(*args, **kwargs):
+      time.sleep(state_duration_ms / 1000.0)
+
+    op.dofn_runner.process_user_timer = mock_process_user_timer
+
+    mock_timer_data = Mock()
+    mock_timer_data.windows = [Mock()]
+    mock_timer_data.user_key = Mock()
+    mock_timer_data.fire_timestamp = Mock()
+    mock_timer_data.paneinfo = Mock()
+    mock_timer_data.dynamic_timer_tag = Mock()
 
-    self.assertIsNot(
-        op.scoped_timer_processing_state, statesampler.NOOP_SCOPED_STATE)
+    sampler.start()
+    op.process_timer('timer_id', mock_timer_data)
+    sampler.stop()
+    sampler.commit_counters()
+
+    if not statesampler.FAST_SAMPLER:
+      return
+
+    expected_counter_name = CounterName(
+        'process-timers-msecs', step_name='test_op', stage_name='test_stage')
+
+    found_counter = None
+    for counter in counter_factory.get_counters():
+      if counter.name == expected_counter_name:
+        found_counter = counter
+        break
+
+    self.assertIsNotNone(
+        found_counter,
+        f"The expected counter '{expected_counter_name}' was not created.")
+
+    actual_value = found_counter.value()
+    expected_value = state_duration_ms
+    self.assertGreater(
+        actual_value,
+        expected_value * (1.0 - margin_of_error),
+        "The timer metric was lower than expected.")
+    self.assertLess(
+        actual_value,
+        expected_value * (1.0 + margin_of_error),
+        "The timer metric was higher than expected.")
 
-  def test_do_operation_without_sampler(self):
+  def test_do_operation_process_timer_metric_with_exception(self):
     """
-    Tests that a DoOperation without a state_sampler correctly uses the
-    NOOP_SCOPED_STATE for timer processing.
+    Tests that the 'process-timers-msecs' metric is still recorded
+    when a timer callback in a DoOperation raises an exception.
     """
+    counter_factory = CounterFactory()
+    sampler = statesampler.StateSampler(
+        'test_stage', counter_factory, sampling_period_ms=1)
+
     mock_spec = operation_specs.WorkerDoFn(
         serialized_fn=pickler.dumps((core.DoFn(), None, None, None, None)),
         output_tags=[],
         input=None,
         side_inputs=[],
         output_coders=[])
 
-    # 1. Create the operation WITHOUT the unexpected keyword argument
-    op = operations.create_operation(
-        name_context='test_op',
+    op = operations.DoOperation(
+        name=common.NameContext('test_op'),
         spec=mock_spec,
-        counter_factory=CounterFactory(),
-        state_sampler=None)
+        counter_factory=counter_factory,
+        sampler=sampler)
+
+    op.dofn_runner = Mock()
+    op.timer_specs = {'timer_id': Mock()}
+    state_duration_ms = 200
+    margin_of_error = 0.75
+
+    def mock_process_user_timer(*args, **kwargs):
+      time.sleep(state_duration_ms / 1000.0)
+      raise ValueError("Test Exception")
+
+    op.dofn_runner.process_user_timer = mock_process_user_timer
+
+    mock_timer_data = Mock()
+    mock_timer_data.windows = [Mock()]
+    mock_timer_data.user_key = Mock()
+    mock_timer_data.fire_timestamp = Mock()
+    mock_timer_data.paneinfo = Mock()
+    mock_timer_data.dynamic_timer_tag = Mock()
+
+    sampler.start()
+    with self.assertRaises(ValueError):
+      op.process_timer('timer_id', mock_timer_data)
+    sampler.stop()
+    sampler.commit_counters()
+
+    if not statesampler.FAST_SAMPLER:
+      return
+
+    expected_counter_name = CounterName(
+        'process-timers-msecs', step_name='test_op', stage_name='test_stage')
+
+    found_counter = None
+    for counter in counter_factory.get_counters():
+      if counter.name == expected_counter_name:
+        found_counter = counter
+        break
 
-    self.assertIs(
-        op.scoped_timer_processing_state, statesampler.NOOP_SCOPED_STATE)
+    self.assertIsNotNone(
+        found_counter,
+        f"The expected counter '{expected_counter_name}' was not created.")
+
+    actual_value = found_counter.value()
+    expected_value = state_duration_ms
+    self.assertGreater(
+        actual_value,
+        expected_value * (1.0 - margin_of_error),
+        "The timer metric was lower than expected.")
+    self.assertLess(
+        actual_value,
+        expected_value * (1.0 + margin_of_error),
+        "The timer metric was higher than expected.")
 
 
 if __name__ == '__main__':