Add a new goodput elastic utils function record_slice_state()

Author: YixuanWang-99

src/maxtext/common/goodput.py

@@ -66,18 +66,31 @@ 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,
     )
-    goodput_monitor = monitoring.GoodputMonitor(
-        job_name=config.run_name,
-        logger_name=f"goodput_{config.run_name}",
-        tensorboard_dir=config.tensorboard_dir,
-        upload_interval=config.goodput_upload_interval_seconds,
-        monitoring_enabled=True,
-        pathway_enabled=config.enable_pathways_goodput,
-        include_badput_breakdown=True,
-        include_step_deviation=config.monitor_step_time_deviation,
-        step_deviation_interval_seconds=config.step_deviation_interval_seconds,
-        gcp_options=gcp_options,
-    )
+    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 = {
+        "job_name": config.run_name,
+        "logger_name": f"goodput_{config.run_name}",
+        "tensorboard_dir": config.tensorboard_dir,
+        "upload_interval": config.goodput_upload_interval_seconds,
+        "monitoring_enabled": True,
+        "include_badput_breakdown": True,
+        "include_step_deviation": config.monitor_step_time_deviation,
+        "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.start_goodput_uploader()
     max_logging.log("Started Goodput upload to Tensorboard & GCM in the background!")
     yield
@@ -121,8 +134,23 @@ def create_goodput_recorder(config):
     if config.enable_goodput_recording and jax.process_index() == 0:
       max_logging.log("[GOODPUT NO-OP] recorder skipped (decoupled stub).")
     return None
-  if config.enable_goodput_recording:
-    logger_name = f"goodput_{config.run_name}"
-    recorder = goodput.GoodputRecorder(config.run_name, logger_name, jax.process_index() == 0)
-    return recorder
-  return None
+
+  if not config.enable_goodput_recording:
+    return None
+
+  logger_name = f"goodput_{config.run_name}"
+
+  # 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
+
+  return goodput.GoodputRecorder(config.run_name, logger_name, jax.process_index() == 0)

src/maxtext/utils/elastic_utils.py

@@ -28,13 +28,37 @@
 pending_elastic_event_type = None
 
 
+def record_slice_state(recorder, active_slices_override: int | None = None) -> None:
+  """Records live slice counts and logs them to the GoodputRecorder."""
+  if recorder is None or not hasattr(recorder, "record_elastic_slice_counts"):
+    return
+
+  if not pathwaysutils.is_pathways_backend_used() or elastic_manager is None:
+    available_slices = 1
+    active_slices = 1 if active_slices_override is None else active_slices_override
+    total_slices = 1
+  else:
+    available_slices = len(pathwaysutils.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()))
+
+  recorder.record_elastic_slice_counts(
+      available_slices=available_slices,
+      active_slices=active_slices,
+      total_slices=total_slices,
+  )
+
+
 def record_elastic_event_start(recorder, config) -> None:
   """Records start of an elastic scale up event."""
   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:
-    recorder.record_custom_badput_event_start_time(custom_badput_event_type=event_type)
+    recorder.record_elastic_wait_start_time(event_type=event_type)
+    record_slice_state(recorder, active_slices_override=0)
 
 
 def record_elastic_wait_end_and_reinit_start(recorder) -> None:
@@ -45,16 +69,18 @@ def record_elastic_wait_end_and_reinit_start(recorder) -> None:
   event_type = pending_elastic_event_type
   pending_elastic_event_type = None
   if recorder:
-    recorder.record_custom_badput_event_end_time(custom_badput_event_type=event_type)
-    recorder.record_custom_badput_event_start_time(custom_badput_event_type="elastic_reinitialization")
+    recorder.record_elastic_wait_end_time(event_type=event_type)
+    recorder.record_elastic_reinit_start_time()
+    record_slice_state(recorder)
   pending_reinit_recorder = recorder
 
 
 def record_elastic_reinit_end() -> None:
   """Records end of elastic reinitialization event."""
   global pending_reinit_recorder
   if pending_reinit_recorder is not None:
-    pending_reinit_recorder.record_custom_badput_event_end_time(custom_badput_event_type="elastic_reinitialization")
+    pending_reinit_recorder.record_elastic_reinit_end_time()
+    record_slice_state(pending_reinit_recorder)
     pending_reinit_recorder = None
 
 
@@ -175,6 +201,10 @@ def elastic_retry(config, callback_fn=None, pre_callback_fn=None):
   Returns:
     A decorator for elastic retry.
   """
+  if not pathwaysutils.is_pathways_backend_used() and config.elastic_enabled:
+    max_logging.log("Elastic Retry disabled (simulated logging mode). Returning identity decorator.")
+    return lambda f: f
+
   if not elastic_enabled(config):
     msg = (
         "Elastic training requires the Pathways backend, and elastic_enabled"

tests/unit/elastic_utils_test.py

@@ -70,6 +70,12 @@ def setUp(self):
 
     # 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 = {
+        0: [FakeDevice()],
+        1: [FakeDevice()],
+    }
     self.fake_jax.process_index.return_value = 0
 
     # Inject fakes into elastic_utils namespace
@@ -97,9 +103,28 @@ def tearDown(self):
     elastic_utils.pending_elastic_event_type = None
     super().tearDown()
 
+  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}
+
+    fake_recorder = Mock()
+    fake_recorder.record_elastic_slice_counts = Mock()
+
+    elastic_utils.record_slice_state(fake_recorder)
+
+    fake_recorder.record_elastic_slice_counts.assert_called_once_with(available_slices=2, active_slices=1, total_slices=2)
+
+    fake_recorder.record_elastic_slice_counts.reset_mock()
+    elastic_utils.record_slice_state(fake_recorder, active_slices_override=0)
+    fake_recorder.record_elastic_slice_counts.assert_called_once_with(available_slices=2, active_slices=0, total_slices=2)
+
   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))
 
@@ -144,6 +169,8 @@ 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]
@@ -168,6 +195,8 @@ 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 = (
@@ -319,28 +348,28 @@ def test_record_elastic_event_start(self):
     """Tests recording an elastic slice down start."""
     elastic_utils.elastic_manager = self.fake_manager
     self.fake_manager.new_slice_event.is_set.return_value = False
+    self.fake_manager.slice_to_devices = {0: [FakeDevice()], 1: [FakeDevice()]}
     fake_recorder = Mock()
     config = FakeConfig()
 
     elastic_utils.record_elastic_event_start(fake_recorder, config)
 
-    fake_recorder.record_custom_badput_event_start_time.assert_called_once_with(
-        custom_badput_event_type="elastic_slice_down"
-    )
+    fake_recorder.record_elastic_wait_start_time.assert_called_once_with(event_type="elastic_slice_down")
+    fake_recorder.record_elastic_slice_counts.assert_called_once()
     self.assertEqual(elastic_utils.pending_elastic_event_type, "elastic_slice_down")
 
   def test_record_elastic_event_start_scale_up(self):
     """Tests recording an elastic slice scale up start."""
     elastic_utils.elastic_manager = self.fake_manager
     self.fake_manager.new_slice_event.is_set.return_value = True
+    self.fake_manager.slice_to_devices = {0: [FakeDevice()], 1: [FakeDevice()]}
     fake_recorder = Mock()
     config = FakeConfig()
 
     elastic_utils.record_elastic_event_start(fake_recorder, config)
 
-    fake_recorder.record_custom_badput_event_start_time.assert_called_once_with(
-        custom_badput_event_type="elastic_scale_up"
-    )
+    fake_recorder.record_elastic_wait_start_time.assert_called_once_with(event_type="elastic_scale_up")
+    fake_recorder.record_elastic_slice_counts.assert_called_once()
 
   def test_record_elastic_wait_end_and_reinit_start_noop_on_first_attempt(self):
     """Tests recording elastic event end and elastic reinit start."""
@@ -349,36 +378,39 @@ def test_record_elastic_wait_end_and_reinit_start_noop_on_first_attempt(self):
 
     elastic_utils.record_elastic_wait_end_and_reinit_start(fake_recorder)
 
-    fake_recorder.record_custom_badput_event_end_time.assert_not_called()
-    fake_recorder.record_custom_badput_event_start_time.assert_not_called()
+    fake_recorder.record_elastic_wait_end_time.assert_not_called()
+    fake_recorder.record_elastic_reinit_start_time.assert_not_called()
     self.assertIsNone(elastic_utils.pending_reinit_recorder)
 
   def test_record_elastic_wait_end_and_reinit_start(self):
     """Test recording end of slice down and start of reinit."""
     elastic_utils.pending_elastic_event_type = "elastic_slice_down"
+    elastic_utils.elastic_manager = self.fake_manager
+    self.fake_manager.active_slice_indices = {0}
+    self.fake_manager.slice_to_devices = {0: [FakeDevice()], 1: [FakeDevice()]}
+    elastic_utils.pending_elastic_event_type = "elastic_slice_down"
     fake_recorder = Mock()
 
     elastic_utils.record_elastic_wait_end_and_reinit_start(fake_recorder)
 
-    fake_recorder.record_custom_badput_event_end_time.assert_called_once_with(
-        custom_badput_event_type="elastic_slice_down"
-    )
-    fake_recorder.record_custom_badput_event_start_time.assert_called_once_with(
-        custom_badput_event_type="elastic_reinitialization"
-    )
+    fake_recorder.record_elastic_wait_end_time.assert_called_once_with(event_type="elastic_slice_down")
+    fake_recorder.record_elastic_reinit_start_time.assert_called_once()
+    fake_recorder.record_elastic_slice_counts.assert_called_once()
     self.assertIs(elastic_utils.pending_reinit_recorder, fake_recorder)
     self.assertIsNone(elastic_utils.pending_elastic_event_type)
 
   def test_record_elastic_reinit_end(self):
     """Tests recording end of elastic reinit."""
     fake_recorder = Mock()
     elastic_utils.pending_reinit_recorder = fake_recorder
+    elastic_utils.elastic_manager = self.fake_manager
+    self.fake_manager.active_slice_indices = {0}
+    self.fake_manager.slice_to_devices = {0: [FakeDevice()], 1: [FakeDevice()]}
 
     elastic_utils.record_elastic_reinit_end()
 
-    fake_recorder.record_custom_badput_event_end_time.assert_called_once_with(
-        custom_badput_event_type="elastic_reinitialization"
-    )
+    fake_recorder.record_elastic_reinit_end_time.assert_called_once()
+    fake_recorder.record_elastic_slice_counts.assert_called_once()
     self.assertIsNone(elastic_utils.pending_reinit_recorder)
 
   def test_record_elastic_reinit_end_on_cold_start(self):
@@ -401,6 +433,8 @@ 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)