Docker for ML Engineers Who Just Want to Ship a Model

You know Python. You can train a model, write a serving script, and get predictions out of an API endpoint. But someone keeps saying you need to "containerize it," and the Docker documentation reads like it was written for infrastructure engineers.

This article skips the container philosophy and goes straight to: here's how to put your model in a Docker image so other people can run it.

The minimum viable Dockerfile

Say you have a directory with these files:

my-model/
├── serve.py
├── model.pkl
└── requirements.txt

Your serve.py is a FastAPI app:

import pickle
from fastapi import FastAPI

app = FastAPI()

with open("model.pkl", "rb") as f:
    model = pickle.load(f)

@app.get("/health")
def health():
    return {"status": "ok"}

@app.post("/predict")
def predict(features: list[float]):
    result = model.predict([features])
    return {"prediction": int(result[0])}

Here's the Dockerfile:

FROM python:3.11-slim

WORKDIR /app

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

COPY serve.py model.pkl ./

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

Build and run:

docker build -t my-model:v1 .
docker run -p 8000:8000 my-model:v1

That's it. Hit http://localhost:8000/health and you should get a response.

What each line does

  • FROM python:3.11-slim — start from a base image that already has Python installed. The slim variant is a stripped-down Debian image (~150 MB instead of ~900 MB).
  • WORKDIR /app — set the working directory inside the container. All subsequent commands run from here.
  • COPY requirements.txt . then RUN pip install — install dependencies first, before copying the rest of the code. Docker caches each layer, so if your code changes but your dependencies don't, pip install won't re-run.
  • COPY serve.py model.pkl ./ — copy your application code and model into the container.
  • CMD [...] — the command that runs when the container starts.

The layer caching trick

Docker builds images in layers. Each COPY or RUN instruction creates a layer. If a layer's input hasn't changed since the last build, Docker reuses the cached version.

This is why you copy requirements.txt and install dependencies before copying your code. If you change serve.py, Docker reuses the cached pip install layer. If you did it all in one step, changing any file would re-trigger a full pip install.

# Good: dependencies are cached separately from code
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY serve.py model.pkl ./

# Bad: any file change re-installs all dependencies
COPY . .
RUN pip install --no-cache-dir -r requirements.txt

This matters when your dependencies include PyTorch (2+ GB). You don't want to re-download it every time you fix a typo in your serving code.

Keeping images small

ML images get big fast. PyTorch alone is over 2 GB. Add CUDA, your model weights, and some data processing libraries and you're looking at 10+ GB images. Big images mean slow pulls, slow deploys, and wasted disk space.

Multi-stage builds

A multi-stage build uses one image to install/build things and a second image to run them. Only the second image ships.

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

WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir --prefix=/install -r requirements.txt

# Stage 2: runtime image
FROM python:3.11-slim

WORKDIR /app
COPY --from=builder /install /usr/local
COPY serve.py model.pkl ./

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

The builder stage installs packages into /install. The runtime stage copies only the installed packages — no pip cache, no build tools, no compiler artifacts. The final image is smaller because it doesn't contain anything needed only at build time.

Other size-reduction tactics

Use --no-cache-dir with pip. Without it, pip caches downloaded packages inside the image. These caches serve no purpose in a container.

Don't install dev dependencies. If your requirements.txt includes pytest, jupyter, or matplotlib but you don't need them at serving time, create a separate requirements-serve.txt with only the serving dependencies.

Use .dockerignore. Without a .dockerignore file, COPY . . copies everything — including your .git directory, __pycache__, .venv, and data files. Create one:

.git
__pycache__
*.pyc
.venv
data/
notebooks/
.env

Consider distroless or Alpine. If you don't need a shell or package manager at runtime, Google's distroless Python images are even smaller than slim. Alpine is another option but can cause issues with Python packages that have C extensions.

GPU images with NVIDIA

If your model needs GPU inference (deep learning, large models), you need NVIDIA's CUDA runtime in the container. This means using a different base image.

FROM nvidia/cuda:12.1.1-runtime-ubuntu22.04

# Install Python (NVIDIA images don't include it)
RUN apt-get update && \
    apt-get install -y --no-install-recommends python3 python3-pip && \
    rm -rf /var/lib/apt/lists/* && \
    ln -s /usr/bin/python3 /usr/bin/python

WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY serve.py model/ ./

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

NVIDIA provides several image variants:

  • base — CUDA driver libraries only. The smallest option.
  • runtime — adds CUDA runtime libraries. Use this for inference.
  • devel — adds compilers and headers. Use this if you need to compile CUDA code (most people don't for serving).

For inference, runtime is almost always what you want.

Running GPU containers

You need two things on the host:

  1. NVIDIA drivers installed on the host machine (not in the container).
  2. NVIDIA Container Toolkit (nvidia-ctk) installed, so Docker can pass GPU devices into the container.
docker run --gpus all -p 8000:8000 my-model-gpu:v1

The --gpus all flag makes all host GPUs available inside the container. You can also specify --gpus '"device=0"' for a specific GPU.

On Kubernetes, you'd request GPUs in the Pod spec instead:

resources:
  limits:
    nvidia.com/gpu: 1

The NVIDIA device plugin handles the rest.

PyTorch-specific tip

PyTorch distributes separate packages for CPU and GPU. The GPU version includes CUDA bindings and is much larger. Use the right one:

# GPU (includes CUDA)
--extra-index-url https://download.pytorch.org/whl/cu121
torch==2.1.0

# CPU only (much smaller)
--extra-index-url https://download.pytorch.org/whl/cpu
torch==2.1.0

If your serving code runs on CPU, use the CPU wheel. Your image will be gigabytes smaller.

Pushing to a registry

A Docker image on your laptop isn't useful to anyone else. Push it to a container registry so your deployment pipeline can pull it.

Docker Hub

docker tag my-model:v1 yourusername/my-model:v1
docker push yourusername/my-model:v1

Docker Hub is public by default (free tier). Fine for open-source, not for proprietary models.

Private registry (GitHub Container Registry)

echo $GITHUB_TOKEN | docker login ghcr.io -u USERNAME --password-stdin
docker tag my-model:v1 ghcr.io/yourorg/my-model:v1
docker push ghcr.io/yourorg/my-model:v1

Private registry (self-hosted)

If you're on-prem, you might run your own registry (Harbor, or the basic Docker registry):

docker tag my-model:v1 registry.internal.company.com/ml/my-model:v1
docker push registry.internal.company.com/ml/my-model:v1

Tagging strategy

Don't use :latest. It's mutable — every push overwrites it. You won't know which version is running in production.

Use one of these:

  • Version tags: :v1, :v2, :v3. Simple and clear.
  • Git SHA: :sha-a1b2c3d. Ties the image to a specific commit.
  • Date-based: :2025-07-11. Useful for models that retrain on a schedule.

Tag every image with something immutable. When something breaks in production, you need to know exactly which image is running.

The full example

Putting it all together — a multi-stage Dockerfile for a model serving endpoint with health checks:

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

WORKDIR /build
COPY requirements.txt .
RUN pip install --no-cache-dir --prefix=/install -r requirements.txt

# Stage 2: runtime
FROM python:3.11-slim

WORKDIR /app

# Copy installed packages from builder
COPY --from=builder /install /usr/local

# Copy application code
COPY serve.py .
COPY model/ model/

# Non-root user for security
RUN useradd --create-home appuser
USER appuser

EXPOSE 8000
HEALTHCHECK --interval=30s --timeout=5s --retries=3 \
  CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')"

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

Build, tag, push:

docker build -t registry.example.com/ml/my-model:v1 .
docker push registry.example.com/ml/my-model:v1

Deploy to Kubernetes:

apiVersion: apps/v1
kind: Deployment
metadata:
  name: my-model
spec:
  replicas: 2
  selector:
    matchLabels:
      app: my-model
  template:
    metadata:
      labels:
        app: my-model
    spec:
      containers:
      - name: serve
        image: registry.example.com/ml/my-model:v1
        ports:
        - containerPort: 8000
        livenessProbe:
          httpGet:
            path: /health
            port: 8000
        readinessProbe:
          httpGet:
            path: /health
            port: 8000
        resources:
          requests:
            memory: "512Mi"
            cpu: "250m"
          limits:
            memory: "1Gi"
            cpu: "500m"

The model is now running in a container, behind a health check, with resource limits, and can be scaled horizontally by changing replicas. Anyone with access to the registry can deploy it without knowing anything about Python versions, pip, or your machine's environment.

Containerizing ML workloads is a core part of what we set up at deploying.ai. If you need help building reliable ML deployment pipelines, reach out.