From 03d05b094d04a445732281ea9f7ff952fdd9ad1d Mon Sep 17 00:00:00 2001
From: yangdx <yangdx@znipower.com>
Date: Thu, 27 Feb 2025 14:13:42 +0800
Subject: [PATCH] Improve Gunicorn support and cleanup shared storage
 initialization
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

• Move Gunicorn check before other startup
• Improve startup flow organization
---
 lightrag/api/lightrag_server.py | 20 +++++++-------------
 lightrag/kg/shared_storage.py   |  2 +-
 2 files changed, 8 insertions(+), 14 deletions(-)

diff --git a/lightrag/api/lightrag_server.py b/lightrag/api/lightrag_server.py
index a9c9ab04..9f162290 100644
--- a/lightrag/api/lightrag_server.py
+++ b/lightrag/api/lightrag_server.py
@@ -470,8 +470,13 @@ def configure_logging():
 
 
 def main():
-    from multiprocessing import freeze_support
+    # Check if running under Gunicorn
+    if 'GUNICORN_CMD_ARGS' in os.environ:
+        # If started with Gunicorn, return directly as Gunicorn will call get_application
+        print("Running under Gunicorn - worker management handled by Gunicorn")
+        return
 
+    from multiprocessing import freeze_support
     freeze_support()
 
     args = parse_args()
@@ -482,18 +487,7 @@ def main():
     configure_logging()
 
     display_splash_screen(args)
-
-    # Check if running under Gunicorn
-    if 'GUNICORN_CMD_ARGS' in os.environ:
-        # If started with Gunicorn, return directly as Gunicorn will call get_application
-        print("Running under Gunicorn - worker management handled by Gunicorn")
-        return
-    
-    # If not running under Gunicorn, initialize shared data here
-    from lightrag.kg.shared_storage import initialize_share_data
-    print("Starting in single-process mode")
-    initialize_share_data(1)  # Force single process mode
-    
+        
     # Create application instance directly instead of using factory function
     app = create_app(args)
     
diff --git a/lightrag/kg/shared_storage.py b/lightrag/kg/shared_storage.py
index 8dc9e1a9..8956d995 100644
--- a/lightrag/kg/shared_storage.py
+++ b/lightrag/kg/shared_storage.py
@@ -52,7 +52,7 @@ def initialize_share_data(workers: int = 1):
     # Check if already initialized
     if _initialized and _initialized.value:
         is_multiprocess = _is_multiprocess.value
-        direct_log(f"Process {os.getpid()} storage data already initialized (multiprocess={_is_multiprocess.value})!")
+        direct_log(f"Process {os.getpid()} storage data already initialized (multiprocess={_is_multiprocess.value})")
         return
     
     _manager = Manager()