Skip to content

Commit 96dbbd5

Browse files
itaruncopybara-github
authored andcommitted
Implements the kv_save_blocks_total counter metric (tpu_raiden_kv_save_blocks_total{dp_rank, tier, outcome}) to track block save outcomes across KV cache offload tiers:
- Declares `kKvSaveBlocksTotal` (`MetricType::kCounter`) along with `outcome` label values `kOutcomeRejected` (`"rejected"`) and `kOutcomeFailed` (`"failed"`) in `metrics_backend.h`. - Instruments `KVCacheStore::Insert` (`tier="local_dram"`, `outcome="rejected"` for the whole batch when `InsertOption::kInsertAndPin` fails, either for lack of unpinned space or because an existing hash could not be pinned), `KVCacheStore::SaveLocal` (`tier="local_dram"`; host block allocation failure counts `outcome="rejected"` on `RESOURCE_EXHAUSTED` and `"failed"` on any other error; storage write-through counts `"ok"` / `"failed"` when the offload settles), `HostOffloadBackend::Save` (`tier="local_dram"`, `outcome="ok"` / `"failed"`), and `HostOffloadBackend::BeginWriteRemote` (`tier="remote_dram"`, `outcome="ok"`, `"rejected"`, and `"failed"` across peer RPC and transfer settle paths). - Adds unit tests verifying `kv_save_blocks_total` across `local_dram` (including storage write-through) and `remote_dram` tiers. PiperOrigin-RevId: 999559866
1 parent b001246 commit 96dbbd5

8 files changed

Lines changed: 407 additions & 15 deletions

File tree

‎tpu_sync/kv_cache/BUILD‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -998,6 +998,7 @@ cc_test(
998998
":kv_cache_store_backend_factory",
999999
":kv_cache_store_client",
10001000
":kv_cache_store_server",
1001+
":kv_cache_telemetry",
10011002
"//tpu_sync/common:raiden_id",
10021003
"//tpu_sync/core:kv_manager_holder",
10031004
"//tpu_sync/core:raiden_transfer_endpoint",

‎tpu_sync/kv_cache/host_offload_backend.cc‎

Lines changed: 65 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,25 @@ void MarkRecallFailed(BlockTracker* absl_nonnull load_tracker,
112112
load_tracker->MarkFailed(block_hashes);
113113
}
114114

115+
// Settle a save batch: count `block_hashes` in `kv_save_blocks_total`
116+
// under `tier` (`ok` for Done, `failed` for Failed), then mark them in
117+
// `save_tracker`.
118+
void MarkSaveDone(BlockTracker* absl_nonnull save_tracker,
119+
absl::string_view dp_rank, absl::string_view tier,
120+
absl::Span<const std::string> block_hashes) {
121+
IncrementKvSaveBlocks(dp_rank, tier, telemetry::metric_labels::kOutcomeOk,
122+
block_hashes.size());
123+
save_tracker->MarkDone(block_hashes);
124+
}
125+
126+
void MarkSaveFailed(BlockTracker* absl_nonnull save_tracker,
127+
absl::string_view dp_rank, absl::string_view tier,
128+
absl::Span<const std::string> block_hashes) {
129+
IncrementKvSaveBlocks(dp_rank, tier, telemetry::metric_labels::kOutcomeFailed,
130+
block_hashes.size());
131+
save_tracker->MarkFailed(block_hashes);
132+
}
133+
115134
} // namespace
116135

117136
HostOffloadBackend::HostOffloadBackend(
@@ -793,11 +812,25 @@ HostOffloadBackend::BeginWriteRemote(
793812
return absl::InvalidArgumentError(
794813
"src_host_block_ids must have one entry per block hash");
795814
}
815+
const auto fail_save = [&](absl::string_view outcome, absl::Status status) {
816+
IncrementKvSaveBlocks(dp_rank_str_,
817+
telemetry::metric_labels::kTierRemoteDram, outcome,
818+
block_hashes.size());
819+
return status;
820+
};
821+
if (raiden_controller_ == nullptr) {
822+
return fail_save(
823+
telemetry::metric_labels::kOutcomeFailed,
824+
absl::FailedPreconditionError("RaidenController is not initialized"));
825+
}
796826

797827
// Resolves the peer through the global registry. A missing registry
798828
// client fails here; there is no separate precondition check.
799-
ABSL_ASSIGN_OR_RETURN(std::shared_ptr<KVCacheStoreClient> client,
800-
GetKVCacheStoreClient(dst_raiden_id));
829+
absl::StatusOr<std::shared_ptr<KVCacheStoreClient>> client =
830+
GetKVCacheStoreClient(dst_raiden_id);
831+
if (!client.ok()) {
832+
return fail_save(telemetry::metric_labels::kOutcomeFailed, client.status());
833+
}
801834

802835
auto hold_expiry = std::make_shared<absl::Time>(absl::Now() + hold_window);
803836

@@ -806,7 +839,7 @@ HostOffloadBackend::BeginWriteRemote(
806839
// below are dropped, and the sweep's eviction skips pinned blocks -- it
807840
// would free nothing and give up the episode. Every verdict is recorded
808841
// here instead, after the release.
809-
auto call = client->WriteRemote(
842+
auto call = (*client)->WriteRemote(
810843
raiden_controller_->unit(), block_hashes, src_host_block_ids,
811844
BuildLocalWorkerEndpoints(raiden_controller_),
812845
absl::ToInt64Milliseconds(requested_deadline), hold_window,
@@ -820,7 +853,10 @@ HostOffloadBackend::BeginWriteRemote(
820853
if (IsTransportError(response.status())) {
821854
InvalidateStoreClient(dst_raiden_id);
822855
}
823-
return response.status();
856+
return fail_save(absl::IsResourceExhausted(response.status())
857+
? telemetry::metric_labels::kOutcomeRejected
858+
: telemetry::metric_labels::kOutcomeFailed,
859+
response.status());
824860
}
825861

826862
RemoteWriteAck ack;
@@ -837,14 +873,18 @@ HostOffloadBackend::BeginWriteRemote(
837873
// by the verdict must find the blocks already unpinned.
838874
Release(block_hashes);
839875
Release(block_hashes);
840-
save_tracker->MarkDone(block_hashes);
876+
MarkSaveDone(save_tracker, dp_rank_str_,
877+
telemetry::metric_labels::kTierRemoteDram, block_hashes);
841878
return ack;
842879
case ::tpu_raiden::kv_cache::proto::WRITE_PARTIAL_EXIST:
843880
ack.existing_hashes.assign(response->existing_hashes().begin(),
844881
response->existing_hashes().end());
845882
// Only the transfer's hold; a failed offer keeps the caller's pin for
846883
// the caller to drop.
847884
Release(block_hashes);
885+
IncrementKvSaveBlocks(
886+
dp_rank_str_, telemetry::metric_labels::kTierRemoteDram,
887+
telemetry::metric_labels::kOutcomeFailed, block_hashes.size());
848888
save_tracker->MarkFailedWithExisting(block_hashes, ack.existing_hashes);
849889
return ack;
850890
default:
@@ -853,8 +893,10 @@ HostOffloadBackend::BeginWriteRemote(
853893
if (ack.operation_id == 0) {
854894
// Neither an existence answer nor an accepted operation: a protocol
855895
// error, not a silent no-op.
856-
return absl::InternalError(
857-
"Destination accepted the offer but returned no operation id.");
896+
return fail_save(
897+
telemetry::metric_labels::kOutcomeFailed,
898+
absl::InternalError(
899+
"Destination accepted the offer but returned no operation id."));
858900
}
859901
std::vector<std::string> hashes(block_hashes.begin(), block_hashes.end());
860902
call.result.OnReady([this, lifetime = lifetime_, dst_raiden_id, save_tracker,
@@ -887,6 +929,11 @@ HostOffloadBackend::BeginWriteRemote(
887929
// A successful save also consumes the caller's pin.
888930
Release(hashes);
889931
}
932+
IncrementKvSaveBlocks(
933+
dp_rank_str_, telemetry::metric_labels::kTierRemoteDram,
934+
succeeded ? telemetry::metric_labels::kOutcomeOk
935+
: telemetry::metric_labels::kOutcomeFailed,
936+
hashes.size());
890937
switch (result.state()) {
891938
case proto::PollWriteRemoteResponse::COMMITTED:
892939
case proto::PollWriteRemoteResponse::ALL_EXIST:
@@ -915,7 +962,8 @@ HostOffloadBackend::BeginWriteRemote(
915962
if (op_id == 0 || remaining_hold <= absl::ZeroDuration()) {
916963
// No operation to chase, or the destination's hold has lapsed.
917964
Release(hashes);
918-
save_tracker->MarkFailed(hashes);
965+
MarkSaveFailed(save_tracker, dp_rank_str_,
966+
telemetry::metric_labels::kTierRemoteDram, hashes);
919967
return;
920968
}
921969
// Started while holding the lock, so `this` cannot be destroyed under us.
@@ -938,6 +986,11 @@ HostOffloadBackend::BeginWriteRemote(
938986
if (poll_succeeded) {
939987
Release(hashes);
940988
}
989+
IncrementKvSaveBlocks(
990+
dp_rank_str_, telemetry::metric_labels::kTierRemoteDram,
991+
poll_succeeded ? telemetry::metric_labels::kOutcomeOk
992+
: telemetry::metric_labels::kOutcomeFailed,
993+
hashes.size());
941994
if (!resp.ok()) {
942995
save_tracker->MarkFailed(hashes);
943996
return;
@@ -1403,13 +1456,15 @@ tsl::Future<> HostOffloadBackend::Save(
14031456
if (!update_hashes.empty()) {
14041457
Insert(update_hashes, update_slices);
14051458
}
1406-
save_tracker->MarkDone(hashes);
1459+
MarkSaveDone(save_tracker, dp_rank_str_,
1460+
telemetry::metric_labels::kTierLocalDram, hashes);
14071461
save_promise.Set(absl::OkStatus());
14081462
} else {
14091463
if (raiden_controller_ != nullptr) {
14101464
(void)raiden_controller_->DeallocateBlockIds(dst_host_ids);
14111465
}
1412-
save_tracker->MarkFailed(hashes);
1466+
MarkSaveFailed(save_tracker, dp_rank_str_,
1467+
telemetry::metric_labels::kTierLocalDram, hashes);
14131468
save_promise.Set(status);
14141469
}
14151470
});

‎tpu_sync/kv_cache/kv_cache_store.cc‎

Lines changed: 36 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@
5151
#include "tpu_sync/core/controller/raiden_controller.h"
5252
#include "tpu_sync/core/host_memory_allocator.h"
5353
#include "tpu_sync/kv_cache/backends/backend.h"
54+
#include "tpu_sync/kv_cache/block_tracker.h"
5455
#include "tpu_sync/kv_cache/global_registry/global_registry_client.h"
5556
#include "tpu_sync/kv_cache/kv_cache_metadata.h"
5657
#include "tpu_sync/kv_cache/kv_cache_metadata_shm.h"
@@ -213,6 +214,25 @@ absl::StatusOr<BackendConfig> ResolveSecondaryBackendTopology(
213214
return sec_cfg;
214215
}
215216

217+
// Settle a save batch: count `block_hashes` in `kv_save_blocks_total`
218+
// under `tier` (`ok` for Done, `failed` for Failed), then mark them in
219+
// `save_tracker`.
220+
void MarkSaveDone(BlockTracker& save_tracker, absl::string_view dp_rank,
221+
absl::string_view tier,
222+
absl::Span<const std::string> block_hashes) {
223+
IncrementKvSaveBlocks(dp_rank, tier, telemetry::metric_labels::kOutcomeOk,
224+
block_hashes.size());
225+
save_tracker.MarkDone(block_hashes);
226+
}
227+
228+
void MarkSaveFailed(BlockTracker& save_tracker, absl::string_view dp_rank,
229+
absl::string_view tier,
230+
absl::Span<const std::string> block_hashes) {
231+
IncrementKvSaveBlocks(dp_rank, tier, telemetry::metric_labels::kOutcomeFailed,
232+
block_hashes.size());
233+
save_tracker.MarkFailed(block_hashes);
234+
}
235+
216236
} // namespace
217237

218238
// --- KVCacheStore Implementation ---
@@ -1068,6 +1088,9 @@ absl::Status KVCacheStore::Insert(const std::vector<std::string>& block_hashes,
10681088

10691089
if (!backends_[0]->Insert(block_hashes, slices,
10701090
InsertOption::kInsertAndPin)) {
1091+
IncrementKvSaveBlocks(
1092+
dp_rank_str_, telemetry::metric_labels::kTierLocalDram,
1093+
telemetry::metric_labels::kOutcomeRejected, block_hashes.size());
10711094
return absl::ResourceExhaustedError(
10721095
"Insert refused: not enough unpinned space, or a hash could not be "
10731096
"pinned");
@@ -1151,8 +1174,15 @@ absl::Status KVCacheStore::SaveLocal(
11511174
save_tracker_.AddPending(block_hashes);
11521175
}
11531176

1154-
auto host_blocks = AllocateBlockIds(block_hashes.size());
1177+
absl::StatusOr<std::vector<int>> host_blocks =
1178+
AllocateBlockIds(block_hashes.size());
11551179
if (!host_blocks.ok()) {
1180+
IncrementKvSaveBlocks(dp_rank_str_,
1181+
telemetry::metric_labels::kTierLocalDram,
1182+
absl::IsResourceExhausted(host_blocks.status())
1183+
? telemetry::metric_labels::kOutcomeRejected
1184+
: telemetry::metric_labels::kOutcomeFailed,
1185+
block_hashes.size());
11561186
save_tracker_.RemovePending(block_hashes);
11571187
return host_blocks.status();
11581188
}
@@ -1193,7 +1223,9 @@ absl::Status KVCacheStore::SaveLocal(
11931223
absl::MutexLock lock(lifetime->mu);
11941224
if (KVCacheStore* self = lifetime->store; self != nullptr) {
11951225
self->DeallocateBlockIds(host_block_ids);
1196-
self->save_tracker_.MarkFailed(block_hashes);
1226+
MarkSaveFailed(self->save_tracker_, self->dp_rank_str_,
1227+
telemetry::metric_labels::kTierLocalDram,
1228+
block_hashes);
11971229
}
11981230
}
11991231
if (host_ram_backend != nullptr)
@@ -1216,7 +1248,8 @@ absl::Status KVCacheStore::SaveLocal(
12161248
{
12171249
absl::MutexLock lock(lifetime->mu);
12181250
if (KVCacheStore* self = lifetime->store; self != nullptr) {
1219-
self->save_tracker_.MarkDone(block_hashes);
1251+
MarkSaveDone(self->save_tracker_, self->dp_rank_str_,
1252+
telemetry::metric_labels::kTierLocalDram, block_hashes);
12201253
}
12211254
}
12221255
std::vector<global_registry::Registration> write_through_regs;

0 commit comments

Comments
 (0)