fix: Correct XPK DAG final status reporting issue (GoogleCloudPlatform#1304)

Correct DAG Status Reporting via Airflow Setup/Teardown API
__The Problem:__ By default, Airflow evaluates the final DAG run status based on the success of its leaf nodes. In our workload runs, `clean_up_workload` acts as a leaf node that is typically triggered with `all_done` rules. When a training workload failed, the cleanup task still completed successfully, causing the Airflow DAG to be incorrectly reported as SUCCESS.
__The Solution:__ We transitioned the cleanup logic to use Airflow's Setup/Teardown API:
Marked `clean_up_workload` as a teardown task linked to `launch_workload` via `.as_teardown(setups=launch_workload)`. This ensures that successful cleanup runs do not mask upstream failures (the DAG run status will now be determined by the actual work task, `wait_for_workload_completion`).
Configured `on_failure_fail_dagrun=True` inside `as_teardown(...)`. This ensures that if the cleanup task itself fails (e.g., if we fail to delete GKE resources, risking resource leaks), the DAG run will be flagged as FAILED.

Author: ryanliuCienet

xlml/apis/task.py

@@ -227,6 +227,10 @@ def run(
       A task group with the following task : run_model.
     """
     with TaskGroup(group_id=self.test_cfg.benchmark_id) as group:
+      dummy_op_for_teardown = EmptyOperator(
+          task_id="dummy_op_for_teardown"
+      ).as_setup()
+
       update_image_tag_cmd = axlearn.update_image_tag_cmd.override(
           owner=self.test_cfg.task_owner
       )(
@@ -293,6 +297,8 @@ def run(
           zone=self.gcp_cfg.zone,
           cluster_name=self.test_cfg.cluster_name,
           xpk_branch=xpk.MAIN_BRANCH,
+      ).as_teardown(
+          setups=dummy_op_for_teardown, on_failure_fail_dagrun=True
       )
 
       # flow1: The launcher task (Kubernetes Pod) that runs the workload.
@@ -302,7 +308,12 @@ def run(
       flow1 = run_workload
       flow2 = wait_for_workload_start >> wait_for_workload_completion >> cleanup
 
-      _ = update_image_tag_cmd >> gen_cmds >> [flow1, flow2]
+      _ = (
+          update_image_tag_cmd
+          >> gen_cmds
+          >> dummy_op_for_teardown
+          >> [flow1, flow2]
+      )
 
     return group
 
@@ -471,6 +482,10 @@ def run_model_with_node_interruption(
             self.task_test_config.benchmark_id,
         )
 
+      dummy_op_for_teardown = EmptyOperator(
+          task_id="dummy_op_for_teardown"
+      ).as_setup()
+
       launch_workload_and_wait_for_reach_step = (
           self.launch_workload_with_node_reach_to_step(
               workload_id,
@@ -512,10 +527,11 @@ def run_model_with_node_interruption(
           zone=self.task_gcp_config.zone,
           cluster_name=self.task_test_config.cluster_name,
           xpk_branch=xpk_branch,
-      )
+      ).as_teardown(setups=dummy_op_for_teardown, on_failure_fail_dagrun=True)
 
       _ = (
           (workload_id, gcs_path)
+          >> dummy_op_for_teardown
           >> launch_workload_and_wait_for_reach_step
           >> run_node_interruption
           >> wait_for_workload_completion
@@ -731,6 +747,11 @@ def run_model(
             self.task_test_config.gcs_subfolder,
             self.task_test_config.benchmark_id,
         )
+
+      dummy_op_for_teardown = EmptyOperator(
+          task_id="dummy_op_for_teardown"
+      ).as_setup()
+
       launch_workload = self.launch_workload(
           workload_id,
           gcs_path,
@@ -756,10 +777,11 @@ def run_model(
           zone=self.task_gcp_config.zone,
           cluster_name=self.task_test_config.cluster_name,
           xpk_branch=xpk_branch,
-      )
+      ).as_teardown(setups=dummy_op_for_teardown, on_failure_fail_dagrun=True)
 
       _ = (
           (workload_id, gcs_path)
+          >> dummy_op_for_teardown
           >> launch_workload
           >> wait_for_workload_completion
           >> clean_up_workload