Implement checkpoint-based elasticity using set-based slice tracking in MaxText.

PiperOrigin-RevId: 948165481

Author: lukebaumann

src/maxtext/utils/elastic_utils.py

@@ -211,7 +211,7 @@ def is_scale_up_event(config) -> bool:
   if elastic_enabled(config):
     ensure_elastic_manager_initialized(config)
     assert elastic_manager is not None
-    return elastic_manager.new_slice_event.is_set()
+    return bool(elastic_manager.available_inactive_slices)
 
   return False
 

tests/unit/elastic_utils_test.py

@@ -16,6 +16,7 @@
 
 import unittest
 from unittest.mock import create_autospec, Mock
+from absl.testing import parameterized
 
 
 from maxtext.utils import elastic_utils
@@ -46,7 +47,7 @@ def __init__(self):
     self.elastic_min_slice_count = 1
 
 
-class ElasticUtilsTest(unittest.TestCase):
+class ElasticUtilsTest(parameterized.TestCase):
   """Unit tests for Elastic Training utility functions."""
 
   def setUp(self):
@@ -66,7 +67,7 @@ def setUp(self):
     self.fake_logging = create_autospec(self.original_max_logging)
     self.fake_jax = create_autospec(self.original_jax)
     self.fake_manager = create_autospec(self.original_manager_class, instance=True)
-    self.fake_manager.new_slice_event = Mock()
+    self.fake_manager.available_inactive_slices = set()
 
     # Configure default behaviors if needed
     self.fake_pathwaysutils.is_pathways_backend_used.return_value = True
@@ -314,7 +315,7 @@ def wait_until_finished(self):
     cm = FakeCheckpointManager()
 
     elastic_utils.elastic_manager = self.fake_manager
-    self.fake_manager.new_slice_event.is_set.return_value = True
+    self.fake_manager.available_inactive_slices = {1}
 
     with self.assertRaises(ScaleUpSignalError):
       elastic_utils.maybe_elastic_scale_up(config, cm)
@@ -359,7 +360,7 @@ def test_elastic_retry_pre_callback_forwarded(self):
   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.available_inactive_slices = set()
     fake_recorder = Mock()
     config = FakeConfig()
 
@@ -373,7 +374,7 @@ def test_record_elastic_event_start(self):
   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.available_inactive_slices = {1}
     fake_recorder = Mock()
     config = FakeConfig()
 
@@ -447,6 +448,26 @@ def __setattr__(self, name, value):
     elastic_utils.ensure_elastic_manager_initialized(config)
     self.assertEqual(elastic_utils.elastic_manager, self.fake_manager)
 
+  @parameterized.parameters(
+      # Positive cases
+      ({1}, True),
+      ({0}, True),
+      ({1, 2}, True),
+      ({0, 3, 6}, True),
+      ({10, 25}, True),
+      # Negative cases
+      (set(), False),
+  )
+  def test_is_scale_up_event_with_set(self, available_inactive_slices, expected):
+    config = FakeConfig()
+    config.elastic_enabled = True
+    elastic_utils.elastic_manager = self.fake_manager
+
+    self.fake_manager.available_inactive_slices = available_inactive_slices
+    self.assertEqual(elastic_utils.is_scale_up_event(config), expected)
+
+
+
 
 if __name__ == "__main__":
   unittest.main()