diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b3126cad..11c6aeb6 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -239,6 +239,16 @@ jobs: python scripts/test_tensor_api.py -n 1 shell: bash + - name: Run RPC Communicator Bandwidth Test + run: | + source test_env/bin/activate + python mooncake-transfer-engine/tests/rpc_communicator_test.py server --url 127.0.0.1:9004 --data-size 1 & + SERVER_PID=$! + sleep 5 + timeout 10 python mooncake-transfer-engine/tests/rpc_communicator_test.py client --url 127.0.0.1:9004 --threads 2 --data-size 1 || true + kill $SERVER_PID 2>/dev/null || true + wait $SERVER_PID 2>/dev/null || true + - name: Test Mooncake EP Backend (CPU Only) run: | source test_env/bin/activate diff --git a/mooncake-integration/CMakeLists.txt b/mooncake-integration/CMakeLists.txt index bf46581e..b88423b9 100644 --- a/mooncake-integration/CMakeLists.txt +++ b/mooncake-integration/CMakeLists.txt @@ -36,6 +36,14 @@ if (WITH_TE) set_target_properties(engine PROPERTIES INSTALL_RPATH "$ORIGIN" ) + + target_link_libraries(engine PRIVATE + $ + ) + + find_package(Python3 COMPONENTS Interpreter Development REQUIRED) + target_link_libraries(engine PRIVATE ${Python3_LIBRARIES}) + target_include_directories(engine PRIVATE ${Python3_INCLUDE_DIRS}) if (USE_ASCEND_DIRECT) target_link_libraries(engine PUBLIC ascendcl diff --git a/mooncake-integration/transfer_engine/transfer_engine_py.cpp b/mooncake-integration/transfer_engine/transfer_engine_py.cpp index 7ab1046e..9fef5a0d 100644 --- a/mooncake-integration/transfer_engine/transfer_engine_py.cpp +++ b/mooncake-integration/transfer_engine/transfer_engine_py.cpp @@ -19,6 +19,7 @@ #include #include +#include "transport/rpc_communicator/rpc_interface.h" #ifdef USE_MNNVL #include "transport/nvlink_transport/nvlink_transport.h" @@ -775,4 +776,7 @@ PYBIND11_MODULE(engine, m) { py::class_>( m, "InnerTransferEngine"); + + // Bind RpcInterface + mooncake::bind_rpc_interface(m); } diff --git a/mooncake-transfer-engine/include/config.h b/mooncake-transfer-engine/include/config.h index d914fc8d..c006d1b3 100644 --- a/mooncake-transfer-engine/include/config.h +++ b/mooncake-transfer-engine/include/config.h @@ -60,6 +60,13 @@ struct GlobalConfig { EndpointStoreType endpoint_store_type = EndpointStoreType::SIEVE; }; +struct RpcCommunicatorConfig { + std::string listen_address; + size_t thread_count = 0; + size_t timeout_seconds = 30; + size_t pool_size = 10; +}; + void loadGlobalConfig(GlobalConfig &config); void dumpGlobalConfig(); diff --git a/mooncake-transfer-engine/include/transport/rpc_communicator/rpc_communicator.h b/mooncake-transfer-engine/include/transport/rpc_communicator/rpc_communicator.h new file mode 100644 index 00000000..4b39a4af --- /dev/null +++ b/mooncake-transfer-engine/include/transport/rpc_communicator/rpc_communicator.h @@ -0,0 +1,81 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "config.h" + +namespace mooncake { + +struct TensorInfo { + void* data_ptr = nullptr; + std::vector shape; + std::string dtype; + size_t total_bytes = 0; +}; + +struct RpcResult { + int code = 0; + std::string err_msg; +}; + +class RpcCommunicator { + public: + RpcCommunicator(); + ~RpcCommunicator(); + + bool initialize(const RpcCommunicatorConfig& config); + bool startServerImpl(bool is_async = true); + bool startServer(); + bool startServerAsync(); + void stopServer(); + + int sendData(const std::string& target_address, const void* data, + size_t data_size); + async_simple::coro::Lazy sendDataAsync( + const std::string& target_address, const void* data, size_t data_size); + + int sendTensor(const std::string& target_address, + const pybind11::object& tensor); + async_simple::coro::Lazy sendTensorAsync( + const std::string& target_address, const TensorInfo& tensor); + + int receiveData(const std::string& source_address, void* buffer, + size_t buffer_size, int timeout_ms = -1); + async_simple::coro::Lazy receiveDataAsync( + const std::string& source_address, int timeout_ms = -1); + + void setDataReceiveCallback( + std::function callback); + + private: + // Handler methods for RPC calls + void handleDataTransfer(coro_rpc::context context, + std::string_view data); + void handleTensorTransfer(coro_rpc::context context); + void handleDataTransferWithAttachment(coro_rpc::context context, + std::string_view data); + void handleTensorTransferWithAttachment(coro_rpc::context context); + + // Implementation members + RpcCommunicatorConfig config_; + bool is_server_started_ = false; + std::unique_ptr server_; + std::function + data_receive_callback_; + pybind11::handle py_callback_; + std::shared_ptr> + client_pools_; +}; + +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-transfer-engine/include/transport/rpc_communicator/rpc_interface.h b/mooncake-transfer-engine/include/transport/rpc_communicator/rpc_interface.h new file mode 100644 index 00000000..eb6e19e5 --- /dev/null +++ b/mooncake-transfer-engine/include/transport/rpc_communicator/rpc_interface.h @@ -0,0 +1,91 @@ +#pragma once + +#include +#include +#include +#include +#include + +namespace mooncake { + +class RpcInterface { + public: + struct ReceivedData { + std::string source_address; + std::string data; + size_t data_size = 0; + + pybind11::bytes getBytes() const { return pybind11::bytes(data); } + pybind11::memoryview getMemoryView() const { + return pybind11::memoryview::from_memory( + const_cast(static_cast(data.data())), + data.size(), true); + } + }; + + struct ReceivedTensor { + std::string source_address; + std::string data; + std::vector shape; + std::string dtype; + size_t total_bytes = 0; + size_t getDataSize() const { return data.size(); } + pybind11::bytes getDataAsBytes() const { return pybind11::bytes(data); } + pybind11::memoryview getMemoryView() const { + return pybind11::memoryview::from_memory( + const_cast(static_cast(data.data())), + data.size(), true); + } + }; + + class Impl; + + RpcInterface(); + ~RpcInterface(); + + bool initialize(const std::string& listen_address = "", + size_t thread_count = 0, size_t timeout_seconds = 30, + size_t pool_size = 10); + + // Convenience methods for common use cases + bool initializeClient(size_t pool_size = 10, size_t timeout_seconds = 30); + bool initializeServer(const std::string& listen_address, + size_t thread_count = 8, size_t timeout_seconds = 30); + + bool startServer(); + bool startServerAsync(); + bool startServerImpl(bool is_async = true); + void stopServer(); + + int sendData(const std::string& target_address, pybind11::handle data); + pybind11::object sendDataAsync(const std::string& target_address, + pybind11::handle data, + pybind11::handle loop); + + int sendTensor(const std::string& target_address, pybind11::handle tensor); + pybind11::object sendTensorAsync(const std::string& target_address, + pybind11::handle tensor, + pybind11::handle loop); + + void setDataReceiveCallback(pybind11::function callback); + void setTensorReceiveCallback(pybind11::function callback); + + void handleIncomingData(std::string_view source_address, + std::string_view data); + void handleIncomingTensor(std::string_view source_address, + std::string_view data, + const std::vector& shape, + std::string_view dtype); + + private: + std::unique_ptr impl_; +}; + +std::unique_ptr createRpcClient(uint64_t local_rank = 0, + uint64_t world_size = 1); +std::unique_ptr createRpcServer(uint64_t local_rank = 0, + uint64_t world_size = 1); + +void bind_rpc_interface(pybind11::module_& m); + +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-transfer-engine/src/transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/CMakeLists.txt index 137202f3..d2a4746c 100644 --- a/mooncake-transfer-engine/src/transport/CMakeLists.txt +++ b/mooncake-transfer-engine/src/transport/CMakeLists.txt @@ -1,7 +1,8 @@ file(GLOB XPORT_SOURCES "*.cpp") add_subdirectory(rdma_transport) -add_library(transport OBJECT ${XPORT_SOURCES} $) +add_subdirectory(rpc_communicator) +add_library(transport OBJECT ${XPORT_SOURCES} $ $) target_link_libraries(transport PRIVATE JsonCpp::JsonCpp yalantinglibs::yalantinglibs glog::glog pthread) if (USE_TCP) diff --git a/mooncake-transfer-engine/src/transport/rpc_communicator/CMakeLists.txt b/mooncake-transfer-engine/src/transport/rpc_communicator/CMakeLists.txt new file mode 100644 index 00000000..841a36aa --- /dev/null +++ b/mooncake-transfer-engine/src/transport/rpc_communicator/CMakeLists.txt @@ -0,0 +1,41 @@ +file(GLOB CORO_RPC_SOURCES "*.cpp") + +find_package(Python3 COMPONENTS Interpreter Development REQUIRED) + +add_library(rpc_communicator OBJECT ${CORO_RPC_SOURCES}) + +target_link_libraries(rpc_communicator + PRIVATE + JsonCpp::JsonCpp + yalantinglibs::yalantinglibs + glog::glog + pthread + ${Python3_LIBRARIES} +) + +target_include_directories(rpc_communicator + PRIVATE + ${Python3_INCLUDE_DIRS} +) + +if(TARGET pybind11::headers) + target_link_libraries(rpc_communicator PRIVATE pybind11::headers) +elseif(DEFINED pybind11_INCLUDE_DIR) + target_include_directories(rpc_communicator PRIVATE ${pybind11_INCLUDE_DIR}) +else() + set(_PYBIND11_DIR "") + if(EXISTS "${CMAKE_SOURCE_DIR}/extern/pybind11/include/pybind11/pybind11.h") + set(_PYBIND11_DIR "${CMAKE_SOURCE_DIR}/extern/pybind11/include") + else() + get_filename_component(_PROJECT_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/../../../../" ABSOLUTE) + set(_PYBIND11_DIR "${_PROJECT_ROOT}/extern/pybind11/include") + if(NOT EXISTS "${_PYBIND11_DIR}/pybind11/pybind11.h") + set(_PYBIND11_DIR "") + endif() + endif() + if(_PYBIND11_DIR) + target_include_directories(rpc_communicator PRIVATE "${_PYBIND11_DIR}") + else() + message(FATAL_ERROR "Could not find pybind11") + endif() +endif() \ No newline at end of file diff --git a/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_communicator.cpp b/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_communicator.cpp new file mode 100644 index 00000000..682ec56d --- /dev/null +++ b/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_communicator.cpp @@ -0,0 +1,419 @@ +#include "transport/rpc_communicator/rpc_communicator.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "async_simple/coro/SyncAwait.h" + +namespace py = pybind11; + +namespace mooncake { + +class py_rpc_context { + public: + void response_msg(py::buffer msg, py::object done) { + py::buffer_info info = msg.request(); + const char* data = static_cast(info.ptr); + context_.get_context_info()->set_response_attachment( + std::string_view(data, info.size)); + context_.get_context_info()->set_complete_handler( + [done](const std::error_code& ec, std::size_t) { + py::gil_scoped_acquire acquire; + done(!ec); + }); + context_.response_msg(); + } + + coro_rpc::context context_; +}; + +RpcCommunicator::RpcCommunicator() {} + +RpcCommunicator::~RpcCommunicator() { stopServer(); } + +void RpcCommunicator::setDataReceiveCallback( + std::function callback) { + LOG(INFO) << "Setting data receive callback..."; + data_receive_callback_ = callback; + LOG(INFO) << "Data receive callback set successfully"; +} + +bool RpcCommunicator::initialize(const RpcCommunicatorConfig& config) { + config_ = config; + easylog::set_min_severity(easylog::Severity::WARNING); // Set log level + // to WARNING + + // Initialize client pools with proper configuration + coro_io::client_pool::pool_config pool_conf{}; + const char* value = std::getenv("MC_RPC_PROTOCOL"); + if (value && std::string_view(value) == "rdma") { + pool_conf.client_config.socket_config = + coro_io::ib_socket_t::config_t{}; + } + client_pools_ = + std::make_shared>( + pool_conf); + + LOG(INFO) << "create coro_rpc_client_pool with " << config.pool_size + << " threads"; + if (!config.listen_address.empty()) { + LOG(INFO) << "Initializing server on " << config.listen_address; + + server_ = std::make_unique( + config.thread_count, config.listen_address, + std::chrono::seconds(config.timeout_seconds)); + + if (value && std::string_view(value) == "rdma") { + if (server_) { + try { + server_->init_ibv(); + LOG(INFO) << "RDMA initialized successfully"; + } catch (const std::exception& e) { + LOG(ERROR) << "RDMA initialization failed: " << e.what(); + LOG(WARNING) << "Falling back to TCP mode"; + // Continue without RDMA - the server will use TCP + } catch (...) { + LOG(ERROR) + << "RDMA initialization failed with unknown error"; + LOG(WARNING) << "Falling back to TCP mode"; + // Continue without RDMA - the server will use TCP + } + } else { + LOG(ERROR) << "Server pointer is null, cannot initialize RDMA"; + LOG(WARNING) << "Falling back to TCP mode"; + } + } + + server_->register_handler<&RpcCommunicator::handleDataTransfer, + &RpcCommunicator::handleTensorTransfer>(this); + } + LOG(INFO) << "Environment variable MC_RPC_PROTOCOL is set to " + << (value ? value : "not set"); + if (value && std::string_view(value) == "rdma") { + LOG(INFO) << "Using RDMA transport for RPC communication"; + } else { + LOG(INFO) << "Using TCP transport for RPC communication"; + } + + LOG(INFO) << "Communicator initialized with client pool support"; + return true; +} + +bool RpcCommunicator::startServerImpl(bool is_async) { + if (is_async) { + return this->startServerAsync(); + } else { + return this->startServer(); + } +} + +bool RpcCommunicator::startServer() { + if (!server_ || config_.listen_address.empty()) return false; + + try { + auto ec = server_->start(); + if (ec.val() == 0) { + is_server_started_ = true; + LOG(INFO) << "Server started on " << config_.listen_address; + return true; + } else { + LOG(ERROR) << "Failed to start server: " << ec.message(); + return false; + } + } catch (const std::exception& e) { + LOG(ERROR) << "Failed to start server: " << e.what(); + return false; + } +} + +bool RpcCommunicator::startServerAsync() { + if (!server_ || config_.listen_address.empty()) return false; + + try { + auto ec = server_->async_start(); + if (!ec.hasResult()) { + is_server_started_ = true; + LOG(INFO) << "Server started asynchronously on " + << config_.listen_address; + return true; + } else { + LOG(ERROR) << "Failed to start server asynchronously"; + return false; + } + } catch (const std::exception& e) { + LOG(ERROR) << "Failed to start server asynchronously: " << e.what(); + return false; + } +} + +void RpcCommunicator::stopServer() { + if (is_server_started_ && server_) { + server_->stop(); + is_server_started_ = false; + LOG(INFO) << "Server stopped"; + } +} + +int RpcCommunicator::sendData(const std::string& target_address, + const void* data, size_t data_size) { + auto result = async_simple::coro::syncAwait( + sendDataAsync(target_address, data, data_size)); + return result.code; +} + +async_simple::coro::Lazy RpcCommunicator::sendDataAsync( + const std::string& target_address, const void* data, size_t data_size) { + std::string_view data_view(static_cast(data), data_size); + + // For large data, use attachment to avoid copying + const size_t ATTACHMENT_THRESHOLD = 1024; // Use attachment for data > 1KB + + auto rpc_result = co_await client_pools_->send_request( + target_address, + [data_view, data_size](coro_rpc::coro_rpc_client& client) + -> async_simple::coro::Lazy { + if (data_size > ATTACHMENT_THRESHOLD) { + // Use attachment for large data - zero copy + client.set_req_attachment(data_view); + // Send empty data parameter, actual data in attachment + auto result = + co_await client.call<&RpcCommunicator::handleDataTransfer>( + std::string_view{}); + if (!result.has_value()) { + LOG(ERROR) << "RPC call failed: " << result.error().msg; + } + } else { + // Use regular parameter for small data + auto result = + co_await client.call<&RpcCommunicator::handleDataTransfer>( + data_view); + if (!result.has_value()) { + LOG(ERROR) << "RPC call failed: " << result.error().msg; + } + } + }); + + if (!rpc_result.has_value()) { + LOG(ERROR) << "RPC send request failed"; + co_return RpcResult{-1, "RPC call failed"}; + } + RpcResult res; + res.code = 0; + co_return res; +} + +int RpcCommunicator::sendTensor(const std::string& target_address, + const pybind11::object& tensor) { + try { + TensorInfo tensor_info; + { + pybind11::gil_scoped_acquire acquire; + pybind11::object tensor_obj = + pybind11::reinterpret_borrow(tensor); + + // Validate tensor type using duck typing - check for required + // attributes + if (!pybind11::hasattr(tensor_obj, "data_ptr") || + !pybind11::hasattr(tensor_obj, "numel") || + !pybind11::hasattr(tensor_obj, "element_size") || + !pybind11::hasattr(tensor_obj, "shape") || + !pybind11::hasattr(tensor_obj, "dtype")) { + LOG(ERROR) << "Input is not a valid tensor object (missing " + "required attributes)"; + return -1; + } + + // Extract tensor properties + uintptr_t data_ptr = + tensor_obj.attr("data_ptr")().cast(); + size_t numel = tensor_obj.attr("numel")().cast(); + size_t element_size = + tensor_obj.attr("element_size")().cast(); + size_t tensor_size = numel * element_size; + + // Get tensor shape + pybind11::object shape_obj = tensor_obj.attr("shape"); + pybind11::tuple shape_tuple = + pybind11::cast(shape_obj); + const size_t shape_size = shape_tuple.size(); + std::vector shape; + shape.reserve(shape_size); // Pre-allocate to avoid reallocations + for (size_t i = 0; i < shape_size; ++i) { + shape.push_back(shape_tuple[i].cast()); + } + + // Get tensor dtype string + pybind11::object dtype_obj = tensor_obj.attr("dtype"); + std::string dtype = pybind11::str(dtype_obj).cast(); + + // Fill TensorInfo structure (no data copying, just metadata) + tensor_info.data_ptr = reinterpret_cast(data_ptr); + tensor_info.total_bytes = tensor_size; + tensor_info.shape = std::move(shape); + tensor_info.dtype = std::move(dtype); + + // Format shape as string for logging - optimize string building + // with reserve (glog LOG macro handles log level checking) + std::string shape_str; + // Estimate size: each number ~10 chars + ", " = ~12 chars per + // dimension + shape_str.reserve(tensor_info.shape.size() * 12); + shape_str = "["; + for (size_t i = 0; i < tensor_info.shape.size(); ++i) { + if (i > 0) shape_str += ", "; + shape_str += std::to_string(tensor_info.shape[i]); + } + shape_str += "]"; + LOG(INFO) << "Sending tensor with shape: " << shape_str + << ", dtype: " << tensor_info.dtype + << ", tensor size: " << tensor_info.total_bytes + << " bytes"; + } + + // Use the async version which supports zero-copy via attachments + pybind11::gil_scoped_release release; + auto result = async_simple::coro::syncAwait( + sendTensorAsync(target_address, tensor_info)); + return result; + + } catch (const std::exception& e) { + LOG(ERROR) << "Send tensor error: " << e.what(); + return -1; + } +} + +async_simple::coro::Lazy RpcCommunicator::sendTensorAsync( + const std::string& target_address, const TensorInfo& tensor) { + auto rpc_result = co_await client_pools_->send_request( + target_address, + [tensor](coro_rpc::coro_rpc_client& client) + -> async_simple::coro::Lazy { + client.set_req_attachment(std::string_view( + static_cast(tensor.data_ptr), tensor.total_bytes)); + + auto result = + co_await client.call<&RpcCommunicator::handleTensorTransfer>(); + + if (!result.has_value()) { + LOG(ERROR) << "Tensor RPC call failed: " << result.error().msg; + } + }); + if (!rpc_result.has_value()) { + LOG(ERROR) << "Tensor RPC send request failed"; + co_return -1; + } + co_return 0; +} + +int RpcCommunicator::receiveData(const std::string& source_address, + void* buffer, size_t buffer_size, + int timeout_ms) { + auto result = async_simple::coro::syncAwait( + receiveDataAsync(source_address, timeout_ms)); + return 0; +} + +async_simple::coro::Lazy RpcCommunicator::receiveDataAsync( + const std::string& source_address, int timeout_ms) { + // For attachment-based data reception, we should use a different approach + // This method is typically called from the handler when data is received + // The actual data reception is handled by the registered handlers + co_return std::string(); +} // Data reception is handled via context and attachment in handlers + +void RpcCommunicator::handleDataTransfer(coro_rpc::context context, + std::string_view data) { + // Check if there's an attachment for large data + auto ctx_info = context.get_context_info(); + auto attachment = ctx_info->get_request_attachment(); + + LOG(INFO) << "Handling data transfer - Data: " << data.size() + << " bytes, Attachment: " << attachment.size() << " bytes"; + // Call the data receive callback if set + if (data_receive_callback_) { + LOG(INFO) << "Calling data receive callback..."; + // Note: coro_rpc context doesn't provide get_remote_endpoint() + // Using empty string as placeholder - can be enhanced if needed + std::string_view source_address = ""; + + // Use attachment if available (for large data), otherwise use data + // parameter + if (!attachment.empty()) { + // Use attachment data directly without copying - zero copy approach + std::string_view attachment_view = attachment; + data_receive_callback_(source_address, attachment_view); + } else { + // For small data, use the regular data parameter + data_receive_callback_(source_address, data); + } + } else { + LOG(INFO) << "No data receive callback set!"; + } + + // Echo back the attachment for response (zero-copy) + if (!attachment.empty()) { + ctx_info->set_response_attachment(std::string_view("ok")); + } + + context.response_msg(); +} + +void RpcCommunicator::handleTensorTransfer(coro_rpc::context context) { + auto ctx_info = context.get_context_info(); + auto attachment = ctx_info->get_request_attachment(); + + LOG(INFO) << "Handling tensor transfer: " << attachment.size() << " bytes"; + + // Call the data receive callback if set (tensor data is received via + // attachment) + if (data_receive_callback_) { + LOG(INFO) << "Calling data receive callback for tensor..."; + // Note: coro_rpc context doesn't provide get_remote_endpoint() + // Using empty string as placeholder - can be enhanced if needed + std::string_view source_address = ""; + + // Pass the attachment data to the callback + data_receive_callback_(source_address, attachment); + } else { + LOG(INFO) << "No data receive callback set for tensor!"; + } + + ctx_info->set_response_attachment(attachment); + context.response_msg(); +} + +void RpcCommunicator::handleDataTransferWithAttachment( + coro_rpc::context context, std::string_view data) { + py_rpc_context t{}; + t.context_ = std::move(context); + py::gil_scoped_acquire acquire; + auto view = + py::memoryview::from_buffer(data.data(), {data.size()}, {sizeof(char)}); + + py_callback_(std::move(t), view); +} + +void RpcCommunicator::handleTensorTransferWithAttachment( + coro_rpc::context context) { + py_rpc_context t{}; + + // Get the attachment before moving the context + auto ctx_info = context.get_context_info(); + auto attachment = ctx_info->get_request_attachment(); + + t.context_ = std::move(context); + py::gil_scoped_acquire acquire; + + auto view = py::memoryview::from_buffer( + attachment.data(), {attachment.size()}, {sizeof(int8_t)}); + + py_callback_(std::move(t), view); +} + +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_interface.cpp b/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_interface.cpp new file mode 100644 index 00000000..13d44f6f --- /dev/null +++ b/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_interface.cpp @@ -0,0 +1,589 @@ +#include "transport/rpc_communicator/rpc_interface.h" +#include "transport/rpc_communicator/rpc_communicator.h" +#include "config.h" +#include +#include +#include +#include +#include +#include "async_simple/coro/SyncAwait.h" + +namespace mooncake { + +static constexpr size_t MAX_TENSOR_DIMS = 4; +// Tensor metadata size: dtype (4 bytes) + ndim (4 bytes) + shape dimensions +// (MAX_TENSOR_DIMS * 8 bytes) +static constexpr size_t TENSOR_METADATA_SIZE = 4 + 4 + MAX_TENSOR_DIMS * 8; + +// Implementation class +class RpcInterface::Impl { + public: + std::unique_ptr communicator; + pybind11::function data_receive_callback; + pybind11::function tensor_receive_callback; +}; + +// Constructor +RpcInterface::RpcInterface() : impl_(std::make_unique()) {} + +// Destructor +RpcInterface::~RpcInterface() = default; + +// Initialize +bool RpcInterface::initialize(const std::string& local_address, + size_t thread_count, size_t timeout_seconds, + size_t pool_size) { + RpcCommunicatorConfig config; + config.listen_address = local_address; + config.thread_count = thread_count; + config.timeout_seconds = timeout_seconds; + config.pool_size = pool_size; + + impl_->communicator = std::make_unique(); + return impl_->communicator->initialize(config); +} + +// Convenience method for client initialization +bool RpcInterface::initializeClient(size_t pool_size, size_t timeout_seconds) { + return initialize("", 0, timeout_seconds, pool_size); +} + +// Convenience method for server initialization +bool RpcInterface::initializeServer(const std::string& listen_address, + size_t thread_count, + size_t timeout_seconds) { + return initialize(listen_address, thread_count, timeout_seconds, 4); +} + +bool RpcInterface::startServer() { + if (!impl_->communicator) return false; + return impl_->communicator->startServer(); +} + +bool RpcInterface::startServerAsync() { + if (!impl_->communicator) return false; + return impl_->communicator->startServerAsync(); +} + +bool RpcInterface::startServerImpl(bool is_async) { + if (!impl_->communicator) return false; + return impl_->communicator->startServerImpl(is_async); +} + +void RpcInterface::stopServer() { + if (impl_->communicator) { + impl_->communicator->stopServer(); + } +} + +int RpcInterface::sendData(const std::string& target_address, + pybind11::handle data) { + if (!impl_->communicator) return -1; + + pybind11::gil_scoped_acquire acquire; + + // Extract data from handle directly + std::string data_str; + try { + // Try to get direct string view from bytes object + pybind11::bytes data_bytes = + pybind11::reinterpret_borrow(data); + data_str = static_cast(data_bytes); + } catch (...) { + // Fallback: convert to bytes and then get string + pybind11::bytes data_bytes = pybind11::cast(data); + data_str = static_cast(data_bytes); + } + + pybind11::gil_scoped_release release; + return impl_->communicator->sendData(target_address, data_str.data(), + data_str.size()); +} + +pybind11::object RpcInterface::sendDataAsync(const std::string& target_address, + pybind11::handle data, + pybind11::handle loop) { + pybind11::gil_scoped_acquire acquire; + + auto future_module = pybind11::module_::import("asyncio"); + auto future_obj = future_module.attr("Future")(); + + if (!impl_->communicator) { + auto exc_type = + pybind11::module_::import("builtins").attr("RuntimeError"); + auto exc = exc_type(pybind11::str("Communicator not initialized")); + future_obj.attr("set_exception")(exc); + return future_obj; + } + + auto communicator = impl_->communicator.get(); + std::string target_addr = target_address; + + // Extract data from handle directly without creating intermediate bytes + // object + std::string_view data_view; + std::string data_str; // Only used if we can't get direct view + + try { + // Try to get direct string view from bytes object + pybind11::bytes data_bytes = + pybind11::reinterpret_borrow(data); + data_view = data_bytes; + // Store a copy for lambda capture since string_view might not be valid + // after GIL release + data_str = std::string(data_view); + } catch (...) { + // Fallback: convert to bytes and then to string + pybind11::bytes data_bytes = pybind11::cast(data); + data_str = data_bytes; + } + + // Release GIL before starting coroutine + pybind11::gil_scoped_release release; + + auto coro_lambda = [communicator, target_addr, data_str, future_obj, + loop]() -> async_simple::coro::Lazy { + try { + auto result_struct = co_await communicator->sendDataAsync( + target_addr, data_str.data(), data_str.size()); + int result = result_struct.code; + + auto call_soon_threadsafe = [future_obj, loop, result]() { + pybind11::gil_scoped_acquire acquire; + if (result >= 0) { + future_obj.attr("set_result")(result); + } else { + auto runtime_error = pybind11::module_::import("builtins") + .attr("RuntimeError"); + future_obj.attr("set_exception")( + runtime_error(pybind11::str("Send data failed"))); + } + }; + + auto callback = pybind11::cpp_function(call_soon_threadsafe); + loop.attr("call_soon_threadsafe")(callback); + } catch (const std::exception& e) { + auto call_soon_threadsafe = [future_obj, loop, e]() { + pybind11::gil_scoped_acquire acquire; + auto runtime_error = + pybind11::module_::import("builtins").attr("RuntimeError"); + future_obj.attr("set_exception")(runtime_error(pybind11::str( + std::string("Send data error: ") + e.what()))); + }; + + auto callback = pybind11::cpp_function(call_soon_threadsafe); + loop.attr("call_soon_threadsafe")(callback); + } + }; + + auto lazy = coro_lambda(); + lazy.start([](auto&& result) { + if (result.hasError()) { + LOG(ERROR) << "Coroutine completed with error"; + } + }); + + return future_obj; +} + +int RpcInterface::sendTensor(const std::string& target_address, + pybind11::handle tensor) { + if (!impl_->communicator) return -1; + + try { + TensorInfo tensor_info; + { + pybind11::gil_scoped_acquire acquire; + pybind11::object tensor_obj = + pybind11::reinterpret_borrow(tensor); + + // Validate tensor type using duck typing - check for required + // attributes + if (!pybind11::hasattr(tensor_obj, "data_ptr") || + !pybind11::hasattr(tensor_obj, "numel") || + !pybind11::hasattr(tensor_obj, "element_size") || + !pybind11::hasattr(tensor_obj, "shape") || + !pybind11::hasattr(tensor_obj, "dtype")) { + LOG(ERROR) << "Input is not a valid tensor object (missing " + "required attributes)"; + return -1; + } + + // Extract tensor properties - zero copy, just get pointers and + // metadata + uintptr_t data_ptr = + tensor_obj.attr("data_ptr")().cast(); + size_t numel = tensor_obj.attr("numel")().cast(); + size_t element_size = + tensor_obj.attr("element_size")().cast(); + size_t tensor_size = numel * element_size; + + // Get tensor shape + pybind11::object shape_obj = tensor_obj.attr("shape"); + pybind11::tuple shape_tuple = + pybind11::cast(shape_obj); + std::vector shape; + for (size_t i = 0; i < shape_tuple.size(); i++) { + shape.push_back(shape_tuple[i].cast()); + } + + // Get tensor dtype string + pybind11::object dtype_obj = tensor_obj.attr("dtype"); + std::string dtype = dtype_obj.attr("__str__")().cast(); + + // Fill TensorInfo structure (no data copying, just metadata) + tensor_info.data_ptr = reinterpret_cast(data_ptr); + tensor_info.total_bytes = tensor_size; + tensor_info.shape = std::move(shape); + tensor_info.dtype = std::move(dtype); + + // Format shape as string for logging + std::string shape_str = "["; + for (size_t i = 0; i < tensor_info.shape.size(); i++) { + shape_str += std::to_string(tensor_info.shape[i]); + if (i < tensor_info.shape.size() - 1) shape_str += ", "; + } + shape_str += "]"; + LOG(INFO) << "Sending tensor with shape: " << shape_str + << ", dtype: " << tensor_info.dtype + << ", tensor size: " << tensor_info.total_bytes + << " bytes"; + } + + // Use the async version which supports zero-copy via attachments + pybind11::gil_scoped_release release; + auto result = async_simple::coro::syncAwait( + impl_->communicator->sendTensorAsync(target_address, tensor_info)); + return result; + + } catch (const std::exception& e) { + LOG(ERROR) << "Send tensor error: " << e.what(); + return -1; + } +} + +pybind11::object RpcInterface::sendTensorAsync( + const std::string& target_address, pybind11::handle tensor, + pybind11::handle loop) { + pybind11::gil_scoped_acquire acquire; + + auto future_module = pybind11::module_::import("asyncio"); + auto future_obj = future_module.attr("Future")(); + + if (!impl_->communicator) { + auto exc_type = + pybind11::module_::import("builtins").attr("RuntimeError"); + auto exc = exc_type(pybind11::str("Communicator not initialized")); + future_obj.attr("set_exception")(exc); + return future_obj; + } + + auto communicator = impl_->communicator.get(); + std::string target_addr = target_address; + + // Extract tensor info + pybind11::object tensor_obj = + pybind11::reinterpret_borrow(tensor); + uintptr_t data_ptr = tensor_obj.attr("data_ptr")().cast(); + size_t numel = tensor_obj.attr("numel")().cast(); + size_t element_size = tensor_obj.attr("element_size")().cast(); + size_t tensor_size = numel * element_size; + + // Get tensor shape and dtype + pybind11::object shape_obj = tensor_obj.attr("shape"); + pybind11::tuple shape_tuple = pybind11::cast(shape_obj); + std::vector shape; + for (size_t i = 0; i < shape_tuple.size(); i++) { + shape.push_back(shape_tuple[i].cast()); + } + + pybind11::object dtype_obj = tensor_obj.attr("dtype"); + std::string dtype = dtype_obj.attr("__str__")().cast(); + + // Create TensorInfo on stack (no need for shared_ptr) + TensorInfo tensor_info; + tensor_info.data_ptr = reinterpret_cast(data_ptr); + tensor_info.total_bytes = tensor_size; + tensor_info.shape = std::move(shape); + tensor_info.dtype = std::move(dtype); + + // Keep a reference to the tensor object to avoid it being garbage collected + pybind11::object py_tensor_obj = tensor_obj; + + pybind11::gil_scoped_release release; + + auto coro_lambda = [communicator, target_addr, tensor_info, future_obj, + loop, + py_tensor_obj]() -> async_simple::coro::Lazy { + try { + auto result = co_await communicator->sendTensorAsync(target_addr, + tensor_info); + + auto call_soon_threadsafe = [future_obj, loop, result]() { + pybind11::gil_scoped_acquire acquire; + if (result >= 0) { + future_obj.attr("set_result")(result); + } else { + auto runtime_error = pybind11::module_::import("builtins") + .attr("RuntimeError"); + future_obj.attr("set_exception")( + runtime_error(pybind11::str("Send tensor failed"))); + } + }; + + auto callback = pybind11::cpp_function(call_soon_threadsafe); + loop.attr("call_soon_threadsafe")(callback); + } catch (const std::exception& e) { + auto call_soon_threadsafe = [future_obj, loop, e]() { + pybind11::gil_scoped_acquire acquire; + auto runtime_error = + pybind11::module_::import("builtins").attr("RuntimeError"); + future_obj.attr("set_exception")(runtime_error(pybind11::str( + std::string("Send tensor error: ") + e.what()))); + }; + + auto callback = pybind11::cpp_function(call_soon_threadsafe); + loop.attr("call_soon_threadsafe")(callback); + } + }; + + // Start the coroutine in a detached manner (fire and forget) + auto lazy = coro_lambda(); + lazy.start([](auto&& result) { + // This callback will be called when the coroutine completes + // We don't need to do anything here since the result is handled in + // the coroutine itself + if (result.hasError()) { + // Log error if needed + LOG(ERROR) << "Tensor coroutine completed with error"; + } + }); + + return future_obj; +} + +void RpcInterface::setDataReceiveCallback(pybind11::function callback) { + pybind11::gil_scoped_acquire acquire; + impl_->data_receive_callback = callback; + if (impl_->communicator) { + // Capture this pointer - safe because RpcInterface is managed by Python + // and will remain valid as long as the callback is set + auto interface_ptr = this; + impl_->communicator->setDataReceiveCallback( + [interface_ptr](std::string_view source, std::string_view data) { + interface_ptr->handleIncomingData(source, data); + }); + } +} + +void RpcInterface::setTensorReceiveCallback(pybind11::function callback) { + pybind11::gil_scoped_acquire acquire; + impl_->tensor_receive_callback = callback; + + // Note: Tensor data is received through the regular data callback + // The handleIncomingData function will detect tensor data and route it + // to handleIncomingTensor automatically +} + +void RpcInterface::handleIncomingData(std::string_view source, + std::string_view data) { + LOG(INFO) << "RpcInterface::handleIncomingData called with " << data.size() + << " bytes"; + + // C++ tensor rebuilding + if (data.size() >= TENSOR_METADATA_SIZE) { + // Read the first few bytes to check if it looks like tensor metadata + const uint32_t* header = reinterpret_cast(data.data()); + uint32_t dtype = header[0]; + uint32_t ndim = header[1]; + + LOG(INFO) << "Checking tensor metadata: dtype=" << dtype + << ", ndim=" << ndim; + + // Basic validation: check if dtype and ndim are in reasonable ranges + if (dtype > 0 && dtype <= 9 && ndim <= 4) { + LOG(INFO) + << "Data recognized as tensor, calling handleIncomingTensor"; + + // This looks like tensor data, handle it as such + std::vector shape; + const int64_t* shape_data = + reinterpret_cast(data.data() + 8); + if (data.size() < TENSOR_METADATA_SIZE || + data.size() < 8 + ndim * sizeof(int64_t)) { + LOG(WARNING) << "Data too small for claimed tensor metadata"; + return; + } + for (int i = 0; i < static_cast(ndim); i++) { + if (shape_data[i] > 0 && shape_data[i] < (1LL << 48)) { + shape.push_back(static_cast(shape_data[i])); + } else if (shape_data[i] > 0) { + LOG(WARNING) << "Shape dimension " << i + << " too large: " << shape_data[i]; + return; + } + } + + // Get dtype name based on dtype ID + std::string_view dtype_name; + switch (dtype) { + case 1: + dtype_name = "float16"; + break; + case 2: + dtype_name = "float32"; + break; + case 3: + dtype_name = "float64"; + break; + case 4: + dtype_name = "int8"; + break; + case 5: + dtype_name = "int16"; + break; + case 6: + dtype_name = "int32"; + break; + case 7: + dtype_name = "int64"; + break; + case 8: + dtype_name = "uint8"; + break; + case 9: + dtype_name = "bool"; + break; + default: + dtype_name = "unknown"; + break; + } + + // Call tensor handler instead of data handler + handleIncomingTensor(source, data, shape, dtype_name); + return; + } + } + + // Handle as regular data if not tensor data + if (!impl_->data_receive_callback) return; + + try { + pybind11::gil_scoped_acquire acquire; + pybind11::dict received; + received["source"] = + std::string(source); // Convert to string for Python + received["data"] = pybind11::bytes( + std::string(data)); // Convert to string for pybind11::bytes + + impl_->data_receive_callback(received); + } catch (const std::exception& e) { + LOG(ERROR) << "Error in data receive callback: " << e.what(); + } +} + +void RpcInterface::handleIncomingTensor(std::string_view source, + std::string_view data, + const std::vector& shape, + std::string_view dtype) { + LOG(INFO) << "RpcInterface::handleIncomingTensor called" + << " - source: " << source << ", data size: " << data.size() + << ", dtype: " << dtype << ", shape size: " << shape.size(); + + if (!impl_->tensor_receive_callback) { + LOG(WARNING) << "No tensor receive callback set!"; + return; + } + + LOG(INFO) << "Calling Python tensor receive callback..."; + + try { + pybind11::gil_scoped_acquire acquire; + + ReceivedTensor received; + received.source_address = std::string(source); + received.data = std::string(data); + received.shape = shape; + received.dtype = std::string(dtype); + + impl_->tensor_receive_callback(received); + } catch (const std::exception& e) { + LOG(ERROR) << "Error in tensor receive callback: " << e.what(); + } +} + +// Factory functions for creating RPC client and server +std::unique_ptr createRpcClient(uint64_t local_rank, + uint64_t world_size) { + auto client = std::make_unique(); + // Initialize client with default settings + client->initialize("", 0, 30, 10); + return client; +} + +std::unique_ptr createRpcServer(uint64_t local_rank, + uint64_t world_size) { + auto server = std::make_unique(); + // Initialize server with default settings + server->initialize("0.0.0.0:8080", 0, 30, 10); + return server; +} + +// Python binding implementation +void bind_rpc_interface(pybind11::module_& m) { + namespace py = pybind11; + using namespace mooncake; + + // Bind RpcInterface::ReceivedData + py::class_(m, "ReceivedData") + .def_readonly("source_address", + &RpcInterface::ReceivedData::source_address) + .def_readonly("data_size", &RpcInterface::ReceivedData::data_size) + .def("get_bytes", &RpcInterface::ReceivedData::getBytes) + .def("get_memory_view", &RpcInterface::ReceivedData::getMemoryView); + + // Bind RpcInterface::ReceivedTensor + py::class_(m, "ReceivedTensor") + .def_readonly("source_address", + &RpcInterface::ReceivedTensor::source_address) + .def_readonly("shape", &RpcInterface::ReceivedTensor::shape) + .def_readonly("dtype", &RpcInterface::ReceivedTensor::dtype) + .def_readonly("total_bytes", &RpcInterface::ReceivedTensor::total_bytes) + .def("get_data_size", &RpcInterface::ReceivedTensor::getDataSize) + .def("get_data_as_bytes", &RpcInterface::ReceivedTensor::getDataAsBytes) + .def("get_memory_view", &RpcInterface::ReceivedTensor::getMemoryView); + + // Bind RpcInterface + py::class_(m, "RpcInterface") + .def(py::init<>()) + .def("initialize", &RpcInterface::initialize, + py::arg("listen_address") = "", py::arg("thread_count") = 0, + py::arg("timeout_seconds") = 30, py::arg("pool_size") = 10) + .def("initialize_client", &RpcInterface::initializeClient, + py::arg("pool_size") = 10, py::arg("timeout_seconds") = 30) + .def("initialize_server", &RpcInterface::initializeServer, + py::arg("listen_address"), py::arg("thread_count") = 8, + py::arg("timeout_seconds") = 30) + .def("start_server", &RpcInterface::startServer) + .def("start_server_async", &RpcInterface::startServerAsync) + .def("stop_server", &RpcInterface::stopServer) + .def("send_data", &RpcInterface::sendData, py::arg("target_address"), + py::arg("data")) + .def("send_data_async", &RpcInterface::sendDataAsync, + py::arg("target_address"), py::arg("data"), py::arg("loop")) + .def("send_tensor", &RpcInterface::sendTensor, + py::arg("target_address"), py::arg("tensor")) + .def("send_tensor_async", &RpcInterface::sendTensorAsync, + py::arg("target_address"), py::arg("tensor"), py::arg("loop")) + .def("set_data_receive_callback", &RpcInterface::setDataReceiveCallback) + .def("set_tensor_receive_callback", + &RpcInterface::setTensorReceiveCallback); + + // Bind factory functions + m.def("create_rpc_client", &createRpcClient, py::arg("local_rank") = 0, + py::arg("world_size") = 1); + m.def("create_rpc_server", &createRpcServer, py::arg("local_rank") = 0, + py::arg("world_size") = 1); +} + +} // namespace mooncake diff --git a/mooncake-transfer-engine/tests/rpc_communicator_test.py b/mooncake-transfer-engine/tests/rpc_communicator_test.py new file mode 100644 index 00000000..aec20877 --- /dev/null +++ b/mooncake-transfer-engine/tests/rpc_communicator_test.py @@ -0,0 +1,134 @@ +import time +import threading +import argparse +import mooncake.engine as engine + +class AtomicCounter: + def __init__(self, initial=0): + self._value = initial + self._lock = threading.Lock() + + def inc(self, num=1): + with self._lock: + self._value += num + return self._value + + def dec(self, num=1): + with self._lock: + self._value -= num + return self._value + + def get_and_reset(self): + with self._lock: + r = self._value + self._value = 0 + return r + +counter = AtomicCounter() + +# Global variable to store data size +data_size = 1024 * 1024 # Default 1MB +test_data = None + +def print_qps(): + while True: + time.sleep(1) + val = counter.get_and_reset() + if val == 0: + continue + print("bandwidth:", 8 * val * data_size / (1024*1024*1024), "GB/s") + +def send_data(client, target_url): + while True: + try: + result = client.send_data(target_url, test_data) + if result < 0: + print(f"Warning: send_data returned error code {result}") + time.sleep(0.01) # Brief delay on error to avoid tight loop + else: + counter.inc() + except Exception as e: + print(f"Error sending data to {target_url}: {e}") + time.sleep(0.1) # Delay on exception to avoid rapid retry + +def run_server(bind_url, data_size_mb=1): + """Run server mode""" + global data_size, test_data + data_size = data_size_mb * 1024 * 1024 + test_data = b'\x00' * data_size + + print(f"Starting server on {bind_url} with {data_size_mb}MB data packets") + + RpcInterface = engine.RpcInterface + server = RpcInterface() + server.initialize_server(bind_url, thread_count=8) + server.start_server_async() + + # Start QPS statistics thread + thread = threading.Thread(target=print_qps) + thread.daemon = True + thread.start() + + print(f"Server started on {bind_url}, press Ctrl+C to stop") + try: + while True: + time.sleep(1) + except KeyboardInterrupt: + print("\nServer stopping...") + server.stop_server() + +def run_client(target_url, num_threads=8, data_size_mb=1): + """Run client mode""" + global data_size, test_data + data_size = data_size_mb * 1024 * 1024 + test_data = b'\x00' * data_size + + print(f"Starting client, connecting to {target_url} with {num_threads} threads, {data_size_mb}MB data packets") + + RpcInterface = engine.RpcInterface + client = RpcInterface() + client.initialize_client(pool_size=100) + + # Start QPS statistics thread + qps_thread = threading.Thread(target=print_qps) + qps_thread.daemon = True + qps_thread.start() + + # Start multiple sending threads + threads = [] + for i in range(num_threads): + thread = threading.Thread(target=send_data, args=(client, target_url)) + thread.daemon = True + thread.start() + threads.append(thread) + + print(f"Client started with {num_threads} threads, press Ctrl+C to stop") + try: + while True: + time.sleep(1) + except KeyboardInterrupt: + print("\nClient stopping...") + +def main(): + parser = argparse.ArgumentParser(description='Mooncake Communication Bandwidth Test Tool') + parser.add_argument('mode', choices=['server', 'client'], + help='Run mode: server or client') + parser.add_argument('--url', default='127.0.0.1:9004', + help='URL address (default: 127.0.0.1:9004)') + parser.add_argument('--threads', type=int, default=8, + help='Number of client threads (default: 8)') + parser.add_argument('--data-size', type=int, default=1, + help='Data packet size in MB (default: 1)') + + args = parser.parse_args() + + if args.mode == 'server': + # Server mode, URL as bind address + bind_url = f"0.0.0.0:{args.url.split(':')[-1]}" if ':' in args.url else f"0.0.0.0:{args.url}" + run_server(bind_url, args.data_size) + else: + # Client mode, URL as target address + run_client(args.url, args.threads, args.data_size) + +if __name__ == "__main__": + main()