How Ai21 Manages Its Gpu Fleet With Kueue

How AI21 manages its GPU fleet with Kueue

Skip to Main Menu

Skip to Main Content

Skip to Footer

Back to Blog

-->

Back to Blog

Back to Blog

-->

Google Cloud recently published our story on AI Hypercomputer , including an 83% reduction in time-to-start for high-priority workloads. Below is the full technical breakdown of how we got there with Kueue on GKE.

In brief

Sharing a large GPU fleet across many teams is a recurring challenge for anyone running AI training workloads on Kubernetes. In this blog, we share a case study from one of our shared GKE clusters, around 10,000 GPUs used across multiple teams, on how we moved from manual coordination over internal messaging channels to Kueue, an open-source, Kubernetes-native job queueing system that admits, prioritizes, and preempts workloads automatically. What began as adopting Kueue turned into a close collaboration with the team behind it at Google, with our real-world production needs feeding directly into Kueue’s roadmap. We’ll walk through what we tried, the results – a dramatic drop in metrics like manual intervention rate and the time a critical job spends “starved” – and the lessons from running it in production that fed back into Kueue itself, including new features such as Admission Fair Sharing (AFS), and is now available to all users.

A sharing problem

It’s no secret that GPUs are a scarce resource. Anyone working on model training – which requires significant GPU usage over a long period of time – is familiar with the problem up close. Their finite nature means managing GPU distribution across teams of researchers, MLOps, engineers, and alignment becomes particularly acute during training, raising challenging questions like:

How do you fairly share a finite, expensive resource across multiple teams?

Who gets prioritized when demand exceeds capacity?

Should you interrupt a running workload to make room for a higher-priority one?

How do you reconcile teams with different priorities and different definitions of “urgent”?

The “before”: Negotiating resources in team channels

At the beginning, we were navigating these questions in an internal messaging channel, called #gpu-resources . When someone needed a GPU type that was fully utilized, they asked in the channel and hoped someone would free capacity.

Figure 1: a typical day on #gpu-resources — manual requests, ad-hoc coordination, idle capacity in some pockets while others starve.

This kind of channel-based negotiation worked while the cluster had headroom. Once utilization stayed pinned near 100% – which is the goal for an expensive reserved GPU fleet – every incoming request became a zero-sum eviction conversation, and the channel hit a wall.

Left to manual negotiation, the cracks in this system quickly appeared:

Real engineering time wasted on triage and negotiation

Unfair queueing

Idle capacity in some pockets while others starved

Reliance on ad hoc “magic commands” to free up sufficient workload capacity

We needed to build something better. We started by devising the design principles that should inform an ideal GPU management system.

The four pillars of an ideal GPU management system

Fairness: Scheduling decisions are made systematically and in accordance with account usage, rather than operating on a first-come-first-served or loudest-voice-wins basis.

Order: Workload queuing reflects business priorities, not just timestamp.

Efficiency: Developers are relieved from manually deleting their idle or non-critical workloads.

Technical integrity: Multi-pod (gang) jobs don’t start until all resources are ready, preventing partial deployments.

From design principles to implementation: Building a GPU management system

Our building blocks

For this blog, we’ll zoom in on one of our shared GKE clusters on Google Cloud (GCP) – the one where we rolled this system out. It pools on the order of 10,000 GPUs across multiple teams.

We run it as a single shared pool rather than a cluster per team on purpose: Kubernetes makes it easy to spin up many clusters, but hard to share capacity between them. Keeping everything in one pool is what keeps utilization high – and it’s exactly what puts the pressure on how we schedule inside it.

Here’s the setup on that cluster:

We have three types of machine capacity – reserved, spot, and Dynamic Workload Scheduler’s flex-start mode – that we need to run three different types of training R&D workloads:

Debug pods: Developers may want to run debugging remotely on a real production environment rather than locally on their laptops

Training jobs: Batch workloads which can scale from 1→100s of workloads within a second. These can appear as single-node jobs or as multi-node jobs, which require multiple machines to run in parallel. We use Indexed Jobs to manage parallelism.

Inference models: Used by our Reinforcement Learning training jobs and managed as Kubernetes deployments, scaled up or down using HPA .

Each of these workloads can be further classified by their ability to be preempted:

Non-preemptible: Interrupting them mid-run is an automatic red light, as it is expensive, wasteful, and potentially damaging to time-critical workloads.

Preemptible: Interrupting them mid-run has minor to negligible consequences.

Searching for a fix: Evaluating vanilla K8s

Before evaluating third-party tools, we checked what Kubernetes natively ships with. Here’s what we saw:

Primitive What it offers Where it falls short PriorityClass Critical pods can preempt non-critical pods. Efficiency and Order Operates at the pod level, not the workload level, with no “gang size” awareness, so entire lower priority runs can be preempted even when just 1 out of several GPUs is needed for a higher priority job. ResourceQuota Hard cap on CPU/memory/GPU/object count per namespace. Fairness No borrowing across teams when capacity is idle; over-quota pods are rejected at admission, not queued. LimitRange Per-container default/min/max requests inside a namespace. Fairness Acts as a guardrail, not a scheduler. Job / Indexed Job Workload abstraction with completions/parallelism. Technical integrity No all-or-nothing mentality. If 3 of 4 pods schedule and the 4th can’t fit, the first 3 still run and burn GPUs.

We needed to turn to other batch scheduler options to find one that would meet all four of our criteria. As...