Keep the peer store client across WriteRemote refusals.

tpu_sync/api/jax/kv_cache_store.py

@@ -162,6 +162,21 @@ def __init__(
     keyword -- there is no value for either that is valid to fall into by
     omission.
 
+    The store monitor (heartbeats to the global registry) and its evict
+    sweep (demoting cold blocks to higher-tier peer stores under memory
+    pressure) are configured through environment variables, not arguments:
+      RAIDEN_ENABLE_STORE_MONITOR       "true"/"1" runs the monitor
+      RAIDEN_ENABLE_EVICT_SWEEP         "true"/"1" runs the evict sweep
+                                        (requires the monitor)
+      RAIDEN_STORE_MONITOR_HEARTBEAT_S  heartbeat period, whole seconds
+      RAIDEN_EVICT_SWEEP_PERIOD_S       sweep fallback period, whole seconds
+      RAIDEN_EVICT_LOW_WATERMARK        free-block ratio in (0, 1) below
+                                        which the sweep starts demoting
+      RAIDEN_EVICT_HIGH_WATERMARK       free-block ratio at which it stops
+    Unset numeric values mean the built-in defaults. The switches only take
+    effect when global_registry_address is set; a registry-less store
+    ignores them, so one fleet-wide environment block is safe.
+
     Args:
       num_shards: Shard count for this store's RaidenController; must be
         >= 1.

tpu_sync/api/torch/kv_cache_store.py

@@ -180,6 +180,21 @@ def __init__(
     keyword -- there is no value for either that is valid to fall into by
     omission.
 
+    The store monitor (heartbeats to the global registry) and its evict
+    sweep (demoting cold blocks to higher-tier peer stores under memory
+    pressure) are configured through environment variables, not arguments:
+      RAIDEN_ENABLE_STORE_MONITOR       "true"/"1" runs the monitor
+      RAIDEN_ENABLE_EVICT_SWEEP         "true"/"1" runs the evict sweep
+                                        (requires the monitor)
+      RAIDEN_STORE_MONITOR_HEARTBEAT_S  heartbeat period, whole seconds
+      RAIDEN_EVICT_SWEEP_PERIOD_S       sweep fallback period, whole seconds
+      RAIDEN_EVICT_LOW_WATERMARK        free-block ratio in (0, 1) below
+                                        which the sweep starts demoting
+      RAIDEN_EVICT_HIGH_WATERMARK       free-block ratio at which it stops
+    Unset numeric values mean the built-in defaults. The switches only take
+    effect when global_registry_address is set; a registry-less store
+    ignores them, so one fleet-wide environment block is safe.
+
     Args:
       num_shards: Shard count for this store's RaidenController; must be
         >= 1.

tpu_sync/kv_cache/BUILD

@@ -137,6 +137,7 @@ cc_library(
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/time",
     ],
 )
 
@@ -148,9 +149,12 @@ cc_test(
     deps = [
         ":kv_cache_store",
         ":kv_cache_store_backend",
+        ":kv_cache_store_backend_factory",
         ":kv_cache_store_wrapper",
         ":raiden_id",
+        "//tpu_sync/kv_cache/global_registry:test_util",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/time",
         "@com_google_googletest//:gtest",
         "@com_google_googletest//:gtest_main",
     ],
@@ -378,6 +382,7 @@ cc_library(
         ":raiden_id",
         "//tpu_sync/kv_cache/global_registry:global_registry_cc_proto",
         "//tpu_sync/kv_cache/global_registry:global_registry_client_cc",
+        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/synchronization",

tpu_sync/kv_cache/host_offload_backend.cc

@@ -81,6 +81,12 @@ BuildLocalWorkerEndpoints(controller::RaidenController* ctrl) {
   return result;
 }
 
+// True for the codes an RPC fails with when the channel itself is the
+// problem, as opposed to the peer's handler answering with an error.
+bool IsTransportError(const absl::Status& status) {
+  return absl::IsUnavailable(status) || absl::IsDeadlineExceeded(status);
+}
+
 }  // namespace
 
 HostOffloadBackend::HostOffloadBackend(
@@ -757,9 +763,15 @@ HostOffloadBackend::BeginWriteRemote(
                         absl::ToInt64Milliseconds(requested_deadline))
           .Await();
   if (!response_or.ok()) {
-    // 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(dst_raiden_id);
+    // On a transport error 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. An application answer -- e.g. the
+    // RESOURCE_EXHAUSTED refusing a batch the peer cannot fit -- proves the
+    // peer is alive on this channel, and refusals cluster exactly when a
+    // reconnect is most wasteful: under memory pressure.
+    if (IsTransportError(response_or.status())) {
+      InvalidateStoreClient(dst_raiden_id);
+    }
     return response_or.status();
   }
 
@@ -795,7 +807,10 @@ HostOffloadBackend::PollWriteRemote(const RaidenId& dst_raiden_id,
 
   auto response_or = client->PollWriteRemote(operation_id).Await();
   if (!response_or.ok()) {
-    InvalidateStoreClient(dst_raiden_id);
+    // Same rule as BeginWriteRemote: only a suspect channel is dropped.
+    if (IsTransportError(response_or.status())) {
+      InvalidateStoreClient(dst_raiden_id);
+    }
     return response_or.status();
   }