feat(common): add processed chunk id checkpoint

kag/builder/default_chain.py

@@ -9,7 +9,6 @@
 # Unless required by applicable law or agreed to in writing, software distributed under the License
 # is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 # or implied.
-
 import logging
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from kag.interface import (
@@ -22,7 +21,6 @@
     SinkWriterABC,
     KAGBuilderChain,
 )
-
 from kag.common.utils import generate_hash_id
 
 logger = logging.getLogger(__name__)
@@ -155,16 +153,28 @@ def run_extract(chunk):
                 if node is None:
                     continue
                 flow_data = execute_node(node, flow_data, key=input_key)
-            return flow_data
+            return {input_key: flow_data[0]}
 
         reader_output = self.reader.invoke(input_data, key=generate_hash_id(input_data))
         splitter_output = []
+
         for chunk in reader_output:
             splitter_output.extend(self.splitter.invoke(chunk, key=chunk.hash_key))
 
+        processed_chunk_keys = kwargs.get("processed_chunk_keys", set())
+        filtered_chunks = []
+        processed = 0
+        for chunk in splitter_output:
+            if chunk.hash_key not in processed_chunk_keys:
+                filtered_chunks.append(chunk)
+            else:
+                processed += 1
+        logger.debug(
+            f"Total chunks: {len(splitter_output)}. Checkpointed: {processed}, Pending: {len(filtered_chunks)}."
+        )
         result = []
         with ThreadPoolExecutor(max_workers) as executor:
-            futures = [executor.submit(run_extract, chunk) for chunk in splitter_output]
+            futures = [executor.submit(run_extract, chunk) for chunk in filtered_chunks]
 
             from tqdm import tqdm
 
@@ -176,5 +186,5 @@ def run_extract(chunk):
                 leave=False,
             ):
                 ret = inner_future.result()
-                result.extend(ret)
+                result.append(ret)
         return result

kag/builder/runner.py

@@ -108,6 +108,14 @@ def __init__(
                 "world_size": self.scanner.sharding_info.get_world_size(),
             }
         )
+        self.processed_chunks = CheckpointerManager.get_checkpointer(
+            {
+                "type": "zodb",
+                "ckpt_dir": os.path.join(self.ckpt_dir, "chain"),
+                "rank": self.scanner.sharding_info.get_rank(),
+                "world_size": self.scanner.sharding_info.get_world_size(),
+            }
+        )
         self._local = threading.local()
 
     def invoke(self, input):
@@ -135,7 +143,11 @@ def invoke(self, input):
 
         def process(data, data_id, data_abstract):
             try:
-                result = self.chain.invoke(data, max_workers=self.num_threads_per_chain)
+                result = self.chain.invoke(
+                    data,
+                    max_workers=self.num_threads_per_chain,
+                    processed_chunk_keys=self.processed_chunks.keys(),
+                )
                 return data, data_id, data_abstract, result
             except Exception:
                 traceback.print_exc()
@@ -177,6 +189,15 @@ def process(data, data_id, data_abstract):
                                 num_nodes += len(item.nodes)
                                 num_edges += len(item.edges)
                                 num_subgraphs += 1
+                            elif isinstance(item, dict):
+
+                                for k, v in item.items():
+                                    self.processed_chunks.write_to_ckpt(k, k)
+                                    if isinstance(v, SubGraph):
+                                        num_nodes += len(v.nodes)
+                                        num_edges += len(v.edges)
+                                        num_subgraphs += 1
+
                         info = {
                             "num_nodes": num_nodes,
                             "num_edges": num_edges,

kag/common/checkpointer/base.py

@@ -111,6 +111,16 @@ def exists(self, key):
         """
         raise NotImplementedError("close not implemented yet.")
 
+    def keys(self):
+        """
+        Returns the key set contained in the checkpoint file.
+
+        Returns:
+            set:  The key set contained in the checkpoint.
+        """
+
+        raise NotImplementedError("keys not implemented yet.")
+
     def size(self):
         """
         Return the number of records in the checkpoint file.

kag/common/checkpointer/bin_checkpointer.py

@@ -92,6 +92,9 @@ def size(self):
 
         return len(self._ckpt)
 
+    def keys(self):
+        return set(self._ckpt.keys())
+
 
 @CheckPointer.register("zodb")
 class ZODBCheckPointer(CheckPointer):
@@ -207,3 +210,8 @@ def size(self):
         with self._lock:
             with self._ckpt.transaction() as conn:
                 return len(conn.root.data)
+
+    def keys(self):
+        with self._lock:
+            with self._ckpt.transaction() as conn:
+                return set(conn.root.data.keys())

kag/common/checkpointer/txt_checkpointer.py

@@ -87,3 +87,6 @@ def _close(self):
 
     def size(self):
         return len(self._ckpt)
+
+    def keys(self):
+        return set(self._ckpt.keys())