fix: register storage actors in Driver

Author: ChenZiHong-Gavin

graphgen/engine.py

@@ -12,7 +12,7 @@
 
 from graphgen.bases import Config, Node
 from graphgen.utils import logger
-from graphgen.common import init_llm
+from graphgen.common import init_llm, init_storage
 
 load_dotenv()
 
@@ -25,6 +25,7 @@ def __init__(
         self.functions = functions
         self.datasets: Dict[str, ray.data.Dataset] = {}
         self.llm_actors = {}
+        self.storage_actors = {}
 
         ctx = DataContext.get_current()
         ctx.enable_rich_progress_bars = False
@@ -54,11 +55,57 @@ def __init__(
             logger.info("Ray Dashboard URL: %s", context.dashboard_url)
 
         self._init_llms()
+        self._init_storage()
 
     def _init_llms(self):
         self.llm_actors["synthesizer"] = init_llm("synthesizer")
         self.llm_actors["trainee"] = init_llm("trainee")
 
+    def _init_storage(self):
+        kv_namespaces, graph_namespaces = self._scan_storage_requirements()
+        working_dir = self.global_params["working_dir"]
+
+        for node_id in kv_namespaces:
+            proxy = init_storage(self.global_params["kv_backend"], working_dir, node_id)
+            self.storage_actors[f"kv_{node_id}"] = proxy
+            logger.info("Create KV Storage Actor: namespace=%s", node_id)
+    
+        for ns in graph_namespaces:
+            proxy = init_storage(self.global_params["graph_backend"], working_dir, ns)
+            self.storage_actors[f"graph_{ns}"] = proxy
+            logger.info("Create Graph Storage Actor: namespace=%s", ns)
+
+    def _scan_storage_requirements(self) -> tuple[set[str], set[str]]:
+        kv_namespaces = set()
+        graph_namespaces = set()
+
+        # TODO: Temporarily hard-coded; node storage will be centrally managed later.
+        for node in self.config.nodes:
+            op_name = node.op_name
+            if self._function_needs_param(op_name, "kv_backend"):
+                kv_namespaces.add(op_name)
+            if self._function_needs_param(op_name, "graph_backend"):
+                graph_namespaces.add("graph")
+        return kv_namespaces, graph_namespaces
+    
+    def _function_needs_param(self, op_name: str, param_name: str) -> bool:
+        if op_name not in self.functions:
+            return False
+        
+        func = self.functions[op_name]
+
+        if inspect.isclass(func):
+            try:
+                sig = inspect.signature(func.__init__)
+                return param_name in sig.parameters
+            except (ValueError, TypeError):
+                return False
+
+        try:
+            sig = inspect.signature(func)
+            return param_name in sig.parameters
+        except (ValueError, TypeError):
+            return False
 
     @staticmethod
     def _topo_sort(nodes: List[Node]) -> List[Node]: