Resolve circular dep in goodput monitor code with minimal change

GitOrigin-RevId: 7ee0b01

Author: wangkuiyi

axlearn/cloud/gcp/measurement.py

@@ -30,16 +30,16 @@
 from ml_goodput_measurement import monitoring as goodput_monitoring
 
 from axlearn.cloud.common.utils import parse_kv_flags, to_bool
-from axlearn.common import measurement
+from axlearn.common import measurement_base
 from axlearn.common.config import REQUIRED, Required, config_class, maybe_set_config
 
 
-@measurement.register_recorder("goodput")
-class GoodputRecorder(measurement.Recorder):
+@measurement_base.register_recorder("goodput")
+class GoodputRecorder(measurement_base.Recorder):
     """Records overall training goodput."""
 
     @config_class
-    class Config(measurement.Recorder.Config):
+    class Config(measurement_base.Recorder.Config):
         """Configures GoodputRecorder.
 
         Attributes:
@@ -77,7 +77,7 @@ def from_flags(cls, fv: flags.FlagValues) -> "GoodputRecorder":
         - jax_backend: The type of jax backend.
         - enable_monitoring: Boolean to enable/disable goodput monitoring (default: true).
         """
-        cfg: measurement.Recorder.Config = cls.default_config()
+        cfg: measurement_base.Recorder.Config = cls.default_config()
         parsed_flags = parse_kv_flags(fv.recorder_spec, delimiter="=")
         if "upload_interval" in parsed_flags:
             parsed_flags["upload_interval"] = int(parsed_flags["upload_interval"])
@@ -100,7 +100,7 @@ def __init__(self, cfg):
         self._logger_name = f"goodput_logger_{cfg.name}"
 
     @contextlib.contextmanager
-    def record_event(self, event: measurement.EventType, *args, **kwargs):
+    def record_event(self, event: measurement_base.EventType, *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:
@@ -225,7 +225,7 @@ def monitor_goodput():
 
         return monitor_goodput()
 
-    def record(self, event: measurement.Event, *args, **kwargs):
+    def record(self, event: measurement_base.Event, *args, **kwargs):
         """Deprecated: `record()` is not used in GoodputRecorder.
         Use the record_event context manager instead.
         """
@@ -238,27 +238,27 @@ def record(self, event: measurement.Event, *args, **kwargs):
                 logging_enabled=(jax.process_index() == 0),
             )
 
-        if event == measurement.Event.START_JOB:
+        if event == measurement_base.Event.START_JOB:
             self._recorder.record_job_start_time(*args, **kwargs)
-        elif event == measurement.Event.END_JOB:
+        elif event == measurement_base.Event.END_JOB:
             self._recorder.record_job_end_time(*args, **kwargs)
-        elif event == measurement.Event.START_STEP:
+        elif event == measurement_base.Event.START_STEP:
             self._recorder.record_step_start_time(*args, **kwargs)
-        elif event == measurement.Event.START_ACCELERATOR_INIT:
+        elif event == measurement_base.Event.START_ACCELERATOR_INIT:
             self._recorder.record_tpu_init_start_time(*args, **kwargs)
-        elif event == measurement.Event.END_ACCELERATOR_INIT:
+        elif event == measurement_base.Event.END_ACCELERATOR_INIT:
             self._recorder.record_tpu_init_end_time(*args, **kwargs)
-        elif event == measurement.Event.START_TRAINING_PREPARATION:
+        elif event == measurement_base.Event.START_TRAINING_PREPARATION:
             self._recorder.record_training_preparation_start_time(*args, **kwargs)
-        elif event == measurement.Event.END_TRAINING_PREPARATION:
+        elif event == measurement_base.Event.END_TRAINING_PREPARATION:
             self._recorder.record_training_preparation_end_time(*args, **kwargs)
-        elif event == measurement.Event.START_DATA_LOADING:
+        elif event == measurement_base.Event.START_DATA_LOADING:
             self._recorder.record_data_loading_start_time(*args, **kwargs)
-        elif event == measurement.Event.END_DATA_LOADING:
+        elif event == measurement_base.Event.END_DATA_LOADING:
             self._recorder.record_data_loading_end_time(*args, **kwargs)
-        elif event == measurement.Event.START_CUSTOM_BADPUT_EVENT:
+        elif event == measurement_base.Event.START_CUSTOM_BADPUT_EVENT:
             self._recorder.record_custom_badput_event_start_time(*args, **kwargs)
-        elif event == measurement.Event.END_CUSTOM_BADPUT_EVENT:
+        elif event == measurement_base.Event.END_CUSTOM_BADPUT_EVENT:
             self._recorder.record_custom_badput_event_end_time(*args, **kwargs)
         else:
             logging.log_first_n(

axlearn/common/measurement.py

@@ -2,169 +2,31 @@
 
 """A library to measure e2e metrics like goodput."""
 
-import contextlib
-import enum
 import importlib
-from typing import Optional, TypeVar
+from typing import Optional
 
 from absl import flags, logging
 
-from axlearn.common.config import REQUIRED, Configurable, Required, config_class
-
-
-class Event(enum.Enum):
-    """Event to be recorded (Legacy).
-
-    Attributes:
-        START_JOB: Start of job.
-        END_JOB: End of job.
-        START_STEP: Start of a training step. Should be recorded with `step` as a positional arg.
-        START_ACCELERATOR_INIT: Start of accelerator mesh initialization.
-        END_ACCELERATOR_INIT: End of accelerator mesh initialization.
-        START_TRAINING_PREPARATION: Start of training preparation.
-        END_TRAINING_PREPARATION: End of training preparation.
-        START_DATA_LOADING: Start of data loading.
-        END_DATA_LOADING: End of data loading.
-        START_CUSTOM_BADPUT_EVENT: Start of custom badput event.
-        END_CUSTOM_BADPUT_EVENT: End of custom badput event.
-    """
-
-    START_JOB = "START_JOB"
-    END_JOB = "END_JOB"
-    START_STEP = "START_STEP"
-    START_ACCELERATOR_INIT = "START_ACCELERATOR_INIT"
-    END_ACCELERATOR_INIT = "END_ACCELERATOR_INIT"
-    START_TRAINING_PREPARATION = "START_TRAINING_PREPARATION"
-    END_TRAINING_PREPARATION = "END_TRAINING_PREPARATION"
-    START_DATA_LOADING = "START_DATA_LOADING"
-    END_DATA_LOADING = "END_DATA_LOADING"
-    START_CUSTOM_BADPUT_EVENT = "START_CUSTOM_BADPUT_EVENT"
-    END_CUSTOM_BADPUT_EVENT = "END_CUSTOM_BADPUT_EVENT"
-
-
-class EventType(enum.Enum):
-    """Event to be recorded.
-
-    Attributes:
-        JOB: Start and end of the job.
-        STEP: Start of a training step. Should be recorded with `step` as a positional arg.
-        ACCELERATOR_INIT: Start and end of accelerator mesh initialization.
-        TRAINING_PREPARATION: Start and end of training preparation.
-        DATA_LOADING: Start and end of data loading.
-        CUSTOM_BADPUT_EVENT: Start and end of custom badput events.
-    """
-
-    JOB = "job"
-    STEP = "step"
-    ACCELERATOR_INIT = "tpu_init"
-    TRAINING_PREPARATION = "training_preparation"
-    DATA_LOADING = "data_loading"
-    CUSTOM_BADPUT_EVENT = "custom_badput_event"
-
-
-class Recorder(Configurable):
-    """The base interface for collecting e2e metrics."""
-
-    @config_class
-    class Config(Configurable.Config):
-        """Configures Recorder.
-
-        Attributes:
-            name: Name of the recorder.
-        """
-
-        name: Required[str] = REQUIRED
-
-    @classmethod
-    def from_flags(cls, fv: Optional[flags.FlagValues]) -> "Recorder":
-        """Converts flags to a recorder."""
-        raise NotImplementedError(cls)
-
-    def record(self, event: Event, *args, **kwargs):
-        """Records a single, instantaneous event.
-
-        Note:
-            This method is maintained for backward compatibility.
-            New child recorder implementations should prioritize implementing the
-            `record_event` context manager instead.
-        """
-        raise NotImplementedError(type(self))
-
-    def start_monitoring(self, **kwargs):
-        """Starts computing and uploading metrics in the background.
-
-        Note:
-            This method is maintained for backward compatibility. New child
-            recorders should prioritize implementing the `maybe_monitor_all`
-            context manager, which provides a clearer lifecycle for monitoring.
-        """
-        raise NotImplementedError(type(self))
-
-    @contextlib.contextmanager
-    def record_event(self, event: EventType, *args, **kwargs):
-        """A context manager to record the start and end of an event.
-
-        This is the preferred method for recording events. Child classes
-        should implement this context manager to handle timed operations.
-
-        Example:
-            with recorder.record_event(EventType.ACCELERATOR_INIT):
-                # Device initialization.
-        """
-        # pylint: disable=unnecessary-pass
-        # pylint: disable=unused-argument
-        try:
-            yield
-        finally:
-            pass
-
-    @contextlib.contextmanager
-    def maybe_monitor_all(self):
-        """Context manager to start and stop computing and monitoring metrics.
-
-        This is the preferred method for monitoring events. Child classes
-        should implement this context manager to manage all types of monitoring.
-
-        Example:
-            with recorder.maybe_monitor_all():
-                # Train
-        """
-        yield
-
-
-_recorders: dict[str, type] = {}
-_T = TypeVar("_T")
-
-
-def register_recorder(name: str):
-    def fn(cls: _T) -> _T:
-        """Registers a recorder class for `get_recorder_config`."""
-        if name in _recorders:
-            raise ValueError(f"Recorder {name} is already registered.")
-        _recorders[name] = cls
-        return cls
-
-    return fn
-
-
-def define_flags(**kwargs):
-    """Common measurement flags."""
-
-    flags.DEFINE_string(
-        "recorder_type",
-        None,
-        "The recorder type. It can be a recorder name, e.g. `my_recorder`, or "
-        "a module paired with a recorder name, e.g. `my.module:my_recorder`.",
-        **kwargs,
-    )
-    flags.DEFINE_multi_string(
-        "recorder_spec",
-        [],
-        "Recorder spec provided as key=value. "
-        "Refer to each recorders's `from_flags` method docstring for details.",
-        **kwargs,
-    )
-
+from axlearn.common.measurement_base import (
+    Event,
+    EventType,
+    Recorder,
+    _recorders,
+    define_flags,
+    register_recorder,
+)
+
+__all__ = [
+    "Event",
+    "EventType",
+    "Recorder",
+    "define_flags",
+    "register_recorder",
+    "global_recorder",
+    "initialize",
+    "record_event",
+    "start_monitoring",
+]
 
 global_recorder: Optional[Recorder] = None