Add graceful shutdown handling for LightRAG server
- Setup signal handlers for SIGINT/SIGTERM - Clean up shared resources on shutdown - Finalize shared data storage
This commit is contained in:
parent
790abf148b
commit
fdf051c234
1 changed files with 27 additions and 1 deletions
|
|
@ -7,6 +7,8 @@ import asyncio
|
|||
import os
|
||||
import logging
|
||||
import logging.config
|
||||
import signal
|
||||
import sys
|
||||
import uvicorn
|
||||
import pipmaster as pm
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
|
|
@ -28,7 +30,6 @@ from .config import (
|
|||
get_default_host,
|
||||
)
|
||||
from lightrag.utils import get_env_value
|
||||
import sys
|
||||
from lightrag import LightRAG, __version__ as core_version
|
||||
from lightrag.api import __api_version__
|
||||
from lightrag.types import GPTKeywordExtractionFormat
|
||||
|
|
@ -53,6 +54,7 @@ from lightrag.kg.shared_storage import (
|
|||
get_pipeline_status_lock,
|
||||
initialize_pipeline_status,
|
||||
cleanup_keyed_lock,
|
||||
finalize_share_data,
|
||||
)
|
||||
from fastapi.security import OAuth2PasswordRequestForm
|
||||
from lightrag.api.auth import auth_handler
|
||||
|
|
@ -74,6 +76,24 @@ config.read("config.ini")
|
|||
auth_configured = bool(auth_handler.accounts)
|
||||
|
||||
|
||||
def setup_signal_handlers():
|
||||
"""Setup signal handlers for graceful shutdown"""
|
||||
|
||||
def signal_handler(sig, frame):
|
||||
print(f"\n\nReceived signal {sig}, shutting down gracefully...")
|
||||
print(f"Process ID: {os.getpid()}")
|
||||
|
||||
# Release shared resources
|
||||
finalize_share_data()
|
||||
|
||||
# Exit with success status
|
||||
sys.exit(0)
|
||||
|
||||
# Register signal handlers
|
||||
signal.signal(signal.SIGINT, signal_handler) # Ctrl+C
|
||||
signal.signal(signal.SIGTERM, signal_handler) # kill command
|
||||
|
||||
|
||||
def create_app(args):
|
||||
# Setup logging
|
||||
logger.setLevel(args.log_level)
|
||||
|
|
@ -159,6 +179,9 @@ def create_app(args):
|
|||
# Clean up database connections
|
||||
await rag.finalize_storages()
|
||||
|
||||
# Clean up shared data
|
||||
finalize_share_data()
|
||||
|
||||
# Initialize FastAPI
|
||||
app_kwargs = {
|
||||
"title": "LightRAG Server API",
|
||||
|
|
@ -737,6 +760,9 @@ def main():
|
|||
update_uvicorn_mode_config()
|
||||
display_splash_screen(global_args)
|
||||
|
||||
# Setup signal handlers for graceful shutdown
|
||||
setup_signal_handlers()
|
||||
|
||||
# Create application instance directly instead of using factory function
|
||||
app = create_app(global_args)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue