rename and add initial call

Author: YixuanWang-99

src/maxtext/common/goodput.py

@@ -143,7 +143,9 @@ def create_goodput_recorder(config):
     if use_elastic:
       try:
         from ml_goodput_measurement import goodput_elastic
-        recorder = goodput_elastic.ElasticGoodputRecorder(config.run_name, logger_name, jax.process_index() == 0)
+        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
 

src/maxtext/utils/elastic_utils.py

@@ -49,7 +49,7 @@ def record_elastic_event_start(recorder, config) -> None:
   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_elastic_wait_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)
 
 
@@ -61,8 +61,8 @@ 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_elastic_wait_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
 
@@ -71,9 +71,7 @@ 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