How to Write a Production Dockerfile for PyTorch Inference

Most Dockerfiles you find online are written for development. Root user, bloated base images, no health checks, no signal handling. Here's a Dockerfile that actually belongs in production.

The Full Dockerfile

FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime

WORKDIR /app

# Install system deps first (rarely changes)
RUN apt-get update && apt-get install -y --no-install-recommends \
    libgl1-mesa-glx \
    libglib2.0-0 \
    curl \
    && rm -rf /var/lib/apt/lists/*

# Install Python deps before code (layer caching)
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy app code
COPY . .

# Create non-root user
RUN groupadd -r appuser && useradd -r -g appuser -u 1000 appuser \
    && chown -R appuser:appuser /app
USER appuser

EXPOSE 8000

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

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

Why Each Decision Matters

Runtime Image, Not Devel

pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime is ~2.5GB. The -devel variant is ~5GB. The devel image includes nvcc, headers, and build tools you don't need at inference time. Smaller image means faster pulls, faster cold starts, and less attack surface.

If you need to compile something at build time (e.g., flash-attention from source), use a multi-stage build: compile in the devel image, copy the wheel to the runtime image.

Layer Ordering

Docker caches layers. If you change a line in app.py, every layer after COPY . . gets rebuilt. By placing COPY requirements.txt and RUN pip install before COPY . ., you avoid reinstalling Python packages every time you change application code.

In practice, this turns a 5-minute rebuild into a 10-second rebuild for code changes.

Non-Root User

Running as root means any code execution vulnerability gives the attacker root access to the container. With UID 1000 (the standard non-root user across most platforms), an exploit is limited to what appuser can do.

Some platforms (OpenShift, certain K8s security policies) reject containers that run as root. Using UID 1000 avoids this.

Health Check

Without a health check, your orchestrator doesn't know if the app is actually working. It only knows if the process is alive. A process can be alive and broken — e.g., the model failed to load, but the FastAPI server is running and returning 500s.

The health check here calls /health every 30 seconds. Your app should implement this:

@app.get("/health")
async def health():
    return {
        "status": "healthy",
        "model_loaded": True,
        "gpu_available": torch.cuda.is_available(),
    }

If /health returns non-200 or times out after 5 seconds, 3 times in a row, the container is marked unhealthy and the orchestrator replaces it.

Graceful Shutdown

FastAPI/Uvicorn handles SIGTERM by default, but your app needs to handle cleanup. On shutdown:

import signal

@app.on_event("shutdown")
def shutdown():
    if hasattr(app.state, 'model'):
        del app.state.model
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

Without this, GPU memory can leak across restarts. After 3-4 deploys, you hit OOM.

What Not to Include

• Supervisor/process managers: Container should run one process. If you need multiple processes (e.g., nginx + app), use a sidecar container or a multi-container pod.
• Debug tools: No vim, htop, strace. These increase image size and attack surface. Debug locally, not in production.
• Development dependencies: No pytest, black, ipython. Use multi-stage builds or a separate dev Dockerfile.
• Hardcoded secrets: No ENV API_KEY=xxx. Use runtime environment injection — Kubernetes secrets, cloud parameter stores, or docker run -e.

Multi-Stage Build for Python Dependencies

If you have compiled Python packages (e.g., pyarrow, onnxruntime-gpu), use a builder stage:

FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-devel AS builder

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

FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime

COPY --from=builder /deps /usr/local/lib/python3.11/site-packages/
# ... rest of the Dockerfile

This keeps the final image small while allowing compilation.

Testing the Dockerfile

Before deploying, test locally:

docker build -t inference-app .
docker run -p 8000:8000 inference-app

# In another terminal:
curl http://localhost:8000/health
curl -X POST http://localhost:8000/predict \
  -H "Content-Type: application/json" \
  -d '{"text": "test input"}'

If the health check returns 200 and predictions work, the Dockerfile is production-ready.

We generate these Dockerfiles automatically at Roptal — the analysis layer reads your repo, detects the framework and dependencies, and produces a production Dockerfile you can review before deploying.