How to Run ML Training Jobs on Kubernetes With GPU Node Pools: A DevOps Engineer's Guide - TZoneLabs How to Run ML Training Jobs on Kubernetes With GPU Node Pools: A DevOps Engineer's Guide

By KP | TZoneLabs | DevOps & Cloud Engineering

Run training as a Kubernetes Job on a dedicated, tainted GPU node pool. Request nvidia.com/gpu: 1 in the Job, tolerate the pool’s taint, set a deadline so it can’t run forever, and let the autoscaler bring a GPU node up when the Job starts and remove it when the Job ends. The GPU bill then only runs while training does.

This is Part 3 of our MLOps series from a DevOps point of view. Part 1 mapped the ML lifecycle onto the DevOps stages we already run, and Part 2 covered versioning data and experiments. This one covers the training pipeline itself.

Each concept below is lined up against the DevOps task it mirrors, since reasoning from something you already run day to day is the fastest way into an unfamiliar workflow.

What Maps to What

ML Training Task	DevOps Equivalent
Kubernetes Job running `train.py`	A one-shot CI job in a throwaway container
GPU taint and toleration	Tagging a CI runner so only matching jobs land on it
NVIDIA device plugin	A DaemonSet agent such as node-exporter, one per node
`nvidia.com/gpu` limit	A CPU or memory resource request
GPU pool that scales from zero	A runner fleet that shrinks to nothing when the queue is empty
IAM role for the training pod	The deploy role a CI job assumes

Step 1: Create a Dedicated GPU Node Group

GPU instances cost many times what a general node does, so they get their own pool. A label lets training Jobs select it, and a taint keeps everything else off it. Here it is in Terraform, which fits the same module and remote state layout from our Terraform modules post:

resource "aws_eks_node_group" "gpu_training" {
  cluster_name    = aws_eks_cluster.main.name
  node_group_name = "gpu-training"
  node_role_arn   = aws_iam_role.node.arn
  subnet_ids      = var.private_subnet_ids

  ami_type       = "AL2023_x86_64_NVIDIA"
  instance_types = ["g5.xlarge"]

  scaling_config {
    min_size     = 0
    desired_size = 0
    max_size     = 4
  }

  labels = {
    workload = "gpu-training"
  }

  taint {
    key    = "nvidia.com/gpu"
    value  = "true"
    effect = "NO_SCHEDULE"
  }
}

Two settings matter most. min_size = 0 means the pool holds no nodes while nothing is training. The NVIDIA AMI type ships with the GPU drivers already installed, so the nodes don’t need a driver install step. Pick the AMI type that matches your Kubernetes version.

Step 2: Install the NVIDIA Device Plugin

Kubernetes doesn’t know what a GPU is until something tells it. The device plugin runs on each GPU node and advertises the nvidia.com/gpu resource to the scheduler:

helm repo add nvdp https://nvidia.github.io/k8s-device-plugin
helm repo update
helm upgrade -i nvdp nvdp/nvidia-device-plugin \
  --namespace nvidia-device-plugin --create-namespace

Once a GPU node is up, confirm it reports the resource:

kubectl describe node <gpu-node-name> | grep nvidia.com/gpu

If the line is missing, the scheduler will never place a GPU Job on that node, and the Job will sit in Pending with no obvious error.

Step 3: Write the Training Job

apiVersion: batch/v1
kind: Job
metadata:
  name: train-IMAGE_TAG
spec:
  backoffLimit: 1
  activeDeadlineSeconds: 14400
  ttlSecondsAfterFinished: 86400
  template:
    spec:
      restartPolicy: Never
      serviceAccountName: ml-training
      nodeSelector:
        workload: gpu-training
      tolerations:
        - key: nvidia.com/gpu
          operator: Exists
          effect: NoSchedule
      containers:
        - name: train
          image: registry.example.com/ml/train:IMAGE_TAG
          args: ["--data-version", "$(DATA_VERSION)"]
          env:
            - name: DATA_VERSION
              value: "a1b2c3d4"
            - name: MLFLOW_TRACKING_URI
              value: http://mlflow.mlops.svc:5000
          resources:
            requests:
              cpu: "4"
              memory: 16Gi
            limits:
              nvidia.com/gpu: 1

Each setting has a reason:

activeDeadlineSeconds kills the Job after four hours. A hung training run on a GPU node costs money every minute, so the deadline is a spending cap.
backoffLimit: 1 allows one retry. The default of six would rerun a broken job six times on expensive hardware.
ttlSecondsAfterFinished removes the finished Job and its pod after a day, so completed runs don’t pile up in the namespace.
The GPU limit is a whole number. Kubernetes doesn’t split a GPU between pods by default, and a GPU in limits doubles as its request.

Step 4: Launch It From the Pipeline

The Job manifest takes the place of the shell command in a normal CI stage. The pipeline from Part 1 starts it, waits for it, and fails if training fails:

train:
  stage: train
  script:
    - sed "s/IMAGE_TAG/$CI_COMMIT_SHORT_SHA/g" k8s/train-job.yaml | kubectl apply -f -
    - kubectl wait --for=condition=complete --timeout=4h job/train-$CI_COMMIT_SHORT_SHA
    - kubectl logs job/train-$CI_COMMIT_SHORT_SHA

One caveat: kubectl wait --for=condition=complete doesn’t notice a failed Job, so a crash holds the pipeline until the timeout. Teams that care about fast feedback poll for the Failed condition as well and exit as soon as either one shows up.

Step 5: Give the Pod Access to Data

The training pod has to read the dataset from S3 and write to the MLflow artifact bucket. On EKS the clean way is an IAM role bound to a Kubernetes service account, so no access keys end up in the Job:

apiVersion: v1
kind: ServiceAccount
metadata:
  name: ml-training
  annotations:
    eks.amazonaws.com/role-arn: arn:aws:iam::123456789012:role/ml-training

Scope that role to the one data prefix and the one artifact prefix the job needs. If a run fails with an access error, the steps in debugging IAM denials apply unchanged.

Mistakes That Show Up Later, Not Immediately

No taint on the GPU pool. Ordinary pods land on the GPU nodes and keep them running, so the pool never scales back to zero.
No activeDeadlineSeconds. One stuck run holds a GPU node overnight.
A multi-gigabyte image pulled fresh on every run. The GPU node bills while it waits for the pull. Keep the base image lean and put the model code on a shared base layer.
Spot GPU nodes with no checkpointing. A reclaimed node throws away hours of training. Write checkpoints to S3 so a retry resumes where it stopped.
Setting the GPU in requests only. Extended resources like nvidia.com/gpu have to be set in limits.
A Pending Job nobody looks at. A missing device plugin, a taint mismatch and a quota limit all look the same from the outside. Check kubectl describe pod for the scheduler’s reason first.

Key Lessons

Treat training as a batch job, not a service.
A Kubernetes Job runs to completion and exits, which matches how training behaves.
Isolate GPU nodes with a taint and a label.
The toleration and node selector together keep expensive hardware reserved for the work that needs it.
Scale the pool from zero.
GPU nodes that only exist during a run turn an always-on cost into a per-run cost.
Set a deadline on every Job.
activeDeadlineSeconds caps the worst case for any single run.
Use a service account role instead of keys.
The pod gets exactly the access it needs, and nothing sensitive sits in the manifest.

Summary

Need	Setting	Where
Reserve GPU nodes	Taint plus `workload` label	Node group
Expose the GPU to the scheduler	NVIDIA device plugin	Helm chart
Claim a GPU	`nvidia.com/gpu: 1` in limits	Job spec
Cap runaway cost	`activeDeadlineSeconds`	Job spec
Reach S3 without keys	Role on a service account	ServiceAccount

Start with one GPU node group at zero, one Job, and one deadline. Everything else on this list can come after the first run completes.

How to Run ML Training Jobs on Kubernetes With GPU Node Pools: A DevOps Engineer’s Guide