- Introduces PinRemoteBlocks and UnpinRemoteBlocks RPCs for fault-tolerant blocking TTL leases.

tpu_raiden/api/jax/kv_cache_store_e2e_test.py

@@ -180,6 +180,7 @@ def setUpClass(cls):
 
   def setUp(self):
     super().setUp()
+    os.environ["RAIDEN_LOCAL_IP"] = "127.0.0.1"
     start_servers()
     try:
       self.devices = jax.devices("tpu")
@@ -239,6 +240,7 @@ def _run_e2e_test(self, enable_multi_numa: bool):
       )
 
     os.environ["ENABLE_MULTI_NUMA"] = "1" if enable_multi_numa else "0"
+    os.environ["RAIDEN_LOCAL_IP"] = "127.0.0.1"
 
     tpu_sharding = self.setup_shardings()
     num_blocks = 2
@@ -272,7 +274,7 @@ def _run_e2e_test(self, enable_multi_numa: bool):
     # 4. Create KVCacheManager (Worker)
     manager = kv_cache_manager.KVCacheManager(
         kv_caches=[tpu_cache],
-        local_control_port=0,
+        local_control_port=-1,
         max_blocks=num_blocks,
         num_slots=2,
         unsafe_skip_buffer_lock=self.skip_lock,
@@ -404,6 +406,7 @@ def _run_remote_read_e2e_test(
       )
 
     os.environ["ENABLE_MULTI_NUMA"] = "1" if enable_multi_numa else "0"
+    os.environ["RAIDEN_LOCAL_IP"] = "127.0.0.1"
 
     if len(self.devices) < 1:
       self.skipTest(
@@ -451,7 +454,7 @@ def _run_remote_read_e2e_test(
     )
     manager_a = kv_cache_manager.KVCacheManager(
         kv_caches=[tpu_cache_a],
-        local_control_port=0,
+        local_control_port=-1,
         max_blocks=num_blocks,
         num_slots=2,
         unsafe_skip_buffer_lock=self.skip_lock,
@@ -475,7 +478,7 @@ def _run_remote_read_e2e_test(
     )
     manager_b = kv_cache_manager.KVCacheManager(
         kv_caches=[tpu_cache_b],
-        local_control_port=0,
+        local_control_port=-1,
         max_blocks=num_blocks,
         num_slots=2,
         unsafe_skip_buffer_lock=self.skip_lock,
@@ -669,7 +672,7 @@ def test_remote_read_e2e_source_missing_block_fails(self):
     if len(self.devices) < 1:
       self.skipTest("Requires at least 1 device")
     os.environ["ENABLE_MULTI_NUMA"] = "0"
-
+    os.environ["RAIDEN_LOCAL_IP"] = "127.0.0.1"
     sharding = self.setup_sharding_for_devices(self.devices)
     num_blocks = 2
     shape = (num_blocks, 128, 8, 8, 128)
@@ -701,7 +704,7 @@ def test_remote_read_e2e_source_missing_block_fails(self):
     )
     manager_a = kv_cache_manager.KVCacheManager(
         kv_caches=[tpu_cache_a],
-        local_control_port=0,
+        local_control_port=-1,
         max_blocks=num_blocks,
         num_slots=2,
         unsafe_skip_buffer_lock=self.skip_lock,
@@ -723,7 +726,7 @@ def test_remote_read_e2e_source_missing_block_fails(self):
     )
     manager_b = kv_cache_manager.KVCacheManager(
         kv_caches=[tpu_cache_b],
-        local_control_port=0,
+        local_control_port=-1,
         max_blocks=num_blocks,
         num_slots=2,
         unsafe_skip_buffer_lock=self.skip_lock,
@@ -768,7 +771,7 @@ def test_remote_read_e2e_source_wrong_status_fails(self):
     if len(self.devices) < 1:
       self.skipTest("Requires at least 1 device")
     os.environ["ENABLE_MULTI_NUMA"] = "0"
-
+    os.environ["RAIDEN_LOCAL_IP"] = "127.0.0.1"
     sharding = self.setup_sharding_for_devices(self.devices)
     num_blocks = 2
     shape = (num_blocks, 128, 8, 8, 128)
@@ -798,7 +801,7 @@ def test_remote_read_e2e_source_wrong_status_fails(self):
     )
     manager_a = kv_cache_manager.KVCacheManager(
         kv_caches=[tpu_cache_a],
-        local_control_port=0,
+        local_control_port=-1,
         max_blocks=num_blocks,
         num_slots=2,
         unsafe_skip_buffer_lock=self.skip_lock,
@@ -835,7 +838,7 @@ def test_remote_read_e2e_source_wrong_status_fails(self):
     )
     manager_b = kv_cache_manager.KVCacheManager(
         kv_caches=[tpu_cache_b],
-        local_control_port=0,
+        local_control_port=-1,
         max_blocks=num_blocks,
         num_slots=2,
         unsafe_skip_buffer_lock=self.skip_lock,

tpu_raiden/core/BUILD

@@ -730,3 +730,15 @@ cc_test(
         "@com_google_googletest//:gtest_main",
     ],
 )
+
+cc_test(
+    name = "kv_manager_holder_test",
+    srcs = ["kv_manager_holder_test.cc"],
+    deps = [
+        ":kv_manager_holder",
+        ":raiden_transfer_endpoint",
+        ":raw_transfer_core",
+        "//tpu_raiden/core/controller:test_util",
+        "@com_google_googletest//:gtest_main",
+    ],
+)

tpu_raiden/core/controller/BUILD

@@ -177,6 +177,7 @@ cc_library(
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/synchronization",
+        "@com_google_absl//absl/time",
         "@com_google_absl//absl/types:span",
         "@xla//xla/tsl/concurrency:future",
     ],
@@ -258,6 +259,7 @@ cc_test(
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/time",
         "@com_google_absl//absl/types:span",
         "@com_google_googletest//:gtest",
         "@com_google_googletest//:gtest_main",

tpu_raiden/core/controller/controller_service.cc

@@ -147,12 +147,7 @@ grpc::Status RaidenControllerServiceImpl::ReadRemote(
     }
   }
 
-  tsl::Future<> future = (*cb)(src_buffers, dst_buffers);
-  absl::Status status = future.Await();
-  if (!status.ok()) {
-    return grpc::Status(grpc::StatusCode::INTERNAL,
-                        std::string(status.message()));
-  }
+  return grpc::Status::OK;
 
   return grpc::Status::OK;
 }
@@ -171,6 +166,138 @@ std::shared_ptr<WorkerRegistry> RaidenControllerServiceImpl::worker_registry()
   return worker_registry_;
 }
 
+RaidenControllerServiceImpl::~RaidenControllerServiceImpl() {
+  {
+    absl::MutexLock lock(&mutex_);
+    sweeper_running_ = false;
+    sweeper_cv_.SignalAll();
+  }
+  if (sweeper_thread_ && sweeper_thread_->joinable()) {
+    sweeper_thread_->join();
+  }
+}
+
+void RaidenControllerServiceImpl::StartSweeperIfNecessary() {
+  if (!sweeper_running_) {
+    sweeper_running_ = true;
+    sweeper_thread_ =
+        std::make_unique<std::thread>([this]() { SweeperLoop(); });
+  }
+}
+
+void RaidenControllerServiceImpl::SweeperLoop() {
+  absl::MutexLock lock(&mutex_);
+  while (sweeper_running_) {
+    if (active_pins_.empty()) {
+      sweeper_cv_.Wait(&mutex_);
+      continue;
+    }
+    absl::Time now = absl::Now();
+    absl::Time next_wakeup = absl::InfiniteFuture();
+    for (auto it = active_pins_.begin(); it != active_pins_.end();) {
+      if (it->expiration <= now) {
+        if (unpin_cb_) {
+          (*unpin_cb_)(it->block_hashes);
+        }
+        it = active_pins_.erase(it);
+      } else {
+        if (it->expiration < next_wakeup) {
+          next_wakeup = it->expiration;
+        }
+        ++it;
+      }
+    }
+    if (next_wakeup != absl::InfiniteFuture() && sweeper_running_) {
+      sweeper_cv_.WaitWithDeadline(&mutex_, next_wakeup);
+    }
+  }
+}
+
+grpc::Status RaidenControllerServiceImpl::PinRemoteBlocks(
+    grpc::ServerContext* context,
+    const ::tpu_raiden::proto::PinRemoteBlocksRequest* request,
+    ::tpu_raiden::proto::PinRemoteBlocksResponse* response) {
+  std::shared_ptr<const ValidateAndPinCallback> validate_cb;
+  std::shared_ptr<WorkerRegistry> registry;
+  {
+    absl::MutexLock lock(&mutex_);
+    validate_cb = validate_and_pin_cb_;
+    registry = worker_registry_;
+  }
+  if (!validate_cb) {
+    return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION,
+                        "ReadRemote hooks not set.");
+  }
+  std::vector<std::string> block_hashes(request->block_hashes().begin(),
+                                        request->block_hashes().end());
+  absl::StatusOr<std::vector<int32_t>> ids_or = (*validate_cb)(block_hashes);
+  if (!ids_or.ok()) {
+    return grpc::Status(static_cast<grpc::StatusCode>(ids_or.status().code()),
+                        std::string(ids_or.status().message()));
+  }
+
+  if (registry) {
+    for (const auto& reg : registry->GetRegisteredWorkers()) {
+      auto* proto_endpoint_group = response->add_src_worker_endpoints();
+      proto_endpoint_group->set_node_id(reg.node_id);
+      proto_endpoint_group->set_worker_id(reg.worker_id);
+      for (const auto& ep : reg.raiden_transfer_endpoints) {
+        auto* ep_proto = proto_endpoint_group->add_endpoints();
+        ep_proto->set_endpoint(ep.endpoint);
+        for (int64_t shard : ep.shards) {
+          ep_proto->add_shards(shard);
+        }
+      }
+    }
+  }
+
+  {
+    absl::MutexLock lock(&mutex_);
+    PinLease lease;
+    lease.block_hashes = block_hashes;
+    lease.block_ids = *ids_or;
+    if (request->ttl_seconds() > 0) {
+      lease.expiration = absl::Now() + absl::Seconds(request->ttl_seconds());
+    } else {
+      lease.expiration = absl::Now() + absl::Seconds(60);
+    }
+    active_pins_.push_back(std::move(lease));
+    StartSweeperIfNecessary();
+    sweeper_cv_.Signal();
+  }
+
+  response->mutable_src_host_block_ids()->Assign(ids_or->begin(),
+                                                 ids_or->end());
+  return grpc::Status::OK;
+}
+
+grpc::Status RaidenControllerServiceImpl::UnpinRemoteBlocks(
+    grpc::ServerContext* context,
+    const ::tpu_raiden::proto::UnpinRemoteBlocksRequest* request,
+    ::tpu_raiden::proto::UnpinRemoteBlocksResponse* response) {
+  std::shared_ptr<const UnpinCallback> cb;
+  std::vector<std::string> hashes_to_unpin;
+  {
+    absl::MutexLock lock(&mutex_);
+    cb = unpin_cb_;
+    std::vector<int32_t> request_ids(request->src_host_block_ids().begin(),
+                                     request->src_host_block_ids().end());
+    for (auto it = active_pins_.begin(); it != active_pins_.end(); ++it) {
+      if (it->block_ids == request_ids) {
+        hashes_to_unpin = std::move(it->block_hashes);
+        active_pins_.erase(it);
+        break;
+      }
+    }
+  }
+
+  if (cb && !hashes_to_unpin.empty()) {
+    (*cb)(hashes_to_unpin);
+  }
+
+  return grpc::Status::OK;
+}
+
 }  // namespace controller
 }  // namespace core
 }  // namespace tpu_raiden

tpu_raiden/core/controller/controller_service.h

@@ -17,8 +17,10 @@
 
 #include <grpcpp/grpcpp.h>
 
+#include <chrono>
 #include <cstdint>
 #include <string>
+#include <thread>
 #include <vector>
 
 #include "absl/base/thread_annotations.h"
@@ -27,6 +29,7 @@
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "absl/synchronization/mutex.h"
+#include "absl/time/time.h"
 #include "absl/types/span.h"
 #include "grpcpp/server_context.h"
 #include "grpcpp/support/status.h"
@@ -47,7 +50,7 @@ class RaidenControllerServiceImpl final
  public:
   explicit RaidenControllerServiceImpl(
       std::shared_ptr<WorkerRegistry> worker_registry = nullptr);
-  ~RaidenControllerServiceImpl() override = default;
+  ~RaidenControllerServiceImpl() override;
 
   // Disallow copy and assign
   RaidenControllerServiceImpl(const RaidenControllerServiceImpl&) = delete;
@@ -65,6 +68,16 @@ class RaidenControllerServiceImpl final
       const ::tpu_raiden::proto::ReadRemoteRequest* request,
       ::tpu_raiden::proto::ReadRemoteResponse* response) override;
 
+  grpc::Status PinRemoteBlocks(
+      grpc::ServerContext* context,
+      const ::tpu_raiden::proto::PinRemoteBlocksRequest* request,
+      ::tpu_raiden::proto::PinRemoteBlocksResponse* response) override;
+
+  grpc::Status UnpinRemoteBlocks(
+      grpc::ServerContext* context,
+      const ::tpu_raiden::proto::UnpinRemoteBlocksRequest* request,
+      ::tpu_raiden::proto::UnpinRemoteBlocksResponse* response) override;
+
   using TransferBuffersCallback = absl::AnyInvocable<tsl::Future<>(
       absl::Span<const Buffer> src_buffers,
       absl::Span<const Buffer> dst_buffers) const>;
@@ -107,6 +120,19 @@ class RaidenControllerServiceImpl final
   std::shared_ptr<const ValidateAndPinCallback> validate_and_pin_cb_
       ABSL_GUARDED_BY(mutex_);
   std::shared_ptr<const UnpinCallback> unpin_cb_ ABSL_GUARDED_BY(mutex_);
+
+  struct PinLease {
+    std::vector<std::string> block_hashes;
+    absl::Time expiration;
+    std::vector<int32_t> block_ids;
+  };
+  std::vector<PinLease> active_pins_ ABSL_GUARDED_BY(mutex_);
+  bool sweeper_running_ ABSL_GUARDED_BY(mutex_) = false;
+  std::unique_ptr<std::thread> sweeper_thread_;
+  absl::CondVar sweeper_cv_;
+
+  void StartSweeperIfNecessary() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
+  void SweeperLoop();
 };
 
 }  // namespace controller