Repository: bentoml/BentoML
Pull request: feat: cascade health checks to runner
- Added `is_ready` method to `RunnerHandle`'s `init` file, then implemented it in the local and remote runners
- Added new method to `Runner` to check that `RunnerHandle` is ready
- Added new method to HTTP service app's `on_startup` method to check the Runner in loop until it returns true

Problem statement:
feat: Cascade readiness and liveness checks to runners in the same container
Health check endpoints, e.g. `/readyz`, `/livez`, on the API server and runner are limited to checking the health of the processes. Ideally, health check to the API server should be cascaded to the runners deployed in the same bento container.

Hints from issue discussion:

Commit message: fix type of async_run_method

Patch to review:
diff --git a/src/bentoml/_internal/configuration/containers.py b/src/bentoml/_internal/configuration/containers.py
index fa3b5fc128d..1be47509aeb 100644
--- a/src/bentoml/_internal/configuration/containers.py
+++ b/src/bentoml/_internal/configuration/containers.py
@@ -171,6 +171,7 @@ def _is_ip_address(addr: str) -> bool:
                 "max_message_length": Or(int, None),
                 "maximum_concurrent_rpcs": Or(int, None),
             },
+            "health_probe": {"check_runners": bool},
         },
         "runners": {
             **RUNNER_CFG_SCHEMA,
diff --git a/src/bentoml/_internal/configuration/default_configuration.yaml b/src/bentoml/_internal/configuration/default_configuration.yaml
index dd959304362..af3371d1ffd 100644
--- a/src/bentoml/_internal/configuration/default_configuration.yaml
+++ b/src/bentoml/_internal/configuration/default_configuration.yaml
@@ -38,6 +38,8 @@ api_server:
     metrics:
       host: 0.0.0.0
       port: 3001
+  health_probe:
+    check_runners: True
 
 runners:
   batching:
diff --git a/src/bentoml/_internal/runner/runner.py b/src/bentoml/_internal/runner/runner.py
index 0b931f3e69c..90f15bc5066 100644
--- a/src/bentoml/_internal/runner/runner.py
+++ b/src/bentoml/_internal/runner/runner.py
@@ -240,6 +240,9 @@ def init_client(self):
     def destroy(self):
         object.__setattr__(self, "_runner_handle", DummyRunnerHandle())
 
+    async def runner_handle_is_ready(self) -> bool:
+        return await self._runner_handle.is_ready()
+
     @property
     def scheduled_worker_count(self) -> int:
         return self.scheduling_strategy.get_worker_count(
diff --git a/src/bentoml/_internal/runner/runner_handle/__init__.py b/src/bentoml/_internal/runner/runner_handle/__init__.py
index 92ffd46a4d3..8da0faea5b2 100644
--- a/src/bentoml/_internal/runner/runner_handle/__init__.py
+++ b/src/bentoml/_internal/runner/runner_handle/__init__.py
@@ -24,6 +24,10 @@ class RunnerHandle(ABC):
     def __init__(self, runner: Runner) -> None:
         ...
 
+    @abstractmethod
+    async def is_ready(self) -> bool:
+        return True
+
     @abstractmethod
     def run_method(
         self,
@@ -49,6 +53,9 @@ def __init__(  # pylint: disable=super-init-not-called
     ) -> None:
         pass
 
+    async def is_ready(self) -> bool:
+        return True
+
     def run_method(
         self,
         __bentoml_method: RunnerMethod[t.Any, t.Any, t.Any],
diff --git a/src/bentoml/_internal/runner/runner_handle/local.py b/src/bentoml/_internal/runner/runner_handle/local.py
index 1ddd9246ed1..d8e2b197a66 100644
--- a/src/bentoml/_internal/runner/runner_handle/local.py
+++ b/src/bentoml/_internal/runner/runner_handle/local.py
@@ -25,6 +25,9 @@ def __init__(self, runner: Runner) -> None:  # pylint: disable=super-init-not-ca
         self._runnable = runner.runnable_class(**runner.runnable_init_params)  # type: ignore
         self._limiter = None
 
+    async def is_ready(self) -> bool:
+        return True
+
     def run_method(
         self,
         __bentoml_method: RunnerMethod[t.Any, P, R],
diff --git a/src/bentoml/_internal/runner/runner_handle/remote.py b/src/bentoml/_internal/runner/runner_handle/remote.py
index d557cb5a84c..db553263d32 100644
--- a/src/bentoml/_internal/runner/runner_handle/remote.py
+++ b/src/bentoml/_internal/runner/runner_handle/remote.py
@@ -130,7 +130,7 @@ async def async_run_method(
         __bentoml_method: RunnerMethod[t.Any, P, R],
         *args: P.args,
         **kwargs: P.kwargs,
-    ) -> R:
+    ) -> R | tuple[R, ...]:
         from ...runner.container import AutoContainer
 
         inp_batch_dim = __bentoml_method.config.batch_dim[0]
@@ -216,5 +216,24 @@ def run_method(
             ),
         )
 
+    async def is_ready(self) -> bool:
+        import aiohttp
+
+        # default kubernetes probe timeout is also 1s; see
+        # https://kubernetes.io/docs/tasks/configure-pod-container/configure-liveness-readiness-startup-probes/#configure-probes
+        timeout = aiohttp.ClientTimeout(total=1)
+        async with self._client.get(
+            f"{self._addr}/readyz",
+            headers={
+                "Bento-Name": component_context.bento_name,
+                "Bento-Version": component_context.bento_version,
+                "Runner-Name": self._runner.name,
+                "Yatai-Bento-Deployment-Name": component_context.yatai_bento_deployment_name,
+                "Yatai-Bento-Deployment-Namespace": component_context.yatai_bento_deployment_namespace,
+            },
+            timeout=timeout,
+        ) as resp:
+            return resp.status == 200
+
     def __del__(self) -> None:
         self._close_conn()
diff --git a/src/bentoml/_internal/server/grpc_app.py b/src/bentoml/_internal/server/grpc_app.py
index c14756929bf..2423f90ee84 100644
--- a/src/bentoml/_internal/server/grpc_app.py
+++ b/src/bentoml/_internal/server/grpc_app.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import typing as t
+import asyncio
 import logging
 from typing import TYPE_CHECKING
 from functools import partial
@@ -19,7 +20,7 @@
     from ..service import Service
     from .grpc.servicer import Servicer
 
-    OnStartup = list[t.Callable[[], None | t.Coroutine[t.Any, t.Any, None]]]
+    OnStartup = list[t.Callable[[], t.Union[None, t.Coroutine[t.Any, t.Any, None]]]]
 
 
 class GRPCAppFactory:
@@ -42,6 +43,35 @@ def __init__(
         self.bento_service = bento_service
         self.enable_metrics = enable_metrics
 
+    @inject
+    async def wait_for_runner_ready(
+        self,
+        *,
+        check_interval: int = 5,
+    ):
+        if BentoMLContainer.api_server_config.health_probe.check_runners.get():
+            logger.info(
+                "Waiting for runners %r to be ready...", self.bento_service.runners
+            )
+
+            while True:
+                try:
+                    runner_statuses = (
+                        runner.runner_handle_is_ready()
+                        for runner in self.bento_service.runners
+                    )
+
+                    runners_ready = all(await asyncio.gather(*runner_statuses))
+
+                    if runners_ready:
+                        break
+                except ConnectionError as e:
+                    logger.debug("[%s] Retrying ...", e)
+
+                await asyncio.sleep(check_interval)
+
+            logger.info("All runners ready; continuing initialization.")
+
     @property
     def on_startup(self) -> OnStartup:
         on_startup: OnStartup = [self.bento_service.on_grpc_server_startup]
@@ -52,6 +82,7 @@ def on_startup(self) -> OnStartup:
             for runner in self.bento_service.runners:
                 on_startup.append(runner.init_client)
 
+        on_startup.append(self.wait_for_runner_ready)
         return on_startup
 
     @property
diff --git a/src/bentoml/_internal/server/http_app.py b/src/bentoml/_internal/server/http_app.py
index 3e4e33ac9a2..910a96a1ad4 100644
--- a/src/bentoml/_internal/server/http_app.py
+++ b/src/bentoml/_internal/server/http_app.py
@@ -10,6 +10,8 @@
 
 from simple_di import inject
 from simple_di import Provide
+from starlette.responses import PlainTextResponse
+from starlette.exceptions import HTTPException
 
 from ..context import trace_context
 from ..context import InferenceApiContext as Context
@@ -267,6 +269,19 @@ def on_startup(self) -> list[t.Callable[[], None]]:
         on_startup.extend(super().on_startup)
         return on_startup
 
+    async def readyz(self, _: "Request") -> "Response":
+        if BentoMLContainer.api_server_config.health_probe.check_runners.get():
+            runner_statuses = (
+                runner.runner_handle_is_ready() for runner in self.bento_service.runners
+            )
+
+            runners_ready = all(await asyncio.gather(*runner_statuses))
+
+            if not runners_ready:
+                raise HTTPException(status_code=503, detail="Runners are not ready.")
+
+        return PlainTextResponse("\n", status_code=200)
+
     @property
     def on_shutdown(self) -> list[t.Callable[[], None]]:
         on_shutdown = [self.bento_service.on_asgi_app_shutdown]
diff --git a/src/bentoml_cli/worker/http_api_server.py b/src/bentoml_cli/worker/http_api_server.py
index e7cdebca063..aecdf90b913 100644
--- a/src/bentoml_cli/worker/http_api_server.py
+++ b/src/bentoml_cli/worker/http_api_server.py
@@ -113,6 +113,11 @@ def main(
     component_context.component_index = worker_id
     configure_server_logging()
 
+    if worker_id is None:
+        # worker ID is not set; this server is running in standalone mode
+        # and should not be concerned with the status of its runners
+        BentoMLContainer.config.health_probe.check_runners.set(False)
+
     BentoMLContainer.development_mode.set(False)
     if prometheus_dir is not None:
         BentoMLContainer.prometheus_multiproc_dir.set(prometheus_dir)
diff --git a/tests/e2e/bento_server_http/tests/test_meta.py b/tests/e2e/bento_server_http/tests/test_meta.py
index 1761a18cb8c..10fe0f4af49 100644
--- a/tests/e2e/bento_server_http/tests/test_meta.py
+++ b/tests/e2e/bento_server_http/tests/test_meta.py
@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import time
+import asyncio
 from pathlib import Path
 
 import pytest
@@ -25,8 +27,6 @@ async def test_api_server_meta(host: str) -> None:
     assert b'{"Hello":"World"}' == body
     status, _, _ = await async_request("GET", f"http://{host}/docs.json")
     assert status == 200
-    status, _, _ = await async_request("GET", f"http://{host}/readyz")
-    assert status == 200
     status, _, body = await async_request("GET", f"http://{host}/metrics")
     assert status == 200
     assert body
@@ -34,6 +34,19 @@ async def test_api_server_meta(host: str) -> None:
     assert status == 404
 
 
+@pytest.mark.asyncio
+async def test_runner_readiness(host: str) -> None:
+    timeout = 300
+    start_time = time.time()
+    status = ""
+    while (time.time() - start_time) < timeout:
+        status, _, _ = await async_request("GET", f"http://{host}/readyz")
+        await asyncio.sleep(20)
+        if status == 200:
+            break
+    assert status == 200
+
+
 @pytest.mark.asyncio
 async def test_cors(host: str, server_config_file: str) -> None:
     ORIGIN = "http://bentoml.ai"


Code review comments:
diff_hunk: @@ -42,6 +43,35 @@ def __init__(
         self.bento_service = bento_service
         self.enable_metrics = enable_metrics
 
+    @inject
+    async def wait_for_runner_ready(
+        self,
+        *,
+        check_interval: int = 5,
+    ):
path: src/bentoml/_internal/server/grpc_app.py
text: @user1:
```suggestion
    async def wait_for_runner_ready(self, *, check_interval: int = 5):
```

diff_hunk: @@ -38,6 +38,8 @@ api_server:
     metrics:
       host: 0.0.0.0
       port: 3001
+  health_probe:
+    check_runners: True
path: src/bentoml/_internal/configuration/default_configuration.yaml
text: @user1:
very nit pick:
```yaml
```suggestion
  probe:
    health:
      check_runners: true
```
```

@author:
Seoyeon and I discussed and went a different route. Let us know if you have any questions.

diff_hunk: @@ -113,6 +113,11 @@ def main(
     component_context.component_index = worker_id
     configure_server_logging()
 
+    if worker_id is None:
+        # worker ID is not set; this server is running in standalone mode
+        # and should not be concerned with the status of its runners
+        BentoMLContainer.config.health_probe.check_runners.set(False)
path: src/bentoml_cli/worker/http_api_server.py
text: @user1:
Should also include this for grpc_api_server.py