feat(airflow): Add semantic inference verification for MaxText E2E DAGs

dags/multipod/maxtext_e2e_tpu_post_training.py

@@ -36,6 +36,7 @@
 from dags.common.quarantined_tests import safe_get_from_variable
 from dags.common.vm_resource import GkeClusters
 from dags.multipod.configs import gke_config
+from dags.multipod.util.validation_util import validate_semantic_inference_output, generate_timestamp
 
 # HF token retrieved from Airflow Variables for secure credential management
 HF_TOKEN = safe_get_from_variable("HF_TOKEN", None)
@@ -248,6 +249,10 @@ def poke(self, context, session=None):
               "core_count", test_config.get("core_count", 8)
           )
           mode_short_name = "multim" if mode == "multimodal_sft" else mode
+          start_time = generate_timestamp.override(
+              task_id="generate_start_time"
+          )()
+
           training_task = gke_config.get_gke_config(
               time_out_in_min=60,
               num_slices=1,
@@ -264,6 +269,8 @@ def poke(self, context, session=None):
               use_gcluster=True,
           ).run(skip_post_process=True)
 
+          end_time = generate_timestamp.override(task_id="generate_end_time")()
+
           to_hf_flags = mode_test_config.get("to_hf_flags", "false true")
           to_hf_script = test_config["to_huggingface"]
           convert_to_huggingface_cmd = (
@@ -284,8 +291,24 @@ def poke(self, context, session=None):
               mounts="/dev/shm;/mnt/shm;rw",
           ).run(skip_post_process=True)
 
+          validation_task = validate_semantic_inference_output.override(
+              task_id="validate_inference"
+          )(
+              project_id=GkeClusters.TPU_V5P_BODABORG_NAP_CLUSTER.project,
+              location=GkeClusters.TPU_V5P_BODABORG_NAP_CLUSTER.zone,
+              cluster_name=GkeClusters.TPU_V5P_BODABORG_NAP_CLUSTER.name,
+              namespace=GkeClusters.TPU_V5P_BODABORG_NAP_CLUSTER.namespace,
+              pod_pattern=f"{mode_short_name}-v5p-.*",
+              container_name="workload-container",
+              start_time=start_time,
+              end_time=end_time,
+          )
+
           chain(
               wait_for_conversion,
+              start_time,
               training_task,
+              end_time,
+              validation_task,
               convert_to_huggingface_task,
           )

dags/multipod/util/validation_util.py

@@ -0,0 +1,130 @@
+import re
+from datetime import datetime, timezone, timedelta
+from typing import Optional
+from absl import logging
+
+from airflow.decorators import task
+from airflow.exceptions import AirflowFailException
+from google.cloud import logging as logging_api
+
+
+def list_log_entries(
+    project_id: str,
+    location: str,
+    cluster_name: str,
+    namespace: str = "default",
+    pod_pattern: str = ".*",
+    container_name: Optional[str] = None,
+    text_filter: Optional[str] = None,
+    start_time: Optional[datetime] = None,
+    end_time: Optional[datetime] = None,
+) -> list[logging_api.LogEntry]:
+  """
+  List log entries for the specified Google Cloud project.
+  """
+  logging_client = logging_api.Client(project=project_id)
+
+  # Set the time window for log retrieval:
+  # default to last 12 hours if not provided
+  if end_time is None:
+    end_time = datetime.now(timezone.utc)
+  if start_time is None:
+    start_time = end_time - timedelta(hours=12)
+
+  # Format times as RFC3339 UTC "Zulu" format required by the Logging API
+  start_time_str = start_time.strftime("%Y-%m-%dT%H:%M:%SZ")
+  end_time_str = end_time.strftime("%Y-%m-%dT%H:%M:%SZ")
+
+  conditions = [
+      f'resource.labels.project_id="{project_id}"',
+      f'resource.labels.location="{location}"',
+      f'resource.labels.cluster_name="{cluster_name}"',
+      f'resource.labels.namespace_name="{namespace}"',
+      f'resource.labels.pod_name=~"{pod_pattern}"',
+      "severity>=DEFAULT",
+      f'timestamp>="{start_time_str}"',
+      f'timestamp<="{end_time_str}"',
+  ]
+
+  if container_name:
+    conditions.append(f'resource.labels.container_name="{container_name}"')
+  if text_filter:
+    conditions.append(f"{text_filter}")
+
+  log_filter = " AND ".join(conditions)
+
+  logging.info(f"Log filter constructed: {log_filter}")
+  return list(logging_client.list_entries(filter_=log_filter))
+
+
+@task
+def validate_semantic_inference_output(
+    project_id: str,
+    location: str,
+    cluster_name: str,
+    namespace: str = "default",
+    pod_pattern: str = ".*",
+    container_name: Optional[str] = None,
+    text_filter: Optional[str] = None,
+    start_time: Optional[datetime] = None,
+    end_time: Optional[datetime] = None,
+) -> None:
+  """Validates the log output for semantic inference quality."""
+  entries = list_log_entries(
+      project_id=project_id,
+      location=location,
+      cluster_name=cluster_name,
+      namespace=namespace,
+      pod_pattern=pod_pattern,
+      container_name=container_name,
+      text_filter=text_filter,
+      start_time=start_time,
+      end_time=end_time,
+  )
+
+  if not entries:
+    raise AirflowFailException(
+        "The log history is empty! Cannot validate inference."
+    )
+
+  buffer = ""
+  for entry in entries:
+    message = None
+    if isinstance(entry, logging_api.TextEntry):
+      message = entry.payload
+    elif isinstance(entry, logging_api.StructEntry):
+      message = entry.payload.get("message")
+
+    if message:
+      buffer += message + "\\n"
+
+  # Isolate the post-training inference which is the last output in the script
+  # The script runs pre-train inference -> train -> post-train inference
+  inference_blocks = buffer.split("Generated text:")
+  if len(inference_blocks) > 1:
+    eval_buffer = inference_blocks[-1]
+    logging.info("Isolated the final 'Generated text:' block for validation.")
+  else:
+    eval_buffer = buffer
+    logging.warning(
+        "Could not find 'Generated text:' delimiter, validating entire log."
+    )
+
+  if "\ufffd" in eval_buffer:
+    raise AirflowFailException(
+        "Semantic validation failed: Non-UTF-8 characters detected (\\ufffd) in post-training inference."
+    )
+
+  if not re.search(
+      r"(?i)(big ben|london eye|tower of london|buckingham palace)", eval_buffer
+  ):
+    raise AirflowFailException(
+        "Semantic validation failed: Expected London landmark keywords not found in post-training inference output."
+    )
+
+  logging.info("Semantic inference verification passed successfully.")
+
+
+@task
+def generate_timestamp():
+  return datetime.now(timezone.utc)