Fix format/linting errors and apply changes to DAGs that're affected

Author: ryanliuCienet

dags/examples/maxtext_aqtp_version_sweep_gke_example_dag.py

@@ -98,4 +98,4 @@
 
   # Run jobs
   for test in tests:
-    test.run_with_run_name_generation()
+    test.to_name_gen_and_quarantine_task().run()

dags/examples/maxtext_profile_namegen_example_dag.py

@@ -14,7 +14,7 @@
 
 """
 An example DAG to extract profile metrics from pretraining mixtral-8x7b model on 1xv4-128.
-Profile extraction can be easily integrated with gke_config + run_with_run_name_generation.
+Profile extraction can be easily integrated with gke_config + (to_name_gen_and_quarantine_task + run).
 """
 
 import datetime
@@ -88,7 +88,8 @@
 ) as dag:
   for run_name, test_scripts_details in test_models_tpu.items():
     for image in docker_image.keys():
-      # file_location: pass in base_output_directory, will be altered in `run_with_run_name_generation`
+      # file_location: pass in base_output_directory, will be altered in
+      # XpkNameGenAndQuarantineTask.run_with_run_name_generation
       job_metric_config = metric_config.MetricConfig()
       # optionally, add tensorboard metrics
       job_metric_config.tensorboard_summary = metric_config.SummaryConfig(
@@ -105,13 +106,17 @@
       if "not_add_profile_config" in test_scripts_details:
         job_metric_config.profile = None
 
-      tpu_task = gke_config.get_gke_config(
-          num_slices=1,
-          time_out_in_min=test_scripts_details["time_out_in_min"],
-          test_name=f"maxtext_{image}_{run_name}",
-          run_model_cmds=test_scripts_details["train_command"],
-          docker_image=docker_image[image],
-          test_owner=test_owner.SHUNING_J,
-          cluster=test_scripts_details["cluster"],
-          user_specified_job_metric_config=job_metric_config,  # customize config
-      ).run_with_run_name_generation(run_name_env="RUN_NAME")
+      tpu_task = (
+          gke_config.get_gke_config(
+              num_slices=1,
+              time_out_in_min=test_scripts_details["time_out_in_min"],
+              test_name=f"maxtext_{image}_{run_name}",
+              run_model_cmds=test_scripts_details["train_command"],
+              docker_image=docker_image[image],
+              test_owner=test_owner.SHUNING_J,
+              cluster=test_scripts_details["cluster"],
+              user_specified_job_metric_config=job_metric_config,  # customize config
+          )
+          .to_name_gen_and_quarantine_task(run_name_env="RUN_NAME")
+          .run()
+      )

dags/examples/maxtext_profile_sweep_example_dag.py

@@ -38,8 +38,7 @@ def dict_to_arg(param_dict):
 
 docker_image = {
     "stable": (
-        "gcr.io/tpu-prod-env-multipod/"
-        "maxtext_jax_stable_stack:2025-05-20"
+        "gcr.io/tpu-prod-env-multipod/" "maxtext_jax_stable_stack:2025-05-20"
     ),
 }
 
@@ -140,22 +139,20 @@ def dict_to_arg(param_dict):
       # generate run_name and tensorboard/profile location for extraction
       num_slices = [1]
       sweep_params = {}
-      maxtext_sweep_gke_test = (
-          maxtext_sweep_gke_config.get_maxtext_sweep_gke_config(
-              test_owner=test_owner.SHUNING_J,
-              dataset_project=Project.CLOUD_ML_AUTO_SOLUTIONS.value,
-              composer_project=Project.CLOUD_ML_AUTO_SOLUTIONS.value,
-              dataset_name=metric_config.DatasetOption.XLML_DATASET,
-              cluster=test_scripts_details["cluster"],
-              time_out_in_min=test_scripts_details["time_out_in_min"],
-              base_output_directory=BASE_OUTPUT_PATH,
-              num_slices=num_slices,
-              docker_image=img_val,
-              run_name_prefix=f"maxtext_{image}_{run_name}",
-              base_run_model_cmds=test_scripts_details["train_command"],
-              sweep_params=sweep_params,
-              enable_profile_config=True,  # add flag to enable profile extraction
-          )
+      maxtext_sweep_gke_test = maxtext_sweep_gke_config.get_maxtext_sweep_gke_config(
+          test_owner=test_owner.SHUNING_J,
+          dataset_project=Project.CLOUD_ML_AUTO_SOLUTIONS.value,
+          composer_project=Project.CLOUD_ML_AUTO_SOLUTIONS.value,
+          dataset_name=metric_config.DatasetOption.XLML_DATASET,
+          cluster=test_scripts_details["cluster"],
+          time_out_in_min=test_scripts_details["time_out_in_min"],
+          base_output_directory=BASE_OUTPUT_PATH,
+          num_slices=num_slices,
+          docker_image=img_val,
+          run_name_prefix=f"maxtext_{image}_{run_name}",
+          base_run_model_cmds=test_scripts_details["train_command"],
+          sweep_params=sweep_params,
+          enable_profile_config=True,  # add flag to enable profile extraction
       )
 
       for test in maxtext_sweep_gke_test:

dags/examples/maxtext_sweep_gke_example_dag.py

@@ -57,4 +57,4 @@
 
   # Run jobs
   for test in maxtext_sweep_gke_test:
-    test.run_with_run_name_generation()
+    test.to_name_gen_and_quarantine_task().run()

dags/multipod/maxtext_convergence.py

@@ -87,16 +87,20 @@
     # The grain dataset takes longer to run, so we give it a longer timeout. The other tests are expected to complete within 5 hours.
     timeout_in_min = 360 if test_name == "maxtext-convergence-grain" else 300
 
-    test_task = gke_config.get_gke_config(
-        cluster=XpkClusters.TPU_V6E_256_MLPERF_CLUSTER,
-        time_out_in_min=timeout_in_min,
-        test_name=test_name,
-        run_model_cmds=run_command,
-        docker_image=DockerImage.MAXTEXT_TPU_JAX_STABLE.value,
-        test_owner=test_owner.MATT_D,
-        base_output_directory=base_output_directory,
-        metric_aggregation_strategy=metric_config.AggregationStrategy.LAST,
-    ).run_with_run_name_generation()
+    test_task = (
+        gke_config.get_gke_config(
+            cluster=XpkClusters.TPU_V6E_256_MLPERF_CLUSTER,
+            time_out_in_min=timeout_in_min,
+            test_name=test_name,
+            run_model_cmds=run_command,
+            docker_image=DockerImage.MAXTEXT_TPU_JAX_STABLE.value,
+            test_owner=test_owner.MATT_D,
+            base_output_directory=base_output_directory,
+            metric_aggregation_strategy=metric_config.AggregationStrategy.LAST,
+        )
+        .to_name_gen_and_quarantine_task()
+        .run()
+    )
     if test_name not in parallel_test_names:
       sequential_tests.append(test_task)
 

xlml/apis/task.py

@@ -769,7 +769,7 @@ def run(
   ) -> DAGNode:
     """Run a test job with run name generation and quarantine."""
     test_name = self.task_test_config.benchmark_id
-    if QuarantineTests.is_quarantined(test_name):
+    if QuarantineTests.is_quarantined(test_name) and self.quarantine_task_group:
       with self.quarantine_task_group:
         return self.run_with_run_name_generation(
             use_pathways=use_pathways,