diff --git a/api/ragflow_server.py b/api/ragflow_server.py index d78990360..ac05dadda 100644 --- a/api/ragflow_server.py +++ b/api/ragflow_server.py @@ -28,6 +28,7 @@ import sys import time import traceback from concurrent.futures import ThreadPoolExecutor +import threading from werkzeug.serving import run_simple from api import settings @@ -42,15 +43,21 @@ from api.versions import get_ragflow_version from api.utils import show_configs from rag.settings import print_rag_settings +stop_event = threading.Event() def update_progress(): - while True: - time.sleep(6) + while not stop_event.is_set(): try: DocumentService.update_progress() + stop_event.wait(6) except Exception: logging.exception("update_progress exception") +def signal_handler(sig, frame): + logging.info("Received interrupt signal, shutting down...") + stop_event.set() + time.sleep(1) + sys.exit(0) if __name__ == '__main__': logging.info(r""" @@ -96,6 +103,9 @@ if __name__ == '__main__': RuntimeConfig.init_env() RuntimeConfig.init_config(JOB_SERVER_HOST=settings.HOST_IP, HTTP_PORT=settings.HOST_PORT) + signal.signal(signal.SIGINT, signal_handler) + signal.signal(signal.SIGTERM, signal_handler) + thread = ThreadPoolExecutor(max_workers=1) thread.submit(update_progress) @@ -112,4 +122,6 @@ if __name__ == '__main__': ) except Exception: traceback.print_exc() + stop_event.set() + time.sleep(1) os.kill(os.getpid(), signal.SIGKILL) diff --git a/rag/svr/task_executor.py b/rag/svr/task_executor.py index 0dd239742..0d9eac99b 100644 --- a/rag/svr/task_executor.py +++ b/rag/svr/task_executor.py @@ -754,12 +754,33 @@ def main(): if TRACE_MALLOC_ENABLED: start_tracemalloc_and_snapshot(None, None) + # Create an event to signal the background thread to exit + stop_event = threading.Event() + background_thread = threading.Thread(target=report_status) background_thread.daemon = True background_thread.start() + + # Handle SIGINT (Ctrl+C) + def signal_handler(sig, frame): + logging.info("Received Ctrl+C, shutting down gracefully...") + stop_event.set() + # Give the background thread time to clean up + if background_thread.is_alive(): + background_thread.join(timeout=5) + logging.info("Exiting...") + sys.exit(0) - while True: - handle_task() + signal.signal(signal.SIGINT, signal_handler) + + try: + while not stop_event.is_set(): + handle_task() + except KeyboardInterrupt: + logging.info("Interrupted by keyboard, shutting down...") + stop_event.set() + if background_thread.is_alive(): + background_thread.join(timeout=5) if __name__ == "__main__": main()