ModelRefs / Docker & Containerization for AI — Tutorial
Docker & Containerization for AI — Tutorial
Reproducible ML environments, GPU-enabled containers, multi-stage builds, and Docker Compose stacks. Covers Why containers for ML.
Overview
Reproducible ML environments, GPU-enabled containers, multi-stage builds, and Docker Compose stacks
Level: Advanced. Estimated reading time: 35 minutes.
Why containers for ML
"It works on my machine" is the original ML reproducibility problem. Model behaviour depends on Python version, CUDA version, library versions, and system libraries — all of which vary between development laptops, CI servers, and GPU clusters.
Docker packages your code, dependencies, and runtime into an immutable image. Every environment — dev, CI, staging, production — runs the exact same container. No dependency drift, no "I forgot to install sentencepiece".
For ML specifically: GPU containers: NVIDIA provides base images (nvidia/cuda:12.x-cudnn9-runtime-ubuntu22.04) with the correct CUDA/cuDNN versions pre-installed. Your model code just runs — no manual CUDA setup.
Reproducibility: tag your image with git SHA or model version. You can reproduce any past inference run by pulling that exact image.
Portability: push to any registry (Docker Hub, ECR, GCR, GHCR). Pull and run on any NVIDIA GPU machine with the container runtime installed.
Writing a production Dockerfile for ML
A good ML Dockerfile has three concerns: correct base image, efficient layer caching, and minimal image size.
Base image: use NVIDIA CUDA base for GPU models; python:3.12-slim for CPU-only.
Layer caching: Docker caches each layer. Copy and install requirements first (rarely changes), then copy application code (changes often). This way pip install only reruns when requirements.txt changes.
Multi-stage builds: use a build stage to install heavy build dependencies (gcc, cmake), then copy only the resulting Python packages to a slim runtime stage. Reduces GPU image size from 8GB to 3GB.
.dockerignore: exclude __pycache__, .git, *.pyc, *.ipynb, .venv, data/, checkpoints/ — don't ship training data or model weights in the image (mount them as volumes).
Non-root user: run as a non-root user for security. Add user myapp and switch with USER myapp before the CMD.
Health check: HEALTHCHECK CMD curl --fail http://localhost:8000/health || exit 1 — Docker will restart the container if the health check fails.
Docker Compose for the full ML stack
A production ML stack needs more than just the inference server: a model store, a cache, a monitoring stack, and often a message queue for async jobs.
Docker Compose orchestrates multi-container stacks on a single machine (or small cluster). For ML dev/staging:
Services: api (FastAPI inference), redis (result caching, job queue), prometheus (metrics), grafana (dashboards), minio (S3-compatible model store).
GPU access in Compose: add the deploy section with resources.reservations.devices targeting nvidia GPUs.
Volumes: mount model weights from the host (./models:/app/models) — don't rebuild the image when weights change.
Networking: all services join a default bridge network, addressable by service name (redis://redis:6379 from the api container).
Prod vs dev overrides: docker-compose.yml (base) + docker-compose.override.yml (dev with volume mounts, hot reload) + docker-compose.prod.yml (replicas, resource limits).
Continue your research
Use these connected ModelRefs sections to compare alternatives, inspect implementation paths, and review the evidence and governance boundaries relevant to Docker & Containerization for AI — Tutorial.