Squashed commits for GoodputRecorder slice state logic

Author: YixuanWang-99

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:
+  """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)
+
+  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)
+    record_slice_state(recorder, active_slices_override=0)
 
 
 def record_elastic_wait_end_and_reinit_start(recorder) -> None:
@@ -46,6 +62,8 @@ def record_elastic_wait_end_and_reinit_start(recorder) -> None:
   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')
+    record_slice_state(recorder)
     recorder.record_custom_badput_event_start_time(custom_badput_event_type="elastic_reinitialization")
   pending_reinit_recorder = recorder
 
@@ -54,6 +72,10 @@ 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'
+    )
+    record_slice_state(pending_reinit_recorder)
     pending_reinit_recorder.record_custom_badput_event_end_time(custom_badput_event_type="elastic_reinitialization")
     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,6 +353,7 @@ 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()
 
@@ -327,12 +362,15 @@ 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_elastic_slice_counts.assert_called_once()
+    self.assertEqual(elastic_utils.pending_elastic_event_type, 'elastic_slice_down')
     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()
 
@@ -341,6 +379,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_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."""
@@ -355,6 +394,10 @@ def test_record_elastic_wait_end_and_reinit_start_noop_on_first_attempt(self):
 
   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()
 
@@ -366,19 +409,24 @@ 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_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_slice_counts.assert_called_once()
     self.assertIsNone(elastic_utils.pending_reinit_recorder)
 
   def test_record_elastic_reinit_end_on_cold_start(self):
@@ -401,6 +449,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)