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 <bdashore3@proton.me>
This commit is contained in:
kingbri 2024-09-19 23:31:29 -04:00
parent b4cda78bcc
commit d5e4285346
2 changed files with 13 additions and 6 deletions

View file

@ -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."""