Add Self-Patching Agent for the Checkpoint Validation Pipeline

[olufiyin19]

Description

This PR introduces the automated Fixer Agent sidecar (agent_sidecar/) built natively on the Google GenAI Agent Development Kit (ADK). This is the fifth and final PR in the Checkpoint Validation Agent stack (building upon PR #4462, PR #4726, #4727, and #4728), completing the autonomous self-healing checkpoint conversion and validation loop.

When checkpoint conversion or validation pipelines fail in Airflow across varied model architectures (e.g. Gemma, Llama, Deepseek, Mistral, Qwen), failures often stem from diagnosable root causes such as remat_policy out-of-memory errors during evaluation, tensor shape transpositions, or missing configuration overrides. Previously, engineers had to manually inspect logs, formulate patches, lint code, re-trigger pipelines, and open pull requests by hand.

This PR provides a containerized, continuously running sidecar service that monitors Google Cloud Storage for validation failure reports, autonomously diagnoses root causes using specialized inspection tools, applies fixes conforming to domain constraints, verifies code against strict repository linting standards (10.00/10 pylint), triggers Airflow validation DAGs, and opens clean GitHub pull requests without human intervention.

Key Implementation Details

1. Google GenAI ADK Multi-Stage Agent Architecture

• Implements adk_agent.py using the native Agent Development Kit (ADK), structured around a multi-stage prompt workflow (01_diagnose.txt, 02_patch.txt, 03_verify.txt, and meta_agent.txt).
• Encodes critical MaxText domain constraints into agent instructions, including:
  * Distinguishing between evaluation vs. training tasks when configuring remat_policy.
  * Using remat_policy=minimal as the standard fallback for evaluation out-of-memory failures.
  * Escaping JSON braces in diagnostic prompts to prevent prompt injection or formatting collisions.

2. Specialized Autonomous Tool Suite (agent_sidecar/fixer/tools/)

• analyze_shapes.py: Compares expected vs. actual checkpoint tensor shapes to isolate parameter dimension mismatches.
• analyze_layer_activations.py: Integrates with layer_metrics.py to localize per-layer numerical drift or NaN/Inf contamination.
• run_linters.py: Runs automated pylint and pyink checks on AI-generated patches, enforcing a strict 10.00/10 pylint score before committing code.
• trigger_airflow_dag.py: Programmatically re-triggers Airflow DAG runs to verify that code patches resolve the original failure end-to-end.
• github_branch_manager.py & create_pull_request.py: Automates branch creation, git staging, and GitHub PR submission with detailed root-cause summaries.

3. GCS Monitoring, State Tracking & Cloud Run Infrastructure

• Implements gcs_poller.py to continuously poll GCS buckets for FAILED validation JSON reports, coordinated via state_manager.py to prevent duplicate repair attempts.
• Includes containerization and serverless deployment scripts (Dockerfile & deploy_to_cloud_run.sh) for deployment to Google Cloud Run.

Stack

5th PR in stack. Depends on the preceding upstream PR: #4728

Checklist

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

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[github-actions]

🤖 Hi @entrpn, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This PR introduces an automated Fixer Agent sidecar built on the Google GenAI Agent Development Kit (ADK) to establish a serverless self-patching and validation pipeline for converted model checkpoints. The overall architectural design is robust and highly innovative, integrating GCS polling, multi-agent reasoning loops (Analyst, Fixer, Verifier), standard repo formatting, and automated Airflow re-triggering to resolve validation failures without human intervention.

🔍 General Feedback

• Positive Highlights: The code structure is incredibly clean, modular, and extremely well-organized. Using native Google GenAI SDK chats with automatic function calling provides a highly scalable framework for orchestration.
• Git State Management: The current git init setup inside the Docker container is a critical structural issue because it creates a disconnected repository with unrelated history compared to origin, causing git push and PR creation to fail. This is easily corrected by initializing the repo and doing a soft reset to align its commit graph with origin.
• Environment and Path Configuration: A couple of hardcoded developer laptop absolute paths (/Users/fiyinbenstowe/...) exist in alerter.py and the prompt template 03_verify.txt. Converting these to relative paths or configuration environment variables ensures seamless operation in Google Cloud Run.
• JSON Compatibility: To prevent downstream parse failures in Google Cloud, avoiding non-standard JSON representations like raw NaN values is highly recommended.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[olufiyin19]

@gemini-cli /review

[github-actions]

🤖 Hi @olufiyin19, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This Pull Request introduces the automated Fixer Agent sidecar built natively on the Google GenAI Agent Development Kit (ADK) to complete the checkpoint validation loop. While the overall multi-stage architecture and autonomous self-healing concepts are highly innovative and well-documented, the implementation currently contains multiple critical correctness bugs and missing files that will prevent it from running successfully in production.

🔍 General Feedback

• High Technical Debt & Fragility: The use of runtime regex-based string manipulation and dynamic exec() to monkeypatch central model restoration routines is extremely fragile. Fixing the underlying bugs directly inside the local model_creation_utils.py module is a much more maintainable approach.
• Robust Exception and Resource Handling: Ensure subprocesses are clean, linting checks are enforced natively, and credentials/environment variables are verified before the container initializes.
• Unused/Superfluous Disables: There are several unnecessary # pylint: disable directives in the files (e.g. pylint: disable=syntax-error) which should be cleaned up to maintain high code-quality scores.