Add more elasticity goodput tests

src/maxtext/common/goodput.py

@@ -42,6 +42,26 @@ class GoodputEvent(Enum):
 RECORD_JOB_END_TIME = f"record_{GoodputEvent.JOB.value}_end_time"
 
 
+def _resolve_goodput_monitor(config):
+  """Resolves the GoodputMonitor class to use and its class-specific kwargs.
+
+  Returns (monitor_class, monitor_kwargs), where monitor_kwargs holds only
+  the constructor arguments specific to that class.
+  """
+  if config.elastic_enabled:
+    try:
+      from maxtext.utils import elastic_utils  # pylint: disable=import-outside-toplevel
+
+      if elastic_utils.should_use_elastic(config):
+        from ml_goodput_measurement import monitoring_elastic  # pylint: disable=import-outside-toplevel
+
+        return monitoring_elastic.ElasticGoodputMonitor, {"include_slice_efficiency": True}
+    except Exception as e:  # pylint: disable=broad-exception-caught
+      max_logging.log(f"Goodput: could not resolve elastic goodput monitor, falling back to base monitor: {e}")
+
+  return monitoring.GoodputMonitor, {"pathway_enabled": config.enable_pathways_goodput}
+
+
 @contextlib.contextmanager
 def maybe_monitor_goodput(config):
   """Monitor cumulative goodput if enabled on the lead host.
@@ -66,16 +86,9 @@ def maybe_monitor_goodput(config):
         enable_gcp_goodput_metrics=config.enable_gcp_goodput_metrics,
         enable_gcp_step_deviation_metrics=config.enable_gcp_step_deviation_metrics,
     )
-    monitor_class = monitoring.GoodputMonitor
-
-    if config.elastic_enabled:
-      try:
-        from ml_goodput_measurement import monitoring_elastic  # pylint: disable=import-outside-toplevel
-
-        monitor_class = monitoring_elastic.ElasticGoodputMonitor
-      except ImportError:
-        max_logging.log("Elastic monitor failed!")
 
+    monitor_class, monitor_kwargs = _resolve_goodput_monitor(config)
+    max_logging.log(f"Goodput: using {monitor_class.__name__} for job: {config.run_name}")
     kwargs = {
         "job_name": config.run_name,
         "logger_name": f"goodput_{config.run_name}",
@@ -86,9 +99,8 @@ def maybe_monitor_goodput(config):
         "include_step_deviation": config.monitor_step_time_deviation,
         "step_deviation_interval_seconds": config.step_deviation_interval_seconds,
         "gcp_options": gcp_options,
+        **monitor_kwargs,
     }
-    if monitor_class == monitoring.GoodputMonitor:
-      kwargs["pathway_enabled"] = config.enable_pathways_goodput
 
     goodput_monitor = monitor_class(**kwargs)
     goodput_monitor.start_goodput_uploader()
@@ -143,14 +155,18 @@ def create_goodput_recorder(config):
   # Detect if we should use the elastic-aware recorder
   if config.elastic_enabled:
     try:
-      from ml_goodput_measurement import goodput_elastic  # pylint: disable=import-outside-toplevel
       from maxtext.utils import elastic_utils  # pylint: disable=import-outside-toplevel
 
-      recorder = goodput_elastic.ElasticGoodputRecorder(config.run_name, logger_name, jax.process_index() == 0)
-      elastic_utils.record_slice_state(recorder)
-    except ImportError as e:
-      max_logging.log(f"Could not create elastic goodput recorder: {e}")
-    else:
-      return recorder
+      if elastic_utils.should_use_elastic(config):
+        from ml_goodput_measurement import goodput_elastic  # pylint: disable=import-outside-toplevel
+
+        recorder = goodput_elastic.ElasticGoodputRecorder(config.run_name, logger_name, jax.process_index() == 0)
+        elastic_utils.record_slice_state(recorder)
+        max_logging.log(f"Goodput: created ElasticGoodputRecorder for job: {config.run_name}")
+        return recorder
+    except Exception as e:  # pylint: disable=broad-exception-caught
+      max_logging.log(f"Goodput: could not create elastic goodput recorder, falling back to base recorder: {e}")
 
-  return goodput.GoodputRecorder(config.run_name, logger_name, jax.process_index() == 0)
+  recorder = goodput.GoodputRecorder(config.run_name, logger_name, jax.process_index() == 0)
+  max_logging.log(f"Goodput: created base GoodputRecorder for job: {config.run_name}")
+  return recorder

src/maxtext/utils/elastic_utils.py

@@ -57,7 +57,7 @@ def record_elastic_event_start(recorder, config) -> None:
   global pending_elastic_event_type
   event_type = "elastic_scale_up" if is_scale_up_event(config) else "elastic_slice_down"
   pending_elastic_event_type = event_type
-  if recorder:
+  if recorder and hasattr(recorder, "record_elastic_wait_start_time"):
     recorder.record_elastic_wait_start_time(event_type=event_type)
     record_slice_state(recorder, active_slices_override=0)
 
@@ -69,7 +69,7 @@ def record_elastic_wait_end_and_reinit_start(recorder) -> None:
     return
   event_type = pending_elastic_event_type
   pending_elastic_event_type = None
-  if recorder:
+  if recorder and hasattr(recorder, "record_elastic_wait_end_time"):
     recorder.record_elastic_wait_end_time(event_type=event_type)
     recorder.record_elastic_reinit_start_time()
     record_slice_state(recorder)
@@ -79,10 +79,10 @@ def record_elastic_wait_end_and_reinit_start(recorder) -> None:
 def record_elastic_reinit_end() -> None:
   """Records end of elastic reinitialization event."""
   global pending_reinit_recorder
-  if pending_reinit_recorder is not None:
+  if pending_reinit_recorder is not None and hasattr(pending_reinit_recorder, "record_elastic_reinit_end_time"):
     pending_reinit_recorder.record_elastic_reinit_end_time()
     record_slice_state(pending_reinit_recorder)
-    pending_reinit_recorder = None
+  pending_reinit_recorder = None
 
 
 def elastic_enabled(config) -> bool:

tests/unit/elastic_utils_test.py

@@ -465,6 +465,36 @@ def test_record_elastic_reinit_end_on_cold_start(self):
 
     elastic_utils.record_elastic_reinit_end()
 
+  def test_record_elastic_event_start_non_elastic_recorder_noop(self):
+    """A recorder lacking the elastic API (e.g. the ImportError fallback) must not raise."""
+    elastic_utils.elastic_manager = self.fake_manager
+    self.fake_manager.available_inactive_slices = set()
+    non_elastic_recorder = Mock(spec=[])  # No record_elastic_* attributes.
+    config = FakeConfig()
+
+    elastic_utils.record_elastic_event_start(non_elastic_recorder, config)  # Must not raise.
+
+    self.assertEqual(elastic_utils.pending_elastic_event_type, "elastic_slice_down")
+
+  def test_record_elastic_wait_end_and_reinit_start_non_elastic_recorder_noop(self):
+    """A recorder lacking the elastic API must not raise, but is still tracked as pending."""
+    elastic_utils.pending_elastic_event_type = "elastic_slice_down"  # pyrefly: ignore[bad-assignment]
+    non_elastic_recorder = Mock(spec=[])
+
+    elastic_utils.record_elastic_wait_end_and_reinit_start(non_elastic_recorder)  # Must not raise.
+
+    self.assertIs(elastic_utils.pending_reinit_recorder, non_elastic_recorder)
+    self.assertIsNone(elastic_utils.pending_elastic_event_type)
+
+  def test_record_elastic_reinit_end_non_elastic_recorder_noop(self):
+    """A recorder lacking the elastic API must not raise, and pending state is still cleared."""
+    non_elastic_recorder = Mock(spec=[])
+    elastic_utils.pending_reinit_recorder = non_elastic_recorder
+
+    elastic_utils.record_elastic_reinit_end()  # Must not raise.
+
+    self.assertIsNone(elastic_utils.pending_reinit_recorder)
+
   def test_ensure_elastic_manager_initialized_readonly_config(self):
     """Tests that ensure_elastic_manager_initialized works with read-only config."""
 

tests/unit/goodput_utils_test.py

@@ -24,11 +24,16 @@
     GoodputEvent,
     RECORD_JOB_END_TIME,
     RECORD_JOB_START_TIME,
+    _resolve_goodput_monitor,
     create_goodput_recorder,
     maybe_monitor_goodput,
     maybe_record_goodput,
     record_goodput,
 )
+from ml_goodput_measurement import goodput
+from ml_goodput_measurement import goodput_elastic
+from ml_goodput_measurement import monitoring
+from ml_goodput_measurement import monitoring_elastic
 from tests.utils.test_helpers import get_test_config_path, get_test_base_output_directory
 
 pytestmark = [pytest.mark.external_training]
@@ -93,6 +98,118 @@ def test_job_recording_constants(self):
     self.assertEqual(RECORD_JOB_START_TIME, "record_job_start_time")
     self.assertEqual(RECORD_JOB_END_TIME, "record_job_end_time")
 
+  def test_resolve_goodput_monitor_non_elastic(self):
+    """A McJAX config must resolve to the base monitor."""
+    self.assertFalse(self.config.elastic_enabled)
+
+    monitor_class, monitor_kwargs = _resolve_goodput_monitor(self.config)
+
+    self.assertIs(monitor_class, monitoring.GoodputMonitor)
+    self.assertIn("pathway_enabled", monitor_kwargs)
+    self.assertNotIn("include_slice_efficiency", monitor_kwargs)
+
+  @mock.patch("maxtext.utils.elastic_utils.should_use_elastic")
+  def test_resolve_goodput_monitor_elastic(self, mock_should_use_elastic):
+    """elastic_enabled=True on an actual Pathways run must resolve to the elastic monitor."""
+    mock_should_use_elastic.return_value = True
+    self.config.elastic_enabled = True
+
+    monitor_class, monitor_kwargs = _resolve_goodput_monitor(self.config)
+
+    self.assertIs(monitor_class, monitoring_elastic.ElasticGoodputMonitor)
+    self.assertEqual(monitor_kwargs, {"include_slice_efficiency": True})
+    self.assertNotIn("pathway_enabled", monitor_kwargs)
+
+  @mock.patch("maxtext.utils.elastic_utils.should_use_elastic")
+  def test_resolve_goodput_monitor_elastic_enabled_not_pathways_falls_back(self, mock_should_use_elastic):
+    """elastic_enabled=True but not actually on Pathways must fall back to the base monitor."""
+    mock_should_use_elastic.return_value = False
+    self.config.elastic_enabled = True
+
+    monitor_class, monitor_kwargs = _resolve_goodput_monitor(self.config)
+
+    self.assertIs(monitor_class, monitoring.GoodputMonitor)
+    self.assertIn("pathway_enabled", monitor_kwargs)
+
+  @mock.patch("ml_goodput_measurement.monitoring_elastic.ElasticGoodputMonitor", side_effect=RuntimeError("boom"))
+  @mock.patch("maxtext.utils.elastic_utils.should_use_elastic")
+  def test_resolve_goodput_monitor_elastic_construction_failure_falls_back(
+      self, mock_should_use_elastic, unused_mock_elastic_monitor_cls
+  ):
+    """A failure constructing the elastic monitor must fall back rather than propagate."""
+    mock_should_use_elastic.return_value = True
+    self.config.elastic_enabled = True
+
+    monitor_class, _ = _resolve_goodput_monitor(self.config)
+
+    self.assertIs(monitor_class, monitoring.GoodputMonitor)
+
+  @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
+  ):
+    """maybe_monitor_goodput actually starts/stops an ElasticGoodputMonitor when elastic is active."""
+    mock_should_use_elastic.return_value = True
+    mock_start_goodput_uploader.return_value = mock.MagicMock()
+    self.config.elastic_enabled = True
+
+    with maybe_monitor_goodput(self.config):
+      mock_start_goodput_uploader.assert_called()
+    mock_stop_goodput_uploader.assert_called()
+
+  @mock.patch("google.cloud.logging.Client")
+  def test_create_goodput_recorder_non_elastic(self, mock_cloud_logger):
+    """Regular (non-Pathways/McJAX) config must get the base recorder."""
+    mock_cloud_logger.return_value = mock.MagicMock()
+    self.assertFalse(self.config.elastic_enabled)
+
+    recorder = create_goodput_recorder(self.config)
+
+    self.assertNotIsInstance(recorder, goodput_elastic.ElasticGoodputRecorder)
+
+  @mock.patch("maxtext.utils.elastic_utils.should_use_elastic")
+  @mock.patch("google.cloud.logging.Client")
+  def test_create_goodput_recorder_elastic(self, mock_cloud_logger, mock_should_use_elastic):
+    """elastic_enabled=True on an actual Pathways run must get the elastic recorder."""
+    mock_cloud_logger.return_value = mock.MagicMock()
+    mock_should_use_elastic.return_value = True
+    self.config.elastic_enabled = True
+
+    recorder = create_goodput_recorder(self.config)
+
+    self.assertIsInstance(recorder, goodput_elastic.ElasticGoodputRecorder)
+
+  @mock.patch("maxtext.utils.elastic_utils.should_use_elastic")
+  @mock.patch("google.cloud.logging.Client")
+  def test_create_goodput_recorder_elastic_enabled_not_pathways_falls_back(
+      self, mock_cloud_logger, mock_should_use_elastic
+  ):
+    """elastic_enabled=True but not actually on Pathways (e.g. McJAX) must get the base recorder."""
+    mock_cloud_logger.return_value = mock.MagicMock()
+    mock_should_use_elastic.return_value = False
+    self.config.elastic_enabled = True
+
+    recorder = create_goodput_recorder(self.config)
+
+    self.assertNotIsInstance(recorder, goodput_elastic.ElasticGoodputRecorder)
+
+  @mock.patch("ml_goodput_measurement.goodput_elastic.ElasticGoodputRecorder", side_effect=RuntimeError("boom"))
+  @mock.patch("maxtext.utils.elastic_utils.should_use_elastic")
+  @mock.patch("google.cloud.logging.Client")
+  def test_create_goodput_recorder_elastic_construction_failure_falls_back(
+      self, mock_cloud_logger, mock_should_use_elastic, unused_mock_elastic_recorder_cls
+  ):
+    """A failure constructing the elastic recorder must fall back rather than propagate."""
+    mock_cloud_logger.return_value = mock.MagicMock()
+    mock_should_use_elastic.return_value = True
+    self.config.elastic_enabled = True
+
+    recorder = create_goodput_recorder(self.config)  # Must not raise.
+
+    self.assertIsInstance(recorder, goodput.GoodputRecorder)
+
   @mock.patch("ml_goodput_measurement.goodput.GoodputRecorder.record_job_end_time")
   @mock.patch("ml_goodput_measurement.goodput.GoodputRecorder.record_job_start_time")
   @mock.patch("google.cloud.logging.Client")