Internal change

PiperOrigin-RevId: 964278984

Author: RecML authors

recml/core/utils/keras_utils.py

@@ -74,13 +74,15 @@ def __init__(
       checkpoint_dir: str,
       max_to_keep: int = 5,
       save_interval_epochs: int = 1,
+      choose_store_cell: bool = True,
   ):
     """Initializes a KerasOrbaxCheckpointManager.
 
     Args:
       checkpoint_dir: The directory to save checkpoints to.
       max_to_keep: The maximum number of checkpoints to keep.
       save_interval_epochs: The interval (in epochs) to save checkpoints.
+      choose_store_cell: Whether to dynamically select the CNS2 store cell.
     """
     if keras.backend.backend() != "jax":
       raise ValueError(
@@ -92,6 +94,11 @@ def __init__(
         options=ocp.CheckpointManagerOptions(
             save_interval_steps=save_interval_epochs,
             max_to_keep=max_to_keep,
+            file_options=ocp.options.FileOptions(
+                cns2_storage_options=ocp.options.Cns2StorageOptions(
+                    choose_store_cell=choose_store_cell,
+                ),
+            ),
         ),
     )
 
@@ -430,20 +437,27 @@ def __init__(
       checkpoint_dir: str,
       max_to_keep: int = 5,
       save_interval_epochs: int = 1,
+      choose_store_cell: bool = True,
   ):
     """Initializes a KerasOrbaxCheckpointManager.
 
     Args:
       checkpoint_dir: The directory to save checkpoints to.
       max_to_keep: The maximum number of checkpoints to keep.
       save_interval_epochs: The interval (in epochs) to save checkpoints.
+      choose_store_cell: Whether to dynamically select the CNS2 store cell.
     """
     super().__init__(
         directory=checkpoint_dir,
         checkpointers=ocp.AsyncCheckpointer(ocp.PyTreeCheckpointHandler()),
         options=ocp.CheckpointManagerOptions(
             save_interval_steps=save_interval_epochs,
             max_to_keep=max_to_keep,
+            file_options=ocp.options.FileOptions(
+                cns2_storage_options=ocp.options.Cns2StorageOptions(
+                    choose_store_cell=choose_store_cell,
+                ),
+            ),
         ),
     )