Integrate AXLearn with Goodput v12

dipannita08 · SujeethJinesh · commit 563f9e62200c · 2025-07-11T20:09:13.000Z
diff --git a/axlearn/cloud/gcp/measurement.py b/axlearn/cloud/gcp/measurement.py
@@ -4,7 +4,7 @@
 
     Example:
 
-    # Enable Goodput when launching an AxLearn training job
+    # Enable Goodput when launching an AXLearn training job
     axlearn gcp launch run --instance_type=tpu-v5litepod-16 \
         --bundler_type=artifactregistry --bundler_spec=image=tpu \
         --bundler_spec=dockerfile=Dockerfile \
@@ -13,7 +13,6 @@
         --recorder_spec=name=my-run-with-goodput \
         --recorder_spec=upload_dir=my-output-directory/summaries \
         --recorder_spec=upload_interval=30 \
-        --recorder_spec=enable_rolling_window_goodput_monitoring=True \
         --recorder_spec=rolling_window_size=86400,259200,432000
 
 """
@@ -43,22 +42,13 @@ class Config(measurement.Recorder.Config):
         Attributes:
             upload_dir: Directory to store metrics for the monitor.
             upload_interval: Time interval (seconds) for monitoring uploads.
-            enable_gcp_goodput_metrics: Whether to upload metrics to Google Cloud Monitoring.
-            enable_pathways_goodput: Whether to enable goodput calculations specific to Pathways.
-            include_badput_breakdown: Whether to include a detailed breakdown of badput sources.
-            enable_rolling_window_goodput_monitoring: Enables goodput/badput monitoring over
-                rolling time windows.
             rolling_window_size: A sequence of integers defining the rolling window sizes in
                 seconds.
         """
 
         upload_dir: Required[str] = REQUIRED
         upload_interval: Required[int] = REQUIRED
-        enable_gcp_goodput_metrics: bool = True
-        enable_pathways_goodput: bool = False
-        include_badput_breakdown: bool = True
-        enable_rolling_window_goodput_monitoring: bool = False
-        rolling_window_size: Sequence[int] = ()
+        rolling_window_size: Sequence[int] = []
 
     @classmethod
     def from_flags(cls, fv: flags.FlagValues) -> "GoodputRecorder":
@@ -70,15 +60,8 @@ def from_flags(cls, fv: flags.FlagValues) -> "GoodputRecorder":
         - upload_dir: The directory to write Tensorboard data to.
         - upload_interval: The time interval in seconds at which to query and upload data
             to Tensorboard.
-        - enable_rolling_window_goodput_monitoring: Whether to enable rolling window Goodput
-            monitoring.
         - rolling_window_size: Comma-separated list of integers representing rolling window
             sizes in seconds.
-        - enable_gcp_goodput_metrics: Whether to upload Goodput metrics to GCM.
-        - enable_pathways_goodput: Whether to enable Pathways-specific Goodput
-        calculations.
-        - include_badput_breakdown: Whether to include a detailed breakdown of
-            badput events in the monitoring.
         """
         cfg: measurement.Recorder.Config = cls.default_config()
         parsed_flags = parse_kv_flags(fv.recorder_spec, delimiter="=")
@@ -90,8 +73,7 @@ def from_flags(cls, fv: flags.FlagValues) -> "GoodputRecorder":
             parsed_flags["rolling_window_size"] = [
                 int(x) for x in parsed_flags["rolling_window_size"].split(",")
             ]
-        cfg = maybe_set_config(cfg, **parsed_flags)
-        return cfg.instantiate()
+        return maybe_set_config(cfg, **parsed_flags).instantiate()
 
     def __init__(self, cfg):
         super().__init__(cfg)
@@ -121,9 +103,9 @@ def record_event(self, event: measurement.Event, *args, **kwargs):
         try:
             if record_event_start:
                 record_event_start(*args, **kwargs)
-        except (TypeError, ValueError, RuntimeError) as e:
+        except RuntimeError as e:
             logging.warning(
-                "Failed to record start of event %s. Error: %s", event.name, e, exc_info=True
+                "Failed to record start of event %s. Error: %s", event.value, e, exc_info=True
             )
 
         try:
@@ -132,9 +114,9 @@ def record_event(self, event: measurement.Event, *args, **kwargs):
             try:
                 if record_event_end:
                     record_event_end(*args, **kwargs)
-            except (TypeError, ValueError, RuntimeError) as e:
+            except RuntimeError as e:
                 logging.warning(
-                    "Failed to record end of event %s. Error: %s", event.name, e, exc_info=True
+                    "Failed to record end of event %s. Error: %s", event.value, e, exc_info=True
                 )
 
     @contextlib.contextmanager
@@ -157,18 +139,17 @@ def maybe_monitor_goodput(self, *args, **kwargs):
             return
         try:
             if self._monitor is None:
-                gcp_options = goodput_monitoring.GCPOptions(
-                    enable_gcp_goodput_metrics=self.config.enable_gcp_goodput_metrics,
+                pathways_enabled = (
+                    hasattr(flags.FLAGS, "jax_backend") and flags.FLAGS.jax_backend == "proxy"
                 )
                 self._monitor = goodput_monitoring.GoodputMonitor(
                     job_name=self.config.name,
                     logger_name=f"goodput_logger_{self.config.name}",
                     tensorboard_dir=self.config.upload_dir,
                     upload_interval=self.config.upload_interval,
                     monitoring_enabled=True,
-                    pathway_enabled=self.config.enable_pathways_goodput,
-                    include_badput_breakdown=self.config.include_badput_breakdown,
-                    gcp_options=gcp_options,
+                    pathway_enabled=pathways_enabled,
+                    include_badput_breakdown=True,
                 )
 
             self._monitor.start_goodput_uploader(*args, **kwargs)
@@ -182,11 +163,14 @@ 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.enable_rolling_window_goodput_monitoring or jax.process_index() != 0:
+        if not self.config.rolling_window_size or jax.process_index() != 0:
             yield
             return
         try:
             if self._rolling_window_monitor is None:
+                pathways_enabled = (
+                    hasattr(flags.FLAGS, "jax_backend") and flags.FLAGS.jax_backend == "proxy"
+                )
                 rolling_window_tensorboard_dir = os.path.join(
                     self.config.upload_dir, f"rolling_window_{self.config.name}"
                 )
@@ -196,7 +180,7 @@ def maybe_monitor_rolling_window_goodput(self):
                     tensorboard_dir=rolling_window_tensorboard_dir,
                     upload_interval=self.config.upload_interval,
                     monitoring_enabled=True,
-                    pathway_enabled=self.config.enable_pathways_goodput,
+                    pathway_enabled=pathways_enabled,
                     include_badput_breakdown=True,
                 )
             self._rolling_window_monitor.start_rolling_window_goodput_uploader(
diff --git a/axlearn/cloud/gcp/measurement_test.py b/axlearn/cloud/gcp/measurement_test.py
@@ -16,59 +16,90 @@
 class GoodputRecorderTest(parameterized.TestCase):
     """Tests GoodputRecorder."""
 
-    def test_from_flags_with_spec(self):
-        """Tests that flags are correctly parsed."""
-        fv = flags.FlagValues()
-        measurement.define_flags(flag_values=fv)
-        fv.set_default(
-            "recorder_spec",
-            [
+    def setUp(self):
+        super().setUp()
+        self.mock_flags = mock.MagicMock(spec=flags.FlagValues)
+        self.mock_flags.jax_backend = None
+        self.patch_flags_global = mock.patch.object(flags, "FLAGS", new=self.mock_flags)
+        self.patch_flags_global.start()
+
+    def tearDown(self):
+        super().tearDown()
+        self.patch_flags_global.stop()
+
+    @parameterized.parameters(
+        dict(
+            recorder_spec=[
+                "name=test-name",
+                "upload_dir=/test/path",
+                "upload_interval=15",
+            ],
+            expected_rolling_window_size=[],
+        ),
+        dict(
+            recorder_spec=[
                 "name=test-name",
                 "upload_dir=/test/path",
                 "upload_interval=15",
-                "enable_rolling_window_goodput_monitoring=True",
                 "rolling_window_size=1,2,3",
             ],
-        )
-        fv.mark_as_parsed()
-        recorder = GoodputRecorder.from_flags(fv)
+            expected_rolling_window_size=[1, 2, 3],
+        ),
+    )
+    def test_from_flags(
+        self,
+        recorder_spec,
+        expected_rolling_window_size,
+    ):
+        """Tests that flags are correctly parsed into the config."""
+        mock_fv = mock.MagicMock(spec=flags.FlagValues)
+        mock_fv.recorder_spec = recorder_spec
+        mock_fv.jax_backend = "tpu"
+
+        recorder = GoodputRecorder.from_flags(mock_fv)
+
         self.assertEqual("test-name", recorder.config.name)
         self.assertEqual("/test/path", recorder.config.upload_dir)
         self.assertEqual(15, recorder.config.upload_interval)
-        self.assertTrue(recorder.config.enable_rolling_window_goodput_monitoring)
-        self.assertEqual([1, 2, 3], recorder.config.rolling_window_size)
+        self.assertEqual(expected_rolling_window_size, recorder.config.rolling_window_size)
 
     def test_from_flags_missing_required(self):
         """Tests that missing required flags raise an error."""
-        fv = flags.FlagValues()
-        measurement.define_flags(flag_values=fv)
-        fv.set_default("recorder_spec", ["name=test-name"])  # Missing upload_dir/interval
-        fv.mark_as_parsed()
+        mock_fv = mock.MagicMock(spec=flags.FlagValues)
+        mock_fv.recorder_spec = ["name=test-name"]  # Missing upload_dir/interval
+        mock_fv.jax_backend = "tpu"
         with self.assertRaisesRegex(RequiredFieldMissingError, "upload_dir"):
-            GoodputRecorder.from_flags(fv)
+            GoodputRecorder.from_flags(mock_fv)
 
     @mock.patch("jax.process_index", return_value=0)
     def test_record_event_context_manager(self, _):
         """Tests the record_event context manager."""
-        recorder = GoodputRecorder(GoodputRecorder.default_config().set(name="test"))
-        with mock.patch("ml_goodput_measurement.goodput.GoodputRecorder") as mock_recorder:
-            mock_instance = mock_recorder.return_value
+        cfg = GoodputRecorder.default_config().set(
+            name="test",
+            upload_dir="/tmp/test",
+            upload_interval=1,
+        )
+        recorder = GoodputRecorder(cfg)
+        with mock.patch("ml_goodput_measurement.goodput.GoodputRecorder") as mock_recorder_cls:
+            mock_instance = mock_recorder_cls.return_value
             with recorder.record_event(measurement.Event.JOB):
                 pass
-            mock_recorder.assert_called_once()
+            mock_recorder_cls.assert_called_once()
             mock_instance.record_job_start_time.assert_called_once()
             mock_instance.record_job_end_time.assert_called_once()
 
+    @parameterized.parameters(
+        dict(is_pathways_job=False, mock_jax_backend="tpu"),
+        dict(is_pathways_job=True, mock_jax_backend="proxy"),
+    )
     @mock.patch("jax.process_index", return_value=0)
-    def test_maybe_monitor_goodput(self, _):
+    def test_maybe_monitor_goodput(self, _, is_pathways_job, mock_jax_backend):
         """Tests the maybe_monitor_goodput context manager."""
+        self.mock_flags.jax_backend = mock_jax_backend
         cfg = GoodputRecorder.default_config().set(
             name="test-monitor",
             upload_dir="/test",
             upload_interval=30,
-            enable_gcp_goodput_metrics=True,
-            enable_pathways_goodput=False,
-            include_badput_breakdown=True,
         )
         recorder = GoodputRecorder(cfg)
 
@@ -84,47 +115,82 @@ def test_maybe_monitor_goodput(self, _):
                 tensorboard_dir="/test",
                 upload_interval=30,
                 monitoring_enabled=True,
-                pathway_enabled=False,
+                pathway_enabled=is_pathways_job,
                 include_badput_breakdown=True,
-                gcp_options=mock.ANY,
             )
-            # Verify the start and stop methods were called.
             mock_monitor_instance.start_goodput_uploader.assert_called_once()
             mock_monitor_instance.stop_goodput_uploader.assert_called_once()
 
+    @parameterized.parameters(
+        dict(
+            is_rolling_window_enabled=True,
+            rolling_window_size=[10, 20],
+            is_pathways_job=False,
+            mock_jax_backend="tpu",
+        ),
+        dict(
+            is_rolling_window_enabled=False,
+            rolling_window_size=[],
+            is_pathways_job=False,
+            mock_jax_backend="tpu",
+        ),
+        dict(
+            is_rolling_window_enabled=True,
+            rolling_window_size=[50],
+            is_pathways_job=True,
+            mock_jax_backend="proxy",
+        ),
+    )
     @mock.patch("jax.process_index", return_value=0)
-    def test_maybe_monitor_rolling_window(self, _):
+    def test_maybe_monitor_rolling_window(
+        self,
+        mock_process_index,
+        is_rolling_window_enabled,
+        rolling_window_size,
+        is_pathways_job,
+        mock_jax_backend,
+    ):  # pylint: disable=unused-argument
         """Tests the rolling window monitoring context manager."""
+        self.mock_flags.jax_backend = mock_jax_backend
         cfg = GoodputRecorder.default_config().set(
             name="test-rolling",
             upload_dir="/test",
             upload_interval=30,
-            enable_rolling_window_goodput_monitoring=True,
-            rolling_window_size=[10, 20],
+            rolling_window_size=rolling_window_size,
         )
         recorder = GoodputRecorder(cfg)
 
         with mock.patch("ml_goodput_measurement.monitoring.GoodputMonitor") as mock_monitor_cls:
             mock_monitor_instance = mock_monitor_cls.return_value
+            if not is_rolling_window_enabled:
+                with recorder.maybe_monitor_rolling_window_goodput():
+                    pass
+                mock_monitor_cls.assert_not_called()
+                return
             with recorder.maybe_monitor_rolling_window_goodput():
                 pass
 
-            # Verify that GoodputMonitor was instantiated for rolling window.
-            mock_monitor_cls.assert_called_once()
-            self.assertEqual(
-                "/test/rolling_window_test-rolling",
-                mock_monitor_cls.call_args.kwargs["tensorboard_dir"],
+            mock_monitor_cls.assert_called_once_with(
+                job_name="test-rolling",
+                logger_name="goodput_logger_test-rolling",
+                tensorboard_dir="/test/rolling_window_test-rolling",
+                upload_interval=30,
+                monitoring_enabled=True,
+                pathway_enabled=is_pathways_job,
+                include_badput_breakdown=True,
             )
 
-            # Verify the correct the start and stop methods were called.
-            mock_monitor_instance.start_rolling_window_goodput_uploader.assert_called_with([10, 20])
+            mock_monitor_instance.start_rolling_window_goodput_uploader.assert_called_with(
+                rolling_window_size
+            )
             mock_monitor_instance.stop_rolling_window_goodput_uploader.assert_called_once()
 
     @mock.patch("jax.process_index", return_value=1)
     def test_non_zero_process_index_skips_monitoring(
         self, mock_process_index
     ):  # pylint: disable=unused-argument
         """Tests that monitoring is skipped on non-zero process indices."""
+        self.mock_flags.jax_backend = "tpu"
         cfg = GoodputRecorder.default_config().set(
             name="test", upload_dir="/test", upload_interval=30
         )
@@ -136,8 +202,13 @@ def test_non_zero_process_index_skips_monitoring(
                 pass
             mock_monitor_cls.assert_not_called()
 
-            # Test maybe_monitor_rolling_window_goodput
-            recorder.config.enable_rolling_window_goodput_monitoring = True
-            with recorder.maybe_monitor_rolling_window_goodput():
+            cfg_rolling = GoodputRecorder.default_config().set(
+                name="test-rolling-skip",
+                upload_dir="/test",
+                upload_interval=30,
+                rolling_window_size=[10, 20],
+            )
+            recorder_rolling = GoodputRecorder(cfg_rolling)
+            with recorder_rolling.maybe_monitor_rolling_window_goodput():
                 pass
             mock_monitor_cls.assert_not_called()
diff --git a/docs/05-Goodput-Monitoring.md b/docs/05-Goodput-Monitoring.md
@@ -1,5 +1,5 @@
 # ML Goodput Monitoring
-AxLearn supports automatic measurement and upload of a wide range of workload
+AXLearn supports automatic measurement and upload of a wide range of workload
 metrics using the **ML Goodput Measurement** library. This includes:
 * **Goodput** and **Badput Breakdown**
 * **Step Metrics** (Ideal Step Time, Step Time Deviation, Last Productive Step etc.)
@@ -97,7 +97,7 @@ Goodput/Badput metrics of a previous workload along with your current workload.
 
 ### How to Monitor Cumulative Goodput Metrics
 
-To enable Goodput recording and monitoring on AxLearn, follow the example below.
+To enable Goodput recording and monitoring on AXLearn, follow the example below.
 
 
 ```bash
@@ -133,11 +133,11 @@ axlearn gcp launch run --instance_type=tpu-v5litepod-16 \
 ### Visualize on Tensorboard
 
 1. Requires packages: `tensorboard-plugin-profile`, `tensorflow` and `tensorboard`.
-2. Use the Tensorboard URL on AxLearn logs to view all metrics in one location.
+2. Use the Tensorboard URL on AXLearn logs to view all metrics in one location.
 
 ### Enabling Google Cloud Monitoring
 
-By default, when Goodput monitoring is enabled via the recorder, AxLearn automatically pushes metrics to Google Cloud Monitoring.
+By default, when Goodput monitoring is enabled via the recorder, AXLearn automatically pushes metrics to Google Cloud Monitoring.
 
 -   **Cumulative Metrics** are enabled by default when you specify the `recorder_type`.
     To disable this, you would need to set `enable_gcp_goodput_metrics` to `False` in
