Add elastic pause/resume functionality to MaxText.

PiperOrigin-RevId: 890594825

Author: abhinavclemson

src/maxtext/configs/base.yml

@@ -1181,3 +1181,8 @@ distill_temperature: 1.0
 # 0.0 value disables this feature.
 distill_beta: 0.0
 distill_layer_indices: None
+
+
+##### Elastic training parameters
+elastic_pause_resume: false
+elastic_timeout: 300

src/maxtext/configs/types.py

@@ -1516,6 +1516,15 @@ class Goodput(BaseModel):
   enable_gcp_step_deviation_metrics: bool = Field(True, description="Enable GCP step deviation metrics.")
 
 
+class ElasticTraining(BaseModel):
+  """Configuration for elastic training and fault tolerance."""
+
+  elastic_pause_resume: bool = Field(
+      False, description="Whether to enable elastic pause and resume functionality."
+  )
+  elastic_timeout: int = Field(3600, description="The timeout in seconds for elastic training operations.")
+
+
 class GcpMonitoring(BaseModel):
   """Configuration for GCP-specific workload monitoring."""
 
@@ -1897,6 +1906,7 @@ class MaxTextConfig(
     Checkpointing,
     OrbaxStorage,
     EmergencyCheckpointing,
+    ElasticTraining,
     # Data Types and Quantization
     DataTypes,
     Quantization,

src/maxtext/trainers/pre_train/train.py

@@ -41,6 +41,7 @@
 from maxtext.configs import pyconfig
 from maxtext.common.common_types import ShardMode
 from maxtext.utils.globals import EPS
+import maxtext.utils.elastic_utils as elastic_utils
 # Placeholder: internal
 
 # pylint: disable=too-many-positional-arguments
@@ -679,9 +680,18 @@ def run(config, recorder, diagnostic_config):
 def main(argv: Sequence[str]) -> None:
   config, recorder, diagnostic_config = initialize(argv)
   record_goodput(recorder, RECORD_JOB_START_TIME)
-  with maybe_monitor_goodput(config):
+
+  def train_func():
+    config, recorder, diagnostic_config = initialize(argv)
     run(config, recorder, diagnostic_config)
 
+  if config.elastic_pause_resume:
+    max_logging.log("Elastic Pause and Resume Enabled.")
+    train_func = elastic_utils.elastic_pause_resume(config)(train_func)
+
+  with maybe_monitor_goodput(config):
+    train_func()
+
 
 if __name__ == "__main__":
   app.run(main)

src/maxtext/utils/elastic_utils.py

@@ -0,0 +1,84 @@
+# Copyright 2023–2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utility functions for Elastic Training."""
+
+import functools
+import re
+import subprocess
+
+import maxtext.utils.max_logging as max_logging
+import pathwaysutils
+from pathwaysutils.elastic import manager
+
+
+elastic_manager: manager.Manager | None = None
+
+
+def elastic_mode_enabled(config) -> bool:
+  """Returns whether elastic mode is enabled."""
+  return (pathwaysutils.is_pathways_backend_used() and
+          config.elastic_pause_resume)
+
+
+def clean_up_checkpoints(checkpoint_dir: str):
+  """Cleans up incomplete checkpoints after an elastic event."""
+  max_logging.log(f"Elastic utils: Checking for incomplete checkpoint after an elastic event...")
+  checkpoint_dir = f"{checkpoint_dir}"
+
+  # 1. List the directory
+  result = subprocess.run(['gsutil', 'ls', checkpoint_dir], capture_output=True, text=True)
+
+  if result.returncode != 0:
+    max_logging.log("Failed to inspect checkpoint dir. Continuing")
+    return
+
+  # 2. Filter for directories ending in numbers/ (equivalent to your grep and sort)
+  checkpoints = [line for line in result.stdout.splitlines() if re.search(r'/\d+/$', line)]
+
+  if not checkpoints:
+    max_logging.log("Found no existing checkpoints. Continuing")
+    return
+
+  # Sort naturally (Version sort) and get the last one
+  checkpoints.sort(key=lambda x: [int(c) if c.isdigit() else c for c in re.split(r'(\d+)', x)])
+  latest_checkpoint = checkpoints[-1]
+
+  max_logging.log(f"Checking latest checkpoint: {latest_checkpoint}")
+
+  # 3. Check for commit_success file
+  # gsutil -q stat returns 0 if found, non-zero if not
+  stat_check = subprocess.run(['gsutil', '-q', 'stat', f"{latest_checkpoint}commit_success*"])
+
+  if stat_check.returncode != 0:
+    max_logging.log(f"No commit_success file found. Deleting {latest_checkpoint}...")
+    subprocess.run(['gsutil', '-m', 'rm', '-rf', latest_checkpoint])
+  else:
+    max_logging.log(f"Found commit_success file. Keeping {latest_checkpoint}.")
+
+
+def elastic_pause_resume(config, callback_fn=None):
+  """Pauses and resumes elastic training."""
+  cleanup_partial = functools.partial(
+      clean_up_checkpoints, config.checkpoint_dir
+  )
+  callback_fn = cleanup_partial if callback_fn is None else callback_fn
+  return elastic_manager.elastic_retry(
+      max_retries=10,
+      poll_interval=10,
+      timeout=config.elastic_timeout,
+      on_elastic_event_callback=callback_fn,
+  )
+
+

src/maxtext/utils/maxtext_utils.py

@@ -44,6 +44,7 @@
 from maxtext.utils import gcs_utils
 from maxtext.utils import max_logging
 from maxtext.utils import max_utils
+from maxtext.utils import elastic_utils
 from maxtext.utils import sharding
 
 OVERWRITE_WITH_GRADIENT = "_overwrite_with_gradient"
@@ -1323,7 +1324,7 @@ def add_config_to_summary_writer(config, summary_writer):
 def create_device_mesh(config, devices=None):
   """Creates a device mesh with each slice in its own data parallel group. If there is only one slice, uses two replicas"""
   if devices is None:
-    devices = jax.devices()
+    devices = elastic_utils.live_devices() if not elastic_utils.elastic_mode_enabled(config)  else jax.devices()
   if config.subslice_shape and config.enable_single_controller and config.num_slices == 1:
     max_logging.log(f"Trying to create a subslice with shape: {config.subslice_shape}")
     subslice_shape = tuple(int(x) for x in config.subslice_shape.split(","))