Tutorial 15: Docker for Data Science and Machine Learning

Table of Contents

Introduction

Reproducibility is one of the greatest challenges in data science and machine learning. The notorious "it works on my machine" problem is amplified when projects depend on specific versions of Python, CUDA, cuDNN, system libraries, and dozens of Python packages. A model that trains perfectly on your laptop may fail spectacularly on a colleague's workstation or in a cloud environment.

Docker solves this by packaging your code, dependencies, and runtime environment into a portable, self-contained unit called a container. With Docker, you can guarantee that your ML pipeline runs identically everywhere — from local development to CI/CD to production serving.

This tutorial covers everything you need to containerize ML workflows: writing efficient Dockerfiles, leveraging multi-stage builds, enabling GPU acceleration, orchestrating multi-service stacks with Docker Compose, and integrating containers into CI/CD pipelines.

Prerequisites

• Docker Engine 24.0 or later installed
• Docker Compose v2 installed
• Basic familiarity with Linux command line
• Python 3.9+ and pip
• (Optional) NVIDIA GPU with drivers installed for GPU sections

# Verify Docker installation
docker --version
docker compose version

Docker Fundamentals for ML Engineers

Key Concepts

Why Docker Matters for ML

| Challenge | Docker Solution |

|-----------|----------------|

| Dependency conflicts | Isolated environments per project |

| Python version mismatches | Specify exact Python version in image |

| CUDA/cuDNN compatibility | Use NVIDIA base images with matching versions |

| Reproducible training | Same image produces same results anywhere |

| Deployment consistency | Dev, staging, and production use the same image |

Writing Dockerfiles for ML Projects

Basic ML Dockerfile

# Use an official Python runtime as the base image
FROM python:3.11-slim

Set environment variables
ENV PYTHONDONTWRITEBYTECODE=1 \
    PYTHONUNBUFFERED=1 \
    PIPNOCACHEDIR=1 \

    PIPDISABLEPIPVERSIONCHECK=1


Set the working directory
WORKDIR /app

Install system dependencies (common for ML libraries)
RUN apt-get update && apt-get install -y --no-install-recommends \
    build-essential \
    libgomp1 \
    libglib2.0-0 \
    libsm6 \
    libxext6 \
    libxrender-dev \
    libgl1-mesa-glx \
    curl \
    && rm -rf /var/lib/apt/lists/


Copy and install Python dependencies first (for layer caching)
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

Copy the application code
COPY . .

Expose the port for the API
EXPOSE 8000

Default command
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]

# requirements.txt
numpy==1.26.4
pandas==2.2.0
scikit-learn==1.4.0
xgboost==2.0.3
lightgbm==4.2.0
fastapi==0.109.0
uvicorn==0.27.0
pydantic==2.5.3
joblib==1.3.2
mlflow==2.10.0

# Build the image
docker build -t ml-inference-api:v1.0 .

Run the container
docker run -d \
    --name ml-api \
    -p 8000:8000 \
    -v $(pwd)/models:/app/models \
    ml-inference-api:v1.0

Check logs
docker logs ml-api

Test the API
curl http://localhost:8000/health

# ============================================
Stage 1: Builder - install dependencies
============================================
FROM python:3.11-slim AS builder

WORKDIR /build

Install build dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
    build-essential \
    gcc \
    g++ \
    gfortran \
    libopenblas-dev \
    && rm -rf /var/lib/apt/lists/

Create a virtual environment
RUN python -m venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"

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

============================================
Stage 2: Training - run model training
============================================
FROM builder AS trainer

WORKDIR /train
COPY . .
RUN python train.py --output /train/models/model.pkl

============================================
Stage 3: Production - minimal runtime image
============================================
FROM python:3.11-slim AS production

Install only runtime system dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
    libgomp1 \
    libglib2.0-0 \
    curl \
    && rm -rf /var/lib/apt/lists/


Copy the virtual environment from builder
COPY --from=builder /opt/venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"

Copy application code
WORKDIR /app
COPY app.py .
COPY utils/ ./utils/

Copy trained model from trainer stage
COPY --from=trainer /train/models/model.pkl ./models/model.pkl

Create a non-root user for security
RUN useradd --create-home --shell /bin/bash mluser
USER mluser

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

EXPOSE 8000
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000", "--workers", "4"]

# Single-stage build (typically 2-3 GB for ML)
docker images ml-api-single

Multi-stage build (typically 500 MB - 1 GB)
docker images ml-api-multi

# Install NVIDIA Container Toolkit
distribution=$(. /etc/os-release; echo $ID$VERSIONID)
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | \
    sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
curl -s -L "https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list" | \
    sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
    sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

sudo apt-get update
sudo apt-get install -y nvidia-container-toolkit
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker

Verify GPU access inside container
docker run --rm --gpus all nvidia/cuda:12.3.1-base-ubuntu22.04 nvidia-smi

# Use NVIDIA CUDA base image
FROM nvidia/cuda:12.3.1-cudnn9-runtime-ubuntu22.04

Set environment variables
ENV DEBIANFRONTEND=noninteractive \

    PYTHONDONTWRITEBYTECODE=1 \
    PYTHONUNBUFFERED=1

Install Python and system dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3.11 \
    python3.11-venv \
    python3-pip \
    && ln -sf /usr/bin/python3.11 /usr/bin/python \
    && rm -rf /var/lib/apt/lists/

WORKDIR /app

Install PyTorch with CUDA support
RUN pip install --no-cache-dir \
    torch==2.2.0 \
    torchvision==0.17.0 \
    --index-url https://download.pytorch.org/whl/cu121

Install other ML dependencies
COPY requirements-gpu.txt .
RUN pip install --no-cache-dir -r requirements-gpu.txt

COPY . .

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

Running with GPU Access

# Run with all GPUs
docker run --gpus all -v $(pwd)/data:/app/data ml-gpu-training:v1.0

Run with specific GPUs
docker run --gpus '"device=0,1"' ml-gpu-training:v1.0

Run with limited GPU memory
docker run --gpus all --shm-size=8g ml-gpu-training:v1.0

Docker Compose for ML Stacks

A typical ML system involves multiple services. Docker Compose orchestrates them together.

version: '3.8'

services:
  # ML API Service
  ml-api:
    build:
      context: .
      dockerfile: Dockerfile
      target: production
    ports:
"8000:8000"
    volumes:

model-store:/app/models
./logs:/app/logs
    environment:

MODELPATH=/app/models/model.pkl

LOGLEVEL=INFO
REDISURL=redis://redis:6379/0

DBURL=postgresql://mluser:mlpass@postgres:5432/mldb
    dependson:

      redis:
        condition: servicehealthy
      postgres:
        condition: servicehealthy

    deploy:
      resources:
        limits:
          cpus: '4.0'
          memory: 8G
    restart: unless-stopped
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:8000/health"]
      interval: 30s
      timeout: 10s
      retries: 3

  # Redis for caching and feature store
  redis:
    image: redis:7-alpine
    ports:
"6379:6379"
    volumes:

redis-data:/data
    command: redis-server --maxmemory 2gb --maxmemory-policy allkeys-lru

    healthcheck:
      test: ["CMD", "redis-cli", "ping"]
      interval: 10s
      timeout: 5s
      retries: 3

  # PostgreSQL for predictions and metadata
  postgres:
    image: postgres:16-alpine
    environment:
      POSTGRESUSER: mluser
      POSTGRESPASSWORD: mlpass

      POSTGRESDB: mldb
    ports:
"5432:5432"
    volumes:

postgres-data:/var/lib/postgresql/data
./init.sql:/docker-entrypoint-initdb.d/init.sql
    healthcheck:

      test: ["CMD-SHELL", "pgisready -U mluser -d mldb"]

      interval: 10s
      timeout: 5s
      retries: 5

  # MLflow for experiment tracking
  mlflow:
    image: ghcr.io/mlflow/mlflow:v2.10.0
    ports:
"5000:5000"
    volumes:

mlflow-data:/mlflow
    command: >

      mlflow server
      --backend-store-uri sqlite:///mlflow/mlflow.db
      --default-artifact-root /mlflow/artifacts
      --host 0.0.0.0
      --port 5000
    dependson:
postgres

  # Jupyter Notebook for experimentation
  jupyter:
    build:
      context: .
      dockerfile: Dockerfile.jupyter
    ports:
"8888:8888"
    volumes:

./notebooks:/app/notebooks
./data:/app/data
model-store:/app/models
    environment:

JUPYTERTOKEN=mlsecret123

    command: >

      jupyter lab --ip=0.0.0.0 --port=8888
      --no-browser --allow-root
      --NotebookApp.token=${JUPYTERTOKEN:-mlsecret123}

  # Prometheus for monitoring
  prometheus:
    image: prom/prometheus:v2.48.0
    ports:
"9090:9090"
    volumes:

./prometheus.yml:/etc/prometheus/prometheus.yml
prometheus-data:/prometheus

  # Grafana for dashboards
  grafana:
    image: grafana/grafana:10.2.0
    ports:
"3000:3000"
    volumes:

grafana-data:/var/lib/grafana
    environment:

GFSECURITYADMINPASSWORD=admin123


volumes:
  model-store:
  redis-data:
  postgres-data:
  mlflow-data:
  prometheus-data:
  grafana-data:

# Start all services
docker compose up -d

Check status
docker compose ps

View logs for a specific service
docker compose logs -f ml-api

Scale the API service
docker compose up -d --scale ml-api=3

Stop everything
docker compose down

Stop and remove all data
docker compose down -v

# Bind mount: map host directory to container path
Best for: development, source code, datasets
docker run -v /host/data:/container/data ml-training

Named volume: Docker-managed persistent storage
Best for: databases, model artifacts, caches
docker volume create model-store
docker run -v model-store:/app/models ml-api

tmpfs mount: in-memory filesystem
Best for: sensitive data, temporary caches
docker run --tmpfs /tmp:rw,size=2g ml-training

# app.py - Handle model loading from mounted volume
import os
import joblib
from pathlib import Path

MODELDIR = Path(os.getenv('MODELDIR', '/app/models'))

DATADIR = Path(os.getenv('DATADIR', '/app/data'))


def loadlatestmodel():

    """Load the most recently modified model file."""
    modelfiles = sorted(
        MODELDIR.glob('.pkl'),

        key=lambda p: p.stat().stmtime,
        reverse=True
    )
    if not modelfiles:

        raise FileNotFoundError(f"No model files found in {MODELDIR}")

    modelpath = modelfiles[0]
    print(f"Loading model: {modelpath}")

    return joblib.load(modelpath)

def savemodel(model, version):

    """Save a model with versioned filename."""
    MODELDIR.mkdir(parents=True, existok=True)

    path = MODELDIR / f"modelv{version}.pkl"

    joblib.dump(model, path)
    print(f"Model saved to {path}")
    return path

# .env
MODELVERSION=v2.1.0
MODELPATH=/app/models/modelv2.1.0.pkl
LOGLEVEL=INFO

REDISURL=redis://redis:6379/0
APIWORKERS=4

MAXBATCHSIZE=64

# docker-compose.yml
services:
  ml-api:
    envfile:
.env
    environment:

EXTRAVAR=overridevalue  # Overrides .env if same key

# docker-compose.override.yml (automatically loaded, for development)
version: '3.8'
services:
  ml-api:
    build:
      context: .
      target: builder  # Use development stage
    volumes:
./src:/app/src  # Hot reload source code
    environment:

DEBUG=1
LOGLEVEL=DEBUG

    command: uvicorn app:app --host 0.0.0.0 --reload


docker-compose.prod.yml (production overrides)
version: '3.8'
services:
  ml-api:
    build:
      target: production
    deploy:
      replicas: 3
      resources:
        limits:
          cpus: '4.0'
          memory: 8G
    environment:
DEBUG=0
LOGLEVEL=WARNING

# Development (uses docker-compose.yml + docker-compose.override.yml)
docker compose up

Production
docker compose -f docker-compose.yml -f docker-compose.prod.yml up -d

# .github/workflows/ml-pipeline.yml
name: ML Pipeline CI/CD

on:
  push:
    branches: [main]
  pullrequest:

    branches: [main]

jobs:
  test:
    runs-on: ubuntu-latest
    steps:
uses: actions/checkout@v4
name: Build test image
        run: docker build --target builder -t ml-test .

name: Run unit tests
        run: |

          docker run --rm ml-test pytest tests/ -v --cov=src --cov-report=xml
name: Run model validation
        run: |

          docker run --rm \
            -v ${{ github.workspace }}/data:/app/data \
            ml-test python validatemodel.py

  build-and-push:
    needs: test
    runs-on: ubuntu-latest
    if: github.ref == 'refs/heads/main'
    steps:
uses: actions/checkout@v4
name: Login to Container Registry
        uses: docker/login-action@v3

        with:
          registry: ghcr.io
          username: ${{ github.actor }}
          password: ${{ secrets.GITHUBTOKEN }}

name: Build and push production image
        uses: docker/build-push-action@v5

        with:
          context: .
          target: production
          push: true
          tags: |
            ghcr.io/${{ github.repository }}:latest
            ghcr.io/${{ github.repository }}:${{ github.sha }}
          cache-from: type=gha
          cache-to: type=gha,mode=max

  deploy:
    needs: build-and-push
    runs-on: ubuntu-latest
    steps:
name: Deploy to production
        run: |

          echo "Deploying ghcr.io/${{ github.repository }}:${{ github.sha }}"
          # Add deployment commands here (kubectl, docker stack, etc.)

# Rarely changes - cached
FROM python:3.11-slim
RUN apt-get update && apt-get install -y libgomp1

Changes occasionally - rebuild from here when deps change
COPY requirements.txt .
RUN pip install -r requirements.txt

Changes frequently - only this layer rebuilds
COPY src/ ./src/

# .dockerignore
.git
.github
pycache
.pyc
.env
.venv
data/raw/
.egg-info
.pytestcache
.mypycache

notebooks/
.ipynbcheckpoints

3. Pin All Dependency Versions

FROM python:3.11.7-slim  # Pin the exact Python version

Pin system packages
RUN apt-get update && apt-get install -y --no-install-recommends \
    libgomp1=12.2.0-14 \
    && rm -rf /var/lib/apt/lists/*

4. Use Non-Root Users

RUN useradd --create-home --shell /bin/bash --uid 1000 mluser
USER mluser

5. Implement Health Checks

HEALTHCHECK --interval=30s --timeout=10s --start-period=15s --retries=3 \
    CMD python -c "import requests; requests.get('http://localhost:8000/health').raiseforstatus()"

6. Keep Images Small

• Use slim or alpine base images when possible.
• Combine RUN commands to reduce layers.
• Remove build tools after compilation.
• Use --no-cache-dir with pip.
• Leverage multi-stage builds to exclude build-only dependencies.

7. Use Build Arguments for Flexibility

ARG PYTHONVERSION=3.11
ARG MODELVERSION=v1.0.0

FROM python:${PYTHONVERSION}-slim


ENV MODELVERSION=${MODELVERSION}

docker build --build-arg MODELVERSION=v2.0.0 -t ml-api:v2 .

Conclusion

Docker transforms how data scientists and ML engineers build, share, and deploy their work. In this tutorial, you learned how to:

• Write efficient Dockerfiles tailored for ML projects with proper layer caching
• Use multi-stage builds to create lean production images
• Enable GPU acceleration with NVIDIA Container Toolkit
• Orchestrate complete ML stacks using Docker Compose with API, database, cache, monitoring, and experiment tracking
• Manage data and models with volumes and bind mounts
• Handle multiple environments (development, staging, production)
• Integrate Docker into CI/CD pipelines for automated testing and deployment

By containerizing your ML workflows, you eliminate the reproducibility problem, speed up onboarding, and create a clear path from experimentation to production. Start with a simple Dockerfile, iterate on it as your project grows, and invest in Docker Compose when your system involves multiple services.