yangdx
commited on
Commit
·
bf7cb9a
1
Parent(s):
b2f0891
refactor: use shared manager from main process for storage implementations.
Browse files
lightrag/kg/faiss_impl.py
CHANGED
|
@@ -3,13 +3,12 @@ import time
|
|
| 3 |
import asyncio
|
| 4 |
from typing import Any, final
|
| 5 |
import threading
|
| 6 |
-
from multiprocessing import Manager
|
| 7 |
-
|
| 8 |
import json
|
| 9 |
import numpy as np
|
| 10 |
|
| 11 |
from dataclasses import dataclass
|
| 12 |
import pipmaster as pm
|
|
|
|
| 13 |
|
| 14 |
from lightrag.utils import (
|
| 15 |
logger,
|
|
@@ -22,7 +21,7 @@ from lightrag.base import (
|
|
| 22 |
if not pm.is_installed("faiss"):
|
| 23 |
pm.install("faiss")
|
| 24 |
|
| 25 |
-
import faiss
|
| 26 |
|
| 27 |
# Global variables for shared memory management
|
| 28 |
_init_lock = threading.Lock()
|
|
@@ -37,7 +36,7 @@ def _get_manager():
|
|
| 37 |
with _init_lock:
|
| 38 |
if _manager is None:
|
| 39 |
try:
|
| 40 |
-
_manager =
|
| 41 |
_shared_indices = _manager.dict()
|
| 42 |
_shared_meta = _manager.dict()
|
| 43 |
except Exception as e:
|
|
|
|
| 3 |
import asyncio
|
| 4 |
from typing import Any, final
|
| 5 |
import threading
|
|
|
|
|
|
|
| 6 |
import json
|
| 7 |
import numpy as np
|
| 8 |
|
| 9 |
from dataclasses import dataclass
|
| 10 |
import pipmaster as pm
|
| 11 |
+
from lightrag.api.utils_api import manager as main_process_manager
|
| 12 |
|
| 13 |
from lightrag.utils import (
|
| 14 |
logger,
|
|
|
|
| 21 |
if not pm.is_installed("faiss"):
|
| 22 |
pm.install("faiss")
|
| 23 |
|
| 24 |
+
import faiss # type: ignore
|
| 25 |
|
| 26 |
# Global variables for shared memory management
|
| 27 |
_init_lock = threading.Lock()
|
|
|
|
| 36 |
with _init_lock:
|
| 37 |
if _manager is None:
|
| 38 |
try:
|
| 39 |
+
_manager = main_process_manager
|
| 40 |
_shared_indices = _manager.dict()
|
| 41 |
_shared_meta = _manager.dict()
|
| 42 |
except Exception as e:
|
lightrag/kg/json_doc_status_impl.py
CHANGED
|
@@ -2,7 +2,6 @@ from dataclasses import dataclass
|
|
| 2 |
import os
|
| 3 |
from typing import Any, Union, final
|
| 4 |
import threading
|
| 5 |
-
from multiprocessing import Manager
|
| 6 |
|
| 7 |
from lightrag.base import (
|
| 8 |
DocProcessingStatus,
|
|
@@ -14,6 +13,7 @@ from lightrag.utils import (
|
|
| 14 |
logger,
|
| 15 |
write_json,
|
| 16 |
)
|
|
|
|
| 17 |
|
| 18 |
# Global variables for shared memory management
|
| 19 |
_init_lock = threading.Lock()
|
|
@@ -27,7 +27,7 @@ def _get_manager():
|
|
| 27 |
with _init_lock:
|
| 28 |
if _manager is None:
|
| 29 |
try:
|
| 30 |
-
_manager =
|
| 31 |
_shared_doc_status_data = _manager.dict()
|
| 32 |
except Exception as e:
|
| 33 |
logger.error(f"Failed to initialize shared memory manager: {e}")
|
|
|
|
| 2 |
import os
|
| 3 |
from typing import Any, Union, final
|
| 4 |
import threading
|
|
|
|
| 5 |
|
| 6 |
from lightrag.base import (
|
| 7 |
DocProcessingStatus,
|
|
|
|
| 13 |
logger,
|
| 14 |
write_json,
|
| 15 |
)
|
| 16 |
+
from lightrag.api.utils_api import manager as main_process_manager
|
| 17 |
|
| 18 |
# Global variables for shared memory management
|
| 19 |
_init_lock = threading.Lock()
|
|
|
|
| 27 |
with _init_lock:
|
| 28 |
if _manager is None:
|
| 29 |
try:
|
| 30 |
+
_manager = main_process_manager
|
| 31 |
_shared_doc_status_data = _manager.dict()
|
| 32 |
except Exception as e:
|
| 33 |
logger.error(f"Failed to initialize shared memory manager: {e}")
|
lightrag/kg/json_kv_impl.py
CHANGED
|
@@ -3,7 +3,6 @@ import os
|
|
| 3 |
from dataclasses import dataclass
|
| 4 |
from typing import Any, final
|
| 5 |
import threading
|
| 6 |
-
from multiprocessing import Manager
|
| 7 |
|
| 8 |
from lightrag.base import (
|
| 9 |
BaseKVStorage,
|
|
@@ -13,6 +12,7 @@ from lightrag.utils import (
|
|
| 13 |
logger,
|
| 14 |
write_json,
|
| 15 |
)
|
|
|
|
| 16 |
|
| 17 |
# Global variables for shared memory management
|
| 18 |
_init_lock = threading.Lock()
|
|
@@ -26,7 +26,7 @@ def _get_manager():
|
|
| 26 |
with _init_lock:
|
| 27 |
if _manager is None:
|
| 28 |
try:
|
| 29 |
-
_manager =
|
| 30 |
_shared_kv_data = _manager.dict()
|
| 31 |
except Exception as e:
|
| 32 |
logger.error(f"Failed to initialize shared memory manager: {e}")
|
|
|
|
| 3 |
from dataclasses import dataclass
|
| 4 |
from typing import Any, final
|
| 5 |
import threading
|
|
|
|
| 6 |
|
| 7 |
from lightrag.base import (
|
| 8 |
BaseKVStorage,
|
|
|
|
| 12 |
logger,
|
| 13 |
write_json,
|
| 14 |
)
|
| 15 |
+
from lightrag.api.utils_api import manager as main_process_manager
|
| 16 |
|
| 17 |
# Global variables for shared memory management
|
| 18 |
_init_lock = threading.Lock()
|
|
|
|
| 26 |
with _init_lock:
|
| 27 |
if _manager is None:
|
| 28 |
try:
|
| 29 |
+
_manager = main_process_manager
|
| 30 |
_shared_kv_data = _manager.dict()
|
| 31 |
except Exception as e:
|
| 32 |
logger.error(f"Failed to initialize shared memory manager: {e}")
|
lightrag/kg/nano_vector_db_impl.py
CHANGED
|
@@ -4,14 +4,13 @@ from typing import Any, final
|
|
| 4 |
from dataclasses import dataclass
|
| 5 |
import numpy as np
|
| 6 |
import threading
|
| 7 |
-
from multiprocessing import Manager
|
| 8 |
-
|
| 9 |
import time
|
| 10 |
|
| 11 |
from lightrag.utils import (
|
| 12 |
logger,
|
| 13 |
compute_mdhash_id,
|
| 14 |
)
|
|
|
|
| 15 |
import pipmaster as pm
|
| 16 |
from lightrag.base import (
|
| 17 |
BaseVectorStorage,
|
|
@@ -34,7 +33,7 @@ def _get_manager():
|
|
| 34 |
with _init_lock:
|
| 35 |
if _manager is None:
|
| 36 |
try:
|
| 37 |
-
_manager =
|
| 38 |
_shared_vector_clients = _manager.dict()
|
| 39 |
except Exception as e:
|
| 40 |
logger.error(f"Failed to initialize shared memory manager: {e}")
|
|
|
|
| 4 |
from dataclasses import dataclass
|
| 5 |
import numpy as np
|
| 6 |
import threading
|
|
|
|
|
|
|
| 7 |
import time
|
| 8 |
|
| 9 |
from lightrag.utils import (
|
| 10 |
logger,
|
| 11 |
compute_mdhash_id,
|
| 12 |
)
|
| 13 |
+
from lightrag.api.utils_api import manager as main_process_manager
|
| 14 |
import pipmaster as pm
|
| 15 |
from lightrag.base import (
|
| 16 |
BaseVectorStorage,
|
|
|
|
| 33 |
with _init_lock:
|
| 34 |
if _manager is None:
|
| 35 |
try:
|
| 36 |
+
_manager = main_process_manager
|
| 37 |
_shared_vector_clients = _manager.dict()
|
| 38 |
except Exception as e:
|
| 39 |
logger.error(f"Failed to initialize shared memory manager: {e}")
|
lightrag/kg/networkx_impl.py
CHANGED
|
@@ -2,14 +2,13 @@ import os
|
|
| 2 |
from dataclasses import dataclass
|
| 3 |
from typing import Any, final
|
| 4 |
import threading
|
| 5 |
-
from multiprocessing import Manager
|
| 6 |
-
|
| 7 |
import numpy as np
|
| 8 |
|
| 9 |
from lightrag.types import KnowledgeGraph, KnowledgeGraphNode, KnowledgeGraphEdge
|
| 10 |
from lightrag.utils import (
|
| 11 |
logger,
|
| 12 |
)
|
|
|
|
| 13 |
|
| 14 |
from lightrag.base import (
|
| 15 |
BaseGraphStorage,
|
|
@@ -37,7 +36,7 @@ def _get_manager():
|
|
| 37 |
with _init_lock:
|
| 38 |
if _manager is None:
|
| 39 |
try:
|
| 40 |
-
_manager =
|
| 41 |
_shared_graphs = _manager.dict()
|
| 42 |
except Exception as e:
|
| 43 |
logger.error(f"Failed to initialize shared memory manager: {e}")
|
|
|
|
| 2 |
from dataclasses import dataclass
|
| 3 |
from typing import Any, final
|
| 4 |
import threading
|
|
|
|
|
|
|
| 5 |
import numpy as np
|
| 6 |
|
| 7 |
from lightrag.types import KnowledgeGraph, KnowledgeGraphNode, KnowledgeGraphEdge
|
| 8 |
from lightrag.utils import (
|
| 9 |
logger,
|
| 10 |
)
|
| 11 |
+
from lightrag.api.utils_api import manager as main_process_manager
|
| 12 |
|
| 13 |
from lightrag.base import (
|
| 14 |
BaseGraphStorage,
|
|
|
|
| 36 |
with _init_lock:
|
| 37 |
if _manager is None:
|
| 38 |
try:
|
| 39 |
+
_manager = main_process_manager
|
| 40 |
_shared_graphs = _manager.dict()
|
| 41 |
except Exception as e:
|
| 42 |
logger.error(f"Failed to initialize shared memory manager: {e}")
|