Escalate XPK clean-up failure to infrastructure team

[olufiyin19]

Escalate XPK clean-up failure to infrastructure team

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[gemini-code-assist]

Code Review

This pull request introduces an automated model onboarding and verification pipeline, featuring a sidecar agent for remediation, multiple validation scripts, and comprehensive training engine parity tests. The code changes include the addition of the agent's core logic, prompt templates, and various utility tools for shape inspection, PR management, and Airflow orchestration. Feedback from the review includes a suggestion to add an explicit handler for the 'escalate' remediation level in the agent to prevent incorrect patching attempts, a recommendation to narrow exception handling in the Airflow polling tool to avoid masking unexpected errors, and a note to remove a redundant print statement in the PR creation tool.

[gemini-code-assist]

The Analyst subagent prompt defines remediation_level as potentially being escalate (e.g., for infrastructure or XPK clean-up failures). However, the agent currently lacks an explicit handler for the escalate case, causing it to fall through to the else block and attempt Phase 3 code patching. Adding an explicit elif remediation_level == "escalate": block to trigger an alert email and halt the workflow will prevent unnecessary patching attempts for infrastructure-level failures.

  if remediation_level == "level_1_config":
    logger.info("Level 1 Config Repair detected. Short-circuiting Phase 3 code patching and git branch creation.")
    fixer_response_text = f"Level 1 Config Repair: Overrides identified = {json.dumps(config_overrides)}"
    new_branch = maxtext_branch
  elif remediation_level == "escalate":
    logger.info("Escalation requested by Analyst. Sending alert email and halting workflow.")
    subject = f"CRITICAL: Validation Pipeline Escalation for {model_name}"
    body = (
        f"The Analyst subagent has requested escalation for run {run_id}.\n\n"
        f"Diagnosis: {plan_json.get('diagnosis')}\n"
        f"Evidence: {plan_json.get('evidence')}"
    )
    send_alert_email(subject, body)
    return f"Escalated to infrastructure team: {plan_json.get('diagnosis')}"
  else:

[gemini-code-assist]

Catching a broad Exception and silently passing (pass) is an anti-pattern that can hide unexpected programming errors (such as NameError, TypeError, or AttributeError), making debugging extremely difficult. It is highly recommended to only catch expected network exceptions like requests.RequestException to allow other unexpected errors to propagate and fail fast.

Suggested change

	except (requests.RequestException, Exception) as e:
	# Log warning but don't fail, allowing subsequent poll iterations to retry
	pass
	except requests.RequestException as e:
	# Log warning but don't fail, allowing subsequent poll iterations to retry
	pass

[gemini-code-assist]

There is a duplicate print statement here. Removing the redundant line will clean up the output logs.

Suggested change

	print("Syncing modified files to git repository...")
	print("Syncing modified files to git repository...")
	print("Syncing modified files to git repository...")