PR3 coro_rpc_communicator python bindings #1106

JasonZhang517 · 2025-11-25T03:49:33Z

Description

The python binding layer of coro_rpc_communicator. Refers to #1104

Type of Change

Types
- [ x ] Bug fix
- New feature
  - Transfer Engine
  - Mooncake Store
  - Mooncake EP
  - Integration
  - P2P Store
  - Python Wheel
- Breaking change
- CI/CD
- Documentation update
- Other

How Has This Been Tested?

Checklist

I have performed a self-review of my own code.
I have updated the documentation.
I have added tests to prove my changes are effective.

- Add cororpc_communicator and cororpc_interface headers and implementation - Implement async RPC communication layer using yalantinglibs coro_rpc - Add CMake configuration for coro_rpc_connector module - Add bandwidth test script for performance validation - Support both TCP and RDMA transports via MC_RPC_PROTOCOL env var This introduces a new transport layer based on coro_rpc for high-performance asynchronous communication in the transfer engine. The implementation uses zero-copy techniques via attachments and provides both synchronous and asynchronous APIs.

Transfer Engine Python bindings: - Add coro_rpc interface integration to transfer_engine_py.cpp - Support transfer operations with notify mechanism - Add getNotifies() method for retrieving transfer notifications - Improve batch transfer sync APIs with notify support Store Python bindings: - Add batch_put_from_multi_buffers() for efficient batch operations - Add batch_get_into_multi_buffers() for batch retrieval - Enhance setup() with additional configuration options - Improve error handling and parameter validation Testing and utilities: - Add comprehensive test_tensor_api.py test script - Update mooncake_config.py with new configuration options - Update CMakeLists.txt for Python integration modules Note: This PR depends on: - PR1 (coro_rpc_connector) for transfer_engine_py.cpp - PR2 (Store enhancements) for store_py.cpp

gemini-code-assist · 2025-11-25T03:49:50Z

Summary of Changes

Hello @JasonZhang517, 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 significantly enhances the interoperability of the Mooncake project by integrating Python bindings for the coro_rpc_communicator. This enables Python-based components to seamlessly utilize the existing C++ RPC communication infrastructure, broadening the scope for development and integration within the ecosystem.

Highlights

Python Bindings for CoroRPCCommunicator: This pull request introduces a new Python binding layer for the coro_rpc_communicator, allowing Python applications to directly interact with the C++ CoroRPCInterface for remote procedure calls.
CMake Build System Updates: The CMake configuration has been updated to correctly build and link the new Python bindings. This includes finding Python3 development components, adding necessary include directories, and conditionally linking against the coro_rpc_connector library.
Exposed RPC Functionality: The Python bindings expose a comprehensive set of RPC functionalities, including methods for initializing clients and servers, sending and receiving data and tensors, and setting up asynchronous receive callbacks. This provides a full-featured Pythonic interface to the underlying C++ RPC system.

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.

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 Python bindings for the coro_rpc_communicator. The changes primarily involve updating CMakeLists.txt to incorporate the new dependencies and adding the necessary pybind11 wrapper code in transfer_engine_py.cpp. The Python bindings themselves are well-implemented. My review focuses on a couple of areas for improvement in the CMake script to remove redundancy and prevent potential build issues.

gemini-code-assist · 2025-11-25T03:50:47Z

mooncake-integration/CMakeLists.txt

+if (TARGET coro_rpc_connector)
+    target_link_libraries(engine PUBLIC coro_rpc_connector)
+endif()


The object files from coro_rpc_connector are included on line 41 via $<TARGET_OBJECTS>. Linking the target again here is either redundant (for static libraries) or an error (for object libraries). Please remove this block to avoid redundancy and potential build failures.

gemini-code-assist · 2025-11-25T03:50:47Z

mooncake-integration/CMakeLists.txt

+if (TARGET transfer_engine)
+    target_link_libraries(engine PUBLIC transfer_engine)
 endif()


The transfer_engine target is already linked to the engine target on line 51. This if block is redundant and can be removed to simplify the script.

…unicator config 3. removed comment lines

Copilot

Pull request overview

This PR attempts to add Python bindings for the coro_rpc_communicator component to the Mooncake transfer engine. However, the implementation has several critical issues that prevent it from compiling. The changes add new binding code in transfer_engine_py.cpp and update the CMakeLists.txt build configuration, but reference files and functions that don't exist in the codebase.

Key issues:

References a non-existent header file (cororpc_interface.h) and undefined classes/functions
Contains duplicate CMake target definitions that will cause build failures
Uses an undefined CMake variable (ENGINE_MODULE_SOURCES)

Reviewed changes

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

File	Description
mooncake-integration/transfer_engine/transfer_engine_py.cpp	Adds Python bindings for CoroRPCInterface, ReceivedData, and ReceivedTensor classes, plus factory functions. However, the included header and referenced symbols don't exist.
mooncake-integration/CMakeLists.txt	Updates build configuration with duplicate pybind11_add_module definitions, redundant Python3 package finding, and conditional linking to the coro_rpc_connector target.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

mooncake-integration/CMakeLists.txt

Copilot · 2025-12-17T06:10:06Z

mooncake-integration/transfer_engine/transfer_engine_py.cpp

+    m.def("create_rpc_client", &createRPCClient, "local_rank"_a,
+          "world_size"_a);
+    m.def("create_rpc_server", &createRPCServer, "local_rank"_a,
+          "world_size"_a);


The factory function names createRPCClient and createRPCServer use inconsistent casing compared to the existing pattern in the codebase. The existing rpc_interface.h uses createRpcClient and createRpcServer (with 'Rpc' in camelCase, not 'RPC' in all caps). For consistency, these should be named createCoroRpcClient and createCoroRpcServer to match the existing naming convention.

mooncake-integration/transfer_engine/transfer_engine_py.cpp

Copilot · 2025-12-17T06:10:07Z

mooncake-integration/transfer_engine/transfer_engine_py.cpp

+    m.def("create_rpc_client", &createRPCClient, "local_rank"_a,
+          "world_size"_a);
+    m.def("create_rpc_server", &createRPCServer, "local_rank"_a,
+          "world_size"_a);


The functions createRPCClient and createRPCServer are referenced but not declared or defined anywhere in the codebase. These factory functions need to be either: 1) declared and defined in the cororpc_interface.h header (similar to how createRpcClient and createRpcServer are defined in rpc_interface.h), or 2) removed from these bindings if they don't exist. Without proper declarations, this code will not compile.

Copilot · 2025-12-17T06:10:07Z

mooncake-integration/transfer_engine/transfer_engine_py.cpp

+// Forward declaration for coro_rpc_interface binding function
+void bind_coro_rpc_interface(py::module_ &m);


The forward declaration of bind_coro_rpc_interface is immediately followed by its implementation in the same file. This forward declaration is unnecessary since the function is defined before it's used in the PYBIND11_MODULE. Consider removing the forward declaration and keeping only the implementation.

Copilot · 2025-12-17T06:10:07Z

mooncake-integration/CMakeLists.txt

    endif()
 endif()

+pybind11_add_module(engine ${ENGINE_MODULE_SOURCES})


The variable ENGINE_MODULE_SOURCES is used but never defined in this CMakeLists.txt file. This will cause the pybind11_add_module command to create an empty module or fail. Either define ENGINE_MODULE_SOURCES before this line, or replace it with the actual source files needed for the engine module (likely transfer_engine/transfer_engine_py.cpp).

Copilot · 2025-12-17T06:10:07Z

mooncake-integration/CMakeLists.txt

    endif()
 endif()

+pybind11_add_module(engine ${ENGINE_MODULE_SOURCES})


The pybind11_add_module for 'engine' is defined twice in this file - once at line 35 (inside the WITH_TE conditional block) and again at line 65 (unconditionally). This creates a duplicate target definition which will cause a CMake configuration error. Remove one of these duplicate definitions. Since the one at line 65 appears to be the newer addition for coro_rpc support, you likely need to either merge the two definitions or ensure only one is active based on the build configuration.

stmatengss · 2025-12-17T06:52:01Z

mooncake-integration/CMakeLists.txt

 include_directories("./")
+include_directories("../mooncake-transfer-engine/include")

+find_package(Python3 COMPONENTS Interpreter Development REQUIRED)


Please see line 46! It's duplicated.

stmatengss · 2025-12-17T06:55:57Z

mooncake-integration/transfer_engine/transfer_engine_py.cpp

+// Include coro_rpc_interface headers
+#include "transport/coro_rpc_connector/cororpc_interface.h"
+
+using namespace pybind11::literals;


py::literals are already accessible.

stmatengss · 2025-12-18T06:07:08Z

Please address the issues above.

…python-bindings

mooncake-integration/transfer_engine/transfer_engine_py.cpp

stmatengss · 2025-12-29T12:43:20Z

Please fix compilation errors for a quick merge. THX

stmatengss · 2026-01-04T02:57:13Z

code format....

codecov-commenter · 2026-01-04T03:22:33Z

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 0% with 4 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...integration/transfer_engine/transfer_engine_py.cpp	0.00%	3 Missing ⚠️
...rc/transport/rpc_communicator/rpc_communicator.cpp	0.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

stmatengss · 2026-01-04T03:26:02Z

mooncake-transfer-engine/src/transport/rpc_communicator/rpc_communicator.cpp

 namespace py = pybind11;

 namespace mooncake {
+namespace py = pybind11;


JasonZhang517 · 2026-01-04T13:02:16Z

https://github.com/kvcache-ai/Mooncake/blob/main/mooncake-transfer-engine/tests/rpc_communicator_test.py
here is the tests

stmatengss · 2026-01-05T02:13:57Z

mooncake-transfer-engine/src/transport/rpc_communicator/rpc_interface.cpp

-        .def_readonly("data_size", &RpcInterface::ReceivedData::data_size)
-        .def("get_bytes", &RpcInterface::ReceivedData::getBytes)
-        .def("get_memory_view", &RpcInterface::ReceivedData::getMemoryView);
+    safe_register_class([&]() {


Don't need it.

luketong777 and others added 6 commits November 21, 2025 16:34

Merge branch 'kvcache-ai:main' into pr1-coro-rpc-core

0a1846b

merged to main and solved the conflicts

e98ee15

Merge remote-tracking branch 'upstream/main' into pr3-python-bindings

b5ac4fd

Merge branch 'kvcache-ai:main' into pr1-coro-rpc-core

0eb7263

gemini-code-assist bot reviewed Nov 25, 2025

View reviewed changes

luketong777 added 21 commits November 26, 2025 17:58

fixed issues: 1.moved config to config.h 2. renamed config to RPCComm…

30d1a50

…unicator config 3. removed comment lines

fixed the names

9491366

fixed compilation errors

f8a1935

fixed some compilation errors

6b96c3c

1. updated names 2. updated CI

d94e502

reformat the code

e6a7dee

1.removed useless pybind dependencies 2. renamed test file

c147359

Merge remote-tracking branch 'upstream/main' into pr1-coro-rpc-core

cc48c16

fixed AI assisted bugs

d0bb308

fixed compile errors

08fd2ed

fixed cmakelists.txt

54722fd

fixed some bugs

6b006e4

removed useless files

3875616

Merge remote-tracking branch 'upstream/main' into pr1-coro-rpc-core

65d8d0b

fixed hardcoding of cmakelists.txt

07cbd4a

updated cmakelists.txt

9cd2e46

updated pybind

734d65e

update pybind

adfaa83

update cmakelists.txt pybind11

d4caead

set cmakelists.txt

fa90a08

Merge remote-tracking branch 'upstream/main' into pr3-python-bindings

a50aa47

Copilot AI review requested due to automatic review settings December 17, 2025 06:05

github-actions bot added the run-ci label Dec 17, 2025

Copilot started reviewing on behalf of JasonZhang517 December 17, 2025 06:06 View session

Copilot AI reviewed Dec 17, 2025

View reviewed changes

stmatengss reviewed Dec 17, 2025

View reviewed changes

luketong777 added 2 commits December 17, 2025 15:24

Merge branch 'pr1-coro-rpc-core' into pr3-python-bindings

d61bc25

updated transfer_engine_py.cpp about the class name

225e39d

JasonZhang517 requested review from chestnut-Q and doujiang24 as code owners December 17, 2025 08:30

github-actions bot added the Transfer Engine label Dec 17, 2025

re-format the codes

e9920fe

JasonZhang517 and others added 5 commits December 25, 2025 19:20

fixed cmakelists.txt

c493206

Merge branch 'kvcache-ai:main' into pr3-python-bindings

83aa96f

Merge remote-tracking branch 'upstream/pr3-python-bindings' into pr3-…

604c58b

…python-bindings

fixed pybind import

af2c0be

reformat the codes

a380f7b

stmatengss reviewed Dec 28, 2025

View reviewed changes

mooncake-integration/transfer_engine/transfer_engine_py.cpp Outdated Show resolved Hide resolved

add safe register class

b963438

JasonZhang517 force-pushed the pr3-python-bindings branch from feca777 to b963438 Compare January 4, 2026 02:12

reformat the codes

1cc7861

stmatengss reviewed Jan 4, 2026

View reviewed changes

JasonZhang517 added 2 commits January 4, 2026 15:45

removed duplicated codes

55bdccf

reformat the code

fbc67b5

stmatengss reviewed Jan 5, 2026

View reviewed changes

remove useless register

71763d8

		// Forward declaration for coro_rpc_interface binding function
		void bind_coro_rpc_interface(py::module_ &m);

PR3 coro_rpc_communicator python bindings #1106

Are you sure you want to change the base?

PR3 coro_rpc_communicator python bindings #1106

Uh oh!

Conversation

JasonZhang517 commented Nov 25, 2025

Description

Type of Change

How Has This Been Tested?

Checklist

Uh oh!

gemini-code-assist bot commented Nov 25, 2025

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 Nov 25, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 25, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Copilot AI Dec 17, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Copilot AI Dec 17, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Dec 17, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Dec 17, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Dec 17, 2025

Choose a reason for hiding this comment

Uh oh!

stmatengss Dec 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

stmatengss Dec 17, 2025

Choose a reason for hiding this comment

Uh oh!

stmatengss commented Dec 18, 2025

Uh oh!

Uh oh!

stmatengss commented Dec 29, 2025

Uh oh!

stmatengss commented Jan 4, 2026

Uh oh!

codecov-commenter commented Jan 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

stmatengss Jan 4, 2026

Choose a reason for hiding this comment

Uh oh!

JasonZhang517 commented Jan 4, 2026

Uh oh!

stmatengss Jan 5, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

stmatengss Dec 17, 2025 •

edited

Loading

codecov-commenter commented Jan 4, 2026 •

edited

Loading