{"title":"Docker for AI Engineers","description":null,"content":"You train a model on your laptop. It works perfectly. You hand it off to your colleague. It crashes.\n\n\"It works on my machine\" is the unofficial motto of software development. In AI, this problem is worse. Your environment has CUDA 11.8, but the server has 12.1. You installed PyTorch with pip, someone else used conda. Your transformers library is version 4.35, theirs is 4.28.\n\nTraining notebooks that ran fine last month now fail because a dependency updated silently. Production deployments break because the inference server has different package versions than your development machine.\n\nThis is where Docker comes in.\n\n---\n\n# What is Docker?\n\nDocker packages your application and all its dependencies into a single, portable unit called a container.\n\nThink of it like shipping containers in the real world. Before standardized containers, loading a ship was chaos. Every item had different dimensions, weights, and handling requirements. The invention of standard shipping containers revolutionized global trade because any container fits on any ship, any truck, any crane.\n\nDocker does the same for software. Your AI application, with all its dependencies, CUDA libraries, Python packages, and model weights, gets packaged into a standardized container that runs identically everywhere.\n\n\n```shell\nWithout Docker: With Docker:\n+-----------------+ +-----------------+\n| Your Machine | | Container |\n+-----------------+ +-----------------+\n| Python 3.11 | | Python 3.11 |\n| CUDA 11.8 | Varies | CUDA 11.8 | Identical\n| PyTorch 2.1 | per machine | PyTorch 2.1 | everywhere\n| transformers | | transformers |\n+-----------------+ +-----------------+\n | |\n v v\n+-----------------+ +-----------------+\n| Server | | Server |\n+-----------------+ +-----------------+\n| Python 3.9 X | | Runs container |\n| CUDA 12.1 X | Different | exactly as-is | Same\n| PyTorch 1.13 X | versions | | behavior\n| Missing libs X | | |\n+-----------------+ +-----------------+\n```\n\n\n---\n\n# Why Docker Matters for AI\n\nAI projects have dependency problems that traditional software does not face.\n\n### 1. Complex Dependency Chains\n\nA simple LLM inference script might need:\n\n- Python 3.10 or 3.11 (not 3.9, not 3.12)\n- PyTorch compiled for your specific CUDA version\n- CUDA toolkit and cuDNN libraries\n- Specific versions of transformers, tokenizers, safetensors\n- Model weights (potentially gigabytes)\n\nEach of these has its own dependencies. PyTorch 2.1 needs specific CUDA versions. Transformers requires specific tokenizers versions. One mismatch and nothing works.\n\n### 2. GPU Driver Complexity\n\nCUDA adds another layer. Your code needs:\n\n- NVIDIA drivers on the host machine\n- CUDA toolkit in your environment\n- cuDNN libraries matching your CUDA version\n- PyTorch compiled for that exact CUDA version\n\nGetting these aligned is tedious. Docker simplifies this by bundling everything except the base drivers.\n\n### 3. Reproducibility Requirements\n\nAI research demands reproducibility. \"We used PyTorch\" is not enough. You need exact versions of every library to reproduce results. Docker captures this automatically.\n\n### 4. Deployment Consistency\n\nThe model that works in your Jupyter notebook needs to work in production. Docker ensures the inference environment matches your development environment exactly.\n\n---\n\n# Docker Fundamentals\n\nBefore diving into AI-specific usage, let us cover Docker basics.\n\n### Images and Containers\n\nA Docker **image** is a template. It contains your operating system, dependencies, and application code. Images are immutable. You build them once and reuse them.\n\nA Docker **container** is a running instance of an image. You can run multiple containers from the same image. Containers are isolated from each other and from the host system.\n\n\n```mermaid\nflowchart TD\n I[Docker Image
blueprint]:::purple\n\n I -->|\"docker run\"| C1[Container 1
running]:::primary\n I -->|\"docker run\"| C2[Container 2
running]:::primary\n\n C1 --> T1[Inference on
request batch 1]:::orange\n C2 --> T2[Training on
dataset A]:::orange\n\n classDef primary fill:#00ceff,stroke:#000,color:#000\n classDef purple fill:#9775fa,stroke:#000,color:#000\n classDef orange fill:#ffa94d,stroke:#000,color:#000\n```\n\n\n### Dockerfiles\n\nA Dockerfile is a recipe for building an image. It specifies the base image, dependencies to install, files to copy, and commands to run.\n\n\n```dockerfile\n# Base image with Python\nFROM python:3.11-slim\n\n# Set working directory\nWORKDIR /app\n\n# Copy requirements first (for caching)\nCOPY requirements.txt .\n\n# Install dependencies\nRUN pip install --no-cache-dir -r requirements.txt\n\n# Copy application code\nCOPY . .\n\n# Command to run when container starts\nCMD [\"python\", \"inference.py\"]\n```\n\n\n### Basic Commands\n\n\n```shell\n# Build an image from a Dockerfile\ndocker build -t my-ai-app .\n\n# Run a container from an image\ndocker run my-ai-app\n\n# Run interactively with terminal\ndocker run -it my-ai-app bash\n\n# List running containers\ndocker ps\n\n# Stop a container\ndocker stop \n\n# List images\ndocker images\n```\n\n\n---\n\n# Docker for GPU Workloads\n\nAI workloads typically need GPU access. Docker handles this through the NVIDIA Container Toolkit.\n\n### How GPU Access Works\n\n\n```mermaid\nflowchart TD\n subgraph Container[\"Docker Container\"]\n A1[Your AI Application]:::primary\n A2[PyTorch]:::secondary\n A3[Model weights]:::secondary\n A4[CUDA Toolkit + cuDNN
in image]:::orange\n\n A1 --> A2\n A1 --> A3\n A2 --> A4\n end\n\n Container -->|\"nvidia-docker
runtime\"| Host\n\n subgraph Host[\"Host Machine\"]\n H1[NVIDIA Drivers]:::purple\n H2[Physical GPU]:::green\n\n H1 --> H2\n end\n\n classDef primary fill:#00ceff,stroke:#000,color:#000\n classDef secondary fill:#38d9a9,stroke:#000,color:#000\n classDef orange fill:#ffa94d,stroke:#000,color:#000\n classDef purple fill:#9775fa,stroke:#000,color:#000\n classDef green fill:#69db7c,stroke:#000,color:#000\n```\n\n\nThe NVIDIA drivers stay on the host machine. Everything else, CUDA toolkit, cuDNN, your application, goes in the container.\n\n### Setting Up GPU Support\n\nFirst, install the NVIDIA Container Toolkit on your host:\n\n\n```shell\n# Add NVIDIA package repository\ndistribution=$(. /etc/os-release;echo $ID$VERSION_ID)\ncurl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -\ncurl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | \\\n sudo tee /etc/apt/sources.list.d/nvidia-docker.list\n\n# Install toolkit\nsudo apt-get update\nsudo apt-get install -y nvidia-container-toolkit\n\n# Restart Docker\nsudo systemctl restart docker\n```\n\n\n### Running GPU Containers\n\nUse the `--gpus` flag:\n\n\n```shell\n# Use all GPUs\ndocker run --gpus all my-ai-app\n\n# Use specific GPUs\ndocker run --gpus '\"device=0,1\"' my-ai-app\n\n# Use a specific number of GPUs\ndocker run --gpus 2 my-ai-app\n```\n\n\n### NVIDIA Base Images\n\nNVIDIA provides pre-built images with CUDA and cuDNN:\n\n\n```dockerfile\n# Start from NVIDIA's CUDA image\nFROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04\n\n# Install Python\nRUN apt-get update && apt-get install -y python3 python3-pip\n\n# Install PyTorch with CUDA support\nRUN pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118\n```\n\n\nThese images handle the CUDA/cuDNN complexity for you. Choose the version that matches your needs:\n\n- `runtime`: For inference (smaller image)\n- `devel`: For training and compilation (includes headers)\n\n---\n\n# Building AI Application Images\n\nLet us walk through building Docker images for common AI workflows.\n\n### Inference Server\n\nA typical inference server needs the model, tokenizer, and serving framework:\n\n\n```dockerfile\nFROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04\n\n# Install system dependencies\nRUN apt-get update && apt-get install -y \\\n python3 \\\n python3-pip \\\n && rm -rf /var/lib/apt/lists/*\n\nWORKDIR /app\n\n# Install Python dependencies\nCOPY requirements.txt .\nRUN pip3 install --no-cache-dir -r requirements.txt\n\n# Copy model and application code\nCOPY model/ ./model/\nCOPY app.py .\n\n# Expose port for API\nEXPOSE 8000\n\n# Run the server\nCMD [\"python3\", \"app.py\"]\n```\n\n\n**requirements.txt:**\n\n\n```shell\ntorch==2.1.0\ntransformers==4.35.0\nfastapi==0.104.0\nuvicorn==0.24.0\n```\n\n\n### Training Environment\n\nTraining images need more tools and often mount data from the host:\n\n\n```dockerfile\nFROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04\n\nRUN apt-get update && apt-get install -y \\\n python3 \\\n python3-pip \\\n git \\\n && rm -rf /var/lib/apt/lists/*\n\nWORKDIR /workspace\n\n# Install training dependencies\nCOPY requirements.txt .\nRUN pip3 install --no-cache-dir -r requirements.txt\n\n# Copy training code\nCOPY train.py .\nCOPY config/ ./config/\n\n# Data will be mounted at runtime\nVOLUME /data\n\nCMD [\"python3\", \"train.py\"]\n```\n\n\nRun with data mounted:\n\n\n```shell\ndocker run --gpus all -v /path/to/data:/data my-training-image\n```\n\n\n### Multi-Stage Builds\n\nUse multi-stage builds to keep final images small:\n\n\n```dockerfile\n# Stage 1: Build dependencies\nFROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 AS builder\n\nRUN apt-get update && apt-get install -y python3 python3-pip\nCOPY requirements.txt .\nRUN pip3 install --no-cache-dir -r requirements.txt --target=/install\n\n# Stage 2: Runtime image\nFROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04\n\nRUN apt-get update && apt-get install -y python3 && rm -rf /var/lib/apt/lists/*\n\n# Copy only installed packages\nCOPY --from=builder /install /usr/local/lib/python3.10/dist-packages\n\nWORKDIR /app\nCOPY app.py .\n\nCMD [\"python3\", \"app.py\"]\n```\n\n\nThe first stage compiles dependencies. The second stage copies only what is needed, resulting in a smaller final image.\n\n---\n\n# Managing Model Weights\n\nLarge models create challenges. A 7B parameter model might be 14GB. You do not want to rebuild images every time you update weights.\n\n### Option 1: Mount Weights at Runtime\n\nKeep weights outside the image and mount them:\n\n\n```dockerfile\nFROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04\n# ... install dependencies ...\n\n# Model will be mounted here\nVOLUME /models\n\nENV MODEL_PATH=/models/llama-7b\n\nCMD [\"python3\", \"serve.py\"]\n```\n\n\n\n```shell\ndocker run --gpus all -v /path/to/models:/models my-inference-server\n```\n\n\n**Pros:** Fast image builds, easy weight updates, shared storage across containers **Cons:** Requires managing files separately, less portable\n\n### Option 2: Download at Startup\n\nDownload weights when the container starts:\n\n\n```python\n# entrypoint.py\nimport os\nfrom huggingface_hub import snapshot_download\n\nMODEL_ID = os.environ.get(\"MODEL_ID\", \"meta-llama/Llama-2-7b-hf\")\nMODEL_PATH = \"/models/current\"\n\nif not os.path.exists(MODEL_PATH):\n snapshot_download(repo_id=MODEL_ID, local_dir=MODEL_PATH)\n\n# Then start the server\nos.system(\"python3 serve.py\")\n```\n\n\n**Pros:** Always gets latest weights, simpler deployment **Cons:** Slow startup, needs network access, download on every new container\n\n### Option 3: Bake into Image\n\nFor production with stable weights, include them in the image:\n\n\n```dockerfile\nFROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04\n# ... install dependencies ...\n\n# Copy model weights into image\nCOPY models/llama-7b /models/llama-7b\n\nCMD [\"python3\", \"serve.py\"]\n```\n\n\n**Pros:** Self-contained, no external dependencies at runtime **Cons:** Huge images, slow builds, every weight change requires rebuild\n\n### Recommendation\n\nUse mounted volumes for development and downloaded weights for ephemeral cloud deployments. Bake weights into images only for stable production deployments where image size is not a concern.\n\n---\n\n# Docker Compose for ML Pipelines\n\nReal ML systems have multiple services. Docker Compose orchestrates them together.\n\n### Example: RAG System\n\nA RAG system might have:\n\n- Vector database (Qdrant/Milvus)\n- Embedding service\n- LLM inference service\n- API gateway\n\n\n```yaml\n# docker-compose.yml\nversion: '3.8'\n\nservices:\n vector-db:\n image: qdrant/qdrant:latest\n ports:\n - \"6333:6333\"\n volumes:\n - qdrant_data:/qdrant/storage\n\n embedding-service:\n build: ./embedding-service\n environment:\n - MODEL_NAME=sentence-transformers/all-MiniLM-L6-v2\n deploy:\n resources:\n reservations:\n devices:\n - driver: nvidia\n count: 1\n capabilities: [gpu]\n\n llm-service:\n build: ./llm-service\n environment:\n - MODEL_PATH=/models/llama-7b\n volumes:\n - ./models:/models\n deploy:\n resources:\n reservations:\n devices:\n - driver: nvidia\n count: 1\n capabilities: [gpu]\n\n api:\n build: ./api\n ports:\n - \"8000:8000\"\n depends_on:\n - vector-db\n - embedding-service\n - llm-service\n\nvolumes:\n qdrant_data:\n```\n\n\nStart everything with:\n\n\n```shell\ndocker compose up\n```\n\n\n### GPU Allocation in Compose\n\nSpecify GPU requirements per service:\n\n\n```yaml\nservices:\n training:\n build: .\n deploy:\n resources:\n reservations:\n devices:\n - driver: nvidia\n count: 2\n capabilities: [gpu]\n```\n\n\n---\n\n# Best Practices\n\n### 1. Pin All Versions\n\nNever use `latest` tags or unpinned packages:\n\n\n```dockerfile\n# Bad\nFROM python:3\nRUN pip install torch transformers\n\n# Good\nFROM python:3.11.6-slim\nRUN pip install torch==2.1.0 transformers==4.35.0\n```\n\n\n### 2. Order Layers by Change Frequency\n\nDocker caches layers. Put rarely-changing steps first:\n\n\n```dockerfile\n# Good: Install system deps first (rarely change)\nRUN apt-get update && apt-get install -y python3 python3-pip\n\n# Then Python deps (change sometimes)\nCOPY requirements.txt .\nRUN pip install -r requirements.txt\n\n# Finally app code (changes often)\nCOPY app.py .\n```\n\n\n### 3. Use .dockerignore\n\nExclude unnecessary files from the build context:\n\n\n```shell\n# .dockerignore\n.git\n__pycache__\n*.pyc\n.env\ndata/\n*.pt\n*.safetensors\nwandb/\n```\n\n\n### 4. Set Memory Limits\n\nAI workloads can consume all available memory. Set limits:\n\n\n```shell\ndocker run --gpus all --memory=32g --memory-swap=32g my-training-image\n```\n\n\n### 5. Use Health Checks\n\nEnsure your service is actually ready:\n\n\n```dockerfile\nHEALTHCHECK --interval=30s --timeout=10s --start-period=60s \\\n CMD curl -f http://localhost:8000/health || exit 1\n```\n\n\n---\n\n# Common Pitfalls\n\n\n| ##### Pitfall | ##### Why It Happens | ##### How to Avoid |\n| --- | --- | --- |\n| CUDA version mismatch | Base image CUDA differs from PyTorch build | Match CUDA versions explicitly |\n| Image too large | Including training data or unnecessary files | Use .dockerignore, multi-stage builds |\n| Slow builds | Not leveraging layer caching | Order Dockerfile by change frequency |\n| Out of GPU memory | No memory limits set | Set explicit memory constraints |\n| Model not found at runtime | Weights not included or mounted | Verify mount paths, use absolute paths |\n| Permission errors | Running as root vs. non-root user | Match user IDs between host and container |\n| Network timeouts downloading models | No retry logic in entrypoint | Add retry logic or pre-download weights |\n\n\n---\n\n# Development vs. Production\n\nDifferent stages need different configurations:\n\n### Development\n\n\n```dockerfile\n# dev.Dockerfile\nFROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04\n\n# Include dev tools\nRUN apt-get update && apt-get install -y \\\n python3 python3-pip git vim curl\n\n# Install with dev dependencies\nCOPY requirements-dev.txt .\nRUN pip install -r requirements-dev.txt\n\n# Mount code at runtime for hot reload\nWORKDIR /app\nCMD [\"uvicorn\", \"app:app\", \"--reload\", \"--host\", \"0.0.0.0\"]\n```\n\n\nRun with code mounted:\n\n\n```shell\ndocker run --gpus all -v $(pwd):/app -p 8000:8000 my-dev-image\n```\n\n\n### Production\n\n\n```dockerfile\n# Dockerfile\nFROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04\n\n# Minimal dependencies\nRUN apt-get update && apt-get install -y python3 python3-pip \\\n && rm -rf /var/lib/apt/lists/*\n\n# Production deps only\nCOPY requirements.txt .\nRUN pip install --no-cache-dir -r requirements.txt\n\n# Copy code (not mounted)\nCOPY app.py .\n\n# Non-root user\nRUN useradd -m appuser\nUSER appuser\n\nCMD [\"uvicorn\", \"app:app\", \"--host\", \"0.0.0.0\", \"--workers\", \"4\"]\n```\n\n\n---\n\n# Useful Tools\n\n### 1. NVIDIA Container Toolkit\n\nEssential for GPU access. Install it once on any machine that runs GPU containers.\n\n### 2. Docker BuildKit\n\nFaster builds with better caching:\n\n\n```shell\nDOCKER_BUILDKIT=1 docker build -t my-image .\n```\n\n\n### 3. dive\n\nAnalyze image layers and find bloat:\n\n\n```shell\ndive my-image:latest\n```\n\n\n### 4. Hadolint\n\nLint your Dockerfiles:\n\n\n```shell\nhadolint Dockerfile\n```\n\n\n---\n\n# Key Takeaways\n\n1. **Docker solves \"works on my machine\" for AI**. Package your environment with your code.\n2. **GPUs work in Docker with NVIDIA Container Toolkit**. The toolkit bridges containers to host GPU drivers.\n3. **Use NVIDIA base images**. They handle CUDA and cuDNN complexity.\n4. **Keep models separate from images**. Mount weights or download at startup for large models.\n5. **Pin everything**. Exact versions ensure reproducibility.\n6. **Layer order matters**. Put stable layers first for faster builds.\n7. **Use Compose for multi-service systems**. RAG pipelines, training clusters, and inference services benefit from orchestration.\n\n---\n\n# References\n\n- [Docker Documentation](undefined) - Official Docker guides and reference\n- [NVIDIA Container Toolkit](undefined) - GPU container setup and configuration\n- [Best Practices for Writing Dockerfiles](undefined) - Official Docker optimization guide\n- [PyTorch Docker Images](undefined) - Official PyTorch containers\n- [Hugging Face Docker Spaces](undefined) - Running HF models in containers","pageType":"ai-engineering"}

Docker for AI Engineers

Ashish Pratap Singh

Get Premium