Merge pull request #78 from fal-ai/add-health-checks

Add a health check endpoint

src/isolate/server/health/__init__.py

@@ -0,0 +1,8 @@
+from isolate.server.health.health_pb2 import (
+    HealthCheckRequest,
+    HealthCheckResponse,
+)
+from isolate.server.health.health_pb2_grpc import HealthServicer, HealthStub
+from isolate.server.health.health_pb2_grpc import (
+    add_HealthServicer_to_server as register_health,
+)

src/isolate/server/health/health.proto

@@ -0,0 +1,23 @@
+syntax = "proto3";
+
+package grpc.health.v1;
+
+message HealthCheckRequest {
+  string service = 1;
+}
+
+message HealthCheckResponse {
+  enum ServingStatus {
+    UNKNOWN = 0;
+    SERVING = 1;
+    NOT_SERVING = 2;
+    SERVICE_UNKNOWN = 3;  // Used only by the Watch method.
+  }
+  ServingStatus status = 1;
+}
+
+service Health {
+  rpc Check(HealthCheckRequest) returns (HealthCheckResponse);
+
+  rpc Watch(HealthCheckRequest) returns (stream HealthCheckResponse);
+}

src/isolate/server/health/health_pb2.pyi

@@ -0,0 +1,75 @@
+"""
+@generated by mypy-protobuf.  Do not edit manually!
+isort:skip_file
+"""
+import builtins
+import google.protobuf.descriptor
+import google.protobuf.internal.enum_type_wrapper
+import google.protobuf.message
+import sys
+import typing
+
+if sys.version_info >= (3, 10):
+    import typing as typing_extensions
+else:
+    import typing_extensions
+
+DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
+
+@typing_extensions.final
+class HealthCheckRequest(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    SERVICE_FIELD_NUMBER: builtins.int
+    service: builtins.str
+    def __init__(
+        self,
+        *,
+        service: builtins.str = ...,
+    ) -> None: ...
+    def ClearField(
+        self, field_name: typing_extensions.Literal["service", b"service"]
+    ) -> None: ...
+
+global___HealthCheckRequest = HealthCheckRequest
+
+@typing_extensions.final
+class HealthCheckResponse(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    class _ServingStatus:
+        ValueType = typing.NewType("ValueType", builtins.int)
+        V: typing_extensions.TypeAlias = ValueType
+
+    class _ServingStatusEnumTypeWrapper(
+        google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[
+            HealthCheckResponse._ServingStatus.ValueType
+        ],
+        builtins.type,
+    ):  # noqa: F821
+        DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+        UNKNOWN: HealthCheckResponse._ServingStatus.ValueType  # 0
+        SERVING: HealthCheckResponse._ServingStatus.ValueType  # 1
+        NOT_SERVING: HealthCheckResponse._ServingStatus.ValueType  # 2
+        SERVICE_UNKNOWN: HealthCheckResponse._ServingStatus.ValueType  # 3
+        """Used only by the Watch method."""
+
+    class ServingStatus(_ServingStatus, metaclass=_ServingStatusEnumTypeWrapper): ...
+    UNKNOWN: HealthCheckResponse.ServingStatus.ValueType  # 0
+    SERVING: HealthCheckResponse.ServingStatus.ValueType  # 1
+    NOT_SERVING: HealthCheckResponse.ServingStatus.ValueType  # 2
+    SERVICE_UNKNOWN: HealthCheckResponse.ServingStatus.ValueType  # 3
+    """Used only by the Watch method."""
+
+    STATUS_FIELD_NUMBER: builtins.int
+    status: global___HealthCheckResponse.ServingStatus.ValueType
+    def __init__(
+        self,
+        *,
+        status: global___HealthCheckResponse.ServingStatus.ValueType = ...,
+    ) -> None: ...
+    def ClearField(
+        self, field_name: typing_extensions.Literal["status", b"status"]
+    ) -> None: ...
+
+global___HealthCheckResponse = HealthCheckResponse

src/isolate/server/health/health_pb2_grpc.py

@@ -0,0 +1,124 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+
+from isolate.server.health import health_pb2 as health__pb2
+
+
+class HealthStub(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def __init__(self, channel):
+        """Constructor.
+
+        Args:
+            channel: A grpc.Channel.
+        """
+        self.Check = channel.unary_unary(
+            "/grpc.health.v1.Health/Check",
+            request_serializer=health__pb2.HealthCheckRequest.SerializeToString,
+            response_deserializer=health__pb2.HealthCheckResponse.FromString,
+        )
+        self.Watch = channel.unary_stream(
+            "/grpc.health.v1.Health/Watch",
+            request_serializer=health__pb2.HealthCheckRequest.SerializeToString,
+            response_deserializer=health__pb2.HealthCheckResponse.FromString,
+        )
+
+
+class HealthServicer(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def Check(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details("Method not implemented!")
+        raise NotImplementedError("Method not implemented!")
+
+    def Watch(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details("Method not implemented!")
+        raise NotImplementedError("Method not implemented!")
+
+
+def add_HealthServicer_to_server(servicer, server):
+    rpc_method_handlers = {
+        "Check": grpc.unary_unary_rpc_method_handler(
+            servicer.Check,
+            request_deserializer=health__pb2.HealthCheckRequest.FromString,
+            response_serializer=health__pb2.HealthCheckResponse.SerializeToString,
+        ),
+        "Watch": grpc.unary_stream_rpc_method_handler(
+            servicer.Watch,
+            request_deserializer=health__pb2.HealthCheckRequest.FromString,
+            response_serializer=health__pb2.HealthCheckResponse.SerializeToString,
+        ),
+    }
+    generic_handler = grpc.method_handlers_generic_handler(
+        "grpc.health.v1.Health", rpc_method_handlers
+    )
+    server.add_generic_rpc_handlers((generic_handler,))
+
+
+# This class is part of an EXPERIMENTAL API.
+class Health(object):
+    """Missing associated documentation comment in .proto file."""
+
+    @staticmethod
+    def Check(
+        request,
+        target,
+        options=(),
+        channel_credentials=None,
+        call_credentials=None,
+        insecure=False,
+        compression=None,
+        wait_for_ready=None,
+        timeout=None,
+        metadata=None,
+    ):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            "/grpc.health.v1.Health/Check",
+            health__pb2.HealthCheckRequest.SerializeToString,
+            health__pb2.HealthCheckResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+        )
+
+    @staticmethod
+    def Watch(
+        request,
+        target,
+        options=(),
+        channel_credentials=None,
+        call_credentials=None,
+        insecure=False,
+        compression=None,
+        wait_for_ready=None,
+        timeout=None,
+        metadata=None,
+    ):
+        return grpc.experimental.unary_stream(
+            request,
+            target,
+            "/grpc.health.v1.Health/Watch",
+            health__pb2.HealthCheckRequest.SerializeToString,
+            health__pb2.HealthCheckResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+        )

src/isolate/server/health_server.py

@@ -0,0 +1,40 @@
+import asyncio
+from dataclasses import dataclass
+from typing import AsyncIterator
+
+from grpc.aio import ServicerContext
+
+from isolate.server import health
+
+
+@dataclass
+class HealthServicer(health.HealthServicer):
+    def __post_init__(self):
+        self._state = {
+            # Empty refers to the whole server
+            "": health.HealthCheckResponse.ServingStatus.SERVING,
+            "isolate": health.HealthCheckResponse.ServingStatus.SERVING,
+        }
+
+    def _get_status(
+        self, service: str
+    ) -> health.HealthCheckResponse.ServingStatus.ValueType:
+        status = self._state.get(
+            service,
+            health.HealthCheckResponse.ServingStatus.SERVICE_UNKNOWN,
+        )
+        return status
+
+    def Check(
+        self, request: health.HealthCheckRequest, context: ServicerContext
+    ) -> health.HealthCheckResponse:
+        return health.HealthCheckResponse(status=self._get_status(request.service))
+
+    async def Watch(
+        self,
+        request: health.HealthCheckRequest,
+        context: ServicerContext,
+    ) -> AsyncIterator[health.HealthCheckResponse]:
+        while True:
+            yield health.HealthCheckResponse(status=self._get_status(request.service))
+            await asyncio.sleep(2)

src/isolate/server/server.py

@@ -27,7 +27,8 @@
 from isolate.connections.grpc import AgentError, LocalPythonGRPC
 from isolate.connections.grpc.configuration import get_default_options
 from isolate.logs import Log, LogLevel, LogSource
-from isolate.server import definitions
+from isolate.server import definitions, health
+from isolate.server.health_server import HealthServicer
 from isolate.server.interface import from_grpc, to_grpc
 
 # Whether to inherit all the packages from the current environment or not.
@@ -292,6 +293,7 @@ def main() -> None:
     )
     with BridgeManager() as bridge_manager:
         definitions.register_isolate(IsolateServicer(bridge_manager), server)
+        health.register_health(HealthServicer(), server)
 
         server.add_insecure_port(f"[::]:50001")
         print("Started listening at localhost:50001")