Docker ML Models — Fixing numpy Version Drift in Serving

ML models are environment-sensitive in ways that web applications are not. A model trained with numpy 1.23 can silently produce different floating-point results on numpy 2.0. A CUDA version mismatch between training and serving causes either crashes or silent CPU fallback that tanks inference latency by 50x. These are not hypothetical — they are the leading causes of 'it works on my machine' failures in ML deployments.

Docker eliminates environment drift by packaging the entire runtime — OS libraries, Python version, pip packages, CUDA toolkit, model weights, and serving logic — into a single versioned image. That image runs identically on your laptop, your CI pipeline, a Kubernetes cluster, and an edge device.

The gap between a Jupyter notebook that produces great metrics and a model that reliably serves predictions in production is wider than most teams expect. Docker closes that gap by making the environment a constant, not a variable. This guide covers the patterns that separate production-grade ML containers from fragile ones.

Why Docker ML Models Is About Dependency Isolation, Not Just Containers

Docker ML models is the practice of packaging a trained machine learning model together with its exact runtime dependencies — Python version, system libraries, and every pip package — into a container image. The core mechanic is that the container becomes a self-contained serving unit: the model artifact, the inference code, and the environment are frozen together at build time. This eliminates the most common failure in ML serving: version drift between training and production environments.

In practice, the container image pins every dependency to a specific version. For example, numpy 1.21.0 compiled against OpenBLAS 0.3.13 is not the same as numpy 1.24.0 compiled against OpenBLAS 0.3.21 — matrix multiplication results can differ in the 5th decimal place, which cascades into different predictions. Docker ensures the exact same binary is used in training, CI, staging, and production. The image is immutable; you never install a newer numpy on top of an existing container.

Use Docker ML models whenever your model’s inference path includes any compiled library (numpy, scipy, TensorFlow, PyTorch, ONNX Runtime). The cost of not doing it is silent correctness failures: the model passes unit tests but produces different outputs in production because a minor version of a linear algebra library changed its rounding behavior. For teams serving models at scale, this is the difference between a reproducible deployment and a debugging nightmare.

Multi-Stage Builds for ML — Separating Training from Serving

The most common mistake in ML Dockerfiles is shipping the training environment as the serving image. A training image includes Jupyter, gcc, test frameworks, debugging tools, and development dependencies — none of which are needed in production. This bloats the image to 5-10GB, increases attack surface, and slows deployments.

Multi-stage builds solve this by using a heavy 'builder' stage with all training and build dependencies, then copying only the trained model and runtime dependencies into a minimal 'serving' stage. The final image contains Python, the serving framework, and the model — nothing else.

Why this matters for ML specifically: ML images are uniquely large because they include CUDA toolkit (2-3GB), PyTorch/TensorFlow (1-2GB), and model weights (1-10GB). A single-stage image that includes training tools, CUDA development headers, and model weights can easily exceed 10GB. A multi-stage serving image with quantized weights can be under 2GB.

Layer caching for ML: Model weights change rarely (only after retraining). Dependencies change occasionally. Application code changes frequently. Order your Dockerfile: base image + CUDA first, dependencies second, model weights third, application code last. This ensures code changes do not trigger a re-download of PyTorch or a re-copy of multi-gigabyte weights.

# ─── STAGE 1: Build environment (training deps, compilers) ───
FROM python:3.10.12-slim-bookworm AS builder

WORKDIR /build

# Install build dependencies (not in final image)
RUN apt-get update && apt-get install -y --no-install-recommends \
    gcc g++ libpq-dev && \
    rm -rf /var/lib/apt/lists/*

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

# Simulate model training artifact (in practice, this comes from
# a training pipeline or model registry)
COPY models/ ./models/

# ─── STAGE 2: Serving runtime (minimal) ───
FROM python:3.10.12-slim-bookworm AS serving

WORKDIR /app

# Install only runtime dependencies
COPY requirements-serving.txt .
RUN pip install --no-cache-dir -r requirements-serving.txt

# Copy trained model weights from builder
COPY --from=builder /build/models/ ./models/

# Copy serving application code
COPY src/serving/ ./src/serving/

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

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

EXPOSE 8080

CMD ["python", "-m", "uvicorn", "src.serving.api:app", "--host", "0.0.0.0", "--port", "8080"]

GPU Access with NVIDIA Container Toolkit

ML inference on CPU is 10-100x slower than on GPU. Docker does not expose GPU devices to containers by default — you need the NVIDIA Container Toolkit and the --gpus flag.

The NVIDIA Container Toolkit (formerly nvidia-docker2) installs a Docker runtime that automatically mounts the GPU device drivers and libraries into containers. Without it, containers see no GPU devices even when the host has GPUs available.

Installation and verification: 1. Install nvidia-container-toolkit on the host 2. Configure Docker to use the nvidia runtime: sudo nvidia-ctk runtime configure --runtime=docker 3. Restart Docker: sudo systemctl restart docker 4. Verify: docker run --gpus all nvidia/cuda:11.8.0-base-ubuntu22.04 nvidia-smi

GPU allocation strategies: - --gpus all: expose all GPUs to the container - --gpus 1: expose one GPU (Docker picks which) - --gpus '\"device=0,2\"': expose specific GPUs by index - NVIDIA_VISIBLE_DEVICES=0,2: set via environment variable (useful in Compose)

Failure scenario — silent CPU fallback: If the NVIDIA Container Toolkit is not installed or --gpus is not passed, PyTorch and TensorFlow silently fall back to CPU. The model loads successfully, inference works, but latency is 50x slower than expected. There is no error — torch.cuda.is_available() returns False, but many serving frameworks do not check this. The fix: always add a startup assertion that verifies GPU availability.

import sys
import logging

logger = logging.getLogger(__name__)

def verify_gpu_availability(required_gpus: int = 1) -> None:
    """Startup assertion: fail fast if GPU is not available.
    
    Call this at application startup before loading the model.
    If GPU is required but not available, exit immediately
    rather than silently falling back to CPU.
    """
    try:
        import torch
        
        available = torch.cuda.is_available()
        device_count = torch.cuda.device_count()
        
        if not available:
            logger.error(
                "GPU not available. torch.cuda.is_available() returned False. "
                "Ensure NVIDIA Container Toolkit is installed and "
                "container is started with --gpus flag."
            )
            sys.exit(1)
        
        if device_count < required_gpus:
            logger.error(
                f"Insufficient GPUs: required={required_gpus}, "
                f"available={device_count}. "
                f"Adjust --gpus flag or reduce requirement."
            )
            sys.exit(1)
        
        gpu_name = torch.cuda.get_device_name(0)
        gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1e9
        logger.info(
            f"GPU verified: {gpu_name}, "
            f"{gpu_memory:.1f}GB VRAM, "
            f"{device_count} device(s) available"
        )
        
    except ImportError:
        logger.error("PyTorch not installed. Cannot verify GPU availability.")
        sys.exit(1)

Model Serving Patterns — FastAPI, TorchServe, and Triton

There are three common patterns for serving ML models in Docker containers. The right choice depends on your latency requirements, model complexity, and operational maturity.

Pattern 1: FastAPI + direct model loading. Load the model at startup, expose a /predict endpoint. Simple, full control over the inference pipeline, easy to customize. Best for single-model serving with custom pre/post-processing. The model is loaded into the application process — startup time equals model load time.

Pattern 2: TorchServe / TF Serving. Purpose-built serving frameworks with built-in batching, model versioning, and A/B testing. More operational overhead but better for multi-model serving and high-throughput scenarios. TorchServe runs a separate model server process — the Docker container wraps the TorchServe binary.

Pattern 3: NVIDIA Triton Inference Server. GPU-optimized serving with dynamic batching, model ensemble pipelines, and multi-framework support (PyTorch, TensorFlow, ONNX, TensorRT). Highest throughput but most complex configuration. Best for latency-critical production workloads with multiple models.

Health check patterns: A health check that only verifies the server is running is insufficient for ML serving. The health check must verify that the model is loaded and can produce a valid output. A /health endpoint should run a dummy inference with a known input and verify the output shape matches expectations.

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List
import numpy as np
import torch
import logging

from io.thecodeforge.ml_serving.startup_check import verify_gpu_availability

logger = logging.getLogger(__name__)

app = FastAPI(title="ML Model Serving API")

# Global model reference — loaded once at startup
model = None
model_device = None


class PredictionRequest(BaseModel):
    features: List[float]


class PredictionResponse(BaseModel):
    prediction: float
    model_version: str
    device: str


@app.on_event("startup")
async def load_model() -> None:
    """Load model once at startup — not on every request."""
    global model, model_device

    # Fail fast if GPU is not available
    verify_gpu_availability(required_gpus=1)

    model_device = torch.device("cuda:0")
    model_path = "./models/production_model.pt"

    logger.info(f"Loading model from {model_path}...")
    model = torch.jit.load(model_path, map_location=model_device)
    model.eval()

    # Warmup inference — ensures CUDA kernels are compiled
    dummy_input = torch.randn(1, 128, device=model_device)
    with torch.no_grad():
        _ = model(dummy_input)

    logger.info("Model loaded and warmed up successfully.")


@app.get("/health")
async def health_check() -> dict:
    """Health check that verifies model is loaded and functional."""
    if model is None:
        raise HTTPException(status_code=503, detail="Model not loaded")

    # Dummy inference to verify the model actually works
    try:
        dummy_input = torch.randn(1, 128, device=model_device)
        with torch.no_grad():
            output = model(dummy_input)
        return {
            "status": "healthy",
            "model_loaded": True,
            "output_shape": list(output.shape),
            "device": str(model_device),
        }
    except Exception as e:
        raise HTTPException(
            status_code=503,
            detail=f"Model health check failed: {str(e)}"
        )


@app.post("/predict", response_model=PredictionResponse)
async def predict(request: PredictionRequest) -> PredictionResponse:
    """Run inference on the loaded model."""
    if model is None:
        raise HTTPException(status_code=503, detail="Model not loaded")

    try:
        input_tensor = torch.tensor(
            [request.features], dtype=torch.float32, device=model_device
        )

        with torch.no_grad():
            output = model(input_tensor)

        return PredictionResponse(
            prediction=output.item(),
            model_version="v1.0.0",
            device=str(model_device),
        )
    except Exception as e:
        logger.error(f"Inference failed: {e}")
        raise HTTPException(status_code=500, detail=f"Inference error: {str(e)}")

Volume Strategies for Model Weights — Baking vs Mounting vs Pulling

Model weights are the largest component of an ML serving image. A production NLP model can be 2-10GB. A large language model can be 50-200GB. How you deliver these weights to the container has a major impact on deployment speed, storage costs, and operational flexibility.

Strategy 1: Bake weights into the image (COPY). Simplest approach. The weights are part of the image layer. Every deployment pulls the full image including weights. Pros: self-contained, no external dependencies. Cons: every model update requires a full image rebuild and multi-gigabyte pull. Not practical for models >1GB.

Strategy 2: Mount weights as a named volume. Weights are stored in a Docker volume, mounted into the container at runtime. The image stays small (just the framework and serving code). Pros: image is small and fast to pull. Cons: weights must be pre-populated in the volume. Requires a separate weight management process.

Strategy 3: Pull weights from object storage at startup. The container downloads weights from S3/GCS/Azure Blob at startup. Pros: always gets the latest version, no pre-population needed, works across environments. Cons: adds startup latency (5-60 seconds depending on model size and network), requires credentials management, adds a failure mode (network timeout during download).

Strategy 4: Hybrid — framework in image, weights in registry. Use a model registry (MLflow, Weights & Biases, SageMaker Model Registry) to version and store weights. The serving image contains the framework and a startup script that pulls the correct model version from the registry. This is the most operationally mature approach — it decouples model updates from image updates.

import os
import logging
from pathlib import Path

logger = logging.getLogger(__name__)


def load_model_weights(model_name: str, version: str) -> Path:
    """Load model weights using the appropriate strategy.
    
    Strategy is determined by environment variable MODEL_SOURCE:
    - 'baked': weights are in the image (COPY in Dockerfile)
    - 'volume': weights are in a mounted volume
    - 's3': weights are downloaded from S3 at startup
    """
    source = os.environ.get("MODEL_SOURCE", "baked")
    
    if source == "baked":
        # Weights were COPY'd into the image during build
        model_path = Path(f"./models/{model_name}/{version}/model.pt")
        if not model_path.exists():
            raise FileNotFoundError(
                f"Baked model not found at {model_path}. "
                f"Ensure the model was copied during docker build."
            )
        logger.info(f"Loaded baked model from {model_path}")
        return model_path
    
    elif source == "volume":
        # Weights are in a mounted Docker volume
        volume_path = os.environ.get("MODEL_VOLUME_PATH", "/data/models")
        model_path = Path(f"{volume_path}/{model_name}/{version}/model.pt")
        if not model_path.exists():
            raise FileNotFoundError(
                f"Model not found in volume at {model_path}. "
                f"Ensure the volume is mounted and contains the model."
            )
        logger.info(f"Loaded model from volume: {model_path}")
        return model_path
    
    elif source == "s3":
        # Download from S3 at startup
        import boto3
        
        bucket = os.environ["MODEL_S3_BUCKET"]
        key = f"models/{model_name}/{version}/model.pt"
        local_path = Path(f"/tmp/models/{model_name}/{version}/model.pt")
        local_path.parent.mkdir(parents=True, exist_ok=True)
        
        logger.info(f"Downloading model from s3://{bucket}/{key}...")
        s3 = boto3.client("s3")
        s3.download_file(bucket, key, str(local_path))
        logger.info(f"Downloaded model to {local_path}")
        return local_path
    
    else:
        raise ValueError(f"Unknown MODEL_SOURCE: {source}")

Production Deployment Patterns — Health Checks, Graceful Shutdown, and Resource Limits

Deploying ML models in production requires patterns that go beyond basic containerization. Three patterns separate production-grade deployments from fragile ones.

1. Health checks that verify inference capability. A /health endpoint must verify that the model is loaded and can produce valid output. Run a dummy inference at startup and on every health check. Kubernetes readiness probes use this to route traffic only to ready containers.

2. Graceful shutdown for in-flight requests. When a container is stopped (docker stop, Kubernetes pod termination), it receives SIGTERM. The serving framework must stop accepting new requests, complete in-flight requests, and exit cleanly. Default stop timeout is 10 seconds — increase it with --stop-timeout or terminationGracePeriodSeconds if inference takes longer.

3. Resource limits to prevent GPU and memory contention. Without resource limits, one container can consume all GPU memory or host memory, crashing other services. Set --memory limits for RAM. Use NVIDIA_MPS or CUDA_VISIBLE_DEVICES for GPU isolation. In Kubernetes, use resource requests and limits for both CPU/memory and nvidia.com/gpu.

4. Model warmup to avoid cold-start latency. The first inference on a GPU model is slow because CUDA kernels must be compiled. Run a dummy inference at startup to warm up the GPU. This ensures the first real request has the same latency as subsequent requests.

# Kubernetes deployment for ML model serving
apiVersion: apps/v1
kind: Deployment
metadata:
  name: ml-model-serving
  namespace: production
spec:
  replicas: 3
  selector:
    matchLabels:
      app: ml-model-serving
  template:
    metadata:
      labels:
        app: ml-model-serving
    spec:
      containers:
        - name: model-server
          image: io.thecodeforge/ml-model:v1.0.0
          ports:
            - containerPort: 8080
          resources:
            requests:
              memory: "2Gi"
              cpu: "1"
              nvidia.com/gpu: "1"
            limits:
              memory: "4Gi"
              cpu: "2"
              nvidia.com/gpu: "1"
          readinessProbe:
            httpGet:
              path: /health
              port: 8080
            initialDelaySeconds: 60  # Model loading takes time
            periodSeconds: 10
            timeoutSeconds: 5
            failureThreshold: 3
          livenessProbe:
            httpGet:
              path: /health
              port: 8080
            initialDelaySeconds: 120
            periodSeconds: 30
            timeoutSeconds: 5
            failureThreshold: 5
          env:
            - name: MODEL_SOURCE
              value: "s3"
            - name: MODEL_S3_BUCKET
              valueFrom:
                secretKeyRef:
                  name: ml-model-secrets
                  key: s3-bucket
          lifecycle:
            preStop:
              exec:
                command: ["/bin/sh", "-c", "sleep 10"]
      terminationGracePeriodSeconds: 60  # Allow in-flight requests to complete

The Model Registry Trap — Why Your Docker Tags Are Lying to You

You've built a solid pipeline. Images versioned, containers reproducible. But six months from now, when that production model starts drifting and you need the exact weights from build #217, your tag-based registry will fail you. Why? Because Docker tags are mutable pointers, not immutable identifiers. Someone rebuilds the same tag for a hotfix, and now your "v2.3.1" image contains different weights, different dependencies, and a different model architecture than what shipped to production.

The fix isn't more tags. It's digest pinning. Every image push produces a digest — a SHA256 hash of the manifest that uniquely identifies that exact image. When you pull from your model registry, pin the digest, not the tag. Store the digest alongside your deployment manifest, your metrics, and your rollback plan. If you're using any model registry like MLflow or DVC, push both the model artifact and the container digest as paired artifacts.

This matters because ML systems fail asymmetrically. A web app crashing is obvious; a model serving slightly wrong predictions for three weeks is not. When you do need to bisect that regression, immutable digests give you a single source of truth. Tags are for humans. Digests are for machines. Trust the machine.

// io.thecodeforge — ml-ai tutorial

import subprocess, json

# Get the digest of your pushed image
image = "registry.thecodeforge.io/ner-model:2.3.1"
result = subprocess.run(
    ["docker", "image", "inspect", image, "--format", "{{json .RepoDigests}}"],
    capture_output=True, text=True
)
digests = json.loads(result.stdout.strip())
print(f"Pinned digest: {digests[0]}")
# Output: registry.thecodeforge.io/ner-model@sha256:a1b2c3d4e5f6...

# Deploy with digest, not tag
compose_manifest = f"""
services:
  inference:
    image: {digests[0]}
    ports: ["8080:8080"]
"""
with open("docker-compose.prod.yml", "w") as f:
    f.write(compose_manifest)
print("Deploying with digest-based immutable reference")

Debugging ML Containers Without Hating Your Life — The Exec vs Entrypoint Gamble

Your container builds fine locally, but on the GPU node it crashes with a cryptic CUDA error. Your first instinct? Cargo cult an interactive shell and poke around. But your Dockerfile ends with CMD ["python", "serve.py"], which means docker run executes the model server, not a shell. You override the entrypoint, get a bash prompt, and now you're hunting for missing shared libraries in an environment that wasn't designed for exploration.

Stop fighting the framework. Use a debug image pattern. Have a Dockerfile.debug or a target in your multi-stage build that adds common debugging tools — curl, strace, nvidia-smi, python3-gdb. Build it separately and mount your model weights as a volume. The key insight: your production image should be minimal, but your debug image should be a full forensic toolkit. Keep them as sibling builds from the same base to guarantee compatibility.

I've debugged a silent NAN issue in a PyTorch model for six hours with nothing but strace and a debug image. The production image was 400MB; the debug image was 1.2GB. I threw it away afterward. That's fine. Debug images are disposable by design — they save you from rebuilding your entire pipeline every time you need to inspect a layer cache miss.

// io.thecodeforge — ml-ai tutorial

# Dockerfile.debug — extends production image with tooling
# Build: docker build -f Dockerfile.debug -t ner-model:debug .

# FROM production-base retains exact GPU/cuDNN versions
# Don't rebuild from scratch — extend your production stage

# In your CI:
# docker build --target production -t ner-model:prod .
# docker build --target debug -t ner-model:debug .

# Run interactive debug:
# docker run --gpus all -v $(pwd)/weights:/models/weights \
#   -it --entrypoint bash ner-model:debug

# Inside container:
# nvidia-smi  # verify GPU visibility
# strace -f -e trace=open,openat python serve.py 2>&1 | grep "No such file"
# python3 -c "import torch; print(torch.cuda.is_available())"

The Cold Start Hell of GPU Containers — Preloading CUDA Kernels

Your model container starts, health checks pass, but the first inference request hangs for 12 seconds. Then it runs at 2ms per request. The difference? The CUDA runtime is JIT-compiling kernels on first use — inside a container that has no warm cache. This isn't a bug; it's the architecture. Every CUDA context creation triggers driver compilation, kernel loading, and context initialization that your dev machine cached months ago.

You have two pragmatic responses. First, use CUDA graphs and persistence of work: pre-run a dummy inference during your container startup health check. That forces the JIT compilation and kernel caching before any real traffic hits. Add a --warmup flag to your entrypoint that runs 3 iterations with a tiny batch and waits for GPU sync. Second, use the CUDA_CACHE_PATH environment variable to point to a writable volume. This persists compiled kernels across container restarts — crucial when Kubernetes reschedules your pod and you don't want to recompile the entire graph.

A team I consulted cut their p99 latency from 8.4s to 47ms with nothing but a Dockerfile ENV CUDA_CACHE_PATH=/cache/cuda and a 15-line warmup script. Their ML engineers had spent two months blaming the network. We fixed it in two hours by understanding what CUDA actually does during container initialization.

// io.thecodeforge — ml-ai tutorial

import torch, time, sys

def warmup(model, device, iterations=5):
    """Precompile CUDA kernels and warm GPU cache."""
    dummy = torch.randn(1, 3, 224, 224).to(device)
    for _ in range(iterations):
        _ = model(dummy)
        torch.cuda.synchronize(device)  # Force completion
        print(f"Warmup iteration {_+1} complete", file=sys.stderr)
    print("GPU warmup done — ready for traffic", file=sys.stderr)

if __name__ == "__main__":
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = torch.jit.load("/models/traced_resnet.pt").to(device)
    
    if "--warmup" in sys.argv:
        warmup(model, device, iterations=3)
    
    # Proceed to start FastAPI/Triton server
    print(f"Model loaded on {device}. Starting server...")