Use helper functions for CheckpointManager operations.

src/maxtext/checkpoint_conversion/reshard_checkpoint.py

@@ -117,7 +117,7 @@ def main(argv: Sequence[str]) -> None:
     save_ckpt_path = os.path.join(save_ckpt_directory, str(step_number), "items")
     max_logging.log(f"Saved checkpoint: {save_ckpt_path}")
     # Upon preemption, exit when and only when all ongoing saves are complete.
-    checkpoint_manager.wait_until_finished()
+    checkpointing.wait_until_finished(checkpoint_manager)
 
   max_logging.log(f"Elapse for checkpoint save: {(time.time() - start) / 60:.2f} min")
   print_peak_memory()

src/maxtext/checkpoint_conversion/standalone_scripts/convert_gemma2_chkpt.py

@@ -301,8 +301,8 @@ def astype_fn(x):
     if checkpointing.save_checkpoint(checkpoint_manager, 0, state_new):
       max_logging.log("saved a checkpoint at step 0")
     # Upon preemption, exit when and only when all ongoing saves are complete.
-    if checkpoint_manager.reached_preemption(0):
-      checkpoint_manager.wait_until_finished()
+    if checkpointing.reached_preemption(checkpoint_manager, 0):
+      checkpointing.wait_until_finished(checkpoint_manager)
       sys.exit()
 
 

src/maxtext/checkpoint_conversion/standalone_scripts/convert_gemma3_chkpt.py

@@ -241,8 +241,8 @@ def astype_fn(x):
       max_logging.log("saved a checkpoint at step 0")
       max_logging.log(f"Checkpoint saved to: {args.maxtext_model_path}")
     # Upon preemption, exit when and only when all ongoing saves are complete.
-    if checkpoint_manager.reached_preemption(0):
-      checkpoint_manager.wait_until_finished()
+    if checkpointing.reached_preemption(checkpoint_manager, 0):
+      checkpointing.wait_until_finished(checkpoint_manager)
       sys.exit()
 
 

src/maxtext/checkpoint_conversion/standalone_scripts/convert_gemma_chkpt.py

@@ -175,8 +175,8 @@ def astype_fn(x):
     if checkpointing.save_checkpoint(checkpoint_manager, 0, state_new):
       max_logging.log("saved a checkpoint at step 0")
     # Upon preemption, exit when and only when all ongoing saves are complete.
-    if checkpoint_manager.reached_preemption(0):
-      checkpoint_manager.wait_until_finished()
+    if checkpointing.reached_preemption(checkpoint_manager, 0):
+      checkpointing.wait_until_finished(checkpoint_manager)
       sys.exit()
 
 

src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py

@@ -306,8 +306,8 @@ def map_fn(key_path, value):
   if checkpointing.save_checkpoint(checkpoint_manager, step_value, converted_state):
     max_logging.log(f"saved a checkpoint at step {step_value}")
   # Upon preemption, exit when and only when all ongoing saves are complete.
-  if checkpoint_manager.reached_preemption(step_value):
-    checkpoint_manager.wait_until_finished()
+  if checkpointing.reached_preemption(checkpoint_manager, step_value):
+    checkpointing.wait_until_finished(checkpoint_manager)
     sys.exit()
 
   max_logging.log(f"Peak cpu memory in a single process: {fmt_size(memory_metrics['max_cpu_bytes'])}")

src/maxtext/checkpoint_conversion/standalone_scripts/llama_ckpt_conversion_inference_only.py

@@ -314,8 +314,8 @@ def checkpoint_device_put(arr):
     if checkpointing.save_checkpoint(checkpoint_manager, step_number_to_save_new_ckpt, state_new):
       max_logging.log(f"saved a checkpoint at step {step_number_to_save_new_ckpt}")
     # Upon preemption, exit when and only when all ongoing saves are complete.
-    if checkpoint_manager.reached_preemption(0):
-      checkpoint_manager.wait_until_finished()
+    if checkpointing.reached_preemption(checkpoint_manager, 0):
+      checkpointing.wait_until_finished(checkpoint_manager)
       sys.exit()
 
 

src/maxtext/checkpoint_conversion/utils/utils.py

@@ -1246,7 +1246,7 @@ def save_weights_to_checkpoint(
   if checkpointing.save_checkpoint(checkpoint_manager, step_number_to_save_new_ckpt, state_new, config=config):
     max_logging.log(f"saved a checkpoint at step {step_number_to_save_new_ckpt}")
   # Upon preemption, exit when and only when all ongoing saves are complete.
-  checkpoint_manager.wait_until_finished()
+  checkpointing.wait_until_finished(checkpoint_manager)
 
   max_logging.log(f"Elapse for checkpoint save: {(time.time() - start) / 60:.2f} min")
 

src/maxtext/common/checkpointing.py

@@ -43,7 +43,6 @@
 from orbax.checkpoint._src.checkpoint_managers import preservation_policy as preservation_policy_lib
 from orbax.checkpoint._src.checkpoint_managers import save_decision_policy as save_decision_policy_lib
 
-
 CheckpointManagerOptions = ocp.CheckpointManagerOptions
 Composite = ocp.args.Composite
 PyTreeCheckpointHandler = ocp.PyTreeCheckpointHandler
@@ -394,6 +393,21 @@ def print_save_message(step, async_checkpointing):
     max_logging.log(f"Saved a checkpoint at step {step}.")
 
 
+def latest_step(checkpoint_manager):
+  """Latest saved step or None, across the v0 emergency manager and the v1 Checkpointer."""
+  return checkpoint_manager.latest_step()
+
+
+def wait_until_finished(checkpoint_manager):
+  """Blocks until pending saves finish, across the v0 emergency manager and the v1 Checkpointer."""
+  return checkpoint_manager.wait_until_finished()
+
+
+def reached_preemption(checkpoint_manager, step: int) -> bool:
+  """Whether a preemption sync point has been reached at `step`, across the v0 emergency manager and the v1 Checkpointer."""
+  return checkpoint_manager.reached_preemption(step)
+
+
 def load_state_if_possible(
     checkpoint_manager: CheckpointManager | None,
     data_iterator: MultiHostDataLoadIterator | list[MultiHostDataLoadIterator] | None,
@@ -446,7 +460,7 @@ def load_state_if_possible(
   if checkpoint_manager is not None:
     max_logging.log("checkpoint manager exists so trying to load this run's existing checkpoint")
 
-    step = checkpoint_manager.latest_step() if step < 0 else step  # pyrefly: ignore[bad-assignment]
+    step = latest_step(checkpoint_manager) if step < 0 else step  # pyrefly: ignore[bad-assignment]
     if step is not None:
       max_logging.log(f"restoring from this run's directory step {step}")
 
@@ -690,16 +704,18 @@ def _should_save_checkpoint_at_step(checkpoint_manager, step, config, force):
   else:
     base_checkpoint_due = step % config.checkpoint_period == 0
   local_checkpoint_due = _uses_local_checkpoint_period(config) and step % config.local_checkpoint_period == 0
-  autocheckpoint_due = config.enable_autocheckpoint and checkpoint_manager.reached_preemption(step)
+  autocheckpoint_due = config.enable_autocheckpoint and reached_preemption(checkpoint_manager, step)
   return base_checkpoint_due or local_checkpoint_due or autocheckpoint_due
 
 
 def _handle_post_checkpoint_preemption(checkpoint_manager, step, force_ckpt_save):
   """Waits on final/preemption saves and raises if preempted."""
-  reached_preemption = checkpoint_manager.reached_preemption(step)
-  if force_ckpt_save or reached_preemption:
-    checkpoint_manager.wait_until_finished()
-  if reached_preemption:
+  # Named is_preempted (not reached_preemption) so it doesn't shadow the module-level
+  # reached_preemption dispatcher we call below.
+  is_preempted = reached_preemption(checkpoint_manager, step)
+  if force_ckpt_save or is_preempted:
+    wait_until_finished(checkpoint_manager)
+  if is_preempted:
     raise exceptions.StopTraining("Job is preempted.")
 
 
@@ -733,7 +749,7 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step
     _handle_post_checkpoint_preemption(checkpoint_manager, actual_step, force_ckpt_save)
     return
 
-  if checkpoint_manager.latest_step() == actual_step:
+  if latest_step(checkpoint_manager) == actual_step:
     max_logging.log(f"Checkpoint for step {actual_step} already exists, skipping save.")
     return
 
@@ -782,7 +798,7 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator=
         force
         or (step % config.checkpoint_period == 0 and not config.enable_continuous_checkpointing)
         or (_uses_local_checkpoint_period(config) and step % config.local_checkpoint_period == 0)
-        or (config.enable_autocheckpoint and checkpoint_manager.reached_preemption(step))
+        or (config.enable_autocheckpoint and reached_preemption(checkpoint_manager, step))
     ):
       blocking_until_ready_start = time.time()
       max_logging.log(f"Waiting for step {step} to finish before checkpoint...")

src/maxtext/experimental/rl/grpo_trainer.py

@@ -1238,7 +1238,7 @@ def generation_worker_fn(
         checkpointing.maybe_save_checkpoint(checkpoint_manager, state_to_save, config, data_iterator)
       elif checkpoint_manager is not None:
         # in case the last checkpoint_period checkpoint is still in progress
-        checkpoint_manager.wait_until_finished()
+        checkpointing.wait_until_finished(checkpoint_manager)
     _job_completed_gracefully = True
   except exceptions.StopTraining as e:
     prof.deactivate()

src/maxtext/trainers/pre_train/train.py

@@ -901,7 +901,7 @@ def train_loop(config, recorder, state=None):
 
     if checkpoint_manager is not None:
       # in case the last checkpoint_period checkpoint is still in progress
-      checkpoint_manager.wait_until_finished()
+      checkpointing.wait_until_finished(checkpoint_manager)
     _job_completed_gracefully = True
   except exceptions.StopTraining as e:
     prof.deactivate()

src/maxtext/utils/elastic_utils.py

@@ -224,7 +224,13 @@ def maybe_elastic_scale_up(config, checkpoint_manager):
         " checkpoint to finish before interrupting."
     )
     if checkpoint_manager is not None:
-      checkpoint_manager.wait_until_finished()
+      # The v1 Checkpointer exposes `.wait()`, the v0 emergency/replicator
+      # managers expose `.wait_until_finished()`; this module cannot import
+      # `checkpointing`'s dispatcher (checkpointing imports elastic_utils).
+      if hasattr(checkpoint_manager, "wait"):
+        checkpoint_manager.wait()
+      else:
+        checkpoint_manager.wait_until_finished()
     max_logging.log("Checkpoint save completed. Interrupting")
     raise manager.ScaleUpSignalError()
 

src/maxtext/utils/generate_param_only_checkpoint.py

@@ -290,7 +290,7 @@ def _save_decode_checkpoint(config, state, checkpoint_manager):
   if checkpoint_manager is not None:
     if checkpointing.save_checkpoint(checkpoint_manager, 0, decode_state):
       max_logging.log(f"saved an decode checkpoint at {config.checkpoint_dir}")
-  checkpoint_manager.wait_until_finished()
+  checkpointing.wait_until_finished(checkpoint_manager)
 
 
 def _save_decode_checkpoint_nnx(config, state, checkpoint_manager):
@@ -318,7 +318,7 @@ def _wrap_value(node):
   if checkpoint_manager is not None:
     if checkpointing.save_checkpoint(checkpoint_manager, 0, bf16_model):
       max_logging.log(f"saved an NNX decode checkpoint at {config.checkpoint_dir}")
-    checkpoint_manager.wait_until_finished()
+    checkpointing.wait_until_finished(checkpoint_manager)
 
 
 def _possibly_unroll_lora_params_nnx(config, lora_state, lora_state_annotations, mesh):
@@ -402,7 +402,7 @@ def _save_lora_decode_checkpoint_nnx(config, lora_state, checkpoint_manager):
   if checkpoint_manager is not None:
     if checkpointing.save_checkpoint(checkpoint_manager, 0, decode_state):
       max_logging.log(f"saved a LoRA decode checkpoint at {config.checkpoint_dir}")
-    checkpoint_manager.wait_until_finished()
+    checkpointing.wait_until_finished(checkpoint_manager)
 
 
 def _generate_lora_decode_checkpoints_nnx(config, mesh):

src/maxtext/utils/rampup_batch.py

@@ -17,6 +17,8 @@
 
 import math
 
+from maxtext.common import checkpointing
+
 
 class RampupBatchManager:
   """
@@ -106,7 +108,7 @@ def create_rampup_manager(config, checkpoint_manager):
 
   # Current step default as -1 if no checkpoint exists
   current_step = -1
-  if checkpoint_manager and checkpoint_manager.latest_step():
-    current_step = checkpoint_manager.latest_step()
+  if checkpoint_manager and checkpointing.latest_step(checkpoint_manager):
+    current_step = checkpointing.latest_step(checkpoint_manager)
 
   return RampupBatchManager(config, current_step)

src/maxtext/utils/standalone_checkpointer.py

@@ -96,7 +96,7 @@ def init_state_fn():
       jax.experimental.multihost_utils.sync_global_devices("Barrier before save")
       state_to_save = train_state_nnx.to_linen_checkpoint_dict(state.to_pure_dict()) if config.pure_nnx else state
       if checkpointing.save_checkpoint(checkpoint_manager, int(step), state_to_save):
-        checkpoint_manager.wait_until_finished()
+        checkpointing.wait_until_finished(checkpoint_manager)
         end_time = datetime.datetime.now()
         if jax.process_index() == 0:
           max_logging.log(

src/maxtext/utils/train_utils.py

@@ -251,7 +251,7 @@ def create_train_state_fn():
       init_state_fn = partial(maxtext_utils.init_initial_state, model, tx, config, is_training, init_rng)
     checkpoint_manager = create_checkpoint_manager(config, mesh, init_state_fn)
     if checkpoint_manager is not None:
-      checkpoint_step = checkpoint_manager.latest_step()
+      checkpoint_step = checkpointing.latest_step(checkpoint_manager)
       if checkpoint_step is not None:
         validate_completed_steps(checkpoint_step + 1, config.steps)
 

tests/post_training/unit/lora_utils_test.py

@@ -421,7 +421,7 @@ def test_save_and_restore_metadata_integration(self):
       # Use save_checkpoint wrapper with a simple state
       dummy_state = {"weight": jnp.array([1.0, 2.0])}
       checkpointing.save_checkpoint(manager, step=0, state=dummy_state, config=cfg_save)
-      manager.wait_until_finished()
+      checkpointing.wait_until_finished(manager)
 
       # Now verify that the saved checkpoint contains metadata on disk
       checkpoint_dir = epath.Path(tmpdir) / "0"

tests/unit/checkpointing_nnx_missing_param_test.py

@@ -125,7 +125,7 @@ def _save_model(self):
     model = _Model(nnx.Rngs(0))
     state = train_state_nnx.TrainStateNNX(model, nnx.Optimizer(model, _TX, wrt=nnx.Param))
     checkpointing.maybe_save_checkpoint(manager, nnx.state(state), _config(), data_iterator=None, step=1)
-    manager.wait_until_finished()
+    checkpointing.wait_until_finished(manager)
     return manager
 
   def _load_overlay(self, manager, model_cls):

tests/unit/checkpointing_nnx_roundtrip_test.py

@@ -158,7 +158,7 @@ def _trained(self, model):
 
   def _save(self, manager, state, step=1):
     checkpointing.maybe_save_checkpoint(manager, nnx.state(state), _config(), data_iterator=None, step=step)
-    manager.wait_until_finished()
+    checkpointing.wait_until_finished(manager)
 
   def _init_state(self, model_cls, seed=123):
     """A fresh concrete init state (real weights/rng, optimizer zeros) for `model_cls`."""

tests/unit/elastic_utils_test.py

@@ -310,7 +310,7 @@ class FakeCheckpointManager:
       def __init__(self):
         self.wait_called = False
 
-      def wait_until_finished(self):
+      def wait(self):
         self.wait_called = True
 
     cm = FakeCheckpointManager()

tests/unit/setup_initial_state_nnx_test.py

@@ -158,7 +158,7 @@ def _train_and_save(self, manager, model_cls, config, seed=0):
     ts.apply_gradients(grads)
     saved = nnx.state(ts).to_pure_dict()
     checkpointing.maybe_save_checkpoint(manager, nnx.state(ts), config, data_iterator=None, step=1)
-    manager.wait_until_finished()
+    checkpointing.wait_until_finished(manager)
     return saved
 
   def test_restores_full_state_via_overlay(self):