Fail loudly when transfer staging is exhausted

A producer that could not seat a pull-serve marked the send done without
transferring, so the consumer learned of the failure only by waiting out
its transfer timeout; a consumer that could not seat a read failed it
silently.

The producer now waits for staging to free, bounded by the transfer
deadline, and reports a transfer failure when a request still cannot be
seated. A request larger than the staging can ever seat -- a fixed slot,
or the whole per-transfer pool -- fails at once. The misses while waiting
are not logged; the one failure message on each side carries the request,
the pages asked, and the pages free, so an exhausted pool is attributed
where it happens.

Measured on the setup of the previous commit: at 4 fixed slots the load
exhausted staging 1112 times across 139 requests, each visible only as a
consumer-side failure. With staging on demand none occurred.

Memory and performance results for on-demand staging itself are in the
commit that introduces it, "Stage disaggregation transfers on demand
instead of in fixed slots" (github.com//pull/722).

During shutdown a producer waiting for staging stops instead of sleeping
out its deadline, and the manager's destructor waits for every pull-serve
worker before tearing down the state they read.

Author: yinlin-google

tpu_sync/core/kv_cache_manager_with_transfer.cc

@@ -620,6 +620,14 @@ KVCacheManagerWithTransfer::KVCacheManagerWithTransfer(
 
 KVCacheManagerWithTransfer::~KVCacheManagerWithTransfer() {
   StopControlServer();
+  // Pull-serve workers read this object's state; nothing may be torn down
+  // while one is still running.
+  shutting_down_.store(true, std::memory_order_relaxed);
+  {
+    absl::MutexLock lock(pull_workers_mu_);
+    pull_workers_mu_.Await(absl::Condition(
+        +[](int* active) { return *active == 0; }, &active_pull_workers_));
+  }
   push_pool_.reset();
   pull_pool_.reset();
   if (host_block_manager_ && !all_slots_.empty()) {
@@ -1556,6 +1564,13 @@ void KVCacheManagerWithTransfer::StartRead(
     if (!staged.has_value()) {
       // Request larger than the staging pool can seat: surface as a recv
       // failure (the connector can recompute) rather than throwing.
+      LOG(ERROR) << "StartRead: cannot stage " << unique_local_bids.size()
+                 << " blocks for req_id=" << req_id
+                 << " (dynamic=" << dynamic_host_staging_
+                 << ", free_host_blocks="
+                 << host_block_manager_->num_free_blocks()
+                 << ", free_slots=" << free_slots_.size()
+                 << ", max_blocks=" << max_blocks_ << ")";
       failed_recving_.insert(req_id);
       return;
     }
@@ -1923,9 +1938,13 @@ KVCacheManagerWithTransfer::AcquireRecvStagingLocked(int64_t num_blocks,
     return blocks;
   }
   // Lock the pages so the pool's LRU cannot evict staging that is mid-flight.
+  // A miss is reported to the caller only: the producer retries it until its
+  // deadline and the consumer fails it at once, and each logs its own verdict.
   auto allocated = host_block_manager_->Allocate(static_cast<int>(num_blocks),
                                                  /*lock=*/true);
-  if (!allocated.ok()) return std::nullopt;
+  if (!allocated.ok()) {
+    return std::nullopt;
+  }
   entry->staged_host_blocks = *allocated;
   std::vector<int64_t> blocks;
   blocks.reserve(allocated->size());
@@ -2319,10 +2338,16 @@ void KVCacheManagerWithTransfer::ProcessPullStream(
     }
   }
 
+  {
+    absl::MutexLock lock(pull_workers_mu_);
+    ++active_pull_workers_;
+  }
   std::thread([this, uuid = req.uuid, remote_data_endpoints, src_block_ids,
                dst_block_ids]() {
     StartPushInternal(uuid, remote_data_endpoints, src_block_ids,
                       dst_block_ids);
+    absl::MutexLock lock(pull_workers_mu_);
+    --active_pull_workers_;
   }).detach();
 }
 
@@ -2335,24 +2360,66 @@ void KVCacheManagerWithTransfer::StartPushInternal(
   // pool. Writing D2H straight to host[src_block_id] overflows the host buffer
   // once a device block id exceeds num_host_blocks.
   std::vector<int64_t> host_block_ids;
-  {
-    absl::MutexLock lock(mu_);
-    auto it = send_entries_.find(uuid);
-    if (it == send_entries_.end()) {
-      return;
+  const auto stage_deadline = DeadlineFromNow();
+  while (true) {
+    if (shutting_down_.load(std::memory_order_relaxed)) {
+      return;  // the manager is being destroyed; its state goes with it
     }
-    RecvEntry staging;
-    auto staged = AcquireRecvStagingLocked(
-        static_cast<int64_t>(src_block_ids.size()), &staging);
-    if (!staged.has_value()) {
-      done_sending_.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()) {
+        return;  // request cancelled while waiting for staging
+      }
+      // Staging that can never seat this request fails it now rather than
+      // after the deadline: a fixed slot holds max_blocks_ pages, the
+      // per-transfer pool holds total_blocks() pages.
+      const int64_t capacity = dynamic_host_staging_
+                                   ? host_block_manager_->total_blocks()
+                                   : max_blocks_;
+      if (static_cast<int64_t>(src_block_ids.size()) > capacity) {
+        LOG(ERROR) << "StartPushInternal: request " << it->second->req_id
+                   << " needs " << src_block_ids.size() << " blocks but "
+                   << (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;
+      }
+      RecvEntry staging;
+      auto staged = AcquireRecvStagingLocked(
+          static_cast<int64_t>(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);
+        break;
+      }
+    }
+    // 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);
+      }
       return;
     }
-    it->second->slot_idx = staging.slot_idx;
-    it->second->staged_host_blocks = std::move(staging.staged_host_blocks);
-    host_block_ids = std::move(*staged);
+    std::this_thread::sleep_for(std::chrono::milliseconds(1));
   }
 
   // Coalesce contiguous (device,host) block runs into a few large copies. With

tpu_sync/core/kv_cache_manager_with_transfer.h

@@ -470,6 +470,12 @@ class KVCacheManagerWithTransfer : public kv_cache::KVCacheManagerBase {
   // fixed slots. A short request then costs its own pages rather than a
   // whole slot, so the same pool seats more transfers at once.
   bool dynamic_host_staging_ = false;
+
+  // Pull-serve workers launched by ProcessPullStream. The destructor waits
+  // for them, and shutting_down_ ends a worker's staging wait early.
+  std::atomic<bool> shutting_down_{false};
+  absl::Mutex pull_workers_mu_;
+  int active_pull_workers_ ABSL_GUARDED_BY(pull_workers_mu_) = 0;
   double timeout_s_ = 120.0;
   bool unsafe_skip_buffer_lock_ = true;