Add a flag to measurement.py to disable goodput monitoring by default

GitOrigin-RevId: e701de09e8b02bc95d21d2fed9bd445ed2d195f9

Author: ampolloreno

axlearn/cloud/gcp/measurement.py

@@ -29,7 +29,7 @@
 from ml_goodput_measurement import goodput
 from ml_goodput_measurement import monitoring as goodput_monitoring
 
-from axlearn.cloud.common.utils import parse_kv_flags
+from axlearn.cloud.common.utils import parse_kv_flags, to_bool
 from axlearn.common import measurement
 from axlearn.common.config import REQUIRED, Required, config_class, maybe_set_config
 
@@ -52,12 +52,17 @@ class Config(measurement.Recorder.Config):
                 See "How to Monitor Rolling Window Goodput Metrics" in
                 docs/05-Goodput-Monitoring.md for more details.
             jax_backend: Jax backend type to infer Pathways environment.
+            enable_monitoring: Whether to enable goodput monitoring/uploading.
         """
 
         upload_dir: Required[str] = REQUIRED
         upload_interval: Required[int] = REQUIRED
         rolling_window_size: Sequence[int] = []
         jax_backend: Optional[str] = None
+        # Disabled by default because of performance degradation. This doesn't disable goodput
+        # recording.
+        # TODO (apolloreno): once the performance degradation is fixed, will change default to True
+        enable_monitoring: bool = False
 
     @classmethod
     def from_flags(cls, fv: flags.FlagValues) -> "GoodputRecorder":
@@ -72,6 +77,7 @@ def from_flags(cls, fv: flags.FlagValues) -> "GoodputRecorder":
         - rolling_window_size: Comma-separated list of integers representing rolling window
             sizes in seconds.
         - jax_backend: The type of jax backend.
+        - enable_monitoring: Boolean to enable/disable goodput monitoring (default: false).
         """
         cfg: measurement.Recorder.Config = cls.default_config()
         parsed_flags = parse_kv_flags(fv.recorder_spec, delimiter="=")
@@ -83,6 +89,8 @@ def from_flags(cls, fv: flags.FlagValues) -> "GoodputRecorder":
             parsed_flags["rolling_window_size"] = [
                 int(x) for x in parsed_flags["rolling_window_size"].split(",")
             ]
+        if "enable_monitoring" in parsed_flags:
+            parsed_flags["enable_monitoring"] = to_bool(parsed_flags["enable_monitoring"])
         return maybe_set_config(cfg, **parsed_flags).instantiate()
 
     def __init__(self, cfg):
@@ -149,7 +157,7 @@ def _maybe_monitor_goodput(self, *args, **kwargs):
         Default behavior is to push metrics to Google Cloud Monitoring.
         This behavior can be overridden by configuring `goodput_monitoring.GCPOptions`
         """
-        if jax.process_index() != 0:
+        if not self.config.enable_monitoring or jax.process_index() != 0:
             yield
             return
         try:
@@ -175,7 +183,11 @@ def _maybe_monitor_goodput(self, *args, **kwargs):
     @contextlib.contextmanager
     def _maybe_monitor_rolling_window_goodput(self):
         """Monitor rolling window goodput if enabled."""
-        if not self.config.rolling_window_size or jax.process_index() != 0:
+        if (
+            not self.config.enable_monitoring
+            or not self.config.rolling_window_size
+            or jax.process_index() != 0
+        ):
             yield
             return
         try:

axlearn/cloud/gcp/measurement_test.py

@@ -37,12 +37,35 @@ class GoodputRecorderTest(parameterized.TestCase):
             expected_rolling_window_size=[1, 2, 3],
             expected_jax_backend="proxy",
         ),
+        dict(
+            recorder_spec=[
+                "name=test-name",
+                "upload_dir=/test/path",
+                "upload_interval=15",
+                "enable_monitoring=true",
+            ],
+            expected_rolling_window_size=[],
+            expected_jax_backend=None,
+            expected_enable_monitoring=True,
+        ),
+        dict(
+            recorder_spec=[
+                "name=test-name",
+                "upload_dir=/test/path",
+                "upload_interval=15",
+                "enable_monitoring=false",
+            ],
+            expected_rolling_window_size=[],
+            expected_jax_backend=None,
+            expected_enable_monitoring=False,
+        ),
     )
     def test_from_flags(
         self,
         recorder_spec,
         expected_rolling_window_size,
         expected_jax_backend,
+        expected_enable_monitoring=False,
     ):
         """Tests that flags are correctly parsed into the config."""
         mock_fv = mock.MagicMock(spec=flags.FlagValues)
@@ -56,6 +79,7 @@ def test_from_flags(
         self.assertEqual(15, recorder.config.upload_interval)
         self.assertEqual(expected_rolling_window_size, recorder.config.rolling_window_size)
         self.assertEqual(expected_jax_backend, recorder.config.jax_backend)
+        self.assertEqual(expected_enable_monitoring, recorder.config.enable_monitoring)
 
     def test_from_flags_missing_required(self):
         """Tests that missing required flags raise an error."""
@@ -188,6 +212,7 @@ def test_maybe_monitor_goodput(self, _, is_pathways_job, mock_jax_backend):
             upload_dir="/test",
             upload_interval=30,
             jax_backend=mock_jax_backend,
+            enable_monitoring=True,
         )
         recorder = GoodputRecorder(cfg)
 
@@ -245,6 +270,7 @@ def test_maybe_monitor_rolling_window(
             upload_interval=30,
             rolling_window_size=rolling_window_size,
             jax_backend=mock_jax_backend,
+            enable_monitoring=True,
         )
         recorder = GoodputRecorder(cfg)
 
@@ -279,7 +305,7 @@ def test_non_zero_process_index_skips_monitoring(
     ):  # pylint: disable=unused-argument
         """Tests that monitoring is skipped on non-zero process indices."""
         cfg = GoodputRecorder.default_config().set(
-            name="test", upload_dir="/test", upload_interval=30
+            name="test", upload_dir="/test", upload_interval=30, enable_monitoring=True
         )
         recorder = GoodputRecorder(cfg)
 
@@ -294,6 +320,7 @@ def test_non_zero_process_index_skips_monitoring(
                 upload_dir="/test",
                 upload_interval=30,
                 rolling_window_size=[10, 20],
+                enable_monitoring=True,
             )
             recorder_rolling = GoodputRecorder(cfg_rolling)
             with recorder_rolling._maybe_monitor_rolling_window_goodput():
@@ -347,6 +374,7 @@ def test_maybe_monitor_all(
             upload_interval=30,
             rolling_window_size=rolling_window_size,
             jax_backend=jax_backend,
+            enable_monitoring=True,
         )
         recorder = GoodputRecorder(cfg)
 
@@ -373,3 +401,118 @@ def test_maybe_monitor_all(
             else:
                 mock_monitor_instance.start_rolling_window_goodput_uploader.assert_not_called()
                 mock_monitor_instance.stop_rolling_window_goodput_uploader.assert_not_called()
+
+    @mock.patch("jax.process_index", return_value=0)
+    def test_enable_monitoring_disabled_by_default(self, _):
+        """Tests that monitoring is disabled by default (enable_monitoring=False)."""
+        cfg = GoodputRecorder.default_config().set(
+            name="test-disabled",
+            upload_dir="/test",
+            upload_interval=30,
+            # enable_monitoring defaults to False
+            rolling_window_size=[10, 20],
+        )
+        recorder = GoodputRecorder(cfg)
+
+        # Verify the flag defaults to False
+        self.assertFalse(recorder.config.enable_monitoring)
+
+        with mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor") as mock_monitor_cls:
+            # Test that cumulative goodput monitoring is skipped
+            with recorder._maybe_monitor_goodput():
+                pass
+            mock_monitor_cls.assert_not_called()
+
+            # Test that rolling window monitoring is skipped
+            with recorder._maybe_monitor_rolling_window_goodput():
+                pass
+            mock_monitor_cls.assert_not_called()
+
+            # Test that maybe_monitor_all is skipped
+            with recorder.maybe_monitor_all():
+                pass
+            mock_monitor_cls.assert_not_called()
+
+    @mock.patch("jax.process_index", return_value=0)
+    def test_enable_monitoring_explicitly_disabled(self, _):
+        """Tests that monitoring is disabled when enable_monitoring=False."""
+        cfg = GoodputRecorder.default_config().set(
+            name="test-explicitly-disabled",
+            upload_dir="/test",
+            upload_interval=30,
+            enable_monitoring=False,  # Explicitly disabled
+            rolling_window_size=[10, 20],
+        )
+        recorder = GoodputRecorder(cfg)
+
+        with mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor") as mock_monitor_cls:
+            # Test cumulative goodput monitoring is skipped
+            with recorder._maybe_monitor_goodput():
+                pass
+            mock_monitor_cls.assert_not_called()
+
+            # Test rolling window monitoring is skipped
+            with recorder._maybe_monitor_rolling_window_goodput():
+                pass
+            mock_monitor_cls.assert_not_called()
+
+            # Test maybe_monitor_all is skipped
+            with recorder.maybe_monitor_all():
+                pass
+            mock_monitor_cls.assert_not_called()
+
+    @mock.patch("jax.process_index", return_value=0)
+    def test_enable_monitoring_explicitly_enabled(self, _):
+        """Tests that monitoring works when enable_monitoring=True."""
+        cfg = GoodputRecorder.default_config().set(
+            name="test-enabled",
+            upload_dir="/test",
+            upload_interval=30,
+            enable_monitoring=True,  # Explicitly enabled
+            rolling_window_size=[10, 20],
+        )
+        recorder = GoodputRecorder(cfg)
+
+        with mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor") as mock_monitor_cls:
+            mock_monitor_instance = mock_monitor_cls.return_value
+
+            # Test cumulative goodput monitoring works
+            with recorder._maybe_monitor_goodput():
+                pass
+
+            # Should be called once for cumulative monitoring
+            self.assertEqual(mock_monitor_cls.call_count, 1)
+            mock_monitor_instance.start_goodput_uploader.assert_called_once()
+            mock_monitor_instance.stop_goodput_uploader.assert_called_once()
+
+    @mock.patch("jax.process_index", return_value=0)
+    def test_record_event_works_with_monitoring_disabled(self, _):
+        """Tests that record_event still works when monitoring is disabled."""
+        cfg = GoodputRecorder.default_config().set(
+            name="test-recording-only",
+            upload_dir="/test",
+            upload_interval=30,
+            enable_monitoring=False,  # Monitoring disabled
+        )
+        recorder = GoodputRecorder(cfg)
+
+        # Verify that goodput recording still works (not monitoring/uploading)
+        with mock.patch("ml_goodput_measurement.goodput.GoodputRecorder") as mock_recorder_cls:
+            mock_instance = mock_recorder_cls.return_value
+            mock_instance.record_job_start_time = mock.MagicMock()
+            mock_instance.record_job_end_time = mock.MagicMock()
+
+            # Record event should work
+            with recorder.record_event(measurement.EventType.JOB):
+                pass
+
+            # Verify goodput recording happened
+            mock_recorder_cls.assert_called_once()
+            mock_instance.record_job_start_time.assert_called_once()
+            mock_instance.record_job_end_time.assert_called_once()
+
+        # Verify no monitoring/uploading happened
+        with mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor") as mock_monitor_cls:
+            with recorder.maybe_monitor_all():
+                pass
+            mock_monitor_cls.assert_not_called()