import os import sys import asyncio from multiprocessing.synchronize import Lock as ProcessLock from multiprocessing import Manager from typing import Any, Dict, Optional, Union # Define a direct print function for critical logs that must be visible in all processes def direct_log(message, level="INFO"): """ Log a message directly to stderr to ensure visibility in all processes, including the Gunicorn master process. """ print(f"{level}: {message}", file=sys.stderr, flush=True) LockType = Union[ProcessLock, asyncio.Lock] is_multiprocess = None _workers = None _manager = None _initialized = None _global_lock: Optional[LockType] = None # shared data for storage across processes _shared_dicts: Optional[Dict[str, Any]] = None _init_flags: Optional[Dict[str, bool]] = None # namespace -> initialized _update_flags: Optional[Dict[str, bool]] = None # namespace -> updated def initialize_share_data(workers: int = 1): """ Initialize shared storage data for single or multi-process mode. When used with Gunicorn's preload feature, this function is called once in the master process before forking worker processes, allowing all workers to share the same initialized data. In single-process mode, this function is called during LightRAG object initialization. The function determines whether to use cross-process shared variables for data storage based on the number of workers. If workers=1, it uses thread locks and local dictionaries. If workers>1, it uses process locks and shared dictionaries managed by multiprocessing.Manager. Args: workers (int): Number of worker processes. If 1, single-process mode is used. If > 1, multi-process mode with shared memory is used. """ global \ _manager, \ _workers, \ is_multiprocess, \ is_multiprocess, \ _global_lock, \ _shared_dicts, \ _init_flags, \ _initialized, \ _update_flags # Check if already initialized if _initialized: direct_log( f"Process {os.getpid()} Shared-Data already initialized (multiprocess={is_multiprocess})" ) return _manager = Manager() _workers = workers if workers > 1: is_multiprocess = True _global_lock = _manager.Lock() _shared_dicts = _manager.dict() _init_flags = _manager.dict() _update_flags = _manager.dict() direct_log( f"Process {os.getpid()} Shared-Data created for Multiple Process (workers={workers})" ) else: is_multiprocess = False _global_lock = asyncio.Lock() _shared_dicts = {} _init_flags = {} _update_flags = {} direct_log(f"Process {os.getpid()} Shared-Data created for Single Process") # Mark as initialized _initialized = True # Initialize pipeline status for document indexing control pipeline_namespace = get_namespace_data("pipeline_status") # Create a shared list object for history_messages history_messages = _manager.list() if is_multiprocess else [] pipeline_namespace.update( { "busy": False, # Control concurrent processes "job_name": "Default Job", # Current job name (indexing files/indexing texts) "job_start": None, # Job start time "docs": 0, # Total number of documents to be indexed "batchs": 0, # Number of batches for processing documents "cur_batch": 0, # Current processing batch "request_pending": False, # Flag for pending request for processing "latest_message": "", # Latest message from pipeline processing "history_messages": history_messages, # 使用共享列表对象 } ) async def get_update_flags(namespace: str): """ Create a updated flags of a specific namespace. Caller must store the dict object locally and use it to determine when to update the storage. """ global _update_flags if _update_flags is None: raise ValueError("Try to create namespace before Shared-Data is initialized") if is_multiprocess: with _global_lock: if namespace not in _update_flags: if _manager is not None: _update_flags[namespace] = _manager.list() direct_log(f"Process {os.getpid()} initialized updated flags for namespace: [{namespace}]") if _manager is not None: new_update_flag = _manager.Value('b', False) _update_flags[namespace].append(new_update_flag) return new_update_flag else: async with _global_lock: if namespace not in _update_flags: _update_flags[namespace] = [] direct_log(f"Process {os.getpid()} initialized updated flags for namespace: [{namespace}]") new_update_flag = False _update_flags[namespace].append(new_update_flag) return new_update_flag async def set_update_flag(namespace: str): """Set all update flag of namespace to indicate storage needs updating""" global _update_flags if _update_flags is None: raise ValueError("Try to create namespace before Shared-Data is initialized") if is_multiprocess: with _global_lock: if namespace not in _update_flags: raise ValueError(f"Namespace {namespace} not found in update flags") # Update flags for multiprocess mode for i in range(len(_update_flags[namespace])): _update_flags[namespace][i].value = True else: async with _global_lock: if namespace not in _update_flags: raise ValueError(f"Namespace {namespace} not found in update flags") # Update flags for single process mode for i in range(len(_update_flags[namespace])): _update_flags[namespace][i] = True def try_initialize_namespace(namespace: str) -> bool: """ Try to initialize a namespace. Returns True if the current process gets initialization permission. Uses atomic operations on shared dictionaries to ensure only one process can successfully initialize. """ global _init_flags, _manager if _init_flags is None: raise ValueError("Try to create nanmespace before Shared-Data is initialized") if namespace not in _init_flags: _init_flags[namespace] = True direct_log( f"Process {os.getpid()} ready to initialize storage namespace: [{namespace}]" ) return True direct_log( f"Process {os.getpid()} storage namespace already to initialized: [{namespace}]" ) return False def get_storage_lock() -> LockType: """return storage lock for data consistency""" return _global_lock async def get_namespace_data(namespace: str) -> Dict[str, Any]: """get storage space for specific storage type(namespace)""" if _shared_dicts is None: direct_log( f"Error: try to getnanmespace before Shared-Data is initialized, pid={os.getpid()}", level="ERROR", ) raise ValueError("Shared dictionaries not initialized") if is_multiprocess: with _global_lock: if namespace not in _shared_dicts: if _manager is not None: _shared_dicts[namespace] = _manager.dict() else: async with _global_lock: if namespace not in _shared_dicts: _shared_dicts[namespace] = {} return _shared_dicts[namespace] def finalize_share_data(): """ Release shared resources and clean up. This function should be called when the application is shutting down to properly release shared resources and avoid memory leaks. In multi-process mode, it shuts down the Manager and releases all shared objects. In single-process mode, it simply resets the global variables. """ global \ _manager, \ is_multiprocess, \ _global_lock, \ _shared_dicts, \ _init_flags, \ _initialized # Check if already initialized if not _initialized: direct_log( f"Process {os.getpid()} storage data not initialized, nothing to finalize" ) return direct_log( f"Process {os.getpid()} finalizing storage data (multiprocess={is_multiprocess})" ) # In multi-process mode, shut down the Manager if is_multiprocess and _manager is not None: try: # Clear shared dictionaries first if _shared_dicts is not None: _shared_dicts.clear() if _init_flags is not None: _init_flags.clear() # Shut down the Manager _manager.shutdown() direct_log(f"Process {os.getpid()} Manager shutdown complete") except Exception as e: direct_log( f"Process {os.getpid()} Error shutting down Manager: {e}", level="ERROR" ) # Reset global variables _manager = None _initialized = None is_multiprocess = None _shared_dicts = None _init_flags = None _global_lock = None direct_log(f"Process {os.getpid()} storage data finalization complete")