diff --git a/.github/workflows/dag-check.yml b/.github/workflows/dag-check.yml index 96fcb76bf..ac3982d0e 100644 --- a/.github/workflows/dag-check.yml +++ b/.github/workflows/dag-check.yml @@ -3,7 +3,6 @@ name: DAG Check on: pull_request: - branches: [master] types: [opened, synchronize, edited] push: diff --git a/.github/workflows/pyink-check.yml b/.github/workflows/pyink-check.yml index ea7a8f250..c5cb1cd83 100644 --- a/.github/workflows/pyink-check.yml +++ b/.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 diff --git a/.github/workflows/pylint-check.yml b/.github/workflows/pylint-check.yml index 5e23d2812..02f4cb31e 100644 --- a/.github/workflows/pylint-check.yml +++ b/.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 diff --git a/.github/workflows/require-checklist.yml b/.github/workflows/require-checklist.yml index d15d19d99..4da288575 100644 --- a/.github/workflows/require-checklist.yml +++ b/.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 \ 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 4f1723cb2..b771ca7e9 100644 --- a/.github/workflows/unit-test.yml +++ b/.github/workflows/unit-test.yml @@ -3,7 +3,6 @@ name: Unit Test on: pull_request: - branches: [master] types: [opened, synchronize, edited] push: diff --git a/dags/examples/maxtext_aqtp_version_sweep_gke_example_dag.py b/dags/examples/maxtext_aqtp_version_sweep_gke_example_dag.py index deaeb3a17..aa484a44c 100644 --- 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( @@ -98,4 +102,4 @@ # Run jobs for test in tests: - test.run_with_run_name_generation() + test.run() diff --git a/dags/examples/maxtext_profile_namegen_example_dag.py b/dags/examples/maxtext_profile_namegen_example_dag.py index bf108062e..71f3ac1e4 100644 --- a/dags/examples/maxtext_profile_namegen_example_dag.py +++ b/dags/examples/maxtext_profile_namegen_example_dag.py @@ -12,15 +12,17 @@ # 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 @@ -28,36 +30,50 @@ 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, @@ -65,13 +81,14 @@ 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() diff --git a/dags/examples/maxtext_profile_sweep_example_dag.py b/dags/examples/maxtext_profile_sweep_example_dag.py index d6d6dc1d1..28e5ff05b 100644 --- a/dags/examples/maxtext_profile_sweep_example_dag.py +++ b/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() diff --git a/dags/examples/maxtext_sweep_gke_example_dag.py b/dags/examples/maxtext_sweep_gke_example_dag.py index 10d76babd..2d11d9e79 100644 --- 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 @@ -57,4 +63,4 @@ # Run jobs for test in maxtext_sweep_gke_test: - test.run_with_run_name_generation() + test.run() diff --git a/dags/multipod/configs/gke_config.py b/dags/multipod/configs/gke_config.py index da705b630..1876856fb 100644 --- a/dags/multipod/configs/gke_config.py +++ b/dags/multipod/configs/gke_config.py @@ -14,13 +14,13 @@ """Utilities to construct configs for maxtext DAG on GKE.""" -from dags.common import test_owner +import datetime +from typing import Any, Iterable + +from dags import gcs_bucket +from dags.common.vm_resource import Project, XpkClusters from xlml.apis import gcp_config, metric_config, task, test_config from xlml.apis.xpk_cluster_config import XpkClusterConfig -from dags import gcs_bucket -from dags.common.vm_resource import TpuVersion, Project, XpkClusters, GpuVersion, CpuVersion -from typing import Iterable -import datetime def get_gke_config( @@ -31,7 +31,9 @@ def get_gke_config( run_model_cmds: Iterable[str], cluster: XpkClusterConfig = XpkClusters.TPU_V4_8_MAXTEXT_CLUSTER, num_slices: int = 1, - dataset_name: metric_config.DatasetOption = metric_config.DatasetOption.XLML_DATASET, + dataset_name: metric_config.DatasetOption = ( + metric_config.DatasetOption.XLML_DATASET + ), dataset_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, composer_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, base_output_directory: str = None, @@ -81,6 +83,138 @@ def get_gke_config( ) +def get_gke_config_with_interrupt( + time_out_in_min: int, + test_name: str, + docker_image: str, + test_owner: str, + run_model_cmds: Iterable[str], + expect_reach_to_step: int, + cluster: XpkClusterConfig = XpkClusters.TPU_V4_8_MAXTEXT_CLUSTER, + num_slices: int = 1, + dataset_name: metric_config.DatasetOption = ( + metric_config.DatasetOption.XLML_DATASET + ), + dataset_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, + composer_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, + base_output_directory: str = None, + metric_aggregation_strategy: metric_config.AggregationStrategy = None, + user_specified_job_metric_config: metric_config.MetricConfig = None, + last_node: bool = False, + check_file_exists: bool = False, +) -> task.XpkNodeInterruptionTask: + job_gcp_config = gcp_config.GCPConfig( + project_name=cluster.project, + zone=cluster.zone, + dataset_name=dataset_name, + dataset_project=dataset_project, + composer_project=composer_project, + ) + + job_test_config = test_config.TpuGkeTest( + test_config.Tpu( + version=cluster.device_version, + cores=cluster.core_count, + ), + test_name=test_name, + run_model_cmds=run_model_cmds, + set_up_cmds=None, + timeout=datetime.timedelta(minutes=time_out_in_min), + task_owner=test_owner, + num_slices=num_slices, + cluster_name=cluster.name, + docker_image=docker_image, + ) + job_metric_config = user_specified_job_metric_config + if job_metric_config is None: + job_metric_config = ( + metric_config.MetricConfig( + tensorboard_summary=metric_config.SummaryConfig( + file_location=base_output_directory, + aggregation_strategy=metric_aggregation_strategy, + use_regex_file_location=True, + ), + ) + if base_output_directory and metric_aggregation_strategy + else None + ) + + return task.XpkNodeInterruptionTask( + task_test_config=job_test_config, + task_gcp_config=job_gcp_config, + task_metric_config=job_metric_config, + expect_reach_to_step=expect_reach_to_step, + last_node=last_node, + check_file_exists=check_file_exists, + ) + + +def get_gke_config_with_name_gen_and_quarantine( + time_out_in_min: int, + test_name: str, + docker_image: str, + test_owner: str, + run_model_cmds: Iterable[str], + cluster: XpkClusterConfig = XpkClusters.TPU_V4_8_MAXTEXT_CLUSTER, + num_slices: int = 1, + dataset_name: metric_config.DatasetOption = ( + metric_config.DatasetOption.XLML_DATASET + ), + dataset_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, + composer_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, + base_output_directory: str = None, + metric_aggregation_strategy: metric_config.AggregationStrategy = None, + user_specified_job_metric_config: metric_config.MetricConfig = None, + quarantine_task_group: Any = None, + run_name_env: str = "M_RUN_NAME", + nested_run_name_in_tb_file_location: bool = True, +) -> task.XpkNameGenAndQuarantineTask: + job_gcp_config = gcp_config.GCPConfig( + project_name=cluster.project, + zone=cluster.zone, + dataset_name=dataset_name, + dataset_project=dataset_project, + composer_project=composer_project, + ) + + job_test_config = test_config.TpuGkeTest( + test_config.Tpu( + version=cluster.device_version, + cores=cluster.core_count, + ), + test_name=test_name, + run_model_cmds=run_model_cmds, + set_up_cmds=None, + timeout=datetime.timedelta(minutes=time_out_in_min), + task_owner=test_owner, + num_slices=num_slices, + cluster_name=cluster.name, + docker_image=docker_image, + ) + job_metric_config = user_specified_job_metric_config + if job_metric_config is None: + job_metric_config = ( + metric_config.MetricConfig( + tensorboard_summary=metric_config.SummaryConfig( + file_location=base_output_directory, + aggregation_strategy=metric_aggregation_strategy, + use_regex_file_location=True, + ), + ) + if base_output_directory and metric_aggregation_strategy + else None + ) + + return task.XpkNameGenAndQuarantineTask( + task_test_config=job_test_config, + task_gcp_config=job_gcp_config, + task_metric_config=job_metric_config, + quarantine_task_group=quarantine_task_group, + run_name_env=run_name_env, + nested_run_name_in_tb_file_location=nested_run_name_in_tb_file_location, + ) + + def get_gke_maxtext_nightly_config( time_out_in_min: int, test_name: str, @@ -88,7 +222,9 @@ def get_gke_maxtext_nightly_config( test_owner: str, cluster: XpkClusterConfig = XpkClusters.TPU_V4_8_MAXTEXT_CLUSTER, num_slices: int = 1, - dataset_name: metric_config.DatasetOption = metric_config.DatasetOption.XLML_DATASET, + dataset_name: metric_config.DatasetOption = ( + metric_config.DatasetOption.XLML_DATASET + ), dataset_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, composer_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, ) -> task.XpkTask: @@ -106,18 +242,25 @@ def get_gke_maxtext_nightly_config( base_output_directory = ( f"{gcs_bucket.BASE_OUTPUT_DIR}/maxtext/nightly/automated/{current_date}" ) - run_name = f"{num_slices}slice-V{cluster.device_version.value}_{cluster.core_count}-maxtext-nightly-{current_datetime}" + run_name = ( + f"{num_slices}slice-V{cluster.device_version.value}_{cluster.core_count}" + f"-maxtext-nightly-{current_datetime}" + ) run_model_cmds = ( "bash src/dependencies/scripts/preflight.sh PLATFORM=GKE", ( "JAX_PLATFORM_NAME=TPU XLA_FLAGS='--xla_dump_to=/tmp/xla_dump/'" " ENABLE_PJRT_COMPATIBILITY=true" - f" python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml run_name={run_name}" - f" base_output_directory={base_output_directory}" + " python3 -m maxtext.trainers.pre_train.train" + " src/maxtext/configs/base.yml" + f" run_name={run_name} base_output_directory={base_output_directory}" " dataset_path=gs://max-datasets-rogue dataset_type=synthetic" - " model_name=llama3-8b per_device_batch_size=12 reuse_example_batch=1 metrics_file='metrics.txt'" - " steps=50 enable_checkpointing=false profiler=xplane upload_all_profiler_results=true skip_first_n_steps_for_profiler=10 profiler_steps=10 gcs_metrics=true" + " model_name=llama3-8b per_device_batch_size=12 reuse_example_batch=1" + " metrics_file='metrics.txt' steps=50 enable_checkpointing=false" + " profiler=xplane upload_all_profiler_results=true" + " skip_first_n_steps_for_profiler=10 profiler_steps=10" + " gcs_metrics=true" ), ) @@ -188,7 +331,9 @@ def get_gke_gpt3_6b_nightly_config( test_owner: str, cluster: XpkClusterConfig = XpkClusters.TPU_V4_8_MAXTEXT_CLUSTER, num_slices: int = 1, - dataset_name: metric_config.DatasetOption = metric_config.DatasetOption.XLML_DATASET, + dataset_name: metric_config.DatasetOption = ( + metric_config.DatasetOption.XLML_DATASET + ), dataset_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, composer_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, ) -> task.XpkTask: @@ -206,18 +351,26 @@ def get_gke_gpt3_6b_nightly_config( base_output_directory = ( f"{gcs_bucket.BASE_OUTPUT_DIR}/maxtext/nightly/automated/{current_date}" ) - run_name = f"{num_slices}slice-V{cluster.device_version.value}_{cluster.core_count}-gpt3-6b-nightly-{current_datetime}" + run_name = ( + f"{num_slices}slice-V{cluster.device_version.value}_{cluster.core_count}" + f"-gpt3-6b-nightly-{current_datetime}" + ) run_model_cmds = ( "bash src/dependencies/scripts/preflight.sh PLATFORM=GKE", ( "JAX_PLATFORM_NAME=TPU XLA_FLAGS='--xla_dump_to=/tmp/xla_dump/'" " ENABLE_PJRT_COMPATIBILITY=true" - f" python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml run_name={run_name} model_name=gpt3-6b" + " python3 -m maxtext.trainers.pre_train.train" + " src/maxtext/configs/base.yml" + f" run_name={run_name} model_name=gpt3-6b" f" base_output_directory={base_output_directory}" " dataset_path=gs://max-datasets-rogue dataset_type=synthetic" - " per_device_batch_size=12 reuse_example_batch=1 global_parameter_scale=1 metrics_file='metrics.txt'" - " steps=50 enable_checkpointing=false profiler=xplane upload_all_profiler_results=true skip_first_n_steps_for_profiler=10 profiler_steps=10 gcs_metrics=true" + " per_device_batch_size=12 reuse_example_batch=1" + " global_parameter_scale=1 metrics_file='metrics.txt' steps=50" + " enable_checkpointing=false profiler=xplane" + " upload_all_profiler_results=true skip_first_n_steps_for_profiler=10" + " profiler_steps=10 gcs_metrics=true" ), ) @@ -251,7 +404,9 @@ def get_maxtext_cpu_end_to_end_gke_config( cluster: XpkClusterConfig = XpkClusters.CPU_N2_STANDARD_64_CLUSTER, machine_count: int = 1, num_slices: int = 1, - dataset_name: metric_config.DatasetOption = metric_config.DatasetOption.XLML_DATASET, + dataset_name: metric_config.DatasetOption = ( + metric_config.DatasetOption.XLML_DATASET + ), dataset_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, composer_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, base_output_directory: str = None, diff --git a/dags/multipod/configs/maxtext_sweep_gke_config.py b/dags/multipod/configs/maxtext_sweep_gke_config.py index 8cdeb69a0..61182aad2 100644 --- a/dags/multipod/configs/maxtext_sweep_gke_config.py +++ b/dags/multipod/configs/maxtext_sweep_gke_config.py @@ -16,7 +16,7 @@ import datetime from xlml.apis import gcp_config, metric_config, task, test_config -from dags.common.vm_resource import TpuVersion, XpkClusterConfig +from dags.common.vm_resource import XpkClusterConfig import itertools from typing import List, Iterable, Dict, Any @@ -49,12 +49,17 @@ def get_maxtext_sweep_gke_config( docker_image: str, base_output_directory: str, base_run_model_cmds: Iterable[str], - dataset_name: metric_config.DatasetOption = metric_config.DatasetOption.BENCHMARK_DATASET, - metric_aggregation_strategy: metric_config.AggregationStrategy = metric_config.AggregationStrategy.MEDIAN, + dataset_name: metric_config.DatasetOption = ( + metric_config.DatasetOption.BENCHMARK_DATASET + ), + metric_aggregation_strategy: metric_config.AggregationStrategy = ( + metric_config.AggregationStrategy.MEDIAN + ), dataset_project: str = None, composer_project: str = None, enable_profile_config: bool = False, -) -> List[task.XpkTask]: + quarantine_task_group: Any = None, +) -> List[task.XpkNameGenAndQuarantineTask]: if not dataset_project: dataset_project = cluster.project if not composer_project: @@ -76,7 +81,8 @@ def get_maxtext_sweep_gke_config( for param, values in sweep_params.items(): sweep_params_list.append([(param, val) for val in values]) - # Generate all combinations of sweep param configurations and create a XpkTask for each one + # Generate all combinations of sweep param configurations + # and create a XpkTask for each one xpk_task_list = [] for idx, config in enumerate(itertools.product(*sweep_params_list)): config_dict = {key: value for (key, value) in config} @@ -122,10 +128,11 @@ def get_maxtext_sweep_gke_config( file_location=base_output_directory, ) - xpk_task = task.XpkTask( + xpk_task = task.XpkNameGenAndQuarantineTask( task_test_config=job_test_config, task_gcp_config=job_gcp_config, task_metric_config=job_metric_config, + quarantine_task_group=quarantine_task_group, ) xpk_task_list.append(xpk_task) diff --git a/dags/multipod/maxtext_convergence.py b/dags/multipod/maxtext_convergence.py index fcb7d9f50..5340bda6a 100644 --- a/dags/multipod/maxtext_convergence.py +++ b/dags/multipod/maxtext_convergence.py @@ -84,10 +84,11 @@ 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 = gke_config.get_gke_config( + test_task = gke_config.get_gke_config_with_name_gen_and_quarantine( cluster=XpkClusters.TPU_V6E_256_MLPERF_CLUSTER, time_out_in_min=timeout_in_min, test_name=test_name, @@ -96,7 +97,7 @@ test_owner=test_owner.MATT_D, base_output_directory=base_output_directory, metric_aggregation_strategy=metric_config.AggregationStrategy.LAST, - ).run_with_run_name_generation() + ).run() if test_name not in parallel_test_names: sequential_tests.append(test_task) diff --git a/dags/multipod/maxtext_e2e_tpu_post_training.py b/dags/multipod/maxtext_e2e_tpu_post_training.py index 352504448..659604fba 100644 --- a/dags/multipod/maxtext_e2e_tpu_post_training.py +++ b/dags/multipod/maxtext_e2e_tpu_post_training.py @@ -123,7 +123,7 @@ def get_workload_name(model, mode, length=6): run_model_cmds=training_cmd, docker_image="{{ params.docker_image }}", test_owner=test_owner.SURBHI_J, - ).run_model( + ).run( use_pathways=True, ) diff --git a/dags/orbax/maxtext_emc_restore_gcs.py b/dags/orbax/maxtext_emc_restore_gcs.py index 28420171b..1edaadcca 100644 --- a/dags/orbax/maxtext_emc_restore_gcs.py +++ b/dags/orbax/maxtext_emc_restore_gcs.py @@ -120,7 +120,7 @@ task_id="generate_start_time" )() - maxtext_chkpt_run_test = gke_config.get_gke_config( + maxtext_chkpt_run_test = gke_config.get_gke_config_with_interrupt( num_slices=slice_num, cluster=test_config.cluster, time_out_in_min=60, @@ -128,12 +128,12 @@ run_model_cmds=workload_command, docker_image=image.value, test_owner=test_owner.DEPP_L, - ).run_with_node_interruption( + expect_reach_to_step=step_to_interrupt, + last_node=True, + ).run( ramdisk_directory=test_config_util.DEFAULT_RAM_DISK, mtc_enabled=True, skip_post_process=True, - last_node=True, - expect_reach_to_step=step_to_interrupt, max_restart=15, ) diff --git a/dags/orbax/maxtext_emc_restore_local.py b/dags/orbax/maxtext_emc_restore_local.py index 7e3c04b3d..a1372fac9 100644 --- a/dags/orbax/maxtext_emc_restore_local.py +++ b/dags/orbax/maxtext_emc_restore_local.py @@ -118,7 +118,7 @@ task_id="generate_start_time" )() - maxtext_chkpt_run_test = gke_config.get_gke_config( + maxtext_chkpt_run_test = gke_config.get_gke_config_with_interrupt( num_slices=slice_num, cluster=test_config.cluster, time_out_in_min=60, @@ -126,11 +126,11 @@ run_model_cmds=workload_command, docker_image=image.value, test_owner=test_owner.DEPP_L, - ).run_with_node_interruption( + expect_reach_to_step=step_to_interrupt, + ).run( ramdisk_directory=test_config_util.DEFAULT_RAM_DISK, mtc_enabled=True, skip_post_process=True, - expect_reach_to_step=step_to_interrupt, max_restart=15, ) diff --git a/dags/orbax/maxtext_mtc_restore_local.py b/dags/orbax/maxtext_mtc_restore_local.py index 66e8f445f..5548b6f80 100644 --- a/dags/orbax/maxtext_mtc_restore_local.py +++ b/dags/orbax/maxtext_mtc_restore_local.py @@ -121,7 +121,7 @@ task_id="generate_start_time" )() - maxtext_chkpt_run_test = gke_config.get_gke_config( + maxtext_chkpt_run_test = gke_config.get_gke_config_with_interrupt( num_slices=slice_num, cluster=test_config.cluster, time_out_in_min=60, @@ -129,11 +129,11 @@ run_model_cmds=workload_command, docker_image=image.value, test_owner=test_owner.DEPP_L, - ).run_with_node_interruption( + expect_reach_to_step=step_to_interrupt, + ).run( ramdisk_directory=test_config_util.DEFAULT_RAM_DISK, mtc_enabled=True, skip_post_process=True, - expect_reach_to_step=step_to_interrupt, max_restart=15, ) diff --git a/dags/orbax/maxtext_reg_restore_gcs_with_node_disruption.py b/dags/orbax/maxtext_reg_restore_gcs_with_node_disruption.py index e94cbdbed..5e9ef528d 100644 --- a/dags/orbax/maxtext_reg_restore_gcs_with_node_disruption.py +++ b/dags/orbax/maxtext_reg_restore_gcs_with_node_disruption.py @@ -130,7 +130,7 @@ def generate_workload_checkpoints_location(gcs_ckpt_location: str) -> str: task_id="generate_start_time" )() - maxtext_chkpt_run_test = gke_config.get_gke_config( + maxtext_chkpt_run_test = gke_config.get_gke_config_with_interrupt( num_slices=slice_num, cluster=test_config.cluster, time_out_in_min=60, @@ -138,13 +138,13 @@ def generate_workload_checkpoints_location(gcs_ckpt_location: str) -> str: run_model_cmds=workload_command, docker_image=image.value, test_owner=test_owner.SHARON_Y, - ).run_with_node_interruption( + expect_reach_to_step=step_to_interrupt, + check_file_exists=True, + ).run( gcs_location=gcs_location, xpk_branch=MAIN_BRANCH, skip_post_process=True, - expect_reach_to_step=step_to_interrupt, max_restart=15, - check_file_exists=True, ) end_time = validation_util.generate_timestamp.override( diff --git a/scripts/code-style.sh b/scripts/code-style.sh index e02e53c4e..5ade7a620 100755 --- a/scripts/code-style.sh +++ b/scripts/code-style.sh @@ -22,14 +22,37 @@ FOLDERS_TO_FORMAT=("dags" "xlml") echo "[code-style] Running Semgrep static analysis..." semgrep --config scripts/semgrep-rules.yaml --error -for folder in "${FOLDERS_TO_FORMAT[@]}" -do - pyink "$folder" --pyink-indentation=2 --pyink-use-majority-quotes --line-length=80 --check --diff -done - -for folder in "${FOLDERS_TO_FORMAT[@]}" -do - pylint "./$folder" --fail-under=9.6 -done +HEAD_SHA="$(git rev-parse HEAD)" +BASE_BRANCH="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." diff --git a/xlml/apis/task.py b/xlml/apis/task.py index 9877f2bf3..8b875834a 100644 --- a/xlml/apis/task.py +++ b/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: @@ -335,7 +338,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 ) @@ -363,117 +367,51 @@ def run( post_process. """ with TaskGroup(group_id=self.task_test_config.benchmark_id) as group: - run_model, gcs_path = self.run_model( - gcs_location, - use_vertex_tensorboard, - use_pathways, - ramdisk_directory, - mtc_enabled, - xpk_branch, - max_restart, - ) - if not skip_post_process: - _ = run_model >> self.post_process(gcs_path) - - return group - - def run_with_node_interruption( - self, - *, - gcs_location: Optional[airflow.XComArg] = None, - use_vertex_tensorboard: bool = False, - expect_reach_to_step: int, - use_pathways: bool = False, - skip_post_process: bool = False, - ramdisk_directory: str = "", - mtc_enabled: bool = False, - xpk_branch: str = xpk.MAIN_BRANCH, - last_node: bool = False, - max_restart: int = 0, - check_file_exists: bool = False, - ) -> DAGNode: - """Run a test job within a docker image. - - Will run a workload with an injected interruption of a GKE node. - Then is expected to automatically restart and continuing running. - - Attributes: - gcs_location: GCS path for all artifacts of the test. - use_vertex_tensorboard: Set to True to view workload data on - Vertex AI Tensorboard. - expect_reach_to_step: The training step at which the node interruption - should be triggered. - use_pathways: Set to True to use the Pathways execution framework. - skip_post_process: If True, the post processing step will be skipped. - ramdisk_directory: The directory for enabling emergency checkpointing. - mtc_enabled: Set to True to enable Multi-tier Checkpointing (MTC). - xpk_branch: The specific git branch of the xpk tool to use. - last_node: If True, the interruption will target the last node in the - workload; otherwise, it targets the first node. - max_restart: By default, this is 0. - This will restart the job with flag "--max-restarts" - check_file_exists: By default, this is False. If set to True, - task branch task_path_decider will be performed. - Returns: - A task group with the following tasks chained: run_model and - post_process. - """ - with TaskGroup(group_id=self.task_test_config.benchmark_id) as group: - run_model, gcs_path = self.run_model_with_node_interruption( + run_model, gcs_path = self._run_model( gcs_location=gcs_location, use_vertex_tensorboard=use_vertex_tensorboard, - expect_reach_to_step=expect_reach_to_step, use_pathways=use_pathways, ramdisk_directory=ramdisk_directory, mtc_enabled=mtc_enabled, xpk_branch=xpk_branch, - last_node=last_node, max_restart=max_restart, - check_file_exists=check_file_exists, ) if not skip_post_process: - _ = run_model >> self.post_process(gcs_path) + _ = run_model >> self._post_process(gcs_path) + return group - def run_model_with_node_interruption( + def _run_model( self, - *, gcs_location: Optional[airflow.XComArg] = None, use_vertex_tensorboard: bool = False, - expect_reach_to_step: int, use_pathways: bool = False, ramdisk_directory: str = "", mtc_enabled: bool = False, xpk_branch: str = xpk.MAIN_BRANCH, - last_node: bool = False, max_restart: int = 0, - check_file_exists: bool = False, ) -> DAGNode: - """Run the TPU/GPU test in `task_test_config` using xpk. + """Run a test job within a docker image. - Different behavior for testing node interruption. + Will run a workload with an injected interruption of a GKE node. + Then is expected to automatically restart and continuing running. Attributes: gcs_location: GCS path for all artifacts of the test. use_vertex_tensorboard: Set to True to view workload data on Vertex AI Tensorboard. - expect_reach_to_step: The training step at which the node interruption - should be triggered. use_pathways: Set to True to use the Pathways execution framework. ramdisk_directory: The directory for enabling emergency checkpointing. mtc_enabled: Set to True to enable Multi-tier Checkpointing (MTC). xpk_branch: The specific git branch of the xpk tool to use. - last_node: If True, the interruption will target the last node in the - workload; otherwise, it targets the first node. max_restart: By default, this is 0. This will restart the job with flag "--max-restarts" - check_file_exists: By default, this is False. If set to True, - task branch task_path_decider will be performed. Returns: A DAG node that executes the model test. """ with TaskGroup(group_id="run_model") as group: workload_id = xpk.generate_workload_id(self.task_test_config.benchmark_id) + if gcs_location: gcs_path = gcs_location else: @@ -486,74 +424,68 @@ def run_model_with_node_interruption( task_id="dummy_op_for_teardown" ).as_setup() - launch_workload_and_wait_for_reach_step = ( - self.launch_workload_with_node_reach_to_step( - workload_id, - gcs_path, - expect_reach_to_step, - use_vertex_tensorboard, - use_pathways, - ramdisk_directory, - mtc_enabled, - xpk_branch, - max_restart, - check_file_exists, - ) + launch_workload = self._launch_workload( + workload_id, + gcs_path, + use_vertex_tensorboard, + use_pathways, + ramdisk_directory, + mtc_enabled, + xpk_branch, + max_restart, ) - run_node_interruption = xpk.delete_node.override( - owner=self.task_test_config.task_owner, trigger_rule="none_failed" - )( - project=self.task_gcp_config.project_name, - zone=self.task_gcp_config.zone, - cluster_name=self.task_test_config.cluster_name, - workload_id=workload_id, - dry_run=False, - last_node=last_node, + wait_for_workload_completion = self._get_wait_for_workload_completion( + workload_id ) - wait_for_workload_completion = xpk.wait_for_workload_completion.override( - timeout=int(self.task_test_config.timeout.total_seconds()), - )( - workload_id=workload_id, - project_id=self.task_gcp_config.project_name, - region=gke.zone_to_region(self.task_gcp_config.zone), - cluster_name=self.task_test_config.cluster_name, - ) - - clean_up_workload = xpk.clean_up_workload( - workload_id=workload_id, - project_id=self.task_gcp_config.project_name, - zone=self.task_gcp_config.zone, - cluster_name=self.task_test_config.cluster_name, - xpk_branch=xpk_branch, + clean_up_workload = self._get_clean_up_workload( + workload_id, xpk_branch ).as_teardown(setups=dummy_op_for_teardown, on_failure_fail_dagrun=True) _ = ( (workload_id, gcs_path) >> dummy_op_for_teardown - >> launch_workload_and_wait_for_reach_step - >> run_node_interruption + >> launch_workload >> wait_for_workload_completion >> clean_up_workload ) - return group, gcs_path + return group, gcs_path - def launch_workload_with_node_reach_to_step( + def _get_wait_for_workload_completion(self, workload_id: str) -> DAGNode: + return xpk.wait_for_workload_completion.override( + timeout=int(self.task_test_config.timeout.total_seconds()), + )( + workload_id=workload_id, + project_id=self.task_gcp_config.project_name, + region=gke.zone_to_region(self.task_gcp_config.zone), + cluster_name=self.task_test_config.cluster_name, + ) + + def _get_clean_up_workload( + self, workload_id: str, xpk_branch: str + ) -> DAGNode: + return xpk.clean_up_workload( + workload_id=workload_id, + project_id=self.task_gcp_config.project_name, + zone=self.task_gcp_config.zone, + cluster_name=self.task_test_config.cluster_name, + xpk_branch=xpk_branch, + ) + + def _launch_workload( self, workload_id: str, gcs_path: str, - expect_reach_to_step: int, use_vertex_tensorboard: bool, use_pathways: bool = False, ramdisk_directory: str = "", mtc_enabled: bool = False, xpk_branch: str = xpk.MAIN_BRANCH, max_restart: int = 0, - check_file_exists: bool = False, ) -> DAGNode: """Create the workload and wait for it to provision.""" - with TaskGroup(group_id="launch_workload_with_node_reach_to_step") as group: + with TaskGroup(group_id="launch_workload") as group: run_workload = xpk.run_workload.override( owner=self.task_test_config.task_owner )( @@ -583,6 +515,136 @@ def launch_workload_with_node_reach_to_step( region=gke.zone_to_region(self.task_gcp_config.zone), cluster_name=self.task_test_config.cluster_name, ) + _ = run_workload >> wait_for_workload_start + return group + + def _post_process(self, result_location: Optional[str] = None) -> DAGNode: + """Process metrics and metadata, and insert them into BigQuery tables. + + Returns: + A DAG node that executes the post process. + """ + with TaskGroup(group_id="post_process") as group: + process_id = metric.generate_process_id.override(retries=0)() + post_process_metrics = metric.process_metrics.override(retries=0)( + process_id, + self.task_test_config, + self.task_metric_config, + self.task_gcp_config, + folder_location=result_location, + ) + + if self.task_metric_config and self.task_metric_config.profile: + self.task_metric_config.profile.metrics = ( + metric.xplane_to_metrics.override(retries=0)( + self.task_metric_config.profile.file_location + ) + ) + _ = ( + process_id + >> self.task_metric_config.profile.metrics + >> post_process_metrics + ) + else: + _ = process_id >> post_process_metrics + + return group + + +@dataclasses.dataclass +class XpkNodeInterruptionTask(XpkTask): + """Task for running XPK workloads with node interruption.""" + + expect_reach_to_step: int = 0 + last_node: bool = False + check_file_exists: bool = False + + def _run_model( + self, + gcs_location: Optional[airflow.XComArg] = None, + use_vertex_tensorboard: bool = False, + use_pathways: bool = False, + ramdisk_directory: str = "", + mtc_enabled: bool = False, + xpk_branch: str = xpk.MAIN_BRANCH, + max_restart: int = 0, + ) -> DAGNode: + """Run a test job within a docker image. + + Will run a workload with an injected interruption of a GKE node. + Then is expected to automatically restart and continuing running. + + Attributes: + gcs_location: GCS path for all artifacts of the test. + use_vertex_tensorboard: Set to True to view workload data on + Vertex AI Tensorboard. + use_pathways: Set to True to use the Pathways execution framework. + ramdisk_directory: The directory for enabling emergency checkpointing. + mtc_enabled: Set to True to enable Multi-tier Checkpointing (MTC). + xpk_branch: The specific git branch of the xpk tool to use. + max_restart: By default, this is 0. + This will restart the job with flag "--max-restarts" + Returns: + A DAG node that executes the model test. + """ + with TaskGroup(group_id="run_model") as group: + workload_id = xpk.generate_workload_id(self.task_test_config.benchmark_id) + + if gcs_location: + gcs_path = gcs_location + else: + gcs_path = name_format.generate_gcs_folder_location( + self.task_test_config.gcs_subfolder, + self.task_test_config.benchmark_id, + ) + + dummy_op_for_teardown = EmptyOperator( + task_id="dummy_op_for_teardown" + ).as_setup() + + launch_workload = self._launch_workload( + workload_id, + gcs_path, + use_vertex_tensorboard, + use_pathways, + ramdisk_directory, + mtc_enabled, + xpk_branch, + max_restart, + ) + + wait_for_workload_reach_step_and_interruption = ( + self._wait_for_workload_reach_step_and_interruption( + workload_id, gcs_path + ) + ) + + wait_for_workload_completion = self._get_wait_for_workload_completion( + workload_id + ) + + clean_up_workload = self._get_clean_up_workload( + workload_id, xpk_branch + ).as_teardown(setups=dummy_op_for_teardown, on_failure_fail_dagrun=True) + + _ = ( + dummy_op_for_teardown + >> launch_workload + >> wait_for_workload_reach_step_and_interruption + >> wait_for_workload_completion + >> clean_up_workload + ) + return group, gcs_path + + def _wait_for_workload_reach_step_and_interruption( + self, + workload_id: str, + gcs_path: str, + ) -> DAGNode: + """Wait for workload to reach specific step and trigger node deletion.""" + with TaskGroup( + group_id="wait_for_workload_reach_step_and_interruption" + ) as group: wait_for_workload_to_reach_step = ( xpk.wait_for_workload_reach_step.override( task_id="wait_for_workload_reach_step" @@ -591,7 +653,7 @@ def launch_workload_with_node_reach_to_step( project_id=self.task_gcp_config.project_name, region=gke.zone_to_region(self.task_gcp_config.zone), cluster_name=self.task_test_config.cluster_name, - expect_reach_to_step=str(expect_reach_to_step), + expect_reach_to_step=str(self.expect_reach_to_step), ) ) @@ -599,15 +661,17 @@ 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(self.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}" @@ -615,59 +679,64 @@ def task_path_decider(check_file_exists: bool = False) -> str: # Conditional checks: depending on the `check_file_exists` argument # specified by the upper-level caller. - maybe_check_file_exists = task_path_decider(check_file_exists) + maybe_check_file_exists = task_path_decider(self.check_file_exists) - _ = ( - run_workload - >> wait_for_workload_start - >> wait_for_workload_to_reach_step - >> maybe_check_file_exists + run_node_interruption = xpk.delete_node.override( + owner=self.task_test_config.task_owner, trigger_rule="none_failed" + )( + project=self.task_gcp_config.project_name, + zone=self.task_gcp_config.zone, + cluster_name=self.task_test_config.cluster_name, + workload_id=workload_id, + dry_run=False, + last_node=self.last_node, ) + + _ = wait_for_workload_to_reach_step >> maybe_check_file_exists _ = maybe_check_file_exists >> [wait_for_file_to_exist, do_nothing] + _ = [wait_for_file_to_exist, do_nothing] >> run_node_interruption - return group + return group + + +@dataclasses.dataclass +class XpkNameGenAndQuarantineTask(XpkTask): + """Task for running XPK workloads with name generation and quarantine.""" + + quarantine_task_group: Any = None + run_name_env: str = "M_RUN_NAME" + nested_run_name_in_tb_file_location: bool = True - def run_with_name_gen_and_quarantine( + def run( self, - quarantine_task_group, use_pathways: bool = False, xpk_branch: str = xpk.MAIN_BRANCH, - run_name_env: str = "M_RUN_NAME", - nested_run_name_in_tb_file_location: bool = True, ) -> DAGNode: + """Generate a unique run name, tensorboard file location, + and profile file location (if metric config has profile), + then run a test job within a docker image. + + Returns: + A task group with the following tasks chained: generate_run_name, + generate_tb_file_location, generate_profile_file_location (optional), + run provision, run_model, post_process. + """ test_name = self.task_test_config.benchmark_id if QuarantineTests.is_quarantined(test_name): - with quarantine_task_group: - return self.run_with_run_name_generation( + with self.quarantine_task_group: + return self._run_with_run_name_generation( use_pathways, xpk_branch, - run_name_env, - nested_run_name_in_tb_file_location, ) else: - return self.run_with_run_name_generation( + return self._run_with_run_name_generation( use_pathways, xpk_branch, - run_name_env, - nested_run_name_in_tb_file_location, ) - def run_with_run_name_generation( - self, - use_pathways: bool = False, - xpk_branch: str = xpk.MAIN_BRANCH, - run_name_env: str = "M_RUN_NAME", - nested_run_name_in_tb_file_location: bool = True, + def _run_with_run_name_generation( + self, use_pathways: bool = False, xpk_branch: str = xpk.MAIN_BRANCH ) -> DAGNode: - """Generate a unique run name, tensorboard file location, - and profile file location (if metric config has profile), - then run a test job within a docker image. - - Returns: - A task group with the following tasks chained: generate_run_name, - generate_tb_file_location, generate_profile_file_location (optional), - run provision, run_model, post_process. - """ with TaskGroup( group_id=self.task_test_config.benchmark_id, prefix_group_id=True ) as group: @@ -677,11 +746,11 @@ def run_with_run_name_generation( tb_file_location = name_format.generate_tb_file_location( run_name, self.task_metric_config.tensorboard_summary.file_location, - nested_run_name_in_tb_file_location, + self.nested_run_name_in_tb_file_location, ) # Set run_name in run_model_cmds - new_run_model_cmds = [f"export {run_name_env}={run_name}"] + new_run_model_cmds = [f"export {self.run_name_env}={run_name}"] for cmd in self.task_test_config.run_model_cmds: new_run_model_cmds.append(cmd) self.task_test_config.run_model_cmds = new_run_model_cmds @@ -697,171 +766,26 @@ def run_with_run_name_generation( run_name, self.task_metric_config.profile.file_location ) self.task_metric_config.profile.file_location = profile_file_location - run_model, gcs_path = self.run_model( + + run_model, gcs_path = self._run_model( use_pathways=use_pathways, xpk_branch=xpk_branch ) _ = ( run_name >> (tb_file_location, profile_file_location) >> run_model - >> self.post_process(gcs_path) + >> self._post_process(gcs_path) ) else: - run_model, gcs_path = self.run_model( + run_model, gcs_path = self._run_model( use_pathways=use_pathways, xpk_branch=xpk_branch ) _ = ( run_name >> tb_file_location >> run_model - >> self.post_process(gcs_path) - ) - return group - - def run_model( - self, - gcs_location: Optional[airflow.XComArg] = None, - use_vertex_tensorboard: bool = False, - use_pathways: bool = False, - ramdisk_directory: str = "", - mtc_enabled: bool = False, - xpk_branch: str = xpk.MAIN_BRANCH, - max_restart: int = 0, - ) -> DAGNode: - """Run the TPU/GPU test in `task_test_config` using xpk. - - Attributes: - gcs_location: GCS path for all artifacts of the test. - use_vertex_tensorboard: Set to True to view workload data on - Vertex AI Tensorboard. - - Returns: - A DAG node that executes the model test. - """ - with TaskGroup(group_id="run_model") as group: - workload_id = xpk.generate_workload_id(self.task_test_config.benchmark_id) - if gcs_location: - gcs_path = gcs_location - else: - gcs_path = name_format.generate_gcs_folder_location( - self.task_test_config.gcs_subfolder, - self.task_test_config.benchmark_id, - ) - - dummy_op_for_teardown = EmptyOperator( - task_id="dummy_op_for_teardown" - ).as_setup() - - launch_workload = self.launch_workload( - workload_id, - gcs_path, - use_vertex_tensorboard, - use_pathways, - ramdisk_directory, - mtc_enabled, - xpk_branch, - max_restart, - ) - wait_for_workload_completion = xpk.wait_for_workload_completion.override( - timeout=int(self.task_test_config.timeout.total_seconds()), - )( - workload_id=workload_id, - project_id=self.task_gcp_config.project_name, - region=gke.zone_to_region(self.task_gcp_config.zone), - cluster_name=self.task_test_config.cluster_name, - ) - - clean_up_workload = xpk.clean_up_workload( - workload_id=workload_id, - project_id=self.task_gcp_config.project_name, - zone=self.task_gcp_config.zone, - cluster_name=self.task_test_config.cluster_name, - xpk_branch=xpk_branch, - ).as_teardown(setups=dummy_op_for_teardown, on_failure_fail_dagrun=True) - - _ = ( - (workload_id, gcs_path) - >> dummy_op_for_teardown - >> launch_workload - >> wait_for_workload_completion - >> clean_up_workload - ) - return group, gcs_path - - def launch_workload( - self, - workload_id: str, - gcs_path: str, - use_vertex_tensorboard: bool, - use_pathways: bool = False, - ramdisk_directory: str = "", - mtc_enabled: bool = False, - xpk_branch: str = xpk.MAIN_BRANCH, - max_restart: int = 0, - ) -> DAGNode: - """Create the workload and wait for it to provision.""" - with TaskGroup(group_id="launch_workload") as group: - run_workload = xpk.run_workload.override( - owner=self.task_test_config.task_owner - )( - task_id="run_workload", - cluster_project=self.task_gcp_config.project_name, - zone=self.task_gcp_config.zone, - cluster_name=self.task_test_config.cluster_name, - benchmark_id=self.task_test_config.benchmark_id, - workload_id=workload_id, - gcs_path=gcs_path, - docker_image=self.task_test_config.docker_image, - accelerator_type=self.task_test_config.accelerator.name, - run_cmds=self.task_test_config.test_script, - num_slices=self.task_test_config.num_slices, - use_vertex_tensorboard=use_vertex_tensorboard, - use_pathways=use_pathways, - ramdisk_directory=ramdisk_directory, - mtc_enabled=mtc_enabled, - xpk_branch=xpk_branch, - max_restart=max_restart, - ) - wait_for_workload_start = xpk.wait_for_workload_start.override( - timeout=self.workload_provision_timeout.total_seconds() - )( - workload_id=workload_id, - project_id=self.task_gcp_config.project_name, - region=gke.zone_to_region(self.task_gcp_config.zone), - cluster_name=self.task_test_config.cluster_name, - ) - _ = run_workload >> wait_for_workload_start - return group - - def post_process(self, result_location: Optional[str] = None) -> DAGNode: - """Process metrics and metadata, and insert them into BigQuery tables. - - Returns: - A DAG node that executes the post process. - """ - with TaskGroup(group_id="post_process") as group: - process_id = metric.generate_process_id.override(retries=0)() - post_process_metrics = metric.process_metrics.override(retries=0)( - process_id, - self.task_test_config, - self.task_metric_config, - self.task_gcp_config, - folder_location=result_location, - ) - - if self.task_metric_config and self.task_metric_config.profile: - self.task_metric_config.profile.metrics = ( - metric.xplane_to_metrics.override(retries=0)( - self.task_metric_config.profile.file_location - ) + >> self._post_process(gcs_path) ) - _ = ( - process_id - >> self.task_metric_config.profile.metrics - >> post_process_metrics - ) - else: - _ = process_id >> post_process_metrics return group