Orbax auto/dag01 save local - #203
Closed
camiloCienet wants to merge 4 commits into
Closed
Conversation
camiloCienet
requested review from
andrewyct,
severus-ho and
xibinliu
as code owners
August 13, 2025 01:51
camiloCienet
force-pushed
the
orbax-auto/dag01-save-local
branch
16 times, most recently
from
August 18, 2025 01:22
31f2833 to
f2675c1
Compare
xibinliu
force-pushed
the
orbax-auto/dag01-save-local
branch
from
August 18, 2025 01:35
f2675c1 to
7f42073
Compare
camiloCienet
force-pushed
the
orbax-auto/dag01-save-local
branch
7 times, most recently
from
August 19, 2025 06:48
ef50d9d to
d0fefd6
Compare
camiloCienet
force-pushed
the
orbax-auto/dag01-save-local
branch
3 times, most recently
from
August 25, 2025 06:46
90d8b46 to
8d82f17
Compare
alfredyu-cienet
requested changes
Aug 25, 2025
camiloCienet
force-pushed
the
orbax-auto/dag01-save-local
branch
3 times, most recently
from
August 25, 2025 08:52
76b0753 to
4197b32
Compare
camiloCienet
force-pushed
the
orbax-auto/dag01-save-local
branch
6 times, most recently
from
August 26, 2025 08:30
e08a623 to
338a7a7
Compare
alfredyu-cienet
requested changes
Aug 27, 2025
camiloCienet
force-pushed
the
orbax-auto/dag01-save-local
branch
2 times, most recently
from
August 27, 2025 09:47
bc105d6 to
f935d74
Compare
camiloCienet
force-pushed
the
orbax-auto/dag01-save-local
branch
from
August 28, 2025 03:26
f935d74 to
523aa82
Compare
…save locally with MTC and EMC enabled.
Fix: Refactor cpc dict into a CheckpointConfiguration class, that holds relevant info about cpc metadata. Fix: Change name of DAG to be mtc and emc. mainly because both of these test are being tested. Fix: Fix some try and catch exceptions that seems to be unnecessary. Fix: Refactor whole Cpc delete task into two task initiate_cpc and wait_to_delete. wait_to_delete task is a sensor task that will check every 10 s if the delition already happened. Fix: Two Dags testing mtc and emc (with replicator=true and replicator=false) Fix: Change loggin.py library to logging_mtc.py names conflict. Fix: Use abhinav-mtc xpk branch instead of main branch since xpk version > 0.4.1 crash airflow pod. b/437817546 Fix: Correct name of DAG to emtc (Emergency multitier Checkpointer). Fix: Change name_id of the test to be related to Orbax. Create new tag with Orbax. Remove --restart-on-user-code-failure flag since is already outdated in both main and develop branch Fix: Backup interval to longer. Only testing local saves. GCS testing is in DAG002 mend
camiloCienet
force-pushed
the
orbax-auto/dag01-save-local
branch
4 times, most recently
from
August 28, 2025 07:14
d932978 to
2126ca4
Compare
alfredyu-cienet
approved these changes
Aug 28, 2025
Fix: Change name logging_mtc.py to validate_util.py. Since validate_util as the filename make more sense, we are validating the logs. Fix: checkpoint_steps --> to Optional. Since this value exist but is not use in this test. Fix: Change name mx-sv --> max-sv-loc. Helps with clarity in bucket checkpoints. Fixes for code review
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This change adds a new DAG which test and verifies the Orbax Emergency Checkpointer and the Multi-tier Checkpointing saving functionality. This DAG automate the installation of the mtc driver.,but requires an existing cluster.
Tests
maxtext_emtc_orbax_sav01_save_localAirflow/Composer
erniechang-test(under GCP project:cloud-ml-auto-solutions)2.13.1Bucket with HNS (Hierarchical Name Space)
Checklist
Before submitting this PR, please make sure (put X in square brackets):