Skip to content

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

Draft
chiajunglien wants to merge 1 commit into
GoogleCloudPlatform:masterfrom
CIeNET-International:maxtext/user/emma/add-verify-main
Draft

chiajunglien wants to merge 1 commit into
GoogleCloudPlatform:masterfrom
CIeNET-International:maxtext/user/emma/add-verify-main

Conversation

@chiajunglien

@chiajunglien chiajunglien commented Sep 8, 2026

Copy link
Copy Markdown
Contributor

Description

Issue:

Currently, our Airflow E2E DAGs (e.g., maxtext_e2e_tpu_post_training.py) rely solely on exit 0 to validate job success. However, during recent model testing (e.g., Llama3.1-70B and Gemma3-4B), we discovered "silent failures" where the inference process completes successfully but generates invalid text, nonsensical repetitions, or produces completely empty output.

Changes:

  1. Semantic Inference Validation Task: Implemented a decoupled validate_semantic_inference_output
    task that securely queries logs from Cloud Logging for the workload-container.
  2. Robust Quality Checks: Added regex parsing to isolate the final Generated text: output block,
    detect non-UTF-8 corruptions (\ufffd), and ensure expected semantic keywords (e.g., London
    landmarks) are present.
  3. Directory Restructuring: Moved the utility functions to dags/multipod/util/validation_util.py to
    maintain a clean directory structure.

Tests

https://06ba93284e31466eb62067a2f46710a7-dot-us-east1.composer.googleusercontent.com/dags/test_inference_validation/grid?search=test_inference_validation

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run one-shot tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant