Merge pull request #5 from code0-tech/3-implement-health-service

Implement health service and proper shutdown

Author: Taucher2003

main.py

@@ -1,5 +1,6 @@
 import concurrent.futures
 import os
+import signal
 import sys
 import time
 
@@ -24,6 +25,7 @@
 import tucana.generated.velorum.info_pb2_grpc as info_pb2_grpc
 import tucana.generated.velorum.info_pb2 as info_pb2
 from grpc_reflection.v1alpha import reflection
+from grpc_health.v1 import health, health_pb2, health_pb2_grpc
 
 from src.endpoint.info.info_endpoint import InfoService
 from src.endpoint.generation.generate_endpoint import GenerateService
@@ -37,6 +39,12 @@
     generate_pb2_grpc.add_GenerateServiceServicer_to_server(GenerateService(), server)
     info_pb2_grpc.add_InfoServiceServicer_to_server(InfoService(), server)
 
+
+    health_servicer = health.HealthServicer()
+    health_servicer.set('liveness', health_pb2.HealthCheckResponse.SERVING)
+    health_servicer.set('readiness', health_pb2.HealthCheckResponse.SERVING)
+    health_pb2_grpc.add_HealthServicer_to_server(health_servicer, server)
+
     port = os.getenv("HOST", "0.0.0.0") + ":" + os.getenv("PORT", "50051")
     server.add_insecure_port(port)
     log.success(f"Velorum listening on {port}")  # type: ignore[attr-defined]
@@ -50,9 +58,13 @@
 
     server.start()
 
-    try:
-        while True:
-            time.sleep(86400)
-    except KeyboardInterrupt:
+
+    def shutdown(signum, frame):
         log.info("Shutting down gracefully...")
-        server.stop(0)
+        health_servicer.enter_graceful_shutdown()
+        server.stop(5).wait()
+
+    signal.signal(signal.SIGTERM, shutdown)
+    signal.signal(signal.SIGINT, shutdown)
+
+    server.wait_for_termination()

pyproject.toml

@@ -7,6 +7,7 @@ requires-python = ">=3.10,<3.13"
 dependencies = [
     "apscheduler>=3.11.2",
     "grpcio>=1.80.0",
+    "grpcio-health-checking>=1.80.0",
     "grpcio-reflection>=1.80.0",
     "instructor>=1.15.1",
     "PyJWT>=2.4.0",

src/interceptor/auth_interceptor.py

@@ -13,6 +13,10 @@
 class AuthInterceptor(grpc.ServerInterceptor):
 
     def intercept_service(self, continuation, handler_call_details):
+        method = handler_call_details.method or ''
+        if method.startswith('/grpc.health.v1.Health/'):
+            return continuation(handler_call_details)
+
         metadata = dict(handler_call_details.invocation_metadata or [])
         auth_token = metadata.get('authorization', '')