Fix priority queue implementation and enhance testing #135
Conversation
There was a problem hiding this comment.
Summary of Changes
Hello @pan-x-c, 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 overhauls the internal queue management system, particularly focusing on the priority queue implementation and its integration with Ray actors. The changes introduce a more modular and robust queue architecture, improve how queue capacity is handled for priority queues, and refine the experience data model's serialization. These foundational improvements are supported by enhanced testing to ensure correctness and reliability.
Highlights
- Queue Architecture Refactor: The core queue implementation has been significantly refactored. An abstract
QueueBufferinterface was introduced, with concrete implementationsAsyncQueue(for standard FIFO queues) andAsyncPriorityQueue(for priority-based queues). TheAsyncPriorityQueueclass was moved from its dedicated file (trinity/buffer/priority_queue.py, which is now removed) intotrinity/buffer/queue.pyfor better consolidation. - Ray Actor Wrapper Redesign: The
QueueActorclass, previously responsible for wrapping queue logic as a Ray actor, has been replaced by a newQueueWrapperclass intrinity/buffer/ray_wrapper.py. This new wrapper now manages the underlyingQueueBufferimplementations, handles reference counting, integrates with experience writers (SQL/JSON), and provides robust batch retrieval with an internal experience pool (exp_pool) and astopped()mechanism for graceful queue exhaustion. - Priority Queue Capacity Logic: The
AsyncPriorityQueue's effective capacity is now dynamically determined as the minimum of the configured capacity and2 * read_batch_size. This ensures the priority queue's size is appropriately scaled relative to the consumer's read rate, as reflected in the updated test cases. - Improved Queue Exhaustion Handling: The
get_batchmethod in the newQueueWrappernow uses astopped()method on the underlying queue and an internalexp_poolto manage experience retrieval. This changes the behavior for an empty queue from raising aTimeoutErrorto raisingStopAsyncIterationwhen the queue is explicitly closed and empty, providing a more Pythonic signal for exhaustion. - Experience Data Model Update: Minor adjustments were made to the
Experiencedata model's serialization. Specifically,chosenandrejectedfields are now serialized aschosen_textandrejected_textrespectively, removing the.tolist()calls, which suggests a change in their underlying data type or representation (likely from tensors to direct string/list of string values). - Enhanced Priority Queue Testing: The
test_priority_queue_capacityintests/buffer/queue_test.pyhas been significantly expanded to cover the new capacity logic and exhaustion behavior, including writing more experiences and asserting the queue length and retrieved experience details more thoroughly.
Using Gemini Code Assist
The 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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.
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 or fill out our survey to provide feedback.
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
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request introduces a refactoring of the queue implementation, replacing the QueueActor with a QueueWrapper and QueueBuffer abstraction. The introduction of AsyncPriorityQueue and AsyncQueue inheriting from a common base class improves modularity, and tests have been enhanced. However, there are critical issues related to potential deadlocks and data loss, as well as minor correctness issues.
|
/unittest-diff |
|
/unittest-all |
Summary
Tests
Github Test Reporter by CTRF 💚 |
Description
As the title says
Checklist
Please check the following items before code is ready to be reviewed.