Fix linting errors

Author: ryanliuCienet

dags/examples/maxtext_profile_sweep_example_dag.py

@@ -12,16 +12,18 @@
 # 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 + (to_name_gen_and_quarantine_task + run).
+"""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.common.vm_resource import XpkClusters, Project
 from dags.multipod.configs import maxtext_sweep_gke_config
 from xlml.apis import metric_config
 
@@ -35,7 +37,10 @@ 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 +50,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 +91,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,25 +135,27 @@ 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(
-          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=docker_image[image],
-          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:

xlml/apis/task.py

@@ -18,7 +18,7 @@
 import dataclasses
 import datetime
 import shlex
-from typing import Optional, Tuple, Union
+from typing import Any, Optional, Tuple, Union
 
 import airflow
 from airflow.models.taskmixin import DAGNode
@@ -42,7 +42,7 @@ def run(self) -> DAGNode:
     Returns:
       A DAG node that executes this test.
     """
-    ...
+    pass
 
   def run_with_quarantine(self, quarantine_task_group):
     """Run a test job. If the test job is flaky, wrap it in a special task grop.
@@ -67,7 +67,7 @@ def run_queued_resource_test(
     tpu_name_env_var: bool = False,
     all_workers: bool = True,
     skip_post_process: bool = False,
-    custom_env: dict[str, str] = {},
+    custom_env: Optional[dict[str, str]] = None,
 ):
   """This is a class to set up tasks for TPU provisioned by Queued Resource.
 
@@ -94,6 +94,9 @@ def run_queued_resource_test(
       post_process and clean_up.
   """
 
+  if custom_env is None:
+    custom_env = {}
+
   with TaskGroup(
       group_id=task_test_config.benchmark_id, prefix_group_id=True
   ) as test:
@@ -324,7 +327,8 @@ class XpkTask(BaseTask):
   task_gcp_config: gcp_config.GCPConfig
   task_metric_config: Optional[metric_config.MetricConfig] = None
   workload_provision_timeout: datetime.timedelta = datetime.timedelta(
-      # Set the provision timeout from 300 to 60 minutes for decreasing the duration of failed tasks
+      # Set the provision timeout from 300 to 60 minutes for decreasing the
+      # duration of failed tasks
       minutes=60
   )
 
@@ -586,7 +590,10 @@ def run_model_with_node_interruption(
       xpk_branch: str = xpk.MAIN_BRANCH,
       max_restart: int = 0,
   ) -> DAGNode:
-    """Run the TPU/GPU test in `task_test_config` using xpk with interruption."""
+    """Run the TPU/GPU test in `task_test_config` using xpk.
+
+    with interruption.
+    """
     with TaskGroup(group_id="run_model") as group:
       workload_id = xpk.generate_workload_id(self.task_test_config.benchmark_id)
       if gcs_location:
@@ -709,14 +716,18 @@ def launch_workload_with_node_reach_to_step(
       wait_for_file_to_exist = gcs.wait_for_file_to_exist.override(
           task_id=task_id_wait_file_exist
       )(
-          file_path=f"{gcs_path}/{str(expect_reach_to_step)}/commit_success.txt",
+          file_path=(
+              f"{gcs_path}/{str(expect_reach_to_step)}/commit_success.txt"
+          ),
       )
       task_id_do_nothing = "do_nothing"
       do_nothing = EmptyOperator(task_id=task_id_do_nothing)
 
       @task.branch
       def task_path_decider(check_file_exists: bool = False) -> str:
-        """Dynamically route the workflow depending on the `check_file_exists`."""
+        """Dynamically route the workflow depending on the
+        `check_file_exists`.
+        """
         if check_file_exists:
           return f"{group.group_id}.{task_id_wait_file_exist}"
         return f"{group.group_id}.{task_id_do_nothing}"