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

[olufiyin19]

Description

This PR concludes the Agent Checkpoint Validation pipeline stack with external tooling and automation wrappers necessary to safely operate the system within Cloud Run and Airflow.

1. Airflow DAG Triggering & Monitoring

• Introduced wait_for_airflow_run.py to seamlessly bridge the Cloud Run Agent with Composer/Airflow DAG runs. The script polls the REST API sequentially, allowing the agent to continuously monitor asynchronous forward_pass and decode job metrics without timing out the Cloud Run sidecar.

2. Automated Email Alerts Delivery

• Built send_email.py as a standalone utility hook. When the validator detects an unresolvable numerical divergence or the agent optimization loop hits the maximum iteration count, it fires a structured Markdown email payload to alert the designated code owners directly.

3. Base Configuration & Lifecycle Tooling

• PyConfig Logging: Explicitly restored logger.setLevel(os.environ.get("LOGLEVEL", "INFO")) inside src/maxtext/configs/pyconfig.py (which had been inadvertently stripped from main). This ensures diagnostic execution traces are properly emitted to stdout for Airflow log capture.
• Ephemeral Branch Cleanup: Included branch_cleanup.py and connected GitHub Actions scripts to safely purge intermediate PR stacks triggered dynamically by the AI agent, avoiding CI/CD branch pollution.

Stack

6th PR in stack. Depends on the preceding upstream PR: #4729

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]

Code Review

This pull request introduces an automated model onboarding and verification pipeline (the Overwatch Agent) to validate converted model checkpoints, including shape validation, checkpoint inspection, forward pass logit verification, and decoding checks. It also adds a validation check in pyconfig.py to prevent string fields from being incorrectly set to None. Feedback on these changes highlights a critical bug in layer_metrics.py where formatting None values (due to NaN or Inf activations) will cause a TypeError and crash the pipeline. Additionally, reliability and portability improvements are recommended, specifically adding a timeout to urllib.request.urlopen in adk_agent.py and replacing the platform-dependent sed -i command in deploy_to_cloud_run.sh with a more portable alternative.

[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.

[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 completes the checkpoint validation pipeline stack by introducing several critical operational components, including an Airflow poller, automated email alerts via GCP Pub/Sub, and an Overwatch agent sidecar. While the high-level architecture and implementation are well-designed for a fast-failing automated validation system, there are several critical bugs, including a syntax error in decode_validator.py and a severe file overwrite risk in create_pull_request.py, that must be addressed before this branch is safe to merge.

🔍 General Feedback

• Well-Structured Automation Flow: The orchestration of shape checking, logit comparison, and automated patching is very cohesive, modular, and directly prevents expensive TPU compute waste on corrupt or mismatching checkpoints.
• Critical Syntactic and File-Handling Vulnerabilities: The syntax error in decode_validator.py makes that component instantly crash, and the recursive file copy in create_pull_request.py introduces a significant regression risk by potentially reverting other files in the main branch to outdated build-time states.
• Incomplete Test Assertions & Mocking: Several test suites (such as the decode and forward-pass validators) are either empty tautological tests or mock incorrect functions, leaving the core operational logic completely untested and uncovered. Improving test completeness is highly recommended to prevent automated agent failures.

[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 completes the Agent Checkpoint Validation pipeline stack by introducing robust Airflow polling, automated Pub/Sub alert hooks, and ephemeral branch cleanup tooling. The overall architecture is highly modular and designed to prevent TPU compute waste by enforcing fast-fail validation checks on model checkpoints.

🔍 General Feedback

• Subprocess Robustness: Improved handling of subprocess output streams using concurrent reading prevents deadlocks and timeout bypasses if the underlying validation processes freeze.
• On-Device Array Conversions: Safe detaching and copying of PyTorch on-device tensors ensures metrics calculation works perfectly across CPU, GPU, and TPU environments without raising conversion exceptions.
• CI/CD Cleanup: Improved Docker and Google Cloud Build setup scripts with strict traps ensures cleaner workspace states and avoids polluting temporary build artifacts.

[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 successfully rounds out the Overwatch Checkpoint Validation automated pipeline stack by integrating robust Airflow DAG polling, alert notification dispatch mechanisms, branch pruning automation, and log capture improvements. While the structural changes and pipeline operational scripts are well-implemented and highly useful, there are critical configuration and schema validation regressions related to DeepSeek V4 that must be resolved before merging to avoid production failures.

🔍 General Feedback

• Strong Pipeline Architecture: The modular separation of shape matching, logit verification, and decoding stages is an excellent design choice for a fast-fail workflow.
• Robust Error Handling & Fallbacks: The use of runtime monkeypatching and structured JSON report logging in the validators is highly commendable for stability.
• Merge Hygiene Caution: Please ensure that PR branches are rebased carefully onto main to prevent inadvertent deletion or reversion of recently merged features (such as DeepSeek V4 load balancing and logging attributes).

[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 successfully introduces the operational tooling, Airflow integrations, pyconfig overrides, and cleanup routines for the Agent Checkpoint Validation pipeline. The overall codebase is modular, clean, and establishes solid automated orchestration capabilities. However, a few critical issues regarding typing-conversion logic and polling resilience must be addressed to ensure the pipeline is stable and production-ready.

🔍 General Feedback

• Well-Structured Architecture: The modular separation between the shape, compiling, forward-pass, and decode-validation stages is clean and enforces a highly effective fast-fail system.
• Improved Observability: Restoring logging levels and outputting HF layer-by-layer details significantly improves diagnostic logging for automated analysis.
• Robust Cleaning: Including automated branch pruning and scheduled pruning GitHub Actions workflows maintains hygiene across the codebase.
• Robustness Overwork: Dynamic runtime monkeypatching (especially source string rewrites) inside forward_pass_validator.py introduces fragility. In the future, resolving upstream bugs inside model_creation_utils.py directly is highly encouraged.