Deploy AI Models with Docker: Production Best Practices

Why Docker for ML?

Docker solves the "works on my machine" problem that plagues ML deployments. Benefits include:

Reproducibility: Exact same environment everywhere
Portability: Deploy to any cloud provider
Isolation: No dependency conflicts
Scalability: Easy horizontal scaling
Version control: Tag and rollback deployments

💡 What you'll learn

How to containerize ML models with GPU support, optimize image sizes, and deploy to production on GPUBrazil instances.

Prerequisites

On your GPUBrazil instance:

# Verify Docker is installed
docker --version

# Verify NVIDIA Container Toolkit
nvidia-smi
docker run --rm --gpus all nvidia/cuda:12.1-base nvidia-smi

Basic ML Dockerfile

Start with NVIDIA's CUDA base images:

# Dockerfile
FROM nvidia/cuda:12.1-cudnn8-runtime-ubuntu22.04

# Set environment variables
ENV PYTHONDONTWRITEBYTECODE=1
ENV PYTHONUNBUFFERED=1
ENV DEBIAN_FRONTEND=noninteractive

# Install Python
RUN apt-get update && apt-get install -y \
    python3.10 \
    python3-pip \
    && rm -rf /var/lib/apt/lists/*

# Set working directory
WORKDIR /app

# Install Python dependencies
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy application code
COPY . .

# Expose port
EXPOSE 8000

# Run application
CMD ["python3", "serve.py"]

Multi-Stage Builds for Smaller Images

ML images can be huge. Multi-stage builds reduce size significantly:

# Dockerfile.multistage
# Stage 1: Build dependencies
FROM nvidia/cuda:12.1-cudnn8-devel-ubuntu22.04 as builder

RUN apt-get update && apt-get install -y \
    python3.10 python3.10-dev python3-pip \
    build-essential cmake git

WORKDIR /build

COPY requirements.txt .
RUN pip install --user --no-cache-dir -r requirements.txt

# Stage 2: Runtime image (much smaller)
FROM nvidia/cuda:12.1-cudnn8-runtime-ubuntu22.04

RUN apt-get update && apt-get install -y \
    python3.10 python3-pip \
    && rm -rf /var/lib/apt/lists/*

# Copy only installed packages from builder
COPY --from=builder /root/.local /root/.local
ENV PATH=/root/.local/bin:$PATH

WORKDIR /app
COPY . .

EXPOSE 8000
CMD ["python3", "serve.py"]

🚀 Size Reduction

Multi-stage builds can reduce image size from 15GB to 5GB by excluding build tools from the final image.

Optimized PyTorch Dockerfile

# Dockerfile.pytorch
FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime

ENV PYTHONDONTWRITEBYTECODE=1
ENV PYTHONUNBUFFERED=1

WORKDIR /app

# Install additional dependencies
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy model weights (or download at runtime)
COPY models/ /app/models/

# Copy application
COPY src/ /app/src/

EXPOSE 8000

# Health check
HEALTHCHECK --interval=30s --timeout=10s --start-period=60s \
    CMD curl -f http://localhost:8000/health || exit 1

CMD ["python", "-m", "src.serve"]

Model Serving with FastAPI

# src/serve.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import os

app = FastAPI()

# Load model on startup
model = None
tokenizer = None

@app.on_event("startup")
async def load_model():
    global model, tokenizer
    model_path = os.getenv("MODEL_PATH", "/app/models/llama-3.1-8b")
    
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        torch_dtype=torch.float16,
        device_map="auto"
    )

class GenerateRequest(BaseModel):
    prompt: str
    max_tokens: int = 256
    temperature: float = 0.7

class GenerateResponse(BaseModel):
    text: str
    tokens_generated: int

@app.get("/health")
async def health():
    return {"status": "healthy", "gpu": torch.cuda.is_available()}

@app.post("/generate", response_model=GenerateResponse)
async def generate(request: GenerateRequest):
    try:
        inputs = tokenizer(request.prompt, return_tensors="pt").to("cuda")
        
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=request.max_tokens,
                temperature=request.temperature,
                do_sample=True
            )
        
        text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        tokens = len(outputs[0]) - len(inputs.input_ids[0])
        
        return GenerateResponse(text=text, tokens_generated=tokens)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

Docker Compose for Development

# docker-compose.yml
version: '3.8'

services:
  inference:
    build:
      context: .
      dockerfile: Dockerfile.pytorch
    ports:
      - "8000:8000"
    volumes:
      - ./models:/app/models:ro
      - ./src:/app/src:ro  # Hot reload in dev
    environment:
      - MODEL_PATH=/app/models/llama-3.1-8b
      - CUDA_VISIBLE_DEVICES=0
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:8000/health"]
      interval: 30s
      timeout: 10s
      retries: 3
      start_period: 120s

  # Optional: Redis for caching
  redis:
    image: redis:7-alpine
    ports:
      - "6379:6379"
    volumes:
      - redis_data:/data

volumes:
  redis_data:

Run with GPU

# Start services
docker compose up -d

# Check logs
docker compose logs -f inference

# Scale horizontally (if you have multiple GPUs)
docker compose up -d --scale inference=2

Production Dockerfile with vLLM

# Dockerfile.vllm
FROM vllm/vllm-openai:latest

# Install additional dependencies
RUN pip install prometheus-client

# Copy custom configuration
COPY vllm_config.yaml /app/config.yaml

# Environment variables
ENV MODEL_NAME=meta-llama/Llama-3.1-8B-Instruct
ENV MAX_MODEL_LEN=8192
ENV GPU_MEMORY_UTILIZATION=0.9

EXPOSE 8000

# Health check
HEALTHCHECK --interval=30s --timeout=10s --start-period=120s \
    CMD curl -f http://localhost:8000/health || exit 1

CMD ["python", "-m", "vllm.entrypoints.openai.api_server", \
     "--model", "${MODEL_NAME}", \
     "--max-model-len", "${MAX_MODEL_LEN}", \
     "--gpu-memory-utilization", "${GPU_MEMORY_UTILIZATION}"]

Baking Models into Images

For faster cold starts, include model weights in the image:

# Dockerfile.with-model
FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime

WORKDIR /app

# Download model during build
RUN pip install huggingface_hub
RUN python -c "from huggingface_hub import snapshot_download; \
    snapshot_download('meta-llama/Llama-3.1-8B-Instruct', \
    local_dir='/app/models/llama-3.1-8b')"

# Install dependencies
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY src/ /app/src/

ENV MODEL_PATH=/app/models/llama-3.1-8b
EXPOSE 8000

CMD ["python", "-m", "src.serve"]

⚠️ Image Size Warning

Baking models creates large images (10-50GB+). Use a container registry close to your deployment region to minimize pull times.

Download Models at Runtime

Alternative: Download models on container start:

# entrypoint.sh
#!/bin/bash
set -e

MODEL_PATH=${MODEL_PATH:-"/app/models"}
MODEL_NAME=${MODEL_NAME:-"meta-llama/Llama-3.1-8B-Instruct"}

# Download model if not present
if [ ! -d "$MODEL_PATH/$(basename $MODEL_NAME)" ]; then
    echo "Downloading model $MODEL_NAME..."
    python -c "from huggingface_hub import snapshot_download; \
        snapshot_download('$MODEL_NAME', local_dir='$MODEL_PATH/$(basename $MODEL_NAME)')"
fi

# Start server
exec python -m src.serve

# Dockerfile with entrypoint
FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY entrypoint.sh /entrypoint.sh
RUN chmod +x /entrypoint.sh

COPY src/ /app/src/
WORKDIR /app

# Mount volume for models
VOLUME /app/models

EXPOSE 8000
ENTRYPOINT ["/entrypoint.sh"]

Security Best Practices

Don't Run as Root

# Create non-root user
RUN useradd -m -u 1000 appuser
USER appuser

WORKDIR /home/appuser/app
COPY --chown=appuser:appuser . .

Scan for Vulnerabilities

# Build image
docker build -t mymodel:latest .

# Scan with Trivy
docker run --rm -v /var/run/docker.sock:/var/run/docker.sock \
    aquasec/trivy image mymodel:latest

Use Secrets Properly

# docker-compose.yml with secrets
services:
  inference:
    build: .
    secrets:
      - hf_token
    environment:
      - HF_TOKEN_FILE=/run/secrets/hf_token

secrets:
  hf_token:
    file: ./secrets/hf_token.txt

CI/CD Pipeline

# .github/workflows/deploy.yml
name: Build and Deploy

on:
  push:
    branches: [main]

jobs:
  build:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v3
      
      - name: Login to Registry
        uses: docker/login-action@v2
        with:
          registry: ghcr.io
          username: ${{ github.actor }}
          password: ${{ secrets.GITHUB_TOKEN }}
      
      - name: Build and Push
        uses: docker/build-push-action@v4
        with:
          context: .
          push: true
          tags: ghcr.io/${{ github.repository }}:${{ github.sha }}
          cache-from: type=gha
          cache-to: type=gha,mode=max
  
  deploy:
    needs: build
    runs-on: ubuntu-latest
    steps:
      - name: Deploy to GPUBrazil
        uses: appleboy/ssh-action@master
        with:
          host: ${{ secrets.GPU_HOST }}
          username: ubuntu
          key: ${{ secrets.SSH_KEY }}
          script: |
            docker pull ghcr.io/${{ github.repository }}:${{ github.sha }}
            docker stop inference || true
            docker run -d --rm --gpus all \
              --name inference \
              -p 8000:8000 \
              ghcr.io/${{ github.repository }}:${{ github.sha }}

Deploy Your Models Today

Get a GPU instance on GPUBrazil and deploy in minutes.

Get $5 Free Credit →

Performance Tips

1. Layer Caching

Order Dockerfile instructions from least to most frequently changing:

# Good: Dependencies cached unless requirements.txt changes
COPY requirements.txt .
RUN pip install -r requirements.txt
COPY src/ /app/src/  # This layer rebuilds often

2. Use .dockerignore

# .dockerignore
.git
__pycache__
*.pyc
.env
.venv
models/  # If downloading at runtime
*.log
.pytest_cache

3. Shared Memory for PyTorch

# Increase shared memory for DataLoader workers
docker run --shm-size=2g --gpus all mymodel:latest

Monitoring

# Add Prometheus metrics
from prometheus_client import Counter, Histogram, start_http_server

REQUEST_COUNT = Counter('requests_total', 'Total requests')
REQUEST_LATENCY = Histogram('request_latency_seconds', 'Request latency')

@app.on_event("startup")
async def start_metrics():
    start_http_server(9090)  # Metrics on port 9090

@app.post("/generate")
@REQUEST_LATENCY.time()
async def generate(request: GenerateRequest):
    REQUEST_COUNT.inc()
    # ... generation logic

Conclusion

Docker is essential for reproducible ML deployments. Key takeaways:

Use multi-stage builds to reduce image size
Start from NVIDIA base images for GPU support
Include health checks for production readiness
Consider baking vs runtime model downloads based on your needs
Implement proper security practices

Deploy your containerized models on GPUBrazil for cost-effective GPU inference at scale.