feat: use storage actor

Author: ChenZiHong-Gavin

graphgen/common/init_storage.py

@@ -1,27 +1,257 @@
-from graphgen.models import JsonKVStorage, NetworkXStorage
+from typing import Any, Dict, Union
+
+import ray
+
+from graphgen.bases.base_storage import BaseGraphStorage, BaseKVStorage
+
+
+class KVStorageActor:
+    def __init__(self, backend: str, working_dir: str, namespace: str):
+        if backend == "json_kv":
+            from graphgen.models import JsonKVStorage
+
+            self.kv = JsonKVStorage(working_dir, namespace)
+        elif backend == "rocksdb":
+            from graphgen.models import RocksDBKVStorage
+
+            self.kv = RocksDBKVStorage(working_dir, namespace)
+        else:
+            raise ValueError(f"Unknown KV backend: {backend}")
+
+    def data(self) -> Dict[str, Dict]:
+        return self.kv.data
+
+    def all_keys(self) -> list[str]:
+        return self.kv.all_keys()
+
+    def index_done_callback(self):
+        return self.kv.index_done_callback()
+
+    def get_by_id(self, id: str) -> Dict:
+        return self.kv.get_by_id(id)
+
+    def get_by_ids(self, ids: list[str], fields=None) -> list:
+        return self.kv.get_by_ids(ids, fields)
+
+    def get_all(self) -> Dict[str, Dict]:
+        return self.kv.get_all()
+
+    def filter_keys(self, data: list[str]) -> set[str]:
+        return self.kv.filter_keys(data)
+
+    def upsert(self, data: dict) -> dict:
+        return self.kv.upsert(data)
+
+    def drop(self):
+        return self.kv.drop()
+
+    def reload(self):
+        return self.kv.reload()
+
+
+class GraphStorageActor:
+    def __init__(self, backend: str, working_dir: str, namespace: str):
+        if backend == "networkx":
+            from graphgen.models import NetworkXStorage
+
+            self.graph = NetworkXStorage(working_dir, namespace)
+        else:
+            raise ValueError(f"Unknown Graph backend: {backend}")
+
+    def index_done_callback(self):
+        return self.graph.index_done_callback()
+
+    def has_node(self, node_id: str) -> bool:
+        return self.graph.has_node(node_id)
+
+    def has_edge(self, source_node_id: str, target_node_id: str):
+        return self.graph.has_edge(source_node_id, target_node_id)
+
+    def node_degree(self, node_id: str) -> int:
+        return self.graph.node_degree(node_id)
+
+    def edge_degree(self, src_id: str, tgt_id: str) -> int:
+        return self.graph.edge_degree(src_id, tgt_id)
+
+    def get_node(self, node_id: str) -> Any:
+        return self.graph.get_node(node_id)
+
+    def update_node(self, node_id: str, node_data: dict[str, str]):
+        return self.graph.update_node(node_id, node_data)
+
+    def get_all_nodes(self) -> Any:
+        return self.graph.get_all_nodes()
+
+    def get_edge(self, source_node_id: str, target_node_id: str):
+        return self.graph.get_edge(source_node_id, target_node_id)
+
+    def update_edge(
+        self, source_node_id: str, target_node_id: str, edge_data: dict[str, str]
+    ):
+        return self.graph.update_edge(source_node_id, target_node_id, edge_data)
+
+    def get_all_edges(self) -> Any:
+        return self.graph.get_all_edges()
+
+    def get_node_edges(self, source_node_id: str) -> Any:
+        return self.graph.get_node_edges(source_node_id)
+
+    def upsert_node(self, node_id: str, node_data: dict[str, str]):
+        return self.graph.upsert_node(node_id, node_data)
+
+    def upsert_edge(
+        self, source_node_id: str, target_node_id: str, edge_data: dict[str, str]
+    ):
+        return self.graph.upsert_edge(source_node_id, target_node_id, edge_data)
+
+    def delete_node(self, node_id: str):
+        return self.graph.delete_node(node_id)
+
+    def reload(self):
+        return self.graph.reload()
+
+
+def get_actor_handle(name: str):
+    try:
+        return ray.get_actor(name)
+    except ValueError as exc:
+        raise RuntimeError(
+            f"Actor {name} not found. Make sure it is created before accessing."
+        ) from exc
+
+
+class RemoteKVStorageProxy(BaseKVStorage):
+    def __init__(self, namespace: str):
+        super().__init__()
+        self.namespace = namespace
+        self.actor_name = f"Actor_KV_{namespace}"
+        self.actor = get_actor_handle(self.actor_name)
+
+    def data(self) -> Dict[str, Any]:
+        return ray.get(self.actor.data.remote())
+
+    def all_keys(self) -> list[str]:
+        return ray.get(self.actor.all_keys.remote())
+
+    def index_done_callback(self):
+        return ray.get(self.actor.index_done_callback.remote())
+
+    def get_by_id(self, id: str) -> Union[Any, None]:
+        return ray.get(self.actor.get_by_id.remote(id))
+
+    def get_by_ids(self, ids: list[str], fields=None) -> list[Any]:
+        return ray.get(self.actor.get_by_ids.remote(ids, fields))
+
+    def get_all(self) -> Dict[str, Any]:
+        return ray.get(self.actor.get_all.remote())
+
+    def filter_keys(self, data: list[str]) -> set[str]:
+        return ray.get(self.actor.filter_keys.remote(data))
+
+    def upsert(self, data: Dict[str, Any]):
+        return ray.get(self.actor.upsert.remote(data))
+
+    def drop(self):
+        return ray.get(self.actor.drop.remote())
+
+    def reload(self):
+        return ray.get(self.actor.reload.remote())
+
+
+class RemoteGraphStorageProxy(BaseGraphStorage):
+    def __init__(self, namespace: str):
+        super().__init__()
+        self.namespace = namespace
+        self.actor_name = f"Actor_Graph_{namespace}"
+        self.actor = get_actor_handle(self.actor_name)
+
+    def index_done_callback(self):
+        return ray.get(self.actor.index_done_callback.remote())
+
+    def has_node(self, node_id: str) -> bool:
+        return ray.get(self.actor.has_node.remote(node_id))
+
+    def has_edge(self, source_node_id: str, target_node_id: str):
+        return ray.get(self.actor.has_edge.remote(source_node_id, target_node_id))
+
+    def node_degree(self, node_id: str) -> int:
+        return ray.get(self.actor.node_degree.remote(node_id))
+
+    def edge_degree(self, src_id: str, tgt_id: str) -> int:
+        return ray.get(self.actor.edge_degree.remote(src_id, tgt_id))
+
+    def get_node(self, node_id: str) -> Any:
+        return ray.get(self.actor.get_node.remote(node_id))
+
+    def update_node(self, node_id: str, node_data: dict[str, str]):
+        return ray.get(self.actor.update_node.remote(node_id, node_data))
+
+    def get_all_nodes(self) -> Any:
+        return ray.get(self.actor.get_all_nodes.remote())
+
+    def get_edge(self, source_node_id: str, target_node_id: str):
+        return ray.get(self.actor.get_edge.remote(source_node_id, target_node_id))
+
+    def update_edge(
+        self, source_node_id: str, target_node_id: str, edge_data: dict[str, str]
+    ):
+        return ray.get(
+            self.actor.update_edge.remote(source_node_id, target_node_id, edge_data)
+        )
+
+    def get_all_edges(self) -> Any:
+        return ray.get(self.actor.get_all_edges.remote())
+
+    def get_node_edges(self, source_node_id: str) -> Any:
+        return ray.get(self.actor.get_node_edges.remote(source_node_id))
+
+    def upsert_node(self, node_id: str, node_data: dict[str, str]):
+        return ray.get(self.actor.upsert_node.remote(node_id, node_data))
+
+    def upsert_edge(
+        self, source_node_id: str, target_node_id: str, edge_data: dict[str, str]
+    ):
+        return ray.get(
+            self.actor.upsert_edge.remote(source_node_id, target_node_id, edge_data)
+        )
+
+    def delete_node(self, node_id: str):
+        return ray.get(self.actor.delete_node.remote(node_id))
+
+    def reload(self):
+        return ray.get(self.actor.reload.remote())
 
 
 class StorageFactory:
     """
     Factory class to create storage instances based on backend.
-    Supported backends:
-        kv_storage(key-value storage):
-            - json_kv: JsonKVStorage
-        graph_storage:
-            - networkx: NetworkXStorage (graph storage)
     """
 
     @staticmethod
     def create_storage(backend: str, working_dir: str, namespace: str):
-        if backend == "json_kv":
-            return JsonKVStorage(working_dir, namespace=namespace)
-
-        if backend == "networkx":
-            return NetworkXStorage(working_dir, namespace=namespace)
-
-        raise NotImplementedError(
-            f"Storage backend '{backend}' is not implemented yet."
-        )
+        if backend in ["json_kv", "rocksdb"]:
+            actor_name = f"Actor_KV_{namespace}"
+            try:
+                ray.get_actor(actor_name)
+            except ValueError:
+                ray.remote(KVStorageActor).options(
+                    name=actor_name,
+                    lifetime="detached",
+                    get_if_exists=True,
+                ).remote(backend, working_dir, namespace)
+            return RemoteKVStorageProxy(namespace)
+        if backend in ["networkx"]:
+            actor_name = f"Actor_Graph_{namespace}"
+            try:
+                ray.get_actor(actor_name)
+            except ValueError:
+                ray.remote(GraphStorageActor).options(
+                    name=actor_name,
+                    lifetime="detached",
+                    get_if_exists=True,
+                ).remote(backend, working_dir, namespace)
+            return RemoteGraphStorageProxy(namespace)
+        raise ValueError(f"Unknown storage backend: {backend}")
 
 
 def init_storage(backend: str, working_dir: str, namespace: str):

graphgen/engine.py

@@ -8,6 +8,7 @@
 import ray.data
 
 from graphgen.bases import Config, Node
+from graphgen.utils import logger
 
 
 class Engine:
@@ -26,7 +27,7 @@ def __init__(
                 log_to_driver=True,
                 **ray_init_kwargs,
             )
-            print(f"Ray Dashboard URL: {context.dashboard_url}")
+            logger.info("Ray Dashboard URL: %s", context.dashboard_url)
 
     @staticmethod
     def _topo_sort(nodes: List[Node]) -> List[Node]:

graphgen/models/__init__.py

@@ -32,5 +32,5 @@
 from .searcher.web.bing_search import BingSearch
 from .searcher.web.google_search import GoogleSearch
 from .splitter import ChineseRecursiveTextSplitter, RecursiveCharacterSplitter
-from .storage import JsonKVStorage, NetworkXStorage, RocksDBCache
+from .storage import JsonKVStorage, NetworkXStorage, RocksDBCache, RocksDBKVStorage
 from .tokenizer import Tokenizer

graphgen/models/storage/__init__.py

@@ -1,4 +1,5 @@
 from graphgen.models.storage.graph.networkx_storage import NetworkXStorage
 from graphgen.models.storage.kv.json_storage import JsonKVStorage
+from graphgen.models.storage.kv.rocksdb_storage import RocksDBKVStorage
 
 from .rocksdb_cache import RocksDBCache

graphgen/models/storage/kv/rocksdb_storage.py

@@ -7,7 +7,6 @@
 from rocksdict import Rdict
 
 from graphgen.bases.base_storage import BaseKVStorage
-from graphgen.utils import logger
 
 
 @dataclass
@@ -18,7 +17,9 @@ class RocksDBKVStorage(BaseKVStorage):
     def __post_init__(self):
         self._db_path = os.path.join(self.working_dir, f"{self.namespace}.db")
         self._db = Rdict(self._db_path)
-        logger.info("Load KV (RocksDB) %s at %s", self.namespace, self._db_path)
+        print(
+            f"RocksDBKVStorage initialized for namespace '{self.namespace}' at '{self._db_path}'"
+        )
 
     @property
     def data(self):
@@ -29,7 +30,7 @@ def all_keys(self) -> List[str]:
 
     def index_done_callback(self):
         self._db.flush()
-        logger.info("RocksDB flushed for %s", self.namespace)
+        print(f"RocksDB flushed for {self.namespace}")
 
     def get_by_id(self, id: str) -> Any:
         return self._db.get(id, None)
@@ -63,7 +64,6 @@ def upsert(self, data: Dict[str, Any]):
         if left_data:
             for k, v in left_data.items():
                 self._db[k] = v
-
             # if left_data is very large, it is recommended to use self._db.write_batch() for optimization
 
         return left_data
@@ -72,8 +72,14 @@ def drop(self):
         self._db.close()
         Rdict.destroy(self._db_path)
         self._db = Rdict(self._db_path)
-        logger.info("Dropped RocksDB %s", self.namespace)
+        print(f"Dropped RocksDB {self.namespace}")
 
     def close(self):
         if self._db:
             self._db.close()
+
+    def reload(self):
+        if self._db:
+            self._db.close()
+        self._db = Rdict(self._db_path)
+        print(f"Reloaded RocksDB {self.namespace}")

graphgen/operators/chunk/chunk_service.py

@@ -47,7 +47,7 @@ def __init__(self, working_dir: str = "cache", **chunk_kwargs):
         tokenizer_model = os.getenv("TOKENIZER_MODEL", "cl100k_base")
         self.tokenizer_instance: Tokenizer = Tokenizer(model_name=tokenizer_model)
         self.chunk_storage = init_storage(
-            backend="json_kv",
+            backend="rocksdb",
             working_dir=working_dir,
             namespace="chunk",
         )

graphgen/operators/partition/partition_service.py

@@ -26,7 +26,7 @@ def __init__(self, working_dir: str = "cache", **partition_kwargs):
             namespace="graph",
         )
         self.chunk_storage: BaseKVStorage = init_storage(
-            backend="json_kv",
+            backend="rocksdb",
             working_dir=working_dir,
             namespace="chunk",
         )

graphgen/operators/quiz/quiz_service.py

@@ -23,7 +23,7 @@ def __init__(
         )
         # { _quiz_id: { "description": str, "quizzes": List[Tuple[str, str]] } }
         self.quiz_storage: BaseKVStorage = init_storage(
-            backend="json_kv", working_dir=working_dir, namespace="quiz"
+            backend="rocksdb", working_dir=working_dir, namespace="quiz"
         )
         self.generator = QuizGenerator(self.llm_client)
         self.concurrency_limit = concurrency_limit