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.