Something went wrong. Try again.
My dotfiles for my nixos machines and infra
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215From f0c483c4df52f82cc5795433c9e9332fb3e8aa21 Mon Sep 17 00:00:00 2001From: hmirin <hmirin@users.noreply.github.com>Date: Sun, 2 Aug 2026 12:55:10 +0900Subject: [PATCH] rpc: avoid serializing buffers from other servers
Only include remote buffer pointers when the buffer belongs to the RPC socket receiving the graph. Add a two-server regression test for cross-server tensor serialization.
Assisted-by: Codex--- ggml/src/ggml-rpc/ggml-rpc.cpp | 25 +++++++++++------- tests/test-rpc-multi-server.cpp | 47 +++++++++++++++++++++++++++++++++ tests/test-rpc-multi-server.sh | 43 ++++++++++++++++++++++++++++++ tools/rpc/CMakeLists.txt | 12 +++++++++ 4 files changed, 117 insertions(+), 10 deletions(-) create mode 100644 tests/test-rpc-multi-server.cpp create mode 100644 tests/test-rpc-multi-server.sh
diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cppindex 17c53a5f049e..894c2cd0ac6b 100644--- a/ggml/src/ggml-rpc/ggml-rpc.cpp+++ b/ggml/src/ggml-rpc/ggml-rpc.cpp@@ -412,7 +412,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 socket_ptr & sock = nullptr) { rpc_tensor result; if (!tensor) { memset(&result, 0, sizeof(result));@@ -424,8 +424,13 @@ 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<uint64_t>(tensor->data);+ if (ctx != nullptr && (sock == nullptr || ctx->sock == sock)) {+ result.buffer = ctx->remote_ptr;+ result.data = reinterpret_cast<uint64_t>(tensor->data);+ } else {+ result.buffer = 0;+ result.data = 0;+ } } else { result.buffer = 0; result.data = 0;@@ -675,7 +680,7 @@ static void ggml_backend_rpc_synchronize(ggml_backend_t backend) { // this is no-op because we don't have any async operations } -static void add_tensor(ggml_tensor * tensor, std::vector<rpc_tensor> & tensors, std::unordered_set<ggml_tensor*> & visited) {+static void add_tensor(ggml_tensor * tensor, const socket_ptr & sock, std::vector<rpc_tensor> & tensors, std::unordered_set<ggml_tensor*> & visited) { if (tensor == nullptr) { return; }@@ -684,18 +689,18 @@ static void add_tensor(ggml_tensor * tensor, std::vector<rpc_tensor> & tensors, } visited.insert(tensor); for (int i = 0; i < GGML_MAX_SRC; i++) {- add_tensor(tensor->src[i], tensors, visited);+ add_tensor(tensor->src[i], sock, tensors, visited); }- add_tensor(tensor->view_src, tensors, visited);- tensors.push_back(serialize_tensor(tensor));+ add_tensor(tensor->view_src, sock, tensors, visited);+ tensors.push_back(serialize_tensor(tensor, sock)); } -static void serialize_graph(uint32_t device, const ggml_cgraph * cgraph, std::vector<uint8_t> & output) {+static void serialize_graph(uint32_t device, const ggml_cgraph * cgraph, const socket_ptr & sock, std::vector<uint8_t> & output) { uint32_t n_nodes = cgraph->n_nodes; std::vector<rpc_tensor> tensors; std::unordered_set<ggml_tensor*> visited; for (uint32_t i = 0; i < n_nodes; i++) {- add_tensor(cgraph->nodes[i], tensors, visited);+ add_tensor(cgraph->nodes[i], sock, 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)) |@@ -733,8 +738,8 @@ static enum ggml_status ggml_backend_rpc_graph_compute(ggml_backend_t backend, g } else { rpc_dev_ctx->last_graph_uid = cgraph->uid; std::vector<uint8_t> input;- serialize_graph(rpc_ctx->device, cgraph, input); auto sock = get_socket(rpc_ctx->endpoint);+ serialize_graph(rpc_ctx->device, cgraph, sock, input); bool status = send_rpc_cmd(sock, RPC_CMD_GRAPH_COMPUTE, input.data(), input.size()); RPC_STATUS_ASSERT(status); }diff --git a/tests/test-rpc-multi-server.cpp b/tests/test-rpc-multi-server.cppnew file mode 100644index 000000000000..4502e2ce71fa--- /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.shnew file mode 100644index 000000000000..a8c72316202a--- /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+fidiff --git a/tools/rpc/CMakeLists.txt b/tools/rpc/CMakeLists.txtindex 0eee9a922e77..2891c7d034cf 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+ $<TARGET_FILE:ggml-rpc-server>+ $<TARGET_FILE:test-rpc-multi-server>)+ set_property(TEST test-rpc-multi-server PROPERTY LABELS main)+endif()+ if(LLAMA_TOOLS_INSTALL) install(TARGETS ${TARGET} RUNTIME) endif()