Create a TPU v4 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 v4 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 v4.

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:

Verify that you have a TPU-compatible version of GKE. For more information, see Validate TPU availability in GKE.
Enable the Kubernetes Engine API.
Install and initialize the Google Cloud CLI.
Set up Cluster Toolkit.
Verify that you have sufficient quota for Cloud TPU v4 in your target region.

Required roles

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

Editor (roles/editor) on your project
Kubernetes Engine Cluster Admin (roles/container.clusterAdmin) on your project
Service Account Admin (roles/iam.serviceAccountAdmin) on your project
To consume a single-project reservation: Compute Instance Admin (v1) (roles/compute.instanceAdmin.v1) on your project
To consume a shared reservation: Compute Instance Admin (v1) (roles/compute.instanceAdmin.v1) on both the owner project and any consumer project where you want to consume the reservation

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:

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.

Enable versioning on the bucket:

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

Open the examples/gke-tpu-v4/gke-tpu-v4-deployment.yaml file.
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
 spot: USE_SPOT_VMS
 reservation: RESERVATION_NAME
 authorized_cidr: AUTHORIZED_CIDR
```
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-central2.
- ZONE: the Google Cloud zone for this deployment, such as us-central2-b.
- NUM_SLICES: the number of independent Cloud TPU slices to create.
- MACHINE_TYPE: the machine type for your Cloud TPU nodes, such as ct4p-hightpu-4t.
- TPU_TOPOLOGY: the physical arrangement of the Cloud TPU chips in a slice, such as 2x2x2.
- USE_SPOT_VMS: set to true to use Spot VMs for the TPU node pool, or false to use on-demand instances or reservations.
- RESERVATION_NAME: the name of the Compute Engine reservation of Cloud TPU v4 nodes. If you set spot: true, then you can comment out this line.
- 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.
Cluster Toolkit automatically calculates the exact number of nodes that are required based on your selected tpu_topology and machine_type values.
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-v4/gke-tpu-v4-deployment.yaml \
  examples/gke-tpu-v4/gke-tpu-v4.yaml

Schedule workloads with Kueue and run a sample JAX JobSet

This blueprint installs and configures Kueue and JobSet by default to manage Cloud TPU quotas and queue workload submissions. You can use the examples/gke-tpu-v4/tpu-kueue-jax-sample.yaml file to create a Kubernetes JobSet 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.

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

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-v4.
- REGION: the Google Cloud region used for this deployment.
- PROJECT_ID: your Google Cloud project ID.
Open the examples/gke-tpu-v4/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-v4-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 2x2x2.
Submit the workload to GKE. The JobSet manifest includes the kueue.x-k8s.io/queue-name: user-queue label, which automatically routes the JobSet through Kueue:
```
kubectl apply -f examples/gke-tpu-v4/tpu-kueue-jax-sample.yaml
```
Verify that Kueue admitted and queued the workload:
```
kubectl get workloads
```
Monitor the JobSet and Pod execution: sh kubectl get jobset kubectl get pods -l jobset.sigs.k8s.io/jobset-name=tpu-v4-kueue-jax-sample
View the logs of the Pods to confirm that JAX detected the Cloud TPU devices:
```
kubectl logs -l jobset.sigs.k8s.io/jobset-name=tpu-v4-kueue-jax-sample
```
When the workload succeeds, the output from both Pods is similar to the following:
```
Global device count: 8
```

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-v4/gke-tpu-v4.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-v4 \
  --compute-type v4-8 \
  --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-v4 \
  --compute-type v4-8 \
  --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:

Track the execution status and view the workload logs by using the gcluster job logs command:
```
./gcluster job logs pathways-job-v4
```
When the workload succeeds, the end of the log output is similar to the following:
```
JAX Device count: 4
```
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-v4
```

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-v4/gke-tpu-v4-deployment.yaml file.

Create a TPU v4 GKE cluster Stay organized with collections Save and categorize content based on your preferences.