diff --git a/tpu_sync/frameworks/jax/kv_cache_manager.cc b/tpu_sync/frameworks/jax/kv_cache_manager.cc index 85a7b4644..ba42c292a 100644 --- a/tpu_sync/frameworks/jax/kv_cache_manager.cc +++ b/tpu_sync/frameworks/jax/kv_cache_manager.cc @@ -641,36 +641,58 @@ NumaAwareKVCacheManager::CompleteReadRaw() { int num_subs = static_cast(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), diff --git a/tpu_sync/frameworks/jax/kv_cache_manager.h b/tpu_sync/frameworks/jax/kv_cache_manager.h index 242df2a30..174f409b0 100644 --- a/tpu_sync/frameworks/jax/kv_cache_manager.h +++ b/tpu_sync/frameworks/jax/kv_cache_manager.h @@ -286,8 +286,8 @@ class NumaAwareKVCacheManager { std::map done_sending_counts_; std::map done_recving_counts_; + std::map failed_recving_counts_; std::map req_expected_counts_; - std::set failed_recving_set_; std::shared_ptr metrics_collector_; }; diff --git a/tpu_sync/frameworks/jax/kv_cache_manager_wrapper_test.cc b/tpu_sync/frameworks/jax/kv_cache_manager_wrapper_test.cc index dac3f05c9..6a08d4197 100644 --- a/tpu_sync/frameworks/jax/kv_cache_manager_wrapper_test.cc +++ b/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(); auto sub1 = std::make_unique(); MockSubManager* ptr0 = sub0.get(); @@ -209,18 +209,152 @@ TEST(KVCacheManagerWrapperTest, TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr 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{"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{"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{"req_fail"}); + EXPECT_TRUE(f3.empty()); +} + +TEST(KVCacheManagerWrapperTest, + FailureWhenAllSubManagersFailEmitsFailedRecvingOnce) { + auto sub0 = std::make_unique(); + auto sub1 = std::make_unique(); + MockSubManager* ptr0 = sub0.get(); + MockSubManager* ptr1 = sub1.get(); + + std::vector> subs; + subs.push_back(std::move(sub0)); + subs.push_back(std::move(sub1)); + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr 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{"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(); + auto sub1 = std::make_unique(); + MockSubManager* ptr0 = sub0.get(); + MockSubManager* ptr1 = sub1.get(); + + std::vector> subs; + subs.push_back(std::move(sub0)); + subs.push_back(std::move(sub1)); + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr 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{"req_fail"}); +} + +TEST(KVCacheManagerWrapperTest, + SendFailureWaitsForSiblingSubManagersBeforeEmittingFailedRecving) { + auto sub0 = std::make_unique(); + auto sub1 = std::make_unique(); + MockSubManager* ptr0 = sub0.get(); + MockSubManager* ptr1 = sub1.get(); + + std::vector> subs; + subs.push_back(std::move(sub0)); + subs.push_back(std::move(sub1)); + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr 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{"req_send_fail"}); +} + +TEST(KVCacheManagerWrapperTest, + SingleSubManagerFailureEmitsImmediatelyAndCleansUp) { + auto sub0 = std::make_unique(); + MockSubManager* ptr0 = sub0.get(); + + std::vector> subs; + subs.push_back(std::move(sub0)); + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr 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{"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{"req_fail"}); EXPECT_TRUE(f2.empty()); }