Add SaveDecisionPolicy and PreservationPolicy to v1, use CheckpointMetadata instead of CheckpointInfo.

PiperOrigin-RevId: 751099231

Author: BlaziusMaximus

checkpoint/orbax/checkpoint/_src/checkpoint_managers/BUILD

@@ -5,18 +5,13 @@ py_library(
     srcs = ["save_decision_policy.py"],
     deps = [
         "//checkpoint/orbax/checkpoint:options",
-        "//checkpoint/orbax/checkpoint/_src/metadata:checkpoint_info",
         "//checkpoint/orbax/checkpoint/_src/multihost",
     ],
 )
 
 py_library(
     name = "preservation_policy",
     srcs = ["preservation_policy.py"],
-    deps = [
-        "//checkpoint/orbax/checkpoint:options",
-        "//checkpoint/orbax/checkpoint/_src/metadata:checkpoint_info",
-    ],
 )
 
 py_test(

checkpoint/orbax/checkpoint/_src/checkpoint_managers/preservation_policy.py

@@ -18,7 +18,6 @@
 import datetime
 from typing import Any, Callable, Dict, Protocol, Sequence, Set
 import numpy as np
-from orbax.checkpoint._src.metadata import checkpoint_info
 
 
 NestedDict = Dict[str, Any]
@@ -36,7 +35,7 @@ class PreservationPolicy(Protocol):
 
   def should_preserve(
       self,
-      checkpoints: Sequence[checkpoint_info.CheckpointInfo],
+      checkpoints,
       *,
       context: PreservationContext,
   ) -> Sequence[bool]:
@@ -52,7 +51,7 @@ class LatestN(PreservationPolicy):
 
   def should_preserve(
       self,
-      checkpoints: Sequence[checkpoint_info.CheckpointInfo],
+      checkpoints,
       *,
       context: PreservationContext,
   ) -> Sequence[bool]:
@@ -69,7 +68,7 @@ class EveryNSeconds(PreservationPolicy):
 
   def should_preserve(
       self,
-      checkpoints: Sequence[checkpoint_info.CheckpointInfo],
+      checkpoints,
       *,
       context: PreservationContext,
   ) -> Sequence[bool]:
@@ -96,7 +95,7 @@ class EveryNSteps(PreservationPolicy):
 
   def should_preserve(
       self,
-      checkpoints: Sequence[checkpoint_info.CheckpointInfo],
+      checkpoints,
       *,
       context: PreservationContext,
   ) -> Sequence[bool]:
@@ -118,7 +117,7 @@ def __post_init__(self, steps_init: Sequence[int]):
 
   def should_preserve(
       self,
-      checkpoints: Sequence[checkpoint_info.CheckpointInfo],
+      checkpoints,
       *,
       context: PreservationContext,
   ) -> Sequence[bool]:
@@ -133,7 +132,7 @@ class AnyPreservationPolicy(PreservationPolicy):
 
   def should_preserve(
       self,
-      checkpoints: Sequence[checkpoint_info.CheckpointInfo],
+      checkpoints,
       *,
       context: PreservationContext,
   ) -> Sequence[bool]:
@@ -154,7 +153,7 @@ class BestN(PreservationPolicy):
 
   def should_preserve(
       self,
-      checkpoints: Sequence[checkpoint_info.CheckpointInfo],
+      checkpoints,
       *,
       context: PreservationContext,
   ) -> Sequence[bool]:

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

@@ -19,7 +19,6 @@
 import typing
 from typing import Container, Protocol, Sequence
 from orbax.checkpoint import options as options_lib
-from orbax.checkpoint._src.metadata import checkpoint_info
 from orbax.checkpoint._src.multihost import multihost
 
 
@@ -41,11 +40,7 @@ class SaveDecisionPolicy(Protocol):
   """
 
   def should_save(
-      self,
-      step: checkpoint_info.CheckpointInfo,
-      previous_steps: Sequence[checkpoint_info.CheckpointInfo],
-      *,
-      context: DecisionContext
+      self, step, previous_steps, *, context: DecisionContext
   ) -> bool:
     ...
 
@@ -57,11 +52,7 @@ class FixedIntervalPolicy(SaveDecisionPolicy):
   interval: int
 
   def should_save(
-      self,
-      step: checkpoint_info.CheckpointInfo,
-      previous_steps: Sequence[checkpoint_info.CheckpointInfo],
-      *,
-      context: DecisionContext
+      self, step, previous_steps, *, context: DecisionContext
   ) -> bool:
     del previous_steps
     del context
@@ -75,11 +66,7 @@ class SpecificStepsPolicy(SaveDecisionPolicy):
   steps: Container[int]
 
   def should_save(
-      self,
-      step: checkpoint_info.CheckpointInfo,
-      previous_steps: Sequence[checkpoint_info.CheckpointInfo],
-      *,
-      context: DecisionContext
+      self, step, previous_steps, *, context: DecisionContext
   ) -> bool:
     del previous_steps
     del context
@@ -93,11 +80,7 @@ class ContinuousCheckpointingPolicy(SaveDecisionPolicy):
   minimum_interval_secs: int | None = None
 
   def should_save(
-      self,
-      step: checkpoint_info.CheckpointInfo,
-      previous_steps: Sequence[checkpoint_info.CheckpointInfo],
-      *,
-      context: DecisionContext
+      self, step, previous_steps, *, context: DecisionContext
   ) -> bool:
     if context.is_saving_in_progress:
       return False
@@ -122,11 +105,7 @@ class PreemptionCheckpointingPolicy(SaveDecisionPolicy):
   """Save a checkpoint when a preemption is detected."""
 
   def should_save(
-      self,
-      step: checkpoint_info.CheckpointInfo,
-      previous_steps: Sequence[checkpoint_info.CheckpointInfo],
-      *,
-      context: DecisionContext
+      self, step, previous_steps, *, context: DecisionContext
   ) -> bool:
     del step
     del previous_steps
@@ -137,11 +116,7 @@ class InitialSavePolicy(SaveDecisionPolicy):
   """Checkpoint as soon as possible if no checkpoints already exist."""
 
   def should_save(
-      self,
-      step: checkpoint_info.CheckpointInfo,
-      previous_steps: Sequence[checkpoint_info.CheckpointInfo],
-      *,
-      context: DecisionContext
+      self, step, previous_steps, *, context: DecisionContext
   ) -> bool:
     del step
     del context
@@ -159,11 +134,7 @@ class AnySavePolicy(SaveDecisionPolicy):
   policies: Sequence[SaveDecisionPolicy]
 
   def should_save(
-      self,
-      step: checkpoint_info.CheckpointInfo,
-      previous_steps: Sequence[checkpoint_info.CheckpointInfo],
-      *,
-      context: DecisionContext
+      self, step, previous_steps, *, context: DecisionContext
   ) -> bool:
     return any(
         policy.should_save(step, previous_steps=previous_steps, context=context)

checkpoint/orbax/checkpoint/checkpoint_manager.py

@@ -169,8 +169,8 @@ class _ShouldSaveFnPolicy(save_decision_policy_lib.SaveDecisionPolicy):
 
   def should_save(
       self,
-      step: checkpoint_info.CheckpointInfo,
-      previous_steps: Sequence[checkpoint_info.CheckpointInfo],
+      step,
+      previous_steps,
       *,
       context: save_decision_policy_lib.DecisionContext,
   ) -> bool:

checkpoint/orbax/checkpoint/experimental/v1/_src/synchronization/multihost.py

@@ -16,9 +16,10 @@
 
 import threading
 import time
-from typing import Collection
+from typing import Collection, Optional
 from absl import logging
 import jax
+from jax.experimental import multihost_utils
 from orbax.checkpoint.experimental.v1._src.synchronization import signaling_client
 
 # Default timeout in seconds.
@@ -139,3 +140,10 @@ def process_index() -> int:
   # global_state.process_id. We rely on the latter to work with barriers over a
   # subset of processes.
   return jax._src.distributed.global_state.process_id  # pylint: disable=protected-access
+
+
+def broadcast_one_to_all(in_tree, is_source: Optional[bool] = None):
+  """Broadcast data from a source host to all other hosts."""
+  if is_source is None:
+    is_source = process_index() == 0
+  return multihost_utils.broadcast_one_to_all(in_tree, is_source=is_source)

checkpoint/orbax/checkpoint/experimental/v1/_src/training/BUILD

@@ -3,7 +3,28 @@ package(default_visibility = ["//visibility:public"])
 py_library(
     name = "save_decision_policies",
     srcs = ["save_decision_policies.py"],
-    deps = ["//checkpoint/orbax/checkpoint/_src/checkpoint_managers:save_decision_policy"],
+    deps = [
+        "//checkpoint/orbax/checkpoint/_src/checkpoint_managers:save_decision_policy",
+        "//orbax/checkpoint/experimental/v1/_src/training/metadata:types",
+    ],
+)
+
+py_library(
+    name = "preservation_policy",
+    srcs = ["preservation_policy.py"],
+    deps = [
+        "//orbax/checkpoint/_src/checkpoint_managers:preservation_policy",
+        "//orbax/checkpoint/experimental/v1/_src/training/metadata:types",
+    ],
+)
+
+py_test(
+    name = "preservation_policy_test",
+    srcs = ["preservation_policy_test.py"],
+    deps = [
+        ":preservation_policy",
+        "//orbax/checkpoint/experimental/v1/_src/training/metadata:types",
+    ],
 )
 
 py_library(

checkpoint/orbax/checkpoint/experimental/v1/_src/training/metadata/types.py

@@ -15,6 +15,7 @@
 """Metadata for `training.Checkpointer`."""
 
 import dataclasses
+import datetime
 from orbax.checkpoint._src.metadata import checkpoint_info
 from orbax.checkpoint.experimental.v1._src.metadata import types as metadata_types
 from orbax.checkpoint.experimental.v1._src.tree import types as tree_types
@@ -70,6 +71,18 @@ def step(self) -> int:
   def metrics(self) -> tree_types.JsonType | None:
     return self._metrics
 
+  @metrics.setter
+  def metrics(self, new_metrics):
+    if new_metrics is not None and not isinstance(new_metrics, (list, dict)):
+      raise ValueError("Metrics must be a JSON-serializable object.")
+    self._metrics = new_metrics
+
+  @property
+  def time(self) -> datetime.datetime:
+    return datetime.datetime.fromtimestamp(
+        self.commit_timestamp_nsecs / 1e9, tz=datetime.timezone.utc
+    )
+
 
 @dataclasses.dataclass(frozen=True, kw_only=True)
 class RootMetadata:

checkpoint/orbax/checkpoint/experimental/v1/_src/training/preservation_policy.py

@@ -0,0 +1,44 @@
+# Copyright 2024 The Orbax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Defines policies for when a checkpoint is preserved."""
+
+from typing import Any, Dict, Protocol, Sequence
+from orbax.checkpoint._src.checkpoint_managers import preservation_policy as preservation_policy_lib
+from orbax.checkpoint.experimental.v1._src.training.metadata import types
+
+
+NestedDict = Dict[str, Any]
+PyTree = Any
+
+PreservationContext = preservation_policy_lib.PreservationContext
+LatestN = preservation_policy_lib.LatestN
+EveryNSeconds = preservation_policy_lib.EveryNSeconds
+EveryNSteps = preservation_policy_lib.EveryNSteps
+CustomSteps = preservation_policy_lib.CustomSteps
+AnyPreservationPolicy = preservation_policy_lib.AnyPreservationPolicy
+BestN = preservation_policy_lib.BestN
+
+
+class PreservationPolicy(Protocol):
+  """A policy that defines when checkpoints should be preserved."""
+
+  def should_preserve(
+      self,
+      checkpoints: Sequence[types.CheckpointMetadata],
+      *,
+      context: PreservationContext,
+  ) -> Sequence[bool]:
+    """Indicates which checkpoints should be preserved.."""
+    ...