[Feature] KVCache: add Pris connector #1303

DwyaneShi · 2025-07-17T19:24:52Z

Pull Request Description

Add Pris Connector

Related Issues

Resolves: #[Insert issue number(s)]

Important: Before submitting, please complete the description above and review the checklist below.

Contribution Guidelines (Expand for Details)

We appreciate your contribution to aibrix! To ensure a smooth review process and maintain high code quality, please adhere to the following guidelines:

Pull Request Title Format

Your PR title should start with one of these prefixes to indicate the nature of the change:

[Bug]: Corrections to existing functionality
[CI]: Changes to build process or CI pipeline
[Docs]: Updates or additions to documentation
[API]: Modifications to aibrix's API or interface
[CLI]: Changes or additions to the Command Line Interface
[Misc]: For changes not covered above (use sparingly)

Note: For changes spanning multiple categories, use multiple prefixes in order of importance.

Submission Checklist

PR title includes appropriate prefix(es)
Changes are clearly explained in the PR description
New and existing tests pass successfully
Code adheres to project style and best practices
Documentation updated to reflect changes (if applicable)
Thorough testing completed, no regressions introduced

By submitting this PR, you confirm that you've read these guidelines and your changes align with the project's contribution standards.

gemini-code-assist

Summary of Changes

Hello @DwyaneShi, 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 introduces a new Pris connector to the KVCache system, expanding its capabilities to integrate with Pris key-value stores. This feature allows aibrix_kvcache to leverage Pris for distributed key-value operations, configurable through new environment variables and integrated seamlessly into the existing connector factory.

Highlights

New Pris Connector: I've added a new PrisConnector to the aibrix_kvcache system. This connector enables aibrix_kvcache to interact with a Pris key-value store, supporting operations like get, put, exists, delete, and their multi-key variants (mget, mput). It also handles memory registration for RDMA operations.
Configuration via Environment Variables: New environment variables have been introduced in aibrix_kvcache/envs.py to configure the Pris connector. These include AIBRIX_KV_CACHE_OL_PRIS_REMOTE_ADDR, AIBRIX_KV_CACHE_OL_PRIS_REMOTE_PORT, AIBRIX_KV_CACHE_OL_PRIS_USE_MPUT_MGET, and AIBRIX_KV_CACHE_OL_PRIS_PASSWORD, allowing flexible setup of the Pris connection.
Connector Factory Integration: The create function in aibrix_kvcache/l2/connectors/connector.py has been updated to include 'PRIS' as a recognized backend. This allows the system to dynamically instantiate the PrisConnector when 'PRIS' is specified as the desired backend.
New Dependency: The pris library (version >= 0.0.4) has been added as a new dependency in requirements/core.txt, as it provides the underlying client for interacting with the Pris key-value store.

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.

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 new PrisConnector for the L2 cache. The review identifies a potential UnboundLocalError in register_slabs when handling empty lists, critical bugs in _mget and _mput where the error handling logic can lead to returning lists of an incorrect size, and a potential misuse of the delete_keys API in the _delete method.

gemini-code-assist · 2025-07-17T19:27:15Z

python/aibrix_kvcache/aibrix_kvcache/l2/connectors/pris.py

+                return Status(StatusCodes.INVALID)
+            desc = PrisRegisterDescriptor(reg_buf)
+            self._register_cache[addr] = desc
+        return Status.ok(desc)


If the input slabs list is empty, the for loop on line 137 will not execute, and desc will not be defined, causing an UnboundLocalError. To fix this, return Status.ok() before the loop if slabs is empty.

if not slabs: return Status.ok() for slab in slabs:

yapple · 2025-07-18T02:16:27Z

python/aibrix_kvcache/aibrix_kvcache/l2/connectors/pris.py

+        else:
+            for i, s in enumerate(details):
+                if s != 0:
+                    op_status[i] = Status(StatusCodes.ERROR)


need we keep the detail error code from the sdk?

Having detailed error messages would be helpful, we can enhance it later.

yapple · 2025-07-18T02:17:10Z

python/aibrix_kvcache/aibrix_kvcache/l2/connectors/pris.py

+        else:
+            for i, s in enumerate(details):
+                if s != 0:
+                    op_status[i] = Status(StatusCodes.ERROR)


need we keep the detail error code from the sdk?

yapple

Perhaps we should create a mock version of the Priskv SDK to facilitate unit testing.

DwyaneShi · 2025-07-18T03:17:34Z

Perhaps we should create a mock version of the Priskv SDK to facilitate unit testing.

Exactly

Xuanjinlong · 2025-07-18T02:09:01Z

python/aibrix_kvcache/aibrix_kvcache/l2/connectors/pris.py

+            if not desc_status.is_ok():
+                for j in range(i, len(keys)):
+                    op_status[j] = Status(desc_status)
+                break


这里如果分配不到，break的话，下面的mget应该还是执行，这里应该return？

Xuanjinlong · 2025-07-18T02:09:56Z

python/aibrix_kvcache/aibrix_kvcache/l2/connectors/pris.py

+            if not desc_status.is_ok():
+                for j in range(i, len(keys)):
+                    op_status[j] = Status(desc_status)
+                break


这里break应该也有问题

python/aibrix_kvcache/aibrix_kvcache/envs.py

sovmatrixac · 2025-07-18T06:49:28Z

python/aibrix_kvcache/aibrix_kvcache/l2/connectors/pris.py

+            desc_status = self._get_register_descriptor(mr)
+            if not desc_status.is_ok():
+                for j in range(i, len(keys)):
+                    op_status[j] = Status(desc_status)


Here, if error occurs when we register memory, it seems we still go to self.conn.mget?

python/aibrix_kvcache/aibrix_kvcache/envs.py

Signed-off-by: Haiyang Shi <[email protected]>

Jeffwan · 2025-07-19T06:50:00Z

@sovmatrixac @yapple @Xuanjinlong any further feedback?

yapple

LGTM

DwyaneShi requested a review from Jeffwan July 17, 2025 19:24

gemini-code-assist bot reviewed Jul 17, 2025

View reviewed changes

DwyaneShi force-pushed the haiyang/kvcache-prisdb branch 2 times, most recently from f88409b to 0d91a8a Compare July 17, 2025 20:18

vllm-project deleted a comment from gemini-code-assist bot Jul 17, 2025

yapple reviewed Jul 18, 2025

View reviewed changes

Xuanjinlong reviewed Jul 18, 2025

View reviewed changes

sovmatrixac reviewed Jul 18, 2025

View reviewed changes

python/aibrix_kvcache/aibrix_kvcache/envs.py Show resolved Hide resolved

DwyaneShi force-pushed the haiyang/kvcache-prisdb branch from 0d91a8a to b0d993c Compare July 18, 2025 21:09

[Feature] KVCache: add Pris connector

5908177

Signed-off-by: Haiyang Shi <[email protected]>

DwyaneShi force-pushed the haiyang/kvcache-prisdb branch from b0d993c to 5908177 Compare July 18, 2025 21:29

yapple approved these changes Jul 19, 2025

View reviewed changes

[Feature] KVCache: add Pris connector #1303

Are you sure you want to change the base?

[Feature] KVCache: add Pris connector #1303

Conversation

DwyaneShi commented Jul 17, 2025

Pull Request Description

Related Issues

Pull Request Title Format

Submission Checklist

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Jul 17, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yapple left a comment

Choose a reason for hiding this comment

Uh oh!

DwyaneShi commented Jul 18, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Jeffwan commented Jul 19, 2025

Uh oh!

yapple left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!