diff --git a/scripts/start_grpc_runtime_server.py b/scripts/start_grpc_runtime_server.py index 732fa42..1281825 100755 --- a/scripts/start_grpc_runtime_server.py +++ b/scripts/start_grpc_runtime_server.py @@ -34,6 +34,7 @@ import argparse import asyncio +import threading import logging import signal import sys @@ -565,7 +566,7 @@ async def _serve(args: argparse.Namespace) -> int: _LOG.info("kakeya gRPC RuntimeService listening on %s", args.bind) http_server = None - http_task = None + http_thread = None if args.network_http_port: if registry is None or prefill_store is None: raise SystemExit( @@ -594,7 +595,12 @@ async def _serve(args: argparse.Namespace) -> int: port=args.network_http_port, log_level=args.log_level.lower(), )) - http_task = asyncio.create_task(http_server.serve()) + http_thread = threading.Thread( + target=http_server.run, + name="kakeya-network-http", + daemon=True, + ) + http_thread.start() _LOG.info( "inference-network dashboard listening on http://%s:%d/network", args.network_http_host, @@ -641,8 +647,8 @@ def _on_signal(sig: int) -> None: pass if http_server is not None: http_server.should_exit = True - if http_task is not None: - await http_task + if http_thread is not None: + await asyncio.to_thread(http_thread.join, args.shutdown_grace_s) if prefill_hook is not None: prefill_hook.close() await server.stop(grace=args.shutdown_grace_s) diff --git a/tests/inference_engine/bridge/test_runtime_http_thread.py b/tests/inference_engine/bridge/test_runtime_http_thread.py new file mode 100644 index 0000000..c5c5249 --- /dev/null +++ b/tests/inference_engine/bridge/test_runtime_http_thread.py @@ -0,0 +1,14 @@ +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[3] +RUNTIME = ROOT / "scripts" / "start_grpc_runtime_server.py" + + +def test_network_http_runs_outside_blocking_grpc_event_loop(): + source = RUNTIME.read_text() + assert 'name="kakeya-network-http"' in source + assert "target=http_server.run" in source + assert "http_thread.start()" in source + assert "asyncio.create_task(http_server.serve())" not in source + assert "asyncio.to_thread(http_thread.join" in source