add a new record_slice_state helper that queries live slice counts and logs them to the GoodputRecorder

Author: YixuanWang-99

src/maxtext/utils/elastic_utils.py

@@ -28,13 +28,30 @@
 pending_elastic_event_type = None
 
 
+def record_slice_state(recorder, active_slices_override: int | None = None) -> None:
+  """Queries live slice counts and logs them to the GoodputRecorder."""
+  if recorder is None 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(elastic_manager.slice_to_devices)
+
+  if recorder:
+    recorder.record_slice_state(
+        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)
+    record_slice_state(recorder, active_slices_override=0)
 
 
 def record_elastic_wait_end_and_reinit_start(recorder) -> None:
@@ -47,6 +64,7 @@ def record_elastic_wait_end_and_reinit_start(recorder) -> 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')
+    record_slice_state(recorder)
   pending_reinit_recorder = recorder
 
 
@@ -57,6 +75,7 @@ def record_elastic_reinit_end() -> None:
     pending_reinit_recorder.record_custom_badput_event_end_time(
         custom_badput_event_type='elastic_reinitialization'
     )
+    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_slice_state = Mock()
+    
+    elastic_utils.record_slice_state(fake_recorder)
+    
+    fake_recorder.record_slice_state.assert_called_once_with(
+        available_slices=2,
+        active_slices=1,
+        total_slices=2
+    )
+
+    fake_recorder.record_slice_state.reset_mock()
+    elastic_utils.record_slice_state(fake_recorder, active_slices_override=0)
+    fake_recorder.record_slice_state.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]
@@ -156,6 +188,8 @@ def test_live_devices_pathways(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 = (
@@ -315,6 +349,7 @@ def test_record_elastic_event_start(self):
     fake_recorder.record_custom_badput_event_start_time.assert_called_once_with(
         custom_badput_event_type='elastic_slice_down'
     )
+    fake_recorder.record_slice_state.assert_called_once()
     self.assertEqual(elastic_utils.pending_elastic_event_type, 'elastic_slice_down')
 
   def test_record_elastic_event_start_scale_up(self):
@@ -329,6 +364,7 @@ def test_record_elastic_event_start_scale_up(self):
     fake_recorder.record_custom_badput_event_start_time.assert_called_once_with(
         custom_badput_event_type='elastic_scale_up'
     )
+    fake_recorder.record_slice_state.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."""
@@ -354,6 +390,7 @@ def test_record_elastic_wait_end_and_reinit_start(self):
     fake_recorder.record_custom_badput_event_start_time.assert_called_once_with(
         custom_badput_event_type='elastic_reinitialization'
     )
+    fake_recorder.record_slice_state.assert_called_once()
     self.assertIs(elastic_utils.pending_reinit_recorder, fake_recorder)
     self.assertIsNone(elastic_utils.pending_elastic_event_type)
 
@@ -367,6 +404,7 @@ def test_record_elastic_reinit_end(self):
     fake_recorder.record_custom_badput_event_end_time.assert_called_once_with(
         custom_badput_event_type='elastic_reinitialization'
     )
+    fake_recorder.record_slice_state.assert_called_once()
     self.assertIsNone(elastic_utils.pending_reinit_recorder)
 
   def test_record_elastic_reinit_end_on_cold_start(self):
@@ -389,6 +427,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)