Skip to content

[Transfer Engine]Coro rpc communicator #795

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
wants to merge 63 commits into
base: main
Choose a base branch
from
Open

[Transfer Engine]Coro rpc communicator #795

wants to merge 63 commits into from

Conversation

JasonZhang517
Copy link
Contributor

  • add a communicator including server and client pool
  • add pybind interface of the communicator
  • add test files

stmatengss
stmatengss reviewed Sep 1, 2025
View reviewed changes
mooncake-integration/CMakeLists.txt Outdated
pybind11_add_module(engine ${SOURCES} ${CACHE_ALLOCATOR_SOURCES}
transfer_engine/transfer_engine_py.cpp
../mooncake-transfer-engine/src/transport/coro_rpc_connector/cororpc_interface.cpp
../mooncake-transfer-engine/src/transport/coro_rpc_connector/cororpc_communicator.cpp
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why need to add these files here?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

if we want to use coro_rpc_connector as a submodule of mooncake transfer engine, we have to compile it here

poor-circle
poor-circle reviewed Sep 2, 2025
View reviewed changes
mooncake-transfer-engine/src/transport/coro_rpc_connector/cororpc_communicator.cpp Outdated
return false;
}

int CoroRPCCommunicator::sendData(const std::string& target_address,
Copy link

@poor-circle poor-circle Sep 2, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sendData should be a wrapper of sendDataAsync, dont rewrite same code twice

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

use syncAwait(Lazy{}) as impl

qicosmos
qicosmos reviewed Sep 3, 2025
View reviewed changes
mooncake-transfer-engine/src/transport/coro_rpc_connector/cororpc_interface.cpp Outdated
auto future_ptr = std::make_shared<pybind11::object>(future_obj);
auto loop_ptr = std::make_shared<pybind11::object>(pybind11::reinterpret_borrow<pybind11::object>(loop));

auto task_func = std::make_shared<std::function<void()>>(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

send_request or
auto lazy = [](int i) -> async_simple::coro::Lazy {
co_await c->send_data_async();
}

lazy(i).start({});

mooncake-transfer-engine/src/transport/coro_rpc_connector/cororpc_interface.cpp Outdated

auto task_func = std::make_shared<std::function<void()>>(
[communicator, target_addr, data_holder, future_ptr, loop_ptr]() {
int result = communicator->sendData(*target_addr, data_holder->data(), data_holder->size());
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

co_await communicator->sendDataAsync

mooncake-transfer-engine/src/transport/coro_rpc_connector/cororpc_interface.cpp Outdated
return future_obj;
}

int CoroRPCInterface::sendTensor(const std::string& target_address, pybind11::handle tensor) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

according async_send_tensor

mooncake-transfer-engine/src/transport/coro_rpc_connector/cororpc_interface.cpp
return future_obj;
}

void CoroRPCInterface::setDataReceiveCallback(pybind11::function callback) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

use asyncio future

@stmatengss stmatengss requested a review from qicosmos September 4, 2025 03:45
@JasonZhang517 JasonZhang517 changed the title Coro rpc communicator [Transfer Engine]Coro rpc communicator Sep 4, 2025
qicosmos
qicosmos reviewed Sep 5, 2025
View reviewed changes
@stmatengss
Copy link
Collaborator

@JasonZhang517 Do you happen to have any updates for this PR?

stmatengss
stmatengss reviewed Sep 14, 2025
View reviewed changes
stmatengss
stmatengss reviewed Sep 14, 2025
View reviewed changes
stmatengss
stmatengss reviewed Sep 14, 2025
View reviewed changes
stmatengss
stmatengss reviewed Sep 14, 2025
View reviewed changes
@JasonZhang517 JasonZhang517 force-pushed the coro_rpc_communicator branch 2 times, most recently from 77b6863 to abd15a1 Compare September 23, 2025 07:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@stmatengss stmatengss stmatengss left review comments

@poor-circle poor-circle Awaiting requested review from poor-circle

+1 more reviewer

@qicosmos qicosmos qicosmos left review comments

Reviewers whose approvals may not affect merge requirements

At least 1 approving review is required to merge this pull request.

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@JasonZhang517 @stmatengss @qicosmos @poor-circle