Add custom route for sticky session

[Yunnglin]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

Write the detail information belongs to this PR.

Experiment results

Paste your experiment result here(if needed).

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! 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 implements a "sticky session" routing strategy for LoRA adapters within the Ray Serve framework. The primary goal is to enhance performance and resource efficiency by ensuring that subsequent requests related to a specific LoRA adapter are consistently directed to the same model replica that initially loaded it. This is achieved through a custom request router that considers both existing multiplexed sessions and the available LoRA capacity on each replica, supported by backend state management and client-side header injection.

Highlights

• Custom Sticky Session Router: Introduced StickyLoraRequestRouter to intelligently route requests to replicas based on multiplexed model IDs and available LoRA capacity, ensuring requests for the same adapter are directed to the same replica.
• Replica Capacity Management: Implemented mechanisms in ModelManager and ServerStateProxy to track the number of LoRA adapters loaded on each replica, allowing the router to identify replicas with available capacity.
• Deployment Configuration Updates: Modified ModelManagement deployment to utilize the new sticky router and register its LoRA capacity, along with increasing replica counts and per_token_model_limit in the configuration.
• Client-Side Header Injection: Updated the client-side patching logic to automatically inject serve_multiplexed_model_id and authentication headers into requests, which are essential for the sticky routing mechanism.

🧠 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

• cookbook/client/tinker/megatron/server_config_7b.yaml
  • per_token_model_limit was increased from 1 to 3.
  • max_loras configuration was added for adapter management.
  • min_replicas and max_replicas for ModelManagement were updated from 1 to 2.
• src/twinkle/model/multi_lora.py
  • The error message for unavailable LoRA slots was enhanced to include the max_loras limit.
• src/twinkle/server/tinker/common/router.py
  • A new StickyLoraRequestRouter class was added to provide custom replica selection logic.
• src/twinkle/server/tinker/model.py
  • The ModelManagement deployment was configured to use the StickyLoraRequestRouter.
  • Replica ID and max_loras were added to the ModelManagement initialization.
  • The replica's LoRA capacity was registered with the server state during initialization.
  • A new _on_request_start method was introduced to centralize token extraction and sticky session setup for API endpoints.
  • All API methods were updated to use the new _on_request_start for token retrieval.
  • The register_model call was updated to pass the replica_id.
• src/twinkle/server/utils/state/model_manager.py
  • Internal dictionaries _replica_models and _replica_max_loras were added to track replica capacity.
  • Methods register_replica, unregister_replica, and get_available_replica_ids were implemented for replica capacity management.
  • Model add and remove operations were updated to manage replica ownership.
  • A new helper method _cleanup_ownership was introduced to streamline ownership cleanup.
• src/twinkle/server/utils/state/models.py
  • An optional replica_id field was added to the ModelRecord class.
• src/twinkle/server/utils/state/server_state.py
  • The register_model method was updated to accept and pass replica_id.
  • Proxy methods for replica management (register_replica, unregister_replica, get_available_replica_ids) were added to ServerStateProxy.
• src/twinkle_client/utils/patch_tinker.py
  • The patching logic for ServiceClient.__init__ was refactored into a separate function.
  • The patched ServiceClient.__init__ was updated to inject serve_multiplexed_model_id and authentication headers into client requests.

Activity

• The pull request introduces a new feature for sticky sessions.
• Configuration files were updated to reflect new limits and replica counts.
• Core routing and state management logic were added and modified.
• Client-side changes were implemented to support the new routing.

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 by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

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 pull request 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. 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

1. 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. ↩

[gemini-code-assist]

Code Review

This pull request introduces a custom router for sticky sessions, which is a key feature for handling stateful LoRA adapters efficiently. The implementation is comprehensive, touching configuration, client-side patching, server-side routing, and state management. The logic for tracking replica capacity and using it for routing decisions is well-thought-out. My feedback focuses on improving the maintainability and robustness of the new code, specifically by replacing print statements with a proper logger and refactoring duplicated logic within the ModelManagement class. Overall, this is a solid implementation of a complex feature.

[Copilot]

Pull request overview

This pull request implements sticky session routing for LoRA adapters by introducing a custom Ray Serve request router that tracks replica capacity and routes requests to replicas with available LoRA slots. The implementation adds replica capacity tracking to the server state management system and modifies the model management service to register replicas and associate models with their hosting replicas.

Changes:

• Adds StickyLoraRequestRouter that routes requests based on both Ray Serve's multiplexing mechanism and replica LoRA capacity
• Implements replica registration and capacity tracking in ServerState and ModelManager
• Modifies model creation to track which replica hosts each model for improved routing decisions
• Refactors client-side tinker patching for better code organization

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 11 comments.

Show a summary per file

File	Description
src/twinkle/server/tinker/common/router.py	New custom router implementing capacity-aware sticky session routing for LoRA adapters
src/twinkle/server/utils/state/server_state.py	Adds replica management methods (register/unregister/query capacity) to ServerState and proxy
src/twinkle/server/utils/state/models.py	Adds replica_id field to ModelRecord for tracking model-replica associations
src/twinkle/server/utils/state/model_manager.py	Implements replica tracking with capacity limits and cleanup logic
src/twinkle/server/tinker/model.py	Registers replica on startup, tracks models per replica, extracts token via helper function, implements multiplexed sticky entry
src/twinkle_client/utils/patch_tinker.py	Refactors ServiceClient patching into separate function for better maintainability
src/twinkle/model/multi_lora.py	Improves error message to include max_loras limit
cookbook/client/tinker/megatron/server_config_7b.yaml	Updates config to scale to 2 replicas with 1 LoRA per replica and increases per-token model limit