diff --git a/mlem/contrib/fastapi.py b/mlem/contrib/fastapi.py index 09a628fb..04859169 100644 --- a/mlem/contrib/fastapi.py +++ b/mlem/contrib/fastapi.py @@ -13,6 +13,7 @@ from fastapi import FastAPI, UploadFile from fastapi.datastructures import Default from fastapi.responses import RedirectResponse +from prometheus_fastapi_instrumentator import Instrumentator from pydantic import BaseModel, create_model, parse_obj_as from pydantic.typing import get_args from starlette.responses import JSONResponse, Response, StreamingResponse @@ -58,6 +59,8 @@ class FastAPIServer(Server, LibRequirementsMixin): """Network interface to use""" port: int = 8080 """Port to use""" + metrics: bool = False + """Generate /metrics endpoint to be scraped by Prometheus""" @classmethod def _create_handler_executor( @@ -197,6 +200,12 @@ def app_init(self, interface: Interface): response_class=response_class or Default(JSONResponse), ) + if self.metrics: + + @app.on_event("startup") + async def startup(): + Instrumentator().instrument(app).expose(app) + return app def serve(self, interface: Interface): diff --git a/setup.py b/setup.py index f3b3754f..41821820 100644 --- a/setup.py +++ b/setup.py @@ -73,7 +73,7 @@ "catboost": ["catboost"], "xgboost": ["xgboost"], "lightgbm": ["lightgbm"], - "fastapi": ["uvicorn", "fastapi"], + "fastapi": ["uvicorn", "fastapi", "prometheus-fastapi-instrumentator"], "streamlit": ["uvicorn", "fastapi", "streamlit", "streamlit_pydantic"], "sagemaker": ["docker", "boto3", "sagemaker"], "torch": ["torch"],