add healthcheck cli

README.md

@@ -16,6 +16,21 @@ pip install vllm-tgis-adapter
 python -m vllm_tgis_adapter
 ```
 
+### HealthCheck CLI
+
+Installing the adapter also install a grpc healthcheck cli that can be used to monitor the status of the grpc server:
+
+```console
+$ grpc_healtheck
+health check...status: SERVING
+```
+
+See usage with
+
+```bash
+grpc_healthcheck --help
+```
+
 ## Build
 
 ```bash

pyproject.toml

@@ -40,6 +40,9 @@ dependencies = [
 Issues = "https://github.com/dtrifiro/vllm_tgis_adapter/issues"
 Source = "https://github.com/dtrifiro/vllm_tgis_adapter"
 
+[project.scripts]
+grpc_healthcheck = "vllm_tgis_adapter.healthcheck:cli"
+
 [project.optional-dependencies]
 tests = [
   "pytest==8.2.0",
@@ -73,7 +76,7 @@ vllm_tgis_adapter = [
 ]
 
 [tool.pytest.ini_options]
-addopts = "-ra"
+addopts = "-ra -W error::grpc.experimental.ExperimentalApiWarning"
 
 [tool.coverage.run]
 branch = true
@@ -158,6 +161,7 @@ select = ["ALL"]
   "N802",  # name should be lowercase
   "ERA001"  # Found commented-out code
 ]
+"src/vllm_tgis_adapter/healthcheck.py" = ["T201"]
 "src/vllm_tgis_adapter/_version.py" = ["ALL"]
 "tests/**" = ["S", "ARG001", "ARG002", "ANN"]
 "setup.py" = [

src/vllm_tgis_adapter/healthcheck.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+import argparse
+import sys
+import warnings
+
+import grpc
+import grpc.experimental  # this is required for grpc_health
+from grpc_health.v1.health_pb2 import HealthCheckRequest
+from grpc_health.v1.health_pb2_grpc import Health
+
+from vllm_tgis_adapter.grpc.grpc_server import TextGenerationService
+
+warnings.simplefilter(
+    action="ignore", category=grpc.experimental.ExperimentalApiWarning
+)
+
+
+def health_check(
+    *,
+    server_url: str = "localhost:8033",
+    service: str | None = None,
+    insecure: bool = True,
+    timeout: float = 1,
+) -> bool:
+    print("health check...", end="")
+    request = HealthCheckRequest(service=service)
+    try:
+        response = Health.Check(
+            request=request,
+            target=server_url,
+            timeout=timeout,
+            insecure=insecure,
+        )
+    except grpc.RpcError as e:
+        print(f"Health.Check failed: code={e.code()}, details={e.details()}")
+        return False
+
+    print(str(response).strip())
+    return response.status == response.SERVING
+
+
+def cli() -> None:
+    args = parse_args()
+    if not health_check(
+        server_url=args.server_url,
+        service=args.service_name,
+        insecure=args.insecure,
+        timeout=args.timeout,
+    ):
+        sys.exit(1)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.formatter_class = argparse.ArgumentDefaultsHelpFormatter
+    group = parser.add_mutually_exclusive_group(required=False)
+    group.add_argument(
+        "--insecure",
+        dest="insecure",
+        action="store_true",
+        help="Use an insecure connection",
+    )
+    group.add_argument(
+        "--secure",
+        dest="secure",
+        action="store_true",
+        help="Use a secure connection",
+    )
+    group.set_defaults(insecure=True, secure=False)
+    parser.add_argument(
+        "--server-url",
+        type=str,
+        help="grpc server url (`host:port`)",
+        default="localhost:8033",
+    )
+    parser.add_argument(
+        "--timeout",
+        type=float,
+        help="Timeout for healthcheck request",
+        default=1,
+    )
+    parser.add_argument(
+        "--service-name",
+        type=str,
+        help="Name of the service to check",
+        required=False,
+        default=TextGenerationService.SERVICE_NAME,
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    cli()

tests/conftest.py

@@ -3,14 +3,13 @@
 import threading
 
 import pytest
-from grpc_health.v1.health_pb2 import HealthCheckRequest
-from grpc_health.v1.health_pb2_grpc import Health
 from vllm import AsyncLLMEngine
 from vllm.config import DeviceConfig
 from vllm.engine.async_llm_engine import _AsyncLLMEngine
 from vllm.entrypoints.openai.cli_args import make_arg_parser
 
 from vllm_tgis_adapter.grpc.grpc_server import TextGenerationService, start_grpc_server
+from vllm_tgis_adapter.healthcheck import health_check
 from vllm_tgis_adapter.tgis_utils.args import (
     EnvVarArgumentParser,
     add_tgis_args,
@@ -69,17 +68,13 @@ def grpc_server_url(grpc_server_thread_port):
 def grpc_server(engine, args, grpc_server_url):
     """Spins up grpc server in a background thread."""
 
-    def health_check():
-        print("Waiting for server to be up...")  # noqa: T201
-        request = HealthCheckRequest(service=TextGenerationService.SERVICE_NAME)
-        resp = Health.Check(
-            request=request,
-            target=grpc_server_url,
-            timeout=1,
+    def _health_check():
+        assert health_check(
+            server_url=grpc_server_url,
             insecure=True,
+            timeout=1,
+            service=TextGenerationService.SERVICE_NAME,
         )
-        assert resp.status == resp.SERVING
-        print("Server is up")  # noqa: T201
 
     global server  # noqa: PLW0602