add health check endpoints (#64)

* Add probes

* add tests

* add tests and checks to ci

.github/workflows/build-vector-serve.yml

@@ -25,8 +25,25 @@ defaults:
     working-directory: ./vector-serve/
 
 jobs:
+  tests:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python 3.11.1
+      uses: actions/setup-python@v5
+      with:
+        python-version: 3.11.1
+    - name: Setup
+      run: make setup
+    - name: Init Model Cache
+      run: make download.models
+    - name: Lints
+      run: make check
+    - name: Tests
+      run: make test
   build_and_push:
     name: Build and push images
+    needs: tests
     runs-on:
       - self-hosted
       - dind

vector-serve/Makefile

@@ -1,5 +1,5 @@
 SOURCE_OBJECTS=app
-
+POETRY_VERSION:=1.7.1
 
 format:
 	poetry run black ${SOURCE_OBJECTS}
@@ -14,3 +14,14 @@ run: download.models
 run.docker:
 	docker build -t vector-serve .
 	docker run -p 3000:3000 vector-serve
+
+test:
+	poetry run pytest
+
+setup:
+	curl -sSL https://install.python-poetry.org | POETRY_VERSION=${POETRY_VERSION} python3 -
+	poetry install
+
+check:
+	poetry run ruff check ${SOURCE_OBJECTS}
+	poetry run black --check ${SOURCE_OBJECTS}

vector-serve/app/app.py

@@ -6,9 +6,10 @@
 
 from app.routes.transform import router as transform_router
 from app.routes.info import router as info_router
+from app.routes.health import router as health_router
+
 from app.models import load_model_cache
 
-import logging
 
 logging.basicConfig(level=logging.DEBUG)
 
@@ -20,8 +21,10 @@
     allow_methods=["*"],
     allow_headers=["*"],
 )
+
 app.include_router(transform_router)
 app.include_router(info_router)
+app.include_router(health_router)
 
 
 def start_app_handler(app: FastAPI) -> Callable:

vector-serve/app/routes/health.py

@@ -0,0 +1,29 @@
+from fastapi import APIRouter
+from pydantic import BaseModel
+import logging
+
+router = APIRouter(tags=["health"])
+
+
+class ReadyResponse(BaseModel):
+    ready: bool
+
+
+class AliveResponse(BaseModel):
+    alive: bool
+
+
+@router.get("/ready", response_model=ReadyResponse)
+def ready() -> ReadyResponse:
+    logging.debug("Health check")
+    return ReadyResponse(
+        ready=True,
+    )
+
+
+@router.get("/alive", response_model=AliveResponse)
+def alive() -> AliveResponse:
+    logging.debug("Health check")
+    return AliveResponse(
+        alive=True,
+    )

vector-serve/pyproject.toml

@@ -17,6 +17,8 @@ torch = "^2.1.2"
 [tool.poetry.group.dev.dependencies]
 ruff = "^0.1.13"
 black = "*"
+pytest = "^8.0.2"
+httpx = "^0.27.0"
 
 [build-system]
 requires = ["poetry-core"]

vector-serve/tests/conftest.py

@@ -0,0 +1,9 @@
+import pytest
+from starlette.testclient import TestClient
+
+from app.app import app
+
+@pytest.fixture()
+def test_client():
+    with TestClient(app) as test_client:
+        yield test_client

vector-serve/tests/test_endpoints.py

@@ -0,0 +1,16 @@
+from fastapi.testclient import TestClient
+from fastapi import FastAPI
+
+def test_ready_endpoint(test_client):
+    response = test_client.get("/ready")
+    assert response.status_code == 200
+    assert response.json() == {"ready": True}
+
+def test_alive_endpoint(test_client):
+    response = test_client.get("/alive")
+    assert response.status_code == 200
+    assert response.json() == {"alive": True}
+
+def test_model_info(test_client):
+    response = test_client.get("/v1/info", params={"model_name": "all-MiniLM-L12-v2"})
+    assert response.status_code == 200