Expose step from orbax.checkpoint.experimental.v1.path.

PiperOrigin-RevId: 889348008

Author: angel-core

checkpoint/CHANGELOG.md

@@ -11,6 +11,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - #v1 Add `use_load_and_broadcast` option.
 - Add PyTorch DCP (Distributed Checkpoint) to the benchmark suite.
+- #v1 Add `DeletionOptions` to configure V1 Checkpointer's checkpoint deletion
+behavior.
+- #v1 Add `cleanup_tmp_directories` setting to V1 checkpointer's `FileOptions`
+to manage temporary directory cleanup behavior.
+- #v1 Add `lightweight_initialize` which allows users to specify whether
+temporary directories should be cleaned up upon Checkpointer initialization
 
 ### Removed
 

checkpoint/orbax/checkpoint/checkpoint_manager.py

@@ -364,12 +364,11 @@ class CheckpointManagerOptions:
     supposed to be created per process. This is used to support async
     directory creation. If True, `multiprocessing_options.primary_host` must be
     None.
-  lightweight_initialize: If True, checkpoint step metadata is not
-  read on
+  lightweight_initialize: If True, checkpoint step metadata is not read on
     CheckpointManager initialization during checkpoint info loading. This is
-    useful to improve init performance
-    when there are O(1k) or more existing checkpoint step present and checkpoint
-    info properties like `time` and `metrics` are not needed.
+    useful to improve init performance when there are O(1k) or more existing
+    checkpoint step present and checkpoint info properties like `time` and
+    `metrics` are not needed.
   """
 
   save_interval_steps: int = 1

checkpoint/orbax/checkpoint/experimental/v1/_src/context/context.py

@@ -115,6 +115,7 @@ def __init__(
       checkpointables_options: options_lib.CheckpointablesOptions | None = None,
       pathways_options: options_lib.PathwaysOptions | None = None,
       checkpoint_layout: options_lib.CheckpointLayout | None = None,
+      deletion_options: options_lib.DeletionOptions | None = None,
   ):
     self._pytree_options = pytree_options or (
         context.pytree_options if context else options_lib.PyTreeOptions()
@@ -146,6 +147,9 @@ def __init__(
         if context
         else options_lib.CheckpointLayout.ORBAX
     )
+    self._deletion_options = deletion_options or (
+        context.deletion_options if context else options_lib.DeletionOptions()
+    )
 
   @property
   def pytree_options(self) -> options_lib.PyTreeOptions:
@@ -179,6 +183,10 @@ def pathways_options(self) -> options_lib.PathwaysOptions:
   def checkpoint_layout(self) -> options_lib.CheckpointLayout:
     return self._checkpoint_layout
 
+  @property
+  def deletion_options(self) -> options_lib.DeletionOptions:
+    return self._deletion_options
+
   def operation_id(self) -> str:
     return synchronization.OperationIdGenerator.get_current_operation_id()
 

checkpoint/orbax/checkpoint/experimental/v1/_src/context/options.py

@@ -494,6 +494,45 @@ class PathwaysOptions:
   checkpointing_impl: pathways_types.CheckpointingImpl | None = None
 
 
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class DeletionOptions:
+  """Options used to configure checkpoint deletion behavior.
+
+  Attributes:
+    gcs_deletion_options: Deletion options specific to GCS.
+  """
+
+  @dataclasses.dataclass(frozen=True, kw_only=True)
+  class GcsDeletionOptions:
+    """Deletion options specific to GCS.
+
+    Attributes:
+      todelete_full_path: Specifies a path relative to the bucket root for
+        "soft-deleting" checkpoints on Google Cloud Storage (GCS). Instead of
+        being permanently removed, checkpoints are moved to this new location
+        within the same bucket. This is useful if direct deletion on GCS is
+        time-consuming, as it allows an external component to
+        manage the actual removal.
+
+        This option gathers all "deleted" items in a centralized path at the
+        bucket level for future cleanup.
+
+        For instance, if a checkpoint is in
+        gs://my-bucket/experiments/run1/, providing the value 'trash' will move
+        a deleted step to gs://my-bucket/trash/<step_id>. Useful when direct
+        deletion is time consuming. It gathers all deleted items in a
+        centralized path for future cleanup. 
+    """
+
+    todelete_full_path: str | None = None
+
+
+  gcs_deletion_options: GcsDeletionOptions = dataclasses.field(
+      default_factory=GcsDeletionOptions
+  )
+
+
+
 class CheckpointLayout(enum.Enum):
   """The layout of the checkpoint.
 

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

@@ -87,6 +87,8 @@ def __init__(
           path_step_lib.NameFormat[path_step_lib.Metadata] | None
       ) = None,
       custom_metadata: tree_types.JsonType | None = None,
+      cleanup_tmp_directories: bool = False,
+      lightweight_initialize: bool = False,
   ):
     """Initializes a Checkpointer.
 
@@ -150,6 +152,13 @@ def __init__(
       custom_metadata: A JSON dictionary representing user-specified custom
         metadata. This should be information that is relevant to the entire
         sequence of checkpoints, rather than to any single checkpoint.
+      cleanup_tmp_directories: If True, cleans up any existing temporary
+        directories on Checkpointer creation.
+      lightweight_initialize: If True, checkpoint step metadata is not read on
+        Checkpointer initialization during checkpoint info loading. This is
+        useful to improve init performance when there are O(1k) or more existing
+        checkpoint step present and checkpoint info properties like `time` and
+        `metrics` are not needed.
     """
     context = context_lib.get_context()
 
@@ -169,9 +178,10 @@ def __init__(
         save_decision_policy=save_decision_policy,
         preservation_policy=preservation_policy,
         step_name_format=step_name_format,
+        cleanup_tmp_directories=cleanup_tmp_directories,
+        lightweight_initialize=lightweight_initialize,
         max_to_keep=None,  # Unlimited.
-        # TODO(b/401541834) Configure todelete_subdir.
-        # TODO(b/401541834) Enable background deletion.
+        todelete_full_path=context.deletion_options.gcs_deletion_options.todelete_full_path,
         async_options=context.async_options.v0(),
         file_options=context.file_options.v0(),
         multiprocessing_options=context.multiprocessing_options.v0(),

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

@@ -669,3 +669,48 @@ def test_preservation_metrics(self, policy, expected_steps):
           [all_metrics[step] for step in expected_steps],
       )
       checkpointer.close()
+
+    def test_gcs_deletion_options(self):
+      deletion_options = ocp.options.DeletionOptions(
+          gcs_deletion_options=ocp.options.DeletionOptions.GcsDeletionOptions(
+              todelete_full_path='gs://bucket/trash'
+          )
+      )
+      self.enter_context(ocp.Context(deletion_options=deletion_options))
+      checkpointer = Checkpointer(self.directory)
+      self.assertEqual(
+          checkpointer._manager._options.todelete_full_path, 'gs://bucket/trash'
+      )
+
+
+    @parameterized.parameters(
+        (True, True),
+        (False, False),
+    )
+    def test_cleanup_tmp_directories(
+        self, cleanup_tmp_directories, expected_cleanup_tmp_directories
+    ):
+      checkpointer = Checkpointer(
+          self.directory, cleanup_tmp_directories=cleanup_tmp_directories
+      )
+      self.assertEqual(
+          checkpointer._manager._options.cleanup_tmp_directories,
+          expected_cleanup_tmp_directories,
+      )
+      checkpointer.close()
+
+    @parameterized.parameters(
+        (True, True),
+        (False, False),
+    )
+    def test_lightweight_initialize(
+        self, lightweight_initialize, expected_lightweight_initialize
+    ):
+      checkpointer = Checkpointer(
+          self.directory, lightweight_initialize=lightweight_initialize
+      )
+      self.assertEqual(
+          checkpointer._manager._options.lightweight_initialize,
+          expected_lightweight_initialize,
+      )
+      checkpointer.close()

checkpoint/orbax/checkpoint/experimental/v1/path.py

@@ -21,3 +21,4 @@
     PathLike,
     PathAwaitingCreation,
 )
+from orbax.checkpoint.experimental.v1._src.path import step