Add DeletionOptions to Orbax v1 Context.

PiperOrigin-RevId: 886843908

Author: angel-core

checkpoint/CHANGELOG.md

@@ -11,6 +11,8 @@ 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.
 
 ### Removed
 

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

@@ -101,6 +101,7 @@ class Context(epy.ContextManager):
     checkpointables_options: Options for controlling checkpointables behavior.
     pathways_options: Options for Pathways checkpointing.
     checkpoint_layout: The layout of the checkpoint. Defaults to ORBAX.
+    deletion_options: Options for controlling deletion behavior.
   """
 
   def __init__(
@@ -115,6 +116,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 +148,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 +184,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,44 @@ 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: 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

@@ -170,8 +170,7 @@ def __init__(
         preservation_policy=preservation_policy,
         step_name_format=step_name_format,
         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,16 @@ 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='trash'
+          )
+      )
+      with ocp.Context(deletion_options=deletion_options):
+        checkpointer = Checkpointer(self.directory)
+        self.assertEqual(
+            checkpointer._manager._options.todelete_full_path, 'trash'
+        )
+