How to serve a scikit-learn model with FastAPI

A scikit-learn model becomes easier to test from other applications when it responds through an HTTP endpoint instead of a Python prompt. Serving the persisted estimator with FastAPI gives frontend code, workers, and integration tests a small JSON scoring interface around the same prediction object.

The app loads a trusted joblib artifact at startup, validates each request with a Pydantic model, and returns a named prediction from the Iris classifier. FastAPI handles the request body and response serialization, while Uvicorn runs the local ASGI server.

Use the joblib loading pattern only for artifacts from a trusted training environment, because pickle-based persistence can execute code while loading and requires matching package versions. For untrusted artifacts or non-Python serving targets, prefer a safer persistence path such as skops.io or ONNX.

1. Create an isolated Python environment for the API project.

   $ python3 -m venv .venv

2. Install scikit-learn, joblib, FastAPI, and Uvicorn in the environment.

   $ .venv/bin/python -m pip install --upgrade pip scikit-learn joblib fastapi "uvicorn[standard]"

   The quoted uvicorn[standard] extra installs the recommended server dependencies without relying on shell-specific bracket handling.

3. Create train_model.py in the project directory.

   train_model.py

   from joblib import dump
   from sklearn.datasets import load_iris
   from sklearn.linear_model import LogisticRegression
   from sklearn.pipeline import make_pipeline
   from sklearn.preprocessing import StandardScaler
    
   iris = load_iris()
   model = make_pipeline(
       StandardScaler(),
       LogisticRegression(max_iter=200, random_state=0),
   )
   model.fit(iris.data, iris.target)
    
   dump(
       {
           "model": model,
           "target_names": iris.target_names.tolist(),
       },
       "iris_model.joblib",
   )
    
   print("saved iris_model.joblib")
   print("classes: setosa, versicolor, virginica")

   The pipeline keeps scaling and classification together so the API uses the same preprocessing path during prediction.

4. Run the training script to write the model artifact.

   $ .venv/bin/python train_model.py
   saved iris_model.joblib
   classes: setosa, versicolor, virginica

5. Create main.py beside the model artifact.

   main.py

   from fastapi import FastAPI
   from joblib import load
   from pydantic import BaseModel, Field
    
   artifact = load("iris_model.joblib")
   model = artifact["model"]
   target_names = artifact["target_names"]
    
   app = FastAPI(title="Iris model API")
    
    
   class PredictionRequest(BaseModel):
       features: list[float] = Field(..., min_length=4, max_length=4)
    
    
   class PredictionResponse(BaseModel):
       prediction: str
    
    
   @app.get("/health")
   def health() -> dict[str, str]:
       return {"status": "ok"}
    
    
   @app.post("/predict", response_model=PredictionResponse)
   def predict(request: PredictionRequest) -> PredictionResponse:
       predicted_class = int(model.predict([request.features])[0])
       return PredictionResponse(prediction=target_names[predicted_class])

   Load only model files created by a trusted training process. joblib uses pickle-based loading, so a malicious artifact can run code during import.

6. Start the local FastAPI server with Uvicorn.

   $ .venv/bin/uvicorn main:app --host 127.0.0.1 --port 8000
   INFO:     Started server process [12345]
   INFO:     Waiting for application startup.
   INFO:     Application startup complete.
   INFO:     Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)

   Keep this terminal open while testing the endpoint.

7. Check the health endpoint from another terminal.

   $ curl --silent http://127.0.0.1:8000/health
   {"status":"ok"}

8. Send a JSON feature vector to the prediction endpoint.

   $ curl --silent --request POST http://127.0.0.1:8000/predict \
     --header 'Content-Type: application/json' \
     --data '{"features":[5.1,3.5,1.4,0.2]}'
   {"prediction":"setosa"}

   The four numbers match the Iris feature order: sepal length, sepal width, petal length, and petal width.

9. Stop the local server with Ctrl+C after the smoke test.

   ^C
   INFO:     Shutting down
   INFO:     Waiting for application shutdown.
   INFO:     Application shutdown complete.
   INFO:     Finished server process [12345]