minimal fastapi prom metrics

backend/main.py

@@ -27,6 +27,7 @@
 from oasst_shared.exceptions import OasstError, OasstErrorCode
 from oasst_shared.schemas import protocol as protocol_schema
 from oasst_shared.utils import utcnow
+from prometheus_fastapi_instrumentator import Instrumentator
 from pydantic import BaseModel
 from sqlmodel import Session, select
 from starlette.middleware.cors import CORSMiddleware
@@ -100,6 +101,13 @@ def create_official_web_api_client():
                 )
 
 
+if settings.ENABLE_PROM_METRICS:
+
+    @app.on_event("startup")
+    async def enable_prom_metrics():
+        Instrumentator().instrument(app).expose(app)
+
+
 if settings.RATE_LIMIT:
 
     @app.on_event("startup")

backend/oasst_backend/config.py

@@ -197,6 +197,8 @@ class Settings(BaseSettings):
 
     ROOT_TOKENS: List[str] = ["1234"]  # supply a string that can be parsed to a json list
 
+    ENABLE_PROM_METRICS: bool = True  # enable prometheus metrics at /metrics
+
     @validator("DATABASE_URI", pre=True)
     def assemble_db_connection(cls, v: Optional[str], values: Dict[str, Any]) -> Any:
         if isinstance(v, str):

backend/requirements.txt

@@ -6,6 +6,7 @@ fastapi-limiter==0.1.5
 fastapi-utils==0.2.1
 loguru==0.6.0
 numpy==1.22.4
+prometheus-fastapi-instrumentator==5.9.1
 psycopg2-binary==2.9.5
 pydantic==1.9.1
 pydantic[email]==1.9.1

inference/server/main.py

@@ -9,10 +9,18 @@
 from fastapi.middleware.cors import CORSMiddleware
 from loguru import logger
 from oasst_shared.schemas import inference, protocol
+from prometheus_fastapi_instrumentator import Instrumentator
 from sse_starlette.sse import EventSourceResponse
 
 app = fastapi.FastAPI()
 
+
+# add prometheus metrics at /metrics
+@app.on_event("startup")
+async def enable_prom_metrics():
+    Instrumentator().instrument(app).expose(app)
+
+
 # Allow CORS
 app.add_middleware(
     CORSMiddleware,

inference/server/requirements.txt

@@ -1,5 +1,6 @@
 fastapi[all]
 loguru
+prometheus-fastapi-instrumentator==5.9.1
 pydantic
 redis
 sse-starlette