Filter caller RaidenId and prune invalid entries in Global Registry Lookup in TPU Raiden.

PiperOrigin-RevId: 971683436

Author: juncgu-google

tpu_sync/kv_cache/global_registry/global_registry.proto

@@ -151,6 +151,10 @@ message LookupRequest {
   // first miss.
   // `bytes` for the reason given on RegisterEntry.prefix_hash.
   repeated bytes prefix_hashes = 1;
+
+  // Optional. If set, the server skips this caller and offers other holders
+  // for the same hash instead of returning the caller to itself.
+  tpu_sync.rpc.RaidenIdProto client_raiden_id = 2;
 }
 
 message LookupResponse {

tpu_sync/kv_cache/global_registry/global_registry_client.cc

@@ -132,12 +132,16 @@ absl::Status GlobalRegistryClient::Register(
 }
 
 absl::StatusOr<std::vector<KVBlockMetadata>> GlobalRegistryClient::Lookup(
-    const std::vector<std::string>& prefix_hashes) {
+    const std::vector<std::string>& prefix_hashes,
+    const RaidenId& client_raiden_id) {
   LookupRequest request;
   request.mutable_prefix_hashes()->Reserve(prefix_hashes.size());
   for (const auto& hash : prefix_hashes) {
     request.add_prefix_hashes(hash);
   }
+  if (!client_raiden_id.empty()) {
+    ToProto(client_raiden_id, request.mutable_client_raiden_id());
+  }
 
   LookupResponse response;
   grpc::ClientContext context;

tpu_sync/kv_cache/global_registry/global_registry_client.h

@@ -81,11 +81,11 @@ class GlobalRegistryClient {
   // (a hash with no active registrations). All subsequent prefix hashes in the
   // input vector are treated as misses and are omitted from the response.
   // The returned vector is aligned in order with the input `prefix_hashes` (the
-  // i-th element of the returned vector corresponds to the i-th input hash).
-  // The size of the returned vector will be equal to the number of sequential
-  // hits before the first miss.
+  // If `client_raiden_id` is set, the server excludes it when returning
+  // holders.
   absl::StatusOr<std::vector<KVBlockMetadata>> Lookup(
-      const std::vector<std::string>& prefix_hashes);
+      const std::vector<std::string>& prefix_hashes,
+      const RaidenId& client_raiden_id = {});
 
   // Unregisters a batch of KV cache entries for a raiden id asynchronously.
   // See RegisterAsync for what the future resolves to, and for the meaning of

tpu_sync/kv_cache/global_registry/global_registry_server.cc

@@ -163,6 +163,9 @@ grpc::Status GlobalRegistryServiceImpl::Register(grpc::ServerContext* context,
 grpc::Status GlobalRegistryServiceImpl::Lookup(grpc::ServerContext* context,
                                                const LookupRequest* request,
                                                LookupResponse* response) {
+  const RaidenId caller = FromProto(request->client_raiden_id());
+  const bool filter_caller = !caller.empty();
+
   absl::MutexLock lock(mutex_);
   absl::Time now = absl::Now();
 
@@ -172,23 +175,31 @@ grpc::Status GlobalRegistryServiceImpl::Lookup(grpc::ServerContext* context,
       break;
     }
 
-    const auto& entries = it->second;
-    std::vector<RegistryEntry> valid_entries;
-    valid_entries.reserve(entries.size());
+    auto& entries = it->second;
     for (const auto& entry : entries) {
-      if (entry.expire_time > now) {
-        valid_entries.push_back(entry);
+      if (entry.expire_time <= now ||
+          (filter_caller && entry.raiden_id == caller)) {
+        EraseFromOwnerIndex(entry.raiden_id, hash);
       }
     }
+    entries.erase(
+        std::remove_if(entries.begin(), entries.end(),
+                       [now, filter_caller, &caller](const RegistryEntry& entry) {
+                         return entry.expire_time <= now ||
+                                (filter_caller && entry.raiden_id == caller);
+                       }),
+        entries.end());
 
-    if (valid_entries.empty()) {
+    if (entries.empty()) {
+      registry_.erase(it);
+      round_robin_indices_.erase(hash);
       break;
     }
 
     // Round-robin selection
     size_t& idx = round_robin_indices_[hash];
-    idx = idx % valid_entries.size();
-    const auto& picked = valid_entries[idx];
+    idx = idx % entries.size();
+    const auto& picked = entries[idx];
     idx++;  // Increment for next lookup
 
     auto* meta = response->add_results();

tpu_sync/kv_cache/global_registry/global_registry_test.cc

@@ -200,6 +200,96 @@ TEST_F(GlobalRegistryTest, MultiRegistrationAndRoundRobinLookup) {
             host1.job_replica_id);
 }
 
+TEST_F(GlobalRegistryTest, LookupDoesNotNameTheCallerToItself) {
+  const std::string hash = "hash1";
+  const RaidenId caller = {"job1", "replica1", "data1", 0};
+
+  ASSERT_TRUE(client_->Register({{hash, caller, 42}}).ok());
+
+  // Caller is skipped and its stale registration is deleted.
+  auto self = client_->Lookup({hash}, caller);
+  ASSERT_TRUE(self.ok()) << self.status().ToString();
+  EXPECT_TRUE(self->empty());
+
+  // Since caller was the only holder, the registry entry was purged.
+  const RaidenId peer = {"job1", "replica2", "data1", 1};
+  auto from_peer = client_->Lookup({hash}, peer);
+  ASSERT_TRUE(from_peer.ok()) << from_peer.status().ToString();
+  EXPECT_TRUE(from_peer->empty());
+}
+
+TEST_F(GlobalRegistryTest, LookupPrunesExpiredEntriesEagerly) {
+  const std::string hash = "expired_lookup_hash";
+  const RaidenId host = {"job1", "replica1", "data1", 0};
+
+  ASSERT_TRUE(client_->Register({{hash, host, 42, absl::Seconds(1)}}).ok());
+  absl::SleepFor(absl::Seconds(2));
+
+  auto res = client_->Lookup({hash});
+  ASSERT_TRUE(res.ok()) << res.status().ToString();
+  EXPECT_TRUE(res->empty());
+}
+
+TEST_F(GlobalRegistryTest, LookupOffersAnotherHolderWhenTheCallerIsOne) {
+  const std::string hash1 = "hash1";
+  const std::string hash2 = "hash2";
+  const RaidenId caller = {"job1", "replica1", "data1", 0};
+  const RaidenId peer = {"job1", "replica2", "data1", 1};
+
+  ASSERT_TRUE(client_
+                  ->Register({{hash1, caller, 42},
+                              {hash1, peer, 43},
+                              {hash2, peer, 44}})
+                  .ok());
+
+  // Skipping caller on hash1 offers peer instead, continuing the prefix walk.
+  auto res = client_->Lookup({hash1, hash2}, caller);
+  ASSERT_TRUE(res.ok()) << res.status().ToString();
+  ASSERT_EQ(res->size(), 2);
+  EXPECT_EQ((*res)[0].raiden_id().job_replica_id(), peer.job_replica_id);
+  EXPECT_EQ((*res)[0].block_id(), 43);
+  EXPECT_EQ((*res)[1].raiden_id().job_replica_id(), peer.job_replica_id);
+  EXPECT_EQ((*res)[1].block_id(), 44);
+}
+
+TEST_F(GlobalRegistryTest, LookupRoundRobinsOverTheHoldersLeftAfterTheSkip) {
+  const std::string hash = "hash1";
+  const RaidenId caller = {"job1", "replica1", "data1", 0};
+  const RaidenId peer_a = {"job1", "replica2", "data1", 1};
+  const RaidenId peer_b = {"job1", "replica3", "data1", 2};
+
+  ASSERT_TRUE(client_
+                  ->Register({{hash, caller, 42},
+                              {hash, peer_a, 43},
+                              {hash, peer_b, 44}})
+                  .ok());
+
+  // Round-robin cycles through remaining peers.
+  auto first = client_->Lookup({hash}, caller);
+  ASSERT_TRUE(first.ok()) << first.status().ToString();
+  ASSERT_EQ(first->size(), 1);
+  EXPECT_EQ((*first)[0].raiden_id().job_replica_id(), peer_a.job_replica_id);
+
+  auto second = client_->Lookup({hash}, caller);
+  ASSERT_TRUE(second.ok()) << second.status().ToString();
+  ASSERT_EQ(second->size(), 1);
+  EXPECT_EQ((*second)[0].raiden_id().job_replica_id(), peer_b.job_replica_id);
+}
+
+TEST_F(GlobalRegistryTest, LookupWithoutACallerIdIsUnfiltered) {
+  const std::string hash = "hash1";
+  const RaidenId host = {"job1", "replica1", "data1", 0};
+
+  ASSERT_TRUE(client_->Register({{hash, host, 42}}).ok());
+
+  // Unfiltered when caller is unset.
+  auto res = client_->Lookup({hash});
+  ASSERT_TRUE(res.ok()) << res.status().ToString();
+  ASSERT_EQ(res->size(), 1);
+  EXPECT_EQ((*res)[0].raiden_id().job_replica_id(), host.job_replica_id);
+  EXPECT_EQ((*res)[0].block_id(), 42);
+}
+
 TEST_F(GlobalRegistryTest, OverwriteRegistrationSameHost) {
   std::string hash = "hash1";
   RaidenId host = {"job1", "replica1", "data1", 0};

tpu_sync/kv_cache/host_offload_backend.cc

@@ -216,7 +216,7 @@ absl::StatusOr<BlockSliceList> HostOffloadBackend::Lookup(
   // at the first one.
   std::vector<global_registry::KVBlockMetadata> remote_hits;
   if (!missing_hashes.empty() && options.enable_global && client != nullptr) {
-    auto global_res_or = client->Lookup(missing_hashes);
+    auto global_res_or = client->Lookup(missing_hashes, local_id);
     if (global_res_or.ok()) {
       remote_hits = std::move(global_res_or).value();
     } else {
@@ -263,9 +263,7 @@ absl::StatusOr<BlockSliceList> HostOffloadBackend::Lookup(
           .data_replica_idx = proto_id.data_replica_idx(),
       };
       if (remote_id == local_id) {
-        // The registry claims this block lives on this node, but it wasn't
-        // present during our sweep (or was evicted since). Returning this as
-        // an unverified HOST hit is unsafe, so this is a miss.
+        // Unverified HOST hit from a registry that did not filter the caller.
         break;
       }
       results.emplace_back(