Fix resource allocation for post-training notebook DAGs (GoogleCloudPlatform#1262)

1. Support v5e and v6e TPU versions in the DAG.
2. Read YAML configurations from a GCS bucket to set the TPU version and zone for the run.

Author: ooops678

dags/post_training/maxtext_rl_notebook.py

@@ -8,7 +8,6 @@
 import datetime
 
 from airflow import models
-from airflow.models.baseoperator import chain
 
 from dags import composer_env
 from dags.common import test_owner
@@ -76,23 +75,26 @@
   test_run_name = "llama31_rl_notebook"
   current_datetime = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
 
+  # Load configurations from GCS
+  config_arg = notebook_util.load_notebook_config_from_gcs_yaml(
+      gcs_path=notebook_util.NOTEBOOK_CONFIG_GCS_PATH,
+      dag_name=DAG_TEST_NAME,
+  )
+  config = notebook_util.NotebookConfig(config_arg)
+
   # Setup commands for MaxText environment
   setup_script = notebook_util.build_maxtext_setup_script()
 
-  notebook_rl_tests = []
-  # Test both GRPO and GSPO algorithms
+  previous_task = config_arg
   for loss_algo in loss_algos:
-    rl_notebook_test = notebook_util.initialize_notebook_test(
-        test_name=f"{DAG_TEST_NAME}_rl_{loss_algo.value}",
+    previous_task = notebook_util.run_notebook_tests(
         dag_name=DAG_TEST_NAME,
+        task_id_prefix=f"rl_{loss_algo.value}",
         notebook_path="src/maxtext/examples/rl_llama3_demo.ipynb",
         set_up_script=setup_script,
         parameters={"LOSS_ALGO": loss_algo.loss_name},
         task_owner=test_owner.DEPP_L,
+        hf_token=HF_TOKEN_LLAMA31,
+        config=config,
+        previous_task=previous_task,
     )
-
-    notebook_rl_tests.append(
-        notebook_util.run_training(rl_notebook_test, HF_TOKEN_LLAMA31)
-    )
-
-  chain(*notebook_rl_tests)

dags/post_training/maxtext_sft_notebook.py

@@ -72,17 +72,24 @@
       f"{gcs_bucket.ORBAX_AUTOMATION_BUCKET_EUROPE_WEST4}/post-training/sft/"
   )
 
+  # Load configurations from GCS
+  config_arg = notebook_util.load_notebook_config_from_gcs_yaml(
+      gcs_path=notebook_util.NOTEBOOK_CONFIG_GCS_PATH,
+      dag_name=DAG_TEST_NAME,
+  )
+  config = notebook_util.NotebookConfig(config_arg)
+
   # Setup commands for MaxText environment
   setup_script = notebook_util.build_maxtext_setup_script()
 
-  # Test SFT training
-  sft_notebook_test = notebook_util.initialize_notebook_test(
-      test_name=f"{DAG_TEST_NAME}_sft",
+  notebook_util.run_notebook_tests(
       dag_name=DAG_TEST_NAME,
+      task_id_prefix="sft",
       notebook_path="src/maxtext/examples/sft_llama3_demo_tpu.ipynb",
       set_up_script=setup_script,
       parameters={"BASE_OUTPUT_DIRECTORY": BASE_OUTPUT_DIRECTORY},
       task_owner=test_owner.DEPP_L,
+      hf_token=HF_TOKEN_LLAMA31,
+      config=config,
+      previous_task=config_arg,
   )
-
-  notebook_util.run_training(sft_notebook_test, HF_TOKEN_LLAMA31)

dags/post_training/util/notebook_util.py

@@ -2,8 +2,18 @@
 
 import datetime
 import inspect
+import logging
 import textwrap
+import airflow
+import json
+import re
+
+from airflow.decorators import task as airflow_task
 from airflow.models.taskmixin import DAGNode
+from airflow.models.baseoperator import chain
+from airflow.operators.empty import EmptyOperator
+from airflow.utils.trigger_rule import TriggerRule
+from airflow.utils.task_group import TaskGroup
 
 from dags.common.vm_resource import (
     Project,
@@ -14,7 +24,11 @@
     Zone,
 )
 from dags.post_training.util import test_config_util
-from xlml.apis import gcp_config, metric_config, task, test_config
+from xlml.apis import gcp_config, gcs, metric_config, task, test_config
+
+NOTEBOOK_CONFIG_GCS_PATH = (
+    "gs://ml-auto-solutions-dag-configs/post-training/notebook_dag_configs.yaml"
+)
 
 
 def build_maxtext_setup_script() -> str:
@@ -78,9 +92,6 @@ def _run_parameter_injection(
       parameters: Dict of {key: value} for literal injection.
       env_params: List of keys to be injected as `os.getenv("KEY")`.
   """
-  import json
-  import re
-
   with open(notebook_path, encoding="utf-8") as f:
     nb = json.load(f)
 
@@ -252,6 +263,7 @@ def initialize_notebook_test(
     set_up_script: str,
     parameters: dict[str, any],
     task_owner: str,
+    tpu_version: TpuVersion,
 ) -> test_config.TpuVmTest:
   """Creates a TpuVmTest configuration for notebook execution."""
   notebook_execution = build_notebook_execution_command(
@@ -262,7 +274,7 @@ def initialize_notebook_test(
   )
   return test_config.TpuVmTest(
       test_config.Tpu(
-          version=TpuVersion.V5E,
+          version=tpu_version,
           cores=8,
           runtime_version=RuntimeVersion.V2_ALPHA_TPUV6.value,
           reserved=False,
@@ -279,14 +291,143 @@ def initialize_notebook_test(
   )
 
 
-def run_training(config: test_config.TpuVmTest, hf_token: str) -> DAGNode:
+class NotebookConfig:
+  """A simple container holding dynamic XComArg fields."""
+
+  def __init__(self, config_arg: airflow.XComArg) -> None:
+    self.zone = config_arg["zone"]
+    self.tpu_version = config_arg["tpu_version"]
+
+
+@airflow_task(multiple_outputs=True)
+def load_notebook_config_from_gcs_yaml(
+    gcs_path: str, dag_name: str
+) -> dict[str, str]:
+  """Loads and parses TPU version and zone configs from GCS yaml config."""
+  config = gcs.load_yaml_from_gcs(gcs_path)
+  dag_cfg = config.get("dag", {}).get(dag_name, {})
+
+  tpu_version = dag_cfg.get("tpu_version")
+  zone = dag_cfg.get("zone")
+
+  # Validate that GCS config values correspond to valid Enum values
+  def assert_is_valid_enum(value: str, enum_class) -> None:
+    try:
+      enum_class(value)
+    except ValueError as e:
+      raise ValueError(f"Config Validation Error: {e}") from e
+
+  assert_is_valid_enum(zone, Zone)
+  assert_is_valid_enum(tpu_version, TpuVersion)
+
+  logging.info(
+      f"Loaded configuration: tpu_version='{tpu_version}', zone='{zone}'."
+  )
+
+  return {"tpu_version": tpu_version, "zone": zone}
+
+
+def run_training(
+    config: test_config.TpuVmTest, hf_token: str, zone: str | airflow.XComArg
+) -> DAGNode:
   return task.run_queued_resource_test(
       task_test_config=config,
       task_gcp_config=gcp_config.GCPConfig(
           project_name=Project.CLOUD_ML_AUTO_SOLUTIONS.value,
-          zone=Zone.EUROPE_WEST4_B.value,
+          zone=zone,
           dataset_name=metric_config.DatasetOption.XLML_DATASET,
       ),
       skip_post_process=True,
       custom_env={"HF_TOKEN": hf_token},
   )
+
+
+def run_notebook_tests(
+    dag_name: str,
+    task_id_prefix: str,
+    notebook_path: str,
+    set_up_script: str,
+    parameters: dict[str, any],
+    task_owner: str,
+    hf_token: str,
+    config: NotebookConfig,
+    previous_task: DAGNode | None = None,
+) -> TaskGroup:
+  """Creates and chains branched notebook tests for all TPU versions.
+
+  Args:
+    dag_name: Name of the DAG.
+    task_id_prefix: Prefix for task and operator IDs (e.g. "rl_grpo" or "sft").
+    notebook_path: Path to the notebook to run.
+    set_up_script: Setup script for MaxText environment.
+    parameters: Dict of parameters to inject in the notebook.
+    task_owner: Owner of the task.
+    hf_token: HuggingFace access token.
+    config: A `NotebookConfig` wrapper containing zone and tpu_version.
+    previous_task: Optional task/DAGNode to chain *before* the branches.
+
+  Returns:
+    A TaskGroup representing the entire branched notebook test workflow.
+  """
+  with TaskGroup(
+      group_id=f"{task_id_prefix}_tests", prefix_group_id=False
+  ) as group:
+    # 1. Initialize and create the V5E test and task group
+    notebook_test_v5e = initialize_notebook_test(
+        test_name=f"{dag_name}_{task_id_prefix}",
+        dag_name=dag_name,
+        notebook_path=notebook_path,
+        set_up_script=set_up_script,
+        parameters=parameters,
+        task_owner=task_owner,
+        tpu_version=TpuVersion.V5E,
+    )
+    run_task_v5e = run_training(notebook_test_v5e, hf_token, zone=config.zone)
+
+    # 2. Initialize and create the TRILLIUM/V6E test and task group
+    notebook_test_v6e = initialize_notebook_test(
+        test_name=f"{dag_name}_{task_id_prefix}",
+        dag_name=dag_name,
+        notebook_path=notebook_path,
+        set_up_script=set_up_script,
+        parameters=parameters,
+        task_owner=task_owner,
+        tpu_version=TpuVersion.TRILLIUM,
+    )
+    run_task_v6e = run_training(notebook_test_v6e, hf_token, zone=config.zone)
+
+    # 3. Create skipped fallback empty operator task
+    skipped = EmptyOperator(task_id=f"skipped_{task_id_prefix}")
+
+    # 4. Define central Task-decorated Branch Operator accepting dynamic parameters
+    @airflow_task.branch(
+        task_id=f"task_path_decider_{task_id_prefix}",
+        trigger_rule=TriggerRule.ALL_DONE,
+        retries=0,
+    )
+    def task_path_decider(tpu_version: str) -> str:
+      logging.info(f"Configured active TPU version: '{tpu_version}'")
+      active_tpu_version = TpuVersion(tpu_version)
+
+      match active_tpu_version:
+        case TpuVersion.V5E:
+          decided_task_id = run_task_v5e.group_id
+        case TpuVersion.TRILLIUM:
+          decided_task_id = run_task_v6e.group_id
+        case _:
+          decided_task_id = skipped.task_id
+
+      logging.info(f"running task_id: {decided_task_id}")
+      return decided_task_id
+
+    # 5. Instantiate branch decider task, passing XComArg directly for Airflow auto-resolution
+    task_decider = task_path_decider(tpu_version=config.tpu_version)
+
+    # 6. Chain previous tasks to the decider if exist
+    if previous_task:
+      chain(previous_task, task_decider)
+
+    # 7. Connect branch decider to target branches
+    chain(task_decider, [run_task_v5e, run_task_v6e, skipped])
+
+  return group

xlml/apis/gcp_config.py

@@ -15,7 +15,11 @@
 """Config file for Google Cloud Project (GCP)."""
 
 import dataclasses
+from typing import Any
 
+from airflow.exceptions import AirflowException
+from airflow.models.xcom_arg import XComArg
+from airflow.operators.python import get_current_context
 from dags.common.vm_resource import Project
 from xlml.apis import metric_config
 
@@ -33,7 +37,32 @@ class GCPConfig:
   """
 
   project_name: str
-  zone: str
+  zone: str | XComArg
   dataset_name: metric_config.DatasetOption
   dataset_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value
   composer_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value
+
+  def __getattribute__(self, name: str) -> Any:
+    # First obtain the underlying value of the attribute.
+    value = super().__getattribute__(name)
+
+    # Actively resolve an XComArg under execution phase.
+    match value:
+      case XComArg():
+        try:
+          context = get_current_context()
+        except AirflowException:
+          # AirflowException means we are not in execution phase yet,
+          # return the XComArg as-is
+          return value
+
+        resolved_value = value.resolve(context)
+
+        # Store the resolved result back so that
+        # further references won't have to resolve again.
+        super().__setattr__(name, resolved_value)
+
+        return resolved_value
+
+      case _:
+        return value

xlml/utils/tpu.py

@@ -23,6 +23,8 @@
 
 from absl import logging
 import airflow
+
+from airflow.utils.trigger_rule import TriggerRule
 from airflow.decorators import task, task_group
 from airflow.utils.task_group import TaskGroup
 from airflow.operators.python import get_current_context
@@ -253,7 +255,7 @@ def delete_queued_resource(qualified_name: airflow.XComArg):
       resource.
   """
 
-  @task(trigger_rule='all_done')
+  @task(trigger_rule=TriggerRule.NONE_SKIPPED)
   def delete_tpu_nodes_request(qualified_name: str):
     # TODO(wcromar): Find a less repetitive way to manage the TPU client.
     creds, _ = google.auth.default()
@@ -299,7 +301,7 @@ def wait_for_tpu_deletion(qualified_name: str):
     logging.info(f'TPU Nodes: {qr.tpu.node_spec}')
     return False
 
-  @task(trigger_rule='all_done')
+  @task(trigger_rule=TriggerRule.NONE_SKIPPED)
   def delete_queued_resource_request(qualified_name: str) -> Optional[str]:
     creds, _ = google.auth.default()
     client = tpu_api.TpuClient(credentials=creds)