feat: MaxText Checkpoint Validation Agent Shape and Mock Tensor sub DAGs

[olufiyin19]

Description

This PR introduces the Airflow DAG infrastructure for the first two phases of the Automated Checkpoint Validation Pipeline. It establishes a master orchestrator that chains validation tasks sequentially, ensuring we "fast fail" on structural issues before spinning up expensive compute resources for tensor math.

Changes included in this PR:

• maxtext_validation_master_dag.py: Added a master orchestrator DAG that sequentially triggers our validation Sub-DAGs.

• dag_verify_checkpoint_shape.py: Added Sub-DAG A (Checkpoint Inspection). This performs a fast, low-compute structural check to ensure checkpoint tensor shapes and layer counts match the expected model configuration.

• dag_verify_mock_tensor.py: Added Sub-DAG B (Mock Tensor Validation). This performs a dry-run forward pass using randomized mock tensors to verify memory allocation and syntax without verifying logit mathematics.

• lib/utils.py: Added get_maxtext_validation_config() helper to abstract boilerplate TPU Kubernetes configuration. It dynamically builds the bash execution environment, injects maxtext_commit_hash overrides via Airflow Jinja templating, and creates the TpuGkeTest object so the Sub-DAGs remain clean.

Tests

Tests Performed:

• I deployed these DAGs to a remote Airflow instance via Cloud VM and triggered the maxtext_validation_master_dag using a custom JSON DEFAULT_PARAMS payload (testing the qwen3-8b architecture).

• Verified that the Master DAG correctly passed the runtime configuration down to the Sub-DAGs using TriggerDagRunOperator.

• Verified that Sub-DAG A correctly pulled the checkpoint shapes and executed.

• Verified that the Master DAG correctly halted until Sub-DAG A completed, and then successfully triggered Sub-DAG B.

Instruction and/or command lines to reproduce your tests:

• Deploy the DAG files to your Composer environment's GCS bucket using bash scripts/upload-tests.sh

• Open the GCP Console, navigate to the tpu-prod-env-multipod project, and go to "Managed Airflow".

• Open the Airflow Webserver for the maxtext-validation-agent-dev environment.

• Search for maxtext_validation_master_dag, click on it, and select "Trigger DAG w/ config".

• Provide a test JSON payload and trigger. You can monitor the sub-tasks via the "Triggered DAG" hyperlinks in the execution details. (See attached screenshots).

List links for your tests (use go/shortn-gen for any internal link): ...

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.

[Rohan-Bierneni]

Thanks for this change!

The code looks good but please make sure to align with any changes/updates made in the maxtext pr. Lets first merge the maxtext pr, then have a quick e2e test run before merging this pr.