Checkpoint Validation Pipeline Operations: Airflow polling integrations, PyConfig overrides, and pipeline lifecycle tooling

.dockerignore

@@ -1,2 +1,4 @@
 .git
 maxtext_venv
+.venv
+Qwen3-4B-Weights*

.github/workflows/cleanup_abandoned_agent_branches.yml

@@ -0,0 +1,33 @@
+name: Cleanup Abandoned Agent Branches
+
+on:
+  schedule:
+    # Run daily at 00:00 UTC
+    - cron: '0 0 * * *'
+  workflow_dispatch:
+    inputs:
+      days:
+        description: 'Inactivity threshold in days'
+        required: false
+        default: '14'
+
+jobs:
+  cleanup:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: write
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+
+      - name: Run branch cleanup script
+        run: |
+          python3 src/maxtext/experimental/agent/ckpt_validation_pipeline/agent_sidecar/monitor/branch_cleanup.py \
+            --days "${{ github.event.inputs.days || '14' }}"

src/maxtext/configs/base.yml

@@ -1086,6 +1086,7 @@ profile_power_events: false # Set to true to enable TPU-specific power/thermal p
 
 log_config: true # Prints the config (after defaults have been set by pyconfig logic)
 debug_sharding: false # Prints model weights sharding info
+debug_tensors: false # Captures intermediate tensors during forward pass using NNX sow
 
 # Checkpoint Structured logging
 enable_checkpoint_cloud_logger: false

src/maxtext/configs/pyconfig.py

@@ -38,6 +38,10 @@
 from maxtext.utils import max_logging
 
 logger = logging.getLogger(__name__)
+try:
+  logger.setLevel(os.environ.get("LOGLEVEL", "INFO").upper())
+except ValueError:
+  logger.setLevel(logging.INFO)
 
 _BASE_CONFIG_ATTR = "base_config"
 _MAX_PREFIX = "M_"
@@ -272,7 +276,24 @@ def _prepare_for_pydantic(raw_keys: dict[str, Any], config_class: type[Any] = ty
 
     new_value = value
     if isinstance(new_value, str) and new_value.lower() == "none":
-      new_value = None
+      field_info = valid_fields.get(key)
+      if field_info:
+        ann = field_info.annotation
+        import typing
+        import types as python_types
+
+        def _allows_none(annotation) -> bool:
+          if annotation is None or annotation is type(None) or annotation is typing.Any:
+            return True
+          origin = typing.get_origin(annotation)
+          if origin in (typing.Union, getattr(python_types, "UnionType", None)):
+            return any(arg is type(None) or arg is None or arg is typing.Any for arg in typing.get_args(annotation))
+          return False
+
+        if _allows_none(ann):
+          new_value = None
+      else:
+        new_value = None
 
     # Pydantic validates enums from their values, so string is fine.
     # It also handles type coercion for simple types.

src/maxtext/configs/types.py

@@ -311,6 +311,7 @@ class RunInfo(BaseModel):
       description="If True, prints the final configuration after initialization.",
   )
   debug_sharding: bool = Field(False, description="If True, print model weight sharding details.")
+  debug_tensors: bool = Field(False, description="Captures intermediate tensors during forward pass using NNX sow")
   base_output_directory: PathStr = Field("", description="Base directory for all outputs, typically a GCS path.")
   sharding_strategy: None | Literal["experimental"] = Field(
       None,

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"]