Polyaxon v3 is coming →
Google Cloud TPU

Run Polyaxon workloads on Google Cloud TPUs

GKE exposes TPU chips as `google.com/tpu` and imposes topology-specific scheduling rules. Polyaxon can request that native resource, target the matching TPU pool, and keep the framework image, workload, logs, and artifacts in the same run model used for GPU jobs.

Outcome

A TPU readiness job that verifies allocation and JAX device discovery before a training or serving recipe is promoted.

Start the guideScaling documentation

Infrastructure responsibilities

Kubernetes discovers

Drivers and a vendor device integration expose healthy accelerator capacity to the scheduler.

Polyaxon places

A component or preset requests the custom resource and targets the compatible node pool and image.

The workload verifies

A smoke test proves device access before training or serving is promoted to that hardware pool.

Prerequisites

  • A GKE Standard or Autopilot environment configured for the selected TPU version and topology.
  • A compatible JAX, PyTorch/XLA, Optimum TPU, or serving image.
  • A non-TPU pool for system workloads plus quota, storage, network, and Polyaxon Agent access.

Validate the infrastructure

  1. 1

    Request a TPU slice

    Select the TPU accelerator and topology exposed by GKE, then request all chips available to the pod on that slice node.

    version: 1.1kind: componentname: tpu-readinessrun:  kind: job  environment:    nodeSelector:      cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice      cloud.google.com/gke-tpu-topology: 2x2  container:    image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:jax0.7.2-rev1    command: ["python", "-c"]    args: ["import jax; print(jax.devices())"]    resources:      requests:        google.com/tpu: "4"      limits:        google.com/tpu: "4"

    TPU labels and valid chip counts depend on GKE version, TPU generation, machine type, and topology. Replace this example with the exact current values for your pool.

  2. 2

    Submit the operation

    Run the component through the target Polyaxon project, queue, preset, and approval path.

    polyaxon run -f tpu-readiness.yaml
  3. 3

    Confirm the TPU topology

    Verify JAX sees the expected devices and that pod replicas, requested chips, node topology, storage, and logs match the selected GKE configuration.

    polyaxon ops dashboard

Production checklist

Pin a compatible driver, device integration, container runtime, framework image, and firmware combination.
Publish accelerator selection through centrally managed Polyaxon presets instead of user-specific labels.
Track allocatable devices, hardware health, utilization, memory, thermals, and workload failures.
Validate one device, one node, and multi-node communication separately before production adoption.
For multi-host TPU slices, validate the framework's distributed initialization and require coordinated scheduling of every expected pod.

Troubleshooting

The resource does not appear

Inspect the node driver, firmware, device integration, kubelet registration, node labels, and allocatable resources.

The container cannot access the device

Verify the requested custom resource, runtime class, permissions, device mounts, and image compatibility.

Performance is unexpectedly low

Check topology, NUMA placement, clocks, memory pressure, host-device transfer, communication libraries, and framework build.

Sources

Official platform, library, model, and Polyaxon references used by this guide.

Continue

Compare AI platformsExplore GPU orchestration