diff --git a/tpu_sync/kv_cache/BUILD b/tpu_sync/kv_cache/BUILD index d2763078..5656efe2 100644 --- a/tpu_sync/kv_cache/BUILD +++ b/tpu_sync/kv_cache/BUILD @@ -382,7 +382,6 @@ 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", diff --git a/tpu_sync/kv_cache/kv_cache_store.cc b/tpu_sync/kv_cache/kv_cache_store.cc index 5ba700fa..7aafbba2 100644 --- a/tpu_sync/kv_cache/kv_cache_store.cc +++ b/tpu_sync/kv_cache/kv_cache_store.cc @@ -15,7 +15,6 @@ #include "tpu_sync/kv_cache/kv_cache_store.h" #include -#include #include #include #include @@ -56,8 +55,8 @@ #include "tpu_sync/kv_cache/kv_cache_store_backend_factory.h" #include "tpu_sync/kv_cache/raiden_id.h" #include "tpu_sync/kv_cache/reshard/reshard_service.h" -#include "tpu_sync/kv_cache/store_monitor.h" #include "tpu_sync/rpc/raiden_service.pb.h" +#include "tpu_sync/kv_cache/store_monitor.h" namespace tpu_raiden { namespace kv_cache { @@ -70,19 +69,6 @@ namespace { // re-register. constexpr int kRegistrationTtlInHeartbeats = 3; -// Evict-sweep defaults. The watermarks can sit this low because detection is -// not periodic: the allocation path nudges the sweep the moment free blocks -// dip below the low watermark, so the cushion only has to cover the demotion -// ramp-up, not a polling interval. -constexpr double kDefaultEvictLowWatermark = 0.03; -constexpr double kDefaultEvictHighWatermark = 0.05; -// Cap on one demotion batch. The batch is sized by what the high watermark -// still needs; this cap only keeps a single sweep step short, since Stop() -// and an interleaved heartbeat wait for at most one step. -constexpr int kMaxEvictBatchBlocks = 128; -// Placement targets requested per pressure episode. -constexpr int32_t kMaxPlacementTargets = 4; - // Bind-and-advertise address for this store's RaidenController. // // Empty ip preserves the legacy behaviour (RaidenController binds the wildcard @@ -112,7 +98,8 @@ absl::Status ValidateConstructionRules(absl::string_view store_server_ip, "store's services. Same-host use: \"127.0.0.1\"."); } if (store_server_ip == "[::]" || store_server_ip == "::" || - store_server_ip == "0.0.0.0" || store_server_ip == "0:0:0:0:0:0:0:0") { + store_server_ip == "0.0.0.0" || + store_server_ip == "0:0:0:0:0:0:0:0") { return absl::InvalidArgumentError(absl::StrCat( "store_server_ip may not be a wildcard (got \"", store_server_ip, "\"): a wildcard binds but cannot be published or dialled.")); @@ -170,9 +157,9 @@ MakeRaidenController(const RaidenId& raiden_id, size_t capacity, int num_shards, absl::StatusOr> KVCacheStore::Create( absl::Span backend_configs, size_t capacity, absl::string_view global_registry_address, RaidenId raiden_id, - int num_shards, int64_t shard_size_bytes, absl::string_view store_server_ip, - int raiden_controller_port, std::optional metadata, - int expected_worker_count) { + int num_shards, int64_t shard_size_bytes, + absl::string_view store_server_ip, int raiden_controller_port, + std::optional metadata, int expected_worker_count) { if (backend_configs.empty()) { return absl::InvalidArgumentError("backend_configs must not be empty"); } @@ -252,33 +239,11 @@ absl::StatusOr> KVCacheStore::Create( // these attributes. store->kv_pool_group_ = effective_config0.kv_pool_group; store->evict_tier_ = effective_config0.evict_tier; - StoreMonitorConfig& monitor_config = store->monitor_config_; - monitor_config = effective_config0.monitor_config; - if (monitor_config.heartbeat_period <= absl::ZeroDuration()) { - monitor_config.heartbeat_period = StoreMonitor::kDefaultHeartbeatPeriod; - } - if (monitor_config.evict_low_watermark <= 0.0) { - monitor_config.evict_low_watermark = kDefaultEvictLowWatermark; - } - if (monitor_config.evict_high_watermark <= 0.0) { - monitor_config.evict_high_watermark = kDefaultEvictHighWatermark; - } - if (monitor_config.enable_evict_sweep) { - if (!monitor_config.enable) { - return absl::FailedPreconditionError( - "StoreMonitorConfig.enable_evict_sweep requires the monitor " - "(StoreMonitorConfig.enable): the sweep runs on the store " - "monitor's schedule."); - } - if (monitor_config.evict_high_watermark < - monitor_config.evict_low_watermark || - monitor_config.evict_high_watermark > 1.0) { - return absl::FailedPreconditionError(absl::StrCat( - "evict watermarks must satisfy low <= high <= 1, got low=", - monitor_config.evict_low_watermark, - " high=", monitor_config.evict_high_watermark)); - } - } + store->enable_store_monitor_ = effective_config0.enable_store_monitor; + store->store_monitor_heartbeat_period_ = + effective_config0.store_monitor_heartbeat_period > absl::ZeroDuration() + ? effective_config0.store_monitor_heartbeat_period + : StoreMonitor::kDefaultHeartbeatPeriod; if (store->raiden_controller_ != nullptr) { RETURN_IF_ERROR( @@ -308,26 +273,19 @@ absl::StatusOr> KVCacheStore::Create( RETURN_IF_ERROR(host_backend->RegisterKVTransferSpecFromWorkers()); } - if (monitor_config.enable) { + if (store->enable_store_monitor_) { // The flag promises heartbeats; a store that never registered has // nothing to heartbeat, so this configuration is a contradiction, not a // degraded mode. if (!store->registered_in_global_registry_) { return absl::FailedPreconditionError( - "StoreMonitorConfig.enable requires a global_registry_address: the " + "enable_store_monitor requires a global_registry_address: the " "monitor's heartbeats refresh this store's registration there."); } - StoreMonitor::Options monitor_options; - monitor_options.heartbeat_period = monitor_config.heartbeat_period; - if (monitor_config.evict_sweep_period > absl::ZeroDuration()) { - monitor_options.sweep_period = monitor_config.evict_sweep_period; - } - StoreMonitor::SweepFn sweep_fn; - if (monitor_config.enable_evict_sweep) { - sweep_fn = [s = store.get()] { return s->SweepOnce(); }; - } store->store_monitor_ = std::make_unique( - monitor_options, store->registry_client_, store->raiden_id_, + StoreMonitor::Options{ + .heartbeat_period = store->store_monitor_heartbeat_period_}, + store->registry_client_, store->raiden_id_, /*status_fn=*/ [s = store.get()] { global_registry::StoreStatus status; @@ -344,13 +302,11 @@ absl::StatusOr> KVCacheStore::Create( s->raiden_controller_ != nullptr ? s->raiden_controller_->controller_address() : "", - kRegistrationTtlInHeartbeats * - s->monitor_config_.heartbeat_period, + kRegistrationTtlInHeartbeats * s->store_monitor_heartbeat_period_, !s->kv_pool_group_.empty() ? s->kv_pool_group_ : s->raiden_id_.job_name, s->evict_tier_); - }, - std::move(sweep_fn)); + }); store->store_monitor_->Start(); } @@ -360,13 +316,14 @@ absl::StatusOr> KVCacheStore::Create( absl::StatusOr> KVCacheStore::Create( const BackendConfig& config, size_t capacity, absl::string_view global_registry_address, RaidenId raiden_id, - int num_shards, int64_t shard_size_bytes, absl::string_view store_server_ip, - int raiden_controller_port, std::optional metadata, - int expected_worker_count) { - return KVCacheStore::Create( - absl::MakeConstSpan(&config, 1), capacity, global_registry_address, - raiden_id, num_shards, shard_size_bytes, store_server_ip, - raiden_controller_port, std::move(metadata), expected_worker_count); + int num_shards, int64_t shard_size_bytes, + absl::string_view store_server_ip, int raiden_controller_port, + std::optional metadata, int expected_worker_count) { + return KVCacheStore::Create(absl::MakeConstSpan(&config, 1), capacity, + global_registry_address, raiden_id, num_shards, + shard_size_bytes, store_server_ip, + raiden_controller_port, std::move(metadata), + expected_worker_count); } absl::StatusOr> KVCacheStore::CreateReshardStore( @@ -390,13 +347,14 @@ absl::StatusOr> KVCacheStore::CreateReshardStore( // num_shards=1 gives the controller an initial partition; workers dynamically // register their actual shard assignments via WorkerService. - ASSIGN_OR_RETURN(auto store, KVCacheStore::Create( - cfg, /*capacity=*/1, - /*global_registry_address=*/"", raiden_id, - /*num_shards=*/1, /*shard_size_bytes=*/0, - store_server_ip, raiden_controller_port, - /*metadata=*/std::nullopt, - /*expected_worker_count=*/0)); + ASSIGN_OR_RETURN(auto store, + KVCacheStore::Create( + cfg, /*capacity=*/1, + /*global_registry_address=*/"", raiden_id, + /*num_shards=*/1, /*shard_size_bytes=*/0, + store_server_ip, raiden_controller_port, + /*metadata=*/std::nullopt, + /*expected_worker_count=*/0)); // Initialize ReshardService with WorkerDelivery::Mode::kController. reshard::ReshardService::Options reshard_opts; @@ -515,7 +473,8 @@ KVCacheStore::KVCacheStore( : raiden_id_(raiden_id), store_server_ip_(store_server_ip), write_through_pool_(std::make_unique<::tpu_raiden::NumaThreadPool>(4)) { - if (absl::Status v = ValidateConstructionRules(store_server_ip, num_shards); + if (absl::Status v = ValidateConstructionRules(store_server_ip, + num_shards); !v.ok()) { LOG(FATAL) << "KVCacheStore construction validation failed: " << v.message() << " Use KVCacheStore::Create() for a recoverable error."; @@ -575,7 +534,8 @@ KVCacheStore::KVCacheStore( std::unique_ptr<::tpu_raiden::controller::RaidenController> raiden_controller, absl::string_view global_registry_address, RaidenId raiden_id, - std::optional metadata, absl::string_view store_server_ip) + std::optional metadata, + absl::string_view store_server_ip) : raiden_id_(raiden_id), raiden_controller_(std::move(raiden_controller)), store_server_ip_(store_server_ip), @@ -583,8 +543,7 @@ KVCacheStore::KVCacheStore( if (absl::Status v = ValidateConstructionRules(store_server_ip, /*num_shards=*/1); !v.ok()) { - LOG(FATAL) << "KVCacheStore construction validation failed: " - << v.message(); + LOG(FATAL) << "KVCacheStore construction validation failed: " << v.message(); } if (raiden_controller_ == nullptr) { LOG(FATAL) << "KVCacheStore requires a RaidenController; the " @@ -684,9 +643,10 @@ absl::Status KVCacheStore::EnsureStoreServerAndRegister( if (!status.ok()) { store_server_ = nullptr; owned_store_server_.reset(); - return absl::Status(status.code(), - absl::StrCat("Failed to start KVCacheStoreServer on ", - bind_address, ": ", status.message())); + return absl::Status( + status.code(), + absl::StrCat("Failed to start KVCacheStoreServer on ", bind_address, + ": ", status.message())); } // The server is the single source of truth for its own address -- it @@ -709,16 +669,17 @@ absl::Status KVCacheStore::EnsureStoreServerAndRegister( raiden_id_, store_server_address_, raiden_controller_ != nullptr ? raiden_controller_->controller_address() : "", - monitor_config_.enable - ? kRegistrationTtlInHeartbeats * monitor_config_.heartbeat_period + enable_store_monitor_ + ? kRegistrationTtlInHeartbeats * store_monitor_heartbeat_period_ : absl::ZeroDuration(), !kv_pool_group_.empty() ? kv_pool_group_ : raiden_id_.job_name, evict_tier_); if (!register_status.ok()) { return absl::Status( register_status.code(), - absl::StrCat("Failed to publish store address ", store_server_address_, - " to the global registry: ", register_status.message())); + absl::StrCat("Failed to publish store address ", + store_server_address_, " to the global registry: ", + register_status.message())); } registered_in_global_registry_ = true; LOG(INFO) << "KVCacheStore published at " << store_server_address_; @@ -1556,164 +1517,11 @@ size_t KVCacheStore::Evict(const std::vector& block_hashes) { return host_ids_to_deallocate.size(); } -bool KVCacheStore::SweepOnce() { - const int free_blocks = - raiden_controller_->block_manager()->num_free_blocks(); - const int total_blocks = raiden_controller_->block_manager()->total_blocks(); - if (total_blocks <= 0) { - return false; - } - const double free_ratio = static_cast(free_blocks) / total_blocks; - // Ends the current pressure episode; the next one re-fetches targets. - auto end_episode = [this] { - sweep_active_ = false; - placement_targets_.clear(); - return false; - }; - - if (!sweep_active_) { - // Idle and enough free blocks: nothing to do. - if (free_ratio >= monitor_config_.evict_low_watermark) { - return false; - } - // Free blocks fell below the low watermark: a pressure episode begins - // (it ends when they recover to the high watermark). The targets are - // fetched once here and reused for the whole episode, so even a long - // drain costs one registry read. If fleets of stores hitting pressure - // together ever make these per-episode reads a load problem for the - // registry, the fetch could be decoupled from the sweep into its own - // periodic task maintaining a local target list -- at the cost of - // staler placement data. - sweep_active_ = true; - placement_targets_.clear(); - auto targets_or = - registry_client_->GetPlacementTargets(raiden_id_, kMaxPlacementTargets); - if (targets_or.ok()) { - for (const auto& info : *targets_or) { - placement_targets_.push_back(RaidenId{ - info.raiden_id().job_name(), info.raiden_id().job_replica_id(), - info.raiden_id().data_name(), - static_cast(info.raiden_id().data_replica_idx())}); - } - } else { - LOG(WARNING) << "Evict sweep could not fetch placement targets: " - << targets_or.status() - << ". Dropping cold blocks locally instead."; - } - } else if (free_ratio >= monitor_config_.evict_high_watermark) { - // Recovered to the high watermark: the episode is over. - return end_episode(); - } - - // One batch per step, sized to what the high watermark still needs and - // capped so a single step stays short. - const int deficit = - static_cast( - std::ceil(monitor_config_.evict_high_watermark * total_blocks)) - - free_blocks; - std::vector batch; - { - absl::MutexLock lock(mutex_); - batch = - backend()->GetEvictableKeys(std::min(deficit, kMaxEvictBatchBlocks)); - } - if (batch.empty()) { - // Everything left is pinned; nothing more this episode can free. - return end_episode(); - } - - // Offer the batch to the episode's targets in order. Every iteration - // either sends the batch or drops one target, so this ends. - while (!placement_targets_.empty()) { - const RaidenId dst = placement_targets_.front(); - absl::Status offered = WriteRemote(batch, dst); - if (!offered.ok()) { - // Refused (peer out of free blocks) or unreachable: this target is - // out for the rest of the episode; try the batch on the next one. - LOG(INFO) << "Evict sweep target " << dst - << " declined a demotion batch: " << offered; - placement_targets_.erase(placement_targets_.begin()); - continue; - } - const BatchWriteResult result = WaitForBatchWriteResult(batch); - const size_t evicted = Evict(result.freeable); - if (result.transfer_failed) { - // The peer accepted but a transfer failed; the failed blocks stay - // local. Drop the target so the next step retries them elsewhere. - placement_targets_.erase(placement_targets_.begin()); - } else if (evicted == 0) { - // The peer holds the whole batch, yet nothing could be freed locally: - // readers pinned every block between the offer and here. Re-offering - // the same keys could spin without raising the free count, so end the - // episode; the next wake re-evaluates. - return end_episode(); - } - return true; - } - - // No target will take the batch (none exist, all dropped, or the registry - // was unreachable): drop it locally -- the same discard AllocateBlockIds - // would be forced into later, done early enough to keep absorbing writes. - if (Evict(batch) == 0) { - // The whole batch got pinned since it was picked: end the episode - // rather than spin on the same keys. - return end_episode(); - } - return true; -} - -KVCacheStore::BatchWriteResult KVCacheStore::WaitForBatchWriteResult( - const std::vector& batch) { - // The sweep is the only WriteRemote caller, so every drained result - // belongs to this batch. WriteRemote's HOLD deadline guarantees each - // block eventually leaves `pending`. - absl::flat_hash_set pending(batch.begin(), batch.end()); - BatchWriteResult result; - while (!pending.empty()) { - auto [done, failed, still_pending, existing, unregistered] = - PollRemoteWriteStatus(); - // Completed transfers: the peer holds these blocks now. - for (const std::string& hash : done) { - if (pending.erase(hash) > 0) { - result.freeable.push_back(hash); - } - } - for (const std::string& hash : failed) { - pending.erase(hash); - } - // `existing` annotates refusals whose bytes the peer already holds -- - // as freeable as a completed transfer. Everything else in `failed` - // (including `unregistered`: landed but unfindable there) keeps its - // local copy. - for (const std::string& hash : existing) { - result.freeable.push_back(hash); - } - if (failed.size() > existing.size()) { - result.transfer_failed = true; - } - if (!pending.empty()) { - absl::SleepFor(absl::Milliseconds(10)); - } - } - return result; -} - absl::StatusOr> KVCacheStore::AllocateBlockIds(int needed) { std::vector hashes_to_deallocate; - bool request_sweep = false; { absl::MutexLock lock(mutex_); int free_count = raiden_controller_->block_manager()->num_free_blocks(); - // Wake the evict sweep the moment this allocation dips below its low - // watermark, instead of leaving detection to the sweep's fallback - // period. The drop-evict below stays the last-ditch fallback for - // allocations the sweep has not made room for. - if (monitor_config_.enable_evict_sweep && store_monitor_ != nullptr) { - const int total = raiden_controller_->block_manager()->total_blocks(); - request_sweep = - total > 0 && - free_count - needed < monitor_config_.evict_low_watermark * total; - } int to_free = needed - free_count; if (to_free > 0) { hashes_to_deallocate = backend()->GetEvictableKeys(to_free); @@ -1727,10 +1535,6 @@ absl::StatusOr> KVCacheStore::AllocateBlockIds(int needed) { } } - if (request_sweep) { - store_monitor_->RequestSweep(); - } - if (!hashes_to_deallocate.empty()) { Evict(hashes_to_deallocate); } @@ -2243,6 +2047,7 @@ void KVCacheStore::PollRemoteReadsInternal( } } + void KVCacheStore::PollFuturesInternal() { std::vector ready_saves; std::vector ready_loads; diff --git a/tpu_sync/kv_cache/kv_cache_store.h b/tpu_sync/kv_cache/kv_cache_store.h index 731f31b6..140be2d7 100644 --- a/tpu_sync/kv_cache/kv_cache_store.h +++ b/tpu_sync/kv_cache/kv_cache_store.h @@ -591,23 +591,6 @@ class KVCacheStore { // service holds it in a pointer it cannot re-seat. void ShutdownBackendStoreServers(KVCacheStoreServer* already_shut); - // One bounded step of the evict sweep: checks the free-block watermarks - // and demotes at most one batch of cold blocks to a placement target (or - // drops the batch locally when no target will take it). Returns true when - // the sweep should run again right away. The store monitor's thread is - // the only caller, which is why the episode state below needs no lock. - bool SweepOnce(); - - // What became of one offered batch, polled until every block is terminal. - struct BatchWriteResult { - // Blocks the peer now holds; their local copies are safe to free. - std::vector freeable; - // At least one block's transfer failed outright; its local copy stays. - bool transfer_failed = false; - }; - BatchWriteResult WaitForBatchWriteResult( - const std::vector& batch); - mutable absl::Mutex mutex_; std::vector> backends_; std::shared_ptr registry_client_; @@ -631,19 +614,11 @@ class KVCacheStore { // True once this store published itself, so teardown knows to unpublish. bool registered_in_global_registry_ = false; - // Registration identity from the tier-0 BackendConfig. + // Registration attributes from the tier-0 BackendConfig. std::string kv_pool_group_; int32_t evict_tier_ = 0; - // Monitor and sweep knobs from the tier-0 BackendConfig, with zeros - // resolved to the built-in defaults in Create. - StoreMonitorConfig monitor_config_; - // Pressure-episode state, touched only by SweepOnce on the monitor's - // thread. One episode = the run of sweep steps from free blocks falling - // below the low watermark until they recover to the high watermark (or - // nothing more can be freed); the placement targets are fetched once per - // episode and reused until it ends or every target has been dropped. - bool sweep_active_ = false; - std::vector placement_targets_; + bool enable_store_monitor_ = false; + absl::Duration store_monitor_heartbeat_period_ = absl::ZeroDuration(); // Constructed and started at the end of Create when enabled; stopped first // thing in the destructor, before anything its callbacks touch goes away. std::unique_ptr store_monitor_; diff --git a/tpu_sync/kv_cache/kv_cache_store_backend_factory.h b/tpu_sync/kv_cache/kv_cache_store_backend_factory.h index bcb67fb7..37088991 100644 --- a/tpu_sync/kv_cache/kv_cache_store_backend_factory.h +++ b/tpu_sync/kv_cache/kv_cache_store_backend_factory.h @@ -55,32 +55,6 @@ struct KVTransferSpecConfig { int num_workers = 0; }; -// The knobs behind a store's StoreMonitor: the heartbeat it sends and the -// evict sweep it schedules. Zero values mean the built-in defaults. -struct StoreMonitorConfig { - // Runs a StoreMonitor thread that heartbeats the store's status to the - // global registry; the store's registration then carries a TTL and expires - // when heartbeats stop. - bool enable = false; - // Heartbeat period. Zero means the StoreMonitor default. With the monitor - // enabled this also sets the registration TTL, a fixed multiple of the - // period. - absl::Duration heartbeat_period = absl::ZeroDuration(); - // Demotes cold blocks to a peer store on a higher evict_tier whenever free - // blocks fall below evict_low_watermark. Requires `enable`: the sweep runs - // on the store monitor's schedule. - bool enable_evict_sweep = false; - // Fallback period between sweep pressure checks; allocation pressure wakes - // the sweep immediately. Zero means the StoreMonitor default. - absl::Duration evict_sweep_period = absl::ZeroDuration(); - // Free-block ratio (free / total) below which the sweep starts demoting. - // Zero means the default. - double evict_low_watermark = 0.0; - // Free-block ratio at which an active sweep stops; must be >= the low - // watermark. Zero means the default. - double evict_high_watermark = 0.0; -}; - struct BackendConfig { std::string type; size_t capacity = 0; @@ -95,7 +69,14 @@ struct BackendConfig { std::string kv_pool_group; // Placement tier the store registers under (see StoreInfo.evict_tier). int32_t evict_tier = 0; - StoreMonitorConfig monitor_config; + // Runs a StoreMonitor thread that heartbeats the store's status to the + // global registry; the store's registration then carries a TTL and expires + // when heartbeats stop. + bool enable_store_monitor = false; + // StoreMonitor heartbeat period. Zero means the StoreMonitor default. With + // the monitor enabled this also sets the registration TTL, a fixed multiple + // of the period. + absl::Duration store_monitor_heartbeat_period = absl::ZeroDuration(); std::string GetProperty(absl::string_view key, absl::string_view default_val = "") const; diff --git a/tpu_sync/kv_cache/kv_cache_store_test.cc b/tpu_sync/kv_cache/kv_cache_store_test.cc index 88650578..1c15d2f8 100644 --- a/tpu_sync/kv_cache/kv_cache_store_test.cc +++ b/tpu_sync/kv_cache/kv_cache_store_test.cc @@ -3526,8 +3526,8 @@ TEST_F(StoreDiscoveryTest, StoreMonitorHeartbeatsTheRegistration) { config.global_registry_address = registry_address_; config.kv_pool_group = "groupA"; config.evict_tier = 1; - config.monitor_config.enable = true; - config.monitor_config.heartbeat_period = absl::Milliseconds(300); + config.enable_store_monitor = true; + config.store_monitor_heartbeat_period = absl::Milliseconds(300); auto store_or = KVCacheStore::Create(config, /*capacity=*/16, registry_address_, rid, @@ -3572,7 +3572,7 @@ TEST_F(StoreDiscoveryTest, StoreMonitorWithoutARegistryIsAnError) { config.type = "HostOffloadBackend"; config.capacity = 16; config.raiden_id = rid; - config.monitor_config.enable = true; + config.enable_store_monitor = true; auto store_or = KVCacheStore::Create(config, /*capacity=*/16, /*global_registry_address=*/"", rid, @@ -3972,9 +3972,7 @@ class RemoteWriteSourceTest : public StoreDiscoveryTest { // Stands the fake up and publishes it under `dst` so the source resolves it // through the registry exactly as it would a real peer. - void StartFakeDestination(const RaidenId& dst, - absl::string_view kv_pool_group = "", - int32_t evict_tier = 0) { + void StartFakeDestination(const RaidenId& dst) { grpc::ServerBuilder builder; int port = 0; builder.AddListeningPort("127.0.0.1:0", grpc::InsecureServerCredentials(), @@ -3984,9 +3982,7 @@ class RemoteWriteSourceTest : public StoreDiscoveryTest { ASSERT_NE(fake_destination_server_, nullptr); ASSERT_TRUE(client_ ->RegisterStore(dst, "127.0.0.1:" + std::to_string(port), - /*controller_address=*/"", - /*ttl=*/absl::ZeroDuration(), kv_pool_group, - evict_tier) + /*controller_address=*/"") .ok()); } @@ -4485,212 +4481,6 @@ TEST_F(RemoteWriteSourceTest, AStaleButRegisteredDestinationFailsPromptly) { << "offering to a dead peer should fail rather than hang"; } -class EvictSweepTest : public RemoteWriteSourceTest { - protected: - // A store with the sweep on. With the 0.5/0.75 watermarks, pressure - // starts below capacity/2 free blocks and relief comes at 3/4 free. - absl::StatusOr> MakeSweepStore( - const RaidenId& id, absl::string_view kv_pool_group, - size_t capacity = kCapacity) { - BackendConfig config; - config.type = "HostOffloadBackend"; - config.capacity = capacity; - config.raiden_id = id; - config.global_registry_address = registry_address_; - config.kv_pool_group = std::string(kv_pool_group); - config.monitor_config.enable = true; - config.monitor_config.enable_evict_sweep = true; - config.monitor_config.evict_sweep_period = absl::Milliseconds(200); - config.monitor_config.evict_low_watermark = 0.5; - config.monitor_config.evict_high_watermark = 0.75; - return KVCacheStore::Create(config, capacity, registry_address_, id, - /*num_shards=*/1, - /*shard_size_bytes=*/1024, - /*store_server_ip=*/"127.0.0.1"); - } - - // Fills `store` with cold (unpinned) host blocks whose ids really come from - // the block manager, so the free-block count the sweep watches drops and - // eviction raises it back. - void PopulateCold(KVCacheStore& store, const RaidenId& id, - const std::vector& hashes) { - auto ids_or = store.raiden_controller()->AllocateBlockIds(hashes.size()); - ASSERT_TRUE(ids_or.ok()) << ids_or.status().ToString(); - std::vector slices; - for (size_t i = 0; i < hashes.size(); ++i) { - slices.push_back(RaidenBlockID(id, (*ids_or)[i], BlockStatus::HOST)); - } - ASSERT_TRUE(store.InsertAndLock(hashes, slices, /*on_host=*/true)); - store.Release(hashes); - } - - // The sweep runs on the monitor's thread; wait until it has raised the free - // count to `expected_free`. Free blocks, not cache size, because eviction - // erases the cache entry a beat before the block id is deallocated. - void AwaitSweepFreed(KVCacheStore& store, int expected_free) { - for (int i = 0; i < 1000; ++i) { - if (store.raiden_controller()->block_manager()->num_free_blocks() == - expected_free) { - break; - } - absl::SleepFor(absl::Milliseconds(10)); - } - EXPECT_EQ(store.raiden_controller()->block_manager()->num_free_blocks(), - expected_free); - } - - // How many of `hashes` the store still resolves locally. - static size_t CountResident(KVCacheStore& store, - const std::vector& hashes) { - size_t resident = 0; - for (const std::string& hash : hashes) { - auto found = store.backend()->Lookup({hash}); - if (found.ok() && found->size() == 1) { - ++resident; - } - } - return resident; - } -}; - -// The main path: pressure starts an episode, the sweep fetches a same-group -// higher-tier target from the registry and demotes LRU-cold batches to it -// until the high watermark, freeing the local copies. -TEST_F(EvictSweepTest, DemotesColdBlocksToAPlacementTarget) { - RaidenId src{"sweep_src", "0", "kv", 0}; - RaidenId dst{"sweep_dst", "0", "kv", 0}; - // Large enough that reaching the high watermark takes several batches of - // the built-in batch cap (128). - auto store_or = MakeSweepStore(src, "sweepgroup", /*capacity=*/600); - ASSERT_TRUE(store_or.ok()) << store_or.status().ToString(); - KVCacheStore& store = **store_or; - ASSERT_EQ(store.raiden_controller()->block_manager()->total_blocks(), 600); - - StartFakeDestination(dst, "sweepgroup", /*evict_tier=*/1); - proto::PollWriteRemoteResponse verdict; - verdict.set_state(proto::PollWriteRemoteResponse::COMMITTED); - fake_destination_.SetPollResponse(verdict); - - // 500 of 600 blocks in use: free ratio 1/6, below the 0.5 low watermark. - std::vector hashes; - for (int i = 0; i < 500; ++i) { - hashes.push_back(absl::StrCat("h", i)); - } - PopulateCold(store, src, hashes); - - // The deficit to the 0.75 high watermark is 350 blocks: batches of - // 128 + 128 + 94, each demoted and freed. - AwaitSweepFreed(store, 450); - EXPECT_EQ(fake_destination_.write_calls(), 3); - EXPECT_EQ(store.backend()->GetSize(), 150); - EXPECT_EQ(CountResident(store, hashes), 150); -} - -// A target that cannot be reached is dropped for the episode and the same -// batch goes to the next target in the ranking. -TEST_F(EvictSweepTest, SkipsAnUnreachableTargetForTheNextOne) { - RaidenId src{"sweep_src_skip", "0", "kv", 0}; - RaidenId dead{"sweep_dead", "0", "kv", 0}; - RaidenId dst{"sweep_dst_skip", "0", "kv", 0}; - auto store_or = MakeSweepStore(src, "skipgroup"); - ASSERT_TRUE(store_or.ok()) << store_or.status().ToString(); - KVCacheStore& store = **store_or; - - // A dead peer, ranked first: registered at tier 1 with the most reported - // free blocks, but nothing listens on its address. - ASSERT_TRUE(client_ - ->RegisterStore(dead, "127.0.0.1:1", - /*controller_address=*/"", - /*ttl=*/absl::ZeroDuration(), "skipgroup", - /*evict_tier=*/1) - .ok()); - global_registry::StoreStatus roomy; - roomy.set_free_blocks(1000); - ASSERT_TRUE(client_->Heartbeat(dead, roomy).ok()); - StartFakeDestination(dst, "skipgroup", /*evict_tier=*/1); - proto::PollWriteRemoteResponse verdict; - verdict.set_state(proto::PollWriteRemoteResponse::COMMITTED); - fake_destination_.SetPollResponse(verdict); - - PopulateCold(store, src, {"a", "b", "c", "d", "e", "f"}); - - // The offer to the dead peer fails, it is dropped, and the batch lands on - // the live target instead. - AwaitSweepFreed(store, 6); - EXPECT_EQ(fake_destination_.write_calls(), 1); - EXPECT_EQ(store.backend()->GetSize(), 2); -} - -// A target that accepts but then fails the transfer is dropped too; with no -// target left, the next step falls back to dropping the blocks locally. -TEST_F(EvictSweepTest, ATransferFailureFallsBackToLocalDrop) { - RaidenId src{"sweep_src_fail", "0", "kv", 0}; - RaidenId dst{"sweep_dst_fail", "0", "kv", 0}; - auto store_or = MakeSweepStore(src, "failgroup"); - ASSERT_TRUE(store_or.ok()) << store_or.status().ToString(); - KVCacheStore& store = **store_or; - - StartFakeDestination(dst, "failgroup", /*evict_tier=*/1); - proto::PollWriteRemoteResponse verdict; - verdict.set_state(proto::PollWriteRemoteResponse::FAILED); - fake_destination_.SetPollResponse(verdict); - - PopulateCold(store, src, {"a", "b", "c", "d", "e", "f"}); - - // One offer is accepted and fails; the target is dropped and the free - // count still recovers, through local drops. - AwaitSweepFreed(store, 6); - EXPECT_EQ(fake_destination_.write_calls(), 1); - EXPECT_EQ(store.backend()->GetSize(), 2); -} - -// The bottom tier has no placement targets; under pressure the sweep drops -// cold blocks locally instead -- the same discard the allocation path would -// be forced into later, done early. -TEST_F(EvictSweepTest, DropsLocallyWhenThereAreNoTargets) { - RaidenId src{"sweep_src_bottom", "0", "kv", 0}; - auto store_or = MakeSweepStore(src, "sweepgroup_bottom"); - ASSERT_TRUE(store_or.ok()) << store_or.status().ToString(); - KVCacheStore& store = **store_or; - - PopulateCold(store, src, {"a", "b", "c", "d", "e", "f"}); - - AwaitSweepFreed(store, 6); - EXPECT_EQ(store.backend()->GetSize(), 2); - EXPECT_EQ(fake_destination_.write_calls(), 0); -} - -// The sweep flag rides on the monitor; asking for one without the other is a -// contradiction, as are watermarks that never let the sweep stop. -TEST_F(EvictSweepTest, SweepConfigContradictionsAreErrors) { - RaidenId id{"sweep_misconfigured", "0", "kv", 0}; - BackendConfig config; - config.type = "HostOffloadBackend"; - config.capacity = kCapacity; - config.raiden_id = id; - config.global_registry_address = registry_address_; - config.monitor_config.enable_evict_sweep = true; - - auto no_monitor = KVCacheStore::Create(config, /*capacity=*/kCapacity, - registry_address_, id, - /*num_shards=*/1, - /*shard_size_bytes=*/1024, - /*store_server_ip=*/"127.0.0.1"); - EXPECT_TRUE(absl::IsFailedPrecondition(no_monitor.status())) - << no_monitor.status().ToString(); - - config.monitor_config.enable = true; - config.monitor_config.evict_low_watermark = 0.5; - config.monitor_config.evict_high_watermark = 0.25; - auto inverted = KVCacheStore::Create(config, /*capacity=*/kCapacity, - registry_address_, id, - /*num_shards=*/1, - /*shard_size_bytes=*/1024, - /*store_server_ip=*/"127.0.0.1"); - EXPECT_TRUE(absl::IsFailedPrecondition(inverted.status())) - << inverted.status().ToString(); -} - // --------------------------------------------------------------------------- // Construction rules (kv_cache_store_construction_rules.md): Create() rejects // a missing/wildcard store_server_ip and a controller-less configuration with diff --git a/tpu_sync/kv_cache/store_monitor.cc b/tpu_sync/kv_cache/store_monitor.cc index 318c5ae8..4c48aaff 100644 --- a/tpu_sync/kv_cache/store_monitor.cc +++ b/tpu_sync/kv_cache/store_monitor.cc @@ -14,16 +14,13 @@ #include "tpu_sync/kv_cache/store_monitor.h" -#include + #include #include // NOLINT(build/c++11) #include #include "absl/log/log.h" #include "absl/status/status.h" -#include "absl/synchronization/mutex.h" -#include "absl/time/clock.h" -#include "absl/time/time.h" #include "tpu_sync/kv_cache/global_registry/global_registry_client.h" #include "tpu_sync/kv_cache/raiden_id.h" @@ -33,70 +30,31 @@ namespace kv_cache { StoreMonitor::StoreMonitor( const Options& options, std::shared_ptr registry_client, - RaidenId raiden_id, StatusFn status_fn, ReregisterFn reregister_fn, - SweepFn sweep_fn) + RaidenId raiden_id, StatusFn status_fn, ReregisterFn reregister_fn) : options_(options), registry_client_(std::move(registry_client)), raiden_id_(std::move(raiden_id)), status_fn_(std::move(status_fn)), - reregister_fn_(std::move(reregister_fn)), - sweep_fn_(std::move(sweep_fn)) {} + reregister_fn_(std::move(reregister_fn)) {} StoreMonitor::~StoreMonitor() { Stop(); } -void StoreMonitor::Start() { thread_ = std::thread(&StoreMonitor::Loop, this); } +void StoreMonitor::Start() { + thread_ = std::thread(&StoreMonitor::Loop, this); +} void StoreMonitor::Stop() { - { - absl::MutexLock lock(mu_); - stop_ = true; - cv_.Signal(); + if (!stop_.HasBeenNotified()) { + stop_.Notify(); } if (thread_.joinable()) { thread_.join(); } } -void StoreMonitor::RequestSweep() { - if (sweep_fn_ == nullptr) { - return; - } - absl::MutexLock lock(mu_); - sweep_requested_ = true; - cv_.Signal(); -} - void StoreMonitor::Loop() { - const bool sweeping = sweep_fn_ != nullptr; - absl::Time next_heartbeat = absl::Now() + options_.heartbeat_period; - absl::Time next_sweep = absl::Now() + options_.sweep_period; - // True while a sweep request is pending or the last sweep step reported more - // work; the loop then skips the wait, but still passes the heartbeat - // check between steps -- so a long drain delays a heartbeat by at most - // one bounded step. - bool sweep_now = false; - while (true) { - { - absl::MutexLock lock(mu_); - const absl::Time wake = - sweeping ? std::min(next_heartbeat, next_sweep) : next_heartbeat; - while (!stop_ && !sweep_requested_ && !sweep_now && absl::Now() < wake) { - cv_.WaitWithDeadline(&mu_, wake); - } - if (stop_) { - return; - } - sweep_now = sweep_now || sweep_requested_; - sweep_requested_ = false; - } - if (absl::Now() >= next_heartbeat) { - HeartbeatOnce(); - next_heartbeat = absl::Now() + options_.heartbeat_period; - } - if (sweeping && (sweep_now || absl::Now() >= next_sweep)) { - sweep_now = sweep_fn_(); - next_sweep = absl::Now() + options_.sweep_period; - } + while (!stop_.WaitForNotificationWithTimeout(options_.heartbeat_period)) { + HeartbeatOnce(); } } diff --git a/tpu_sync/kv_cache/store_monitor.h b/tpu_sync/kv_cache/store_monitor.h index 4b147fe9..64c8f987 100644 --- a/tpu_sync/kv_cache/store_monitor.h +++ b/tpu_sync/kv_cache/store_monitor.h @@ -19,9 +19,8 @@ #include #include // NOLINT(build/c++11) -#include "absl/base/thread_annotations.h" #include "absl/status/status.h" -#include "absl/synchronization/mutex.h" +#include "absl/synchronization/notification.h" #include "absl/time/time.h" #include "tpu_sync/kv_cache/global_registry/global_registry.pb.h" #include "tpu_sync/kv_cache/global_registry/global_registry_client.h" @@ -37,47 +36,34 @@ namespace kv_cache { // registration (registry restart or TTL expiry); heartbeats carry no // coordinates, so the monitor re-registers through the callback instead. // -// With a SweepFn, the same thread also schedules the store's evict sweep -- -// one thread total, a tradeoff to balance the number of threads raiden is -// managing. The sweep runs once per sweep period, immediately when -// NudgeSweep() is called, and step by step while the callback keeps -// returning true; the heartbeat check is interleaved between steps, so -// sweeping delays a heartbeat by at most one bounded step. -// // Owned by KVCacheStore; deliberately talks to the store only through the -// callbacks so it never sees store internals. +// two callbacks so it never sees store internals. class StoreMonitor { public: static constexpr absl::Duration kDefaultHeartbeatPeriod = absl::Seconds(300); - static constexpr absl::Duration kDefaultSweepPeriod = absl::Seconds(30); struct Options { absl::Duration heartbeat_period = kDefaultHeartbeatPeriod; - absl::Duration sweep_period = kDefaultSweepPeriod; }; // Snapshots the store's current status for one heartbeat. using StatusFn = std::function; // Re-publishes the store's full registration (coordinates and TTL). using ReregisterFn = std::function; - // One bounded step of evict-sweep work. Returns true if more work remains, - // in which case the monitor runs the next step right away; false parks the - // sweep until the next period or request. - using SweepFn = std::function; StoreMonitor(const Options& options, std::shared_ptr registry_client, RaidenId raiden_id, StatusFn status_fn, - ReregisterFn reregister_fn, SweepFn sweep_fn = nullptr); + ReregisterFn reregister_fn); - // Stops the threads; the callbacks must outlive this call, not the object. + // Stops the thread; the callbacks must outlive this call, not the object. ~StoreMonitor(); StoreMonitor(const StoreMonitor&) = delete; StoreMonitor& operator=(const StoreMonitor&) = delete; - // Starts the monitor thread. The first heartbeat fires one period from + // Starts the heartbeat thread. The first heartbeat fires one period from // now: the caller registers before starting the monitor, so the registry // is already fresh. Call at most once. void Start(); @@ -86,12 +72,6 @@ class StoreMonitor { // cannot be restarted. void Stop(); - // Request to run the sweep now instead of at its next period. Called from - // the store's allocation path when free blocks dip below the sweep's low - // watermark, so pressure is acted on immediately rather than discovered a - // sweep period later. No-op without a SweepFn or after Stop(). - void RequestSweep(); - private: void Loop(); void HeartbeatOnce(); @@ -102,18 +82,12 @@ class StoreMonitor { const RaidenId raiden_id_; const StatusFn status_fn_; const ReregisterFn reregister_fn_; - const SweepFn sweep_fn_; - - absl::Mutex mu_; - // Signaled on Stop() and RequestSweep(); the loop waits on it with a - // deadline, so a signal is its only early wake-up. - absl::CondVar cv_; - bool stop_ ABSL_GUARDED_BY(mu_) = false; - bool sweep_requested_ ABSL_GUARDED_BY(mu_) = false; + + absl::Notification stop_; std::thread thread_; }; } // namespace kv_cache } // namespace tpu_raiden -#endif // THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_KV_CACHE_STORE_MONITOR_H_ +#endif // THIRD_PARTY_TPU_RAIDEN_KV_CACHE_STORE_MONITOR_H_ diff --git a/tpu_sync/kv_cache/store_monitor_test.cc b/tpu_sync/kv_cache/store_monitor_test.cc index dc9b6bca..2be875ae 100644 --- a/tpu_sync/kv_cache/store_monitor_test.cc +++ b/tpu_sync/kv_cache/store_monitor_test.cc @@ -148,79 +148,6 @@ TEST_F(StoreMonitorTest, ReportedStatusFeedsThePlacementRanking) { EXPECT_EQ((*targets)[1].raiden_id().job_name(), "crowded"); } -TEST_F(StoreMonitorTest, SweepRunsOnItsPeriod) { - RaidenId id = {"sweeping", "r0", "dataS", 0}; - std::atomic sweep_calls{0}; - StoreMonitor monitor( - StoreMonitor::Options{.heartbeat_period = absl::Hours(1), - .sweep_period = absl::Milliseconds(200)}, - client_, id, ReportFreeBlocks(1), - /*reregister_fn=*/[] { return absl::OkStatus(); }, - /*sweep_fn=*/[&sweep_calls] { - ++sweep_calls; - return false; - }); - monitor.Start(); - absl::SleepFor(absl::Seconds(1)); - EXPECT_GE(sweep_calls.load(), 2); - monitor.Stop(); -} - -TEST_F(StoreMonitorTest, ARequestWakesTheSweepBeforeItsPeriod) { - RaidenId id = {"sweeping", "r0", "dataS", 0}; - std::atomic sweep_calls{0}; - StoreMonitor monitor( - StoreMonitor::Options{.heartbeat_period = absl::Hours(1), - .sweep_period = absl::Hours(1)}, - client_, id, ReportFreeBlocks(1), - /*reregister_fn=*/[] { return absl::OkStatus(); }, - /*sweep_fn=*/[&sweep_calls] { - ++sweep_calls; - return false; - }); - monitor.Start(); - absl::SleepFor(absl::Milliseconds(200)); - ASSERT_EQ(sweep_calls.load(), 0); // A period away from the first sweep. - - monitor.RequestSweep(); - for (int i = 0; i < 100 && sweep_calls.load() == 0; ++i) { - absl::SleepFor(absl::Milliseconds(10)); - } - EXPECT_EQ(sweep_calls.load(), 1); - monitor.Stop(); -} - -TEST_F(StoreMonitorTest, SweepKeepsSteppingWhileItReportsMoreWork) { - RaidenId id = {"sweeping", "r0", "dataS", 0}; - std::atomic sweep_calls{0}; - StoreMonitor monitor( - StoreMonitor::Options{.heartbeat_period = absl::Hours(1), - .sweep_period = absl::Hours(1)}, - client_, id, ReportFreeBlocks(1), - /*reregister_fn=*/[] { return absl::OkStatus(); }, - // Three steps of pending work, then done: one wake-up must run all - // four calls back to back. - /*sweep_fn=*/[&sweep_calls] { return ++sweep_calls < 4; }); - monitor.Start(); - monitor.RequestSweep(); - for (int i = 0; i < 100 && sweep_calls.load() < 4; ++i) { - absl::SleepFor(absl::Milliseconds(10)); - } - EXPECT_EQ(sweep_calls.load(), 4); - monitor.Stop(); -} - -TEST_F(StoreMonitorTest, ARequestWithoutASweepIsANoOp) { - RaidenId id = {"monitored", "r0", "dataS", 0}; - StoreMonitor monitor( - StoreMonitor::Options{.heartbeat_period = absl::Hours(1)}, client_, id, - ReportFreeBlocks(1), - /*reregister_fn=*/[] { return absl::OkStatus(); }); - monitor.Start(); - monitor.RequestSweep(); - monitor.Stop(); -} - TEST_F(StoreMonitorTest, DestructorStopsAnUnstoppedMonitor) { RaidenId id = {"monitored", "r0", "dataS", 0}; {