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

dags/common/scheduling_helper/dag_integrity_whitelist.yaml

@@ -178,4 +178,8 @@ whitelisted_dags:
 - trt_llm_mlperf_v41
 - trtllm_bench_inference
 - xlml_to_buganizer
-
+- maxtext_validation_master_dag 
+- dag_verify_checkpoint_shape # maxtext validation sub dag
+- dag_verify_forward_compile # maxtext validation sub dag
+- dag_verify_forward_pass # maxtext validation sub dag
+- dag_verify_decoding # maxtext validation sub dag

dags/maxtext_validation_agent/.airflowignore

@@ -0,0 +1 @@
+tests/

dags/maxtext_validation_agent/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

dags/maxtext_validation_agent/commands.txt

@@ -0,0 +1,46 @@
+<1> to get the airflow running with your json file from maxtext repo
+bash scripts/upload-tests.sh gs://us-central1-maxtext-validat-16478791-bucket/dags
+
+<2> to check available clusters
+gcloud container clusters list --project=tpu-prod-env-multipod
+
+<3> Please remember you're here to be critical, help me, and not just agree with me
+
+<4> Do not start task except i ask you to
+
+<5> aound line 1073 in ~/ml-auto-solutions/xlml/utils/metric.py, REMOVE the try-except block and simply
+replace it with the commented out original line when you want to make your PR
+
+<6> If you paste your personal Hugging Face token into the Airflow Web UI for your dev environment (maxtext-validation-agent-dev-fiyinbenstowe), only DAGs running in your specific dev sandbox will use it.
+However, here is the reality of how this works when you move to production:
+If you delete it later: Your Mock Tensor DAG will instantly break again with the exact same KeyError: 'Variable HF_TOKEN does not exist' error.
+The Production Solution: You will never use your personal token in the main ml-auto-solutions production environment. When your PR is merged, the team will create a shared "Service Account" or "Bot" token on Hugging Face (e.g., a generic Google/MaxText bot account). An admin will save that token into the Production Airflow Variables.
+For right now, using your personal token in your isolated dev environment is the standard way to test. Just remember to revoke it on your Hugging Face account later if you ever feel it was compromised.
+
+<7> What happens when Snehal merges to main?
+Your intuition is 100% correct. Hardcoding "git clone -b dsv4-conversion-pr" is technical debt.
+If Snehal merges his code to main and deletes his PR branch tomorrow, your Airflow DAG will immediately fail with a fatal: Remote branch not found error.
+The Engineering Standard:
+It is totally fine to hardcode the branch right now while you are building in a dev sandbox (maxtext-validation-agent-dev-fiyinbenstowe). However, before you submit your Pull Request to ml-auto-solutions, you must parameterize it.
+You will change your code to pull the branch name dynamically from an Airflow DAG run configuration, defaulting to main:
+f"git clone -b {{{{ dag_run.conf.get('maxtext_branch', 'main') }}}} [https://github.com/](https://github.com/)..."
+This way, production runs main automatically, but developers can trigger the DAG manually and pass in their own feature branches to test them. For today, stick to the hardcoded branch just to get the pod passing.
+
+<8> rsync commands for mac:
+rsync -avz /Users/fiyinbenstowe/Desktop/project/ml-auto-solutions/ jfacevedo-v6e-8:~/ml-auto-solutions/
+rsync -avz /Users/fiyinbenstowe/Desktop/project/maxtext/ jfacevedo-v6e-8:~/maxtext/
+
+<9> rsync commands for cloudtop:
+rsync -avz ~/Desktop/Project/ml-auto-solutions/ jfacevedo-v6e-8:~/ml-auto-solutions/
+rsync -avz ~/Desktop/Project/maxtext/ jfacevedo-v6e-8:~/maxtext/
+
+<10> for maxtext linting:
+pyink --pyink-indentation=2 --line-length=122 src/maxtext/experimental/agent/ckpt_validation_pipeline
+pylint src/maxtext/experimental/agent/ckpt_validation_pipeline
+
+<11> In maxtext, change the recipient email from mine to specifc user's
+
+<12> To edit and retest the agent, run in maxtext directory on local terminal
+bash src/maxtext/experimental/agent/ckpt_validation_pipeline/agent_sidecar/deploy_to_cloud_run.sh
+
+<13>

dags/maxtext_validation_agent/dag_verify_checkpoint_shape.py

@@ -0,0 +1,73 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""DAG to automate MaxText Checkpoint Structural Shape Validation."""
+
+# pylint: disable=line-too-long
+
+import datetime
+from airflow import models
+from dags.maxtext_validation_agent.lib import utils
+from dags.maxtext_validation_agent.lib.utils import trigger_agent_on_failure
+
+
+DEFAULT_PARAMS = {
+    "checkpoint_gcs_path": "",
+    "hf_config_url": "",
+    "hf_model_path": "",
+    "hf_ref_code_url": "",
+    "hf_token": "",
+    "max_kl_div": "",
+    "maxtext_branch": "",
+    "maxtext_commit_hash": "",
+    "maxtext_model_name": "",
+    "forward_pass_maxtext_overrides": {
+        "attention": "",
+        "scan_layers": "",
+        "weight_dtype": "",
+        "tokenizer_path": "",
+        "tokenizer_type": "",
+    },
+    "report_gcs_dir": "",
+    "run_name": "",
+    "xpk_cluster_name": "",
+    "xpk_project": "",
+    "xpk_zone": "",
+}
+
+with models.DAG(
+    dag_id="dag_verify_checkpoint_shape",
+    schedule=None,
+    tags=["maxtext", "checkpoint", "validation"],
+    start_date=datetime.datetime(2026, 6, 26),
+    catchup=False,
+    params=DEFAULT_PARAMS,
+    default_args={
+        "retries": 0,
+        "on_failure_callback": trigger_agent_on_failure,
+    },
+) as dag:
+    # Looks for keys in runtime conf first (from manual JSON or Master DAG),
+    # falls back to defaults if run is standalone.
+
+    checkpoint_task = utils.get_checkpoint_shape_validation_task(
+        dag=dag,
+        model_name="{{ dag_run.conf.get('maxtext_model_name', params['maxtext_model_name']) }}",
+        checkpoint_gcs_path="{{ dag_run.conf.get('checkpoint_gcs_path', params['checkpoint_gcs_path']) }}",
+        scan_layers="{{ dag_run.conf.get('forward_pass_maxtext_overrides', params['forward_pass_maxtext_overrides']).get('scan_layers', False) | lower }}",
+    )
+
+    # Execute Task A
+    check_task = utils.get_upstream_failure_validator_task(dag)
+    checkpoint_task >> check_task

dags/maxtext_validation_agent/dag_verify_decoding.py

@@ -0,0 +1,77 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""DAG to automate MaxText Checkpoint Decoding Validation (Sub-DAG D)."""
+
+# pylint: disable=line-too-long
+
+import datetime
+from airflow import models
+from dags.maxtext_validation_agent.lib import utils
+from dags.maxtext_validation_agent.lib.utils import trigger_agent_on_failure
+from dags.common.vm_resource import XpkClusters, TpuVersion
+
+DEFAULT_PARAMS = {
+    "email": "",
+    "xpk_cluster_name": "",
+    "xpk_project": "",
+    "xpk_zone": "",
+    "checkpoint_gcs_path": "",
+    "decode_maxtext_overrides": {
+        "attention": "",
+        "per_device_batch_size": "",
+        "scan_layers": "",
+        "tokenizer_path": "",
+        "tokenizer_type": "",
+        "weight_dtype": "",
+    },
+    "hf_config_url": "",
+    "hf_model_path": "",
+    "hf_ref_code_url": "",
+    "hf_token": "",
+    "max_kl_div": "",
+    "maxtext_branch": "",
+    "maxtext_commit_hash": "",
+    "maxtext_model_name": "",
+    "report_gcs_dir": "",
+    "run_name": "",
+}
+
+with models.DAG(
+    dag_id="dag_verify_decoding",
+    schedule=None,
+    tags=["maxtext", "checkpoint", "decoding", "validation"],
+    start_date=datetime.datetime(2026, 6, 26),
+    catchup=False,
+    params=DEFAULT_PARAMS,
+    default_args={
+        "retries": 0,
+        "on_failure_callback": trigger_agent_on_failure,
+    },
+) as dag:
+    cluster_param = DEFAULT_PARAMS.get("xpk_cluster_name", "v4-8-maxtext")
+    cluster_name = (
+        cluster_param.default
+        if hasattr(cluster_param, "default")
+        else str(cluster_param)
+    )
+    cluster_config = utils.get_cluster_config(cluster_name)
+
+    decoding_task = utils.get_decoding_validation_task(
+        cluster_config=cluster_config,
+        time_out_in_min=45,
+    ).run(skip_post_process=True)
+
+    check_task = utils.get_upstream_failure_validator_task(dag)
+    decoding_task >> check_task

dags/maxtext_validation_agent/dag_verify_forward_compile.py

@@ -0,0 +1,71 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""DAG to automate MaxText Checkpoint Inspection (Task B)."""
+
+# pylint: disable=line-too-long
+
+import datetime
+from airflow import models
+from dags.maxtext_validation_agent.lib import utils
+from dags.maxtext_validation_agent.lib.utils import trigger_agent_on_failure
+
+
+DEFAULT_PARAMS = {
+    "email": "",
+    "xpk_cluster_name": "",
+    "xpk_project": "",
+    "xpk_zone": "",
+    "checkpoint_gcs_path": "",
+    "forward_pass_maxtext_overrides": {
+        "attention": "",
+        "per_device_batch_size": "",
+        "scan_layers": "",
+        "tokenizer_path": "",
+        "tokenizer_type": "",
+        "weight_dtype": "",
+    },
+    "hf_config_url": "",
+    "hf_model_path": "",
+    "hf_ref_code_url": "",
+    "hf_token": "",
+    "max_kl_div": "",
+    "maxtext_branch": "",
+    "maxtext_commit_hash": "",
+    "maxtext_model_name": "",
+    "report_gcs_dir": "",
+    "run_name": "",
+}
+
+with models.DAG(
+    dag_id="dag_verify_forward_compile",
+    schedule=None,
+    tags=["maxtext", "checkpoint", "inspection"],
+    start_date=datetime.datetime(2026, 6, 26),
+    catchup=False,
+    params=DEFAULT_PARAMS,
+    default_args={
+        "retries": 0,
+        "on_failure_callback": trigger_agent_on_failure,
+    },
+) as dag:
+    # Looks for keys in runtime conf first (from manual JSON or Master DAG),
+    # falls back to defaults if run standalone.
+    forward_compile_task = utils.get_forward_compile_validation_task(
+        dag=dag,
+    )
+
+    # Execute Task B
+    check_task = utils.get_upstream_failure_validator_task(dag)
+    forward_compile_task >> check_task