Skip to content

Commit

Permalink
Signals: Split signal handler between sync and async
Browse files Browse the repository at this point in the history
Asyncio requires a closure of the event loop while sync can use SystemExit
to kill the program.

Signed-off-by: kingbri <bdashore3@proton.me>
  • Loading branch information
kingbri1 committed Sep 20, 2024
1 parent b4cda78 commit d5e4285
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 6 deletions.
18 changes: 12 additions & 6 deletions common/signals.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,21 +22,27 @@ 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)

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."""
Expand Down
1 change: 1 addition & 0 deletions main.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import pathlib
import platform
import signal
import time
from loguru import logger
from typing import Optional

Expand Down

0 comments on commit d5e4285

Please sign in to comment.