The Practical Guide to Running ML Workloads on Kubernetes
Kubernetes has become the default orchestration platform for ML workloads, but running machine learning on K8s is fundamentally different from running web applications. GPU scheduling, large model artifacts, long-running training jobs, and bursty inference traffic all require specific configurations that standard K8s tutorials don't cover. Here's what actually works.
GPU Scheduling: The Basics Most Teams Get Wrong
The NVIDIA device plugin for Kubernetes makes GPUs available as a schedulable resource. But the default behavior has significant limitations:
- GPUs are not shareable by default. If a pod requests 1 GPU, it gets exclusive access to that entire GPU, even if it only uses 20% of the compute. For inference workloads that don't need a full GPU, this is enormously wasteful.
- GPU memory isn't a schedulable resource. Kubernetes doesn't natively understand GPU memory. A pod requesting a GPU might land on a 16GB A10 when it needs a 40GB A100, or vice versa.
- Multi-GPU training requires careful topology awareness. For distributed training, you want GPUs on the same node (or at least same rack) to minimize communication overhead.
Solutions
- NVIDIA MIG (Multi-Instance GPU): Supported on A100 and H100 GPUs. Allows partitioning a single GPU into isolated instances. Configure MIG profiles in your device plugin config to offer fractional GPUs as resources.
- Time-slicing: For less critical workloads, NVIDIA's time-slicing allows multiple pods to share a GPU by time-multiplexing. Lower isolation than MIG but works on more GPU models.
- Node labels and affinity rules: Label nodes by GPU type (gpu-type: a100-40gb) and use node affinity in pod specs to ensure workloads land on appropriate hardware.
- Topology-aware scheduling: Use the Topology Manager with the single-numa-node policy for training workloads that benefit from GPU locality.
Training Jobs: Not Your Standard Deployment
ML training jobs are fundamentally different from web services:
- They run for hours or days, not indefinitely
- They're compute-intensive and bursty
- They need checkpointing to survive preemption
- They often require coordinated multi-node execution
Use Jobs, Not Deployments
This seems obvious, but we've seen teams run training workloads as Deployments with restartPolicy: Always. Use Kubernetes Jobs (or CronJobs for scheduled retraining) with appropriate completion and backoff settings.
Checkpointing
Training jobs should checkpoint model state to persistent storage (S3, GCS, or a PersistentVolume) at regular intervals. If a node goes down or a spot instance is reclaimed, the job can resume from the last checkpoint rather than starting over. Configure your training framework to save checkpoints every N epochs or N minutes.
Spot/Preemptible Instances
GPU instances are expensive. Spot instances can save 60-70% but can be reclaimed with minimal warning. Combine spot instances with checkpointing and use a mix of on-demand (for critical jobs) and spot (for fault-tolerant training) nodes. Karpenter or Cluster Autoscaler can manage this node provisioning automatically.
Model Serving: Where Performance Matters
Serving ML models in production has different requirements than training:
Serving Frameworks
- Triton Inference Server: NVIDIA's server supports multiple frameworks (TensorFlow, PyTorch, ONNX, TensorRT) with dynamic batching, model ensemble, and GPU sharing. Our default choice for most deployments.
- TorchServe: Good for pure PyTorch models with simpler requirements.
- vLLM: Purpose-built for LLM serving with PagedAttention for efficient memory management. If you're serving large language models, vLLM significantly outperforms general-purpose serving frameworks.
Autoscaling for Inference
Standard Kubernetes HPA (Horizontal Pod Autoscaler) based on CPU doesn't work well for ML inference. Instead:
- Scale on GPU utilization metrics (exposed via DCGM exporter)
- Scale on request queue depth (number of pending inference requests)
- Scale on custom latency metrics (p99 response time exceeding SLA)
- Use KEDA for event-driven autoscaling from message queues
Set scale-down stabilization to at least 5 minutes — ML model loading takes time, and you don't want to thrash.
Storage Patterns
ML workloads are data-hungry. Get storage wrong and everything bottlenecks:
- Model artifacts: Store in object storage (S3/GCS). Use init containers to pull models into local storage before the serving container starts. For large models (10GB+), consider ReadWriteMany PersistentVolumes backed by EFS/Filestore so multiple pods can share the same model files.
- Training data: For large datasets, avoid pulling data through the network every training run. Use persistent volumes or a shared filesystem. Consider a data caching layer for frequently accessed training data.
- Experiment tracking: Use MLflow, Weights & Biases, or similar tools to track experiments, metrics, and model versions. Store artifacts in object storage, not on node-local disks.
Monitoring ML-Specific Metrics
Beyond standard Kubernetes monitoring (CPU, memory, pod health), ML workloads need:
- GPU metrics: Utilization, memory usage, temperature, power draw via DCGM exporter + Prometheus
- Inference metrics: Request latency (p50, p95, p99), throughput, batch sizes, queue depth
- Model metrics: Prediction distribution drift, feature drift, accuracy degradation over time
- Cost metrics: GPU-hours per model, cost per inference, spot vs. on-demand spend
We typically deploy a Grafana dashboard with these four categories visible at a glance, with alerts on latency SLA breaches, GPU memory exhaustion, and model drift thresholds.
Our Standard Stack
For most enterprise ML deployments on Kubernetes, our default stack is: EKS or GKE with managed node groups, Karpenter for autoscaling, Triton or vLLM for serving, Argo Workflows for training pipelines, MLflow for experiment tracking, Prometheus + Grafana for monitoring, and Istio for traffic management and canary deployments. This stack handles everything from small single-model deployments to multi-model platforms serving millions of predictions per day.
Need help with your next project?
Book a free 30-minute discovery session with our senior engineers to discuss your specific challenges.