Merge pull request #179 from djwu563/release/v0.2.2

Add support for setting a domain for workers

Author: zhang_lu

omagent-core/src/omagent_core/clients/devices/cli/client.py

@@ -42,18 +42,19 @@ def __init__(
         self._config_path = config_path
         self._workers = workers
         self._input_prompt = input_prompt
+        self._task_to_domain = {}
 
     def start_interactor(self):
         workflow_instance_id = None
         try:
             absolute_path = Path(self._config_path).resolve()
             worker_config = build_from_file(self._config_path)
             self._task_handler_interactor = TaskHandler(
-                worker_config=worker_config, workers=self._workers
+                worker_config=worker_config, workers=self._workers, task_to_domain=self._task_to_domain
             )
             self._task_handler_interactor.start_processes()
             workflow_instance_id = self._interactor.start_workflow_with_input(
-                workflow_input={}
+                workflow_input={}, task_to_domain=self._task_to_domain
             )
 
             stream_name = f"{workflow_instance_id}_output"
@@ -180,11 +181,11 @@ def start_processor(self):
         try:
             worker_config = build_from_file(self._config_path)
             self._task_handler_processor = TaskHandler(
-                worker_config=worker_config, workers=self._workers
+                worker_config=worker_config, workers=self._workers, task_to_domain=self._task_to_domain
             )
             self._task_handler_processor.start_processes()
             workflow_instance_id = self._processor.start_workflow_with_input(
-                workflow_input={}
+                workflow_input={}, task_to_domain=self._task_to_domain
             )
             user_input = input(
                 f"{Fore.GREEN}Please input a folder path of images:(WaitPress Enter to finish the entire input.):\n>>>{Style.RESET_ALL}"

omagent-core/src/omagent_core/clients/devices/programmatic/client.py

@@ -22,21 +22,22 @@ def __init__(
         self._config_path = config_path
         self._workers = workers
         self._task_handler_processor = None
+        self._task_to_domain = {}
 
     def start_processor(self):
         worker_config = build_from_file(self._config_path)
         self._task_handler_processor = TaskHandler(
             worker_config=worker_config, workers=self._workers
         )
         self._task_handler_processor.start_processes()
-        self._processor.start_workflow_with_input(workflow_input={})
+        self._processor.start_workflow_with_input(workflow_input={}, task_to_domain=self._task_to_domain)
 
     def start_processor_with_input(self, workflow_input: dict):
         try:
             if self._task_handler_processor is None:
                 worker_config = build_from_file(self._config_path)
                 self._task_handler_processor = TaskHandler(
-                    worker_config=worker_config, workers=self._workers
+                    worker_config=worker_config, workers=self._workers, task_to_domain=self._task_to_domain
                 )
                 self._task_handler_processor.start_processes()
             return self._process_workflow(self._processor, workflow_input)
@@ -46,8 +47,9 @@ def start_processor_with_input(self, workflow_input: dict):
     def start_batch_processor(self, workflow_input_list: list[dict], max_tasks: int = 10):
         results = [None] * len(workflow_input_list)
         worker_config = build_from_file(self._config_path)
-        self._task_handler_processor = TaskHandler(worker_config=worker_config, workers=self._workers)
-        self._task_handler_processor.start_processes()
+        if self._task_handler_processor is None:
+            self._task_handler_processor = TaskHandler(worker_config=worker_config, workers=self._workers, task_to_domain=self._task_to_domain)
+            self._task_handler_processor.start_processes()
 
         result_queue = multiprocessing.Queue()
         active_processes = []
@@ -87,7 +89,7 @@ def _process_workflow(self, workflow: ConductorWorkflow, workflow_input: dict):
         workflow_instance_id = None
         try:
             workflow_instance_id = workflow.start_workflow_with_input(
-                workflow_input=workflow_input
+                workflow_input=workflow_input, task_to_domain=self._task_to_domain
             )
             while True:
                 status = workflow.get_workflow(workflow_id=workflow_instance_id).status

omagent-core/src/omagent_core/clients/devices/webpage/client.py

@@ -36,6 +36,8 @@ def __init__(
         self._config_path = config_path
         self._workers = workers
         self._workflow_instance_id = None
+        self._worker_config = build_from_file(self._config_path)
+        self._task_to_domain = {}
         self._incomplete_message = ""
         self._custom_css = """
             #OmAgent {
@@ -70,9 +72,8 @@ def __init__(
         """
 
     def start_interactor(self):
-        worker_config = build_from_file(self._config_path)
         self._task_handler_interactor = TaskHandler(
-            worker_config=worker_config, workers=self._workers
+            worker_config=self._worker_config, workers=self._workers, task_to_domain=self._task_to_domain
         )
         self._task_handler_interactor.start_processes()
         try:
@@ -113,9 +114,8 @@ def stop_interactor(self):
         self._task_handler_interactor.stop_processes()
 
     def start_processor(self):
-        worker_config = build_from_file(self._config_path)
         self._task_handler_processor = TaskHandler(
-            worker_config=worker_config, workers=self._workers
+            worker_config=self._worker_config, workers=self._workers, task_to_domain=self._task_to_domain
         )
         self._task_handler_processor.start_processes()
 
@@ -161,7 +161,7 @@ def stop_processor(self):
     def add_message(self, history, message):
         if self._workflow_instance_id is None:
             self._workflow_instance_id = self._interactor.start_workflow_with_input(
-                workflow_input={}
+                workflow_input={}, task_to_domain=self._task_to_domain
             )
         contents = []
         for x in message["files"]:
@@ -184,7 +184,7 @@ def add_message(self, history, message):
     def add_processor_message(self, history, message):
         if self._workflow_instance_id is None:
             self._workflow_instance_id = self._processor.start_workflow_with_input(
-                workflow_input={}
+                workflow_input={}, task_to_domain=self._task_to_domain
             )
         image_items = []
         for idx, x in enumerate(message["files"]):

omagent-core/src/omagent_core/engine/automator/task_handler.py

@@ -54,6 +54,7 @@ def __init__(
         workers: List[BaseWorker] = [],
         metrics_settings: MetricsSettings = None,
         import_modules: List[str] = None,
+        task_to_domain: dict[str, str] = None,
     ):
         """Initialize a new TaskHandler instance.
 
@@ -95,7 +96,8 @@ def __init__(
             )
             workers.extend([worker_cls(**config) for _ in range(concurrency)])
         for worker in workers:
-            worker.task_definition_name = worker.id or worker.name
+            if task_to_domain is not None:
+                task_to_domain[worker.task_definition_name] = worker.domain
         self.__create_task_runner_processes(
             workers, container.conductor_config, metrics_settings
         )

omagent-core/src/omagent_core/engine/automator/task_runner.py

@@ -96,6 +96,8 @@ def __poll_task(self) -> Task:
                 else:
                     task = None
             else:
+                if domain is not None:
+                    params["domain"] = domain
                 task = self.task_client.poll(tasktype=task_definition_name, **params)
 
             finish_time = time.time()

omagent-core/src/omagent_core/engine/worker/base.py

@@ -57,7 +57,7 @@ class BaseWorker(BotBase, ABC):
     _task_type: Optional[str] = None
 
     def model_post_init(self, __context: Any) -> None:
-        self.task_definition_name = self.name
+        self.task_definition_name = self.id or self.name
         self.next_task_index = 0
         self._task_definition_name_cache = None