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 = 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 = 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 = 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 = 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 = 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}")