From a7cc83bbae43e548df42c0af0df68f391315aa77 Mon Sep 17 00:00:00 2001 From: hmirin Date: Mon, 31 Aug 2026 02:26:16 +0900 Subject: [PATCH] rpc: avoid serializing buffers from other servers (#26500) * rpc: avoid serializing buffers from other servers Only include remote buffer pointers when the buffer belongs to the RPC dispatcher receiving the graph. Add a two-server regression test for cross-server tensor serialization. Assisted-by: Codex * cont : add ref --------- Co-authored-by: Georgi Gerganov --- ggml/src/ggml-rpc/ggml-rpc.cpp | 26 +++++++++++------- tests/test-rpc-multi-server.cpp | 47 +++++++++++++++++++++++++++++++++ tests/test-rpc-multi-server.sh | 43 ++++++++++++++++++++++++++++++ tools/rpc/CMakeLists.txt | 12 +++++++++ 4 files changed, 118 insertions(+), 10 deletions(-) create mode 100644 tests/test-rpc-multi-server.cpp create mode 100755 tests/test-rpc-multi-server.sh diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index 58a8a030c..a97db24e6 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -625,7 +625,7 @@ static bool ggml_backend_buffer_is_rpc(ggml_backend_buffer_t buffer) { return buffer->iface.free_buffer == ggml_backend_rpc_buffer_free_buffer; } -static rpc_tensor serialize_tensor(const ggml_tensor * tensor) { +static rpc_tensor serialize_tensor(const ggml_tensor * tensor, const std::shared_ptr & dispatcher = nullptr) { rpc_tensor result; if (!tensor) { memset(&result, 0, sizeof(result)); @@ -637,8 +637,14 @@ static rpc_tensor serialize_tensor(const ggml_tensor * tensor) { if (tensor->buffer && ggml_backend_buffer_is_rpc(tensor->buffer)) { ggml_backend_buffer_t buffer = tensor->buffer; ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context; - result.buffer = ctx != nullptr ? ctx->remote_ptr : 0; - result.data = reinterpret_cast(tensor->data); + // ref: https://github.com/ggml-org/llama.cpp/pull/26500 + if (ctx != nullptr && (dispatcher == nullptr || ctx->dispatcher == dispatcher)) { + result.buffer = ctx->remote_ptr; + result.data = reinterpret_cast(tensor->data); + } else { + result.buffer = 0; + result.data = 0; + } } else { result.buffer = 0; result.data = 0; @@ -958,7 +964,7 @@ static void ggml_backend_rpc_synchronize(ggml_backend_t backend) { rpc_ctx->dispatcher->synchronize(); } -static void add_tensor(ggml_tensor * tensor, const ggml_cgraph * cgraph, std::vector & tensors, std::unordered_set & visited) { +static void add_tensor(ggml_tensor * tensor, const ggml_cgraph * cgraph, const std::shared_ptr & dispatcher, std::vector & tensors, std::unordered_set & visited) { if (tensor == nullptr) { return; } @@ -967,10 +973,10 @@ static void add_tensor(ggml_tensor * tensor, const ggml_cgraph * cgraph, std::ve } visited.insert(tensor); for (int i = 0; i < GGML_MAX_SRC; i++) { - add_tensor(tensor->src[i], cgraph, tensors, visited); + add_tensor(tensor->src[i], cgraph, dispatcher, tensors, visited); } - add_tensor(tensor->view_src, cgraph, tensors, visited); - rpc_tensor result = serialize_tensor(tensor); + add_tensor(tensor->view_src, cgraph, dispatcher, tensors, visited); + rpc_tensor result = serialize_tensor(tensor, dispatcher); const size_t hash_pos = ggml_hash_find(&cgraph->visited_hash_set, tensor); if (hash_pos != GGML_HASHSET_FULL && ggml_bitset_get(cgraph->visited_hash_set.used, hash_pos)) { result.use_count = cgraph->use_counts[hash_pos]; @@ -978,12 +984,12 @@ static void add_tensor(ggml_tensor * tensor, const ggml_cgraph * cgraph, std::ve tensors.push_back(result); } -static uint8_t * serialize_graph(uint32_t device, const ggml_cgraph * cgraph, size_t * output_size) { +static uint8_t * serialize_graph(uint32_t device, const ggml_cgraph * cgraph, const std::shared_ptr & dispatcher, size_t * output_size) { uint32_t n_nodes = cgraph->n_nodes; std::vector tensors; std::unordered_set visited; for (uint32_t i = 0; i < n_nodes; i++) { - add_tensor(cgraph->nodes[i], cgraph, tensors, visited); + add_tensor(cgraph->nodes[i], cgraph, dispatcher, tensors, visited); } // serialization format: // | device (4 bytes) | n_nodes (4 bytes) | nodes (n_nodes * sizeof(uint64_t) | n_tensors (4 bytes) | tensors (n_tensors * sizeof(rpc_tensor)) | @@ -1020,7 +1026,7 @@ static enum ggml_status ggml_backend_rpc_graph_compute(ggml_backend_t backend, g } else { rpc_dev_ctx->last_graph_uid = cgraph->uid; size_t input_size = 0; - uint8_t * input = serialize_graph(rpc_ctx->device, cgraph, &input_size); + uint8_t * input = serialize_graph(rpc_ctx->device, cgraph, rpc_ctx->dispatcher, &input_size); std::shared_ptr input_ptr(input, std::default_delete()); rpc_ctx->dispatcher->send_async(RPC_CMD_GRAPH_COMPUTE, input_ptr, input_size); } diff --git a/tests/test-rpc-multi-server.cpp b/tests/test-rpc-multi-server.cpp new file mode 100644 index 000000000..4502e2ce7 --- /dev/null +++ b/tests/test-rpc-multi-server.cpp @@ -0,0 +1,47 @@ +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml-impl.h" +#include "ggml-rpc.h" +#include "ggml.h" + +int main(int argc, char ** argv) { + GGML_ASSERT(argc == 3); + ggml_backend_load_all(); + + const char * endpoint_a = argv[1]; + const char * endpoint_b = argv[2]; + + ggml_backend_t backend_a = ggml_backend_rpc_init(endpoint_a, 0); + ggml_backend_t backend_b = ggml_backend_rpc_init(endpoint_b, 0); + GGML_ASSERT(backend_a != nullptr); + GGML_ASSERT(backend_b != nullptr); + + ggml_init_params params = { + /* .mem_size = */ ggml_tensor_overhead() + ggml_graph_overhead_custom(1, false), + /* .mem_buffer = */ nullptr, + /* .no_alloc = */ true, + }; + ggml_context * ctx = ggml_init(params); + GGML_ASSERT(ctx != nullptr); + + ggml_tensor * tensor = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend_a); + GGML_ASSERT(buffer != nullptr); + + // A remote pointer allocated by server A is not meaningful to server B. + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 1, false); + graph->nodes[0] = tensor; + graph->n_nodes = 1; + + GGML_ASSERT(ggml_backend_graph_compute(backend_b, graph) == GGML_STATUS_SUCCESS); + // Wait for server B to finish the graph before the script checks its log. + size_t free_mem; + size_t total_mem; + ggml_backend_rpc_get_device_memory(endpoint_b, 0, &free_mem, &total_mem); + GGML_ASSERT(total_mem > 0); + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend_b); + ggml_backend_free(backend_a); + return 0; +} diff --git a/tests/test-rpc-multi-server.sh b/tests/test-rpc-multi-server.sh new file mode 100755 index 000000000..a8c723162 --- /dev/null +++ b/tests/test-rpc-multi-server.sh @@ -0,0 +1,43 @@ +#!/usr/bin/env bash +set -euo pipefail + +server=$1 +client=$2 +port_a=$((40000 + $$ % 10000)) +port_b=$((port_a + 1)) +endpoint_a="127.0.0.1:${port_a}" +endpoint_b="127.0.0.1:${port_b}" +test_dir=$(mktemp -d) + +cleanup() { + kill "${pid_a:-}" "${pid_b:-}" 2>/dev/null || true + rm -rf "$test_dir" +} +trap cleanup EXIT + +wait_for_port() { + local port=$1 + for _ in {1..600}; do + if (exec 3<>"/dev/tcp/127.0.0.1/$port") 2>/dev/null; then + exec 3>&- + exec 3<&- + return 0 + fi + sleep 0.05 + done + return 1 +} + +"$server" --device CPU --host 127.0.0.1 --port "$port_a" >"$test_dir/server-a.log" 2>&1 & +pid_a=$! +"$server" --device CPU --host 127.0.0.1 --port "$port_b" >"$test_dir/server-b.log" 2>&1 & +pid_b=$! +wait_for_port "$port_a" +wait_for_port "$port_b" + +"$client" "$endpoint_a" "$endpoint_b" + +if grep -q "invalid data ptr" "$test_dir/server-b.log"; then + cat "$test_dir/server-b.log" + exit 1 +fi diff --git a/tools/rpc/CMakeLists.txt b/tools/rpc/CMakeLists.txt index 0eee9a922..2891c7d03 100644 --- a/tools/rpc/CMakeLists.txt +++ b/tools/rpc/CMakeLists.txt @@ -3,6 +3,18 @@ add_executable(${TARGET} rpc-server.cpp) target_link_libraries(${TARGET} PRIVATE ggml) target_compile_features(${TARGET} PRIVATE cxx_std_17) +if (LLAMA_BUILD_TESTS AND UNIX AND NOT GGML_BACKEND_DL) + add_executable(test-rpc-multi-server ${PROJECT_SOURCE_DIR}/tests/test-rpc-multi-server.cpp) + target_link_libraries(test-rpc-multi-server PRIVATE ggml ggml-rpc) + target_include_directories(test-rpc-multi-server PRIVATE ${PROJECT_SOURCE_DIR}/ggml/src) + add_test( + NAME test-rpc-multi-server + COMMAND bash ${PROJECT_SOURCE_DIR}/tests/test-rpc-multi-server.sh + $ + $) + set_property(TEST test-rpc-multi-server PROPERTY LABELS main) +endif() + if(LLAMA_TOOLS_INSTALL) install(TARGETS ${TARGET} RUNTIME) endif()