Shrinking a Python ML Docker Image from 8GB to 900MB with Multi-Stage Builds

Your ML Docker image is 8GB. Your CI pipeline spends 14 minutes pushing it to ECR. Here's how to get it to 900MB.

You’re not alone. Docker is used by 57% of professional developers (Stack Overflow 2025), and a huge chunk of them are wrestling with the same bloated container problem. The good news? Multi-stage builds reduce average image size by 60–80% vs single-stage (Docker Hub analysis, 2025). This isn't magic; it's just deleting the mess you made in the first stage before anyone sees the final product.

Let’s turn your shipping container into a courier envelope.

Why ML Docker Images Balloon: Layer-by-Layer Audit with Dive

First, you need to see the enemy. Your Dockerfile isn't a monolith; it's a stack of immutable layers. Every RUN, COPY, and ADD creates a new one. The problem with ML images is the sequence: you copy your entire codebase, install system dependencies (CUDA, OpenCV libraries), then run pip install -r requirements.txt, which pulls down half of PyPI, including torch and tensorflow. Those packages are massive, and they live forever in your image history.

Don't guess. Use Dive to perform an autopsy.

brew install dive

# Build your image with a tag
docker build -t my-ml-monstrosity:fat .

# Analyze it
dive my-ml-monstrosity:fat

Dive's TUI will show you a tree of your image layers. You'll likely see a single RUN pip install layer consuming 3-4GB. You'll also see your __pycache__ directories, .git, and test files, all baked in. This is the foundation of your problem: the runtime environment is polluted with build-time artifacts.

Choosing the Right Base: python:3.12-slim vs nvidia/cuda:12-cudnn-runtime

Your base image choice sets the floor. The default python:3.12 is over 1.1GB. For CPU inference, python:3.12-slim (130MB) is your best friend. But ML often means GPU. The temptation is nvidia/cuda:12.4.1-devel-ubuntu22.04—a 4GB behemoth containing everything needed to compile CUDA code.

You don't need to compile. You need to run. Swap the -devel for -runtime.

# FROM nvidia/cuda:12.4.1-devel-ubuntu22.04  # ~4GB - NO
FROM nvidia/cuda:12.4.1-runtime-ubuntu22.04  # ~1.5GB - YES

Even better, combine with Python slim. The official Python images provide CUDA variants.

FROM python:3.12-slim  # For CPU, ~130MB

# OR for CUDA 12.1 runtime
FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04  # Check latest tag

This one decision can save you 2.5GB before you write another line.

Multi-Stage Build: Separate Builder Stage from Runtime Stage

This is the core maneuver. You create a disposable "builder" stage where you install every tool and dependency you need to assemble your application. Then, you create a clean, final stage and copy only the necessary artifacts from the builder. The builder stage and all its intermediate layers are discarded, leaving a pristine final image.

Here’s the skeleton:

# Stage 1: The Builder (The messy workshop)
FROM python:3.12-slim as builder

WORKDIR /app

# Copy requirements first for better layer caching
COPY requirements.txt .
# Install ALL dependencies, including dev ones, into a local directory
RUN pip install --user --no-warn-script-location -r requirements.txt

# Copy application code
COPY . .

# Stage 2: The Runtime (The clean showroom)
FROM python:3.12-slim

WORKDIR /app

# Copy *only* the installed packages from the builder's user directory
COPY --from=builder /root/.local /root/.local
# Copy *only* your application code, not the cache or temp files
COPY --from=builder /app /app

# Ensure scripts in .local are usable
ENV PATH=/root/.local/bin:$PATH

# Your application runs here
CMD ["python", "app/main.py"]

The magic is in the COPY --from=builder. It selectively plucks the site-packages folder and your code, leaving behind the pip cache, build tools, and downloaded .whl files.

BuildKit Cache Mounts for pip: Rebuilds in 12s Instead of 8min

Even with multi-stage, rebuilding because you changed one line of code means re-running pip install, which is slow. Enter BuildKit cache mounts. With DOCKER_BUILDKIT=1, you can mount a persistent cache directory for pip (or apt) across builds. BuildKit is 2–5x faster than legacy Docker build on cache-heavy projects (Docker, 2025).

Enable BuildKit and use cache mounts:

# syntax=docker/dockerfile:1.4
FROM python:3.12-slim as builder

WORKDIR /app

COPY requirements.txt .

# This mount persists the pip cache across builds
RUN --mount=type=cache,target=/root/.cache/pip \
    pip install --user --no-warn-script-location -r requirements.txt

COPY . .

Now, rebuild after a code change. The pip install step will show Using cache and finish in seconds because it's reusing downloaded packages. To use this, build with:

# Set BuildKit as the default (or prepend DOCKER_BUILDKIT=1 to commands)
export DOCKER_BUILDKIT=1
docker build -t my-app:lean .

Real Error Fix: If you get Error response from daemon: network not found during a docker-compose up after a docker-compose down, the default network was removed. Fix: Re-create the network with docker network create myapp_default or just use docker-compose up --force-recreate.

.dockerignore: The 10 Lines That Save 2GB Every Build

Your COPY . . is a Trojan horse. Without a .dockerignore file, you're copying your .git folder (hundreds of MB of history), virtual environments (venv/), IDE settings (.vscode/), local datasets, and logs. This bloats the build context sent to the Docker daemon, making every build slower and potentially leaking secrets.

Create a .dockerignore in your project root:

# Python
__pycache__/
*.py[cod]
*.so
.Python
venv/
env/
.pytest_cache/
.coverage

# Data & Logs
data/
*.csv
*.parquet
logs/
*.log

# Git & IDE
.git/
.gitignore
.vscode/
.idea/

# Docker
Dockerfile
docker-compose.yml

# OS
.DS_Store
Thumbs.db

This file is your first line of defense. It prevents gigabytes of irrelevant files from ever becoming a layer.

Removing Dev Dependencies, Tests, and pycache in Runtime Stage

The builder stage might have installed pytest, black, jupyter, and ipython. Your production container doesn't need them. Use a multi-requirements.txt strategy.

requirements.txt:

torch==2.3.0
transformers==4.40.0
numpy==1.26.0
# ... production deps

requirements-dev.txt:

-r requirements.txt
pytest==8.1.0
black==24.4.0
jupyter==1.0.0
ipython==8.22.0

In your builder stage, install requirements-dev.txt. In your final stage, copy only the packages installed from requirements.txt. A more surgical approach is to use pip install --no-deps for your own packages or tools like pip-chill to generate a minimal production list.

Also, ensure no cache sneaks through. Add a cleanup in your final stage:

FROM python:3.12-slim

# ... copy artifacts ...

# Remove pip cache and find/remove __pycache__ directories
RUN rm -rf /root/.cache/pip && \
    find /app -type d -name __pycache__ -exec rm -rf {} + 2>/dev/null || true

# Optional: If you know you don't need the package installer at runtime
RUN apt-get purge -y python3-pip && apt-get autoremove -y

CMD ["python", "app/main.py"]

Real Error Fix: Building on an Apple Silicon Mac for an AMD64 cloud server? You'll hit OCI runtime exec failed: exec format error. Fix: Build with the explicit platform flag: docker build --platform linux/amd64 -t my-app:amd64 ..

Final Benchmark: Layer-by-Layer Before vs After Comparison

Let’s put hard numbers to it. Assume a typical ML app with torch, transformers, scikit-learn, and pandas.

*These layers are discarded and not present in the final image.

The final Dockerfile looks like this:

# syntax=docker/dockerfile:1.4
# STAGE 1: BUILDER
FROM python:3.12-slim as builder

WORKDIR /app
COPY requirements.txt .
RUN --mount=type=cache,target=/root/.cache/pip \
    pip install --user --no-warn-script-location -r requirements.txt

COPY . .

# STAGE 2: RUNTIME
FROM python:3.12-slim

WORKDIR /app
# Copy only the installed packages and our code
COPY --from=builder /root/.local /root/.local
COPY --from=builder /app /app

# Clean up
RUN find /app -type d -name __pycache__ -exec rm -rf {} + 2>/dev/null || true

ENV PATH=/root/.local/bin:$PATH
ENV PYTHONUNBUFFERED=1

CMD ["python", "app/main.py"]

Next Steps: From Lean Image to Robust Pipeline

You've shrunk the image. Now harden the process.

1. Security Scan: Run trivy image my-app:lean to find CVEs in your slimmed-down image. You'll have fewer false positives now.
2. Registry Hygiene: Use Docker Scout to generate a software bill of materials (SBOM) and set policy gates in your CI/CD. Docker Hub serves 13M+ image pulls per day as of Q1 2026; don't let yours be the one that causes a breach.
3. Orchestrate: In docker-compose.yml, define your app and a model cache volume. Use the platform: key to avoid format errors.

   version: '3.8'
   services:
     ml-api:
       platform: linux/amd64
       build: .
       ports:
         - "8000:8000"
       volumes:
         - model-cache:/app/models
   volumes:
     model-cache:
   

4. Automate Updates: For internal apps, consider Watchtower to automatically update running containers with new builds.
5. Monitor: Use lazydocker or Portainer to keep an eye on your running containers' resource usage—your now-lean images will start faster and use less memory.

The goal isn't just a small image. It's a fast, secure, and repeatable build process that doesn't make your infrastructure weep. With 84% of containerized workloads running Docker as the container runtime (CNCF Annual Survey 2025), doing this right isn't a niche skill—it's core engineering. Stop pushing gigabytes of waste and start shipping the application.