Support empty folder deletion for HNS bucket with Orbax #1286

deepikarajani24 · 2025-07-07T03:51:31Z

Orbax made a change to support empty folders deletion which require enable_hns_rmtree flag to be True in CheckpointManagerOptions.

Tested orbax regular checkpoint with HNS bucket with keep_last_n as 3, and confirmed only latest 3 checkpoint folders were left and rest were deleted. Also, without this change confirmed all the checkpoint folders would exist and none would be deleted.
This change would help with improved listing for HNS bucket which gets slow if empty folders are present

Ethanlm

Thanks. Left some comment

Ethanlm · 2025-07-07T17:17:35Z

Dockerfile

@@ -83,14 +83,15 @@ ENTRYPOINT ["/opt/apache/beam/boot"]

 FROM base AS tpu

-ARG EXTRAS=
+ARG EXTRAS=orbax


Why do we need to change this?

Removed it, made orbax dependency change in pyproject.toml

Ethanlm · 2025-07-07T17:17:45Z

Dockerfile


 ENV UV_FIND_LINKS=https://storage.googleapis.com/jax-releases/libtpu_releases.html
 # Ensure we install the TPU version, even if building locally.
 # Jax will fallback to CPU when run on a machine without TPU.
 RUN uv pip install --prerelease=allow .[core,tpu] && uv cache clean
 RUN if [ -n "$EXTRAS" ]; then uv pip install .[$EXTRAS] && uv cache clean; fi
 COPY . .
+RUN pip install -U "orbax-checkpoint @ git+https://github.com/google/orbax.git@refs/pull/2074/head#subdirectory=checkpoint"


Why is this needed

Suggest moving this change to pyproject.toml,

# Orbax checkpointing. orbax = [ "humanize==4.10.0", "orbax-checkpoint@git+https://github.com/google/orbax.git@refs/pull/2074/head#subdirectory=checkpoint", ]

We can't use something from an unmerged PR

I've updated this PR to point to the HEAD of the Orbax repository as the orbax PR is now merged. Should we merge this now, or should we wait for the next official Orbax release and use that version instead ?

Also, moved the change to pyproject.yaml

And this change is needed because it relies on the enable_hns_rmtree flag, which was recently added to the Orbax library.

Ethanlm · 2025-07-07T17:18:00Z

axlearn/common/checkpointer_orbax.py

@@ -242,6 +242,7 @@ def save_fn_with_summaries(step: int, last_saved_step: Optional[int]) -> bool:
                should_save_fn=save_fn_with_summaries,
                enable_background_delete=True,
                async_options=ocp.options.AsyncOptions(timeout_secs=cfg.async_timeout_secs),
+                enable_hns_rmtree=True,


Can you please add some test results in the PR summary to help with reviewers?

I have added the details. Kindly review and I'm happy to provide any additional details as needed.

What happens if the GCS is not HNS? Is there any regression in terms of functionality and performance?

for non HNS bucket, there is existing path.rmtree() in orbax which works to delete objects in non HNS. But that does not work for HNS bucket as HNS bucket have true folder structure and path.rmtree() only removes objects, leaving all the empty parent folders intact. So, orbax is now manually deleting the empty folders recursively.

_rmtree function in Orbax https://github.com/google/orbax/blob/1769e61f1a380f975d7094f6b8c6ecff035ac5db/checkpoint/orbax/checkpoint/_src/path/deleter.py#L140

And Axlearn uses enable_background_delete=True in CheckpointManagerOptions CheckpointManagerOptionshttps://github.com/apple/axlearn/blob/89991e862f2889641dd705040106f1706cd8db5c/axlearn/common/checkpointer_orbax.py#L243C16 which makes the deletion in background and should not cause regression in performance

andersensam · 2025-07-07T19:36:01Z

pyproject.toml

@@ -195,4 +196,4 @@ junit_family="xunit2"

 [tool.isort]
 line_length = 100
-profile = "black"
+profile = "black"


Fix extra line -- these changes will be rejected by pre-commit

We should revert the change here

andersensam · 2025-07-07T19:36:15Z

Dockerfile

@@ -112,4 +113,4 @@ COPY . .
 # Final target spec.                                                           #
 ################################################################################

-FROM ${TARGET} AS final
+FROM ${TARGET} AS final


Remove extra line

We should revert the change in this file

Ethanlm · 2025-07-08T23:13:23Z

Dockerfile

@@ -112,4 +113,4 @@ COPY . .
 # Final target spec.                                                           #
 ################################################################################

-FROM ${TARGET} AS final
+FROM ${TARGET} AS final


We should revert the change in this file

Ethanlm · 2025-07-08T23:13:40Z

pyproject.toml

@@ -195,4 +196,4 @@ junit_family="xunit2"

 [tool.isort]
 line_length = 100
-profile = "black"
+profile = "black"


We should revert the change here

Ethanlm · 2025-07-08T23:13:58Z

pyproject.toml

@@ -120,7 +121,7 @@ vertexai_tensorboard = [
    "setuptools==65.7.0",
    # Pin version to fix Tensorboard uploader TypeError: can only concatenate str (not "NoneType") to str
    # https://github.com/googleapis/python-aiplatform/commit/4f982ab254b05fe44a9d2ed959fca2793961b56c
-    "google-cloud-aiplatform[tensorboard]==1.61.0",
+    "google-cloud-aiplatform[tensorboard]==1.92.0",


Is this a required change?

Yes, as we updated the version of google-cloud-storage to recognize buckets with hierarchical namespace enabled, google-cloud-aiplatform is also updated correspondingly.

Ethanlm · 2025-07-08T23:14:23Z

axlearn/common/checkpointer_orbax.py

@@ -242,6 +242,7 @@ def save_fn_with_summaries(step: int, last_saved_step: Optional[int]) -> bool:
                should_save_fn=save_fn_with_summaries,
                enable_background_delete=True,
                async_options=ocp.options.AsyncOptions(timeout_secs=cfg.async_timeout_secs),
+                enable_hns_rmtree=True,


What happens if the GCS is not HNS? Is there any regression in terms of functionality and performance?

Ethanlm · 2025-07-08T23:15:35Z

pyproject.toml

@@ -154,7 +155,7 @@ mmau = [
 # Orbax checkpointing.
 orbax = [
    "humanize==4.10.0",
-    "orbax-checkpoint==0.11.15",
+    "orbax-checkpoint @ git+https://github.com/google/orbax.git#subdirectory=checkpoint",


When will the release be available? yes it is preferred to use a released version

Checked with Orbax team it will be available likely tomorrow. I will update the release version here.

using the latest release version now.

Support empty folder deletion for HNS bucket with Orbax

3669e5a

deepikarajani24 requested review from ruomingp, markblee and a team as code owners July 7, 2025 03:51

Ethanlm requested changes Jul 7, 2025

View reviewed changes

andersensam reviewed Jul 7, 2025

View reviewed changes

deepikarajani24 added 2 commits July 8, 2025 19:48

update the orbax version to head as the PR is merged

6378138

Merge branch 'main' of https://github.com/apple/axlearn

13511b4

Ethanlm reviewed Jul 8, 2025

View reviewed changes

deepikarajani24 added 4 commits July 9, 2025 01:21

remove extra line in Dockerfile

2b4eca2

revert the last change to Dockerfile

b1e07dc

add extra newline at the end as the code seems to to resolve the diff

9666470

use orbax release 0.11.19

48b2692

deepikarajani24 force-pushed the main branch from 9a84b2d to 48b2692 Compare July 10, 2025 00:41

deepikarajani24 requested review from Ethanlm and andersensam July 11, 2025 19:36

Support empty folder deletion for HNS bucket with Orbax #1286

Are you sure you want to change the base?

Support empty folder deletion for HNS bucket with Orbax #1286

Conversation

deepikarajani24 commented Jul 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Ethanlm left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

deepikarajani24 commented Jul 7, 2025 •

edited

Loading