Create a TPU v5e GKE cluster

This document shows you how to use a Cluster Toolkit blueprint to automate the deployment of a Google Kubernetes Engine (GKE) cluster that has a dedicated Cloud TPU v5e node pool. This blueprint lets you rapidly provision repeatable, scalable infrastructure optimized to train and serve large-scale AI models.

For more information about the architecture of this TPU, see TPU v5e.

For more information about TPUs in GKE, see How TPUs in GKE work.

Before you begin

Before you begin, verify that you have completed the following tasks:

Required roles

To get the permissions that you need to deploy the GKE Cloud TPU v5e cluster, ask your administrator to grant you the following IAM roles:

For more information about granting roles, see Manage access to projects, folders, and organizations.

You might also be able to get the required permissions through custom roles or other predefined roles.

Set up the cluster infrastructure

To set up the cluster infrastructure that is required for the blueprint deployment, do the following:

  1. Create a Cloud Storage bucket to store the state of the Terraform deployment:

    gcloud storage buckets create gs://BUCKET_NAME \
  --default-storage-class=STANDARD \
  --location=COMPUTE_REGION \
  --uniform-bucket-level-access

    Replace the following:

    • BUCKET_NAME: the name of the new Cloud Storage bucket.
    • COMPUTE_REGION: the compute region where you want to store the Terraform state.

  2. Enable versioning on the bucket:

    gcloud storage buckets update gs://BUCKET_NAME \
  --versioning

  3. Open the examples/gke-tpu-v5e/gke-tpu-v5e-deployment.yaml file.

  4. In the terraform_backend_defaults and vars sections, replace the placeholders to match your deployment:

    terraform_backend_defaults:
 type: gcs
 configuration:
   bucket: BUCKET_NAME

vars:
 project_id: PROJECT_ID
 deployment_name: DEPLOYMENT_NAME
 region: REGION
 zone: ZONE
 num_slices: NUM_SLICES
 machine_type: MACHINE_TYPE
 tpu_topology: TPU_TOPOLOGY
 authorized_cidr: AUTHORIZED_CIDR
 reservation: RESERVATION_NAME

    Replace the following:

    • BUCKET_NAME: the Cloud Storage bucket that stores the Terraform state.
    • PROJECT_ID: your Google Cloud project ID.
    • DEPLOYMENT_NAME: the name of your deployment.
    • REGION: the Google Cloud region for this deployment, such as us-east5.
    • ZONE: the Google Cloud zone for this deployment, such as us-east5-a.
    • NUM_SLICES: the number of independent Cloud TPU slices to create.
    • MACHINE_TYPE: the machine type for your Cloud TPU nodes, such as ct5lp-hightpu-4t.
    • TPU_TOPOLOGY: the physical arrangement of the Cloud TPU chips in a slice, such as 2x2 or 2x4.
    • AUTHORIZED_CIDR: the CIDR block that contains the IP address of the machine that calls Terraform. To allow all IP addresses, use 0.0.0.0/0. Identity and Access Management (IAM) restrictions are still enforced. To allow only your IP address, use your IP address followed by /32.
    • RESERVATION_NAME: the name of the Compute Engine reservation of Cloud TPU v5e nodes.

    Cluster Toolkit automatically calculates the exact number of nodes that are required based on your selected tpu_topology and machine_type values.

  5. Generate Application Default Credentials (ADC) to provide access to Terraform:

    gcloud auth application-default login

Deploy the blueprint

After you have set up the cluster infrastructure, deploy the blueprint to provision the GKE infrastructure by using the gcluster deploy command:

cd ~/cluster-toolkit
./gcluster deploy -d \
  examples/gke-tpu-v5e/gke-tpu-v5e-deployment.yaml \
  examples/gke-tpu-v5e/gke-tpu-v5e.yaml

Schedule workloads with Kueue and run a sample JAX Job

This blueprint installs and configures Kueue and JobSet by default to manage Cloud TPU quotas and queue Job submissions. You can use the examples/gke-tpu-v5e/tpu-kueue-jax-sample.yaml file to create a Kubernetes Job that integrates routing through Kueue queues and validates the Cloud TPU device count by using JAX.

By default, Cluster Toolkit dynamically applies an embedded Kueue configuration based on your Pathways and dynamic slicing settings:

  • If you explicitly disable both Pathways and dynamic slicing, then Cluster Toolkit still installs the Kueue engine and controllers, but leaves them unconfigured without creating default queues or resource flavors.
  • If you want to override the default embedded configurations or apply a custom configuration when both settings are disabled, then uncomment the config_path field in the kueue section of the workload-manager-install module in your blueprint. For more information about default configurations and variables, see the kubectl-apply module and the kubectl-apply module inputs on GitHub.

Understand Kueue quota

The blueprint automatically calculates and sets a google.com/tpu quota in the ClusterQueue resource in Kueue. The blueprint derives the node count from your machine_type and tpu_topology settings, and calculates the quota by using the num_slices * (total_chips_in_topology / chips_per_machine) * chips_per_machine formula.

Submit and verify the sample JAX Job

To submit and verify the sample JAX Job through Kueue, do the following:

  1. Connect to your cluster:

    gcloud container clusters get-credentials DEPLOYMENT_NAME \
  --region=REGION \
  --project=PROJECT_ID

    Replace the following:

    • DEPLOYMENT_NAME: the name of your deployment, such as gke-tpu-v5e.
    • REGION: the Google Cloud region used for this deployment.
    • PROJECT_ID: your Google Cloud project ID.

  2. Open the examples/gke-tpu-v5e/tpu-kueue-jax-sample.yaml file and verify that the nodeSelector values match the Cloud TPU accelerator and topology that you configured in your deployment:

    nodeSelector:
  cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
  cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY

    Replace TPU_TOPOLOGY with the topology that you configured for your Cloud TPU slice, such as 2x2.

  3. Submit the workload to GKE. The Job manifest includes the kueue.x-k8s.io/queue-name: user-queue label, which automatically routes the Job through Kueue:

    kubectl apply -f examples/gke-tpu-v5e/tpu-kueue-jax-sample.yaml

  4. Verify that Kueue admitted and queued the workload:

    kubectl get workloads

    Monitor the Job and Pod execution: sh kubectl get jobs

  5. View the logs of the Pod to confirm that JAX detected the Cloud TPU devices:

    kubectl logs -l job-name=tpu-kueue-jax-sample -c jax-tpu

    The output is similar to the following:

    Global device count: 4

    The global device count depends on your topology. For example, there are four chips for a 2x2 topology.

Run Pathways workloads

This blueprint supports Pathways on Cloud orchestration, which lets you run JAX workloads that are distributed across remote Cloud TPU workers and coordinated by a CPU-based head node. The following sections show you how to deploy and verify a sample workload that prints the active Cloud TPU device count across the coordinated Pathways network.

For more information about submitting workloads, see the gcluster job submit reference documentation and the gcluster job submission guide on GitHub.

Enable Pathways in the blueprint

Before you deploy the cluster, verify that Pathways is enabled in the examples/gke-tpu-v5e/gke-tpu-v5e.yaml file:

vars:
  enable_pathways_for_tpus: true

When you enable Pathways, the blueprint provisions a cpu-np node pool that uses the n4-standard-64 machine type by default to host the coordinator head Pod. GKE autoscaling dynamically manages the size of this node pool between 0 and 100 nodes based on the workload.

Preview the JobSet manifest

Before you submit the workload to your live GKE cluster, we recommend that you generate and inspect the Kubernetes JobSet manifest without submitting the workload. To generate the manifest, use the gcluster job submit command:

./gcluster job submit \
  --name pathways-job \
  --compute-type v5litepod-4 \
  --pathways \
  --pathways-gcs-location gs://COORDINATION_BUCKET/pathways-scratch \
  --image us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest \
  --command "pip install pathwaysutils && python -c 'import pathwaysutils; pathwaysutils.initialize(); import jax; print(\"JAX Device count:\", jax.device_count())'" \
  --dry-run-out my-manifest.yaml

Replace COORDINATION_BUCKET with the name of the Cloud Storage bucket that the coordinator and worker Pods use to synchronize network state and parameters.

This command generates a local my-manifest.yaml file that contains the complete Kubernetes JobSet manifest. Inspect the my-manifest.yaml file and verify the following:

  • The coordinator Pod is assigned to the CPU node pool by using the cloud.google.com/gke-nodepool: cpu-np node selector.
  • The worker Pods target the Cloud TPU node pool.
  • The pathways-proxy and pathways-rm sidecar containers in the coordinator Pod and the pathways-worker container in the worker Pods are configured with the required ports and environment variables.

Submit the workload

After you verify the JobSet manifest, submit the workload to your live GKE cluster without the --dry-run-out flag:

./gcluster job submit \
  --name pathways-job \
  --compute-type v5litepod-4 \
  --pathways \
  --pathways-gcs-location gs://COORDINATION_BUCKET/pathways-scratch \
  --image us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest \
  --command "pip install pathwaysutils && python -c 'import pathwaysutils; pathwaysutils.initialize(); import jax; print(\"JAX Device count:\", jax.device_count())'"

Monitor and manage the workload

To monitor the workload logs and clean up resources, do the following:

  1. Track the execution status and view the workload logs by using the gcluster job logs command:

    ./gcluster job logs pathways-job

    When the workload succeeds, the end of the log output is similar to the following:

    JAX Device count: 4

  2. Optional: To cancel a running workload or clean up resources after the workload finishes, use the gcluster job cancel command:

    ./gcluster job cancel pathways-job

Delete resources

To avoid recurring charges for the resources used on this page, delete the resources provisioned by Cluster Toolkit, including the Virtual Private Cloud (VPC) networks and GKE cluster, by using the gcluster destroy command:

./gcluster destroy DEPLOYMENT_NAME/

Replace DEPLOYMENT_NAME with the name of your deployment. You can find this name in the Set up the cluster infrastructure section, where you defined the deployment_name variable in the examples/gke-tpu-v5e/gke-tpu-v5e-deployment.yaml file.

What's next