Deploy a Machine Learning Model Using Sklearn, FastAPI, and Docker
Learn how to build a REST API to deploy scikit-learn models using FastAPI and Docker. From training to a production-ready containerized deployment in minutes.
Training a machine learning model is only half the job. The other 50% — the part many data scientists put off — is making that model available for other systems to consume. A Jupyter notebook is not an API. A .pkl file on your desktop is not in production. This article shows you how to close that gap: you will train a scikit-learn classifier, expose it as a validated REST API with FastAPI, and package everything into a Docker container you can deploy anywhere.
Why You Need to Deploy ML Models
When your model lives in a notebook, you are the only one who can use it, and only manually. To integrate it with a web application, a mobile app, a data pipeline, or any other service, you need to expose it through a standard interface. A REST API is the most universal contract: any system that can make an HTTP request can consume your model, regardless of the language it is written in.
FastAPI offers concrete advantages over alternatives like Flask for this use case:
- Automatic data validation with Pydantic
- Automatically generated interactive documentation (Swagger UI at
/docs) - Native async support, useful when the model is slow to respond
- Static typing that prevents bugs in production
Training and Saving the Model with joblib
We start by training an Iris classifier and saving it to disk. joblib is more efficient than pickle for objects that contain large NumPy arrays (like sklearn's decision trees). If the training step itself is new to you, Your First Machine Learning Model with Scikit-learn walks through it in detail.
# train_model.py
import joblib
import numpy as np
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
# Load data and train
iris = load_iris()
X_train, X_test, y_train, y_test = train_test_split(
iris.data, iris.target, test_size=0.2, random_state=42
)
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train, y_train)
# Verify performance
accuracy = accuracy_score(y_test, model.predict(X_test))
print(f"Test accuracy: {accuracy:.4f}")
# Save model and class metadata
joblib.dump(model, "model/classifier.joblib")
joblib.dump(iris.target_names.tolist(), "model/class_names.joblib")
print("Model saved to model/classifier.joblib")
Run this script first:
mkdir model
python train_model.py
# Test accuracy: 1.0000
# Model saved to model/classifier.joblib
Building the FastAPI App
With the model saved, we build the API. The key is to define input and output schemas with Pydantic: FastAPI uses them to automatically validate every request and generate the documentation.
# app/main.py
from contextlib import asynccontextmanager
from typing import List
import joblib
import numpy as np
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
# Input schema: the four measurements of an iris flower
class IrisFeatures(BaseModel):
sepal_length: float = Field(..., gt=0, description="Sepal length in cm")
sepal_width: float = Field(..., gt=0, description="Sepal width in cm")
petal_length: float = Field(..., gt=0, description="Petal length in cm")
petal_width: float = Field(..., gt=0, description="Petal width in cm")
model_config = {
"json_schema_extra": {
"example": {
"sepal_length": 5.1,
"sepal_width": 3.5,
"petal_length": 1.4,
"petal_width": 0.2
}
}
}
# Output schema
class PredictionResponse(BaseModel):
prediction: int
class_name: str
probabilities: List[float]
# Load model at application startup
ml_models = {}
@asynccontextmanager
async def lifespan(app: FastAPI):
ml_models["classifier"] = joblib.load("model/classifier.joblib")
ml_models["class_names"] = joblib.load("model/class_names.joblib")
yield
ml_models.clear()
app = FastAPI(
title="Iris Classifier API",
description="Classifies iris flowers using a Random Forest",
version="1.0.0",
lifespan=lifespan
)
@app.get("/health")
def health_check():
return {"status": "healthy", "model_loaded": "classifier" in ml_models}
@app.post("/predict", response_model=PredictionResponse)
def predict(features: IrisFeatures):
if "classifier" not in ml_models:
raise HTTPException(status_code=503, detail="Model unavailable")
# Prepare features as a 2D array
X = np.array([[
features.sepal_length,
features.sepal_width,
features.petal_length,
features.petal_width
]])
model = ml_models["classifier"]
prediction = int(model.predict(X)[0])
probabilities = model.predict_proba(X)[0].tolist()
class_name = ml_models["class_names"][prediction]
return PredictionResponse(
prediction=prediction,
class_name=class_name,
probabilities=probabilities
)
@app.post("/predict/batch", response_model=List[PredictionResponse])
def predict_batch(features_list: List[IrisFeatures]):
if "classifier" not in ml_models:
raise HTTPException(status_code=503, detail="Model unavailable")
X = np.array([[f.sepal_length, f.sepal_width, f.petal_length, f.petal_width]
for f in features_list])
model = ml_models["classifier"]
predictions = model.predict(X).tolist()
probabilities = model.predict_proba(X).tolist()
class_names = ml_models["class_names"]
return [
PredictionResponse(
prediction=pred,
class_name=class_names[pred],
probabilities=probs
)
for pred, probs in zip(predictions, probabilities)
]
Testing the API Locally with uvicorn
Install the dependencies and start the server:
pip install fastapi uvicorn joblib scikit-learn numpy
uvicorn app.main:app --reload --host 0.0.0.0 --port 8000
Open http://localhost:8000/docs for the interactive documentation. You can also test with curl:
curl -X POST "http://localhost:8000/predict" \
-H "Content-Type: application/json" \
-d '{"sepal_length": 5.1, "sepal_width": 3.5, "petal_length": 1.4, "petal_width": 0.2}'
Expected response:
{
"prediction": 0,
"class_name": "setosa",
"probabilities": [0.97, 0.02, 0.01]
}
Writing the Dockerfile
To containerize the API we need a Dockerfile. We use python:3.12-slim to keep the image small:
# Dockerfile
FROM python:3.12-slim
# Create a non-root user for security
RUN useradd --create-home appuser
WORKDIR /app
# Install dependencies first (leverages Docker layer cache)
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
# Copy application code and model
COPY app/ ./app/
COPY model/ ./model/
# Switch to non-root user
USER appuser
# Expose port
EXPOSE 8000
# Start command
CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000"]
The requirements.txt file:
fastapi==0.115.6
uvicorn[standard]==0.34.0
scikit-learn==1.6.1
joblib==1.5.0
numpy==2.2.3
pydantic==2.11.3
Building and Running the Docker Container
# Build the image
docker build -t iris-classifier:v1 .
# Run the container
docker run -d \
--name iris-api \
-p 8000:8000 \
iris-classifier:v1
# Verify it is running
docker ps
docker logs iris-api
Testing the Dockerized API
With the container running, the same tests work:
# Health check
curl http://localhost:8000/health
# {"status":"healthy","model_loaded":true}
# Single prediction
curl -X POST "http://localhost:8000/predict" \
-H "Content-Type: application/json" \
-d '{"sepal_length": 6.7, "sepal_width": 3.0, "petal_length": 5.2, "petal_width": 2.3}'
# {"prediction":2,"class_name":"virginica","probabilities":[0.0,0.05,0.95]}
Or using Python requests:
import requests
url = "http://localhost:8000/predict/batch"
payload = [
{"sepal_length": 5.1, "sepal_width": 3.5, "petal_length": 1.4, "petal_width": 0.2},
{"sepal_length": 6.7, "sepal_width": 3.0, "petal_length": 5.2, "petal_width": 2.3},
]
response = requests.post(url, json=payload)
for result in response.json():
print(f"{result['class_name']}: {max(result['probabilities']):.2%} confidence")
# setosa: 97.00% confidence
# virginica: 95.00% confidence
Production Considerations
The code above is functional, but a real production environment requires a few additional adjustments.
CORS: if your API is consumed from a frontend on another domain, configure CORS:
from fastapi.middleware.cors import CORSMiddleware
app.add_middleware(
CORSMiddleware,
allow_origins=["https://your-domain.com"],
allow_methods=["POST", "GET"],
allow_headers=["*"],
)
Authentication: for non-public APIs, add an API key header:
import os
from fastapi import Security
from fastapi.security.api_key import APIKeyHeader
api_key_header = APIKeyHeader(name="X-API-Key")
@app.post("/predict")
def predict(features: IrisFeatures, api_key: str = Security(api_key_header)):
if api_key != os.environ["API_KEY"]:
raise HTTPException(status_code=403, detail="Invalid key")
# ... rest of the handler
Scaling: run uvicorn with multiple workers or behind gunicorn (add gunicorn to requirements.txt):
CMD ["gunicorn", "app.main:app", "-w", "4", "-k", "uvicorn.workers.UvicornWorker", "--bind", "0.0.0.0:8000"]
Conclusion
Going from a trained model to a containerized REST API takes less than an hour following this workflow: train with sklearn, save with joblib, deploy with FastAPI, package with Docker. The structure is simple enough to iterate quickly and robust enough to deploy to any container orchestrator (Kubernetes, Azure Container Apps, AWS ECS). The next step is adding a CI/CD pipeline that rebuilds the image every time you retrain the model — retraining you can automate with Prefect — and watching the deployed model for drift: ML Model Monitoring in Production with Evidently.
About the author
Rodrigo Arenas is a software architect and machine learning engineer in Medellín, Colombia. He builds products, platforms, and open-source software for AI, Data Engineering, and Machine Learning — creator of Ciaren, sklearn-genetic-opt, and PyWorkforce.
Recommended reading
Rodrigo Arenas — ML Engineer & Open Source Maintainer
Working on something similar? I can help you take it to production.