diff --git a/.github/workflows/dag-check.yml b/.github/workflows/dag-check.yml index 96fcb76bf..c10eb9a89 100644 --- a/.github/workflows/dag-check.yml +++ b/.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: diff --git a/.github/workflows/pyink-check.yml b/.github/workflows/pyink-check.yml index ea7a8f250..0254abeb9 100644 --- a/.github/workflows/pyink-check.yml +++ b/.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 diff --git a/.github/workflows/pylint-check.yml b/.github/workflows/pylint-check.yml index 5e23d2812..c13899a02 100644 --- a/.github/workflows/pylint-check.yml +++ b/.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 diff --git a/.github/workflows/require-checklist.yml b/.github/workflows/require-checklist.yml index d15d19d99..507a07381 100644 --- a/.github/workflows/require-checklist.yml +++ b/.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 \ No newline at end of file + requireChecklist: false # If this is true and there are no checklists detected, the action will fail diff --git a/.github/workflows/unit-test.yml b/.github/workflows/unit-test.yml index e2ba42c08..a86cb0df1 100644 --- a/.github/workflows/unit-test.yml +++ b/.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: diff --git a/dags/sparsity_diffusion_devx/maxtext_moe_tpu_e2e.py b/dags/sparsity_diffusion_devx/maxtext_moe_tpu_e2e.py index 9a1a8951a..23d7ec750 100644 --- a/dags/sparsity_diffusion_devx/maxtext_moe_tpu_e2e.py +++ b/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 diff --git a/scripts/code-style.sh b/scripts/code-style.sh index 940f80e8f..51f38bd78 100755 --- a/scripts/code-style.sh +++ b/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."