Skip to content

Commit 7cbe10e

Browse files
ryanliuCienetryanliuCienet
committed
Fix format/linting errors
1 parent 04e1f25 commit 7cbe10e

5 files changed

Lines changed: 82 additions & 64 deletions

File tree

‎dags/orbax/maxtext_emc_restore_gcs.py‎

Lines changed: 20 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -120,22 +120,26 @@
120120
task_id="generate_start_time"
121121
)()
122122

123-
maxtext_chkpt_run_test = gke_config.get_gke_config(
124-
num_slices=slice_num,
125-
cluster=test_config.cluster,
126-
time_out_in_min=60,
127-
test_name=f"{test_config.short_id}-emc",
128-
run_model_cmds=workload_command,
129-
docker_image=image.value,
130-
test_owner=test_owner.DEPP_L,
131-
).to_node_interruption_task(
132-
expect_reach_to_step=step_to_interrupt,
133-
last_node=True,
134-
).run(
135-
ramdisk_directory=test_config_util.DEFAULT_RAM_DISK,
136-
mtc_enabled=True,
137-
skip_post_process=True,
138-
max_restart=15,
123+
maxtext_chkpt_run_test = (
124+
gke_config.get_gke_config(
125+
num_slices=slice_num,
126+
cluster=test_config.cluster,
127+
time_out_in_min=60,
128+
test_name=f"{test_config.short_id}-emc",
129+
run_model_cmds=workload_command,
130+
docker_image=image.value,
131+
test_owner=test_owner.DEPP_L,
132+
)
133+
.to_node_interruption_task(
134+
expect_reach_to_step=step_to_interrupt,
135+
last_node=True,
136+
)
137+
.run(
138+
ramdisk_directory=test_config_util.DEFAULT_RAM_DISK,
139+
mtc_enabled=True,
140+
skip_post_process=True,
141+
max_restart=15,
142+
)
139143
)
140144

141145
end_time = validation_util.generate_timestamp.override(

‎dags/orbax/maxtext_emc_restore_local.py‎

Lines changed: 19 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -118,21 +118,25 @@
118118
task_id="generate_start_time"
119119
)()
120120

121-
maxtext_chkpt_run_test = gke_config.get_gke_config(
122-
num_slices=slice_num,
123-
cluster=test_config.cluster,
124-
time_out_in_min=60,
125-
test_name=f"{test_config.short_id}-emc",
126-
run_model_cmds=workload_command,
127-
docker_image=image.value,
128-
test_owner=test_owner.DEPP_L,
129-
).to_node_interruption_task(
130-
expect_reach_to_step=step_to_interrupt,
131-
).run(
132-
ramdisk_directory=test_config_util.DEFAULT_RAM_DISK,
133-
mtc_enabled=True,
134-
skip_post_process=True,
135-
max_restart=15,
121+
maxtext_chkpt_run_test = (
122+
gke_config.get_gke_config(
123+
num_slices=slice_num,
124+
cluster=test_config.cluster,
125+
time_out_in_min=60,
126+
test_name=f"{test_config.short_id}-emc",
127+
run_model_cmds=workload_command,
128+
docker_image=image.value,
129+
test_owner=test_owner.DEPP_L,
130+
)
131+
.to_node_interruption_task(
132+
expect_reach_to_step=step_to_interrupt,
133+
)
134+
.run(
135+
ramdisk_directory=test_config_util.DEFAULT_RAM_DISK,
136+
mtc_enabled=True,
137+
skip_post_process=True,
138+
max_restart=15,
139+
)
136140
)
137141

138142
end_time = validation_util.generate_timestamp.override(

‎dags/orbax/maxtext_mtc_restore_local.py‎

Lines changed: 19 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -121,21 +121,25 @@
121121
task_id="generate_start_time"
122122
)()
123123

124-
maxtext_chkpt_run_test = gke_config.get_gke_config(
125-
num_slices=slice_num,
126-
cluster=test_config.cluster,
127-
time_out_in_min=60,
128-
test_name=f"{test_config.short_id}-mtc",
129-
run_model_cmds=workload_command,
130-
docker_image=image.value,
131-
test_owner=test_owner.DEPP_L,
132-
).to_node_interruption_task(
133-
expect_reach_to_step=step_to_interrupt,
134-
).run(
135-
ramdisk_directory=test_config_util.DEFAULT_RAM_DISK,
136-
mtc_enabled=True,
137-
skip_post_process=True,
138-
max_restart=15,
124+
maxtext_chkpt_run_test = (
125+
gke_config.get_gke_config(
126+
num_slices=slice_num,
127+
cluster=test_config.cluster,
128+
time_out_in_min=60,
129+
test_name=f"{test_config.short_id}-mtc",
130+
run_model_cmds=workload_command,
131+
docker_image=image.value,
132+
test_owner=test_owner.DEPP_L,
133+
)
134+
.to_node_interruption_task(
135+
expect_reach_to_step=step_to_interrupt,
136+
)
137+
.run(
138+
ramdisk_directory=test_config_util.DEFAULT_RAM_DISK,
139+
mtc_enabled=True,
140+
skip_post_process=True,
141+
max_restart=15,
142+
)
139143
)
140144

141145
end_time = validation_util.generate_timestamp.override(

‎dags/orbax/maxtext_reg_restore_gcs_with_node_disruption.py‎

Lines changed: 20 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -130,22 +130,26 @@ def generate_workload_checkpoints_location(gcs_ckpt_location: str) -> str:
130130
task_id="generate_start_time"
131131
)()
132132

133-
maxtext_chkpt_run_test = gke_config.get_gke_config(
134-
num_slices=slice_num,
135-
cluster=test_config.cluster,
136-
time_out_in_min=60,
137-
test_name=f"{test_config.short_id}",
138-
run_model_cmds=workload_command,
139-
docker_image=image.value,
140-
test_owner=test_owner.SHARON_Y,
141-
).to_node_interruption_task(
142-
expect_reach_to_step=step_to_interrupt,
143-
check_file_exists=True,
144-
).run(
145-
gcs_location=gcs_location,
146-
xpk_branch=MAIN_BRANCH,
147-
skip_post_process=True,
148-
max_restart=15,
133+
maxtext_chkpt_run_test = (
134+
gke_config.get_gke_config(
135+
num_slices=slice_num,
136+
cluster=test_config.cluster,
137+
time_out_in_min=60,
138+
test_name=f"{test_config.short_id}",
139+
run_model_cmds=workload_command,
140+
docker_image=image.value,
141+
test_owner=test_owner.SHARON_Y,
142+
)
143+
.to_node_interruption_task(
144+
expect_reach_to_step=step_to_interrupt,
145+
check_file_exists=True,
146+
)
147+
.run(
148+
gcs_location=gcs_location,
149+
xpk_branch=MAIN_BRANCH,
150+
skip_post_process=True,
151+
max_restart=15,
152+
)
149153
)
150154

151155
end_time = validation_util.generate_timestamp.override(

‎xlml/apis/task.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -398,7 +398,6 @@ def to_name_gen_and_quarantine_task(
398398
nested_run_name_in_tb_file_location=nested_run_name_in_tb_file_location,
399399
)
400400

401-
402401
def run_model(
403402
self,
404403
gcs_location: Optional[airflow.XComArg] = None,
@@ -540,6 +539,7 @@ def post_process(self, result_location: Optional[str] = None) -> DAGNode:
540539

541540
return group
542541

542+
543543
@dataclasses.dataclass
544544
class XpkNodeInterruptionTask(XpkTask):
545545
"""Task for running XPK workloads with node interruption."""
@@ -811,7 +811,9 @@ def run_with_run_name_generation(
811811
run_model, gcs_path = self.run_model(
812812
use_pathways=use_pathways, xpk_branch=xpk_branch
813813
)
814-
flow = run_name >> (tb_file_location, profile_file_location) >> run_model
814+
flow = (
815+
run_name >> (tb_file_location, profile_file_location) >> run_model
816+
)
815817
if not skip_post_process:
816818
_ = flow >> self.post_process(gcs_path)
817819
else:

0 commit comments

Comments
 (0)