Use helper functions for CheckpointManager operations.

Refactors code to interact with the CheckpointManager through dedicated helper functions instead of direct method calls. This is necessary because:
- The underlying CheckpointManager operations (`.wait_until_finished`, `.latest_step`, `.reached_preemption`) have different semantics in v1 compared to v0.
- We need to ensure the v0 emergency checkpointer functionality remains intact.

This abstraction layer isolates the core logic from the specific CheckpointManager version

PiperOrigin-RevId: 953793325

Author: angel-core

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()