From 22a8bf22d919298f1b257221b74399497dd9aede Mon Sep 17 00:00:00 2001 From: justinlu Date: Sun, 6 Sep 2026 20:29:38 -0700 Subject: [PATCH] Use BlockTracker to track WriteRemote status for KvCacheBackend. PiperOrigin-RevId: 977243211 --- tpu_sync/kv_cache/BUILD | 5 +- tpu_sync/kv_cache/block_tracker.cc | 16 ++ tpu_sync/kv_cache/block_tracker.h | 6 + tpu_sync/kv_cache/host_offload_backend.cc | 110 +++++++++++- tpu_sync/kv_cache/host_offload_backend.h | 29 ++-- tpu_sync/kv_cache/kv_cache_store.cc | 156 ++---------------- tpu_sync/kv_cache/kv_cache_store_client.cc | 132 +++++++++------ tpu_sync/kv_cache/kv_cache_store_client.h | 23 ++- .../kv_cache/kv_cache_store_client_test.cc | 73 ++++++++ 9 files changed, 314 insertions(+), 236 deletions(-) diff --git a/tpu_sync/kv_cache/BUILD b/tpu_sync/kv_cache/BUILD index da928365..134199e2 100644 --- a/tpu_sync/kv_cache/BUILD +++ b/tpu_sync/kv_cache/BUILD @@ -336,6 +336,7 @@ cc_library( hdrs = ["host_offload_backend.h"], visibility = ["//visibility:public"], deps = [ + ":block_tracker", ":kv_cache_metadata", ":kv_cache_store_backend", ":kv_cache_store_backend_factory", @@ -726,7 +727,7 @@ cc_library( hdrs = ["kv_cache_store_client.h"], visibility = ["//visibility:public"], deps = [ - ":completion_executor", + ":block_tracker", "//tpu_sync/proto:kv_cache_store_service_cc_grpc", "//tpu_sync/proto:kv_cache_store_service_cc_proto", "//tpu_sync/proto:worker_service_cc_proto", @@ -748,6 +749,7 @@ cc_test( srcs = ["kv_cache_store_client_test.cc"], local_defines = ["ABSL_DEFINE_UNQUALIFIED_STATUS_TESTING_MACROS"], deps = [ + ":block_tracker", ":kv_cache_store_client", "//tpu_sync/core:raw_transfer_core", "//tpu_sync/proto:kv_cache_store_service_cc_grpc", @@ -756,6 +758,7 @@ cc_test( "@com_github_grpc_grpc//:grpc++", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_matchers", + "@com_google_absl//absl/time", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", "@xla//xla/tsl/concurrency:future", diff --git a/tpu_sync/kv_cache/block_tracker.cc b/tpu_sync/kv_cache/block_tracker.cc index 3754bde7..aead47ce 100644 --- a/tpu_sync/kv_cache/block_tracker.cc +++ b/tpu_sync/kv_cache/block_tracker.cc @@ -146,6 +146,22 @@ void BlockTracker::MarkUnregistered(const std::string& block_hash) { MarkUnregistered(absl::MakeConstSpan(&block_hash, 1)); } +void BlockTracker::MarkFailedWithExisting( + absl::Span failed, + absl::Span existing) { + absl::MutexLock lock(mutex_); + MarkExistingLocked(existing); + MarkFailedLocked(failed); +} + +void BlockTracker::MarkFailedWithUnregistered( + absl::Span failed, + absl::Span unregistered) { + absl::MutexLock lock(mutex_); + MarkUnregisteredLocked(unregistered); + MarkFailedLocked(failed); +} + void BlockTracker::Update(absl::Span done, absl::Span failed) { absl::MutexLock lock(mutex_); diff --git a/tpu_sync/kv_cache/block_tracker.h b/tpu_sync/kv_cache/block_tracker.h index ec6924bb..6b954b14 100644 --- a/tpu_sync/kv_cache/block_tracker.h +++ b/tpu_sync/kv_cache/block_tracker.h @@ -70,6 +70,12 @@ class BlockTracker { void MarkUnregistered(absl::Span block_hashes); void MarkUnregistered(const std::string& block_hash); + // Atomically records failed blocks alongside existing or unregistered blocks. + void MarkFailedWithExisting(absl::Span failed, + absl::Span existing); + void MarkFailedWithUnregistered(absl::Span failed, + absl::Span unregistered); + void Update(absl::Span done, absl::Span failed = {}); diff --git a/tpu_sync/kv_cache/host_offload_backend.cc b/tpu_sync/kv_cache/host_offload_backend.cc index 879935cc..ba7ce5f2 100644 --- a/tpu_sync/kv_cache/host_offload_backend.cc +++ b/tpu_sync/kv_cache/host_offload_backend.cc @@ -34,6 +34,7 @@ #include "absl/strings/str_join.h" #include "absl/strings/string_view.h" #include "absl/synchronization/mutex.h" +#include "absl/time/clock.h" #include "absl/time/time.h" #include "absl/types/span.h" #include "grpcpp/create_channel.h" @@ -42,6 +43,7 @@ #include "tpu_sync/common/raiden_id.h" #include "tpu_sync/core/buffer.h" #include "tpu_sync/core/controller/raiden_controller.h" +#include "tpu_sync/kv_cache/block_tracker.h" #include "tpu_sync/kv_cache/global_registry/global_registry_client.h" #include "tpu_sync/kv_cache/kv_cache_metadata.h" #include "tpu_sync/kv_cache/kv_cache_store_backend.h" @@ -102,6 +104,10 @@ HostOffloadBackend::HostOffloadBackend( registry_client_(std::move(registry_client)) {} HostOffloadBackend::~HostOffloadBackend() { + { + absl::MutexLock lock(lifetime_->mu); + lifetime_->is_alive = false; + } if (server_) { server_->Shutdown(); } @@ -774,9 +780,8 @@ absl::StatusOr HostOffloadBackend::BeginWriteRemote( const RaidenId& dst_raiden_id, absl::Span block_hashes, absl::Span src_host_block_ids, - absl::Duration requested_deadline, - absl::Duration hold_window, - WriteRemoteVerdictCallback on_verdict) { + absl::Duration requested_deadline, absl::Duration hold_window, + BlockTracker* save_tracker) { if (block_hashes.empty()) { return absl::InvalidArgumentError("WriteRemote requires at least one hash"); } @@ -790,11 +795,12 @@ HostOffloadBackend::BeginWriteRemote( ABSL_ASSIGN_OR_RETURN(std::shared_ptr client, GetKVCacheStoreClient(dst_raiden_id)); - auto call = client->WriteRemote(raiden_controller_->unit(), block_hashes, - src_host_block_ids, - BuildLocalWorkerEndpoints(raiden_controller_), - absl::ToInt64Milliseconds(requested_deadline), - hold_window, std::move(on_verdict)); + auto hold_expiry = std::make_shared(absl::Now() + hold_window); + + auto call = client->WriteRemote( + raiden_controller_->unit(), block_hashes, src_host_block_ids, + BuildLocalWorkerEndpoints(raiden_controller_), + absl::ToInt64Milliseconds(requested_deadline), hold_window, save_tracker); auto response = call.ack.Await(); if (!response.ok()) { // On a transport error the peer may have restarted on a new port; drop @@ -811,13 +817,26 @@ HostOffloadBackend::BeginWriteRemote( ack.cancel = std::move(call.cancel); ack.operation_id = response->operation_id(); ack.granted_deadline = absl::Milliseconds(response->granted_deadline_ms()); + if (ack.granted_deadline >= hold_window) { + *hold_expiry = std::max(*hold_expiry, absl::Now() + ack.granted_deadline); + } switch (response->exist_state()) { case ::tpu_raiden::kv_cache::proto::WRITE_ALL_EXIST: ack.all_exist = true; + if (save_tracker != nullptr) { + save_tracker->MarkDone(block_hashes); + } + // Release both the transfer's hold and the caller's pin. + Release(block_hashes); + Release(block_hashes); return ack; case ::tpu_raiden::kv_cache::proto::WRITE_PARTIAL_EXIST: ack.existing_hashes.assign(response->existing_hashes().begin(), response->existing_hashes().end()); + if (save_tracker != nullptr) { + save_tracker->MarkFailedWithExisting(block_hashes, ack.existing_hashes); + } + Release(block_hashes); return ack; default: break; @@ -828,6 +847,81 @@ HostOffloadBackend::BeginWriteRemote( return absl::InternalError( "Destination accepted the offer but returned no operation id."); } + if (save_tracker != nullptr) { + std::vector hashes(block_hashes.begin(), block_hashes.end()); + call.result.OnReady([this, lifetime = lifetime_, dst_raiden_id, + save_tracker, hashes = std::move(hashes), hold_expiry, + op_id = ack.operation_id]( + absl::StatusOr + result_or) { + absl::MutexLock lock(lifetime->mu); + if (!lifetime->is_alive) { + return; + } + if (!result_or.ok() && + result_or.status().code() == absl::StatusCode::kCancelled) { + return; + } + bool succeeded = false; + if (result_or.ok()) { + const auto& result = *result_or; + succeeded = + (result.state() == proto::PollWriteRemoteResponse::COMMITTED || + result.state() == proto::PollWriteRemoteResponse::ALL_EXIST); + } else { + const absl::Duration remaining_hold = *hold_expiry - absl::Now(); + if (op_id != 0 && remaining_hold > absl::ZeroDuration()) { + auto fut = PollWriteRemoteAsync( + dst_raiden_id, op_id, absl::ToInt64Milliseconds(remaining_hold)); + fut.OnReady([this, lifetime, save_tracker, hashes]( + absl::StatusOr resp) { + absl::MutexLock lock(lifetime->mu); + if (!lifetime->is_alive) { + return; + } + bool poll_succeeded = false; + if (resp.ok()) { + switch (resp->state()) { + case proto::PollWriteRemoteResponse::COMMITTED: + case proto::PollWriteRemoteResponse::ALL_EXIST: + save_tracker->MarkDone(hashes); + poll_succeeded = true; + break; + case proto::PollWriteRemoteResponse::PARTIAL_EXIST: + save_tracker->MarkFailedWithExisting( + hashes, + std::vector(resp->existing_hashes().begin(), + resp->existing_hashes().end())); + break; + case proto::PollWriteRemoteResponse::STORED_UNREGISTERED: + save_tracker->MarkFailedWithUnregistered( + hashes, std::vector( + resp->unregistered_hashes().begin(), + resp->unregistered_hashes().end())); + break; + default: + save_tracker->MarkFailed(hashes); + break; + } + } else { + save_tracker->MarkFailed(hashes); + } + Release(hashes); + if (poll_succeeded) { + Release(hashes); + } + }); + return; + } else { + save_tracker->MarkFailed(hashes); + } + } + Release(hashes); + if (succeeded) { + Release(hashes); + } + }); + } return ack; } diff --git a/tpu_sync/kv_cache/host_offload_backend.h b/tpu_sync/kv_cache/host_offload_backend.h index de858bd1..501a25f6 100644 --- a/tpu_sync/kv_cache/host_offload_backend.h +++ b/tpu_sync/kv_cache/host_offload_backend.h @@ -52,6 +52,8 @@ class RaidenController; namespace kv_cache { +class BlockTracker; + class HostOffloadBackend : public KVCacheStoreBackend { public: static absl::StatusOr> Create( @@ -154,8 +156,8 @@ class HostOffloadBackend : public KVCacheStoreBackend { // --- Remote write, source side ------------------------------------------ // - // BeginWriteRemote offers blocks; the verdict arrives later through - // on_verdict, on the same call. PollWriteRemoteAsync is recovery for a + // BeginWriteRemote offers blocks; the verdict arrives later on the same + // call and updates save_tracker. PollWriteRemoteAsync is recovery for a // source that lost that call. KVCacheStore owns the pins. // What the destination decided, before any bytes have moved. @@ -175,23 +177,16 @@ class HostOffloadBackend : public KVCacheStoreBackend { std::shared_ptr cancel; }; - // Reports how the offer's call ended, once an ack has arrived; same shape - // as KVCacheStoreClient::WriteRemoteVerdictCallback. - using WriteRemoteVerdictCallback = std::function result, - uint64_t operation_id)>; // Offers `block_hashes` to `dst_raiden_id` and blocks until the ack (not // the bytes). `requested_deadline` is how long the destination may hold - // its landing blocks; `hold_window` is the call's deadline. `on_verdict` - // runs when the call ends; if the call fails before any ack, the failure - // is the return status and on_verdict never runs. + // its landing blocks; `hold_window` is the call's deadline. If + // `save_tracker` is non-null, its transfer status is updated when the + // operation completes or settles. absl::StatusOr BeginWriteRemote( const RaidenId& dst_raiden_id, absl::Span block_hashes, absl::Span src_host_block_ids, - absl::Duration requested_deadline, - absl::Duration hold_window, - WriteRemoteVerdictCallback on_verdict = nullptr); + absl::Duration requested_deadline, absl::Duration hold_window, + BlockTracker* save_tracker = nullptr); // Asks the destination what became of an accepted offer; recovery for a // source that lost its stream. `wait_ms > 0` asks it to hold the answer @@ -279,6 +274,12 @@ class HostOffloadBackend : public KVCacheStoreBackend { absl::flat_hash_map, RaidenIdHash> store_clients_ ABSL_GUARDED_BY(mutex_); + + struct Lifetime { + absl::Mutex mu; + bool is_alive ABSL_GUARDED_BY(mu) = true; + }; + std::shared_ptr lifetime_ = std::make_shared(); }; } // namespace kv_cache diff --git a/tpu_sync/kv_cache/kv_cache_store.cc b/tpu_sync/kv_cache/kv_cache_store.cc index cf3af1f2..cac45fd6 100644 --- a/tpu_sync/kv_cache/kv_cache_store.cc +++ b/tpu_sync/kv_cache/kv_cache_store.cc @@ -1828,157 +1828,21 @@ absl::Status KVCacheStore::SaveRemote( // settle path, never by rollback. std::move(rollback).Cancel(); - // Runs when the offer's call ends after the ack: with a result, an error, - // or its deadline. (A call that fails before any ack reports through - // BeginWriteRemote's return instead.) Every settle claims the operation - // via TakeRemoteWrite, so exactly one path settles it. - auto on_verdict = [lifetime = lifetime_, op_key, dst_raiden_id]( - absl::Status status, - std::optional result, - uint64_t stream_op_id) { - absl::MutexLock lock(lifetime->mu); - if (lifetime->store == nullptr) { - return; - } - - auto settle_verdict = [](KVCacheStore* store, OperationKey key, - proto::PollWriteRemoteResponse::State state, - std::vector existing_hashes, - std::vector unregistered_hashes) { - auto taken = store->TakeRemoteWrite(key); - if (!taken.has_value()) { - return; - } - switch (state) { - case proto::PollWriteRemoteResponse::COMMITTED: - case proto::PollWriteRemoteResponse::ALL_EXIST: - store->OnWriteRemoteVerdict(std::move(*taken), - /*succeeded=*/true, {}); - break; - case proto::PollWriteRemoteResponse::PARTIAL_EXIST: - store->OnWriteRemoteVerdict(std::move(*taken), /*succeeded=*/false, - std::move(existing_hashes)); - break; - case proto::PollWriteRemoteResponse::STORED_UNREGISTERED: - store->OnWriteRemoteVerdict(std::move(*taken), /*succeeded=*/false, {}, - std::move(unregistered_hashes)); - break; - default: - store->OnWriteRemoteVerdict(std::move(*taken), - /*succeeded=*/false, {}); - break; - } - }; - - if (result.has_value()) { - settle_verdict( - lifetime->store, op_key, result->state(), - std::vector(result->existing_hashes().begin(), - result->existing_hashes().end()), - std::vector(result->unregistered_hashes().begin(), - result->unregistered_hashes().end())); - } else { - // The call ended without a verdict. Decide on the operation's CURRENT - // hold_expiry, not on the call's status: the ack may have extended the - // hold past the call's fixed deadline, in which case the destination - // may still be pulling and we must recover, not settle. - uint64_t op_id = stream_op_id; - absl::Time current_hold_expiry = absl::InfinitePast(); - { - absl::MutexLock lock(lifetime->store->mutex_); - auto it = lifetime->store->active_remote_writes_.find(op_key); - if (it == lifetime->store->active_remote_writes_.end()) { - // Already settled elsewhere; nothing to do. - return; - } - if (op_id == 0) { - op_id = it->second.operation_id; - } - current_hold_expiry = it->second.hold_expiry; - } - const absl::Duration remaining_hold = current_hold_expiry - absl::Now(); - if (op_id == 0 || remaining_hold <= absl::ZeroDuration()) { - // No operation id to ask about, or the hold has fully elapsed: - // settle as failed. - settle_verdict(lifetime->store, op_key, - proto::PollWriteRemoteResponse::FAILED, {}, {}); - return; - } - // Still inside the hold: ask the destination once to hold the answer - // until the operation is terminal. One attempt; the continuation - // always settles. - auto* host_backend = - dynamic_cast(lifetime->store->backends_[0].get()); - if (host_backend != nullptr) { - auto fut = host_backend->PollWriteRemoteAsync( - dst_raiden_id, op_id, absl::ToInt64Milliseconds(remaining_hold)); - fut.OnReady([lifetime, op_key, settle_verdict]( - absl::StatusOr resp) { - CompletionExecutor::Schedule( - [lifetime, op_key, settle_verdict, resp = std::move(resp)]() { - absl::MutexLock lock(lifetime->mu); - if (lifetime->store == nullptr) { - return; - } - if (resp.ok()) { - settle_verdict( - lifetime->store, op_key, resp->state(), - std::vector(resp->existing_hashes().begin(), - resp->existing_hashes().end()), - std::vector(resp->unregistered_hashes().begin(), - resp->unregistered_hashes().end())); - } else { - settle_verdict(lifetime->store, op_key, - proto::PollWriteRemoteResponse::FAILED, {}, {}); - } - }); - }); - return; - } - settle_verdict(lifetime->store, op_key, - proto::PollWriteRemoteResponse::FAILED, {}, {}); - } - }; - - auto ack_res = backend->BeginWriteRemote( - dst_raiden_id, block_hashes, src_host_block_ids, - hold - kRemoteWriteMargin, hold, std::move(on_verdict)); + BlockTracker* save_tracker = + (owner == SaveOwner::kApplication) ? &save_tracker_ : &sweep_tracker_; + auto ack_res = + backend->BeginWriteRemote(dst_raiden_id, block_hashes, src_host_block_ids, + hold - kRemoteWriteMargin, hold, save_tracker); if (!ack_res.ok()) { - // The offer failed before any ack. Undo it completely: release the pin, - // clear the marks, and report only through the return status. - // - // KNOWN GAP: the destination starts pulling before it answers, so if - // only the answer was lost it may still be reading these blocks. We - // release anyway: holding for the full HOLD would stall the sweep - // whenever a peer is simply down, and gRPC reports both cases as - // UNAVAILABLE. - if (auto taken = TakeRemoteWrite(op_key); taken.has_value()) { - backend->Release(taken->block_hashes); - if (taken->owner == SaveOwner::kApplication) { - save_tracker_.RemovePending(taken->block_hashes); - } else { - sweep_tracker_.RemovePending(taken->block_hashes); - } - } + backend->Release(block_hashes); + save_tracker->RemovePending(block_hashes); + TakeRemoteWrite(op_key); return ack_res.status(); } const auto& ack = *ack_res; - if (ack.all_exist) { - // SUCCESS with nothing to wait for. - auto taken = TakeRemoteWrite(op_key); - if (taken.has_value()) { - OnWriteRemoteVerdict(std::move(*taken), /*succeeded=*/true, {}); - } - return absl::OkStatus(); - } - if (!ack.existing_hashes.empty()) { - // FAILURE. The caller gets the list and decides what to re-offer. - auto taken = TakeRemoteWrite(op_key); - if (taken.has_value()) { - OnWriteRemoteVerdict(std::move(*taken), /*succeeded=*/false, - std::move(ack.existing_hashes)); - } + if (ack.all_exist || !ack.existing_hashes.empty()) { + TakeRemoteWrite(op_key); return absl::OkStatus(); } diff --git a/tpu_sync/kv_cache/kv_cache_store_client.cc b/tpu_sync/kv_cache/kv_cache_store_client.cc index 8c900c08..dfe8450e 100644 --- a/tpu_sync/kv_cache/kv_cache_store_client.cc +++ b/tpu_sync/kv_cache/kv_cache_store_client.cc @@ -12,7 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include "tpu_sync/kv_cache/completion_executor.h" #include "tpu_sync/kv_cache/kv_cache_store_client.h" #include @@ -21,6 +20,7 @@ #include #include #include +#include #include "absl/status/status.h" #include "absl/status/statusor.h" @@ -32,6 +32,7 @@ #include "grpcpp/impl/status.h" #include "grpcpp/support/client_callback.h" #include "xla/tsl/concurrency/future.h" +#include "tpu_sync/kv_cache/block_tracker.h" #include "tpu_sync/proto/kv_cache_store_service.grpc.pb.h" #include "tpu_sync/proto/kv_cache_store_service.pb.h" #include "tpu_sync/rpc/raiden_service.pb.h" @@ -124,9 +125,10 @@ KVCacheStoreClient::Fetch( } // The source's end of one WriteRemote call. The ack resolves `ack_promise_`. -// When the call ends, OnDone hands the status and any streamed result to -// `on_verdict_` via the CompletionExecutor, then deletes the reactor. If the -// call fails before any ack, only the ack future reports the error. +// When the stream ends, OnDone fulfills `result_promise_` with the verdict +// or error, then deletes the reactor. If `tracker_` is provided, OnDone +// updates it with the streamed verdict. If the call fails before any ack, both +// the ack and result futures report the error. class WriteRemoteClientReactor : public ::grpc::ClientReadReactor< ::tpu_raiden::kv_cache::proto::WriteRemoteEvent> { @@ -136,11 +138,14 @@ class WriteRemoteClientReactor ::tpu_raiden::kv_cache::proto::WriteRemoteRequest request, std::shared_ptr<::grpc::ClientContext> context, tsl::Promise<::tpu_raiden::kv_cache::proto::WriteRemoteAck> ack_promise, - KVCacheStoreClient::WriteRemoteVerdictCallback on_verdict) + tsl::Promise<::tpu_raiden::kv_cache::proto::WriteRemoteResult> + result_promise, + BlockTracker* tracker) : context_(std::move(context)), request_(std::move(request)), ack_promise_(std::move(ack_promise)), - on_verdict_(std::move(on_verdict)) { + result_promise_(std::move(result_promise)), + tracker_(tracker) { stub->async()->WriteRemote(context_.get(), &request_, this); StartRead(&event_); StartCall(); @@ -153,16 +158,15 @@ class WriteRemoteClientReactor if (event_.has_ack()) { ack_received_ = true; - operation_id_ = event_.ack().operation_id(); ack_promise_.Set(event_.ack()); if (event_.ack().exist_state() == ::tpu_raiden::kv_cache::proto::WRITE_EXIST_STATE_UNSPECIFIED) { StartRead(&event_); } else { // Existence answers (ALL_EXIST / PARTIAL_EXIST) settle synchronously - // inside SaveRemote via the ack promise. Clearing on_verdict_ prevents - // OnDone from scheduling an unneeded verdict that races with SaveRemote. - on_verdict_ = nullptr; + // inside SaveRemote/BeginWriteRemote via the ack promise. Clearing + // tracker_ prevents OnDone from recording an unneeded verdict. + tracker_ = nullptr; } } else if (event_.has_result()) { has_result_ = true; @@ -174,29 +178,47 @@ class WriteRemoteClientReactor void OnDone(const ::grpc::Status& status) override { bool initial_ack_failed = !ack_received_; if (!ack_received_) { - ack_promise_.Set(absl::Status( + absl::Status rpc_status( static_cast(status.error_code()), - status.error_message())); + status.error_message()); + ack_promise_.Set(rpc_status); + result_promise_.Set(rpc_status); ack_received_ = true; + } else if (has_result_) { + result_promise_.Set(result_); + } else if (status.ok()) { + result_promise_.Set( + absl::InternalError("WriteRemote stream closed without result")); + } else { + result_promise_.Set( + absl::Status(static_cast(status.error_code()), + status.error_message())); } - absl::Status rpc_status = - status.ok() - ? absl::OkStatus() - : absl::Status(static_cast(status.error_code()), - status.error_message()); - - std::optional result; - if (has_result_) { - result = std::move(result_); - } - - if (on_verdict_ && !initial_ack_failed) { - CompletionExecutor::Schedule( - [on_verdict = std::move(on_verdict_), rpc_status, - result = std::move(result), op_id = operation_id_]() mutable { - on_verdict(rpc_status, std::move(result), op_id); - }); + if (tracker_ != nullptr && !initial_ack_failed && has_result_) { + std::vector hashes(request_.block_hashes().begin(), + request_.block_hashes().end()); + switch (result_.state()) { + case proto::PollWriteRemoteResponse::COMMITTED: + case proto::PollWriteRemoteResponse::ALL_EXIST: + tracker_->MarkDone(hashes); + break; + case proto::PollWriteRemoteResponse::PARTIAL_EXIST: + tracker_->MarkFailedWithExisting( + hashes, + std::vector(result_.existing_hashes().begin(), + result_.existing_hashes().end())); + break; + case proto::PollWriteRemoteResponse::STORED_UNREGISTERED: + tracker_->MarkFailedWithUnregistered( + hashes, + std::vector(result_.unregistered_hashes().begin(), + result_.unregistered_hashes().end())); + break; + default: + tracker_->MarkFailed(hashes); + break; + } } delete this; @@ -211,10 +233,10 @@ class WriteRemoteClientReactor proto::WriteRemoteEvent event_; // Resolved by the ack, or by the call's error if no ack arrived. tsl::Promise ack_promise_; - // Invoked from OnDone with the call's outcome. May be null. - KVCacheStoreClient::WriteRemoteVerdictCallback on_verdict_; - // Operation id from the ack; 0 until the ack arrives. - uint64_t operation_id_ = 0; + // Resolved when the stream delivers a verdict or ends. + tsl::Promise result_promise_; + // Updated with the streamed verdict if non-null. + BlockTracker* tracker_ = nullptr; // True once ack_promise_ has been set. bool ack_received_ = false; // True when the peer streamed a result; result_ then holds it. @@ -228,31 +250,30 @@ KVCacheStoreClient::WriteRemoteCall KVCacheStoreClient::WriteRemote( absl::Span src_host_block_ids, absl::Span src_worker_endpoints, - int64_t deadline_ms, - absl::Duration hold_window, - WriteRemoteVerdictCallback on_verdict) { + int64_t deadline_ms, absl::Duration hold_window, BlockTracker* tracker) { if (block_hashes.empty()) { + absl::Status status = + absl::InvalidArgumentError("WriteRemote requires at least one hash."); return WriteRemoteCall{ - tsl::Future<::tpu_raiden::kv_cache::proto::WriteRemoteAck>( - absl::InvalidArgumentError( - "WriteRemote requires at least one hash.")), + tsl::Future<::tpu_raiden::kv_cache::proto::WriteRemoteAck>(status), + tsl::Future<::tpu_raiden::kv_cache::proto::WriteRemoteResult>(status), nullptr}; } if (src_host_block_ids.size() != block_hashes.size()) { + absl::Status status = absl::InvalidArgumentError(absl::StrCat( + "Mismatched src_host_block_ids count (", src_host_block_ids.size(), + ") vs block_hashes count (", block_hashes.size(), ").")); return WriteRemoteCall{ - tsl::Future<::tpu_raiden::kv_cache::proto::WriteRemoteAck>( - absl::InvalidArgumentError(absl::StrCat( - "Mismatched src_host_block_ids count (", - src_host_block_ids.size(), ") vs block_hashes count (", - block_hashes.size(), ")."))), + tsl::Future<::tpu_raiden::kv_cache::proto::WriteRemoteAck>(status), + tsl::Future<::tpu_raiden::kv_cache::proto::WriteRemoteResult>(status), nullptr}; } if (deadline_ms <= 0) { + absl::Status status = absl::InvalidArgumentError(absl::StrCat( + "WriteRemote requires a positive deadline_ms, got ", deadline_ms, ".")); return WriteRemoteCall{ - tsl::Future<::tpu_raiden::kv_cache::proto::WriteRemoteAck>( - absl::InvalidArgumentError(absl::StrCat( - "WriteRemote requires a positive deadline_ms, got ", - deadline_ms, "."))), + tsl::Future<::tpu_raiden::kv_cache::proto::WriteRemoteAck>(status), + tsl::Future<::tpu_raiden::kv_cache::proto::WriteRemoteResult>(status), nullptr}; } @@ -270,8 +291,10 @@ KVCacheStoreClient::WriteRemoteCall KVCacheStoreClient::WriteRemote( } request.set_deadline_ms(deadline_ms); - auto [promise, future] = + auto [ack_promise, ack_future] = tsl::MakePromise<::tpu_raiden::kv_cache::proto::WriteRemoteAck>(); + auto [result_promise, result_future] = + tsl::MakePromise<::tpu_raiden::kv_cache::proto::WriteRemoteResult>(); // The call's deadline is the hold window; it is what ends a call that // never gets an answer. @@ -284,21 +307,22 @@ KVCacheStoreClient::WriteRemoteCall KVCacheStoreClient::WriteRemote( // the life of the call. auto cancel = std::make_shared(); { - absl::MutexLock lock(&cancel->state_->mutex); + absl::MutexLock lock(cancel->state_->mutex); cancel->state_->context = context; } // Owns itself until OnDone. new WriteRemoteClientReactor(stub_.get(), std::move(request), - std::move(context), std::move(promise), - std::move(on_verdict)); - return WriteRemoteCall{std::move(future), std::move(cancel)}; + std::move(context), std::move(ack_promise), + std::move(result_promise), tracker); + return WriteRemoteCall{std::move(ack_future), std::move(result_future), + std::move(cancel)}; } void WriteRemoteCancel::TryCancel() { std::shared_ptr<::grpc::ClientContext> context; { - absl::MutexLock lock(&state_->mutex); + absl::MutexLock lock(state_->mutex); context = state_->context.lock(); } // If the call already ended, the weak pointer is empty and there is diff --git a/tpu_sync/kv_cache/kv_cache_store_client.h b/tpu_sync/kv_cache/kv_cache_store_client.h index 406cfe8a..64463946 100644 --- a/tpu_sync/kv_cache/kv_cache_store_client.h +++ b/tpu_sync/kv_cache/kv_cache_store_client.h @@ -23,6 +23,7 @@ #include "absl/base/thread_annotations.h" #include "absl/status/statusor.h" #include "absl/synchronization/mutex.h" +#include "absl/time/time.h" #include "absl/types/span.h" #include "grpcpp/channel.h" #include "grpcpp/client_context.h" @@ -34,6 +35,8 @@ namespace tpu_raiden { namespace kv_cache { +class BlockTracker; + // Cancels an open WriteRemote call from outside it. ~KVCacheStore uses this // to end abandoned offers, so the destination is not left holding a call // nobody will answer. Holds the call's context weakly: the handle does not @@ -74,21 +77,16 @@ class KVCacheStoreClient { absl::Span client_worker_endpoints = {}); - // Reports how the offer's call ended, once an ack has arrived: `result` is - // the streamed verdict if the peer sent one; otherwise `rpc_status` says - // how the call ended. - using WriteRemoteVerdictCallback = std::function result, - uint64_t operation_id)>; // Offers `block_hashes` to the connected peer, on one streaming call whose // deadline is `hold_window`. `ack` resolves when the peer has decided; it - // does not wait for the bytes. `on_verdict` runs on the CompletionExecutor - // when the call ends. If the call fails before any ack, the error goes to - // `ack` and `on_verdict` never runs. `deadline_ms` must be > 0. + // does not wait for the bytes. `result` resolves when the stream ends with + // a verdict or error. If `tracker` is non-null, its status is updated with + // the streamed verdict when the call ends. `deadline_ms` must be > 0. struct WriteRemoteCall { // The peer's decision, or the error that ended the call before it. tsl::Future<::tpu_raiden::kv_cache::proto::WriteRemoteAck> ack; + // The final result streamed by the peer, or the error if the stream failed. + tsl::Future<::tpu_raiden::kv_cache::proto::WriteRemoteResult> result; // Cancels the open call. Null only if this client refused to make it. std::shared_ptr cancel; }; @@ -98,9 +96,8 @@ class KVCacheStoreClient { absl::Span src_host_block_ids, absl::Span src_worker_endpoints, - int64_t deadline_ms, - absl::Duration hold_window, - WriteRemoteVerdictCallback on_verdict = nullptr); + int64_t deadline_ms, absl::Duration hold_window, + BlockTracker* tracker = nullptr); // Asks the peer what became of an accepted operation. `wait_ms > 0` asks // it to hold the answer until the operation is terminal. UNKNOWN means the diff --git a/tpu_sync/kv_cache/kv_cache_store_client_test.cc b/tpu_sync/kv_cache/kv_cache_store_client_test.cc index d4c53d8b..9bf549d2 100644 --- a/tpu_sync/kv_cache/kv_cache_store_client_test.cc +++ b/tpu_sync/kv_cache/kv_cache_store_client_test.cc @@ -23,13 +23,16 @@ #include #include "absl/status/status.h" #include "absl/status/status_matchers.h" +#include "absl/time/time.h" #include "grpcpp/grpcpp.h" #include "grpcpp/security/credentials.h" #include "grpcpp/security/server_credentials.h" #include "grpcpp/support/status.h" +#include "grpcpp/support/sync_stream.h" #include "xla/tsl/concurrency/future.h" #include "xla/tsl/platform/statusor.h" #include "tpu_sync/core/raiden_future.h" +#include "tpu_sync/kv_cache/block_tracker.h" #include "tpu_sync/proto/kv_cache_store_service.grpc.pb.h" #include "tpu_sync/proto/kv_cache_store_service.pb.h" #include "tpu_sync/rpc/raiden_service.pb.h" @@ -60,13 +63,36 @@ class TestKVCacheStoreService return ::grpc::Status::OK; } + ::grpc::Status WriteRemote( + ::grpc::ServerContext* context, + const ::tpu_raiden::kv_cache::proto::WriteRemoteRequest* request, + ::grpc::ServerWriter<::tpu_raiden::kv_cache::proto::WriteRemoteEvent>* + writer) override { + if (fail_write_remote_) { + return ::grpc::Status(::grpc::StatusCode::INTERNAL, + "Simulated WriteRemote RPC error"); + } + ::tpu_raiden::kv_cache::proto::WriteRemoteEvent ack_event; + ack_event.mutable_ack()->set_operation_id(12345); + ack_event.mutable_ack()->set_granted_deadline_ms(5000); + writer->Write(ack_event); + + ::tpu_raiden::kv_cache::proto::WriteRemoteEvent result_event; + result_event.mutable_result()->set_state( + ::tpu_raiden::kv_cache::proto::PollWriteRemoteResponse::COMMITTED); + writer->Write(result_event); + return ::grpc::Status::OK; + } + void SetFailRpc(bool fail) { fail_rpc_ = fail; } + void SetFailWriteRemote(bool fail) { fail_write_remote_ = fail; } const ::tpu_raiden::kv_cache::proto::FetchRequest& last_request() const { return last_request_; } private: bool fail_rpc_ = false; + bool fail_write_remote_ = false; ::tpu_raiden::kv_cache::proto::FetchRequest last_request_; }; @@ -177,6 +203,53 @@ TEST_F(KVCacheStoreClientTest, FetchEmptyHashesReturnsEmptyResponse) { EXPECT_EQ(response.done_block_hashes_size(), 0); } +TEST_F(KVCacheStoreClientTest, WriteRemoteUpdatesBlockTrackerOnSuccess) { + BlockTracker tracker; + std::vector hashes = {"hash_1", "hash_2"}; + std::vector host_ids = {10, 11}; + tracker.AddPending(hashes); + + auto call = client_->WriteRemote( + ::tpu_sync::rpc::RaidenIdProto(), hashes, host_ids, {}, + /*deadline_ms=*/5000, absl::Seconds(10), &tracker); + auto ack = call.ack.Await(); + ABSL_ASSERT_OK(ack); + EXPECT_EQ(ack->operation_id(), 12345); + + auto result = call.result.Await(); + ABSL_ASSERT_OK(result); + EXPECT_EQ(result->state(), + ::tpu_raiden::kv_cache::proto::PollWriteRemoteResponse::COMMITTED); + + auto status = tracker.Poll(); + EXPECT_THAT(status.done, UnorderedElementsAre("hash_1", "hash_2")); + EXPECT_TRUE(status.pending.empty()); +} + +TEST_F(KVCacheStoreClientTest, WriteRemoteValidatesArguments) { + BlockTracker tracker; + // Empty hashes + auto empty_call = + client_->WriteRemote(::tpu_sync::rpc::RaidenIdProto(), {}, {}, {}, 5000, + absl::Seconds(10), &tracker); + EXPECT_THAT(empty_call.ack.Await().status(), + StatusIs(absl::StatusCode::kInvalidArgument)); + + // Mismatched host ids + auto mismatch_call = + client_->WriteRemote(::tpu_sync::rpc::RaidenIdProto(), {"a"}, {1, 2}, {}, + 5000, absl::Seconds(10), &tracker); + EXPECT_THAT(mismatch_call.ack.Await().status(), + StatusIs(absl::StatusCode::kInvalidArgument)); + + // Non-positive deadline + auto deadline_call = + client_->WriteRemote(::tpu_sync::rpc::RaidenIdProto(), {"a"}, {1}, {}, 0, + absl::Seconds(10), &tracker); + EXPECT_THAT(deadline_call.ack.Await().status(), + StatusIs(absl::StatusCode::kInvalidArgument)); +} + } // namespace } // namespace kv_cache } // namespace tpu_raiden