address comments

Author: YixuanWang-99

src/maxtext/common/goodput.py

@@ -22,8 +22,8 @@
 import contextlib
 import jax
 from enum import Enum
+import pathwaysutils
 from maxtext.utils import max_logging
-from maxtext.utils import pathwaysutils
 from maxtext.common.gcloud_stub import goodput_modules
 
 goodput, monitoring, _GOODPUT_STUB = goodput_modules()
@@ -68,12 +68,13 @@ def maybe_monitor_goodput(config):
         enable_gcp_step_deviation_metrics=config.enable_gcp_step_deviation_metrics,
     )
     monitor_class = monitoring.GoodputMonitor
-    use_elastic = getattr(config, 'elastic_enabled', False) and pathwaysutils.is_pathways_backend_used()
-    if use_elastic:
+
+    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(
@@ -138,9 +139,8 @@ def create_goodput_recorder(config):
     logger_name = f"goodput_{config.run_name}"
 
     # Detect if we should use the elastic-aware recorder
-    use_elastic = getattr(config, 'elastic_enabled', False) and pathwaysutils.is_pathways_backend_used()
     recorder = None
-    if use_elastic:
+    if config.elastic_enabled:
       try:
         from ml_goodput_measurement import goodput_elastic
         from maxtext.utils import elastic_utils

src/maxtext/utils/elastic_utils.py

@@ -35,7 +35,7 @@ def record_slice_state(recorder, active_slices_override: int | None = None) -> N
 
   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)
+  total_slices = len(pathwaysutils.elastic.get_slice_to_devices(jax.devices()))
 
   recorder.record_elastic_slice_counts(
       available_slices=available_slices,

tests/unit/elastic_utils_test.py

@@ -359,12 +359,11 @@ def test_record_elastic_event_start(self):
 
     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')
-    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."""
@@ -376,8 +375,8 @@ def test_record_elastic_event_start_scale_up(self):
 
     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()
 
@@ -388,8 +387,8 @@ 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):
@@ -403,12 +402,10 @@ def test_record_elastic_wait_end_and_reinit_start(self):
 
     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)
@@ -423,9 +420,7 @@ def test_record_elastic_reinit_end(self):
 
     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)