Escalate XPK clean-up failure to infrastructure team

src/maxtext/experimental/agent/ckpt_validation_pipeline/README.md

@@ -0,0 +1,116 @@
+# Automated Model Onboarding & Verification Pipeline
+
+This pipeline is used to automate the validation of converted model checkpoints. It is designed to be triggered deterministically by Airflow DAGs to verify the correctness of model checkpoints in a fast-fail architecture, preventing the waste of expensive TPU compute on malformed checkpoints.
+
+If a step fails, the Overwatch Agent analyzes the divergence, attempts to fix the MaxText code, and re-runs the validation step automatically.
+
+## The Pipeline Lifecycle
+
+1. **Task A: Shape Matching (Mock Tensor) - The "Fast Fail"**
+   Validates basic matrix shapes and model architecture acceptance using mock tensors in seconds.
+   Script: `checkpoint_shape_validator.py`
+
+2. **Task B: Checkpoint Inspection**
+   Inspects the structure of the Orbax/MaxText checkpoint to ensure all required files and layers are present in GCS.
+   Script: [`inspect_checkpoint.py`](/src/maxtext/checkpoint_conversion/inspect_checkpoint.py)
+
+3. **Task C: Forward Pass Logit Verification** (WIP)
+   Runs the model on PyTorch and MaxText simultaneously and compares the intermediate layer outputs (using Flax `sow`) to catch the exact layer where a conversion bug exists.
+   Script: `forward_pass_validator.py`
+
+4. **Task D: SFT & Decoding (Caching Logic)** (WIP)
+   *   **SFT**: Tests the backward pass by running training steps to ensure loss decreases without hitting NaNs.
+   *   **Decoding Check**: Tests text generation and autoregressive caching logic (KV Cache) for new models.
+   Script: `decode_validator.py`
+
+## Quick starts
+To begin, you'll need:
+
+1. A valid Google Cloud Storage (GCS) bucket where your converted checkpoint is located (e.g., `gs://my-bucket/converted_ckpt/0/items`).
+2. The corresponding MaxText internal model name (e.g., `qwen3-8b`, `llama3-70b`).
+3. To trigger the pipeline via the Airflow UI using the `maxtext_validation_agent` DAG.
+4. A full run of the pipeline should typically take about 1-2 hours if all stages pass.
+
+## 1. Prepare the inputs (Shape Validation)
+
+The first step of the pipeline (`checkpoint_shape_validator.py`) requires context files about the theoretical MaxText blueprint and the actual Orbax checkpoint layer. You can generate them using the `inspect_checkpoint.py` tool.
+
+* **Theoretical MaxText Blueprint**: Generated on-the-fly dynamically by parsing abstract JAX shapes without executing compute. Following MaxText's architecture transition, this now validates shapes against **NNX** model trees by default. (A legacy Linen `init` fallback is preserved via a custom `inspect_checkpoint.py` specifically to support older models like Deepseekv4).
+* **Actual Orbax Checkpoint Layer**: Generated by reading the `safetensors` or `pth` file headers to extract metadata instantly, avoiding host RAM allocation.
+
+The Airflow DAG automatically generates these `/tmp/ideal_shapes.txt` and `/tmp/actual_shapes.txt` files and passes them to the validator.
+
+## 2. Run the pipeline
+While the primary interaction is via the Airflow UI, you can execute the validation process step-by-step manually.
+
+## Manual Run Instructions (For Debugging)
+
+### Step 1: Shape Validation (No TPU Required)
+
+> **Note on Device Expectations**: Steps 1 and 2 rely on abstract shape tracing (`jax.eval_shape`) and mock tensors. Because they do not execute actual math, they are extremely cheap and **do not require a TPU** (they can run on a standard CPU VM or a TPU VM without locking the chips). In contrast, the subsequent downstream steps (Logit Verification and Decoding) execute the actual model weights and explicitly require TPU hardware (e.g. v4-8) to run.
+
+```bash
+python3 src/maxtext/experimental/agent/ckpt_validation_pipeline/checkpoint_shape_validator.py \
+  --ideal_shapes_path=/tmp/ideal_shapes.txt \
+  --actual_shapes_path=/tmp/actual_shapes.txt \
+  --report_gcs_dir=gs://your-bucket/reports/
+```
+
+### Step 2: Forward Compile Validation (Mock Tensors)
+
+```bash
+python3 src/maxtext/experimental/agent/ckpt_validation_pipeline/forward_compile_validator.py \
+  --checkpoint_gcs_path=gs://your-bucket/checkpoint/0/items \
+  --maxtext_model_name=qwen3-8b \
+  --report_gcs_dir=gs://your-bucket/reports/ \
+  --scan_layers=true
+```
+
+### Step 3: Forward Pass Logit Verification (WIP)
+
+```bash
+python3 src/maxtext/experimental/agent/ckpt_validation_pipeline/forward_pass_validator.py \
+  --checkpoint_gcs_path=gs://your-bucket/checkpoint/0/items \
+  --maxtext_model_name=qwen3-8b \
+  --run_hf_model=true \
+  --hf_model_path=Qwen/Qwen2.5-7B-Instruct \
+  --report_gcs_dir=gs://your-bucket/reports/
+```
+
+### Step 4: Decoding (Caching Logic) Verification (WIP)
+
+```bash
+python3 src/maxtext/experimental/agent/ckpt_validation_pipeline/decode_validator.py \
+  --checkpoint_gcs_path=gs://your-bucket/checkpoint/0/items \
+  --maxtext_model_name=qwen3-8b \
+  --report_gcs_dir=gs://your-bucket/reports/
+```
+
+## Architecture Notes (Linen vs. NNX)
+
+MaxText currently supports two neural network frameworks internally: Flax Linen and the newer Flax NNX. 
+**Going forward, NNX is the only supported architecture.** DeepSeekV4 is officially the last model that will be compatible with Linen. All validation scripts have been migrated to support NNX abstract states natively:
+
+*   **`forward_compile_validator.py` (NNX):** Uses the `create_nnx_abstract_model` abstraction.
+*   **`checkpoint_shape_validator.py` (NNX):** Theoretical inputs are derived from `inspect_checkpoint.py`, which supports extracting the parameter tree from `nnx.State`.
+*   **`decode_validator.py` & `forward_pass_validator.py` (NNX):** Will automatically use NNX models directly without falling back to Linen overrides.
+
+### Reading the JSON Reports
+
+If you specified `--report_gcs_dir=gs://your-bucket/reports/`, each step will upload a JSON file containing the validation results.
+* **Success**: The status will be `"SUCCESS"` and the pipeline proceeds to the next stage.
+* **Failure**: The status will be `"FAILURE"` and the `error_message` or `stderr` field will contain the stack trace. 
+
+## Debugging tips
+
+1. If a validation step fails in Airflow, check the task logs directly in the Airflow UI to see the exact stdout/stderr from the Python script.
+2. If the **Shape Validation** fails, ensure your model configuration matches the checkpoint architecture exactly.
+3. If the **Forward Compile** fails, look for OOMs or distributed check failures that might indicate incorrect batch size or sequence length overrides.
+4. If the **Forward Pass** fails with `401 Unauthorized`, ensure you are using an open HuggingFace model or providing a valid `HF_TOKEN`.
+5. If the **Decoding** step fails, check the KV caching parameters in your model configuration.
+
+## Tests
+Run standard MaxText tests:
+```bash
+python3 -m pytest src/maxtext/experimental/agent/ckpt_validation_pipeline/tests/
+```

src/maxtext/experimental/agent/ckpt_validation_pipeline/__init__.py

@@ -0,0 +1,20 @@
+# Copyright 2023-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 "innovation" 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.
+
+"""
+Checkpoint Validation Agent Package.
+Used to verify and report the status of converted model checkpoints.
+"""
+
+from maxtext.experimental.agent.ckpt_validation_pipeline import layer_metrics

src/maxtext/experimental/agent/ckpt_validation_pipeline/agent_sidecar/Dockerfile

@@ -0,0 +1,33 @@
+FROM python:3.12-slim
+
+WORKDIR /app
+
+# Install git, curl, gpg, and GitHub CLI (gh)
+RUN apt-get update && apt-get install -y git curl ca-certificates gpg && \
+    curl -fsSL https://cli.github.com/packages/githubcli-archive-keyring.gpg | dd of=/usr/share/keyrings/githubcli-archive-keyring.gpg && \
+    chmod go+r /usr/share/keyrings/githubcli-archive-keyring.gpg && \
+    echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/githubcli-archive-keyring.gpg] https://cli.github.com/packages stable main" | tee /etc/apt/sources.list.d/github-cli.list > /dev/null && \
+    echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] http://packages.cloud.google.com/apt cloud-sdk main" | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list && \
+    curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | gpg --dearmor -o /usr/share/keyrings/cloud.google.gpg && \
+    apt-get update && apt-get install -y gh google-cloud-cli && \
+    rm -rf /var/lib/apt/lists/*
+
+# Copy only requirements first to leverage Docker cache for heavy installations
+COPY src/dependencies/requirements/generated_requirements/tpu-requirements.txt /tmp/tpu-requirements.txt
+COPY pyproject.toml /app/pyproject.toml
+
+# Automatically install all MaxText TPU requirements and local package
+RUN pip install --no-cache-dir google-cloud-storage google-genai requests google-auth pyink pylint && \
+    pip install --no-cache-dir -r /tmp/tpu-requirements.txt
+
+# Now copy the full source code (any python code changes will invalidate this layer but skip the pip install)
+COPY . /app
+RUN pip install --no-cache-dir --no-deps -e .
+
+# Set git global identity for commits and associate local repo with origin/main history
+RUN git config --global user.email "overwatch-agent@google.com" && \
+    git config --global user.name "Overwatch Agent" 
+
+ENV PYTHONPATH="/app/src/maxtext/utils:/app/src:/app"
+
+CMD ["python", "src/maxtext/experimental/agent/ckpt_validation_pipeline/agent_sidecar/main.py"]