From 7bf2b07d4c5b5821f2451371bdd8f3ed8a84e207 Mon Sep 17 00:00:00 2001 From: kingbri Date: Fri, 2 Aug 2024 15:10:27 -0400 Subject: [PATCH] Signals: Exit on async cleanup The async signal exit function should be the internal for exiting the program. In addition, prevent the handler from being called twice by adding a boolean. May become an asyncio event later on. In addition, make sure to skip_wait when running model.unload. Signed-off-by: kingbri --- common/signals.py | 21 ++++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/common/signals.py b/common/signals.py index f0b7f19..d4e144c 100644 --- a/common/signals.py +++ b/common/signals.py @@ -7,24 +7,35 @@ from types import FrameType from common import model +SHUTTING_DOWN: bool = False + + def signal_handler(*_): """Signal handler for main function. Run before uvicorn starts.""" + global SHUTTING_DOWN + + if SHUTTING_DOWN: + return + logger.warning("Shutdown signal called. Exiting gracefully.") + SHUTTING_DOWN = True # Run async unloads for model asyncio.ensure_future(signal_handler_async()) - # Exit the program - sys.exit(0) - async def signal_handler_async(*_): + """Internal signal handler. Runs all async code to shut down the program.""" + if model.container: - await model.container.unload() + await model.unload_model(skip_wait=True) if model.embeddings_container: - await model.embeddings_container.unload() + await model.unload_embedding_model() + + # Exit the program + sys.exit(0) def uvicorn_signal_handler(signal_event: signal.Signals):