Docker for AI: Running GPU Workloads in Containers

Containers are the standard way to package and deploy AI workloads on a dedicated GPU server. Docker with NVIDIA GPU support gives you reproducible environments, easy version management, and clean isolation between multiple models running on the same machine. This tutorial covers everything from installing the NVIDIA Container Toolkit to running open-source LLMs in production containers on Ubuntu 22.04 and 24.04.

Install Docker Engine

Install Docker from the official repository (not the Ubuntu snap package):

# Remove old versions
sudo apt remove -y docker docker-engine docker.io containerd runc

# Install dependencies
sudo apt update
sudo apt install -y ca-certificates curl gnupg

# Add Docker GPG key and repository
sudo install -m 0755 -d /etc/apt/keyrings
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor -o /etc/apt/keyrings/docker.gpg
sudo chmod a+r /etc/apt/keyrings/docker.gpg

echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] \
    https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable" | \
    sudo tee /etc/apt/sources.list.d/docker.list > /dev/null

# Install Docker Engine
sudo apt update
sudo apt install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin

# Add your user to the docker group
sudo usermod -aG docker $USER
newgrp docker

# Verify installation
docker --version
docker run hello-world

Install NVIDIA Container Toolkit

The NVIDIA Container Toolkit allows Docker containers to access host GPUs. Ensure you have NVIDIA drivers installed on the host first.

# Add NVIDIA Container Toolkit repository
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | \
    sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg

curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
    sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
    sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

# Install the toolkit
sudo apt update
sudo apt install -y nvidia-container-toolkit

# Configure Docker to use NVIDIA runtime
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker

# Verify GPU access inside a container
docker run --rm --gpus all nvidia/cuda:12.4.0-runtime-ubuntu22.04 nvidia-smi

You should see the same GPU listing inside the container as you do on the host. The container inherits the host’s NVIDIA driver, so no CUDA installation is needed inside the container unless you need the compiler (nvcc).

Running GPU Containers

Use the --gpus flag to assign GPUs to containers:

# Use all available GPUs
docker run --rm --gpus all nvidia/cuda:12.4.0-runtime-ubuntu22.04 nvidia-smi

# Use specific GPUs by index
docker run --rm --gpus '"device=0,1"' nvidia/cuda:12.4.0-runtime-ubuntu22.04 nvidia-smi

# Use a specific number of GPUs
docker run --rm --gpus 2 nvidia/cuda:12.4.0-runtime-ubuntu22.04 nvidia-smi

# Run a PyTorch container with GPU
docker run --rm --gpus all -it pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime \
    python3 -c "import torch; print(f'GPUs: {torch.cuda.device_count()}')"

For persistent model serving, use named volumes to cache downloaded models:

# Create a volume for model storage
docker volume create model-cache

# Run vLLM with model caching
docker run -d --gpus all \
    --name vllm-server \
    -v model-cache:/root/.cache/huggingface \
    -p 8000:8000 \
    vllm/vllm-openai:latest \
    --model meta-llama/Llama-3.1-8B-Instruct \
    --max-model-len 4096

Docker Compose with GPU Support

Docker Compose is ideal for managing multi-service AI deployments. Here is a compose file that runs vLLM with an Nginx reverse proxy:

# docker-compose.yml
services:
  vllm:
    image: vllm/vllm-openai:latest
    command: >
      --model meta-llama/Llama-3.1-8B-Instruct
      --max-model-len 4096
      --gpu-memory-utilization 0.95
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    volumes:
      - model-cache:/root/.cache/huggingface
    ports:
      - "8000:8000"
    restart: unless-stopped

  nginx:
    image: nginx:latest
    volumes:
      - ./nginx.conf:/etc/nginx/conf.d/default.conf:ro
    ports:
      - "80:80"
      - "443:443"
    depends_on:
      - vllm
    restart: unless-stopped

volumes:
  model-cache:

# Start the stack
docker compose up -d

# Check logs
docker compose logs -f vllm

# Scale to multiple replicas (if you have the GPUs)
docker compose up -d --scale vllm=2

For Nginx configuration details, see the Nginx reverse proxy for AI inference guide.

AI-Specific Docker Images

Use purpose-built images for common AI frameworks:

# vLLM — high-throughput LLM serving
docker run --gpus all -p 8000:8000 vllm/vllm-openai:latest \
    --model meta-llama/Llama-3.1-8B-Instruct

# Ollama — simple LLM deployment
docker run -d --gpus all -p 11434:11434 \
    -v ollama-data:/root/.ollama \
    --name ollama ollama/ollama
docker exec ollama ollama pull llama3.1:8b

# PyTorch — training and fine-tuning
docker run --gpus all -it -v $(pwd):/workspace \
    pytorch/pytorch:2.3.0-cuda12.1-cudnn8-devel bash

# TensorFlow — training and serving
docker run --gpus all -it -v $(pwd):/workspace \
    tensorflow/tensorflow:latest-gpu bash

If you plan to run DeepSeek or other large models, check our best GPU for LLM inference guide for VRAM requirements before selecting your container configuration.

Multi-GPU Container Configuration

For models that require multiple GPUs on a multi-GPU cluster, assign specific devices:

# Docker Compose with specific GPU assignment
services:
  model-a:
    image: vllm/vllm-openai:latest
    command: --model meta-llama/Llama-3.1-8B-Instruct
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ['0']
              capabilities: [gpu]

  model-b:
    image: vllm/vllm-openai:latest
    command: --model mistralai/Mistral-7B-Instruct-v0.3
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ['1']
              capabilities: [gpu]

# Verify GPU isolation
docker exec model-a nvidia-smi  # Should show only GPU 0
docker exec model-b nvidia-smi  # Should show only GPU 1

Production Deployment Patterns

Set resource limits and health checks for production stability:

services:
  vllm:
    image: vllm/vllm-openai:latest
    command: >
      --model meta-llama/Llama-3.1-8B-Instruct
      --gpu-memory-utilization 0.95
      --max-model-len 4096
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]
        limits:
          memory: 32g
          cpus: '8'
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:8000/health"]
      interval: 30s
      timeout: 10s
      retries: 3
      start_period: 120s
    logging:
      driver: json-file
      options:
        max-size: "100m"
        max-file: "3"
    restart: unless-stopped

For monitoring GPU usage inside containers, see our GPU monitoring guide. To secure your containerised API with authentication, follow the secure AI inference API tutorial. Find more guides in the tutorials category.