How to Serve a Machine Learning Model with FastAPI

Knowledgebase

FastAPI is a modern, fast Python web framework well-suited for wrapping a machine learning model in a proper HTTP API — letting other applications call your model over the network rather than embedding it directly.

Prerequisites

Python installed (see How to Install Python on Ubuntu & Debian)
A trained model or an AI service (like Ollama) already working locally

Step 1 — Install FastAPI and Uvicorn

pip install fastapi uvicorn --break-system-packages

Step 2 — Create a Basic API

nano main.py

from fastapi import FastAPI
from pydantic import BaseModel
import requests

app = FastAPI()

class Query(BaseModel):
    prompt: str

@app.post("/generate")
def generate(query: Query):
    response = requests.post("http://localhost:11434/api/generate", json={
        "model": "llama3.2",
        "prompt": query.prompt,
        "stream": False
    })
    return {"response": response.json()["response"]}

This example wraps a local Ollama installation, but the same pattern applies to any model or ML library.

Step 3 — Run the API

uvicorn main:app --host 0.0.0.0 --port 8000

Step 4 — Test the Endpoint

curl -X POST http://localhost:8000/generate \
  -H "Content-Type: application/json" \
  -d '{"prompt": "What is a VPS?"}'

Step 5 — Explore the Automatic API Documentation

http://YOUR_SERVER_IP:8000/docs

FastAPI automatically generates interactive API documentation (via Swagger UI) from your code — useful for testing and for other developers integrating with your API.

Step 6 — Run in Production with Gunicorn (Recommended)

pip install gunicorn --break-system-packages
gunicorn main:app -w 4 -k uvicorn.workers.UvicornWorker --bind 0.0.0.0:8000

See How to Deploy a Python Application with Gunicorn and Nginx for the complete production deployment pattern, including process management and reverse proxy configuration.

Step 7 — Add Basic Input Validation

class Query(BaseModel):
    prompt: str

    class Config:
        str_max_length = 2000

Prevents excessively long inputs from causing performance issues or unexpected behavior.

Step 8 — Add API Key Authentication

from fastapi import Header, HTTPException

API_KEY = "your-secret-key"

@app.post("/generate")
def generate(query: Query, x_api_key: str = Header(None)):
    if x_api_key != API_KEY:
        raise HTTPException(status_code=401, detail="Invalid API key")
    # ... rest of the logic

See How to Set Up API Authentication with JWT for a more robust token-based approach appropriate for production use.

Handling Slow AI Responses Gracefully

AI inference can take several seconds or more — consider async request handling and appropriate client-side timeout configuration, and potentially a queue-based approach (see Message Queue Basics: When and Why to Use One) for longer-running generation tasks.

Common Errors

Request times out on longer generations — increase the timeout on both the FastAPI/Uvicorn side and any reverse proxy in front of it (Nginx's default proxy timeout is often too short for AI workloads).

Continue Reading

Browse more articles in AI & Machine Learning on a VPS.

Knowledgebase

Categories

Categories