Save variables with keys in the checkpoint

PiperOrigin-RevId: 756065918

Author: edwardzhou130

recml/core/training/keras_trainer.py

@@ -118,6 +118,7 @@ def __init__(
       max_checkpoints_to_keep: int = 5,
       checkpoint_save_interval_epochs: int = 1,
       rng_seed: int = core.DEFAULT_RNG_SEED,
+      legacy_format: bool = True,
   ):
     """Initializes the instance."""
 
@@ -148,12 +149,14 @@ def __init__(
         model_dir, core.TRAINING_COMPLETE_MARKER_FILE
     )
     self._checkpoint_dir = os.path.join(model_dir, core.CHECKPOINT_DIR)
+    self._legacy_format = legacy_format
 
     if keras.backend.backend() == "jax":
       self._checkpoint_manager = keras_utils.KerasOrbaxCheckpointManager(
           checkpoint_dir=self._checkpoint_dir,
           max_to_keep=max_checkpoints_to_keep,
           save_interval_epochs=checkpoint_save_interval_epochs,
+          legacy_format=self._legacy_format,
       )
       self._train_callbacks = [
           keras_utils.EpochSummaryCallback(
@@ -314,13 +317,18 @@ def __init__(
           self,
           checkpoint_dir: str,
           epoch: int,
+          legacy_format: bool,
       ):
         self._checkpoint_dir = checkpoint_dir
         self._epoch = epoch
+        self._legacy_format = legacy_format
 
       def on_test_begin(self, logs: Mapping[str, Any] | None = None):
         keras_utils.restore_keras_model(
-            model, self._checkpoint_dir, step=self._epoch
+            model,
+            self._checkpoint_dir,
+            step=self._epoch,
+            legacy_format=self._legacy_format,
         )
 
     history = None
@@ -329,7 +337,9 @@ def on_test_begin(self, logs: Mapping[str, Any] | None = None):
         timeout=self._continuous_eval_timeout,
         timeout_fn=timeout_fn,
     ):
-      restore_callback = _RestoreCallback(self._checkpoint_dir, epoch)
+      restore_callback = _RestoreCallback(
+          self._checkpoint_dir, epoch, self._legacy_format
+      )
       [tb_cbk] = [
           cbk
           for cbk in self._eval_callbacks

recml/core/training/keras_trainer_test.py

@@ -94,6 +94,22 @@ def test_keras_task_and_trainer(self, mode: str):
     ):
       self.assertEqual(history.history["num_params/trainable"][0], 2)
 
+  def test_new_checkpoint_format(self):
+    if keras.backend.backend() != "jax":
+      self.skipTest("Only supported on the Jax backend.")
+    trainer = keras_trainer.KerasTrainer(
+        distribution=keras.distribution.DataParallel(),
+        train_steps=5,
+        steps_per_eval=3,
+        steps_per_loop=2,
+        model_dir=self.create_tempdir().full_path,
+        continuous_eval_timeout=5,
+        legacy_format=False,
+    )
+    experiment = core.Experiment(_KerasTask(), trainer)
+    core.run_experiment(experiment, core.Experiment.Mode.TRAIN)
+    core.run_experiment(experiment, core.Experiment.Mode.CONTINUOUS_EVAL)
+
 
 if __name__ == "__main__":
   absltest.main()