feat: Add Airflow DAG for Llama3.1-8B RL training (GoogleCloudPlatform#1127)

This pull request introduces automation for running RL training from a Jupyter notebook on TPU VMs, adds utilities for notebook parameter injection and execution, and updates test configuration and task orchestration to support the new workflow. The main changes are the addition of a new Airflow DAG for automating Llama3.1-8B RL training from a notebook, a utility module for notebook automation, and enhancements to the task orchestration to allow skipping post-processing steps.

Author: RexBearIU

dags/post_training/maxtext_rl.py

@@ -44,17 +44,15 @@
       # RL Training (GRPO/GSPO) MaxText RL Training
 
       ### Overview
-      This DAG runs RL training using GRPO (Group Relative Policy Optimization) 
-      and GSPO algorithms to validate the MaxText reinforcement learning pipeline. 
-      The workflow tests the complete RL training stack including infrastructure setup,
-      model initialization, training execution, and result validation.
+      Runs RL training with GRPO/GSPO to validate MaxText RL pipeline.
+      Tests RL stack: infra, model, training, validation.
 
       ### Execution Flow
-      1. **Job Launch:** Deploy RL training jobs to GKE cluster using Pathways infrastructure
-      2. **Model Loading:** Initialize Llama3.1 70B model with HuggingFace authentication
-      3. **Training Run:** Execute train_rl with JAX proxy/CPU platforms for GRPO and GSPO
-      4. **Log Validation:** Monitor and check for "Post RL Training" completion signal
-      5. **Success/Failure:** Report final status based on log validation and job completion
+      1. **Job Launch:** Deploy RL jobs to GKE using Pathways infra
+      2. **Model Loading:** Init Llama3.1 70B with HF auth
+      3. **Training Run:** Run train_rl with JAX proxy for GRPO/GSPO
+      4. **Log Validation:** Check for 'Post RL Training' signal
+      5. **Success/Failure:** Report status from logs and completion
 
       ### Success Criteria
       The test passes when:
@@ -85,7 +83,7 @@
       ],
   )
   # HF token retrieved from Airflow Variables for secure credential management
-  HF_TOKEN_LLAMA3_1 = models.Variable.get("HF_TOKEN_LLAMA3_1", None)
+  HF_TOKEN_LLAMA3_1 = models.Variable.get("HF_TOKEN_CIENET", None)
 
   for mode, image in test_config_util.POST_TRAINING_DOCKER_IMAGES:
     # TODO: Enable stable mode once a new version of MaxText is available
@@ -104,7 +102,10 @@
         )
 
         with TaskGroup(
-            group_id=f"{loss_algo.value}-{mode.value}-{slice_num}x{training_config.accelerator}"
+            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:
             start_time = validation_util.generate_timestamp.override(

dags/post_training/maxtext_rl_notebook.py

@@ -0,0 +1,154 @@
+"""
+Airflow DAG for automating Llama3.1-8B RL training from Jupyter notebook.
+
+This DAG automates the rl_llama3_demo.ipynb notebook, executing GRPO/GSPO
+training on single-host TPU VMs.
+"""
+
+import datetime
+
+from airflow import models
+
+from dags import composer_env
+from dags.common import test_owner
+from dags.common.vm_resource import (
+    Project,
+    RuntimeVersion,
+    TpuVersion,
+    V6E_GCE_NETWORK,
+    V6E_GCE_SUBNETWORK,
+    Zone,
+)
+from dags.post_training.util import notebook_util, test_config_util
+from xlml.apis import gcp_config, metric_config, task, test_config
+
+SCHEDULE = "0 21 * * *" if composer_env.is_prod_env() else None
+DAG_TEST_NAME = "maxtext_rl_notebook"
+DEFAULT_BUCKET = "gs://rl-automation"
+
+with models.DAG(
+    dag_id=DAG_TEST_NAME,
+    start_date=datetime.datetime(2026, 1, 9),
+    schedule_interval=SCHEDULE,
+    catchup=False,
+    tags=[
+        "maxtext",
+        "post-training",
+        "rl",
+        "grpo",
+        "gspo",
+        "notebook",
+        "TPU",
+        "v6e-8",
+        "nightly",
+    ],
+    description="Automated Llama3.1-8B RL training from Jupyter notebook.",
+    doc_md="""
+      # Llama3.1-8B RL Training (Notebook Automation)
+
+      ### Overview
+      This DAG automates the rl_llama3_demo.ipynb notebook, which demonstrates
+      reinforcement learning training on Llama3.1-8B-Instruct using GRPO
+      (Group Relative Policy Optimization) or GSPO (Group Sequence Policy
+      Optimization) algorithms.
+
+      ### Prerequisites
+      - MaxText checkpoint for Llama3.1-8B-Instruct model
+      - HuggingFace access token with read permissions
+      - Single-host TPU VM (v6e-8 or v5p-8)
+
+      ### Execution Flow
+      1. **TPU Creation:** Create TPU VM with required specifications
+      2. **Environment Setup:** Clone MaxText, install dependencies
+      3. **RL Training:** Execute GRPO/GSPO training with reward model
+      4. **Log Validation:** Verify training completion signals
+      5. **Cleanup:** Delete TPU resources
+
+      ### Success Criteria
+      The test passes when:
+      1. TPU VM is created successfully
+      2. Training completes without errors
+      3. "RL Training Completed Successfully" appears in logs
+      4. Checkpoints are saved to output directory
+    """,
+    concurrency=1,
+) as dag:
+  # Test configuration
+  notebook_config = test_config_util.RLTestConfig(
+      cluster=None,  # Not used for TPU VM tests
+      accelerator="v6e-8",
+      slices=[1],
+      model_name="llama3.1-8b",
+      base_dir=f"{DEFAULT_BUCKET}/llama3.1-8b-Instruct/outputs",
+      tokenizer_path="meta-llama/Llama-3.1-8B-Instruct",
+      load_parameters_path=(
+          f"{DEFAULT_BUCKET}/llama3.1-8b-Instruct/scanned-pathways/0/items"
+      ),
+      loss_algos=[
+          test_config_util.LossAlgo.GRPO,
+          test_config_util.LossAlgo.GSPO,
+      ],
+  )
+
+  # HF token retrieved from Airflow Variables
+  HF_TOKEN_LLAMA31 = models.Variable.get("HF_TOKEN_CIENET", None)
+
+  # Test configuration
+  test_run_name = "llama31_rl_notebook"
+  current_datetime = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
+
+  # Setup commands for MaxText environment
+  setup_script = notebook_util.build_maxtext_setup_script()
+
+  # Path to the RL demo notebook
+  notebook_path = "src/MaxText/examples/rl_llama3_demo.ipynb"
+
+  # Test both GRPO and GSPO algorithms
+  for loss_algo in notebook_config.loss_algos:
+    run_name = f"{loss_algo.value}-{current_datetime}"
+
+    # Parameters to inject into notebook
+    notebook_params = {
+        "MODEL_CHECKPOINT_PATH": notebook_config.load_parameters_path,
+        "OUTPUT_DIRECTORY": notebook_config.base_dir,
+        "LOSS_ALGO": loss_algo.loss_name,
+    }
+
+    # Build notebook execution command
+    notebook_execution = notebook_util.build_notebook_execution_command(
+        notebook_path=notebook_path,
+        parameters=notebook_params,
+        maxtext_path="maxtext",
+        venv_path="maxtext_venv",
+        env_params={"HF_TOKEN": HF_TOKEN_LLAMA31},
+    )
+
+    # Create TPU VM test configuration
+    rl_notebook_test = test_config.TpuVmTest(
+        test_config.Tpu(
+            version=TpuVersion.TRILLIUM,
+            cores=8,
+            runtime_version=RuntimeVersion.V2_ALPHA_TPUV6.value,
+            reserved=False,
+            network=V6E_GCE_NETWORK,
+            subnetwork=V6E_GCE_SUBNETWORK,
+        ),
+        test_name=f"{DAG_TEST_NAME}_{loss_algo.value}",
+        set_up_cmds=[setup_script],
+        run_model_cmds=[notebook_execution],
+        timeout=datetime.timedelta(minutes=180),
+        task_owner=test_owner.JACKY_F,
+        num_slices=1,
+        gcs_subfolder=f"{DEFAULT_BUCKET}/{DAG_TEST_NAME}",
+    )
+
+    # Run the training task
+    training_task = task.run_queued_resource_test(
+        task_test_config=rl_notebook_test,
+        task_gcp_config=gcp_config.GCPConfig(
+            project_name=Project.CLOUD_ML_AUTO_SOLUTIONS.value,
+            zone=Zone.EUROPE_WEST4_A.value,
+            dataset_name=metric_config.DatasetOption.XLML_DATASET,
+        ),
+        skip_post_process=True,
+    )