Merge pull request #135 from zhuzhongshu123/0.6_dev

fix(common): fix zodb writer bug

kag/builder/default_chain.py

@@ -186,7 +186,6 @@ def get_component_with_ckpts(self):
             self.extractor,
             self.vectorizer,
             self.post_processor,
-            self.writer,
         ]
 
     def close_checkpointers(self):

kag/builder/model/sub_graph.py

@@ -10,6 +10,7 @@
 # is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 # or implied.
 import pprint
+import copy
 from typing import Dict, List, Any
 
 from kag.builder.model.spg_record import SPGRecord
@@ -48,7 +49,7 @@ def to_dict(self):
             "id": self.id,
             "name": self.name,
             "label": self.label,
-            "properties": self.properties,
+            "properties": copy.deepcopy(self.properties),
         }
 
     @classmethod
@@ -119,7 +120,7 @@ def to_dict(self):
             "fromType": self.from_type,
             "toType": self.to_type,
             "label": self.label,
-            "properties": self.properties,
+            "properties": copy.deepcopy(self.properties),
         }
 
     @classmethod

kag/common/checkpointer/bin_checkpointer.py

@@ -13,7 +13,7 @@
 import logging
 import transaction
 import threading
-
+import pickle
 import BTrees.OOBTree
 from ZODB import DB
 from ZODB.FileStorage import FileStorage
@@ -135,7 +135,10 @@ def read_from_ckpt(self, key):
         with self._lock:
             with self._ckpt.transaction() as conn:
                 obj = conn.root.data.get(key, None)
-            return obj
+            if obj:
+                return pickle.loads(obj)
+            else:
+                return None
 
     def write_to_ckpt(self, key, value):
         """
@@ -148,7 +151,7 @@ def write_to_ckpt(self, key, value):
         with self._lock:
             try:
                 with self._ckpt.transaction() as conn:
-                    conn.root.data[key] = value
+                    conn.root.data[key] = pickle.dumps(value)
             except Exception as e:
                 logger.warn(f"failed to write checkpoint {key} to db, info: {e}")