FastAPI

import os
from contextlib import asynccontextmanager
from typing import Dict

import joblib
import pandas as pd
from fastapi import FastAPI

model = None
MODEL_PATH = os.getenv("MODEL_PATH", "iris_model.joblib")

def load_model():
    _model = joblib.load(MODEL_PATH)
    return _model

"""
Here we are loading the model in global variable model using 
FastAPI's lifespan context manager (https://fastapi.tiangolo.com/advanced/events/#lifespan-events).
This makes sure we load the model before server is ready to serve requests.
"""
@asynccontextmanager
async def lifespan(app: FastAPI):
    global model
    model = load_model()
    yield


app = FastAPI(lifespan=lifespan, root_path=os.getenv("TFY_SERVICE_ROOT_PATH", ""))

# Add a health check endpoint which will be used for readiness and liveness 
# probe in the deployment
@app.get("/health")
async def health() -> Dict[str, bool]:
    return {"healthy": True}

# This is the main prediction endpoint that is used to serve the model
@app.post("/predict")
def predict(
    sepal_length: float, sepal_width: float, petal_length: float, petal_width: float
):
    global model
    class_names = ["setosa", "versicolor", "virginica"]
    data = dict(
        sepal_length=sepal_length,
        sepal_width=sepal_width,
        petal_length=petal_length,
        petal_width=petal_width,
    )
    prediction = model.predict_proba(pd.DataFrame([data]))[0]
    predictions = []
    for label, confidence in zip(class_names, prediction):
        predictions.append({"label": label, "score": confidence})
    return {"predictions": predictions}