refactor remote instance info dict#14971
refactor remote instance info dict#14971amysaq2023 wants to merge 2 commits intosgl-project:mainfrom
Conversation
Signed-off-by: Anqi Shen <amy.saq@antgroup.com>
Summary of ChangesHello @amysaq2023, 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 refactors the internal handling of remote instance transfer engine information within the SGLang framework. By introducing a dedicated utility function for parsing this data from scheduler outputs and adjusting the data flow in key entrypoints, the changes aim to improve code modularity, reduce redundancy, and streamline the management of configuration details related to remote weight loading. This addresses previous feedback and enhances the maintainability of the system. Highlights
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
|
There was a problem hiding this comment.
Code Review
This pull request refactors the logic for parsing remote_instance_transfer_engine_info by moving it into a dedicated utility function. This is a good change that improves code organization and reduces duplication. The changes in engine.py, http_server.py, and scheduler.py correctly adopt this new utility function. My feedback includes a suggestion to make the new utility function more Pythonic and to remove some commented-out code for better maintainability.
| def parse_remote_instance_transfer_engine_info_from_scheduler_infos(scheduler_infos): | ||
| remote_instance_transfer_engine_info = {} | ||
| for data in scheduler_infos: | ||
| if ( | ||
| "tp_rank" in data | ||
| and "remote_instance_transfer_engine_session_id" in data | ||
| and "remote_instance_transfer_engine_weights_info_dict" in data | ||
| ): | ||
| remote_instance_transfer_engine_info[data["tp_rank"]] = ( | ||
| data["remote_instance_transfer_engine_session_id"], | ||
| data["remote_instance_transfer_engine_weights_info_dict"], | ||
| ) | ||
| return remote_instance_transfer_engine_info |
There was a problem hiding this comment.
The implementation of this function can be made more concise and Pythonic by using a dictionary comprehension. This improves readability and reduces verbosity.
| def parse_remote_instance_transfer_engine_info_from_scheduler_infos(scheduler_infos): | |
| remote_instance_transfer_engine_info = {} | |
| for data in scheduler_infos: | |
| if ( | |
| "tp_rank" in data | |
| and "remote_instance_transfer_engine_session_id" in data | |
| and "remote_instance_transfer_engine_weights_info_dict" in data | |
| ): | |
| remote_instance_transfer_engine_info[data["tp_rank"]] = ( | |
| data["remote_instance_transfer_engine_session_id"], | |
| data["remote_instance_transfer_engine_weights_info_dict"], | |
| ) | |
| return remote_instance_transfer_engine_info | |
| def parse_remote_instance_transfer_engine_info_from_scheduler_infos(scheduler_infos): | |
| required_keys = { | |
| "tp_rank", | |
| "remote_instance_transfer_engine_session_id", | |
| "remote_instance_transfer_engine_weights_info_dict", | |
| } | |
| return { | |
| data["tp_rank"]: ( | |
| data["remote_instance_transfer_engine_session_id"], | |
| data["remote_instance_transfer_engine_weights_info_dict"], | |
| ) | |
| for data in scheduler_infos | |
| if required_keys.issubset(data) | |
| } |
This commit address comments in sgl-project#13125 (comment), sgl-project#13125 (comment) Signed-off-by: Anqi Shen <amy.saq@antgroup.com>
5073fe7 to
1562ee0
Compare
Address comments in #13125 (comment) and #13125 (comment)
Checklist