Use BlockTracker to track WriteRemote status for KvCacheBackend.

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",

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<const std::string> failed,
+    absl::Span<const std::string> existing) {
+  absl::MutexLock lock(mutex_);
+  MarkExistingLocked(existing);
+  MarkFailedLocked(failed);
+}
+
+void BlockTracker::MarkFailedWithUnregistered(
+    absl::Span<const std::string> failed,
+    absl::Span<const std::string> unregistered) {
+  absl::MutexLock lock(mutex_);
+  MarkUnregisteredLocked(unregistered);
+  MarkFailedLocked(failed);
+}
+
 void BlockTracker::Update(absl::Span<const std::string> done,
                           absl::Span<const std::string> failed) {
   absl::MutexLock lock(mutex_);

tpu_sync/kv_cache/block_tracker.h

@@ -70,6 +70,12 @@ class BlockTracker {
   void MarkUnregistered(absl::Span<const std::string> block_hashes);
   void MarkUnregistered(const std::string& block_hash);
 
+  // Atomically records failed blocks alongside existing or unregistered blocks.
+  void MarkFailedWithExisting(absl::Span<const std::string> failed,
+                              absl::Span<const std::string> existing);
+  void MarkFailedWithUnregistered(absl::Span<const std::string> failed,
+                                  absl::Span<const std::string> unregistered);
+
   void Update(absl::Span<const std::string> done,
               absl::Span<const std::string> failed = {});
 

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::RemoteWriteAck>
 HostOffloadBackend::BeginWriteRemote(
     const RaidenId& dst_raiden_id, absl::Span<const std::string> block_hashes,
     absl::Span<const int32_t> 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<KVCacheStoreClient> 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::Time>(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<std::string> 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<proto::WriteRemoteResult>
+                                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<proto::PollWriteRemoteResponse> 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<std::string>(resp->existing_hashes().begin(),
+                                               resp->existing_hashes().end()));
+                  break;
+                case proto::PollWriteRemoteResponse::STORED_UNREGISTERED:
+                  save_tracker->MarkFailedWithUnregistered(
+                      hashes, std::vector<std::string>(
+                                  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;
 }
 

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<std::shared_ptr<KVCacheStoreBackend>> 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<WriteRemoteCancel> cancel;
   };
 
-  // Reports how the offer's call ended, once an ack has arrived; same shape
-  // as KVCacheStoreClient::WriteRemoteVerdictCallback.
-  using WriteRemoteVerdictCallback = std::function<void(
-      absl::Status rpc_status,
-      std::optional<::tpu_raiden::kv_cache::proto::WriteRemoteResult> 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<RemoteWriteAck> BeginWriteRemote(
       const RaidenId& dst_raiden_id, absl::Span<const std::string> block_hashes,
       absl::Span<const int32_t> 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<RaidenId, std::shared_ptr<KVCacheStoreClient>,
                       RaidenIdHash>
       store_clients_ ABSL_GUARDED_BY(mutex_);
+
+  struct Lifetime {
+    absl::Mutex mu;
+    bool is_alive ABSL_GUARDED_BY(mu) = true;
+  };
+  std::shared_ptr<Lifetime> lifetime_ = std::make_shared<Lifetime>();
 };
 
 }  // namespace kv_cache