move run_grpc_server to grpc module

src/vllm_tgis_adapter/__main__.py

@@ -35,7 +35,7 @@
 from vllm.entrypoints.openai.serving_embedding import OpenAIServingEmbedding
 from vllm.usage.usage_lib import UsageContext
 
-from .grpc import start_grpc_server
+from .grpc import run_grpc_server
 from .logging import init_logger
 from .tgis_utils.args import EnvVarArgumentParser, add_tgis_args, postprocess_tgis_args
 
@@ -228,32 +228,6 @@ async def run_http_server(
         await server.shutdown()
 
 
-async def run_grpc_server(
-    engine: AsyncLLMEngine,
-    *,
-    disable_log_stats: bool,
-) -> None:
-    async def _force_log() -> None:
-        while True:
-            await asyncio.sleep(10)
-            await engine.do_log_stats()
-
-    if not disable_log_stats:
-        asyncio.create_task(_force_log())  # noqa: RUF006
-
-    assert args is not None
-
-    server = await start_grpc_server(engine, args)
-
-    try:
-        while True:
-            await asyncio.sleep(60)
-    except asyncio.CancelledError:
-        print("Gracefully stopping gRPC server")  # noqa: T201
-        await server.stop(30)  # TODO configurable grace
-        await server.wait_for_termination()
-
-
 if __name__ == "__main__":
     # convert to our custom env var arg parser
     parser = EnvVarArgumentParser(parser=make_arg_parser())
@@ -322,7 +296,7 @@ async def _force_log() -> None:
         run_http_server(engine, args, model_config)
     )
     grpc_server_task = event_loop.create_task(
-        run_grpc_server(engine, disable_log_stats=engine_args.disable_log_stats)
+        run_grpc_server(engine, args, disable_log_stats=engine_args.disable_log_stats)
     )
 
     def signal_handler() -> None:

src/vllm_tgis_adapter/grpc/__init__.py

@@ -1 +1 @@
-from .grpc_server import start_grpc_server
+from .grpc_server import run_grpc_server

src/vllm_tgis_adapter/grpc/grpc_server.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import inspect
 import time
 import uuid
@@ -804,3 +805,30 @@ async def start_grpc_server(
     logger.info("gRPC Server started at %s", listen_on)
 
     return server
+
+
+async def run_grpc_server(
+    engine: AsyncLLMEngine,
+    args: argparse.Namespace,
+    *,
+    disable_log_stats: bool,
+) -> None:
+    async def _force_log() -> None:
+        while True:
+            await asyncio.sleep(10)
+            await engine.do_log_stats()
+
+    if not disable_log_stats:
+        asyncio.create_task(_force_log())  # noqa: RUF006
+
+    assert args is not None
+
+    server = await start_grpc_server(engine, args)
+
+    try:
+        while True:
+            await asyncio.sleep(60)
+    except asyncio.CancelledError:
+        print("Gracefully stopping gRPC server")  # noqa: T201
+        await server.stop(30)  # TODO configurable grace
+        await server.wait_for_termination()

tests/conftest.py

@@ -13,7 +13,10 @@
 from vllm.entrypoints.openai.serving_completion import OpenAIServingCompletion
 from vllm.entrypoints.openai.serving_embedding import OpenAIServingEmbedding
 
-from vllm_tgis_adapter.grpc.grpc_server import TextGenerationService, start_grpc_server
+import vllm_tgis_adapter
+from vllm_tgis_adapter.__main__ import run_http_server
+from vllm_tgis_adapter.grpc import run_grpc_server
+from vllm_tgis_adapter.grpc.grpc_server import TextGenerationService
 from vllm_tgis_adapter.healthcheck import health_check
 from vllm_tgis_adapter.tgis_utils.args import (
     EnvVarArgumentParser,
@@ -90,7 +93,7 @@ def grpc_server_url(grpc_server_thread_port):
 
 
 @pytest.fixture()
-def grpc_server(engine, args, grpc_server_url):
+def _grpc_server(engine, args, grpc_server_url):
     """Spins up the grpc server in a background thread."""
 
     def _health_check():
@@ -101,34 +104,24 @@ def _health_check():
             service=TextGenerationService.SERVICE_NAME,
         )
 
-    global server  # noqa: PLW0602
-
     loop = asyncio.new_event_loop()
 
-    async def run_server():
-        global server  # noqa: PLW0603
-
-        server = await start_grpc_server(engine, args)
-        while server._server.is_running():  # noqa: SLF001
-            await asyncio.sleep(1)
+    global task  # noqa: PLW0602
 
     def target():
-        loop.run_until_complete(run_server())
+        global task  # noqa: PLW0603
+
+        task = loop.create_task(run_grpc_server(engine, args, disable_log_stats=False))
+        loop.run_until_complete(task)
 
     t = threading.Thread(target=target)
     t.start()
 
-    async def stop():
-        global server  # noqa: PLW0602
-
-        await server.stop(grace=None)
-        await server.wait_for_termination()
-
     try:
         wait_until(_health_check)
-        yield server
+        yield
     finally:
-        loop.create_task(stop())  # noqa: RUF006
+        task.cancel()
         t.join()
 
 

tests/test_grpc_server.py

@@ -1,3 +1,6 @@
-def test_startup(grpc_server):
+import pytest
+
+
+@pytest.mark.usefixtures("_grpc_server")
+def test_startup():
     """Test that the grpc_server fixture starts up properly."""
-    assert grpc_server._server.is_running()  # noqa: SLF001