Add monitoring API to upload step deviation to Tensorboard.

PiperOrigin-RevId: 720236693

Author: dipannita08

ml_goodput_measurement/src/monitoring.py

@@ -16,6 +16,7 @@
 _TENSORBOARD_GCS_SUBDIR = 'goodput'
 _TENSORBOARD_GOODPUT_LABEL = 'goodput'
 _TENSORBOARD_BADPUT_LABEL = 'badput'
+_TENSORBOARD_STEP_DEVIATION_LABEL = 'step_deviation'
 
 logger = logging.getLogger(__name__)
 
@@ -32,6 +33,9 @@ def __init__(
       monitoring_enabled: bool = False,
       pathway_enabled: bool = False,
       include_badput_breakdown=False,
+      include_step_deviation=False,
+      configured_ideal_step_time=None,
+      step_deviation_interval_seconds=10,
   ):
     """Initializes the GoodputMonitor.
 
@@ -47,6 +51,11 @@ def __init__(
         pathway_enabled: Whether the application is using Pathways.
         include_badput_breakdown: Whether to query and upload badput breakdown
           data to Tensorboard.
+        include_step_deviation: Whether to query and upload step deviation data
+          to Tensorboard.
+        step_deviation_interval_seconds: The interval to query step deviation
+          data.
+        configured_ideal_step_time: The optionalideal step time configured by the user.
     """
     if not monitoring_enabled:
       logger.info(
@@ -55,32 +64,48 @@ def __init__(
       )
       return
 
+    # Common configurations.
     self._job_name = job_name
     self._logger_name = logger_name
     self._tensorboard_dir = os.path.join(
         tensorboard_dir, _TENSORBOARD_GCS_SUBDIR
     )
+    # Goodput configurations.
     self._upload_interval = upload_interval
+    self._include_badput_breakdown = include_badput_breakdown
+
+    # Step deviation configurations.
+    self._include_step_deviation = include_step_deviation
+    self._step_deviation_interval_seconds = step_deviation_interval_seconds
+    self._configured_ideal_step_time = configured_ideal_step_time
+
+    # Initialize the GoodputCalculator.
     self._goodput_calculator = GoodputCalculator(
         job_name=self._job_name,
         logger_name=self._logger_name,
         using_pathways=pathway_enabled,
     )
     self._writer = writer.SummaryWriter(self._tensorboard_dir)
-    self._include_badput_breakdown = include_badput_breakdown
 
-    # Flag to signal the daemon thread if it exists when to initate
+    # Goodput uploader flags to signal the daemon thread if it exists when to initate
     # shutdown and wait for termination.
     self._uploader_thread_running = False
     self._goodput_upload_thread = None
     self._termination_event = threading.Event()
     self._termination_event.clear()
 
+    # Step deviation threading flags.
+    self._step_deviation_uploader_thread_running = False
+    self._step_deviation_upload_thread = None
+    self._step_deviation_termination_event = threading.Event()
+    self._step_deviation_termination_event.clear()
+
   def __del__(self):
     if self._uploader_thread_running:
       self.stop_goodput_uploader()
+      self.stop_step_deviation_uploader()
 
-  def _write_to_tensorboard(
+  def _write_goodput_data_to_tensorboard(
       self,
       job_goodput: float,
       badput_breakdown: dict[BadputType, float],
@@ -125,7 +150,7 @@ def _query_and_upload_goodput(self):
                 include_badput_breakdown=self._include_badput_breakdown
             )
         )
-        self._write_to_tensorboard(job_goodput, job_badput_breakdown, last_step)
+        self._write_goodput_data_to_tensorboard(job_goodput, job_badput_breakdown, last_step)
       except Exception as e:  # pylint: disable=broad-exception-caught
         logger.error(
             'Error while querying and uploading goodput to Tensorboard. This'
@@ -165,3 +190,81 @@ def stop_goodput_uploader(self):
         ' be uploaded to Tensorboard.'
     )
     self._uploader_thread_running = False
+
+  def _write_step_deviation_to_tensorboard(
+      self, step_deviation: dict[int, float]
+  ):
+    if self._writer is not None:
+      for step_count, step_deviation in step_deviation.items():
+        self._writer.add_scalar(
+            _TENSORBOARD_STEP_DEVIATION_LABEL,
+            float(step_deviation),
+            step_count,
+        )
+      self._writer.flush()
+
+  def _query_and_upload_step_deviation(self):
+    """Queries and uploads step deviation data to Tensorboard."""
+    while not self._step_deviation_termination_event.is_set():
+      time.sleep(self._step_deviation_interval_seconds)
+      try:
+        step_deviation = self._goodput_calculator.get_step_deviation(
+            self._configured_ideal_step_time
+        )
+      except Exception as e:  # pylint: disable=broad-exception-caught
+        logger.error(
+            'Error while querying step deviation to Tensorboard.'
+            ' This will not impact the workload. Error: %s',
+            e,
+        )
+        continue
+      try:
+        self._write_step_deviation_to_tensorboard(step_deviation)
+      except Exception as e:  # pylint: disable=broad-exception-caught
+        logger.error(
+            'Error while writing step deviation to Tensorboard.'
+            ' This will not impact the workload. Error: %s',
+            e,
+        )
+
+  def start_step_deviation_uploader(self):
+    """Starts the step deviation uploader thread."""
+    if not self._include_step_deviation:
+      logger.info(
+          'Step deviation monitoring is disabled. Returning without initializing'
+          ' step deviation uploader thread.'
+      )
+      return
+
+    if self._step_deviation_uploader_thread_running:
+      raise RuntimeError('Step deviation uploader thread is already running.')
+
+    self._step_deviation_termination_event.clear()
+    self._step_deviation_upload_thread = threading.Thread(
+        target=self._query_and_upload_step_deviation, daemon=True
+    )
+    logger.info(
+        'Starting step deviation query and uploader thread in the background'
+        ' for job: %s and logger: %s',
+        self._job_name,
+        self._logger_name,
+    )
+    self._step_deviation_upload_thread.start()
+    self._step_deviation_uploader_thread_running = True
+
+  def stop_step_deviation_uploader(self):
+    """Stops the step deviation uploader thread."""
+    if not self._step_deviation_uploader_thread_running:
+      raise RuntimeError('Step deviation uploader thread is not running.')
+
+    self._step_deviation_termination_event.set()
+    if self._step_deviation_upload_thread is not None:
+      logger.info(
+          'Waiting for step deviation query and uploader thread to complete.'
+      )
+      self._step_deviation_upload_thread.join()
+    logger.info(
+        'Step deviation query and uploader thread stopped. No more step'
+        ' deviation data will be uploaded to Tensorboard.'
+    )
+    self._step_deviation_uploader_thread_running = False

ml_goodput_measurement/tests/monitoring_test.py

@@ -0,0 +1,217 @@
+"""Tests to validate the monitoring module.
+
+This module tests the GoodputMonitor class and its functionality, specifically
+the uploading of step deviation, goodput and badput data to Tensorboard.
+"""
+
+from unittest import mock
+
+from absl.testing import absltest
+from cloud_goodput.ml_goodput_measurement.src import monitoring
+
+GoodputMonitor = monitoring.GoodputMonitor
+patch = mock.patch
+MagicMock = mock.MagicMock
+
+_TEST_UPLOAD_INTERVAL = 1
+
+
+class GoodputMonitorTests(absltest.TestCase):
+  """Tests for the GoodputMonitor class."""
+
+  def setUp(self):
+    super().setUp()
+    self.job_name = 'test-run'
+    self.logger_name = 'test-logger'
+    self.tensorboard_dir = 'test-dir'
+
+  @patch('tensorboardX.writer.SummaryWriter')
+  @patch('google.cloud.logging.Client')
+  def test_goodput_monitor_init(self, mock_logger_client, mock_summary_writer):
+    mock_summary_writer.return_value = MagicMock()
+    mock_logger_client.return_value = MagicMock()
+    goodput_monitor = GoodputMonitor(
+        self.job_name,
+        self.logger_name,
+        self.tensorboard_dir,
+        upload_interval=_TEST_UPLOAD_INTERVAL,
+        monitoring_enabled=True,
+    )
+    # Objects should be initialized correctly.
+    self.assertIsNotNone(goodput_monitor)
+    self.assertIs(goodput_monitor._writer, mock_summary_writer.return_value)
+    self.assertIsNotNone(goodput_monitor._goodput_calculator)
+
+    # Thread events should be initialized correctly.
+    self.assertIsNotNone(goodput_monitor._step_deviation_termination_event)
+    self.assertFalse(goodput_monitor._step_deviation_termination_event.is_set())
+    self.assertFalse(goodput_monitor._step_deviation_uploader_thread_running)
+    self.assertIsNotNone(goodput_monitor._termination_event)
+    self.assertFalse(goodput_monitor._termination_event.is_set())
+    self.assertFalse(goodput_monitor._uploader_thread_running)
+
+  @patch(
+      'cloud_goodput.ml_goodput_measurement.src.monitoring.GoodputMonitor._write_goodput_to_tensorboard'
+  )
+  @patch('tensorboardX.writer.SummaryWriter')
+  @patch('google.cloud.logging.Client')
+  async def test_goodput_monitor_start_goodput_uploader_success(
+      self, mock_logger_client, mock_summary_writer, mock_goodput_to_tensorboard
+  ):
+    mock_summary_writer.return_value = MagicMock()
+    mock_goodput_to_tensorboard.return_value = MagicMock()
+    mock_logger_client.return_value = MagicMock()
+    goodput_monitor = monitoring.GoodputMonitor(
+        self.job_name,
+        self.logger_name,
+        self.tensorboard_dir,
+        upload_interval=_TEST_UPLOAD_INTERVAL,
+        monitoring_enabled=True,
+    )
+    goodput_monitor.start_goodput_uploader()
+    self.assertTrue(goodput_monitor._uploader_thread_running)
+    self.assertIsNotNone(goodput_monitor._goodput_upload_thread)
+    self.assertFalse(goodput_monitor._termination_event.is_set())
+    mock_goodput_to_tensorboard.assert_called_once()
+    mock_summary_writer.return_value.add_scalar.assert_called_once()
+    goodput_monitor.stop_goodput_uploader()
+    self.assertFalse(goodput_monitor._uploader_thread_running)
+    self.assertIsNone(goodput_monitor._goodput_upload_thread)
+    self.assertTrue(goodput_monitor._termination_event.is_set())
+
+  @patch(
+      'cloud_goodput.ml_goodput_measurement.src.monitoring.GoodputMonitor._write_goodput_to_tensorboard'
+  )
+  @patch('tensorboardX.writer.SummaryWriter')
+  @patch('google.cloud.logging.Client')
+  async def test_goodput_monitor_start_goodput_uploader_failure(
+      self, mock_logger_client, mock_summary_writer, mock_goodput_to_tensorboard
+  ):
+    mock_logger_client.return_value = MagicMock()
+    mock_summary_writer.return_value = MagicMock()
+    mock_goodput_to_tensorboard.side_effect = ValueError('Test Error')
+    goodput_monitor = monitoring.GoodputMonitor(
+        self.job_name,
+        self.logger_name,
+        self.tensorboard_dir,
+        upload_interval=_TEST_UPLOAD_INTERVAL,
+        monitoring_enabled=True,
+    )
+    goodput_monitor.start_goodput_uploader()
+    self.assertTrue(goodput_monitor._uploader_thread_running)
+    self.assertIsNotNone(goodput_monitor._goodput_upload_thread)
+    self.assertFalse(goodput_monitor._termination_event.is_set())
+    mock_goodput_to_tensorboard.assert_called_once()
+    with self.assertRaisesRegex(ValueError, 'Test Error'):
+      goodput_monitor._query_and_upload_goodput()
+    mock_summary_writer.return_value.add_scalar.assert_not_called()
+    goodput_monitor.stop_goodput_uploader()
+    self.assertFalse(goodput_monitor._uploader_thread_running)
+    self.assertIsNone(goodput_monitor._goodput_upload_thread)
+    self.assertTrue(goodput_monitor._termination_event.is_set())
+
+  @patch(
+      'cloud_goodput.ml_goodput_measurement.src.monitoring.GoodputMonitor._write_badput_to_tensorboard'
+  )
+  @patch('tensorboardX.writer.SummaryWriter')
+  @patch('google.cloud.logging.Client')
+  async def test_goodput_monitor_start_badput_uploader_success(
+      self, mock_logger_client, mock_summary_writer, mock_badput_to_tensorboard
+  ):
+    mock_summary_writer.return_value = MagicMock()
+    mock_badput_to_tensorboard.return_value = MagicMock()
+    mock_logger_client.return_value = MagicMock()
+    goodput_monitor = monitoring.GoodputMonitor(
+        self.job_name,
+        self.logger_name,
+        self.tensorboard_dir,
+        upload_interval=_TEST_UPLOAD_INTERVAL,
+        monitoring_enabled=True,
+        include_badput_breakdown=True,
+    )
+
+    goodput_monitor.start_goodput_uploader()
+    self.assertTrue(goodput_monitor._uploader_thread_running)
+    self.assertIsNotNone(goodput_monitor._goodput_upload_thread)
+    self.assertFalse(goodput_monitor._termination_event.is_set())
+    self.assertTrue(goodput_monitor._include_badput_breakdown)
+
+    mock_badput_to_tensorboard.assert_called_once()
+    mock_summary_writer.return_value.add_scalar.assert_called_once()
+
+    goodput_monitor.stop_goodput_uploader()
+    self.assertFalse(goodput_monitor._uploader_thread_running)
+    self.assertIsNone(goodput_monitor._goodput_upload_thread)
+    self.assertTrue(goodput_monitor._termination_event.is_set())
+
+  @patch(
+      'cloud_goodput.ml_goodput_measurement.src.monitoring.GoodputMonitor._write_step_deviation_to_tensorboard'
+  )
+  @patch('tensorboardX.writer.SummaryWriter')
+  @patch('google.cloud.logging.Client')
+  async def test_goodput_monitor_start_step_deviation_uploader_success(
+      self,
+      mock_logger_client,
+      mock_summary_writer,
+      mock_step_deviation_to_tensorboard,
+  ):
+    mock_logger_client.return_value = MagicMock()
+    mock_summary_writer.return_value = MagicMock()
+    mock_step_deviation_to_tensorboard.return_value = MagicMock()
+    goodput_monitor = monitoring.GoodputMonitor(
+        self.job_name,
+        self.logger_name,
+        self.tensorboard_dir,
+        upload_interval=_TEST_UPLOAD_INTERVAL,
+        monitoring_enabled=True,
+        include_step_deviation=True,
+    )
+    goodput_monitor.start_step_deviation_uploader()
+    self.assertTrue(goodput_monitor._step_deviation_uploader_thread_running)
+    self.assertIsNotNone(goodput_monitor._step_deviation_upload_thread)
+    self.assertFalse(goodput_monitor._step_deviation_termination_event.is_set())
+    mock_step_deviation_to_tensorboard.assert_called_once()
+    mock_summary_writer.return_value.add_scalar.assert_called_once()
+    goodput_monitor.stop_step_deviation_uploader()
+    self.assertFalse(goodput_monitor._step_deviation_uploader_thread_running)
+    self.assertIsNone(goodput_monitor._step_deviation_upload_thread)
+    self.assertTrue(goodput_monitor._step_deviation_termination_event.is_set())
+
+  @patch(
+      'cloud_goodput.ml_goodput_measurement.src.monitoring.GoodputMonitor._write_step_deviation_to_tensorboard'
+  )
+  @patch('tensorboardX.writer.SummaryWriter')
+  @patch('google.cloud.logging.Client')
+  async def test_goodput_monitor_start_step_deviation_uploader_failure(
+      self,
+      mock_logger_client,
+      mock_summary_writer,
+      mock_query_and_upload_step_deviation,
+  ):
+    mock_logger_client.return_value = MagicMock()
+    mock_summary_writer.return_value = MagicMock()
+    mock_query_and_upload_step_deviation.side_effect = ValueError('Test Error')
+    goodput_monitor = monitoring.GoodputMonitor(
+        self.job_name,
+        self.logger_name,
+        self.tensorboard_dir,
+        upload_interval=_TEST_UPLOAD_INTERVAL,
+        monitoring_enabled=True,
+        include_step_deviation=True,
+    )
+    goodput_monitor.start_step_deviation_uploader()
+    self.assertTrue(goodput_monitor._step_deviation_uploader_thread_running)
+    self.assertIsNotNone(goodput_monitor._step_deviation_upload_thread)
+    self.assertFalse(goodput_monitor._step_deviation_termination_event.is_set())
+    mock_query_and_upload_step_deviation.assert_called_once()
+    with self.assertRaisesRegex(ValueError, 'Test Error'):
+      goodput_monitor._query_and_upload_step_deviation()
+    mock_summary_writer.return_value.add_scalar.assert_not_called()
+    goodput_monitor.stop_step_deviation_uploader()
+    self.assertFalse(goodput_monitor._step_deviation_uploader_thread_running)
+    self.assertIsNone(goodput_monitor._step_deviation_upload_thread)
+    self.assertTrue(goodput_monitor._step_deviation_termination_event.is_set())
+
+
+if __name__ == '__main__':
+  absltest.main()