fix: Adjust execution schedule for multiple DAGs (GoogleCloudPlatform#1273)

This change adjusts the execution order of Orbax DAGs to isolate a certain issue for further troubleshooting. Will change time schedule from Save -> Resume -> Restore.

Author: Chrisliao0806

dags/orbax/axlearn_checkpoint_regular.py

@@ -27,10 +27,9 @@
 from dags.orbax.util import test_config_util, validation_util
 from xlml.utils.gke import zone_to_region
 from xlml.utils import axlearn
-from xlml.apis.xpk_cluster_config import XpkClusterConfig
 
 
-SCHEDULE = "0 17 * * *" if composer_env.is_prod_env() else None
+SCHEDULE = "0 13 * * *" if composer_env.is_prod_env() else None
 DAG_TEST_NAME = "axlearn_reg_save"
 
 RESERVE_TIME_FOR_OTHERS = datetime.timedelta(minutes=5)

dags/orbax/maxtext_emc_restore_gcs.py

@@ -18,7 +18,7 @@
 from xlml.utils.gke import zone_to_region
 
 DAG_TEST_NAME = "maxtext_emc_orbax_res_gcs"
-SCHEDULE = "0 10 * * *" if composer_env.is_prod_env() else None
+SCHEDULE = "30 18 * * *" if composer_env.is_prod_env() else None
 
 with models.DAG(
     dag_id=DAG_TEST_NAME,
@@ -34,7 +34,10 @@
         "TPU",
         "v5p-128",
     ],
-    description="DAG to verify MaxText's emergency restore from GCS checkpoints after a full cluster interruption.",
+    description=(
+        "DAG to verify MaxText's emergency restore from GCS"
+        " checkpoints after a full cluster interruption."
+    ),
     doc_md="""
       # MaxText Emergency Restore from GCS Validation DAG
 
@@ -108,7 +111,9 @@
             run_name=run_name,
             slice_num=slice_num,
             out_folder="maxtext_emc_orbax_res_gcs",
-            enable_multi_tier_checkpointing=checkpointing.enable_multi_tier_checkpointing,
+            enable_multi_tier_checkpointing=(
+                checkpointing.enable_multi_tier_checkpointing
+            ),
         )
 
         start_time = validation_util.generate_timestamp.override(
@@ -174,18 +179,23 @@
             is_local=False
         )
 
-        validate_saved_checkpoints_steps_gcs = validation_util.validate_gcs_checkpoint_files(
-            bucket_path=(
-                f"{test_config_util.DEFAULT_BUCKET}/{DAG_TEST_NAME}/{run_name}"
-            ),
-            steps_to_validate=gcs_saved_steps_to_validate,
+        validate_saved_checkpoints_steps_gcs = (
+            validation_util.validate_gcs_checkpoint_files(
+                bucket_path=(
+                    f"{test_config_util.DEFAULT_BUCKET}"
+                    f"/{DAG_TEST_NAME}/{run_name}"
+                ),
+                steps_to_validate=gcs_saved_steps_to_validate,
+            )
         )
         # Final CPC cleanup to ensure symmetric start/end
         wait_delete_cpc_final = checkpoint_util.wait_for_cpc_deletion.override(
             trigger_rule=TriggerRule.ALL_DONE,
             task_id="wait_delete_cpc_final",
         )(test_config.cpc_config).as_teardown(setups=apply_cpc)
 
+        # Airflow uses >> for task chaining, which is pointless for pylint.
+        # pylint: disable=pointless-statement
         (
             wait_delete_cpc
             >> apply_cpc
@@ -198,3 +208,4 @@
             >> validate_saved_checkpoints_steps_gcs
             >> wait_delete_cpc_final
         )
+        # pylint: enable=pointless-statement

dags/orbax/maxtext_emc_restore_local.py

@@ -15,7 +15,7 @@
 from xlml.utils.gke import zone_to_region
 
 DAG_TEST_NAME = "maxtext_emc_orbax_res_local"
-SCHEDULE = "45 10 * * *" if composer_env.is_prod_env() else None
+SCHEDULE = "15 19 * * *" if composer_env.is_prod_env() else None
 
 
 with models.DAG(
@@ -32,7 +32,10 @@
         "TPU",
         "v5p-128",
     ],
-    description="DAG to verify MaxText's emergency restore from local checkpoints after a node interruption.",
+    description=(
+        "DAG to verify MaxText's emergency restore from local"
+        " checkpoints after a node interruption."
+    ),
     doc_md="""
       # MaxText Emergency Restore from Local Checkpoint Validation DAG
 
@@ -106,7 +109,9 @@
             run_name=run_name,
             slice_num=slice_num,
             out_folder="maxtext_emc_orbax_res_local",
-            enable_multi_tier_checkpointing=checkpointing.enable_multi_tier_checkpointing,
+            enable_multi_tier_checkpointing=(
+                checkpointing.enable_multi_tier_checkpointing
+            ),
         )
 
         start_time = validation_util.generate_timestamp.override(
@@ -183,6 +188,8 @@
             task_id="wait_delete_cpc_final",
         )(test_config.cpc_config).as_teardown(setups=apply_cpc)
 
+        # Airflow uses >> for task chaining, which is pointless for pylint.
+        # pylint: disable=pointless-statement
         (
             wait_delete_cpc
             >> apply_cpc
@@ -195,3 +202,4 @@
             >> validate_log
             >> wait_delete_cpc_final
         )
+        # pylint: enable=pointless-statement

dags/orbax/maxtext_emc_resume_gcs.py

@@ -21,7 +21,7 @@
 from dags.orbax.util import checkpoint_util
 from xlml.utils.gke import zone_to_region
 
-SCHEDULE = "30 15 * * *" if composer_env.is_prod_env() else None
+SCHEDULE = "15 14 * * *" if composer_env.is_prod_env() else None
 DAG_TEST_NAME = "maxtext_emc_resume_from_gcs"
 
 with models.DAG(
@@ -38,7 +38,10 @@
         "TPU",
         "v5p-128",
     ],
-    description="A DAG to test MaxText Emergency Checkpoint Manager GCS restore functionality.",
+    description=(
+        "A DAG to test MaxText Emergency Checkpoint Manager"
+        " GCS restore functionality."
+    ),
     doc_md="""
       # Emergency Checkpoint Manager GCS Restore Validation DAG
 
@@ -136,7 +139,9 @@
             run_name=run_name,
             slice_num=slice_num,
             out_folder=out_folder,
-            enable_multi_tier_checkpointing=checkpointing.enable_multi_tier_checkpointing,
+            enable_multi_tier_checkpointing=(
+                checkpointing.enable_multi_tier_checkpointing
+            ),
         )
 
         start_time = validation_util.generate_timestamp.override(
@@ -169,7 +174,9 @@
             run_name=run_name,
             slice_num=slice_num,
             out_folder=out_folder,
-            enable_multi_tier_checkpointing=checkpointing.enable_multi_tier_checkpointing,
+            enable_multi_tier_checkpointing=(
+                checkpointing.enable_multi_tier_checkpointing
+            ),
         )
 
         resume_training_run = gke_config.get_gke_config(
@@ -249,6 +256,8 @@
             task_id="wait_delete_cpc_final",
         )(test_config.cpc_config).as_teardown(setups=apply_first_cpc)
 
+        # Airflow uses >> for task chaining, which is pointless for pylint.
+        # pylint: disable=pointless-statement
         (
             wait_delete_cpc
             >> apply_first_cpc
@@ -264,3 +273,4 @@
             >> validate_saved_checkpoints_steps_gcs
             >> wait_delete_cpc_final
         )
+        # pylint: enable=pointless-statement

dags/orbax/maxtext_emc_save_gcs.py

@@ -14,12 +14,12 @@
 from dags.common.vm_resource import XpkClusters
 from dags.multipod.configs import gke_config
 from dags.orbax.util import checkpoint_util
+from dags.orbax.util import test_config_util
 from dags.orbax.util import validation_util
 from xlml.utils.gke import zone_to_region
-from dags.orbax.util import test_config_util
 
 
-SCHEDULE = "45 12 * * *" if composer_env.is_prod_env() else None
+SCHEDULE = "15 11 * * *" if composer_env.is_prod_env() else None
 DAG_TEST_NAME = "maxtext_emc_save_gcs"
 
 
@@ -37,7 +37,10 @@
         "TPU",
         "v5p-128",
     ],
-    description="DAG that verifies the orbax multi-tier checkpointing saving functionality with replicator to GCS bucket",
+    description=(
+        "DAG that verifies the orbax multi-tier checkpointing saving"
+        " functionality with replicator to GCS bucket"
+    ),
     doc_md="""
       # Multi-tier Checkpoint Validation DAG
 
@@ -101,7 +104,9 @@
             checkpoint_dir=test_config_util.DEFAULT_RAM_DISK,
             run_name=run_name,
             out_folder="maxtext_emc_orbax_save_gcs",
-            enable_multi_tier_checkpointing=checkpointing.enable_multi_tier_checkpointing,
+            enable_multi_tier_checkpointing=(
+                checkpointing.enable_multi_tier_checkpointing
+            ),
             slice_num=slice_num,
         )
 
@@ -133,16 +138,21 @@
             location=zone_to_region(test_config.cluster.zone),
             cluster_name=test_config.cluster.name,
             ram_disk="gcs",
-            pod_pattern=f"{test_config.short_id}-emc.*-0-\d+-",
+            pod_pattern=rf"{test_config.short_id}-emc.*-0-\d+-",
             start_time=start_time,
             end_time=end_time,
             steps_to_validate=steps_to_validate,
         )
 
         # Validate that GCS restore happened during the second training run
-        validate_checkpoints_steps_gcs = validation_util.validate_gcs_checkpoint_files(
-            bucket_path=f"{test_config_util.DEFAULT_BUCKET}/maxtext_emc_orbax_save_gcs/{run_name}",
-            steps_to_validate=steps_to_validate,
+        validate_checkpoints_steps_gcs = (
+            validation_util.validate_gcs_checkpoint_files(
+                bucket_path=(
+                    f"{test_config_util.DEFAULT_BUCKET}"
+                    f"/maxtext_emc_orbax_save_gcs/{run_name}"
+                ),
+                steps_to_validate=steps_to_validate,
+            )
         )
 
         # Final CPC cleanup to ensure symmetric start/end
@@ -151,6 +161,8 @@
             task_id="wait_delete_cpc_final",
         )(test_config.cpc_config).as_teardown(setups=apply_cpc)
 
+        # Airflow uses >> for task chaining, which is pointless for pylint.
+        # pylint: disable=pointless-statement
         (
             wait_delete_cpc
             >> apply_cpc
@@ -162,3 +174,4 @@
             >> validate_checkpoints_steps_gcs
             >> wait_delete_cpc_final
         )
+        # pylint: enable=pointless-statement

dags/orbax/maxtext_mtc_emergency_save_local.py

@@ -25,7 +25,7 @@
 from xlml.utils.gke import zone_to_region
 
 
-SCHEDULE = "0 12 * * *" if composer_env.is_prod_env() else None
+SCHEDULE = "45 11 * * *" if composer_env.is_prod_env() else None
 DAG_TEST_NAME = "maxtext_emc_and_mtc_orbax_save_local"
 
 

dags/orbax/maxtext_mtc_restore_local.py

@@ -18,7 +18,7 @@
 from xlml.utils.gke import zone_to_region
 
 DAG_TEST_NAME = "maxtext_mtc_orbax_res_local"
-SCHEDULE = "45 13 * * *" if composer_env.is_prod_env() else None
+SCHEDULE = "30 21 * * *" if composer_env.is_prod_env() else None
 
 with models.DAG(
     dag_id=DAG_TEST_NAME,
@@ -35,7 +35,10 @@
         "TPU",
         "v5p-128",
     ],
-    description="DAG to verify MaxText's multi-tier restore from local checkpoints after a node interruption.",
+    description=(
+        "DAG to verify MaxText's multi-tier restore from local"
+        " checkpoints after a node interruption."
+    ),
     doc_md="""
       # MaxText Multi-tier Restore from Local Checkpoint Validation DAG
 
@@ -109,7 +112,9 @@
             run_name=run_name,
             slice_num=slice_num,
             out_folder="maxtext_mtc_orbax_res_local",
-            enable_multi_tier_checkpointing=checkpointing.enable_multi_tier_checkpointing,
+            enable_multi_tier_checkpointing=(
+                checkpointing.enable_multi_tier_checkpointing
+            ),
         )
 
         start_time = validation_util.generate_timestamp.override(
@@ -188,6 +193,8 @@
             task_id="wait_delete_cpc_final",
         )(test_config.cpc_config).as_teardown(setups=apply_cpc)
 
+        # Airflow uses >> for task chaining, which is pointless for pylint.
+        # pylint: disable=pointless-statement
         (
             wait_delete_cpc
             >> apply_cpc
@@ -200,3 +207,4 @@
             >> validate_local_saved_steps
             >> wait_delete_cpc_final
         )
+        # pylint: enable=pointless-statement

dags/orbax/maxtext_mtc_resume_gcs.py

@@ -21,7 +21,7 @@
 from dags.orbax.util import checkpoint_util
 from xlml.utils.gke import zone_to_region
 
-SCHEDULE = "45 18 * * *" if composer_env.is_prod_env() else None
+SCHEDULE = "45 15 * * *" if composer_env.is_prod_env() else None
 DAG_TEST_NAME = "maxtext_mtc_resume_from_gcs"
 
 with models.DAG(
@@ -39,7 +39,10 @@
         "TPU",
         "v5p-128",
     ],
-    description="A DAG to test MaxText Multi-tier Checkpointing (MTC) GCS restore functionality.",
+    description=(
+        "A DAG to test MaxText Multi-tier Checkpointing (MTC)"
+        " GCS restore functionality."
+    ),
     doc_md="""
       # Multi-tier Checkpointing (MTC) GCS Restore Validation DAG
 
@@ -139,7 +142,9 @@
             run_name=run_name,
             slice_num=slice_num,
             out_folder=out_folder,
-            enable_multi_tier_checkpointing=checkpointing.enable_multi_tier_checkpointing,
+            enable_multi_tier_checkpointing=(
+                checkpointing.enable_multi_tier_checkpointing
+            ),
         )
 
         start_time = validation_util.generate_timestamp.override(
@@ -175,7 +180,9 @@
             run_name=run_name,
             slice_num=slice_num,
             out_folder=out_folder,
-            enable_multi_tier_checkpointing=checkpointing.enable_multi_tier_checkpointing,
+            enable_multi_tier_checkpointing=(
+                checkpointing.enable_multi_tier_checkpointing
+            ),
         )
 
         resume_training_run = gke_config.get_gke_config(
@@ -239,7 +246,8 @@
             )
         )
 
-        # Validate that MTC checkpoint files exist in GCS bucket with correct backup folder structure
+        # Validate that MTC checkpoint files exist in GCS bucket with
+        # correct backup folder structure
         validate_mtc_gcs_files = (
             validation_util.validate_gcs_checkpoint_files.override(
                 task_id="validate_mtc_gcs_files"
@@ -256,6 +264,8 @@
             task_id="wait_delete_cpc_final",
         )(test_config.cpc_config).as_teardown(setups=apply_cpc)
 
+        # Airflow uses >> for task chaining, which is pointless for pylint.
+        # pylint: disable=pointless-statement
         (
             wait_delete_cpc
             >> apply_cpc
@@ -272,3 +282,4 @@
             >> validate_mtc_gcs_files
             >> wait_delete_cpc_final
         )
+        # pylint: enable=pointless-statement