Use rolling window to track goodput metrics over time

src/maxtext/common/gcloud_stub.py

@@ -314,6 +314,12 @@ def start_goodput_uploader(self):
     def start_step_deviation_uploader(self):
       print("[DECOUPLED NO-OP] goodput step deviation uploader skipped.")
 
+    def start_rolling_window_goodput_uploader(self, *_a, **_k):
+      print("[DECOUPLED NO-OP] rolling window goodput uploader skipped.")
+
+    def stop_rolling_window_goodput_uploader(self, *_a, **_k):
+      print("[DECOUPLED NO-OP] rolling window goodput uploader stop skipped.")
+
   monitoring_ns = SimpleNamespace(GCPOptions=_StubMonitoringOptions, GoodputMonitor=_StubGoodputMonitor)
   goodput_ns = SimpleNamespace(GoodputRecorder=_StubGoodputRecorder)
   return goodput_ns, monitoring_ns, True

src/maxtext/common/goodput.py

@@ -109,12 +109,42 @@ def maybe_monitor_goodput(config):
 
     goodput_monitor = _construct_goodput_monitor(config, common_kwargs)
     goodput_monitor.start_goodput_uploader()
-    max_logging.log("Started Goodput upload to Tensorboard & GCM in the background!")
-    yield
+    max_logging.log("Started Cumulative Goodput upload to Tensorboard & GCM in the background!")
+    with maybe_monitor_rolling_window_goodput(goodput_monitor, config):
+      yield
   finally:
     if goodput_monitor:
       goodput_monitor.stop_goodput_uploader()
-      max_logging.log("Flushed final metrics and safe exited from Goodput monitoring.")
+      max_logging.log("Flushed final metrics and safe exited from cumulative Goodput monitoring.")
+
+
+@contextlib.contextmanager
+def maybe_monitor_rolling_window_goodput(goodput_monitor, config):
+  """Monitor rolling window goodput if enabled on the lead host.
+
+  Failures to start or stop the rolling window uploader are logged and
+  swallowed so that monitoring degrades gracefully and never interrupts
+  training. Exceptions raised by the wrapped block are not caught here.
+  """
+  if _GOODPUT_STUB or not goodput_monitor or not config.enable_rolling_window_goodput:
+    yield
+    return
+  started = False
+  try:
+    goodput_monitor.start_rolling_window_goodput_uploader(config.rolling_windows_seconds)
+    started = True
+    max_logging.log("Started Rolling Window Goodput upload to Tensorboard & GCM in the background!")
+  except Exception as e:  # pylint: disable=broad-exception-caught
+    max_logging.log(f"Goodput: failed to start rolling window goodput uploader, continuing without it: {e}")
+  try:
+    yield
+  finally:
+    if started:
+      try:
+        goodput_monitor.stop_rolling_window_goodput_uploader()
+        max_logging.log("Flushed final metrics and safe exited from rolling window Goodput monitoring.")
+      except Exception as e:  # pylint: disable=broad-exception-caught
+        max_logging.log(f"Goodput: failed to stop rolling window goodput uploader cleanly: {e}")
 
 
 @contextlib.contextmanager

src/maxtext/configs/base.yml

@@ -1191,6 +1191,8 @@ monitor_step_time_deviation: true
 step_deviation_interval_seconds: 30
 enable_gcp_goodput_metrics: true
 enable_gcp_step_deviation_metrics: true
+rolling_windows_seconds: [3600, 86400, 259200, 432000]
+enable_rolling_window_goodput: true
 
 # GCP workload monitoring
 report_heartbeat_metric_for_gcp_monitoring: false

src/maxtext/configs/types.py

@@ -2863,6 +2863,12 @@ class Goodput(BaseModel):
   step_deviation_interval_seconds: int = Field(30, description="Interval to check step time deviation.")
   enable_gcp_goodput_metrics: bool = Field(True, description="Enable GCP goodput metrics.")
   enable_gcp_step_deviation_metrics: bool = Field(True, description="Enable GCP step deviation metrics.")
+  rolling_windows_seconds: list[PositiveInt] = Field(
+      [3600, 86400, 259200, 432000],
+      min_length=1,
+      description="Rolling window intervals in seconds. Must contain at least one window.",
+  )
+  enable_rolling_window_goodput: bool = Field(True, description="Enable rolling window goodput.")
 
 
 class ElasticTraining(BaseModel):

tests/unit/goodput_utils_test.py

@@ -100,15 +100,127 @@ class TestException(BaseException):
     mock_record_job_start_time.assert_called_once()
     mock_record_job_end_time.assert_not_called()
 
+  @mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor.stop_rolling_window_goodput_uploader")
+  @mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor.start_rolling_window_goodput_uploader")
   @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_stop_goodput_uploader):
+  def test_monitor_goodput(self, mock_start_goodput_uploader, mock_stop_goodput_uploader, *unused_rolling_mocks):
     mock_start_goodput_uploader.return_value = mock.MagicMock()
 
     with maybe_monitor_goodput(self.config):
       mock_start_goodput_uploader.assert_called()
     mock_stop_goodput_uploader.assert_called()
 
+  @mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor.stop_rolling_window_goodput_uploader")
+  @mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor.start_rolling_window_goodput_uploader")
+  @mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor.stop_goodput_uploader")
+  @mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor.start_goodput_uploader")
+  def test_monitor_rolling_window_goodput(
+      self, mock_start_goodput, mock_stop_goodput, mock_start_rolling, mock_stop_rolling
+  ):
+    """Both uploaders start; on exit the rolling window stops before the cumulative one."""
+    windows = [3600, 86400]
+    config = _ConfigOverride(self.config, enable_rolling_window_goodput=True, rolling_windows_seconds=windows)
+    calls = mock.Mock()
+    calls.attach_mock(mock_stop_rolling, "stop_rolling")
+    calls.attach_mock(mock_stop_goodput, "stop_goodput")
+
+    with maybe_monitor_goodput(config):
+      mock_start_goodput.assert_called_once()
+      mock_start_rolling.assert_called_once_with(windows)
+      mock_stop_rolling.assert_not_called()
+
+    self.assertEqual(calls.mock_calls, [mock.call.stop_rolling(), mock.call.stop_goodput()])
+
+  @mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor.stop_rolling_window_goodput_uploader")
+  @mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor.start_rolling_window_goodput_uploader")
+  @mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor.stop_goodput_uploader")
+  @mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor.start_goodput_uploader")
+  def test_monitor_rolling_window_goodput_disabled(
+      self, mock_start_goodput, mock_stop_goodput, mock_start_rolling, mock_stop_rolling
+  ):
+    """enable_rolling_window_goodput=False only runs the cumulative uploader."""
+    config = _ConfigOverride(self.config, enable_rolling_window_goodput=False)
+
+    with maybe_monitor_goodput(config):
+      mock_start_goodput.assert_called_once()
+
+    mock_stop_goodput.assert_called_once()
+    mock_start_rolling.assert_not_called()
+    mock_stop_rolling.assert_not_called()
+
+  @mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor.stop_rolling_window_goodput_uploader")
+  @mock.patch(
+      "ml_goodput_measurement.monitoring.GoodputMonitor.start_rolling_window_goodput_uploader",
+      side_effect=RuntimeError("start boom"),
+  )
+  @mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor.stop_goodput_uploader")
+  @mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor.start_goodput_uploader")
+  def test_monitor_rolling_window_goodput_start_failure_degrades_gracefully(
+      self, unused_mock_start_goodput, mock_stop_goodput, mock_start_rolling, mock_stop_rolling
+  ):
+    """A failing rolling window start is swallowed; the body still runs and nothing is stopped twice."""
+    config = _ConfigOverride(self.config, enable_rolling_window_goodput=True)
+    body_ran = False
+
+    with maybe_monitor_goodput(config):  # Must not raise.
+      body_ran = True
+
+    self.assertTrue(body_ran)
+    mock_start_rolling.assert_called_once()
+    mock_stop_rolling.assert_not_called()  # Never started, so never stopped.
+    mock_stop_goodput.assert_called_once()  # Cumulative monitoring is unaffected.
+
+  @mock.patch(
+      "ml_goodput_measurement.monitoring.GoodputMonitor.stop_rolling_window_goodput_uploader",
+      side_effect=RuntimeError("stop boom"),
+  )
+  @mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor.start_rolling_window_goodput_uploader")
+  @mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor.stop_goodput_uploader")
+  @mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor.start_goodput_uploader")
+  def test_monitor_rolling_window_goodput_stop_failure_degrades_gracefully(
+      self, unused_mock_start_goodput, mock_stop_goodput, unused_mock_start_rolling, mock_stop_rolling
+  ):
+    """A failing rolling window stop is swallowed and the cumulative uploader is still stopped."""
+    config = _ConfigOverride(self.config, enable_rolling_window_goodput=True)
+
+    with maybe_monitor_goodput(config):  # Must not raise on exit.
+      pass
+
+    mock_stop_rolling.assert_called_once()
+    mock_stop_goodput.assert_called_once()
+
+  @mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor.stop_rolling_window_goodput_uploader")
+  @mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor.start_rolling_window_goodput_uploader")
+  @mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor.stop_goodput_uploader")
+  @mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor.start_goodput_uploader")
+  def test_monitor_rolling_window_goodput_propagates_body_exception(
+      self, unused_mock_start_goodput, mock_stop_goodput, unused_mock_start_rolling, mock_stop_rolling
+  ):
+    """Graceful degradation must not swallow errors raised by training itself."""
+    config = _ConfigOverride(self.config, enable_rolling_window_goodput=True)
+
+    class TrainingError(Exception):
+      pass
+
+    with self.assertRaises(TrainingError):
+      with maybe_monitor_goodput(config):
+        raise TrainingError()
+
+    mock_stop_rolling.assert_called_once()
+    mock_stop_goodput.assert_called_once()
+
+  def test_rolling_windows_seconds_rejects_empty_list(self):
+    """An empty window list would spawn an idle uploader process, so config validation rejects it."""
+    with self.assertRaisesRegex(Exception, "rolling_windows_seconds"):
+      pyconfig.initialize(
+          [None, get_test_config_path()],
+          base_output_directory=get_test_base_output_directory(),
+          run_name="runner_test",
+          enable_checkpointing=False,
+          rolling_windows_seconds=[],
+      )
+
   def test_job_recording_constants(self):
     """Constants must map to the recorder method names."""
     self.assertEqual(RECORD_JOB_START_TIME, "record_job_start_time")
@@ -183,11 +295,13 @@ def test_construct_goodput_monitor_elastic_construction_failure_falls_back(
     # here - that name is itself patched to a Mock for this test, so it isn't a usable type.
     self.assertIsInstance(monitor, monitoring.GoodputMonitor)
 
+  @mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor.stop_rolling_window_goodput_uploader")
+  @mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor.start_rolling_window_goodput_uploader")
   @mock.patch("ml_goodput_measurement.monitoring_elastic.ElasticGoodputMonitor.stop_goodput_uploader")
   @mock.patch("ml_goodput_measurement.monitoring_elastic.ElasticGoodputMonitor.start_goodput_uploader")
   @mock.patch("maxtext.utils.elastic_utils.should_use_elastic")
   def test_monitor_goodput_elastic(
-      self, mock_should_use_elastic, mock_start_goodput_uploader, mock_stop_goodput_uploader
+      self, mock_should_use_elastic, mock_start_goodput_uploader, mock_stop_goodput_uploader, *unused_rolling_mocks
   ):
     """maybe_monitor_goodput actually starts/stops an ElasticGoodputMonitor when elastic is active."""
     mock_should_use_elastic.return_value = True