Apply the changes to the affected DAGs

Author: ryanliuCienet

dags/examples/maxtext_profile_sweep_example_dag.py

@@ -14,7 +14,7 @@
 
 """
 An example DAG to extract profile metrics from pretraining mixtral-8x7b model on v6-256.
-Profile extraction can be seamlessly integrated with maxtext_sweep_gke_config + run_with_name_gen_and_quarantine.
+Profile extraction can be seamlessly integrated with maxtext_sweep_gke_config + (to_name_gen_and_quarantine_task + run).
 """
 
 import datetime
@@ -142,4 +142,4 @@ def dict_to_arg(param_dict):
       )
 
       for test in maxtext_sweep_gke_test:
-        test.run_with_name_gen_and_quarantine(quarantine_task_group)
+        test.to_name_gen_and_quarantine_task(quarantine_task_group).run()

dags/orbax/maxtext_emc_restore_gcs.py

@@ -128,12 +128,13 @@
             run_model_cmds=workload_command,
             docker_image=image.value,
             test_owner=test_owner.DEPP_L,
-        ).run_with_node_interruption(
+        ).to_node_interruption_task(
+            expect_reach_to_step=step_to_interrupt,
+            last_node=True,
+        ).run(
             ramdisk_directory=test_config_util.DEFAULT_RAM_DISK,
             mtc_enabled=True,
             skip_post_process=True,
-            last_node=True,
-            expect_reach_to_step=step_to_interrupt,
             max_restart=15,
         )
 

dags/orbax/maxtext_emc_restore_local.py

@@ -126,11 +126,12 @@
             run_model_cmds=workload_command,
             docker_image=image.value,
             test_owner=test_owner.DEPP_L,
-        ).run_with_node_interruption(
+        ).to_node_interruption_task(
+            expect_reach_to_step=step_to_interrupt,
+        ).run(
             ramdisk_directory=test_config_util.DEFAULT_RAM_DISK,
             mtc_enabled=True,
             skip_post_process=True,
-            expect_reach_to_step=step_to_interrupt,
             max_restart=15,
         )
 

dags/orbax/maxtext_mtc_restore_local.py

@@ -129,11 +129,12 @@
             run_model_cmds=workload_command,
             docker_image=image.value,
             test_owner=test_owner.DEPP_L,
-        ).run_with_node_interruption(
+        ).to_node_interruption_task(
+            expect_reach_to_step=step_to_interrupt,
+        ).run(
             ramdisk_directory=test_config_util.DEFAULT_RAM_DISK,
             mtc_enabled=True,
             skip_post_process=True,
-            expect_reach_to_step=step_to_interrupt,
             max_restart=15,
         )
 

dags/orbax/maxtext_reg_restore_gcs_with_node_disruption.py

@@ -138,13 +138,14 @@ def generate_workload_checkpoints_location(gcs_ckpt_location: str) -> str:
             run_model_cmds=workload_command,
             docker_image=image.value,
             test_owner=test_owner.SHARON_Y,
-        ).run_with_node_interruption(
+        ).to_node_interruption_task(
+            expect_reach_to_step=step_to_interrupt,
+            check_file_exists=True,
+        ).run(
             gcs_location=gcs_location,
             xpk_branch=MAIN_BRANCH,
             skip_post_process=True,
-            expect_reach_to_step=step_to_interrupt,
             max_restart=15,
-            check_file_exists=True,
         )
 
         end_time = validation_util.generate_timestamp.override(