Machine Learning course menu
Serving a Machine Learning Model with FastAPI
Turn a trained model into a web API with FastAPI: loading the model at startup, validating input with Pydantic, a prediction endpoint, testing with TestClient, batch predictions, health checks, and running the server with uvicorn.
On this page
A saved model file is not useful to a website or mobile app on its own. The common way to make predictions available is a small web API: other programs send input data as JSON over HTTP and receive predictions back. FastAPI is a popular Python framework for this.
pip install fastapi uvicornA prediction API#
This example trains and saves a small model first, so it runs on its own. In a real project, training happens in a separate script and the API only loads the saved file.
import joblib
from fastapi import FastAPI
from pydantic import BaseModel, Field
from sklearn.datasets import load_iris
from sklearn.linear_model import LogisticRegression
# training step (normally a separate script)
iris = load_iris()
joblib.dump(LogisticRegression(max_iter=1000).fit(iris.data, iris.target), "iris_model.joblib")
# the API
app = FastAPI(title="Iris classifier")
model = joblib.load("iris_model.joblib") # load once at startup, not per request
SPECIES = ["setosa", "versicolor", "virginica"]
class Flower(BaseModel):
sepal_length: float = Field(gt=0, lt=20)
sepal_width: float = Field(gt=0, lt=20)
petal_length: float = Field(gt=0, lt=20)
petal_width: float = Field(gt=0, lt=20)
@app.get("/health")
def health():
return {"status": "ok"}
@app.post("/predict")
def predict(flower: Flower):
row = [[flower.sepal_length, flower.sepal_width, flower.petal_length, flower.petal_width]]
probabilities = model.predict_proba(row)[0]
best = int(probabilities.argmax())
return {"species": SPECIES[best], "probability": round(float(probabilities[best]), 3)}Key points:
- The model loads once when the server starts. Loading it inside
predictwould make every request slow. - The
Flowerclass (Pydantic) defines the expected JSON and validates it: missing fields, wrong types or values outside the allowed range are rejected automatically with a clear error. /healthlets monitoring tools check the service is running.
Testing the API#
FastAPI's TestClient sends requests to the app without starting a real server, ideal for automated tests. Continuing from the code above:
from fastapi.testclient import TestClient
client = TestClient(app)
print(client.get("/health").json())
response = client.post("/predict", json={"sepal_length": 5.1, "sepal_width": 3.5,
"petal_length": 1.4, "petal_width": 0.2})
print(response.status_code, response.json())
response = client.post("/predict", json={"sepal_length": 6.7, "sepal_width": 3.0,
"petal_length": 5.6, "petal_width": 2.3})
print(response.status_code, response.json())
bad = client.post("/predict", json={"sepal_length": 5.1, "sepal_width": 3.5, "petal_length": -1})
print(bad.status_code, [(e["loc"][-1], e["type"]) for e in bad.json()["detail"]]){'status': 'ok'}
200 {'species': 'setosa', 'probability': 0.982}
200 {'species': 'virginica', 'probability': 0.974}
422 [('petal_length', 'greater_than'), ('petal_width', 'missing')]The invalid request (a negative petal length and a missing field) is rejected with status 422 before it ever reaches the model, and the error lists exactly which fields are wrong.
Batch predictions#
Sending many rows in one request is far more efficient than one request per row:
from fastapi.testclient import TestClient
@app.post("/predict-batch")
def predict_batch(flowers: list[Flower]):
rows = [[f.sepal_length, f.sepal_width, f.petal_length, f.petal_width] for f in flowers]
return {"species": [SPECIES[i] for i in model.predict(rows)]}
client = TestClient(app)
batch = [{"sepal_length": 5.0, "sepal_width": 3.4, "petal_length": 1.5, "petal_width": 0.2},
{"sepal_length": 5.9, "sepal_width": 2.8, "petal_length": 4.3, "petal_width": 1.3}]
print(client.post("/predict-batch", json=batch).json()){'species': ['setosa', 'versicolor']}Running the server#
Save the API code as main.py and start it:
uvicorn main:app --reloadThe API is then available at http://127.0.0.1:8000, with interactive documentation, generated automatically from your Pydantic models, at http://127.0.0.1:8000/docs. --reload restarts on code changes and is for development only.
For production, the API typically runs in a Docker container on a cloud service, behind HTTPS, with several worker processes. Managed platforms (AWS SageMaker, Google Vertex AI, Azure ML) offer hosted model serving as an alternative to running your own API.
Common mistakes#
Practice#
Practice exercise
Add a GET /model-info endpoint to the iris API that returns the model's class names and number of features, and test it with TestClient.
Show hint
Return a dictionary built from SPECIES and model.n_features_in_.
Show solution
from fastapi.testclient import TestClient
@app.get("/model-info")
def model_info():
return {"classes": SPECIES, "n_features": int(model.n_features_in_)}
print(TestClient(app).get("/model-info").json()){'classes': ['setosa', 'versicolor', 'virginica'], 'n_features': 4}Key takeaway#
Next: MLOps: Models in Production.
Finished this lesson?