Add GKE TPU v5e example#5840
Conversation
agrawalkhushi18
added
the
release-key-new-features
Summary of ChangesHello, I'm Gemini Code Assist1! 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 introduces a new blueprint example for deploying GKE TPU v5e clusters using the Cluster Toolkit. It provides comprehensive configuration files, including infrastructure deployment manifests, Kueue integration for job scheduling, and sample workloads to verify TPU connectivity. Additionally, the PR updates the Google provider version constraints to ensure compatibility across the repository. Highlights
Using Gemini Code AssistThe 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
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 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
|
![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 a new GKE TPU v5e blueprint example, including deployment configurations, a Kueue setup, and sample workloads, alongside a minor downgrade of the maximum Google provider version constraint in pkg/config/expand.go and its golden copies. Feedback on the changes includes registering the new example in the main examples/README.md index, providing a default value for tpu_topology to match the default machine type, and declaring the spot variable in the blueprint to prevent validation errors when configuring spot instances.
|
|
||
| 1. Clone the Cluster Toolkit from the git repository: | ||
|
|
||
| ```sh |
There was a problem hiding this comment.
Shall we update this to binary installation and not git clone.
|
|
||
| ## Running Pathways Workloads | ||
|
|
||
| This blueprint supports **Pathways-on-Cloud** orchestration, allowing you to run JAX workloads distributed across remote TPU workers coordinated by a CPU-based head node. The sections below guide you through deploying and verifying a sample job that prints the active TPU device count over the coordinated Pathways network. |
There was a problem hiding this comment.
Shall we also hyperlink to the pathways GitHub repo?
| * `region`: the compute region for the cluster. | ||
| * `zone`: the compute zone for the TPUs. | ||
| * `num_slices`: the number of TPU slices to create. | ||
| * `machine_type`: the machine type of the TPU (e.g., `ct5lp-hightpu-8t`). |
There was a problem hiding this comment.
Shall we update the example in ct51p-hightpu-8t to the machine types relevant to tpu v5e
4 participants
This PR adds a new GKE blueprint and workload templates optimized specifically for provisioning and running jobs on Cloud TPU v5e clusters.
Key Changes
examples/gke-tpu-v5e/):gke-tpu-v5e.yaml: Adds a blueprint that sets up a GKE cluster with a single VPC network and pre-configures Kueue, JobSet, and Cert-Manager.gke-tpu-v5e-deployment.yaml: Adds the deployment settings file, pre-configured with recommended TPU v5e variables and GCS backend configurations.kueue-job-sample.yaml: An example Kueue-enabled JobSet workload that demonstrates how to submit jobs to the Kueue queue and request/limit TPU resources (google.com/tpu) properly on the v5e cluster.kueue-configuration.yaml.tftpl: Added a Kueue configuration template that defines resource flavors, cluster queues, and local queues dynamically mapped to the TPU v5e accelerator type (tpu-v5-lite-podslice) and total chip quota.tpu-singleslice.yaml: A standard, lightweight KubernetesJobconfigured for single-slice (2x2) JAX workloads.kueue-job-sample.yaml: A sample workload pre-configured for Kueue job queue integration.README.md: A guide covering deployment steps and teardown.Verification
tpu-singleslice.yamlto verify successful TPU hardware initialization and JAX device detection (Global device count: 4).Submission Checklist
NOTE: Community submissions can take up to 2 weeks to be reviewed.
Please take the following actions before submitting this pull request.