diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index 6855d75744c..57b34945425 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,12 @@ enum rpc_cmd { RPC_CMD_GRAPH_RECOMPUTE, RPC_CMD_MEMSET_TENSOR, 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, }; @@ -88,12 +95,20 @@ struct rpc_msg_hello_req { }; // Server feature flags, carried in the byte that used to be pure padding in the HELLO response. -// Advertised here rather than by bumping the protocol minor: a client rejects any server whose -// minor exceeds its own, so a bump locks out every already-deployed older client even when the new -// command is purely additive and such a client would never send it. This byte is fixed size, -// already on the wire, and read as padding by existing clients, which see zero. +// 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 and such a client would never send them. This byte +// is fixed size, already on the wire, and ignored by every existing client, which reads it as +// padding and sees zero. +// +// One bit per feature, and the bits are not interchangeable: batched get and peer copy were +// developed on separate branches and each started life at bit 0. Keeping both at bit 0 would make a +// server that only implements batched get look, to a client that knows about peer copy, like it +// supports peer copy too, and that client would then send it RPC_CMD_COPY_TENSOR_TO. Batched get +// keeps bit 0 because it is the value already published on this branch. enum rpc_srv_flag { RPC_SRV_FLAG_BATCHED_GET = 1 << 0, // supports RPC_CMD_GET_TENSORS + RPC_SRV_FLAG_PEER_COPY = 1 << 1, // supports RPC_CMD_COPY_TENSOR_TO and RPC_CMD_PEER_BARRIER }; struct rpc_msg_hello_rsp { @@ -212,6 +227,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; }; @@ -348,6 +386,9 @@ 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"; + case RPC_CMD_PEER_LINK: return "PEER_LINK"; default: return "?"; } } @@ -491,13 +532,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) { @@ -527,7 +574,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); @@ -550,7 +597,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()); @@ -1041,6 +1088,82 @@ static bool rpc_supports_batched_get(const socket_ptr & sock) { return !disabled && (sock->conn.server_flags & RPC_SRV_FLAG_BATCHED_GET); } +// 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_flags & RPC_SRV_FLAG_PEER_COPY); +} + +// 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); + + // 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); + 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 socket_ptr tensor_socket(const ggml_tensor * tensor) { ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; if (buf == nullptr || !ggml_backend_buffer_is_rpc(buf)) { @@ -1087,7 +1210,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); - 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; } @@ -1379,10 +1505,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(); @@ -1400,6 +1538,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); @@ -1411,32 +1550,63 @@ 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); - 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, 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; + // 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) { response.major = RPC_PROTO_MAJOR_VERSION; response.minor = RPC_PROTO_MINOR_VERSION; response.patch = RPC_PROTO_PATCH_VERSION; - response.srv_flags = RPC_SRV_FLAG_BATCHED_GET; + // this server implements both, so it advertises both + response.srv_flags = RPC_SRV_FLAG_BATCHED_GET | 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) { + std::lock_guard lock(shared.mtx); + uint32_t dev_id = request.device; if (dev_id >= backends.size()) { return false; @@ -1477,6 +1647,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; @@ -1490,7 +1662,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); } @@ -1498,6 +1671,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; @@ -1510,6 +1685,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; @@ -1522,9 +1699,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 (owned_buffers.find(buffer) == owned_buffers.end()) { GGML_LOG_ERROR("[%s] buffer not found\n", __func__); return false; } @@ -1534,21 +1713,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 (owned_buffers.find(buffer) == owned_buffers.end()) { GGML_LOG_ERROR("[%s] buffer not found\n", __func__); return false; } @@ -1557,6 +1741,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, @@ -1602,7 +1788,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); @@ -1628,7 +1815,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 && buffers.find(result->buffer) == 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; } @@ -1653,6 +1845,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; @@ -1670,7 +1864,15 @@ 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. + // + // 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; @@ -1726,6 +1928,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; @@ -1767,6 +1971,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, @@ -1802,6 +2008,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, @@ -1838,6 +2046,8 @@ bool rpc_server::get_tensor(const rpc_msg_get_tensor_req & request, std::vector< 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; } @@ -1921,6 +2131,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, @@ -1961,6 +2173,156 @@ 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; + } + + // 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) { + return fail("failed to connect"); + } + if (!(sock->conn.server_flags & RPC_SRV_FLAG_PEER_COPY)) { + 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; +} + +// 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); + + // 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 + { + 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; + } + + // 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); + } + + 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, @@ -2018,6 +2380,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)) { @@ -2092,6 +2456,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; @@ -2107,6 +2473,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; @@ -2121,14 +2489,30 @@ 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); } } +// 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 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; @@ -2364,6 +2748,44 @@ 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_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))) { @@ -2475,17 +2897,42 @@ 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) { 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); - 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); + rpc_live_conns.fetch_sub(1, std::memory_order_relaxed); + 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 5ec15dc80c0..757b21c8458 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; @@ -674,10 +705,10 @@ 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))); + 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 diff --git a/ggml/src/ggml-rpc/transport.h b/ggml/src/ggml-rpc/transport.h index aae0d920291..004183c4952 100644 --- a/ggml/src/ggml-rpc/transport.h +++ b/ggml/src/ggml-rpc/transport.h @@ -41,6 +41,8 @@ struct rpc_conn_state { std::unordered_map last_graph_uid; uint32_t server_minor = 0; + + // server feature flags from the HELLO response uint8_t server_flags = 0; // lock order: mtx_defer before mtx_send, never the reverse