Why this matters

You already know Docker. The gap when moving to AI workloads is not container basics — it is the operational realities that make naive Docker usage painful: model weights that are gigabytes, Python dependencies that take ten minutes to install, GPU drivers that conflict, and images that bloat to 15 GB because someone ran pip install torch in the wrong layer.

This lesson covers the patterns that keep AI containers fast to build, small to ship, and correct to run.

Core concepts

The dependency problem

A standard Flask API image might be 200 MB. An LLM application that pulls in torch, transformers, sentence-transformers, and a vector client easily reaches 6–10 GB before you add your own code. That has concrete consequences:

• Cold start time on serverless or new nodes: 3–8 minutes pulling the image
• CI pipeline time dominated by layer pulls
• Storage costs at scale
• Dependency conflicts between CUDA versions and library expectations

The fix is to treat AI dependencies as a separate concern from application code.

Multi-stage builds for AI applications

Multi-stage builds let you separate the heavy installation phase from the final runtime image. The canonical pattern for an LLM application:

# Stage 1: dependency builder
FROM python:3.11-slim AS builder

WORKDIR /build
COPY requirements.txt .

# Install build tools in this stage only
RUN apt-get update && apt-get install -y --no-install-recommends \
    build-essential git \
    && rm -rf /var/lib/apt/lists/*

RUN pip install --no-cache-dir --prefix=/install -r requirements.txt

# Stage 2: runtime image
FROM python:3.11-slim AS runtime

# Copy installed packages from builder
COPY --from=builder /install /usr/local

WORKDIR /app
COPY src/ ./src/

# Non-root user for production
RUN useradd --create-home appuser
USER appuser

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

The builder stage installs everything. The runtime stage copies only the installed packages — no build tools, no apt cache, no pip cache. This typically saves 300–500 MB.

GPU containers

If your workload runs local inference (vLLM, Ollama, llama.cpp), you need a CUDA base image and GPU access at runtime. The key decisions:

CUDA version pinning. The CUDA version in your image must be compatible with the driver version on your host. Check with nvidia-smi — the output shows the maximum supported CUDA version. Mismatches cause silent runtime failures.

# For PyTorch inference with CUDA 12.1
FROM nvidia/cuda:12.1.0-cudnn8-runtime-ubuntu22.04 AS runtime

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

# torch with CUDA support — must match the base image CUDA version
RUN pip install torch==2.1.0+cu121 --index-url https://download.pytorch.org/whl/cu121

Docker Compose for local GPU development:

services:
  inference:
    build: .
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]
    environment:
      - CUDA_VISIBLE_DEVICES=0

Model weight management

Model weights should not be in your Docker image. A 7B parameter model in float16 is ~14 GB. That would make your image effectively unshippable and would rebuild from scratch every time your application code changes.

The correct pattern is to separate model weights from application code entirely:

Option 1: Volume mount at runtime. Pull weights to a host directory or network volume, mount at /models. The container expects MODEL_PATH=/models/llama-3-8b from the environment.

Option 2: Sidecar init container. A small init container downloads the model at pod startup to a shared volume. The main container waits until the model file exists.

Option 3: Model server as a separate service. Run Ollama or vLLM as their own service. Your application calls their HTTP API. This is the cleanest boundary for most production deployments.

import os
from pathlib import Path


def get_model_path() -> Path:
    '''Resolve model path from environment, with a local fallback for development.'''
    model_path = Path(os.environ.get("MODEL_PATH", "/models/default"))
    if not model_path.exists():
        raise RuntimeError(
            f"Model not found at {model_path}. "
            "Set MODEL_PATH or mount the models volume."
        )
    return model_path

Requirements file hygiene

AI requirements files grow organic complexity fast. Three practices that keep them manageable:

1. Pin exact versions for stability. torch==2.1.0 not torch>=2. Provider SDKs change APIs between minor versions.

2. Split by concern. Keep requirements.txt for production dependencies and requirements-dev.txt for testing tools. This stops pytest and black from inflating your production image.

3. Use a .dockerignore that excludes model files. If you have any model weights in your repo (for tests), make sure they are in .dockerignore.

# .dockerignore
__pycache__/
*.pyc
.pytest_cache/
.env
*.gguf
*.bin
models/

Working example

A complete multi-stage Dockerfile for an LLM API that calls an external provider (no local model, no GPU requirement):

FROM python:3.11-slim AS builder
WORKDIR /build
COPY requirements.txt .
RUN pip install --no-cache-dir --prefix=/install -r requirements.txt

FROM python:3.11-slim AS runtime
COPY --from=builder /install /usr/local

# Security: no root at runtime
RUN useradd --create-home --shell /bin/bash appuser
WORKDIR /app
COPY src/ ./src/
RUN chown -R appuser:appuser /app
USER appuser

# Health check endpoint
HEALTHCHECK --interval=30s --timeout=10s --start-period=10s --retries=3 \
    CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')"

EXPOSE 8000
CMD ["python", "-m", "uvicorn", "src.main:app", "--host", "0.0.0.0", "--port", "8000"]

Build with docker build --target runtime -t myapp:latest . to ensure only the runtime stage is tagged and shipped.

Common mistakes

Putting model weights in the image. Every image rebuild redownloads gigabytes. The image cannot be deployed to environments without the model baked in. Use volumes or a separate model service.

Running as root. Many LLM libraries default-assume they can write to home directories. Running as a non-root user with explicit write permissions to the paths you need is a small step that removes a whole class of security exposure.

Not pinning the CUDA version. A CUDA version mismatch between image and host produces cryptic errors or silent CPU-only fallback. Always pin and document the required driver version in your README.

Using the full nvidia/cuda:*-devel image in production. The devel image includes compilers and headers needed to build CUDA extensions. At runtime you only need the runtime image — it is significantly smaller.

Installing dev tools in the production layer. pip install pytest mypy black in your production Dockerfile adds 200+ MB with no runtime value.

Try it yourself

1. Take an existing Python project and write a multi-stage Dockerfile that separates dependency installation from the final runtime image. Measure the size difference with docker images.
2. Add a .dockerignore file and verify that your build context drops to under 1 MB using docker build --no-cache 2>&1 | grep "Sending build context".
3. Add a HEALTHCHECK instruction to your Dockerfile and test that docker inspect reports health status after startup.