Skip to content

Commit 65d3df0

Browse files
khatwanimohitkhatwanimohit
authored
change paths for e2e tests (#1162)
1 parent 5ee6a23 commit 65d3df0

File tree

9 files changed

+31
-27
lines changed

9 files changed

+31
-27
lines changed

‎dags/multipod/configs/jax_tests_gke_config.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ def get_jax_distributed_initialize_config(
2727
num_slices: int = 1,
2828
):
2929
run_model_cmds = [
30-
"bash end_to_end/test_jdi.sh",
30+
"bash tests/end_to_end/test_jdi.sh",
3131
]
3232

3333
return gke_config.get_gke_config(

‎dags/multipod/configs/maxtext_gce_config.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,9 @@ def get_maxtext_end_to_end_test_config(
119119

120120
test_platform = common.Platform.GCE
121121
set_up_cmds = common.setup_maxtext(test_mode, test_platform)
122-
run_model_cmds = (f"cd /tmp/maxtext && bash end_to_end/tpu/{test_script}.sh",)
122+
run_model_cmds = (
123+
f"cd /tmp/maxtext && bash tests/end_to_end/tpu/{test_script}.sh",
124+
)
123125

124126
job_test_config = test_config.TpuVmTest(
125127
test_config.Tpu(

‎dags/multipod/maxtext_checkpointing.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@
5656
async_checkpointing = chkpt_mode == "async"
5757
run_name = f"checkpointing-{mode.value}-{chkpt_mode}-{current_datetime}"
5858
command = (
59-
"bash end_to_end/test_checkpointing.sh "
59+
"bash tests/end_to_end/test_checkpointing.sh "
6060
f" {run_name} {base_output_directory} {dataset_path}"
6161
f" true tfds autoselected {async_checkpointing}",
6262
)
@@ -79,7 +79,7 @@
7979
time_out_in_min=60,
8080
test_name=f"maxtext-checkpoint-resharding-{mode.value}",
8181
run_model_cmds=(
82-
f"bash end_to_end/tpu/test_checkpoint_resharding.sh checkpoint-resharding-{mode.value} {base_output_directory} {dataset_path}",
82+
f"bash tests/end_to_end/tpu/test_checkpoint_resharding.sh checkpoint-resharding-{mode.value} {base_output_directory} {dataset_path}",
8383
),
8484
docker_image=image.value,
8585
test_owner=test_owner.SURBHI_J,

‎dags/multipod/maxtext_convergence.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@
5353
per_device_batch_size = 2.0 # 256 chips * 2 pdb = 512 gbs.
5454

5555
base_convergence_command = (
56-
"bash end_to_end/tpu/test_convergence_1b_params.sh"
56+
"bash tests/end_to_end/tpu/test_convergence_1b_params.sh"
5757
f" OUTPUT_PATH={base_output_directory} DATASET_PATH={dataset_path}"
5858
f" LOSS_THRESHOLD={loss_threshold} STEPS={steps}"
5959
f" PER_DEVICE_BATCH_SIZE={per_device_batch_size}"

‎dags/multipod/maxtext_end_to_end.py‎

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -50,40 +50,42 @@
5050
test_models_tpu = {
5151
"llama2-7b": {
5252
"owner": test_owner.MOHIT_K,
53-
"commands": ["bash end_to_end/tpu/llama2/7b/test_llama2_7b.sh"],
53+
"commands": ["bash tests/end_to_end/tpu/llama2/7b/test_llama2_7b.sh"],
5454
},
5555
"mistral-7b": {
5656
"owner": test_owner.MOHIT_K,
57-
"commands": ["bash end_to_end/tpu/mistral/7b/test_mistral-7b.sh"],
57+
"commands": [
58+
"bash tests/end_to_end/tpu/mistral/7b/test_mistral-7b.sh"
59+
],
5860
},
5961
"gemma-2b": {
6062
"owner": test_owner.MOHIT_K,
61-
"commands": ["bash end_to_end/tpu/gemma/2b/test_gemma.sh"],
63+
"commands": ["bash tests/end_to_end/tpu/gemma/2b/test_gemma.sh"],
6264
},
6365
"gemma2-2b": {
6466
"owner": test_owner.HENGTAO_G,
6567
"commands": [
66-
"bash end_to_end/tpu/gemma2/2b/test_gemma2_to_mt.sh",
67-
"bash end_to_end/tpu/gemma2/2b/test_gemma2_to_hf.sh",
68+
"bash tests/end_to_end/tpu/gemma2/2b/test_gemma2_to_mt.sh",
69+
"bash tests/end_to_end/tpu/gemma2/2b/test_gemma2_to_hf.sh",
6870
],
6971
},
7072
"gemma3-4b": {
7173
"owner": test_owner.HENGTAO_G,
7274
"commands": [
73-
"bash end_to_end/tpu/gemma3/4b/test_gemma3_to_mt.sh",
74-
"bash end_to_end/tpu/gemma3/4b/test_gemma3_to_hf.sh",
75+
"bash tests/end_to_end/tpu/gemma3/4b/test_gemma3_to_mt.sh",
76+
"bash tests/end_to_end/tpu/gemma3/4b/test_gemma3_to_hf.sh",
7577
],
7678
},
7779
"qwen3-4b": {
7880
"owner": test_owner.HENGTAO_G,
7981
"commands": [
80-
"bash end_to_end/tpu/qwen3/4b/test_qwen3_to_mt.sh",
81-
"bash end_to_end/tpu/qwen3/4b/test_qwen3_to_hf.sh",
82+
"bash tests/end_to_end/tpu/qwen3/4b/test_qwen3_to_mt.sh",
83+
"bash tests/end_to_end/tpu/qwen3/4b/test_qwen3_to_hf.sh",
8284
],
8385
},
8486
"gpt3": {
8587
"owner": test_owner.MOHIT_K,
86-
"commands": ["bash end_to_end/tpu/test_gpt3.sh"],
88+
"commands": ["bash tests/end_to_end/tpu/test_gpt3.sh"],
8789
},
8890
}
8991

@@ -160,7 +162,7 @@ def convert_checkpoint_and_run_training(
160162
time_out_in_min=test_scripts_details[0]["time_out_in_min"],
161163
test_name=test_name,
162164
run_model_cmds=(
163-
f"export BASE_OUTPUT_PATH=$GCS_OUTPUT; bash end_to_end/{test_scripts_details[0]['script_name']}.sh",
165+
f"export BASE_OUTPUT_PATH=$GCS_OUTPUT; bash tests/end_to_end/{test_scripts_details[0]['script_name']}.sh",
164166
),
165167
docker_image=docker_image,
166168
test_owner=test_owner.ANISHA_M,
@@ -170,7 +172,7 @@ def convert_checkpoint_and_run_training(
170172
time_out_in_min=test_scripts_details[1]["time_out_in_min"],
171173
test_name=test_name,
172174
run_model_cmds=(
173-
f"export BASE_OUTPUT_PATH=$GCS_OUTPUT; bash end_to_end/{test_scripts_details[1]['script_name']}.sh",
175+
f"export BASE_OUTPUT_PATH=$GCS_OUTPUT; bash tests/end_to_end/{test_scripts_details[1]['script_name']}.sh",
174176
),
175177
docker_image=docker_image,
176178
test_owner=test_owner.ANISHA_M,

‎dags/multipod/maxtext_gpu_end_to_end.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -168,31 +168,31 @@ def run_maxtext_tests(dag: models.DAG):
168168
),
169169
"generate-param-only-checkpoint": (
170170
"XLA_PYTHON_CLIENT_MEM_FRACTION=0.65 TF_FORCE_GPU_ALLOW_GROWTH=true "
171-
f"bash end_to_end/test_generate_param_only_checkpoint.sh -r runner-{timestamp}-8 "
171+
f"bash tests/end_to_end/test_generate_param_only_checkpoint.sh -r runner-{timestamp}-8 "
172172
"-o gs://runner-maxtext-logs -d gs://maxtext-dataset -i 4 -a dot_product",
173173
1,
174174
),
175175
"generate-param-only-checkpoint-int8": (
176176
"XLA_PYTHON_CLIENT_MEM_FRACTION=0.65 TF_FORCE_GPU_ALLOW_GROWTH=true "
177-
f"bash end_to_end/test_generate_param_only_checkpoint.sh -r runner-{timestamp}-9 "
177+
f"bash tests/end_to_end/test_generate_param_only_checkpoint.sh -r runner-{timestamp}-9 "
178178
"-o gs://runner-maxtext-logs -d gs://maxtext-dataset -i 4 -q int8 -a dot_product",
179179
1,
180180
),
181181
"grain-checkpoint-determinism": (
182182
"XLA_PYTHON_CLIENT_MEM_FRACTION=0.65 TF_FORCE_GPU_ALLOW_GROWTH=true "
183-
"bash end_to_end/test_checkpointing.sh runner gs://runner-maxtext-logs "
183+
"bash tests/end_to_end/test_checkpointing.sh runner gs://runner-maxtext-logs "
184184
"gs://maxtext-dataset False c4-array_record dot_product",
185185
1,
186186
),
187187
"checkpoint-compatibility": (
188188
"XLA_PYTHON_CLIENT_MEM_FRACTION=0.65 TF_FORCE_GPU_ALLOW_GROWTH=true "
189-
"bash end_to_end/test_checkpoint_compatibility.sh runner "
189+
"bash tests/end_to_end/test_checkpoint_compatibility.sh runner "
190190
"gs://runner-maxtext-logs gs://maxtext-dataset dot_product",
191191
1,
192192
),
193193
"llama2-7b-train-1node": ("bash MaxText/configs/a3/llama_2_7b/1vm.sh", 1),
194194
"llama2-7b-train-2node": ("bash MaxText/configs/a3/llama_2_7b/2vm.sh", 2),
195-
"llama2-7b": ("bash end_to_end/gpu/a3/test_llama2_7b.sh", 1),
195+
"llama2-7b": ("bash tests/end_to_end/gpu/a3/test_llama2_7b.sh", 1),
196196
}
197197

198198
quarantine_task_group = TaskGroup(

‎dags/multipod/maxtext_multi_tier_checkpointing.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@
6060
for accelerator, slices in test_configs.items():
6161
for slice_num in slices:
6262
command = (
63-
"bash end_to_end/test_multi_tier_checkpointing.sh"
63+
"bash tests/end_to_end/test_multi_tier_checkpointing.sh"
6464
f" multi_tier_checkpointing-{slice_num}x-{accelerator}"
6565
f" {base_output_directory} {dataset_path}",
6666
)

‎dags/multipod/maxtext_sft_trainer.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@
5757
f'export BASE_OUTPUT_DIRECTORY={base_output_directory}',
5858
'export STEPS=1000',
5959
'export PROMPT="Suggest some famous landmarks in London."',
60-
'bash end_to_end/tpu/test_sft_trainer.sh',
60+
'bash tests/end_to_end/tpu/test_sft_trainer.sh',
6161
)
6262
maxtext_v4_configs_test = gke_config.get_gke_config(
6363
cluster=XpkClusters.TPU_V5P_8_CLUSTER,

‎dags/sparsity_diffusion_devx/maxtext_moe_tpu_e2e.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,7 @@
8888
time_out_in_min=test_scripts_details["time_out_in_min"],
8989
test_name=f"{test_name_prefix}_{image}_{model}",
9090
run_model_cmds=(
91-
f"export HF_TOKEN={HF_TOKEN}; export BASE_OUTPUT_PATH=$GCS_OUTPUT; bash end_to_end/{test_scripts_details['script_name']}.sh",
91+
f"export HF_TOKEN={HF_TOKEN}; export BASE_OUTPUT_PATH=$GCS_OUTPUT; bash tests/end_to_end/{test_scripts_details['script_name']}.sh",
9292
),
9393
docker_image=docker_image[image],
9494
test_owner=test_owner.SHUNING_J,
@@ -148,7 +148,7 @@ def convert_checkpoint_and_run_training(
148148
time_out_in_min=test_scripts_details[0]["time_out_in_min"],
149149
test_name=test_name,
150150
run_model_cmds=(
151-
f"export BASE_OUTPUT_PATH=$GCS_OUTPUT; bash end_to_end/{test_scripts_details[0]['script_name']}.sh",
151+
f"export BASE_OUTPUT_PATH=$GCS_OUTPUT; bash tests/end_to_end/{test_scripts_details[0]['script_name']}.sh",
152152
),
153153
docker_image=docker_image,
154154
test_owner=test_owner.SHUNING_J,
@@ -158,7 +158,7 @@ def convert_checkpoint_and_run_training(
158158
time_out_in_min=test_scripts_details[1]["time_out_in_min"],
159159
test_name=test_name,
160160
run_model_cmds=(
161-
f"export BASE_OUTPUT_PATH=$GCS_OUTPUT; bash end_to_end/{test_scripts_details[1]['script_name']}.sh",
161+
f"export BASE_OUTPUT_PATH=$GCS_OUTPUT; bash tests/end_to_end/{test_scripts_details[1]['script_name']}.sh",
162162
),
163163
docker_image=docker_image,
164164
test_owner=test_owner.SHUNING_J,

0 commit comments

Comments
 (0)