From 45e4d8837defd62cf144f558da6dd5c36169d473 Mon Sep 17 00:00:00 2001 From: Yin Lin Date: Fri, 28 Aug 2026 16:03:17 +0000 Subject: [PATCH] Cut fixed per-request costs out of the reshard control plane Every Stage-3 coordination paid four avoidable fixed costs on its framed RPCs (coordinate, GET_METADATA, receiver arm): 1. The 4-byte length prefix and the body went out as two separate send() calls with Nagle enabled, on requests and responses alike, exposing every hop to the delayed-ACK stall (tens of ms on small RPCs). 2. The destination controller was asked for every registered unit's full pool manifest on every request, although work units register once per engine lifetime. 3. Client sockets carried no keepalive, so a black-holed peer was only detected at the full receive timeout. 4. The framed server never reaped its per-connection threads: one std::thread handle and stack per request, held until shutdown. Changes: - framed_rpc: single-buffer framing on both directions; TCP_NODELAY on client and accepted sockets; SO_KEEPALIVE plus TCP_USER_TIMEOUT bounded by the call's I/O timeout on client sockets; the accept loop joins finished connection threads. - reshard_coordinator: destination metadata is cached per controller address. Staleness (engine replacement) surfaces as a plan-build or receiver-arm failure. A failed attempt that used the cache drops the entry; while no receiver has acknowledged its arm the attempt is side-effect-free beyond the abandoned claim and is replayed once on fresh metadata. Once any receiver has acknowledged, the failure is returned as is and the next request re-queries. Validation: reshard package builds and reshard_service_test passes in the OSS bazel build (ml-build container, clang-18), including the new RemoteMetadataCachedAndRefreshedOnStaleFailure test (cache hit on the second request, exactly one refetch after a fingerprint-mismatch replay) and PartialReceiverArmFailureDropsCacheWithoutReplay test (one of two receivers refuses its arm: no replay, the cache entry is dropped, and the next request re-queries). --- tpu_sync/kv_cache/reshard/framed_rpc.cc | 68 +++++++-- tpu_sync/kv_cache/reshard/framed_rpc.h | 11 +- .../kv_cache/reshard/reshard_coordinator.cc | 64 ++++++-- .../kv_cache/reshard/reshard_coordinator.h | 19 +++ .../kv_cache/reshard/reshard_service_test.cc | 137 +++++++++++++++++- 5 files changed, 272 insertions(+), 27 deletions(-) diff --git a/tpu_sync/kv_cache/reshard/framed_rpc.cc b/tpu_sync/kv_cache/reshard/framed_rpc.cc index cb004b1d..2b080d02 100644 --- a/tpu_sync/kv_cache/reshard/framed_rpc.cc +++ b/tpu_sync/kv_cache/reshard/framed_rpc.cc @@ -17,12 +17,15 @@ #include #include #include +#include #include #include +#include #include #include #include +#include #include #include // NOLINT(build/c++11) #include @@ -60,6 +63,33 @@ bool SendAll(int fd, const char* data, size_t n) { return true; } +// Sends the 4-byte length prefix and the body as ONE buffer. Two separate +// send() calls under Nagle stall the body until the peer ACKs the prefix +// segment (delayed-ACK interaction), putting a tens-of-ms floor on small +// RPCs. +bool SendFramed(int fd, absl::string_view payload) { + std::string framed; + framed.reserve(sizeof(uint32_t) + payload.size()); + uint32_t net_len = htonl(static_cast(payload.size())); + framed.append(reinterpret_cast(&net_len), sizeof(net_len)); + framed.append(payload.data(), payload.size()); + return SendAll(fd, framed.data(), framed.size()); +} + +// Control-plane RPCs are short request/response exchanges: disable Nagle so +// each frame goes out immediately, and enable keepalive plus +// TCP_USER_TIMEOUT so a black-holed peer surfaces as a socket error within +// the I/O timeout instead of only at the full receive deadline. +void ConfigureControlSocket(int fd, absl::Duration io_timeout) { + int one = 1; + setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &one, sizeof(one)); + setsockopt(fd, SOL_SOCKET, SO_KEEPALIVE, &one, sizeof(one)); + unsigned int user_timeout_ms = static_cast( + absl::ToInt64Milliseconds(io_timeout)); + setsockopt(fd, IPPROTO_TCP, TCP_USER_TIMEOUT, &user_timeout_ms, + sizeof(user_timeout_ms)); +} + // Splits "host:port" at the last colon; strips IPv6 brackets, mirroring // raiden_controller.connect_socket. absl::Status SplitAddress(absl::string_view address, std::string* host, @@ -104,6 +134,7 @@ int TryConnectOnce(const std::string& host, int port, tv.tv_usec = 0; setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv)); setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv)); + ConfigureControlSocket(fd, io_timeout); if (connect(fd, res->ai_addr, res->ai_addrlen) == 0) break; close(fd); fd = -1; @@ -140,10 +171,7 @@ absl::StatusOr SocketFramedTransport::Call( std::string response; { - uint32_t net_len = htonl(static_cast(payload.size())); - if (!SendAll(fd, reinterpret_cast(&net_len), - sizeof(net_len)) || - !SendAll(fd, payload.data(), payload.size())) { + if (!SendFramed(fd, payload)) { close(fd); return absl::UnavailableError( absl::StrCat("Failed to send framed payload to ", address, ": ", @@ -242,10 +270,10 @@ void FramedServer::Stop() { server_fd_ = -1; } if (accept_thread_.joinable()) accept_thread_.join(); - for (std::thread& t : connection_threads_) { - if (t.joinable()) t.join(); + for (const std::unique_ptr& connection : connections_) { + if (connection->thread.joinable()) connection->thread.join(); } - connection_threads_.clear(); + connections_.clear(); } void FramedServer::AcceptLoop() { @@ -262,12 +290,29 @@ void FramedServer::AcceptLoop() { close(client_fd); break; } - connection_threads_.emplace_back(&FramedServer::ServeConnection, this, - client_fd); + connections_.erase( + std::remove_if(connections_.begin(), connections_.end(), + [](const std::unique_ptr& connection) { + if (!connection->done.load()) return false; + if (connection->thread.joinable()) { + connection->thread.join(); + } + return true; + }), + connections_.end()); + auto connection = std::make_unique(); + Connection* raw = connection.get(); + raw->thread = std::thread([this, raw, client_fd]() { + ServeConnection(client_fd); + raw->done.store(true); + }); + connections_.push_back(std::move(connection)); } } void FramedServer::ServeConnection(int client_fd) { + int one = 1; + setsockopt(client_fd, IPPROTO_TCP, TCP_NODELAY, &one, sizeof(one)); uint32_t net_len = 0; if (!ReadExactly(client_fd, reinterpret_cast(&net_len), sizeof(net_len))) { @@ -281,10 +326,7 @@ void FramedServer::ServeConnection(int client_fd) { return; } std::string response = handler_(request); - uint32_t resp_net_len = htonl(static_cast(response.size())); - SendAll(client_fd, reinterpret_cast(&resp_net_len), - sizeof(resp_net_len)); - SendAll(client_fd, response.data(), response.size()); + SendFramed(client_fd, response); close(client_fd); } diff --git a/tpu_sync/kv_cache/reshard/framed_rpc.h b/tpu_sync/kv_cache/reshard/framed_rpc.h index de307b98..51833d9f 100644 --- a/tpu_sync/kv_cache/reshard/framed_rpc.h +++ b/tpu_sync/kv_cache/reshard/framed_rpc.h @@ -17,6 +17,7 @@ #include #include +#include #include #include // NOLINT(build/c++11) #include @@ -79,6 +80,14 @@ class FramedServer final { int port() const { return port_; } private: + // One handler thread per accepted connection; `done` lets the accept loop + // reap finished threads so a long-lived server does not accumulate one + // un-joined thread (and its stack) per request. + struct Connection { + std::thread thread; + std::atomic done{false}; + }; + void AcceptLoop(); void ServeConnection(int client_fd); @@ -88,7 +97,7 @@ class FramedServer final { Handler handler_; std::atomic stopping_{false}; std::thread accept_thread_; - std::vector connection_threads_; + std::vector> connections_; }; } // namespace reshard diff --git a/tpu_sync/kv_cache/reshard/reshard_coordinator.cc b/tpu_sync/kv_cache/reshard/reshard_coordinator.cc index c6e80d5e..acaee71e 100644 --- a/tpu_sync/kv_cache/reshard/reshard_coordinator.cc +++ b/tpu_sync/kv_cache/reshard/reshard_coordinator.cc @@ -328,22 +328,65 @@ absl::Status ReshardCoordinator::ExecutePoolReshard( "(the destination-side relay is retired)"); } - const int64_t controller_start_ns = MonotonicNs(); - const int64_t plan_build_start_ns = controller_start_ns; - - auto src_metadata = directory_->LocalMetadata(args.src_units); - if (!src_metadata.ok()) return src_metadata.status(); std::vector dst_metadata; + bool used_cache = false; if (!args.dst_controller_address.empty()) { - auto remote = QueryRemoteMetadata(args.dst_controller_address); - if (!remote.ok()) return remote.status(); - dst_metadata = *std::move(remote); + { + absl::MutexLock lock(metadata_cache_mu_); + auto it = remote_metadata_cache_.find(args.dst_controller_address); + if (it != remote_metadata_cache_.end()) { + dst_metadata = it->second; + used_cache = true; + } + } + if (!used_cache) { + auto remote = QueryRemoteMetadata(args.dst_controller_address); + if (!remote.ok()) return remote.status(); + dst_metadata = *std::move(remote); + absl::MutexLock lock(metadata_cache_mu_); + remote_metadata_cache_[args.dst_controller_address] = dst_metadata; + } } else { auto local = directory_->LocalMetadata(args.dst_units); if (!local.ok()) return local.status(); dst_metadata = *std::move(local); } + bool receiver_armed = false; + absl::Status status = + ExecutePoolReshardAttempt(args, dst_metadata, &receiver_armed); + if (status.ok() || !used_cache) return status; + // Cached destination metadata can be stale after an engine replacement, + // so a failed attempt drops the entry and the next request re-queries. + // Planning and arming are side-effect-free beyond the abandoned claim + // until a receiver acknowledges its arm, so the attempt is replayed on + // fresh metadata only while no receiver has acknowledged. + { + absl::MutexLock lock(metadata_cache_mu_); + remote_metadata_cache_.erase(args.dst_controller_address); + } + if (receiver_armed) return status; + auto remote = QueryRemoteMetadata(args.dst_controller_address); + if (!remote.ok()) return remote.status(); + dst_metadata = *std::move(remote); + { + absl::MutexLock lock(metadata_cache_mu_); + remote_metadata_cache_[args.dst_controller_address] = dst_metadata; + } + receiver_armed = false; + return ExecutePoolReshardAttempt(args, dst_metadata, &receiver_armed); +} + +absl::Status ReshardCoordinator::ExecutePoolReshardAttempt( + const PoolReshardArgs& args, + const std::vector& dst_metadata, + bool* receiver_armed) { + const int64_t controller_start_ns = MonotonicNs(); + const int64_t plan_build_start_ns = controller_start_ns; + + auto src_metadata = directory_->LocalMetadata(args.src_units); + if (!src_metadata.ok()) return src_metadata.status(); + int64_t uuid = args.uuid; if (uuid <= 0) { // Python: random.randint(1, 2**63 - 1) when the wire carries no uuid. @@ -355,7 +398,7 @@ absl::Status ReshardCoordinator::ExecutePoolReshard( plan_request.src_units = args.src_units; plan_request.dst_units = args.dst_units; plan_request.src_metadata = *std::move(src_metadata); - plan_request.dst_metadata = std::move(dst_metadata); + plan_request.dst_metadata = dst_metadata; plan_request.req_id = args.req_id; plan_request.uuid = uuid; plan_request.dst_device_block_ids = args.dst_device_block_ids; @@ -401,6 +444,9 @@ absl::Status ReshardCoordinator::ExecutePoolReshard( }); } for (std::thread& t : armers) t.join(); + for (const absl::Status& status : arm_status) { + if (status.ok()) *receiver_armed = true; + } for (const absl::Status& status : arm_status) { if (!status.ok()) { registry_->AbandonClaim(args.req_id, uuid, claim_owner); diff --git a/tpu_sync/kv_cache/reshard/reshard_coordinator.h b/tpu_sync/kv_cache/reshard/reshard_coordinator.h index 6c595845..e19bf54f 100644 --- a/tpu_sync/kv_cache/reshard/reshard_coordinator.h +++ b/tpu_sync/kv_cache/reshard/reshard_coordinator.h @@ -112,6 +112,15 @@ class ReshardCoordinator { private: absl::Status ExecutePoolReshard(const PoolReshardArgs& args); + // One planning/arming/dispatch pass over the given destination metadata. + // Sets *receiver_armed once any destination acknowledges its arm; past + // that point the transfer has side effects beyond the abandoned claim and + // must not be replayed. + absl::Status ExecutePoolReshardAttempt( + const PoolReshardArgs& args, + const std::vector& dst_metadata, + bool* receiver_armed); + // GET_METADATA against the destination controller (dst_controller_address // path), recorded shape-identical to Python's _query_remote_metadata. absl::StatusOr> @@ -124,6 +133,16 @@ class ReshardCoordinator { mutable absl::Mutex status_mu_; std::map transfer_status_ ABSL_GUARDED_BY(status_mu_); + + // Destination work-unit metadata rarely changes (units register once per + // engine lifetime), so the per-request GET_METADATA round trip is served + // from this per-address cache. Staleness surfaces as a plan-build or + // receiver-arm failure; a failed attempt drops the entry, and + // ExecutePoolReshard replays once on fresh metadata while no receiver + // has acknowledged its arm. + mutable absl::Mutex metadata_cache_mu_; + std::map> + remote_metadata_cache_ ABSL_GUARDED_BY(metadata_cache_mu_); }; } // namespace reshard diff --git a/tpu_sync/kv_cache/reshard/reshard_service_test.cc b/tpu_sync/kv_cache/reshard/reshard_service_test.cc index a18d4a94..763f5cbe 100644 --- a/tpu_sync/kv_cache/reshard/reshard_service_test.cc +++ b/tpu_sync/kv_cache/reshard/reshard_service_test.cc @@ -17,6 +17,7 @@ #include #include #include +#include #include #include #include @@ -91,6 +92,9 @@ class FakeTransport final : public FramedTransport { absl::MutexLock lock(mu_); calls_.emplace_back(std::string(address), std::string(payload)); } + if (!metadata_addr_.empty() && metadata_addr_ == address) { + return metadata_responder_(); + } if (fail_addr_ == address) { tpu_sync::rpc::ControlResponse failed; failed.set_success(false); @@ -104,9 +108,17 @@ class FakeTransport final : public FramedTransport { void FailFor(const std::string& addr) { fail_addr_ = addr; } + void ServeMetadata(const std::string& addr, + std::function responder) { + metadata_addr_ = addr; + metadata_responder_ = std::move(responder); + } + absl::Mutex mu_; std::vector> calls_; std::string fail_addr_; + std::string metadata_addr_; + std::function metadata_responder_; }; class ReshardStackTest : public ::testing::Test { @@ -147,7 +159,8 @@ class ReshardStackTest : public ::testing::Test { // Registers rank units 0..7 plus the decode unit(s) through the framed // surface (byte-level, like the real facade would). void RegisterAllUnits(int num_src, int64_t live, int64_t stride, - int64_t num_blocks, int num_dst = 1) { + int64_t num_blocks, int num_dst = 1, + const std::string& fingerprint = "fp1") { for (int rank = 0; rank < num_src; ++rank) { tpu_sync::rpc::ControlRequest req; req.set_command( @@ -158,7 +171,7 @@ class ReshardStackTest : public ::testing::Test { reg->set_control_plane_rpc_address( absl::StrCat("10.0.0.1:", 9100 + 2 * rank)); *reg->add_pools() = MakePool("fa", live, stride, num_blocks); - reg->set_layout_fingerprint("fp1"); + reg->set_layout_fingerprint(fingerprint); reg->set_page_tokens(512); reg->set_transfer_parallelism(num_src); reg->set_transfer_rank(rank); @@ -166,7 +179,7 @@ class ReshardStackTest : public ::testing::Test { ASSERT_TRUE(resp.success()) << resp.message(); } for (int idx = 0; idx < num_dst; ++idx) { - RegisterDstUnit(idx, num_src, live, stride, num_blocks); + RegisterDstUnit(idx, num_src, live, stride, num_blocks, fingerprint); } } @@ -314,7 +327,7 @@ class ReshardStackTest : public ::testing::Test { tpu_sync::rpc::ControllerResponse Coordinate( const std::string& req_id, int64_t uuid, int num_src, std::vector dst_blocks, std::vector dst_skip = {}, - int num_dst = 1) { + int num_dst = 1, const std::string& dst_controller_address = "") { tpu_sync::rpc::ControllerRequest req; req.set_command( tpu_sync::rpc::ControllerRequest::COMMAND_COORDINATE_TRANSFER); @@ -337,6 +350,9 @@ class ReshardStackTest : public ::testing::Test { for (int64_t skip : dst_skip) { coord->add_dst_skip_bytes(skip); } + if (!dst_controller_address.empty()) { + coord->set_dst_controller_address(dst_controller_address); + } return HandleController(req.SerializeAsString()); } @@ -590,6 +606,119 @@ TEST_F(ReshardStackTest, ArmFailureAbandonsClaimAndSkipsSenders) { EXPECT_EQ(transport_.calls_.size(), 2u); } +TEST_F(ReshardStackTest, RemoteMetadataCachedAndRefreshedOnStaleFailure) { + RegisterAllUnits(/*num_src=*/2, /*live=*/1024, /*stride=*/1024, + /*num_blocks=*/16, /*num_dst=*/0); + const std::string dst_ctrl = "10.7.7.7:28000"; + std::string dst_fp = "fp1"; + auto dst_metadata = [&dst_fp]() { + tpu_sync::rpc::ControlResponse resp; + resp.set_success(true); + auto* meta = resp.mutable_get_metadata_response()->add_metadata(); + *meta->mutable_unit() = RaidenIdToProto(DstUnit()); + meta->add_shards("10.0.0.2:9400"); + meta->set_control_plane_rpc_address("10.0.0.2:9600"); + *meta->add_pools() = MakePool("fa", 1024, 1024, 16); + meta->set_layout_fingerprint(dst_fp); + meta->set_page_tokens(4096); + meta->set_transfer_parallelism(2); + meta->set_transfer_rank(0); + return resp.SerializeAsString(); + }; + transport_.ServeMetadata(dst_ctrl, dst_metadata); + auto metadata_queries = [&]() { + int count = 0; + for (const auto& call : transport_.calls_) { + if (call.first == dst_ctrl) ++count; + } + return count; + }; + + RegisterSpans(0, "req-m1", 61, 1024, 3, 0, 0, 1024); + RegisterSpans(1, "req-m1", 61, 1024, 5, 1, 0, 512); + ASSERT_TRUE(Coordinate("req-m1", 61, 2, {7, 9}, {}, /*num_dst=*/1, dst_ctrl) + .success()); + EXPECT_EQ(metadata_queries(), 1); + + // Second request is served from the metadata cache: no new query. + RegisterSpans(0, "req-m2", 62, 1024, 4, 0, 0, 1024); + RegisterSpans(1, "req-m2", 62, 1024, 6, 1, 0, 512); + ASSERT_TRUE(Coordinate("req-m2", 62, 2, {11, 12}, {}, /*num_dst=*/1, dst_ctrl) + .success()); + EXPECT_EQ(metadata_queries(), 1); + + // The peer re-admits with a new layout: sources move to fp2 and the + // remote now reports fp2. The fp1 cache entry fails plan build once, + // is invalidated, refetched, and the replay succeeds. + RegisterAllUnits(/*num_src=*/2, /*live=*/1024, /*stride=*/1024, + /*num_blocks=*/16, /*num_dst=*/0, "fp2"); + dst_fp = "fp2"; + RegisterSpans(0, "req-m3", 63, 1024, 8, 0, 0, 1024); + RegisterSpans(1, "req-m3", 63, 1024, 9, 1, 0, 512); + tpu_sync::rpc::ControllerResponse resp = + Coordinate("req-m3", 63, 2, {13, 14}, {}, /*num_dst=*/1, dst_ctrl); + ASSERT_TRUE(resp.success()) << resp.message(); + EXPECT_EQ(metadata_queries(), 2); +} + +TEST_F(ReshardStackTest, PartialReceiverArmFailureDropsCacheWithoutReplay) { + RegisterAllUnits(/*num_src=*/2, /*live=*/1024, /*stride=*/1024, + /*num_blocks=*/16, /*num_dst=*/0); + const std::string dst_ctrl = "10.7.7.7:28000"; + auto dst_metadata = []() { + tpu_sync::rpc::ControlResponse resp; + resp.set_success(true); + for (int idx = 0; idx < 2; ++idx) { + auto* meta = resp.mutable_get_metadata_response()->add_metadata(); + *meta->mutable_unit() = RaidenIdToProto(DstUnit(idx)); + meta->add_shards(absl::StrCat("10.0.0.2:", 9400 + idx)); + meta->set_control_plane_rpc_address( + absl::StrCat("10.0.0.2:", 9600 + idx)); + *meta->add_pools() = MakePool("fa", 1024, 1024, 16); + meta->set_layout_fingerprint("fp1"); + meta->set_page_tokens(4096); + meta->set_transfer_parallelism(2); + meta->set_transfer_rank(0); + } + return resp.SerializeAsString(); + }; + transport_.ServeMetadata(dst_ctrl, dst_metadata); + auto metadata_queries = [&]() { + int count = 0; + for (const auto& call : transport_.calls_) { + if (call.first == dst_ctrl) ++count; + } + return count; + }; + + RegisterSpans(0, "req-p1", 71, 1024, 3, 0, 0, 1024); + RegisterSpans(1, "req-p1", 71, 1024, 5, 1, 0, 512); + ASSERT_TRUE(Coordinate("req-p1", 71, 2, {7, 9}, {}, /*num_dst=*/2, dst_ctrl) + .success()); + EXPECT_EQ(metadata_queries(), 1); + + // One of the two receivers refuses its arm while the other acknowledges: + // the request fails with that refusal and is not replayed (a replay would + // re-arm the acknowledged receiver), but the cache entry is dropped. + transport_.FailFor("10.0.0.2:9601"); + RegisterSpans(0, "req-p2", 72, 1024, 4, 0, 0, 1024); + RegisterSpans(1, "req-p2", 72, 1024, 6, 1, 0, 512); + tpu_sync::rpc::ControllerResponse refused = + Coordinate("req-p2", 72, 2, {11, 12}, {}, /*num_dst=*/2, dst_ctrl); + ASSERT_FALSE(refused.success()); + EXPECT_THAT(refused.message(), HasSubstr("injected arm refusal")); + EXPECT_EQ(metadata_queries(), 1); + + // The next request re-queries the destination before planning. + transport_.FailFor(""); + RegisterSpans(0, "req-p3", 73, 1024, 8, 0, 0, 1024); + RegisterSpans(1, "req-p3", 73, 1024, 9, 1, 0, 512); + ASSERT_TRUE( + Coordinate("req-p3", 73, 2, {13, 14}, {}, /*num_dst=*/2, dst_ctrl) + .success()); + EXPECT_EQ(metadata_queries(), 2); +} + TEST_F(ReshardStackTest, CancelTombstoneBlocksLateRegistration) { RegisterAllUnits(/*num_src=*/1, /*live=*/1024, /*stride=*/1024, /*num_blocks=*/16);