From a4f335eeba4047ccf387c822c5a60f1e0ef7d4f7 Mon Sep 17 00:00:00 2001 From: Yin Date: Fri, 4 Sep 2026 23:00:13 +0000 Subject: [PATCH] Use sliding window to track goodput metrics over time. This will allow for more accurate measurement of performance and help identify trends in data throughput. --- src/maxtext/common/gcloud_stub.py | 6 ++ src/maxtext/common/goodput.py | 36 ++++++++- src/maxtext/configs/base.yml | 2 + src/maxtext/configs/types.py | 6 ++ tests/unit/goodput_utils_test.py | 118 +++++++++++++++++++++++++++++- 5 files changed, 163 insertions(+), 5 deletions(-) diff --git a/src/maxtext/common/gcloud_stub.py b/src/maxtext/common/gcloud_stub.py index 52c39a8abb..cec3ccb43f 100644 --- a/src/maxtext/common/gcloud_stub.py +++ b/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 diff --git a/src/maxtext/common/goodput.py b/src/maxtext/common/goodput.py index f40c0ef42a..2ba5fd23ca 100644 --- a/src/maxtext/common/goodput.py +++ b/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 diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 61b3f19869..ba2d6ca4f6 100644 --- a/src/maxtext/configs/base.yml +++ b/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 diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 380d24f4ab..3fa3668e74 100644 --- a/src/maxtext/configs/types.py +++ b/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): diff --git a/tests/unit/goodput_utils_test.py b/tests/unit/goodput_utils_test.py index ceca74bb07..e17398b861 100644 --- a/tests/unit/goodput_utils_test.py +++ b/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