From 4a6171e6a0fcc4970e8ba76a06e81f242a5938c1 Mon Sep 17 00:00:00 2001
From: "jinli.yl" <jinli.yl@alibaba-inc.com>
Date: Fri, 2 Aug 2024 19:48:36 +0800
Subject: [PATCH] add not parallel method

---
 memoryscope/core/config/demo_config.yaml        |  1 +
 .../worker/backend/update_insight_worker.py     | 10 +++++++---
 memoryscope/core/worker/base_worker.py          | 17 +++++++++++++----
 3 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/memoryscope/core/config/demo_config.yaml b/memoryscope/core/config/demo_config.yaml
index 58e792d4..3d597958 100644
--- a/memoryscope/core/config/demo_config.yaml
+++ b/memoryscope/core/config/demo_config.yaml
@@ -148,6 +148,7 @@ worker:
     rank_model: rank_model
     embedding_model: embedding_model
     update_insight_threshold: 0.01
+    enable_parallel: false
   long_contra_repeat:
     class: core.worker.backend.long_contra_repeat_worker
     generation_model: generation_model
diff --git a/memoryscope/core/worker/backend/update_insight_worker.py b/memoryscope/core/worker/backend/update_insight_worker.py
index 5d2b69ea..7bee8aa0 100644
--- a/memoryscope/core/worker/backend/update_insight_worker.py
+++ b/memoryscope/core/worker/backend/update_insight_worker.py
@@ -212,7 +212,8 @@ class UpdateInsightWorker(MemoryBaseWorker):
 
         # Process active insight nodes with corresponding not updated nodes
         for node in insight_nodes:
-            time.sleep(1)
+            if self.enable_parallel:
+                time.sleep(1)
             if node.action_status == ActionStatusEnum.NEW.value:
                 self.submit_thread_task(fn=self.filter_obs_nodes,
                                         insight_node=node,
@@ -233,8 +234,11 @@ class UpdateInsightWorker(MemoryBaseWorker):
 
         # Submit tasks to update insights for the top nodes
         for insight_node, filtered_nodes, _ in result_sorted:
-            time.sleep(1)
-            self.submit_thread_task(fn=self.update_insight, insight_node=insight_node, filtered_nodes=filtered_nodes)
+            if self.enable_parallel:
+                time.sleep(1)
+            self.submit_thread_task(fn=self.update_insight,
+                                    insight_node=insight_node,
+                                    filtered_nodes=filtered_nodes)
 
         # Gather the final results from all update tasks
         for _ in self.gather_thread_result():
diff --git a/memoryscope/core/worker/base_worker.py b/memoryscope/core/worker/base_worker.py
index 40056a35..6a07c26e 100644
--- a/memoryscope/core/worker/base_worker.py
+++ b/memoryscope/core/worker/base_worker.py
@@ -43,6 +43,7 @@ class BaseWorker(metaclass=ABCMeta):
         self.raise_exception: bool = raise_exception
         self.is_multi_thread: bool = is_multi_thread
         self.thread_pool: ThreadPoolExecutor = thread_pool
+        self.enable_parallel: bool = True
         self.kwargs: dict = kwargs
 
         self.continue_run: bool = True
@@ -113,7 +114,10 @@ class BaseWorker(metaclass=ABCMeta):
             args: Positional arguments for the function.
             kwargs: Keyword arguments for the function.
         """
-        self.thread_task_list.append(self.thread_pool.submit(fn, *args, **kwargs))
+        if self.enable_parallel:
+            self.thread_task_list.append(self.thread_pool.submit(fn, *args, **kwargs))
+        else:
+            self.thread_task_list.append(fn(*args, **kwargs))
 
     def gather_thread_result(self):
         """
@@ -122,9 +126,14 @@ class BaseWorker(metaclass=ABCMeta):
         Yields:
             The result of each completed task.
         """
-        for future in as_completed(self.thread_task_list):
-            yield future.result()
-        self.thread_task_list.clear()
+        if self.enable_parallel:
+            for future in as_completed(self.thread_task_list):
+                yield future.result()
+            self.thread_task_list.clear()
+        else:
+            for future in self.thread_task_list:
+                yield future
+            self.thread_task_list.clear()
 
     @abstractmethod
     def _run(self):