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