Running 40 Training Jobs a Day on Kubernetes: Patterns That Stuck

When we first moved our ML training workloads to Kubernetes, we envisioned seamless scalability. The reality, as it often is, was a grind. We faced critical challenges with GPU resource management and job orchestration that nearly derailed us by week three. This is the story of how my team and I tackled these issues to reliably run over 40 distinct machine learning training jobs per day on our Kubernetes clusters.

The Early Days: Raw Pods and GPU Fragmentation

Our initial approach was simple: each training job was a Kubernetes Pod, requesting a GPU. We quickly learned that even with only 5 NVIDIA A100 GPUs per Kubernetes node, packing more than 2 jobs per node was tricky due to memory and thermal issues. We routinely saw GPU fragmentation where a node might have 4 GPUs free, but each was requested by a job needing a full A100's memory capacity. At our peak, prior to implementing a scheduler, we had over 30 jobs waiting in pending state during peak hours.

Kueue vs. Volcano: The Scheduler Showdown

To combat fragmentation and manage fair sharing, we evaluated Kubernetes batch schedulers. Our initial testing with Volcano showed promise, particularly its bin-packing and gang scheduling features. However, its installation and integration proved cumbersome with our existing Argo Workflows setup. After 6 weeks of struggling, we pivoted to Kueue. Kueue's integration with the Kubernetes API, using ResourceFlavor and ClusterQueue concepts, was a much better fit for our team. We configured 3 ClusterQueues, each with a total capacity of 15 A100-80GB GPUs, reflecting our three main ML product teams.

The Argo Workflows Pattern That Stuck

While the scheduler handled resource allocation, we still needed an orchestration layer. Argo Workflows became our backbone. We evolved a pattern where each training job was encapsulated in a single Argo Workflow template. The central piece was a TrainingJob custom resource definition (CRD) that our platform team built, acting as an abstraction layer for the ML engineers.

apiVersion: ml.example.com/v1
kind: TrainingJob
metadata:
  generateName: sentiment-model-v2-
spec:
  gitRepo: "git@github.com:myorg/ml-models.git"
  trainingScript: "src/train.py"
  hardware:
    gpuType: "a100-80gb"
    gpuCount: 1
    memoryGb: 128
    cpuCores: 8
  datasetVersion: "v1.2.0"
  hyperparameters:
    learning_rate: 0.0001
    batch_size: 64

Dynamic GPU Slicing with MIG and Kueue

Even with Kueue, simply requesting a full A100 GPU still led to underutilization for smaller tasks. NVIDIA's Multi-Instance GPU (MIG) capability on our A100s was the answer. We configured our nodes to expose MIG-sliced GPUs (e.g., 1g.5gb, 2g.10gb). This meant an ML engineer could request a gpuType: "a100-1g.5gb" for a lightweight experiment, sharing a physical A100 with 6 other similar jobs. This alone boosted our overall GPU utilization by nearly 45%, moving from ~35% utilization to over 80% on average during peak hours.

Observability and Cost Management

With 40+ jobs per day, monitoring became crucial. We integrated Prometheus and Grafana, collecting metrics from Kueue (e.g., queue length, workload status) and Kubernetes. We implemented custom exporters to get GPU utilization statistics from dcgm-exporter. This data was invaluable for showing our leadership the efficiency gains. We also developed a simple billing-per-GPU-hour system, tagging each Kubernetes Pod with team and project labels, which fed into our FinOps dashboard.

What We Removed and What I'd Do Differently

Initially, we tried to build our own custom scheduler for GPU time-sharing, which was quickly abandoned after 4 weeks. It was an over-engineered solution trying to solve a problem that Kueue already had robust solutions for. We also started with a very complex GitOps setup for our ML models. We simplified this to only rebuild images on base dependency changes, reducing our image build times and registry storage costs by 70GB per week.

If I were to approach this again, I'd push for MIG from day one. The initial overhead of configuring MIG and educating teams on its use cases would have paid dividends much earlier.

Honest Limitations

Despite our successes, challenges remain. Pre-emption for low-priority jobs is still a work-in-progress. While Kueue supports pre-emption, integrating it seamlessly with Argo Workflows and ensuring proper checkpointing for ML models is a complex problem we haven't fully solved. Currently, if a high-priority job needs to run, we still have to manually drain nodes or terminate lower-priority jobs.