Merge pull request #163 from zhuzhongshu123/0.6_dev

feat(bridge): disable ckpt for spg server runner
OpenSPG · Dec 25, 2024 · 69ce1f4 · 69ce1f4
2 parents 2ecb439 + 6055c46
commit 69ce1f4
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/kag/bridge/spg_server_bridge.py b/kag/bridge/spg_server_bridge.py
@@ -26,7 +26,7 @@ def run_reader(self, config, input_data):
         reader = interface.ReaderABC.from_config(reader_config)
         chunks = []
         for data in scanner.generate(input_data):
-            chunks += reader.invoke(data)
+            chunks += reader.invoke(data, write_ckpt=False)
         return [x.to_dict() for x in chunks]
 
     def run_component(self, component_name, component_config, input_data):
@@ -35,4 +35,6 @@ def run_component(self, component_name, component_config, input_data):
 
         cls = getattr(interface, component_name)
         instance = cls.from_config(component_config)
-        return [x.to_dict() for x in instance.invoke(input_data)]
+        if hasattr(instance.input_types, "from_dict"):
+            input_data = instance.input_types.from_dict(input_data)
+        return [x.to_dict() for x in instance.invoke(input_data, write_ckpt=False)]
diff --git a/kag/interface/builder/base.py b/kag/interface/builder/base.py
@@ -97,7 +97,8 @@ def _invoke(self, input: Input, **kwargs) -> List[Output]:
         )
 
     def invoke(self, input: Input, **kwargs) -> List[Output]:
-        if self.checkpointer:
+        write_ckpt = kwargs.get("write_ckpt", True)
+        if write_ckpt and self.checkpointer:
             input_key = kwargs.get("key")
             # found existing data in checkpointer
             if input_key and self.checkpointer.exists(input_key):