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 the "Overwatch" automated model onboarding and verification pipeline for MaxText, which automates checkpoint validation, includes a Cloud Run sidecar agent driven by Gemini for auto-remediation, and adds a comprehensive training engine parity verification suite. The review feedback focuses on enhancing the robustness, performance, and style of the newly added scripts. Key recommendations include adding repository cloning fallbacks in the PR tool, enforcing timeouts on network requests, implementing strict type checks on parsed JSON configurations to prevent runtime crashes, logging caught exceptions during polling, limiting GCS blob listings to avoid performance degradation, avoiding fragile runtime source code modifications, and resolving a PEP 8 naming convention violation.

[gemini-code-assist]

If /tmp/maxtext_repo does not exist (e.g., if this tool is executed first or the directory was cleaned up), subprocess.run with cwd=repo_dir will raise a FileNotFoundError and crash. Clone the repository if it does not exist to ensure robustness.

Suggested change

	repo_dir = "/tmp/maxtext_repo"
	fork_branch = args.fix_branch or f"fix/agent-remediation-{int(time.time())}"
	print(f"1. Checking out fix branch '{fork_branch}'...")
	repo_dir = "/tmp/maxtext_repo"
	fork_branch = args.fix_branch or f"fix/agent-remediation-{int(time.time())}"
	if not os.path.exists(repo_dir):
	print(f"Cloning repository into {repo_dir}...")
	subprocess.run(
	["git", "clone", "https://github.com/AI-Hypercomputer/maxtext.git", repo_dir], check=True
	)
	print(f"1. Checking out fix branch '{fork_branch}'...")
	subprocess.run(["git", "checkout", fork_branch], cwd=repo_dir, check=False)

[gemini-code-assist]

The urllib.request.urlopen function is called without a timeout. If the remote server hangs or takes too long to respond, this call will block indefinitely, which can stall the entire agent execution. It is highly recommended to specify a timeout.

Suggested change

	with urllib.request.urlopen(req) as response:
	with urllib.request.urlopen(req, timeout=30) as response:

[gemini-code-assist]

If the LLM response is not a JSON object (e.g., if it returns a list or a string), calling .get() on plan_json will raise an AttributeError and crash the agent. Adding a type check ensures robust error handling.

    parsed = json.loads(raw_text.strip())
    if not isinstance(parsed, dict):
      raise ValueError("Analyst response is not a JSON object")
    plan_json = parsed
    logger.info(f"Analyst diagnosis: {plan_json.get('diagnosis')}")

[gemini-code-assist]

If original_conf is not a dictionary (e.g., if json.loads returned a list or string), calling .get() will raise an AttributeError. Similarly, if clean_conf is not a dictionary, "dag_conf" in clean_conf can raise a TypeError. Adding type checks prevents runtime crashes.

Suggested change

	clean_conf = original_conf.get("dag_conf", original_conf)
	# Just in case we are dealing with an already nested config from before this fix,
	# gracefully un-nest it.
	while "dag_conf" in clean_conf:
	clean_conf = clean_conf["dag_conf"]
	conf_dict.update(clean_conf)
	if isinstance(original_conf, dict):
	clean_conf = original_conf.get("dag_conf", original_conf)
	# Just in case we are dealing with an already nested config from before this fix,
	# gracefully un-nest it.
	while isinstance(clean_conf, dict) and "dag_conf" in clean_conf:
	clean_conf = clean_conf["dag_conf"]
	if isinstance(clean_conf, dict):
	conf_dict.update(clean_conf)

[gemini-code-assist]

Silently catching and passing on all exceptions (including Exception) is dangerous because it can mask critical bugs (like NameError or TypeError) and cause the loop to run unnecessarily until the 2-hour timeout. Print the warning to sys.stderr so that failures are visible in the logs.

Suggested change

	except (requests.RequestException, Exception) as e:
	# Log warning but don't fail, allowing subsequent poll iterations to retry
	pass
	except (requests.RequestException, Exception) as e:
	# Log warning but don't fail, allowing subsequent poll iterations to retry
	print(f"Warning: encountered exception during poll: {e}", file=sys.stderr)

[gemini-code-assist]

PEP 8 style guide violation. Never use the character 'l' (lowercase letter el) as a single-character variable name because it is indistinguishable from the numeral one ('1') or uppercase letter 'I' in many fonts.

Suggested change

	for logger_name in [None, "absl"]:
	l = logging.getLogger(logger_name)
	for h in l.handlers:
	for logger_name in [None, "absl"]:
	log_obj = logging.getLogger(logger_name)
	for h in log_obj.handlers:

References

1. PEP 8: Never use the characters 'l' (lowercase letter el), 'O' (uppercase letter oh), or 'I' (uppercase letter eye) as single character variable names. ^(link)

[gemini-code-assist]

Relying on inspect.getsource to dynamically retrieve, regex-replace, and exec library code at runtime is highly fragile. If the source code is not available on disk (e.g., zipped/compiled environments), it will raise an OSError. Furthermore, any minor upstream change to _fix_restore_args_for_shape_mismatch will break the regex replacement and crash the validation pipeline. Consider defining a complete custom version of _fix_restore_args_for_shape_mismatch or refactoring the monkeypatch to avoid string manipulation of the library's source code.

[gemini-code-assist]

Listing all blobs in the bucket without any limit or prefix can cause severe performance degradation and high memory usage as the bucket grows over time. Consider using max_results or page through the blobs to limit the search space.

Suggested change

	blobs = list(bucket.list_blobs())
	blobs = list(bucket.list_blobs(max_results=1000))