feat: Implement strict serial execution

.github/workflows/dag-check.yml

@@ -3,7 +3,8 @@ name: DAG Check
 
 on:
   pull_request:
-    branches: [master]
+    branch:
+      - maxtext/dev/**
     types: [opened, synchronize, edited]
 
   push:

.github/workflows/pyink-check.yml

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

.github/workflows/pylint-check.yml

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

.github/workflows/require-checklist.yml

@@ -1,11 +1,16 @@
 name: Require Checklist
 on:
   pull_request:
+    branch:
+      - maxtext/dev/**
     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,8 @@ name: Unit Test
 
 on:
   pull_request:
-    branches: [master]
+    branch:
+      - maxtext/dev/**
     types: [opened, synchronize, edited]
 
   push:

dags/sparsity_diffusion_devx/maxtext_moe_tpu_e2e.py

@@ -1,5 +1,4 @@
 # Copyright 2024 Google LLC
-#
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
@@ -10,25 +9,28 @@
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""A DAG to run end-to-end MoE tests."""
+# limitations under the License."""A DAG to run end-to-end MoE tests."""
 
+"""A DAG to run end-to-end moe tests."""
 
 import datetime
+from datetime import timedelta
 from airflow import models
 from airflow.utils.task_group import TaskGroup
 from dags import composer_env
-from dags.common.quarantined_tests import QuarantineTests
 from dags.common import test_owner
 from dags.common.vm_resource import XpkClusters, DockerImage
 from dags.multipod.configs import gke_config
 from xlml.utils import name_format
 
+
 # Run once a day at 1 am UTC (5 pm PST)
 SCHEDULED_TIME = "0 4 * * *" if composer_env.is_prod_env() else None
 HF_TOKEN = models.Variable.get("HF_TOKEN", None)
 
+DEFAULT_RETRIES = 3
+DEFAULT_RETRY_DELAY = timedelta(minutes=5)
+
 
 with models.DAG(
     dag_id="maxtext_moe_tpu_e2e",
@@ -46,8 +48,12 @@
     ],
     start_date=datetime.datetime(2024, 11, 14),
     catchup=False,
+    default_args={
+        "retries": DEFAULT_RETRIES,
+        "retry_delay": DEFAULT_RETRY_DELAY,
+    },
 ) as dag:
-  test_name_prefix = "maxtext"
+  TEST_NAME_PREFIX = "maxtext"
   quarantine_task_group = TaskGroup(
       group_id="Quarantine", dag=dag, prefix_group_id=False
   )
@@ -56,8 +62,6 @@
       "nightly": DockerImage.MAXTEXT_TPU_STABLE_STACK_NIGHTLY_JAX.value,
   }
 
-  # Unchained tests
-  # TODO(ranran): add back ckpt conversation after b/384580048
   test_models_tpu = {
       "mixtral-8x22b": {
           "script_name": "tpu/mixtral/8x22b/2_test_mixtral",
@@ -70,31 +74,30 @@
           "time_out_in_min": 90,
       },
   }
-
   unchained_tests = []
   for model, test_scripts_details in test_models_tpu.items():
-    for image in docker_image.keys():
-      # TODO(shuningjin): remove this once stable image is upgraded
-      # gpt-oss with flash attention requires jax>=0.7.2, skip stable image
-      if model.startswith("gpt-oss") and image == "stable":
+    for image_name, image_value in docker_image.items():
+      if model.startswith("gpt-oss") and image_name == "stable":
         continue
       training_tpu = gke_config.get_gke_config(
           time_out_in_min=test_scripts_details["time_out_in_min"],
-          test_name=f"{test_name_prefix}_{image}_{model}",
+          test_name=f"{TEST_NAME_PREFIX}_{image_name}_{model}",
           run_model_cmds=(
-              f"export HF_TOKEN={HF_TOKEN}; export BASE_OUTPUT_PATH=$GCS_OUTPUT; bash end_to_end/{test_scripts_details['script_name']}.sh",
+              f"export HF_TOKEN={HF_TOKEN}; "
+              "export BASE_OUTPUT_PATH=$GCS_OUTPUT; "
+              "bash end_to_end/{test_scripts_details['script_name']}.sh",
           ),
-          docker_image=docker_image[image],
+          docker_image=image_value,
           test_owner=test_owner.SHUNING_J,
           cluster=test_scripts_details["cluster"],
       ).run_with_quarantine(quarantine_task_group)
       unchained_tests.append(training_tpu)
 
-  # stable_tpu >> nightly_tpu
   for i in range(len(unchained_tests) - 1):
-    unchained_tests[i] >> unchained_tests[i + 1]
+    downstream_task = unchained_tests[i + 1]
+    downstream_task.trigger_rule = "all_done"
+    _ = unchained_tests[i] >> downstream_task
 
-  # Chained tests
   multicluster_test_models = {
       "mixtral-8x7b": [
           {
@@ -123,72 +126,86 @@
   }
 
   def convert_checkpoint_and_run_training(
-      test_group_id,
-      test_name_prefix,
-      image,
-      docker_image,
-      model,
-      test_scripts_details,
+      run_config, scripts_config, upstream_task=None
   ):
-    with TaskGroup(group_id=test_group_id, prefix_group_id=False) as group:
-      test_name = f"{test_name_prefix}_{image}_{model}"
-      shared_gcs_location = name_format.generate_gcs_folder_location.override(
-          task_id=f"{test_group_id}_generate_gcs_folder_location"
-      )(
-          gcs_subfolder,
-          test_group_id,
-      )
-      conversion_cpu = gke_config.get_maxtext_cpu_end_to_end_gke_config(
-          time_out_in_min=test_scripts_details[0]["time_out_in_min"],
-          test_name=test_name,
-          run_model_cmds=(
-              f"export BASE_OUTPUT_PATH=$GCS_OUTPUT; bash end_to_end/{test_scripts_details[0]['script_name']}.sh",
-          ),
-          docker_image=docker_image,
-          test_owner=test_owner.SHUNING_J,
-          cluster=test_scripts_details[0]["cluster"],
-      ).run(gcs_location=shared_gcs_location)
-      training_tpu = gke_config.get_gke_config(
-          time_out_in_min=test_scripts_details[1]["time_out_in_min"],
-          test_name=test_name,
-          run_model_cmds=(
-              f"export BASE_OUTPUT_PATH=$GCS_OUTPUT; bash end_to_end/{test_scripts_details[1]['script_name']}.sh",
-          ),
-          docker_image=docker_image,
-          test_owner=test_owner.SHUNING_J,
-          cluster=test_scripts_details[1]["cluster"],
-      ).run(gcs_location=shared_gcs_location)
-      return conversion_cpu, training_tpu
+    """Orchestrates the CPU conversion and TPU training tasks for a given model.
 
-  tests = []
-  for model, test_scripts_details in multicluster_test_models.items():
-    gcs_subfolder = (
-        f"{test_owner.Team.JAX_MODELS_AND_PERFORMANCE.value}/maxtext"
+    Creates a shared GCS location, configures the CPU conversion task and the
+    subsequent TPU training task, and enforces sequential execution if an
+    upstream task is provided.
+    """
+
+    cpu_details = scripts_config[0]
+    tpu_details = scripts_config[1]
+
+    group_id = run_config["group_id"]
+    test_name = run_config["test_name"]
+    image_uri = run_config["image_uri"]
+
+    test_id_suffix = group_id.replace("chained_tests_", "")
+    gcs_location_task_id = f"generate_gcs_location_{test_id_suffix}"
+
+    trigger_rule = "all_done" if upstream_task is not None else "all_success"
+
+    shared_gcs_location = name_format.generate_gcs_folder_location.override(
+        task_id=gcs_location_task_id, trigger_rule=trigger_rule
+    )(
+        GCS_SUBFOLDER,
+        group_id,
     )
-    for image in docker_image.keys():
-      test_group_id = "chained_tests" + "_" + model + "_" + image
-      if QuarantineTests.is_quarantined(test_group_id):
-        with quarantine_task_group:
-          mode_cpu, mode_tpu = convert_checkpoint_and_run_training(
-              test_group_id,
-              test_name_prefix,
-              image,
-              docker_image[image],
-              model,
-              test_scripts_details,
-          )
-      else:
-        mode_cpu, mode_tpu = convert_checkpoint_and_run_training(
-            test_group_id,
-            test_name_prefix,
-            image,
-            docker_image[image],
-            model,
-            test_scripts_details,
-        )
-      tests.append(mode_cpu)
-      tests.append(mode_tpu)
-
-    # stable_cpu >> stable_tpu >> nightly_cpu >> nightly_tpu
-    for i in range(len(tests) - 1):
-      tests[i] >> tests[i + 1]
+
+    if upstream_task is not None:
+      _ = upstream_task >> shared_gcs_location
+
+    conversion_cpu = gke_config.get_maxtext_cpu_end_to_end_gke_config(
+        time_out_in_min=cpu_details["time_out_in_min"],
+        test_name=test_name,
+        run_model_cmds=(
+            "export BASE_OUTPUT_PATH=$GCS_OUTPUT; "
+            f"bash end_to_end/{cpu_details['script_name']}.sh",
+        ),
+        docker_image=image_uri,
+        test_owner=test_owner.SHUNING_J,
+        cluster=cpu_details["cluster"],
+    ).run(gcs_location=shared_gcs_location)
+
+    _ = shared_gcs_location >> conversion_cpu
+
+    tpu_train_task = gke_config.get_gke_config(
+        time_out_in_min=tpu_details["time_out_in_min"],
+        test_name=test_name,
+        run_model_cmds=(
+            "export BASE_OUTPUT_PATH=$GCS_OUTPUT; "
+            f"bash end_to_end/{tpu_details['script_name']}.sh",
+        ),
+        docker_image=image_uri,
+        test_owner=test_owner.SHUNING_J,
+        cluster=tpu_details["cluster"],
+    ).run(gcs_location=shared_gcs_location)
+
+    _ = shared_gcs_location >> training_tpu
+    _ = conversion_cpu >> tpu_train_task
+
+    return tpu_train_task
+
+  last_task = None
+  GCS_SUBFOLDER = f"{test_owner.Team.JAX_MODELS_AND_PERFORMANCE.value}/maxtext"
+
+  for model, test_scripts_details in multicluster_test_models.items():
+    for image_key, image_value in docker_image.items():
+      current_group_id = f"chained_tests_{model}_{image_key}"
+      current_test_name = f"{TEST_NAME_PREFIX}_{image_key}_{model}"
+
+      config = {
+          "group_id": current_group_id,
+          "test_name": current_test_name,
+          "image_uri": image_value,
+      }
+
+      mode_tpu = convert_checkpoint_and_run_training(
+          run_config=config,
+          scripts_config=test_scripts_details,
+          upstream_task=last_task,
+      )
+
+      last_task = mode_tpu

scripts/code-style.sh

@@ -19,14 +19,47 @@ set -e
 
 FOLDERS_TO_FORMAT=("dags" "xlml")
 
-for folder in "${FOLDERS_TO_FORMAT[@]}"
-do
-  pyink "$folder" --pyink-indentation=2 --pyink-use-majority-quotes --line-length=80
-done
-
-for folder in "${FOLDERS_TO_FORMAT[@]}"
-do
-  pylint "./$folder" --fail-under=9.6
-done
+# for folder in "${FOLDERS_TO_FORMAT[@]}"
+# do
+#   pyink "$folder" --pyink-indentation=2 --pyink-use-majority-quotes --line-length=80
+# done
+#
+# for folder in "${FOLDERS_TO_FORMAT[@]}"
+# do
+#   pylint "./$folder" --fail-under=9.6
+# done
+
+HEAD_SHA="$(git rev-parse HEAD)"
+BASE_BRANCH="maxtext/dev"
+
+if ! git rev-parse --verify "$BASE_BRANCH" >/dev/null 2>&1; then
+  git fetch origin "$BASE_BRANCH":"$BASE_BRANCH" || {
+    echo "[code-style] base branch '$BASE_BRANCH' not found, skip diff-based check."
+    exit 0
+  }
+fi
+
+CHANGED_PY_FILES="$(
+  git diff --name-only --diff-filter=ACM "${BASE_BRANCH}" "${HEAD_SHA}" \
+    | grep '\.py$' \
+    | while read -r f; do
+        for folder in "${FOLDERS_TO_FORMAT[@]}"; do
+          if [[ "$f" == "$folder/"* ]]; then
+            echo "$f"
+            break
+          fi
+        done
+      done \
+    | sort -u
+)"
+
+if [[ -z "${CHANGED_PY_FILES}" ]]; then
+  echo "[pre-push hook] no changed files detected between ${HEAD_SHA} and ${BASE_BRANCH}"
+  exit 1
+fi
+
+pyink ${CHANGED_PY_FILES} --pyink-indentation=2 --pyink-use-majority-quotes --line-length=80 --check --diff
+
+pylint ${CHANGED_PY_FILES} --fail-under=9.6 --disable=E1123
 
 echo "Successfully clean up all codes."