Use the new preservation policy in CheckpointManager.

PiperOrigin-RevId: 750914697

Author: abhishek002002

checkpoint/orbax/checkpoint/BUILD

@@ -139,6 +139,7 @@ py_library(
         "//checkpoint/orbax/checkpoint/_src/path:step",
         "//checkpoint/orbax/checkpoint/_src/path:utils",
         "//orbax/checkpoint/_src:threading",
+        "//orbax/checkpoint/_src/checkpoint_managers:preservation_policy",
     ],
 )
 
@@ -353,5 +354,6 @@ py_library(
         ":abstract_checkpoint_manager",
         ":checkpoint_manager",
         "//checkpoint/orbax/checkpoint/_src/checkpoint_managers:save_decision_policy",
+        "//orbax/checkpoint/_src/checkpoint_managers:preservation_policy",
     ],
 )

checkpoint/orbax/checkpoint/checkpoint_manager.py

@@ -35,6 +35,7 @@
 from orbax.checkpoint import options as options_lib
 from orbax.checkpoint import utils
 from orbax.checkpoint._src import threading as threading_lib
+from orbax.checkpoint._src.checkpoint_managers import preservation_policy as preservation_policy_lib
 from orbax.checkpoint._src.checkpoint_managers import save_decision_policy as save_decision_policy_lib
 from orbax.checkpoint._src.checkpointers import abstract_checkpointer
 from orbax.checkpoint._src.checkpointers import async_checkpointer
@@ -206,6 +207,60 @@ def _get_default_save_decision_policy(
   return save_decision_policy_lib.AnySavePolicy(save_interval_policies)
 
 
+@dataclasses.dataclass
+class _ShouldKeepFnPolicy(preservation_policy_lib.PreservationPolicy):
+  """Return true based on a provided function of the step."""
+  should_keep_fn: Callable[[int], bool]
+
+  def should_preserve(
+      self,
+      checkpoints: Sequence[checkpoint_info.CheckpointInfo],
+      *,
+      context: preservation_policy_lib.PreservationContext,
+  ) -> list[bool]:
+    return [self.should_keep_fn(ckpt.step) for ckpt in checkpoints]
+
+
+def _get_default_preservation_policy(
+    options: CheckpointManagerOptions,
+) -> preservation_policy_lib.PreservationPolicy:
+  """Returns a default preservation policy."""
+  # Must have set max_to_keep in order to remove any checkpoints.
+  preservation_policies = []
+  if options.keep_period is not None:
+    preservation_policies.append(
+        preservation_policy_lib.EveryNSteps(options.keep_period)
+    )
+  if options.should_keep_fn is not None:
+    preservation_policies.append(
+        _ShouldKeepFnPolicy(
+            should_keep_fn=options.should_keep_fn
+        )
+    )
+  if options.keep_time_interval is not None:
+    total_seconds = int(options.keep_time_interval.total_seconds())
+    preservation_policies.append(
+        preservation_policy_lib.EveryNSeconds(
+            interval_secs=total_seconds
+        )
+    )
+  if options.best_fn is not None:
+    preservation_policies.append(
+        preservation_policy_lib.BestN(
+            best_fn=options.best_fn,
+            reverse=(options.best_mode == 'min'),
+            n=options.max_to_keep,
+        )
+    )
+  else:
+    preservation_policies.append(
+        preservation_policy_lib.LatestN(n=options.max_to_keep)
+    )
+  return preservation_policy_lib.AnyPreservationPolicy(
+      preservation_policies
+  )
+
+
 # TODO(b/268051457) Clean up when no longer depended upon by internal users.
 def is_async_checkpointer(checkpointer: AbstractCheckpointer):
   return isinstance(
@@ -319,6 +374,12 @@ class CheckpointManagerOptions:
     is the sole means of determining when a checkpoint should be saved. If not
     provided, these other options are used instead. Prefer to use this option
     over others.
+  preservation_policy: An object used to determine which checkpoints to
+    preserve. If provided, overrides any other options dealing with this
+    subject, including `max_to_keep`, `keep_time_interval`, `keep_period`, and
+    `should_keep_fn`, `best_fn`, and is the sole means of determining which
+    checkpoints to preserve. If not provided, these other options are used 
+    instead. Prefer to use this option over others.
   """
 
   save_interval_steps: int = 1
@@ -351,6 +412,9 @@ class CheckpointManagerOptions:
   save_decision_policy: Optional[
       save_decision_policy_lib.SaveDecisionPolicy
   ] = None
+  preservation_policy: Optional[
+      preservation_policy_lib.PreservationPolicy
+  ] = None
 
   def __post_init__(self):
     step_name_format_single_host_load_and_broadcast = (
@@ -632,6 +696,10 @@ def __init__(
         self._options.save_decision_policy
         or _get_default_save_decision_policy(self._options)
     )
+    self._preservation_policy = (
+        self._options.preservation_policy
+        or _get_default_preservation_policy(self._options)
+    )
 
     if self._options.best_mode not in ['min', 'max']:
       raise ValueError('`best_mode` must be one of: "min", "max"')
@@ -1151,17 +1219,12 @@ def delete(self, step: int):
 
     Args:
       step: The step to delete.
-
-    Raises:
-      FileNotFoundError: If the step does not exist.
     """
     if self._options.read_only:
       logging.warning('%s is read only, delete will be skipped', self.directory)
       return
     if step not in self.all_steps():
-      raise FileNotFoundError(
-          f'Requested deleting a non-existent step: {step}.'
-      )
+      raise ValueError(f'Requested deleting a non-existent step: {step}.')
     self._checkpoint_deleter.delete(step)
     multihost.sync_global_processes(
         multihost.unique_barrier_key(
@@ -1704,22 +1767,6 @@ def build_checkpoint_info(step_metadata):
       )
       return checkpoint_infos
 
-  def _get_interval_preserved_checkpoints(
-      self, checkpoints: checkpoint_info.CheckpointInfos
-  ) -> List[CheckpointInfo]:
-    """Gets which checkpoints should be kept based on keep_time_interval."""
-    if checkpoints.empty():
-      return []
-    interval_preserved_checkpoints = [checkpoints[0]]
-    if self._options.keep_time_interval is not None:
-      for info in checkpoints[1:]:
-        if info.time >= (
-            interval_preserved_checkpoints[-1].time
-            + self._options.keep_time_interval
-        ):
-          interval_preserved_checkpoints.append(info)
-    return interval_preserved_checkpoints
-
   def _add_checkpoint_info(self, step: int, metrics: Optional[PyTree]):
     self._checkpoints.append(
         CheckpointInfo(
@@ -1866,102 +1913,21 @@ def _cleanup_tmp_directories(self):
 
   def _get_old_steps_to_remove(self) -> List[int]:
     """Returns checkpoints that should be deleted."""
-    # Must have set max_to_keep in order to remove any checkpoints.
-    if self._options.max_to_keep is None:
-      return []
-    # Not enough checkpoints accumulated to consider deletion.
-    if self._checkpoints.size() <= self._options.max_to_keep:
-      return []
-
     # This isn't a duration but there isn't a general counter that we can use so
     # we abuse a duration metric to count the number of steps examined.
     jax.monitoring.record_event_duration_secs(
         '/jax/checkpoint/write/old_steps_examined_count',
         self._checkpoints.size(),
     )
-
-    if self._track_best:
-      # Best steps (to keep) are at the end, after sorting.
-      (
-          checkpoints_without_metrics,
-          sorted_checkpoints,
-      ) = self._sort_checkpoints_by_metrics(self._checkpoints)
-    else:
-      # checkpoints already sorted by ascending step
-      checkpoints_without_metrics = []
-      sorted_checkpoints = [info for info in self._checkpoints]
-
-    keep = int(self._options.max_to_keep)
-    if self._options.keep_checkpoints_without_metrics:
-      maybe_delete = (
-          sorted_checkpoints[:-keep] if keep > 0 else sorted_checkpoints
-      )
-      active_checkpoints = set(
-          checkpoints_without_metrics
-          + (sorted_checkpoints[-keep:] if keep > 0 else [])
-      )
-    else:
-      all_checkpoints = checkpoints_without_metrics + sorted_checkpoints
-      maybe_delete = all_checkpoints[:-keep] if keep > 0 else sorted_checkpoints
-      active_checkpoints = set(all_checkpoints[-keep:] if keep > 0 else [])
-
-    interval_preserved_checkpoints = self._get_interval_preserved_checkpoints(
-        self._checkpoints
+    preservation_result = self._preservation_policy.should_preserve(
+        [info for info in self._checkpoints],
+        context=preservation_policy_lib.PreservationContext(),
     )
-    kept_checkpoints = set()
-    for info in maybe_delete:
-      if (
-          self._options.keep_time_interval is not None
-          and interval_preserved_checkpoints
-      ):
-        if info in interval_preserved_checkpoints:
-          logging.info(
-              'Preserving %s: (Reason: older falling on keep_time_interval).',
-              info,
-          )
-          kept_checkpoints.add(info)
-          continue
-        elif info.time >= (
-            interval_preserved_checkpoints[-1].time
-            + self._options.keep_time_interval
-        ):
-          interval_preserved_checkpoints.append(info)
-          logging.info(
-              'Preserving %s: (Reason: latest falling on keep_time_interval).',
-              info,
-          )
-          kept_checkpoints.add(info)
-          continue
-      if (
-          self._options.should_keep_fn is not None
-          and self._options.should_keep_fn(info.step)
-      ):
-        logging.info(
-            'Preserving %s: (Reason: on should_keep_fn callback).', info
-        )
-        kept_checkpoints.add(info)
-        continue
-      if (
-          self._options.keep_period is not None
-          and info.step % self._options.keep_period == 0
-      ):
-        logging.info(
-            'Preserving %s: (Reason: on keep_period=%s).',
-            info,
-            self._options.keep_period,
-        )
-        kept_checkpoints.add(info)
-        continue
-
-    kept_checkpoints.update(active_checkpoints)
-
-    steps_to_remove = []
-    for info in self._checkpoints:
-      if info not in kept_checkpoints:
-        reason = 'worse metric' if self._track_best else 'old checkpoint'
-        logging.info('Deleting %s: (Reason: %s).', info, reason)
-        steps_to_remove.append(info.step)
-    return steps_to_remove
+    result = []
+    for i in range(len(self._checkpoints)):
+      if not preservation_result[i]:
+        result.append(self._checkpoints[i].step)
+    return result
 
   def _wait_for_checkpointers(self):
     if is_async_checkpointer(self._checkpointer):

checkpoint/orbax/checkpoint/checkpoint_managers.py

@@ -26,6 +26,18 @@
     AnySavePolicy,
 )
 
+from orbax.checkpoint._src.checkpoint_managers import preservation_policy
+from orbax.checkpoint._src.checkpoint_managers.preservation_policy import (
+    PreservationPolicy,
+    LatestN,
+    EveryNSeconds,
+    EveryNSteps,
+    CustomSteps,
+    AnyPreservationPolicy,
+    BestN,
+)
+
+
 from orbax.checkpoint.checkpoint_manager import CheckpointManagerOptions
 from orbax.checkpoint.checkpoint_manager import CheckpointManager