Allow configuring the minimum interval seconds for continuous checkpointing.

checkpoint/orbax/checkpoint/_src/checkpoint_managers/save_decision_policy.py

@@ -15,6 +15,7 @@
 """Defines policies for when a checkpoint is saved."""
 
 import dataclasses
+import datetime
 import typing
 from typing import Container, Protocol, Sequence
 
@@ -23,6 +24,7 @@
 class StepInfo:
   """Relevant information about a checkpoint step."""
   step: int
+  time: datetime.datetime
 
 
 @dataclasses.dataclass(kw_only=True)
@@ -87,19 +89,26 @@ def should_save(
     return step.step in self.steps
 
 
+@dataclasses.dataclass(kw_only=True)
 class ContinuousCheckpointingPolicy(SaveDecisionPolicy):
   """Checkpoint as often as possible, as long as a save is not ongoing."""
 
+  minimum_interval_secs: int | None = None
+
   def should_save(
       self,
       step: StepInfo,
       previous_steps: Sequence[StepInfo],
       *,
       context: DecisionContext
   ) -> bool:
-    del step
-    del previous_steps
-    return not context.is_saving_in_progress
+    if context.is_saving_in_progress:
+      return False
+    if not previous_steps or self.minimum_interval_secs is None:
+      return True
+    return step.time - previous_steps[-1].time >= datetime.timedelta(
+        seconds=self.minimum_interval_secs
+    )
 
 
 class PreemptionCheckpointingPolicy(SaveDecisionPolicy):

checkpoint/orbax/checkpoint/checkpoint_manager.py

@@ -1084,10 +1084,12 @@ def should_save(self, step: int) -> bool:
     is_saving_in_progress = self.is_saving_in_progress()
     reached_preemption = self.reached_preemption(step)
     previous_step_infos = [
-        save_decision_policy_lib.StepInfo(step=ckpt.step)
+        save_decision_policy_lib.StepInfo(step=ckpt.step, time=ckpt.time)
         for ckpt in self._checkpoints
     ]
-    current_step_info = save_decision_policy_lib.StepInfo(step=step)
+    current_step_info = save_decision_policy_lib.StepInfo(
+        step=step, time=datetime.datetime.now(tz=datetime.timezone.utc),
+    )
     context = save_decision_policy_lib.DecisionContext(
         is_saving_in_progress=is_saving_in_progress,
         reached_preemption=reached_preemption,