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 _construct_goodput_monitor(config, common_kwargs):
+  """Constructs a GoodputMonitor, preferring the elastic-aware monitor when applicable."""
+  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
+
+        monitor = monitoring_elastic.ElasticGoodputMonitor(include_slice_efficiency=True, **common_kwargs)
+        max_logging.log(f"Goodput: using ElasticGoodputMonitor for job: {config.run_name}")
+        return monitor
+    except Exception as e:  # pylint: disable=broad-exception-caught
+      max_logging.log(f"Goodput: could not create elastic goodput monitor, falling back to base monitor: {e}")
+
+  monitor = monitoring.GoodputMonitor(pathway_enabled=config.enable_pathways_goodput, **common_kwargs)
+  max_logging.log(f"Goodput: using GoodputMonitor for job: {config.run_name}")
+  return monitor
+
+
 @contextlib.contextmanager
 def maybe_monitor_goodput(config):
   """Monitor cumulative goodput if enabled on the lead host.
@@ -66,17 +86,8 @@ 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!")
 
-    kwargs = {
+    common_kwargs = {
         "job_name": config.run_name,
         "logger_name": f"goodput_{config.run_name}",
         "tensorboard_dir": config.tensorboard_dir,
@@ -87,10 +98,8 @@ def maybe_monitor_goodput(config):
         "step_deviation_interval_seconds": config.step_deviation_interval_seconds,
         "gcp_options": gcp_options,
     }
-    if monitor_class == monitoring.GoodputMonitor:
-      kwargs["pathway_enabled"] = config.enable_pathways_goodput
 
-    goodput_monitor = monitor_class(**kwargs)
+    goodput_monitor = _construct_goodput_monitor(config, common_kwargs)
     goodput_monitor.start_goodput_uploader()
     max_logging.log("Started Goodput upload to Tensorboard & GCM in the background!")
     yield
@@ -143,14 +152,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

@@ -22,6 +22,7 @@
 from maxtext.utils import gcs_utils
 from maxtext.utils import max_logging
 import pathwaysutils
+from pathwaysutils.elastic import elastic
 from pathwaysutils.elastic import manager
 
 elastic_manager: manager.Manager | None = None
@@ -39,11 +40,11 @@ def record_slice_state(recorder, active_slices_override: int | None = None) -> N
   ):
     return
 
-  available_slices = len(pathwaysutils.elastic.get_active_slice_indices())
+  available_slices = len(elastic.get_active_slice_indices())
   active_slices = (
       active_slices_override if active_slices_override is not None else len(elastic_manager.active_slice_indices)
   )
-  total_slices = len(pathwaysutils.elastic.get_slice_to_devices(jax.devices()))
+  total_slices = len(elastic.get_slice_to_devices(jax.devices()))
 
   recorder.record_elastic_slice_counts(
       available_slices=available_slices,
@@ -57,7 +58,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 +70,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 +80,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

@@ -61,6 +61,7 @@ def setUp(self):
     self.original_jax = elastic_utils.jax
     self.original_gcs_utils = elastic_utils.gcs_utils
     self.original_max_logging = elastic_utils.max_logging
+    self.original_elastic = elastic_utils.elastic
     self.original_manager_class = pathwaysutils.elastic.manager.Manager
     self.original_scale_up_signal_error = getattr(pathwaysutils.elastic.manager, "ScaleUpSignalError", None)
 
@@ -72,12 +73,12 @@ def setUp(self):
     self.fake_jax = create_autospec(self.original_jax)
     self.fake_manager = create_autospec(self.original_manager_class, instance=True)
     self.fake_manager.available_inactive_slices = set()
+    self.fake_elastic = create_autospec(self.original_elastic)
 
     # Configure default behaviors if needed
     self.fake_pathwaysutils.is_pathways_backend_used.return_value = True
-    self.fake_pathwaysutils.elastic = Mock()
-    self.fake_pathwaysutils.elastic.get_active_slice_indices.return_value = [0, 1]
-    self.fake_pathwaysutils.elastic.get_slice_to_devices.return_value = {
+    self.fake_elastic.get_active_slice_indices.return_value = [0, 1]
+    self.fake_elastic.get_slice_to_devices.return_value = {
         0: [FakeDevice()],
         1: [FakeDevice()],
     }
@@ -89,6 +90,7 @@ def setUp(self):
     self.fake_jax.errors.JaxRuntimeError = MockJaxRuntimeError
     elastic_utils.gcs_utils = self.fake_gcs_utils
     elastic_utils.max_logging = self.fake_logging
+    elastic_utils.elastic = self.fake_elastic
 
     # Hook up pathwaysutils.elastic.manager.Manager to return our fake_manager
     pathwaysutils.elastic.manager.Manager = lambda *args, **kwargs: self.fake_manager  # pyrefly: ignore[bad-assignment]
@@ -102,6 +104,7 @@ def tearDown(self):
     elastic_utils.jax = self.original_jax
     elastic_utils.gcs_utils = self.original_gcs_utils
     elastic_utils.max_logging = self.original_max_logging
+    elastic_utils.elastic = self.original_elastic
     pathwaysutils.elastic.manager.Manager = self.original_manager_class
     pathwaysutils.elastic.manager.ScaleUpSignalError = (  # pyrefly: ignore[bad-assignment]
         self.original_scale_up_signal_error,
@@ -115,7 +118,7 @@ def test_record_slice_state(self):
     elastic_utils.elastic_manager = self.fake_manager
     self.fake_manager.active_slice_indices = {0}
     self.fake_manager.slice_to_devices = {0: [FakeDevice()], 1: [FakeDevice()]}
-    self.fake_pathwaysutils.elastic.get_active_slice_indices.return_value = {0, 1}
+    self.fake_elastic.get_active_slice_indices.return_value = {0, 1}
 
     fake_recorder = Mock()
     fake_recorder.record_elastic_slice_counts = Mock()
@@ -131,8 +134,6 @@ def test_record_slice_state(self):
   def test_elastic_enabled(self):
     config = FakeConfig()
     self.fake_pathwaysutils.is_pathways_backend_used.return_value = True
-    self.fake_pathwaysutils.elastic = Mock()
-    self.fake_pathwaysutils.elastic.get_active_slice_indices.return_value = [0, 1]
     config.elastic_enabled = True
     self.assertTrue(elastic_utils.elastic_enabled(config))
 
@@ -177,8 +178,6 @@ def test_live_devices_no_pathways(self):
   def test_live_devices_pathways(self):
     """Tests live_devices when pathways is used."""
     self.fake_pathwaysutils.is_pathways_backend_used.return_value = True
-    self.fake_pathwaysutils.elastic = Mock()
-    self.fake_pathwaysutils.elastic.get_active_slice_indices.return_value = [0, 1]
     device0 = FakeDevice(slice_index=0)
     device1 = FakeDevice(slice_index=1)
     self.fake_jax.devices.return_value = [device0, device1]
@@ -203,8 +202,6 @@ def test_live_devices_disabled(self):
   def test_elastic_retry_disabled(self):
     """Tests elastic_retry when disabled but pathways is used."""
     self.fake_pathwaysutils.is_pathways_backend_used.return_value = True
-    self.fake_pathwaysutils.elastic = Mock()
-    self.fake_pathwaysutils.elastic.get_active_slice_indices.return_value = [0, 1]
     config = FakeConfig()
     config.elastic_enabled = False
     msg = (
@@ -465,6 +462,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."""
 
@@ -479,8 +506,6 @@ def __setattr__(self, name, value):
 
     config = ReadOnlyConfig()
     self.fake_pathwaysutils.is_pathways_backend_used.return_value = True
-    self.fake_pathwaysutils.elastic = Mock()
-    self.fake_pathwaysutils.elastic.get_active_slice_indices.return_value = [0, 1]
 
     # Should not raise ValueError
     elastic_utils.ensure_elastic_manager_initialized(config)