Add max_training_steps option to Kauldron evaluations

PiperOrigin-RevId: 914132571

Author: jrruijli

kauldron/evals/eval_impl.py

@@ -176,8 +176,9 @@ def _preemptable_iter_new_checkpoints(
     state: train_step.TrainState,
 ) -> Iterator[train_step.TrainState]:
   """Yields the new checkpoints."""
-  # Skip the `iter_new_checkpoints` for eval-only jobs.
-  if trainer.setup.eval_only:
+  # Skip the `iter_new_checkpoints` for eval-only jobs ONLY IF
+  # max_training_steps is not set.
+  if trainer.setup.eval_only and trainer.setup.max_training_steps is None:
     return
 
   trainer_ckpt = trainer.checkpointer
@@ -196,6 +197,17 @@ def _preemptable_iter_new_checkpoints(
       # state might have been donated, we should not access it after this point.
       # Eval is done, remove the duplicated checkpoint
       eval_ckpt.delete(step)
+      # Exit if resumed step meets or exceeds limit
+      if (
+          trainer.setup.max_training_steps is not None
+          and step >= trainer.setup.max_training_steps
+      ):
+        logging.info(
+            f'Resumed step {step} reaches max_training_steps'
+            f' {trainer.setup.max_training_steps}. Exiting.'
+        )
+        return
+
     for step in trainer_ckpt.iter_new_checkpoints(
         min_interval_secs=10,
         timeout=10,
@@ -222,6 +234,17 @@ def _preemptable_iter_new_checkpoints(
       if trainer.setup.preemptable_eval:
         eval_ckpt.delete(step)
 
+      # Simply break the loop if we have evaluated the max step
+      if (
+          trainer.setup.max_training_steps is not None
+          and step >= trainer.setup.max_training_steps
+      ):
+        logging.info(
+            f'Reached max_training_steps {trainer.setup.max_training_steps}.'
+            ' Stopping.'
+        )
+        break
+
 
 def _restore_checkpoint(
     *,

kauldron/train/setup_utils.py

@@ -50,6 +50,7 @@ class Setup:
       `Writer`, to ensure all dashboards are written to the same collection.
     eval_only: Whether the job is a eval-only job.
     preemptable_eval: If `True`, makes eval preemptible.
+    max_training_steps: If set, limits the evaluation to only evaluate checkpoints up to this step.
   """
   # pyformat: enable
 
@@ -60,6 +61,7 @@ class Setup:
   eval_only: bool = False
   # TODO(epot): Not the best place for this.
   preemptable_eval: bool = False
+  max_training_steps: int | None = None
 
   def __post_init__(self):
     # Normalize tags to a list.

kauldron/xm/_src/kauldron_utils.py

@@ -215,6 +215,10 @@ def _resolve_run_konfig(
       # Cannot use `dataclasses.replace` as it interact with the `merge_utils`
       run = type(run)(**init_kwargs)
 
+    max_training_steps = kontext.get_by_path(
+        self._cfg, "setup.max_training_steps", None
+    )
+
     if self._is_eval_only:  # Eval only, normalize run config.
       if isinstance(run, run_strategies.AlongTrain):
         run = run_strategies.StandaloneLastCheckpoint(
@@ -226,10 +230,16 @@ def _resolve_run_konfig(
         init_kwargs = dict(run._kxm_init_kwargs)  # pylint: disable=protected-access
         init_kwargs.pop("job_group", None)  # Avoid duplicated kwarg.
 
-        run = run_strategies.StandaloneLastCheckpoint(
-            job_group=run.job_group,
-            **init_kwargs,
-        )
+        if max_training_steps is not None:
+          run = run_strategies.StandaloneEveryCheckpoint(
+              job_group=run.job_group,
+              **init_kwargs,
+          )
+        else:
+          run = run_strategies.StandaloneLastCheckpoint(
+              job_group=run.job_group,
+              **init_kwargs,
+          )
       else:
         raise TypeError(
             f"Unexpected run strategy for {eval_name}. Got: {type(run)}."