google
Refactor GCS Operations to Use TensorFlow's File System
Summary:
This PR decouples Orbax from direct GCS API dependencies by delegating all filesystem operations to the underlying TensorFlow gfile implementation. This change improves testability and makes Orbax less susceptible to regressions from GCS client changes.
Key Changes:
Dependency Removal: All direct GCS API calls have been removed from Orbax and are now handled by TensorFlow's GCS filesystem.
HNS Checkpoint Renaming: Orbax now implicitly supports fast checkpoint renames on HNS-enabled buckets. Users can enable this feature by configuring _todelete_full_path, which triggers a tf.io.gfile.rename call. The underlying library handles HNS detection automatically.
Simplified Listing Logic: The custom logic for includeFoldersAsPrefixes has been removed, as this is now handled by default in TensorFlow's listdir implementation.
