From d109c9c010ac09471822a80c4df079c3160d24d5 Mon Sep 17 00:00:00 2001 From: Yin Lin Date: Mon, 24 Aug 2026 00:27:44 +0000 Subject: [PATCH 1/2] Hold transfer staging until every accepted operation drains A transfer that settles (timeout, a failed copy or push, or completion) now publishes its outcome only when every accepted copy, push, and incoming payload against its staging has finished, so neither its host blocks nor the device pages they mirror are handed to another transfer while late work can still touch them. A failed transfer retires its uuid before its blocks return, so a late payload resolves nothing. Destruction admits no new handlers, counts handlers and workers, unblocks accepted sockets, cancels queued pushes exactly once, stops the data transport, and waits for the last callback before tearing state down. --- .../core/kv_cache_manager_with_transfer.cc | 751 ++++++++++++++---- .../core/kv_cache_manager_with_transfer.h | 72 ++ ...manager_with_transfer_pool_reshard_test.cc | 36 + tpu_sync/core/raiden_manager_base.cc | 18 + tpu_sync/core/raiden_manager_base.h | 6 + tpu_sync/kv_cache/kv_cache_manager_base.cc | 46 +- tpu_sync/kv_cache/kv_cache_manager_base.h | 17 + tpu_sync/transport/block_transport.cc | 55 ++ tpu_sync/transport/block_transport.h | 8 + tpu_sync/transport/block_transport_delegate.h | 8 + 10 files changed, 848 insertions(+), 169 deletions(-) diff --git a/tpu_sync/core/kv_cache_manager_with_transfer.cc b/tpu_sync/core/kv_cache_manager_with_transfer.cc index 502d1eaa..80ff16e6 100644 --- a/tpu_sync/core/kv_cache_manager_with_transfer.cc +++ b/tpu_sync/core/kv_cache_manager_with_transfer.cc @@ -619,15 +619,39 @@ KVCacheManagerWithTransfer::KVCacheManagerWithTransfer( } KVCacheManagerWithTransfer::~KVCacheManagerWithTransfer() { - StopControlServer(); - // Pull-serve workers read this object's state; nothing may be torn down - // while one is still running. + // Refuse new work first: once this is visible, no pull handler, push + // worker, or H2H launch starts. shutting_down_.store(true, std::memory_order_relaxed); + StopControlServer(); + // Pull-serve handlers and workers read this object's state; nothing may + // be torn down while one is still running. { absl::MutexLock lock(pull_workers_mu_); pull_workers_mu_.Await(absl::Condition( +[](int* active) { return *active == 0; }, &active_pull_workers_)); } + { + // A drained worker's last D2H callback may still be handing an H2H push + // to the transport; the hand-off must finish before the transport stops. + absl::MutexLock lock(mu_); + mu_.Await(absl::Condition( + +[](int* pending) { return *pending == 0; }, &pending_h2h_launches_)); + } + // Break in-flight pushes so their completion callbacks fire, then stop + // the transport while this object is still whole: its handler threads + // call back into the delegate for chunk resolution, payload leases, and + // layer completion. + if (auto* transport = transport_server_if_started()) { + transport->AbortActiveSends(); + } + StopTransportServer(); + { + // Every accepted operation holds a lease on the staging it uses; the + // last release frees that staging and reaps its entry. + absl::MutexLock lock(mu_); + mu_.Await(absl::Condition( + +[](int* ops) { return *ops == 0; }, &total_outstanding_ops_)); + } push_pool_.reset(); pull_pool_.reset(); if (host_block_manager_ && !all_slots_.empty()) { @@ -669,10 +693,22 @@ int64_t KVCacheManagerWithTransfer::NotifyForRead( entry->deadline = DeadlineFromNow(); entry->register_start = register_start; + SettleActions replaced; { absl::MutexLock lock(mu_); + if (send_entries_.contains(uuid)) { + // A live entry re-registered under the same uuid is failed and + // drained first; silently replacing it would orphan its staging and + // strand its callbacks' leases. + LOG(ERROR) << "NotifyForRead: uuid " << uuid + << " re-registered while still active; failing the old entry"; + SettleSendLocked(uuid, /*failed=*/true, &replaced); + } send_entries_[uuid] = entry; } + ApplySettleActions(&replaced); + // A fresh registration under a previously failed uuid serves again. + ReviveTransferUuid(uuid); cv_.SignalAll(); std::ostringstream timing; @@ -696,6 +732,7 @@ absl::Status KVCacheManagerWithTransfer::RegisterActivePlan( return absl::AlreadyExistsError( absl::StrCat("Plan with UUID ", uuid, " is already registered!")); } + SettleActions replace_actions; const uint64_t generation = ++plan_generation_counter_; // Under demand staging a plan's device blocks are staged in host blocks // allocated for the plan, so the host mirror no longer has to span the @@ -795,6 +832,11 @@ absl::Status KVCacheManagerWithTransfer::RegisterActivePlan( ReleaseRecvStagingLocked(&recv_entry); } if (total_blocks > 0) { + if (active_recv_entries_.contains(uuid)) { + // Only a plan-less receive can collide here (a planned one is + // refused above); it is failed and drained rather than replaced. + SettleRecvLocked(uuid, /*failed=*/true, &replace_actions); + } active_recv_entries_[uuid] = std::move(recv_entry); LOG(INFO) << "RegisterActivePlan (Receiver): Populated " "active_recv_entries_ for UUID " @@ -804,6 +846,8 @@ absl::Status KVCacheManagerWithTransfer::RegisterActivePlan( } } + ApplySettleActions(&replace_actions); + // Publish the plan last: pushes resolve through it, so everything they // may touch exists by the time it is visible. absl::Status registered = kv_cache::KVCacheManagerBase::RegisterActivePlan( @@ -828,6 +872,21 @@ absl::Status KVCacheManagerWithTransfer::RegisterActivePlan( absl::Status KVCacheManagerWithTransfer::RegisterRecv( uint64_t uuid, const std::string& req_id, int64_t expected_block_count) { + SettleActions replaced; + { + absl::MutexLock lock(mu_); + if (active_recv_entries_.contains(uuid)) { + // A live entry re-registered under the same uuid is failed and + // drained first; silently replacing it would orphan its staging and + // strand its callbacks' leases. + LOG(ERROR) << "RegisterRecv: uuid " << uuid + << " re-registered while still active; failing the old entry"; + SettleRecvLocked(uuid, /*failed=*/true, &replaced); + } + } + ApplySettleActions(&replaced); + // A fresh registration under a previously failed uuid serves again. + ReviveTransferUuid(uuid); absl::MutexLock lock(mu_); RecvEntry recv_entry; recv_entry.req_id = req_id; @@ -1358,13 +1417,20 @@ absl::Status KVCacheManagerWithTransfer::PoolReshardPush( } raiden::PjRtCopyFuture future = std::move(future_or).value(); state->d2h_futures.push_back(future); + { + absl::MutexLock lock(mu_); + ++total_outstanding_ops_; + } future.OnReady([this, uuid = static_cast(plan.uuid()), pool_idx](auto status_or) { if (!status_or.ok()) { FinishPoolReshardSend(uuid, status_or.status()); - return; + } else { + StartPoolReshardPush(uuid, pool_idx); } - StartPoolReshardPush(uuid, pool_idx); + // The callback reads manager state; destruction waits for it. + absl::MutexLock lock(mu_); + --total_outstanding_ops_; }); } return absl::OkStatus(); @@ -1812,13 +1878,20 @@ void KVCacheManagerWithTransfer::StartRead( LOG(ERROR) << "Raiden consumer error during Hybrid Bridge StartRead connect: " << e.what(); - absl::MutexLock lock(mu_); - failed_recving_.insert(req_id); - auto it = active_recv_entries_.find(uuid); - if (it != active_recv_entries_.end()) { - ReleaseRecvStagingLocked(&it->second); - active_recv_entries_.erase(it); + // The producer may have accepted the request before the failure and + // may already be pushing; the receive settles like any other failure, + // so its staging outlives whatever that push still lands. + SettleActions actions; + { + absl::MutexLock lock(mu_); + auto it = active_recv_entries_.find(uuid); + if (it != active_recv_entries_.end()) { + SettleRecvLocked(uuid, /*failed=*/true, &actions); + } else { + failed_recving_.insert(req_id); + } } + ApplySettleActions(&actions); } }); } @@ -1830,22 +1903,23 @@ KVCacheManagerWithTransfer::CompleteReadRaw() { std::vector done_recving; std::vector failed_recving; std::vector> settled_plans; + std::vector settle_actions; { absl::MutexLock lock(mu_); const auto now = std::chrono::steady_clock::now(); - for (auto it = send_entries_.begin(); it != send_entries_.end();) { - const auto& entry = it->second; + std::vector expired_sends; + for (const auto& [uuid, entry] : send_entries_) { if (entry->deadline <= now) { - // Nothing pulled this entry within its deadline; the bytes were - // never sent, so the transfer failed. - failed_recving_.insert(entry->req_id); - ReleaseEntrySlotLocked(entry); - settled_plans.emplace_back(it->first, 0); - it = send_entries_.erase(it); - } else { - ++it; + expired_sends.push_back(uuid); } } + for (uint64_t uuid : expired_sends) { + // Nothing pulled this entry within its deadline; the bytes were + // never sent, so the transfer failed. + settle_actions.emplace_back(); + SettleSendLocked(uuid, /*failed=*/true, &settle_actions.back()); + settled_plans.emplace_back(uuid, 0); + } for (auto it = active_pool_reshard_sends_.begin(); it != active_pool_reshard_sends_.end();) { const auto& entry = it->second; @@ -1862,9 +1936,8 @@ KVCacheManagerWithTransfer::CompleteReadRaw() { // died or never finished pushing). Without this the entry and its host // staging slot leak forever, eventually exhausting the slot pool. Surface // the timeout as a recv failure so the connector can recompute the blocks. - for (auto it = active_recv_entries_.begin(); - it != active_recv_entries_.end();) { - auto& entry = it->second; + std::vector> recv_outcomes; + for (auto& [uuid, entry] : active_recv_entries_) { if (entry.network_completed || entry.num_completed_layers == num_layers()) { bool all_h2d_done = true; @@ -1877,24 +1950,28 @@ KVCacheManagerWithTransfer::CompleteReadRaw() { if (all_h2d_done) { LOG(INFO) << "CompleteReadRaw (polling completion): req_id=" << entry.req_id; - done_recving_.insert(entry.req_id); - ReleaseRecvStagingLocked(&entry); - if (entry.unregister_on_settle) { - settled_plans.emplace_back(it->first, entry.plan_generation); - } - active_recv_entries_.erase(it++); + recv_outcomes.emplace_back(uuid, /*failed=*/false); continue; } } if (entry.deadline <= now) { - failed_recving_.insert(entry.req_id); - ReleaseRecvStagingLocked(&entry); - settled_plans.emplace_back(it->first, entry.plan_generation); - active_recv_entries_.erase(it++); - } else { - ++it; + recv_outcomes.emplace_back(uuid, /*failed=*/true); + } + } + for (const auto& [uuid, failed] : recv_outcomes) { + if (failed) { + // Settlement itself unregisters only plans that travel with the + // receive; a timed-out pool or fixed-slot receive still drops its + // plan here, as before. + auto eit = active_recv_entries_.find(uuid); + if (eit != active_recv_entries_.end() && + !eit->second.unregister_on_settle) { + settled_plans.emplace_back(uuid, eit->second.plan_generation); + } } + settle_actions.emplace_back(); + SettleRecvLocked(uuid, failed, &settle_actions.back()); } done_sending.assign(done_sending_.begin(), done_sending_.end()); done_recving.assign(done_recving_.begin(), done_recving_.end()); @@ -1903,6 +1980,9 @@ KVCacheManagerWithTransfer::CompleteReadRaw() { done_recving_.clear(); failed_recving_.clear(); } + for (auto& actions : settle_actions) { + ApplySettleActions(&actions); + } // Unregistering drops the plan and its transport receive-progress counters // (ForgetPushProgress), so a settled uuid is reusable. for (const auto& [uuid, generation] : settled_plans) { @@ -2139,6 +2219,221 @@ void KVCacheManagerWithTransfer::ReleaseEntrySlotLocked( entry->slot_released = true; } +uint64_t KVCacheManagerWithTransfer::EnsureOpTokenLocked(uint64_t* op_token) { + if (*op_token == 0) { + *op_token = ++op_token_counter_; + } + return *op_token; +} + +void KVCacheManagerWithTransfer::AddSendOpLocked( + const std::shared_ptr& entry) { + EnsureOpTokenLocked(&entry->op_token); + ++entry->outstanding_ops; + ++total_outstanding_ops_; +} + +void KVCacheManagerWithTransfer::AddRecvOpLocked(RecvEntry* entry) { + EnsureOpTokenLocked(&entry->op_token); + ++entry->outstanding_ops; + ++total_outstanding_ops_; +} + +void KVCacheManagerWithTransfer::ReapSendLocked( + const std::shared_ptr& entry) { + // A failed transfer's uuid is retired before its staging returns, so a + // late pull or payload cannot resolve chunks into blocks that may + // already belong to someone else. + if (entry->failed) { + RetireTransferUuid(entry->uuid); + } + ReleaseEntrySlotLocked(entry); +} + +void KVCacheManagerWithTransfer::SettleSendLocked(uint64_t uuid, bool failed, + SettleActions* actions) { + auto it = send_entries_.find(uuid); + if (it == send_entries_.end()) { + return; // already settled on another path + } + std::shared_ptr entry = it->second; + send_entries_.erase(it); + entry->failed = failed; + if (failed) { + failed_recving_.insert(entry->req_id); + } else { + done_sending_.insert(entry->req_id); + } + actions->uuid = uuid; + actions->is_sender = true; + actions->op_token = EnsureOpTokenLocked(&entry->op_token); + // Futures beyond the chained prefix have no completion callback; hand + // them drain-only ones so the leases they hold can fall. + for (size_t i = entry->chained_d2h_layers; + i < entry->d2h_layer_futures.size(); ++i) { + actions->drain_futures.push_back(entry->d2h_layer_futures[i]); + } + entry->chained_d2h_layers = entry->d2h_layer_futures.size(); + if (entry->outstanding_ops == 0) { + ReapSendLocked(entry); + return; + } + const uint64_t token = entry->op_token; + entry->draining = true; + draining_sends_[token] = std::move(entry); +} + +void KVCacheManagerWithTransfer::FinishSendOpLocked(uint64_t uuid, + uint64_t op_token, + SettleActions* actions) { + std::shared_ptr entry; + if (auto it = send_entries_.find(uuid); + it != send_entries_.end() && it->second->op_token == op_token) { + entry = it->second; + } else if (auto dit = draining_sends_.find(op_token); + dit != draining_sends_.end()) { + entry = dit->second; + } + if (!entry) { + return; + } + --entry->outstanding_ops; + --total_outstanding_ops_; + if (entry->draining && entry->outstanding_ops == 0) { + draining_sends_.erase(entry->op_token); + actions->uuid = uuid; + actions->is_sender = true; + ReapSendLocked(entry); + } +} + +void KVCacheManagerWithTransfer::SettleRecvLocked(uint64_t uuid, bool failed, + SettleActions* actions) { + auto it = active_recv_entries_.find(uuid); + if (it == active_recv_entries_.end()) { + return; // already settled on another path + } + RecvEntry& entry = it->second; + entry.settled_failed = failed; + actions->uuid = uuid; + actions->is_sender = false; + if (entry.outstanding_ops == 0) { + // The outcome becomes visible only here, when neither the staging nor + // the device pages it names have accepted work left against them, and + // a failed uuid is retired before any of its blocks can be reused. + (failed ? failed_recving_ : done_recving_).insert(entry.req_id); + if (failed) { + RetireTransferUuid(uuid); + } + ReleaseRecvStagingLocked(&entry); + if (entry.unregister_on_settle) { + actions->unregister_plan = {uuid, entry.plan_generation}; + } + active_recv_entries_.erase(it); + return; + } + entry.draining = true; + entry.uuid = uuid; + const uint64_t token = EnsureOpTokenLocked(&entry.op_token); + actions->op_token = token; + draining_recvs_[token] = std::move(entry); + active_recv_entries_.erase(it); +} + +void KVCacheManagerWithTransfer::FinishRecvOpLocked(uint64_t uuid, + uint64_t op_token, + SettleActions* actions) { + RecvEntry* entry = nullptr; + bool in_draining = false; + if (auto dit = draining_recvs_.find(op_token); + dit != draining_recvs_.end()) { + entry = &dit->second; + in_draining = true; + } else if (auto it = active_recv_entries_.find(uuid); + it != active_recv_entries_.end() && + it->second.op_token == op_token) { + entry = &it->second; + } + if (entry == nullptr) { + return; + } + --entry->outstanding_ops; + --total_outstanding_ops_; + if (in_draining && entry->outstanding_ops == 0) { + (entry->settled_failed ? failed_recving_ : done_recving_) + .insert(entry->req_id); + if (entry->settled_failed) { + RetireTransferUuid(entry->uuid); + } + ReleaseRecvStagingLocked(entry); + actions->uuid = uuid; + actions->is_sender = false; + if (entry->unregister_on_settle) { + actions->unregister_plan = {uuid, entry->plan_generation}; + } + draining_recvs_.erase(op_token); + } +} + +void KVCacheManagerWithTransfer::ApplySettleActions(SettleActions* actions) { + for (auto& future : actions->drain_futures) { + future.OnReady([this, uuid = actions->uuid, token = actions->op_token, + is_sender = actions->is_sender](auto status_or) { + (void)status_or; + DrainOp(uuid, token, is_sender); + }); + } + actions->drain_futures.clear(); + if (actions->unregister_plan.has_value()) { + UnregisterSettledPlan(actions->unregister_plan->first, + actions->unregister_plan->second); + actions->unregister_plan.reset(); + } +} + +void KVCacheManagerWithTransfer::DrainOp(uint64_t uuid, uint64_t op_token, + bool is_sender) { + SettleActions actions; + { + absl::MutexLock lock(mu_); + if (is_sender) { + FinishSendOpLocked(uuid, op_token, &actions); + } else { + FinishRecvOpLocked(uuid, op_token, &actions); + } + } + ApplySettleActions(&actions); +} + +uint64_t KVCacheManagerWithTransfer::BeginPayloadResolution(uint64_t uuid) { + absl::MutexLock lock(mu_); + if (auto it = active_recv_entries_.find(uuid); + it != active_recv_entries_.end()) { + // The payload's bytes land in this entry's staging; the lease keeps + // that staging owned until EndPayloadResolution. + AddRecvOpLocked(&it->second); + return it->second.op_token; + } + // A settled entry that is still draining keeps its staging until every + // accepted operation ends; a payload resolving through its plan is one. + for (auto& [token, entry] : draining_recvs_) { + if (entry.uuid == uuid) { + ++entry.outstanding_ops; + ++total_outstanding_ops_; + return token; + } + } + return 0; +} + +void KVCacheManagerWithTransfer::EndPayloadResolution(uint64_t uuid, + uint64_t token) { + if (token == 0) { + return; + } + DrainOp(uuid, token, /*is_sender=*/false); +} + std::shared_ptr KVCacheManagerWithTransfer::CreateStagingReadiness(int64_t slot_idx, int64_t num_blocks) { @@ -2306,6 +2601,14 @@ void KVCacheManagerWithTransfer::StopControlServer() { shutdown(control_fd_, SHUT_RDWR); close(control_fd_); } + // Unblock handlers parked in reads on accepted connections; each handler + // still closes its own socket when it returns. + { + absl::MutexLock lock(pull_workers_mu_); + for (int fd : accepted_control_fds_) { + shutdown(fd, SHUT_RDWR); + } + } if (control_thread_.joinable()) { control_thread_.join(); } @@ -2330,8 +2633,33 @@ void KVCacheManagerWithTransfer::ControlServerLoop() { } std::optional source_node = assigned_numa_node(); + { + absl::MutexLock lock(pull_workers_mu_); + if (shutting_down_.load(std::memory_order_relaxed)) { + close(client_fd); + continue; + } + accepted_control_fds_.insert(client_fd); + } pull_pool_->Schedule(source_node, [this, client_fd]() { + // Handler admission is counted so destruction waits for handlers, not + // only for the workers they spawn; a handler admitted after shutdown + // began must not touch manager state at all. + { + absl::MutexLock lock(pull_workers_mu_); + if (shutting_down_.load(std::memory_order_relaxed)) { + accepted_control_fds_.erase(client_fd); + close(client_fd); + return; + } + ++active_pull_workers_; + } HandleControlConnection(client_fd); + { + absl::MutexLock lock(pull_workers_mu_); + accepted_control_fds_.erase(client_fd); + --active_pull_workers_; + } close(client_fd); }); } @@ -2393,6 +2721,25 @@ void KVCacheManagerWithTransfer::ProcessPullStream( std::vector dst_block_ids = ReadBlockIds(fd, req.num_blocks); ValidateRequestedBlocks(*entry, src_block_ids); + // Admission comes before the acknowledgement: a duplicate pull, a settled + // transfer, or a stopping manager is refused here instead of being told + // to expect a payload that will never come. + { + absl::MutexLock lock(mu_); + auto it = send_entries_.find(req.uuid); + if (it == send_entries_.end() || it->second != entry) { + throw std::runtime_error("pull request refused: transfer settled"); + } + if (it->second->pull_started) { + throw std::runtime_error( + "pull request refused: transfer is already being served"); + } + if (shutting_down_.load(std::memory_order_relaxed)) { + throw std::runtime_error("pull request refused: manager is stopping"); + } + it->second->pull_started = true; + } + // Acknowledge acceptance to consumer immediately ControlResponseHeader response; response.magic = kResponseMagic; @@ -2460,19 +2807,13 @@ void KVCacheManagerWithTransfer::ProcessPullStream( << (remote_data_endpoints.empty() ? "" : remote_data_endpoints[0]) << (remote_data_endpoints.size() > 1 ? " and others" : ""); - { - absl::MutexLock lock(mu_); - if (auto it = send_entries_.find(req.uuid); it != send_entries_.end()) { - if (it->second->pull_started) { - VLOG(1) << "StartPushInternal already running for UUID: " << req.uuid; - return; - } - it->second->pull_started = true; - } - } - { absl::MutexLock lock(pull_workers_mu_); + if (shutting_down_.load(std::memory_order_relaxed)) { + // The manager is being destroyed; the consumer sees its connection + // drop and fails the transfer on its own deadline. + return; + } ++active_pull_workers_; } std::thread([this, uuid = req.uuid, remote_data_endpoints, src_block_ids, @@ -2495,6 +2836,8 @@ bool KVCacheManagerWithTransfer::AcquireSendStagingWithRetry( if (shutting_down_.load(std::memory_order_relaxed)) { return false; // the manager is being destroyed; its state goes with it } + bool never_fits = false; + SettleActions actions; { absl::MutexLock lock(mu_); auto it = send_entries_.find(uuid); @@ -2514,40 +2857,45 @@ bool KVCacheManagerWithTransfer::AcquireSendStagingWithRetry( << (dynamic_host_staging_ ? "the host staging pool holds " : "a staging slot holds ") << capacity; - failed_recving_.insert(it->second->req_id); - ReleaseEntrySlotLocked(it->second); - send_entries_.erase(it); - return false; - } - RecvEntry staging; - auto staged = AcquireRecvStagingLocked( - static_cast(src_block_ids.size()), &staging); - if (staged.has_value()) { - it->second->slot_idx = staging.slot_idx; - it->second->staged_host_blocks = std::move(staging.staged_host_blocks); - *host_block_ids = std::move(*staged); - return true; + SettleSendLocked(uuid, /*failed=*/true, &actions); + never_fits = true; + } else { + RecvEntry staging; + auto staged = AcquireRecvStagingLocked( + static_cast(src_block_ids.size()), &staging); + if (staged.has_value()) { + it->second->slot_idx = staging.slot_idx; + it->second->staged_host_blocks = + std::move(staging.staged_host_blocks); + *host_block_ids = std::move(*staged); + return true; + } } } + if (never_fits) { + ApplySettleActions(&actions); + return false; + } // Staging exhausted: wait for in-flight sends to hand blocks back // instead of reporting a send that never happened. The consumer's own // deadline still bounds the total wait. if (std::chrono::steady_clock::now() >= stage_deadline) { - absl::MutexLock lock(mu_); - auto it = send_entries_.find(uuid); - if (it != send_entries_.end()) { - LOG(ERROR) << "StartPushInternal: staging exhausted serving " - << it->second->req_id << " (" << src_block_ids.size() - << " blocks; free_host_blocks=" - << host_block_manager_->num_free_blocks() - << ", total_host_blocks=" - << host_block_manager_->total_blocks() - << ", free_slots=" << free_slots_.size() - << "); reporting transfer failure"; - failed_recving_.insert(it->second->req_id); - ReleaseEntrySlotLocked(it->second); - send_entries_.erase(it); + { + absl::MutexLock lock(mu_); + auto it = send_entries_.find(uuid); + if (it != send_entries_.end()) { + LOG(ERROR) << "StartPushInternal: staging exhausted serving " + << it->second->req_id << " (" << src_block_ids.size() + << " blocks; free_host_blocks=" + << host_block_manager_->num_free_blocks() + << ", total_host_blocks=" + << host_block_manager_->total_blocks() + << ", free_slots=" << free_slots_.size() + << "); reporting transfer failure"; + SettleSendLocked(uuid, /*failed=*/true, &actions); + } } + ApplySettleActions(&actions); return false; } std::this_thread::sleep_for(std::chrono::milliseconds(1)); @@ -2587,13 +2935,27 @@ void KVCacheManagerWithTransfer::StartPushInternal( } CopySpec d2h_copy = BuildCoalescedCopySpec(src_block_ids, host_block_ids); - entry->d2h_layer_futures.reserve(num_layers()); + { + absl::MutexLock lock(mu_); + entry->d2h_layer_futures.reserve(num_layers()); + } // 1. Issue D2H copies layer-by-layer! for (size_t l = 0; l < num_layers(); ++l) { LOG(INFO) << "StartPushInternal (D2H start) layer " << l << ": uuid=" << uuid << ", numa=" << assigned_numa_node().value_or(-1); + // The copy about to be issued reads the entry's staging; the lease + // taken here keeps that staging owned even if a timeout settles the + // entry mid-loop. + { + absl::MutexLock lock(mu_); + auto it = send_entries_.find(uuid); + if (it == send_entries_.end() || it->second != entry) { + return; // settled while issuing; leases cover the copies so far + } + AddSendOpLocked(entry); + } auto future_or = D2hSyncDispatch(d2h_copy.src_offsets, d2h_copy.dst_offsets, d2h_copy.sizes, /*slot_idx=*/std::nullopt, /*layer_idx=*/l); @@ -2602,16 +2964,34 @@ void KVCacheManagerWithTransfer::StartPushInternal( // thread has no caller for an exception to reach. LOG(ERROR) << "StartPushInternal: failed to issue D2H for layer " << l << ": " << future_or.status(); + SettleActions actions; + { + absl::MutexLock lock(mu_); + SettleSendLocked(uuid, /*failed=*/true, &actions); + FinishSendOpLocked(uuid, entry->op_token, &actions); + } + ApplySettleActions(&actions); + return; + } + bool published = false; + { absl::MutexLock lock(mu_); auto it = send_entries_.find(uuid); - if (it != send_entries_.end()) { - failed_recving_.insert(it->second->req_id); - ReleaseEntrySlotLocked(it->second); - send_entries_.erase(it); - } + if (it != send_entries_.end() && it->second == entry) { + entry->d2h_layer_futures.push_back(future_or.value()); + published = true; + } + } + if (!published) { + // The entry settled between the lease and this publication; nothing + // will chain this copy, so it hands its lease back itself. + future_or.value().OnReady( + [this, uuid, token = entry->op_token](auto status_or) { + (void)status_or; + DrainOp(uuid, token, /*is_sender=*/true); + }); return; } - entry->d2h_layer_futures.push_back(std::move(future_or.value())); } entry->remote_data_endpoints = remote_data_endpoints; @@ -2631,6 +3011,10 @@ void KVCacheManagerWithTransfer::SendNextLayer(uint64_t uuid, size_t l) { return; } entry = it->second; + if (l < entry->d2h_layer_futures.size() && + l >= entry->chained_d2h_layers) { + entry->chained_d2h_layers = l + 1; + } } if (l >= num_layers()) { @@ -2639,56 +3023,80 @@ void KVCacheManagerWithTransfer::SendNextLayer(uint64_t uuid, size_t l) { return; } - entry->d2h_layer_futures[l].OnReady([this, uuid, l](auto status_or) { + const uint64_t op_token = entry->op_token; + entry->d2h_layer_futures[l].OnReady([this, uuid, l, + op_token](auto status_or) { if (!status_or.ok()) { LOG(ERROR) << "StartPushInternal: D2H copy failed for layer " << l << ", status: " << status_or.status().ToString(); + SettleActions actions; + { + absl::MutexLock lock(mu_); + SettleSendLocked(uuid, /*failed=*/true, &actions); + FinishSendOpLocked(uuid, op_token, &actions); + } + ApplySettleActions(&actions); + return; + } + + bool launch = false; + SettleActions actions; + { absl::MutexLock lock(mu_); auto it = send_entries_.find(uuid); - if (it != send_entries_.end()) { - failed_recving_.insert(it->second->req_id); - ReleaseEntrySlotLocked(it->second); - send_entries_.erase(it); - } + if (it != send_entries_.end() && it->second->op_token == op_token) { + // The push about to be queued reads the entry's staging; it takes a + // lease of its own, released by its completion callback. + AddSendOpLocked(it->second); + ++pending_h2h_launches_; + launch = true; + } + // The completed copy's own lease. + FinishSendOpLocked(uuid, op_token, &actions); + } + ApplySettleActions(&actions); + if (!launch) { return; } - push_pool_->Schedule([this, uuid, l]() { + push_pool_->Schedule([this, uuid, l, op_token]() { std::shared_ptr entry; + SettleActions actions; { absl::MutexLock lock(mu_); auto it = send_entries_.find(uuid); - if (it == send_entries_.end()) { - return; + if (it != send_entries_.end() && it->second->op_token == op_token && + !shutting_down_.load(std::memory_order_relaxed)) { + entry = it->second; + } else { + // Settled or shutting down between queueing and launch: the push + // never starts, so its lease falls here. + --pending_h2h_launches_; + FinishSendOpLocked(uuid, op_token, &actions); } - entry = it->second; + } + ApplySettleActions(&actions); + if (!entry) { + return; } LOG(INFO) << "StartPushInternal (H2H start layer " << l << "): uuid=" << uuid << ", numa=" << assigned_numa_node().value_or(-1); H2hWriteDirectAsync( entry->remote_data_endpoints, entry->src_ints, entry->dst_ints, uuid, - l, [this, uuid, l](absl::StatusOr> push_res) { - std::shared_ptr entry; - { - absl::MutexLock lock(mu_); - auto it = send_entries_.find(uuid); - if (it != send_entries_.end()) { - entry = it->second; - } - } - if (!entry) return; - + l, + [this, uuid, l, + op_token](absl::StatusOr> push_res) { + SettleActions actions; if (!push_res.ok()) { LOG(ERROR) << "H2hWrite failed for layer " << l << ": " << push_res.status().ToString(); - absl::MutexLock lock(mu_); - if (auto it = send_entries_.find(uuid); - it != send_entries_.end()) { - failed_recving_.insert(entry->req_id); - ReleaseEntrySlotLocked(entry); - send_entries_.erase(it); + { + absl::MutexLock lock(mu_); + SettleSendLocked(uuid, /*failed=*/true, &actions); + FinishSendOpLocked(uuid, op_token, &actions); } + ApplySettleActions(&actions); return; } @@ -2696,21 +3104,36 @@ void KVCacheManagerWithTransfer::SendNextLayer(uint64_t uuid, size_t l) { << "): uuid=" << uuid << ", numa=" << assigned_numa_node().value_or(-1); - if (entry->remaining_h2h_layers.fetch_sub(1) == 1) { - LOG(INFO) << "StartPushInternal (All H2H complete): uuid=" - << uuid; + std::shared_ptr entry; + { absl::MutexLock lock(mu_); - if (auto it = send_entries_.find(uuid); - it != send_entries_.end()) { - done_sending_.insert(entry->req_id); - ReleaseEntrySlotLocked(entry); - send_entries_.erase(it); + auto it = send_entries_.find(uuid); + if (it != send_entries_.end() && + it->second->op_token == op_token) { + entry = it->second; } } + const bool all_layers_done = + entry != nullptr && + entry->remaining_h2h_layers.fetch_sub(1) == 1; + { + absl::MutexLock lock(mu_); + if (all_layers_done) { + LOG(INFO) << "StartPushInternal (All H2H complete): uuid=" + << uuid; + SettleSendLocked(uuid, /*failed=*/false, &actions); + } + FinishSendOpLocked(uuid, op_token, &actions); + } + ApplySettleActions(&actions); }); + { + absl::MutexLock lock(mu_); + --pending_h2h_launches_; + } - // Immediately queue the next layer's push without waiting for this one to - // finish + // Immediately queue the next layer's push without waiting for this one + // to finish SendNextLayer(uuid, l + 1); }); }); @@ -2737,7 +3160,18 @@ absl::Status KVCacheManagerWithTransfer::WaitForPendingWork() { } if (recv_pending) break; } - if (!recv_pending && active_pool_reshard_sends_.empty()) { + bool send_pending = false; + for (const auto& [uuid, entry] : send_entries_) { + (void)uuid; + if (entry->pull_started) { + send_pending = true; + break; + } + } + if (!recv_pending && !send_pending && + active_pool_reshard_sends_.empty() && draining_sends_.empty() && + draining_recvs_.empty() && total_outstanding_ops_ == 0 && + pending_h2h_launches_ == 0) { break; } const absl::Duration elapsed = absl::Now() - start; @@ -2876,15 +3310,12 @@ absl::Status KVCacheManagerWithTransfer::OnBlocksReceived( << uuid << ", received blocks count: " << block_ids.size(); std::string req_id; - int64_t recv_slot = -1; - std::vector recv_staged; CopySpec h2d_copy; absl::flat_hash_map host_to_chip; bool found = false; - bool unregister_plan = false; - uint64_t plan_generation = 0; std::vector accumulated_host_blocks; + SettleActions actions; std::chrono::steady_clock::time_point start_time; bool should_record_duration = false; @@ -2909,7 +3340,6 @@ absl::Status KVCacheManagerWithTransfer::OnBlocksReceived( it->second.total_blocks * num_layers()) { it->second.network_completed = true; req_id = it->second.req_id; - recv_slot = it->second.slot_idx; if (metrics_collector_) { metrics_collector_->RecordLastPacket(uuid); } @@ -2921,10 +3351,7 @@ absl::Status KVCacheManagerWithTransfer::OnBlocksReceived( metrics_collector_->RecordEnd(uuid); } found = true; - recv_staged = std::move(it->second.staged_host_blocks); - unregister_plan = it->second.unregister_on_settle; - plan_generation = it->second.plan_generation; - active_recv_entries_.erase(it); + SettleRecvLocked(uuid, /*failed=*/false, &actions); } } else { VLOG(1) << "OnBlocksReceived: Partial blocks received for uuid " << uuid @@ -2945,12 +3372,7 @@ absl::Status KVCacheManagerWithTransfer::OnBlocksReceived( return RaidenManagerBase::OnBlocksReceived(block_ids, uuid); } - { - absl::MutexLock lock(mu_); - done_recving_.insert(req_id); - ReleaseStagingLocked(recv_slot, &recv_staged); - } - if (unregister_plan) UnregisterSettledPlan(uuid, plan_generation); + ApplySettleActions(&actions); LOG(INFO) << "OnBlocksReceived (Network + H2D complete): req_id=" << req_id << ", uuid=" << uuid @@ -3031,10 +3453,17 @@ void KVCacheManagerWithTransfer::LaunchEligiblePoolH2ds(uint64_t uuid) { continue; } raiden::PjRtCopyFuture future = std::move(future_or).value(); + { + absl::MutexLock lock(mu_); + ++total_outstanding_ops_; + } future.OnReady([this, uuid, pool_idx = pool_idx](auto status_or) { FinishPoolReshardRecvPool( uuid, pool_idx, status_or.ok() ? absl::OkStatus() : status_or.status()); + // The callback reads manager state; destruction waits for it. + absl::MutexLock lock(mu_); + --total_outstanding_ops_; }); { absl::MutexLock lock(mu_); @@ -3102,7 +3531,7 @@ absl::Status KVCacheManagerWithTransfer::OnLayerReceived(size_t layer_idx, uint64_t uuid) { CopySpec h2d_copy; std::string req_id; - int64_t recv_slot; + uint64_t op_token = 0; bool trigger_enqueue = false; { absl::MutexLock lock(mu_); @@ -3113,7 +3542,10 @@ absl::Status KVCacheManagerWithTransfer::OnLayerReceived(size_t layer_idx, auto& entry = it->second; h2d_copy = entry.h2d_copy; req_id = entry.req_id; - recv_slot = entry.slot_idx; + // The copy about to be issued writes into the entry's staging; the + // lease taken here keeps that staging owned until the copy completes. + AddRecvOpLocked(&entry); + op_token = entry.op_token; if (!entry.h2d_started) { entry.h2d_started = true; trigger_enqueue = true; @@ -3131,39 +3563,32 @@ absl::Status KVCacheManagerWithTransfer::OnLayerReceived(size_t layer_idx, h2d_copy.sizes, /*slot_idx=*/std::nullopt, /*layer_idx=*/layer_idx); if (!future_or.ok()) { - bool unregister_plan = false; - uint64_t plan_generation = 0; + SettleActions actions; { absl::MutexLock lock(mu_); - failed_recving_.insert(req_id); - auto it = active_recv_entries_.find(uuid); - if (it != active_recv_entries_.end()) { - ReleaseRecvStagingLocked(&it->second); - unregister_plan = it->second.unregister_on_settle; - plan_generation = it->second.plan_generation; - active_recv_entries_.erase(it); - } + SettleRecvLocked(uuid, /*failed=*/true, &actions); + FinishRecvOpLocked(uuid, op_token, &actions); } - if (unregister_plan) UnregisterSettledPlan(uuid, plan_generation); + ApplySettleActions(&actions); return future_or.status(); } auto future = future_or.value(); - future.OnReady([this, uuid, layer_idx, recv_slot, req_id, + future.OnReady([this, uuid, layer_idx, op_token, req_id, metrics_collector = metrics_collector_](auto status_or) { - bool unregister_plan = false; - uint64_t plan_generation = 0; + SettleActions actions; { absl::MutexLock lock(mu_); auto it = active_recv_entries_.find(uuid); - if (it == active_recv_entries_.end()) { - return; - } - auto& entry = it->second; - if (status_or.ok()) { + if (it == active_recv_entries_.end() || + it->second.op_token != op_token) { + // Settled while the copy ran; only the copy's lease falls here. + FinishRecvOpLocked(uuid, op_token, &actions); + } else if (status_or.ok()) { LOG(INFO) << "OnLayerReceived (H2D copy complete) layer " << layer_idx << ": req_id=" << req_id << ", numa=" << assigned_numa_node().value_or(-1); + auto& entry = it->second; entry.num_completed_layers++; if (entry.num_completed_layers == num_layers()) { // TODO: Find a way to optimize this by moving out of the mutex. @@ -3177,24 +3602,18 @@ absl::Status KVCacheManagerWithTransfer::OnLayerReceived(size_t layer_idx, if (metrics_collector) { metrics_collector->RecordEnd(uuid); } - done_recving_.insert(req_id); - ReleaseRecvStagingLocked(&entry); - unregister_plan = entry.unregister_on_settle; - plan_generation = entry.plan_generation; - active_recv_entries_.erase(uuid); + SettleRecvLocked(uuid, /*failed=*/false, &actions); } + FinishRecvOpLocked(uuid, op_token, &actions); } else { LOG(ERROR) << "OnLayerReceived (H2D copy failed) layer " << layer_idx << " for req_id: " << req_id << ", error: " << status_or.status().ToString(); - failed_recving_.insert(req_id); - ReleaseRecvStagingLocked(&entry); - unregister_plan = entry.unregister_on_settle; - plan_generation = entry.plan_generation; - active_recv_entries_.erase(uuid); + SettleRecvLocked(uuid, /*failed=*/true, &actions); + FinishRecvOpLocked(uuid, op_token, &actions); } } - if (unregister_plan) UnregisterSettledPlan(uuid, plan_generation); + ApplySettleActions(&actions); }); { diff --git a/tpu_sync/core/kv_cache_manager_with_transfer.h b/tpu_sync/core/kv_cache_manager_with_transfer.h index f17ceed3..f6878ad2 100644 --- a/tpu_sync/core/kv_cache_manager_with_transfer.h +++ b/tpu_sync/core/kv_cache_manager_with_transfer.h @@ -36,6 +36,7 @@ #include #include "absl/container/flat_hash_map.h" +#include "absl/container/flat_hash_set.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/synchronization/mutex.h" @@ -281,6 +282,18 @@ class KVCacheManagerWithTransfer : public kv_cache::KVCacheManagerBase { std::vector src_ints; std::vector dst_ints; std::atomic remaining_h2h_layers{0}; + // Identity of this entry in draining bookkeeping, assigned with its + // first accepted operation; a reused uuid cannot confuse late callbacks. + uint64_t op_token = 0; + // Accepted asynchronous operations (issued D2H copies, H2H pushes) that + // still read this entry's staging. + int outstanding_ops = 0; + // Layers whose D2H future has a completion callback attached; an early + // settle hands the remaining futures drain-only callbacks. + size_t chained_d2h_layers = 0; + // The entry has settled (its outcome is published) and waits only for + // outstanding operations to drain before its staging is released. + bool draining = false; }; struct StagingLayerReady { @@ -368,6 +381,31 @@ class KVCacheManagerWithTransfer : public kv_cache::KVCacheManagerBase { std::vector* staged_host_blocks); void ReleaseEntrySlotLocked(const std::shared_ptr& entry); + // Deferred reclamation: a settle publishes the entry's outcome at once, + // but its staging is released only when every accepted operation has + // drained. Work that must run after mu_ is dropped (attaching drain + // callbacks, retiring uuids, unregistering settled plans) is handed back + // in a SettleActions the caller applies. + struct SettleActions { + uint64_t uuid = 0; + uint64_t op_token = 0; + bool is_sender = false; + std::vector drain_futures; + std::optional> unregister_plan; + }; + uint64_t EnsureOpTokenLocked(uint64_t* op_token); + void AddSendOpLocked(const std::shared_ptr& entry); + void AddRecvOpLocked(RecvEntry* entry); + void ReapSendLocked(const std::shared_ptr& entry); + void SettleSendLocked(uint64_t uuid, bool failed, SettleActions* actions); + void FinishSendOpLocked(uint64_t uuid, uint64_t op_token, + SettleActions* actions); + void SettleRecvLocked(uint64_t uuid, bool failed, SettleActions* actions); + void FinishRecvOpLocked(uint64_t uuid, uint64_t op_token, + SettleActions* actions); + void ApplySettleActions(SettleActions* actions); + void DrainOp(uint64_t uuid, uint64_t op_token, bool is_sender); + void StartControlServer(); void StopControlServer(); void ControlServerLoop(); @@ -380,6 +418,8 @@ class KVCacheManagerWithTransfer : public kv_cache::KVCacheManagerBase { size_t layer_idx, size_t shard_idx, int block_id, uint64_t uuid, transport::BlockTransportDelegate::HostBlockReadyCallback cb) override; void ScheduleAsyncTask(std::function task) override; + uint64_t BeginPayloadResolution(uint64_t uuid) override; + void EndPayloadResolution(uint64_t uuid, uint64_t token) override; std::shared_ptr CreateStagingReadiness( int64_t slot_idx, int64_t num_blocks); void MarkStagingLayerReady( @@ -438,8 +478,28 @@ class KVCacheManagerWithTransfer : public kv_cache::KVCacheManagerBase { // Multi-tag plans: each pool uploads only its own group's destination // block ids (the flat chip_block_ids list concatenates all groups). std::map> pool_dst_block_ids; + // The transfer this entry belongs to; set when the entry starts + // draining, so a payload lease can still find it. + uint64_t uuid = 0; + // Identity of this entry in draining bookkeeping, assigned with its + // first accepted operation; a reused uuid cannot confuse late callbacks. + uint64_t op_token = 0; + // Accepted asynchronous operations (issued H2D copies, payload reads) + // that still write into this entry's staging. + int outstanding_ops = 0; + // The entry has settled (its outcome is published) and waits only for + // outstanding operations to drain before its staging is released. + bool draining = false; + // Whether the published outcome was a failure; a drained failed entry + // also retires its uuid. + bool settled_failed = false; }; absl::flat_hash_map active_recv_entries_; + // Entries that settled while accepted operations still hold their staging, + // keyed by op_token; reaped when the last operation drains. + absl::flat_hash_map> draining_sends_ + ABSL_GUARDED_BY(mu_); + absl::flat_hash_map draining_recvs_ ABSL_GUARDED_BY(mu_); struct PoolReshardSendEntry { std::string req_id; @@ -509,6 +569,10 @@ class KVCacheManagerWithTransfer : public kv_cache::KVCacheManagerBase { std::atomic shutting_down_{false}; absl::Mutex pull_workers_mu_; int active_pull_workers_ ABSL_GUARDED_BY(pull_workers_mu_) = 0; + // Control sockets whose handler is queued or running; shut down on stop + // so a handler blocked in a read unwinds and its pool can join. + absl::flat_hash_set accepted_control_fds_ + ABSL_GUARDED_BY(pull_workers_mu_); double timeout_s_ = 120.0; bool unsafe_skip_buffer_lock_ = true; @@ -532,6 +596,14 @@ class KVCacheManagerWithTransfer : public kv_cache::KVCacheManagerBase { active_producer_blocks_; absl::Mutex mu_; absl::CondVar cv_; + // Sum of outstanding operations across all entries, active and draining, + // plus pool copy callbacks; destruction waits for it to reach zero. + int total_outstanding_ops_ ABSL_GUARDED_BY(mu_) = 0; + // H2H pushes between their queueing decision and their hand-off to the + // transport; destruction stops the transport only once this is zero. + int pending_h2h_launches_ ABSL_GUARDED_BY(mu_) = 0; + // Source of entry op_token values. + uint64_t op_token_counter_ ABSL_GUARDED_BY(mu_) = 0; int control_fd_ = -1; std::atomic stopping_{false}; std::thread control_thread_; diff --git a/tpu_sync/core/kv_cache_manager_with_transfer_pool_reshard_test.cc b/tpu_sync/core/kv_cache_manager_with_transfer_pool_reshard_test.cc index 42ef597b..70720086 100644 --- a/tpu_sync/core/kv_cache_manager_with_transfer_pool_reshard_test.cc +++ b/tpu_sync/core/kv_cache_manager_with_transfer_pool_reshard_test.cc @@ -734,6 +734,42 @@ TEST(SendDeadlineTest, ExpiredSendEntryFailsInsteadOfReportingDone) { EXPECT_THAT(failed_recving, Contains("expired_send_req")); } +TEST(DrainingTest, PayloadLeaseDefersOutcomeStagingAndPlanUntilItEnds) { + TestManager manager(/*timeout_s=*/0.05); + manager.EnableDemandStaging(); + auto* pool = manager.host_block_manager(); + const int free_before = pool->num_free_blocks(); + ASSERT_TRUE(manager + .RegisterActivePlan( + 41, BlockPlan(41, {0, 1}, {2, 3}, MEMORY_TYPE_HBM), + /*is_sender=*/false) + .ok()); + EXPECT_EQ(pool->num_free_blocks(), free_before - 2); + transport::BlockTransportDelegate* delegate = &manager; + const uint64_t token = delegate->BeginPayloadResolution(41); + ASSERT_GT(token, 0u); + + // The transfer times out while the payload lease is open: no outcome is + // published, and the staging and the plan stay owned. + absl::SleepFor(absl::Milliseconds(120)); + auto during = manager.CompleteReadRaw(); + EXPECT_THAT(std::get<2>(during), IsEmpty()); + EXPECT_EQ(pool->num_free_blocks(), free_before - 2); + EXPECT_TRUE(manager.HasActivePlan(41)); + + // Ending the lease drains the transfer: the failure is published, the + // staging returns, the plan is gone, and a late payload resolves nothing. + delegate->EndPayloadResolution(41, token); + auto after = manager.CompleteReadRaw(); + EXPECT_THAT(std::get<2>(after), Contains("block_plan_req_41")); + EXPECT_EQ(pool->num_free_blocks(), free_before); + EXPECT_FALSE(manager.HasActivePlan(41)); + const int64_t dst_block = 2; + EXPECT_TRUE( + manager.GetBlockChunks(0, 0, absl::MakeConstSpan(&dst_block, 1), 16, 41) + .empty()); +} + TEST(DemandStagingTest, SenderPlanReturnsStagingOnUnregister) { TestManager manager; manager.EnableDemandStaging(); diff --git a/tpu_sync/core/raiden_manager_base.cc b/tpu_sync/core/raiden_manager_base.cc index 40862e34..95c1678c 100644 --- a/tpu_sync/core/raiden_manager_base.cc +++ b/tpu_sync/core/raiden_manager_base.cc @@ -97,6 +97,24 @@ RaidenManagerBase::~RaidenManagerBase() { } } +tpu_raiden::transport::BlockTransport* +RaidenManagerBase::transport_server_if_started() { + absl::MutexLock lock(server_init_mu_); + return server_.get(); +} + +void RaidenManagerBase::StopTransportServer() { + std::unique_ptr server; + { + absl::MutexLock lock(server_init_mu_); + server = std::move(server_); + } + // The destructor joins transport workers; it runs unlocked so a worker + // that lazily consults the server pointer on its way out cannot deadlock + // against a stop in progress. + server.reset(); +} + std::vector RaidenManagerBase::GetHostNics() const { return GetLocalHostNicAddresses(); } diff --git a/tpu_sync/core/raiden_manager_base.h b/tpu_sync/core/raiden_manager_base.h index 4936dc95..f341d09c 100644 --- a/tpu_sync/core/raiden_manager_base.h +++ b/tpu_sync/core/raiden_manager_base.h @@ -146,6 +146,12 @@ class RaidenManagerBase : public tpu_raiden::transport::BlockTransportDelegate { std::vector local_ips_; tpu_raiden::transport::BlockTransport* InitTransportServer(); + // The data transport, or nullptr when no transfer has started one. + tpu_raiden::transport::BlockTransport* transport_server_if_started(); + // Stops the data transport and joins its workers, so transport threads + // stop calling into the delegate before the state they use is torn down. + // Idempotent; a later transfer would lazily start a fresh transport. + void StopTransportServer(); virtual std::vector GetHostNics() const; void DetectAndAssignNumaNode( diff --git a/tpu_sync/kv_cache/kv_cache_manager_base.cc b/tpu_sync/kv_cache/kv_cache_manager_base.cc index f9f56a1b..964bf3f5 100644 --- a/tpu_sync/kv_cache/kv_cache_manager_base.cc +++ b/tpu_sync/kv_cache/kv_cache_manager_base.cc @@ -2338,6 +2338,8 @@ absl::Status KVCacheManagerBase::RegisterActivePlan( } } absl::MutexLock l(plans_mu_); + // A new registration supersedes any earlier retirement of this uuid. + ReviveTransferUuidLocked(uuid); if (auto [it, inserted] = active_plans_.try_emplace( uuid, std::make_shared(RegisteredPlan{ request, is_sender, std::move(host_block_of), @@ -2352,6 +2354,34 @@ absl::Status KVCacheManagerBase::RegisterActivePlan( return absl::OkStatus(); } +void KVCacheManagerBase::ReviveTransferUuidLocked(uint64_t uuid) { + if (retired_transfer_uuids_.erase(uuid) == 0) { + return; + } + auto order_it = std::find(retired_transfer_uuid_order_.begin(), + retired_transfer_uuid_order_.end(), uuid); + if (order_it != retired_transfer_uuid_order_.end()) { + retired_transfer_uuid_order_.erase(order_it); + } +} + +void KVCacheManagerBase::ReviveTransferUuid(uint64_t uuid) { + absl::MutexLock l(plans_mu_); + ReviveTransferUuidLocked(uuid); +} + +void KVCacheManagerBase::RetireTransferUuid(uint64_t uuid) { + absl::MutexLock l(plans_mu_); + if (!retired_transfer_uuids_.insert(uuid).second) { + return; + } + retired_transfer_uuid_order_.push_back(uuid); + while (retired_transfer_uuid_order_.size() > kMaxRetiredTransferUuids) { + retired_transfer_uuids_.erase(retired_transfer_uuid_order_.front()); + retired_transfer_uuid_order_.pop_front(); + } +} + absl::Status KVCacheManagerBase::UnregisterActivePlan(uint64_t uuid) { { absl::MutexLock l(plans_mu_); @@ -2384,13 +2414,23 @@ KVCacheManagerBase::GetBlockChunks(size_t layer_idx, size_t shard_idx, absl::string_view peer, int64_t src_block_id, int64_t dst_block_id) { std::shared_ptr plan_snapshot; + bool retired = false; { absl::MutexLock l(plans_mu_); - auto it = active_plans_.find(uuid); - if (it != active_plans_.end()) { - plan_snapshot = it->second; + retired = retired_transfer_uuids_.contains(uuid); + if (!retired) { + auto it = active_plans_.find(uuid); + if (it != active_plans_.end()) { + plan_snapshot = it->second; + } } } + // A retired transfer's late payloads resolve nothing at all — neither + // through a still-registered plan nor through the planless identity + // fallback: the blocks they name may already belong to someone else. + if (retired) { + return {}; + } const bool has_plan = plan_snapshot != nullptr; // Resolve addressing geometry. With explicit pools the wire index is a pool diff --git a/tpu_sync/kv_cache/kv_cache_manager_base.h b/tpu_sync/kv_cache/kv_cache_manager_base.h index 9419495e..450b5025 100644 --- a/tpu_sync/kv_cache/kv_cache_manager_base.h +++ b/tpu_sync/kv_cache/kv_cache_manager_base.h @@ -20,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -387,6 +388,15 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase { virtual absl::Status UnregisterActivePlan(uint64_t uuid); + // Refuses planless payload resolution for a transfer that settled + // abnormally: a late push for its uuid must not land at identity-addressed + // blocks that may have new owners. Registering the uuid again lifts the + // refusal. The set is bounded; the oldest retirements fall off first. + void RetireTransferUuid(uint64_t uuid); + // Lifts an earlier retirement when a transfer legitimately reuses the + // uuid without registering a plan. + void ReviveTransferUuid(uint64_t uuid); + // Whether a transfer plan is currently registered under `uuid`. bool HasActivePlan(uint64_t uuid) const { absl::MutexLock l(plans_mu_); @@ -602,6 +612,13 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase { // previous copy semantics exactly. absl::flat_hash_map> active_plans_ ABSL_GUARDED_BY(plans_mu_); + void ReviveTransferUuidLocked(uint64_t uuid) + ABSL_EXCLUSIVE_LOCKS_REQUIRED(plans_mu_); + static constexpr size_t kMaxRetiredTransferUuids = 4096; + // Uuids RetireTransferUuid() has retired, and their retirement order. + absl::flat_hash_set retired_transfer_uuids_ + ABSL_GUARDED_BY(plans_mu_); + std::deque retired_transfer_uuid_order_ ABSL_GUARDED_BY(plans_mu_); // An asynchronous FFI task item representing a queued H2D or D2H copy // request. Bundles the work lambda with the XLA promise that signals Python diff --git a/tpu_sync/transport/block_transport.cc b/tpu_sync/transport/block_transport.cc index 0c6276e8..38d0af49 100644 --- a/tpu_sync/transport/block_transport.cc +++ b/tpu_sync/transport/block_transport.cc @@ -241,6 +241,23 @@ BlockTransport::~BlockTransport() { for (auto& t : socket_workers_) { if (t.joinable()) t.join(); } + // Workers are joined, so anything still queued will never run; cancelling + // each task fires its completion exactly once and no caller waits forever. + std::vector> orphaned_tasks; + { + absl::MutexLock lock(scheduler_mu_); + for (auto& queue_entry : peer_queues_) { + for (auto& task : queue_entry.second.tasks) { + orphaned_tasks.push_back(std::move(task)); + } + queue_entry.second.tasks.clear(); + } + } + for (auto& task : orphaned_tasks) { + if (task && task->cancel) { + task->cancel(absl::CancelledError("transport shutting down")); + } + } { absl::MutexLock lock(active_sends_mu_); for (const auto& [uuid, state] : active_sends_) { @@ -253,6 +270,15 @@ BlockTransport::~BlockTransport() { } } +void BlockTransport::AbortActiveSends() { + absl::MutexLock lock(active_sends_mu_); + for (const auto& [uuid, state] : active_sends_) { + if (state && state->client_fd >= 0) { + shutdown(state->client_fd, SHUT_RDWR); + } + } +} + void BlockTransport::SocketWorkerLoop() { while (!scheduler_stopping_) { std::unique_ptr task; @@ -328,6 +354,13 @@ absl::Status BlockTransport::HandleCustomRequest( absl::Status BlockTransport::HandleIncomingPush( int client_fd, const lib::ChunkHeader& header) { ASSIGN_OR_RETURN(MajorOrder major_order, ParseMajorOrder(header.flags)); + // One lease spans the whole stream from admission: a stream blocked in + // any read below cannot outlive its transfer's staging unnoticed. + const uint64_t stream_token = + block_delegate_->BeginPayloadResolution(header.uuid); + absl::Cleanup end_stream_lease = [&] { + block_delegate_->EndPayloadResolution(header.uuid, stream_token); + }; std::vector target_layers; if (header.local_id == 0xFFFFFFFF) { target_layers.resize(block_delegate_->num_block_arrays()); @@ -404,6 +437,14 @@ absl::Status BlockTransport::HandleIncomingPush( uint8_t size_buf[lib::kChunkSizeFieldSize]; RETURN_IF_ERROR(ReadExact(client_fd, size_buf, sizeof(size_buf))); const uint32_t sender_size = lib::DeserializeChunkSize(size_buf); + // The delegate may reclaim the memory behind resolved chunks when + // the transfer settles; the lease taken here keeps this payload's + // destination alive until its bytes have landed or the stream fails. + const uint64_t payload_token = + block_delegate_->BeginPayloadResolution(header.uuid); + absl::Cleanup end_payload_lease = [&] { + block_delegate_->EndPayloadResolution(header.uuid, payload_token); + }; const int64_t block_id_val = dst_id; int64_t src_bid = -1; @@ -849,6 +890,20 @@ void BlockTransport::AsyncPush( task->stream_idx = i; task->peer = remote_peer; task->run = std::move(task_run); + task->cancel = [i, statuses, remaining_workers, + on_complete](const absl::Status& status) { + (*statuses)[i] = status; + if (remaining_workers->fetch_sub(1) == 1) { + absl::Status final_status = absl::OkStatus(); + for (const auto& s : *statuses) { + if (!s.ok()) { + final_status = s; + break; + } + } + on_complete(final_status); + } + }; { absl::MutexLock lock(scheduler_mu_); diff --git a/tpu_sync/transport/block_transport.h b/tpu_sync/transport/block_transport.h index f2997163..962e36c2 100644 --- a/tpu_sync/transport/block_transport.h +++ b/tpu_sync/transport/block_transport.h @@ -64,6 +64,11 @@ class BlockTransport final { // Destructor closes all sockets and joins all threads. ~BlockTransport(); + // Fails every started push by shutting down its socket, without waiting. + // Each push reports through its completion callback as usual; a push still + // queued behind a worker runs later against its own socket. + void AbortActiveSends(); + // Return the TCP listening socket port. int local_port() const { return raw_transport_.local_port(); } @@ -144,6 +149,9 @@ class BlockTransport final { int stream_idx; std::string peer; std::function run; + // Fails this stream's share of the push without running it; fires the + // aggregated completion when it is the last share to finish. + std::function cancel; }; struct PeerQueue { diff --git a/tpu_sync/transport/block_transport_delegate.h b/tpu_sync/transport/block_transport_delegate.h index a142125a..3d985f97 100644 --- a/tpu_sync/transport/block_transport_delegate.h +++ b/tpu_sync/transport/block_transport_delegate.h @@ -115,6 +115,14 @@ class BlockTransportDelegate : public lib::RawBufferTransportDelegate { return result; } + // Brackets one incoming payload's use of the chunk pointers GetBlockChunks + // resolves for it: Begin runs before resolution, End after the payload's + // last byte has landed or its stream has failed. A delegate that recycles + // the memory behind its chunks keeps it alive between the two. End receives + // Begin's return value unchanged. + virtual uint64_t BeginPayloadResolution(uint64_t uuid) { return 0; } + virtual void EndPayloadResolution(uint64_t uuid, uint64_t token) {} + virtual absl::StatusOr> AllocateBlocks( size_t num_blocks, uint64_t uuid = 0) = 0; From 0b2e9d61fb9c3ab171cb8a5956b5e772bca4d66f Mon Sep 17 00:00:00 2001 From: Yin Lin Date: Mon, 24 Aug 2026 00:27:52 +0000 Subject: [PATCH 2/2] Register pool reshard plans in one generation-scoped step Pool reshard registration now arms the send or receive state first and publishes the plan last under the plan lifecycle lock, carrying a generation like block plans. Sender and receiver settlement, timeout cleanup, and each copy-to-push hand-off act only on their own registration, a failed pool receive drains like any other receive, and a failed pool send retires its uuid, so stale cleanup can no longer drop a reused plan or its transfer progress. --- .../core/kv_cache_manager_with_transfer.cc | 208 +++++++++++------- .../core/kv_cache_manager_with_transfer.h | 10 +- ...manager_with_transfer_pool_reshard_test.cc | 35 ++- 3 files changed, 172 insertions(+), 81 deletions(-) diff --git a/tpu_sync/core/kv_cache_manager_with_transfer.cc b/tpu_sync/core/kv_cache_manager_with_transfer.cc index 80ff16e6..b3392751 100644 --- a/tpu_sync/core/kv_cache_manager_with_transfer.cc +++ b/tpu_sync/core/kv_cache_manager_with_transfer.cc @@ -1346,9 +1346,6 @@ absl::Status KVCacheManagerWithTransfer::PoolReshardPush( } InitTransportServer(); - TF_RETURN_IF_ERROR(kv_cache::KVCacheManagerBase::RegisterActivePlan( - plan.uuid(), plan, /*is_sender=*/true)); - auto state = std::make_shared(); state->req_id = plan.req_id(); state->uuid = plan.uuid(); @@ -1358,13 +1355,30 @@ absl::Status KVCacheManagerWithTransfer::PoolReshardPush( state->plan = plan; state->deadline = DeadlineFromNow(); { - absl::MutexLock lock(mu_); - if (active_pool_reshard_sends_.contains(plan.uuid())) { - (void)kv_cache::KVCacheManagerBase::UnregisterActivePlan(plan.uuid()); - return absl::AlreadyExistsError( - absl::StrCat("pool reshard send UUID already active: ", plan.uuid())); + // One lifecycle step, like block plans: the send state is armed first + // and the plan published last under the lifecycle lock, with a + // generation that scopes every later cleanup to this registration. + absl::MutexLock lifecycle(plan_lifecycle_mu_); + if (kv_cache::KVCacheManagerBase::HasActivePlan(plan.uuid())) { + return absl::AlreadyExistsError(absl::StrCat( + "Plan with UUID ", plan.uuid(), " is already registered!")); + } + state->plan_generation = ++plan_generation_counter_; + { + absl::MutexLock lock(mu_); + if (active_pool_reshard_sends_.contains(plan.uuid())) { + return absl::AlreadyExistsError(absl::StrCat( + "pool reshard send UUID already active: ", plan.uuid())); + } + active_pool_reshard_sends_[plan.uuid()] = state; + } + absl::Status registered = kv_cache::KVCacheManagerBase::RegisterActivePlan( + plan.uuid(), plan, /*is_sender=*/true, {}, state->plan_generation); + if (!registered.ok()) { + absl::MutexLock lock(mu_); + active_pool_reshard_sends_.erase(plan.uuid()); + return registered; } - active_pool_reshard_sends_[plan.uuid()] = state; } // Multi-tag plans scope each pool's staging and pushes to its group's @@ -1405,14 +1419,16 @@ absl::Status KVCacheManagerWithTransfer::PoolReshardPush( // per-peer completion slots instead of failing the plan. The // receiver's expected pushes count only senders with scheduled pairs. for (size_t peer_idx = 0; peer_idx < peers.size(); ++peer_idx) { - FinishPoolReshardSend(plan.uuid(), absl::OkStatus()); + FinishPoolReshardSend(plan.uuid(), state->plan_generation, + absl::OkStatus()); } continue; } } auto future_or = D2hPoolBlocks(pool_idx, pool_src_block_ids); if (!future_or.ok()) { - FinishPoolReshardSend(plan.uuid(), future_or.status()); + FinishPoolReshardSend(plan.uuid(), state->plan_generation, + future_or.status()); return future_or.status(); } raiden::PjRtCopyFuture future = std::move(future_or).value(); @@ -1422,11 +1438,12 @@ absl::Status KVCacheManagerWithTransfer::PoolReshardPush( ++total_outstanding_ops_; } future.OnReady([this, uuid = static_cast(plan.uuid()), + generation = state->plan_generation, pool_idx](auto status_or) { if (!status_or.ok()) { - FinishPoolReshardSend(uuid, status_or.status()); + FinishPoolReshardSend(uuid, generation, status_or.status()); } else { - StartPoolReshardPush(uuid, pool_idx); + StartPoolReshardPush(uuid, pool_idx, generation); } // The callback reads manager state; destruction waits for it. absl::MutexLock lock(mu_); @@ -1437,7 +1454,8 @@ absl::Status KVCacheManagerWithTransfer::PoolReshardPush( } void KVCacheManagerWithTransfer::StartPoolReshardPush(uint64_t uuid, - size_t pool_idx) { + size_t pool_idx, + uint64_t generation) { std::shared_ptr state; { absl::MutexLock lock(mu_); @@ -1445,6 +1463,9 @@ void KVCacheManagerWithTransfer::StartPoolReshardPush(uint64_t uuid, if (it == active_pool_reshard_sends_.end()) return; state = it->second; } + // A copy completing for an earlier registration of this uuid must not + // launch pushes against the current one's plan. + if (state->plan_generation != generation) return; auto schedule_it = state->plan.shard_push_schedules().find(0); if (schedule_it == state->plan.shard_push_schedules().end()) { @@ -1474,44 +1495,50 @@ void KVCacheManagerWithTransfer::StartPoolReshardPush(uint64_t uuid, } } - transport::BlockTransport* transport_server = nullptr; { + // The pushes are queued while the transport pointer is held under its + // lock, so a concurrent transport stop cannot destroy it mid-queue. absl::MutexLock lock(server_init_mu_); - transport_server = server_.get(); - } - if (transport_server == nullptr) { - FinishPoolReshardSend( - uuid, absl::FailedPreconditionError("transport server is not running")); - return; - } - - for (const auto& [peer, transfers] : transfers_by_peer) { - std::vector src_ids; - std::vector dst_ids; - src_ids.reserve(transfers.size()); - dst_ids.reserve(transfers.size()); - for (const auto& [src_id, dst_id] : transfers) { - src_ids.push_back(src_id); - dst_ids.push_back(dst_id); + if (server_ != nullptr) { + for (const auto& [peer, transfers] : transfers_by_peer) { + std::vector src_ids; + std::vector dst_ids; + src_ids.reserve(transfers.size()); + dst_ids.reserve(transfers.size()); + for (const auto& [src_id, dst_id] : transfers) { + src_ids.push_back(src_id); + dst_ids.push_back(dst_id); + } + server_->AsyncPush( + {peer}, src_ids, dst_ids, state->parallelism, + transport::MajorOrder::kLayerMajor, uuid, + static_cast(pool_idx), + [this, uuid, generation = state->plan_generation]( + absl::StatusOr> result) { + FinishPoolReshardSend( + uuid, generation, + result.ok() ? absl::OkStatus() : result.status()); + }); + } + return; } - transport_server->AsyncPush( - {peer}, src_ids, dst_ids, state->parallelism, - transport::MajorOrder::kLayerMajor, uuid, static_cast(pool_idx), - [this, uuid](absl::StatusOr> result) { - FinishPoolReshardSend( - uuid, result.ok() ? absl::OkStatus() : result.status()); - }); } + FinishPoolReshardSend( + uuid, state->plan_generation, + absl::FailedPreconditionError("transport server is not running")); } void KVCacheManagerWithTransfer::FinishPoolReshardSend( - uint64_t uuid, const absl::Status& status) { + uint64_t uuid, uint64_t generation, const absl::Status& status) { bool finished = false; { absl::MutexLock lock(mu_); auto it = active_pool_reshard_sends_.find(uuid); if (it == active_pool_reshard_sends_.end()) return; auto& state = *it->second; + // Completion for an earlier registration of this uuid must not touch + // the current one's progress. + if (state.plan_generation != generation) return; if (state.finalizing) return; if (!status.ok()) { LOG(ERROR) << "Pool reshard send failed uuid=" << uuid @@ -1525,16 +1552,17 @@ void KVCacheManagerWithTransfer::FinishPoolReshardSend( } } if (finished) { - absl::Status unregister = UnregisterActivePlan(uuid); - if (!unregister.ok() && !absl::IsNotFound(unregister)) { - LOG(ERROR) << "Failed to unregister pool reshard sender plan " << uuid - << ": " << unregister; + UnregisterSettledPlan(uuid, generation); + if (!status.ok()) { + // A failed pool send retires its uuid so a late chunk lookup cannot + // resolve into the mirror once the plan is gone. + RetireTransferUuid(uuid); } absl::MutexLock lock(mu_); auto it = active_pool_reshard_sends_.find(uuid); if (it == active_pool_reshard_sends_.end()) return; - if (it->second->failed || - (!unregister.ok() && !absl::IsNotFound(unregister))) { + if (it->second->plan_generation != generation) return; + if (it->second->failed) { failed_recving_.insert(it->second->req_id); } else { done_sending_.insert(it->second->req_id); @@ -1559,16 +1587,6 @@ absl::Status KVCacheManagerWithTransfer::PoolReshardRegisterRecv( return absl::InvalidArgumentError( "pool reshard receiver requires dst_mem_type=HBM"); } - { - absl::MutexLock lock(mu_); - if (active_recv_entries_.contains(plan.uuid())) { - return absl::AlreadyExistsError( - absl::StrCat("pool reshard recv UUID already active: ", plan.uuid())); - } - } - - TF_RETURN_IF_ERROR(kv_cache::KVCacheManagerBase::RegisterActivePlan( - plan.uuid(), plan, /*is_sender=*/false)); RecvEntry recv_entry; recv_entry.req_id = plan.req_id(); recv_entry.is_pool_reshard = true; @@ -1593,8 +1611,32 @@ absl::Status KVCacheManagerWithTransfer::PoolReshardRegisterRecv( } } { - absl::MutexLock lock(mu_); - active_recv_entries_[plan.uuid()] = std::move(recv_entry); + // One lifecycle step, like block plans: the receive state is armed + // first and the plan published last under the lifecycle lock, so an + // early inbound push can never observe the plan without its receiver, + // and a generation scopes every later cleanup to this registration. + absl::MutexLock lifecycle(plan_lifecycle_mu_); + if (kv_cache::KVCacheManagerBase::HasActivePlan(plan.uuid())) { + return absl::AlreadyExistsError(absl::StrCat( + "Plan with UUID ", plan.uuid(), " is already registered!")); + } + const uint64_t generation = ++plan_generation_counter_; + recv_entry.plan_generation = generation; + { + absl::MutexLock lock(mu_); + if (active_recv_entries_.contains(plan.uuid())) { + return absl::AlreadyExistsError(absl::StrCat( + "pool reshard recv UUID already active: ", plan.uuid())); + } + active_recv_entries_[plan.uuid()] = std::move(recv_entry); + } + absl::Status registered = kv_cache::KVCacheManagerBase::RegisterActivePlan( + plan.uuid(), plan, /*is_sender=*/false, {}, generation); + if (!registered.ok()) { + absl::MutexLock lock(mu_); + active_recv_entries_.erase(plan.uuid()); + return registered; + } } return absl::OkStatus(); } @@ -1925,7 +1967,7 @@ KVCacheManagerWithTransfer::CompleteReadRaw() { const auto& entry = it->second; if (entry->deadline <= now) { failed_recving_.insert(entry->req_id); - settled_plans.emplace_back(it->first, 0); + settled_plans.emplace_back(it->first, entry->plan_generation); auto erase_it = it++; active_pool_reshard_sends_.erase(erase_it); } else { @@ -3416,11 +3458,13 @@ absl::Status KVCacheManagerWithTransfer::OnPoolReceived(size_t pool_idx, void KVCacheManagerWithTransfer::LaunchEligiblePoolH2ds(uint64_t uuid) { std::vector>> to_launch; + uint64_t generation = 0; { absl::MutexLock lock(mu_); auto it = active_recv_entries_.find(uuid); if (it == active_recv_entries_.end()) return; RecvEntry& entry = it->second; + generation = entry.plan_generation; if (entry.reshard_finalizing) return; for (size_t pool_idx : entry.started_pool_indices) { if (entry.h2d_launched_pools.count(pool_idx)) continue; @@ -3449,7 +3493,8 @@ void KVCacheManagerWithTransfer::LaunchEligiblePoolH2ds(uint64_t uuid) { for (auto& [pool_idx, chip_block_ids] : to_launch) { auto future_or = H2dPoolBlocks(pool_idx, chip_block_ids); if (!future_or.ok()) { - FinishPoolReshardRecvPool(uuid, pool_idx, future_or.status()); + FinishPoolReshardRecvPool(uuid, pool_idx, generation, + future_or.status()); continue; } raiden::PjRtCopyFuture future = std::move(future_or).value(); @@ -3457,9 +3502,10 @@ void KVCacheManagerWithTransfer::LaunchEligiblePoolH2ds(uint64_t uuid) { absl::MutexLock lock(mu_); ++total_outstanding_ops_; } - future.OnReady([this, uuid, pool_idx = pool_idx](auto status_or) { + future.OnReady([this, uuid, pool_idx = pool_idx, + generation](auto status_or) { FinishPoolReshardRecvPool( - uuid, pool_idx, + uuid, pool_idx, generation, status_or.ok() ? absl::OkStatus() : status_or.status()); // The callback reads manager state; destruction waits for it. absl::MutexLock lock(mu_); @@ -3476,13 +3522,17 @@ void KVCacheManagerWithTransfer::LaunchEligiblePoolH2ds(uint64_t uuid) { } void KVCacheManagerWithTransfer::FinishPoolReshardRecvPool( - uint64_t uuid, size_t pool_idx, const absl::Status& status) { + uint64_t uuid, size_t pool_idx, uint64_t generation, + const absl::Status& status) { bool finished = false; { absl::MutexLock lock(mu_); auto it = active_recv_entries_.find(uuid); if (it == active_recv_entries_.end()) return; RecvEntry& entry = it->second; + // Completion for an earlier registration of this uuid must not touch + // the current one's progress. + if (entry.plan_generation != generation) return; if (entry.reshard_finalizing) return; if (!status.ok()) { entry.reshard_finalizing = true; @@ -3502,24 +3552,26 @@ void KVCacheManagerWithTransfer::FinishPoolReshardRecvPool( std::chrono::steady_clock::time_point start_time; bool should_record_duration = false; if (finished) { - absl::Status unregister = UnregisterActivePlan(uuid); - if (!unregister.ok() && !absl::IsNotFound(unregister)) { - LOG(ERROR) << "Failed to unregister pool reshard receiver plan " << uuid - << ": " << unregister; - } - absl::MutexLock lock(mu_); - auto it = active_recv_entries_.find(uuid); - if (it == active_recv_entries_.end()) return; - if (!status.ok() || (!unregister.ok() && !absl::IsNotFound(unregister))) { - failed_recving_.insert(it->second.req_id); - active_recv_entries_.erase(it); - } else { - start_time = it->second.start_time; - should_record_duration = true; + UnregisterSettledPlan(uuid, generation); + SettleActions actions; + { + absl::MutexLock lock(mu_); + auto it = active_recv_entries_.find(uuid); + if (it == active_recv_entries_.end()) return; + if (it->second.plan_generation != generation) return; + if (!status.ok()) { + // A failed pool receive settles like any other receive: an open + // payload lease keeps the entry parked until its stream lets go. + SettleRecvLocked(uuid, /*failed=*/true, &actions); + } else { + start_time = it->second.start_time; + should_record_duration = true; - it->second.network_completed = true; - done_recving_.insert(it->second.req_id); + it->second.network_completed = true; + done_recving_.insert(it->second.req_id); + } } + ApplySettleActions(&actions); } if (should_record_duration) { RecordTransferDuration( diff --git a/tpu_sync/core/kv_cache_manager_with_transfer.h b/tpu_sync/core/kv_cache_manager_with_transfer.h index f6878ad2..34321ea6 100644 --- a/tpu_sync/core/kv_cache_manager_with_transfer.h +++ b/tpu_sync/core/kv_cache_manager_with_transfer.h @@ -504,6 +504,9 @@ class KVCacheManagerWithTransfer : public kv_cache::KVCacheManagerBase { struct PoolReshardSendEntry { std::string req_id; uint64_t uuid = 0; + // Generation of the plan this send belongs to; settlement cleanup only + // touches that registration. + uint64_t plan_generation = 0; int parallelism = 8; int remaining_pool_peer_pushes = 0; bool failed = false; @@ -540,9 +543,12 @@ class KVCacheManagerWithTransfer : public kv_cache::KVCacheManagerBase { // corruption; the arming worker validates for itself. absl::Status ValidatePoolReshardReceiverCoverage( const ::tpu_sync::rpc::StartTransferRequest& plan); - void StartPoolReshardPush(uint64_t uuid, size_t pool_idx); - void FinishPoolReshardSend(uint64_t uuid, const absl::Status& status); + void StartPoolReshardPush(uint64_t uuid, size_t pool_idx, + uint64_t generation); + void FinishPoolReshardSend(uint64_t uuid, uint64_t generation, + const absl::Status& status); void FinishPoolReshardRecvPool(uint64_t uuid, size_t pool_idx, + uint64_t generation, const absl::Status& status); // Launches H2D uploads for every wire-complete pool whose order-rank // prerequisites (all lower-rank pools uploaded) are satisfied. diff --git a/tpu_sync/core/kv_cache_manager_with_transfer_pool_reshard_test.cc b/tpu_sync/core/kv_cache_manager_with_transfer_pool_reshard_test.cc index 70720086..69235b9d 100644 --- a/tpu_sync/core/kv_cache_manager_with_transfer_pool_reshard_test.cc +++ b/tpu_sync/core/kv_cache_manager_with_transfer_pool_reshard_test.cc @@ -699,7 +699,9 @@ TEST(PoolReshardRecvTest, FinishPoolReshardRecvRecordsDurationMetric) { manager.PoolReshardRegisterRecv(plan, std::vector{0}).ok()); // Simulate pool completion - manager.FinishPoolReshardRecvPool(3001, /*pool_idx=*/0, absl::OkStatus()); + manager.FinishPoolReshardRecvPool(3001, /*pool_idx=*/0, + manager.ActivePlanGeneration(3001).value(), + absl::OkStatus()); } TEST(PoolReshardRecvTest, FinishPoolReshardRecvDoesNotRecordMetricOnFailure) { @@ -719,10 +721,41 @@ TEST(PoolReshardRecvTest, FinishPoolReshardRecvDoesNotRecordMetricOnFailure) { // Simulate pool failure manager.FinishPoolReshardRecvPool(3002, /*pool_idx=*/0, + manager.ActivePlanGeneration(3002).value(), absl::InternalError("simulated failure")); } +TEST(PoolReshardRecvTest, CleanupIsScopedToItsOwnRegistration) { + TestManager manager; + ASSERT_TRUE(manager.RegisterPools({DensePool("fa")}).ok()); + manager.AttachPlaceholderDeviceHold(); + + StartTransferRequest plan = ValidPlan(/*uuid=*/3003); + ASSERT_TRUE( + manager.PoolReshardRegisterRecv(plan, std::vector{0}).ok()); + const uint64_t generation = manager.ActivePlanGeneration(3003).value(); + ASSERT_GT(generation, 0u); + + // Cleanup carrying another registration's generation must leave this + // registration and its receive state untouched. + manager.FinishPoolReshardRecvPool(3003, /*pool_idx=*/0, generation + 1, + absl::InternalError("stale cleanup")); + EXPECT_TRUE(manager.HasActivePlan(3003)); + EXPECT_EQ( + manager.PoolReshardRegisterRecv(plan, std::vector{0}).code(), + absl::StatusCode::kAlreadyExists); + + // Cleanup for this registration settles it; the uuid is registrable + // again and the new registration carries a fresh generation. + manager.FinishPoolReshardRecvPool(3003, /*pool_idx=*/0, generation, + absl::InternalError("real failure")); + EXPECT_FALSE(manager.HasActivePlan(3003)); + ASSERT_TRUE( + manager.PoolReshardRegisterRecv(plan, std::vector{0}).ok()); + EXPECT_NE(manager.ActivePlanGeneration(3003).value(), generation); +} + TEST(SendDeadlineTest, ExpiredSendEntryFailsInsteadOfReportingDone) { TestManager manager(/*timeout_s=*/0.05); ASSERT_GT(manager.NotifyForRead("expired_send_req", 31, {0, 1}), 0);