Barrier-wait failed transfers and clean up completion state in NumaAwareKVCacheManager.

tpu_sync/frameworks/jax/kv_cache_manager.cc

@@ -641,36 +641,58 @@ NumaAwareKVCacheManager::CompleteReadRaw() {
   int num_subs = static_cast<int>(sub_managers_.size());
   if (num_subs == 0) return {};
 
+  auto process_completion = [&](const std::string& req) {
+    int expected = num_subs;
+    if (auto exp_it = req_expected_counts_.find(req);
+        exp_it != req_expected_counts_.end()) {
+      expected = exp_it->second;
+    }
+    int s_count = 0;
+    if (auto it = done_sending_counts_.find(req);
+        it != done_sending_counts_.end()) {
+      s_count = it->second;
+    }
+    int r_count = 0;
+    if (auto it = done_recving_counts_.find(req);
+        it != done_recving_counts_.end()) {
+      r_count = it->second;
+    }
+    int f_count = 0;
+    if (auto it = failed_recving_counts_.find(req);
+        it != failed_recving_counts_.end()) {
+      f_count = it->second;
+    }
+
+    if (s_count + r_count + f_count >= expected) {
+      if (f_count > 0) {
+        failed_recving.push_back(req);
+      } else if (r_count >= expected) {
+        done_recving.push_back(req);
+      } else if (s_count >= expected) {
+        done_sending.push_back(req);
+      } else {
+        failed_recving.push_back(req);
+      }
+      done_sending_counts_.erase(req);
+      done_recving_counts_.erase(req);
+      failed_recving_counts_.erase(req);
+      req_expected_counts_.erase(req);
+    }
+  };
+
   for (auto& sub : sub_managers_) {
     auto [s, r, f] = sub->CompleteReadRaw();
     for (const auto& req : s) {
       done_sending_counts_[req]++;
-      if (done_sending_counts_[req] == num_subs) {
-        done_sending.push_back(req);
-        done_sending_counts_.erase(req);
-      }
+      process_completion(req);
     }
     for (const auto& req : r) {
-      auto [cnt_it, inserted] = done_recving_counts_.emplace(req, 0);
-      cnt_it->second++;
-      int expected = num_subs;
-      if (auto exp_it = req_expected_counts_.find(req);
-          exp_it != req_expected_counts_.end()) {
-        expected = exp_it->second;
-      }
-      if (cnt_it->second >= expected) {
-        done_recving_counts_.erase(cnt_it);
-        req_expected_counts_.erase(req);
-        if (failed_recving_set_.find(req) == failed_recving_set_.end()) {
-          done_recving.push_back(req);
-        }
-        failed_recving_set_.erase(req);
-      }
+      done_recving_counts_[req]++;
+      process_completion(req);
     }
     for (const auto& req : f) {
-      if (failed_recving_set_.insert(req).second) {
-        failed_recving.push_back(req);
-      }
+      failed_recving_counts_[req]++;
+      process_completion(req);
     }
   }
   return {std::move(done_sending), std::move(done_recving),

tpu_sync/frameworks/jax/kv_cache_manager.h

@@ -286,8 +286,8 @@ class NumaAwareKVCacheManager {
 
   std::map<std::string, int> done_sending_counts_;
   std::map<std::string, int> done_recving_counts_;
+  std::map<std::string, int> failed_recving_counts_;
   std::map<std::string, int> req_expected_counts_;
-  std::set<std::string> failed_recving_set_;
   std::shared_ptr<MetricsCollector> metrics_collector_;
 };
 

tpu_sync/frameworks/jax/kv_cache_manager_wrapper_test.cc

@@ -196,7 +196,7 @@ TEST(KVCacheManagerWrapperTest,
 }
 
 TEST(KVCacheManagerWrapperTest,
-     AnyFailureImmediatelyTriggersFailedRecvingAndSuppressesFutureAcks) {
+     FailureWaitsForSiblingSubManagersBeforeEmittingFailedRecving) {
   auto sub0 = std::make_unique<MockSubManager>();
   auto sub1 = std::make_unique<MockSubManager>();
   MockSubManager* ptr0 = sub0.get();
@@ -209,18 +209,152 @@ TEST(KVCacheManagerWrapperTest,
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<KVCacheManager> mgr,
                           KVCacheManager::Create(std::move(subs)));
 
-  // 1. sub0 reports failure for req_fail
+  // 1. sub0 reports failure for req_fail, but sibling sub1 still has DMAs
+  // in flight. CompleteReadRaw must barrier-wait and not emit prematurely.
   ptr0->failed.push_back("req_fail");
   auto [s1, r1, f1] = mgr->CompleteReadRaw();
   EXPECT_TRUE(s1.empty());
   EXPECT_TRUE(r1.empty());
-  EXPECT_EQ(f1, std::vector<std::string>{"req_fail"});
+  EXPECT_TRUE(f1.empty());
+
+  // 2. Later sub1 finishes recving for req_fail. Now that all sub-managers
+  // have reached terminal status, failed_recving is emitted and done_recving
+  // is suppressed.
+  ptr1->recving.push_back("req_fail");
+  auto [s2, r2, f2] = mgr->CompleteReadRaw();
+  EXPECT_TRUE(s2.empty());
+  EXPECT_TRUE(r2.empty());
+  EXPECT_EQ(f2, std::vector<std::string>{"req_fail"});
 
-  // 2. Later sub1 finishes recving for req_fail (should be suppressed)
+  // 3. Verifies state cleanup & req_id reuse: a subsequent transfer reusing
+  // "req_fail" succeeds cleanly without being suppressed by the previous
+  // failure.
+  ptr0->recving.push_back("req_fail");
   ptr1->recving.push_back("req_fail");
+  auto [s3, r3, f3] = mgr->CompleteReadRaw();
+  EXPECT_TRUE(s3.empty());
+  EXPECT_EQ(r3, std::vector<std::string>{"req_fail"});
+  EXPECT_TRUE(f3.empty());
+}
+
+TEST(KVCacheManagerWrapperTest,
+     FailureWhenAllSubManagersFailEmitsFailedRecvingOnce) {
+  auto sub0 = std::make_unique<MockSubManager>();
+  auto sub1 = std::make_unique<MockSubManager>();
+  MockSubManager* ptr0 = sub0.get();
+  MockSubManager* ptr1 = sub1.get();
+
+  std::vector<std::unique_ptr<KVCacheManagerWithTransfer>> subs;
+  subs.push_back(std::move(sub0));
+  subs.push_back(std::move(sub1));
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<KVCacheManager> mgr,
+                          KVCacheManager::Create(std::move(subs)));
+
+  // 1. sub0 reports failure for req_fail.
+  ptr0->failed.push_back("req_fail");
+  auto [s1, r1, f1] = mgr->CompleteReadRaw();
+  EXPECT_TRUE(s1.empty());
+  EXPECT_TRUE(r1.empty());
+  EXPECT_TRUE(f1.empty());
+
+  // 2. sub1 also reports failure for req_fail.
+  ptr1->failed.push_back("req_fail");
   auto [s2, r2, f2] = mgr->CompleteReadRaw();
   EXPECT_TRUE(s2.empty());
   EXPECT_TRUE(r2.empty());
+  EXPECT_EQ(f2, std::vector<std::string>{"req_fail"});
+
+  // 3. Subsequent poll returns nothing.
+  auto [s3, r3, f3] = mgr->CompleteReadRaw();
+  EXPECT_TRUE(s3.empty());
+  EXPECT_TRUE(r3.empty());
+  EXPECT_TRUE(f3.empty());
+}
+
+TEST(KVCacheManagerWrapperTest,
+     RecvFirstThenFailureWaitsForSubManagersAndEmitsFailure) {
+  auto sub0 = std::make_unique<MockSubManager>();
+  auto sub1 = std::make_unique<MockSubManager>();
+  MockSubManager* ptr0 = sub0.get();
+  MockSubManager* ptr1 = sub1.get();
+
+  std::vector<std::unique_ptr<KVCacheManagerWithTransfer>> subs;
+  subs.push_back(std::move(sub0));
+  subs.push_back(std::move(sub1));
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<KVCacheManager> mgr,
+                          KVCacheManager::Create(std::move(subs)));
+
+  // 1. sub0 succeeds recving first.
+  ptr0->recving.push_back("req_fail");
+  auto [s1, r1, f1] = mgr->CompleteReadRaw();
+  EXPECT_TRUE(s1.empty());
+  EXPECT_TRUE(r1.empty());
+  EXPECT_TRUE(f1.empty());
+
+  // 2. sub1 fails. Overall transfer must emit failure, not success.
+  ptr1->failed.push_back("req_fail");
+  auto [s2, r2, f2] = mgr->CompleteReadRaw();
+  EXPECT_TRUE(s2.empty());
+  EXPECT_TRUE(r2.empty());
+  EXPECT_EQ(f2, std::vector<std::string>{"req_fail"});
+}
+
+TEST(KVCacheManagerWrapperTest,
+     SendFailureWaitsForSiblingSubManagersBeforeEmittingFailedRecving) {
+  auto sub0 = std::make_unique<MockSubManager>();
+  auto sub1 = std::make_unique<MockSubManager>();
+  MockSubManager* ptr0 = sub0.get();
+  MockSubManager* ptr1 = sub1.get();
+
+  std::vector<std::unique_ptr<KVCacheManagerWithTransfer>> subs;
+  subs.push_back(std::move(sub0));
+  subs.push_back(std::move(sub1));
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<KVCacheManager> mgr,
+                          KVCacheManager::Create(std::move(subs)));
+
+  // 1. sub0 succeeds sending first.
+  ptr0->sending.push_back("req_send_fail");
+  auto [s1, r1, f1] = mgr->CompleteReadRaw();
+  EXPECT_TRUE(s1.empty());
+  EXPECT_TRUE(r1.empty());
+  EXPECT_TRUE(f1.empty());
+
+  // 2. sub1 fails sending. Overall transfer must emit failure, suppressing
+  // done_sending.
+  ptr1->failed.push_back("req_send_fail");
+  auto [s2, r2, f2] = mgr->CompleteReadRaw();
+  EXPECT_TRUE(s2.empty());
+  EXPECT_TRUE(r2.empty());
+  EXPECT_EQ(f2, std::vector<std::string>{"req_send_fail"});
+}
+
+TEST(KVCacheManagerWrapperTest,
+     SingleSubManagerFailureEmitsImmediatelyAndCleansUp) {
+  auto sub0 = std::make_unique<MockSubManager>();
+  MockSubManager* ptr0 = sub0.get();
+
+  std::vector<std::unique_ptr<KVCacheManagerWithTransfer>> subs;
+  subs.push_back(std::move(sub0));
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<KVCacheManager> mgr,
+                          KVCacheManager::Create(std::move(subs)));
+
+  // 1. sub0 reports failure for req_fail. With 1 sub-manager, emitted
+  // immediately.
+  ptr0->failed.push_back("req_fail");
+  auto [s1, r1, f1] = mgr->CompleteReadRaw();
+  EXPECT_TRUE(s1.empty());
+  EXPECT_TRUE(r1.empty());
+  EXPECT_EQ(f1, std::vector<std::string>{"req_fail"});
+
+  // 2. Reuse req_fail for healthy transfer; state was erased so it succeeds.
+  ptr0->recving.push_back("req_fail");
+  auto [s2, r2, f2] = mgr->CompleteReadRaw();
+  EXPECT_TRUE(s2.empty());
+  EXPECT_EQ(r2, std::vector<std::string>{"req_fail"});
   EXPECT_TRUE(f2.empty());
 }