Merge pull request #3235 from AI-Hypercomputer:pw-elastic-training-gp-fixes

PiperOrigin-RevId: 878677943

Author: Google-ML-Automation

.github/CODEOWNERS

@@ -1,4 +1,4 @@
-* @gobbleturk @khatwanimohit @bvandermoon @vipannalla @RissyRan @richjames0 @gagika @shralex @SurbhiJainUSC @hengtaoguo @A9isha @aireenmei @NuojCheng @jiangjy1982 @suexu1025 @NicoGrande @jesselu-google
+* @gobbleturk @khatwanimohit @bvandermoon @vipannalla @RissyRan @richjames0 @gagika @shralex @SurbhiJainUSC @hengtaoguo @A9isha @aireenmei @NuojCheng @jiangjy1982 @suexu1025 @NicoGrande @jesselu-google @dipannita08 @igorts-git
 
 # Model bring-up
 src/MaxText/assets @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande

src/maxtext/common/goodput.py

@@ -1,4 +1,4 @@
-# Copyright 2023–2025 Google LLC
+# Copyright 2023–2026 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -36,6 +36,12 @@ class GoodputEvent(Enum):
   STEP = "step"
 
 
+# Recorder method name constants for explicit job start/end recording.
+# Derived from the enum so they stay in sync if the value ever changes.
+RECORD_JOB_START_TIME = f"record_{GoodputEvent.JOB.value}_start_time"
+RECORD_JOB_END_TIME = f"record_{GoodputEvent.JOB.value}_end_time"
+
+
 @contextlib.contextmanager
 def maybe_monitor_goodput(config):
   """Monitor cumulative goodput if enabled on the lead host.
@@ -83,18 +89,22 @@ def maybe_monitor_goodput(config):
 
 @contextlib.contextmanager
 def maybe_record_goodput(recorder, event_name, *args):
-  """Record goodput if `enable_goodput_recording=True`."""
+  """Record goodput if `enable_goodput_recording=True`.
+
+  The end-time event is only recorded when the wrapped block exits without
+  raising an exception (i.e. the event truly completed).  Callers that need
+  explicit end-time control — e.g. GoodputEvent.JOB under elastic training
+  where the elastic manager may suppress the JAX exception internally —
+  should call record_goodput directly rather than using this context manager.
+  """
+  record_goodput(recorder, f"record_{event_name.value}_start_time", *args)
+  completed = False
   try:
-    start_event_name = f"record_{event_name.value}_start_time"
-    record_goodput(recorder, start_event_name, *args)
     yield
-  except BaseException:  # pylint: disable=W0706
-    raise
-  else:
-    end_event_name = f"record_{event_name.value}_end_time"
-    record_goodput(recorder, end_event_name, *args)
+    completed = True
   finally:
-    pass
+    if completed:
+      record_goodput(recorder, f"record_{event_name.value}_end_time", *args)
 
 
 def record_goodput(recorder, event_name, *args):

src/maxtext/experimental/rl/grpo_trainer.py

@@ -1,4 +1,4 @@
-# Copyright 2023–2025 Google LLC
+# Copyright 2023–2026 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -74,9 +74,12 @@
 from maxtext.common.data_loader import DataLoader
 from maxtext.common.goodput import (
     GoodputEvent,
+    RECORD_JOB_END_TIME,
+    RECORD_JOB_START_TIME,
     create_goodput_recorder,
     maybe_monitor_goodput,
     maybe_record_goodput,
+    record_goodput,
 )
 from maxtext.experimental.rl import grpo_input_pipeline
 from maxtext.experimental.rl import grpo_utils
@@ -795,6 +798,7 @@ def generation_worker_fn(
   )
   generation_thread.start()
 
+  _job_completed_gracefully = False
   try:
     last_step_completion = datetime.datetime.now()
     for step in np.arange(start_step, config.steps):
@@ -881,9 +885,13 @@ def generation_worker_fn(
       elif checkpoint_manager is not None:
         # in case the last checkpoint_period checkpoint is still in progress
         checkpoint_manager.wait_until_finished()
+    _job_completed_gracefully = True
   except exceptions.StopTraining as e:
     max_logging.log(f"Training stopped: {str(e)}")
+    _job_completed_gracefully = True
   finally:
+    if _job_completed_gracefully:
+      record_goodput(recorder, RECORD_JOB_END_TIME)
     metric_logger.flush_metrics_and_cleanup()
     max_logging.log("Training loop finished or exited. Signaling generation worker to stop.")
     stop_event.set()
@@ -955,8 +963,9 @@ def main(argv: Sequence[str]) -> None:
   )
   diagnostic_config = diagnostic_configuration.DiagnosticConfig(debug_config)
 
+  record_goodput(recorder, RECORD_JOB_START_TIME)
   with diagnostic.diagnose(diagnostic_config):
-    with maybe_record_goodput(recorder, GoodputEvent.JOB), maybe_monitor_goodput(config):
+    with maybe_monitor_goodput(config):
       train_loop(config, config_inference, recorder)
 
 

src/maxtext/trainers/post_train/sft/train_sft.py

@@ -53,9 +53,12 @@
 from maxtext.trainers.pre_train.train import loss_fn
 from maxtext.common.goodput import (
     GoodputEvent,
+    RECORD_JOB_END_TIME,
+    RECORD_JOB_START_TIME,
     create_goodput_recorder,
     maybe_monitor_goodput,
     maybe_record_goodput,
+    record_goodput,
 )
 from maxtext.optimizers import optimizers
 from maxtext.trainers.post_train.sft import hooks
@@ -181,7 +184,13 @@ def train(mt_config, goodput_recorder=None):
     goodput_recorder: An optional GoodputRecorder to record performance metrics.
   """
   trainer, mesh = setup_trainer_state(mt_config, goodput_recorder)
-  trainer = train_model(mt_config, trainer, mesh)
+  _job_completed_gracefully = False
+  try:
+    trainer = train_model(mt_config, trainer, mesh)
+    _job_completed_gracefully = True
+  finally:
+    if _job_completed_gracefully:
+      record_goodput(goodput_recorder, RECORD_JOB_END_TIME)
   return trainer, mesh
 
 
@@ -198,8 +207,8 @@ def main(argv: Sequence[str]) -> None:
   max_utils.print_system_information()
 
   goodput_recorder = create_goodput_recorder(mt_config)
-
-  with maybe_record_goodput(goodput_recorder, GoodputEvent.JOB), maybe_monitor_goodput(mt_config):
+  record_goodput(goodput_recorder, RECORD_JOB_START_TIME)
+  with maybe_monitor_goodput(mt_config):
     train(mt_config, goodput_recorder)
 
 

src/maxtext/trainers/post_train/sft/train_sft_deprecated.py

@@ -1,4 +1,4 @@
-# Copyright 2023–2025 Google LLC
+# Copyright 2023–2026 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -37,9 +37,12 @@
 from maxtext.common.data_loader import DataLoader
 from maxtext.common.goodput import (
     GoodputEvent,
+    RECORD_JOB_END_TIME,
+    RECORD_JOB_START_TIME,
     create_goodput_recorder,
     maybe_monitor_goodput,
     maybe_record_goodput,
+    record_goodput,
 )
 from maxtext.common.metric_logger import MetricLogger
 from maxtext.utils import exceptions
@@ -90,6 +93,7 @@ def train_loop(config, recorder, state=None):
   # Write train config params, num model params, and XLA flags to tensorboard
   metric_logger.write_setup_info_to_tensorboard(state.params)
 
+  _job_completed_gracefully = False
   try:
     last_step_completion = datetime.datetime.now()
     for step in np.arange(start_step, config.steps):
@@ -147,9 +151,13 @@ def train_loop(config, recorder, state=None):
     elif checkpoint_manager is not None:
       # in case the last checkpoint_period checkpoint is still in progress
       checkpoint_manager.wait_until_finished()
+    _job_completed_gracefully = True
   except exceptions.StopTraining as e:
     max_logging.log(f"Training stopped: {str(e)}")
+    _job_completed_gracefully = True
   finally:
+    if _job_completed_gracefully:
+      record_goodput(recorder, RECORD_JOB_END_TIME)
     metric_logger.flush_metrics_and_cleanup()
 
   return state
@@ -172,7 +180,8 @@ def main(argv: Sequence[str]) -> None:
   os.environ["TFDS_DATA_DIR"] = config.dataset_path
 
   recorder = create_goodput_recorder(config)
-  with maybe_record_goodput(recorder, GoodputEvent.JOB), maybe_monitor_goodput(config):
+  record_goodput(recorder, RECORD_JOB_START_TIME)
+  with maybe_monitor_goodput(config):
     train_loop(config, recorder)
 
 

src/maxtext/trainers/pre_train/train.py

@@ -1,4 +1,4 @@
-# Copyright 2023–2025 Google LLC
+# Copyright 2023–2026 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -48,9 +48,12 @@
 from maxtext.common import checkpointing, profiler
 from maxtext.common.goodput import (
     GoodputEvent,
+    RECORD_JOB_END_TIME,
+    RECORD_JOB_START_TIME,
     create_goodput_recorder,
     maybe_monitor_goodput,
     maybe_record_goodput,
+    record_goodput,
 )
 from maxtext.common.gcloud_stub import cloud_diagnostics as _cloud_diag, is_decoupled
 from maxtext.common.gcloud_stub import vertex_tensorboard_modules
@@ -493,6 +496,7 @@ def train_loop(config, recorder, state=None):
   # Write train config params, num model params, and XLA flags to tensorboard
   metric_logger.write_setup_info_to_tensorboard(state.params)
 
+  _job_completed_gracefully = False
   try:
     last_step_completion = datetime.datetime.now()
     for step in np.arange(start_step, config.steps):
@@ -558,9 +562,13 @@ def train_loop(config, recorder, state=None):
     if checkpoint_manager is not None:
       # in case the last checkpoint_period checkpoint is still in progress
       checkpoint_manager.wait_until_finished()
+    _job_completed_gracefully = True
   except exceptions.StopTraining as e:
     max_logging.log(f"Training stopped: {str(e)}")
+    _job_completed_gracefully = True
   finally:
+    if _job_completed_gracefully:
+      record_goodput(recorder, RECORD_JOB_END_TIME)
     metric_logger.flush_metrics_and_cleanup()
 
   return state
@@ -623,14 +631,14 @@ def run(config, recorder, diagnostic_config):
 
   with (
       diagnostics_context,
-      maybe_record_goodput(recorder, GoodputEvent.JOB),
       max_utils.maybe_get_transformer_engine_context(config),
   ):
     train_loop(config, recorder)
 
 
 def main(argv: Sequence[str]) -> None:
   config, recorder, diagnostic_config = initialize(argv)
+  record_goodput(recorder, RECORD_JOB_START_TIME)
   with maybe_monitor_goodput(config):
     run(config, recorder, diagnostic_config)
 

tests/unit/goodput_utils_test.py

@@ -1,4 +1,4 @@
-# Copyright 2023–2025 Google LLC
+# Copyright 2023–2026 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -20,7 +20,15 @@
 import pytest
 
 from maxtext.configs import pyconfig
-from maxtext.common.goodput import create_goodput_recorder, maybe_monitor_goodput, maybe_record_goodput, GoodputEvent
+from maxtext.common.goodput import (
+    GoodputEvent,
+    RECORD_JOB_END_TIME,
+    RECORD_JOB_START_TIME,
+    create_goodput_recorder,
+    maybe_monitor_goodput,
+    maybe_record_goodput,
+    record_goodput,
+)
 from tests.utils.test_helpers import get_test_config_path, get_test_base_output_directory
 
 pytestmark = [pytest.mark.external_training]
@@ -80,6 +88,59 @@ def test_monitor_goodput(self, mock_start_goodput_uploader, mock_stop_goodput_up
       mock_start_goodput_uploader.assert_called()
     mock_stop_goodput_uploader.assert_called()
 
+  def test_job_recording_constants(self):
+    """Constants must map to the recorder method names."""
+    self.assertEqual(RECORD_JOB_START_TIME, "record_job_start_time")
+    self.assertEqual(RECORD_JOB_END_TIME, "record_job_end_time")
+
+  @mock.patch("ml_goodput_measurement.goodput.GoodputRecorder.record_job_end_time")
+  @mock.patch("ml_goodput_measurement.goodput.GoodputRecorder.record_job_start_time")
+  @mock.patch("google.cloud.logging.Client")
+  def test_explicit_job_recording_graceful_completion(
+      self, mock_cloud_logger, mock_record_job_start_time, mock_record_job_end_time
+  ):
+    """Both start and end are recorded when the job completes gracefully."""
+    mock_cloud_logger.return_value = mock.MagicMock()
+    recorder = create_goodput_recorder(self.config)
+
+    record_goodput(recorder, RECORD_JOB_START_TIME)
+    _job_completed_gracefully = False
+    try:
+      _job_completed_gracefully = True
+    finally:
+      if _job_completed_gracefully:
+        record_goodput(recorder, RECORD_JOB_END_TIME)
+
+    mock_record_job_start_time.assert_called_once()
+    mock_record_job_end_time.assert_called_once()
+
+  @mock.patch("ml_goodput_measurement.goodput.GoodputRecorder.record_job_end_time")
+  @mock.patch("ml_goodput_measurement.goodput.GoodputRecorder.record_job_start_time")
+  @mock.patch("google.cloud.logging.Client")
+  def test_explicit_job_recording_elastic_restart(
+      self, mock_cloud_logger, mock_record_job_start_time, mock_record_job_end_time
+  ):
+    """Only start is recorded when the elastic manager handles the error internally.
+
+    This simulates the elastic-restart scenario: the manager catches the JAX
+    exception inside train_loop, so the loop exits without raising.  The
+    _job_completed_gracefully flag is never set, so record_job_end_time must
+    not be called.
+    """
+    mock_cloud_logger.return_value = mock.MagicMock()
+    recorder = create_goodput_recorder(self.config)
+
+    record_goodput(recorder, RECORD_JOB_START_TIME)
+    _job_completed_gracefully = False
+    try:
+      pass  # Elastic manager caught and suppressed the exception.
+    finally:
+      if _job_completed_gracefully:
+        record_goodput(recorder, RECORD_JOB_END_TIME)
+
+    mock_record_job_start_time.assert_called_once()
+    mock_record_job_end_time.assert_not_called()
+
 
 if __name__ == "__main__":
   unittest.main()