Refactor XpkTask using OOP subclasses

.github/workflows/dag-check.yml

@@ -3,7 +3,6 @@ name: DAG Check
 
 on:
   pull_request:
-    branches: [master]
     types: [opened, synchronize, edited]
 
   push:

.github/workflows/pyink-check.yml

@@ -2,11 +2,12 @@ name: Formatter
 
 on:
   pull_request:
-    branches: [master]
     types: [opened, synchronize, edited]
   push:
     branches: [master]
 
+  workflow_dispatch: {}
+
 jobs:
   format_check:
     runs-on: ubuntu-latest

.github/workflows/pylint-check.yml

@@ -2,12 +2,13 @@ name: Linter
 
 on:
   pull_request:
-    branches: [master]
     types: [opened, synchronize, edited]
 
   push:
     branches: [master]
 
+  workflow_dispatch: {}
+
 jobs:
   linting_check:
     runs-on: ubuntu-latest

.github/workflows/require-checklist.yml

@@ -2,10 +2,13 @@ name: Require Checklist
 on:
   pull_request:
     types: [opened, edited, synchronize]
+
+  workflow_dispatch: {}
+
 jobs:
   check_pr_body:
     runs-on: ubuntu-latest
     steps:
       - uses: mheap/require-checklist-action@v2
         with:
-          requireChecklist: false # If this is true and there are no checklists detected, the action will fail
+          requireChecklist: false # If this is true and there are no checklists detected, the action will fail

.github/workflows/unit-test.yml

@@ -3,7 +3,6 @@ name: Unit Test
 
 on:
   pull_request:
-    branches: [master]
     types: [opened, synchronize, edited]
 
   push:

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(
@@ -98,4 +102,4 @@
 
   # Run jobs
   for test in tests:
-    test.run_with_run_name_generation()
+    test.run()

dags/examples/maxtext_profile_namegen_example_dag.py

@@ -12,66 +12,83 @@
 # 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.
-Profile extraction can be easily integrated with gke_config + run_with_run_name_generation.
+"""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,8 +104,9 @@
     concurrency=2,
 ) 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`
+    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()
       # optionally, add tensorboard metrics
       job_metric_config.tensorboard_summary = metric_config.SummaryConfig(
@@ -105,13 +123,16 @@
       if "not_add_profile_config" in test_scripts_details:
         job_metric_config.profile = None
 
-      tpu_task = gke_config.get_gke_config(
+      tpu_task = gke_config.get_gke_config_with_name_gen_and_quarantine(
           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],
+          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
-      ).run_with_run_name_generation(run_name_env="RUN_NAME")
+          user_specified_job_metric_config=(
+              job_metric_config
+          ),  # customize config
+          run_name_env="RUN_NAME",
+      ).run()

dags/examples/maxtext_profile_sweep_example_dag.py

@@ -12,17 +12,19 @@
 # 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 v6-256.
-Profile extraction can be seamlessly integrated with maxtext_sweep_gke_config + run_with_name_gen_and_quarantine.
+"""An example DAG to extract profile metrics.
+
+Pretraining mixtral-8x7b model on v6-256.
+Profile extraction can be seamlessly integrated with
+maxtext_sweep_gke_config + (to_name_gen_and_quarantine_task + run).
 """
 
 import datetime
 from airflow import models
 from airflow.utils.task_group import TaskGroup
 from dags.common import test_owner
-from dags.common.vm_resource import XpkClusters, DockerImage, Project
-from dags.multipod.configs import maxtext_sweep_gke_config
+from dags.common.vm_resource import XpkClusters, Project
+from dags.multipod.configs import maxtext_sweep_gke_config as sweep_config
 from xlml.apis import metric_config
 
 SCHEDULED_TIME = None
@@ -35,7 +37,9 @@ def dict_to_arg(param_dict):
 
 
 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"
+    ),
 }
 
 # https://github.com/AI-Hypercomputer/maxtext/blob/main/benchmarks/maxtext_trillium_model_configs.py
@@ -45,9 +49,14 @@ def dict_to_arg(param_dict):
         "cluster": XpkClusters.TPU_V6E_256_MLPERF_CLUSTER,
         "train_command": [
             f"export BASE_OUTPUT_PATH={BASE_OUTPUT_PATH} && "
-            "python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml base_output_directory=${BASE_OUTPUT_PATH} model_name=mixtral-8x7b "
-            # add profiler config: ensure steps > skip_first_n_steps_for_profiler + profiler_steps
-            "steps=10 profiler=xplane skip_first_n_steps_for_profiler=5 profiler_steps=3 "
+            "python3 -m maxtext.trainers.pre_train.train "
+            "src/maxtext/configs/base.yml "
+            "base_output_directory=${BASE_OUTPUT_PATH} "
+            "model_name=mixtral-8x7b "
+            # add profiler config: ensure steps >
+            # skip_first_n_steps_for_profiler + profiler_steps
+            "steps=10 profiler=xplane skip_first_n_steps_for_profiler=5 "
+            "profiler_steps=3 "
             + dict_to_arg({
                 "per_device_batch_size": 12,
                 "ici_fsdp_parallelism": -1,
@@ -81,9 +90,14 @@ def dict_to_arg(param_dict):
         "base_output_directory": "gs://runner-maxtext-logs",
         "train_command": [
             f"export BASE_OUTPUT_PATH={BASE_OUTPUT_PATH} && "
-            "python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml base_output_directory=${BASE_OUTPUT_PATH} model_name=mixtral-8x7b "
-            # add profiler config: ensure steps > skip_first_n_steps_for_profiler + profiler_steps
-            "steps=10 profiler=xplane skip_first_n_steps_for_profiler=5 profiler_steps=3 "
+            "python3 -m maxtext.trainers.pre_train.train "
+            "src/maxtext/configs/base.yml "
+            "base_output_directory=${BASE_OUTPUT_PATH} "
+            "model_name=mixtral-8x7b "
+            # add profiler config: ensure steps >
+            # skip_first_n_steps_for_profiler + profiler_steps
+            "steps=10 profiler=xplane skip_first_n_steps_for_profiler=5 "
+            "profiler_steps=3 "
             + dict_to_arg({
                 "per_device_batch_size": 12,
                 "ici_fsdp_parallelism": -1,
@@ -120,12 +134,12 @@ def dict_to_arg(param_dict):
   )
 
   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():
       # sweep num_slices and other training params
       # 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,
@@ -134,12 +148,13 @@ def dict_to_arg(param_dict):
           time_out_in_min=test_scripts_details["time_out_in_min"],
           base_output_directory=BASE_OUTPUT_PATH,
           num_slices=num_slices,
-          docker_image=docker_image[image],
+          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
+          quarantine_task_group=quarantine_task_group,
       )
 
       for test in maxtext_sweep_gke_test:
-        test.run_with_name_gen_and_quarantine(quarantine_task_group)
+        test.run()

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
@@ -57,4 +63,4 @@
 
   # Run jobs
   for test in maxtext_sweep_gke_test:
-    test.run_with_run_name_generation()
+    test.run()