feat: Add Vertex AI integration for log uploads (GoogleCloudPlatform#1169)

This change adds support for uploading training metrics to Vertex AI TensorBoard for both RL and SFT post-training DAGs, improves run name generation for better traceability, and enhances multi-slice training support. The changes include new utility functions, configuration structures, and updates to training command generation and orchestration.

Author: RexBearIU

dags/post_training/maxtext_rl.py

@@ -18,6 +18,7 @@
 from dags.common.vm_resource import XpkClusters
 from dags.multipod.configs import gke_config
 from dags.post_training.util import validation_util, test_config_util
+from dags.post_training.util.test_config_util import VertexAI
 from xlml.utils.xpk import MAIN_BRANCH
 from xlml.utils.gke import zone_to_region
 
@@ -52,21 +53,23 @@
       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
+      5. **Vertex AI Upload:** Execute script to upload metrics to TensorBoard
+      6. **Success/Failure:** Report status from logs and completion
 
       ### Success Criteria
       The test passes when:
       1. Training jobs complete successfully without errors
       2. "Post RL Training" log message appears in jax-tpu container logs
       3. No infrastructure failures or container launch issues occur
       4. All training steps execute within expected parameters
+      5. Metrics are successfully synced to Vertex AI TensorBoard
     """,
     concurrency=2,
 ) as dag:
   training_config = test_config_util.RLTestConfig(
       cluster=XpkClusters.TPU_V5P_128_CLUSTER,
       accelerator="v5p-128",
-      slices=[1],  # Single slice for RL training
+      slices=[1, 2],  # Multi-slice support
       model_name="llama3.1-70b",
       base_dir=(
           f"{test_config_util.DEFAULT_BUCKET}/llama3.1-70b-Instruct/outputs"
@@ -92,13 +95,17 @@
 
     for loss_algo in training_config.loss_algos:
       for slice_num in training_config.slices:
-        current_datetime = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
-        run_name = f"{loss_algo.value}-{mode.value}-{current_datetime}"
+        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(
@@ -166,7 +173,13 @@
                 )
             )
 
-          chain(
-              training_group,
-              validation_group,
+          upload_to_vertex_ai = validation_util.upload_to_vertex_ai(
+              project_id=VertexAI.POST_TRAINING.project_id,
+              region=VertexAI.POST_TRAINING.region,
+              tensorboard_id=VertexAI.POST_TRAINING.tensorboard_id,
+              logdir=f"{training_config.base_dir}/{run_name}/tensorboard/",
+              experiment_name=loss_algo.value,
+              run_name_prefix=run_name,
           )
+
+          chain(training_group, [validation_group, upload_to_vertex_ai])

dags/post_training/maxtext_sft.py

@@ -20,6 +20,7 @@
 from dags.common.vm_resource import DockerImage, XpkClusters
 from dags.multipod.configs import gke_config
 from dags.post_training.util import validation_util, test_config_util
+from dags.post_training.util.test_config_util import VertexAI
 from xlml.utils.xpk import MAIN_BRANCH
 from xlml.utils.gke import zone_to_region
 
@@ -128,7 +129,7 @@ def validate_training(
   training_config = test_config_util.SFTTestConfig(
       cluster=XpkClusters.TPU_V5P_128_CLUSTER,
       accelerator="v5p-128",
-      slices=[1],  # Single slice for SFT training
+      slices=[2],
       model_name="llama3.1-70b",
       steps=training_steps,
       short_id="msft",
@@ -151,8 +152,11 @@ def validate_training(
       continue  # Skip stable for SFT training tests
 
     for slice_num in training_config.slices:
-      current_datetime = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
-      run_name = f"sft-{mode.value}-{current_datetime}"
+      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,
@@ -168,4 +172,13 @@ def validate_training(
             training_config, training_steps, start_time, end_time
         )
 
-        chain(training_group, validation_group)
+        upload_to_vertex_ai = validation_util.upload_to_vertex_ai(
+            project_id=VertexAI.POST_TRAINING.project_id,
+            region=VertexAI.POST_TRAINING.region,
+            tensorboard_id=VertexAI.POST_TRAINING.tensorboard_id,
+            logdir=f"{training_config.base_dir}/{run_name}/tensorboard/",
+            experiment_name="sft",
+            run_name_prefix=run_name,
+        )
+
+        chain(training_group, [validation_group, upload_to_vertex_ai])

dags/post_training/util/test_config_util.py

@@ -4,10 +4,34 @@
 from enum import Enum
 
 from dags import gcs_bucket
-from dags.common.vm_resource import XpkClusters, DockerImage
+from dags.common.vm_resource import (
+    XpkClusters,
+    DockerImage,
+    Project,
+    Region,
+)
 from dags.multipod.configs.common import SetupMode
 
 
+@dataclass(frozen=True)
+class VertexAIConfig:
+  """Configuration for Vertex AI TensorBoard instance."""
+
+  project_id: str
+  region: str
+  tensorboard_id: str
+
+
+class VertexAI:
+  """Vertex AI TensorBoard instances for post-training dashboards."""
+
+  POST_TRAINING = VertexAIConfig(
+      project_id=Project.CLOUD_TPU_MULTIPOD_DEV.value,
+      region=Region.US_CENTRAL1.value,
+      tensorboard_id="9105049192543289344",
+  )
+
+
 DEFAULT_BUCKET = gcs_bucket.RL_AUTOMATION_BUCKET
 
 # Docker images for post-training
@@ -107,40 +131,56 @@ def __init__(
     self.rl_config_path = rl_config_path
 
   def generate_rl_training_command(
-      self, loss_algo: LossAlgo, run_name: str, hf_token: str
+      self,
+      loss_algo: LossAlgo,
+      run_name: str,
+      hf_token: str,
+      num_slices: int = 1,
   ) -> tuple[str]:
-    """Generates the RL training command as a tuple for GKE compatibility.
+    """Generates the RL training command as a tuple for GKE compatibility."""
+    command_params = [
+        f"run_name={run_name}",
+        f"model_name={self.model_name}",
+        f"tokenizer_path={self.tokenizer_path}",
+        f"load_parameters_path={self.load_parameters_path}",
+        f"base_output_directory={self.base_dir}",
+        f"rl.loss_algo={loss_algo.loss_name}",
+    ]
 
-    Args:
-      loss_algo: The loss algorithm to use (e.g., LossAlgo.GRPO).
-      run_name: The run name for the training job.
-      hf_token: The HuggingFace token for authentication.
+    num_trainer = 1
+    num_samplers = max(1, num_slices - num_trainer)
+
+    if num_slices > 1:
+      command_params.extend([
+          f"num_trainer_slices={num_trainer}",
+          f"num_samplers_slices={num_samplers}",
+          f"rollout_data_parallelism={num_samplers * 2}",
+      ])
 
-    Returns:
-      A tuple containing the RL training command string.
-    """
     rl_command = (
         "python -m src.MaxText.rl.train_rl "
-        f"{self.rl_config_path} run_name={run_name} "
-        f"model_name={self.model_name} "
-        f"tokenizer_path={self.tokenizer_path} "
-        f"load_parameters_path={self.load_parameters_path} "
-        f"base_output_directory={self.base_dir} "
-        f"rl.loss_algo={loss_algo.loss_name}"
+        f"{self.rl_config_path} " + " ".join(command_params)
     )
-    command = " && ".join([
+
+    environment_variables = [
         f"export HF_TOKEN={hf_token}",
         "export TPU_MIN_LOG_LEVEL=0",
         "export TF_CPP_MIN_LOG_LEVEL=0",
         "export TPU_STDERR_LOG_LEVEL=0",
         "export JAX_PLATFORMS=proxy,cpu",
         "export JAX_BACKEND_TARGET=grpc://127.0.0.1:29000",
         "export ENABLE_PATHWAYS_PERSISTENCE='1'",
-        rl_command,
-    ])
+    ]
 
-    # Return as tuple for k8s yaml compatibility.
-    return (command,)
+    if num_slices > 1:
+      environment_variables.extend([
+          f"export NUM_SLICES={num_samplers}",
+          "export JAX_RANDOM_WEIGHTS=true",
+          "export VLLM_ENABLE_V1_MULTIPROCESSING=0",
+      ])
+
+    full_command = " && ".join(environment_variables + [rl_command])
+    return (full_command,)
 
 
 @dataclass

dags/post_training/util/validation_util.py

@@ -4,13 +4,78 @@
 workflows, reusing generic utilities from the orbax module where applicable.
 """
 
+from absl import logging
+from datetime import datetime, timezone
+
+from airflow import AirflowException
+from airflow.decorators import task
+from google.cloud import aiplatform
+from google.cloud import storage
+
 # Re-export commonly used validation functions from orbax
 from dags.orbax.util.validation_util import (
     generate_timestamp,
     validate_log_exist,
 )
 
+
+def _ensure_tfevents_exist(logdir: str) -> None:
+  """Ensures the provided GCS logdir contains TensorBoard event files."""
+  if not logdir.startswith("gs://"):
+    raise AirflowException("Vertex AI upload only supports GCS paths.")
+
+  gcs_path = logdir[len("gs://") :]
+  bucket_name, _, prefix = gcs_path.partition("/")
+  prefix = prefix.lstrip("/")
+  storage_client = storage.Client()
+  blobs = storage_client.list_blobs(bucket_name, prefix=prefix)
+  has_event_files = any(".tfevents." in blob.name for blob in blobs)
+  if not has_event_files:
+    raise AirflowException(
+        "No TensorBoard event files found at logdir: "
+        f"{logdir}. Ensure the path contains .tfevents.* files."
+    )
+
+
+@task
+def generate_run_name(
+    prefix: str,
+    mode: str,
+    num_slices: int,
+) -> str:
+  """Generates a unique run name with a timestamp."""
+  current_datetime = datetime.now(timezone.utc).strftime("%Y-%m-%d-%H-%M-%S")
+  return f"{prefix}-{mode}-{num_slices}x-{current_datetime}"
+
+
+@task
+def upload_to_vertex_ai(
+    project_id: str,
+    region: str,
+    tensorboard_id: str,
+    logdir: str,
+    experiment_name: str,
+    run_name_prefix: str,
+) -> None:
+  """Uploads TensorBoard logs to Vertex AI.
+
+  This task uses the Vertex AI Python SDK to sync logs from a GCS directory
+  to a specific TensorBoard experiment.
+  """
+  _ensure_tfevents_exist(logdir)
+
+  aiplatform.init(project=project_id, location=region)
+  aiplatform.upload_tb_log(
+      tensorboard_id=tensorboard_id,
+      tensorboard_experiment_name=experiment_name,
+      logdir=logdir,
+      run_name_prefix=run_name_prefix,
+  )
+  logging.info("Upload completed successfully.")
+
+
 __all__ = [
     "generate_timestamp",
     "validate_log_exist",
+    "upload_to_vertex_ai",
 ]