From 961c87f2c7947a831d2190cee770f1206ec5b415 Mon Sep 17 00:00:00 2001 From: Ihar Hrachyshka Date: Fri, 7 Mar 2025 22:01:08 +0000 Subject: [PATCH] fix: Don't confuse CancelledError from implementation shutdown If implementation raises CancelledError (e.g. when it runs its own async loop for jobs), the main server shutdown handler gets confused and doesn't attempt to shut down the main loop tasks. Handle it by separating exception handling for implementations and for the main loop. While at it, also fixing the following failure when this happens: UnboundLocalError: cannot access local variable 'loop' where it is not associated with a value Signed-off-by: Ihar Hrachyshka --- llama_stack/distribution/server/server.py | 30 +++++++++++------------ 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/llama_stack/distribution/server/server.py b/llama_stack/distribution/server/server.py index c4ef79a69..b85c463ae 100644 --- a/llama_stack/distribution/server/server.py +++ b/llama_stack/distribution/server/server.py @@ -143,23 +143,23 @@ def handle_signal(app, signum, _) -> None: logger.info(f"Received signal {signame} ({signum}). Exiting gracefully...") async def shutdown(): - try: - # Gracefully shut down implementations - for impl in app.__llama_stack_impls__.values(): - impl_name = impl.__class__.__name__ - logger.info("Shutting down %s", impl_name) - try: - if hasattr(impl, "shutdown"): - await asyncio.wait_for(impl.shutdown(), timeout=5) - else: - logger.warning("No shutdown method for %s", impl_name) - except asyncio.TimeoutError: - logger.exception("Shutdown timeout for %s ", impl_name, exc_info=True) - except Exception as e: - logger.exception("Failed to shutdown %s: %s", impl_name, {e}) + # Gracefully shut down implementations + for impl in app.__llama_stack_impls__.values(): + impl_name = impl.__class__.__name__ + logger.info("Shutting down %s", impl_name) + try: + if hasattr(impl, "shutdown"): + await asyncio.wait_for(impl.shutdown(), timeout=5) + else: + logger.warning("No shutdown method for %s", impl_name) + except asyncio.TimeoutError: + logger.exception("Shutdown timeout for %s ", impl_name, exc_info=True) + except (Exception, asyncio.CancelledError) as e: + logger.exception("Failed to shutdown %s: %s", impl_name, {e}) + loop = asyncio.get_running_loop() + try: # Gather all running tasks - loop = asyncio.get_running_loop() tasks = [task for task in asyncio.all_tasks(loop) if task is not asyncio.current_task()] # Cancel all tasks