Skip to main content
    Scalexa — Senior Engineering & AI Solutions
    Engineering

    The Practical Guide to Running ML Workloads on Kubernetes

    Aurelien DuarteMarch 14, 202612 min read

    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.

    Get Started

    Ready to Get Started?

    Book a free 30-minute discovery session with our senior engineers to identify quick wins and show you what's possible.

    View Our Work