An Introduction to Docker for Machine Learning

What is Docker?

Docker is an open-source platform that enables you to package, distribute, and run applications inside lightweight, isolated environments called containers. A container packages your code alongside all its dependencies — system libraries, Python runtimes, CUDA drivers, and configuration files — ensuring deterministic behavior regardless of the host environment.

Key Distinction: Unlike traditional Virtual Machines (VMs) that virtualize entire hardware stacks and boot independent guest operating systems, Docker containers share the host operating system kernel. This makes them orders of magnitude lighter, faster to initialize (milliseconds vs. minutes), and significantly more memory-efficient.

+-----------------------------------+     +-----------------------------------+
|            Docker Container       |     |          Virtual Machine          |
|  +-----------------------------+  |     |  +-----------------------------+  |
|  | App + Python + CUDA Deps    |  |     |  | App + Python + CUDA Deps    |  |
|  +-----------------------------+  |     |  +-----------------------------+  |
|  | Shared Host OS Kernel       |  |     |  | Guest OS (Full Kernel)      |  |
|  +-----------------------------+  |     |  +-----------------------------+  |
|  | Host Hardware (GPU / CPU)   |  |     |  | Hypervisor + Host OS        |  |
+-----------------------------------+     +-----------------------------------+

In Machine Learning, containerization is indispensable. A model trained on Python 3.10 with CUDA 12.1 and specific pinned versions of torch and transformers will execute identically on your local workstation, an on-prem GPU cluster, or a cloud instance.

Core Concepts

Before building containers, it helps to understand four foundational building blocks:

Running Your First Container

Assuming Docker is already installed and running on your system, you can verify your installation:

# Pull and run the official hello-world verification image
docker run hello-world

For a more practical service example using Nginx:

# Run Nginx in detached mode (-d), mapping host port 8080 to container port 80
docker run -d -p 8080:80 --name web-server nginx

Navigate to http://localhost:8080 in your browser to confirm the server is live. Essential container management commands include:

# List active containers
docker ps

# List all containers (including stopped ones)
docker ps -a

# Stream live container logs
docker logs -f web-server

# Open an interactive Bash shell inside a running container
docker exec -it web-server bash

# Gracefully stop and delete the container
docker stop web-server && docker rm web-server

Writing an Optimized Dockerfile for Python ML

A Dockerfile is a text document containing instructions to assemble an image. Docker builds images as a series of cached layers. Ordering instructions from least frequently changed to most frequently changed drastically accelerates build times.

Here is a standard, production-ready Dockerfile for a Python ML workload:

# 1. Use an official lightweight Python runtime
FROM python:3.11-slim

# 2. Set environment variables to prevent Python from writing .pyc files and buffer stdout
ENV PYTHONDONTWRITEBYTECODE=1 \
    PYTHONUNBUFFERED=1

# 3. Establish working directory
WORKDIR /app

# 4. Install system dependencies if required (e.g., git, curl, build tools)
RUN apt-get update && apt-get install -y --no-install-recommends \
    build-essential \
    && rm -rf /var/lib/apt/lists/*

# 5. Copy requirements FIRST to leverage Docker layer caching
COPY requirements.txt .

# 6. Install Python packages without caching wheels inside the image
RUN pip install --no-cache-dir --upgrade pip && \
    pip install --no-cache-dir -r requirements.txt

# 7. Copy source code (changes frequently)
COPY . .

# 8. Expose documentation port
EXPOSE 8000

# 9. Default command
CMD ["python", "train.py"]

Pro Tip (Layer Caching): Always copy requirements.txt and run pip install before copying the application source code (COPY . .). This ensures Docker reuses the installed packages layer across builds unless your dependencies explicitly change.

Build the image and execute it while mounting a local dataset directory:

# Build and tag image
docker build -t ml-pipeline:latest .

# Run interactively, mounting local data into the container
docker run -it --rm \
  -v $(pwd)/data:/app/data \
  -v $(pwd)/checkpoints:/app/checkpoints \
  ml-pipeline:latest

The -v parameter creates a bind mount: changes made in /app/checkpoints inside the container immediately persist on your host disk, preventing model loss if the container terminates.

GPU Acceleration for Deep Learning

When training Deep Learning models with PyTorch or TensorFlow, you need GPU acceleration.

1. The Recommended Approach: Official Framework Images

Instead of installing CUDA drivers from scratch on a raw Ubuntu image, use the official PyTorch base images. They come with compatible CUDA, cuDNN, and Python runtime binaries pre-configured:

# Official PyTorch image with CUDA 12.1 and cuDNN 8 pre-installed
FROM pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime

WORKDIR /app

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

COPY . .

CMD ["python", "train.py"]

2. The Host Prerequisite: NVIDIA Container Toolkit

To expose physical GPUs to Docker containers, install the NVIDIA Container Toolkit on your host machine. Once installed, pass the --gpus all flag at runtime:

# Verify GPU visibility inside the container
docker run --gpus all --rm pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime nvidia-smi

3. Critical ML Gotcha: Shared Memory (--shm-size)

When training PyTorch models with multi-process data loading (DataLoader(..., num_workers > 0)), worker processes exchange tensors via shared memory (/dev/shm).

By default, Docker allocates only 64 MB to shared memory. If your batch size or worker count exceeds this limit, training crashes with a cryptic error:

RuntimeError: DataLoader worker (pid 42) is killed by signal: Bus error.

The Solution: Always allocate sufficient shared memory (e.g., 8 GB or more) when launching training containers:

docker run --gpus all --shm-size="8g" -it --rm \
  -v $(pwd)/data:/app/data \
  -v $(pwd)/checkpoints:/app/checkpoints \
  ml-pipeline:gpu

Multi-Service Systems with Docker Compose

Modern Machine Learning architectures rarely operate in isolation. A realistic production system often coordinates:

1. A Training Worker that processes datasets and outputs model artifacts to a shared volume.
2. A FastAPI Serving API that loads the trained weights to perform real-time inference.
3. A Redis Cache / Queue to manage request throughput and background job scheduling.

Docker Compose orchestrates all these components through a declarative docker-compose.yml file:

services:

  # ── Training Service (GPU accelerated) ────────────────────────
  trainer:
    build:
      context: .
      dockerfile: Dockerfile.train
    shm_size: '8gb'
    volumes:
      - ./data:/app/data:ro           # Read-only dataset mount
      - ./checkpoints:/app/checkpoints # Persist trained weights
    environment:
      EPOCHS: "100"
      BATCH_SIZE: "64"
      LEARNING_RATE: "1e-4"
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]

  # ── Model Serving API (FastAPI) ──────────────────────────────
  inference-api:
    build:
      context: .
      dockerfile: Dockerfile.serve
    ports:
      - "8000:8000"
    volumes:
      - ./checkpoints:/app/checkpoints:ro
    environment:
      MODEL_PATH: "/app/checkpoints/best_model.pt"
      REDIS_HOST: "redis-cache"
    depends_on:
      - redis-cache
    restart: unless-stopped

  # ── In-Memory Cache & Message Broker ─────────────────────────
  redis-cache:
    image: redis:7-alpine
    ports:
      - "6379:6379"
    volumes:
      - redis_data:/data

volumes:
  redis_data:

Essential Docker Compose Workflow

# Build images and start all services in the background
docker compose up -d --build

# Monitor live logs for the API service
docker compose logs -f inference-api

# Execute an ad-hoc evaluation script inside the trainer environment
docker compose run --rm trainer python evaluate.py

# Gracefully shut down containers while keeping persistent volumes intact
docker compose down

# Shut down containers AND remove named data volumes
docker compose down -v

Cheat Sheet: Essential Docker Commands

Summary & Best Practices

Adopting containerization in your AI workflows eliminates environmental entropy. To maximize reliability:

1. Leverage Framework Images: Use pytorch/pytorch or tensorflow/tensorflow base images to avoid CUDA/cuDNN mismatch headaches.
2. Optimize Layer Caching: Separate dependency installation (requirements.txt) from application code copies.
3. Decouple Data & Artifacts: Never bake datasets or multi-gigabyte model weights into image layers; always use volumes or cloud object storage (S3/GCS).
4. Configure Shared Memory: Set --shm-size whenever using multi-worker data loaders in PyTorch.
5. Orchestrate with Compose: Define multi-service pipelines (training, inference, caching, tracking) in reproducible declarative configuration files.