Checkpoint Validation Pipeline Operations: Airflow polling integrations, PyConfig overrides, and pipeline lifecycle tooling - #4800
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
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.
93482fb to
aa0f3ab
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
|
@gemini-cli /review |
|
🤖 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. |
|
@gemini-cli /review |
|
🤖 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[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
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.
dd145f1 to
8488751
Compare
|
@gemini-cli /review |
|
🤖 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. |
There was a problem hiding this comment.
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
mainto prevent inadvertent deletion or reversion of recently merged features (such as DeepSeek V4 load balancing and logging attributes).
3378f32 to
8488751
Compare
8488751 to
fb12c29
Compare
|
@gemini-cli /review |
|
🤖 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. |
There was a problem hiding this comment.
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.pyintroduces fragility. In the future, resolving upstream bugs insidemodel_creation_utils.pydirectly is highly encouraged.
fb12c29 to
a0edc81
Compare
a0edc81 to
9bd506a
Compare
4279cf1 to
25724ad
Compare
| # Note: In a production environment, the Airflow Webserver URL and auth token would be injected via environment variables. | ||
| AIRFLOW_URL = os.environ.get( | ||
| "AIRFLOW_WEBSERVER_URL", | ||
| "https://4bae0a6de8f94e92aa8ee3a6ffc8b278-dot-us-central1.composer.googleusercontent.com", |
There was a problem hiding this comment.
I dont' think we should expose google urls, even if internal only.
|
|
||
| # Compile the final Markdown body with Github & Airflow links | ||
| pr_url = state_entry.get("pr_url", "https://github.com/AI-Hypercomputer/maxtext/pulls") | ||
| log_url = f"https://console.cloud.google.com/run/jobs/executions/list?project=tpu-prod-env-multipod" |
There was a problem hiding this comment.
if the project is hardcoded here, what happens if the job runs in a different project?
| except Exception as e: | ||
| logger.error("Error checking GCS for failures: %s", e) | ||
|
|
||
| return None, None |
There was a problem hiding this comment.
instead create a report_data = None, blob = None and return that instead.
|
|
||
| # Ensure temporary Dockerfile in root is cleaned up on exit | ||
| trap "rm -f ./Dockerfile .gcloudignore.tmp" EXIT | ||
| PROJECT_ID="tpu-prod-env-multipod" |
There was a problem hiding this comment.
it would be nice if these are configurable through env variables.
| 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`). |
There was a problem hiding this comment.
has qwen3-8b and llama3-70b been tested in this pipeline? If not, remove from readme and add models that have been tested e2e.
| @@ -1,2 +1,4 @@ | |||
| .git | |||
| maxtext_venv | |||
| .venv | |||
There was a problem hiding this comment.
Remove these lines, keep them in your own environment.
2 participants
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.
wait_for_airflow_run.pyto 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.send_email.pyas 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.logger.setLevel(os.environ.get("LOGLEVEL", "INFO"))insidesrc/maxtext/configs/pyconfig.py(which had been inadvertently stripped frommain). This ensures diagnostic execution traces are properly emitted to stdout for Airflow log capture.branch_cleanup.pyand 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):
gemini-reviewlabel.