Conversation
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 addresses a critical bug in the Twinkle server's handling of the 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. 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. 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 addresses an issue with serve_multiplexed_model_id by correcting the HTTP header name to the canonical Serve-Multiplexed-Model-Id format and introducing a patch for Ray Serve to handle header normalization from proxies. The changes also improve robustness by using .get() for dictionary access and add a request start hook to several endpoints. My review focuses on improving maintainability by refactoring duplicated code in the endpoint handlers. The core fix and the patching mechanism appear solid.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a significant set of changes, primarily focused on fixing an HTTP header normalization issue with serve_multiplexed_model_id by patching Ray Serve. It also includes a major and beneficial refactoring of the server state management and resource lifecycle from synchronous to asynchronous, which improves the architecture.
My review has identified a critical bug in the new lifecycle management for processors that will cause a TypeError at runtime. I've also pointed out a high-severity issue regarding a potential resource leak in the sampler application due to the removal of adapter lifecycle management. Additionally, I've suggested a medium-severity refactoring to reduce code duplication in the model endpoint handlers. Please address the critical and high-severity issues before merging.
| def _on_processor_expired(self, processor_id: str) -> None: | ||
| """Hook called when a processor's session expires. | ||
|
|
||
| Must be overridden by inheriting classes. | ||
|
|
||
| Raises: | ||
| NotImplementedError: If not overridden. | ||
| """ | ||
| raise NotImplementedError(f'_on_processor_expired must be implemented by {self.__class__.__name__}') |
There was a problem hiding this comment.
The abstract method _on_processor_expired is defined as a synchronous function, but it is being awaited in the base SessionResourceMixin's cleanup loop (await self._on_resource_expired(resource_id) which in turn calls await self._on_processor_expired(resource_id)). This will lead to a TypeError: object NoneType can't be awaited at runtime when a processor expires. The implementation in src/twinkle/server/processor/app.py is also synchronous. Both the abstract method definition here and its implementation should be async def.
| def _on_processor_expired(self, processor_id: str) -> None: | |
| """Hook called when a processor's session expires. | |
| Must be overridden by inheriting classes. | |
| Raises: | |
| NotImplementedError: If not overridden. | |
| """ | |
| raise NotImplementedError(f'_on_processor_expired must be implemented by {self.__class__.__name__}') | |
| async def _on_processor_expired(self, processor_id: str) -> None: | |
| """Hook called when a processor's session expires. | |
| Must be overridden by inheriting classes. | |
| Raises: | |
| NotImplementedError: If not overridden. | |
| """ | |
| raise NotImplementedError(f'_on_processor_expired must be implemented by {self.__class__.__name__}') |
|
|
||
|
|
||
| class SamplerManagement(TaskQueueMixin, AdapterManagerMixin): | ||
| class SamplerManagement(TaskQueueMixin): |
There was a problem hiding this comment.
The AdapterManagerMixin has been removed from SamplerManagement, and the corresponding _init_adapter_manager and _on_adapter_expired logic has also been removed. This means that LoRA adapters added to the sampler via the /twinkle/add_adapter_to_sampler endpoint are no longer session-bound or automatically cleaned up upon session expiration. This could lead to a resource leak on the sampler replicas as adapters accumulate over time without being unloaded. Was this intentional? If not, AdapterManagerMixin (or the new SessionResourceMixin) should be restored to manage the adapter lifecycle on the sampler.
| @app.post('/twinkle/create', response_model=types.CreateResponse) | ||
| async def create(request: Request, body: types.CreateRequest, | ||
| self: ModelManagement = Depends(self_fn)) -> types.CreateResponse: | ||
| await self._on_request_start(request) |
There was a problem hiding this comment.
The call await self._on_request_start(request) is repeated at the beginning of almost every endpoint in this file. This introduces code duplication and makes the handlers more verbose. To improve maintainability, you could refactor this into a FastAPI dependency. This would centralize the request startup logic and clean up the endpoint definitions.
For example:
async def startup_dependency(request: Request, self: ModelManagement = Depends(self_fn)):
await self._on_request_start(request)
return self
@app.post('/twinkle/create', response_model=types.CreateResponse)
async def create(request: Request, body: types.CreateRequest,
self: ModelManagement = Depends(startup_dependency)) -> types.CreateResponse:
# No need for the explicit await here anymore
return types.CreateResponse()1 participant
PR type
PR information
Write the detail information belongs to this PR.
Experiment results
Paste your experiment result here(if needed).