Add tpu_raiden_storage_metadata_check_total counter for POSIX secondary storage lookups

Register and instrument `storage_metadata_check_total` (exported as `tpu_raiden_storage_metadata_check_total`) to count secondary storage metadata existence checks in `PosixKVCacheStoreBackend::Lookup`:

- **Metric Metadata (`metrics_backend.h`)**: Register `kStorageMetadataCheckTotal` (`MetricType::kCounter`) with labels `{dp_rank, result}` where `result` is one of `cache_hit` (`kResultCacheHit`), `stat_hit` (`kResultStatHit`), or `stat_miss` (`kResultStatMiss`). Host/pod IP is not duplicated in metric labels because Prometheus/GMP target scraping attaches `instance`/`pod` automatically, and `dp_rank` is populated from `RaidenId::data_replica_idx`.
- **POSIX Lookup Instrumentation (`posix_backend.{h,cc}`, `tds_backend.cc`, `kv_cache_store.cc`)**:
  - Propagate `effective_raiden_id` to secondary backend configs in `KVCacheStore::Create` when `effective_config.raiden_id.empty()`, and pass `absl::StrCat(config.raiden_id.data_replica_idx)` to `PosixKVCacheStoreBackend` in the `"posix"` and `"tds"` backend factory registrations (defaulting to `"0"`).
  - Batch shard-level metadata check counts per `PosixKVCacheStoreBackend::Lookup()` call across Phase 1 (`CachedFresh` positive-only TTL cache hits -> `cache_hit`) and Phase 2 (`ProbeExists` filesystem `stat()` hits/misses -> `stat_hit` / `stat_miss`), emitting at most three `IncrementCounter` calls per `Lookup()` when `metric_store.HasBackends()` is true and `count > 0` outside `cache_mu_`.
- **Unit Tests**: Update lockstep registration tests in `metrics_api_test.cc` and `telemetry_binding_test.py`, and add unit tests in `posix_backend_test.cc` and `kv_cache_store_test.cc` covering disabled-cache `stat_hit`/`stat_miss`, enabled-cache `cache_hit` + TTL expiry + `Delete()` invalidation, and custom `dp_rank` propagation.

PiperOrigin-RevId: 1013600513

Author: Tarun-Arora

tpu_sync/kv_cache/BUILD

@@ -1123,6 +1123,7 @@ cc_test(
         "//tpu_sync/kv_cache/global_registry:global_registry_server_lib",
         "//tpu_sync/kv_cache/global_registry:test_util",
         "//tpu_sync/proto:kv_cache_store_service_cc_grpc",
+        "//tpu_sync/telemetry:metrics_api",
         "//tpu_sync/telemetry:metrics_backend",
         "//tpu_sync/telemetry:mock_metrics_backend",
         "@com_github_grpc_grpc//:grpc++",

tpu_sync/kv_cache/backends/storage/BUILD

@@ -37,6 +37,8 @@ cc_library(
         "//tpu_sync/kv_cache:kv_cache_store_backend_factory",
         "//tpu_sync/kv_cache:lru_cache",
         "//tpu_sync/kv_cache/backends:backend",
+        "//tpu_sync/telemetry:metrics_api",
+        "//tpu_sync/telemetry:metrics_backend",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/base:nullability",
@@ -76,6 +78,9 @@ cc_test(
         "//tpu_sync/kv_cache:kv_cache_store_backend_factory",
         "//tpu_sync/kv_cache/backends:backend",
         "//tpu_sync/proto:worker_service_cc_proto",
+        "//tpu_sync/telemetry:metrics_api",
+        "//tpu_sync/telemetry:metrics_backend",
+        "//tpu_sync/telemetry:mock_metrics_backend",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/status",

tpu_sync/kv_cache/backends/storage/posix_backend.cc

@@ -60,6 +60,8 @@
 #include "tpu_sync/kv_cache/backends/backend.h"
 #include "tpu_sync/kv_cache/kv_cache_store_backend.h"
 #include "tpu_sync/kv_cache/kv_cache_store_backend_factory.h"
+#include "tpu_sync/telemetry/metrics_api.h"
+#include "tpu_sync/telemetry/metrics_backend.h"
 
 namespace tpu_raiden {
 namespace kv_cache {
@@ -780,13 +782,14 @@ MetadataCacheOptions MetadataCacheOptions::FromPosixOptions(
 PosixKVCacheStoreBackend::PosixKVCacheStoreBackend(
     std::shared_ptr<KVBackend> storage_backend, std::string name,
     size_t capacity_bytes, size_t lookup_batch_size,
-    MetadataCacheOptions cache_options)
+    MetadataCacheOptions cache_options, std::string dp_rank)
     : storage_backend_(std::move(storage_backend)),
       name_(std::move(name)),
       capacity_bytes_(capacity_bytes),
       lookup_batch_size_(lookup_batch_size > 0 ? lookup_batch_size
                                                : kDefaultLookupBatchSize),
       cache_options_(cache_options),
+      dp_rank_(dp_rank.empty() ? "0" : std::move(dp_rank)),
       cache_(cache_options.max_entries) {
   if (cache_options_.enabled()) {
     LOG(INFO) << "PosixKVCacheStoreBackend '" << name_
@@ -968,12 +971,19 @@ absl::StatusOr<BlockSliceList> PosixKVCacheStoreBackend::Lookup(
     size_t block;
     size_t shard;
   };
+  uint64_t cache_hits = 0;
   std::vector<ShardRef> misses;
   for (size_t i = 0; i < shard_keys.size(); ++i) {
     for (size_t j = 0; j < shard_keys[i].size(); ++j) {
-      if (!present[i][j]) misses.push_back({i, j});
+      if (present[i][j]) {
+        ++cache_hits;
+      } else {
+        misses.push_back({i, j});
+      }
     }
   }
+  uint64_t stat_hits = 0;
+  uint64_t stat_misses = 0;
   std::vector<const BlockKey*> confirmed;
   for (size_t offset = 0; offset < misses.size();
        offset += lookup_batch_size_) {
@@ -992,8 +1002,10 @@ absl::StatusOr<BlockSliceList> PosixKVCacheStoreBackend::Lookup(
       if (exists[k]) {
         present[miss.block][miss.shard] = true;
         confirmed.push_back(&shard_keys[miss.block][miss.shard]);
+        ++stat_hits;
       } else {
         chunk_has_absent = true;
+        ++stat_misses;
       }
     }
     if (options.match_mode == LookupMatchMode::kPrefix && chunk_has_absent) {
@@ -1002,6 +1014,23 @@ absl::StatusOr<BlockSliceList> PosixKVCacheStoreBackend::Lookup(
   }
   CacheInsert(confirmed);
 
+  telemetry::RaidenMetricStore& metric_store =
+      telemetry::RaidenMetricStore::GetGlobalMetricStore();
+  if (metric_store.HasBackends()) {
+    auto record = [&](absl::string_view result, uint64_t count) {
+      if (count == 0) return;
+      const telemetry::MetricLabel labels[] = {
+          {telemetry::metric_labels::kDpRank, dp_rank_},
+          {telemetry::metric_labels::kResult, result},
+      };
+      metric_store.IncrementCounter(
+          telemetry::metric_names::kStorageMetadataCheckTotal, labels, count);
+    };
+    record(telemetry::metric_labels::kResultCacheHit, cache_hits);
+    record(telemetry::metric_labels::kResultStatHit, stat_hits);
+    record(telemetry::metric_labels::kResultStatMiss, stat_misses);
+  }
+
   // Phase 3 (merge).
   auto all_present = [](const std::vector<bool>& shards) {
     return !shards.empty() && absl::c_all_of(shards, [](bool shard_present) {
@@ -1127,5 +1156,6 @@ REGISTER_KV_CACHE_STORE_BACKEND(
           std::move(backend), backend_name, options.capacity_bytes,
           options.lookup_batch_size,
           ::tpu_raiden::kv_cache::backends::storage::MetadataCacheOptions::
-              FromPosixOptions(options));
+              FromPosixOptions(options),
+          absl::StrCat(config.raiden_id.data_replica_idx));
     });

tpu_sync/kv_cache/backends/storage/posix_backend.h

@@ -272,10 +272,13 @@ class PosixKVCacheStoreBackend : public KVCacheStoreBackend {
                            std::string name = std::string(kPosixBackendName),
                            size_t capacity_bytes = 0,
                            size_t lookup_batch_size = kDefaultLookupBatchSize,
-                           MetadataCacheOptions cache_options = {});
+                           MetadataCacheOptions cache_options = {},
+                           std::string dp_rank = "0");
 
   std::string name() const override { return name_; }
 
+  const std::string& dp_rank() const { return dp_rank_; }
+
   size_t lookup_batch_size() const { return lookup_batch_size_; }
 
   const MetadataCacheOptions& metadata_cache_options() const {
@@ -354,6 +357,7 @@ class PosixKVCacheStoreBackend : public KVCacheStoreBackend {
   size_t capacity_bytes_ = 0;
   size_t lookup_batch_size_ = 32;
   const MetadataCacheOptions cache_options_;
+  const std::string dp_rank_ = "0";
 
   // Never held across storage I/O or key mapping.
   mutable absl::Mutex cache_mu_;

tpu_sync/kv_cache/backends/storage/posix_backend_test.cc

@@ -59,6 +59,9 @@
 #include "tpu_sync/kv_cache/kv_cache_store_backend.h"
 #include "tpu_sync/kv_cache/kv_cache_store_backend_factory.h"
 #include "tpu_sync/proto/worker_service.pb.h"
+#include "tpu_sync/telemetry/metrics_api.h"
+#include "tpu_sync/telemetry/metrics_backend.h"
+#include "tpu_sync/telemetry/mock_metrics_backend.h"
 
 namespace tpu_raiden {
 namespace kv_cache {
@@ -2709,6 +2712,185 @@ TEST_F(MetadataCacheTest, UnmappableHashTruncatesPrefix) {
             1);
 }
 
+auto MetadataCheckLabels(absl::string_view dp_rank, absl::string_view result) {
+  return testing::ElementsAre(
+      testing::FieldsAre(telemetry::metric_labels::kDpRank,
+                         std::string(dp_rank)),
+      testing::FieldsAre(telemetry::metric_labels::kResult,
+                         std::string(result)));
+}
+
+TEST_F(MetadataCacheTest, EmitsStatHitAndStatMissWhenCacheDisabled) {
+  Env env = CreateEnv(MetadataCacheOptions{}, /*batch_size=*/8, /*tp_size=*/2);
+  env.CreateAllShards("a");
+  env.CreateBlock("b", /*tp_rank=*/0);  // Rank 1 shard is missing.
+
+  auto mock_backend = std::make_unique<telemetry::MockMetricsBackend>();
+  auto* raw_backend = mock_backend.get();
+  EXPECT_CALL(
+      *raw_backend,
+      IncrementCounter(
+          telemetry::metric_names::kStorageMetadataCheckTotal,
+          MetadataCheckLabels("0", telemetry::metric_labels::kResultStatHit),
+          3));
+  EXPECT_CALL(
+      *raw_backend,
+      IncrementCounter(
+          telemetry::metric_names::kStorageMetadataCheckTotal,
+          MetadataCheckLabels("0", telemetry::metric_labels::kResultStatMiss),
+          3));
+
+  telemetry::ScopedMetricsBackendReset scoped_reset(std::move(mock_backend));
+  // 6 shards in one chunk: a_r0, a_r1, b_r0 exist (3 stat_hits); b_r1, c_r0,
+  // c_r1 are absent (3 stat_misses).
+  EXPECT_EQ(env.LookupCount({"a", "b", "c"}), 1);
+}
+
+TEST_F(MetadataCacheTest,
+       EmitsCacheHitStatHitAndStatMissWithTtlExpiryAndDelete) {
+  Env env =
+      CreateEnv(Enabled(/*max_entries=*/1000, kShortTtl), /*batch_size=*/4);
+  env.CreateBlock("a");
+  env.CreateBlock("b");
+
+  // 1. Cold lookup of {"a", "b"}: 2 stat_hit, 0 cache_hit, 0 stat_miss.
+  {
+    auto mock_backend = std::make_unique<telemetry::MockMetricsBackend>();
+    EXPECT_CALL(
+        *mock_backend,
+        IncrementCounter(
+            telemetry::metric_names::kStorageMetadataCheckTotal,
+            MetadataCheckLabels("0", telemetry::metric_labels::kResultStatHit),
+            2));
+    telemetry::ScopedMetricsBackendReset scoped_reset(std::move(mock_backend));
+    EXPECT_EQ(env.LookupCount({"a", "b"}), 2);
+  }
+
+  // 2. Warm lookup of {"a", "b", "c"} where "c" is missing: 2 cache_hit,
+  // 1 stat_miss.
+  {
+    auto mock_backend = std::make_unique<telemetry::MockMetricsBackend>();
+    EXPECT_CALL(
+        *mock_backend,
+        IncrementCounter(
+            telemetry::metric_names::kStorageMetadataCheckTotal,
+            MetadataCheckLabels("0", telemetry::metric_labels::kResultCacheHit),
+            2));
+    EXPECT_CALL(
+        *mock_backend,
+        IncrementCounter(
+            telemetry::metric_names::kStorageMetadataCheckTotal,
+            MetadataCheckLabels("0", telemetry::metric_labels::kResultStatMiss),
+            1));
+    telemetry::ScopedMetricsBackendReset scoped_reset(std::move(mock_backend));
+    EXPECT_EQ(env.LookupCount({"a", "b", "c"}), 2);
+  }
+
+  // 3. Create "c" on disk and lookup {"a", "b", "c"}: 2 cache_hit, 1 stat_hit.
+  env.CreateBlock("c");
+  {
+    auto mock_backend = std::make_unique<telemetry::MockMetricsBackend>();
+    EXPECT_CALL(
+        *mock_backend,
+        IncrementCounter(
+            telemetry::metric_names::kStorageMetadataCheckTotal,
+            MetadataCheckLabels("0", telemetry::metric_labels::kResultCacheHit),
+            2));
+    EXPECT_CALL(
+        *mock_backend,
+        IncrementCounter(
+            telemetry::metric_names::kStorageMetadataCheckTotal,
+            MetadataCheckLabels("0", telemetry::metric_labels::kResultStatHit),
+            1));
+    telemetry::ScopedMetricsBackendReset scoped_reset(std::move(mock_backend));
+    EXPECT_EQ(env.LookupCount({"a", "b", "c"}), 3);
+  }
+
+  // 4. Delete "b" from cache and disk, then lookup {"a", "b", "c"}: "a" and "c"
+  // hit in cache (2 cache_hit), "b" is re-probed on disk (1 stat_miss).
+  EXPECT_EQ(env.store->Delete({"b"}), 1);
+  env.RemoveBlock("b");
+  {
+    auto mock_backend = std::make_unique<telemetry::MockMetricsBackend>();
+    EXPECT_CALL(
+        *mock_backend,
+        IncrementCounter(
+            telemetry::metric_names::kStorageMetadataCheckTotal,
+            MetadataCheckLabels("0", telemetry::metric_labels::kResultCacheHit),
+            2));
+    EXPECT_CALL(
+        *mock_backend,
+        IncrementCounter(
+            telemetry::metric_names::kStorageMetadataCheckTotal,
+            MetadataCheckLabels("0", telemetry::metric_labels::kResultStatMiss),
+            1));
+    telemetry::ScopedMetricsBackendReset scoped_reset(std::move(mock_backend));
+    EXPECT_EQ(env.LookupCount({"a", "b", "c"}), 1);
+  }
+
+  // 5. After TTL expiry, "a" is re-probed on disk (1 stat_hit, 0 cache_hit).
+  absl::SleepFor(kShortTtl + absl::Milliseconds(50));
+  {
+    auto mock_backend = std::make_unique<telemetry::MockMetricsBackend>();
+    EXPECT_CALL(
+        *mock_backend,
+        IncrementCounter(
+            telemetry::metric_names::kStorageMetadataCheckTotal,
+            MetadataCheckLabels("0", telemetry::metric_labels::kResultStatHit),
+            1));
+    telemetry::ScopedMetricsBackendReset scoped_reset(std::move(mock_backend));
+    EXPECT_EQ(env.LookupCount({"a"}), 1);
+  }
+}
+
+TEST_F(MetadataCacheTest, PropagatesCustomDpRankViaConstructorAndFactory) {
+  auto backend = std::make_shared<CountingPosixKVBackend>(
+      "posix", absl::flat_hash_map<std::string, std::string>{});
+  auto mapper = std::make_shared<PosixPathMapper>(
+      scratch_dir_, "model_test",
+      ParallelismConfig{.tp_size = 1, .tp_rank = 0});
+  backend->set_mapper(mapper);
+
+  // Empty dp_rank falls back to "0".
+  PosixKVCacheStoreBackend default_store(backend, "posix", /*capacity=*/0,
+                                         /*lookup_batch_size=*/4,
+                                         MetadataCacheOptions{},
+                                         /*dp_rank=*/"");
+  EXPECT_EQ(default_store.dp_rank(), "0");
+
+  // Factory propagates config.raiden_id.data_replica_idx to dp_rank.
+  BackendConfig config;
+  config.type = "posix";
+  config.SetProperty("root_dir", scratch_dir_);
+  config.SetProperty("model_name", "model_test");
+  config.parallelism = {.tp_size = 1, .tp_rank = 0};
+  config.raiden_id = {"test_job", "0", "kv_cache", /*data_replica_idx=*/3};
+
+  TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<KVCacheStoreBackend> created,
+                       KVCacheStoreBackendFactory::Instance().CreateBackend(
+                           config, /*controller=*/nullptr));
+  auto factory_store =
+      std::dynamic_pointer_cast<PosixKVCacheStoreBackend>(created);
+  ASSERT_NE(factory_store, nullptr);
+  EXPECT_EQ(factory_store->dp_rank(), "3");
+
+  auto mock_backend = std::make_unique<telemetry::MockMetricsBackend>();
+  auto* raw_backend = mock_backend.get();
+  EXPECT_CALL(
+      *raw_backend,
+      IncrementCounter(
+          telemetry::metric_names::kStorageMetadataCheckTotal,
+          MetadataCheckLabels("3", telemetry::metric_labels::kResultStatMiss),
+          1));
+  EXPECT_CALL(*raw_backend,
+              SetGauge(telemetry::metric_names::kStorageFilesystemBytes,
+                       testing::_, testing::_))
+      .Times(testing::AnyNumber());
+  telemetry::ScopedMetricsBackendReset scoped_reset(std::move(mock_backend));
+  TF_ASSERT_OK_AND_ASSIGN(BlockSliceList result,
+                       factory_store->Lookup({"missing_block"}));
+  EXPECT_TRUE(result.empty());
+}
 }  // namespace
 }  // namespace storage
 }  // namespace backends

tpu_sync/kv_cache/kv_cache_store.cc

@@ -379,18 +379,16 @@ absl::StatusOr<std::unique_ptr<KVCacheStore>> KVCacheStore::Create(
     const auto& config = effective_configs[i];
     BackendConfig effective_config = (i == 0) ? effective_config0 : config;
 
-    // Apply global registry address and raiden_id to tier 0
-    // (HostOffloadBackend).
+    // Apply global registry address to tier 0 (HostOffloadBackend), and
+    // propagate the store's raiden_id to every tier that did not set its own.
     // TODO: Revisit secondary tier registration with global registry.
-    if (i == 0) {
-      if (effective_config.global_registry_address.empty() &&
-          !global_registry_address.empty()) {
-        effective_config.global_registry_address =
-            std::string(global_registry_address);
-      }
-      if (effective_config.raiden_id.empty() && !raiden_id.empty()) {
-        effective_config.raiden_id = raiden_id;
-      }
+    if (i == 0 && effective_config.global_registry_address.empty() &&
+        !global_registry_address.empty()) {
+      effective_config.global_registry_address =
+          std::string(global_registry_address);
+    }
+    if (effective_config.raiden_id.empty()) {
+      effective_config.raiden_id = effective_raiden_id;
     }
 
     ABSL_ASSIGN_OR_RETURN(auto backend,