Add read-only request block status probe in reshard stack

tpu_sync/api/torch/reshard_client.py

@@ -48,6 +48,18 @@
 # coordinate_transfer. Absent on older shims, so consumers must getattr.
 SUPPORTS_DST_SKIP_BYTES = bool(
     getattr(_impl, "reshard_client_supports_dst_skip_bytes", False))
+# True iff get_request_block_status (read-only registry lifecycle probe)
+# is available; absent on older shims, so consumers must getattr.
+SUPPORTS_REQUEST_BLOCK_STATUS = bool(
+    getattr(_impl, "reshard_client_supports_request_block_status", False)
+)
+
+# GetRequestBlockStatusResponse.Status values returned by
+# get_request_block_status.
+REQUEST_BLOCK_STATUS_UNKNOWN = 1
+REQUEST_BLOCK_STATUS_REGISTERED = 2
+REQUEST_BLOCK_STATUS_CLAIMED = 3
+REQUEST_BLOCK_STATUS_CANCELLED = 4
 
 
 def _unit_tuple(unit: Any) -> tuple:
@@ -191,6 +203,20 @@ def cancel_request_blocks_if_unclaimed(self, req_id: str,
     return bool(
         self._impl.cancel_request_blocks_if_unclaimed(str(req_id), int(uuid)))
 
+  def get_request_block_status(
+      self, keys: typing.Sequence[tuple[str, int]]
+  ) -> list[int]:
+    """Read-only registry lifecycle probe: one REQUEST_BLOCK_STATUS_* value
+
+    per (req_id, uuid) key, in order.
+    """
+    return [
+        int(status)
+        for status in self._impl.get_request_block_status(
+            [(str(req_id), int(uuid)) for req_id, uuid in keys]
+        )
+    ]
+
   def coordinate_transfer(
       self,
       src_units: list,

tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc

@@ -1083,6 +1083,15 @@ NB_MODULE(_tpu_raiden_torch, m) {
              ::tpu_raiden::kv_cache::reshardpb2::ThrowIfError(result.status());
              return *result;
            })
+      .def("get_request_block_status",
+           [](::tpu_raiden::kv_cache::reshardpb2::ClientHandle& self,
+              const std::vector<std::pair<std::string, int64_t>>& keys) {
+             nb::gil_scoped_release release;
+             absl::StatusOr<std::vector<int32_t>> result =
+                 self.client->GetRequestBlockStatus(keys);
+             ::tpu_raiden::kv_cache::reshardpb2::ThrowIfError(result.status());
+             return *result;
+           })
       .def("start_transfer",
            [](::tpu_raiden::kv_cache::reshardpb2::ClientHandle& self,
               const std::vector<::tpu_raiden::kv_cache::reshardpb2::UnitTuple>&
@@ -1163,6 +1172,9 @@ NB_MODULE(_tpu_raiden_torch, m) {
   // Capability marker for the vLLM connector's version-skew probe: present
   // and true iff start_transfer accepts the dst_skip_bytes clip.
   m.attr("reshard_client_supports_dst_skip_bytes") = true;
+  // Present and true iff get_request_block_status (the read-only registry
+  // lifecycle probe) is bound.
+  m.attr("reshard_client_supports_request_block_status") = true;
 
   // NOTE: KVCacheStoreWrapper is already bound above as "KVCacheStore";
   // nanobind keys class bindings by C++ type, so a second nb::class_ for the

tpu_sync/kv_cache/reshard/request_block_registry.cc

@@ -404,6 +404,37 @@ absl::StatusOr<bool> RequestBlockRegistry::CancelIfUnclaimed(
   return true;
 }
 
+std::vector<RequestBlockRegistry::RequestBlockStatus>
+RequestBlockRegistry::Status(const std::vector<RequestBlockKey>& keys) {
+  const double now = clock_();
+  absl::MutexLock lock(*mu_);
+  PurgeExpiredRequestBlocksLocked(now);
+  PurgeExpiredLifecycleLocked(now);
+  std::vector<RequestBlockStatus> statuses;
+  statuses.reserve(keys.size());
+  for (const RequestBlockKey& key : keys) {
+    const LifecycleKey lifecycle_key{key.first, key.second};
+    if (cancelled_.find(lifecycle_key) != cancelled_.end()) {
+      statuses.push_back(RequestBlockStatus::kCancelled);
+      continue;
+    }
+    if (claimed_.find(lifecycle_key) != claimed_.end()) {
+      statuses.push_back(RequestBlockStatus::kClaimed);
+      continue;
+    }
+    bool registered = false;
+    for (const auto& [block_key, registration] : request_blocks_) {
+      if (block_key.first == key.first && registration.uuid == key.second) {
+        registered = true;
+        break;
+      }
+    }
+    statuses.push_back(registered ? RequestBlockStatus::kRegistered
+                                  : RequestBlockStatus::kUnknown);
+  }
+  return statuses;
+}
+
 absl::StatusOr<std::map<RaidenId, RequestBlockRegistration,
                         RequestBlockRegistry::RaidenIdLess>>
 RequestBlockRegistry::LookupAndClaim(const std::string& req_id, int64_t uuid,

tpu_sync/kv_cache/reshard/request_block_registry.h

@@ -79,6 +79,20 @@ class RequestBlockRegistry {
   absl::StatusOr<bool> CancelIfUnclaimed(const std::string& req_id,
                                          int64_t uuid);
 
+  enum class RequestBlockStatus {
+    kUnknown = 0,
+    kRegistered = 1,
+    kClaimed = 2,
+    kCancelled = 3,
+  };
+  using RequestBlockKey = std::pair<std::string, int64_t>;
+
+  // Read-only lifecycle probe (get_request_block_status): one status per
+  // key, in order. Expired state is purged first so a TTL-lapsed row reads
+  // as kUnknown. Precedence: cancelled tombstone > claim > registered row.
+  std::vector<RequestBlockStatus> Status(
+      const std::vector<RequestBlockKey>& keys);
+
   // _lookup_request_blocks: the claim linearization point. claim_owner is
   // an opaque identity pointer (Python: `object()`).
   absl::StatusOr<std::map<RaidenId, RequestBlockRegistration, RaidenIdLess>>

tpu_sync/kv_cache/reshard/reshard_client.cc

@@ -14,6 +14,7 @@
 
 #include "tpu_sync/kv_cache/reshard/reshard_client.h"
 
+#include <cstddef>
 #include <cstdint>
 #include <memory>
 #include <string>
@@ -248,6 +249,20 @@ tpu_sync::rpc::ControllerRequest ReshardClient::BuildGetTransferStatus(
   return req;
 }
 
+tpu_sync::rpc::ControllerRequest ReshardClient::BuildGetRequestBlockStatus(
+    const std::vector<std::pair<std::string, int64_t>>& keys) {
+  tpu_sync::rpc::ControllerRequest req;
+  req.set_command(
+      tpu_sync::rpc::ControllerRequest::COMMAND_GET_REQUEST_BLOCK_STATUS);
+  auto* status = req.mutable_get_request_block_status_request();
+  for (const auto& [req_id, uuid] : keys) {
+    auto* key = status->add_keys();
+    key->set_req_id(req_id);
+    key->set_uuid(uuid);
+  }
+  return req;
+}
+
 tpu_sync::rpc::ControlRequest ReshardClient::BuildGetMetadata() {
   tpu_sync::rpc::ControlRequest req;
   req.set_command(tpu_sync::rpc::ControlRequest::COMMAND_GET_METADATA);
@@ -348,6 +363,23 @@ absl::StatusOr<int32_t> ReshardClient::GetTransferStatus(
       response->get_transfer_status_response().status());
 }
 
+absl::StatusOr<std::vector<int32_t>> ReshardClient::GetRequestBlockStatus(
+    const std::vector<std::pair<std::string, int64_t>>& keys) {
+  absl::StatusOr<tpu_sync::rpc::ControllerResponse> response =
+      CallController(BuildGetRequestBlockStatus(keys));
+  if (!response.ok()) return response.status();
+  const auto& body = response->get_request_block_status_response();
+  if (static_cast<size_t>(body.statuses_size()) != keys.size()) {
+    return absl::InternalError(
+        absl::StrCat("Remote Controller Server returned ", body.statuses_size(),
+                     " request block statuses for ", keys.size(), " keys"));
+  }
+  std::vector<int32_t> statuses;
+  statuses.reserve(keys.size());
+  for (int status : body.statuses()) statuses.push_back(status);
+  return statuses;
+}
+
 absl::StatusOr<std::vector<std::string>> ReshardClient::GetMetadata() {
   absl::StatusOr<tpu_sync::rpc::ControlResponse> response =
       CallRaiden(BuildGetMetadata());

tpu_sync/kv_cache/reshard/reshard_client.h

@@ -31,6 +31,7 @@
 #include <memory>
 #include <optional>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "absl/status/status.h"
@@ -162,6 +163,8 @@ class ReshardClient {
       const StartTransferArgs& args);
   static tpu_sync::rpc::ControllerRequest BuildGetTransferStatus(
       const std::string& req_id, int64_t uuid);
+  static tpu_sync::rpc::ControllerRequest BuildGetRequestBlockStatus(
+      const std::vector<std::pair<std::string, int64_t>>& keys);
   static tpu_sync::rpc::ControlRequest BuildGetMetadata();
   static tpu_sync::rpc::ControlRequest BuildShutdown();
 
@@ -179,6 +182,9 @@ class ReshardClient {
   absl::StatusOr<bool> StartTransfer(const StartTransferArgs& args);
   absl::StatusOr<int32_t> GetTransferStatus(const std::string& req_id,
                                             int64_t uuid);
+  // GetRequestBlockStatusResponse::Status values, parallel to keys.
+  absl::StatusOr<std::vector<int32_t>> GetRequestBlockStatus(
+      const std::vector<std::pair<std::string, int64_t>>& keys);
   // Serialized RegisterWorkUnitRequest payloads; the Python shim parses
   // them back into pb2 objects so callers see the facade's return shape.
   absl::StatusOr<std::vector<std::string>> GetMetadata();

tpu_sync/kv_cache/reshard/reshard_service.cc

@@ -117,7 +117,10 @@ std::string ReshardService::HandleFrame(const std::string& request_bytes) {
        req.has_complete_request_blocks_request()) ||
       (req.command() == tpu_sync::rpc::ControllerRequest::
                             COMMAND_CANCEL_REQUEST_BLOCKS_IF_UNCLAIMED &&
-       req.has_cancel_request_blocks_if_unclaimed_request());
+       req.has_cancel_request_blocks_if_unclaimed_request()) ||
+      (req.command() ==
+           tpu_sync::rpc::ControllerRequest::COMMAND_GET_REQUEST_BLOCK_STATUS &&
+       req.has_get_request_block_status_request());
   if (is_controller_command) {
     return HandleControllerCommand(req);
   }
@@ -274,6 +277,35 @@ std::string ReshardService::HandleControllerCommand(
       resp.set_success(true);
       break;
     }
+    case tpu_sync::rpc::ControllerRequest::COMMAND_GET_REQUEST_BLOCK_STATUS: {
+      const auto& status_req = req.get_request_block_status_request();
+      std::vector<RequestBlockRegistry::RequestBlockKey> keys;
+      keys.reserve(status_req.keys_size());
+      for (const auto& key : status_req.keys()) {
+        keys.emplace_back(key.req_id(), key.uuid());
+      }
+      auto* out = resp.mutable_get_request_block_status_response();
+      for (RequestBlockRegistry::RequestBlockStatus status :
+           registry_->Status(keys)) {
+        using ProtoStatus = tpu_sync::rpc::GetRequestBlockStatusResponse;
+        switch (status) {
+          case RequestBlockRegistry::RequestBlockStatus::kRegistered:
+            out->add_statuses(ProtoStatus::STATUS_REGISTERED);
+            break;
+          case RequestBlockRegistry::RequestBlockStatus::kClaimed:
+            out->add_statuses(ProtoStatus::STATUS_CLAIMED);
+            break;
+          case RequestBlockRegistry::RequestBlockStatus::kCancelled:
+            out->add_statuses(ProtoStatus::STATUS_CANCELLED);
+            break;
+          case RequestBlockRegistry::RequestBlockStatus::kUnknown:
+            out->add_statuses(ProtoStatus::STATUS_UNKNOWN);
+            break;
+        }
+      }
+      resp.set_success(true);
+      break;
+    }
     default:
       resp.set_message("COMMAND_UNSPECIFIED");
       break;

tpu_sync/kv_cache/reshard/reshard_service_test.cc

@@ -620,6 +620,72 @@ TEST_F(ReshardStackTest, CancelTombstoneBlocksLateRegistration) {
                         "req_id=req-5, uuid=46"));
 }
 
+TEST_F(ReshardStackTest, RequestBlockStatusProbeTracksLifecycle) {
+  using ProtoStatus = tpu_sync::rpc::GetRequestBlockStatusResponse;
+  RegisterAllUnits(/*num_src=*/2, /*live=*/1024, /*stride=*/1024,
+                   /*num_blocks=*/16);
+  auto probe = [&](std::vector<std::pair<std::string, int64_t>> keys) {
+    tpu_sync::rpc::ControllerRequest req;
+    req.set_command(
+        tpu_sync::rpc::ControllerRequest::COMMAND_GET_REQUEST_BLOCK_STATUS);
+    for (const auto& [req_id, uuid] : keys) {
+      auto* key = req.mutable_get_request_block_status_request()->add_keys();
+      key->set_req_id(req_id);
+      key->set_uuid(uuid);
+    }
+    tpu_sync::rpc::ControllerResponse resp =
+        HandleController(req.SerializeAsString());
+    EXPECT_TRUE(resp.success()) << resp.message();
+    std::vector<int> statuses;
+    for (int status : resp.get_request_block_status_response().statuses()) {
+      statuses.push_back(status);
+    }
+    return statuses;
+  };
+  auto cancel = [&](const std::string& req_id, int64_t uuid) {
+    tpu_sync::rpc::ControllerRequest req;
+    req.set_command(tpu_sync::rpc::ControllerRequest::
+                        COMMAND_CANCEL_REQUEST_BLOCKS_IF_UNCLAIMED);
+    req.mutable_cancel_request_blocks_if_unclaimed_request()->set_req_id(
+        req_id);
+    req.mutable_cancel_request_blocks_if_unclaimed_request()->set_uuid(uuid);
+    return HandleController(req.SerializeAsString());
+  };
+
+  // Never registered.
+  EXPECT_EQ(probe({{"req-s", 60}}),
+            std::vector<int>{ProtoStatus::STATUS_UNKNOWN});
+
+  // Registered (one rank is enough) and unclaimed.
+  RegisterSpans(0, "req-s", 60, 1024, 3, 0, 0, 1024);
+  EXPECT_EQ(probe({{"req-s", 60}}),
+            std::vector<int>{ProtoStatus::STATUS_REGISTERED});
+
+  // The consumer's release-only cancel: the row is gone, the tombstone
+  // reports it — this is what the producer worker polls for.
+  tpu_sync::rpc::ControllerResponse cancel_resp = cancel("req-s", 60);
+  ASSERT_TRUE(cancel_resp.success());
+  EXPECT_EQ(cancel_resp.response_data(), "true");
+  EXPECT_EQ(probe({{"req-s", 60}}),
+            std::vector<int>{ProtoStatus::STATUS_CANCELLED});
+
+  // Claimed registrations refuse cancellation and read as claimed; the
+  // batch keeps key order.
+  RegisterSpans(0, "req-c", 61, 1024, 5, 0, 0, 1024);
+  RegisterSpans(1, "req-c", 61, 1024, 7, 1, 0, 512);
+  ASSERT_TRUE(Coordinate("req-c", 61, 2, {7, 9}).success());
+  EXPECT_EQ(probe({{"req-c", 61}, {"req-s", 60}, {"req-none", 1}}),
+            (std::vector<int>{ProtoStatus::STATUS_CLAIMED,
+                              ProtoStatus::STATUS_CANCELLED,
+                              ProtoStatus::STATUS_UNKNOWN}));
+  EXPECT_EQ(cancel("req-c", 61).response_data(), "false");
+
+  // Tombstones lapse with the TTL and read as unknown.
+  now_ += 601.0;
+  EXPECT_EQ(probe({{"req-s", 60}}),
+            std::vector<int>{ProtoStatus::STATUS_UNKNOWN});
+}
+
 TEST_F(ReshardStackTest, CompletionVotesRetireAfterClaim) {
   RegisterAllUnits(/*num_src=*/2, /*live=*/1024, /*stride=*/1024,
                    /*num_blocks=*/16);

tpu_sync/rpc/controller_service.proto

@@ -148,6 +148,36 @@ message GetTransferStatusRequest {
   int64 uuid = 2;
 }
 
+message RequestBlockKey {
+  string req_id = 1;
+  int64 uuid = 2;
+}
+
+// Read-only lifecycle probe over the request-block registry. Producer
+// workers poll it for registrations that have not reached a native
+// terminal, so a consumer-side cancel_request_blocks_if_unclaimed (which
+// only retires the registry row) becomes observable before the TTL sweep.
+message GetRequestBlockStatusRequest {
+  repeated RequestBlockKey keys = 1;
+}
+
+message GetRequestBlockStatusResponse {
+  enum Status {
+    STATUS_UNSPECIFIED = 0;
+    // No row, claim, or tombstone for the key (never registered, retired
+    // after completion, or purged by TTL).
+    STATUS_UNKNOWN = 1;
+    // At least one rank row is registered and unclaimed.
+    STATUS_REGISTERED = 2;
+    // Claimed by a transfer plan (in flight).
+    STATUS_CLAIMED = 3;
+    // Cancelled while unclaimed; tombstone still live.
+    STATUS_CANCELLED = 4;
+  }
+  // Parallel to GetRequestBlockStatusRequest.keys.
+  repeated Status statuses = 1;
+}
+
 message GetTransferStatusResponse {
   enum Status {
     STATUS_UNSPECIFIED = 0;
@@ -168,6 +198,7 @@ message ControllerRequest {
     COMMAND_RELEASE_REQUEST_BLOCKS = 4;
     COMMAND_COMPLETE_REQUEST_BLOCKS = 5;
     COMMAND_CANCEL_REQUEST_BLOCKS_IF_UNCLAIMED = 6;
+    COMMAND_GET_REQUEST_BLOCK_STATUS = 7;
   }
   Command command = 1;
 
@@ -185,6 +216,7 @@ message ControllerRequest {
   CompleteRequestBlocksRequest complete_request_blocks_request = 10;
   CancelRequestBlocksIfUnclaimedRequest
       cancel_request_blocks_if_unclaimed_request = 11;
+  GetRequestBlockStatusRequest get_request_block_status_request = 14;
 }
 
 message ControllerResponse {
@@ -193,4 +225,5 @@ message ControllerResponse {
   GetTransferStatusResponse get_transfer_status_response = 3;
   // Command-specific scalar result (e.g. cancellation verdict, retire count).
   string response_data = 4;
+  GetRequestBlockStatusResponse get_request_block_status_response = 5;
 }