GCP ACE: GKE and AI — Training, Serving, and Cost Management for ML Workloads
This closes the “Manage Scalable Workloads in GKE” module. Everything in this series so far scaled general-purpose web workloads. AI/ML training and serving workloads have different needs — specialized hardware that’s often scarce, jobs that need all-or-nothing resource allocation, and cost profiles where the difference between the right and wrong scheduling choice is enormous.
Why AI/ML workloads don’t just fit the patterns already covered
A regular HPA (from the autoscaling article) reacts to CPU/memory metrics and adds or removes interchangeable replicas one at a time. A distributed training job doesn’t work that way — it typically needs all of its worker Pods scheduled together before training can start at all (there’s no useful “half the workers” state), and it needs specific hardware (GPUs, TPUs) that’s frequently in short supply across a region. This is the gap GKE’s AI/ML-specific tooling closes.
GPU and TPU support
GKE supports both accelerator types for both training and inference:
- GPUs — general-purpose accelerators; GKE includes a faster-startup node provisioning path for GPU workloads, cutting node boot time by up to 80% compared to standard node startup, which matters when you’re trying to scale up capacity for a burst of inference traffic quickly.
- TPUs — Google’s custom accelerators, particularly suited to large-scale training. TPU multislice links several TPU slices together over the data center network so they can train one large model cooperatively, rather than being limited to what fits on a single slice.
Kueue — job queueing for batch and training workloads
Kueue is described as “a Kubernetes-native job queueing system.” It sits above the Job/CronJob primitives (from the earlier workloads article in this series) and adds what a raw Job on its own doesn’t have:
- Resource allocation and quota enforcement — Kueue decides which queued jobs get to actually run based on available quota, rather than every submitted Job competing for the scheduler’s attention immediately and unpredictably.
- Priority — higher-priority training or batch jobs can be scheduled ahead of lower-priority ones, rather than strict first-in-first-out.
- All-or-nothing scheduling for a job’s Pods — this is the piece a plain Job/Deployment doesn’t give you, and it’s exactly what distributed training needs: a job’s workers either all get scheduled together, or the job waits, rather than partially starting and stalling.
Dynamic Workload Scheduler (DWS)
GPUs and TPUs are frequently capacity-constrained — you can request one and simply not get it immediately when a zone runs short. Dynamic Workload Scheduler, exposed as the flex-start provisioning mode, exists specifically to “improve your ability to secure scarce GPU and TPU accelerators” for jobs that are short-duration and can tolerate waiting briefly for capacity to become available, in exchange for a meaningfully better chance of actually getting that hardware (and often a better price) than trying to provision it on-demand, immediately, every time.
The practical pairing worth remembering: Kueue decides what should run and in what order; DWS decides how to actually go get the scarce hardware that decision needs. They’re commonly used together — Kueue manages the queue and priority, DWS handles the underlying capacity acquisition for GPU/TPU-hungry jobs in that queue.
Custom Compute Classes — defining fallback hardware preferences
A custom ComputeClass is a Kubernetes custom resource defining a profile of node attributes GKE should use when autoscaling to meet a workload’s needs — with an explicit, ordered fallback list.
apiVersion: cloud.google.com/v1
kind: ComputeClass
metadata:
name: fallback-class
spec:
nodePoolAutoCreation:
enabled: true
priorities:
- machineFamily: n4
spot: true
minCores: 16
- machineFamily: n4
spot: false
minCores: 16
- machineFamily: c3
spot: false
minCores: 16
whenUnsatisfiable: ScaleUpAnyway
GKE tries each entry in priorities, in order: Spot N4 VMs first (cheapest), falling back to on-demand N4, then on-demand C3 as a last resort. whenUnsatisfiable: ScaleUpAnyway means if none of the preferred options can be satisfied, GKE still scales up with whatever it can get, rather than leaving the workload unscheduled (the alternative, DoNotScaleUp, would block scaling instead). DoNotScaleUp is the one to pick when the hardware is non-negotiable, for example a training job that must land on a specific TPU or GPU type.
A Pod opts into a ComputeClass via nodeSelector:
spec:
nodeSelector:
cloud.google.com/compute-class: fallback-class
One detail worth remembering for the exam: ComputeClasses influence autoscaling decisions — what kind of new node GKE provisions — but they are not considered by kube-scheduler itself when placing Pods onto already-existing nodes. The fallback logic is about what gets provisioned, not about where an existing Pod lands.
Cost management, pulling it together
For a GPU/TPU-heavy training workload, the realistic cost-management stack looks like this:
- Kueue queues the job and enforces quota/priority, so you’re not paying for idle capacity waiting on jobs that shouldn’t run yet.
- DWS (flex-start) improves the odds of actually securing scarce accelerator capacity, often at better pricing than a rigid on-demand request.
- A custom ComputeClass defines a cost-aware fallback order — try Spot first (cheapest), then on-demand of the same family, then a different family entirely — rather than defaulting straight to the most expensive guaranteed option.
- Know which Autopilot billing model applies. Autopilot’s Pod-based billing (from the earlier GKE architecture article) covers general-purpose Pods only. Pods that select specific hardware, such as GPUs, TPUs, or a machine series through a ComputeClass, use node-based billing: Compute Engine pricing for the node plus an Autopilot management premium. For accelerator-heavy training, that makes bin-packing and scale-down (and Spot fallbacks) the real cost levers.
Exam Takeaways
- Distributed training needs all-or-nothing Pod scheduling — a capability plain Jobs/Deployments don’t provide on their own, which is exactly what Kueue adds.
- Kueue handles queueing, quota, and priority for batch/training jobs; it’s a layer above the Job/CronJob primitives, not a replacement for them.
- Dynamic Workload Scheduler (flex-start mode) improves access to scarce GPU/TPU capacity for jobs that can tolerate a short wait — pair it with Kueue for the full queue-then-provision pipeline.
- TPU multislice links multiple TPU slices together over the network to train one model cooperatively, beyond what a single slice could handle alone.
- A custom ComputeClass defines an ordered fallback list of node attributes (machine family, Spot vs. on-demand, core count) for autoscaling — tried top to bottom.
whenUnsatisfiable: ScaleUpAnywayvs.DoNotScaleUpdecides whether GKE scales up with a worse option or simply doesn’t scale when nothing in the priority list can be satisfied.- ComputeClasses affect what gets provisioned during autoscaling — they don’t affect kube-scheduler’s placement decisions on already-existing nodes.
- On Autopilot, accelerator Pods (and any Pod selecting specific hardware) are billed per node, not per Pod request.