Skip to content

Refactor kv_cache to conform to TPU Raiden status handling conventions. #869

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
Sep 7, 2026
Merged

Refactor kv_cache to conform to TPU Raiden status handling conventions. #869

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion .bazelrc
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,13 @@
# Allow empty globs across all Bazel commands (required for external deps e.g. Eigen / XLA under Bazel 8).
common --incompatible_disallow_empty_glob=false

build:oss --experimental_repo_remote_exec
build:oss --nocheck_visibility
build:oss --copt=-DMLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir.
build:oss --override_module=torch_tpu=%workspace%/third_party/torch_tpu_stub
build:oss --action_env=CC=clang-18
build:oss --action_env=CXX=clang++-18
build:oss --repo_env=CC=clang-18
build:oss --incompatible_disallow_empty_glob=false

build:ci --google_default_credentials=true
build:ci --remote_cache=grpcs://remotebuildexecution.googleapis.com
Expand Down
19 changes: 19 additions & 0 deletions tpu_sync/kv_cache/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,9 +70,11 @@ cc_test(
deps = [
":logical_block_manager",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/status:statusor",
"@com_google_googletest//:gtest",
"@com_google_googletest//:gtest_main",
"@xla//xla/tsl/platform:statusor",
],
)

Expand All @@ -83,6 +85,7 @@ cc_library(
visibility = ["//visibility:public"],
deps = [
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_macros",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:string_view",
Expand All @@ -96,10 +99,12 @@ cc_test(
deps = [
":kv_cache_metadata",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/types:span",
"@com_google_googletest//:gtest",
"@com_google_googletest//:gtest_main",
"@xla//xla/tsl/platform:statusor",
],
)

Expand Down Expand Up @@ -155,10 +160,12 @@ cc_test(
":kv_cache_store_wrapper",
"//tpu_sync/common:raiden_id",
"//tpu_sync/kv_cache/global_registry:test_util",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/time",
"@com_google_googletest//:gtest",
"@com_google_googletest//:gtest_main",
"@xla//xla/tsl/platform:statusor",
],
)

Expand All @@ -168,10 +175,12 @@ cc_test(
deps = [
":kv_cache_metadata",
":kv_cache_metadata_shm",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:string_view",
"@com_google_googletest//:gtest",
"@com_google_googletest//:gtest_main",
"@xla//xla/tsl/platform:statusor",
],
)

Expand All @@ -187,6 +196,7 @@ cc_library(
deps = [
"//tpu_sync/rpc:raiden_service_cc_proto",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_macros",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:span",
Expand All @@ -207,9 +217,11 @@ cc_test(
"//tpu_sync/rpc:raiden_service_cc_proto",
"//tpu_sync/transport:block_transport_delegate",
"@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/platform:statusor",
],
)

Expand All @@ -229,6 +241,7 @@ cc_test(
"@com_google_absl//absl/status:statusor",
"@com_google_googletest//:gtest",
"@com_google_googletest//:gtest_main",
"@xla//xla/tsl/platform:statusor",
],
)

Expand Down Expand Up @@ -427,10 +440,12 @@ cc_test(
"//tpu_sync/kv_cache/global_registry:global_registry_server_lib",
"//tpu_sync/kv_cache/global_registry:test_util",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/time",
"@com_google_googletest//:gtest",
"@com_google_googletest//:gtest_main",
"@xla//xla/tsl/platform:statusor",
],
)

Expand Down Expand Up @@ -575,6 +590,7 @@ cc_test(
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_macros",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
Expand Down Expand Up @@ -743,6 +759,7 @@ cc_test(
"@com_google_googletest//:gtest",
"@com_google_googletest//:gtest_main",
"@xla//xla/tsl/concurrency:future",
"@xla//xla/tsl/platform:statusor",
],
)

Expand Down Expand Up @@ -771,13 +788,15 @@ cc_test(
"//tpu_sync/kv_cache/global_registry:test_util",
"//tpu_sync/proto:kv_cache_store_service_cc_grpc",
"@com_github_grpc_grpc//:grpc++",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/synchronization",
"@com_google_absl//absl/time",
"@com_google_absl//absl/types:span",
"@com_google_googletest//:gtest",
"@com_google_googletest//:gtest_main",
"@xla//xla/tsl/concurrency:future",
"@xla//xla/tsl/platform:statusor",
],
)

Expand Down
77 changes: 36 additions & 41 deletions tpu_sync/kv_cache/host_offload_backend.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -217,13 +217,13 @@ 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, local_id);
if (global_res_or.ok()) {
remote_hits = std::move(global_res_or).value();
} else {
LOG(WARNING) << "Global registry lookup failed: "
<< global_res_or.status().message();
}
auto lookup_remote = [&]() -> absl::Status {
ABSL_ASSIGN_OR_RETURN(
remote_hits, client->Lookup(missing_hashes, local_id),
_.LogWarning() << "Global registry lookup failed: ");
return absl::OkStatus();
};
(void)lookup_remote();
}

// Phases 3 and 4, under one lock: walk the request in order taking each hash
Expand Down Expand Up @@ -580,11 +580,7 @@ absl::StatusOr<size_t> HostOffloadBackend::RecoverFromLocalManifest() {
for (const KVCacheMetadata::Entry* entry : recoverable) {
block_ids.push_back(entry->block_id);
}
absl::Status allocate_status =
raiden_controller_->AllocateTargetBlockIds(block_ids);
if (!allocate_status.ok()) {
return allocate_status;
}
ABSL_RETURN_IF_ERROR(raiden_controller_->AllocateTargetBlockIds(block_ids));

uint64_t max_seq = 0;
for (const KVCacheMetadata::Entry* entry : recoverable) {
Expand Down Expand Up @@ -799,29 +795,29 @@ HostOffloadBackend::BeginWriteRemote(
BuildLocalWorkerEndpoints(raiden_controller_),
absl::ToInt64Milliseconds(requested_deadline),
hold_window, std::move(on_verdict));
auto response_or = call.ack.Await();
if (!response_or.ok()) {
auto response = call.ack.Await();
if (!response.ok()) {
// On a transport error the peer may have restarted on a new port; drop
// the cached client so the next attempt re-resolves it. An application
// error (e.g. RESOURCE_EXHAUSTED) means the peer is alive, so the
// channel is kept.
if (IsTransportError(response_or.status())) {
if (IsTransportError(response.status())) {
InvalidateStoreClient(dst_raiden_id);
}
return response_or.status();
return response.status();
}

RemoteWriteAck ack;
ack.cancel = std::move(call.cancel);
ack.operation_id = response_or->operation_id();
ack.granted_deadline = absl::Milliseconds(response_or->granted_deadline_ms());
switch (response_or->exist_state()) {
ack.operation_id = response->operation_id();
ack.granted_deadline = absl::Milliseconds(response->granted_deadline_ms());
switch (response->exist_state()) {
case ::tpu_raiden::kv_cache::proto::WRITE_ALL_EXIST:
ack.all_exist = true;
return ack;
case ::tpu_raiden::kv_cache::proto::WRITE_PARTIAL_EXIST:
ack.existing_hashes.assign(response_or->existing_hashes().begin(),
response_or->existing_hashes().end());
ack.existing_hashes.assign(response->existing_hashes().begin(),
response->existing_hashes().end());
return ack;
default:
break;
Expand All @@ -839,11 +835,11 @@ tsl::Future<proto::PollWriteRemoteResponse>
HostOffloadBackend::PollWriteRemoteAsync(const RaidenId& dst_raiden_id,
uint64_t operation_id,
int64_t wait_ms) {
auto client_or = GetKVCacheStoreClient(dst_raiden_id);
if (!client_or.ok()) {
return tsl::Future<proto::PollWriteRemoteResponse>(client_or.status());
auto client = GetKVCacheStoreClient(dst_raiden_id);
if (!client.ok()) {
return tsl::Future<proto::PollWriteRemoteResponse>(client.status());
}
return (*client_or)->PollWriteRemote(operation_id, wait_ms);
return (*client)->PollWriteRemote(operation_id, wait_ms);
}

std::vector<std::string> HostOffloadBackend::AlreadyPresentHostResident(
Expand Down Expand Up @@ -991,46 +987,45 @@ tsl::Future<> HostOffloadBackend::Load(
tsl::Future<> HostOffloadBackend::LoadRemoteBlocks(
const RaidenId& remote_id, absl::Span<const std::string> block_hashes,
absl::Span<const int32_t> device_block_ids) {
auto client_or = GetKVCacheStoreClient(remote_id);
if (!client_or.ok()) {
return tsl::Future<>(client_or.status());
auto client = GetKVCacheStoreClient(remote_id);
if (!client.ok()) {
return tsl::Future<>(client.status());
}
std::shared_ptr<KVCacheStoreClient> client = std::move(client_or.value());
std::shared_ptr<KVCacheStoreClient> client_ptr = *std::move(client);

auto host_blocks_or =
raiden_controller_->AllocateBlockIds(block_hashes.size());
if (!host_blocks_or.ok()) {
return tsl::Future<>(host_blocks_or.status());
auto host_blocks = raiden_controller_->AllocateBlockIds(block_hashes.size());
if (!host_blocks.ok()) {
return tsl::Future<>(host_blocks.status());
}
std::vector<int32_t> dst_host_block_ids(host_blocks_or.value().begin(),
host_blocks_or.value().end());
std::vector<int32_t> dst_host_block_ids(host_blocks->begin(),
host_blocks->end());

auto [load_promise, load_future] = tsl::MakePromise<>();

::tpu_sync::rpc::RaidenIdProto client_raiden_id = raiden_controller_->unit();
std::vector<::tpu_sync::proto::RaidenWorkerEndpointsProto>
client_worker_endpoints = BuildLocalWorkerEndpoints(raiden_controller_);
tsl::Future<::tpu_raiden::kv_cache::proto::FetchResponse> fetch_future =
client->Fetch(block_hashes, device_block_ids, dst_host_block_ids,
client_raiden_id, client_worker_endpoints);
client_ptr->Fetch(block_hashes, device_block_ids, dst_host_block_ids,
client_raiden_id, client_worker_endpoints);

fetch_future.OnReady(
[this, remote_id, dst_host_block_ids,
dev_ids_vec = std::vector<int32_t>(device_block_ids.begin(),
device_block_ids.end()),
load_promise = std::move(load_promise)](
const absl::StatusOr<::tpu_raiden::kv_cache::proto::FetchResponse>&
response_or) mutable {
if (!response_or.ok()) {
fetch_response) mutable {
if (!fetch_response.ok()) {
(void)raiden_controller_->DeallocateBlockIds(dst_host_block_ids);
// The peer may have restarted on a new port; drop the cached client
// so the next attempt re-resolves instead of redialling a dead one.
InvalidateStoreClient(remote_id);
load_promise.Set(response_or.status());
load_promise.Set(fetch_response.status());
return;
}

const auto& response = response_or.value();
const auto& response = *fetch_response;
if (!response.failed_block_hashes().empty()) {
(void)raiden_controller_->DeallocateBlockIds(dst_host_block_ids);
std::string err_msg = response.error_message().empty()
Expand Down
Loading
Loading