feat: Add health check to API and client

Add health check endpoint that checks the health of the database, object
store, and queue, returning a JSON status report. Leveraging this, add a
`health_check` method to the client to verify service health, along with
an "is_healthy" boolean method for convenience.

Signed-off-by: Phoevos Kalemkeris <[email protected]>

Author: phoevos

client/cogstack_model_gateway_client/client.py

@@ -312,6 +312,43 @@ async def deploy_model(
         resp = await self._request("POST", url, json=data)
         return resp.json()
 
+    @require_client
+    async def health_check(self):
+        """Check if the Gateway and its components are healthy and responsive.
+
+        Returns:
+            dict: Health status information with 'status' and 'components' fields.
+                  Status will be 'healthy' or 'unhealthy'.
+
+        Raises:
+            httpx.HTTPStatusError: For HTTP errors other than 503 Service Unavailable.
+            Other exceptions: For network errors, timeouts, etc.
+        """
+        url = f"{self.base_url}/health"
+        try:
+            resp = await self._request("GET", url)
+            return resp.json()
+        except httpx.HTTPStatusError as e:
+            if e.response.status_code == 503:
+                try:
+                    return e.response.json()
+                except Exception:
+                    return {"status": "unhealthy", "error": "Service unavailable"}
+            raise
+
+    @require_client
+    async def is_healthy(self):
+        """Check if the Gateway and its components are healthy.
+
+        Returns:
+            bool: True if overall status is 'healthy', False otherwise.
+        """
+        try:
+            health_data = await self.health_check()
+            return health_data.get("status") == "healthy"
+        except Exception:
+            return False
+
 
 class GatewayClientSync:
     """A simplified synchronous wrapper around the async GatewayClient.
@@ -574,3 +611,24 @@ def deploy_model(
                 ttl=ttl,
             )
         )
+
+    def health_check(self):
+        """Check if the Gateway and its components are healthy and responsive.
+
+        Returns:
+            dict: Health status information with 'status' and 'components' fields.
+                  Status will be 'healthy' or 'unhealthy'.
+
+        Raises:
+            httpx.HTTPStatusError: For HTTP errors other than 503 Service Unavailable.
+            Other exceptions: For network errors, timeouts, etc.
+        """
+        return asyncio.run(self._client.health_check())
+
+    def is_healthy(self):
+        """Check if the Gateway and its components are healthy.
+
+        Returns:
+            bool: True if overall status is 'healthy', False otherwise.
+        """
+        return asyncio.run(self._client.is_healthy())

cogstack_model_gateway/common/db.py

@@ -2,7 +2,7 @@
 from contextlib import contextmanager
 
 from sqlalchemy.exc import ProgrammingError
-from sqlmodel import Session, SQLModel, create_engine
+from sqlmodel import Session, SQLModel, create_engine, text
 
 log = logging.getLogger("cmg.common")
 
@@ -42,3 +42,13 @@ def get_session(self):
         """Get a database session."""
         with Session(self.engine) as session:
             yield session
+
+    def health_check(self) -> bool:
+        """Perform a health check by executing a simple query."""
+        try:
+            with self.get_session() as session:
+                session.exec(text("SELECT 1"))
+            return True
+        except Exception as e:
+            log.error("Health check failed for database: %s", e)
+            return False

cogstack_model_gateway/common/object_store.py

@@ -88,3 +88,12 @@ def get_object_url(
         """Get a presigned URL for the specified object."""
         bucket_name = bucket_name if bucket_name else self.default_bucket
         return self.client.presigned_get_object(bucket_name, object_key, expires=expires)
+
+    def health_check(self) -> bool:
+        """Perform a health check by listing objects in the configured default bucket."""
+        try:
+            self.client.list_objects(self.default_bucket, extra_query_params={"max-keys": "1"})
+            return True
+        except Exception as e:
+            log.error("Health check failed for bucket '%s': %s", self.default_bucket, e)
+            return False

cogstack_model_gateway/common/queue.py

@@ -178,3 +178,11 @@ def is_queue_empty(self):
             return queue.method.message_count == 0
         except pika.exceptions.ChannelClosed:
             return True
+
+    @with_connection
+    def health_check(self) -> bool:
+        """Perform a health check by verifying a connection can be established."""
+        healthy = self.connection and not self.connection.is_closed
+        if not healthy:
+            log.error("Health check failed for queue '%s'", self.queue_name)
+        return healthy

cogstack_model_gateway/gateway/main.py

@@ -3,10 +3,10 @@
 from contextlib import asynccontextmanager
 
 import urllib3
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from prometheus_client import CollectorRegistry, make_asgi_app, multiprocess
 
-from cogstack_model_gateway.common.config import load_config
+from cogstack_model_gateway.common.config import get_config, load_config
 from cogstack_model_gateway.common.db import DatabaseManager
 from cogstack_model_gateway.common.logging import configure_logging
 from cogstack_model_gateway.common.object_store import ObjectStoreManager
@@ -99,3 +99,41 @@ async def prometheus_request_counter(request, call_next):
 async def root():
     """Root endpoint for the gateway API."""
     return {"message": "Enter the cult... I mean, the API."}
+
+
+@app.get("/health")
+async def health_check():
+    """Health check endpoint that verifies the status of critical components."""
+    try:
+        config = get_config()
+
+        components_to_check = {
+            "database": config.database_manager,
+            "task_object_store": config.task_object_store_manager,
+            "results_object_store": config.results_object_store_manager,
+            "queue": config.queue_manager,
+        }
+
+        component_status = {
+            name: "healthy" if manager.health_check() else "unhealthy"
+            for name, manager in components_to_check.items()
+        }
+
+        overall_status = (
+            "healthy"
+            if all(status == "healthy" for status in component_status.values())
+            else "unhealthy"
+        )
+
+        health_status = {"status": overall_status, "components": component_status}
+
+        if overall_status == "unhealthy":
+            raise HTTPException(status_code=503, detail=health_status)
+
+        return health_status
+
+    except Exception as e:
+        raise HTTPException(
+            status_code=503,
+            detail={"status": "unhealthy", "error": f"Failed to perform health check: {str(e)}"},
+        )

docker-compose.yaml

@@ -128,6 +128,11 @@ services:
       - gateway
       - cogstack-model-serve_cms
       - observability
+    healthcheck:
+      test: ["CMD", "curl", "-f", "http://localhost:8000/health"]
+      interval: 60s
+      timeout: 20s
+      retries: 3
 
   scheduler:
     image: cogstacksystems/cogstack-model-gateway-scheduler:${CMG_IMAGE_TAG:-latest}

tests/integration/test_api.py

@@ -91,6 +91,24 @@ def test_root(client: TestClient):
     assert response.json() == {"message": "Enter the cult... I mean, the API."}
 
 
+def test_health_check(client: TestClient):
+    response = client.get("/health")
+    assert response.status_code == 200
+
+    health_data = response.json()
+    assert "status" in health_data
+    assert "components" in health_data
+
+    expected_components = ["database", "task_object_store", "results_object_store", "queue"]
+    for component in expected_components:
+        assert component in health_data["components"]
+
+    # In a properly configured test environment, all components should be healthy
+    assert health_data["status"] == "healthy"
+    for component in expected_components:
+        assert health_data["components"][component] == "healthy"
+
+
 def test_get_tasks(client: TestClient):
     response = client.get("/tasks/")
     assert response.status_code == 403