Squashed commits: update GoodputRecorder and ElasticRecorder

Author: YixuanWang-99

src/maxtext/common/goodput.py

@@ -22,6 +22,7 @@
 import contextlib
 import jax
 from enum import Enum
+import pathwaysutils
 from maxtext.utils import max_logging
 from maxtext.common.gcloud_stub import goodput_modules
 
@@ -66,18 +67,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(
+    monitor_class = monitoring.GoodputMonitor
+
+    if config.elastic_enabled:
+      try:
+        from ml_goodput_measurement import monitoring_elastic
+        monitor_class = monitoring_elastic.ElasticGoodputMonitor
+      except ImportError:
+        max_logging.log("Elastic monitor failed!")
+        pass
+
+    kwargs = dict(
         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,
     )
+    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
@@ -123,6 +137,19 @@ def create_goodput_recorder(config):
     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)
+    
+    # Detect if we should use the elastic-aware recorder
+    recorder = None
+    if config.elastic_enabled:
+      try:
+        from ml_goodput_measurement import goodput_elastic
+        from maxtext.utils import elastic_utils
+        recorder = goodput_elastic.ElasticGoodputRecorder(config.run_name, logger_name, jax.process_index() == 0)      
+        elastic_utils.record_slice_state(recorder)
+      except ImportError:
+        pass
+    
+    if recorder is None:
+      recorder = goodput.GoodputRecorder(config.run_name, logger_name, jax.process_index() == 0)
     return recorder
   return None

src/maxtext/utils/elastic_utils.py

@@ -28,13 +28,29 @@
 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') or not pathwaysutils.is_pathways_backend_used() or elastic_manager is None:
+    return
+
+  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 +61,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
 
 

tests/unit/elastic_utils_test.py

@@ -70,6 +70,8 @@ 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_jax.process_index.return_value = 0
 
     # Inject fakes into elastic_utils namespace
@@ -97,9 +99,37 @@ 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 +174,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 +200,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 +353,32 @@ 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"
     )
-    self.assertEqual(elastic_utils.pending_elastic_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 +387,41 @@ 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 +444,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)