This PR decouples Orbax from HNS rename API dependencies by delegating the filesystem operations to the underlying TensorFlow gfile implementation. This change improves testability and makes Orbax less susceptible to regressions from this api client changes.

PiperOrigin-RevId: 833996717

Author: Orbax Authors

checkpoint/orbax/checkpoint/_src/path/deleter.py

@@ -21,6 +21,7 @@
 import threading
 import time
 from typing import Optional, Protocol, Sequence
+from urllib import parse
 
 from absl import logging
 from etils import epath
@@ -31,6 +32,7 @@
 from orbax.checkpoint._src.path import step as step_lib
 
 
+urlparse = parse.urlparse
 PurePosixPath = pathlib.PurePosixPath
 
 _THREADED_DELETE_DURATION = (
@@ -183,7 +185,8 @@ def delete(self, step: int) -> None:
       # Attempt to rename using GCS HNS API if configured.
       if self._todelete_full_path is not None:
         if gcs_utils.is_gcs_path(self._directory):
-          self._rename_gcs_step_with_hns(step, delete_target)
+          # This is recommended for GCS buckets with HNS enabled.
+          self._gcs_rename_step(step, delete_target)
         else:
           raise NotImplementedError()
       # Attempt to rename to local subdirectory using `todelete_subdir`
@@ -204,88 +207,55 @@ def delete(self, step: int) -> None:
           time.time() - start,
       )
 
-  def _rename_gcs_step_with_hns(
+  def _gcs_rename_step(
       self, step: int, delete_target: epath.Path
   ):
-    """Renames a GCS directory using the Storage Control API.
+    """Renames a GCS directory to a temporary location for deletion.
+
+    This method renames the directory using the
+    underlying `tf.io.gfile.rename` method. This underlying
+    implementation will automatically detect if the bucket is HNS-enabled
+    and use a fast atomic rename, or fall back to a legacy
+    copy/delete rename if it is not.
 
     Args:
       step: The checkpoint step number.
       delete_target: The path to the directory to be renamed.
-
-    Raises:
-      ValueError: If the GCS bucket is not HNS-enabled, as this is a
-        hard requirement for this operation.
     """
-    logging.info(
-        'Condition: GCS path with `todelete_full_path` set. Checking for HNS.'
-    )
-    bucket_name, _ = gcs_utils.parse_gcs_path(self._directory)
-    if not gcs_utils.is_hierarchical_namespace_enabled(self._directory):
-      raise ValueError(
-          f'Bucket "{bucket_name}" does not have Hierarchical Namespace'
-          ' enabled, which is required when _todelete_full_path is set.'
-      )
-
-    logging.info('HNS bucket detected. Attempting to rename step %d.', step)
-    # pylint: disable=g-import-not-at-top
-    from google.api_core import exceptions as google_exceptions  # pytype: disable=import-error
     try:
-      from google.cloud import storage_control_v2  # pytype: disable=import-error
-      import google.auth  # pytype: disable=import-error
-
-      # Use default credentials, but without a quota project to avoid
-      # quota issues with this API.
-      credentials, _ = google.auth.default()
-      creds_without_quota_project = credentials.with_quota_project(None)
-      client = storage_control_v2.StorageControlClient(
-          credentials=creds_without_quota_project
-      )
-      # Destination parent is the absolute path to the bucket.
-      destination_parent_dir_str = (
+      # Get the bucket name from the source path
+      bucket_name = urlparse(str(delete_target)).netloc
+      if not bucket_name:
+        raise ValueError(
+            f'Could not parse bucket name from path: {delete_target}'
+        )
+
+      # Construct the destination path inside the `_todelete_full_path` dir.
+      destination_parent_path = epath.Path(
           f'gs://{bucket_name}/{self._todelete_full_path}'
       )
-      destination_parent_path = PurePosixPath(destination_parent_dir_str)
-      logging.info(
-          'Ensuring destination parent folder exists via HNS API: %s',
-          destination_parent_dir_str,
-      )
-      try:
-        parent_folder_id = str(
-            destination_parent_path.relative_to(f'gs://{bucket_name}')
-        )
-        bucket_resource_name = f'projects/_/buckets/{bucket_name}'
-        client.create_folder(
-            request=storage_control_v2.CreateFolderRequest(
-                parent=bucket_resource_name,
-                folder_id=parent_folder_id,
-                recursive=True,
-            )
-        )
-        logging.info('HNS parent folder creation request sent.')
-      except google_exceptions.AlreadyExists:
-        logging.info('HNS parent folder already exists, proceeding.')
+      destination_parent_path.mkdir(parents=True, exist_ok=True)
 
+      # Create a unique name for the destination to avoid collisions.
       now = datetime.datetime.now()
       timestamp_str = now.strftime('%Y%m%d-%H%M%S-%f')
       new_name_with_timestamp = f'{delete_target.name}-{timestamp_str}'
       dest_path = destination_parent_path / new_name_with_timestamp
-      source_folder_id = str(delete_target.relative_to(f'gs://{bucket_name}'))
-      destination_folder_id = str(dest_path.relative_to(f'gs://{bucket_name}'))
-      source_resource_name = (
-          f'projects/_/buckets/{bucket_name}/folders/{source_folder_id}'
-      )
-      logging.info('Rename API call: Source: %s', source_resource_name)
-      logging.info('Rename API call: Destination ID: %s', destination_folder_id)
-      request = storage_control_v2.RenameFolderRequest(
-          name=source_resource_name,
-          destination_folder_id=destination_folder_id,
+
+      logging.info(
+          'Executing filesystem-aware rename: Source=`%s`, Destination=`%s`',
+          delete_target,
+          dest_path,
       )
-      op = client.rename_folder(request=request)
-      op.result()
+
+      # Call the high-level rename method.
+      # This will be fast on HNS and slow (but functional) on non-HNS.
+      delete_target.rename(dest_path)
       logging.info('Successfully renamed step %d to %s', step, dest_path)
-    except google_exceptions.GoogleAPIError as e:
-      logging.error('HNS rename failed for step %d. Error: %s', step, e)
+
+    except Exception as e:
+      logging.error('Rename failed for step %d. Error: %s', step, e)
+      raise
 
   def _rename_step_to_subdir(self, step: int, delete_target: epath.Path):
     """Renames a step directory to its corresponding todelete_subdir."""

checkpoint/orbax/checkpoint/_src/path/step.py

@@ -321,13 +321,11 @@ class _StandardNameFormat(NameFormat[Metadata]):
     single_host_load_and_broadcast: If True, the jax process=0 will list all
       steps and broadcast them to all other processes. NOTE: Ignored if jax
       backend is not multi controller.
-    enable_hns: Enables HNS-specific path logic.
   """
 
   step_prefix: Optional[str] = None
   step_format_fixed_length: Optional[int] = None
   single_host_load_and_broadcast: bool = False
-  enable_hns: bool = False
 
   def __str__(self):
     return f'StandardNameFormat("{self.build_name(1234)}")'
@@ -375,7 +373,7 @@ def _glob_step_paths(self, base_path: epath.PathLike) -> list[epath.Path]:
     """Returns step paths under `base_path`."""
     base_path = epath.Path(base_path)
     # <step_prefix>_?<0 padding>?*
-    if self.enable_hns and gcs_utils.is_hierarchical_namespace_enabled(
+    if gcs_utils.is_hierarchical_namespace_enabled(
         base_path
     ):
       logging.vlog(
@@ -560,7 +558,6 @@ def standard_name_format(
     step_prefix: Optional[str] = None,
     step_format_fixed_length: Optional[int] = None,
     single_host_load_and_broadcast: bool = False,
-    enable_hns: bool = False,
 ) -> NameFormat[Metadata]:
   """Returns NameFormat for 'standard' steps for common Orbax use cases.
 
@@ -580,13 +577,11 @@ def standard_name_format(
     single_host_load_and_broadcast: If True, the jax process=0 will list all
       steps and broadcast them to all other processes. NOTE: Ignored if jax
       backend is not multi controller.
-    enable_hns: Enables HNS-specific path logic.
   """
   return _StandardNameFormat(
       step_prefix=step_prefix,
       step_format_fixed_length=step_format_fixed_length,
       single_host_load_and_broadcast=single_host_load_and_broadcast,
-      enable_hns=enable_hns,
   )
 
 

checkpoint/orbax/checkpoint/checkpoint_manager.py

@@ -318,8 +318,6 @@ class CheckpointManagerOptions:
     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.
-  enable_hns: If True, enables HNS-specific path manipulation logic.
-    Experimental feature.
   enable_background_delete: If True, old checkpoint deletions will be done in a
     background thread, otherwise, it will be done at the end of each save.  When
     it's enabled, make sure to call CheckpointManager.close() or use context to
@@ -386,7 +384,6 @@ class CheckpointManagerOptions:
   single_host_load_and_broadcast: bool = False
   todelete_subdir: Optional[str] = None
   todelete_full_path: Optional[str] = None
-  enable_hns: bool = False
   enable_background_delete: bool = False
   read_only: bool = False
   enable_async_checkpointing: bool = True
@@ -874,7 +871,6 @@ def __init__(
             single_host_load_and_broadcast=(
                 self._options.single_host_load_and_broadcast
             ),
-            enable_hns=self._options.enable_hns,
         )
     )
 
@@ -905,7 +901,6 @@ def __init__(
             primary_host=self._multiprocessing_options.primary_host,
             todelete_subdir=self._options.todelete_subdir,
             todelete_full_path=self._options.todelete_full_path,
-            enable_hns=self._options.enable_hns,
             enable_background_delete=self._options.enable_background_delete,
         )
     )