[dev] modify workflow kwargs & worker init params

This commit is contained in:
jinli.yl 2024-07-09 16:31:18 +08:00
parent 523b752a3d
commit a8fcfeeaa2
11 changed files with 29 additions and 67 deletions

View file

@ -69,12 +69,13 @@ class CliJob(object):
def run(self, config: str):
self.load_config(config)
with G_CONTEXT.thread_pool:
memory_chat = list(G_CONTEXT.memory_chat_dict.values())[0]
memory_chat.run()
# with G_CONTEXT.thread_pool:
memory_chat = list(G_CONTEXT.memory_chat_dict.values())[0]
memory_chat.run()
G_CONTEXT.memory_store.close()
G_CONTEXT.monitor.close()
G_CONTEXT.thread_pool.shutdown()
if __name__ == "__main__":

View file

@ -12,9 +12,15 @@ class BaseOperation(metaclass=ABCMeta):
self.description: str = description
self.kwargs: dict = kwargs
def init_workflow(self):
def init_workflow(self, **kwargs):
pass
@abstractmethod
def run_operation(self, **kwargs):
raise NotImplementedError
def run_operation_backend(self):
pass
def stop_operation_backend(self):
pass

View file

@ -82,7 +82,7 @@ class BaseWorkflow(object):
continue
self.logger.info(f"----- print_workflow_{self.name}_end -----")
def init_workers(self):
def init_workers(self, is_backend: bool = False, **kwargs):
for name in list(self.worker_dict.keys()):
if name not in G_CONTEXT.worker_config:
raise RuntimeError(f"worker={name} is not exists in worker_config!")
@ -91,9 +91,10 @@ class BaseWorkflow(object):
config=G_CONTEXT.worker_config[name],
suffix_name="worker",
name=name,
is_multi_thread=self.worker_dict[name],
is_multi_thread=is_backend or self.worker_dict[name],
context=self.context,
context_lock=self.context_lock)
context_lock=self.context_lock,
**kwargs)
def _run_sub_workflow(self, worker_list: List[str]) -> bool:
for name in worker_list:
@ -113,8 +114,7 @@ class BaseWorkflow(object):
else:
t_list = []
for sub_workflow in workflow_part:
t_list.append(G_CONTEXT.thread_pool.submit(
self._run_sub_workflow, sub_workflow))
t_list.append(G_CONTEXT.thread_pool.submit(self._run_sub_workflow, sub_workflow))
flag = True
for future in as_completed(t_list):

View file

@ -21,8 +21,8 @@ class ReadMemory(BaseWorkflow, BaseOperation):
self.chat_messages: List[Message] = chat_messages
self.his_msg_count: int = his_msg_count
def init_workflow(self):
self.init_workers()
def init_workflow(self, **kwargs):
self.init_workers(**kwargs)
def run_operation(self, **kwargs):
max_count = 1 + self.his_msg_count

View file

@ -11,8 +11,8 @@ class SummaryMemory(BaseWorkflow, BaseBackendOperation):
super().__init__(**kwargs)
BaseBackendOperation.__init__(self, **kwargs)
def init_workflow(self):
self.init_workers()
def init_workflow(self, **kwargs):
self.init_workers(is_backend=True, **kwargs)
def _run_operation(self, **kwargs):
self.context[CHAT_KWARGS] = kwargs

View file

@ -35,8 +35,8 @@ class WriteMemory(BaseWorkflow, BaseBackendOperation):
for msg in self.chat_messages:
msg.memorized = True
def init_workflow(self):
self.init_workers()
def init_workflow(self, **kwargs):
self.init_workers(is_backend=True, **kwargs)
def _run_operation(self, **kwargs):
self.context[CHAT_KWARGS] = kwargs

View file

@ -33,7 +33,7 @@ class BaseMemoryService(metaclass=ABCMeta):
def add_messages(self, messages: List[Message] | Message):
raise NotImplementedError
def start_service(self):
def start_service(self, **kwargs):
pass
@abstractmethod

View file

@ -37,9 +37,9 @@ class ChatMemoryService(BaseMemoryService):
for _ in range(gap_size):
self.chat_messages.pop(0)
def start_service(self):
def start_service(self, **kwargs):
for _, operation in self._operation_dict.items():
operation.init_workflow()
operation.init_workflow(**kwargs)
if operation.operation_type == "backend":
operation.run_operation_backend()

View file

@ -32,14 +32,11 @@ class BaseWorker(metaclass=ABCMeta):
raise RuntimeError(f"async_task is not allowed in multi_thread condition")
self.task_list.append((fn, args, kwargs))
async def _async_gather(self):
return await asyncio.gather(*[fn(*args, **kwargs) for fn, args, kwargs in self.task_list])
def gather_async_result(self):
if self.is_multi_thread:
raise RuntimeError(f"async_task is not allowed in multi_thread condition")
async def async_gather():
return await asyncio.gather(*[fn(*args, **kwargs) for fn, args, kwargs in self.task_list])
results = asyncio.run(async_gather())
results = asyncio.run(self._async_gather())
self.task_list.clear()
return results

11
test.py
View file

@ -1,11 +0,0 @@
from memory_scope.cli import CliJob
def main(config_path: str):
job = CliJob()
job.run(config=config_path)
if __name__ == "__main__":
# fire.Fire(main)
main("config/config.yaml")

31
tt.py
View file

@ -1,31 +0,0 @@
import asyncio
class TT(object):
def __init__(self):
self.task_list = []
async def async_func(self, i: int):
await asyncio.sleep(i) # 模拟异步操作
print(f"函数{i}的结果")
def submit_async_task(self, fn, *args, **kwargs):
self.task_list.append((fn, args, kwargs))
def gather_async_result(self):
async def async_gather():
return await asyncio.gather(*[fn(*args, **kwargs) for fn, args, kwargs in self.task_list])
results = asyncio.run(async_gather())
self.task_list.clear()
return results
def run(self):
self.submit_async_task(self.async_func, i=1)
self.submit_async_task(self.async_func, i=2)
self.submit_async_task(self.async_func, i=3)
self.gather_async_result()
TT().run()