Upgrades version of Orbax/TF/TFIO to support HNS atomic folder rename

[deepikarajani24]

Summary

This PR updates the TensorFlow and Orbax dependencies to support the new HNS-native RenameFolder API.

Changes

To leverage this feature for HNS buckets, you'd need to configure ocp.CheckpointManagerOptions with todelete_full_path="_trash".

• Impact: When max_to_keep is exceeded, old checkpoints are now atomically moved to a _trash subdirectory instead of being deleted.

Context & Motivation

• TensorFlow Support: TensorFlow has added support for the HNS RenameFolder API, allowing for recursive, atomic directory moves.
• Orbax Integration: Orbax now exposes a todelete_full_path option in CheckpointManagerOptions. When enabled, Orbax delegates to tf.io.gfile.rename to move old checkpoints to a trash directory rather than performing a slow, object-by-object deletion.
• Performance: On HNS buckets, renaming a folder is significantly faster than standard deletion.

Validation

Scale testing was conducted on Axlearn workloads using this configuration. Results confirmed that the rename operations were significantly faster than the previous deletion mechanism, reducing overhead during checkpoint rotation.

Configuration Snippet

options=ocp.CheckpointManagerOptions(
    create=True,
    max_to_keep=cfg.keep_last_n,
    enable_async_checkpointing=True,
    step_name_format=self._name_format,
    should_save_fn=save_fn_with_summaries,
    enable_background_delete=True,
    async_options=ocp.options.AsyncOptions(timeout_secs=cfg.async_timeout_secs),
    # New HNS optimization:
    todelete_full_path="_trash",
)

[github-actions]

This pull request has been automatically marked as stale because it has been inactive for 60 days. It will be closed in 7 days if no further activity occurs. If you would like to continue working on this, please remove the stale label or leave a comment.

[github-actions]

This pull request was closed because it has been inactive for more than 7 days since being marked as stale. Please feel free to reopen it if you would like to continue.