Containerizing a Machine Learning Model with Docker

Why Containerize an ML Model?

A trained model is useless in production if it only runs on the data scientist’s laptop. Containerization bundles your model artifact, inference code, system libraries, and pinned Python dependencies into a single immutable image that runs identically on any machine with a container runtime.

Key benefits for MLOps:

• Reproducibility — the exact same environment everywhere, eliminating “works on my machine”.
• Portability — the same image runs on a laptop, CI runner, or Kubernetes cluster.
• Isolation — dependencies don’t conflict with the host or other services.
• Scalability — images are the unit of deployment for orchestrators like Kubernetes.

Anatomy of an ML Container

A typical inference container includes:

1. A base image (e.g. python:3.11-slim).
2. System dependencies (e.g. libgomp1 for some ML libs).
3. Python dependencies from a pinned requirements.txt.
4. The model artifact (e.g. model.pkl, model.joblib, or an ONNX file).
5. Inference code — usually an API server such as FastAPI or Flask.

Writing the Dockerfile

# 1. Pin the base image for reproducibility
FROM python:3.11-slim

# 2. Set a working directory
WORKDIR /app

# 3. Install Python deps first to leverage layer caching
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# 4. Copy code and the trained model artifact
COPY app/ ./app/
COPY model.joblib .

# 5. Run as a non-root user for security
RUN useradd --create-home appuser
USER appuser

# 6. Document the port and define the start command
EXPOSE 8000
CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000"]

Why this order matters

Docker builds images in layers and caches each one. Because dependencies change far less often than your code, copying and installing requirements.txt before copying source code means a code change only rebuilds the cheap final layers — not the expensive pip install.

A Minimal Inference App

# app/main.py
import joblib
from fastapi import FastAPI
from pydantic import BaseModel

model = joblib.load("model.joblib")
app = FastAPI()

class Features(BaseModel):
    values: list[float]

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

@app.post("/predict")
def predict(f: Features):
    pred = model.predict([f.values])
    return {"prediction": pred.tolist()}

Building and Running

# Build and tag the image (semantic + git-sha tags are good practice)
docker build -t my-model:1.0.0 .

# Run it, mapping container port 8000 to host 8000
docker run -p 8000:8000 my-model:1.0.0

# Test the endpoint
curl -X POST http://localhost:8000/predict \
  -H 'Content-Type: application/json' \
  -d '{"values": [5.1, 3.5, 1.4, 0.2]}'

Optimizing the Image

Use a .dockerignore

Exclude virtual envs, data, notebooks, and .git to keep the build context small and avoid leaking secrets.

.git
__pycache__/
*.ipynb
data/
venv/

Multi-stage builds

When you need build tools (compilers, wheels) but don’t want them in the final image, use a builder stage and copy only the artifacts forward.

FROM python:3.11-slim AS builder
WORKDIR /app
COPY requirements.txt .
RUN pip wheel --no-cache-dir --wheel-dir /wheels -r requirements.txt

FROM python:3.11-slim
WORKDIR /app
COPY --from=builder /wheels /wheels
RUN pip install --no-cache-dir /wheels/*
COPY app/ ./app/
COPY model.joblib .
EXPOSE 8000
CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000"]

Best Practices Checklist

• Pin everything: base image tag and every dependency version.
• Use slim or distroless base images to reduce attack surface and size.
• Run as non-root.
• Add a HEALTHCHECK or expose a /health endpoint for orchestrators.
• Tag with both a semantic version and the git SHA for traceability.
• Don’t bake secrets into the image; inject them at runtime via env vars or secret managers.
• Decide on model storage: small models can be baked in; large models are often pulled from object storage / a registry at startup.