feat: integrate GKE-managed ML Diagnostics

[AdarshK15]

Summary

This PR integrates GKE-managed Machine Learning Diagnostics support to the gke-cluster module. It also refactors the mldiagnostics module (which installs ML Diagnostics helm charts) to support older GKE versions (< 1.35.0-gke.3065000). Validations are added to ensure only one installation method is used at a time to prevent duplicate or conflicting installations.

Changes Made

• Added the boolean variable enable_ml_diagnostics to the gke-cluster module to support GKE's managed_machine_learning_diagnostics_config for GKE version >= 1.35.0-gke.3065000.
• Moved workload namespace labelling (managed-mldiagnostics-gke: "true") to the gke-cluster module after namespace creation.
• Added preflight metadata validation to ensure Workload Identity is enabled when enable_ml_diagnostics is set to true.
• Added a validation in mldiagnostics module to ensure only one ML Diagnostics installation method is used in the blueprint.
• Updated both gke-tpu-7x and gke-tpu-v6e blueprints to use GKE managed ML Diagnostics.

Documentation

• updated blueprint documentation:
  • GKE TPU v6e README
  • GKE TPU 7x README

Usage Example

To enable GKE managed ML Diagnostics, set enable_managed_ml_diagnostics to true in gke-cluster module:

  - id: gke-tpu-v6e-cluster
    source: modules/scheduler/gke-cluster
    settings:
      configure_workload_identity_sa: true
      namespace: $(vars.user_namespace)
      enable_managed_ml_diagnostics: true

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request integrates native GKE-managed Machine Learning Diagnostics support into the gke-cluster module. It refactors the existing mldiagnostics module to maintain backward compatibility for older GKE versions while introducing strict validation to prevent conflicting installation methods. The update also migrates namespace labeling to the gke-cluster module and updates the gke-tpu-7x and gke-tpu-v6e blueprints to adopt the managed diagnostics approach, accompanied by refreshed documentation.

Highlights

• GKE-managed ML Diagnostics Integration: Added the enable_ml_diagnostics boolean variable to the gke-cluster module to support native GKE-managed ML Diagnostics for supported versions.
• Mutual Exclusion Validation: Implemented a precondition in the mldiagnostics module to prevent concurrent use of managed and manual installation methods, ensuring configuration consistency.
• Blueprint Updates: Updated gke-tpu-7x and gke-tpu-v6e blueprints to leverage the new managed diagnostics feature and removed the manual mldiagnostics module dependency.
• Namespace Management: Refactored workload namespace labeling logic, moving it directly into the gke-cluster module to streamline cluster configuration.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request transitions Google Cloud ML Diagnostics from a manual Helm-based installation to a native GKE-managed feature in the gke-cluster module for GKE versions 1.35.0-gke.3065000 and higher. The manual module remains as an optional fallback for older GKE versions, with validation added to prevent both methods from being active simultaneously. Feedback on the changes highlights several key improvement opportunities: preventing silent failures on unsupported GKE versions by using strict validation to fail fast, adding robust validation checks in the gke-cluster module to ensure GKE version and Workload Identity requirements are met, correcting a variable name mismatch in the metadata validation error message, and removing the unused helm provider and its requirement from the gke-cluster module.

[LAVEEN]

LGTM

[SwarnaBharathiMantena]

LGTM! Please ensure at least one of the TPU 7x or TPU v6e tests are running successfully.

[AdarshK15]

Test failures

• Test PR-test-gke-tpu-v6e failed due to a known timeout issue at Run CHS suites task, ML Diagnostics related tasks completed successfully.
• Other GKE tests passed.