Simplest MLflow serving with FastAPI

serving.py

import os

import mlflow
import uvicorn
from fastapi import FastAPI, HTTPException


class MLService:
    def __init__(self, model):
        self.model = model
        self.app = FastAPI()

        @self.app.post("/invocations/")
        async def invocations(payload: dict):
            try:
                return int(self.model.predict(payload))
            except Exception as e:
                raise HTTPException(
                    status_code=400,
                    detail=str(e),
                )


if __name__ == "__main__":
    pyfunc_model = mlflow.pyfunc.load_model(os.environ["MODEL_PATH"])
    ml_service = MLService(pyfunc_model)

    uvicorn.run(ml_service.app, host="0.0.0.0", port=5000, log_level="info")