Prasanta Kumar Behera
AI-Assisted Infrastructure Engineering›13 · GPU scheduling for AI workloads
Learning Hub / AI-Assisted Engineering / AI-Assisted Infrastructure Engineering

Lesson 13 of 14 · Modules

GPU scheduling for AI workloads

Run AI workloads on Kubernetes GPUs: how GPUs become schedulable resources (device plugin, GPU Operator), requesting them, sharing expensive GPUs with MIG partitions or time-slicing (and their isolation trade-offs), keeping GPU nodes for GPU work, and queueing batch jobs fairly.

Advanced
Key wordsGPUs on KubernetesNVIDIA device pluginGPU Operatornvidia.com/gpuMIGtime-slicingMPSDynamic Resource Allocationnode labelstaintsKueue

GPUs as schedulable resources

Kubernetes doesn't know about GPUs natively. A device plugin (e.g. NVIDIA's) discovers GPUs on each node and advertises them as an extended resource such as nvidia.com/gpu. Pods request them in limits:

apiVersion: v1
kind: Pod
metadata:
  name: cuda-check
spec:
  restartPolicy: Never
  tolerations:
    - key: nvidia.com/gpu
      operator: Exists
      effect: NoSchedule
  containers:
    - name: cuda
      image: nvcr.io/nvidia/cuda:12.4.1-base-ubuntu22.04   # pick a tag matching your driver
      command: [ "nvidia-smi" ]
      resources:
        limits:
          nvidia.com/gpu: 1

The NVIDIA GPU Operator installs and manages the pieces on each GPU node: driver (or uses a preinstalled one), container toolkit, device plugin, GPU Feature Discovery (node labels such as nvidia.com/gpu.product), and the DCGM exporter for Prometheus metrics. Other vendors have equivalent plugins/operators.

GPUs are the school's few expensive science microscopes. The device plugin is the lab technician who counts them and hands them out. A child must book one (request in limits). To let more children use them, you can split a microscope into separate small stations (MIG, each with its own lens), or let children take turns quickly at one microscope (time-slicing), which is cheaper but they can bump into each other.

Sharing expensive GPUs

Method Isolation Good for
Whole GPU Full Training, heavy inference
MIG (Multi-Instance GPU, on supported data-centre GPUs) Hardware-partitioned compute and memory Multiple inference services with predictable performance
Time-slicing None (shared memory, no fault isolation) Dev/test, light or bursty inference
MPS Partial (process-level sharing with limits) Many small concurrent processes
  • MIG profiles appear as resources like nvidia.com/mig-1g.10gb (with the "mixed" strategy), configured through the GPU Operator.
  • Time-slicing is configured in the device plugin (a number of replicas per GPU); pods still request nvidia.com/gpu: 1 but may share a physical GPU.
  • Dynamic Resource Allocation (DRA), which has matured in recent Kubernetes releases, provides richer device requests (attributes, sharing) through DeviceClasses and ResourceClaims; vendor drivers are adopting it. Check your Kubernetes and driver versions.

Placement and fairness

  • Taint GPU nodes (nvidia.com/gpu=present:NoSchedule or similar) and add tolerations to GPU workloads.
  • Use node affinity on GPU product labels for workloads that need specific models.
  • For batch/training queues, Kueue (a Kubernetes-native job queueing system) provides quotas, fair sharing and admission for GPU jobs across teams.
  • Monitor with DCGM exporter metrics (utilisation, memory, temperature, errors); idle GPUs are expensive.

Try it: GPUs (or the concepts without them)

  1. If you have a GPU node (or a cloud GPU instance), install the NVIDIA GPU Operator with Helm and run the cuda-check pod; read nvidia-smi output in its logs.
  2. Label and taint the GPU node; confirm a normal pod doesn't land there and the GPU pod does.
  3. Enable time-slicing with 4 replicas in the device plugin config; run 4 pods requesting one GPU each and watch them share it.
  4. Scrape DCGM exporter metrics into Prometheus and build a utilisation panel.
  5. Without hardware: read the GPU Operator and Kueue docs and write a short design for sharing 8 GPUs among 3 teams (MIG vs time-slicing vs whole GPUs, quotas).

Going deeper: GPU platforms

  • Plan drivers carefully: driver, CUDA and container image compatibility; upgrade GPU nodes like any other node image (see Edge Kubernetes & Zero-Touch Provisioning, lesson 06).
  • Chargeback GPU time per team; utilisation reports change behaviour quickly.
  • For large training jobs, look at gang scheduling and topology-aware placement (NVLink/network), via queueing systems and schedulers built for batch.

Recap

  • GPUs are extended resources advertised by a device plugin; request them in limits.
  • The GPU Operator manages drivers, toolkit, plugin, feature discovery and DCGM metrics.
  • Share GPUs with MIG (isolated partitions), time-slicing (no isolation) or MPS; watch DRA for richer allocation.
  • Taint GPU nodes, select by product labels, queue fairly with Kueue, and monitor utilisation.

This site is a public version of my personal engineering knowledge hub. It intentionally excludes confidential company information and internal operational details.