Autoscaling using TPU instances

The Cast AI Autoscaler can scale GKE clusters using TPU-optimized instances. This guide describes the steps needed to configure your cluster so that TPU nodes can join it and workloads that request TPU resources are scheduled correctly.

Supported providers

📘

Note

TPU instances are only supported on GCP GKE. AWS EKS and Azure AKS are not supported.

How does it work?

Once a pending pod requests google.com/tpu resources, Cast AI's Autoscaler detects the workload and provisions a TPU node that satisfies the request.

Cast AI supports single-host TPU slices only. Multi-host (pod slice) topologies are not supported.

To enable provisioning of TPU nodes, you need:

• A pod that requests google.com/tpu resources
• A toleration for the google.com/tpu=true:NoSchedule taint
• Node selectors that specify the TPU version and topology

Cast AI automatically selects the correct TPU instance type based on the node selectors in your pod specification — no device plugins or additional configuration are required. GKE handles TPU driver installation automatically on TPU nodes.

TPU node labels

Cast AI applies the following labels to TPU nodes, which you can use in node selectors or affinities:

Supported TPU versions and accelerator labels

Supported single-host topologies

TPU v5e (2D topologies)

TPU v5p (3D topologies)

When multiple topologies are valid for a given chip count, Cast AI prefers the most balanced (closest to cube/square) topology.

Tainting of TPU nodes

TPU nodes will have the google.com/tpu=true:NoSchedule taint applied automatically. Your TPU workloads must include a matching toleration; without it, the pod will not be scheduled onto a TPU node, and Cast AI will not provision one.

Workload configuration examples

The following examples show common patterns for configuring TPU workloads.

TPU v5e — 4 chips (2x2 topology)

spec:
  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
    cloud.google.com/gke-tpu-topology: 2x2
  tolerations:
    - key: "google.com/tpu"
      operator: Exists
      effect: NoSchedule
  containers:
    - image: my-tpu-image
      name: tpu-workload
      resources:
        requests:
          cpu: 4
          memory: 8Gi
          google.com/tpu: 4
        limits:
          cpu: 4
          memory: 8Gi
          google.com/tpu: 4

TPU v5p — 4 chips (2x2x1 topology)

spec:
  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v5p-slice
    cloud.google.com/gke-tpu-topology: 2x2x1
  tolerations:
    - key: "google.com/tpu"
      operator: Exists
      effect: NoSchedule
  containers:
    - image: my-tpu-image
      name: tpu-workload
      resources:
        requests:
          cpu: 4
          memory: 8Gi
          google.com/tpu: 4
        limits:
          cpu: 4
          memory: 8Gi
          google.com/tpu: 4

TPU v5e — 8 chips (2x4 topology)

spec:
  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
    cloud.google.com/gke-tpu-topology: 2x4
  tolerations:
    - key: "google.com/tpu"
      operator: Exists
      effect: NoSchedule
  containers:
    - image: my-tpu-image
      name: tpu-workload
      resources:
        requests:
          cpu: 8
          memory: 16Gi
          google.com/tpu: 8
        limits:
          cpu: 8
          memory: 16Gi
          google.com/tpu: 8

Required pod fields

The following fields in your pod specification are required for Cast AI to correctly provision a TPU node:

📘

Note

The google.com/tpu resource must be set in both requests and limits, and both values must be equal.

No device plugin required

Unlike GPU instances, TPU nodes on GKE do not require an additional device plugin. GKE manages TPU drivers and the google.com/tpu extended resource automatically on TPU nodes. Cast AI does not validate for any TPU device plugin before autoscaling.

Further reading

• GKE TPU documentation
• Plan your TPU configuration on GKE