From 5f6d1cd606bf129333206529cbf8a964819446b5 Mon Sep 17 00:00:00 2001 From: Bhargav Dodla Date: Sat, 17 Aug 2024 09:04:45 -0700 Subject: [PATCH] feat: Add health check service to registry server Signed-off-by: Bhargav Dodla --- sdk/python/feast/registry_server.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/sdk/python/feast/registry_server.py b/sdk/python/feast/registry_server.py index 53acb9f625c..2ac968bed8e 100644 --- a/sdk/python/feast/registry_server.py +++ b/sdk/python/feast/registry_server.py @@ -3,6 +3,8 @@ import grpc from google.protobuf.empty_pb2 import Empty +from grpc_health.v1 import health, health_pb2, health_pb2_grpc +from grpc_reflection.v1alpha import reflection from feast import FeatureStore from feast.data_source import DataSource @@ -353,6 +355,18 @@ def start_server(store: FeatureStore, port: int): RegistryServer_pb2_grpc.add_RegistryServerServicer_to_server( RegistryServer(store.registry), server ) + # Add health check service to server + health_servicer = health.HealthServicer() + health_pb2_grpc.add_HealthServicer_to_server(health_servicer, server) + health_servicer.set("", health_pb2.HealthCheckResponse.SERVING) + + service_names_available_for_reflection = ( + RegistryServer_pb2.DESCRIPTOR.services_by_name["RegistryServer"].full_name, + health_pb2.DESCRIPTOR.services_by_name["Health"].full_name, + reflection.SERVICE_NAME, + ) + reflection.enable_server_reflection(service_names_available_for_reflection, server) + server.add_insecure_port(f"[::]:{port}") server.start() server.wait_for_termination()