diff --git a/tpu_sync/api/torch/reshard_client.py b/tpu_sync/api/torch/reshard_client.py index 9bf5274a..b8c1e2bb 100644 --- a/tpu_sync/api/torch/reshard_client.py +++ b/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, diff --git a/tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc b/tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc index e11707f2..151e9f6f 100644 --- a/tpu_sync/frameworks/torch/tpu_raiden_torch_module.cc +++ b/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>& keys) { + nb::gil_scoped_release release; + absl::StatusOr> 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 diff --git a/tpu_sync/kv_cache/reshard/request_block_registry.cc b/tpu_sync/kv_cache/reshard/request_block_registry.cc index e900f48f..c6d31934 100644 --- a/tpu_sync/kv_cache/reshard/request_block_registry.cc +++ b/tpu_sync/kv_cache/reshard/request_block_registry.cc @@ -404,6 +404,37 @@ absl::StatusOr RequestBlockRegistry::CancelIfUnclaimed( return true; } +std::vector +RequestBlockRegistry::Status(const std::vector& keys) { + const double now = clock_(); + absl::MutexLock lock(*mu_); + PurgeExpiredRequestBlocksLocked(now); + PurgeExpiredLifecycleLocked(now); + std::vector 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> RequestBlockRegistry::LookupAndClaim(const std::string& req_id, int64_t uuid, diff --git a/tpu_sync/kv_cache/reshard/request_block_registry.h b/tpu_sync/kv_cache/reshard/request_block_registry.h index 966c2014..e27ed88b 100644 --- a/tpu_sync/kv_cache/reshard/request_block_registry.h +++ b/tpu_sync/kv_cache/reshard/request_block_registry.h @@ -79,6 +79,20 @@ class RequestBlockRegistry { absl::StatusOr CancelIfUnclaimed(const std::string& req_id, int64_t uuid); + enum class RequestBlockStatus { + kUnknown = 0, + kRegistered = 1, + kClaimed = 2, + kCancelled = 3, + }; + using RequestBlockKey = std::pair; + + // 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 Status( + const std::vector& keys); + // _lookup_request_blocks: the claim linearization point. claim_owner is // an opaque identity pointer (Python: `object()`). absl::StatusOr> diff --git a/tpu_sync/kv_cache/reshard/reshard_client.cc b/tpu_sync/kv_cache/reshard/reshard_client.cc index a5079182..18053d95 100644 --- a/tpu_sync/kv_cache/reshard/reshard_client.cc +++ b/tpu_sync/kv_cache/reshard/reshard_client.cc @@ -14,6 +14,7 @@ #include "tpu_sync/kv_cache/reshard/reshard_client.h" +#include #include #include #include @@ -248,6 +249,20 @@ tpu_sync::rpc::ControllerRequest ReshardClient::BuildGetTransferStatus( return req; } +tpu_sync::rpc::ControllerRequest ReshardClient::BuildGetRequestBlockStatus( + const std::vector>& 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 ReshardClient::GetTransferStatus( response->get_transfer_status_response().status()); } +absl::StatusOr> ReshardClient::GetRequestBlockStatus( + const std::vector>& keys) { + absl::StatusOr response = + CallController(BuildGetRequestBlockStatus(keys)); + if (!response.ok()) return response.status(); + const auto& body = response->get_request_block_status_response(); + if (static_cast(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 statuses; + statuses.reserve(keys.size()); + for (int status : body.statuses()) statuses.push_back(status); + return statuses; +} + absl::StatusOr> ReshardClient::GetMetadata() { absl::StatusOr response = CallRaiden(BuildGetMetadata()); diff --git a/tpu_sync/kv_cache/reshard/reshard_client.h b/tpu_sync/kv_cache/reshard/reshard_client.h index 55dcc5e4..485b2047 100644 --- a/tpu_sync/kv_cache/reshard/reshard_client.h +++ b/tpu_sync/kv_cache/reshard/reshard_client.h @@ -31,6 +31,7 @@ #include #include #include +#include #include #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>& keys); static tpu_sync::rpc::ControlRequest BuildGetMetadata(); static tpu_sync::rpc::ControlRequest BuildShutdown(); @@ -179,6 +182,9 @@ class ReshardClient { absl::StatusOr StartTransfer(const StartTransferArgs& args); absl::StatusOr GetTransferStatus(const std::string& req_id, int64_t uuid); + // GetRequestBlockStatusResponse::Status values, parallel to keys. + absl::StatusOr> GetRequestBlockStatus( + const std::vector>& keys); // Serialized RegisterWorkUnitRequest payloads; the Python shim parses // them back into pb2 objects so callers see the facade's return shape. absl::StatusOr> GetMetadata(); diff --git a/tpu_sync/kv_cache/reshard/reshard_service.cc b/tpu_sync/kv_cache/reshard/reshard_service.cc index 468928e5..f5a902bb 100644 --- a/tpu_sync/kv_cache/reshard/reshard_service.cc +++ b/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 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; diff --git a/tpu_sync/kv_cache/reshard/reshard_service_test.cc b/tpu_sync/kv_cache/reshard/reshard_service_test.cc index a18d4a94..2f180aa5 100644 --- a/tpu_sync/kv_cache/reshard/reshard_service_test.cc +++ b/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> 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 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{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{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{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{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{ProtoStatus::STATUS_UNKNOWN}); +} + TEST_F(ReshardStackTest, CompletionVotesRetireAfterClaim) { RegisterAllUnits(/*num_src=*/2, /*live=*/1024, /*stride=*/1024, /*num_blocks=*/16); diff --git a/tpu_sync/rpc/controller_service.proto b/tpu_sync/rpc/controller_service.proto index 7672c764..de5f3e11 100644 --- a/tpu_sync/rpc/controller_service.proto +++ b/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; }