diff --git a/tpu_sync/kv_cache/reshard_send_session.cc b/tpu_sync/kv_cache/reshard_send_session.cc index f77052a91..602843986 100644 --- a/tpu_sync/kv_cache/reshard_send_session.cc +++ b/tpu_sync/kv_cache/reshard_send_session.cc @@ -204,6 +204,13 @@ absl::StatusOr ReshardSendSession::AcquireStagingSlot(size_t pool_idx) { absl::StatusOr ReshardSendSession::LookupStagingSlot( size_t pool_idx) const { + // A lease outlives the session's finalizing_ flip: a push whose completion + // started finalizing may still be addressing another pool's host bytes. + // SettleLocked releases every lease, so a done session never resolves one. + if (std::optional slot = staging_layer_manager_.Lookup(pool_idx); + slot.has_value()) { + return *slot; + } { absl::MutexLock lock(mu_); if (done_ || finalizing_) { @@ -211,12 +218,8 @@ absl::StatusOr ReshardSendSession::LookupStagingSlot( : absl::CancelledError("Send session is draining"); } } - std::optional slot = staging_layer_manager_.Lookup(pool_idx); - if (!slot.has_value()) { - return absl::FailedPreconditionError( - absl::StrCat("Pool ", pool_idx, " has no leased staging slot")); - } - return *slot; + return absl::FailedPreconditionError( + absl::StrCat("Pool ", pool_idx, " has no leased staging slot")); } absl::Status ReshardSendSession::ValidatePushPreconditions() const { diff --git a/tpu_sync/kv_cache/reshard_send_session.h b/tpu_sync/kv_cache/reshard_send_session.h index 8fc1e1d2b..ae8d54976 100644 --- a/tpu_sync/kv_cache/reshard_send_session.h +++ b/tpu_sync/kv_cache/reshard_send_session.h @@ -103,8 +103,10 @@ class ReshardSendSession // has none and waiting until deadline when every slot is leased. absl::StatusOr AcquireStagingSlot(size_t pool_idx); - // Returns the host staging buffer currently leased to |pool_idx|, or - // kFailedPrecondition if |pool_idx| holds no lease. + // Returns the host staging buffer currently leased to |pool_idx|. Fails + // with the session's status (or kCancelled) once the session is draining + // and holds no lease for |pool_idx|, and kFailedPrecondition when a live + // session holds none. absl::StatusOr LookupStagingSlot(size_t pool_idx) const; private: diff --git a/tpu_sync/kv_cache/reshard_send_session_test.cc b/tpu_sync/kv_cache/reshard_send_session_test.cc index 4d46b4273..71399505a 100644 --- a/tpu_sync/kv_cache/reshard_send_session_test.cc +++ b/tpu_sync/kv_cache/reshard_send_session_test.cc @@ -56,6 +56,13 @@ struct ReshardSendSessionTestPeer { return session.staging_layer_manager_.Release(pool_idx); } + // Issues `pool_idx`'s peer pushes as the D2H completion callback does. + static void StartPoolPush(ReshardSendSession& session, + KVCacheManagerWithTransfer& manager, + size_t pool_idx) { + session.StartPoolPush(manager, pool_idx); + } + // Counts one (pool, peer) push of `session` in flight on `manager`, as // ExecutePush and StartPoolPush do before issuing it. The executor's device // copies are device-only, so deviceless tests drive the op accounting here. @@ -531,5 +538,124 @@ TEST(ReshardSendSessionTest, BoundedStagingSlotAcquisitionAndRelease) { StatusIs(absl::StatusCode::kDeadlineExceeded)); } +// A precondition that needs no waiting fails ExecutePush synchronously, so +// the START_TRANSFER reply carries it, and the send settles with its plan +// dropped rather than waiting for the deadline sweep. +TEST(ReshardSendSessionTest, ExecutePushFailsFastWithoutTransportServer) { + KVCacheManagerWithTransfer manager = MakeTestManager(); + manager.base()->AttachPlaceholderDeviceHoldForTest(); + std::unique_ptr allocator = + CreateStagingBlockAllocator(manager.base(), /*num_slots=*/0, + /*max_blocks=*/0); + StartTransferRequest plan = ValidPlan(/*uuid=*/1401, /*dtype_tags=*/{""}); + TF_ASSERT_OK_AND_ASSIGN( + std::shared_ptr session, + ReshardSendSession::Create( + manager.base(), allocator.get(), std::vector{0}, + /*parallelism=*/1, + std::chrono::steady_clock::now() + std::chrono::seconds(10), plan)); + ASSERT_TRUE(manager.base()->HasActivePlan(plan.uuid())); + + manager.base()->StopTransportServer(); + EXPECT_THAT(session->ExecutePush(manager, std::vector{0}), + StatusIs(absl::StatusCode::kFailedPrecondition, + HasSubstr("transport server"))); + EXPECT_TRUE(session->Done()); + EXPECT_THAT(session->GetStatus(), + StatusIs(absl::StatusCode::kFailedPrecondition)); + EXPECT_FALSE(manager.base()->HasActivePlan(plan.uuid())); + // A finished send refuses further pushes with its terminal status. + EXPECT_THAT(session->ExecutePush(manager, std::vector{0}), + StatusIs(absl::StatusCode::kFailedPrecondition)); +} + +// A failure on the pipeline thread (here: the staging slot wait running +// into the session deadline) is reported out of band: ExecutePush returns +// OK, the send settles with the error, every lease is released, and the plan +// is dropped. +TEST(ReshardSendSessionTest, PipelineSlotWaitTimeoutSettlesSendOutOfBand) { + KVCacheManagerWithTransfer manager( + /*num_layers=*/2, /*num_shards=*/1, /*slice_byte_size=*/128, + /*local_port=*/std::nullopt, + /*host_blocks_to_allocate=*/std::make_optional(4)); + ASSERT_EQ(manager.base()->num_host_layers(), 2u); + std::vector pools = { + DensePool("fa", 128), + DensePool("state", 128), + }; + pools[1].storage_index = 1; + ABSL_ASSERT_OK(manager.base()->RegisterPools(pools)); + manager.base()->AttachPlaceholderDeviceHoldForTest(); + std::unique_ptr allocator = + CreateStagingBlockAllocator(manager.base(), /*num_slots=*/0, + /*max_blocks=*/0); + StartTransferRequest plan = ValidPlan( + /*uuid=*/1402, /*dtype_tags=*/{"bf16", "bf16"}, + /*transferred_pools=*/{0, 1}); + TF_ASSERT_OK_AND_ASSIGN( + std::shared_ptr session, + ReshardSendSession::Create( + manager.base(), allocator.get(), std::vector{0}, + /*parallelism=*/1, + std::chrono::steady_clock::now() + std::chrono::milliseconds(200), + plan)); + + // Park every slot on units the pipeline never releases, so pool 0's + // AcquireStagingSlot can only end at the deadline. + constexpr size_t kParkedUnits[] = {7, 8}; + EXPECT_THAT(session->AcquireStagingSlot(kParkedUnits[0]), IsOkAndHolds(0u)); + EXPECT_THAT(session->AcquireStagingSlot(kParkedUnits[1]), IsOkAndHolds(1u)); + + ABSL_ASSERT_OK(session->ExecutePush(manager, std::vector{0})); + EXPECT_THAT(session->AwaitForDone(), + StatusIs(absl::StatusCode::kDeadlineExceeded)); + EXPECT_TRUE(session->Done()); + EXPECT_FALSE(manager.base()->HasActivePlan(plan.uuid())); + // SettleLocked released the parked leases along with everything else. + for (size_t unit : kParkedUnits) { + EXPECT_THAT(session->LookupStagingSlot(unit), + StatusIs(absl::StatusCode::kDeadlineExceeded)); + } + EXPECT_THAT(session->LookupStagingSlot(0), + StatusIs(absl::StatusCode::kDeadlineExceeded)); +} + +// A pool pushing to several peers keeps its staging slot until the last +// peer push completes, then the send settles clean. +TEST(ReshardSendSessionTest, StagingSlotReleasedAfterLastPeerPush) { + KVCacheManagerWithTransfer manager = MakeTestManager(); + kv_cache::KVCacheManagerBase receiver_a = MakeTestBase(); + kv_cache::KVCacheManagerBase receiver_b = MakeTestBase(); + const std::string peer_a = + absl::StrCat("127.0.0.1:", receiver_a.transport_server()->local_port()); + const std::string peer_b = + absl::StrCat("127.0.0.1:", receiver_b.transport_server()->local_port()); + manager.base()->AttachPlaceholderDeviceHoldForTest(); + std::unique_ptr allocator = + CreateStagingBlockAllocator(manager.base(), /*num_slots=*/0, + /*max_blocks=*/0); + StartTransferRequest plan = ValidPlan(/*uuid=*/1403, /*dtype_tags=*/{""}); + auto* schedule = &(*plan.mutable_shard_push_schedules())[0]; + schedule->mutable_entries(0)->set_dst_peer(peer_a); + *schedule->add_entries() = schedule->entries(0); + schedule->mutable_entries(1)->set_dst_peer(peer_b); + TF_ASSERT_OK_AND_ASSIGN( + std::shared_ptr session, + ReshardSendSession::Create( + manager.base(), allocator.get(), std::vector{0}, + /*parallelism=*/1, + std::chrono::steady_clock::now() + std::chrono::seconds(10), plan)); + + // Stage pool 0 as the pipeline would, then issue its two peer pushes. + EXPECT_THAT(session->AcquireStagingSlot(0), IsOkAndHolds(0u)); + ReshardSendSessionTestPeer::StartPoolPush(*session, manager, /*pool_idx=*/0); + + ABSL_EXPECT_OK(session->AwaitForDone()); + EXPECT_TRUE(session->Done()); + EXPECT_FALSE(manager.base()->HasActivePlan(plan.uuid())); + EXPECT_THAT(session->LookupStagingSlot(0), + StatusIs(absl::StatusCode::kCancelled)); +} + } // namespace } // namespace tpu_raiden