ML Model Deployment with FastAPI
Deploy machine learning models as production REST APIs using FastAPI with input validation, async inference, and health checks.
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.
Related articles
- FastAPI Production Deployment of FastAPI with Docker and Gunicorn
Deploy FastAPI to production with Docker, Gunicorn, Uvicorn workers, health checks, multi-stage builds, and best practices.
- Django Django REST Framework vs FastAPI Compared
A practical comparison of DRF and FastAPI: performance, ORM, validation, async, and how to choose for a new Python service.
- FastAPI FastAPI Deployment with Uvicorn and Gunicorn
Deploy FastAPI to production with Gunicorn managing Uvicorn workers. Cover process counts, timeouts, and health checks.
- Machine Learning ML Ensemble Methods Beyond Random Forests
Master stacking, blending, and voting ensembles to combine multiple ML models for better predictive performance.