Skip to content

Commit 1ace5db

Browse files
SurbhiJainUSCSurbhiJainUSC
authored
Add MaxText E2E tests DAG triggered by Github Actions workflow (GoogleCloudPlatform#1275)
1 parent 580460d commit 1ace5db

4 files changed

Lines changed: 357 additions & 0 deletions

File tree

‎dags/common/scheduling_helper/dag_integrity_whitelist.yaml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,9 @@ whitelisted_dags:
5151
- maxtext_configs_aot_gpu
5252
- maxtext_configs_aot_hybridsim
5353
- maxtext_convergence
54+
- maxtext_e2e_tests
55+
- maxtext_e2e_tpu_post_training
56+
- maxtext_e2e_tpu_pre_training
5457
- maxtext_emc_and_mtc_orbax_save_local
5558
- maxtext_emc_orbax_res_gcs
5659
- maxtext_emc_orbax_res_local

‎dags/multipod/maxtext_e2e_tests.py‎

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""
16+
Parent DAG that triggers MaxText E2E TPU pre-training and post-training child
17+
DAGs in parallel, waits for both to complete, then fires a GitHub
18+
repository_dispatch callback with the aggregated result.
19+
"""
20+
import datetime
21+
import requests
22+
from airflow import models
23+
from airflow.models.param import Param
24+
from airflow.operators.python import PythonOperator
25+
from airflow.operators.trigger_dagrun import TriggerDagRunOperator
26+
from airflow.utils.trigger_rule import TriggerRule
27+
28+
with models.DAG(
29+
dag_id="maxtext_e2e_tests",
30+
schedule=None,
31+
tags=[
32+
"maxtext",
33+
"e2e",
34+
"pre-training",
35+
"post-training",
36+
],
37+
start_date=datetime.datetime(2026, 6, 10),
38+
catchup=False,
39+
params={
40+
"build_mode": Param(
41+
type="string",
42+
description="Build mode: stable or nightly",
43+
),
44+
"maxtext_sha": Param(
45+
type="string",
46+
description="Commit SHA being tested",
47+
),
48+
"github_run_id": Param(
49+
type="string",
50+
description="GitHub Actions run ID of the original build workflow",
51+
),
52+
"github_repo": Param(
53+
type="string",
54+
description="GitHub repository in owner/repo format",
55+
),
56+
"github_callback_token": Param(
57+
type="string",
58+
description="GitHub PAT used to fire the repository_dispatch callback",
59+
),
60+
},
61+
) as dag:
62+
trigger_pre_training = TriggerDagRunOperator(
63+
task_id="trigger_tpu_pre_training",
64+
trigger_dag_id="maxtext_e2e_tpu_pre_training",
65+
conf={
66+
"docker_image": "gcr.io/tpu-prod-env-multipod/maxtext_jax_{{ params.build_mode }}:{{ params.github_run_id }}"
67+
},
68+
wait_for_completion=True,
69+
poke_interval=600, # check every 10 minutes for child DAG completion
70+
)
71+
72+
trigger_post_training = TriggerDagRunOperator(
73+
task_id="trigger_tpu_post_training",
74+
trigger_dag_id="maxtext_e2e_tpu_post_training",
75+
conf={
76+
"docker_image": "gcr.io/tpu-prod-env-multipod/maxtext_post_training_{{ params.build_mode }}:{{ params.github_run_id }}"
77+
},
78+
wait_for_completion=True,
79+
poke_interval=600, # check every 10 minutes for child DAG completion
80+
)
81+
82+
def fire_github_callback(**context):
83+
params = context["params"]
84+
dag_run = context["dag_run"]
85+
86+
response = requests.post(
87+
f"https://api.github.com/repos/{params['github_repo']}/dispatches",
88+
headers={
89+
"Authorization": f"Bearer {params['github_callback_token']}",
90+
"Accept": "application/vnd.github+json",
91+
"X-GitHub-Api-Version": "2022-11-28",
92+
},
93+
json={
94+
"event_type": "airflow-dag-complete",
95+
"client_payload": {
96+
"state": "success",
97+
"dag_id": dag_run.dag_id,
98+
"dag_run_id": dag_run.run_id,
99+
"sha": params["maxtext_sha"],
100+
"github_run_id": params["github_run_id"],
101+
},
102+
},
103+
timeout=30,
104+
)
105+
response.raise_for_status()
106+
107+
github_callback = PythonOperator(
108+
task_id="fire_github_callback",
109+
python_callable=fire_github_callback,
110+
trigger_rule=TriggerRule.ALL_SUCCESS, # Only fire if all upstream tasks succeeded
111+
)
112+
113+
[trigger_pre_training, trigger_post_training] >> github_callback

‎dags/multipod/maxtext_e2e_tpu_post_training.py‎

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""
16+
A DAG to run MaxText E2E TPU Post-Training tests.
17+
"""
18+
import datetime
19+
import hashlib
20+
21+
from click import command
22+
from airflow import models
23+
from airflow.models.param import Param
24+
from airflow.utils.task_group import TaskGroup
25+
from dags.common import test_owner
26+
from dags.common.vm_resource import XpkClusters
27+
from dags.multipod.configs import gke_config
28+
29+
# HF token retrieved from Airflow Variables for secure credential management
30+
HF_TOKEN = models.Variable.get("HF_TOKEN", None)
31+
32+
33+
def get_workload_name(model, mode, length=6):
34+
hex_code = f"{mode}-{hashlib.sha256(model.encode()).hexdigest()}"
35+
return hex_code[:length]
36+
37+
38+
with models.DAG(
39+
dag_id="maxtext_e2e_tpu_post_training",
40+
schedule=None,
41+
tags=[
42+
"maxtext",
43+
"post-training",
44+
"TPU",
45+
],
46+
start_date=datetime.datetime(2026, 6, 10),
47+
catchup=False,
48+
params={
49+
"docker_image": Param(
50+
type="string",
51+
description="Docker image URI for the candidate to test",
52+
),
53+
},
54+
) as dag:
55+
test_models = {
56+
"gemma3-4b": {
57+
"checkpoint_conversion": {
58+
"to_maxtext": "bash tests/end_to_end/tpu/gemma3/4b/test_gemma3_to_mt.sh",
59+
"to_huggingface": "bash tests/end_to_end/tpu/gemma3/4b/test_gemma3_to_hf.sh",
60+
},
61+
"post_training": {
62+
"sft": {
63+
"command": "bash tests/end_to_end/tpu/gemma3/4b/test_gemma3_sft.sh",
64+
"maxtext_ckpt_path": "gs://runner-maxtext-logs/gemma3-4b/sft/{run_name}/checkpoints/5/model_params",
65+
},
66+
"rl": {
67+
"command": "bash tests/end_to_end/tpu/gemma3/4b/test_gemma3_rl.sh",
68+
"maxtext_ckpt_path": "gs://runner-maxtext-logs/gemma3-4b/rl/{run_name}/checkpoints/actor/5/model_params",
69+
},
70+
},
71+
},
72+
}
73+
74+
for model, test_config in test_models.items():
75+
with TaskGroup(group_id=model) as model_group:
76+
run_name = "post-{{ ts_nodash }}"
77+
78+
convert_to_maxtext_cmd = (f"export HF_TOKEN={HF_TOKEN}",) + (
79+
f"{test_config['checkpoint_conversion']['to_maxtext']} {run_name}",
80+
)
81+
convert_to_maxtext_task = gke_config.get_gke_config(
82+
time_out_in_min=60,
83+
test_name=f"convert-to-maxtext",
84+
run_model_cmds=convert_to_maxtext_cmd,
85+
docker_image="{{ params.docker_image }}",
86+
cluster=XpkClusters.TPU_V5P_8_CLUSTER_V2,
87+
test_owner=test_owner.SURBHI_J,
88+
).run(skip_post_process=True)
89+
90+
for mode, mode_test_config in test_config["post_training"].items():
91+
with TaskGroup(group_id=f"{mode}-{model}") as model_group:
92+
environment_variables = [
93+
f"export HF_TOKEN={HF_TOKEN}",
94+
"export TPU_MIN_LOG_LEVEL=0",
95+
"export TF_CPP_MIN_LOG_LEVEL=0",
96+
"export TPU_STDERR_LOG_LEVEL=0",
97+
"export JAX_PLATFORMS=proxy,cpu",
98+
"export JAX_BACKEND_TARGET=grpc://127.0.0.1:29000",
99+
"export ENABLE_PATHWAYS_PERSISTENCE='1'",
100+
]
101+
102+
with TaskGroup(group_id=f"train-{mode}-{model}") as model_group:
103+
command = mode_test_config["command"]
104+
training_cmd = (
105+
" && ".join(
106+
environment_variables + [f"{command} {run_name} true"]
107+
),
108+
)
109+
training_task = gke_config.get_gke_config(
110+
time_out_in_min=60,
111+
num_slices=1,
112+
cluster=XpkClusters.TPU_V5P_128_CLUSTER,
113+
test_name=get_workload_name(model, mode),
114+
run_model_cmds=training_cmd,
115+
docker_image="{{ params.docker_image }}",
116+
test_owner=test_owner.SURBHI_J,
117+
).run_model(
118+
use_pathways=True,
119+
)
120+
121+
model_path = mode_test_config["maxtext_ckpt_path"].format(
122+
run_name=run_name
123+
)
124+
convert_to_huggingface_cmd = (f"export HF_TOKEN={HF_TOKEN}",) + (
125+
f"{test_config['checkpoint_conversion']['to_huggingface']} {run_name} {model_path} false true",
126+
)
127+
convert_to_huggingface_task = gke_config.get_gke_config(
128+
time_out_in_min=60,
129+
test_name="convert-to-huggingface",
130+
run_model_cmds=convert_to_huggingface_cmd,
131+
docker_image="{{ params.docker_image }}",
132+
cluster=XpkClusters.TPU_V5P_8_CLUSTER_V2,
133+
test_owner=test_owner.SURBHI_J,
134+
).run(skip_post_process=True)
135+
136+
(
137+
convert_to_maxtext_task
138+
>> training_task
139+
>> convert_to_huggingface_task
140+
)

‎dags/multipod/maxtext_e2e_tpu_pre_training.py‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""
16+
A DAG to run MaxText E2E TPU Pre-Training tests.
17+
"""
18+
import datetime
19+
from airflow import models
20+
from airflow.models.param import Param
21+
from airflow.utils.task_group import TaskGroup
22+
from dags.common import test_owner
23+
from dags.common.vm_resource import XpkClusters
24+
from dags.multipod.configs import gke_config
25+
26+
HF_TOKEN = models.Variable.get("HF_TOKEN", None)
27+
28+
with models.DAG(
29+
dag_id="maxtext_e2e_tpu_pre_training",
30+
schedule=None,
31+
tags=[
32+
"maxtext",
33+
"pre-training",
34+
"TPU",
35+
],
36+
start_date=datetime.datetime(2026, 6, 10),
37+
catchup=False,
38+
params={
39+
"docker_image": Param(
40+
type="string",
41+
description="Docker image URI for the candidate to test",
42+
),
43+
},
44+
) as dag:
45+
test_models = {
46+
"gemma3-4b": {
47+
"checkpoint_conversion": {
48+
"to_maxtext": "bash tests/end_to_end/tpu/gemma3/4b/test_gemma3_to_mt.sh",
49+
"to_huggingface": "bash tests/end_to_end/tpu/gemma3/4b/test_gemma3_to_hf.sh",
50+
},
51+
"training": {
52+
"command": "bash tests/end_to_end/tpu/gemma3/4b/test_gemma3.sh",
53+
"maxtext_ckpt_path": "gs://runner-maxtext-logs/gemma3-4b/train/{run_name}/checkpoints/4/items",
54+
},
55+
},
56+
}
57+
58+
for model, test_config in test_models.items():
59+
with TaskGroup(group_id=model) as model_group:
60+
run_name = "pre-{{ ts_nodash }}"
61+
62+
convert_to_maxtext_cmd = (f"export HF_TOKEN={HF_TOKEN}",) + (
63+
f"{test_config['checkpoint_conversion']['to_maxtext']} {run_name}",
64+
)
65+
convert_to_maxtext_task = gke_config.get_gke_config(
66+
time_out_in_min=60,
67+
test_name="convert-to-maxtext",
68+
run_model_cmds=convert_to_maxtext_cmd,
69+
docker_image="{{ params.docker_image }}",
70+
cluster=XpkClusters.TPU_V5P_8_CLUSTER_V2,
71+
test_owner=test_owner.SURBHI_J,
72+
).run(skip_post_process=True)
73+
74+
training_cmd = (f"export HF_TOKEN={HF_TOKEN}",) + (
75+
f"{test_config['training']['command']} {run_name}",
76+
)
77+
training_task = gke_config.get_gke_config(
78+
time_out_in_min=60,
79+
test_name="training",
80+
run_model_cmds=training_cmd,
81+
docker_image="{{ params.docker_image }}",
82+
cluster=XpkClusters.TPU_V5P_128_CLUSTER,
83+
test_owner=test_owner.SURBHI_J,
84+
).run(skip_post_process=True)
85+
86+
model_path = test_config["training"]["maxtext_ckpt_path"].format(
87+
run_name=run_name
88+
)
89+
convert_to_huggingface_cmd = (f"export HF_TOKEN={HF_TOKEN}",) + (
90+
f"{test_config['checkpoint_conversion']['to_huggingface']} {run_name} {model_path}",
91+
)
92+
convert_to_huggingface_task = gke_config.get_gke_config(
93+
time_out_in_min=60,
94+
test_name="convert-to-huggingface",
95+
run_model_cmds=convert_to_huggingface_cmd,
96+
docker_image="{{ params.docker_image }}",
97+
cluster=XpkClusters.TPU_V5P_8_CLUSTER_V2,
98+
test_owner=test_owner.SURBHI_J,
99+
).run(skip_post_process=True)
100+
101+
convert_to_maxtext_task >> training_task >> convert_to_huggingface_task

0 commit comments

Comments
 (0)