feature/prefix-aware-routing #546
Conversation
Signed-off-by: BrianPark314 <[email protected]>
Signed-off-by: BrianPark314 <[email protected]>
Signed-off-by: BrianPark314 <[email protected]>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @BrianPark314, 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 implements a new prefix-aware routing mechanism for the inference gateway, designed to intelligently direct requests based on prompt similarity to optimize resource utilization and performance. It includes the core Go logic for this new routing strategy, updates to the build process, and an example Kubernetes configuration for a vLLM runtime.
Highlights
- New Routing Strategy: I've introduced a new
PrefixMatchPickerwithin thegateway-api-inference-extensionto enable prefix-aware routing of inference requests. This picker uses ahashTrieto perform longest-prefix matching on incoming prompts, aiming to route similar prompts to the same backend inference engine (pod) to potentially improve cache hit rates and overall efficiency. - Kubernetes Integration: The new
PrefixMatchPickeris integrated into the build process of thegateway-api-inference-extensionDocker image. Additionally, a new exampleVLLMRuntimeconfiguration (vllm-runtime.yaml) has been added, demonstrating how a vLLM model can be configured within the Kubernetes environment, which would leverage such a routing mechanism. - Scheduler Configuration Update: The
scheduler.patchfile has been updated to include the newPrefixAwarePickeras a potential or default picker in the scheduling configuration, allowing the inference extension to utilize this new routing logic.
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 is currently in preview and 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 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
The pull request introduces a prefix-aware routing mechanism. However, there's a critical issue where multiple pickers are added to the scheduler, which will cause only the last one to be used. Additionally, there's an ignored error in the prefix-aware picker, and a hardcoded copyright year.
# Conflicts: # src/gateway_inference_extension/configs/vllm/vllm-runtime.yaml
BrianPark314
force-pushed
the
feature/prefix-aware-routing
branch
from
db38fc4 to
9cb5ea8
Compare
| return &hashTrie{children: make(map[rune]*hashTrie)} | ||
| } | ||
|
|
||
| func (t *hashTrie) insert(key, endpoint string) { |
There was a problem hiding this comment.
If the key is very long (such as a complete natural language sentence), will it cause excessive memory usage? 🤔
| } | ||
|
|
||
| /*---------------------------- trie implementation ---------------------------*/ | ||
|
|
There was a problem hiding this comment.
Do we need a mechanism just like limiting the cache size or LRU to ensure that our services will not be caused by the infinitely growing prefix cache? 🤔
| return &types.Result{} | ||
| } | ||
|
|
||
| prompt, _ := ctx.RequestBody["prompt"].(string) |
There was a problem hiding this comment.
In chat endpoint, there isn't a prompt in the body, should we consider this kind of scenario too?
| defer t.mu.Unlock() | ||
|
|
||
| node := t | ||
| for _, r := range key { |
There was a problem hiding this comment.
should we do prefix chunk hash instead of treating each character one by one? Like the python version implementation.
Do you mean the python implementation of the routing logics are inconsistent? If so, that should be worked on a seperate PR. |
|
Hi Brian, In the python implementation, the prompt is chunked and hashed, however in this go version, it is per char/token. Maybe chunk and hash would be better and it aligns python implementation. |
Signed-off-by: BrianPark314 <[email protected]>
FILL IN THE PR DESCRIPTION HERE
FIX #xxxx (link existing issues this PR will resolve)
BEFORE SUBMITTING, PLEASE READ THE CHECKLIST BELOW AND FILL IN THE DESCRIPTION ABOVE
-swhen doinggit commit[Bugfix],[Feat], and[CI].Detailed Checklist (Click to Expand)
Thank you for your contribution to production-stack! Before submitting the pull request, please ensure the PR meets the following criteria. This helps us maintain the code quality and improve the efficiency of the review process.
PR Title and Classification
Please try to classify PRs for easy understanding of the type of changes. The PR title is prefixed appropriately to indicate the type of change. Please use one of the following:
[Bugfix]for bug fixes.[CI/Build]for build or continuous integration improvements.[Doc]for documentation fixes and improvements.[Feat]for new features in the cluster (e.g., autoscaling, disaggregated prefill, etc.).[Router]for changes to thevllm_router(e.g., routing algorithm, router observability, etc.).[Misc]for PRs that do not fit the above categories. Please use this sparingly.Note: If the PR spans more than one category, please include all relevant prefixes.
Code Quality
The PR need to meet the following code quality standards:
pre-committo format your code. SeeREADME.mdfor installation.DCO and Signed-off-by
When contributing changes to this project, you must agree to the DCO. Commits must include a
Signed-off-by:header which certifies agreement with the terms of the DCO.Using
-swithgit commitwill automatically add this header.What to Expect for the Reviews
We aim to address all PRs in a timely manner. If no one reviews your PR within 5 days, please @-mention one of YuhanLiu11
, Shaoting-Feng or ApostaC.