Skip to content
Laviyx Learn

Try

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.

Intermediate6 min readUpdated
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.

Terminal
pip install fastapi uvicorn

A 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.

Python
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 predict would make every request slow.
  • The Flower class (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.
  • /health lets 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:

Python
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"]])
Output
{'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:

Python
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())
Output
{'species': ['setosa', 'versicolor']}

Running the server#

Save the API code as main.py and start it:

Terminal
uvicorn main:app --reload

The 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
Python
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())
Output
{'classes': ['setosa', 'versicolor', 'virginica'], 'n_features': 4}

Key takeaway#

Next: MLOps: Models in Production.

Finished this lesson?