How to use FastAPI for machine learning inference, Docker-ized and deployed using Kubernetes
FastAPI is fast to build, iterate, and deploy. It uses Python types to automatically do validation and generate OpenAPI docs that are interactive. For machine learning model inference, it’s a great way to quickly get a REST API running in development. It’s also easy and fast to Dockerize it and deploy it to Kubernetes for production.
From the FastAPI docs:
By declaring parameters with type hints, you get:
And FastAPI uses the same declarations to:
Check our main.py in the demo repo here
for this example of setting up FastAPI with a redirect to the auto-built /docs.
Plus 2 simple health check endpoints, one for the load balancer and one for the app.
import os
from datetime import datetime, timezone
from fastapi import FastAPI
from starlette.responses import RedirectResponse
from app.env import load_env_vars
app = FastAPI()
load_env_vars()
booted_at = {}
@app.on_event("startup")
def startup_event():
booted_at["time"] = datetime.now(timezone.utc)
# Redirect base url to /docs
@app.get("/")
def redirect_to_docs():
return RedirectResponse(url="/docs")
# ...
@app.get("/health_check")
def health_check():
return {
"aws_region": os.environ.get("AWS_REGION"),
"booted_at": booted_at.get("time"),
"health": "OK",
"k8s_env": os.environ.get("K8S_ENV"),
"python_env": os.environ["PYTHON_ENV"],
"titanic_model_version": TITANIC_MODEL_VERSION,
"version": os.environ.get(
"APP_REVISION", "Missing $APP_REVISION env var, not set"
),
}
# Don't render JSON for load balancer health check, just return 200
@app.get("/health_check/load_balancer")
def health_check_load_balancer():
return
Next in main.py we add an inference endpoint, something like:
@app.post(
"/predict_survival_by_passengers", response_model=SurvivalPredictionByPassengerIds
)
def predict_survival_by_passengers(passengers: list[Passenger]):
# Convert passengers to a list of dicts; this is what the model expects
# There must be a better way to do this.
passengers = [passenger.dict() for passenger in passengers]
# Make predictions
results = [] if not passengers else titanic_survival_predict(passengers)
return {
"predictions": results,
"metadata": {
"titanic_model_version": TITANIC_MODEL_VERSION,
},
}
That endpoint simply calls the predict() function (here as titanic_survival_predict()) that loads the model and makes predictions. See demo app for more code
FastAPI handles the type checking, validation, and most errors for us automatically.
I’ve found poetry to be a best way to manage dependencies and you’ll see poetry commands in the Dockerfile below.
Here’s the Dockerfile:
# Versions with defaults. Override with env var to build a different version.
ARG PYTHON_VERSION=3.11
# More args
# For security, set a non-root user. Name is arbitrary.
ARG USER=nonroot
ARG USER_ID=1001
ARG PYTHON_ENV=production
# Use the official lightweight Python image.
# https://hub.docker.com/_/python
FROM python:${PYTHON_VERSION}-slim-buster
# Args needed for this container
ARG USER
ARG USER_ID
ARG PYTHON_ENV
ENV PYTHON_ENV=${PYTHON_ENV}
# Recommended by hadolint
# https://github.com/hadolint/hadolint/wiki/DL4006
SHELL ["/bin/bash", "-o", "pipefail", "-c"]
# Add non-root user
RUN groupadd --gid $USER_ID $USER \
&& useradd --uid $USER_ID --gid $USER --shell /bin/bash --create-home $USER
# Create a directory for the app code
RUN mkdir /app \
&& chown -R $USER:$USER /app
WORKDIR /app
# Copy dependency definitions for production-only requirements
COPY --chown=$USER:$USER pyproject.toml ./pyproject.toml
# Install poetry
RUN pip install poetry
# Don't create a virtualenv, we'll use the container's python so everything else just works
RUN poetry config virtualenvs.create false
# Install production-only dependencies
RUN poetry install --without dev,test
# COPY the app code that is needed to start the server and run the app
# This will be at `app/app` in the container.
COPY --chown=$USER:$USER app app
COPY --chown=$USER:$USER .env.production .env.production
# Set user to non-root $USER
# This needs to be the numeric uid, not the username, for the k8s
# securityContext: runAsNonRoot check to work.
USER $USER_ID
# https://www.uvicorn.org/deployment/
CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "80"]
Assuming you already have an existing Kubernetes cluster, it’s straightforward to deploy this app. Check out the demo app’s kubernetes/ directory for the deployment and service specs.
I always recommend following the NSA/CISA Kubernetes Hardening Guidance as much as possible.
In the demo app, you’ll notice it follows some key security recommendations:
The demo app does the checked ones above and everyone really should. They’re straightforward to implement. See the Dockerfile for “nonroot” USER setup and USER_ID
The other parts are as simple as passing some options in the K8s Deployment spec
spec:
securityContext:
# Must match Dockerfile's USER_ID for User and Group
runAsUser: 1001
runAsGroup: 1001
# Set ownership of mounted volumes to the user running the container
fsGroup: 1001
# https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.21/#container-v1-core
containers:
- name: demo-fastapi-ml-container
image: ghcr.io/briansigafoos/docker-fastapi-ml-inference
securityContext:
allowPrivilegeEscalation: false
readOnlyRootFilesystem: true
runAsNonRoot: true
FastAPI is a great framework for building APIs, including for machine learning model inference. It’s easy to use, has great documentation, and is fast. It’s also easy to Dockerize and to deploy to Kubernetes.
Check out the demo app code and let me know what you think.