Repository: AI-Hypercomputer/maxtext
Pull request #1310: Poc elastic training
Author: login: lukebaumann
Labels: name: pull ready
Created: 2025-02-25T21:47:57Z | Merged: 2025-04-17T22:25:51Z

# Description

Elastic training is enabled for data parallel maxtext training workloads when using Pathways on Cloud.

An elastic manager can be configured to support different host snapshot periods and number, consecutive retries, total slices losses.

Documentation will be shared after Cloud Next for more information about how to deploy with Pathways on Cloud with elasticity enabled.

# Tests

Manually tested. Will add an airflow test after merging.

# Checklist

Before submitting this PR, please make sure (put X in square brackets):
- [X] I have performed a self-review of my code.
- [X] I have necessary comments in my code, particularly in hard-to-understand areas.
- [X] I have run end-to-end tests tests and provided workload links above if applicable.
- [X] I have made or will make corresponding changes to the doc if needed.



Diff:
diff --git a/MaxText/elastic_train.py b/MaxText/elastic_train.py
new file mode 100644
index 0000000000..fe07b0fb16
--- /dev/null
+++ b/MaxText/elastic_train.py
@@ -0,0 +1,515 @@
+# Copyright 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.
+
+"""Training loop for elastic training model.
+
+Elastic training via Pathways on Cloud allows for slices to go down without
+crashing the main program. In this way, elastic events (slices going up and
+down) can be caught or polled for and reacted to within an elastic_handler
+function without causing the whole workload to restart. This involves
+creating a new mesh with the available slices, reinitializing variables,
+and recompiling functions in addition to restoring from a host offloaded
+snapshot of the state.
+
+The purpose of this training loop is to serve as an example for how to
+support elasticity and is actively in development. As such, there are some
+performance optimizations that have yet to be added as well as some features
+not supported.
+
+Current limitations:
+- The host offloaded snapshot is currently blocking
+- Elastic event handling during async checkpointing is not tested
+- Elastic event handling during profiling is not tested
+- Elastic manager configuration values are hard coded
+- Debug logging statements for elasticity are hard coded as enabled
+
+See https://github.com/AI-Hypercomputer/pathways-utils/tree/main/pathwaysutils/elastic
+for more details about the elastic manager.
+"""
+from collections.abc import Sequence
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from absl import app
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+from flax.linen import partitioning as nn_partitioning
+import jax
+from ml_goodput_measurement import monitoring
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+import tensorflow as tf
+
+from MaxText import checkpointing
+from MaxText import max_utils
+from MaxText import maxtext_utils
+from MaxText import max_logging
+from MaxText import profiler
+from MaxText import pyconfig
+from MaxText.gcp_workload_monitor import GCPWorkloadMonitor
+from MaxText.input_pipeline.input_pipeline_interface import create_data_iterator
+from MaxText.metric_logger import MetricLogger
+from MaxText.train import check_example_batch
+from MaxText.train import create_goodput_recorder
+from MaxText.train import get_first_step
+from MaxText.train import load_next_batch
+from MaxText.train import record_goodput
+from MaxText.train import record_scalar_metrics
+from MaxText.train import save_checkpoint
+from MaxText.train import setup_mesh_and_model
+from MaxText.train import setup_train_loop
+from MaxText.train import train_step
+from MaxText.train import validate_train_config
+from MaxText.vertex_tensorboard import VertexTensorboardManager
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.INFO)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    elastic_manager,
+    checkpoint_manager,
+):
+  """Reconfigures the workload onto the currently available slices.
+
+  This is called by the elastic manager's maybe_reshard_up/down
+  functions and is responsible for creating a new mesh,
+  reinitializing the state and any objects that depend on the mesh.
+
+  It returns all of the reinitialized objects.
+
+  maybe_reshard_up/down take this function and its arguments and if
+  there is an elastic event, those functions will call this function
+  and return its returns.
+  """
+  # We reuse setup_mesh_and_model because it contains most of the
+  # reconfiguration. Depending on the configuration, the checkpoint
+  # manager depends on the mesh and must be recreated. Therefore, we
+  # close the previous checkpoint manager and get a new checkpoint
+  # manager from setup_mesh_and_model.
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(elastic_manager.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, elastic_manager.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = elastic_manager.get_resharded_snapshot(mesh)
+
+      # We do not want to restore from the previous checkpoint but instead
+      # restore from the host offloaded snapshot.
+      if checkpoint_manager is not None:
+        latest_step = checkpoint_manager.latest_step()
+
+        # If we checkpointed after the latest snapshot, the checkpoint manager
+        # will try to take another checkpoint and fail because it already
+        # exists. Therefore, we delete the checkpoint and let the checkpoint
+        # manager re-take the checkpoint.
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
+          checkpoint_manager.delete(latest_step)
+
+      state, _, state_mesh_shardings, data_iterator = maxtext_utils.setup_training_state(
+          model,
+          data_iterator,
+          tx,
+          config,
+          jax.random.fold_in(init_rng, step),
+          mesh,
+          checkpoint_manager=None,
+      )
+
+      state = state.replace(**snapshot)
+      state = state.replace(step=state.step.at[None].set(step))
+
+      (
+          functional_train,
+          in_shard_train,
+          out_shard_train,
+          static_argnums_train,
+          donate_argnums_train,
+      ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+      p_train_step = jax.jit(
+          functional_train,
+          in_shardings=in_shard_train,
+          out_shardings=out_shard_train,
+          static_argnums=static_argnums_train,
+          donate_argnums=donate_argnums_train,
+      )
+
+      example_batch = None
+      metric_logger = MetricLogger(writer, config)
+
+      jax.block_until_ready(state)
+
+  return (
+      config,
+      step,
+      state,
+      mesh,
+      checkpoint_manager,
+      data_iterator,
+      p_train_step,
+      example_batch,
+      learning_rate_schedule,
+      metric_logger,
+      writer,
+  )
+
+
+def train_loop(config, elastic_manager, state=None):
+  """Main Training loop.
+  Args:
+    config:
+    state:
+    ckpt_path:
+  Returns:
+  """
+  # Create a GoodputRecorder to log information
+  recorder = create_goodput_recorder(config)
+  record_goodput(recorder, config, recorder.record_job_start_time if recorder else None)
+
+  (
+      init_rng,
+      writer,
+      checkpoint_manager,
+      state_mesh_shardings,
+      model,
+      mesh,
+      learning_rate_schedule,
+      data_iterator,
+      _,
+      state,
+  ) = setup_train_loop(config)
+
+  # pylint: disable=line-too-long
+  (
+      functional_train,
+      in_shard_train,
+      out_shard_train,
+      static_argnums_train,
+      donate_argnums_train,
+  ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+  num_model_parameters = max_utils.calculate_num_params_from_pytree(state.params)
+  max_logging.log(f"number parameters: {num_model_parameters/1e9:.3f} billion")
+  per_device_tflops, _, _ = maxtext_utils.calculate_tflops_training_per_device(config)
+  per_device_tokens = maxtext_utils.calculate_tokens_training_per_device(config)
+
+  # Write train config params, num model params, and XLA flags to tensorboard
+  max_utils.add_text_to_summary_writer("num_model_parameters", str(num_model_parameters), writer)
+  max_utils.add_text_to_summary_writer("libtpu_init_args", os.environ["LIBTPU_INIT_ARGS"], writer)
+  maxtext_utils.add_config_to_summary_writer(config, writer)
+
+  p_train_step = jax.jit(
+      functional_train,
+      in_shardings=in_shard_train,
+      out_shardings=out_shard_train,
+      static_argnums=static_argnums_train,
+      donate_argnums=donate_argnums_train,
+  )
+  running_gcs_metrics = [] if config.gcs_metrics else None
+
+  start_step = get_first_step(state)  # this is the start_step for training
+  prof = profiler.Profiler(config, offset_step=start_step)
+  first_profiling_step = prof.start_initial_profile_step
+  if config.profiler != "" and first_profiling_step >= config.steps:
+    raise ValueError("Profiling requested but initial profiling step set past training final step")
+  last_profiling_step = prof.finished_initial_profile_step
+
+  example_batch = None
+  last_step_completion = datetime.datetime.now()
+
+  performance_metric_queue = None
+  if config.report_heartbeat_metric_for_gcp_monitoring or config.report_performance_metric_for_gcp_monitoring:
+    gcp_workload_monitor = GCPWorkloadMonitor(config.run_name)
+    if config.report_heartbeat_metric_for_gcp_monitoring:
+      gcp_workload_monitor.start_heartbeat_reporting_thread(config.heartbeat_reporting_interval_in_seconds)
+    if config.report_performance_metric_for_gcp_monitoring:
+      performance_metric_queue = queue.Queue()
+      gcp_workload_monitor.start_performance_reporting_thread(performance_metric_queue)
+
+  metric_logger = MetricLogger(writer, config)
+  step = start_step
+
+  elastic_manager.maybe_snapshot(
+      step,
+      snapshot={
+          "params": state.params,
+          "opt_state": state.opt_state,
+      },
+      force=True,
+      block=True,
+  )
+
+  # Using while loop instead of a for loop because with elasticity
+  # the step is restored back to the latest snapshot when a slice is lost
+  while step < config.steps:
+    try:
+      if step == first_profiling_step or prof.should_activate_periodic_profile(step):
+        optional_postfix = f"step_{step}" if config.profile_periodically_period > 0 else ""
+        prof.activate(blocking_object=state, optional_postfix=optional_postfix)
+
+      max_logging.log(f"{step=} {elastic_manager.elastic_down_event_count=} {elastic_manager.good_slice_count=}")
+      with mesh, nn_partitioning.axis_rules(config.logical_axis_rules), jax.default_device(elastic_manager.default_device):
+        with jax.profiler.StepTraceAnnotation("train", step_num=step):
+          record_goodput(recorder, config, recorder.record_data_loading_start_time if recorder else None)
+          try:
+            example_batch = load_next_batch(data_iterator, example_batch, config)
+          except Exception as e:  # pylint: disable=broad-except
+            max_logging.log(f"load_next_batch failed, you may have run out of data. Error message: {e}")
+            break
+          record_goodput(recorder, config, recorder.record_data_loading_end_time if recorder else None)
+          check_example_batch(config, example_batch=example_batch)
+          # pylint: disable=not-callable
+          nextrng = jax.jit(jax.random.fold_in)(init_rng, step)
+          record_goodput(recorder, config, recorder.record_step_start_time if recorder else None, step)
+          state, metrics = p_train_step(state, example_batch, nextrng)
+
+        step_time_delta = datetime.datetime.now() - last_step_completion
+        last_step_completion = datetime.datetime.now()
+        record_scalar_metrics(metrics, step_time_delta, per_device_tflops, learning_rate_schedule(step), per_device_tokens)
+        if performance_metric_queue:
+          performance_metric_queue.put(step_time_delta.total_seconds())
+
+        if checkpoint_manager is not None:
+          state_to_save = state
+          if save_checkpoint(checkpoint_manager, int(step), state_to_save, config.dataset_type, data_iterator, config):
+            checkpointing.print_save_message(step, config.async_checkpointing)
+
+          # Upon preemption, exit when and only when all ongoing saves are complete.
+          if checkpoint_manager.reached_preemption(step):
+            checkpoint_manager.wait_until_finished()
+            sys.exit()
+
+        metric_logger.write_metrics(running_gcs_metrics, metrics, step)
+
+        if step == last_profiling_step or prof.should_deactivate_periodic_profile(step):
+          prof.deactivate(blocking_object=state)
+
+      elastic_manager.maybe_snapshot(
+          step=step,
+          snapshot={
+              "params": state.params,
+              "opt_state": state.opt_state,
+          },
+          block=True,
+      )
+
+      ret = elastic_manager.maybe_reshard_up(
+          step=step,
+          snapshot={
+              "params": state.params,
+              "opt_state": state.opt_state,
+          },
+          elastic_handler=elastic_handler,
+          handler_kwargs={
+              "config": config,
+              "elastic_manager": elastic_manager,
+              "checkpoint_manager": checkpoint_manager,
+          },
+      )
+      if ret is not None:
+        (
+            config,
+            step,
+            state,
+            mesh,
+            checkpoint_manager,
+            data_iterator,
+            p_train_step,
+            example_batch,
+            learning_rate_schedule,
+            metric_logger,
+            writer,
+        ) = ret
+
+      if step == start_step:
+        max_utils.print_mem_stats("After params initialized")
+
+      step += 1
+
+    except jax.errors.JaxRuntimeError as error:
+      ret = elastic_manager.maybe_reshard_down(
+          error=error,
+          elastic_handler=elastic_handler,
+          handler_kwargs={
+              "config": config,
+              "elastic_manager": elastic_manager,
+              "checkpoint_manager": checkpoint_manager,
+          },
+      )
+      if ret is not None:
+        (
+            config,
+            step,
+            state,
+            mesh,
+            checkpoint_manager,
+            data_iterator,
+            p_train_step,
+            example_batch,
+            learning_rate_schedule,
+            metric_logger,
+            writer,
+        ) = ret
+
+  if checkpoint_manager is not None:
+    if (int(state.step) - 1) % config.checkpoint_period != 0:
+      try:
+        state_to_save = state
+        if save_checkpoint(
+            checkpoint_manager,
+            int(state.step) - 1,
+            state_to_save,
+            config.dataset_type,
+            data_iterator,
+            config,
+            force=True,
+        ):
+          checkpointing.print_save_message(int(state.step) - 1, config.async_checkpointing)
+      except Exception:  # pylint: disable=broad-except
+        max_logging.log(f"Checkpoint is already saved for step {int(state.step)-1}.")
+
+    checkpoint_manager.wait_until_finished()
+  metric_logger.write_metrics(running_gcs_metrics, metrics, config.steps - 1)  # final step metrics
+  max_utils.close_summary_writer(writer)
+  record_goodput(recorder, config, recorder.record_job_end_time if recorder else None)
+  with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
+    # pytype: disable=attribute-error
+    compiled = p_train_step.lower(state, example_batch, nextrng).compile()
+    compiled_stats = compiled.memory_analysis()
+    if compiled_stats is not None:
+      max_logging.log(
+          f"Output size: {compiled_stats.output_size_in_bytes}, "
+          f"temp size: {compiled_stats.temp_size_in_bytes}, "
+          f"argument size: {compiled_stats.argument_size_in_bytes}, "
+          f"host temp size: {compiled_stats.host_temp_size_in_bytes}, in bytes."
+      )
+  return state
+
+
+def wait_for_all_slices(elastic_manager: manager.Manager) -> None:
+  elastic_manager.good_slice_indices = elastic_manager.get_slice_availability()
+  while len(elastic_manager.good_slice_indices) < elastic_manager.total_slice_count:
+    max_logging.log(
+        f"Only {elastic_manager.good_slice_count} slices out of {elastic_manager.total_slice_count} available. "
+        "Sleeping for 5 seconds."
+    )
+    time.sleep(5)
+    elastic_manager.good_slice_indices = elastic_manager.get_slice_availability()
+  max_logging.log("All slices are available")
+
+
+def elastic_initialize(devices: Sequence[jax.Device]) -> manager.Manager:
+  """Initializes the elastic manager and pyconfig to support elastic training
+
+  Args:
+    devices: The devices used for training
+
+  Returns:
+    The initialized elastic manager
+  """
+  elastic_manager = manager.Manager(
+      devices,
+      reshard_check_period=1,
+      snapshot_period=5,
+      max_elastic_down_event_count=100,
+      max_reshard_retry_count=3,
+  )
+
+  # Do not start training until all slices are available
+  # TODO: b/408455557 - Migrate to pathwaysutils and make configurable
+  wait_for_all_slices(elastic_manager)
+
+  pyconfig.HyperParameters.global_batch_size_to_train_on = property(
+      lambda self: elastic_manager.scale_by_good_slices(self.get_keys()["global_batch_size_to_train_on"])
+  )
+  pyconfig.HyperParameters.global_batch_size_to_load = property(
+      lambda self: elastic_manager.scale_by_good_slices(self.get_keys()["global_batch_size_to_load"])
+  )
+  pyconfig.HyperParameters.micro_batch_size_to_train_on = property(
+      lambda self: elastic_manager.scale_by_good_slices(self.get_keys()["micro_batch_size_to_train_on"])
+  )
+  pyconfig.HyperParameters.num_slices = property(lambda self: elastic_manager.good_slice_count)
+
+  return elastic_manager
+
+
+def main(argv: Sequence[str]) -> None:
+  pathwaysutils.initialize()
+  jax.config.update("jax_default_prng_impl", "unsafe_rbg")
+  # TF allocates extraneous GPU memory when using TFDS data
+  # this leads to CUDA OOMs. WAR for now is to hide GPUs from TF
+  tf.config.set_visible_devices([], "GPU")
+  os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"
+  if "xla_tpu_spmd_rng_bit_generator_unsafe" not in os.environ.get("LIBTPU_INIT_ARGS", ""):
+    os.environ["LIBTPU_INIT_ARGS"] = os.environ.get("LIBTPU_INIT_ARGS", "") + " --xla_tpu_spmd_rng_bit_generator_unsafe=true"
+
+  elastic_manager = elastic_initialize(jax.devices())
+
+  config = pyconfig.initialize(argv)
+  max_utils.print_system_information()
+  validate_train_config(config)
+  os.environ["TFDS_DATA_DIR"] = config.dataset_path or ""
+  vertex_tensorboard_manager = VertexTensorboardManager()
+  if config.use_vertex_tensorboard or os.environ.get("UPLOAD_DATA_TO_TENSORBOARD"):
+    vertex_tensorboard_manager.configure_vertex_tensorboard(config)
+
+  if config.monitor_goodput and jax.process_index() == 0:
+    logger_name = f"goodput_{config.run_name}"
+    goodput_monitor = monitoring.GoodputMonitor(
+        job_name=config.run_name,
+        logger_name=logger_name,
+        tensorboard_dir=config.tensorboard_dir,
+        upload_interval=config.goodput_upload_interval_seconds,
+        monitoring_enabled=True,
+        pathway_enabled=config.enable_pathways_goodput,
+        include_badput_breakdown=True,
+        include_step_deviation=config.monitor_step_time_deviation,
+        step_deviation_interval_seconds=config.step_deviation_interval_seconds,
+    )
+    goodput_monitor.start_goodput_uploader()
+    max_logging.log("Started Goodput upload to Tensorboard in the background!")
+    if config.monitor_step_time_deviation:
+      goodput_monitor.start_step_deviation_uploader()
+      max_logging.log("Started step time deviation upload to Tensorboard in the background!")
+  debug_config = debug_configuration.DebugConfig(
+      stack_trace_config=stack_trace_configuration.StackTraceConfig(
+          collect_stack_trace=config.collect_stack_trace,
+          stack_trace_to_cloud=config.stack_trace_to_cloud,
+          stack_trace_interval_seconds=config.stack_trace_interval_seconds,
+      )
+  )
+  diagnostic_config = diagnostic_configuration.DiagnosticConfig(debug_config)
+  with diagnostic.diagnose(diagnostic_config):
+    train_loop(config, elastic_manager)
+
+
+if __name__ == "__main__":
+  app.run(main)
diff --git a/MaxText/tests/elastic_train_test.py b/MaxText/tests/elastic_train_test.py
new file mode 100644
index 0000000000..3e3f2f6936
--- /dev/null
+++ b/MaxText/tests/elastic_train_test.py
@@ -0,0 +1,136 @@
+#  Copyright 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.
+"""
+Tests the elastic related functions in elastic_train.py
+"""
+
+import logging
+import time
+from unittest import mock
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import jax
+from MaxText import elastic_train
+from MaxText import max_utils
+from MaxText import pyconfig
+from pathwaysutils.elastic import manager
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.INFO)
+
+
+class ElasticTrainTest(parameterized.TestCase):
+
+  def tearDown(self):
+    """Clean up at the end of the test
+
+    HyperParameters and they must be removed so that they do not impact other unittest
+    """
+    try:
+      del pyconfig.HyperParameters.global_batch_size_to_train_on
+      del pyconfig.HyperParameters.global_batch_size_to_load
+      del pyconfig.HyperParameters.micro_batch_size_to_train_on
+      del pyconfig.HyperParameters.num_slices
+    except AttributeError:
+      pass
+
+  @parameterized.named_parameters(
+      ("ready_after_0_try", [{0, 1}]),
+      ("nothing_available_at_first", [{}, {0, 1}]),
+      ("nothing_available_for_a_few_times", [{}, {}, {}, {0, 1}]),
+      ("back_and_forth", [{}, {1}, {}, {0}, {}, {1}, {0}, {}, {0}, {0, 1}]),
+  )
+  def test_wait_for_all_slices(self, slice_availability_side_effect):
+    mock_manager = mock.create_autospec(manager.Manager, instance=True)
+    mock_manager.total_slice_count = 2
+    mock_manager.get_slice_availability.side_effect = slice_availability_side_effect
+
+    mock_sleep = self.enter_context(mock.patch.object(time, "sleep", create_autospec=True))
+
+    elastic_train.wait_for_all_slices(mock_manager)
+
+    self.assertEqual(mock_sleep.call_count, len(slice_availability_side_effect) - 1)
+
+  @parameterized.named_parameters(
+      ("4_out_of_4_100", {0, 1, 2, 3}, 4, 100, 100),
+      ("3_out_of_4_100", {1, 2, 3}, 4, 100, 75),
+      ("2_out_of_4_100", {0, 3}, 4, 100, 50),
+      ("1_out_of_4_100", {2}, 4, 100, 25),
+      ("0_out_of_4_100", {}, 4, 100, 0),
+      ("3_out_of_3_100", {0, 1, 2}, 3, 63, 63),
+      ("2_out_of_3_100", {0, 2}, 3, 63, 42),
+      ("1_out_of_3_100", {1}, 3, 63, 21),
+      ("0_out_of_3_100", {}, 3, 63, 0),
+  )
+  def test_pyconfig_changes(self, good_slice_indices, total_slice_count, base_number, expected_number):
+
+    # Mock max_utils to report that there are 4 slices
+    # This is used to set config.num_slices
+    self.enter_context(
+        mock.patch.object(
+            max_utils,
+            "get_num_slices",
+            return_value=total_slice_count,
+            create_autospec=True,
+        )
+    )
+
+    # Mock jax.devices to report that there is 1 device
+    # This is used to compute the micro_batch_size_to_load,
+    # micro_batch_size_to_train_on, global_batch_size_to_load,
+    # global_batch_size_to_train_on
+    # All of those will be equal to per_device_batch_size based on base.yml
+    self.enter_context(
+        mock.patch.object(
+            jax,
+            "devices",
+            return_value=[
+                mock.create_autospec(
+                    jax.Device,
+                    instance=True,
+                ),
+            ],
+            create_autospec=True,
+        )
+    )
+
+    # Set checkpoint_period which should be unchanged.
+    # Set enable_single_controller to avoid jax.distribute_initialize
+    config = pyconfig.initialize(
+        argv=[
+            "test",
+            "MaxText/configs/base.yml",
+        ],
+        per_device_batch_size=base_number,
+        checkpoint_period=1234,
+        enable_single_controller=True,
+    )
+
+    # Do not set any devices and instead overwrite the total_slice_count and the
+    # good_slice_indices directly to avoid code paths that would otherwise need
+    # additional mocking
+    elastic_manager = elastic_train.elastic_initialize([])
+    elastic_manager._total_slice_count = total_slice_count  # pylint: disable=protected-access
+    elastic_manager.good_slice_indices = good_slice_indices
+
+    self.assertEqual(config.global_batch_size_to_train_on, expected_number)
+    self.assertEqual(config.global_batch_size_to_load, expected_number)
+    self.assertEqual(config.micro_batch_size_to_train_on, expected_number)
+    self.assertEqual(config.num_slices, len(good_slice_indices))
+    self.assertEqual(config.checkpoint_period, 1234)
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/MaxText/train.py b/MaxText/train.py
index 858f5f15e7..1e8c2c8c54 100644
--- a/MaxText/train.py
+++ b/MaxText/train.py
@@ -569,11 +569,12 @@ def check_example_batch(config, example_batch):
     err.throw()
 
 
-def setup_mesh_and_model(config):
+def setup_mesh_and_model(config, devices=None):
   """Set up the mesh and the model for training
 
   Args:
     config
+    devices
 
   Returns:
     init_rng: RNG key
@@ -590,7 +591,7 @@ def setup_mesh_and_model(config):
   writer = max_utils.initialize_summary_writer(config.tensorboard_dir, config.run_name)
 
   # Mesh definition
-  devices_array = maxtext_utils.create_device_mesh(config)
+  devices_array = maxtext_utils.create_device_mesh(config, devices)
   mesh = Mesh(devices_array, config.mesh_axes)
 
   # Model and Optimizer definition
diff --git a/requirements.txt b/requirements.txt
index 118a8deb22..753540e4ef 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -19,7 +19,6 @@ ml-collections
 ml-goodput-measurement==0.0.8
 numpy
 optax
-pathwaysutils
 protobuf==3.20.3
 pylint
 pytest
@@ -37,5 +36,5 @@ transformers
 mlperf-logging@git+https://github.com/mlperf/logging.git
 google-jetstream@git+https://github.com/AI-Hypercomputer/JetStream.git
 jsonlines
-pathwaysutils@git+https://github.com/google/pathways-utils.git
+pathwaysutils@git+https://github.com/AI-Hypercomputer/pathways-utils.git
 omegaconf
diff --git a/requirements_with_jax_stable_stack.txt b/requirements_with_jax_stable_stack.txt
index 5957029669..2a53c24d6d 100644
--- a/requirements_with_jax_stable_stack.txt
+++ b/requirements_with_jax_stable_stack.txt
@@ -20,7 +20,7 @@ transformers
 mlperf-logging@git+https://github.com/mlperf/logging.git
 google-jetstream@git+https://github.com/AI-Hypercomputer/JetStream.git
 jsonlines
-pathwaysutils@git+https://github.com/google/pathways-utils.git
+pathwaysutils@git+https://github.com/AI-Hypercomputer/pathways-utils.git
 google-cloud-monitoring
 google-api-core
 google-api-python-client


Reviews:
user: login: gagika
body: A few initial comments
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-02-27T18:12:51Z
user: login: lukebaumann
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-04T16:40:11Z
user: login: gobbleturk
body: This looks awesome, thanks Luke!! I am impressed with the modularity and lack of train.py changes
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-04T17:36:43Z
user: login: lukebaumann
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-04T17:39:34Z
user: login: lukebaumann
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-04T17:39:46Z
user: login: lukebaumann
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-04T17:48:23Z
user: login: lukebaumann
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-04T17:48:40Z
user: login: lukebaumann
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-04T17:49:17Z
user: login: lukebaumann
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-04T17:56:43Z
user: login: lukebaumann
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-04T17:59:51Z
user: login: lukebaumann
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-04T18:07:11Z
user: login: lukebaumann
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-04T18:07:32Z
user: login: lukebaumann
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-04T18:14:41Z
user: login: gobbleturk
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-04T18:15:16Z
user: login: gobbleturk
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-08T18:04:03Z
user: login: gobbleturk
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-08T18:04:39Z
user: login: gobbleturk
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-08T18:05:08Z
user: login: gobbleturk
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-08T18:05:27Z
user: login: gobbleturk
body: Approved! There are still requests from the first review to add some comments around tricky sections (or preferrably to make the code super clear and self explanatory, e.g. descriptive method names)
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-08T18:06:38Z
user: login: SurbhiJainUSC
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-08T18:16:25Z
user: login: SurbhiJainUSC
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-08T21:16:33Z
user: login: SurbhiJainUSC
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-08T21:20:26Z
user: login: lukebaumann
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T00:09:48Z
user: login: lukebaumann
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T00:10:08Z
user: login: lukebaumann
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T00:18:47Z
user: login: lukebaumann
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T01:14:48Z
user: login: lukebaumann
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T02:11:35Z
user: login: lukebaumann
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T02:24:25Z
user: login: lukebaumann
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T03:06:55Z
user: login: lukebaumann
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T03:08:31Z
user: login: richjames0
body: cc @gobbleturk 

Thanks for this, @lukebaumann - looks great! 

One bigger picture question I have, which doesn't necessarily need to be addressed now, but could be, is - should this be integrated with train.py, particularly in a world where we're focused on improving code quality and modularizing the codebase. I see that there's about a 30% overlap in LoC (maybe at the bottom of the two files) and I'm wondering whether this could be accomplished by train.py having hooks out to elastic training.

In addition, I've commented in a couple of places about breaking long functions out into sub-functions, which would help with testability as well as comprehensibility. I notice there are no unit tests right now - it would be great to get some, and I'm sure trying to do so would drive some of that decomposition in a practical way (vs. because Rich said so)
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T17:04:40Z
user: login: lukebaumann
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T18:08:08Z
user: login: lukebaumann
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T18:18:13Z
user: login: lukebaumann
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T18:23:11Z
user: login: lukebaumann
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T18:24:30Z
user: login: lukebaumann
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T19:03:33Z
user: login: lukebaumann
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T19:05:59Z
user: login: SurbhiJainUSC
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T20:28:51Z
user: login: richjames0
body: LGTM
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-17T00:18:23Z
user: login: gobbleturk
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-17T00:47:29Z

Review comments:
path: MaxText/train.py
user: login: gagika
body: Could you move it to a separate elastic_train.py file similar to sft_trainer.py: https://github.com/AI-Hypercomputer/maxtext/blob/sft-maxtext/MaxText/sft_trainer.py 
created_at: 2025-02-27T18:07:40Z
author_association: COLLABORATOR
diff_hunk: @@ -767,6 +782,101 @@ def setup_train_loop(config):
   )
 
 
+@elasticutils.timeit
+def reshard_fn(config: pyconfig.HyperParameters):
+  """Reshard function."""
+  while True:
path: MaxText/train.py
user: login: gagika
body: Could you add num of reshards instead of infinite loop?
created_at: 2025-02-27T18:10:09Z
author_association: COLLABORATOR
diff_hunk: @@ -767,6 +782,101 @@ def setup_train_loop(config):
   )
 
 
+@elasticutils.timeit
+def reshard_fn(config: pyconfig.HyperParameters):
+  """Reshard function."""
+  while True:
path: MaxText/train.py
user: login: lukebaumann
body: Out dated. Now the number of retry attempts and total number of slice down events are configurable.
created_at: 2025-04-04T16:40:11Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,452 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
path: MaxText/elastic_train.py
user: login: gobbleturk
body: can you include more description for what elastic training is? E.g. training continues if a slice goes down etc
created_at: 2025-04-04T16:45:43Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,452 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
path: MaxText/elastic_train.py
user: login: gobbleturk
body: what is this comment? Is this necessary?
created_at: 2025-04-04T16:46:02Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,452 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput, record_scalar_metrics,
+    save_checkpoint, setup_mesh_and_model, setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
path: MaxText/elastic_train.py
user: login: gobbleturk
body: I think user should have the power to set these? Perhaps add options in base.yml? Open for discussion though this is okay

We currently plumb jax debuggin options like this in pyconfig

https://github.com/AI-Hypercomputer/maxtext/blob/d2f7baf506b9facb3efcb39a4ae8e5be3b823ef6/MaxText/pyconfig.py#L402-L403
created_at: 2025-04-04T16:49:29Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,452 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput, record_scalar_metrics,
+    save_checkpoint, setup_mesh_and_model, setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger('pathwaysutils.elastic.manager').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.elastic.reshard').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.debug.timing').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.debug.watchdog').setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    em,
path: MaxText/elastic_train.py
user: login: gobbleturk
body: nit: more descriptive variable than em? My initial guess is "engineering manager" but I think that is probably not it
created_at: 2025-04-04T16:50:36Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,452 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput, record_scalar_metrics,
+    save_checkpoint, setup_mesh_and_model, setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger('pathwaysutils.elastic.manager').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.elastic.reshard').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.debug.timing').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.debug.watchdog').setLevel(logging.DEBUG)
+
+
+@timing.timeit
path: MaxText/elastic_train.py
user: login: gobbleturk
body: uh is this leftover from testing? We don't want a timeit wrapping here do we?
created_at: 2025-04-04T16:51:03Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,452 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput, record_scalar_metrics,
+    save_checkpoint, setup_mesh_and_model, setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger('pathwaysutils.elastic.manager').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.elastic.reshard').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.debug.timing').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.debug.watchdog').setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    em,
path: MaxText/elastic_train.py
user: login: gobbleturk
body: Okay I figured out its probably elastic manager?
created_at: 2025-04-04T16:53:03Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,452 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput, record_scalar_metrics,
+    save_checkpoint, setup_mesh_and_model, setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger('pathwaysutils.elastic.manager').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.elastic.reshard').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.debug.timing').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.debug.watchdog').setLevel(logging.DEBUG)
+
+
+@timing.timeit
path: MaxText/elastic_train.py
user: login: gobbleturk
body: Discussed offline, TIL @timing.timeit just wraps time.time() Comment what this function does - I'm surprised to see a time.time() around a method called elastic_handler
created_at: 2025-04-04T16:59:25Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,452 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
path: MaxText/elastic_train.py
user: login: gobbleturk
body: Discussed offline, this is for internal integration, not needed for initial integration
created_at: 2025-04-04T17:03:17Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,444 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput,
+    record_scalar_metrics,
+    save_checkpoint,
+    setup_mesh_and_model,
+    setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.watchdog").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    em,
+    checkpoint_manager,
+):
+  """Reshard function."""
path: MaxText/elastic_train.py
user: login: gobbleturk
body: If this function reshards, is there a better name than elastic handler?
created_at: 2025-04-04T17:22:14Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,444 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput,
+    record_scalar_metrics,
+    save_checkpoint,
+    setup_mesh_and_model,
+    setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.watchdog").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    em,
+    checkpoint_manager,
+):
+  """Reshard function."""
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(em.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, em.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = em.get_resharded_snapshot(mesh)
+
+      if checkpoint_manager is not None:
+        # Confirm this is the right thing to do
path: MaxText/elastic_train.py
user: login: gobbleturk
body: Is this a question for me?
created_at: 2025-04-04T17:24:04Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,444 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput,
+    record_scalar_metrics,
+    save_checkpoint,
+    setup_mesh_and_model,
+    setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.watchdog").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    em,
+    checkpoint_manager,
+):
+  """Reshard function."""
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(em.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, em.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = em.get_resharded_snapshot(mesh)
+
+      if checkpoint_manager is not None:
+        # Confirm this is the right thing to do
+        latest_step = checkpoint_manager.latest_step()
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
path: MaxText/elastic_train.py
user: login: gobbleturk
body: Can you include comments about this logic. Why don't we restore from {latest_step}? Do we believe that checkpoint might be corrupted? Is it too hard to plumb through a checkpoint restore so instead we always do an EM restore?
created_at: 2025-04-04T17:25:43Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,444 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput,
+    record_scalar_metrics,
+    save_checkpoint,
+    setup_mesh_and_model,
+    setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.watchdog").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    em,
+    checkpoint_manager,
+):
+  """Reshard function."""
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(em.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, em.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = em.get_resharded_snapshot(mesh)
+
+      if checkpoint_manager is not None:
+        # Confirm this is the right thing to do
+        latest_step = checkpoint_manager.latest_step()
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
+          checkpoint_manager.delete(latest_step)
+
+      state, _, state_mesh_shardings, data_iterator = max_utils.setup_training_state(
+          model,
+          data_iterator,
+          tx,
+          config,
+          jax.random.fold_in(init_rng, step),
+          mesh,
+          checkpoint_manager=None,
+      )
+
+      state = state.replace(**snapshot)
+      state = state.replace(step=state.step.at[None].set(step))
+
+      (
+          functional_train,
+          in_shard_train,
+          out_shard_train,
+          static_argnums_train,
+          donate_argnums_train,
+      ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+      p_train_step = jax.jit(
+          functional_train,
+          in_shardings=in_shard_train,
+          out_shardings=out_shard_train,
+          static_argnums=static_argnums_train,
+          donate_argnums=donate_argnums_train,
+      )
+
+      example_batch = None
+      metric_logger = MetricLogger(writer, config)
+
+      jax.block_until_ready(state)
+
+  return (
+      config,
+      step,
+      state,
+      mesh,
+      checkpoint_manager,
+      data_iterator,
+      p_train_step,
+      example_batch,
+      learning_rate_schedule,
+      metric_logger,
+      writer,
+  )
+
+
+def train_loop(config, em, state=None):
+  """Main Training loop.
+  Args:
+    config:
+    state:
+    ckpt_path:
+  Returns:
+  """
+  # Create a GoodputRecorder to log information
+  recorder = create_goodput_recorder(config)
+  record_goodput(recorder, config, recorder.record_job_start_time if recorder else None)
+
+  (
+      init_rng,
+      writer,
+      checkpoint_manager,
+      state_mesh_shardings,
+      model,
+      mesh,
+      learning_rate_schedule,
+      data_iterator,
+      _,
+      state,
+  ) = setup_train_loop(config)
+
+  # pylint: disable=line-too-long
+  (
+      functional_train,
+      in_shard_train,
+      out_shard_train,
+      static_argnums_train,
+      donate_argnums_train,
+  ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+  p_train_step = jax.jit(
+      functional_train,
+      in_shardings=in_shard_train,
+      out_shardings=out_shard_train,
+      static_argnums=static_argnums_train,
+      donate_argnums=donate_argnums_train,
+  )
+
+  num_model_parameters = max_utils.calculate_num_params_from_pytree(state.params)
+  max_logging.log(f"number parameters: {num_model_parameters/1e9:.3f} billion")
+  per_device_tflops, _, _ = maxtext_utils.calculate_tflops_training_per_device(config)
+  per_device_tokens = maxtext_utils.calculate_tokens_training_per_device(config)
+
+  # Write train config params, num model params, and XLA flags to tensorboard
+  max_utils.add_text_to_summary_writer("num_model_parameters", str(num_model_parameters), writer)
+  max_utils.add_text_to_summary_writer("libtpu_init_args", os.environ["LIBTPU_INIT_ARGS"], writer)
+  max_utils.add_config_to_summary_writer(config, writer)
+
+  running_gcs_metrics = [] if config.gcs_metrics else None
+
+  start_step = get_first_step(state)  # this is the start_step for training
+  prof = profiler.Profiler(config, offset_step=start_step)
+  first_profiling_step = prof.start_initial_profile_step
+  if config.profiler != "" and first_profiling_step >= config.steps:
+    raise ValueError("Profiling requested but initial profiling step set past training final step")
+  last_profiling_step = prof.finished_initial_profile_step
+
+  example_batch = None
+  last_step_completion = datetime.datetime.now()
+
+  performance_metric_queue = None
+  if config.report_heartbeat_metric_for_gcp_monitoring or config.report_performance_metric_for_gcp_monitoring:
+    gcp_workload_monitor = GCPWorkloadMonitor(config.run_name)
+    if config.report_heartbeat_metric_for_gcp_monitoring:
+      gcp_workload_monitor.start_heartbeat_reporting_thread(config.heartbeat_reporting_interval_in_seconds)
+    if config.report_performance_metric_for_gcp_monitoring:
+      performance_metric_queue = queue.Queue()
+      gcp_workload_monitor.start_performance_reporting_thread(performance_metric_queue)
+
+  metric_logger = MetricLogger(writer, config)
+  step = start_step
+
+  em.maybe_snapshot(
+      step,
+      snapshot={
+          "params": state.params,
+          "opt_state": state.opt_state,
+      },
+      force=True,
+      block=True,
+  )
+
+  while True:
path: MaxText/elastic_train.py
user: login: gobbleturk
body: Can you include a comment why you are using `while True:` Does a while `step <= config.step` not work here?
created_at: 2025-04-04T17:27:45Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,444 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput,
+    record_scalar_metrics,
+    save_checkpoint,
+    setup_mesh_and_model,
+    setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.watchdog").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    em,
+    checkpoint_manager,
+):
+  """Reshard function."""
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(em.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, em.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = em.get_resharded_snapshot(mesh)
+
+      if checkpoint_manager is not None:
+        # Confirm this is the right thing to do
+        latest_step = checkpoint_manager.latest_step()
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
+          checkpoint_manager.delete(latest_step)
+
+      state, _, state_mesh_shardings, data_iterator = max_utils.setup_training_state(
+          model,
+          data_iterator,
+          tx,
+          config,
+          jax.random.fold_in(init_rng, step),
+          mesh,
+          checkpoint_manager=None,
+      )
+
+      state = state.replace(**snapshot)
+      state = state.replace(step=state.step.at[None].set(step))
+
+      (
+          functional_train,
+          in_shard_train,
+          out_shard_train,
+          static_argnums_train,
+          donate_argnums_train,
+      ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+      p_train_step = jax.jit(
+          functional_train,
+          in_shardings=in_shard_train,
+          out_shardings=out_shard_train,
+          static_argnums=static_argnums_train,
+          donate_argnums=donate_argnums_train,
+      )
+
+      example_batch = None
+      metric_logger = MetricLogger(writer, config)
+
+      jax.block_until_ready(state)
+
+  return (
+      config,
+      step,
+      state,
+      mesh,
+      checkpoint_manager,
+      data_iterator,
+      p_train_step,
+      example_batch,
+      learning_rate_schedule,
+      metric_logger,
+      writer,
+  )
+
+
+def train_loop(config, em, state=None):
+  """Main Training loop.
+  Args:
+    config:
+    state:
+    ckpt_path:
+  Returns:
+  """
+  # Create a GoodputRecorder to log information
+  recorder = create_goodput_recorder(config)
+  record_goodput(recorder, config, recorder.record_job_start_time if recorder else None)
+
+  (
+      init_rng,
+      writer,
+      checkpoint_manager,
+      state_mesh_shardings,
+      model,
+      mesh,
+      learning_rate_schedule,
+      data_iterator,
+      _,
+      state,
+  ) = setup_train_loop(config)
+
+  # pylint: disable=line-too-long
+  (
+      functional_train,
+      in_shard_train,
+      out_shard_train,
+      static_argnums_train,
+      donate_argnums_train,
+  ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+  p_train_step = jax.jit(
+      functional_train,
+      in_shardings=in_shard_train,
+      out_shardings=out_shard_train,
+      static_argnums=static_argnums_train,
+      donate_argnums=donate_argnums_train,
+  )
+
+  num_model_parameters = max_utils.calculate_num_params_from_pytree(state.params)
+  max_logging.log(f"number parameters: {num_model_parameters/1e9:.3f} billion")
+  per_device_tflops, _, _ = maxtext_utils.calculate_tflops_training_per_device(config)
+  per_device_tokens = maxtext_utils.calculate_tokens_training_per_device(config)
+
+  # Write train config params, num model params, and XLA flags to tensorboard
+  max_utils.add_text_to_summary_writer("num_model_parameters", str(num_model_parameters), writer)
+  max_utils.add_text_to_summary_writer("libtpu_init_args", os.environ["LIBTPU_INIT_ARGS"], writer)
+  max_utils.add_config_to_summary_writer(config, writer)
+
+  running_gcs_metrics = [] if config.gcs_metrics else None
+
+  start_step = get_first_step(state)  # this is the start_step for training
+  prof = profiler.Profiler(config, offset_step=start_step)
+  first_profiling_step = prof.start_initial_profile_step
+  if config.profiler != "" and first_profiling_step >= config.steps:
+    raise ValueError("Profiling requested but initial profiling step set past training final step")
+  last_profiling_step = prof.finished_initial_profile_step
+
+  example_batch = None
+  last_step_completion = datetime.datetime.now()
+
+  performance_metric_queue = None
+  if config.report_heartbeat_metric_for_gcp_monitoring or config.report_performance_metric_for_gcp_monitoring:
+    gcp_workload_monitor = GCPWorkloadMonitor(config.run_name)
+    if config.report_heartbeat_metric_for_gcp_monitoring:
+      gcp_workload_monitor.start_heartbeat_reporting_thread(config.heartbeat_reporting_interval_in_seconds)
+    if config.report_performance_metric_for_gcp_monitoring:
+      performance_metric_queue = queue.Queue()
+      gcp_workload_monitor.start_performance_reporting_thread(performance_metric_queue)
+
+  metric_logger = MetricLogger(writer, config)
+  step = start_step
+
+  em.maybe_snapshot(
+      step,
+      snapshot={
+          "params": state.params,
+          "opt_state": state.opt_state,
+      },
+      force=True,
+      block=True,
+  )
+
+  while True:
+    with watchdog.watchdog(120):
path: MaxText/elastic_train.py
user: login: gobbleturk
body: cc @SurbhiJainUSC this looks like an alternative to our stack trace package - perhaps instead we can use our stack trace package?
created_at: 2025-04-04T17:28:27Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,444 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput,
+    record_scalar_metrics,
+    save_checkpoint,
+    setup_mesh_and_model,
+    setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.watchdog").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    em,
+    checkpoint_manager,
+):
+  """Reshard function."""
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(em.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, em.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = em.get_resharded_snapshot(mesh)
+
+      if checkpoint_manager is not None:
+        # Confirm this is the right thing to do
+        latest_step = checkpoint_manager.latest_step()
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
+          checkpoint_manager.delete(latest_step)
+
+      state, _, state_mesh_shardings, data_iterator = max_utils.setup_training_state(
+          model,
+          data_iterator,
+          tx,
+          config,
+          jax.random.fold_in(init_rng, step),
+          mesh,
+          checkpoint_manager=None,
+      )
+
+      state = state.replace(**snapshot)
+      state = state.replace(step=state.step.at[None].set(step))
+
+      (
+          functional_train,
+          in_shard_train,
+          out_shard_train,
+          static_argnums_train,
+          donate_argnums_train,
+      ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+      p_train_step = jax.jit(
+          functional_train,
+          in_shardings=in_shard_train,
+          out_shardings=out_shard_train,
+          static_argnums=static_argnums_train,
+          donate_argnums=donate_argnums_train,
+      )
+
+      example_batch = None
+      metric_logger = MetricLogger(writer, config)
+
+      jax.block_until_ready(state)
+
+  return (
+      config,
+      step,
+      state,
+      mesh,
+      checkpoint_manager,
+      data_iterator,
+      p_train_step,
+      example_batch,
+      learning_rate_schedule,
+      metric_logger,
+      writer,
+  )
+
+
+def train_loop(config, em, state=None):
+  """Main Training loop.
+  Args:
+    config:
+    state:
+    ckpt_path:
+  Returns:
+  """
+  # Create a GoodputRecorder to log information
+  recorder = create_goodput_recorder(config)
+  record_goodput(recorder, config, recorder.record_job_start_time if recorder else None)
+
+  (
+      init_rng,
+      writer,
+      checkpoint_manager,
+      state_mesh_shardings,
+      model,
+      mesh,
+      learning_rate_schedule,
+      data_iterator,
+      _,
+      state,
+  ) = setup_train_loop(config)
+
+  # pylint: disable=line-too-long
+  (
+      functional_train,
+      in_shard_train,
+      out_shard_train,
+      static_argnums_train,
+      donate_argnums_train,
+  ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+  p_train_step = jax.jit(
+      functional_train,
+      in_shardings=in_shard_train,
+      out_shardings=out_shard_train,
+      static_argnums=static_argnums_train,
+      donate_argnums=donate_argnums_train,
+  )
+
+  num_model_parameters = max_utils.calculate_num_params_from_pytree(state.params)
+  max_logging.log(f"number parameters: {num_model_parameters/1e9:.3f} billion")
+  per_device_tflops, _, _ = maxtext_utils.calculate_tflops_training_per_device(config)
+  per_device_tokens = maxtext_utils.calculate_tokens_training_per_device(config)
+
+  # Write train config params, num model params, and XLA flags to tensorboard
+  max_utils.add_text_to_summary_writer("num_model_parameters", str(num_model_parameters), writer)
+  max_utils.add_text_to_summary_writer("libtpu_init_args", os.environ["LIBTPU_INIT_ARGS"], writer)
+  max_utils.add_config_to_summary_writer(config, writer)
+
+  running_gcs_metrics = [] if config.gcs_metrics else None
+
+  start_step = get_first_step(state)  # this is the start_step for training
+  prof = profiler.Profiler(config, offset_step=start_step)
+  first_profiling_step = prof.start_initial_profile_step
+  if config.profiler != "" and first_profiling_step >= config.steps:
+    raise ValueError("Profiling requested but initial profiling step set past training final step")
+  last_profiling_step = prof.finished_initial_profile_step
+
+  example_batch = None
+  last_step_completion = datetime.datetime.now()
+
+  performance_metric_queue = None
+  if config.report_heartbeat_metric_for_gcp_monitoring or config.report_performance_metric_for_gcp_monitoring:
+    gcp_workload_monitor = GCPWorkloadMonitor(config.run_name)
+    if config.report_heartbeat_metric_for_gcp_monitoring:
+      gcp_workload_monitor.start_heartbeat_reporting_thread(config.heartbeat_reporting_interval_in_seconds)
+    if config.report_performance_metric_for_gcp_monitoring:
+      performance_metric_queue = queue.Queue()
+      gcp_workload_monitor.start_performance_reporting_thread(performance_metric_queue)
+
+  metric_logger = MetricLogger(writer, config)
+  step = start_step
+
+  em.maybe_snapshot(
+      step,
+      snapshot={
+          "params": state.params,
+          "opt_state": state.opt_state,
+      },
+      force=True,
+      block=True,
+  )
+
+  while True:
+    with watchdog.watchdog(120):
+      try:
+        if step == first_profiling_step or prof.should_activate_periodic_profile(step):
+          optional_postfix = f"step_{step}" if config.profile_periodically_period > 0 else ""
+          prof.activate(blocking_object=state, optional_postfix=optional_postfix)
+
+        if step >= config.steps:
+          break
+
+        max_logging.log(f"{step=} {em.elastic_down_event_count=} {em.good_slice_count=}")
+        with mesh, nn_partitioning.axis_rules(config.logical_axis_rules), jax.default_device(em.default_device):
+          with jax.profiler.StepTraceAnnotation("train", step_num=step):
+            record_goodput(recorder, config, recorder.record_data_loading_start_time if recorder else None)
+            example_batch = load_next_batch(data_iterator, example_batch, config)
+            record_goodput(recorder, config, recorder.record_data_loading_end_time if recorder else None)
+            check_example_batch(config, example_batch=example_batch)
+            # pylint: disable=not-callable
+            nextrng = jax.jit(jax.random.fold_in)(init_rng, step)
+            record_goodput(recorder, config, recorder.record_step_start_time if recorder else None, step)
+            state, metrics = p_train_step(state, example_batch, nextrng)
+
+          step_time_delta = datetime.datetime.now() - last_step_completion
+          last_step_completion = datetime.datetime.now()
+          record_scalar_metrics(metrics, step_time_delta, per_device_tflops, learning_rate_schedule(step), per_device_tokens)
+          if performance_metric_queue:
+            performance_metric_queue.put(step_time_delta.total_seconds())
+
+          if checkpoint_manager is not None:
+            state_to_save = state
+            if save_checkpoint(checkpoint_manager, int(step), state_to_save, config.dataset_type, data_iterator, config):
+              checkpointing.print_save_message(step, config.async_checkpointing)
+
+            # Upon preemption, exit when and only when all ongoing saves are complete.
+            if checkpoint_manager.reached_preemption(step):
+              checkpoint_manager.wait_until_finished()
+              sys.exit()
+
+          metric_logger.write_metrics(running_gcs_metrics, metrics, step)
+
+          if step == last_profiling_step or prof.should_deactivate_periodic_profile(step):
+            prof.deactivate(blocking_object=state)
+
+        em.maybe_snapshot(
path: MaxText/elastic_train.py
user: login: gobbleturk
body: Nit: can we wrap lines 283-348 in its own method called something like "get_next_em_state"? This is a stylistic choice to make code more modular easier to read/maintain
created_at: 2025-04-04T17:29:57Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,444 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput,
+    record_scalar_metrics,
+    save_checkpoint,
+    setup_mesh_and_model,
+    setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.watchdog").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    em,
+    checkpoint_manager,
+):
+  """Reshard function."""
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(em.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, em.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = em.get_resharded_snapshot(mesh)
+
+      if checkpoint_manager is not None:
+        # Confirm this is the right thing to do
+        latest_step = checkpoint_manager.latest_step()
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
+          checkpoint_manager.delete(latest_step)
+
+      state, _, state_mesh_shardings, data_iterator = max_utils.setup_training_state(
+          model,
+          data_iterator,
+          tx,
+          config,
+          jax.random.fold_in(init_rng, step),
+          mesh,
+          checkpoint_manager=None,
+      )
+
+      state = state.replace(**snapshot)
+      state = state.replace(step=state.step.at[None].set(step))
+
+      (
+          functional_train,
+          in_shard_train,
+          out_shard_train,
+          static_argnums_train,
+          donate_argnums_train,
+      ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+      p_train_step = jax.jit(
+          functional_train,
+          in_shardings=in_shard_train,
+          out_shardings=out_shard_train,
+          static_argnums=static_argnums_train,
+          donate_argnums=donate_argnums_train,
+      )
+
+      example_batch = None
+      metric_logger = MetricLogger(writer, config)
+
+      jax.block_until_ready(state)
+
+  return (
+      config,
+      step,
+      state,
+      mesh,
+      checkpoint_manager,
+      data_iterator,
+      p_train_step,
+      example_batch,
+      learning_rate_schedule,
+      metric_logger,
+      writer,
+  )
+
+
+def train_loop(config, em, state=None):
+  """Main Training loop.
+  Args:
+    config:
+    state:
+    ckpt_path:
+  Returns:
+  """
+  # Create a GoodputRecorder to log information
+  recorder = create_goodput_recorder(config)
+  record_goodput(recorder, config, recorder.record_job_start_time if recorder else None)
+
+  (
+      init_rng,
+      writer,
+      checkpoint_manager,
+      state_mesh_shardings,
+      model,
+      mesh,
+      learning_rate_schedule,
+      data_iterator,
+      _,
+      state,
+  ) = setup_train_loop(config)
+
+  # pylint: disable=line-too-long
+  (
+      functional_train,
+      in_shard_train,
+      out_shard_train,
+      static_argnums_train,
+      donate_argnums_train,
+  ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+  p_train_step = jax.jit(
+      functional_train,
+      in_shardings=in_shard_train,
+      out_shardings=out_shard_train,
+      static_argnums=static_argnums_train,
+      donate_argnums=donate_argnums_train,
+  )
+
+  num_model_parameters = max_utils.calculate_num_params_from_pytree(state.params)
+  max_logging.log(f"number parameters: {num_model_parameters/1e9:.3f} billion")
+  per_device_tflops, _, _ = maxtext_utils.calculate_tflops_training_per_device(config)
+  per_device_tokens = maxtext_utils.calculate_tokens_training_per_device(config)
+
+  # Write train config params, num model params, and XLA flags to tensorboard
+  max_utils.add_text_to_summary_writer("num_model_parameters", str(num_model_parameters), writer)
+  max_utils.add_text_to_summary_writer("libtpu_init_args", os.environ["LIBTPU_INIT_ARGS"], writer)
+  max_utils.add_config_to_summary_writer(config, writer)
+
+  running_gcs_metrics = [] if config.gcs_metrics else None
+
+  start_step = get_first_step(state)  # this is the start_step for training
+  prof = profiler.Profiler(config, offset_step=start_step)
+  first_profiling_step = prof.start_initial_profile_step
+  if config.profiler != "" and first_profiling_step >= config.steps:
+    raise ValueError("Profiling requested but initial profiling step set past training final step")
+  last_profiling_step = prof.finished_initial_profile_step
+
+  example_batch = None
+  last_step_completion = datetime.datetime.now()
+
+  performance_metric_queue = None
+  if config.report_heartbeat_metric_for_gcp_monitoring or config.report_performance_metric_for_gcp_monitoring:
+    gcp_workload_monitor = GCPWorkloadMonitor(config.run_name)
+    if config.report_heartbeat_metric_for_gcp_monitoring:
+      gcp_workload_monitor.start_heartbeat_reporting_thread(config.heartbeat_reporting_interval_in_seconds)
+    if config.report_performance_metric_for_gcp_monitoring:
+      performance_metric_queue = queue.Queue()
+      gcp_workload_monitor.start_performance_reporting_thread(performance_metric_queue)
+
+  metric_logger = MetricLogger(writer, config)
+  step = start_step
+
+  em.maybe_snapshot(
+      step,
+      snapshot={
+          "params": state.params,
+          "opt_state": state.opt_state,
+      },
+      force=True,
+      block=True,
+  )
+
+  while True:
+    with watchdog.watchdog(120):
+      try:
+        if step == first_profiling_step or prof.should_activate_periodic_profile(step):
+          optional_postfix = f"step_{step}" if config.profile_periodically_period > 0 else ""
+          prof.activate(blocking_object=state, optional_postfix=optional_postfix)
+
+        if step >= config.steps:
+          break
+
+        max_logging.log(f"{step=} {em.elastic_down_event_count=} {em.good_slice_count=}")
+        with mesh, nn_partitioning.axis_rules(config.logical_axis_rules), jax.default_device(em.default_device):
+          with jax.profiler.StepTraceAnnotation("train", step_num=step):
+            record_goodput(recorder, config, recorder.record_data_loading_start_time if recorder else None)
+            example_batch = load_next_batch(data_iterator, example_batch, config)
+            record_goodput(recorder, config, recorder.record_data_loading_end_time if recorder else None)
+            check_example_batch(config, example_batch=example_batch)
+            # pylint: disable=not-callable
+            nextrng = jax.jit(jax.random.fold_in)(init_rng, step)
+            record_goodput(recorder, config, recorder.record_step_start_time if recorder else None, step)
+            state, metrics = p_train_step(state, example_batch, nextrng)
+
+          step_time_delta = datetime.datetime.now() - last_step_completion
+          last_step_completion = datetime.datetime.now()
+          record_scalar_metrics(metrics, step_time_delta, per_device_tflops, learning_rate_schedule(step), per_device_tokens)
+          if performance_metric_queue:
+            performance_metric_queue.put(step_time_delta.total_seconds())
+
+          if checkpoint_manager is not None:
+            state_to_save = state
+            if save_checkpoint(checkpoint_manager, int(step), state_to_save, config.dataset_type, data_iterator, config):
+              checkpointing.print_save_message(step, config.async_checkpointing)
+
+            # Upon preemption, exit when and only when all ongoing saves are complete.
+            if checkpoint_manager.reached_preemption(step):
+              checkpoint_manager.wait_until_finished()
+              sys.exit()
+
+          metric_logger.write_metrics(running_gcs_metrics, metrics, step)
+
+          if step == last_profiling_step or prof.should_deactivate_periodic_profile(step):
+            prof.deactivate(blocking_object=state)
+
+        em.maybe_snapshot(
+            step=step,
+            snapshot={
+                "params": state.params,
+                "opt_state": state.opt_state,
+            },
+            block=True,
+        )
+
+        ret = em.maybe_reshard_up(
+            step=step,
+            snapshot={
+                "params": state.params,
+                "opt_state": state.opt_state,
+            },
+            elastic_handler=elastic_handler,
+            handler_kwargs={
+                "config": config,
+                "em": em,
+                "checkpoint_manager": checkpoint_manager,
+            },
+        )
+        if ret is not None:
+          (
+              config,
+              step,
+              state,
+              mesh,
+              checkpoint_manager,
+              data_iterator,
+              p_train_step,
+              example_batch,
+              learning_rate_schedule,
+              metric_logger,
+              writer,
+          ) = ret
+
+        if step == start_step:
+          max_utils.print_mem_stats("After params initialized")
+
+        step += 1
+
+      except jax.errors.JaxRuntimeError as error:
+        ret = em.maybe_reshard_down(
+            error=error,
+            elastic_handler=elastic_handler,
+            handler_kwargs={
+                "config": config,
+                "em": em,
+                "checkpoint_manager": checkpoint_manager,
+            },
+        )
+        if ret is not None:
+          (
+              config,
+              step,
+              state,
+              mesh,
+              checkpoint_manager,
+              data_iterator,
+              p_train_step,
+              example_batch,
+              learning_rate_schedule,
+              metric_logger,
+              writer,
+          ) = ret
+
+  if checkpoint_manager is not None:
+    checkpoint_manager.wait_until_finished()
+  metric_logger.write_metrics(running_gcs_metrics, metrics, config.steps - 1)  # final step metrics
+  max_utils.close_summary_writer(writer)
+  record_goodput(recorder, config, recorder.record_job_end_time if recorder else None)
+  with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
+    # pytype: disable=attribute-error
+    compiled = p_train_step.lower(state, example_batch, nextrng).compile()
+    compiled_stats = compiled.memory_analysis()
+    if compiled_stats is not None:
+      max_logging.log(
+          f"Output size: {compiled_stats.output_size_in_bytes}, "
+          f"temp size: {compiled_stats.temp_size_in_bytes}, "
+          f"argument size: {compiled_stats.argument_size_in_bytes}, "
+          f"host temp size: {compiled_stats.host_temp_size_in_bytes}, in bytes."
+      )
+  return state
+
+
+def main(argv: Sequence[str]) -> None:
+  pathwaysutils.initialize()
+  jax.config.update("jax_default_prng_impl", "unsafe_rbg")
+  # TF allocates extraneous GPU memory when using TFDS data
+  # this leads to CUDA OOMs. WAR for now is to hide GPUs from TF
+  tf.config.set_visible_devices([], "GPU")
+  os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"
+  if "xla_tpu_spmd_rng_bit_generator_unsafe" not in os.environ.get("LIBTPU_INIT_ARGS", ""):
+    os.environ["LIBTPU_INIT_ARGS"] = os.environ.get("LIBTPU_INIT_ARGS", "") + " --xla_tpu_spmd_rng_bit_generator_unsafe=true"
+
+  em = manager.Manager(
+      jax.devices(),
+      reshard_check_period=1,
+      snapshot_period=1,
+      max_elastic_down_event_count=100,
+      max_reshard_retry_count=3,
+  )
+
+  em.good_slice_indices = em.get_slice_availability()
+  while len(em.good_slice_indices) < em.total_slice_count:
path: MaxText/elastic_train.py
user: login: gobbleturk
body: Can you include a comment here what is em.total_slice_count so users don't have to dig into source code? Should users have the ability to set this so that e.g. training could start if say only n-1 slices start successfully?
created_at: 2025-04-04T17:31:37Z
author_association: COLLABORATOR
path: MaxText/train.py
user: login: gobbleturk
body: I am amazed by the lack of changes here, great work!
created_at: 2025-04-04T17:32:27Z
author_association: COLLABORATOR
diff_hunk: @@ -37,8 +37,8 @@
 
 # The set of tuning params required for long-running pathways jobs.
 PATHWAYS_LONG_RUN_CHECKPOINTING_TUNING_PARAMS = {
-    "enable_checkpointing": True,
-    "async_checkpointing": True,
+    "enable_checkpointing": False,
path: benchmarks/maxtext_trillium_model_configs.py
user: login: gobbleturk
body: Why is this changed in this PR? This seems orthogonal?
created_at: 2025-04-04T17:33:03Z
author_association: COLLABORATOR
diff_hunk: @@ -405,7 +405,7 @@ def build_user_command(
       'export ENABLE_PATHWAYS_PERSISTENCE=1 &&',
       f'export JAX_PLATFORMS={jax_platforms} &&',
       'export ENABLE_PJRT_COMPATIBILITY=true &&',
-      'python3 MaxText/train.py MaxText/configs/base.yml',
+      'python3 MaxText/elastic_train.py MaxText/configs/base.yml',
path: benchmarks/maxtext_xpk_runner.py
user: login: gobbleturk
body: please revert this
created_at: 2025-04-04T17:33:38Z
author_association: COLLABORATOR
diff_hunk: @@ -405,7 +405,7 @@ def build_user_command(
       'export ENABLE_PATHWAYS_PERSISTENCE=1 &&',
       f'export JAX_PLATFORMS={jax_platforms} &&',
       'export ENABLE_PJRT_COMPATIBILITY=true &&',
-      'python3 MaxText/train.py MaxText/configs/base.yml',
+      'python3 MaxText/elastic_train.py MaxText/configs/base.yml',
path: benchmarks/maxtext_xpk_runner.py
user: login: lukebaumann
body: Reverted all benchmarks/ changes
created_at: 2025-04-04T17:39:34Z
author_association: COLLABORATOR
diff_hunk: @@ -37,8 +37,8 @@
 
 # The set of tuning params required for long-running pathways jobs.
 PATHWAYS_LONG_RUN_CHECKPOINTING_TUNING_PARAMS = {
-    "enable_checkpointing": True,
-    "async_checkpointing": True,
+    "enable_checkpointing": False,
path: benchmarks/maxtext_trillium_model_configs.py
user: login: lukebaumann
body: Reverted all benchmarks/ changes
created_at: 2025-04-04T17:39:46Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,444 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput,
+    record_scalar_metrics,
+    save_checkpoint,
+    setup_mesh_and_model,
+    setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.watchdog").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    em,
+    checkpoint_manager,
+):
+  """Reshard function."""
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(em.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, em.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = em.get_resharded_snapshot(mesh)
+
+      if checkpoint_manager is not None:
+        # Confirm this is the right thing to do
+        latest_step = checkpoint_manager.latest_step()
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
+          checkpoint_manager.delete(latest_step)
+
+      state, _, state_mesh_shardings, data_iterator = max_utils.setup_training_state(
+          model,
+          data_iterator,
+          tx,
+          config,
+          jax.random.fold_in(init_rng, step),
+          mesh,
+          checkpoint_manager=None,
+      )
+
+      state = state.replace(**snapshot)
+      state = state.replace(step=state.step.at[None].set(step))
+
+      (
+          functional_train,
+          in_shard_train,
+          out_shard_train,
+          static_argnums_train,
+          donate_argnums_train,
+      ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+      p_train_step = jax.jit(
+          functional_train,
+          in_shardings=in_shard_train,
+          out_shardings=out_shard_train,
+          static_argnums=static_argnums_train,
+          donate_argnums=donate_argnums_train,
+      )
+
+      example_batch = None
+      metric_logger = MetricLogger(writer, config)
+
+      jax.block_until_ready(state)
+
+  return (
+      config,
+      step,
+      state,
+      mesh,
+      checkpoint_manager,
+      data_iterator,
+      p_train_step,
+      example_batch,
+      learning_rate_schedule,
+      metric_logger,
+      writer,
+  )
+
+
+def train_loop(config, em, state=None):
+  """Main Training loop.
+  Args:
+    config:
+    state:
+    ckpt_path:
+  Returns:
+  """
+  # Create a GoodputRecorder to log information
+  recorder = create_goodput_recorder(config)
+  record_goodput(recorder, config, recorder.record_job_start_time if recorder else None)
+
+  (
+      init_rng,
+      writer,
+      checkpoint_manager,
+      state_mesh_shardings,
+      model,
+      mesh,
+      learning_rate_schedule,
+      data_iterator,
+      _,
+      state,
+  ) = setup_train_loop(config)
+
+  # pylint: disable=line-too-long
+  (
+      functional_train,
+      in_shard_train,
+      out_shard_train,
+      static_argnums_train,
+      donate_argnums_train,
+  ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+  p_train_step = jax.jit(
+      functional_train,
+      in_shardings=in_shard_train,
+      out_shardings=out_shard_train,
+      static_argnums=static_argnums_train,
+      donate_argnums=donate_argnums_train,
+  )
+
+  num_model_parameters = max_utils.calculate_num_params_from_pytree(state.params)
+  max_logging.log(f"number parameters: {num_model_parameters/1e9:.3f} billion")
+  per_device_tflops, _, _ = maxtext_utils.calculate_tflops_training_per_device(config)
+  per_device_tokens = maxtext_utils.calculate_tokens_training_per_device(config)
+
+  # Write train config params, num model params, and XLA flags to tensorboard
+  max_utils.add_text_to_summary_writer("num_model_parameters", str(num_model_parameters), writer)
+  max_utils.add_text_to_summary_writer("libtpu_init_args", os.environ["LIBTPU_INIT_ARGS"], writer)
+  max_utils.add_config_to_summary_writer(config, writer)
+
+  running_gcs_metrics = [] if config.gcs_metrics else None
+
+  start_step = get_first_step(state)  # this is the start_step for training
+  prof = profiler.Profiler(config, offset_step=start_step)
+  first_profiling_step = prof.start_initial_profile_step
+  if config.profiler != "" and first_profiling_step >= config.steps:
+    raise ValueError("Profiling requested but initial profiling step set past training final step")
+  last_profiling_step = prof.finished_initial_profile_step
+
+  example_batch = None
+  last_step_completion = datetime.datetime.now()
+
+  performance_metric_queue = None
+  if config.report_heartbeat_metric_for_gcp_monitoring or config.report_performance_metric_for_gcp_monitoring:
+    gcp_workload_monitor = GCPWorkloadMonitor(config.run_name)
+    if config.report_heartbeat_metric_for_gcp_monitoring:
+      gcp_workload_monitor.start_heartbeat_reporting_thread(config.heartbeat_reporting_interval_in_seconds)
+    if config.report_performance_metric_for_gcp_monitoring:
+      performance_metric_queue = queue.Queue()
+      gcp_workload_monitor.start_performance_reporting_thread(performance_metric_queue)
+
+  metric_logger = MetricLogger(writer, config)
+  step = start_step
+
+  em.maybe_snapshot(
+      step,
+      snapshot={
+          "params": state.params,
+          "opt_state": state.opt_state,
+      },
+      force=True,
+      block=True,
+  )
+
+  while True:
+    with watchdog.watchdog(120):
+      try:
+        if step == first_profiling_step or prof.should_activate_periodic_profile(step):
+          optional_postfix = f"step_{step}" if config.profile_periodically_period > 0 else ""
+          prof.activate(blocking_object=state, optional_postfix=optional_postfix)
+
+        if step >= config.steps:
+          break
+
+        max_logging.log(f"{step=} {em.elastic_down_event_count=} {em.good_slice_count=}")
+        with mesh, nn_partitioning.axis_rules(config.logical_axis_rules), jax.default_device(em.default_device):
+          with jax.profiler.StepTraceAnnotation("train", step_num=step):
+            record_goodput(recorder, config, recorder.record_data_loading_start_time if recorder else None)
+            example_batch = load_next_batch(data_iterator, example_batch, config)
+            record_goodput(recorder, config, recorder.record_data_loading_end_time if recorder else None)
+            check_example_batch(config, example_batch=example_batch)
+            # pylint: disable=not-callable
+            nextrng = jax.jit(jax.random.fold_in)(init_rng, step)
+            record_goodput(recorder, config, recorder.record_step_start_time if recorder else None, step)
+            state, metrics = p_train_step(state, example_batch, nextrng)
+
+          step_time_delta = datetime.datetime.now() - last_step_completion
+          last_step_completion = datetime.datetime.now()
+          record_scalar_metrics(metrics, step_time_delta, per_device_tflops, learning_rate_schedule(step), per_device_tokens)
+          if performance_metric_queue:
+            performance_metric_queue.put(step_time_delta.total_seconds())
+
+          if checkpoint_manager is not None:
+            state_to_save = state
+            if save_checkpoint(checkpoint_manager, int(step), state_to_save, config.dataset_type, data_iterator, config):
+              checkpointing.print_save_message(step, config.async_checkpointing)
+
+            # Upon preemption, exit when and only when all ongoing saves are complete.
+            if checkpoint_manager.reached_preemption(step):
+              checkpoint_manager.wait_until_finished()
+              sys.exit()
+
+          metric_logger.write_metrics(running_gcs_metrics, metrics, step)
+
+          if step == last_profiling_step or prof.should_deactivate_periodic_profile(step):
+            prof.deactivate(blocking_object=state)
+
+        em.maybe_snapshot(
+            step=step,
+            snapshot={
+                "params": state.params,
+                "opt_state": state.opt_state,
+            },
+            block=True,
+        )
+
+        ret = em.maybe_reshard_up(
+            step=step,
+            snapshot={
+                "params": state.params,
+                "opt_state": state.opt_state,
+            },
+            elastic_handler=elastic_handler,
+            handler_kwargs={
+                "config": config,
+                "em": em,
+                "checkpoint_manager": checkpoint_manager,
+            },
+        )
+        if ret is not None:
+          (
+              config,
+              step,
+              state,
+              mesh,
+              checkpoint_manager,
+              data_iterator,
+              p_train_step,
+              example_batch,
+              learning_rate_schedule,
+              metric_logger,
+              writer,
+          ) = ret
+
+        if step == start_step:
+          max_utils.print_mem_stats("After params initialized")
+
+        step += 1
+
+      except jax.errors.JaxRuntimeError as error:
+        ret = em.maybe_reshard_down(
+            error=error,
+            elastic_handler=elastic_handler,
+            handler_kwargs={
+                "config": config,
+                "em": em,
+                "checkpoint_manager": checkpoint_manager,
+            },
+        )
+        if ret is not None:
+          (
+              config,
+              step,
+              state,
+              mesh,
+              checkpoint_manager,
+              data_iterator,
+              p_train_step,
+              example_batch,
+              learning_rate_schedule,
+              metric_logger,
+              writer,
+          ) = ret
+
+  if checkpoint_manager is not None:
+    checkpoint_manager.wait_until_finished()
+  metric_logger.write_metrics(running_gcs_metrics, metrics, config.steps - 1)  # final step metrics
+  max_utils.close_summary_writer(writer)
+  record_goodput(recorder, config, recorder.record_job_end_time if recorder else None)
+  with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
+    # pytype: disable=attribute-error
+    compiled = p_train_step.lower(state, example_batch, nextrng).compile()
+    compiled_stats = compiled.memory_analysis()
+    if compiled_stats is not None:
+      max_logging.log(
+          f"Output size: {compiled_stats.output_size_in_bytes}, "
+          f"temp size: {compiled_stats.temp_size_in_bytes}, "
+          f"argument size: {compiled_stats.argument_size_in_bytes}, "
+          f"host temp size: {compiled_stats.host_temp_size_in_bytes}, in bytes."
+      )
+  return state
+
+
+def main(argv: Sequence[str]) -> None:
+  pathwaysutils.initialize()
+  jax.config.update("jax_default_prng_impl", "unsafe_rbg")
+  # TF allocates extraneous GPU memory when using TFDS data
+  # this leads to CUDA OOMs. WAR for now is to hide GPUs from TF
+  tf.config.set_visible_devices([], "GPU")
+  os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"
+  if "xla_tpu_spmd_rng_bit_generator_unsafe" not in os.environ.get("LIBTPU_INIT_ARGS", ""):
+    os.environ["LIBTPU_INIT_ARGS"] = os.environ.get("LIBTPU_INIT_ARGS", "") + " --xla_tpu_spmd_rng_bit_generator_unsafe=true"
+
+  em = manager.Manager(
+      jax.devices(),
+      reshard_check_period=1,
+      snapshot_period=1,
+      max_elastic_down_event_count=100,
+      max_reshard_retry_count=3,
+  )
+
+  em.good_slice_indices = em.get_slice_availability()
+  while len(em.good_slice_indices) < em.total_slice_count:
path: MaxText/elastic_train.py
user: login: lukebaumann
body: Added a to-do.
created_at: 2025-04-04T17:48:23Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,452 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput, record_scalar_metrics,
+    save_checkpoint, setup_mesh_and_model, setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger('pathwaysutils.elastic.manager').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.elastic.reshard').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.debug.timing').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.debug.watchdog').setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    em,
path: MaxText/elastic_train.py
user: login: lukebaumann
body: Done
created_at: 2025-04-04T17:48:40Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,452 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
path: MaxText/elastic_train.py
user: login: lukebaumann
body: Removed
created_at: 2025-04-04T17:49:17Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,452 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
path: MaxText/elastic_train.py
user: login: lukebaumann
body: Done
created_at: 2025-04-04T17:56:42Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,452 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput, record_scalar_metrics,
+    save_checkpoint, setup_mesh_and_model, setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
path: MaxText/elastic_train.py
user: login: lukebaumann
body: Added a comment that these are hard coded for the moment.
created_at: 2025-04-04T17:59:51Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,452 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput, record_scalar_metrics,
+    save_checkpoint, setup_mesh_and_model, setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger('pathwaysutils.elastic.manager').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.elastic.reshard').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.debug.timing').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.debug.watchdog').setLevel(logging.DEBUG)
+
+
+@timing.timeit
path: MaxText/elastic_train.py
user: login: lukebaumann
body: I want to log how long the function call takes. `timing.timeit` is a decorator included with pathwaysutils to time every call to the function.
created_at: 2025-04-04T18:07:11Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,452 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput, record_scalar_metrics,
+    save_checkpoint, setup_mesh_and_model, setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger('pathwaysutils.elastic.manager').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.elastic.reshard').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.debug.timing').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.debug.watchdog').setLevel(logging.DEBUG)
+
+
+@timing.timeit
path: MaxText/elastic_train.py
user: login: lukebaumann
body: What the function does is discussed at a high level in the module docstring.
created_at: 2025-04-04T18:07:32Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,444 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput,
+    record_scalar_metrics,
+    save_checkpoint,
+    setup_mesh_and_model,
+    setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.watchdog").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    em,
+    checkpoint_manager,
+):
+  """Reshard function."""
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(em.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, em.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = em.get_resharded_snapshot(mesh)
+
+      if checkpoint_manager is not None:
+        # Confirm this is the right thing to do
+        latest_step = checkpoint_manager.latest_step()
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
+          checkpoint_manager.delete(latest_step)
+
+      state, _, state_mesh_shardings, data_iterator = max_utils.setup_training_state(
+          model,
+          data_iterator,
+          tx,
+          config,
+          jax.random.fold_in(init_rng, step),
+          mesh,
+          checkpoint_manager=None,
+      )
+
+      state = state.replace(**snapshot)
+      state = state.replace(step=state.step.at[None].set(step))
+
+      (
+          functional_train,
+          in_shard_train,
+          out_shard_train,
+          static_argnums_train,
+          donate_argnums_train,
+      ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+      p_train_step = jax.jit(
+          functional_train,
+          in_shardings=in_shard_train,
+          out_shardings=out_shard_train,
+          static_argnums=static_argnums_train,
+          donate_argnums=donate_argnums_train,
+      )
+
+      example_batch = None
+      metric_logger = MetricLogger(writer, config)
+
+      jax.block_until_ready(state)
+
+  return (
+      config,
+      step,
+      state,
+      mesh,
+      checkpoint_manager,
+      data_iterator,
+      p_train_step,
+      example_batch,
+      learning_rate_schedule,
+      metric_logger,
+      writer,
+  )
+
+
+def train_loop(config, em, state=None):
+  """Main Training loop.
+  Args:
+    config:
+    state:
+    ckpt_path:
+  Returns:
+  """
+  # Create a GoodputRecorder to log information
+  recorder = create_goodput_recorder(config)
+  record_goodput(recorder, config, recorder.record_job_start_time if recorder else None)
+
+  (
+      init_rng,
+      writer,
+      checkpoint_manager,
+      state_mesh_shardings,
+      model,
+      mesh,
+      learning_rate_schedule,
+      data_iterator,
+      _,
+      state,
+  ) = setup_train_loop(config)
+
+  # pylint: disable=line-too-long
+  (
+      functional_train,
+      in_shard_train,
+      out_shard_train,
+      static_argnums_train,
+      donate_argnums_train,
+  ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+  p_train_step = jax.jit(
+      functional_train,
+      in_shardings=in_shard_train,
+      out_shardings=out_shard_train,
+      static_argnums=static_argnums_train,
+      donate_argnums=donate_argnums_train,
+  )
+
+  num_model_parameters = max_utils.calculate_num_params_from_pytree(state.params)
+  max_logging.log(f"number parameters: {num_model_parameters/1e9:.3f} billion")
+  per_device_tflops, _, _ = maxtext_utils.calculate_tflops_training_per_device(config)
+  per_device_tokens = maxtext_utils.calculate_tokens_training_per_device(config)
+
+  # Write train config params, num model params, and XLA flags to tensorboard
+  max_utils.add_text_to_summary_writer("num_model_parameters", str(num_model_parameters), writer)
+  max_utils.add_text_to_summary_writer("libtpu_init_args", os.environ["LIBTPU_INIT_ARGS"], writer)
+  max_utils.add_config_to_summary_writer(config, writer)
+
+  running_gcs_metrics = [] if config.gcs_metrics else None
+
+  start_step = get_first_step(state)  # this is the start_step for training
+  prof = profiler.Profiler(config, offset_step=start_step)
+  first_profiling_step = prof.start_initial_profile_step
+  if config.profiler != "" and first_profiling_step >= config.steps:
+    raise ValueError("Profiling requested but initial profiling step set past training final step")
+  last_profiling_step = prof.finished_initial_profile_step
+
+  example_batch = None
+  last_step_completion = datetime.datetime.now()
+
+  performance_metric_queue = None
+  if config.report_heartbeat_metric_for_gcp_monitoring or config.report_performance_metric_for_gcp_monitoring:
+    gcp_workload_monitor = GCPWorkloadMonitor(config.run_name)
+    if config.report_heartbeat_metric_for_gcp_monitoring:
+      gcp_workload_monitor.start_heartbeat_reporting_thread(config.heartbeat_reporting_interval_in_seconds)
+    if config.report_performance_metric_for_gcp_monitoring:
+      performance_metric_queue = queue.Queue()
+      gcp_workload_monitor.start_performance_reporting_thread(performance_metric_queue)
+
+  metric_logger = MetricLogger(writer, config)
+  step = start_step
+
+  em.maybe_snapshot(
+      step,
+      snapshot={
+          "params": state.params,
+          "opt_state": state.opt_state,
+      },
+      force=True,
+      block=True,
+  )
+
+  while True:
+    with watchdog.watchdog(120):
+      try:
+        if step == first_profiling_step or prof.should_activate_periodic_profile(step):
+          optional_postfix = f"step_{step}" if config.profile_periodically_period > 0 else ""
+          prof.activate(blocking_object=state, optional_postfix=optional_postfix)
+
+        if step >= config.steps:
+          break
+
+        max_logging.log(f"{step=} {em.elastic_down_event_count=} {em.good_slice_count=}")
+        with mesh, nn_partitioning.axis_rules(config.logical_axis_rules), jax.default_device(em.default_device):
+          with jax.profiler.StepTraceAnnotation("train", step_num=step):
+            record_goodput(recorder, config, recorder.record_data_loading_start_time if recorder else None)
+            example_batch = load_next_batch(data_iterator, example_batch, config)
+            record_goodput(recorder, config, recorder.record_data_loading_end_time if recorder else None)
+            check_example_batch(config, example_batch=example_batch)
+            # pylint: disable=not-callable
+            nextrng = jax.jit(jax.random.fold_in)(init_rng, step)
+            record_goodput(recorder, config, recorder.record_step_start_time if recorder else None, step)
+            state, metrics = p_train_step(state, example_batch, nextrng)
+
+          step_time_delta = datetime.datetime.now() - last_step_completion
+          last_step_completion = datetime.datetime.now()
+          record_scalar_metrics(metrics, step_time_delta, per_device_tflops, learning_rate_schedule(step), per_device_tokens)
+          if performance_metric_queue:
+            performance_metric_queue.put(step_time_delta.total_seconds())
+
+          if checkpoint_manager is not None:
+            state_to_save = state
+            if save_checkpoint(checkpoint_manager, int(step), state_to_save, config.dataset_type, data_iterator, config):
+              checkpointing.print_save_message(step, config.async_checkpointing)
+
+            # Upon preemption, exit when and only when all ongoing saves are complete.
+            if checkpoint_manager.reached_preemption(step):
+              checkpoint_manager.wait_until_finished()
+              sys.exit()
+
+          metric_logger.write_metrics(running_gcs_metrics, metrics, step)
+
+          if step == last_profiling_step or prof.should_deactivate_periodic_profile(step):
+            prof.deactivate(blocking_object=state)
+
+        em.maybe_snapshot(
path: MaxText/elastic_train.py
user: login: lukebaumann
body: This is 2 function calls and one if. It is extra long because there are several variables that need to be updated by the elastic handler after reconfiguring and these are split by line.

I think I might wrap things in a follow up and will make a bug to track it.
created_at: 2025-04-04T18:14:41Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,452 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput, record_scalar_metrics,
+    save_checkpoint, setup_mesh_and_model, setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger('pathwaysutils.elastic.manager').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.elastic.reshard').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.debug.timing').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.debug.watchdog').setLevel(logging.DEBUG)
+
+
+@timing.timeit
path: MaxText/elastic_train.py
user: login: gobbleturk
body: Is there a more appropriate name for this method - perhaps maybe_elastic_reshard?
created_at: 2025-04-04T18:15:16Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,452 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput, record_scalar_metrics,
+    save_checkpoint, setup_mesh_and_model, setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger('pathwaysutils.elastic.manager').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.elastic.reshard').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.debug.timing').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.debug.watchdog').setLevel(logging.DEBUG)
+
+
+@timing.timeit
path: MaxText/elastic_train.py
user: login: gobbleturk
body: Is there a more appropriate name for this method? elastic_handler sounds like a noun instead of a verb, its surprising to see it as a method name
created_at: 2025-04-08T18:04:02Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,444 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput,
+    record_scalar_metrics,
+    save_checkpoint,
+    setup_mesh_and_model,
+    setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.watchdog").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    em,
+    checkpoint_manager,
+):
+  """Reshard function."""
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(em.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, em.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = em.get_resharded_snapshot(mesh)
+
+      if checkpoint_manager is not None:
+        # Confirm this is the right thing to do
path: MaxText/elastic_train.py
user: login: gobbleturk
body: Can you clarify this comment?
created_at: 2025-04-08T18:04:39Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,444 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput,
+    record_scalar_metrics,
+    save_checkpoint,
+    setup_mesh_and_model,
+    setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.watchdog").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    em,
+    checkpoint_manager,
+):
+  """Reshard function."""
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(em.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, em.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = em.get_resharded_snapshot(mesh)
+
+      if checkpoint_manager is not None:
+        # Confirm this is the right thing to do
+        latest_step = checkpoint_manager.latest_step()
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
path: MaxText/elastic_train.py
user: login: gobbleturk
body: Can you include some comments about this logic, see above question?
created_at: 2025-04-08T18:05:08Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,444 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput,
+    record_scalar_metrics,
+    save_checkpoint,
+    setup_mesh_and_model,
+    setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.watchdog").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    em,
+    checkpoint_manager,
+):
+  """Reshard function."""
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(em.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, em.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = em.get_resharded_snapshot(mesh)
+
+      if checkpoint_manager is not None:
+        # Confirm this is the right thing to do
+        latest_step = checkpoint_manager.latest_step()
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
+          checkpoint_manager.delete(latest_step)
+
+      state, _, state_mesh_shardings, data_iterator = max_utils.setup_training_state(
+          model,
+          data_iterator,
+          tx,
+          config,
+          jax.random.fold_in(init_rng, step),
+          mesh,
+          checkpoint_manager=None,
+      )
+
+      state = state.replace(**snapshot)
+      state = state.replace(step=state.step.at[None].set(step))
+
+      (
+          functional_train,
+          in_shard_train,
+          out_shard_train,
+          static_argnums_train,
+          donate_argnums_train,
+      ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+      p_train_step = jax.jit(
+          functional_train,
+          in_shardings=in_shard_train,
+          out_shardings=out_shard_train,
+          static_argnums=static_argnums_train,
+          donate_argnums=donate_argnums_train,
+      )
+
+      example_batch = None
+      metric_logger = MetricLogger(writer, config)
+
+      jax.block_until_ready(state)
+
+  return (
+      config,
+      step,
+      state,
+      mesh,
+      checkpoint_manager,
+      data_iterator,
+      p_train_step,
+      example_batch,
+      learning_rate_schedule,
+      metric_logger,
+      writer,
+  )
+
+
+def train_loop(config, em, state=None):
+  """Main Training loop.
+  Args:
+    config:
+    state:
+    ckpt_path:
+  Returns:
+  """
+  # Create a GoodputRecorder to log information
+  recorder = create_goodput_recorder(config)
+  record_goodput(recorder, config, recorder.record_job_start_time if recorder else None)
+
+  (
+      init_rng,
+      writer,
+      checkpoint_manager,
+      state_mesh_shardings,
+      model,
+      mesh,
+      learning_rate_schedule,
+      data_iterator,
+      _,
+      state,
+  ) = setup_train_loop(config)
+
+  # pylint: disable=line-too-long
+  (
+      functional_train,
+      in_shard_train,
+      out_shard_train,
+      static_argnums_train,
+      donate_argnums_train,
+  ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+  p_train_step = jax.jit(
+      functional_train,
+      in_shardings=in_shard_train,
+      out_shardings=out_shard_train,
+      static_argnums=static_argnums_train,
+      donate_argnums=donate_argnums_train,
+  )
+
+  num_model_parameters = max_utils.calculate_num_params_from_pytree(state.params)
+  max_logging.log(f"number parameters: {num_model_parameters/1e9:.3f} billion")
+  per_device_tflops, _, _ = maxtext_utils.calculate_tflops_training_per_device(config)
+  per_device_tokens = maxtext_utils.calculate_tokens_training_per_device(config)
+
+  # Write train config params, num model params, and XLA flags to tensorboard
+  max_utils.add_text_to_summary_writer("num_model_parameters", str(num_model_parameters), writer)
+  max_utils.add_text_to_summary_writer("libtpu_init_args", os.environ["LIBTPU_INIT_ARGS"], writer)
+  max_utils.add_config_to_summary_writer(config, writer)
+
+  running_gcs_metrics = [] if config.gcs_metrics else None
+
+  start_step = get_first_step(state)  # this is the start_step for training
+  prof = profiler.Profiler(config, offset_step=start_step)
+  first_profiling_step = prof.start_initial_profile_step
+  if config.profiler != "" and first_profiling_step >= config.steps:
+    raise ValueError("Profiling requested but initial profiling step set past training final step")
+  last_profiling_step = prof.finished_initial_profile_step
+
+  example_batch = None
+  last_step_completion = datetime.datetime.now()
+
+  performance_metric_queue = None
+  if config.report_heartbeat_metric_for_gcp_monitoring or config.report_performance_metric_for_gcp_monitoring:
+    gcp_workload_monitor = GCPWorkloadMonitor(config.run_name)
+    if config.report_heartbeat_metric_for_gcp_monitoring:
+      gcp_workload_monitor.start_heartbeat_reporting_thread(config.heartbeat_reporting_interval_in_seconds)
+    if config.report_performance_metric_for_gcp_monitoring:
+      performance_metric_queue = queue.Queue()
+      gcp_workload_monitor.start_performance_reporting_thread(performance_metric_queue)
+
+  metric_logger = MetricLogger(writer, config)
+  step = start_step
+
+  em.maybe_snapshot(
+      step,
+      snapshot={
+          "params": state.params,
+          "opt_state": state.opt_state,
+      },
+      force=True,
+      block=True,
+  )
+
+  while True:
path: MaxText/elastic_train.py
user: login: gobbleturk
body: ^
created_at: 2025-04-08T18:05:27Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,444 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput,
+    record_scalar_metrics,
+    save_checkpoint,
+    setup_mesh_and_model,
+    setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.watchdog").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    em,
+    checkpoint_manager,
+):
+  """Reshard function."""
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(em.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, em.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = em.get_resharded_snapshot(mesh)
+
+      if checkpoint_manager is not None:
+        # Confirm this is the right thing to do
+        latest_step = checkpoint_manager.latest_step()
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
+          checkpoint_manager.delete(latest_step)
+
+      state, _, state_mesh_shardings, data_iterator = max_utils.setup_training_state(
+          model,
+          data_iterator,
+          tx,
+          config,
+          jax.random.fold_in(init_rng, step),
+          mesh,
+          checkpoint_manager=None,
+      )
+
+      state = state.replace(**snapshot)
+      state = state.replace(step=state.step.at[None].set(step))
+
+      (
+          functional_train,
+          in_shard_train,
+          out_shard_train,
+          static_argnums_train,
+          donate_argnums_train,
+      ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+      p_train_step = jax.jit(
+          functional_train,
+          in_shardings=in_shard_train,
+          out_shardings=out_shard_train,
+          static_argnums=static_argnums_train,
+          donate_argnums=donate_argnums_train,
+      )
+
+      example_batch = None
+      metric_logger = MetricLogger(writer, config)
+
+      jax.block_until_ready(state)
+
+  return (
+      config,
+      step,
+      state,
+      mesh,
+      checkpoint_manager,
+      data_iterator,
+      p_train_step,
+      example_batch,
+      learning_rate_schedule,
+      metric_logger,
+      writer,
+  )
+
+
+def train_loop(config, em, state=None):
+  """Main Training loop.
+  Args:
+    config:
+    state:
+    ckpt_path:
+  Returns:
+  """
+  # Create a GoodputRecorder to log information
+  recorder = create_goodput_recorder(config)
+  record_goodput(recorder, config, recorder.record_job_start_time if recorder else None)
+
+  (
+      init_rng,
+      writer,
+      checkpoint_manager,
+      state_mesh_shardings,
+      model,
+      mesh,
+      learning_rate_schedule,
+      data_iterator,
+      _,
+      state,
+  ) = setup_train_loop(config)
+
+  # pylint: disable=line-too-long
+  (
+      functional_train,
+      in_shard_train,
+      out_shard_train,
+      static_argnums_train,
+      donate_argnums_train,
+  ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+  p_train_step = jax.jit(
+      functional_train,
+      in_shardings=in_shard_train,
+      out_shardings=out_shard_train,
+      static_argnums=static_argnums_train,
+      donate_argnums=donate_argnums_train,
+  )
+
+  num_model_parameters = max_utils.calculate_num_params_from_pytree(state.params)
+  max_logging.log(f"number parameters: {num_model_parameters/1e9:.3f} billion")
+  per_device_tflops, _, _ = maxtext_utils.calculate_tflops_training_per_device(config)
+  per_device_tokens = maxtext_utils.calculate_tokens_training_per_device(config)
+
+  # Write train config params, num model params, and XLA flags to tensorboard
+  max_utils.add_text_to_summary_writer("num_model_parameters", str(num_model_parameters), writer)
+  max_utils.add_text_to_summary_writer("libtpu_init_args", os.environ["LIBTPU_INIT_ARGS"], writer)
+  max_utils.add_config_to_summary_writer(config, writer)
+
+  running_gcs_metrics = [] if config.gcs_metrics else None
+
+  start_step = get_first_step(state)  # this is the start_step for training
+  prof = profiler.Profiler(config, offset_step=start_step)
+  first_profiling_step = prof.start_initial_profile_step
+  if config.profiler != "" and first_profiling_step >= config.steps:
+    raise ValueError("Profiling requested but initial profiling step set past training final step")
+  last_profiling_step = prof.finished_initial_profile_step
+
+  example_batch = None
+  last_step_completion = datetime.datetime.now()
+
+  performance_metric_queue = None
+  if config.report_heartbeat_metric_for_gcp_monitoring or config.report_performance_metric_for_gcp_monitoring:
+    gcp_workload_monitor = GCPWorkloadMonitor(config.run_name)
+    if config.report_heartbeat_metric_for_gcp_monitoring:
+      gcp_workload_monitor.start_heartbeat_reporting_thread(config.heartbeat_reporting_interval_in_seconds)
+    if config.report_performance_metric_for_gcp_monitoring:
+      performance_metric_queue = queue.Queue()
+      gcp_workload_monitor.start_performance_reporting_thread(performance_metric_queue)
+
+  metric_logger = MetricLogger(writer, config)
+  step = start_step
+
+  em.maybe_snapshot(
+      step,
+      snapshot={
+          "params": state.params,
+          "opt_state": state.opt_state,
+      },
+      force=True,
+      block=True,
+  )
+
+  while True:
+    with watchdog.watchdog(120):
path: MaxText/elastic_train.py
user: login: SurbhiJainUSC
body: I think `watchdog` monitors file system events and stack trace package would periodically log python traces in addition of logging traces during failures. @lukebaumann - what is the purpose of `watchdog` in this case?
created_at: 2025-04-08T18:16:25Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,472 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model.
+
+Elastic training via Pathways on Cloud allows for slices to go down without
+crashing the main program. In this way, elastic events (slices going up and
+down) can be caught or polled for and reacted to within an elastic_handler
+function without causing the whole workload to restart. This involves
+creating a new mesh with the available slices, reinitializing variables,
+and recompiling functions in addition to restoring from a host offloaded
+snapshot of the state.
+
+The purpose of this training loop is to serve as an example for how to
+support elasticity and is actively in development. As such, there are some
+performance optimizations that have yet to be added as well as some features
+not supported.
+
+Current limitations:
+- The host offloaded snapshot is currently blocking
+- Elastic event handling during async checkpointing is not tested
+- Elastic event handling during profiling is not tested
+- Elastic manager configuration values are hard coded
+- Debug logging statements for elasticity are hard coded as enabled
+
+See https://github.com/AI-Hypercomputer/pathways-utils/tree/main/pathwaysutils/elastic
+for more details about the elastic manager.
+"""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput,
+    record_scalar_metrics,
+    save_checkpoint,
+    setup_mesh_and_model,
+    setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.watchdog").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    elastic_manager,
+    checkpoint_manager,
+):
+  """Reshard function."""
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(elastic_manager.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, elastic_manager.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = elastic_manager.get_resharded_snapshot(mesh)
+
+      if checkpoint_manager is not None:
+        # Confirm this is the right thing to do
+        latest_step = checkpoint_manager.latest_step()
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
+          checkpoint_manager.delete(latest_step)
+
+      state, _, state_mesh_shardings, data_iterator = max_utils.setup_training_state(
+          model,
+          data_iterator,
+          tx,
+          config,
+          jax.random.fold_in(init_rng, step),
+          mesh,
+          checkpoint_manager=None,
+      )
+
+      state = state.replace(**snapshot)
+      state = state.replace(step=state.step.at[None].set(step))
+
+      (
+          functional_train,
+          in_shard_train,
+          out_shard_train,
+          static_argnums_train,
+          donate_argnums_train,
+      ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+      p_train_step = jax.jit(
+          functional_train,
+          in_shardings=in_shard_train,
+          out_shardings=out_shard_train,
+          static_argnums=static_argnums_train,
+          donate_argnums=donate_argnums_train,
+      )
+
+      example_batch = None
+      metric_logger = MetricLogger(writer, config)
+
+      jax.block_until_ready(state)
+
+  return (
+      config,
+      step,
+      state,
+      mesh,
+      checkpoint_manager,
+      data_iterator,
+      p_train_step,
+      example_batch,
+      learning_rate_schedule,
+      metric_logger,
+      writer,
+  )
+
+
+def train_loop(config, elastic_manager, state=None):
+  """Main Training loop.
+  Args:
+    config:
+    state:
+    ckpt_path:
+  Returns:
+  """
+  # Create a GoodputRecorder to log information
+  recorder = create_goodput_recorder(config)
+  record_goodput(recorder, config, recorder.record_job_start_time if recorder else None)
+
+  (
+      init_rng,
+      writer,
+      checkpoint_manager,
+      state_mesh_shardings,
+      model,
+      mesh,
+      learning_rate_schedule,
+      data_iterator,
+      _,
+      state,
+  ) = setup_train_loop(config)
+
+  # pylint: disable=line-too-long
+  (
+      functional_train,
+      in_shard_train,
+      out_shard_train,
+      static_argnums_train,
+      donate_argnums_train,
+  ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+  p_train_step = jax.jit(
+      functional_train,
+      in_shardings=in_shard_train,
+      out_shardings=out_shard_train,
+      static_argnums=static_argnums_train,
+      donate_argnums=donate_argnums_train,
+  )
+
+  num_model_parameters = max_utils.calculate_num_params_from_pytree(state.params)
+  max_logging.log(f"number parameters: {num_model_parameters/1e9:.3f} billion")
+  per_device_tflops, _, _ = maxtext_utils.calculate_tflops_training_per_device(config)
+  per_device_tokens = maxtext_utils.calculate_tokens_training_per_device(config)
+
+  # Write train config params, num model params, and XLA flags to tensorboard
+  max_utils.add_text_to_summary_writer("num_model_parameters", str(num_model_parameters), writer)
+  max_utils.add_text_to_summary_writer("libtpu_init_args", os.environ["LIBTPU_INIT_ARGS"], writer)
+  max_utils.add_config_to_summary_writer(config, writer)
+
+  running_gcs_metrics = [] if config.gcs_metrics else None
+
+  start_step = get_first_step(state)  # this is the start_step for training
+  prof = profiler.Profiler(config, offset_step=start_step)
+  first_profiling_step = prof.start_initial_profile_step
+  if config.profiler != "" and first_profiling_step >= config.steps:
+    raise ValueError("Profiling requested but initial profiling step set past training final step")
+  last_profiling_step = prof.finished_initial_profile_step
+
+  example_batch = None
+  last_step_completion = datetime.datetime.now()
+
+  performance_metric_queue = None
+  if config.report_heartbeat_metric_for_gcp_monitoring or config.report_performance_metric_for_gcp_monitoring:
+    gcp_workload_monitor = GCPWorkloadMonitor(config.run_name)
+    if config.report_heartbeat_metric_for_gcp_monitoring:
+      gcp_workload_monitor.start_heartbeat_reporting_thread(config.heartbeat_reporting_interval_in_seconds)
+    if config.report_performance_metric_for_gcp_monitoring:
+      performance_metric_queue = queue.Queue()
+      gcp_workload_monitor.start_performance_reporting_thread(performance_metric_queue)
+
+  metric_logger = MetricLogger(writer, config)
+  step = start_step
+
+  elastic_manager.maybe_snapshot(
+      step,
+      snapshot={
+          "params": state.params,
+          "opt_state": state.opt_state,
+      },
+      force=True,
+      block=True,
+  )
+
+  while True:
+    with watchdog.watchdog(120):
+      try:
+        if step == first_profiling_step or prof.should_activate_periodic_profile(step):
+          optional_postfix = f"step_{step}" if config.profile_periodically_period > 0 else ""
+          prof.activate(blocking_object=state, optional_postfix=optional_postfix)
+
+        if step >= config.steps:
+          break
+
+        max_logging.log(f"{step=} {elastic_manager.elastic_down_event_count=} {elastic_manager.good_slice_count=}")
+        with mesh, nn_partitioning.axis_rules(config.logical_axis_rules), jax.default_device(elastic_manager.default_device):
+          with jax.profiler.StepTraceAnnotation("train", step_num=step):
+            record_goodput(recorder, config, recorder.record_data_loading_start_time if recorder else None)
+            example_batch = load_next_batch(data_iterator, example_batch, config)
+            record_goodput(recorder, config, recorder.record_data_loading_end_time if recorder else None)
+            check_example_batch(config, example_batch=example_batch)
+            # pylint: disable=not-callable
+            nextrng = jax.jit(jax.random.fold_in)(init_rng, step)
+            record_goodput(recorder, config, recorder.record_step_start_time if recorder else None, step)
+            state, metrics = p_train_step(state, example_batch, nextrng)
+
+          step_time_delta = datetime.datetime.now() - last_step_completion
+          last_step_completion = datetime.datetime.now()
+          record_scalar_metrics(metrics, step_time_delta, per_device_tflops, learning_rate_schedule(step), per_device_tokens)
+          if performance_metric_queue:
+            performance_metric_queue.put(step_time_delta.total_seconds())
+
+          if checkpoint_manager is not None:
+            state_to_save = state
+            if save_checkpoint(checkpoint_manager, int(step), state_to_save, config.dataset_type, data_iterator, config):
+              checkpointing.print_save_message(step, config.async_checkpointing)
+
+            # Upon preemption, exit when and only when all ongoing saves are complete.
+            if checkpoint_manager.reached_preemption(step):
+              checkpoint_manager.wait_until_finished()
+              sys.exit()
+
+          metric_logger.write_metrics(running_gcs_metrics, metrics, step)
+
+          if step == last_profiling_step or prof.should_deactivate_periodic_profile(step):
+            prof.deactivate(blocking_object=state)
+
+        elastic_manager.maybe_snapshot(
+            step=step,
+            snapshot={
+                "params": state.params,
+                "opt_state": state.opt_state,
+            },
+            block=True,
+        )
+
+        ret = elastic_manager.maybe_reshard_up(
+            step=step,
+            snapshot={
+                "params": state.params,
+                "opt_state": state.opt_state,
+            },
+            elastic_handler=elastic_handler,
+            handler_kwargs={
+                "config": config,
+                "elastic_manager": elastic_manager,
+                "checkpoint_manager": checkpoint_manager,
+            },
+        )
+        if ret is not None:
+          (
+              config,
+              step,
+              state,
+              mesh,
+              checkpoint_manager,
+              data_iterator,
+              p_train_step,
+              example_batch,
+              learning_rate_schedule,
+              metric_logger,
+              writer,
+          ) = ret
+
+        if step == start_step:
+          max_utils.print_mem_stats("After params initialized")
+
+        step += 1
+
+      except jax.errors.JaxRuntimeError as error:
+        ret = elastic_manager.maybe_reshard_down(
+            error=error,
+            elastic_handler=elastic_handler,
+            handler_kwargs={
+                "config": config,
+                "elastic_manager": elastic_manager,
+                "checkpoint_manager": checkpoint_manager,
+            },
+        )
+        if ret is not None:
+          (
+              config,
+              step,
+              state,
+              mesh,
+              checkpoint_manager,
+              data_iterator,
+              p_train_step,
+              example_batch,
+              learning_rate_schedule,
+              metric_logger,
+              writer,
+          ) = ret
+
+  if checkpoint_manager is not None:
path: MaxText/elastic_train.py
user: login: SurbhiJainUSC
body: Did you verify if you are able to save the checkpoint for the last executed step. If not, consider adding this: https://github.com/AI-Hypercomputer/maxtext/pull/1443
created_at: 2025-04-08T21:16:33Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,472 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model.
+
+Elastic training via Pathways on Cloud allows for slices to go down without
+crashing the main program. In this way, elastic events (slices going up and
+down) can be caught or polled for and reacted to within an elastic_handler
+function without causing the whole workload to restart. This involves
+creating a new mesh with the available slices, reinitializing variables,
+and recompiling functions in addition to restoring from a host offloaded
+snapshot of the state.
+
+The purpose of this training loop is to serve as an example for how to
+support elasticity and is actively in development. As such, there are some
+performance optimizations that have yet to be added as well as some features
+not supported.
+
+Current limitations:
+- The host offloaded snapshot is currently blocking
+- Elastic event handling during async checkpointing is not tested
+- Elastic event handling during profiling is not tested
+- Elastic manager configuration values are hard coded
+- Debug logging statements for elasticity are hard coded as enabled
+
+See https://github.com/AI-Hypercomputer/pathways-utils/tree/main/pathwaysutils/elastic
+for more details about the elastic manager.
+"""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput,
+    record_scalar_metrics,
+    save_checkpoint,
+    setup_mesh_and_model,
+    setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.watchdog").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    elastic_manager,
+    checkpoint_manager,
+):
+  """Reshard function."""
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(elastic_manager.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, elastic_manager.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = elastic_manager.get_resharded_snapshot(mesh)
+
+      if checkpoint_manager is not None:
+        # Confirm this is the right thing to do
+        latest_step = checkpoint_manager.latest_step()
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
+          checkpoint_manager.delete(latest_step)
+
+      state, _, state_mesh_shardings, data_iterator = max_utils.setup_training_state(
+          model,
+          data_iterator,
+          tx,
+          config,
+          jax.random.fold_in(init_rng, step),
+          mesh,
+          checkpoint_manager=None,
+      )
+
+      state = state.replace(**snapshot)
+      state = state.replace(step=state.step.at[None].set(step))
+
+      (
+          functional_train,
+          in_shard_train,
+          out_shard_train,
+          static_argnums_train,
+          donate_argnums_train,
+      ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+      p_train_step = jax.jit(
+          functional_train,
+          in_shardings=in_shard_train,
+          out_shardings=out_shard_train,
+          static_argnums=static_argnums_train,
+          donate_argnums=donate_argnums_train,
+      )
+
+      example_batch = None
+      metric_logger = MetricLogger(writer, config)
+
+      jax.block_until_ready(state)
+
+  return (
+      config,
+      step,
+      state,
+      mesh,
+      checkpoint_manager,
+      data_iterator,
+      p_train_step,
+      example_batch,
+      learning_rate_schedule,
+      metric_logger,
+      writer,
+  )
+
+
+def train_loop(config, elastic_manager, state=None):
+  """Main Training loop.
+  Args:
+    config:
+    state:
+    ckpt_path:
+  Returns:
+  """
+  # Create a GoodputRecorder to log information
+  recorder = create_goodput_recorder(config)
+  record_goodput(recorder, config, recorder.record_job_start_time if recorder else None)
+
+  (
+      init_rng,
+      writer,
+      checkpoint_manager,
+      state_mesh_shardings,
+      model,
+      mesh,
+      learning_rate_schedule,
+      data_iterator,
+      _,
+      state,
+  ) = setup_train_loop(config)
+
+  # pylint: disable=line-too-long
+  (
+      functional_train,
+      in_shard_train,
+      out_shard_train,
+      static_argnums_train,
+      donate_argnums_train,
+  ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+  p_train_step = jax.jit(
+      functional_train,
+      in_shardings=in_shard_train,
+      out_shardings=out_shard_train,
+      static_argnums=static_argnums_train,
+      donate_argnums=donate_argnums_train,
+  )
+
+  num_model_parameters = max_utils.calculate_num_params_from_pytree(state.params)
+  max_logging.log(f"number parameters: {num_model_parameters/1e9:.3f} billion")
+  per_device_tflops, _, _ = maxtext_utils.calculate_tflops_training_per_device(config)
+  per_device_tokens = maxtext_utils.calculate_tokens_training_per_device(config)
+
+  # Write train config params, num model params, and XLA flags to tensorboard
+  max_utils.add_text_to_summary_writer("num_model_parameters", str(num_model_parameters), writer)
+  max_utils.add_text_to_summary_writer("libtpu_init_args", os.environ["LIBTPU_INIT_ARGS"], writer)
+  max_utils.add_config_to_summary_writer(config, writer)
+
+  running_gcs_metrics = [] if config.gcs_metrics else None
+
+  start_step = get_first_step(state)  # this is the start_step for training
+  prof = profiler.Profiler(config, offset_step=start_step)
+  first_profiling_step = prof.start_initial_profile_step
+  if config.profiler != "" and first_profiling_step >= config.steps:
+    raise ValueError("Profiling requested but initial profiling step set past training final step")
+  last_profiling_step = prof.finished_initial_profile_step
+
+  example_batch = None
+  last_step_completion = datetime.datetime.now()
+
+  performance_metric_queue = None
+  if config.report_heartbeat_metric_for_gcp_monitoring or config.report_performance_metric_for_gcp_monitoring:
+    gcp_workload_monitor = GCPWorkloadMonitor(config.run_name)
+    if config.report_heartbeat_metric_for_gcp_monitoring:
+      gcp_workload_monitor.start_heartbeat_reporting_thread(config.heartbeat_reporting_interval_in_seconds)
+    if config.report_performance_metric_for_gcp_monitoring:
+      performance_metric_queue = queue.Queue()
+      gcp_workload_monitor.start_performance_reporting_thread(performance_metric_queue)
+
+  metric_logger = MetricLogger(writer, config)
+  step = start_step
+
+  elastic_manager.maybe_snapshot(
+      step,
+      snapshot={
+          "params": state.params,
+          "opt_state": state.opt_state,
+      },
+      force=True,
+      block=True,
+  )
+
+  while True:
+    with watchdog.watchdog(120):
+      try:
+        if step == first_profiling_step or prof.should_activate_periodic_profile(step):
+          optional_postfix = f"step_{step}" if config.profile_periodically_period > 0 else ""
+          prof.activate(blocking_object=state, optional_postfix=optional_postfix)
+
+        if step >= config.steps:
+          break
+
+        max_logging.log(f"{step=} {elastic_manager.elastic_down_event_count=} {elastic_manager.good_slice_count=}")
+        with mesh, nn_partitioning.axis_rules(config.logical_axis_rules), jax.default_device(elastic_manager.default_device):
+          with jax.profiler.StepTraceAnnotation("train", step_num=step):
+            record_goodput(recorder, config, recorder.record_data_loading_start_time if recorder else None)
+            example_batch = load_next_batch(data_iterator, example_batch, config)
+            record_goodput(recorder, config, recorder.record_data_loading_end_time if recorder else None)
+            check_example_batch(config, example_batch=example_batch)
+            # pylint: disable=not-callable
+            nextrng = jax.jit(jax.random.fold_in)(init_rng, step)
+            record_goodput(recorder, config, recorder.record_step_start_time if recorder else None, step)
+            state, metrics = p_train_step(state, example_batch, nextrng)
+
+          step_time_delta = datetime.datetime.now() - last_step_completion
+          last_step_completion = datetime.datetime.now()
+          record_scalar_metrics(metrics, step_time_delta, per_device_tflops, learning_rate_schedule(step), per_device_tokens)
+          if performance_metric_queue:
+            performance_metric_queue.put(step_time_delta.total_seconds())
+
+          if checkpoint_manager is not None:
+            state_to_save = state
+            if save_checkpoint(checkpoint_manager, int(step), state_to_save, config.dataset_type, data_iterator, config):
+              checkpointing.print_save_message(step, config.async_checkpointing)
+
+            # Upon preemption, exit when and only when all ongoing saves are complete.
+            if checkpoint_manager.reached_preemption(step):
+              checkpoint_manager.wait_until_finished()
+              sys.exit()
+
+          metric_logger.write_metrics(running_gcs_metrics, metrics, step)
+
+          if step == last_profiling_step or prof.should_deactivate_periodic_profile(step):
+            prof.deactivate(blocking_object=state)
+
+        elastic_manager.maybe_snapshot(
+            step=step,
+            snapshot={
+                "params": state.params,
+                "opt_state": state.opt_state,
+            },
+            block=True,
+        )
+
+        ret = elastic_manager.maybe_reshard_up(
+            step=step,
+            snapshot={
+                "params": state.params,
+                "opt_state": state.opt_state,
+            },
+            elastic_handler=elastic_handler,
+            handler_kwargs={
+                "config": config,
+                "elastic_manager": elastic_manager,
+                "checkpoint_manager": checkpoint_manager,
+            },
+        )
+        if ret is not None:
+          (
+              config,
+              step,
+              state,
+              mesh,
+              checkpoint_manager,
+              data_iterator,
+              p_train_step,
+              example_batch,
+              learning_rate_schedule,
+              metric_logger,
+              writer,
+          ) = ret
+
+        if step == start_step:
+          max_utils.print_mem_stats("After params initialized")
+
+        step += 1
+
+      except jax.errors.JaxRuntimeError as error:
+        ret = elastic_manager.maybe_reshard_down(
+            error=error,
+            elastic_handler=elastic_handler,
+            handler_kwargs={
+                "config": config,
+                "elastic_manager": elastic_manager,
+                "checkpoint_manager": checkpoint_manager,
+            },
+        )
+        if ret is not None:
+          (
+              config,
+              step,
+              state,
+              mesh,
+              checkpoint_manager,
+              data_iterator,
+              p_train_step,
+              example_batch,
+              learning_rate_schedule,
+              metric_logger,
+              writer,
+          ) = ret
+
+  if checkpoint_manager is not None:
+    checkpoint_manager.wait_until_finished()
+  metric_logger.write_metrics(running_gcs_metrics, metrics, config.steps - 1)  # final step metrics
+  max_utils.close_summary_writer(writer)
+  record_goodput(recorder, config, recorder.record_job_end_time if recorder else None)
+  with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
+    # pytype: disable=attribute-error
+    compiled = p_train_step.lower(state, example_batch, nextrng).compile()
+    compiled_stats = compiled.memory_analysis()
+    if compiled_stats is not None:
+      max_logging.log(
+          f"Output size: {compiled_stats.output_size_in_bytes}, "
+          f"temp size: {compiled_stats.temp_size_in_bytes}, "
+          f"argument size: {compiled_stats.argument_size_in_bytes}, "
+          f"host temp size: {compiled_stats.host_temp_size_in_bytes}, in bytes."
+      )
+  return state
+
+
+def main(argv: Sequence[str]) -> None:
+  pathwaysutils.initialize()
+  jax.config.update("jax_default_prng_impl", "unsafe_rbg")
+  # TF allocates extraneous GPU memory when using TFDS data
+  # this leads to CUDA OOMs. WAR for now is to hide GPUs from TF
+  tf.config.set_visible_devices([], "GPU")
+  os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"
+  if "xla_tpu_spmd_rng_bit_generator_unsafe" not in os.environ.get("LIBTPU_INIT_ARGS", ""):
+    os.environ["LIBTPU_INIT_ARGS"] = os.environ.get("LIBTPU_INIT_ARGS", "") + " --xla_tpu_spmd_rng_bit_generator_unsafe=true"
+
+  elastic_manager = manager.Manager(
+      jax.devices(),
+      reshard_check_period=1,
+      snapshot_period=5,
+      max_elastic_down_event_count=100,
+      max_reshard_retry_count=3,
+  )
+
+  # Do not start training until all slices are available
+  # TODO: b/408455557 - Migrate to pathwaysutils and make configurable
+  elastic_manager.good_slice_indices = elastic_manager.get_slice_availability()
+  while len(elastic_manager.good_slice_indices) < elastic_manager.total_slice_count:
+    max_logging.log(
+        f"Only {elastic_manager.good_slice_count} slices out of {elastic_manager.total_slice_count} available. "
+        "Sleeping for 5 seconds."
+    )
+    time.sleep(5)
+    elastic_manager.good_slice_indices = elastic_manager.get_slice_availability()
+  max_logging.log("All slices are available")
+
+  pyconfig.HyperParameters.global_batch_size_to_train_on = property(
+      lambda self: elastic_manager.scale_by_good_slices(self.get_keys()["global_batch_size_to_train_on"])
+  )
+  pyconfig.HyperParameters.global_batch_size_to_load = property(
+      lambda self: elastic_manager.scale_by_good_slices(self.get_keys()["global_batch_size_to_load"])
+  )
+  pyconfig.HyperParameters.micro_batch_size_to_train_on = property(
+      lambda self: elastic_manager.scale_by_good_slices(self.get_keys()["micro_batch_size_to_train_on"])
+  )
+  pyconfig.HyperParameters.num_slices = property(lambda self: elastic_manager.good_slice_count)
+
+  config = pyconfig.initialize(argv)
+  max_utils.print_system_information()
+  validate_train_config(config)
+  os.environ["TFDS_DATA_DIR"] = config.dataset_path
+  vertex_tensorboard_manager = VertexTensorboardManager()
+  if config.use_vertex_tensorboard or os.environ.get("UPLOAD_DATA_TO_TENSORBOARD"):
+    vertex_tensorboard_manager.configure_vertex_tensorboard(config)
+
+  if config.monitor_goodput and jax.process_index() == 0:
+    logger_name = f"goodput_{config.run_name}"
+    goodput_monitor = monitoring.GoodputMonitor(
path: MaxText/elastic_train.py
user: login: SurbhiJainUSC
body: FYI: there is a [PR](https://github.com/AI-Hypercomputer/maxtext/pull/1539) in review for Goodput and those changes might be needed here as well.
created_at: 2025-04-08T21:20:26Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,452 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput, record_scalar_metrics,
+    save_checkpoint, setup_mesh_and_model, setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger('pathwaysutils.elastic.manager').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.elastic.reshard').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.debug.timing').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.debug.watchdog').setLevel(logging.DEBUG)
+
+
+@timing.timeit
path: MaxText/elastic_train.py
user: login: lukebaumann
body: Talked offline. Keeping the name elastic_handler since it is how it is named in pathwaysutils and we can keep it consistent and event handler functions are commonly named this way and this is effectively an event handler for elastic events.
created_at: 2025-04-09T00:09:48Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,444 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput,
+    record_scalar_metrics,
+    save_checkpoint,
+    setup_mesh_and_model,
+    setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.watchdog").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    em,
+    checkpoint_manager,
+):
+  """Reshard function."""
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(em.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, em.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = em.get_resharded_snapshot(mesh)
+
+      if checkpoint_manager is not None:
+        # Confirm this is the right thing to do
path: MaxText/elastic_train.py
user: login: lukebaumann
body: Nice catch!
created_at: 2025-04-09T00:10:08Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,444 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput,
+    record_scalar_metrics,
+    save_checkpoint,
+    setup_mesh_and_model,
+    setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.watchdog").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    em,
+    checkpoint_manager,
+):
+  """Reshard function."""
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(em.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, em.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = em.get_resharded_snapshot(mesh)
+
+      if checkpoint_manager is not None:
+        # Confirm this is the right thing to do
+        latest_step = checkpoint_manager.latest_step()
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
path: MaxText/elastic_train.py
user: login: lukebaumann
body: Added comments explaining this section
created_at: 2025-04-09T00:18:47Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,444 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput,
+    record_scalar_metrics,
+    save_checkpoint,
+    setup_mesh_and_model,
+    setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.watchdog").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    em,
+    checkpoint_manager,
+):
+  """Reshard function."""
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(em.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, em.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = em.get_resharded_snapshot(mesh)
+
+      if checkpoint_manager is not None:
+        # Confirm this is the right thing to do
+        latest_step = checkpoint_manager.latest_step()
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
+          checkpoint_manager.delete(latest_step)
+
+      state, _, state_mesh_shardings, data_iterator = max_utils.setup_training_state(
+          model,
+          data_iterator,
+          tx,
+          config,
+          jax.random.fold_in(init_rng, step),
+          mesh,
+          checkpoint_manager=None,
+      )
+
+      state = state.replace(**snapshot)
+      state = state.replace(step=state.step.at[None].set(step))
+
+      (
+          functional_train,
+          in_shard_train,
+          out_shard_train,
+          static_argnums_train,
+          donate_argnums_train,
+      ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+      p_train_step = jax.jit(
+          functional_train,
+          in_shardings=in_shard_train,
+          out_shardings=out_shard_train,
+          static_argnums=static_argnums_train,
+          donate_argnums=donate_argnums_train,
+      )
+
+      example_batch = None
+      metric_logger = MetricLogger(writer, config)
+
+      jax.block_until_ready(state)
+
+  return (
+      config,
+      step,
+      state,
+      mesh,
+      checkpoint_manager,
+      data_iterator,
+      p_train_step,
+      example_batch,
+      learning_rate_schedule,
+      metric_logger,
+      writer,
+  )
+
+
+def train_loop(config, em, state=None):
+  """Main Training loop.
+  Args:
+    config:
+    state:
+    ckpt_path:
+  Returns:
+  """
+  # Create a GoodputRecorder to log information
+  recorder = create_goodput_recorder(config)
+  record_goodput(recorder, config, recorder.record_job_start_time if recorder else None)
+
+  (
+      init_rng,
+      writer,
+      checkpoint_manager,
+      state_mesh_shardings,
+      model,
+      mesh,
+      learning_rate_schedule,
+      data_iterator,
+      _,
+      state,
+  ) = setup_train_loop(config)
+
+  # pylint: disable=line-too-long
+  (
+      functional_train,
+      in_shard_train,
+      out_shard_train,
+      static_argnums_train,
+      donate_argnums_train,
+  ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+  p_train_step = jax.jit(
+      functional_train,
+      in_shardings=in_shard_train,
+      out_shardings=out_shard_train,
+      static_argnums=static_argnums_train,
+      donate_argnums=donate_argnums_train,
+  )
+
+  num_model_parameters = max_utils.calculate_num_params_from_pytree(state.params)
+  max_logging.log(f"number parameters: {num_model_parameters/1e9:.3f} billion")
+  per_device_tflops, _, _ = maxtext_utils.calculate_tflops_training_per_device(config)
+  per_device_tokens = maxtext_utils.calculate_tokens_training_per_device(config)
+
+  # Write train config params, num model params, and XLA flags to tensorboard
+  max_utils.add_text_to_summary_writer("num_model_parameters", str(num_model_parameters), writer)
+  max_utils.add_text_to_summary_writer("libtpu_init_args", os.environ["LIBTPU_INIT_ARGS"], writer)
+  max_utils.add_config_to_summary_writer(config, writer)
+
+  running_gcs_metrics = [] if config.gcs_metrics else None
+
+  start_step = get_first_step(state)  # this is the start_step for training
+  prof = profiler.Profiler(config, offset_step=start_step)
+  first_profiling_step = prof.start_initial_profile_step
+  if config.profiler != "" and first_profiling_step >= config.steps:
+    raise ValueError("Profiling requested but initial profiling step set past training final step")
+  last_profiling_step = prof.finished_initial_profile_step
+
+  example_batch = None
+  last_step_completion = datetime.datetime.now()
+
+  performance_metric_queue = None
+  if config.report_heartbeat_metric_for_gcp_monitoring or config.report_performance_metric_for_gcp_monitoring:
+    gcp_workload_monitor = GCPWorkloadMonitor(config.run_name)
+    if config.report_heartbeat_metric_for_gcp_monitoring:
+      gcp_workload_monitor.start_heartbeat_reporting_thread(config.heartbeat_reporting_interval_in_seconds)
+    if config.report_performance_metric_for_gcp_monitoring:
+      performance_metric_queue = queue.Queue()
+      gcp_workload_monitor.start_performance_reporting_thread(performance_metric_queue)
+
+  metric_logger = MetricLogger(writer, config)
+  step = start_step
+
+  em.maybe_snapshot(
+      step,
+      snapshot={
+          "params": state.params,
+          "opt_state": state.opt_state,
+      },
+      force=True,
+      block=True,
+  )
+
+  while True:
path: MaxText/elastic_train.py
user: login: lukebaumann
body: Thanks! Changed to `while step < config.step`
created_at: 2025-04-09T01:14:48Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,444 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput,
+    record_scalar_metrics,
+    save_checkpoint,
+    setup_mesh_and_model,
+    setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.watchdog").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    em,
+    checkpoint_manager,
+):
+  """Reshard function."""
path: MaxText/elastic_train.py
user: login: lukebaumann
body: I updated this comment to describe what it does more fully.
created_at: 2025-04-09T02:11:35Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,444 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput,
+    record_scalar_metrics,
+    save_checkpoint,
+    setup_mesh_and_model,
+    setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.watchdog").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    em,
+    checkpoint_manager,
+):
+  """Reshard function."""
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(em.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, em.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = em.get_resharded_snapshot(mesh)
+
+      if checkpoint_manager is not None:
+        # Confirm this is the right thing to do
+        latest_step = checkpoint_manager.latest_step()
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
+          checkpoint_manager.delete(latest_step)
+
+      state, _, state_mesh_shardings, data_iterator = max_utils.setup_training_state(
+          model,
+          data_iterator,
+          tx,
+          config,
+          jax.random.fold_in(init_rng, step),
+          mesh,
+          checkpoint_manager=None,
+      )
+
+      state = state.replace(**snapshot)
+      state = state.replace(step=state.step.at[None].set(step))
+
+      (
+          functional_train,
+          in_shard_train,
+          out_shard_train,
+          static_argnums_train,
+          donate_argnums_train,
+      ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+      p_train_step = jax.jit(
+          functional_train,
+          in_shardings=in_shard_train,
+          out_shardings=out_shard_train,
+          static_argnums=static_argnums_train,
+          donate_argnums=donate_argnums_train,
+      )
+
+      example_batch = None
+      metric_logger = MetricLogger(writer, config)
+
+      jax.block_until_ready(state)
+
+  return (
+      config,
+      step,
+      state,
+      mesh,
+      checkpoint_manager,
+      data_iterator,
+      p_train_step,
+      example_batch,
+      learning_rate_schedule,
+      metric_logger,
+      writer,
+  )
+
+
+def train_loop(config, em, state=None):
+  """Main Training loop.
+  Args:
+    config:
+    state:
+    ckpt_path:
+  Returns:
+  """
+  # Create a GoodputRecorder to log information
+  recorder = create_goodput_recorder(config)
+  record_goodput(recorder, config, recorder.record_job_start_time if recorder else None)
+
+  (
+      init_rng,
+      writer,
+      checkpoint_manager,
+      state_mesh_shardings,
+      model,
+      mesh,
+      learning_rate_schedule,
+      data_iterator,
+      _,
+      state,
+  ) = setup_train_loop(config)
+
+  # pylint: disable=line-too-long
+  (
+      functional_train,
+      in_shard_train,
+      out_shard_train,
+      static_argnums_train,
+      donate_argnums_train,
+  ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+  p_train_step = jax.jit(
+      functional_train,
+      in_shardings=in_shard_train,
+      out_shardings=out_shard_train,
+      static_argnums=static_argnums_train,
+      donate_argnums=donate_argnums_train,
+  )
+
+  num_model_parameters = max_utils.calculate_num_params_from_pytree(state.params)
+  max_logging.log(f"number parameters: {num_model_parameters/1e9:.3f} billion")
+  per_device_tflops, _, _ = maxtext_utils.calculate_tflops_training_per_device(config)
+  per_device_tokens = maxtext_utils.calculate_tokens_training_per_device(config)
+
+  # Write train config params, num model params, and XLA flags to tensorboard
+  max_utils.add_text_to_summary_writer("num_model_parameters", str(num_model_parameters), writer)
+  max_utils.add_text_to_summary_writer("libtpu_init_args", os.environ["LIBTPU_INIT_ARGS"], writer)
+  max_utils.add_config_to_summary_writer(config, writer)
+
+  running_gcs_metrics = [] if config.gcs_metrics else None
+
+  start_step = get_first_step(state)  # this is the start_step for training
+  prof = profiler.Profiler(config, offset_step=start_step)
+  first_profiling_step = prof.start_initial_profile_step
+  if config.profiler != "" and first_profiling_step >= config.steps:
+    raise ValueError("Profiling requested but initial profiling step set past training final step")
+  last_profiling_step = prof.finished_initial_profile_step
+
+  example_batch = None
+  last_step_completion = datetime.datetime.now()
+
+  performance_metric_queue = None
+  if config.report_heartbeat_metric_for_gcp_monitoring or config.report_performance_metric_for_gcp_monitoring:
+    gcp_workload_monitor = GCPWorkloadMonitor(config.run_name)
+    if config.report_heartbeat_metric_for_gcp_monitoring:
+      gcp_workload_monitor.start_heartbeat_reporting_thread(config.heartbeat_reporting_interval_in_seconds)
+    if config.report_performance_metric_for_gcp_monitoring:
+      performance_metric_queue = queue.Queue()
+      gcp_workload_monitor.start_performance_reporting_thread(performance_metric_queue)
+
+  metric_logger = MetricLogger(writer, config)
+  step = start_step
+
+  em.maybe_snapshot(
+      step,
+      snapshot={
+          "params": state.params,
+          "opt_state": state.opt_state,
+      },
+      force=True,
+      block=True,
+  )
+
+  while True:
+    with watchdog.watchdog(120):
path: MaxText/elastic_train.py
user: login: lukebaumann
body: I removed the watchdog since it was for debugging during development.

To answer the question, `watchdog(n)` is a context manager that starts a thread to log the stack traces of all threads every `n` seconds until the code block is complete. If the code block executes in under `n` seconds, no logs are printed. The idea is that if a step is taking longer than expected, the stack can be logged for debugging purposes. This can happen under normal circumstances, e.g., if compilation is needed, the step might be longer than expected and print the stack trace. There is no monitoring of system events.

As I understand the stack trace package, it periodically logs the stack of all active threads regardless of what is going on in the code. This functionality can easily be added to the stack trace package if you think it would be useful.

This was added to pathwaysutils but is not Pathways specific. It is inspired by the watchdog in AXLearn.
created_at: 2025-04-09T02:24:25Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,472 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model.
+
+Elastic training via Pathways on Cloud allows for slices to go down without
+crashing the main program. In this way, elastic events (slices going up and
+down) can be caught or polled for and reacted to within an elastic_handler
+function without causing the whole workload to restart. This involves
+creating a new mesh with the available slices, reinitializing variables,
+and recompiling functions in addition to restoring from a host offloaded
+snapshot of the state.
+
+The purpose of this training loop is to serve as an example for how to
+support elasticity and is actively in development. As such, there are some
+performance optimizations that have yet to be added as well as some features
+not supported.
+
+Current limitations:
+- The host offloaded snapshot is currently blocking
+- Elastic event handling during async checkpointing is not tested
+- Elastic event handling during profiling is not tested
+- Elastic manager configuration values are hard coded
+- Debug logging statements for elasticity are hard coded as enabled
+
+See https://github.com/AI-Hypercomputer/pathways-utils/tree/main/pathwaysutils/elastic
+for more details about the elastic manager.
+"""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput,
+    record_scalar_metrics,
+    save_checkpoint,
+    setup_mesh_and_model,
+    setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.watchdog").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    elastic_manager,
+    checkpoint_manager,
+):
+  """Reshard function."""
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(elastic_manager.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, elastic_manager.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = elastic_manager.get_resharded_snapshot(mesh)
+
+      if checkpoint_manager is not None:
+        # Confirm this is the right thing to do
+        latest_step = checkpoint_manager.latest_step()
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
+          checkpoint_manager.delete(latest_step)
+
+      state, _, state_mesh_shardings, data_iterator = max_utils.setup_training_state(
+          model,
+          data_iterator,
+          tx,
+          config,
+          jax.random.fold_in(init_rng, step),
+          mesh,
+          checkpoint_manager=None,
+      )
+
+      state = state.replace(**snapshot)
+      state = state.replace(step=state.step.at[None].set(step))
+
+      (
+          functional_train,
+          in_shard_train,
+          out_shard_train,
+          static_argnums_train,
+          donate_argnums_train,
+      ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+      p_train_step = jax.jit(
+          functional_train,
+          in_shardings=in_shard_train,
+          out_shardings=out_shard_train,
+          static_argnums=static_argnums_train,
+          donate_argnums=donate_argnums_train,
+      )
+
+      example_batch = None
+      metric_logger = MetricLogger(writer, config)
+
+      jax.block_until_ready(state)
+
+  return (
+      config,
+      step,
+      state,
+      mesh,
+      checkpoint_manager,
+      data_iterator,
+      p_train_step,
+      example_batch,
+      learning_rate_schedule,
+      metric_logger,
+      writer,
+  )
+
+
+def train_loop(config, elastic_manager, state=None):
+  """Main Training loop.
+  Args:
+    config:
+    state:
+    ckpt_path:
+  Returns:
+  """
+  # Create a GoodputRecorder to log information
+  recorder = create_goodput_recorder(config)
+  record_goodput(recorder, config, recorder.record_job_start_time if recorder else None)
+
+  (
+      init_rng,
+      writer,
+      checkpoint_manager,
+      state_mesh_shardings,
+      model,
+      mesh,
+      learning_rate_schedule,
+      data_iterator,
+      _,
+      state,
+  ) = setup_train_loop(config)
+
+  # pylint: disable=line-too-long
+  (
+      functional_train,
+      in_shard_train,
+      out_shard_train,
+      static_argnums_train,
+      donate_argnums_train,
+  ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+  p_train_step = jax.jit(
+      functional_train,
+      in_shardings=in_shard_train,
+      out_shardings=out_shard_train,
+      static_argnums=static_argnums_train,
+      donate_argnums=donate_argnums_train,
+  )
+
+  num_model_parameters = max_utils.calculate_num_params_from_pytree(state.params)
+  max_logging.log(f"number parameters: {num_model_parameters/1e9:.3f} billion")
+  per_device_tflops, _, _ = maxtext_utils.calculate_tflops_training_per_device(config)
+  per_device_tokens = maxtext_utils.calculate_tokens_training_per_device(config)
+
+  # Write train config params, num model params, and XLA flags to tensorboard
+  max_utils.add_text_to_summary_writer("num_model_parameters", str(num_model_parameters), writer)
+  max_utils.add_text_to_summary_writer("libtpu_init_args", os.environ["LIBTPU_INIT_ARGS"], writer)
+  max_utils.add_config_to_summary_writer(config, writer)
+
+  running_gcs_metrics = [] if config.gcs_metrics else None
+
+  start_step = get_first_step(state)  # this is the start_step for training
+  prof = profiler.Profiler(config, offset_step=start_step)
+  first_profiling_step = prof.start_initial_profile_step
+  if config.profiler != "" and first_profiling_step >= config.steps:
+    raise ValueError("Profiling requested but initial profiling step set past training final step")
+  last_profiling_step = prof.finished_initial_profile_step
+
+  example_batch = None
+  last_step_completion = datetime.datetime.now()
+
+  performance_metric_queue = None
+  if config.report_heartbeat_metric_for_gcp_monitoring or config.report_performance_metric_for_gcp_monitoring:
+    gcp_workload_monitor = GCPWorkloadMonitor(config.run_name)
+    if config.report_heartbeat_metric_for_gcp_monitoring:
+      gcp_workload_monitor.start_heartbeat_reporting_thread(config.heartbeat_reporting_interval_in_seconds)
+    if config.report_performance_metric_for_gcp_monitoring:
+      performance_metric_queue = queue.Queue()
+      gcp_workload_monitor.start_performance_reporting_thread(performance_metric_queue)
+
+  metric_logger = MetricLogger(writer, config)
+  step = start_step
+
+  elastic_manager.maybe_snapshot(
+      step,
+      snapshot={
+          "params": state.params,
+          "opt_state": state.opt_state,
+      },
+      force=True,
+      block=True,
+  )
+
+  while True:
+    with watchdog.watchdog(120):
+      try:
+        if step == first_profiling_step or prof.should_activate_periodic_profile(step):
+          optional_postfix = f"step_{step}" if config.profile_periodically_period > 0 else ""
+          prof.activate(blocking_object=state, optional_postfix=optional_postfix)
+
+        if step >= config.steps:
+          break
+
+        max_logging.log(f"{step=} {elastic_manager.elastic_down_event_count=} {elastic_manager.good_slice_count=}")
+        with mesh, nn_partitioning.axis_rules(config.logical_axis_rules), jax.default_device(elastic_manager.default_device):
+          with jax.profiler.StepTraceAnnotation("train", step_num=step):
+            record_goodput(recorder, config, recorder.record_data_loading_start_time if recorder else None)
+            example_batch = load_next_batch(data_iterator, example_batch, config)
+            record_goodput(recorder, config, recorder.record_data_loading_end_time if recorder else None)
+            check_example_batch(config, example_batch=example_batch)
+            # pylint: disable=not-callable
+            nextrng = jax.jit(jax.random.fold_in)(init_rng, step)
+            record_goodput(recorder, config, recorder.record_step_start_time if recorder else None, step)
+            state, metrics = p_train_step(state, example_batch, nextrng)
+
+          step_time_delta = datetime.datetime.now() - last_step_completion
+          last_step_completion = datetime.datetime.now()
+          record_scalar_metrics(metrics, step_time_delta, per_device_tflops, learning_rate_schedule(step), per_device_tokens)
+          if performance_metric_queue:
+            performance_metric_queue.put(step_time_delta.total_seconds())
+
+          if checkpoint_manager is not None:
+            state_to_save = state
+            if save_checkpoint(checkpoint_manager, int(step), state_to_save, config.dataset_type, data_iterator, config):
+              checkpointing.print_save_message(step, config.async_checkpointing)
+
+            # Upon preemption, exit when and only when all ongoing saves are complete.
+            if checkpoint_manager.reached_preemption(step):
+              checkpoint_manager.wait_until_finished()
+              sys.exit()
+
+          metric_logger.write_metrics(running_gcs_metrics, metrics, step)
+
+          if step == last_profiling_step or prof.should_deactivate_periodic_profile(step):
+            prof.deactivate(blocking_object=state)
+
+        elastic_manager.maybe_snapshot(
+            step=step,
+            snapshot={
+                "params": state.params,
+                "opt_state": state.opt_state,
+            },
+            block=True,
+        )
+
+        ret = elastic_manager.maybe_reshard_up(
+            step=step,
+            snapshot={
+                "params": state.params,
+                "opt_state": state.opt_state,
+            },
+            elastic_handler=elastic_handler,
+            handler_kwargs={
+                "config": config,
+                "elastic_manager": elastic_manager,
+                "checkpoint_manager": checkpoint_manager,
+            },
+        )
+        if ret is not None:
+          (
+              config,
+              step,
+              state,
+              mesh,
+              checkpoint_manager,
+              data_iterator,
+              p_train_step,
+              example_batch,
+              learning_rate_schedule,
+              metric_logger,
+              writer,
+          ) = ret
+
+        if step == start_step:
+          max_utils.print_mem_stats("After params initialized")
+
+        step += 1
+
+      except jax.errors.JaxRuntimeError as error:
+        ret = elastic_manager.maybe_reshard_down(
+            error=error,
+            elastic_handler=elastic_handler,
+            handler_kwargs={
+                "config": config,
+                "elastic_manager": elastic_manager,
+                "checkpoint_manager": checkpoint_manager,
+            },
+        )
+        if ret is not None:
+          (
+              config,
+              step,
+              state,
+              mesh,
+              checkpoint_manager,
+              data_iterator,
+              p_train_step,
+              example_batch,
+              learning_rate_schedule,
+              metric_logger,
+              writer,
+          ) = ret
+
+  if checkpoint_manager is not None:
path: MaxText/elastic_train.py
user: login: lukebaumann
body: I did not verify that. I based this off of train.py and do not mean to make this different from train.py. I just rebased this off of train.py at HEAD so this should include #1443. (separately I left a comment on that PR)

As train.py becomes more of a library, I expect this code duplication to be reduced.
created_at: 2025-04-09T03:06:55Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,472 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model.
+
+Elastic training via Pathways on Cloud allows for slices to go down without
+crashing the main program. In this way, elastic events (slices going up and
+down) can be caught or polled for and reacted to within an elastic_handler
+function without causing the whole workload to restart. This involves
+creating a new mesh with the available slices, reinitializing variables,
+and recompiling functions in addition to restoring from a host offloaded
+snapshot of the state.
+
+The purpose of this training loop is to serve as an example for how to
+support elasticity and is actively in development. As such, there are some
+performance optimizations that have yet to be added as well as some features
+not supported.
+
+Current limitations:
+- The host offloaded snapshot is currently blocking
+- Elastic event handling during async checkpointing is not tested
+- Elastic event handling during profiling is not tested
+- Elastic manager configuration values are hard coded
+- Debug logging statements for elasticity are hard coded as enabled
+
+See https://github.com/AI-Hypercomputer/pathways-utils/tree/main/pathwaysutils/elastic
+for more details about the elastic manager.
+"""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput,
+    record_scalar_metrics,
+    save_checkpoint,
+    setup_mesh_and_model,
+    setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.watchdog").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    elastic_manager,
+    checkpoint_manager,
+):
+  """Reshard function."""
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(elastic_manager.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, elastic_manager.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = elastic_manager.get_resharded_snapshot(mesh)
+
+      if checkpoint_manager is not None:
+        # Confirm this is the right thing to do
+        latest_step = checkpoint_manager.latest_step()
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
+          checkpoint_manager.delete(latest_step)
+
+      state, _, state_mesh_shardings, data_iterator = max_utils.setup_training_state(
+          model,
+          data_iterator,
+          tx,
+          config,
+          jax.random.fold_in(init_rng, step),
+          mesh,
+          checkpoint_manager=None,
+      )
+
+      state = state.replace(**snapshot)
+      state = state.replace(step=state.step.at[None].set(step))
+
+      (
+          functional_train,
+          in_shard_train,
+          out_shard_train,
+          static_argnums_train,
+          donate_argnums_train,
+      ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+      p_train_step = jax.jit(
+          functional_train,
+          in_shardings=in_shard_train,
+          out_shardings=out_shard_train,
+          static_argnums=static_argnums_train,
+          donate_argnums=donate_argnums_train,
+      )
+
+      example_batch = None
+      metric_logger = MetricLogger(writer, config)
+
+      jax.block_until_ready(state)
+
+  return (
+      config,
+      step,
+      state,
+      mesh,
+      checkpoint_manager,
+      data_iterator,
+      p_train_step,
+      example_batch,
+      learning_rate_schedule,
+      metric_logger,
+      writer,
+  )
+
+
+def train_loop(config, elastic_manager, state=None):
+  """Main Training loop.
+  Args:
+    config:
+    state:
+    ckpt_path:
+  Returns:
+  """
+  # Create a GoodputRecorder to log information
+  recorder = create_goodput_recorder(config)
+  record_goodput(recorder, config, recorder.record_job_start_time if recorder else None)
+
+  (
+      init_rng,
+      writer,
+      checkpoint_manager,
+      state_mesh_shardings,
+      model,
+      mesh,
+      learning_rate_schedule,
+      data_iterator,
+      _,
+      state,
+  ) = setup_train_loop(config)
+
+  # pylint: disable=line-too-long
+  (
+      functional_train,
+      in_shard_train,
+      out_shard_train,
+      static_argnums_train,
+      donate_argnums_train,
+  ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+  p_train_step = jax.jit(
+      functional_train,
+      in_shardings=in_shard_train,
+      out_shardings=out_shard_train,
+      static_argnums=static_argnums_train,
+      donate_argnums=donate_argnums_train,
+  )
+
+  num_model_parameters = max_utils.calculate_num_params_from_pytree(state.params)
+  max_logging.log(f"number parameters: {num_model_parameters/1e9:.3f} billion")
+  per_device_tflops, _, _ = maxtext_utils.calculate_tflops_training_per_device(config)
+  per_device_tokens = maxtext_utils.calculate_tokens_training_per_device(config)
+
+  # Write train config params, num model params, and XLA flags to tensorboard
+  max_utils.add_text_to_summary_writer("num_model_parameters", str(num_model_parameters), writer)
+  max_utils.add_text_to_summary_writer("libtpu_init_args", os.environ["LIBTPU_INIT_ARGS"], writer)
+  max_utils.add_config_to_summary_writer(config, writer)
+
+  running_gcs_metrics = [] if config.gcs_metrics else None
+
+  start_step = get_first_step(state)  # this is the start_step for training
+  prof = profiler.Profiler(config, offset_step=start_step)
+  first_profiling_step = prof.start_initial_profile_step
+  if config.profiler != "" and first_profiling_step >= config.steps:
+    raise ValueError("Profiling requested but initial profiling step set past training final step")
+  last_profiling_step = prof.finished_initial_profile_step
+
+  example_batch = None
+  last_step_completion = datetime.datetime.now()
+
+  performance_metric_queue = None
+  if config.report_heartbeat_metric_for_gcp_monitoring or config.report_performance_metric_for_gcp_monitoring:
+    gcp_workload_monitor = GCPWorkloadMonitor(config.run_name)
+    if config.report_heartbeat_metric_for_gcp_monitoring:
+      gcp_workload_monitor.start_heartbeat_reporting_thread(config.heartbeat_reporting_interval_in_seconds)
+    if config.report_performance_metric_for_gcp_monitoring:
+      performance_metric_queue = queue.Queue()
+      gcp_workload_monitor.start_performance_reporting_thread(performance_metric_queue)
+
+  metric_logger = MetricLogger(writer, config)
+  step = start_step
+
+  elastic_manager.maybe_snapshot(
+      step,
+      snapshot={
+          "params": state.params,
+          "opt_state": state.opt_state,
+      },
+      force=True,
+      block=True,
+  )
+
+  while True:
+    with watchdog.watchdog(120):
+      try:
+        if step == first_profiling_step or prof.should_activate_periodic_profile(step):
+          optional_postfix = f"step_{step}" if config.profile_periodically_period > 0 else ""
+          prof.activate(blocking_object=state, optional_postfix=optional_postfix)
+
+        if step >= config.steps:
+          break
+
+        max_logging.log(f"{step=} {elastic_manager.elastic_down_event_count=} {elastic_manager.good_slice_count=}")
+        with mesh, nn_partitioning.axis_rules(config.logical_axis_rules), jax.default_device(elastic_manager.default_device):
+          with jax.profiler.StepTraceAnnotation("train", step_num=step):
+            record_goodput(recorder, config, recorder.record_data_loading_start_time if recorder else None)
+            example_batch = load_next_batch(data_iterator, example_batch, config)
+            record_goodput(recorder, config, recorder.record_data_loading_end_time if recorder else None)
+            check_example_batch(config, example_batch=example_batch)
+            # pylint: disable=not-callable
+            nextrng = jax.jit(jax.random.fold_in)(init_rng, step)
+            record_goodput(recorder, config, recorder.record_step_start_time if recorder else None, step)
+            state, metrics = p_train_step(state, example_batch, nextrng)
+
+          step_time_delta = datetime.datetime.now() - last_step_completion
+          last_step_completion = datetime.datetime.now()
+          record_scalar_metrics(metrics, step_time_delta, per_device_tflops, learning_rate_schedule(step), per_device_tokens)
+          if performance_metric_queue:
+            performance_metric_queue.put(step_time_delta.total_seconds())
+
+          if checkpoint_manager is not None:
+            state_to_save = state
+            if save_checkpoint(checkpoint_manager, int(step), state_to_save, config.dataset_type, data_iterator, config):
+              checkpointing.print_save_message(step, config.async_checkpointing)
+
+            # Upon preemption, exit when and only when all ongoing saves are complete.
+            if checkpoint_manager.reached_preemption(step):
+              checkpoint_manager.wait_until_finished()
+              sys.exit()
+
+          metric_logger.write_metrics(running_gcs_metrics, metrics, step)
+
+          if step == last_profiling_step or prof.should_deactivate_periodic_profile(step):
+            prof.deactivate(blocking_object=state)
+
+        elastic_manager.maybe_snapshot(
+            step=step,
+            snapshot={
+                "params": state.params,
+                "opt_state": state.opt_state,
+            },
+            block=True,
+        )
+
+        ret = elastic_manager.maybe_reshard_up(
+            step=step,
+            snapshot={
+                "params": state.params,
+                "opt_state": state.opt_state,
+            },
+            elastic_handler=elastic_handler,
+            handler_kwargs={
+                "config": config,
+                "elastic_manager": elastic_manager,
+                "checkpoint_manager": checkpoint_manager,
+            },
+        )
+        if ret is not None:
+          (
+              config,
+              step,
+              state,
+              mesh,
+              checkpoint_manager,
+              data_iterator,
+              p_train_step,
+              example_batch,
+              learning_rate_schedule,
+              metric_logger,
+              writer,
+          ) = ret
+
+        if step == start_step:
+          max_utils.print_mem_stats("After params initialized")
+
+        step += 1
+
+      except jax.errors.JaxRuntimeError as error:
+        ret = elastic_manager.maybe_reshard_down(
+            error=error,
+            elastic_handler=elastic_handler,
+            handler_kwargs={
+                "config": config,
+                "elastic_manager": elastic_manager,
+                "checkpoint_manager": checkpoint_manager,
+            },
+        )
+        if ret is not None:
+          (
+              config,
+              step,
+              state,
+              mesh,
+              checkpoint_manager,
+              data_iterator,
+              p_train_step,
+              example_batch,
+              learning_rate_schedule,
+              metric_logger,
+              writer,
+          ) = ret
+
+  if checkpoint_manager is not None:
+    checkpoint_manager.wait_until_finished()
+  metric_logger.write_metrics(running_gcs_metrics, metrics, config.steps - 1)  # final step metrics
+  max_utils.close_summary_writer(writer)
+  record_goodput(recorder, config, recorder.record_job_end_time if recorder else None)
+  with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
+    # pytype: disable=attribute-error
+    compiled = p_train_step.lower(state, example_batch, nextrng).compile()
+    compiled_stats = compiled.memory_analysis()
+    if compiled_stats is not None:
+      max_logging.log(
+          f"Output size: {compiled_stats.output_size_in_bytes}, "
+          f"temp size: {compiled_stats.temp_size_in_bytes}, "
+          f"argument size: {compiled_stats.argument_size_in_bytes}, "
+          f"host temp size: {compiled_stats.host_temp_size_in_bytes}, in bytes."
+      )
+  return state
+
+
+def main(argv: Sequence[str]) -> None:
+  pathwaysutils.initialize()
+  jax.config.update("jax_default_prng_impl", "unsafe_rbg")
+  # TF allocates extraneous GPU memory when using TFDS data
+  # this leads to CUDA OOMs. WAR for now is to hide GPUs from TF
+  tf.config.set_visible_devices([], "GPU")
+  os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"
+  if "xla_tpu_spmd_rng_bit_generator_unsafe" not in os.environ.get("LIBTPU_INIT_ARGS", ""):
+    os.environ["LIBTPU_INIT_ARGS"] = os.environ.get("LIBTPU_INIT_ARGS", "") + " --xla_tpu_spmd_rng_bit_generator_unsafe=true"
+
+  elastic_manager = manager.Manager(
+      jax.devices(),
+      reshard_check_period=1,
+      snapshot_period=5,
+      max_elastic_down_event_count=100,
+      max_reshard_retry_count=3,
+  )
+
+  # Do not start training until all slices are available
+  # TODO: b/408455557 - Migrate to pathwaysutils and make configurable
+  elastic_manager.good_slice_indices = elastic_manager.get_slice_availability()
+  while len(elastic_manager.good_slice_indices) < elastic_manager.total_slice_count:
+    max_logging.log(
+        f"Only {elastic_manager.good_slice_count} slices out of {elastic_manager.total_slice_count} available. "
+        "Sleeping for 5 seconds."
+    )
+    time.sleep(5)
+    elastic_manager.good_slice_indices = elastic_manager.get_slice_availability()
+  max_logging.log("All slices are available")
+
+  pyconfig.HyperParameters.global_batch_size_to_train_on = property(
+      lambda self: elastic_manager.scale_by_good_slices(self.get_keys()["global_batch_size_to_train_on"])
+  )
+  pyconfig.HyperParameters.global_batch_size_to_load = property(
+      lambda self: elastic_manager.scale_by_good_slices(self.get_keys()["global_batch_size_to_load"])
+  )
+  pyconfig.HyperParameters.micro_batch_size_to_train_on = property(
+      lambda self: elastic_manager.scale_by_good_slices(self.get_keys()["micro_batch_size_to_train_on"])
+  )
+  pyconfig.HyperParameters.num_slices = property(lambda self: elastic_manager.good_slice_count)
+
+  config = pyconfig.initialize(argv)
+  max_utils.print_system_information()
+  validate_train_config(config)
+  os.environ["TFDS_DATA_DIR"] = config.dataset_path
+  vertex_tensorboard_manager = VertexTensorboardManager()
+  if config.use_vertex_tensorboard or os.environ.get("UPLOAD_DATA_TO_TENSORBOARD"):
+    vertex_tensorboard_manager.configure_vertex_tensorboard(config)
+
+  if config.monitor_goodput and jax.process_index() == 0:
+    logger_name = f"goodput_{config.run_name}"
+    goodput_monitor = monitoring.GoodputMonitor(
path: MaxText/elastic_train.py
user: login: lukebaumann
body: I am not sure if there is a plan to make importable functions so that the different trainers can stay up-to-date but I think that would be very beneficial.
created_at: 2025-04-09T03:08:31Z
author_association: COLLABORATOR
diff_hunk: @@ -37,5 +36,5 @@ transformers
 mlperf-logging@git+https://github.com/mlperf/logging.git
 google-jetstream@git+https://github.com/AI-Hypercomputer/JetStream.git
 jsonlines
-pathwaysutils@git+https://github.com/google/pathways-utils.git
+pathwaysutils@git+https://github.com/AI-Hypercomputer/pathways-utils.git
path: requirements.txt
user: login: richjames0
body: Do we need to make sure this change is made in JStS also? IDK
created_at: 2025-04-09T16:44:43Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,452 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput, record_scalar_metrics,
+    save_checkpoint, setup_mesh_and_model, setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger('pathwaysutils.elastic.manager').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.elastic.reshard').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.debug.timing').setLevel(logging.DEBUG)
+logging.getLogger('pathwaysutils.debug.watchdog').setLevel(logging.DEBUG)
+
+
+@timing.timeit
path: MaxText/elastic_train.py
user: login: richjames0
body: +1 on commenting why there's a timeit here. I also found this surprising
created_at: 2025-04-09T16:51:28Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,452 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model."""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+# Placeholder: internal
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput, record_scalar_metrics,
+    save_checkpoint, setup_mesh_and_model, setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
path: MaxText/elastic_train.py
user: login: richjames0
body: +1 on Matt's comment. One of us is redundant here. 
created_at: 2025-04-09T16:52:04Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,493 @@
+# Copyright 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.
+
+"""Training loop for elastic training model.
+
+Elastic training via Pathways on Cloud allows for slices to go down without
+crashing the main program. In this way, elastic events (slices going up and
+down) can be caught or polled for and reacted to within an elastic_handler
+function without causing the whole workload to restart. This involves
+creating a new mesh with the available slices, reinitializing variables,
+and recompiling functions in addition to restoring from a host offloaded
+snapshot of the state.
+
+The purpose of this training loop is to serve as an example for how to
+support elasticity and is actively in development. As such, there are some
+performance optimizations that have yet to be added as well as some features
+not supported.
+
+Current limitations:
+- The host offloaded snapshot is currently blocking
+- Elastic event handling during async checkpointing is not tested
+- Elastic event handling during profiling is not tested
+- Elastic manager configuration values are hard coded
+- Debug logging statements for elasticity are hard coded as enabled
+
+See https://github.com/AI-Hypercomputer/pathways-utils/tree/main/pathwaysutils/elastic
+for more details about the elastic manager.
+"""
+from collections.abc import Sequence
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from absl import app
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+from flax.linen import partitioning as nn_partitioning
+import jax
+from ml_goodput_measurement import monitoring
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+import tensorflow as tf
+
+from MaxText import checkpointing
+from MaxText import max_utils
+from MaxText import maxtext_utils
+from MaxText import max_logging
+from MaxText import profiler
+from MaxText import pyconfig
+from MaxText.gcp_workload_monitor import GCPWorkloadMonitor
+from MaxText.input_pipeline.input_pipeline_interface import create_data_iterator
+from MaxText.metric_logger import MetricLogger
+from MaxText.train import check_example_batch
+from MaxText.train import create_goodput_recorder
+from MaxText.train import get_first_step
+from MaxText.train import load_next_batch
+from MaxText.train import record_goodput
+from MaxText.train import record_scalar_metrics
+from MaxText.train import save_checkpoint
+from MaxText.train import setup_mesh_and_model
+from MaxText.train import setup_train_loop
+from MaxText.train import train_step
+from MaxText.train import validate_train_config
+from MaxText.vertex_tensorboard import VertexTensorboardManager
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    elastic_manager,
+    checkpoint_manager,
+):
+  """Reconfigures the workload onto the currently available slices.
+
+  This is called by the elastic manager's maybe_reshard_up/down
+  functions and is responsible for creating a new mesh,
+  reinitializing the state and any objects that depend on the mesh.
+
+  It returns all of the reinitialized objects.
+
+  maybe_reshard_up/down take this function and its arguments and if
+  there is an elastic event, those functions will call this function
+  and return its returns.
+  """
+  if checkpoint_manager is not None:
path: MaxText/elastic_train.py
user: login: richjames0
body: A comment on why this is necessary would be nice
created_at: 2025-04-09T16:52:39Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,493 @@
+# Copyright 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.
+
+"""Training loop for elastic training model.
+
+Elastic training via Pathways on Cloud allows for slices to go down without
+crashing the main program. In this way, elastic events (slices going up and
+down) can be caught or polled for and reacted to within an elastic_handler
+function without causing the whole workload to restart. This involves
+creating a new mesh with the available slices, reinitializing variables,
+and recompiling functions in addition to restoring from a host offloaded
+snapshot of the state.
+
+The purpose of this training loop is to serve as an example for how to
+support elasticity and is actively in development. As such, there are some
+performance optimizations that have yet to be added as well as some features
+not supported.
+
+Current limitations:
+- The host offloaded snapshot is currently blocking
+- Elastic event handling during async checkpointing is not tested
+- Elastic event handling during profiling is not tested
+- Elastic manager configuration values are hard coded
+- Debug logging statements for elasticity are hard coded as enabled
+
+See https://github.com/AI-Hypercomputer/pathways-utils/tree/main/pathwaysutils/elastic
+for more details about the elastic manager.
+"""
+from collections.abc import Sequence
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from absl import app
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+from flax.linen import partitioning as nn_partitioning
+import jax
+from ml_goodput_measurement import monitoring
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+import tensorflow as tf
+
+from MaxText import checkpointing
+from MaxText import max_utils
+from MaxText import maxtext_utils
+from MaxText import max_logging
+from MaxText import profiler
+from MaxText import pyconfig
+from MaxText.gcp_workload_monitor import GCPWorkloadMonitor
+from MaxText.input_pipeline.input_pipeline_interface import create_data_iterator
+from MaxText.metric_logger import MetricLogger
+from MaxText.train import check_example_batch
+from MaxText.train import create_goodput_recorder
+from MaxText.train import get_first_step
+from MaxText.train import load_next_batch
+from MaxText.train import record_goodput
+from MaxText.train import record_scalar_metrics
+from MaxText.train import save_checkpoint
+from MaxText.train import setup_mesh_and_model
+from MaxText.train import setup_train_loop
+from MaxText.train import train_step
+from MaxText.train import validate_train_config
+from MaxText.vertex_tensorboard import VertexTensorboardManager
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    elastic_manager,
+    checkpoint_manager,
+):
+  """Reconfigures the workload onto the currently available slices.
+
+  This is called by the elastic manager's maybe_reshard_up/down
+  functions and is responsible for creating a new mesh,
+  reinitializing the state and any objects that depend on the mesh.
+
+  It returns all of the reinitialized objects.
+
+  maybe_reshard_up/down take this function and its arguments and if
+  there is an elastic event, those functions will call this function
+  and return its returns.
+  """
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(elastic_manager.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, elastic_manager.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = elastic_manager.get_resharded_snapshot(mesh)
+
+      # We do not want to restore from the previous checkpoint but instead
+      # restore from the host offloaded snapshot.
+      if checkpoint_manager is not None:
+        latest_step = checkpoint_manager.latest_step()
+
+        # If we checkpointed after the latest snapshot, the checkpoint manager
+        # will try to take another checkpoint and fail because it already
+        # exists. Therefore, we delete the checkpoint and let the checkpoint
+        # manager re-take the checkpoint.
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
+          checkpoint_manager.delete(latest_step)
+
+      state, _, state_mesh_shardings, data_iterator = max_utils.setup_training_state(
+          model,
+          data_iterator,
+          tx,
+          config,
+          jax.random.fold_in(init_rng, step),
+          mesh,
+          checkpoint_manager=None,
+      )
+
+      state = state.replace(**snapshot)
+      state = state.replace(step=state.step.at[None].set(step))
+
+      (
+          functional_train,
+          in_shard_train,
+          out_shard_train,
+          static_argnums_train,
+          donate_argnums_train,
+      ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+      p_train_step = jax.jit(
+          functional_train,
+          in_shardings=in_shard_train,
+          out_shardings=out_shard_train,
+          static_argnums=static_argnums_train,
+          donate_argnums=donate_argnums_train,
+      )
+
+      example_batch = None
+      metric_logger = MetricLogger(writer, config)
+
+      jax.block_until_ready(state)
+
+  return (
+      config,
+      step,
+      state,
+      mesh,
+      checkpoint_manager,
+      data_iterator,
+      p_train_step,
+      example_batch,
+      learning_rate_schedule,
+      metric_logger,
+      writer,
+  )
+
+
+def train_loop(config, elastic_manager, state=None):
+  """Main Training loop.
+  Args:
+    config:
+    state:
+    ckpt_path:
+  Returns:
+  """
+  # Create a GoodputRecorder to log information
+  recorder = create_goodput_recorder(config)
+  record_goodput(recorder, config, recorder.record_job_start_time if recorder else None)
+
+  (
+      init_rng,
+      writer,
+      checkpoint_manager,
+      state_mesh_shardings,
+      model,
+      mesh,
+      learning_rate_schedule,
+      data_iterator,
+      _,
+      state,
+  ) = setup_train_loop(config)
+
+  # pylint: disable=line-too-long
+  (
+      functional_train,
+      in_shard_train,
+      out_shard_train,
+      static_argnums_train,
+      donate_argnums_train,
+  ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+  num_model_parameters = max_utils.calculate_num_params_from_pytree(state.params)
+  max_logging.log(f"number parameters: {num_model_parameters/1e9:.3f} billion")
+  per_device_tflops, _, _ = maxtext_utils.calculate_tflops_training_per_device(config)
+  per_device_tokens = maxtext_utils.calculate_tokens_training_per_device(config)
+
+  # Write train config params, num model params, and XLA flags to tensorboard
+  max_utils.add_text_to_summary_writer("num_model_parameters", str(num_model_parameters), writer)
+  max_utils.add_text_to_summary_writer("libtpu_init_args", os.environ["LIBTPU_INIT_ARGS"], writer)
+  max_utils.add_config_to_summary_writer(config, writer)
+
+  p_train_step = jax.jit(
+      functional_train,
+      in_shardings=in_shard_train,
+      out_shardings=out_shard_train,
+      static_argnums=static_argnums_train,
+      donate_argnums=donate_argnums_train,
+  )
+  running_gcs_metrics = [] if config.gcs_metrics else None
+
+  start_step = get_first_step(state)  # this is the start_step for training
+  prof = profiler.Profiler(config, offset_step=start_step)
+  first_profiling_step = prof.start_initial_profile_step
+  if config.profiler != "" and first_profiling_step >= config.steps:
+    raise ValueError("Profiling requested but initial profiling step set past training final step")
+  last_profiling_step = prof.finished_initial_profile_step
+
+  example_batch = None
+  last_step_completion = datetime.datetime.now()
+
+  performance_metric_queue = None
+  if config.report_heartbeat_metric_for_gcp_monitoring or config.report_performance_metric_for_gcp_monitoring:
+    gcp_workload_monitor = GCPWorkloadMonitor(config.run_name)
+    if config.report_heartbeat_metric_for_gcp_monitoring:
+      gcp_workload_monitor.start_heartbeat_reporting_thread(config.heartbeat_reporting_interval_in_seconds)
+    if config.report_performance_metric_for_gcp_monitoring:
+      performance_metric_queue = queue.Queue()
+      gcp_workload_monitor.start_performance_reporting_thread(performance_metric_queue)
+
+  metric_logger = MetricLogger(writer, config)
+  step = start_step
+
+  elastic_manager.maybe_snapshot(
+      step,
+      snapshot={
+          "params": state.params,
+          "opt_state": state.opt_state,
+      },
+      force=True,
+      block=True,
+  )
+
+  # Using while loop instead of a for loop because with elasticity
+  # the step is restored back to the latest snapshot when a slice is lost
+  while step < config.steps:
+    try:
+      if step == first_profiling_step or prof.should_activate_periodic_profile(step):
+        optional_postfix = f"step_{step}" if config.profile_periodically_period > 0 else ""
+        prof.activate(blocking_object=state, optional_postfix=optional_postfix)
+
+      max_logging.log(f"{step=} {elastic_manager.elastic_down_event_count=} {elastic_manager.good_slice_count=}")
+      with mesh, nn_partitioning.axis_rules(config.logical_axis_rules), jax.default_device(elastic_manager.default_device):
+        with jax.profiler.StepTraceAnnotation("train", step_num=step):
+          record_goodput(recorder, config, recorder.record_data_loading_start_time if recorder else None)
+          try:
+            example_batch = load_next_batch(data_iterator, example_batch, config)
+          except Exception as e:  # pylint: disable=broad-except
+            max_logging.log(f"load_next_batch failed, you may have run out of data. Error message: {e}")
+            break
+          record_goodput(recorder, config, recorder.record_data_loading_end_time if recorder else None)
+          check_example_batch(config, example_batch=example_batch)
+          # pylint: disable=not-callable
+          nextrng = jax.jit(jax.random.fold_in)(init_rng, step)
+          record_goodput(recorder, config, recorder.record_step_start_time if recorder else None, step)
+          state, metrics = p_train_step(state, example_batch, nextrng)
+
+        step_time_delta = datetime.datetime.now() - last_step_completion
+        last_step_completion = datetime.datetime.now()
+        record_scalar_metrics(metrics, step_time_delta, per_device_tflops, learning_rate_schedule(step), per_device_tokens)
+        if performance_metric_queue:
+          performance_metric_queue.put(step_time_delta.total_seconds())
+
+        if checkpoint_manager is not None:
+          state_to_save = state
+          if save_checkpoint(checkpoint_manager, int(step), state_to_save, config.dataset_type, data_iterator, config):
+            checkpointing.print_save_message(step, config.async_checkpointing)
+
+          # Upon preemption, exit when and only when all ongoing saves are complete.
+          if checkpoint_manager.reached_preemption(step):
+            checkpoint_manager.wait_until_finished()
+            sys.exit()
path: MaxText/elastic_train.py
user: login: richjames0
body: I don't love seeing exit calls mid program - they're kinda like a goto and can be confusing when debugging, so if it's possible to not have this without radically changing control flow that would be nice. But if that would be a nightmare then ignore this
created_at: 2025-04-09T16:58:58Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,493 @@
+# Copyright 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.
+
+"""Training loop for elastic training model.
+
+Elastic training via Pathways on Cloud allows for slices to go down without
+crashing the main program. In this way, elastic events (slices going up and
+down) can be caught or polled for and reacted to within an elastic_handler
+function without causing the whole workload to restart. This involves
+creating a new mesh with the available slices, reinitializing variables,
+and recompiling functions in addition to restoring from a host offloaded
+snapshot of the state.
+
+The purpose of this training loop is to serve as an example for how to
+support elasticity and is actively in development. As such, there are some
+performance optimizations that have yet to be added as well as some features
+not supported.
+
+Current limitations:
+- The host offloaded snapshot is currently blocking
+- Elastic event handling during async checkpointing is not tested
+- Elastic event handling during profiling is not tested
+- Elastic manager configuration values are hard coded
+- Debug logging statements for elasticity are hard coded as enabled
+
+See https://github.com/AI-Hypercomputer/pathways-utils/tree/main/pathwaysutils/elastic
+for more details about the elastic manager.
+"""
+from collections.abc import Sequence
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from absl import app
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+from flax.linen import partitioning as nn_partitioning
+import jax
+from ml_goodput_measurement import monitoring
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+import tensorflow as tf
+
+from MaxText import checkpointing
+from MaxText import max_utils
+from MaxText import maxtext_utils
+from MaxText import max_logging
+from MaxText import profiler
+from MaxText import pyconfig
+from MaxText.gcp_workload_monitor import GCPWorkloadMonitor
+from MaxText.input_pipeline.input_pipeline_interface import create_data_iterator
+from MaxText.metric_logger import MetricLogger
+from MaxText.train import check_example_batch
+from MaxText.train import create_goodput_recorder
+from MaxText.train import get_first_step
+from MaxText.train import load_next_batch
+from MaxText.train import record_goodput
+from MaxText.train import record_scalar_metrics
+from MaxText.train import save_checkpoint
+from MaxText.train import setup_mesh_and_model
+from MaxText.train import setup_train_loop
+from MaxText.train import train_step
+from MaxText.train import validate_train_config
+from MaxText.vertex_tensorboard import VertexTensorboardManager
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    elastic_manager,
+    checkpoint_manager,
+):
+  """Reconfigures the workload onto the currently available slices.
+
+  This is called by the elastic manager's maybe_reshard_up/down
+  functions and is responsible for creating a new mesh,
+  reinitializing the state and any objects that depend on the mesh.
+
+  It returns all of the reinitialized objects.
+
+  maybe_reshard_up/down take this function and its arguments and if
+  there is an elastic event, those functions will call this function
+  and return its returns.
+  """
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(elastic_manager.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, elastic_manager.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = elastic_manager.get_resharded_snapshot(mesh)
+
+      # We do not want to restore from the previous checkpoint but instead
+      # restore from the host offloaded snapshot.
+      if checkpoint_manager is not None:
+        latest_step = checkpoint_manager.latest_step()
+
+        # If we checkpointed after the latest snapshot, the checkpoint manager
+        # will try to take another checkpoint and fail because it already
+        # exists. Therefore, we delete the checkpoint and let the checkpoint
+        # manager re-take the checkpoint.
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
+          checkpoint_manager.delete(latest_step)
+
+      state, _, state_mesh_shardings, data_iterator = max_utils.setup_training_state(
+          model,
+          data_iterator,
+          tx,
+          config,
+          jax.random.fold_in(init_rng, step),
+          mesh,
+          checkpoint_manager=None,
+      )
+
+      state = state.replace(**snapshot)
+      state = state.replace(step=state.step.at[None].set(step))
+
+      (
+          functional_train,
+          in_shard_train,
+          out_shard_train,
+          static_argnums_train,
+          donate_argnums_train,
+      ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+      p_train_step = jax.jit(
+          functional_train,
+          in_shardings=in_shard_train,
+          out_shardings=out_shard_train,
+          static_argnums=static_argnums_train,
+          donate_argnums=donate_argnums_train,
+      )
+
+      example_batch = None
+      metric_logger = MetricLogger(writer, config)
+
+      jax.block_until_ready(state)
+
+  return (
+      config,
+      step,
+      state,
+      mesh,
+      checkpoint_manager,
+      data_iterator,
+      p_train_step,
+      example_batch,
+      learning_rate_schedule,
+      metric_logger,
+      writer,
+  )
+
+
+def train_loop(config, elastic_manager, state=None):
path: MaxText/elastic_train.py
user: login: richjames0
body: This is a whole lot of function and is accordingly a bit hard to follow IMO. TBF you're just following the (probably not well thought out) established idioms of train.py but I think it would be nice to break this down for both readability and testability.
created_at: 2025-04-09T17:00:46Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,493 @@
+# Copyright 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.
+
+"""Training loop for elastic training model.
+
+Elastic training via Pathways on Cloud allows for slices to go down without
+crashing the main program. In this way, elastic events (slices going up and
+down) can be caught or polled for and reacted to within an elastic_handler
+function without causing the whole workload to restart. This involves
+creating a new mesh with the available slices, reinitializing variables,
+and recompiling functions in addition to restoring from a host offloaded
+snapshot of the state.
+
+The purpose of this training loop is to serve as an example for how to
+support elasticity and is actively in development. As such, there are some
+performance optimizations that have yet to be added as well as some features
+not supported.
+
+Current limitations:
+- The host offloaded snapshot is currently blocking
+- Elastic event handling during async checkpointing is not tested
+- Elastic event handling during profiling is not tested
+- Elastic manager configuration values are hard coded
+- Debug logging statements for elasticity are hard coded as enabled
+
+See https://github.com/AI-Hypercomputer/pathways-utils/tree/main/pathwaysutils/elastic
+for more details about the elastic manager.
+"""
+from collections.abc import Sequence
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from absl import app
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+from flax.linen import partitioning as nn_partitioning
+import jax
+from ml_goodput_measurement import monitoring
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+import tensorflow as tf
+
+from MaxText import checkpointing
+from MaxText import max_utils
+from MaxText import maxtext_utils
+from MaxText import max_logging
+from MaxText import profiler
+from MaxText import pyconfig
+from MaxText.gcp_workload_monitor import GCPWorkloadMonitor
+from MaxText.input_pipeline.input_pipeline_interface import create_data_iterator
+from MaxText.metric_logger import MetricLogger
+from MaxText.train import check_example_batch
+from MaxText.train import create_goodput_recorder
+from MaxText.train import get_first_step
+from MaxText.train import load_next_batch
+from MaxText.train import record_goodput
+from MaxText.train import record_scalar_metrics
+from MaxText.train import save_checkpoint
+from MaxText.train import setup_mesh_and_model
+from MaxText.train import setup_train_loop
+from MaxText.train import train_step
+from MaxText.train import validate_train_config
+from MaxText.vertex_tensorboard import VertexTensorboardManager
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    elastic_manager,
+    checkpoint_manager,
+):
+  """Reconfigures the workload onto the currently available slices.
+
+  This is called by the elastic manager's maybe_reshard_up/down
+  functions and is responsible for creating a new mesh,
+  reinitializing the state and any objects that depend on the mesh.
+
+  It returns all of the reinitialized objects.
+
+  maybe_reshard_up/down take this function and its arguments and if
+  there is an elastic event, those functions will call this function
+  and return its returns.
+  """
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(elastic_manager.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, elastic_manager.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = elastic_manager.get_resharded_snapshot(mesh)
+
+      # We do not want to restore from the previous checkpoint but instead
+      # restore from the host offloaded snapshot.
+      if checkpoint_manager is not None:
+        latest_step = checkpoint_manager.latest_step()
+
+        # If we checkpointed after the latest snapshot, the checkpoint manager
+        # will try to take another checkpoint and fail because it already
+        # exists. Therefore, we delete the checkpoint and let the checkpoint
+        # manager re-take the checkpoint.
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
+          checkpoint_manager.delete(latest_step)
+
+      state, _, state_mesh_shardings, data_iterator = max_utils.setup_training_state(
+          model,
+          data_iterator,
+          tx,
+          config,
+          jax.random.fold_in(init_rng, step),
+          mesh,
+          checkpoint_manager=None,
+      )
+
+      state = state.replace(**snapshot)
+      state = state.replace(step=state.step.at[None].set(step))
+
+      (
+          functional_train,
+          in_shard_train,
+          out_shard_train,
+          static_argnums_train,
+          donate_argnums_train,
+      ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+      p_train_step = jax.jit(
+          functional_train,
+          in_shardings=in_shard_train,
+          out_shardings=out_shard_train,
+          static_argnums=static_argnums_train,
+          donate_argnums=donate_argnums_train,
+      )
+
+      example_batch = None
+      metric_logger = MetricLogger(writer, config)
+
+      jax.block_until_ready(state)
+
+  return (
+      config,
+      step,
+      state,
+      mesh,
+      checkpoint_manager,
+      data_iterator,
+      p_train_step,
+      example_batch,
+      learning_rate_schedule,
+      metric_logger,
+      writer,
+  )
+
+
+def train_loop(config, elastic_manager, state=None):
+  """Main Training loop.
+  Args:
+    config:
+    state:
+    ckpt_path:
+  Returns:
+  """
+  # Create a GoodputRecorder to log information
+  recorder = create_goodput_recorder(config)
+  record_goodput(recorder, config, recorder.record_job_start_time if recorder else None)
+
+  (
+      init_rng,
+      writer,
+      checkpoint_manager,
+      state_mesh_shardings,
+      model,
+      mesh,
+      learning_rate_schedule,
+      data_iterator,
+      _,
+      state,
+  ) = setup_train_loop(config)
+
+  # pylint: disable=line-too-long
+  (
+      functional_train,
+      in_shard_train,
+      out_shard_train,
+      static_argnums_train,
+      donate_argnums_train,
+  ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+  num_model_parameters = max_utils.calculate_num_params_from_pytree(state.params)
+  max_logging.log(f"number parameters: {num_model_parameters/1e9:.3f} billion")
+  per_device_tflops, _, _ = maxtext_utils.calculate_tflops_training_per_device(config)
+  per_device_tokens = maxtext_utils.calculate_tokens_training_per_device(config)
+
+  # Write train config params, num model params, and XLA flags to tensorboard
+  max_utils.add_text_to_summary_writer("num_model_parameters", str(num_model_parameters), writer)
+  max_utils.add_text_to_summary_writer("libtpu_init_args", os.environ["LIBTPU_INIT_ARGS"], writer)
+  max_utils.add_config_to_summary_writer(config, writer)
+
+  p_train_step = jax.jit(
+      functional_train,
+      in_shardings=in_shard_train,
+      out_shardings=out_shard_train,
+      static_argnums=static_argnums_train,
+      donate_argnums=donate_argnums_train,
+  )
+  running_gcs_metrics = [] if config.gcs_metrics else None
+
+  start_step = get_first_step(state)  # this is the start_step for training
+  prof = profiler.Profiler(config, offset_step=start_step)
+  first_profiling_step = prof.start_initial_profile_step
+  if config.profiler != "" and first_profiling_step >= config.steps:
+    raise ValueError("Profiling requested but initial profiling step set past training final step")
+  last_profiling_step = prof.finished_initial_profile_step
+
+  example_batch = None
+  last_step_completion = datetime.datetime.now()
+
+  performance_metric_queue = None
+  if config.report_heartbeat_metric_for_gcp_monitoring or config.report_performance_metric_for_gcp_monitoring:
+    gcp_workload_monitor = GCPWorkloadMonitor(config.run_name)
+    if config.report_heartbeat_metric_for_gcp_monitoring:
+      gcp_workload_monitor.start_heartbeat_reporting_thread(config.heartbeat_reporting_interval_in_seconds)
+    if config.report_performance_metric_for_gcp_monitoring:
+      performance_metric_queue = queue.Queue()
+      gcp_workload_monitor.start_performance_reporting_thread(performance_metric_queue)
+
+  metric_logger = MetricLogger(writer, config)
+  step = start_step
+
+  elastic_manager.maybe_snapshot(
+      step,
+      snapshot={
+          "params": state.params,
+          "opt_state": state.opt_state,
+      },
+      force=True,
+      block=True,
+  )
+
+  # Using while loop instead of a for loop because with elasticity
+  # the step is restored back to the latest snapshot when a slice is lost
+  while step < config.steps:
path: MaxText/elastic_train.py
user: login: richjames0
body: My comment about breaking out functions applies even more so here with all these nested blocks. If this degree of complexity of control flow is required (and I can totally believe it is) it would be a lot easier to understand imo if the actual code blocks (vs. control flow) were broken out into functions
created_at: 2025-04-09T17:01:53Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,493 @@
+# Copyright 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.
+
+"""Training loop for elastic training model.
+
+Elastic training via Pathways on Cloud allows for slices to go down without
+crashing the main program. In this way, elastic events (slices going up and
+down) can be caught or polled for and reacted to within an elastic_handler
+function without causing the whole workload to restart. This involves
+creating a new mesh with the available slices, reinitializing variables,
+and recompiling functions in addition to restoring from a host offloaded
+snapshot of the state.
+
+The purpose of this training loop is to serve as an example for how to
+support elasticity and is actively in development. As such, there are some
+performance optimizations that have yet to be added as well as some features
+not supported.
+
+Current limitations:
+- The host offloaded snapshot is currently blocking
+- Elastic event handling during async checkpointing is not tested
+- Elastic event handling during profiling is not tested
+- Elastic manager configuration values are hard coded
+- Debug logging statements for elasticity are hard coded as enabled
+
+See https://github.com/AI-Hypercomputer/pathways-utils/tree/main/pathwaysutils/elastic
+for more details about the elastic manager.
+"""
+from collections.abc import Sequence
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from absl import app
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+from flax.linen import partitioning as nn_partitioning
+import jax
+from ml_goodput_measurement import monitoring
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+import tensorflow as tf
+
+from MaxText import checkpointing
+from MaxText import max_utils
+from MaxText import maxtext_utils
+from MaxText import max_logging
+from MaxText import profiler
+from MaxText import pyconfig
+from MaxText.gcp_workload_monitor import GCPWorkloadMonitor
+from MaxText.input_pipeline.input_pipeline_interface import create_data_iterator
+from MaxText.metric_logger import MetricLogger
+from MaxText.train import check_example_batch
+from MaxText.train import create_goodput_recorder
+from MaxText.train import get_first_step
+from MaxText.train import load_next_batch
+from MaxText.train import record_goodput
+from MaxText.train import record_scalar_metrics
+from MaxText.train import save_checkpoint
+from MaxText.train import setup_mesh_and_model
+from MaxText.train import setup_train_loop
+from MaxText.train import train_step
+from MaxText.train import validate_train_config
+from MaxText.vertex_tensorboard import VertexTensorboardManager
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    elastic_manager,
+    checkpoint_manager,
+):
+  """Reconfigures the workload onto the currently available slices.
+
+  This is called by the elastic manager's maybe_reshard_up/down
+  functions and is responsible for creating a new mesh,
+  reinitializing the state and any objects that depend on the mesh.
+
+  It returns all of the reinitialized objects.
+
+  maybe_reshard_up/down take this function and its arguments and if
+  there is an elastic event, those functions will call this function
+  and return its returns.
+  """
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(elastic_manager.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, elastic_manager.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = elastic_manager.get_resharded_snapshot(mesh)
+
+      # We do not want to restore from the previous checkpoint but instead
+      # restore from the host offloaded snapshot.
+      if checkpoint_manager is not None:
+        latest_step = checkpoint_manager.latest_step()
+
+        # If we checkpointed after the latest snapshot, the checkpoint manager
+        # will try to take another checkpoint and fail because it already
+        # exists. Therefore, we delete the checkpoint and let the checkpoint
+        # manager re-take the checkpoint.
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
+          checkpoint_manager.delete(latest_step)
+
+      state, _, state_mesh_shardings, data_iterator = max_utils.setup_training_state(
+          model,
+          data_iterator,
+          tx,
+          config,
+          jax.random.fold_in(init_rng, step),
+          mesh,
+          checkpoint_manager=None,
+      )
+
+      state = state.replace(**snapshot)
+      state = state.replace(step=state.step.at[None].set(step))
+
+      (
+          functional_train,
+          in_shard_train,
+          out_shard_train,
+          static_argnums_train,
+          donate_argnums_train,
+      ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+      p_train_step = jax.jit(
+          functional_train,
+          in_shardings=in_shard_train,
+          out_shardings=out_shard_train,
+          static_argnums=static_argnums_train,
+          donate_argnums=donate_argnums_train,
+      )
+
+      example_batch = None
+      metric_logger = MetricLogger(writer, config)
+
+      jax.block_until_ready(state)
+
+  return (
+      config,
+      step,
+      state,
+      mesh,
+      checkpoint_manager,
+      data_iterator,
+      p_train_step,
+      example_batch,
+      learning_rate_schedule,
+      metric_logger,
+      writer,
+  )
+
+
+def train_loop(config, elastic_manager, state=None):
+  """Main Training loop.
+  Args:
+    config:
+    state:
+    ckpt_path:
+  Returns:
+  """
+  # Create a GoodputRecorder to log information
+  recorder = create_goodput_recorder(config)
+  record_goodput(recorder, config, recorder.record_job_start_time if recorder else None)
+
+  (
+      init_rng,
+      writer,
+      checkpoint_manager,
+      state_mesh_shardings,
+      model,
+      mesh,
+      learning_rate_schedule,
+      data_iterator,
+      _,
+      state,
+  ) = setup_train_loop(config)
+
+  # pylint: disable=line-too-long
+  (
+      functional_train,
+      in_shard_train,
+      out_shard_train,
+      static_argnums_train,
+      donate_argnums_train,
+  ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+  num_model_parameters = max_utils.calculate_num_params_from_pytree(state.params)
+  max_logging.log(f"number parameters: {num_model_parameters/1e9:.3f} billion")
+  per_device_tflops, _, _ = maxtext_utils.calculate_tflops_training_per_device(config)
+  per_device_tokens = maxtext_utils.calculate_tokens_training_per_device(config)
+
+  # Write train config params, num model params, and XLA flags to tensorboard
+  max_utils.add_text_to_summary_writer("num_model_parameters", str(num_model_parameters), writer)
+  max_utils.add_text_to_summary_writer("libtpu_init_args", os.environ["LIBTPU_INIT_ARGS"], writer)
+  max_utils.add_config_to_summary_writer(config, writer)
+
+  p_train_step = jax.jit(
+      functional_train,
+      in_shardings=in_shard_train,
+      out_shardings=out_shard_train,
+      static_argnums=static_argnums_train,
+      donate_argnums=donate_argnums_train,
+  )
+  running_gcs_metrics = [] if config.gcs_metrics else None
+
+  start_step = get_first_step(state)  # this is the start_step for training
+  prof = profiler.Profiler(config, offset_step=start_step)
+  first_profiling_step = prof.start_initial_profile_step
+  if config.profiler != "" and first_profiling_step >= config.steps:
+    raise ValueError("Profiling requested but initial profiling step set past training final step")
+  last_profiling_step = prof.finished_initial_profile_step
+
+  example_batch = None
+  last_step_completion = datetime.datetime.now()
+
+  performance_metric_queue = None
+  if config.report_heartbeat_metric_for_gcp_monitoring or config.report_performance_metric_for_gcp_monitoring:
+    gcp_workload_monitor = GCPWorkloadMonitor(config.run_name)
+    if config.report_heartbeat_metric_for_gcp_monitoring:
+      gcp_workload_monitor.start_heartbeat_reporting_thread(config.heartbeat_reporting_interval_in_seconds)
+    if config.report_performance_metric_for_gcp_monitoring:
+      performance_metric_queue = queue.Queue()
+      gcp_workload_monitor.start_performance_reporting_thread(performance_metric_queue)
+
+  metric_logger = MetricLogger(writer, config)
+  step = start_step
+
+  elastic_manager.maybe_snapshot(
+      step,
+      snapshot={
+          "params": state.params,
+          "opt_state": state.opt_state,
+      },
+      force=True,
+      block=True,
+  )
+
+  # Using while loop instead of a for loop because with elasticity
+  # the step is restored back to the latest snapshot when a slice is lost
+  while step < config.steps:
+    try:
+      if step == first_profiling_step or prof.should_activate_periodic_profile(step):
+        optional_postfix = f"step_{step}" if config.profile_periodically_period > 0 else ""
+        prof.activate(blocking_object=state, optional_postfix=optional_postfix)
+
+      max_logging.log(f"{step=} {elastic_manager.elastic_down_event_count=} {elastic_manager.good_slice_count=}")
+      with mesh, nn_partitioning.axis_rules(config.logical_axis_rules), jax.default_device(elastic_manager.default_device):
+        with jax.profiler.StepTraceAnnotation("train", step_num=step):
+          record_goodput(recorder, config, recorder.record_data_loading_start_time if recorder else None)
+          try:
+            example_batch = load_next_batch(data_iterator, example_batch, config)
+          except Exception as e:  # pylint: disable=broad-except
+            max_logging.log(f"load_next_batch failed, you may have run out of data. Error message: {e}")
+            break
+          record_goodput(recorder, config, recorder.record_data_loading_end_time if recorder else None)
+          check_example_batch(config, example_batch=example_batch)
+          # pylint: disable=not-callable
+          nextrng = jax.jit(jax.random.fold_in)(init_rng, step)
+          record_goodput(recorder, config, recorder.record_step_start_time if recorder else None, step)
+          state, metrics = p_train_step(state, example_batch, nextrng)
+
+        step_time_delta = datetime.datetime.now() - last_step_completion
+        last_step_completion = datetime.datetime.now()
+        record_scalar_metrics(metrics, step_time_delta, per_device_tflops, learning_rate_schedule(step), per_device_tokens)
+        if performance_metric_queue:
+          performance_metric_queue.put(step_time_delta.total_seconds())
+
+        if checkpoint_manager is not None:
+          state_to_save = state
+          if save_checkpoint(checkpoint_manager, int(step), state_to_save, config.dataset_type, data_iterator, config):
+            checkpointing.print_save_message(step, config.async_checkpointing)
+
+          # Upon preemption, exit when and only when all ongoing saves are complete.
+          if checkpoint_manager.reached_preemption(step):
+            checkpoint_manager.wait_until_finished()
+            sys.exit()
+
+        metric_logger.write_metrics(running_gcs_metrics, metrics, step)
+
+        if step == last_profiling_step or prof.should_deactivate_periodic_profile(step):
+          prof.deactivate(blocking_object=state)
+
+      elastic_manager.maybe_snapshot(
+          step=step,
+          snapshot={
+              "params": state.params,
+              "opt_state": state.opt_state,
+          },
+          block=True,
+      )
+
+      ret = elastic_manager.maybe_reshard_up(
+          step=step,
+          snapshot={
+              "params": state.params,
+              "opt_state": state.opt_state,
+          },
+          elastic_handler=elastic_handler,
+          handler_kwargs={
+              "config": config,
+              "elastic_manager": elastic_manager,
+              "checkpoint_manager": checkpoint_manager,
+          },
+      )
+      if ret is not None:
+        (
+            config,
+            step,
+            state,
+            mesh,
+            checkpoint_manager,
+            data_iterator,
+            p_train_step,
+            example_batch,
+            learning_rate_schedule,
+            metric_logger,
+            writer,
+        ) = ret
+
+      if step == start_step:
+        max_utils.print_mem_stats("After params initialized")
+
+      step += 1
+
+    except jax.errors.JaxRuntimeError as error:
+      ret = elastic_manager.maybe_reshard_down(
+          error=error,
+          elastic_handler=elastic_handler,
+          handler_kwargs={
+              "config": config,
+              "elastic_manager": elastic_manager,
+              "checkpoint_manager": checkpoint_manager,
+          },
+      )
+      if ret is not None:
+        (
+            config,
+            step,
+            state,
+            mesh,
+            checkpoint_manager,
+            data_iterator,
+            p_train_step,
+            example_batch,
+            learning_rate_schedule,
+            metric_logger,
+            writer,
+        ) = ret
+
+  if checkpoint_manager is not None:
+    if int(state.step) - 1 % config.checkpoint_period != 0:
+      try:
+        state_to_save = state
+        if save_checkpoint(
+            checkpoint_manager,
+            int(state_to_save) - 1,
+            state_to_save,
+            config.dataset_type,
+            data_iterator,
+            config,
+            force=True,
+        ):
+          checkpointing.print_save_message(int(state_to_save) - 1, config.async_checkpointing)
+      except Exception:  # pylint: disable=broad-except
+        max_logging.log(f"Checkpoint is already saved for step {int(state.step)-1}.")
+
+    checkpoint_manager.wait_until_finished()
+  metric_logger.write_metrics(running_gcs_metrics, metrics, config.steps - 1)  # final step metrics
+  max_utils.close_summary_writer(writer)
+  record_goodput(recorder, config, recorder.record_job_end_time if recorder else None)
+  with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
+    # pytype: disable=attribute-error
+    compiled = p_train_step.lower(state, example_batch, nextrng).compile()
+    compiled_stats = compiled.memory_analysis()
+    if compiled_stats is not None:
+      max_logging.log(
+          f"Output size: {compiled_stats.output_size_in_bytes}, "
+          f"temp size: {compiled_stats.temp_size_in_bytes}, "
+          f"argument size: {compiled_stats.argument_size_in_bytes}, "
+          f"host temp size: {compiled_stats.host_temp_size_in_bytes}, in bytes."
+      )
+  return state
+
+
+def main(argv: Sequence[str]) -> None:
path: MaxText/elastic_train.py
user: login: richjames0
body: Not as significant as in train_step but I would also break this function up a bunch
created_at: 2025-04-09T17:02:28Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,493 @@
+# Copyright 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.
+
+"""Training loop for elastic training model.
+
+Elastic training via Pathways on Cloud allows for slices to go down without
+crashing the main program. In this way, elastic events (slices going up and
+down) can be caught or polled for and reacted to within an elastic_handler
+function without causing the whole workload to restart. This involves
+creating a new mesh with the available slices, reinitializing variables,
+and recompiling functions in addition to restoring from a host offloaded
+snapshot of the state.
+
+The purpose of this training loop is to serve as an example for how to
+support elasticity and is actively in development. As such, there are some
+performance optimizations that have yet to be added as well as some features
+not supported.
+
+Current limitations:
+- The host offloaded snapshot is currently blocking
+- Elastic event handling during async checkpointing is not tested
+- Elastic event handling during profiling is not tested
+- Elastic manager configuration values are hard coded
+- Debug logging statements for elasticity are hard coded as enabled
+
+See https://github.com/AI-Hypercomputer/pathways-utils/tree/main/pathwaysutils/elastic
+for more details about the elastic manager.
+"""
+from collections.abc import Sequence
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from absl import app
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+from flax.linen import partitioning as nn_partitioning
+import jax
+from ml_goodput_measurement import monitoring
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+import tensorflow as tf
+
+from MaxText import checkpointing
+from MaxText import max_utils
+from MaxText import maxtext_utils
+from MaxText import max_logging
+from MaxText import profiler
+from MaxText import pyconfig
+from MaxText.gcp_workload_monitor import GCPWorkloadMonitor
+from MaxText.input_pipeline.input_pipeline_interface import create_data_iterator
+from MaxText.metric_logger import MetricLogger
+from MaxText.train import check_example_batch
+from MaxText.train import create_goodput_recorder
+from MaxText.train import get_first_step
+from MaxText.train import load_next_batch
+from MaxText.train import record_goodput
+from MaxText.train import record_scalar_metrics
+from MaxText.train import save_checkpoint
+from MaxText.train import setup_mesh_and_model
+from MaxText.train import setup_train_loop
+from MaxText.train import train_step
+from MaxText.train import validate_train_config
+from MaxText.vertex_tensorboard import VertexTensorboardManager
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    elastic_manager,
+    checkpoint_manager,
+):
+  """Reconfigures the workload onto the currently available slices.
+
+  This is called by the elastic manager's maybe_reshard_up/down
+  functions and is responsible for creating a new mesh,
+  reinitializing the state and any objects that depend on the mesh.
+
+  It returns all of the reinitialized objects.
+
+  maybe_reshard_up/down take this function and its arguments and if
+  there is an elastic event, those functions will call this function
+  and return its returns.
+  """
+  if checkpoint_manager is not None:
path: MaxText/elastic_train.py
user: login: lukebaumann
body: Done
created_at: 2025-04-09T18:08:07Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,493 @@
+# Copyright 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.
+
+"""Training loop for elastic training model.
+
+Elastic training via Pathways on Cloud allows for slices to go down without
+crashing the main program. In this way, elastic events (slices going up and
+down) can be caught or polled for and reacted to within an elastic_handler
+function without causing the whole workload to restart. This involves
+creating a new mesh with the available slices, reinitializing variables,
+and recompiling functions in addition to restoring from a host offloaded
+snapshot of the state.
+
+The purpose of this training loop is to serve as an example for how to
+support elasticity and is actively in development. As such, there are some
+performance optimizations that have yet to be added as well as some features
+not supported.
+
+Current limitations:
+- The host offloaded snapshot is currently blocking
+- Elastic event handling during async checkpointing is not tested
+- Elastic event handling during profiling is not tested
+- Elastic manager configuration values are hard coded
+- Debug logging statements for elasticity are hard coded as enabled
+
+See https://github.com/AI-Hypercomputer/pathways-utils/tree/main/pathwaysutils/elastic
+for more details about the elastic manager.
+"""
+from collections.abc import Sequence
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from absl import app
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+from flax.linen import partitioning as nn_partitioning
+import jax
+from ml_goodput_measurement import monitoring
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+import tensorflow as tf
+
+from MaxText import checkpointing
+from MaxText import max_utils
+from MaxText import maxtext_utils
+from MaxText import max_logging
+from MaxText import profiler
+from MaxText import pyconfig
+from MaxText.gcp_workload_monitor import GCPWorkloadMonitor
+from MaxText.input_pipeline.input_pipeline_interface import create_data_iterator
+from MaxText.metric_logger import MetricLogger
+from MaxText.train import check_example_batch
+from MaxText.train import create_goodput_recorder
+from MaxText.train import get_first_step
+from MaxText.train import load_next_batch
+from MaxText.train import record_goodput
+from MaxText.train import record_scalar_metrics
+from MaxText.train import save_checkpoint
+from MaxText.train import setup_mesh_and_model
+from MaxText.train import setup_train_loop
+from MaxText.train import train_step
+from MaxText.train import validate_train_config
+from MaxText.vertex_tensorboard import VertexTensorboardManager
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    elastic_manager,
+    checkpoint_manager,
+):
+  """Reconfigures the workload onto the currently available slices.
+
+  This is called by the elastic manager's maybe_reshard_up/down
+  functions and is responsible for creating a new mesh,
+  reinitializing the state and any objects that depend on the mesh.
+
+  It returns all of the reinitialized objects.
+
+  maybe_reshard_up/down take this function and its arguments and if
+  there is an elastic event, those functions will call this function
+  and return its returns.
+  """
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(elastic_manager.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, elastic_manager.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = elastic_manager.get_resharded_snapshot(mesh)
+
+      # We do not want to restore from the previous checkpoint but instead
+      # restore from the host offloaded snapshot.
+      if checkpoint_manager is not None:
+        latest_step = checkpoint_manager.latest_step()
+
+        # If we checkpointed after the latest snapshot, the checkpoint manager
+        # will try to take another checkpoint and fail because it already
+        # exists. Therefore, we delete the checkpoint and let the checkpoint
+        # manager re-take the checkpoint.
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
+          checkpoint_manager.delete(latest_step)
+
+      state, _, state_mesh_shardings, data_iterator = max_utils.setup_training_state(
+          model,
+          data_iterator,
+          tx,
+          config,
+          jax.random.fold_in(init_rng, step),
+          mesh,
+          checkpoint_manager=None,
+      )
+
+      state = state.replace(**snapshot)
+      state = state.replace(step=state.step.at[None].set(step))
+
+      (
+          functional_train,
+          in_shard_train,
+          out_shard_train,
+          static_argnums_train,
+          donate_argnums_train,
+      ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+      p_train_step = jax.jit(
+          functional_train,
+          in_shardings=in_shard_train,
+          out_shardings=out_shard_train,
+          static_argnums=static_argnums_train,
+          donate_argnums=donate_argnums_train,
+      )
+
+      example_batch = None
+      metric_logger = MetricLogger(writer, config)
+
+      jax.block_until_ready(state)
+
+  return (
+      config,
+      step,
+      state,
+      mesh,
+      checkpoint_manager,
+      data_iterator,
+      p_train_step,
+      example_batch,
+      learning_rate_schedule,
+      metric_logger,
+      writer,
+  )
+
+
+def train_loop(config, elastic_manager, state=None):
path: MaxText/elastic_train.py
user: login: lukebaumann
body: I agree. I keep this as close to train.py as possible with the hope that as train.py is updated, the same updates can be brought here without much effort.

I think making the train_loop updated in elastic_train.py and not train.py is would lead divergence that would be more difficult to maintain.

Diffing the train_loop against train.py with ignore white space (because of the indentation from the try...except, addition of the default_device context manager, and moving the mesh context manager to cover most of the loop) shows how the removals are eval and dpo (because they are untested and unrelated cc: @gobbleturk) and the only additions are the maybe_snapshot, maybe_reshard_up, maybe_reshard_down which is about 70 LOC which is mostly comprised of unboxing the returned tuples.

Let me know if you want more from this PR.
created_at: 2025-04-09T18:18:13Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,493 @@
+# Copyright 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.
+
+"""Training loop for elastic training model.
+
+Elastic training via Pathways on Cloud allows for slices to go down without
+crashing the main program. In this way, elastic events (slices going up and
+down) can be caught or polled for and reacted to within an elastic_handler
+function without causing the whole workload to restart. This involves
+creating a new mesh with the available slices, reinitializing variables,
+and recompiling functions in addition to restoring from a host offloaded
+snapshot of the state.
+
+The purpose of this training loop is to serve as an example for how to
+support elasticity and is actively in development. As such, there are some
+performance optimizations that have yet to be added as well as some features
+not supported.
+
+Current limitations:
+- The host offloaded snapshot is currently blocking
+- Elastic event handling during async checkpointing is not tested
+- Elastic event handling during profiling is not tested
+- Elastic manager configuration values are hard coded
+- Debug logging statements for elasticity are hard coded as enabled
+
+See https://github.com/AI-Hypercomputer/pathways-utils/tree/main/pathwaysutils/elastic
+for more details about the elastic manager.
+"""
+from collections.abc import Sequence
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from absl import app
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+from flax.linen import partitioning as nn_partitioning
+import jax
+from ml_goodput_measurement import monitoring
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+import tensorflow as tf
+
+from MaxText import checkpointing
+from MaxText import max_utils
+from MaxText import maxtext_utils
+from MaxText import max_logging
+from MaxText import profiler
+from MaxText import pyconfig
+from MaxText.gcp_workload_monitor import GCPWorkloadMonitor
+from MaxText.input_pipeline.input_pipeline_interface import create_data_iterator
+from MaxText.metric_logger import MetricLogger
+from MaxText.train import check_example_batch
+from MaxText.train import create_goodput_recorder
+from MaxText.train import get_first_step
+from MaxText.train import load_next_batch
+from MaxText.train import record_goodput
+from MaxText.train import record_scalar_metrics
+from MaxText.train import save_checkpoint
+from MaxText.train import setup_mesh_and_model
+from MaxText.train import setup_train_loop
+from MaxText.train import train_step
+from MaxText.train import validate_train_config
+from MaxText.vertex_tensorboard import VertexTensorboardManager
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    elastic_manager,
+    checkpoint_manager,
+):
+  """Reconfigures the workload onto the currently available slices.
+
+  This is called by the elastic manager's maybe_reshard_up/down
+  functions and is responsible for creating a new mesh,
+  reinitializing the state and any objects that depend on the mesh.
+
+  It returns all of the reinitialized objects.
+
+  maybe_reshard_up/down take this function and its arguments and if
+  there is an elastic event, those functions will call this function
+  and return its returns.
+  """
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(elastic_manager.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, elastic_manager.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = elastic_manager.get_resharded_snapshot(mesh)
+
+      # We do not want to restore from the previous checkpoint but instead
+      # restore from the host offloaded snapshot.
+      if checkpoint_manager is not None:
+        latest_step = checkpoint_manager.latest_step()
+
+        # If we checkpointed after the latest snapshot, the checkpoint manager
+        # will try to take another checkpoint and fail because it already
+        # exists. Therefore, we delete the checkpoint and let the checkpoint
+        # manager re-take the checkpoint.
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
+          checkpoint_manager.delete(latest_step)
+
+      state, _, state_mesh_shardings, data_iterator = max_utils.setup_training_state(
+          model,
+          data_iterator,
+          tx,
+          config,
+          jax.random.fold_in(init_rng, step),
+          mesh,
+          checkpoint_manager=None,
+      )
+
+      state = state.replace(**snapshot)
+      state = state.replace(step=state.step.at[None].set(step))
+
+      (
+          functional_train,
+          in_shard_train,
+          out_shard_train,
+          static_argnums_train,
+          donate_argnums_train,
+      ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+      p_train_step = jax.jit(
+          functional_train,
+          in_shardings=in_shard_train,
+          out_shardings=out_shard_train,
+          static_argnums=static_argnums_train,
+          donate_argnums=donate_argnums_train,
+      )
+
+      example_batch = None
+      metric_logger = MetricLogger(writer, config)
+
+      jax.block_until_ready(state)
+
+  return (
+      config,
+      step,
+      state,
+      mesh,
+      checkpoint_manager,
+      data_iterator,
+      p_train_step,
+      example_batch,
+      learning_rate_schedule,
+      metric_logger,
+      writer,
+  )
+
+
+def train_loop(config, elastic_manager, state=None):
+  """Main Training loop.
+  Args:
+    config:
+    state:
+    ckpt_path:
+  Returns:
+  """
+  # Create a GoodputRecorder to log information
+  recorder = create_goodput_recorder(config)
+  record_goodput(recorder, config, recorder.record_job_start_time if recorder else None)
+
+  (
+      init_rng,
+      writer,
+      checkpoint_manager,
+      state_mesh_shardings,
+      model,
+      mesh,
+      learning_rate_schedule,
+      data_iterator,
+      _,
+      state,
+  ) = setup_train_loop(config)
+
+  # pylint: disable=line-too-long
+  (
+      functional_train,
+      in_shard_train,
+      out_shard_train,
+      static_argnums_train,
+      donate_argnums_train,
+  ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+  num_model_parameters = max_utils.calculate_num_params_from_pytree(state.params)
+  max_logging.log(f"number parameters: {num_model_parameters/1e9:.3f} billion")
+  per_device_tflops, _, _ = maxtext_utils.calculate_tflops_training_per_device(config)
+  per_device_tokens = maxtext_utils.calculate_tokens_training_per_device(config)
+
+  # Write train config params, num model params, and XLA flags to tensorboard
+  max_utils.add_text_to_summary_writer("num_model_parameters", str(num_model_parameters), writer)
+  max_utils.add_text_to_summary_writer("libtpu_init_args", os.environ["LIBTPU_INIT_ARGS"], writer)
+  max_utils.add_config_to_summary_writer(config, writer)
+
+  p_train_step = jax.jit(
+      functional_train,
+      in_shardings=in_shard_train,
+      out_shardings=out_shard_train,
+      static_argnums=static_argnums_train,
+      donate_argnums=donate_argnums_train,
+  )
+  running_gcs_metrics = [] if config.gcs_metrics else None
+
+  start_step = get_first_step(state)  # this is the start_step for training
+  prof = profiler.Profiler(config, offset_step=start_step)
+  first_profiling_step = prof.start_initial_profile_step
+  if config.profiler != "" and first_profiling_step >= config.steps:
+    raise ValueError("Profiling requested but initial profiling step set past training final step")
+  last_profiling_step = prof.finished_initial_profile_step
+
+  example_batch = None
+  last_step_completion = datetime.datetime.now()
+
+  performance_metric_queue = None
+  if config.report_heartbeat_metric_for_gcp_monitoring or config.report_performance_metric_for_gcp_monitoring:
+    gcp_workload_monitor = GCPWorkloadMonitor(config.run_name)
+    if config.report_heartbeat_metric_for_gcp_monitoring:
+      gcp_workload_monitor.start_heartbeat_reporting_thread(config.heartbeat_reporting_interval_in_seconds)
+    if config.report_performance_metric_for_gcp_monitoring:
+      performance_metric_queue = queue.Queue()
+      gcp_workload_monitor.start_performance_reporting_thread(performance_metric_queue)
+
+  metric_logger = MetricLogger(writer, config)
+  step = start_step
+
+  elastic_manager.maybe_snapshot(
+      step,
+      snapshot={
+          "params": state.params,
+          "opt_state": state.opt_state,
+      },
+      force=True,
+      block=True,
+  )
+
+  # Using while loop instead of a for loop because with elasticity
+  # the step is restored back to the latest snapshot when a slice is lost
+  while step < config.steps:
path: MaxText/elastic_train.py
user: login: lukebaumann
body: I agree but I think the refactor should happen in train.py.

I know there have been efforts recently to refactor the train_loop into functions. As far as elasticity is concerned, IMO the code blocks have already been moved into functions and are located in `pathwaysutils`. maybe_snapshot, maybe_reshard_down, and maybe_reshard_up are the only elastic related functions/code blocks that are imported from pathwaysutils and the elastic_handler defined in this module is by design its own function.
created_at: 2025-04-09T18:23:11Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,493 @@
+# Copyright 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.
+
+"""Training loop for elastic training model.
+
+Elastic training via Pathways on Cloud allows for slices to go down without
+crashing the main program. In this way, elastic events (slices going up and
+down) can be caught or polled for and reacted to within an elastic_handler
+function without causing the whole workload to restart. This involves
+creating a new mesh with the available slices, reinitializing variables,
+and recompiling functions in addition to restoring from a host offloaded
+snapshot of the state.
+
+The purpose of this training loop is to serve as an example for how to
+support elasticity and is actively in development. As such, there are some
+performance optimizations that have yet to be added as well as some features
+not supported.
+
+Current limitations:
+- The host offloaded snapshot is currently blocking
+- Elastic event handling during async checkpointing is not tested
+- Elastic event handling during profiling is not tested
+- Elastic manager configuration values are hard coded
+- Debug logging statements for elasticity are hard coded as enabled
+
+See https://github.com/AI-Hypercomputer/pathways-utils/tree/main/pathwaysutils/elastic
+for more details about the elastic manager.
+"""
+from collections.abc import Sequence
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from absl import app
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+from flax.linen import partitioning as nn_partitioning
+import jax
+from ml_goodput_measurement import monitoring
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+import tensorflow as tf
+
+from MaxText import checkpointing
+from MaxText import max_utils
+from MaxText import maxtext_utils
+from MaxText import max_logging
+from MaxText import profiler
+from MaxText import pyconfig
+from MaxText.gcp_workload_monitor import GCPWorkloadMonitor
+from MaxText.input_pipeline.input_pipeline_interface import create_data_iterator
+from MaxText.metric_logger import MetricLogger
+from MaxText.train import check_example_batch
+from MaxText.train import create_goodput_recorder
+from MaxText.train import get_first_step
+from MaxText.train import load_next_batch
+from MaxText.train import record_goodput
+from MaxText.train import record_scalar_metrics
+from MaxText.train import save_checkpoint
+from MaxText.train import setup_mesh_and_model
+from MaxText.train import setup_train_loop
+from MaxText.train import train_step
+from MaxText.train import validate_train_config
+from MaxText.vertex_tensorboard import VertexTensorboardManager
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    elastic_manager,
+    checkpoint_manager,
+):
+  """Reconfigures the workload onto the currently available slices.
+
+  This is called by the elastic manager's maybe_reshard_up/down
+  functions and is responsible for creating a new mesh,
+  reinitializing the state and any objects that depend on the mesh.
+
+  It returns all of the reinitialized objects.
+
+  maybe_reshard_up/down take this function and its arguments and if
+  there is an elastic event, those functions will call this function
+  and return its returns.
+  """
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(elastic_manager.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, elastic_manager.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = elastic_manager.get_resharded_snapshot(mesh)
+
+      # We do not want to restore from the previous checkpoint but instead
+      # restore from the host offloaded snapshot.
+      if checkpoint_manager is not None:
+        latest_step = checkpoint_manager.latest_step()
+
+        # If we checkpointed after the latest snapshot, the checkpoint manager
+        # will try to take another checkpoint and fail because it already
+        # exists. Therefore, we delete the checkpoint and let the checkpoint
+        # manager re-take the checkpoint.
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
+          checkpoint_manager.delete(latest_step)
+
+      state, _, state_mesh_shardings, data_iterator = max_utils.setup_training_state(
+          model,
+          data_iterator,
+          tx,
+          config,
+          jax.random.fold_in(init_rng, step),
+          mesh,
+          checkpoint_manager=None,
+      )
+
+      state = state.replace(**snapshot)
+      state = state.replace(step=state.step.at[None].set(step))
+
+      (
+          functional_train,
+          in_shard_train,
+          out_shard_train,
+          static_argnums_train,
+          donate_argnums_train,
+      ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+      p_train_step = jax.jit(
+          functional_train,
+          in_shardings=in_shard_train,
+          out_shardings=out_shard_train,
+          static_argnums=static_argnums_train,
+          donate_argnums=donate_argnums_train,
+      )
+
+      example_batch = None
+      metric_logger = MetricLogger(writer, config)
+
+      jax.block_until_ready(state)
+
+  return (
+      config,
+      step,
+      state,
+      mesh,
+      checkpoint_manager,
+      data_iterator,
+      p_train_step,
+      example_batch,
+      learning_rate_schedule,
+      metric_logger,
+      writer,
+  )
+
+
+def train_loop(config, elastic_manager, state=None):
+  """Main Training loop.
+  Args:
+    config:
+    state:
+    ckpt_path:
+  Returns:
+  """
+  # Create a GoodputRecorder to log information
+  recorder = create_goodput_recorder(config)
+  record_goodput(recorder, config, recorder.record_job_start_time if recorder else None)
+
+  (
+      init_rng,
+      writer,
+      checkpoint_manager,
+      state_mesh_shardings,
+      model,
+      mesh,
+      learning_rate_schedule,
+      data_iterator,
+      _,
+      state,
+  ) = setup_train_loop(config)
+
+  # pylint: disable=line-too-long
+  (
+      functional_train,
+      in_shard_train,
+      out_shard_train,
+      static_argnums_train,
+      donate_argnums_train,
+  ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+  num_model_parameters = max_utils.calculate_num_params_from_pytree(state.params)
+  max_logging.log(f"number parameters: {num_model_parameters/1e9:.3f} billion")
+  per_device_tflops, _, _ = maxtext_utils.calculate_tflops_training_per_device(config)
+  per_device_tokens = maxtext_utils.calculate_tokens_training_per_device(config)
+
+  # Write train config params, num model params, and XLA flags to tensorboard
+  max_utils.add_text_to_summary_writer("num_model_parameters", str(num_model_parameters), writer)
+  max_utils.add_text_to_summary_writer("libtpu_init_args", os.environ["LIBTPU_INIT_ARGS"], writer)
+  max_utils.add_config_to_summary_writer(config, writer)
+
+  p_train_step = jax.jit(
+      functional_train,
+      in_shardings=in_shard_train,
+      out_shardings=out_shard_train,
+      static_argnums=static_argnums_train,
+      donate_argnums=donate_argnums_train,
+  )
+  running_gcs_metrics = [] if config.gcs_metrics else None
+
+  start_step = get_first_step(state)  # this is the start_step for training
+  prof = profiler.Profiler(config, offset_step=start_step)
+  first_profiling_step = prof.start_initial_profile_step
+  if config.profiler != "" and first_profiling_step >= config.steps:
+    raise ValueError("Profiling requested but initial profiling step set past training final step")
+  last_profiling_step = prof.finished_initial_profile_step
+
+  example_batch = None
+  last_step_completion = datetime.datetime.now()
+
+  performance_metric_queue = None
+  if config.report_heartbeat_metric_for_gcp_monitoring or config.report_performance_metric_for_gcp_monitoring:
+    gcp_workload_monitor = GCPWorkloadMonitor(config.run_name)
+    if config.report_heartbeat_metric_for_gcp_monitoring:
+      gcp_workload_monitor.start_heartbeat_reporting_thread(config.heartbeat_reporting_interval_in_seconds)
+    if config.report_performance_metric_for_gcp_monitoring:
+      performance_metric_queue = queue.Queue()
+      gcp_workload_monitor.start_performance_reporting_thread(performance_metric_queue)
+
+  metric_logger = MetricLogger(writer, config)
+  step = start_step
+
+  elastic_manager.maybe_snapshot(
+      step,
+      snapshot={
+          "params": state.params,
+          "opt_state": state.opt_state,
+      },
+      force=True,
+      block=True,
+  )
+
+  # Using while loop instead of a for loop because with elasticity
+  # the step is restored back to the latest snapshot when a slice is lost
+  while step < config.steps:
+    try:
+      if step == first_profiling_step or prof.should_activate_periodic_profile(step):
+        optional_postfix = f"step_{step}" if config.profile_periodically_period > 0 else ""
+        prof.activate(blocking_object=state, optional_postfix=optional_postfix)
+
+      max_logging.log(f"{step=} {elastic_manager.elastic_down_event_count=} {elastic_manager.good_slice_count=}")
+      with mesh, nn_partitioning.axis_rules(config.logical_axis_rules), jax.default_device(elastic_manager.default_device):
+        with jax.profiler.StepTraceAnnotation("train", step_num=step):
+          record_goodput(recorder, config, recorder.record_data_loading_start_time if recorder else None)
+          try:
+            example_batch = load_next_batch(data_iterator, example_batch, config)
+          except Exception as e:  # pylint: disable=broad-except
+            max_logging.log(f"load_next_batch failed, you may have run out of data. Error message: {e}")
+            break
+          record_goodput(recorder, config, recorder.record_data_loading_end_time if recorder else None)
+          check_example_batch(config, example_batch=example_batch)
+          # pylint: disable=not-callable
+          nextrng = jax.jit(jax.random.fold_in)(init_rng, step)
+          record_goodput(recorder, config, recorder.record_step_start_time if recorder else None, step)
+          state, metrics = p_train_step(state, example_batch, nextrng)
+
+        step_time_delta = datetime.datetime.now() - last_step_completion
+        last_step_completion = datetime.datetime.now()
+        record_scalar_metrics(metrics, step_time_delta, per_device_tflops, learning_rate_schedule(step), per_device_tokens)
+        if performance_metric_queue:
+          performance_metric_queue.put(step_time_delta.total_seconds())
+
+        if checkpoint_manager is not None:
+          state_to_save = state
+          if save_checkpoint(checkpoint_manager, int(step), state_to_save, config.dataset_type, data_iterator, config):
+            checkpointing.print_save_message(step, config.async_checkpointing)
+
+          # Upon preemption, exit when and only when all ongoing saves are complete.
+          if checkpoint_manager.reached_preemption(step):
+            checkpoint_manager.wait_until_finished()
+            sys.exit()
path: MaxText/elastic_train.py
user: login: lukebaumann
body: 100% agree. This is copied verbatim from train.py and unrelated to elasticity. The reason for not improving the train_loop is discussed in an above comment.
created_at: 2025-04-09T18:24:30Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,493 @@
+# Copyright 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.
+
+"""Training loop for elastic training model.
+
+Elastic training via Pathways on Cloud allows for slices to go down without
+crashing the main program. In this way, elastic events (slices going up and
+down) can be caught or polled for and reacted to within an elastic_handler
+function without causing the whole workload to restart. This involves
+creating a new mesh with the available slices, reinitializing variables,
+and recompiling functions in addition to restoring from a host offloaded
+snapshot of the state.
+
+The purpose of this training loop is to serve as an example for how to
+support elasticity and is actively in development. As such, there are some
+performance optimizations that have yet to be added as well as some features
+not supported.
+
+Current limitations:
+- The host offloaded snapshot is currently blocking
+- Elastic event handling during async checkpointing is not tested
+- Elastic event handling during profiling is not tested
+- Elastic manager configuration values are hard coded
+- Debug logging statements for elasticity are hard coded as enabled
+
+See https://github.com/AI-Hypercomputer/pathways-utils/tree/main/pathwaysutils/elastic
+for more details about the elastic manager.
+"""
+from collections.abc import Sequence
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from absl import app
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+from flax.linen import partitioning as nn_partitioning
+import jax
+from ml_goodput_measurement import monitoring
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+import tensorflow as tf
+
+from MaxText import checkpointing
+from MaxText import max_utils
+from MaxText import maxtext_utils
+from MaxText import max_logging
+from MaxText import profiler
+from MaxText import pyconfig
+from MaxText.gcp_workload_monitor import GCPWorkloadMonitor
+from MaxText.input_pipeline.input_pipeline_interface import create_data_iterator
+from MaxText.metric_logger import MetricLogger
+from MaxText.train import check_example_batch
+from MaxText.train import create_goodput_recorder
+from MaxText.train import get_first_step
+from MaxText.train import load_next_batch
+from MaxText.train import record_goodput
+from MaxText.train import record_scalar_metrics
+from MaxText.train import save_checkpoint
+from MaxText.train import setup_mesh_and_model
+from MaxText.train import setup_train_loop
+from MaxText.train import train_step
+from MaxText.train import validate_train_config
+from MaxText.vertex_tensorboard import VertexTensorboardManager
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    elastic_manager,
+    checkpoint_manager,
+):
+  """Reconfigures the workload onto the currently available slices.
+
+  This is called by the elastic manager's maybe_reshard_up/down
+  functions and is responsible for creating a new mesh,
+  reinitializing the state and any objects that depend on the mesh.
+
+  It returns all of the reinitialized objects.
+
+  maybe_reshard_up/down take this function and its arguments and if
+  there is an elastic event, those functions will call this function
+  and return its returns.
+  """
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(elastic_manager.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, elastic_manager.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = elastic_manager.get_resharded_snapshot(mesh)
+
+      # We do not want to restore from the previous checkpoint but instead
+      # restore from the host offloaded snapshot.
+      if checkpoint_manager is not None:
+        latest_step = checkpoint_manager.latest_step()
+
+        # If we checkpointed after the latest snapshot, the checkpoint manager
+        # will try to take another checkpoint and fail because it already
+        # exists. Therefore, we delete the checkpoint and let the checkpoint
+        # manager re-take the checkpoint.
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
+          checkpoint_manager.delete(latest_step)
+
+      state, _, state_mesh_shardings, data_iterator = max_utils.setup_training_state(
+          model,
+          data_iterator,
+          tx,
+          config,
+          jax.random.fold_in(init_rng, step),
+          mesh,
+          checkpoint_manager=None,
+      )
+
+      state = state.replace(**snapshot)
+      state = state.replace(step=state.step.at[None].set(step))
+
+      (
+          functional_train,
+          in_shard_train,
+          out_shard_train,
+          static_argnums_train,
+          donate_argnums_train,
+      ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+      p_train_step = jax.jit(
+          functional_train,
+          in_shardings=in_shard_train,
+          out_shardings=out_shard_train,
+          static_argnums=static_argnums_train,
+          donate_argnums=donate_argnums_train,
+      )
+
+      example_batch = None
+      metric_logger = MetricLogger(writer, config)
+
+      jax.block_until_ready(state)
+
+  return (
+      config,
+      step,
+      state,
+      mesh,
+      checkpoint_manager,
+      data_iterator,
+      p_train_step,
+      example_batch,
+      learning_rate_schedule,
+      metric_logger,
+      writer,
+  )
+
+
+def train_loop(config, elastic_manager, state=None):
+  """Main Training loop.
+  Args:
+    config:
+    state:
+    ckpt_path:
+  Returns:
+  """
+  # Create a GoodputRecorder to log information
+  recorder = create_goodput_recorder(config)
+  record_goodput(recorder, config, recorder.record_job_start_time if recorder else None)
+
+  (
+      init_rng,
+      writer,
+      checkpoint_manager,
+      state_mesh_shardings,
+      model,
+      mesh,
+      learning_rate_schedule,
+      data_iterator,
+      _,
+      state,
+  ) = setup_train_loop(config)
+
+  # pylint: disable=line-too-long
+  (
+      functional_train,
+      in_shard_train,
+      out_shard_train,
+      static_argnums_train,
+      donate_argnums_train,
+  ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+  num_model_parameters = max_utils.calculate_num_params_from_pytree(state.params)
+  max_logging.log(f"number parameters: {num_model_parameters/1e9:.3f} billion")
+  per_device_tflops, _, _ = maxtext_utils.calculate_tflops_training_per_device(config)
+  per_device_tokens = maxtext_utils.calculate_tokens_training_per_device(config)
+
+  # Write train config params, num model params, and XLA flags to tensorboard
+  max_utils.add_text_to_summary_writer("num_model_parameters", str(num_model_parameters), writer)
+  max_utils.add_text_to_summary_writer("libtpu_init_args", os.environ["LIBTPU_INIT_ARGS"], writer)
+  max_utils.add_config_to_summary_writer(config, writer)
+
+  p_train_step = jax.jit(
+      functional_train,
+      in_shardings=in_shard_train,
+      out_shardings=out_shard_train,
+      static_argnums=static_argnums_train,
+      donate_argnums=donate_argnums_train,
+  )
+  running_gcs_metrics = [] if config.gcs_metrics else None
+
+  start_step = get_first_step(state)  # this is the start_step for training
+  prof = profiler.Profiler(config, offset_step=start_step)
+  first_profiling_step = prof.start_initial_profile_step
+  if config.profiler != "" and first_profiling_step >= config.steps:
+    raise ValueError("Profiling requested but initial profiling step set past training final step")
+  last_profiling_step = prof.finished_initial_profile_step
+
+  example_batch = None
+  last_step_completion = datetime.datetime.now()
+
+  performance_metric_queue = None
+  if config.report_heartbeat_metric_for_gcp_monitoring or config.report_performance_metric_for_gcp_monitoring:
+    gcp_workload_monitor = GCPWorkloadMonitor(config.run_name)
+    if config.report_heartbeat_metric_for_gcp_monitoring:
+      gcp_workload_monitor.start_heartbeat_reporting_thread(config.heartbeat_reporting_interval_in_seconds)
+    if config.report_performance_metric_for_gcp_monitoring:
+      performance_metric_queue = queue.Queue()
+      gcp_workload_monitor.start_performance_reporting_thread(performance_metric_queue)
+
+  metric_logger = MetricLogger(writer, config)
+  step = start_step
+
+  elastic_manager.maybe_snapshot(
+      step,
+      snapshot={
+          "params": state.params,
+          "opt_state": state.opt_state,
+      },
+      force=True,
+      block=True,
+  )
+
+  # Using while loop instead of a for loop because with elasticity
+  # the step is restored back to the latest snapshot when a slice is lost
+  while step < config.steps:
+    try:
+      if step == first_profiling_step or prof.should_activate_periodic_profile(step):
+        optional_postfix = f"step_{step}" if config.profile_periodically_period > 0 else ""
+        prof.activate(blocking_object=state, optional_postfix=optional_postfix)
+
+      max_logging.log(f"{step=} {elastic_manager.elastic_down_event_count=} {elastic_manager.good_slice_count=}")
+      with mesh, nn_partitioning.axis_rules(config.logical_axis_rules), jax.default_device(elastic_manager.default_device):
+        with jax.profiler.StepTraceAnnotation("train", step_num=step):
+          record_goodput(recorder, config, recorder.record_data_loading_start_time if recorder else None)
+          try:
+            example_batch = load_next_batch(data_iterator, example_batch, config)
+          except Exception as e:  # pylint: disable=broad-except
+            max_logging.log(f"load_next_batch failed, you may have run out of data. Error message: {e}")
+            break
+          record_goodput(recorder, config, recorder.record_data_loading_end_time if recorder else None)
+          check_example_batch(config, example_batch=example_batch)
+          # pylint: disable=not-callable
+          nextrng = jax.jit(jax.random.fold_in)(init_rng, step)
+          record_goodput(recorder, config, recorder.record_step_start_time if recorder else None, step)
+          state, metrics = p_train_step(state, example_batch, nextrng)
+
+        step_time_delta = datetime.datetime.now() - last_step_completion
+        last_step_completion = datetime.datetime.now()
+        record_scalar_metrics(metrics, step_time_delta, per_device_tflops, learning_rate_schedule(step), per_device_tokens)
+        if performance_metric_queue:
+          performance_metric_queue.put(step_time_delta.total_seconds())
+
+        if checkpoint_manager is not None:
+          state_to_save = state
+          if save_checkpoint(checkpoint_manager, int(step), state_to_save, config.dataset_type, data_iterator, config):
+            checkpointing.print_save_message(step, config.async_checkpointing)
+
+          # Upon preemption, exit when and only when all ongoing saves are complete.
+          if checkpoint_manager.reached_preemption(step):
+            checkpoint_manager.wait_until_finished()
+            sys.exit()
+
+        metric_logger.write_metrics(running_gcs_metrics, metrics, step)
+
+        if step == last_profiling_step or prof.should_deactivate_periodic_profile(step):
+          prof.deactivate(blocking_object=state)
+
+      elastic_manager.maybe_snapshot(
+          step=step,
+          snapshot={
+              "params": state.params,
+              "opt_state": state.opt_state,
+          },
+          block=True,
+      )
+
+      ret = elastic_manager.maybe_reshard_up(
+          step=step,
+          snapshot={
+              "params": state.params,
+              "opt_state": state.opt_state,
+          },
+          elastic_handler=elastic_handler,
+          handler_kwargs={
+              "config": config,
+              "elastic_manager": elastic_manager,
+              "checkpoint_manager": checkpoint_manager,
+          },
+      )
+      if ret is not None:
+        (
+            config,
+            step,
+            state,
+            mesh,
+            checkpoint_manager,
+            data_iterator,
+            p_train_step,
+            example_batch,
+            learning_rate_schedule,
+            metric_logger,
+            writer,
+        ) = ret
+
+      if step == start_step:
+        max_utils.print_mem_stats("After params initialized")
+
+      step += 1
+
+    except jax.errors.JaxRuntimeError as error:
+      ret = elastic_manager.maybe_reshard_down(
+          error=error,
+          elastic_handler=elastic_handler,
+          handler_kwargs={
+              "config": config,
+              "elastic_manager": elastic_manager,
+              "checkpoint_manager": checkpoint_manager,
+          },
+      )
+      if ret is not None:
+        (
+            config,
+            step,
+            state,
+            mesh,
+            checkpoint_manager,
+            data_iterator,
+            p_train_step,
+            example_batch,
+            learning_rate_schedule,
+            metric_logger,
+            writer,
+        ) = ret
+
+  if checkpoint_manager is not None:
+    if int(state.step) - 1 % config.checkpoint_period != 0:
+      try:
+        state_to_save = state
+        if save_checkpoint(
+            checkpoint_manager,
+            int(state_to_save) - 1,
+            state_to_save,
+            config.dataset_type,
+            data_iterator,
+            config,
+            force=True,
+        ):
+          checkpointing.print_save_message(int(state_to_save) - 1, config.async_checkpointing)
+      except Exception:  # pylint: disable=broad-except
+        max_logging.log(f"Checkpoint is already saved for step {int(state.step)-1}.")
+
+    checkpoint_manager.wait_until_finished()
+  metric_logger.write_metrics(running_gcs_metrics, metrics, config.steps - 1)  # final step metrics
+  max_utils.close_summary_writer(writer)
+  record_goodput(recorder, config, recorder.record_job_end_time if recorder else None)
+  with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
+    # pytype: disable=attribute-error
+    compiled = p_train_step.lower(state, example_batch, nextrng).compile()
+    compiled_stats = compiled.memory_analysis()
+    if compiled_stats is not None:
+      max_logging.log(
+          f"Output size: {compiled_stats.output_size_in_bytes}, "
+          f"temp size: {compiled_stats.temp_size_in_bytes}, "
+          f"argument size: {compiled_stats.argument_size_in_bytes}, "
+          f"host temp size: {compiled_stats.host_temp_size_in_bytes}, in bytes."
+      )
+  return state
+
+
+def main(argv: Sequence[str]) -> None:
path: MaxText/elastic_train.py
user: login: lukebaumann
body: Thanks! I moved elastic initialization into its own function.
created_at: 2025-04-09T19:03:33Z
author_association: COLLABORATOR
diff_hunk: @@ -37,5 +36,5 @@ transformers
 mlperf-logging@git+https://github.com/mlperf/logging.git
 google-jetstream@git+https://github.com/AI-Hypercomputer/JetStream.git
 jsonlines
-pathwaysutils@git+https://github.com/google/pathways-utils.git
+pathwaysutils@git+https://github.com/AI-Hypercomputer/pathways-utils.git
path: requirements.txt
user: login: lukebaumann
body: Made the change in the stable stack requirements as well. This redirects automatically from google to AI-Hypercomputer but I wanted to clean this up.

The more important change in this file is removing the other pathwaysutils requirement which installs from PyPI but is being overruled by the later one installed from GitHub. We want the GitHub HEAD version only and not to confuse anyone about what is required.
created_at: 2025-04-09T19:05:59Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,472 @@
+"""
+Copyright 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.
+"""
+
+"""Training loop for elastic training model.
+
+Elastic training via Pathways on Cloud allows for slices to go down without
+crashing the main program. In this way, elastic events (slices going up and
+down) can be caught or polled for and reacted to within an elastic_handler
+function without causing the whole workload to restart. This involves
+creating a new mesh with the available slices, reinitializing variables,
+and recompiling functions in addition to restoring from a host offloaded
+snapshot of the state.
+
+The purpose of this training loop is to serve as an example for how to
+support elasticity and is actively in development. As such, there are some
+performance optimizations that have yet to be added as well as some features
+not supported.
+
+Current limitations:
+- The host offloaded snapshot is currently blocking
+- Elastic event handling during async checkpointing is not tested
+- Elastic event handling during profiling is not tested
+- Elastic manager configuration values are hard coded
+- Debug logging statements for elasticity are hard coded as enabled
+
+See https://github.com/AI-Hypercomputer/pathways-utils/tree/main/pathwaysutils/elastic
+for more details about the elastic manager.
+"""
+import datetime
+import logging
+import os
+import sys
+import time
+import queue
+
+from collections.abc import Sequence
+from absl import app
+from flax.linen import partitioning as nn_partitioning
+import jax
+
+import checkpointing
+import max_utils
+import maxtext_utils
+import max_logging
+import profiler
+import pyconfig
+import pathwaysutils
+from pathwaysutils.elastic import manager
+from pathwaysutils.debug import timing
+from pathwaysutils.debug import watchdog
+import tensorflow as tf
+
+from metric_logger import MetricLogger
+
+from vertex_tensorboard import VertexTensorboardManager
+
+from input_pipeline.input_pipeline_interface import create_data_iterator
+
+from gcp_workload_monitor import GCPWorkloadMonitor
+
+from cloud_tpu_diagnostics import diagnostic
+from cloud_tpu_diagnostics.configuration import debug_configuration
+from cloud_tpu_diagnostics.configuration import diagnostic_configuration
+from cloud_tpu_diagnostics.configuration import stack_trace_configuration
+
+from ml_goodput_measurement import monitoring
+
+from train import (
+    check_example_batch,
+    create_goodput_recorder,
+    get_first_step,
+    load_next_batch,
+    record_goodput,
+    record_scalar_metrics,
+    save_checkpoint,
+    setup_mesh_and_model,
+    setup_train_loop,
+    train_step,
+    validate_train_config,
+)
+
+logging.basicConfig()
+logging.getLogger("pathwaysutils.elastic.manager").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.elastic.reshard").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.timing").setLevel(logging.DEBUG)
+logging.getLogger("pathwaysutils.debug.watchdog").setLevel(logging.DEBUG)
+
+
+@timing.timeit
+def elastic_handler(
+    config: pyconfig.HyperParameters,
+    elastic_manager,
+    checkpoint_manager,
+):
+  """Reshard function."""
+  if checkpoint_manager is not None:
+    checkpoint_manager.close()
+
+  with jax.default_device(elastic_manager.default_device):
+    init_rng, writer, checkpoint_manager, mesh, model, learning_rate_schedule, tx = setup_mesh_and_model(
+        config, elastic_manager.good_devices
+    )
+    with mesh:
+      data_iterator, _ = create_data_iterator(config, mesh)
+
+      step, snapshot = elastic_manager.get_resharded_snapshot(mesh)
+
+      if checkpoint_manager is not None:
+        # Confirm this is the right thing to do
+        latest_step = checkpoint_manager.latest_step()
+        if latest_step is not None and latest_step >= step:
+          max_logging.log(f"Deleting checkpoint from step {latest_step} since we are rewinding to step {step}.")
+          checkpoint_manager.delete(latest_step)
+
+      state, _, state_mesh_shardings, data_iterator = max_utils.setup_training_state(
+          model,
+          data_iterator,
+          tx,
+          config,
+          jax.random.fold_in(init_rng, step),
+          mesh,
+          checkpoint_manager=None,
+      )
+
+      state = state.replace(**snapshot)
+      state = state.replace(step=state.step.at[None].set(step))
+
+      (
+          functional_train,
+          in_shard_train,
+          out_shard_train,
+          static_argnums_train,
+          donate_argnums_train,
+      ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+      p_train_step = jax.jit(
+          functional_train,
+          in_shardings=in_shard_train,
+          out_shardings=out_shard_train,
+          static_argnums=static_argnums_train,
+          donate_argnums=donate_argnums_train,
+      )
+
+      example_batch = None
+      metric_logger = MetricLogger(writer, config)
+
+      jax.block_until_ready(state)
+
+  return (
+      config,
+      step,
+      state,
+      mesh,
+      checkpoint_manager,
+      data_iterator,
+      p_train_step,
+      example_batch,
+      learning_rate_schedule,
+      metric_logger,
+      writer,
+  )
+
+
+def train_loop(config, elastic_manager, state=None):
+  """Main Training loop.
+  Args:
+    config:
+    state:
+    ckpt_path:
+  Returns:
+  """
+  # Create a GoodputRecorder to log information
+  recorder = create_goodput_recorder(config)
+  record_goodput(recorder, config, recorder.record_job_start_time if recorder else None)
+
+  (
+      init_rng,
+      writer,
+      checkpoint_manager,
+      state_mesh_shardings,
+      model,
+      mesh,
+      learning_rate_schedule,
+      data_iterator,
+      _,
+      state,
+  ) = setup_train_loop(config)
+
+  # pylint: disable=line-too-long
+  (
+      functional_train,
+      in_shard_train,
+      out_shard_train,
+      static_argnums_train,
+      donate_argnums_train,
+  ) = maxtext_utils.get_functional_train_with_signature(train_step, mesh, state_mesh_shardings, model, config)
+
+  p_train_step = jax.jit(
+      functional_train,
+      in_shardings=in_shard_train,
+      out_shardings=out_shard_train,
+      static_argnums=static_argnums_train,
+      donate_argnums=donate_argnums_train,
+  )
+
+  num_model_parameters = max_utils.calculate_num_params_from_pytree(state.params)
+  max_logging.log(f"number parameters: {num_model_parameters/1e9:.3f} billion")
+  per_device_tflops, _, _ = maxtext_utils.calculate_tflops_training_per_device(config)
+  per_device_tokens = maxtext_utils.calculate_tokens_training_per_device(config)
+
+  # Write train config params, num model params, and XLA flags to tensorboard
+  max_utils.add_text_to_summary_writer("num_model_parameters", str(num_model_parameters), writer)
+  max_utils.add_text_to_summary_writer("libtpu_init_args", os.environ["LIBTPU_INIT_ARGS"], writer)
+  max_utils.add_config_to_summary_writer(config, writer)
+
+  running_gcs_metrics = [] if config.gcs_metrics else None
+
+  start_step = get_first_step(state)  # this is the start_step for training
+  prof = profiler.Profiler(config, offset_step=start_step)
+  first_profiling_step = prof.start_initial_profile_step
+  if config.profiler != "" and first_profiling_step >= config.steps:
+    raise ValueError("Profiling requested but initial profiling step set past training final step")
+  last_profiling_step = prof.finished_initial_profile_step
+
+  example_batch = None
+  last_step_completion = datetime.datetime.now()
+
+  performance_metric_queue = None
+  if config.report_heartbeat_metric_for_gcp_monitoring or config.report_performance_metric_for_gcp_monitoring:
+    gcp_workload_monitor = GCPWorkloadMonitor(config.run_name)
+    if config.report_heartbeat_metric_for_gcp_monitoring:
+      gcp_workload_monitor.start_heartbeat_reporting_thread(config.heartbeat_reporting_interval_in_seconds)
+    if config.report_performance_metric_for_gcp_monitoring:
+      performance_metric_queue = queue.Queue()
+      gcp_workload_monitor.start_performance_reporting_thread(performance_metric_queue)
+
+  metric_logger = MetricLogger(writer, config)
+  step = start_step
+
+  elastic_manager.maybe_snapshot(
+      step,
+      snapshot={
+          "params": state.params,
+          "opt_state": state.opt_state,
+      },
+      force=True,
+      block=True,
+  )
+
+  while True:
+    with watchdog.watchdog(120):
+      try:
+        if step == first_profiling_step or prof.should_activate_periodic_profile(step):
+          optional_postfix = f"step_{step}" if config.profile_periodically_period > 0 else ""
+          prof.activate(blocking_object=state, optional_postfix=optional_postfix)
+
+        if step >= config.steps:
+          break
+
+        max_logging.log(f"{step=} {elastic_manager.elastic_down_event_count=} {elastic_manager.good_slice_count=}")
+        with mesh, nn_partitioning.axis_rules(config.logical_axis_rules), jax.default_device(elastic_manager.default_device):
+          with jax.profiler.StepTraceAnnotation("train", step_num=step):
+            record_goodput(recorder, config, recorder.record_data_loading_start_time if recorder else None)
+            example_batch = load_next_batch(data_iterator, example_batch, config)
+            record_goodput(recorder, config, recorder.record_data_loading_end_time if recorder else None)
+            check_example_batch(config, example_batch=example_batch)
+            # pylint: disable=not-callable
+            nextrng = jax.jit(jax.random.fold_in)(init_rng, step)
+            record_goodput(recorder, config, recorder.record_step_start_time if recorder else None, step)
+            state, metrics = p_train_step(state, example_batch, nextrng)
+
+          step_time_delta = datetime.datetime.now() - last_step_completion
+          last_step_completion = datetime.datetime.now()
+          record_scalar_metrics(metrics, step_time_delta, per_device_tflops, learning_rate_schedule(step), per_device_tokens)
+          if performance_metric_queue:
+            performance_metric_queue.put(step_time_delta.total_seconds())
+
+          if checkpoint_manager is not None:
+            state_to_save = state
+            if save_checkpoint(checkpoint_manager, int(step), state_to_save, config.dataset_type, data_iterator, config):
+              checkpointing.print_save_message(step, config.async_checkpointing)
+
+            # Upon preemption, exit when and only when all ongoing saves are complete.
+            if checkpoint_manager.reached_preemption(step):
+              checkpoint_manager.wait_until_finished()
+              sys.exit()
+
+          metric_logger.write_metrics(running_gcs_metrics, metrics, step)
+
+          if step == last_profiling_step or prof.should_deactivate_periodic_profile(step):
+            prof.deactivate(blocking_object=state)
+
+        elastic_manager.maybe_snapshot(
+            step=step,
+            snapshot={
+                "params": state.params,
+                "opt_state": state.opt_state,
+            },
+            block=True,
+        )
+
+        ret = elastic_manager.maybe_reshard_up(
+            step=step,
+            snapshot={
+                "params": state.params,
+                "opt_state": state.opt_state,
+            },
+            elastic_handler=elastic_handler,
+            handler_kwargs={
+                "config": config,
+                "elastic_manager": elastic_manager,
+                "checkpoint_manager": checkpoint_manager,
+            },
+        )
+        if ret is not None:
+          (
+              config,
+              step,
+              state,
+              mesh,
+              checkpoint_manager,
+              data_iterator,
+              p_train_step,
+              example_batch,
+              learning_rate_schedule,
+              metric_logger,
+              writer,
+          ) = ret
+
+        if step == start_step:
+          max_utils.print_mem_stats("After params initialized")
+
+        step += 1
+
+      except jax.errors.JaxRuntimeError as error:
+        ret = elastic_manager.maybe_reshard_down(
+            error=error,
+            elastic_handler=elastic_handler,
+            handler_kwargs={
+                "config": config,
+                "elastic_manager": elastic_manager,
+                "checkpoint_manager": checkpoint_manager,
+            },
+        )
+        if ret is not None:
+          (
+              config,
+              step,
+              state,
+              mesh,
+              checkpoint_manager,
+              data_iterator,
+              p_train_step,
+              example_batch,
+              learning_rate_schedule,
+              metric_logger,
+              writer,
+          ) = ret
+
+  if checkpoint_manager is not None:
path: MaxText/elastic_train.py
user: login: SurbhiJainUSC
body: Thanks @lukebaumann for catching this bug. This should be fixed in [#1555](https://github.com/AI-Hypercomputer/maxtext/pull/1555).
created_at: 2025-04-09T20:28:50Z
author_association: COLLABORATOR

PR comments:
user: login: gobbleturk
created_at: 2025-04-04T17:47:22Z
body: Can you include a short description + "how to use" for elastic training in our readme in this section 

https://github.com/AI-Hypercomputer/maxtext?tab=readme-ov-file#features-and-diagnostics
author_association: COLLABORATOR
user: login: gobbleturk
created_at: 2025-04-04T22:04:36Z
body: Maybe link https://github.com/AI-Hypercomputer/pathways-utils/tree/main/pathwaysutils/elastic as mention a lot of the logic handled there under the elastic_manager class in maxtext. This code in pathwaysutils is very professional, you should defintely reference it =)

Are there tests you can reference in a different codebase? Are the tests only internal?
author_association: COLLABORATOR
user: login: lukebaumann
created_at: 2025-04-07T17:48:55Z
body: Added the reference. There are internal tests that will be moved to OSS soon. In addition, an airflow XLML test will be written once this is merged.
author_association: COLLABORATOR
user: login: lukebaumann
created_at: 2025-04-07T17:50:16Z
body: > Can you include a short description + "how to use" for elastic training in our readme in this section
> 
> https://github.com/AI-Hypercomputer/maxtext?tab=readme-ov-file#features-and-diagnostics

There will be official documentation published soon that we can link in the readme but it is not linkable yet.
author_association: COLLABORATOR
user: login: lukebaumann
created_at: 2025-04-09T19:17:48Z
body: > cc @gobbleturk
> 
> Thanks for this, @lukebaumann - looks great!
> 
> One bigger picture question I have, which doesn't necessarily need to be addressed now, but could be, is - should this be integrated with train.py, particularly in a world where we're focused on improving code quality and modularizing the codebase. I see that there's about a 30% overlap in LoC (maybe at the bottom of the two files) and I'm wondering whether this could be accomplished by train.py having hooks out to elastic training.
> 
> In addition, I've commented in a couple of places about breaking long functions out into sub-functions, which would help with testability as well as comprehensibility. I notice there are no unit tests right now - it would be great to get some, and I'm sure trying to do so would drive some of that decomposition in a practical way (vs. because Rich said so)

Spoke with @gobbleturk about this and he recommended keeping them separate. For one, elasticity is only for Pathways workloads and while the changes should not impact McJAX workloads, I do not want to risk any unintended impacts to existing users. I also have not tested elasticity with all of the possible configurations (in part because I am not aware of most of the configurations possible in MaxText).

I think this is a great example of what is possible because of the existing modularization. I foresee most of train.py to be moved into a new module train_lib.py and train.py, elastic_train.py, sft_train.py share even less code.

The elasticity unit tests are in pathwaysutils and will be OSSed soon. I will be making an airflow test for elasticity once this is merged. 
author_association: COLLABORATOR
user: login: gobbleturk
created_at: 2025-04-17T21:03:31Z
body: Add pull ready manually, surprisingly it is complaning 2 owners haven't approved despite rich and I are both code owners https://github.com/AI-Hypercomputer/maxtext/actions/runs/14523196283/job/40748720809
author_association: COLLABORATOR