Conversation
schedule strategy
add executor
add utility function for task exec
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a foundational task scheduling framework, enabling more sophisticated management and execution of tasks within the system. It centralizes task state handling, provides flexible scheduling strategies, and refines the event-driven architecture with more precise lifecycle hooks. These changes lay the groundwork for advanced automation and workflow orchestration, ensuring tasks are processed efficiently and reliably according to defined schedules and dependencies. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request primarily refactors task status handling by renaming TaskStatusValue to TaskStatus and making it inherit from str for consistency, along with updating all references across multiple files including aworld/core/common.py, aworld/core/context/amni/contexts.py, aworld/core/context/base.py, aworld/core/event/message_future.py, aworld/core/task.py, aworld/events/util.py, aworld/runners/event_runner.py, aworld/runners/handler/background_task.py, aworld/runners/handler/task.py, aworld/runners/task_manager.py, and aworld/schedule/scheduler.py. It also introduces a new TaskTypeValue class and adds new hook points and corresponding classes (OnStartLLMCallHook, OnFinishedLLMCallHook, OnLLMCallHook, OnSuccessLLMCallHook, OnErrorLLMCallHook, OnStartToolCallHook, OnFinishedToolCallHook, OnToolCallHook, OnSuccessToolCallHook, OnErrorToolCallHook, OnRunHook, OnSuccessHook, OnErrorHook, OnStartHook, OnFinishHook) to aworld/runners/hook/hooks.py, aworld/runners/hook/agent_hooks.py, aworld/runners/hook/task_hooks.py, and aworld/runners/hook/tool_hooks.py, while removing older Pre/Post hook types. Additionally, it removes model_config arguments from @scorer_register decorators in aworld/evaluations/scorers/output_validators.py and removes several unused imports. Review comments highlight a type hint mismatch for the executor's return value in TaskScheduler, an AttributeError risk in execute_schedulable_tasks due to incorrect return type, a bug in create_strategy where AutoStrategy is instantiated prematurely, an issue with PriorityTaskQueue using a min-heap instead of a max-heap for priority, a potential security vulnerability in _merge_by_topic due to reliance on untrusted agent_id from message payload, and a need to update the TaskTypeValue class to inherit from str for consistency, as well as updating a docstring example in TaskManager to reflect the TaskStatus renaming.
aworld/schedule/scheduler.py
Outdated
| return self._executor | ||
|
|
||
| @executor.setter | ||
| def executor(self, executor: Callable[[Task], Awaitable[bool]]): |
There was a problem hiding this comment.
The type hint for the executor's return value is Awaitable[bool], but the execute method expects it to return a TaskResponse object (line 161: result: TaskResponse = await self._executor(task)). The type hint should be updated to Awaitable[TaskResponse].
Additionally, the docstring for the setter should be updated to reflect this change.
- Current docstring:
executor: Async function that takes a Task and returns bool (success) - Suggested docstring:
executor: Async function that takes a Task and returns a TaskResponse object
| def executor(self, executor: Callable[[Task], Awaitable[bool]]): | |
| def executor(self, executor: Callable[[Task], Awaitable[TaskResponse]]): |
aworld/schedule/scheduler.py
Outdated
| async def task_execute(st=scheduled_task): | ||
| st.task_status = TaskStatus.RUNNING | ||
| st.started_at = time.time() | ||
| res = await exec_tasks(tasks=[st]) | ||
| st.task_status = res.get(st.id).status | ||
| st.completed_at = time.time() | ||
| return res |
There was a problem hiding this comment.
The task_execute function returns res, which is a dictionary of {'task_id': TaskResponse}. However, the LocalRuntime's execute method, which consumes this, expects an object with an .id attribute to build its result dictionary. This will cause an AttributeError. The task_execute function should return the TaskResponse object directly.
| async def task_execute(st=scheduled_task): | |
| st.task_status = TaskStatus.RUNNING | |
| st.started_at = time.time() | |
| res = await exec_tasks(tasks=[st]) | |
| st.task_status = res.get(st.id).status | |
| st.completed_at = time.time() | |
| return res | |
| async def task_execute(st=scheduled_task): | |
| st.task_status = TaskStatus.RUNNING | |
| st.started_at = time.time() | |
| res_dict = await exec_tasks(tasks=[st]) | |
| task_response = res_dict.get(st.id) | |
| if task_response: | |
| st.task_status = task_response.status | |
| st.completed_at = time.time() | |
| return task_response |
aworld/schedule/strategy.py
Outdated
|
|
||
| def create_strategy(strategy_type: Optional[str] = None, **kwargs) -> ScheduleStrategy: | ||
| """Create strategy instance.""" | ||
| return STRATEGY_MAP.get(strategy_type, AutoStrategy())(**kwargs) |
There was a problem hiding this comment.
There is a bug in the create_strategy factory. AutoStrategy() creates an instance of the class, and then you are trying to call that instance. The intention is to get the class from the map and then instantiate it with **kwargs. You should retrieve the class AutoStrategy itself as the default value, not an instance of it.
| return STRATEGY_MAP.get(strategy_type, AutoStrategy())(**kwargs) | |
| return STRATEGY_MAP.get(strategy_type, AutoStrategy)(**kwargs) |
aworld/schedule/strategy.py
Outdated
| def push(self, task: ScheduledTask): | ||
| # (priority_value, counter, task) | ||
| priority_value = task.priority | ||
| heapq.heappush(self.heap, (priority_value, self.counter, task)) |
There was a problem hiding this comment.
The PriorityTaskQueue uses heapq, which is a min-heap. Pushing priority_value directly will result in tasks with lower priority values being processed first. This contradicts the common expectation and the implementation in other parts of the code (like TaskManager.get_ready_tasks) where higher priority values are treated as more important. To implement a max-heap for high-priority-first scheduling, you should push the negative of the priority value.
| heapq.heappush(self.heap, (priority_value, self.counter, task)) | |
| heapq.heappush(self.heap, (-priority_value, self.counter, task)) |
| if not agent_id: | ||
| agent_id = data.get('env_content', {}).get('agent_id') |
There was a problem hiding this comment.
In _merge_by_topic, the agent_id is extracted from the message payload (data.get('env_content', {}).get('agent_id')) if it's not already present in the message. Since the payload is the result of a background task, which may be untrusted or compromised, an attacker could manipulate this value to cause the runner to send messages to arbitrary agents. This could lead to unauthorized actions or lateral movement within the agent system. Furthermore, when forwarding messages to parent tasks (lines 105-112), the agent_id is omitted, forcing the receiver to rely on the untrusted payload.
| class TaskTypeValue: | ||
| """Task type constants.""" | ||
| INSTANT = 'instant' | ||
| SCHEDULED = 'scheduled' |
There was a problem hiding this comment.
For consistency with the TaskStatus class, TaskTypeValue should also inherit from str. This makes it clearer that it's a group of string constants and allows for potential use in type hints if needed in the future.
| class TaskTypeValue: | |
| """Task type constants.""" | |
| INSTANT = 'instant' | |
| SCHEDULED = 'scheduled' | |
| class TaskTypeValue(str): | |
| """Task type constants.""" | |
| INSTANT = 'instant' | |
| SCHEDULED = 'scheduled' |
aworld/runners/task_manager.py
Outdated
|
|
||
| # List tasks | ||
| all_tasks = await manager.list() | ||
| pending_tasks = await manager.list(status=TaskStatusValue.INIT) |
There was a problem hiding this comment.
The example in the docstring uses TaskStatusValue, which has been removed in this pull request. It should be updated to use TaskStatus.
| pending_tasks = await manager.list(status=TaskStatusValue.INIT) | |
| pending_tasks = await manager.list(status=TaskStatus.INIT) |
3 participants
No description provided.