Skip to content
Codeloom
Machine Learning

ML Model Deployment with FastAPI

Deploy machine learning models as production REST APIs using FastAPI with input validation, async inference, and health checks.

·6 min read · By Codeloom
Intermediate 12 min read

What you'll learn

  • How to serve ML models behind a FastAPI endpoint
  • Input validation with Pydantic for model predictions
  • Async request handling for better throughput
  • Health checks, versioning, and production best practices

Prerequisites

None — this post is self-contained.

Training a model is half the work. Getting it into production where applications can call it reliably is the other half. FastAPI is a strong choice for serving ML models: it is fast, supports async request handling, has built-in input validation with Pydantic, and generates OpenAPI documentation automatically.

This guide walks through building a production-ready ML serving API from scratch.

Project Structure

A clean project layout separates the API layer from model logic:

ml-api/
  app/
    __init__.py
    main.py          # FastAPI application
    model.py         # Model loading and inference
    schemas.py       # Pydantic request/response models
  models/
    classifier_v1.joblib
  requirements.txt
  Dockerfile

Defining Request and Response Schemas

Start with Pydantic models that validate inputs and document the API contract.

# app/schemas.py
from pydantic import BaseModel, Field

class PredictionRequest(BaseModel):
    features: list[float] = Field(
        ...,
        min_length=1,
        max_length=100,
        description="Feature vector for prediction",
    )
    return_probabilities: bool = Field(
        default=False,
        description="Whether to return class probabilities",
    )

    model_config = {
        "json_schema_extra": {
            "examples": [
                {
                    "features": [5.1, 3.5, 1.4, 0.2],
                    "return_probabilities": True,
                }
            ]
        }
    }

class PredictionResponse(BaseModel):
    prediction: int | str
    confidence: float = Field(ge=0.0, le=1.0)
    probabilities: dict[str, float] | None = None
    model_version: str

class HealthResponse(BaseModel):
    status: str
    model_loaded: bool
    model_version: str
    uptime_seconds: float

Pydantic validates every incoming request before your model code runs. Malformed inputs return a clear 422 error with details about what failed, without ever touching the model.

Loading and Wrapping the Model

Wrap model loading and inference in a class that handles initialization once and provides a clean prediction interface.

# app/model.py
import joblib
import numpy as np
from pathlib import Path

class ModelService:
    def __init__(self):
        self.model = None
        self.version = "unknown"
        self.class_names: list[str] = []

    def load(self, model_path: str, version: str = "v1"):
        """Load model from disk. Called once at startup."""
        path = Path(model_path)
        if not path.exists():
            raise FileNotFoundError(f"Model not found: {model_path}")

        self.model = joblib.load(path)
        self.version = version

        # Extract class names if available
        if hasattr(self.model, "classes_"):
            self.class_names = [str(c) for c in self.model.classes_]

    def predict(self, features: list[float]) -> dict:
        """Run inference on a single feature vector."""
        if self.model is None:
            raise RuntimeError("Model not loaded")

        X = np.array(features).reshape(1, -1)
        prediction = self.model.predict(X)[0]

        # Get probabilities if the model supports them
        probabilities = None
        confidence = 1.0

        if hasattr(self.model, "predict_proba"):
            proba = self.model.predict_proba(X)[0]
            confidence = float(np.max(proba))

            if self.class_names:
                probabilities = {
                    name: float(p)
                    for name, p in zip(self.class_names, proba)
                }
            else:
                probabilities = {
                    str(i): float(p) for i, p in enumerate(proba)
                }

        return {
            "prediction": prediction,
            "confidence": confidence,
            "probabilities": probabilities,
            "model_version": self.version,
        }

# Singleton instance
model_service = ModelService()

Using a singleton avoids reloading the model on every request. The model loads into memory once at startup and stays there.

Building the FastAPI Application

# app/main.py
import time
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException
from app.schemas import PredictionRequest, PredictionResponse, HealthResponse
from app.model import model_service

START_TIME = time.time()

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Load model on startup, clean up on shutdown."""
    model_service.load(
        model_path="models/classifier_v1.joblib",
        version="v1",
    )
    print(f"Model loaded: version {model_service.version}")
    yield
    print("Shutting down, releasing model resources")

app = FastAPI(
    title="ML Prediction API",
    version="1.0.0",
    lifespan=lifespan,
)

@app.get("/health", response_model=HealthResponse)
async def health_check():
    return HealthResponse(
        status="healthy",
        model_loaded=model_service.model is not None,
        model_version=model_service.version,
        uptime_seconds=time.time() - START_TIME,
    )

@app.post("/predict", response_model=PredictionResponse)
async def predict(request: PredictionRequest):
    try:
        result = model_service.predict(request.features)

        if not request.return_probabilities:
            result["probabilities"] = None

        return PredictionResponse(**result)

    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except RuntimeError as e:
        raise HTTPException(status_code=503, detail=str(e))

@app.post("/predict/batch", response_model=list[PredictionResponse])
async def predict_batch(requests: list[PredictionRequest]):
    """Process multiple predictions in a single request."""
    if len(requests) > 100:
        raise HTTPException(
            status_code=400,
            detail="Batch size limited to 100 requests",
        )

    results = []
    for req in requests:
        result = model_service.predict(req.features)
        if not req.return_probabilities:
            result["probabilities"] = None
        results.append(PredictionResponse(**result))

    return results

The lifespan context manager replaces the deprecated @app.on_event pattern. It loads the model before the first request and handles cleanup on shutdown.

Adding Request Logging and Metrics

Track predictions for monitoring and debugging.

# app/middleware.py
import time
import logging
from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware

logger = logging.getLogger("ml_api")

class RequestLoggingMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        start = time.time()
        response = await call_next(request)
        duration_ms = (time.time() - start) * 1000

        logger.info(
            "request",
            extra={
                "method": request.method,
                "path": request.url.path,
                "status": response.status_code,
                "duration_ms": round(duration_ms, 2),
            },
        )

        response.headers["X-Response-Time-Ms"] = str(round(duration_ms, 2))
        return response

# Add to main.py:
# app.add_middleware(RequestLoggingMiddleware)

Containerizing with Docker

FROM python:3.11-slim

WORKDIR /app

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

COPY app/ app/
COPY models/ models/

EXPOSE 8000

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

For production, add a multi-stage build, run as a non-root user, and use gunicorn with uvicorn workers for multi-core utilization:

CMD ["gunicorn", "app.main:app", \
     "-w", "4", \
     "-k", "uvicorn.workers.UvicornWorker", \
     "--bind", "0.0.0.0:8000"]

Set the worker count based on your CPU count and whether inference is CPU-bound or GPU-bound. For GPU models, you typically need just one worker per GPU.

Model Versioning

Support serving multiple model versions behind the same API:

@app.post("/v2/predict", response_model=PredictionResponse)
async def predict_v2(request: PredictionRequest):
    """Serve a newer model version at a separate endpoint."""
    result = model_service_v2.predict(request.features)
    return PredictionResponse(**result)

Alternatively, accept a version parameter in the request and route internally. Either way, never remove an old version endpoint without warning consumers — use deprecation headers and a migration timeline.

Testing the API

# test_api.py
from fastapi.testclient import TestClient
from app.main import app

client = TestClient(app)

def test_health():
    response = client.get("/health")
    assert response.status_code == 200
    data = response.json()
    assert data["status"] == "healthy"
    assert data["model_loaded"] is True

def test_prediction():
    response = client.post("/predict", json={
        "features": [5.1, 3.5, 1.4, 0.2],
        "return_probabilities": True,
    })
    assert response.status_code == 200
    data = response.json()
    assert "prediction" in data
    assert 0.0 <= data["confidence"] <= 1.0

def test_invalid_input():
    response = client.post("/predict", json={"features": []})
    assert response.status_code == 422  # Validation error

Summary

FastAPI provides a solid foundation for serving ML models in production. Define strict input/output schemas with Pydantic. Load models once at startup using the lifespan pattern. Add health check endpoints for orchestrators. Include request logging for observability. Containerize with Docker and use gunicorn with uvicorn workers for multi-core serving. Plan for model versioning from the start so you can upgrade models without breaking consumers.