rpc: avoid serializing buffers from other servers

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<rpc_dispatcher> & 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<uint64_t>(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<uint64_t>(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<rpc_tensor> & tensors, std::unordered_set<ggml_tensor*> & visited) {
+static void add_tensor(ggml_tensor * tensor, const ggml_cgraph * cgraph, const std::shared_ptr<rpc_dispatcher> & dispatcher, std::vector<rpc_tensor> & tensors, std::unordered_set<ggml_tensor*> & visited) {
     if (tensor == nullptr) {
         return;
     }
@@ -967,23 +973,23 @@ 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];
     }
     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<rpc_dispatcher> & dispatcher, size_t * output_size) {
     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], 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<uint8_t> input_ptr(input, std::default_delete<uint8_t[]>());
         rpc_ctx->dispatcher->send_async(RPC_CMD_GRAPH_COMPUTE, input_ptr, input_size);
     }

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;
+}

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

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()