feat: custom readiness hook method (#5337)

* feat: custom readiness hook method

Signed-off-by: Frost Ming <me@frostming.com>

* fix: add is_alive hook

Signed-off-by: Frost Ming <me@frostming.com>

---------

Signed-off-by: Frost Ming <me@frostming.com>

Author: frostming

docs/source/build-with-bentoml/lifecycle-hooks.rst

@@ -148,3 +148,44 @@ Use the ``@bentoml.on_shutdown`` decorator to specify a method as a shutdown hoo
         @bentoml.on_shutdown
         async def async_shutdown(self):
             print("Async cleanup actions on Service shutdown.")
+
+Health check hooks
+^^^^^^^^^^^^^^^^^^
+
+Health check hooks allow you to specify custom logic for determining when your Service is healthy and ready to handle requests.
+This is particularly useful when your Service depends on external resources that need to be checked before the Service can be considered operational.
+
+You can define the following methods in your service class to implement health checks:
+
+- ``__is_alive__``: This method is called to check if the Service is alive. It should return a boolean value indicating the Service's health status. This responds to the ``/livez`` endpoint.
+- ``__is_ready__``: This method is called to check if the Service is ready to handle requests. It should return a boolean value indicating the Service's readiness status. This responds to the ``/readyz`` endpoint.
+
+Both can be asynchronous functions.
+
+For example:
+
+.. code-block:: python
+
+    import bentoml
+
+    @bentoml.service(workers=4)
+    class HookService:
+        def __init__(self) -> None:
+            self.db_connection = None
+            self.cache = None
+
+        @bentoml.on_startup
+        def init_resources(self):
+            self.db_connection = setup_database()
+            self.cache = setup_cache()
+
+        def __is_ready__(self) -> bool:
+            # Check if required resources are available
+            if self.db_connection is None or self.cache is None:
+                return False
+            return self.db_connection.is_connected() and self.cache.is_available()
+
+When you call the ``/readyz`` endpoint, it returns:
+
+- HTTP 200 if the Service is ready (the hook returns ``True``)
+- HTTP 503 if the Service is not ready (the hook returns ``False``)

src/_bentoml_impl/server/app.py

@@ -358,12 +358,48 @@ async def destroy_instance(self, _: Starlette) -> None:
         set_current_service(None)
         await self._result_store.__aexit__(None, None, None)
 
+    async def livez(self, _: Request) -> Response:
+        from starlette.exceptions import HTTPException
+
+        from bentoml._internal.utils import is_async_callable
+
+        if hasattr(self.service.inner, "__is_alive__"):
+            assert self._service_instance is not None, "Service must be initialized"
+            if is_async_callable(self.service.inner.__is_alive__):
+                is_alive = await self._service_instance.__is_alive__()
+            else:
+                is_alive = await anyio.to_thread.run_sync(
+                    self._service_instance.__is_alive__
+                )
+            if not is_alive:
+                raise HTTPException(
+                    status_code=503,
+                    detail="Service is dead because .__is_alive__() returns False.",
+                )
+        return await super().livez(_)
+
     async def readyz(self, _: Request) -> Response:
         from starlette.exceptions import HTTPException
         from starlette.responses import PlainTextResponse
 
+        from bentoml._internal.utils import is_async_callable
+
         from ..client import RemoteProxy
 
+        if hasattr(self.service.inner, "__is_ready__"):
+            assert self._service_instance is not None, "Service must be initialized"
+            if is_async_callable(self.service.inner.__is_ready__):
+                is_ready = await self._service_instance.__is_ready__()
+            else:
+                is_ready = await anyio.to_thread.run_sync(
+                    self._service_instance.__is_ready__
+                )
+            if not is_ready:
+                raise HTTPException(
+                    status_code=503,
+                    detail="Service is not ready because .__is_ready__() returns False.",
+                )
+
         if BentoMLContainer.api_server_config.runner_probe.enabled.get():
             dependency_statuses: list[t.Coroutine[None, None, bool]] = []
             for dependency in self.service.dependencies.values():