Skip to content
Python Central Hub

Building an ML API with Flask/FastAPI

Once your model is saved to disk, the most common way to make it usable by other software is to put it behind an HTTP API. A client sends features as JSON, your server loads the pipeline once at startup, runs .predict().predict(), and sends the result back as JSON. This page builds that endpoint twice — once with FastAPI, once with Flask — and covers the practical details (validation, logging, versioning) that matter once real traffic starts hitting it.

Why APIs are common for ML deployment

An API lets other services call your model:

  • web apps
  • mobile apps
  • internal tools

The minimal contract

Inputs:

  • JSON payload with features

Outputs:

  • prediction (and optionally probability)
diagram Diagram mermaid

FastAPI is popular because:

  • automatic docs
  • type hints
  • validation
app.py (FastAPI)
# Requires: fastapi, uvicorn, joblib
from fastapi import FastAPI
from pydantic import BaseModel
import joblib
 
app = FastAPI()
model = joblib.load("model.joblib")
 
class Features(BaseModel):
    age: int
    income: float
    city: str
    plan: str
 
@app.post("/predict")
def predict(features: Features):
    X = [features.model_dump()]
    pred = model.predict(X)[0]
    return {"prediction": int(pred)}
app.py (FastAPI)
# Requires: fastapi, uvicorn, joblib
from fastapi import FastAPI
from pydantic import BaseModel
import joblib
 
app = FastAPI()
model = joblib.load("model.joblib")
 
class Features(BaseModel):
    age: int
    income: float
    city: str
    plan: str
 
@app.post("/predict")
def predict(features: Features):
    X = [features.model_dump()]
    pred = model.predict(X)[0]
    return {"prediction": int(pred)}

Flask example (minimal)

app.py (Flask)
# Requires: flask, joblib
from flask import Flask, request, jsonify
import joblib
 
app = Flask(__name__)
model = joblib.load("model.joblib")
 
@app.post("/predict")
def predict():
    payload = request.get_json(force=True)
    pred = model.predict([payload])[0]
    return jsonify({"prediction": int(pred)})
app.py (Flask)
# Requires: flask, joblib
from flask import Flask, request, jsonify
import joblib
 
app = Flask(__name__)
model = joblib.load("model.joblib")
 
@app.post("/predict")
def predict():
    payload = request.get_json(force=True)
    pred = model.predict([payload])[0]
    return jsonify({"prediction": int(pred)})

Calling the API with curl

Once either server is running (e.g. uvicorn app:app --reloaduvicorn app:app --reload for FastAPI, or flask --app app runflask --app app run for Flask), you can test it from the command line:

Test the /predict endpoint
curl -X POST http://127.0.0.1:8000/predict \
  -H "Content-Type: application/json" \
  -d '{"age": 35, "income": 50000, "city": "Pune", "plan": "Pro"}'
 
# -> {"prediction":1}
Test the /predict endpoint
curl -X POST http://127.0.0.1:8000/predict \
  -H "Content-Type: application/json" \
  -d '{"age": 35, "income": 50000, "city": "Pune", "plan": "Pro"}'
 
# -> {"prediction":1}

Visualize it

Here’s the round trip a request takes when a client asks your API for a prediction.

diagram ML API request flow mermaid
How a client's JSON request becomes a prediction returned from a Flask or FastAPI endpoint.

Handling bad input gracefully

Pydantic rejects malformed JSON automatically, but it can’t catch everything — a request might be perfectly well-typed and still crash your model (an unseen category, a divide-by-zero in a custom feature, a NaNNaN). Wrap the prediction call so a bad request returns a clean 4xx4xx/5xx5xx response instead of a stack trace leaking to the client:

app.py (error handling)
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import joblib
 
app = FastAPI()
model = joblib.load("model.joblib")
 
class Features(BaseModel):
    age: int
    income: float
 
@app.post("/predict")
def predict(features: Features):
    try:
        X = [features.model_dump()]
        pred = model.predict(X)[0]
    except Exception as exc:
        # never let a raw exception/traceback reach the client
        raise HTTPException(status_code=400, detail=f"prediction failed: {exc}")
    return {"prediction": int(pred)}
app.py (error handling)
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import joblib
 
app = FastAPI()
model = joblib.load("model.joblib")
 
class Features(BaseModel):
    age: int
    income: float
 
@app.post("/predict")
def predict(features: Features):
    try:
        X = [features.model_dump()]
        pred = model.predict(X)[0]
    except Exception as exc:
        # never let a raw exception/traceback reach the client
        raise HTTPException(status_code=400, detail=f"prediction failed: {exc}")
    return {"prediction": int(pred)}

A health check endpoint

Before a load balancer, a Kubernetes pod, or a teammate’s script sends real traffic to your API, it usually wants a cheap way to ask “are you alive, and is the model actually loaded?” A /health/health route that returns instantly (no model inference) is the standard way to answer that:

app.py (health check)
@app.get("/health")
def health():
    return {"status": "ok", "model_loaded": model is not None}
app.py (health check)
@app.get("/health")
def health():
    return {"status": "ok", "model_loaded": model is not None}

Orchestration tools poll endpoints like this on a schedule and stop routing traffic to an instance that fails the check — the same idea Géron describes for TF Serving’s own readiness signals, just implemented by hand here.

Serving more than one model version (canary requests)

Once you have a model_v2.joblibmodel_v2.joblib you want to try cautiously, you don’t have to cut over every client at once. Load both versions and let the caller (or a small percentage of random traffic) opt into the new one — this is exactly the pattern Géron describes for TF Serving, where appending /versions/0002/versions/0002 to a model path lets you “test a new version on a small group of users before releasing it widely (this is called a canary)“:

app.py (canary routing)
import random
import joblib
from fastapi import FastAPI
 
app = FastAPI()
models = {
    "v1": joblib.load("model_v1.joblib"),
    "v2": joblib.load("model_v2.joblib"),
}
 
@app.post("/predict")
def predict(features: dict, version: str | None = None):
    # explicit version wins; otherwise send 10% of traffic to the canary
    chosen = version or ("v2" if random.random() < 0.10 else "v1")
    model = models[chosen]
    pred = model.predict([features])[0]
    return {"prediction": int(pred), "model_version": chosen}
app.py (canary routing)
import random
import joblib
from fastapi import FastAPI
 
app = FastAPI()
models = {
    "v1": joblib.load("model_v1.joblib"),
    "v2": joblib.load("model_v2.joblib"),
}
 
@app.post("/predict")
def predict(features: dict, version: str | None = None):
    # explicit version wins; otherwise send 10% of traffic to the canary
    chosen = version or ("v2" if random.random() < 0.10 else "v1")
    model = models[chosen]
    pred = model.predict([features])[0]
    return {"prediction": int(pred), "model_version": chosen}

If v2v2 starts returning worse predictions, roll back by simply routing 0% of traffic to it — no redeploy, no downtime, and no code change beyond flipping that percentage.

Practical tips

  • validate inputs (types, ranges)
  • log requests (careful with PII)
  • version your model artifact
  • add a /health/health route so load balancers and orchestrators know the service is ready
  • roll out risky model changes as a canary before sending 100% of traffic to them

🧪 Try It Yourself

Exercise 1 – Define the Request Schema with Pydantic

Exercise 2 – Load the Model Once at Startup

Exercise 3 – Shape the JSON Response

Next

Continue to Deploying ML Models to Streamlit — for cases where you want a user-facing demo UI instead of (or in addition to) a raw JSON API.

If this helped you, consider buying me a coffee ☕

Buy me a coffee

Was this page helpful?

Let us know how we did