Kubernetes provides powerful orchestration for AI workloads with GPU scheduling, resource management, and scalability.

GPU Setup

NVIDIA GPU Operator

# Install GPU Operator
helm repo add nvidia https://helm.ngc.nvidia.com/nvidia
helm install gpu-operator nvidia/gpu-operator \
    --namespace gpu-operator \
    --create-namespace \
    --set driver.enabled=true \
    --set mig.strategy=single  # For Ampere+ GPUs

GPU Pod Configuration

apiVersion: v1
kind: Pod
metadata:
  name: gpu-training-job
spec:
  containers:
  - name: trainer
    image: nvcr.io/nvidia/pytorch:24.01-py3
    command: ["python", "train.py"]
    resources:
      limits:
        nvidia.com/gpu: 4      # Request 4 GPUs
        memory: "32Gi"
        cpu: "16"
      requests:
        nvidia.com/gpu: 4
        memory: "32Gi"
    volumeMounts:
    - name: dataset
      mountPath: /data
    - name: checkpoints
      mountPath: /checkpoints
  volumes:
  - name: dataset
    persistentVolumeClaim:
      claimName: dataset-pvc
  - name: checkpoints
    persistentVolumeClaim:
      claimName: checkpoints-pvc
  tolerations:
  - key: nvidia.com/gpu
    operator: Exists
    effect: NoSchedule

MIG (Multi-Instance GPU)

# Request specific MIG slice
resources:
  limits:
    nvidia.com/mig-1g.10gb: 1  # 1/7th of A100
    # Options: 1g.10gb, 2g.20gb, 3g.40gb, 7g.80gb

Node Affinity for GPU Selection

affinity:
  nodeAffinity:
    requiredDuringSchedulingIgnoredDuringExecution:
      nodeSelectorTerms:
      - matchExpressions:
        - key: nvidia.com/gpu.product
          operator: In
          values:
          - NVIDIA-A100-SXM4-80GB
          - NVIDIA-H100-80GB-HBM3