From 9e3ae503060cc685749839630a379c6e63d0eae6 Mon Sep 17 00:00:00 2001 From: Ashwin Bharambe Date: Fri, 18 Jul 2025 10:29:19 -0700 Subject: [PATCH] feat(server): construct the stack in a persistent event loop (#2818) When we call `construct_stack()`, providers are instantiated and `initialize()` is called. This call can end up doing _anything_ at all -- specifically, providers are free to create long running background tasks as part of this. If we wrapped this within a `asyncio.run()` as in the current code, these tasks get canceled when the stack construction finishes. This is not correct. The PR addresses the issue by creating a persistent event loop which is used for both the stack as well as for running the uvicorn server. In other words, the lifetime of the providers (and downstream async code) is now the same as the lifetime of the uvicorn server. ## Test Plan This should not affect any current code since we don't have background tasks created right now. However, https://github.com/meta-llama/llama-stack/pull/2805 will start using this functionality. --- llama_stack/distribution/server/server.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/llama_stack/distribution/server/server.py b/llama_stack/distribution/server/server.py index 974064b58..fb00b8384 100644 --- a/llama_stack/distribution/server/server.py +++ b/llama_stack/distribution/server/server.py @@ -455,6 +455,7 @@ def main(args: argparse.Namespace | None = None): redoc_url="/redoc", openapi_url="/openapi.json", ) + if not os.environ.get("LLAMA_STACK_DISABLE_VERSION_CHECK"): app.add_middleware(ClientVersionMiddleware) @@ -493,7 +494,13 @@ def main(args: argparse.Namespace | None = None): ) try: - impls = asyncio.run(construct_stack(config)) + # Create and set the event loop that will be used for both construction and server runtime + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + # Construct the stack in the persistent event loop + impls = loop.run_until_complete(construct_stack(config)) + except InvalidProviderError as e: logger.error(f"Error: {str(e)}") sys.exit(1) @@ -591,7 +598,8 @@ def main(args: argparse.Namespace | None = None): if ssl_config: uvicorn_config.update(ssl_config) - uvicorn.run(**uvicorn_config) + # Run uvicorn in the existing event loop to preserve background tasks + loop.run_until_complete(uvicorn.Server(uvicorn.Config(**uvicorn_config)).serve()) def extract_path_params(route: str) -> list[str]: