add not parallel method

This commit is contained in:
jinli.yl 2024-08-06 15:05:44 +08:00
commit 924b67fca2
3 changed files with 21 additions and 7 deletions

View file

@ -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

View file

@ -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():

View file

@ -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):