Merge pull request #2687 from AI-Hypercomputer:fix-goodput-v15-integration

PiperOrigin-RevId: 832608784

Author: Google-ML-Automation

src/MaxText/elastic_train.py

@@ -384,8 +384,7 @@ def main(argv: Sequence[str]) -> None:
   if config.use_vertex_tensorboard or os.environ.get("UPLOAD_DATA_TO_TENSORBOARD"):
     vertex_tensorboard_manager.configure_vertex_tensorboard(config)
 
-  # Goodput configurations
-  maybe_monitor_goodput(config)
+  # Create the Goodput recorder
   recorder = create_goodput_recorder(config)
 
   # Stack traces configurations
@@ -399,7 +398,7 @@ def main(argv: Sequence[str]) -> None:
   diagnostic_config = diagnostic_configuration.DiagnosticConfig(debug_config)
 
   with diagnostic.diagnose(diagnostic_config):
-    with maybe_record_goodput(recorder, GoodputEvent.JOB):
+    with maybe_record_goodput(recorder, GoodputEvent.JOB), maybe_monitor_goodput(config):
       train_loop(config, elastic_manager, recorder)
 
 

src/MaxText/experimental/rl/grpo_trainer.py

@@ -946,8 +946,7 @@ def main(argv: Sequence[str]) -> None:
   if config.use_vertex_tensorboard or os.environ.get("UPLOAD_DATA_TO_TENSORBOARD"):
     vertex_tensorboard_manager.configure_vertex_tensorboard(config)
 
-  # Goodput configurations
-  maybe_monitor_goodput(config)
+  # Create the Goodput recorder
   recorder = create_goodput_recorder(config)
 
   # Stack traces configurations
@@ -961,7 +960,7 @@ def main(argv: Sequence[str]) -> None:
   diagnostic_config = diagnostic_configuration.DiagnosticConfig(debug_config)
 
   with diagnostic.diagnose(diagnostic_config):
-    with maybe_record_goodput(recorder, GoodputEvent.JOB):
+    with maybe_record_goodput(recorder, GoodputEvent.JOB), maybe_monitor_goodput(config):
       train_loop(config, config_inference, recorder)
 
 

src/MaxText/sft/sft_trainer.py

@@ -194,10 +194,9 @@ def main(argv: Sequence[str]) -> None:
   mt_config = pyconfig.initialize(argv)
   max_utils.print_system_information()
 
-  maybe_monitor_goodput(mt_config)
   goodput_recorder = create_goodput_recorder(mt_config)
 
-  with maybe_record_goodput(goodput_recorder, GoodputEvent.JOB):
+  with maybe_record_goodput(goodput_recorder, GoodputEvent.JOB), maybe_monitor_goodput(mt_config):
     train(mt_config, goodput_recorder)
 
 

src/MaxText/sft_trainer.py

@@ -168,9 +168,8 @@ def main(argv: Sequence[str]) -> None:
   validate_train_config(config)
   os.environ["TFDS_DATA_DIR"] = config.dataset_path
 
-  maybe_monitor_goodput(config)
   recorder = create_goodput_recorder(config)
-  with maybe_record_goodput(recorder, GoodputEvent.JOB):
+  with maybe_record_goodput(recorder, GoodputEvent.JOB), maybe_monitor_goodput(config):
     train_loop(config, recorder)
 
 

src/MaxText/train.py

@@ -502,8 +502,7 @@ def initialize(argv: Sequence[str]) -> tuple[pyconfig.HyperParameters, Any, Any]
   if config.use_vertex_tensorboard or os.environ.get("UPLOAD_DATA_TO_TENSORBOARD"):
     vertex_tensorboard_manager.configure_vertex_tensorboard(config)
 
-  # Goodput configurations
-  maybe_monitor_goodput(config)
+  # Create the Goodput recorder
   recorder = create_goodput_recorder(config)
 
   # Stack traces configurations
@@ -524,6 +523,7 @@ def run(config, recorder, diagnostic_config):
       diagnostic.diagnose(diagnostic_config),
       maybe_record_goodput(recorder, GoodputEvent.JOB),
       max_utils.maybe_get_transformer_engine_context(config),
+      maybe_monitor_goodput(config),
   ):
     train_loop(config, recorder)
 

src/MaxText/utils/goodput_utils.py

@@ -34,13 +34,14 @@ class GoodputEvent(Enum):
   STEP = "step"
 
 
+@contextlib.contextmanager
 def maybe_monitor_goodput(config):
-  """Monitor goodput if `monitor_goodput=True`."""
-  if config.monitor_goodput and jax.process_index() == 0:
-    # Workload monitoring and Goodput monitoring both uses /workload/performance
-    # GCM metric to publish step_time and step_deviation metrics. For now, we
-    # will disable publishing step deviation metrics to GCM if workload
-    # monitoring is enabled. Will reconcile this in the future.
+  """Monitor cumulative goodput if enabled."""
+  if not config.monitor_goodput or jax.process_index() != 0:
+    yield
+    return
+  goodput_monitor = None
+  try:
     if config.report_performance_metric_for_gcp_monitoring:
       config.enable_gcp_step_deviation_metrics = False
 
@@ -62,10 +63,11 @@ def maybe_monitor_goodput(config):
     )
     goodput_monitor.start_goodput_uploader()
     max_logging.log("Started Goodput upload to Tensorboard & GCM in the background!")
-
-    if config.monitor_step_time_deviation:
-      goodput_monitor.start_step_deviation_uploader()
-      max_logging.log("Started step time deviation upload to Tensorboard & GCM in the background!")
+    yield
+  finally:
+    if goodput_monitor:
+      goodput_monitor.stop_goodput_uploader()
+      max_logging.log("Flushed final metrics and safe exited from Goodput monitoring.")
 
 
 @contextlib.contextmanager
@@ -75,9 +77,13 @@ def maybe_record_goodput(recorder, event_name, *args):
     start_event_name = f"record_{event_name.value}_start_time"
     record_goodput(recorder, start_event_name, *args)
     yield
-  finally:
+  except BaseException:  # pylint: disable=W0706
+    raise
+  else:
     end_event_name = f"record_{event_name.value}_end_time"
     record_goodput(recorder, end_event_name, *args)
+  finally:
+    pass
 
 
 def record_goodput(recorder, event_name, *args):

tests/goodput_utils_test.py

@@ -53,16 +53,27 @@ def test_record_goodput(self, mock_cloud_logger, mock_record_job_start_time, moc
     mock_record_job_start_time.assert_called()
     mock_record_job_end_time.assert_called()
 
-  @mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor.start_step_deviation_uploader")
+    class TestException(BaseException):
+      pass
+
+    mock_record_job_start_time.reset_mock()
+    mock_record_job_end_time.reset_mock()
+    with self.assertRaises(TestException):
+      with maybe_record_goodput(recorder, GoodputEvent.JOB):
+        mock_record_job_start_time.assert_called_once()
+        raise TestException()
+
+    mock_record_job_start_time.assert_called_once()
+    mock_record_job_end_time.assert_not_called()
+
+  @mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor.stop_goodput_uploader")
   @mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor.start_goodput_uploader")
-  def test_monitor_goodput(self, mock_start_goodput_uploader, mock_start_step_deviation_uploader):
+  def test_monitor_goodput(self, mock_start_goodput_uploader, mock_stop_goodput_uploader):
     mock_start_goodput_uploader.return_value = mock.MagicMock()
-    mock_start_step_deviation_uploader.return_value = mock.MagicMock()
-
-    maybe_monitor_goodput(self.config)
 
-    mock_start_goodput_uploader.assert_called()
-    mock_start_step_deviation_uploader.assert_called()
+    with maybe_monitor_goodput(self.config):
+      mock_start_goodput_uploader.assert_called()
+    mock_stop_goodput_uploader.assert_called()
 
 
 if __name__ == "__main__":