Internal change

PiperOrigin-RevId: 940694716

Author: RecML authors

recml/core/training/keras_trainer.py

@@ -122,9 +122,32 @@ def __init__(
       max_checkpoints_to_keep: int = 5,
       checkpoint_save_interval_epochs: int = 1,
       rng_seed: int = core.DEFAULT_RNG_SEED,
-      legacy_checkpoint_format: bool = True,
+      checkpoint_version: keras_utils.CheckpointVersion | str = "v2",
+      legacy_checkpoint_format: bool | None = None,
   ):
-    """Initializes the instance."""
+    """Initializes the instance.
+
+    Args:
+      distribution: The distribution strategy to use.
+      model_dir: The directory to save checkpoints and logs.
+      train_steps: Total number of training steps.
+      steps_per_eval: Number of steps between evaluations.
+      continuous_eval_timeout: Timeout for continuous evaluation.
+      steps_per_loop: Number of steps per training loop.
+      max_checkpoints_to_keep: Maximum number of checkpoints to keep.
+      checkpoint_save_interval_epochs: Interval in epochs to save checkpoints.
+      rng_seed: Random seed.
+      checkpoint_version: The checkpoint version to use. Supported versions:
+        "v1" (legacy V1), "v2" (V2, default), "v3" (V3).
+      legacy_checkpoint_format: Deprecated. Use checkpoint_version instead. If
+        set, True maps to V1, and False maps to V2. TODO(b/542602169): Remove
+        this in v2.
+      resume_training_launcher: Launcher to resume training.
+      enable_xmanager_measurements: Whether to enable XManager measurements.
+      enable_autoxprof: Whether to enable AutoXprof.
+      autoxprof_settings: Settings for AutoXprof.
+      s2_logging_settings: Settings for S2 logging.
+    """
 
     keras.utils.set_random_seed(rng_seed)
 
@@ -153,25 +176,42 @@ def __init__(
     self._checkpoint_dir = os.path.join(model_dir, core.CHECKPOINT_DIR)
     self._max_checkpoints_to_keep = max_checkpoints_to_keep
     self._checkpoint_save_interval_epochs = checkpoint_save_interval_epochs
-    self._legacy_checkpoint_format = legacy_checkpoint_format
+    if legacy_checkpoint_format is not None:
+      logging.warning(
+          "legacy_checkpoint_format is deprecated, use checkpoint_version"
+          " instead."
+      )
+      self._checkpoint_version = (
+          keras_utils.CheckpointVersion.V1
+          if legacy_checkpoint_format
+          else keras_utils.CheckpointVersion.V2
+      )
+    else:
+      self._checkpoint_version = keras_utils.CheckpointVersion(
+          checkpoint_version
+      )
 
   @functools.cached_property
   def train_callbacks(self) -> list[keras.callbacks.Callback]:
     """Returns the training callbacks."""
     if keras.backend.backend() == "jax":
-      if self._legacy_checkpoint_format:
-        checkpoint_manager = keras_utils.KerasOrbaxCheckpointManager(
-            checkpoint_dir=self._checkpoint_dir,
-            max_to_keep=self._max_checkpoints_to_keep,
-            save_interval_epochs=self._checkpoint_save_interval_epochs,
-        )
+      if self._checkpoint_version == keras_utils.CheckpointVersion.V1:
+        manager_cls = keras_utils.KerasOrbaxCheckpointManager
+      elif self._checkpoint_version == keras_utils.CheckpointVersion.V2:
+        manager_cls = keras_utils.KerasOrbaxCheckpointManagerV2
+      elif self._checkpoint_version == keras_utils.CheckpointVersion.V3:
+        manager_cls = keras_utils.KerasOrbaxCheckpointManagerV3
       else:
-        checkpoint_manager = keras_utils.KerasOrbaxCheckpointManagerV2(
-            checkpoint_dir=self._checkpoint_dir,
-            max_to_keep=self._max_checkpoints_to_keep,
-            save_interval_epochs=self._checkpoint_save_interval_epochs,
+        raise ValueError(
+            f"Unsupported checkpoint version: {self._checkpoint_version}"
         )
 
+      checkpoint_manager = manager_cls(
+          checkpoint_dir=self._checkpoint_dir,
+          max_to_keep=self._max_checkpoints_to_keep,
+          save_interval_epochs=self._checkpoint_save_interval_epochs,
+      )
+
       callbacks = [
           keras_utils.EpochSummaryCallback(
               log_dir=os.path.join(self._model_dir, core.LOG_DIR),
@@ -379,7 +419,7 @@ def timeout_fn() -> bool:
     else:
       steps_msg = "running complete evaluation..."
 
-    use_legacy_checkpoint_format = self._legacy_checkpoint_format
+    checkpoint_version = self._checkpoint_version
 
     class _RestoreCallback(keras.callbacks.Callback):
       """Callback for restoring the model from the latest checkpoint."""
@@ -393,14 +433,21 @@ def __init__(
         self._epoch = epoch
 
       def on_test_begin(self, logs: Mapping[str, Any] | None = None):
-        if use_legacy_checkpoint_format:
+        if checkpoint_version == keras_utils.CheckpointVersion.V1:
           keras_utils.restore_keras_model(
               model, self._checkpoint_dir, step=self._epoch
           )
-        else:
+        elif checkpoint_version in (
+            keras_utils.CheckpointVersion.V2,
+            keras_utils.CheckpointVersion.V3,
+        ):
           keras_utils.restore_keras_checkpoint(
               self._checkpoint_dir, model=model, epoch=self._epoch
           )
+        else:
+          raise ValueError(
+              f"Unsupported checkpoint version: {checkpoint_version}"
+          )
 
     history = None
     for epoch in ocp.checkpoint_utils.checkpoints_iterator(

recml/core/training/keras_trainer_test.py

@@ -55,6 +55,7 @@ class KerasTrainerTest(parameterized.TestCase):
 
   def setUp(self):
     super().setUp()
+    keras.backend.clear_session()
     # Workaround to make `create_tempdir` work with pytest.
     if not flags.FLAGS.is_parsed():
       flags.FLAGS.mark_as_parsed()
@@ -63,26 +64,38 @@ def setUp(self):
       {"testcase_name": "train", "mode": core.Trainer.Mode.TRAIN},
       {"testcase_name": "eval", "mode": core.Trainer.Mode.EVAL},
       {
-          "testcase_name": "train_and_eval",
+          "testcase_name": "train_and_eval_v3",
           "mode": core.Trainer.Mode.TRAIN_AND_EVAL,
+          "checkpoint_version": "v3",
       },
       {
-          "testcase_name": "continuous_eval_",
-          "mode": core.Trainer.Mode.CONTINUOUS_EVAL,
+          "testcase_name": "train_and_eval_v2",
+          "mode": core.Trainer.Mode.TRAIN_AND_EVAL,
+          "checkpoint_version": "v2",
       },
       {
-          "testcase_name": "train_and_eval_legacy_checkpoint_format",
+          "testcase_name": "train_and_eval_v1",
           "mode": core.Trainer.Mode.TRAIN_AND_EVAL,
-          "legacy_checkpoint_format": True,
+          "checkpoint_version": "v1",
+      },
+      {
+          "testcase_name": "continuous_eval_v3",
+          "mode": core.Trainer.Mode.CONTINUOUS_EVAL,
+          "checkpoint_version": "v3",
+      },
+      {
+          "testcase_name": "continuous_eval_v2",
+          "mode": core.Trainer.Mode.CONTINUOUS_EVAL,
+          "checkpoint_version": "v2",
       },
       {
-          "testcase_name": "continuous_eval_legacy_checkpoint_format",
+          "testcase_name": "continuous_eval_v1",
           "mode": core.Trainer.Mode.CONTINUOUS_EVAL,
-          "legacy_checkpoint_format": True,
+          "checkpoint_version": "v1",
       },
   )
   def test_keras_task_and_trainer(
-      self, mode: str, legacy_checkpoint_format: bool = False
+      self, mode: str, checkpoint_version: str = "v3"
   ):
     if keras.backend.backend() == "jax":
       distribution = keras.distribution.DataParallel()
@@ -98,13 +111,14 @@ def test_keras_task_and_trainer(
         steps_per_loop=2,
         model_dir=self.create_tempdir().full_path,
         continuous_eval_timeout=5,
-        legacy_checkpoint_format=legacy_checkpoint_format,
+        checkpoint_version=checkpoint_version,
     )
     experiment = core.Experiment(_KerasTask(), trainer)
 
     if mode == core.Trainer.Mode.CONTINUOUS_EVAL:
       # Produce one checkpoint so there is something to evaluate.
       core.run_experiment(experiment, core.Trainer.Mode.TRAIN)
+      keras.backend.clear_session()
 
     history = core.run_experiment(experiment, mode)