[Frontend] Add readiness and liveness endpoints to OpenAI API server

tests/entrypoints/openai/test_basic.py

@@ -59,3 +59,13 @@ async def test_log_metrics(client: openai.AsyncOpenAI):
     response = requests.get(base_url + "/metrics")
 
     assert response.status_code == HTTPStatus.OK
+
+
+@pytest.mark.asyncio
+async def test_get_readiness_ok(client: openai.AsyncOpenAI):
+    """Test the technical route /readiness when the model is fully loaded"""
+    base_url = str(client.base_url)[:-3].strip("/")
+
+    response = requests.get(base_url + "/ready")
+
+    assert response.status_code == HTTPStatus.OK

vllm/entrypoints/openai/api_server.py

@@ -9,7 +9,7 @@
 
 import fastapi
 import uvicorn
-from fastapi import APIRouter, Request
+from fastapi import APIRouter, Request, HTTPException
 from fastapi.exceptions import RequestValidationError
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse, Response, StreamingResponse
@@ -90,6 +90,29 @@ async def health() -> Response:
     return Response(status_code=200)
 
 
+@router.get(
+    "/ready",
+    name="readiness",
+    tags=["technical"],
+)
+async def get_readiness() -> Response:
+    """Readiness probe for k8s"""
+    try :
+        model_executor = openai_serving_chat.engine.engine.model_executor
+        model_runner = model_executor.driver_worker.model_runner
+
+        # check if model weight are loaded in gpu memory
+        model_weights = model_runner.model_memory_usage
+
+        # check if KV cache has been set up
+        num_cpu_blocks = model_runner.num_cpu_blocks
+        num_gpu_blocks = model_runner.num_gpu_blocks
+
+        if model_weights > 0 and num_cpu_blocks > 0  and num_gpu_blocks > 0 :
+            return Response(status_code=200)
+    except: HTTPException(status_code=500, detail="Model not loaded yet or KV cache not setup yet")
+
+
 @router.post("/tokenize")
 async def tokenize(request: TokenizeRequest):
     generator = await openai_serving_tokenization.create_tokenize(request)

vllm/entrypoints/openai/protocol.py

@@ -719,4 +719,4 @@ class DetokenizeRequest(OpenAIBaseModel):
 
 
 class DetokenizeResponse(OpenAIBaseModel):
-    prompt: str
+    prompt: str