Transfer staging leases and generation-scoped pool plan cleanup

tpu_sync/core/kv_cache_manager_with_transfer.h

@@ -36,6 +36,7 @@
 #include <vector>
 
 #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<int> src_ints;
     std::vector<int> dst_ints;
     std::atomic<size_t> 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<int>* staged_host_blocks);
   void ReleaseEntrySlotLocked(const std::shared_ptr<SendEntry>& 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<raiden::PjRtCopyFuture> drain_futures;
+    std::optional<std::pair<uint64_t, uint64_t>> unregister_plan;
+  };
+  uint64_t EnsureOpTokenLocked(uint64_t* op_token);
+  void AddSendOpLocked(const std::shared_ptr<SendEntry>& entry);
+  void AddRecvOpLocked(RecvEntry* entry);
+  void ReapSendLocked(const std::shared_ptr<SendEntry>& 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<void()> task) override;
+  uint64_t BeginPayloadResolution(uint64_t uuid) override;
+  void EndPayloadResolution(uint64_t uuid, uint64_t token) override;
   std::shared_ptr<StagingReadinessState> CreateStagingReadiness(
       int64_t slot_idx, int64_t num_blocks);
   void MarkStagingLayerReady(
@@ -438,12 +478,35 @@ 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<size_t, std::vector<int64_t>> 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<uint64_t, RecvEntry> 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<uint64_t, std::shared_ptr<SendEntry>> draining_sends_
+      ABSL_GUARDED_BY(mu_);
+  absl::flat_hash_map<uint64_t, RecvEntry> draining_recvs_ ABSL_GUARDED_BY(mu_);
 
   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;
@@ -480,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.
@@ -509,6 +575,10 @@ class KVCacheManagerWithTransfer : public kv_cache::KVCacheManagerBase {
   std::atomic<bool> 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<int> accepted_control_fds_
+      ABSL_GUARDED_BY(pull_workers_mu_);
   double timeout_s_ = 120.0;
   bool unsafe_skip_buffer_lock_ = true;
 
@@ -532,6 +602,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<bool> stopping_{false};
   std::thread control_thread_;

tpu_sync/core/kv_cache_manager_with_transfer_pool_reshard_test.cc

@@ -699,7 +699,9 @@ TEST(PoolReshardRecvTest, FinishPoolReshardRecvRecordsDurationMetric) {
       manager.PoolReshardRegisterRecv(plan, std::vector<int64_t>{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<int64_t>{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<int64_t>{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<int64_t>{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);
@@ -734,6 +767,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();

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<tpu_raiden::transport::BlockTransport> 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<HostNicAddress> RaidenManagerBase::GetHostNics() const {
   return GetLocalHostNicAddresses();
 }

tpu_sync/core/raiden_manager_base.h

@@ -146,6 +146,12 @@ class RaidenManagerBase : public tpu_raiden::transport::BlockTransportDelegate {
   std::vector<std::string> 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<HostNicAddress> GetHostNics() const;
 
   void DetectAndAssignNumaNode(

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<const RegisteredPlan>(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<const RegisteredPlan> 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

tpu_sync/kv_cache/kv_cache_manager_base.h

@@ -20,6 +20,7 @@
 #include <cstdint>
 #include <functional>
 #include <memory>
+#include <deque>
 #include <optional>
 #include <queue>
 #include <string>
@@ -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<uint64_t, std::shared_ptr<const RegisteredPlan>>
       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<uint64_t> retired_transfer_uuids_
+      ABSL_GUARDED_BY(plans_mu_);
+  std::deque<uint64_t> 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