Add workload hang monitoring & rolling window goodput support

Author: dipannita08

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,10 +13,15 @@
         --recorder_spec=name=my-run-with-goodput \
         --recorder_spec=upload_dir=my-output-directory/summaries \
         --recorder_spec=upload_interval=30 \
-        --recorder_spec=step_deviation_interval_seconds=30
+        --recorder_spec=enable_rolling_window_goodput_monitoring=True \
+        --recorder_spec=rolling_window_size=86400,259200,432000
 
 """
 
+import contextlib
+import os
+from typing import Optional, Sequence
+
 import jax
 from absl import flags, logging
 from ml_goodput_measurement import goodput
@@ -38,13 +43,22 @@ class Config(measurement.Recorder.Config):
         Attributes:
             upload_dir: Directory to store metrics for the monitor.
             upload_interval: Time interval (seconds) for monitoring uploads.
-            step_deviation_interval_seconds: Time interval (seconds) for step deviation metrics
-            uploads. -1 to disable step deviation 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
-        step_deviation_interval_seconds: int = 30  # Default to 30 seconds
+        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] = ()
 
     @classmethod
     def from_flags(cls, fv: flags.FlagValues) -> "GoodputRecorder":
@@ -53,68 +67,83 @@ def from_flags(cls, fv: flags.FlagValues) -> "GoodputRecorder":
         `fv.recorder_spec` will be interpreted as a list of `key=value` pairs; config names
         corresponding to keys will be set to the corresponding values. A GoodputRecorder can
         additionally take in following Tensorboard configs in the recorder_spec:
-         - 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.
-        - step_deviation_interval_seconds: Time interval (seconds) for step deviation metrics
-        uploads. Set to less than or equal to 0 to disable step deviation uploads.
+        - 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()
-        cfg = maybe_set_config(cfg, **parse_kv_flags(fv.recorder_spec, delimiter="="))
+        parsed_flags = parse_kv_flags(fv.recorder_spec, delimiter="=")
+        if "upload_interval" in parsed_flags:
+            parsed_flags["upload_interval"] = int(parsed_flags["upload_interval"])
+        if "rolling_window_size" in parsed_flags and isinstance(
+            parsed_flags["rolling_window_size"], str
+        ):
+            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()
 
     def __init__(self, cfg):
         super().__init__(cfg)
-        cfg: GoodputRecorder.Config = self.config
-        self._recorder = None
-        self._monitor = None
-
-    def record(self, event: measurement.Event, *args, **kwargs):
-        # Lazily instantiate the recorder. This avoids invoking jax before setup is complete.
+        self._recorder: Optional[goodput.GoodputRecorder] = None
+        self._monitor: Optional[goodput_monitoring.GoodputMonitor] = None
+        self._rolling_window_monitor: Optional[goodput_monitoring.GoodputMonitor] = None
+
+    @contextlib.contextmanager
+    def record_event(self, event: measurement.Event, *args, **kwargs):
+        """Records a goodput event using a context manager."""
+        # Lazily instantiate the recorder if it hasn't been already.
         if self._recorder is None:
-            cfg: GoodputRecorder.Config = self.config
+            if jax.process_index() == 0:
+                logging.info("Lazily instantiating goodput recorder.")
             self._recorder = goodput.GoodputRecorder(
-                job_name=cfg.name,
-                logger_name=f"goodput_logger_{cfg.name}",
+                job_name=self.config.name,
+                logger_name=f"goodput_logger_{self.config.name}",
                 logging_enabled=(jax.process_index() == 0),
             )
 
-        if event == measurement.Event.START_JOB:
-            self._recorder.record_job_start_time(*args, **kwargs)
-        elif event == measurement.Event.END_JOB:
-            self._recorder.record_job_end_time(*args, **kwargs)
-        elif event == measurement.Event.START_STEP:
-            self._recorder.record_step_start_time(*args, **kwargs)
-        elif event == measurement.Event.START_ACCELERATOR_INIT:
-            self._recorder.record_tpu_init_start_time(*args, **kwargs)
-        elif event == measurement.Event.END_ACCELERATOR_INIT:
-            self._recorder.record_tpu_init_end_time(*args, **kwargs)
-        elif event == measurement.Event.START_TRAINING_PREPARATION:
-            self._recorder.record_training_preparation_start_time(*args, **kwargs)
-        elif event == measurement.Event.END_TRAINING_PREPARATION:
-            self._recorder.record_training_preparation_end_time(*args, **kwargs)
-        elif event == measurement.Event.START_DATA_LOADING:
-            self._recorder.record_data_loading_start_time(*args, **kwargs)
-        elif event == measurement.Event.END_DATA_LOADING:
-            self._recorder.record_data_loading_end_time(*args, **kwargs)
-        elif event == measurement.Event.START_CUSTOM_BADPUT_EVENT:
-            self._recorder.record_custom_badput_event_start_time(*args, **kwargs)
-        elif event == measurement.Event.END_CUSTOM_BADPUT_EVENT:
-            self._recorder.record_custom_badput_event_end_time(*args, **kwargs)
-        else:
-            logging.log_first_n(
-                logging.WARNING,
-                "Ignoring unknown event %s",
-                1,
-                event,
+        start_method_name = f"record_{event.value}_start_time"
+        end_method_name = f"record_{event.value}_end_time"
+
+        record_event_start = getattr(self._recorder, start_method_name, None)
+        record_event_end = getattr(self._recorder, end_method_name, None)
+
+        try:
+            if record_event_start:
+                record_event_start(*args, **kwargs)
+        except (TypeError, ValueError, RuntimeError) as e:
+            logging.warning(
+                "Failed to record start of event %s. Error: %s", event.name, e, exc_info=True
             )
 
-    def start_monitoring(self, *args, **kwargs):
-        """Starts Monitoring of Goodput.
+        try:
+            yield
+        finally:
+            try:
+                if record_event_end:
+                    record_event_end(*args, **kwargs)
+            except (TypeError, ValueError, RuntimeError) as e:
+                logging.warning(
+                    "Failed to record end of event %s. Error: %s", event.name, e, exc_info=True
+                )
+
+    @contextlib.contextmanager
+    def maybe_monitor_goodput(self, *args, **kwargs):
+        """Monitor cumulative goodput if enabled.
 
         Instantiate ml-goodput-measurement's GoodputMonitor to asynchronously calculate
-        Goodput and Badput at the upload_interval and upload to the specified TensorBoard
-        directory.
+        Goodput, Badput, Step & Disruption Information at the upload_interval to the
+        specified TensorBoard directory and Google Cloud Monitoring.
         Note: This function requires initialization of distributed JAX before it is called.
         If there are internal GCP errors from querying and uploading data, these will be
         logged without affecting the workload. GoodputMonitor logs will provide further
@@ -123,33 +152,66 @@ def start_monitoring(self, *args, **kwargs):
         Default behavior is to push metrics to Google Cloud Monitoring.
         This behavior can be overridden by configuring `goodput_monitoring.GCPOptions`
         """
-        cfg: GoodputRecorder.Config = self.config
-        include_step_deviation = True
-        if jax.process_index() == 0:
+        if jax.process_index() != 0:
+            yield
+            return
+        try:
             if self._monitor is None:
-                if int(cfg.step_deviation_interval_seconds) <= 0:
-                    include_step_deviation = False
-
                 gcp_options = goodput_monitoring.GCPOptions(
-                    enable_gcp_goodput_metrics=True,
-                    enable_gcp_step_deviation_metrics=include_step_deviation,
+                    enable_gcp_goodput_metrics=self.config.enable_gcp_goodput_metrics,
                 )
                 self._monitor = goodput_monitoring.GoodputMonitor(
-                    job_name=cfg.name,
-                    logger_name=f"goodput_logger_{cfg.name}",
-                    tensorboard_dir=cfg.upload_dir,
-                    upload_interval=int(cfg.upload_interval),
+                    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,
-                    include_badput_breakdown=True,
-                    include_step_deviation=include_step_deviation,
-                    step_deviation_interval_seconds=int(cfg.step_deviation_interval_seconds),
+                    pathway_enabled=self.config.enable_pathways_goodput,
+                    include_badput_breakdown=self.config.include_badput_breakdown,
                     gcp_options=gcp_options,
                 )
 
             self._monitor.start_goodput_uploader(*args, **kwargs)
             logging.info("Started Goodput upload to Tensorboard & GCM in the background!")
-            if include_step_deviation:
-                self._monitor.start_step_deviation_uploader(*args, **kwargs)
+            yield
+        finally:
+            if self._monitor:
+                self._monitor.stop_goodput_uploader()
+                logging.info("Flushed final metrics and safe exited from Goodput monitoring.")
+
+    @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:
+            yield
+            return
+        try:
+            if self._rolling_window_monitor is None:
+                rolling_window_tensorboard_dir = os.path.join(
+                    self.config.upload_dir, f"rolling_window_{self.config.name}"
+                )
+                self._rolling_window_monitor = goodput_monitoring.GoodputMonitor(
+                    job_name=self.config.name,
+                    logger_name=f"goodput_logger_{self.config.name}",
+                    tensorboard_dir=rolling_window_tensorboard_dir,
+                    upload_interval=self.config.upload_interval,
+                    monitoring_enabled=True,
+                    pathway_enabled=self.config.enable_pathways_goodput,
+                    include_badput_breakdown=True,
+                )
+            self._rolling_window_monitor.start_rolling_window_goodput_uploader(
+                self.config.rolling_window_size
+            )
+            logging.info("Started Rolling Window Goodput monitoring in the background!")
+            yield
+        finally:
+            if self._rolling_window_monitor:
+                self._rolling_window_monitor.stop_rolling_window_goodput_uploader()
                 logging.info(
-                    "Started Step Deviation upload to Tensorboard & GCM in the background!"
+                    "Flushed final metrics and safe exited from Rolling Window Goodput monitoring."
                 )
+
+
+def create_goodput_recorder(cfg: GoodputRecorder.Config):
+    """Factory method to create GoodputRecorder."""
+    return GoodputRecorder(cfg)