refactor: use shared manager from main process for storage implementations.
This commit is contained in:
@@ -3,13 +3,12 @@ import time
|
|||||||
import asyncio
|
import asyncio
|
||||||
from typing import Any, final
|
from typing import Any, final
|
||||||
import threading
|
import threading
|
||||||
from multiprocessing import Manager
|
|
||||||
|
|
||||||
import json
|
import json
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
import pipmaster as pm
|
import pipmaster as pm
|
||||||
|
from lightrag.api.utils_api import manager as main_process_manager
|
||||||
|
|
||||||
from lightrag.utils import (
|
from lightrag.utils import (
|
||||||
logger,
|
logger,
|
||||||
@@ -22,7 +21,7 @@ from lightrag.base import (
|
|||||||
if not pm.is_installed("faiss"):
|
if not pm.is_installed("faiss"):
|
||||||
pm.install("faiss")
|
pm.install("faiss")
|
||||||
|
|
||||||
import faiss
|
import faiss # type: ignore
|
||||||
|
|
||||||
# Global variables for shared memory management
|
# Global variables for shared memory management
|
||||||
_init_lock = threading.Lock()
|
_init_lock = threading.Lock()
|
||||||
@@ -37,7 +36,7 @@ def _get_manager():
|
|||||||
with _init_lock:
|
with _init_lock:
|
||||||
if _manager is None:
|
if _manager is None:
|
||||||
try:
|
try:
|
||||||
_manager = Manager()
|
_manager = main_process_manager
|
||||||
_shared_indices = _manager.dict()
|
_shared_indices = _manager.dict()
|
||||||
_shared_meta = _manager.dict()
|
_shared_meta = _manager.dict()
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
@@ -2,7 +2,6 @@ from dataclasses import dataclass
|
|||||||
import os
|
import os
|
||||||
from typing import Any, Union, final
|
from typing import Any, Union, final
|
||||||
import threading
|
import threading
|
||||||
from multiprocessing import Manager
|
|
||||||
|
|
||||||
from lightrag.base import (
|
from lightrag.base import (
|
||||||
DocProcessingStatus,
|
DocProcessingStatus,
|
||||||
@@ -14,6 +13,7 @@ from lightrag.utils import (
|
|||||||
logger,
|
logger,
|
||||||
write_json,
|
write_json,
|
||||||
)
|
)
|
||||||
|
from lightrag.api.utils_api import manager as main_process_manager
|
||||||
|
|
||||||
# Global variables for shared memory management
|
# Global variables for shared memory management
|
||||||
_init_lock = threading.Lock()
|
_init_lock = threading.Lock()
|
||||||
@@ -27,7 +27,7 @@ def _get_manager():
|
|||||||
with _init_lock:
|
with _init_lock:
|
||||||
if _manager is None:
|
if _manager is None:
|
||||||
try:
|
try:
|
||||||
_manager = Manager()
|
_manager = main_process_manager
|
||||||
_shared_doc_status_data = _manager.dict()
|
_shared_doc_status_data = _manager.dict()
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Failed to initialize shared memory manager: {e}")
|
logger.error(f"Failed to initialize shared memory manager: {e}")
|
||||||
|
@@ -3,7 +3,6 @@ import os
|
|||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from typing import Any, final
|
from typing import Any, final
|
||||||
import threading
|
import threading
|
||||||
from multiprocessing import Manager
|
|
||||||
|
|
||||||
from lightrag.base import (
|
from lightrag.base import (
|
||||||
BaseKVStorage,
|
BaseKVStorage,
|
||||||
@@ -13,6 +12,7 @@ from lightrag.utils import (
|
|||||||
logger,
|
logger,
|
||||||
write_json,
|
write_json,
|
||||||
)
|
)
|
||||||
|
from lightrag.api.utils_api import manager as main_process_manager
|
||||||
|
|
||||||
# Global variables for shared memory management
|
# Global variables for shared memory management
|
||||||
_init_lock = threading.Lock()
|
_init_lock = threading.Lock()
|
||||||
@@ -26,7 +26,7 @@ def _get_manager():
|
|||||||
with _init_lock:
|
with _init_lock:
|
||||||
if _manager is None:
|
if _manager is None:
|
||||||
try:
|
try:
|
||||||
_manager = Manager()
|
_manager = main_process_manager
|
||||||
_shared_kv_data = _manager.dict()
|
_shared_kv_data = _manager.dict()
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Failed to initialize shared memory manager: {e}")
|
logger.error(f"Failed to initialize shared memory manager: {e}")
|
||||||
|
@@ -4,14 +4,13 @@ from typing import Any, final
|
|||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import threading
|
import threading
|
||||||
from multiprocessing import Manager
|
|
||||||
|
|
||||||
import time
|
import time
|
||||||
|
|
||||||
from lightrag.utils import (
|
from lightrag.utils import (
|
||||||
logger,
|
logger,
|
||||||
compute_mdhash_id,
|
compute_mdhash_id,
|
||||||
)
|
)
|
||||||
|
from lightrag.api.utils_api import manager as main_process_manager
|
||||||
import pipmaster as pm
|
import pipmaster as pm
|
||||||
from lightrag.base import (
|
from lightrag.base import (
|
||||||
BaseVectorStorage,
|
BaseVectorStorage,
|
||||||
@@ -34,7 +33,7 @@ def _get_manager():
|
|||||||
with _init_lock:
|
with _init_lock:
|
||||||
if _manager is None:
|
if _manager is None:
|
||||||
try:
|
try:
|
||||||
_manager = Manager()
|
_manager = main_process_manager
|
||||||
_shared_vector_clients = _manager.dict()
|
_shared_vector_clients = _manager.dict()
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Failed to initialize shared memory manager: {e}")
|
logger.error(f"Failed to initialize shared memory manager: {e}")
|
||||||
|
@@ -2,14 +2,13 @@ import os
|
|||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from typing import Any, final
|
from typing import Any, final
|
||||||
import threading
|
import threading
|
||||||
from multiprocessing import Manager
|
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
|
||||||
from lightrag.types import KnowledgeGraph, KnowledgeGraphNode, KnowledgeGraphEdge
|
from lightrag.types import KnowledgeGraph, KnowledgeGraphNode, KnowledgeGraphEdge
|
||||||
from lightrag.utils import (
|
from lightrag.utils import (
|
||||||
logger,
|
logger,
|
||||||
)
|
)
|
||||||
|
from lightrag.api.utils_api import manager as main_process_manager
|
||||||
|
|
||||||
from lightrag.base import (
|
from lightrag.base import (
|
||||||
BaseGraphStorage,
|
BaseGraphStorage,
|
||||||
@@ -37,7 +36,7 @@ def _get_manager():
|
|||||||
with _init_lock:
|
with _init_lock:
|
||||||
if _manager is None:
|
if _manager is None:
|
||||||
try:
|
try:
|
||||||
_manager = Manager()
|
_manager = main_process_manager
|
||||||
_shared_graphs = _manager.dict()
|
_shared_graphs = _manager.dict()
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Failed to initialize shared memory manager: {e}")
|
logger.error(f"Failed to initialize shared memory manager: {e}")
|
||||||
|
Reference in New Issue
Block a user