chore: Set maxtext TPU E2E pre/post training DAGs priority to "very-high" (GoogleCloudPlatform#1311)

- Exposed the `priority` argument (defaulting to `"high"`) in `XpkTask` methods (`run`, `run_model`, `launch_workload`) within `xlml/apis/task.py` to allow passing custom scheduling priorities down to `xpk.run_workload`.    
- Explicitly set `priority="very-high"` in both `maxtext_e2e_tpu_pre_training.py` and `maxtext_e2e_tpu_post_training.py` DAGs to expedite these specific test runs.

Author: chiajunglien

dags/multipod/maxtext_e2e_tpu_post_training.py

@@ -94,7 +94,7 @@ def get_workload_name(model, mode, length=6):
           docker_image="{{ params.docker_image }}",
           cluster=XpkClusters.TPU_V5P_8_CLUSTER_V2,
           test_owner=test_owner.SURBHI_J,
-      ).run(skip_post_process=True)
+      ).run(skip_post_process=True, priority="very-high")
 
       for mode, mode_test_config in test_config["post_training"].items():
         with TaskGroup(group_id=f"{mode}-{model}") as model_group:
@@ -125,6 +125,7 @@ def get_workload_name(model, mode, length=6):
                 test_owner=test_owner.SURBHI_J,
             ).run_model(
                 use_pathways=True,
+                priority="very-high",
             )
 
           model_path = mode_test_config["maxtext_ckpt_path"].format(
@@ -140,7 +141,7 @@ def get_workload_name(model, mode, length=6):
               docker_image="{{ params.docker_image }}",
               cluster=XpkClusters.TPU_V5P_8_CLUSTER_V2,
               test_owner=test_owner.SURBHI_J,
-          ).run(skip_post_process=True)
+          ).run(skip_post_process=True, priority="very-high")
 
           (
               convert_to_maxtext_task

dags/multipod/maxtext_e2e_tpu_pre_training.py

@@ -70,7 +70,7 @@
           docker_image="{{ params.docker_image }}",
           cluster=XpkClusters.TPU_V5P_8_CLUSTER_V2,
           test_owner=test_owner.SURBHI_J,
-      ).run(skip_post_process=True)
+      ).run(skip_post_process=True, priority="very-high")
 
       training_cmd = (f"export HF_TOKEN={HF_TOKEN}",) + (
           f"{test_config['training']['command']} {run_name}",
@@ -82,7 +82,7 @@
           docker_image="{{ params.docker_image }}",
           cluster=XpkClusters.TPU_V5P_128_CLUSTER,
           test_owner=test_owner.SURBHI_J,
-      ).run(skip_post_process=True)
+      ).run(skip_post_process=True, priority="very-high")
 
       model_path = test_config["training"]["maxtext_ckpt_path"].format(
           run_name=run_name
@@ -97,6 +97,6 @@
           docker_image="{{ params.docker_image }}",
           cluster=XpkClusters.TPU_V5P_8_CLUSTER_V2,
           test_owner=test_owner.SURBHI_J,
-      ).run(skip_post_process=True)
+      ).run(skip_post_process=True, priority="very-high")
 
       convert_to_maxtext_task >> training_task >> convert_to_huggingface_task

xlml/apis/task.py

@@ -350,6 +350,7 @@ def run(
       mtc_enabled: bool = False,
       xpk_branch: str = xpk.MAIN_BRANCH,
       max_restart: int = 0,
+      priority: str = "high",
   ) -> DAGNode:
     """Run a test job within a docker image.
 
@@ -371,6 +372,7 @@ def run(
           mtc_enabled,
           xpk_branch,
           max_restart,
+          priority=priority,
       )
       if not skip_post_process:
         _ = run_model >> self.post_process(gcs_path)
@@ -727,6 +729,7 @@ def run_model(
       mtc_enabled: bool = False,
       xpk_branch: str = xpk.MAIN_BRANCH,
       max_restart: int = 0,
+      priority: str = "high",
   ) -> DAGNode:
     """Run the TPU/GPU test in `task_test_config` using xpk.
 
@@ -761,6 +764,7 @@ def run_model(
           mtc_enabled,
           xpk_branch,
           max_restart,
+          priority=priority,
       )
       wait_for_workload_completion = xpk.wait_for_workload_completion.override(
           timeout=int(self.task_test_config.timeout.total_seconds()),
@@ -798,6 +802,7 @@ def launch_workload(
       mtc_enabled: bool = False,
       xpk_branch: str = xpk.MAIN_BRANCH,
       max_restart: int = 0,
+      priority: str = "high",
   ) -> DAGNode:
     """Create the workload and wait for it to provision."""
     with TaskGroup(group_id="launch_workload") as group:
@@ -821,6 +826,7 @@ def launch_workload(
           mtc_enabled=mtc_enabled,
           xpk_branch=xpk_branch,
           max_restart=max_restart,
+          priority=priority,
       )
       wait_for_workload_start = xpk.wait_for_workload_start.override(
           timeout=self.workload_provision_timeout.total_seconds()