Skip to content

Demote cold KV blocks to placement targets when free blocks run low. #712

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

Closed
copybara-service wants to merge 1 commit into from
Closed

Demote cold KV blocks to placement targets when free blocks run low. #712

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
1 change: 0 additions & 1 deletion tpu_sync/kv_cache/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
293 changes: 49 additions & 244 deletions tpu_sync/kv_cache/kv_cache_store.cc
View file
Open in desktop

Large diffs are not rendered by default.

31 changes: 3 additions & 28 deletions tpu_sync/kv_cache/kv_cache_store.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<std::string> freeable;
// At least one block's transfer failed outright; its local copy stays.
bool transfer_failed = false;
};
BatchWriteResult WaitForBatchWriteResult(
const std::vector<std::string>& batch);

mutable absl::Mutex mutex_;
std::vector<std::shared_ptr<KVCacheStoreBackend>> backends_;
std::shared_ptr<global_registry::GlobalRegistryClient> registry_client_;
Expand All @@ -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<RaidenId> 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<StoreMonitor> store_monitor_;
Expand Down
35 changes: 8 additions & 27 deletions tpu_sync/kv_cache/kv_cache_store_backend_factory.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down
220 changes: 5 additions & 215 deletions tpu_sync/kv_cache/kv_cache_store_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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(),
Expand All @@ -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());
}

Expand Down Expand Up @@ -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<std::unique_ptr<KVCacheStore>> 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<std::string>& hashes) {
auto ids_or = store.raiden_controller()->AllocateBlockIds(hashes.size());
ASSERT_TRUE(ids_or.ok()) << ids_or.status().ToString();
std::vector<RaidenBlockID> 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<std::string>& 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<std::string> 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
Expand Down
Loading
Loading