Fix maxtext_rl and maxtext_sft DAGs (GoogleCloudPlatform#1211)

Author: ooops678

dags/post_training/maxtext_rl.py

@@ -64,7 +64,7 @@
       4. All training steps execute within expected parameters
       5. Metrics are successfully synced to Vertex AI TensorBoard
     """,
-    concurrency=2,
+    concurrency=1,
 ) as dag:
   training_config = test_config_util.RLTestConfig(
       cluster=XpkClusters.TPU_V5P_128_CLUSTER,
@@ -88,41 +88,41 @@
   # HF token retrieved from Airflow Variables for secure credential management
   HF_TOKEN_LLAMA3_1 = models.Variable.get("HF_TOKEN_CIENET", None)
 
+  tests = []
   for mode, image in test_config_util.POST_TRAINING_DOCKER_IMAGES:
     # TODO: Enable stable mode once a new version of MaxText is available
     if mode == test_config_util.SetupMode.STABLE:
       continue  # Skip stable for RL training tests
 
     for loss_algo in training_config.loss_algos:
       for slice_num in training_config.slices:
-        run_name = validation_util.generate_run_name(
-            prefix=loss_algo.value,
-            mode=mode.value,
-            num_slices=slice_num,
-        )
-
-        rl_training_command = training_config.generate_rl_training_command(
-            loss_algo=loss_algo,
-            run_name=run_name,
-            hf_token=HF_TOKEN_LLAMA3_1,
-            num_slices=slice_num,
-        )
-
         with TaskGroup(
             group_id=(
                 f"{loss_algo.value}-{mode.value}-"
                 f"{slice_num}x{training_config.accelerator}"
             )
         ) as group:
           with TaskGroup(group_id="run_training") as training_group:
+            run_name = validation_util.generate_run_name(
+                prefix=loss_algo.value,
+                mode=mode.value,
+                num_slices=slice_num,
+            )
+
+            rl_training_command = training_config.generate_rl_training_command(
+                loss_algo=loss_algo,
+                run_name=run_name,
+                hf_token=HF_TOKEN_LLAMA3_1,
+                num_slices=slice_num,
+            )
             start_time = validation_util.generate_timestamp.override(
                 task_id="generate_start_time"
             )()
 
             training_task = gke_config.get_gke_config(
+                time_out_in_min=30,
                 num_slices=slice_num,
                 cluster=training_config.cluster,
-                time_out_in_min=30,
                 test_name=loss_algo.value,
                 run_model_cmds=rl_training_command,
                 docker_image=image.value,
@@ -137,6 +137,7 @@
             )()
 
             chain(
+                run_name,
                 start_time,
                 training_task,
                 end_time,
@@ -149,7 +150,7 @@
                 project_id=training_config.cluster.project,
                 location=zone_to_region(training_config.cluster.zone),
                 cluster_name=training_config.cluster.name,
-                text_filter=f'"Config param loss_algo: {loss_algo.loss_name}"',
+                text_filter=f"\"'loss_algo': '{loss_algo.loss_name}'\"",
                 namespace="default",
                 container_name="jax-tpu",
                 pod_pattern=f"{loss_algo.value}.*",
@@ -183,3 +184,6 @@
           )
 
           chain(training_group, [validation_group, upload_to_vertex_ai])
+          tests.append(group)
+
+        chain(*tests)

dags/post_training/maxtext_sft.py

@@ -75,7 +75,7 @@ def validate_training(
         project_id=config.cluster.project,
         location=zone_to_region(config.cluster.zone),
         cluster_name=config.cluster.name,
-        text_filter=f"\"'step': {steps}\"",
+        text_filter=f'"Training: 100%" AND "{steps}/{steps}"',
         namespace="default",
         container_name="jax-tpu",
         pod_pattern=f"{config.short_id}.*",
@@ -122,7 +122,7 @@ def validate_training(
       3. No infrastructure failures or container launch issues occur
       4. All training steps execute within expected parameters
     """,
-    concurrency=2,
+    concurrency=1,
 ) as dag:
   training_steps = 30
 
@@ -152,19 +152,19 @@ def validate_training(
       continue  # Skip stable for SFT training tests
 
     for slice_num in training_config.slices:
-      run_name = validation_util.generate_run_name(
-          prefix="sft",
-          mode=mode.value,
-          num_slices=slice_num,
-      )
-
-      sft_training_command = training_config.generate_sft_training_command(
-          run_name=run_name,
-          hf_token=HF_TOKEN_CIENET,
-      )
       with TaskGroup(
           group_id=f"sft-{mode.value}-{slice_num}x{training_config.accelerator}"
       ) as group:
+        run_name = validation_util.generate_run_name(
+            prefix="sft",
+            mode=mode.value,
+            num_slices=slice_num,
+        )
+
+        sft_training_command = training_config.generate_sft_training_command(
+            run_name=run_name,
+            hf_token=HF_TOKEN_CIENET,
+        )
         training_group, start_time, end_time = run_training(
             slice_num, training_config, sft_training_command, image
         )
@@ -181,4 +181,4 @@ def validate_training(
             run_name_prefix=run_name,
         )
 
-        chain(training_group, [validation_group, upload_to_vertex_ai])
+        chain(run_name, training_group, [validation_group, upload_to_vertex_ai])

dags/post_training/util/test_config_util.py

@@ -158,7 +158,7 @@ def generate_rl_training_command(
       ])
 
     rl_command = (
-        "python -m src.maxtext.rl.train_rl "
+        "python -m src.MaxText.rl.train_rl "
         f"{self.rl_config_path} " + " ".join(command_params)
     )