add health endpoint

dnth · Nov 1, 2024 · 46f5757 · 46f5757
1 parent 7c7eb5d
commit 46f5757
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 342 deletions.
diff --git a/nbs/serving.ipynb b/nbs/serving.ipynb
diff --git a/xinfer/serve.py b/xinfer/serve.py
@@ -1,3 +1,5 @@
+import time
+
 from fastapi import FastAPI
 from loguru import logger
 from pydantic import BaseModel
@@ -25,7 +27,10 @@ def __init__(
         model_id,
         **kwargs,
     ):
-        self.model = create_model(model_id, **kwargs)
+        try:
+            self.model = create_model(model_id, **kwargs)
+        except Exception as e:
+            raise RuntimeError(f"Failed to load model {model_id}: {str(e)}")
 
     @app.post("/infer")
     async def infer(self, request: InferRequest) -> dict:
@@ -45,6 +50,16 @@ async def infer_batch(self, request: InferBatchRequest) -> list[dict]:
         except Exception as e:
             return [{"error": f"An error occurred: {str(e)}"}]
 
+    @app.get("/health")
+    async def health(self):
+        return {
+            "status": "healthy",
+            "timestamp": time.time(),
+            "model_id": self.model.model_id,
+            "device": self.model.device,
+            "dtype": str(self.model.dtype),
+        }
+
 
 def serve_model(
     model_id: str,