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:
yangdx 2025-08-01 10:56:18 +08:00
parent 790abf148b
commit fdf051c234

View file

@ -7,6 +7,8 @@ import asyncio
import os import os
import logging import logging
import logging.config import logging.config
import signal
import sys
import uvicorn import uvicorn
import pipmaster as pm import pipmaster as pm
from fastapi.staticfiles import StaticFiles from fastapi.staticfiles import StaticFiles
@ -28,7 +30,6 @@ from .config import (
get_default_host, get_default_host,
) )
from lightrag.utils import get_env_value from lightrag.utils import get_env_value
import sys
from lightrag import LightRAG, __version__ as core_version from lightrag import LightRAG, __version__ as core_version
from lightrag.api import __api_version__ from lightrag.api import __api_version__
from lightrag.types import GPTKeywordExtractionFormat from lightrag.types import GPTKeywordExtractionFormat
@ -53,6 +54,7 @@ from lightrag.kg.shared_storage import (
get_pipeline_status_lock, get_pipeline_status_lock,
initialize_pipeline_status, initialize_pipeline_status,
cleanup_keyed_lock, cleanup_keyed_lock,
finalize_share_data,
) )
from fastapi.security import OAuth2PasswordRequestForm from fastapi.security import OAuth2PasswordRequestForm
from lightrag.api.auth import auth_handler from lightrag.api.auth import auth_handler
@ -74,6 +76,24 @@ config.read("config.ini")
auth_configured = bool(auth_handler.accounts) 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): def create_app(args):
# Setup logging # Setup logging
logger.setLevel(args.log_level) logger.setLevel(args.log_level)
@ -159,6 +179,9 @@ def create_app(args):
# Clean up database connections # Clean up database connections
await rag.finalize_storages() await rag.finalize_storages()
# Clean up shared data
finalize_share_data()
# Initialize FastAPI # Initialize FastAPI
app_kwargs = { app_kwargs = {
"title": "LightRAG Server API", "title": "LightRAG Server API",
@ -737,6 +760,9 @@ def main():
update_uvicorn_mode_config() update_uvicorn_mode_config()
display_splash_screen(global_args) display_splash_screen(global_args)
# Setup signal handlers for graceful shutdown
setup_signal_handlers()
# Create application instance directly instead of using factory function # Create application instance directly instead of using factory function
app = create_app(global_args) app = create_app(global_args)