Fix format/linting errors and set the quarantine_task_group arg default to None

ryanliuCienet · ryanliuCienet · commit 1e7b72abab9c · 2026-07-08T02:09:50.000Z
diff --git a/dags/examples/maxtext_aqtp_version_sweep_gke_example_dag.py b/dags/examples/maxtext_aqtp_version_sweep_gke_example_dag.py
@@ -21,7 +21,7 @@
 import datetime
 from airflow import models
 from dags.common import test_owner
-from dags.common.vm_resource import TpuVersion, Zone, Project, XpkClusters, DockerImage
+from dags.common.vm_resource import XpkClusters, DockerImage
 from dags.multipod.configs import maxtext_sweep_gke_config
 
 # Set concurrency to number of workers otherwise tasks may time out
@@ -54,7 +54,11 @@
     for model_size in models:
       run_cmds = [
           "pip show aqtp",
-          f"bash src/maxtext/configs/tpu/{tpu}/{model_size}.sh EXECUTABLE=train.py OUTPUT_PATH={base_output_directory} PLATFORM=gke",
+          (
+              f"bash src/maxtext/configs/tpu/{tpu}/{model_size}.sh "
+              f"EXECUTABLE=train.py OUTPUT_PATH={base_output_directory} "
+              "PLATFORM=gke"
+          ),
       ]
 
       tests.extend(
diff --git a/dags/examples/maxtext_profile_namegen_example_dag.py b/dags/examples/maxtext_profile_namegen_example_dag.py
@@ -12,66 +12,82 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""
-An example DAG to extract profile metrics from pretraining mixtral-8x7b model on 1xv4-128.
+"""An example DAG to extract profile metrics.
+
+Pretraining mixtral-8x7b model on 1xv4-128.
 Profile extraction can be easily integrated with gke_config + (to_name_gen_and_quarantine_task + run).
 """
 
 import datetime
 from airflow import models
 from dags.common import test_owner
-from dags.common.vm_resource import XpkClusters, DockerImage
+from dags.common.vm_resource import XpkClusters
 from dags.multipod.configs import gke_config
 from xlml.apis import metric_config
 
 SCHEDULED_TIME = None
 BASE_OUTPUT_PATH = "gs://runner-maxtext-logs"
 
 docker_image = {
-    "stable": "gcr.io/tpu-prod-env-multipod/maxtext_jax_stable_stack:2025-05-20",
+    "stable": (
+        "gcr.io/tpu-prod-env-multipod/" "maxtext_jax_stable_stack:2025-05-20"
+    ),
 }
 
 base_command = (
     f"export BASE_OUTPUT_PATH={BASE_OUTPUT_PATH} && "
-    + "python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml base_output_directory=gs://runner-maxtext-logs run_name=${RUN_NAME} model_name=mixtral-8x7b tokenizer_path=assets/tokenizer.mistral-v1 dataset_path=gs://maxtext-dataset per_device_batch_size=4 enable_checkpointing=false ici_fsdp_parallelism=-1 max_target_length=1024 async_checkpointing=false attention=flash dtype=bfloat16 weight_dtype=bfloat16"
+    "python3 -m maxtext.trainers.pre_train.train "
+    "src/maxtext/configs/base.yml "
+    "base_output_directory=gs://runner-maxtext-logs "
+    "run_name=${RUN_NAME} model_name=mixtral-8x7b "
+    "tokenizer_path=assets/tokenizer.mistral-v1 "
+    "dataset_path=gs://maxtext-dataset per_device_batch_size=4 "
+    "enable_checkpointing=false ici_fsdp_parallelism=-1 "
+    "max_target_length=1024 async_checkpointing=false "
+    "attention=flash dtype=bfloat16 weight_dtype=bfloat16"
 )
 
 test_models_tpu = {
     # use: upload single profile from the first host, extract profile
-    # add profiler config: ensure steps > skip_first_n_steps_for_profiler + profiler_steps
+    # add profiler config: ensure steps >
+    # skip_first_n_steps_for_profiler + profiler_steps
     "mixtral-8x7b_pretraining-megablox_config-true_upload-one": {
         "cluster": XpkClusters.TPU_V4_128_CLUSTER,
         "time_out_in_min": 60,
         "train_command": [
-            base_command
-            + " steps=10 profiler=xplane skip_first_n_steps_for_profiler=5 profiler_steps=3",
+            base_command + " steps=10 profiler=xplane "
+            "skip_first_n_steps_for_profiler=5 profiler_steps=3",
         ],
     },
     # use: upload profiles from all hosts, extract one of the profiles
-    # add profiler config: ensure steps > skip_first_n_steps_for_profiler + profiler_steps
+    # add profiler config: ensure steps >
+    # skip_first_n_steps_for_profiler + profiler_steps
     "mixtral-8x7b_pretraining-megablox_config-true_upload-all": {
         "cluster": XpkClusters.TPU_V4_128_CLUSTER,
         "time_out_in_min": 60,
         "train_command": [
-            base_command
-            + " steps=10 profiler=xplane skip_first_n_steps_for_profiler=5 profiler_steps=3 upload_all_profiler_results=True",
+            base_command + " steps=10 profiler=xplane "
+            "skip_first_n_steps_for_profiler=5 profiler_steps=3 "
+            "upload_all_profiler_results=True",
         ],
     },
-    # testing: handle edge case, attempt to extract, find no match, proceed to post_process without error
+    # testing: handle edge case, attempt to extract, find no match,
+    # proceed to post_process without error
     "testing_config-true_upload-none": {
         "cluster": XpkClusters.TPU_V4_128_CLUSTER,
         "time_out_in_min": 60,
         "train_command": [
             base_command + " steps=10",
         ],
     },
-    # testing: not generate profile location, not extract profile in post_process
+    # testing: not generate profile location, not extract profile in
+    # post_process
     "testing_config-false_upload-one": {
         "cluster": XpkClusters.TPU_V4_128_CLUSTER,
         "time_out_in_min": 60,
         "train_command": [
-            base_command
-            + " steps=10 profiler=xplane skip_first_n_steps_for_profiler=5 profiler_steps=3",
+            base_command + " steps=10 profiler=xplane "
+            "skip_first_n_steps_for_profiler=5 profiler_steps=3",
         ],
         "not_add_profile_config": True,
     },
@@ -87,7 +103,7 @@
     concurrency=2,
 ) as dag:
   for run_name, test_scripts_details in test_models_tpu.items():
-    for image in docker_image.keys():
+    for image, img_val in docker_image.items():
       # file_location: pass in base_output_directory, will be altered in
       # XpkNameGenAndQuarantineTask.run_with_run_name_generation
       job_metric_config = metric_config.MetricConfig()
@@ -112,10 +128,12 @@
               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],
+              docker_image=img_val,
               test_owner=test_owner.SHUNING_J,
               cluster=test_scripts_details["cluster"],
-              user_specified_job_metric_config=job_metric_config,  # customize config
+              user_specified_job_metric_config=(
+                  job_metric_config
+              ),  # customize config
           )
           .to_name_gen_and_quarantine_task(run_name_env="RUN_NAME")
           .run()
diff --git a/dags/examples/maxtext_profile_sweep_example_dag.py b/dags/examples/maxtext_profile_sweep_example_dag.py
@@ -24,7 +24,7 @@
 from airflow.utils.task_group import TaskGroup
 from dags.common import test_owner
 from dags.common.vm_resource import XpkClusters, Project
-from dags.multipod.configs import maxtext_sweep_gke_config
+from dags.multipod.configs import maxtext_sweep_gke_config as sweep_config
 from xlml.apis import metric_config
 
 SCHEDULED_TIME = None
@@ -139,7 +139,7 @@ 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(
+      maxtext_sweep_gke_test = sweep_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,
diff --git a/dags/examples/maxtext_sweep_gke_example_dag.py b/dags/examples/maxtext_sweep_gke_example_dag.py
@@ -21,7 +21,7 @@
 import datetime
 from airflow import models
 from dags.common import test_owner
-from dags.common.vm_resource import TpuVersion, Zone, Project, XpkClusters, DockerImage
+from dags.common.vm_resource import XpkClusters, DockerImage
 from dags.multipod.configs import maxtext_sweep_gke_config
 
 # Set concurrency to number of workers otherwise tasks may time out
@@ -37,7 +37,13 @@
   # MaxText set up and run commands
   base_output_directory = "gs://maxtext-experiments-multipod"
   base_run_model_cmds = [
-      f"python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml base_output_directory={base_output_directory} dataset_path=gs://max-datasets-rogue enable_checkpointing=false global_parameter_scale=16 steps=10",
+      (
+          "python3 -m maxtext.trainers.pre_train.train "
+          "src/maxtext/configs/base.yml "
+          f"base_output_directory={base_output_directory} "
+          "dataset_path=gs://max-datasets-rogue enable_checkpointing=false "
+          "global_parameter_scale=16 steps=10"
+      ),
   ]
 
   # Get list of MaxText GKE XPK jobs
diff --git a/dags/multipod/maxtext_convergence.py b/dags/multipod/maxtext_convergence.py
@@ -84,7 +84,8 @@
 
   sequential_tests = []
   for test_name, run_command in convergence_tests.items():
-    # The grain dataset takes longer to run, so we give it a longer timeout. The other tests are expected to complete within 5 hours.
+    # 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 = (
diff --git a/xlml/apis/task.py b/xlml/apis/task.py
@@ -388,7 +388,7 @@ def to_node_interruption_task(
 
   def to_name_gen_and_quarantine_task(
       self,
-      quarantine_task_group,
+      quarantine_task_group: Any = None,
       run_name_env: str = "M_RUN_NAME",
       nested_run_name_in_tb_file_location: bool = True,
   ) -> "XpkNameGenAndQuarantineTask":
