DEV Community

qing
qing

Posted on

Build a Machine Learning Model API with FastAPI

Build a Machine Learning Model API with FastAPI

Imagine you’ve spent weeks training a machine learning model, tweaking hyperparameters, and watching accuracy curves climb. But now it sits in a .pkl file on your laptop, useless to anyone else. That’s the gap between a cool experiment and a real product: deployment. The fastest, most modern way to bridge that gap is FastAPI, a framework that lets you wrap your model in a production-ready API with minimal boilerplate and maximum speed.

In this guide, you’ll build a complete Machine Learning Model API using FastAPI and Scikit-Learn. By the end, you’ll have a working endpoint that accepts input data, runs a prediction, and returns JSON—ready to be consumed by a frontend, mobile app, or another service. Let’s get your model out of the notebook and into the world.

Why FastAPI for Machine Learning?

FastAPI isn’t just another Python web framework. It’s built on Starlette for the HTTP layer and Pydantic for data validation, giving you three critical advantages for ML:

  • Performance: FastAPI is among the fastest Python frameworks available, thanks to asynchronous support and native JSON handling.
  • Automatic Validation: Pydantic models automatically validate and serialize incoming data, so you never have to write manual type checks.
  • Interactive Docs: FastAPI generates Swagger UI and ReDoc documentation automatically, letting you test endpoints instantly without writing extra test code.

For ML models that might need to serve thousands of requests per second, these features are non-negotiable.

Step 1: Set Up Your Environment

First, create a virtual environment and install the required packages:

python -m venv ml-api-env
source ml-api-env/bin/activate  # On Windows: ml-api-env\Scripts\activate
pip install fastapi uvicorn scikit-learn pydantic
Enter fullscreen mode Exit fullscreen mode
  • fastapi: The core framework.
  • uvicorn: The ASGI server that runs FastAPI.
  • scikit-learn: For our sample ML model.
  • pydantic: For data validation (already included with FastAPI).

Step 2: Create a Dummy ML Model

You don’t need a real trained model to start. Let’s create a simple house price predictor using Scikit-Learn. This will be saved as a .pkl file so FastAPI can load it later.

# train_model.py
from sklearn.linear_model import LinearRegression
import pickle

# Sample data: square feet -> price
X = [[500], [600], [700], [800], [900], [1000]]
y = [100000, 120000, 140000, 160000, 180000, 200000]

model = LinearRegression()
model.fit(X, y)

# Save the model
with open("house_price_model.pkl", "wb") as f:
    pickle.dump(model, f)

print("Model trained and saved as 'house_price_model.pkl'")
Enter fullscreen mode Exit fullscreen mode

Run this script once to generate the model file:

python train_model.py
Enter fullscreen mode Exit fullscreen mode

Step 3: Build the FastAPI Application

Now, create your main API file: api.py. This script will:

  1. Load the saved model.
  2. Define input validation with Pydantic.
  3. Create a /predict endpoint that returns predictions.
# api.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import pickle

# Initialize FastAPI app
app = FastAPI(title="House Price Prediction API", version="1.0.0")

# Load the trained model
model = None

@app.on_event("startup")
def load_model():
    with open("house_price_model.pkl", "rb") as f:
        model = pickle.load(f)

# Define input data model
class HouseInput(BaseModel):
    square_feet: int

# Prediction endpoint
@app.post("/predict")
def predict_price(input_data: HouseInput):
    if model is None:
        raise HTTPException(status_code=500, detail="Model not loaded")

    # Prepare data for model
    X = [[input_data.square_feet]]
    prediction = model.predict(X)[0]

    return {
        "square_feet": input_data.square_feet,
        "predicted_price": round(prediction, 2)
    }

# Health check endpoint
@app.get("/")
def read_root():
    return {"message": "Welcome to the House Price Prediction API"}
Enter fullscreen mode Exit fullscreen mode

Step 4: Run and Test Your API

Start the server with Uvicorn:

uvicorn api:app --reload
Enter fullscreen mode Exit fullscreen mode

The --reload flag enables auto-reloading on code changes, perfect for development.

Once running, open your browser to http://localhost:8000. You’ll see:

  • The root endpoint (/) with a welcome message.
  • Swagger UI at http://localhost:8000/docs, where you can test the /predict endpoint interactively.

Testing with a Real Request

Use curl to send a prediction request:

curl -X POST "http://localhost:8000/predict" \
     -H "Content-Type: application/json" \
     -d '{"square_feet": 750}'
Enter fullscreen mode Exit fullscreen mode

Expected response:

{
  "square_feet": 750,
  "predicted_price": 150000.0
}
Enter fullscreen mode Exit fullscreen mode

You can also test it from Python using requests:

import requests

response = requests.post(
    "http://localhost:8000/predict",
    json={"square_feet": 750}
)
print(response.json())
Enter fullscreen mode Exit fullscreen mode

Step 5: Production Tips for Real ML APIs

While this example uses a simple model and local file loading, real-world deployments require more robustness. Here’s what to add next:

Feature Why It Matters
Async Functions Prevent blocking when the model takes time to compute [6]
Logging & Monitoring Track errors, latency, and prediction volume [7]
Dockerization Ensure consistent environments across dev and production [8]
Model Storage in AWS S3 Avoid storing large models in your app directory [7]
Background Tasks Run heavy preprocessing/postprocessing without blocking the response [7]
Unit Tests Validate that your API returns correct predictions for known inputs [7]

For production, consider using multistage Docker builds to reduce image size and storing models in cloud storage like AWS S3 instead of local files [7][8].

Your Model is Now Live

You’ve just built a fully functional ML API that can be consumed by any client. No more Jupyter notebooks—your model is now a service.

What’s Next?

  • Replace the dummy model with your own trained model (e.g., from TensorFlow, PyTorch, or XGBoost).
  • Add input preprocessing (e.g., scaling, encoding) inside the endpoint [9].
  • Implement authentication and rate limiting for public APIs.
  • Deploy to cloud platforms like AWS, Google Cloud, or Azure using Docker [8].

Call to Action

Try it today: Copy the code above, run it on your machine, and send a prediction request. Then, share your API endpoint in the comments or on Twitter—tag me so I can see what you built!

If you found this helpful, check out my other posts on deploying ML models with Docker and building async ML pipelines. Let’s keep pushing ML from research to real-world impact. 🚀


If you found this helpful, consider buying me a coffee ☕ — it keeps these articles coming!

Also check out my AI tools collection: AI 次元世界 — free AI tools for developers.


💡 Related: **Content Creator Ultimate Bundle (Save 33%)* — $29.99*

Top comments (0)