chore: refactor separate servicer implementation

Signed-off-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>

src/bentoml/_internal/io_descriptors/base.py

@@ -30,6 +30,7 @@
     )
     OpenAPIResponse = dict[str, str | dict[str, MediaType] | dict[str, t.Any]]
 
+
 IO_DESCRIPTOR_REGISTRY: dict[str, type[IODescriptor[t.Any]]] = {}
 
 IOType = t.TypeVar("IOType")

src/bentoml/_internal/server/grpc/servicer/__init__.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+import typing as t
+import logging
+import importlib
+from typing import TYPE_CHECKING
+from inspect import isawaitable
+
+from ....utils import LazyLoader
+from .....grpc.utils import import_generated_stubs
+from .....grpc.utils import LATEST_PROTOCOL_VERSION
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from types import ModuleType
+
+    from grpc import aio
+    from grpc_health.v1 import health
+    from typing_extensions import Self
+
+    from bentoml.grpc.v1 import service_pb2_grpc as services
+    from bentoml.grpc.types import Interceptors
+    from bentoml.grpc.types import AddServicerFn
+    from bentoml.grpc.types import ServicerClass
+
+    from ....service.service import Service
+
+    class ServicerModule(ModuleType):
+        @staticmethod
+        def create_bento_servicer(service: Service) -> services.BentoServiceServicer:
+            ...
+
+else:
+    health = LazyLoader(
+        "health",
+        globals(),
+        "grpc_health.v1.health",
+        exc_msg="'grpcio-health-checking' is required for using health checking endpoints. Install with 'pip install grpcio-health-checking'.",
+    )
+    struct_pb2 = LazyLoader("struct_pb2", globals(), "google.protobuf.struct_pb2")
+
+
+class Servicer:
+    """Create an instance of gRPC Servicer."""
+
+    _cached_module = None
+
+    def __init__(
+        self: Self,
+        service: Service,
+        on_startup: t.Sequence[t.Callable[[], t.Any]] | None = None,
+        on_shutdown: t.Sequence[t.Callable[[], t.Any]] | None = None,
+        mount_servicers: t.Sequence[tuple[ServicerClass, AddServicerFn, list[str]]]
+        | None = None,
+        interceptors: Interceptors | None = None,
+        protocol_version: str = LATEST_PROTOCOL_VERSION,
+    ) -> None:
+        self.bento_service = service
+
+        self.on_startup = [] if not on_startup else list(on_startup)
+        self.on_shutdown = [] if not on_shutdown else list(on_shutdown)
+        self.mount_servicers = [] if not mount_servicers else list(mount_servicers)
+        self.interceptors = [] if not interceptors else list(interceptors)
+        self.protocol_version = protocol_version
+
+        self.loaded = False
+
+    @property
+    def _servicer_module(self) -> ServicerModule:
+        if self._cached_module is None:
+            object.__setattr__(
+                self,
+                "_cached_module",
+                importlib.import_module(f".{self.protocol_version}", package=__name__),
+            )
+        assert self._cached_module is not None
+        return self._cached_module
+
+    def load(self):
+        assert not self.loaded
+
+        pb, _ = import_generated_stubs(self.protocol_version)
+
+        self.interceptors_stack = self.build_interceptors_stack()
+
+        self.bento_servicer = self._servicer_module.create_bento_servicer(
+            self.bento_service
+        )
+
+        # Create a health check servicer. We use the non-blocking implementation
+        # to avoid thread starvation.
+        self.health_servicer = health.aio.HealthServicer()
+
+        self.service_names = tuple(
+            service.full_name for service in pb.DESCRIPTOR.services_by_name.values()
+        ) + (health.SERVICE_NAME,)
+        self.loaded = True
+
+    def build_interceptors_stack(self) -> list[aio.ServerInterceptor]:
+        return list(map(lambda x: x(), self.interceptors))
+
+    async def startup(self):
+        for handler in self.on_startup:
+            out = handler()
+            if isawaitable(out):
+                await out
+
+    async def shutdown(self):
+        for handler in self.on_shutdown:
+            out = handler()
+            if isawaitable(out):
+                await out
+
+    def __bool__(self):
+        return self.loaded