Running AI Workloads on Kubernetes: GPUs, Scheduling, Scaling, and Model Serving

In Part 1 of AI Infrastructure for Cloud Engineers, we explored why Kubernetes is emerging as a crucial foundation for production AI systems. The next question to consider is: what changes when AI workloads are run on Kubernetes? While traditional applications typically schedule around familiar resources like CPU, memory, storage, and network capacity, AI workloads introduce a new factor: GPUs. GPUs are powerful, expensive, and limited resources.

When teams begin running AI workloads at scale, such as model inference, embedding services, fine-tuning jobs, or other AI tasks, merely deploying a container is insufficient. The platform must also determine which workload receives which GPU, how workloads should be scheduled, when inference capacity should scale, and how to prevent idle expensive accelerators. Additionally, models need to be served reliably in production.

Kubernetes scheduling, autoscaling, and model-serving patterns become increasingly important in this context. The following are some key aspects to understand:

1. How Kubernetes views GPUs: Kubernetes traditionally schedules Pods based on resources such as CPU and memory. In a cluster with GPU-enabled nodes, a workload can request a GPU using the nvidia.com/gpu resource limit. Hardware vendors typically expose devices like GPUs to Kubernetes through device plugins.

2. The diversity of GPUs: Not all GPUs are the same. A cluster might include various accelerator types, such as NVIDIA T4, A100, and H100. Different workloads require different GPU types, memory capacities, and sizes. Efficient placement of workloads on appropriate GPUs is essential for optimal performance.

3. Workload-aware scheduling: Traditional Kubernetes scheduling focuses on individual Pods. However, AI workloads might require multiple Pods to run together in a coordinated manner, such as distributed training. This necessitates workload-aware and topology-aware scheduling, where Kubernetes considers the complete workload's requirements and topology for optimal placement.

4. GPU utilization: GPUs represent a significant portion of infrastructure costs in self-hosted AI environments. Low utilization can be expensive. For instance, if a GPU capacity is only 30% utilized due to oversized models, inconsistent traffic, inefficient requests, or insufficient batching, the remaining capacity is still being paid for. AI platforms aim to keep models responsive while using expensive compute efficiently, employing GPU-sharing and accelerator-aware scheduling approaches.

5. Inference vs. training: AI infrastructure discussions often combine training and inference, but they have different operational characteristics. Training typically involves dataset processing, a training job, multiple GPUs, hours or days, and the resulting model. Inference, on the other hand, involves user requests, model servers, GPUs, and generated responses. Inference is sensitive to latency, throughput, availability, queue depth, and concurrent requests, requiring a different scaling approach compared to training.

6. CPU scaling is not always sufficient: For traditional web applications, Kubernetes autoscaling often uses CPU utilization. However, AI inference may require different scaling signals. For example, if an inference server has low CPU utilization but high GPU utilization and a queue of waiting requests, CPU alone may not indicate the need for additional pods. Inference scaling should consider GPU utilization and other relevant metrics to ensure the model can handle concurrent requests effectively.