Demystifying GPU Management: High-Performance AI on Kubernetes#

Key GPU Features#

• Parallelism: Thousands of smaller cores for concurrent operations.
• High Throughput: GPUs can perform billions of floating-point operations per second (FLOPS).
• Specialized Libraries: Ecosystems such as CUDA (for NVIDIA GPUs) and ROCm (for AMD GPUs) provide specialized APIs and libraries.

Because of these capabilities, GPUs have become indispensable to modern AI and high-performance computing tasks. However, the challenge is how to effectively manage and scale these GPU resources across dozens or even hundreds of machines. This is where Kubernetes steps in.

Container Orchestration and AI#

Container orchestration with Kubernetes isn’t just about running any workload; it’s about managing distributed systems in an automated and scalable manner. For AI and ML workloads, which can be resource-intensive, Kubernetes provides advanced scheduling to ensure that containers using GPUs are placed on nodes that have GPU capabilities.

Prerequisites#

1. A functioning Kubernetes cluster: You should already have a cluster set up—either on premises, on a public cloud, or in a hybrid environment.
2. GPU Hardware: At least one node should be equipped with a GPU (e.g., an NVIDIA GPU).
3. Suitable drivers and runtime: For NVIDIA GPUs, you need to install NVIDIA drivers on the node.

Installing GPU Drivers on Your Nodes#

On GPU-equipped nodes, you need to install the appropriate drivers for your hardware. For NVIDIA GPUs:

1. Add the CUDA repository:

   1
   distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
   2
   wget https://developer.download.nvidia.com/compute/cuda/repos/$distribution/x86_64/cuda-keyring.gpg
   3
   sudo rpm --import cuda-keyring.gpg  # or 'apt-key add' for Debian/Ubuntu
   

2. Install the driver packages:

   1
   sudo apt-get update
   2
   sudo apt-get install -y cuda-drivers
   

   or:

   1
   sudo yum clean expire-cache
   2
   sudo yum install -y nvidia-driver-latest-dkms
   

3. Verify:

   1
   nvidia-smi
   

   If successful, you should see information about your GPU model and driver version.

NVIDIA Device Plugin for Kubernetes#

Kubernetes includes support for accelerated computing using plugin-like integrations. NVIDIA provides a device plugin to make the GPU resources discoverable by the kubelet.

1. Deploy using YAML (for a typical setup):

   1
   kubectl create -f https://raw.githubusercontent.com/NVIDIA/k8s-device-plugin/v0.12.3/nvidia-device-plugin.yml
   

   This creates a DaemonSet that runs on every GPU node in your cluster, registering NVIDIA GPUs with Kubernetes.

2. Check the logs to confirm:

   1
   kubectl logs -n kube-system <nvidia-device-plugin-pod-name>
   

   It should report something like “Found X GPU(s)”.

This device plugin fundamentally exposes GPUs as schedulable resources. With the plugin in place, you can start running GPU-accelerated Pods.

GPU Requests and Limits#

Kubernetes uses a resource request/limit model to ensure Pods get the resources they need:

• Requests: The amount of resource (CPU, GPU, memory) that a container is guaranteed.
• Limits: The maximum amount of resource a container is allowed to use.

For GPU-enabled Pods, you specify GPUs in the resources section of your Pod spec:

1
apiVersion: v1
2
kind: Pod
3
metadata:
4
  name: gpu-pod
5
spec:
6
  containers:
7
  - name: gpu-container
8
    image: nvidia/cuda:11.3.0-base
9
    resources:
10
      limits:
11
        nvidia.com/gpu: 1

In this example, Kubernetes will schedule the Pod on a node with at least one available GPU.

Node Labels and GPU Affinity#

Kubernetes node labels help the scheduler decide where to place your workloads. If you have different GPU models across your cluster—for example, some nodes have Tesla V100, and others have Tesla T4—you can label the nodes:

1
kubectl label nodes <node-name> gpu-type=tesla-t4
2
kubectl label nodes <node-name> gpu-type=tesla-v100

Then you can use node affinity in your Pod specification:

1
apiVersion: v1
2
kind: Pod
3
metadata:
4
  name: gpu-affinity-pod
5
spec:
6
  affinity:
7
    nodeAffinity:
8
      requiredDuringSchedulingIgnoredDuringExecution:
9
        nodeSelectorTerms:
10
        - matchExpressions:
11
          - key: gpu-type
12
            operator: In
13
            values:
14
            - tesla-t4
15
  containers:
16
    - name: gpu-container
17
      image: nvidia/cuda:11.3.0-base
18
      resources:
19
        limits:
20
          nvidia.com/gpu: 1

In this manner, you can ensure that your workloads land on nodes that have the correct GPU type.

GPU-Sharing Mechanisms#

By default, Kubernetes schedules GPU resources in an exclusive manner, meaning one container receives the entire physical GPU. However, you can use proprietary or specialized device plugins that enable fractional GPU sharing. This is particularly useful for inference workloads, where a single GPU can be shared among multiple models at once without exhausting GPU memory. Examples of GPU sharing solutions include:

• NVIDIA MPS (Multi-Process Service): Allows multiple processes to share GPU contexts.
• Third-Party Device Plugins: Provide fractional GPU allocation.

However, these are more advanced setups and might require extra configuration in your cluster.

Example: TensorFlow with GPUs#

Here is a minimal YAML file running TensorFlow on GPUs:

1
apiVersion: batch/v1
2
kind: Job
3
metadata:
4
  name: tf-gpu-job
5
spec:
6
  template:
7
    spec:
8
      restartPolicy: Never
9
      containers:
10
      - name: tf-container
11
        image: tensorflow/tensorflow:latest-gpu
12
        command: ["python", "-c"]
13
        args:
14
        - >
15
          import tensorflow as tf;
16
          print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')));
17
        resources:
18
          limits:
19
            nvidia.com/gpu: 1

1. Create the Job:

   1
   kubectl apply -f tf-gpu-job.yaml
   

2. Check the status:

   1
   kubectl get jobs
   

3. Check the logs:

   1
   kubectl logs <the-pod-name>
   

If everything is correct, the output will show the number of GPUs recognized by TensorFlow.

Example: PyTorch with GPUs#

Similarly, for PyTorch:

1
apiVersion: batch/v1
2
kind: Job
3
metadata:
4
  name: pytorch-gpu-job
5
spec:
6
  template:
7
    spec:
8
      restartPolicy: Never
9
      containers:
10
      - name: pytorch-container
11
        image: pytorch/pytorch:latest
12
        command: ["python", "-c"]
13
        args:
14
        - >
15
          import torch;
16
          print("Is CUDA available?", torch.cuda.is_available());
17
          print("GPU Count:", torch.cuda.device_count());
18
        resources:
19
          limits:
20
            nvidia.com/gpu: 1

You should see Is CUDA available? True and the number of GPU devices recognized by PyTorch.

Managing GPU Memory#

Deep learning containers can be memory-hungry, and GPU memory constraints are a common bottleneck. If your GPU is 16GB of VRAM, and your model plus batch size demands more, you’ll run into out-of-memory errors.

• Batch Size Tuning: Find the optimal batch size that fits in GPU memory.
• Gradient Checkpointing: Reduce GPU memory usage by recomputing some intermediate states during backpropagation.
• Mixed-Precision Training: Use half-precision (FP16) or Tensor Cores on supported GPUs to reduce memory usage and improve speed.

Multi-GPU and Distributed Training#

For more advanced scenarios, you may want to run distributed training across multiple GPUs and potentially across multiple nodes. Popular frameworks like PyTorch’s torch.distributed or TensorFlow’s tf.distribute can help.

1
apiVersion: v1
2
kind: Pod
3
metadata:
4
  name: multi-gpu-pod
5
spec:
6
  containers:
7
  - name: multi-gpu-container
8
    image: tensorflow/tensorflow:latest-gpu
9
    resources:
10
      limits:
11
        nvidia.com/gpu: 2

Monitoring and Metrics#

Keeping an eye on GPU utilization is crucial. You can integrate solutions like:

1. Prometheus + Grafana:

   • Deploy the Prometheus node exporter or the NVIDIA GPU exporter.
   • Export metrics like GPU usage, memory usage, temperature, etc.
   • Visualize using Grafana dashboards.

2. NVIDIA DCGM (Data Center GPU Manager):

   • Provides GPU-level metrics and can integrate with Prometheus.

Below is a snippet of a Prometheus configuration for scraping GPU metrics from a node exporter:

1
scrape_configs:
2
- job_name: 'gpu-metrics'
3
  static_configs:
4
  - targets: ['<node-ip>:9100']  # or a service if you have one

GPU Operator and CRDs#

The NVIDIA GPU Operator uses the Operator pattern to simplify GPU management. It automates:

• Driver installation
• Device plugin deployment
• Monitoring components

It leverages Custom Resource Definitions (CRDs) to abstract complex GPU management tasks into higher-level concepts. If you’re scaling up GPU usage across many nodes, or you frequently add new GPU nodes, adopting the GPU Operator can save you a lot of time.

Version Compatibility#

The moving parts in a GPU-enabled Kubernetes cluster include:

• Kubernetes version
• NVIDIA drivers
• CUDA toolkit
• NVIDIA device plugin

Incompatibilities can cause major headaches—for instance, if the device plugin version does not match the GPU driver version. Always check the official compatibility matrix from NVIDIA or your GPU vendor.

Resource Quotas and Avoiding Overcommitment#

Kubernetes allows you to set Resource Quotas in a namespace to define the maximum GPU usage. Overcommitment is generally not applicable to GPUs because they cannot be shared the same way CPU or memory can be overcommitted. If your cluster tries to schedule a GPU workload when none are available, the Pod will be stuck in Pending.

Security and Isolation#

Given GPUs are shared hardware, GPU isolation is crucial when running multi-tenant workloads. By default, the device plugin provides exclusive GPU allocation. However, malicious container processes could still attempt side-channel attacks. To mitigate:

• Use Runtime Security Tools: Tools such as SELinux, AppArmor, or seccomp can provide additional constraint layers.
• Sandboxing: For extremely sensitive workloads, consider using separate GPU nodes or advanced techniques like GPU virtualization.

Upgrades and Maintenance#

Keeping your cluster current is a never-ending task. Regularly updating the Kubernetes version, device plugin versions, and GPU drivers is important but can be disruptive:

• Rollouts: Use rolling upgrades for node updates.
• Version Skew: Validate that your device plugin is forward-compatible with the new driver or Kubernetes release.
• Backup: Always back up your cluster state, especially if using critical data or states in ephemeral volumes.

Real-Time Inference#

For use cases like real-time video analytics or online recommendation engines, GPU-enabled inference can drastically reduce latency. By deploying your inference model in a GPU-enabled container, you can scale horizontally as traffic grows. Tools like Triton Inference Server from NVIDIA provide a standardized way to serve multiple models on the same GPU.

Data Analytics and ETL#

ETL or data-processing tasks can benefit from GPU acceleration. Libraries like RAPIDS accelerate data transformations and analytics in Python using GPUs. Utilizing these libraries within a Kubernetes environment allows large data sets to be processed faster.

Model Training Pipelines#

From data preprocessing to hyperparameter tuning, GPUs accelerate each step of an ML pipeline. Integrating Kubeflow or Argo Workflows allows you to define complex pipelines where certain tasks specifically request GPUs when required, and revert to CPU nodes for lighter tasks.

Demo Architecture#

In this section, we’ll walk through a simple demonstration pipeline:

1. A data preprocessing step (CPU-based).
2. A model training step (GPU-based).
3. An inference service (GPU-based, for demonstration).
4. A monitoring setup to track GPU usage.

Below is a high-level architecture diagram (in text form):

1
       Data Source
2
           |
3
           v
4
    [Preprocessing Pod] --(CPU only)--> [Model Checkpoint]
5
           |
6
           v
7
    [Training Pod] ------(GPU)--------> [Trained Model]
8
           |
9
           v
10
   [Inference Pod] ----(GPU)----> [Service Exposure]
11
           |
12
           v
13
    [Monitoring Setup: Prometheus + Node Exporter + Grafana]

Step-by-Step Guide#

1
apiVersion: batch/v1
2
kind: Job
3
metadata:
4
  name: data-preprocessing-job
5
spec:
6
  template:
7
    spec:
8
      containers:
9
      - name: preprocess
10
        image: python:3.9
11
        command: ["python", "-c"]
12
        args:
13
        - >
14
          import sys;
15
          # Pseudo code for processing:
16
          # read input data, clean it, transform it, save to persistent storage
17
          print("Data preprocessing completed");
18
      restartPolicy: Never

1
apiVersion: batch/v1
2
kind: Job
3
metadata:
4
  name: model-training-job
5
spec:
6
  template:
7
    spec:
8
      containers:
9
      - name: train
10
        image: pytorch/pytorch:latest
11
        command: ["python", "train.py"]
12
        # 'train.py' could be part of the container or mounted via ConfigMap
13
        resources:
14
          limits:
15
            nvidia.com/gpu: 1
16
      restartPolicy: Never

1
import torch
2

3
# Hyperparameters
4
epochs = 5
5
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
6

7
print(f"Using device: {device}")
8

9
# Hypothetical model, data loader, etc.
10
# model = ...
11
# optimizer = ...
12
# data_loader = ...
13

14
for epoch in range(epochs):
15
    # training loop
16
    pass
17

18
# Save the trained model
19
torch.save(model.state_dict(), "model.pt")

1
apiVersion: apps/v1
2
kind: Deployment
3
metadata:
4
  name: inference-deployment
5
spec:
6
  replicas: 2
7
  selector:
8
    matchLabels:
9
      app: inference
10
  template:
11
    metadata:
12
      labels:
13
        app: inference
14
    spec:
15
      containers:
16
      - name: inference
17
        image: python:3.9
18
        command: ["python", "inference_server.py"]
19
        resources:
20
          limits:
21
            nvidia.com/gpu: 1

1
from flask import Flask, request, jsonify
2
import torch
3

4
app = Flask(__name__)
5

6
model = ... # load model from model.pt
7
model.cuda()
8
model.eval()
9

10
@app.route('/predict', methods=['POST'])
11
def predict():
12
    data = request.json
13
    # preprocess input
14
    with torch.no_grad():
15
        output = model(...)  # pass data to model on GPU
16
    return jsonify({"output": output.tolist()})
17

18
if __name__ == '__main__':
19
    app.run(host='0.0.0.0', port=8080)

1
apiVersion: v1
2
kind: Service
3
metadata:
4
  name: inference-service
5
spec:
6
  type: NodePort
7
  selector:
8
    app: inference
9
  ports:
10
    - port: 8080
11
      targetPort: 8080
12
      nodePort: 30002