Allow skipping tmp checkpoint saving during eval.

PiperOrigin-RevId: 729482348

kauldron/evals/eval_impl.py

@@ -169,6 +169,11 @@ def _preemptable_iter_new_checkpoints(
 ) -> Iterator[train_step.TrainState]:
   """Yields the new checkpoints."""
   # Skip the `iter_new_checkpoints` for eval-only jobs.
+  if 'save_tmp_ckpt' in trainer.aux:
+    save_tmp_ckpt = trainer.aux['save_tmp_ckpt']
+  else:
+    save_tmp_ckpt = True
+
   if trainer.setup.eval_only:
     return
 
@@ -206,11 +211,13 @@ def _preemptable_iter_new_checkpoints(
       assert int(state.step) == step
       # Temporarily copy the state to the eval checkpoint, to ensure that
       # it won't be deleted by the train job until the current eval is done.
-      eval_ckpt.save(state, step=step)
+      if save_tmp_ckpt:
+        eval_ckpt.save(state, step=step)
       yield state
       # state might have been donated, we should not access it after this point.
       # Eval is done, remove the duplicated checkpoint
-      eval_ckpt.delete(step)
+      if save_tmp_ckpt:
+        eval_ckpt.delete(step)
 
 
 def _restore_checkpoint(