From d5e4285346e59de714f1e48a2e295f119ec34e59 Mon Sep 17 00:00:00 2001 From: kingbri Date: Thu, 19 Sep 2024 23:31:29 -0400 Subject: [PATCH] Signals: Split signal handler between sync and async Asyncio requires a closure of the event loop while sync can use SystemExit to kill the program. Signed-off-by: kingbri --- common/signals.py | 18 ++++++++++++------ main.py | 1 + 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/common/signals.py b/common/signals.py index 97f595b..1b47778 100644 --- a/common/signals.py +++ b/common/signals.py @@ -22,11 +22,20 @@ def signal_handler(*_): SHUTTING_DOWN = True # Run async unloads for model - asyncio.ensure_future(signal_handler_async()) + # If an event loop doesn't exist (synchronous), exit. + try: + loop = asyncio.get_running_loop() + loop.run_until_complete(signal_handler_async()) + except RuntimeError: + sys.exit(0) -async def signal_handler_async(*_): - """Internal signal handler. Runs all async code to shut down the program.""" +async def signal_handler_async(): + """ + Internal signal handler. Runs all async code to shut down the program. + + asyncio.run will cancel all remaining tasks and close the event loop. + """ if model.container: await model.unload_model(skip_wait=True, shutdown=True) @@ -34,9 +43,6 @@ async def signal_handler_async(*_): if model.embeddings_container: await model.unload_embedding_model() - # Exit the program - sys.exit(0) - def uvicorn_signal_handler(signal_event: signal.Signals): """Overrides uvicorn's signal handler.""" diff --git a/main.py b/main.py index 06db5d5..fd297ef 100644 --- a/main.py +++ b/main.py @@ -5,6 +5,7 @@ import os import pathlib import platform import signal +import time from loguru import logger from typing import Optional