From 1614b745e41ebebe1dc43dfb8139d503b658bee4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 18:57:08 -0700 Subject: [PATCH 1/8] rpc: move a tensor between two servers without passing through the client A layer split over several RPC devices moved every hidden state through the host that runs the scheduler: ggml_backend_sched has no direct path between two RPC buffers on different endpoints, so it fell back to reading the tensor into the client's memory and writing it out again, two transfers and a synchronize per stage boundary. RPC_CMD_COPY_TENSOR_TO tells the server that holds the source tensor to write it into a tensor on another server. The source server opens a connection to the destination with the same HELLO negotiation a client uses (RDMA when both rails allow it, TCP otherwise), pushes the data as an ordinary RPC_CMD_SET_TENSOR so the destination applies its own tensor deserialization and buffer range checks, and waits for RPC_CMD_PEER_BARRIER before answering the client, so the destination cannot compute before the write has landed. Connections to other servers are cached per destination endpoint and closed when the client disconnects. Serving several connections at once is what this needs, so a server now runs one thread per connection over a shared buffer registry and a shared execution mutex; a session still owns and frees only the buffers it allocated. The client uses the command from the RPC backend's cpy_tensor_async when the source and the destination are RPC buffers on different endpoints and both servers report protocol minor 3 or higher; everything else, including two devices of one server, keeps its previous path. GGML_RPC_P2P=0 forces the old path. Protocol minor 2 -> 3, every existing command unchanged. --- ggml/include/ggml-rpc.h | 2 +- ggml/src/ggml-rpc/ggml-rpc.cpp | 325 ++++++++++++++++++++++++++++++-- ggml/src/ggml-rpc/transport.cpp | 2 +- 3 files changed, 310 insertions(+), 19 deletions(-) diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h index 5d4030128246..dc51b73e940e 100644 --- a/ggml/include/ggml-rpc.h +++ b/ggml/include/ggml-rpc.h @@ -7,7 +7,7 @@ extern "C" { #endif #define RPC_PROTO_MAJOR_VERSION 5 -#define RPC_PROTO_MINOR_VERSION 2 +#define RPC_PROTO_MINOR_VERSION 3 #define RPC_PROTO_PATCH_VERSION 0 #ifdef __cplusplus diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index 66bd1afb0041..03ad0f0e2846 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -19,6 +19,7 @@ #include #include #include +#include static const char * RPC_DEBUG = std::getenv("GGML_RPC_DEBUG"); @@ -75,6 +76,8 @@ enum rpc_cmd { RPC_CMD_GRAPH_RECOMPUTE, RPC_CMD_MEMSET_TENSOR, RPC_CMD_GET_TENSORS, + RPC_CMD_COPY_TENSOR_TO, + RPC_CMD_PEER_BARRIER, RPC_CMD_COUNT, }; @@ -197,6 +200,29 @@ struct rpc_msg_copy_tensor_rsp { uint8_t result; }; +// RPC_CMD_COPY_TENSOR_TO tells the server that holds `src` to write it into `dst` on another +// server, without the data passing through the client. The request is +// | rpc_msg_copy_tensor_to_hdr | endpoint_len bytes of the destination endpoint |, and the +// response arrives once the destination has acknowledged the write. +struct rpc_msg_copy_tensor_to_hdr { + rpc_tensor src; + rpc_tensor dst; + uint64_t size; + uint32_t endpoint_len; +}; + +struct rpc_msg_copy_tensor_to_rsp { + uint8_t result; +}; + +// RPC_CMD_PEER_BARRIER has an empty request and a one byte response. A connection is served +// strictly in order, so receiving its response proves that every command sent earlier on the +// same connection has finished. It is what a source server waits on after pushing a +// RPC_CMD_SET_TENSOR to a destination server. +struct rpc_msg_peer_barrier_rsp { + uint8_t result; +}; + struct rpc_msg_get_device_memory_req { uint32_t device; }; @@ -337,6 +363,8 @@ static const char * rpc_cmd_name(int cmd) { case RPC_CMD_GRAPH_RECOMPUTE: return "GRAPH_RECOMPUTE"; case RPC_CMD_MEMSET_TENSOR: return "MEMSET_TENSOR"; case RPC_CMD_GET_TENSORS: return "GET_TENSORS"; + case RPC_CMD_COPY_TENSOR_TO: return "COPY_TENSOR_TO"; + case RPC_CMD_PEER_BARRIER: return "PEER_BARRIER"; default: return "?"; } } @@ -1025,6 +1053,67 @@ static bool rpc_supports_batched_get(const socket_ptr & sock) { return !disabled && sock->conn.server_minor >= 2; } +// true when the server understands RPC_CMD_COPY_TENSOR_TO and RPC_CMD_PEER_BARRIER. +// GGML_RPC_P2P=0 forces the old path, for A/B measurements. +static bool rpc_supports_p2p(const socket_ptr & sock) { + static const char * env = std::getenv("GGML_RPC_P2P"); + static const bool disabled = env != nullptr && std::strcmp(env, "0") == 0; + return !disabled && sock != nullptr && sock->conn.server_minor >= 3; +} + +// Server to server movement of a tensor that crosses a stage boundary of a layer split. The +// client tells the server that holds the source to write it into the destination tensor on +// another server; only the command and its acknowledgement cross the client's host, not the +// data. Two servers of one endpoint keep using the server local RPC_CMD_COPY_TENSOR. +static bool ggml_backend_rpc_cpy_tensor_p2p(ggml_backend_t backend_src, ggml_backend_t backend_dst, + const ggml_tensor * src, ggml_tensor * dst) { + ggml_backend_rpc_context * src_ctx = (ggml_backend_rpc_context *) backend_src->context; + ggml_backend_rpc_context * dst_ctx = (ggml_backend_rpc_context *) backend_dst->context; + + if (src_ctx->endpoint == dst_ctx->endpoint) { + // same server: ggml_backend_rpc_buffer_cpy_tensor does it without leaving the server + return false; + } + if (src->buffer == nullptr || dst->buffer == nullptr || + !ggml_backend_buffer_is_rpc(src->buffer) || !ggml_backend_buffer_is_rpc(dst->buffer)) { + return false; + } + const uint64_t size = (uint64_t) ggml_nbytes(src); + if (size != (uint64_t) ggml_nbytes(dst)) { + return false; + } + + auto src_sock = get_socket(src_ctx->endpoint); + auto dst_sock = get_socket(dst_ctx->endpoint); + if (!rpc_supports_p2p(src_sock) || !rpc_supports_p2p(dst_sock)) { + return false; + } + + // anything queued for the destination has to be on the wire before the peer write lands + rpc_flush_deferred_guarded(dst_sock); + + const std::string & endpoint = dst_ctx->endpoint; + rpc_msg_copy_tensor_to_hdr hdr; + hdr.src = serialize_tensor(src); + hdr.dst = serialize_tensor(dst); + hdr.size = size; + hdr.endpoint_len = (uint32_t) endpoint.size(); + + std::vector input(sizeof(hdr) + endpoint.size()); + memcpy(input.data(), &hdr, sizeof(hdr)); + memcpy(input.data() + sizeof(hdr), endpoint.data(), endpoint.size()); + + // The reply arrives once the destination server has acknowledged the write, so the + // destination cannot compute before the hidden state has landed: the next thing the + // scheduler does on that backend is RPC_CMD_GRAPH_COMPUTE, which is sent after this + // returns. + rpc_msg_copy_tensor_to_rsp response; + bool status = send_rpc_cmd(src_sock, RPC_CMD_COPY_TENSOR_TO, input.data(), input.size(), + &response, sizeof(response)); + RPC_STATUS_ASSERT(status); + return response.result != 0; +} + static void ggml_backend_rpc_get_tensor_async(ggml_backend_t backend, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { ggml_backend_rpc_context * rpc_ctx = (ggml_backend_rpc_context *)backend->context; auto sock = get_socket(rpc_ctx->endpoint); @@ -1059,8 +1148,10 @@ static bool ggml_backend_rpc_cpy_tensor_async(ggml_backend_t backend_src, ggml_b const bool src_is_rpc = ggml_backend_is_rpc(backend_src); const bool dst_is_rpc = ggml_backend_is_rpc(backend_dst); - // RPC to RPC on the same server is handled by the buffer level copy - if (src_is_rpc == dst_is_rpc) { + if (src_is_rpc && dst_is_rpc) { + return ggml_backend_rpc_cpy_tensor_p2p(backend_src, backend_dst, src, dst); + } + if (!src_is_rpc && !dst_is_rpc) { return false; } @@ -1360,10 +1451,22 @@ void ggml_backend_rpc_get_device_memory(const char * endpoint, uint32_t device, // RPC server-side implementation +// State that every connection of one server process shares. +// +// Before peer to peer copies a server served one client at a time, so a session could own its +// buffers privately. A source server now opens a second connection to the destination server +// and writes into a buffer that the coordinator allocated over its own connection, so the live +// buffers have to be visible to every connection, and command execution has to be serialised +// across connections (two connections must never drive the same backend at the same time). +struct rpc_server_shared { + std::mutex mtx; // serialises command execution + std::unordered_set buffers; // every live buffer of this process +}; + class rpc_server { public: - rpc_server(std::vector all_backends, const char * cache_dir) - : backends(std::move(all_backends)), cache_dir(cache_dir) { + rpc_server(std::vector all_backends, const char * cache_dir, rpc_server_shared & shared) + : backends(std::move(all_backends)), cache_dir(cache_dir), shared(shared) { stored_graphs.resize(backends.size()); } ~rpc_server(); @@ -1381,6 +1484,7 @@ class rpc_server { bool get_tensor(const rpc_msg_get_tensor_req & request, std::vector & response); bool get_tensors(const std::vector & input, std::vector & response); bool copy_tensor(const rpc_msg_copy_tensor_req & request, rpc_msg_copy_tensor_rsp & response); + bool copy_tensor_to(const std::vector & input, rpc_msg_copy_tensor_to_rsp & response); bool graph_compute(const std::vector & input); bool graph_recompute(const rpc_msg_graph_recompute_req & request); bool init_tensor(const rpc_msg_init_tensor_req & request); @@ -1401,9 +1505,17 @@ class rpc_server { std::unordered_map & tensor_map); + socket_ptr get_peer_socket(const std::string & endpoint); + std::vector backends; const char * cache_dir; - std::unordered_set buffers; + rpc_server_shared & shared; + // buffers allocated over this connection; freed when it closes + std::unordered_set owned_buffers; + // connections to other servers, one per destination endpoint, closed with this connection + std::unordered_map peer_socks; + // reused staging for RPC_CMD_COPY_TENSOR_TO + std::vector p2p_buf; // store the last computed graph for each backend std::vector stored_graphs; }; @@ -1416,6 +1528,8 @@ void rpc_server::hello(rpc_msg_hello_rsp & response) { } bool rpc_server::get_alloc_size(const rpc_msg_get_alloc_size_req & request, rpc_msg_get_alloc_size_rsp & response) { + std::lock_guard lock(shared.mtx); + uint32_t dev_id = request.device; if (dev_id >= backends.size()) { return false; @@ -1456,6 +1570,8 @@ bool rpc_server::get_alloc_size(const rpc_msg_get_alloc_size_req & request, rpc_ } bool rpc_server::alloc_buffer(const rpc_msg_alloc_buffer_req & request, rpc_msg_alloc_buffer_rsp & response) { + std::lock_guard lock(shared.mtx); + uint32_t dev_id = request.device; if (dev_id >= backends.size()) { return false; @@ -1469,7 +1585,8 @@ bool rpc_server::alloc_buffer(const rpc_msg_alloc_buffer_req & request, rpc_msg_ response.remote_size = buffer->size; LOG_DBG("[%s] device: %d, size: %" PRIu64 " -> remote_ptr: %" PRIx64 ", remote_size: %" PRIu64 "\n", __func__, dev_id, request.size, response.remote_ptr, response.remote_size); - buffers.insert(buffer); + shared.buffers.insert(buffer); + owned_buffers.insert(buffer); } else { LOG_DBG("[%s] device: %d, size: %" PRIu64 " -> failed\n", __func__, dev_id, request.size); } @@ -1477,6 +1594,8 @@ bool rpc_server::alloc_buffer(const rpc_msg_alloc_buffer_req & request, rpc_msg_ } bool rpc_server::get_alignment(const rpc_msg_get_alignment_req & request, rpc_msg_get_alignment_rsp & response) { + std::lock_guard lock(shared.mtx); + uint32_t dev_id = request.device; if (dev_id >= backends.size()) { return false; @@ -1489,6 +1608,8 @@ bool rpc_server::get_alignment(const rpc_msg_get_alignment_req & request, rpc_ms } bool rpc_server::get_max_size(const rpc_msg_get_max_size_req & request, rpc_msg_get_max_size_rsp & response) { + std::lock_guard lock(shared.mtx); + uint32_t dev_id = request.device; if (dev_id >= backends.size()) { return false; @@ -1501,9 +1622,11 @@ bool rpc_server::get_max_size(const rpc_msg_get_max_size_req & request, rpc_msg_ } bool rpc_server::buffer_get_base(const rpc_msg_buffer_get_base_req & request, rpc_msg_buffer_get_base_rsp & response) { + std::lock_guard lock(shared.mtx); + LOG_DBG("[%s] remote_ptr: %" PRIx64 "\n", __func__, request.remote_ptr); ggml_backend_buffer_t buffer = reinterpret_cast(request.remote_ptr); - if (buffers.find(buffer) == buffers.end()) { + if (shared.buffers.find(buffer) == shared.buffers.end()) { GGML_LOG_ERROR("[%s] buffer not found\n", __func__); return false; } @@ -1513,21 +1636,26 @@ bool rpc_server::buffer_get_base(const rpc_msg_buffer_get_base_req & request, rp } bool rpc_server::free_buffer(const rpc_msg_free_buffer_req & request) { + std::lock_guard lock(shared.mtx); + LOG_DBG("[%s] remote_ptr: %" PRIx64 "\n", __func__, request.remote_ptr); ggml_backend_buffer_t buffer = reinterpret_cast(request.remote_ptr); - if (buffers.find(buffer) == buffers.end()) { + if (owned_buffers.find(buffer) == owned_buffers.end()) { GGML_LOG_ERROR("[%s] buffer not found\n", __func__); return false; } ggml_backend_buffer_free(buffer); - buffers.erase(buffer); + owned_buffers.erase(buffer); + shared.buffers.erase(buffer); return true; } bool rpc_server::buffer_clear(const rpc_msg_buffer_clear_req & request) { + std::lock_guard lock(shared.mtx); + LOG_DBG("[%s] remote_ptr: %" PRIx64 ", value: %u\n", __func__, request.remote_ptr, request.value); ggml_backend_buffer_t buffer = reinterpret_cast(request.remote_ptr); - if (buffers.find(buffer) == buffers.end()) { + if (shared.buffers.find(buffer) == shared.buffers.end()) { GGML_LOG_ERROR("[%s] buffer not found\n", __func__); return false; } @@ -1536,6 +1664,8 @@ bool rpc_server::buffer_clear(const rpc_msg_buffer_clear_req & request) { } bool rpc_server::memset_tensor(const rpc_msg_memset_tensor_req & request) { + std::lock_guard lock(shared.mtx); + struct ggml_init_params params { /*.mem_size =*/ ggml_tensor_overhead(), /*.mem_buffer =*/ NULL, @@ -1607,7 +1737,7 @@ ggml_tensor * rpc_server::deserialize_tensor(struct ggml_context * ctx, const rp result->nb[i] = tensor->nb[i]; } result->buffer = reinterpret_cast(tensor->buffer); - if (result->buffer && buffers.find(result->buffer) == buffers.end()) { + if (result->buffer && shared.buffers.find(result->buffer) == shared.buffers.end()) { result->buffer = nullptr; } @@ -1632,6 +1762,8 @@ ggml_tensor * rpc_server::deserialize_tensor(struct ggml_context * ctx, const rp bool rpc_server::set_tensor(const std::vector & input) { + std::lock_guard lock(shared.mtx); + // serialization format: | rpc_tensor | offset (8 bytes) | data (size bytes) | if (input.size() < sizeof(rpc_tensor) + sizeof(uint64_t)) { return false; @@ -1705,6 +1837,8 @@ bool rpc_server::get_cached_file(uint64_t hash, std::vector & data) { bool rpc_server::set_tensor_hash(const rpc_msg_set_tensor_hash_req & request, rpc_msg_set_tensor_hash_rsp & response) { + std::lock_guard lock(shared.mtx); + std::vector cached_file; if (!get_cached_file(request.hash, cached_file)) { response.result = 0; @@ -1746,6 +1880,8 @@ bool rpc_server::set_tensor_hash(const rpc_msg_set_tensor_hash_req & request, rp } bool rpc_server::init_tensor(const rpc_msg_init_tensor_req & request) { + std::lock_guard lock(shared.mtx); + struct ggml_init_params params { /*.mem_size =*/ ggml_tensor_overhead(), /*.mem_buffer =*/ NULL, @@ -1781,6 +1917,8 @@ bool rpc_server::init_tensor(const rpc_msg_init_tensor_req & request) { } bool rpc_server::get_tensor(const rpc_msg_get_tensor_req & request, std::vector & response) { + std::lock_guard lock(shared.mtx); + struct ggml_init_params params { /*.mem_size =*/ ggml_tensor_overhead(), /*.mem_buffer =*/ NULL, @@ -1820,6 +1958,8 @@ bool rpc_server::get_tensor(const rpc_msg_get_tensor_req & request, std::vector< // request order, so one decode step of a backend sampled batch is one round trip instead of one // per sequence and per sampler output. bool rpc_server::get_tensors(const std::vector & input, std::vector & response) { + std::lock_guard lock(shared.mtx); + if (input.size() < sizeof(uint32_t)) { return false; } @@ -1871,6 +2011,8 @@ bool rpc_server::get_tensors(const std::vector & input, std::vector lock(shared.mtx); + struct ggml_init_params params { /*.mem_size =*/ 2*ggml_tensor_overhead(), /*.mem_buffer =*/ NULL, @@ -1911,6 +2053,111 @@ bool rpc_server::copy_tensor(const rpc_msg_copy_tensor_req & request, rpc_msg_co return true; } +socket_ptr rpc_server::get_peer_socket(const std::string & endpoint) { + auto it = peer_socks.find(endpoint); + if (it != peer_socks.end()) { + return it->second; + } + // the same connect and HELLO negotiation a client does, so a server to server link uses + // RDMA whenever both rails allow it and TCP otherwise + auto sock = get_socket(endpoint); + if (sock == nullptr) { + GGML_LOG_ERROR("[%s] failed to connect to %s\n", __func__, endpoint.c_str()); + return nullptr; + } + if (sock->conn.server_minor < 3) { + GGML_LOG_ERROR("[%s] %s does not support peer to peer copies\n", __func__, endpoint.c_str()); + return nullptr; + } + peer_socks[endpoint] = sock; + return sock; +} + +// Writes a tensor of this server into a tensor of another server. The data never reaches the +// client that issued the command. The source region is validated exactly like RPC_CMD_GET_TENSOR +// and is pushed to the destination as an ordinary RPC_CMD_SET_TENSOR, so the destination applies +// its own deserialize_tensor and buffer range checks and nothing about the destination pointer is +// trusted here. RPC_CMD_PEER_BARRIER then tells us that the write has landed. +// +// Recoverable problems (the destination is unreachable or too old, the write was rejected) are +// reported as result = 0, which puts the client back on its previous path; only a malformed +// request or an out of bounds source is a protocol error. +bool rpc_server::copy_tensor_to(const std::vector & input, rpc_msg_copy_tensor_to_rsp & response) { + response.result = 0; + if (input.size() < sizeof(rpc_msg_copy_tensor_to_hdr)) { + return false; + } + rpc_msg_copy_tensor_to_hdr hdr; + memcpy(&hdr, input.data(), sizeof(hdr)); + if (input.size() != sizeof(hdr) + (size_t) hdr.endpoint_len) { + return false; + } + const std::string endpoint((const char *) input.data() + sizeof(hdr), hdr.endpoint_len); + + const size_t msg_size = sizeof(rpc_tensor) + sizeof(uint64_t) + (size_t) hdr.size; + if (p2p_buf.size() < msg_size) { + p2p_buf.resize(msg_size); + } + const uint64_t dst_offset = 0; + memcpy(p2p_buf.data(), &hdr.dst, sizeof(rpc_tensor)); + memcpy(p2p_buf.data() + sizeof(rpc_tensor), &dst_offset, sizeof(dst_offset)); + + // read the source out of this server's device; the peer connection is used without this + // lock, so a ring of servers cannot deadlock on each other's execution mutex + { + std::lock_guard lock(shared.mtx); + + struct ggml_init_params params { + /*.mem_size =*/ ggml_tensor_overhead(), + /*.mem_buffer =*/ NULL, + /*.no_alloc =*/ true, + }; + ggml_context_ptr ctx_ptr { ggml_init(params) }; + GGML_ASSERT(ctx_ptr != nullptr); + ggml_tensor * tensor = deserialize_tensor(ctx_ptr.get(), &hdr.src); + if (tensor == nullptr || tensor->buffer == nullptr) { + GGML_LOG_ERROR("[%s] error deserializing tensor\n", __func__); + return false; + } + + // sanitize tensor->data + const size_t p0 = (size_t) ggml_backend_buffer_get_base(tensor->buffer); + const size_t p1 = p0 + ggml_backend_buffer_get_size(tensor->buffer); + if (hdr.src.data < p0 || hdr.src.data >= p1 || hdr.size > (p1 - hdr.src.data)) { + GGML_LOG_ERROR("[%s] source region (data=0x%" PRIx64 ", size=%" PRIu64 ") out of buffer bounds [0x%zx, 0x%zx)\n", + __func__, hdr.src.data, hdr.size, p0, p1); + return false; + } + if (hdr.size > (uint64_t) ggml_nbytes(tensor)) { + GGML_LOG_ERROR("[%s] source region larger than the tensor\n", __func__); + return false; + } + + ggml_backend_tensor_get(tensor, p2p_buf.data() + sizeof(rpc_tensor) + sizeof(dst_offset), 0, hdr.size); + } + + socket_ptr peer = get_peer_socket(endpoint); + if (peer == nullptr) { + return true; + } + + LOG_DBG("[%s] %" PRIu64 " bytes to %s\n", __func__, hdr.size, endpoint.c_str()); + + if (!send_rpc_cmd(peer, RPC_CMD_SET_TENSOR, p2p_buf.data(), msg_size)) { + GGML_LOG_ERROR("[%s] failed to send to %s\n", __func__, endpoint.c_str()); + peer_socks.erase(endpoint); + return true; + } + rpc_msg_peer_barrier_rsp barrier; + if (!send_rpc_cmd(peer, RPC_CMD_PEER_BARRIER, nullptr, 0, &barrier, sizeof(barrier))) { + GGML_LOG_ERROR("[%s] %s did not acknowledge the write\n", __func__, endpoint.c_str()); + peer_socks.erase(endpoint); + return true; + } + response.result = barrier.result; + return true; +} + ggml_tensor * rpc_server::create_node(uint64_t id, struct ggml_context * ctx, const std::unordered_map & tensor_ptrs, @@ -1968,6 +2215,8 @@ ggml_tensor * rpc_server::create_node(uint64_t id, } bool rpc_server::graph_compute(const std::vector & input) { + std::lock_guard lock(shared.mtx); + // 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)) | if (input.size() < 2*sizeof(uint32_t)) { @@ -2042,6 +2291,8 @@ bool rpc_server::graph_compute(const std::vector & input) { } bool rpc_server::graph_recompute(const rpc_msg_graph_recompute_req & request) { + std::lock_guard lock(shared.mtx); + uint32_t device = request.device; if (device >= backends.size()) { return false; @@ -2057,6 +2308,8 @@ bool rpc_server::graph_recompute(const rpc_msg_graph_recompute_req & request) { } bool rpc_server::get_device_memory(const rpc_msg_get_device_memory_req & request, rpc_msg_get_device_memory_rsp & response) { + std::lock_guard lock(shared.mtx); + uint32_t dev_id = request.device; if (dev_id >= backends.size()) { return false; @@ -2071,14 +2324,18 @@ bool rpc_server::get_device_memory(const rpc_msg_get_device_memory_req & request } rpc_server::~rpc_server() { - for (auto buffer : buffers) { + // the connections this server opened to other servers go away with this connection + peer_socks.clear(); + std::lock_guard lock(shared.mtx); + for (auto buffer : owned_buffers) { + shared.buffers.erase(buffer); ggml_backend_buffer_free(buffer); } } static void rpc_serve_client(const std::vector & backends, const char * cache_dir, - socket_ptr sock) { - rpc_server server(backends, cache_dir); + socket_ptr sock, rpc_server_shared & shared) { + rpc_server server(backends, cache_dir, shared); uint8_t cmd; if (!sock->recv_data(&cmd, 1)) { return; @@ -2314,6 +2571,32 @@ static void rpc_serve_client(const std::vector & backends, const } break; } + case RPC_CMD_COPY_TENSOR_TO: { + std::vector input; + if (!recv_msg(sock, input)) { + return; + } + rpc_msg_copy_tensor_to_rsp response; + if (!server.copy_tensor_to(input, response)) { + return; + } + if (!send_msg(sock, &response, sizeof(response))) { + return; + } + break; + } + case RPC_CMD_PEER_BARRIER: { + if (!recv_msg(sock, nullptr, 0)) { + return; + } + // every command received earlier on this connection has already been served + rpc_msg_peer_barrier_rsp response; + response.result = 1; + if (!send_msg(sock, &response, sizeof(response))) { + return; + } + break; + } case RPC_CMD_COPY_TENSOR: { rpc_msg_copy_tensor_req request; if (!recv_msg(sock, &request, sizeof(request))) { @@ -2425,6 +2708,10 @@ void ggml_backend_rpc_start_server(const char * endpoint, const char * cache_dir fprintf(stderr, "Failed to create server socket\n"); return; } + // One thread per connection: besides the client that drives this server, other servers + // connect to it to deliver the tensors of a layer split directly (RPC_CMD_COPY_TENSOR_TO). + // The connections share one buffer registry and one execution mutex, see rpc_server_shared. + auto shared = std::make_shared(); while (true) { auto client_socket = server_socket->accept(); if (client_socket == nullptr) { @@ -2433,9 +2720,13 @@ void ggml_backend_rpc_start_server(const char * endpoint, const char * cache_dir } printf("Accepted client connection\n"); fflush(stdout); - rpc_serve_client(backends, cache_dir, client_socket); - printf("Client connection closed\n"); - fflush(stdout); + // the state the threads share outlives this function, so a failed accept cannot pull + // it out from under a connection that is still being served + std::thread([backends, cache_dir, client_socket, shared]() { + rpc_serve_client(backends, cache_dir, client_socket, *shared); + printf("Client connection closed\n"); + fflush(stdout); + }).detach(); } rpc_transport_shutdown(); for (auto backend : backends) { diff --git a/ggml/src/ggml-rpc/transport.cpp b/ggml/src/ggml-rpc/transport.cpp index 5ec15dc80c0c..b077d7e6d287 100644 --- a/ggml/src/ggml-rpc/transport.cpp +++ b/ggml/src/ggml-rpc/transport.cpp @@ -674,7 +674,7 @@ socket_ptr socket_t::create_server(const char * host, int port) { if (bind(sockfd, (struct sockaddr *) &serv_addr, sizeof(serv_addr)) < 0) { return nullptr; } - if (listen(sockfd, 1) < 0) { + if (listen(sockfd, 16) < 0) { return nullptr; } return socket_ptr(new socket_t(std::make_unique(sockfd))); From 8b1cbcd871aa5702f7f21c19911c0d7a08577b65 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 8 Sep 2026 22:21:24 -0700 Subject: [PATCH 2/8] rpc: do not abort the source server when a peer negotiation fails get_peer_socket opens a connection from inside a server, and get_socket ends in negotiate_hello, whose failed exchange hits RPC_STATUS_ASSERT and aborts the process. A destination that accepts the connection and then disconnects, or returns a truncated HELLO, therefore took down the source server and every client it was serving. That also contradicted what copy_tensor_to documents: an unreachable or too old destination is a recoverable condition reported as result = 0, which puts the client back on its previous path. negotiate_hello and get_socket now take may_fail, set only for peer connections, and return failure instead of aborting. A client keeps aborting as before. --- ggml/src/ggml-rpc/ggml-rpc.cpp | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index 25e4261f6300..de78ed746021 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -502,13 +502,19 @@ static bool send_rpc_cmd(socket_ptr sock, enum rpc_cmd cmd, const void * input, // Performs HELLO handshake with transport auto-negotiation. // Advertises local capabilities via conn_caps; if the server responds with // matching capabilities, the socket is upgraded transparently. -static bool negotiate_hello(const std::shared_ptr & sock) { +// may_fail is for connections opened from inside a server: a destination that disconnects or +// truncates its HELLO must not abort the source server and every client it is serving, so the +// caller gets false and copy_tensor_to reports result = 0. A client keeps aborting as before. +static bool negotiate_hello(const std::shared_ptr & sock, bool may_fail = false) { rpc_msg_hello_req request = {}; rpc_msg_hello_rsp response = {}; sock->get_caps(request.conn_caps); bool status = send_rpc_cmd(sock, RPC_CMD_HELLO, &request, sizeof(request), &response, sizeof(response)); + if (!status && may_fail) { + return false; + } RPC_STATUS_ASSERT(status); if (response.major != RPC_PROTO_MAJOR_VERSION || response.minor > RPC_PROTO_MINOR_VERSION) { @@ -537,7 +543,7 @@ static std::shared_ptr find_socket(const std::string & endpoint) { return nullptr; } -static std::shared_ptr get_socket(const std::string & endpoint) { +static std::shared_ptr get_socket(const std::string & endpoint, bool may_fail = false) { std::lock_guard lock(g_sockets_mutex); auto it = g_sockets.find(endpoint); @@ -560,7 +566,7 @@ static std::shared_ptr get_socket(const std::string & endpoint) { if (sock == nullptr) { return nullptr; } - if (!negotiate_hello(sock)) { + if (!negotiate_hello(sock, may_fail)) { return nullptr; } LOG_DBG("[%s] connected to %s\n", __func__, endpoint.c_str()); @@ -2035,7 +2041,8 @@ socket_ptr rpc_server::get_peer_socket(const std::string & endpoint) { } // the same connect and HELLO negotiation a client does, so a server to server link uses // RDMA whenever both rails allow it and TCP otherwise - auto sock = get_socket(endpoint); + // may_fail: a destination that is down or restarting is a recoverable condition here + auto sock = get_socket(endpoint, /* may_fail */ true); if (sock == nullptr) { GGML_LOG_ERROR("[%s] failed to connect to %s\n", __func__, endpoint.c_str()); return nullptr; From e2e3c040d17a0bec76dc9d7ea6aee78ca113edfa Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 8 Sep 2026 22:21:24 -0700 Subject: [PATCH 3/8] rpc: size the peer staging buffer only after validating the request copy_tensor_to resized p2p_buf from hdr.size, which a client controls, before any of the checks against the source tensor and its buffer. A large value aborted the server on an uncaught bad_alloc, and a value near SIZE_MAX wrapped msg_size so that the two header copies ran past the end of an undersized vector. The checks already existed further down. The sizing and the header copies now happen after them, where hdr.size is known to fit both the buffer and the tensor. --- ggml/src/ggml-rpc/ggml-rpc.cpp | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index de78ed746021..f85787987858 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -2076,13 +2076,8 @@ bool rpc_server::copy_tensor_to(const std::vector & input, rpc_msg_copy } const std::string endpoint((const char *) input.data() + sizeof(hdr), hdr.endpoint_len); - const size_t msg_size = sizeof(rpc_tensor) + sizeof(uint64_t) + (size_t) hdr.size; - if (p2p_buf.size() < msg_size) { - p2p_buf.resize(msg_size); - } - const uint64_t dst_offset = 0; - memcpy(p2p_buf.data(), &hdr.dst, sizeof(rpc_tensor)); - memcpy(p2p_buf.data() + sizeof(rpc_tensor), &dst_offset, sizeof(dst_offset)); + // sized below, once hdr.size has been checked against the source tensor + size_t msg_size = 0; // read the source out of this server's device; the peer connection is used without this // lock, so a ring of servers cannot deadlock on each other's execution mutex @@ -2115,6 +2110,18 @@ bool rpc_server::copy_tensor_to(const std::vector & input, rpc_msg_copy return false; } + // only now is hdr.size known to fit both the buffer and the tensor. Sizing the staging + // buffer before these checks let a client pick any 64 bit length: a large one aborts the + // server on an uncaught bad_alloc, and one near SIZE_MAX wraps msg_size so that the two + // header copies below run past the end of an undersized vector. + const uint64_t dst_offset = 0; + msg_size = sizeof(rpc_tensor) + sizeof(dst_offset) + (size_t) hdr.size; + if (p2p_buf.size() < msg_size) { + p2p_buf.resize(msg_size); + } + memcpy(p2p_buf.data(), &hdr.dst, sizeof(rpc_tensor)); + memcpy(p2p_buf.data() + sizeof(rpc_tensor), &dst_offset, sizeof(dst_offset)); + ggml_backend_tensor_get(tensor, p2p_buf.data() + sizeof(rpc_tensor) + sizeof(dst_offset), 0, hdr.size); } From 4bc3b0a9bb875a1117d382609c8f598ce589e52d Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 8 Sep 2026 22:21:24 -0700 Subject: [PATCH 4/8] rpc: order the peer write against work already sent to the destination rpc_flush_deferred_guarded only puts queued commands on the wire. It does not wait for them to be served, and the peer write arrives on a different connection, so the two are separated by nothing but the destination's execution mutex. An earlier GRAPH_COMPUTE or SET_TENSOR that has not yet taken that mutex could run after the peer's SET_TENSOR and read or overwrite a split input buffer that had just been reused. RPC_CMD_PEER_BARRIER already means that everything received earlier on a connection has been served, so it is now issued on dst_sock before the copy starts. A failure there is recoverable and falls back to routing the tensor through the client. --- ggml/src/ggml-rpc/ggml-rpc.cpp | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index f85787987858..dcb065819032 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -1081,6 +1081,21 @@ static bool ggml_backend_rpc_cpy_tensor_p2p(ggml_backend_t backend_src, ggml_bac // anything queued for the destination has to be on the wire before the peer write lands rpc_flush_deferred_guarded(dst_sock); + // and it has to have been served, not merely sent. The peer write arrives on a different + // connection, so flushing this one orders the bytes but nothing else: an earlier + // GRAPH_COMPUTE or SET_TENSOR that has not yet taken the destination's execution mutex can + // still run after the peer's SET_TENSOR and read or overwrite a reused split input buffer. + // One connection is served in order, so a command with a reply is the barrier: when this + // returns, everything sent earlier on dst_sock is done. Failure is recoverable, the caller + // falls back to routing the tensor through the client. + { + rpc_msg_peer_barrier_rsp barrier; + if (!send_rpc_cmd(dst_sock, RPC_CMD_PEER_BARRIER, nullptr, 0, &barrier, sizeof(barrier))) { + GGML_LOG_ERROR("[%s] %s did not acknowledge the barrier\n", __func__, dst_ctx->endpoint.c_str()); + return false; + } + } + const std::string & endpoint = dst_ctx->endpoint; rpc_msg_copy_tensor_to_hdr hdr; hdr.src = serialize_tensor(src); From 08ad5b62a42950b2138490a61d7597881329afd5 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 8 Sep 2026 22:39:48 -0700 Subject: [PATCH 5/8] rpc: advertise the peer copy as a server flag instead of a minor bump Raising RPC_PROTO_MINOR_VERSION to 3 locked out every already-deployed client: negotiate_hello rejects any server whose minor exceeds its own, so a 5.2 client could not connect to a 5.3 server for any operation at all, even though the new commands are appended compatibly. The version is back to 2 and support is advertised in the byte of the HELLO response that was previously pure padding. That byte is fixed size, already on the wire, and read as padding by every existing client, so an old client connects exactly as before and simply never sees the flag, while an old server sends zero and is never asked for a peer copy. The two gates that keyed on minor >= 3 now test the flag. Note on why the flag does not live in conn_caps, which would otherwise be the obvious place: update_caps treats any non-zero byte there as "the peer speaks RDMA", and rdma_caps already fills all 24 bytes, so a spare bit would both be unavailable and, if taken, make an RDMA-less peer look RDMA-capable. --- ggml/include/ggml-rpc.h | 2 +- ggml/src/ggml-rpc/ggml-rpc.cpp | 20 ++++++++++++++++---- ggml/src/ggml-rpc/transport.h | 3 +++ 3 files changed, 20 insertions(+), 5 deletions(-) diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h index dc51b73e940e..5d4030128246 100644 --- a/ggml/include/ggml-rpc.h +++ b/ggml/include/ggml-rpc.h @@ -7,7 +7,7 @@ extern "C" { #endif #define RPC_PROTO_MAJOR_VERSION 5 -#define RPC_PROTO_MINOR_VERSION 3 +#define RPC_PROTO_MINOR_VERSION 2 #define RPC_PROTO_PATCH_VERSION 0 #ifdef __cplusplus diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index dcb065819032..e78a5bef37e5 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -90,11 +90,20 @@ struct rpc_msg_hello_req { uint8_t conn_caps[RPC_CONN_CAPS_SIZE]; }; +// Server feature flags, carried in the byte that used to be pure padding in the HELLO response. +// Features are advertised here rather than by bumping the protocol minor because a client rejects +// any server whose minor exceeds its own, so a bump locks out every already-deployed older client +// even when the new commands are purely additive. This byte is fixed size, already on the wire, +// and ignored by every existing client, which reads it as padding and sees zero. +enum rpc_srv_flag { + RPC_SRV_FLAG_PEER_COPY = 1 << 0, // supports RPC_CMD_COPY_TENSOR_TO and RPC_CMD_PEER_BARRIER +}; + struct rpc_msg_hello_rsp { uint8_t major; uint8_t minor; uint8_t patch; - uint8_t padding; + uint8_t srv_flags; uint8_t conn_caps[RPC_CONN_CAPS_SIZE]; }; @@ -524,6 +533,7 @@ static bool negotiate_hello(const std::shared_ptr & sock, bool may_fai } sock->conn.server_minor = response.minor; + sock->conn.server_flags = response.srv_flags; sock->update_caps(response.conn_caps); return true; @@ -1047,7 +1057,7 @@ static bool rpc_supports_batched_get(const socket_ptr & sock) { static bool rpc_supports_p2p(const socket_ptr & sock) { static const char * env = std::getenv("GGML_RPC_P2P"); static const bool disabled = env != nullptr && std::strcmp(env, "0") == 0; - return !disabled && sock != nullptr && sock->conn.server_minor >= 3; + return !disabled && sock != nullptr && (sock->conn.server_flags & RPC_SRV_FLAG_PEER_COPY); } // Server to server movement of a tensor that crosses a stage boundary of a layer split. The @@ -1523,7 +1533,9 @@ void rpc_server::hello(rpc_msg_hello_rsp & response) { response.major = RPC_PROTO_MAJOR_VERSION; response.minor = RPC_PROTO_MINOR_VERSION; response.patch = RPC_PROTO_PATCH_VERSION; - LOG_DBG("[%s] version: %d.%d.%d\n", __func__, response.major, response.minor, response.patch); + response.srv_flags = RPC_SRV_FLAG_PEER_COPY; + LOG_DBG("[%s] version: %d.%d.%d flags: 0x%02x\n", __func__, response.major, response.minor, + response.patch, response.srv_flags); } bool rpc_server::get_alloc_size(const rpc_msg_get_alloc_size_req & request, rpc_msg_get_alloc_size_rsp & response) { @@ -2062,7 +2074,7 @@ socket_ptr rpc_server::get_peer_socket(const std::string & endpoint) { GGML_LOG_ERROR("[%s] failed to connect to %s\n", __func__, endpoint.c_str()); return nullptr; } - if (sock->conn.server_minor < 3) { + if (!(sock->conn.server_flags & RPC_SRV_FLAG_PEER_COPY)) { GGML_LOG_ERROR("[%s] %s does not support peer to peer copies\n", __func__, endpoint.c_str()); return nullptr; } diff --git a/ggml/src/ggml-rpc/transport.h b/ggml/src/ggml-rpc/transport.h index cffe5d3aa642..968455305ac1 100644 --- a/ggml/src/ggml-rpc/transport.h +++ b/ggml/src/ggml-rpc/transport.h @@ -40,6 +40,9 @@ struct rpc_conn_state { std::unordered_map last_graph_uid; + // server feature flags from the HELLO response + uint8_t server_flags = 0; + uint32_t server_minor = 0; // lock order: mtx_defer before mtx_send, never the reverse From 0008f46912777f7129789272b6e04334bf129ed5 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Wed, 9 Sep 2026 00:00:13 -0700 Subject: [PATCH 6/8] rpc: do not leak a descriptor or die on SIGPIPE when a peer write fails Two ways a recoverable peer-copy failure could take down the source server. socket_t::connect() created a descriptor and then returned nullptr from three later failure paths without closing it: TCP_NODELAY, name resolution, and the connect() itself. Until now that cost at most one descriptor per process, because a client that cannot connect aborts anyway. get_peer_socket() reaches the same code once per tensor per boundary crossing, does not cache failed peers, and treats a refused connection as recoverable, so the same leak becomes unbounded and a long decode eventually exhausts the descriptor limit, at which point the server can no longer connect to peers or accept clients. socket_t::impl already closes the descriptor in its destructor, so the fix is to adopt it into the socket_ptr immediately rather than carrying it raw until the last line. accept() and create_server() had the same shape and are fixed the same way; accept()'s is per client connection. Separately, send() was called with no flags. On POSIX, writing to a socket whose peer has gone away raises SIGPIPE, and neither this library nor the rpc-server executable installs a handler, so the default action terminates the process. A destination that was restarted while a source held a cached connection to it is exactly the case that is supposed to end in result = 0; instead it killed the source server and every client it was serving. Sends now pass MSG_NOSIGNAL where it exists, and Apple, which has no MSG_NOSIGNAL, gets SO_NOSIGPIPE set once per descriptor instead. Windows has no SIGPIPE and is unaffected. --- ggml/src/ggml-rpc/transport.cpp | 48 ++++++++++++++++++++++++++++++--- 1 file changed, 44 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-rpc/transport.cpp b/ggml/src/ggml-rpc/transport.cpp index b077d7e6d287..757b21c8458b 100644 --- a/ggml/src/ggml-rpc/transport.cpp +++ b/ggml/src/ggml-rpc/transport.cpp @@ -41,6 +41,30 @@ using ssize_t = __int64; typedef int sockfd_t; #endif +// Writing to a socket whose peer has gone away raises SIGPIPE on POSIX, and neither this library +// nor the rpc-server executable installs a handler, so the default action terminates the whole +// process. That was survivable while every send belonged to a client that aborts on a broken +// connection anyway. It is not survivable now: a source server writes to a cached peer connection, +// and a destination that was restarted is exactly the recoverable case that is supposed to end in +// result = 0, not in the source server dying and taking every client it serves with it. +// +// Linux and most BSDs take MSG_NOSIGNAL per send. Apple has no MSG_NOSIGNAL and needs the +// SO_NOSIGPIPE socket option instead, applied once per descriptor. Windows has no SIGPIPE at all. +#if defined(MSG_NOSIGNAL) +# define RPC_SEND_FLAGS MSG_NOSIGNAL +#else +# define RPC_SEND_FLAGS 0 +#endif + +static void set_no_sigpipe(sockfd_t sockfd) { +#if !defined(MSG_NOSIGNAL) && defined(SO_NOSIGPIPE) + int flag = 1; + setsockopt(sockfd, SOL_SOCKET, SO_NOSIGPIPE, (char *) &flag, sizeof(flag)); +#else + (void) sockfd; +#endif +} + static const char * RPC_DEBUG = std::getenv("GGML_RPC_DEBUG"); #define LOG_DBG(...) \ @@ -490,7 +514,7 @@ bool socket_t::impl::send_data(const void * data, size_t size) { size_t bytes_sent = 0; while (bytes_sent < size) { size_t size_to_send = std::min(size - bytes_sent, MAX_CHUNK_SIZE); - ssize_t n = send(fd, (const char *)data + bytes_sent, size_to_send, 0); + ssize_t n = send(fd, (const char *)data + bytes_sent, size_to_send, RPC_SEND_FLAGS); if (n < 0) { GGML_LOG_ERROR("send failed (bytes_sent=%zu, size_to_send=%zu)\n", bytes_sent, size_to_send); @@ -646,11 +670,15 @@ socket_ptr socket_t::accept() { if (!is_valid_fd(client_socket_fd)) { return nullptr; } + set_no_sigpipe(client_socket_fd); + // Adopt first here too: a TCP_NODELAY failure used to leak the accepted descriptor, and this + // one is reached once per client connection rather than once per process. + socket_ptr sock(new socket_t(std::make_unique(client_socket_fd))); if (!set_no_delay(client_socket_fd)) { GGML_LOG_ERROR("Failed to set TCP_NODELAY\n"); return nullptr; } - return socket_ptr(new socket_t(std::make_unique(client_socket_fd))); + return sock; } socket_ptr socket_t::create_server(const char * host, int port) { @@ -658,6 +686,9 @@ socket_ptr socket_t::create_server(const char * host, int port) { if (!is_valid_fd(sockfd)) { return nullptr; } + set_no_sigpipe(sockfd); + // Same reason as socket_t::connect: adopt first, so the failure paths below cannot leak. + socket_ptr sock(new socket_t(std::make_unique(sockfd))); if (!set_reuse_addr(sockfd)) { GGML_LOG_ERROR("Failed to set SO_REUSEADDR\n"); return nullptr; @@ -677,7 +708,7 @@ socket_ptr socket_t::create_server(const char * host, int port) { if (listen(sockfd, 16) < 0) { return nullptr; } - return socket_ptr(new socket_t(std::make_unique(sockfd))); + return sock; } socket_ptr socket_t::connect(const char * host, int port) { @@ -685,6 +716,15 @@ socket_ptr socket_t::connect(const char * host, int port) { if (!is_valid_fd(sockfd)) { return nullptr; } + set_no_sigpipe(sockfd); + // Adopt the descriptor before anything that can fail, so every early return below closes it. + // socket_t::impl's destructor already calls closesocket(); the leak was purely that the fd + // was carried raw until the last line. This used to cost at most one descriptor per process + // at startup, because a client that cannot connect aborts. A server calls this per peer copy + // through get_peer_socket(), failed peers are not cached, and a destination that refuses the + // connection is a recoverable condition that keeps being retried, so the same leak becomes one + // descriptor per tensor per boundary crossing and a long decode exhausts the limit. + socket_ptr sock(new socket_t(std::make_unique(sockfd))); if (!set_no_delay(sockfd)) { GGML_LOG_ERROR("Failed to set TCP_NODELAY\n"); return nullptr; @@ -701,7 +741,7 @@ socket_ptr socket_t::connect(const char * host, int port) { if (::connect(sockfd, (struct sockaddr *)&addr, sizeof(addr)) < 0) { return nullptr; } - return socket_ptr(new socket_t(std::make_unique(sockfd))); + return sock; } #ifdef _WIN32 From de47be8ba8cf264c038fe0bf9fa202f2f00de98b Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Wed, 9 Sep 2026 00:00:13 -0700 Subject: [PATCH 7/8] rpc: validate tensor buffers against the session that owns them Moving the buffer registry into rpc_server_shared so a peer copy could resolve the destination buffer also widened every other command. deserialize_tensor() validated against shared.buffers, the set of all live buffers in the process, so any connection could name a buffer belonging to another connection and then read it with GET_TENSOR, overwrite it with SET_TENSOR, or reference it from a graph. Before the registry moved, buffers belonged to the connection's rpc_server and the same request simply did not resolve. deserialize_tensor() now validates against owned_buffers, this session's own allocations, and takes an allow_foreign flag for the one path that genuinely needs more. That path is SET_TENSOR: the destination of a peer copy receives it on the connection the source server opened rather than on the connection of the client that allocated the buffer, so it has to resolve a buffer another session owns. The write stays bounded by the existing buffer range checks. GET_TENSOR, GRAPH_COMPUTE, memset_tensor and the local copy path go back to being confined to the session, as do buffer_get_base() and buffer_clear(), which were checking process-wide liveness rather than ownership. free_buffer() already checked ownership. owned_buffers is maintained on allocation, on free, and on session teardown, so it is the same lifetime the pre-existing per-connection registry had. This restores the isolation that existed before peer copies for every command except the peer write itself. It is not authentication: the RPC server has none, and a caller that can open a connection can still address the destination buffer of a peer copy. Narrowing that further needs the destination to be told which writes to expect, which is a protocol change and not attempted here. --- ggml/src/ggml-rpc/ggml-rpc.cpp | 25 +++++++++++++++++++------ 1 file changed, 19 insertions(+), 6 deletions(-) diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index e78a5bef37e5..91f4705eb4cc 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -1507,7 +1507,10 @@ class rpc_server { private: bool get_cached_file(uint64_t hash, std::vector & data); - ggml_tensor * deserialize_tensor(struct ggml_context * ctx, const rpc_tensor * tensor); + // allow_foreign widens buffer validation from this session's buffers to every live buffer + // in the process. Only the peer-copy write path sets it, see set_tensor(). + ggml_tensor * deserialize_tensor(struct ggml_context * ctx, const rpc_tensor * tensor, + bool allow_foreign = false); ggml_tensor * create_node(uint64_t id, struct ggml_context * ctx, const std::unordered_map & tensor_ptrs, @@ -1637,7 +1640,7 @@ bool rpc_server::buffer_get_base(const rpc_msg_buffer_get_base_req & request, rp LOG_DBG("[%s] remote_ptr: %" PRIx64 "\n", __func__, request.remote_ptr); ggml_backend_buffer_t buffer = reinterpret_cast(request.remote_ptr); - if (shared.buffers.find(buffer) == shared.buffers.end()) { + if (owned_buffers.find(buffer) == owned_buffers.end()) { GGML_LOG_ERROR("[%s] buffer not found\n", __func__); return false; } @@ -1666,7 +1669,7 @@ bool rpc_server::buffer_clear(const rpc_msg_buffer_clear_req & request) { LOG_DBG("[%s] remote_ptr: %" PRIx64 ", value: %u\n", __func__, request.remote_ptr, request.value); ggml_backend_buffer_t buffer = reinterpret_cast(request.remote_ptr); - if (shared.buffers.find(buffer) == shared.buffers.end()) { + if (owned_buffers.find(buffer) == owned_buffers.end()) { GGML_LOG_ERROR("[%s] buffer not found\n", __func__); return false; } @@ -1722,7 +1725,8 @@ bool rpc_server::memset_tensor(const rpc_msg_memset_tensor_req & request) { return true; } -ggml_tensor * rpc_server::deserialize_tensor(struct ggml_context * ctx, const rpc_tensor * tensor) { +ggml_tensor * rpc_server::deserialize_tensor(struct ggml_context * ctx, const rpc_tensor * tensor, + bool allow_foreign) { // Validate tensor type before using it if (tensor->type >= GGML_TYPE_COUNT) { GGML_LOG_ERROR("[%s] invalid tensor type received: %u\n", __func__, tensor->type); @@ -1748,7 +1752,12 @@ ggml_tensor * rpc_server::deserialize_tensor(struct ggml_context * ctx, const rp result->nb[i] = tensor->nb[i]; } result->buffer = reinterpret_cast(tensor->buffer); - if (result->buffer && shared.buffers.find(result->buffer) == shared.buffers.end()) { + // Validate against this session's own buffers by default. Before peer copies existed the + // registry was per connection, so naming another connection's buffer simply did not resolve; + // moving the registry into rpc_server_shared made every live buffer in the process reachable + // from every connection, which is a wider grant than the peer-copy write actually needs. + const auto & allowed = allow_foreign ? shared.buffers : owned_buffers; + if (result->buffer && allowed.find(result->buffer) == allowed.end()) { result->buffer = nullptr; } @@ -1792,7 +1801,11 @@ bool rpc_server::set_tensor(const std::vector & input) { ggml_context_ptr ctx_ptr { ggml_init(params) }; GGML_ASSERT(ctx_ptr != nullptr); ggml_context * ctx = ctx_ptr.get(); - ggml_tensor * tensor = deserialize_tensor(ctx, in_tensor); + // The destination of a peer copy receives this command on the connection the source server + // opened, not on the connection of the client that allocated the buffer, so this is the one + // path that has to resolve a buffer belonging to another session. The write is still bounded + // by the buffer range checks in deserialize_tensor and below. + ggml_tensor * tensor = deserialize_tensor(ctx, in_tensor, /* allow_foreign */ true); if (tensor == nullptr || tensor->buffer == nullptr) { GGML_LOG_ERROR("[%s] error deserializing tensor\n", __func__); return false; From 4f54a47d66400291fb19c9e61aea5a7d633f12bd Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Wed, 9 Sep 2026 06:25:30 -0700 Subject: [PATCH 8/8] rpc: restrict foreign writes to declared peer links, back off unreachable peers, bound connections Three review items, each reproduced against this branch before anything was changed. All three are CPU reproducible, so the numbers below are measured rather than argued. Foreign writes. set_tensor() passed allow_foreign = true for every caller, because the destination of a peer copy is the one path that has to resolve a buffer belonging to another session. Ordinary client connections dispatch the same command, so any client could name any live buffer in the process and write it. Demonstrated with two plain connections to one server: client A allocates a buffer and fills it with 0x11, client B allocates nothing at all and sends SET_TENSOR naming A's buffer, and A reads its own buffer back. Before, 4096 of 4096 bytes were B's 0xAA and none were A's 0x11. After, 4096 of 4096 are still 0x11 and none are 0xAA. The widening now applies only to a link that declared itself with RPC_CMD_PEER_LINK, which a source server sends on the connection it opens to a destination and no ordinary client sends. To be clear about what this is: the RPC protocol has no authentication of any kind, so this is isolation rather than authorization and does not make the port safe to expose. What it changes is that the privilege belongs to the one connection that asked for it instead of to every connection. The control matters more than the fix here, because refusing too much would break peer copies silently: a rejected write surfaces as result = 0 and a fallback, not as an error. Driving RPC_CMD_COPY_TENSOR_TO between two servers, both before and after, gives result = 1 and 4096 of 4096 bytes of the source pattern at the destination. Identical on both sides of the change. Unreachable peers. A destination the coordinator can reach but the source server cannot, which is what asymmetric firewall or NAT rules produce, failed in get_peer_socket() on every single copy with nothing remembered. Where packets are dropped rather than refused that is a full TCP timeout per split boundary, which turns a mild fallback into a stall. Failures are now cached with a 1s, 2s, 4s backoff to a 60s cap, cleared on the first success, so a destination that is briefly down is retried promptly while one that is unreachable by routing stops costing anything measurable. Connection threads. Every accepted socket got a detached thread before HELLO was validated, and that thread can block in recv_data() forever because the transport sets no read timeout. Measured by opening 200 connections that connect and then say nothing: before, 200 accepted, 0 refused, thread count +200. After, 64 accepted, 136 refused, thread count +64, and back to baseline once the clients close. The limit is 64 by default and GGML_RPC_MAX_CONNECTIONS raises it. Not covered here: the backoff is reasoned from the code path rather than measured, because provoking a genuinely unreachable destination needs the two-node setup, and the peer path was exercised only between two servers on one host. --- ggml/src/ggml-rpc/ggml-rpc.cpp | 114 +++++++++++++++++++++++++++++++-- 1 file changed, 109 insertions(+), 5 deletions(-) diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index 63b78435787e..57b349454253 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -78,6 +78,10 @@ enum rpc_cmd { RPC_CMD_GET_TENSORS, RPC_CMD_COPY_TENSOR_TO, RPC_CMD_PEER_BARRIER, + // sent by a source server on the link it opened to a destination server, to mark that link as + // server to server. Only such a link may write a buffer it does not own. Appended last so the + // existing command numbers, and RPC_CMD_HELLO in particular, do not move. + RPC_CMD_PEER_LINK, RPC_CMD_COUNT, }; @@ -384,6 +388,7 @@ static const char * rpc_cmd_name(int cmd) { case RPC_CMD_GET_TENSORS: return "GET_TENSORS"; case RPC_CMD_COPY_TENSOR_TO: return "COPY_TENSOR_TO"; case RPC_CMD_PEER_BARRIER: return "PEER_BARRIER"; + case RPC_CMD_PEER_LINK: return "PEER_LINK"; default: return "?"; } } @@ -1545,6 +1550,8 @@ class rpc_server { ggml_cgraph * graph; }; + void set_peer_link() { is_peer_link = true; } + private: bool get_cached_file(uint64_t hash, std::vector & data); // allow_foreign widens buffer validation from this session's buffers to every live buffer @@ -1566,10 +1573,25 @@ class rpc_server { std::unordered_set owned_buffers; // connections to other servers, one per destination endpoint, closed with this connection std::unordered_map peer_socks; + // endpoints this server could not reach, and until when not to try again. See get_peer_socket. + std::unordered_map peer_failed_until; + std::unordered_map peer_backoff; // reused staging for RPC_CMD_COPY_TENSOR_TO std::vector p2p_buf; // store the last computed graph for each backend std::vector stored_graphs; + + // Set only by RPC_CMD_PEER_LINK, which a source server sends on the link it opened to push a + // peer copy. Ordinary coordinator connections never send it and so can never write a buffer + // belonging to another session, which they previously could: SET_TENSOR widened validation for + // every caller because the peer-copy destination is the one path that legitimately needs it. + // + // This is isolation, not authentication. The RPC protocol has no authentication of any kind and + // an untrusted client could send this command too, so it does not make the port safe to expose; + // what it does is stop an ordinary or buggy client from reaching another session's buffers at + // all, which is the difference between every connection holding the privilege and only the one + // that asked for it. + bool is_peer_link = false; }; void rpc_server::hello(rpc_msg_hello_rsp & response) { @@ -1846,7 +1868,11 @@ bool rpc_server::set_tensor(const std::vector & input) { // opened, not on the connection of the client that allocated the buffer, so this is the one // path that has to resolve a buffer belonging to another session. The write is still bounded // by the buffer range checks in deserialize_tensor and below. - ggml_tensor * tensor = deserialize_tensor(ctx, in_tensor, /* allow_foreign */ true); + // + // Only a link that declared itself with RPC_CMD_PEER_LINK gets that widening. An ordinary + // client connection is held to its own buffers, so it can no longer overwrite another + // session's buffer by supplying that buffer's pointer. + ggml_tensor * tensor = deserialize_tensor(ctx, in_tensor, /* allow_foreign */ is_peer_link); if (tensor == nullptr || tensor->buffer == nullptr) { GGML_LOG_ERROR("[%s] error deserializing tensor\n", __func__); return false; @@ -2152,18 +2178,55 @@ socket_ptr rpc_server::get_peer_socket(const std::string & endpoint) { if (it != peer_socks.end()) { return it->second; } + + // A destination the coordinator can reach but this server cannot, which is what asymmetric + // firewall or NAT rules produce, fails here on every single cross-server copy. Nothing about + // that failure was remembered, so each split boundary paid another blocking connect, and when + // the packets are dropped rather than refused that is the full OS TCP timeout, tens of seconds + // each time. The fallback path is meant to be a mild slowdown, not a stall. Remember the + // failure and stop retrying it until the backoff expires. + const auto now = std::chrono::steady_clock::now(); + { + auto bad = peer_failed_until.find(endpoint); + if (bad != peer_failed_until.end()) { + if (now < bad->second) { + return nullptr; + } + peer_failed_until.erase(bad); + } + } + + auto fail = [&](const char * why) -> socket_ptr { + // grows 1s, 2s, 4s ... to a cap, so a destination that is down briefly is retried soon + // while one that is unreachable by routing stops costing anything measurable + auto & backoff = peer_backoff[endpoint]; + backoff = backoff == std::chrono::seconds(0) ? std::chrono::seconds(1) + : std::min(backoff * 2, std::chrono::seconds(60)); + peer_failed_until[endpoint] = now + backoff; + GGML_LOG_ERROR("[%s] %s: %s, not retrying for %llds\n", __func__, endpoint.c_str(), why, + (long long) backoff.count()); + return nullptr; + }; + // the same connect and HELLO negotiation a client does, so a server to server link uses // RDMA whenever both rails allow it and TCP otherwise // may_fail: a destination that is down or restarting is a recoverable condition here auto sock = get_socket(endpoint, /* may_fail */ true); if (sock == nullptr) { - GGML_LOG_ERROR("[%s] failed to connect to %s\n", __func__, endpoint.c_str()); - return nullptr; + return fail("failed to connect"); } if (!(sock->conn.server_flags & RPC_SRV_FLAG_PEER_COPY)) { - GGML_LOG_ERROR("[%s] %s does not support peer to peer copies\n", __func__, endpoint.c_str()); - return nullptr; + return fail("does not support peer to peer copies"); } + // Declare this link server to server before any write goes over it. The destination only + // widens buffer validation for links that have said this, so without it the copy would be + // rejected as a write to a buffer this session does not own. + rpc_msg_peer_barrier_rsp linked; + if (!send_rpc_cmd(sock, RPC_CMD_PEER_LINK, nullptr, 0, &linked, sizeof(linked)) || linked.result == 0) { + return fail("peer link handshake failed"); + } + + peer_backoff.erase(endpoint); peer_socks[endpoint] = sock; return sock; } @@ -2435,6 +2498,18 @@ rpc_server::~rpc_server() { } } +// connections currently being served, and the ceiling on them. See the accept loop. +static std::atomic rpc_live_conns{0}; + +static int rpc_max_conns() { + static const int n = [] { + const char * e = getenv("GGML_RPC_MAX_CONNECTIONS"); + const int v = e != nullptr ? atoi(e) : 0; + return v > 0 ? v : 64; + }(); + return n; +} + static void rpc_serve_client(const std::vector & backends, const char * cache_dir, socket_ptr sock, rpc_server_shared & shared) { rpc_server server(backends, cache_dir, shared); @@ -2699,6 +2774,18 @@ static void rpc_serve_client(const std::vector & backends, const } break; } + case RPC_CMD_PEER_LINK: { + if (!recv_msg(sock, nullptr, 0)) { + return; + } + server.set_peer_link(); + rpc_msg_peer_barrier_rsp response; + response.result = 1; + if (!send_msg(sock, &response, sizeof(response))) { + return; + } + break; + } case RPC_CMD_COPY_TENSOR: { rpc_msg_copy_tensor_req request; if (!recv_msg(sock, &request, sizeof(request))) { @@ -2820,12 +2907,29 @@ void ggml_backend_rpc_start_server(const char * endpoint, const char * cache_dir fprintf(stderr, "Failed to accept client connection\n"); return; } + // Every accepted socket used to get a detached thread before HELLO was validated, and that + // thread can sit in recv_data() indefinitely because the transport sets no read timeout. A + // host that opens connections and then says nothing therefore consumed a thread, its stack + // and a descriptor each time, without ever completing a handshake, until the server ran out + // of one of them. Cap the number of connections served at once and refuse beyond it, which + // costs a well behaved deployment nothing: the coordinator and its peers are few. + if (rpc_live_conns.load(std::memory_order_relaxed) >= rpc_max_conns()) { + fprintf(stderr, "Refusing client connection: already serving %d, limit %d " + "(raise with GGML_RPC_MAX_CONNECTIONS)\n", + rpc_live_conns.load(std::memory_order_relaxed), rpc_max_conns()); + fflush(stderr); + // dropping the last reference closes it, so the peer sees the connection go away + continue; + } + rpc_live_conns.fetch_add(1, std::memory_order_relaxed); + printf("Accepted client connection\n"); fflush(stdout); // the state the threads share outlives this function, so a failed accept cannot pull // it out from under a connection that is still being served std::thread([backends, cache_dir, client_socket, shared]() { rpc_serve_client(backends, cache_dir, client_socket, *shared); + rpc_live_conns.fetch_sub(1, std::memory_order_relaxed); printf("Client connection closed\n"); fflush(stdout); }).detach();