Docker for ML Models: Containerize, Deploy and Scale with Confidence

# ─── 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"]

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)

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)}")

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}")

# 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