Internal change

recml/core/training/keras_trainer.py

@@ -122,7 +122,8 @@ 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: int = 3,
+      legacy_checkpoint_format: bool | None = None,
   ):
     """Initializes the instance."""
 
@@ -153,24 +154,41 @@ 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 = 1 if legacy_checkpoint_format else 3
+    else:
+      self._checkpoint_version = 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:
+      if self._checkpoint_version == 1:
         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,
         )
-      else:
+      elif self._checkpoint_version == 2:
         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,
         )
+      elif self._checkpoint_version == 3:
+        checkpoint_manager = keras_utils.KerasOrbaxCheckpointManagerV3(
+            checkpoint_dir=self._checkpoint_dir,
+            max_to_keep=self._max_checkpoints_to_keep,
+            save_interval_epochs=self._checkpoint_save_interval_epochs,
+        )
+      else:
+        raise ValueError(
+            f"Unsupported checkpoint version: {self._checkpoint_version}"
+        )
 
       callbacks = [
           keras_utils.EpochSummaryCallback(
@@ -379,7 +397,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 +411,18 @@ def __init__(
         self._epoch = epoch
 
       def on_test_begin(self, logs: Mapping[str, Any] | None = None):
-        if use_legacy_checkpoint_format:
+        if checkpoint_version == 1:
           keras_utils.restore_keras_model(
               model, self._checkpoint_dir, step=self._epoch
           )
-        else:
+        elif checkpoint_version in (2, 3):
           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,27 +64,37 @@ 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": 3,
       },
       {
-          "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": 2,
       },
       {
-          "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": 1,
+      },
+      {
+          "testcase_name": "continuous_eval_v3",
+          "mode": core.Trainer.Mode.CONTINUOUS_EVAL,
+          "checkpoint_version": 3,
+      },
+      {
+          "testcase_name": "continuous_eval_v2",
+          "mode": core.Trainer.Mode.CONTINUOUS_EVAL,
+          "checkpoint_version": 2,
       },
       {
-          "testcase_name": "continuous_eval_legacy_checkpoint_format",
+          "testcase_name": "continuous_eval_v1",
           "mode": core.Trainer.Mode.CONTINUOUS_EVAL,
-          "legacy_checkpoint_format": True,
+          "checkpoint_version": 1,
       },
   )
-  def test_keras_task_and_trainer(
-      self, mode: str, legacy_checkpoint_format: bool = False
-  ):
+  def test_keras_task_and_trainer(self, mode: str, checkpoint_version: int = 3):
     if keras.backend.backend() == "jax":
       distribution = keras.distribution.DataParallel()
     else:
@@ -98,13 +109,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)