Skip to content

Commit ab83747

Browse files
committed
[protocol/manager/service] support configurable host-state P2P host count
1 parent f905d96 commit ab83747

7 files changed

Lines changed: 218 additions & 115 deletions

File tree

kv_cache_manager/manager/cache_manager.cc

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4385,7 +4385,8 @@ CacheManager::GetHostCacheState(RequestContext *request_context,
43854385
const std::string &instance_id,
43864386
QueryType query_type,
43874387
const KeyVector &block_cache_keys,
4388-
const std::vector<std::string> &medium_filter) {
4388+
const std::vector<std::string> &medium_filter,
4389+
size_t p2p_host_count) {
43894390
SPAN_TRACER(request_context);
43904391
const std::string &trace_id = request_context->trace_id();
43914392
auto *service_metrics_collector = dynamic_cast<ServiceMetricsCollector *>(request_context->metrics_collector());
@@ -4434,7 +4435,8 @@ CacheManager::GetHostCacheState(RequestContext *request_context,
44344435
use_eagle_pop,
44354436
medium_filter,
44364437
host_matches,
4437-
&request_check_loc_data_exist);
4438+
&request_check_loc_data_exist,
4439+
p2p_host_count);
44384440
break;
44394441
}
44404442
case QueryType::QT_PREFIX_MATCH_WITH_MAMBA: {
@@ -4444,7 +4446,8 @@ CacheManager::GetHostCacheState(RequestContext *request_context,
44444446
medium_filter,
44454447
instance_info->location_spec_groups(),
44464448
host_matches,
4447-
&request_check_loc_data_exist);
4449+
&request_check_loc_data_exist,
4450+
p2p_host_count);
44484451
break;
44494452
}
44504453
default:

kv_cache_manager/manager/cache_manager.h

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -214,7 +214,8 @@ class CacheManager {
214214
const std::string &instance_id,
215215
QueryType query_type,
216216
const KeyVector &block_cache_keys,
217-
const std::vector<std::string> &medium_filter = {});
217+
const std::vector<std::string> &medium_filter = {},
218+
size_t p2p_host_count = 0);
218219
ErrorCode TrimCache(RequestContext *request_context,
219220
const std::string &instance_id,
220221
const proto::meta::TrimStrategy &trim_strategy,

kv_cache_manager/manager/meta_searcher.cc

Lines changed: 105 additions & 99 deletions
Original file line numberDiff line numberDiff line change
@@ -357,17 +357,19 @@ bool IsMediumMatched(const StandardUri &uri, const std::unordered_set<std::strin
357357
using HostToSpecNames = std::map<std::string, std::set<std::string>>;
358358
using KeyToHostSpecNames = std::vector<HostToSpecNames>; // key -> host -> spec names
359359

360-
constexpr size_t kMaxP2PHostCount = 5;
361-
362-
std::vector<size_t> SelectTopHostIndicesByLocal(const std::vector<MetaSearcher::HostCacheMatch> &host_matches) {
360+
std::vector<size_t> SelectTopHostIndicesByLocal(const std::vector<MetaSearcher::HostCacheMatch> &host_matches,
361+
size_t p2p_host_count) {
362+
if (p2p_host_count == 0) {
363+
return {};
364+
}
363365
std::vector<size_t> host_indices;
364366
host_indices.reserve(host_matches.size());
365367
for (size_t i = 0; i < host_matches.size(); ++i) {
366368
if (host_matches[i].local > 0) {
367369
host_indices.push_back(i);
368370
}
369371
}
370-
const size_t selected_count = std::min(kMaxP2PHostCount, host_indices.size());
372+
const size_t selected_count = std::min(p2p_host_count, host_indices.size());
371373
std::partial_sort(host_indices.begin(),
372374
host_indices.begin() + selected_count,
373375
host_indices.end(),
@@ -932,8 +934,7 @@ ErrorCode MetaSearcher::BatchGetBestLocationByBackend(RequestContext *request_co
932934

933935
const auto &vmap = valid_maps[i];
934936
const std::string_view requested_spec_name =
935-
requested_spec_names.empty() ? std::string_view{}
936-
: requested_spec_names[query_to_output_index[i]];
937+
requested_spec_names.empty() ? std::string_view{} : requested_spec_names[query_to_output_index[i]];
937938
std::vector<std::string> vineyard_addrs;
938939

939940
for (const auto &[id, loc] : vmap) {
@@ -987,8 +988,7 @@ ErrorCode MetaSearcher::BatchGetBestLocationByBackend(RequestContext *request_co
987988
// --- Per-key independent selection (WEIGHTED_RANDOM or other non-event-report) ---
988989
for (size_t i = 0; i < query_keys.size(); ++i) {
989990
const std::string_view requested_spec_name =
990-
requested_spec_names.empty() ? std::string_view{}
991-
: requested_spec_names[query_to_output_index[i]];
991+
requested_spec_names.empty() ? std::string_view{} : requested_spec_names[query_to_output_index[i]];
992992
const auto &vmap = valid_maps[i];
993993
CacheLocationMap filtered;
994994
for (const auto &[id, loc] : vmap) {
@@ -1118,7 +1118,8 @@ ErrorCode MetaSearcher::PrefixMatchByHost(RequestContext *request_context,
11181118
bool use_eagle_pop,
11191119
const std::vector<std::string> &medium_filter,
11201120
std::vector<HostCacheMatch> &out_matches,
1121-
const CheckLocDataExistFunc *request_check_loc_data_exist) const {
1121+
const CheckLocDataExistFunc *request_check_loc_data_exist,
1122+
size_t p2p_host_count) const {
11221123
SPAN_TRACER(request_context);
11231124
out_matches.clear();
11241125
if (keys.empty()) {
@@ -1196,34 +1197,36 @@ ErrorCode MetaSearcher::PrefixMatchByHost(RequestContext *request_context,
11961197
}
11971198
});
11981199
if (reduce_ok) {
1199-
const auto top_host_indices = SelectTopHostIndicesByLocal(host_matches);
1200-
reduce_ok = top_host_indices.empty() || meta_indexer_->ParallelForQuery(
1201-
top_host_indices.size(),
1202-
[&top_host_indices,
1203-
&candidate_hosts,
1204-
&keys,
1205-
&key_to_host_spec_names,
1206-
&key_to_vineyard_spec_names,
1207-
&host_matches,
1208-
use_eagle_pop](std::size_t begin, std::size_t end) {
1209-
for (std::size_t selected_index = begin; selected_index < end; ++selected_index) {
1210-
const size_t host_index = top_host_indices[selected_index];
1211-
const auto &host = candidate_hosts[host_index];
1212-
auto p2p_selection = SelectP2PByPrefix(host, key_to_host_spec_names, key_to_vineyard_spec_names);
1213-
std::unordered_set<int64_t> fetched_block_keys;
1214-
for (size_t index : p2p_selection.covered_indices) {
1215-
fetched_block_keys.insert(keys[index]);
1216-
}
1217-
int64_t total_match = host_matches[host_index].local;
1218-
if (!p2p_selection.covered_indices.empty()) {
1219-
auto p2p_specs = MergeHostAndP2PSpecs(
1220-
host, key_to_host_spec_names, key_to_vineyard_spec_names, p2p_selection);
1221-
total_match = ComputePrefixMatchBlocks(p2p_specs, use_eagle_pop);
1222-
}
1223-
host_matches[host_index].p2p_1_fetch = static_cast<int64_t>(fetched_block_keys.size());
1224-
host_matches[host_index].p2p_1_total_match = total_match;
1225-
}
1226-
});
1200+
const auto top_host_indices = SelectTopHostIndicesByLocal(host_matches, p2p_host_count);
1201+
reduce_ok = top_host_indices.empty() ||
1202+
meta_indexer_->ParallelForQuery(
1203+
top_host_indices.size(),
1204+
[&top_host_indices,
1205+
&candidate_hosts,
1206+
&keys,
1207+
&key_to_host_spec_names,
1208+
&key_to_vineyard_spec_names,
1209+
&host_matches,
1210+
use_eagle_pop](std::size_t begin, std::size_t end) {
1211+
for (std::size_t selected_index = begin; selected_index < end; ++selected_index) {
1212+
const size_t host_index = top_host_indices[selected_index];
1213+
const auto &host = candidate_hosts[host_index];
1214+
auto p2p_selection =
1215+
SelectP2PByPrefix(host, key_to_host_spec_names, key_to_vineyard_spec_names);
1216+
std::unordered_set<int64_t> fetched_block_keys;
1217+
for (size_t index : p2p_selection.covered_indices) {
1218+
fetched_block_keys.insert(keys[index]);
1219+
}
1220+
int64_t total_match = host_matches[host_index].local;
1221+
if (!p2p_selection.covered_indices.empty()) {
1222+
auto p2p_specs = MergeHostAndP2PSpecs(
1223+
host, key_to_host_spec_names, key_to_vineyard_spec_names, p2p_selection);
1224+
total_match = ComputePrefixMatchBlocks(p2p_specs, use_eagle_pop);
1225+
}
1226+
host_matches[host_index].p2p_1_fetch = static_cast<int64_t>(fetched_block_keys.size());
1227+
host_matches[host_index].p2p_1_total_match = total_match;
1228+
}
1229+
});
12271230
}
12281231
KVCM_METRICS_COLLECTOR_CHRONO_MARK_END(service_metrics_collector, MetaSearcherHostPrefixReduce);
12291232
if (!reduce_ok) {
@@ -1245,7 +1248,8 @@ ErrorCode MetaSearcher::PrefixMatchWithMambaByHost(RequestContext *request_conte
12451248
const std::vector<std::string> &medium_filter,
12461249
const std::vector<LocationSpecGroup> &location_spec_groups,
12471250
std::vector<HostCacheMatch> &out_matches,
1248-
const CheckLocDataExistFunc *request_check_loc_data_exist) const {
1251+
const CheckLocDataExistFunc *request_check_loc_data_exist,
1252+
size_t p2p_host_count) const {
12491253
SPAN_TRACER(request_context);
12501254
out_matches.clear();
12511255
if (keys.empty()) {
@@ -1343,74 +1347,76 @@ ErrorCode MetaSearcher::PrefixMatchWithMambaByHost(RequestContext *request_conte
13431347
}
13441348
});
13451349
if (reduce_ok) {
1346-
const auto top_host_indices = SelectTopHostIndicesByLocal(host_matches);
1347-
reduce_ok = top_host_indices.empty() || meta_indexer_->ParallelForQuery(
1348-
top_host_indices.size(),
1349-
[&top_host_indices,
1350-
&candidate_hosts,
1351-
&keys,
1352-
&key_to_host_spec_names,
1353-
&key_to_vineyard_spec_names,
1354-
&full_groups,
1355-
&mamba_state_groups,
1356-
&host_matches,
1357-
use_eagle_pop](std::size_t begin, std::size_t end) {
1358-
for (std::size_t selected_index = begin; selected_index < end; ++selected_index) {
1359-
const size_t host_index = top_host_indices[selected_index];
1360-
const auto &host = candidate_hosts[host_index];
1361-
const V6DPeerSelection no_p2p;
1362-
auto combined_specs =
1363-
MergeHostAndP2PSpecs(host, key_to_host_spec_names, key_to_vineyard_spec_names, no_p2p);
1364-
std::unordered_set<int64_t> fetched_block_keys;
1365-
1366-
for (const auto *full_group : full_groups) {
1367-
auto full_selection =
1368-
SelectP2PGroupByPrefix(host, key_to_host_spec_names, key_to_vineyard_spec_names, *full_group);
1369-
for (size_t block_index : full_selection.covered_indices) {
1370-
const auto peer_it =
1371-
key_to_vineyard_spec_names[block_index].find(full_selection.peer_addr);
1372-
if (peer_it == key_to_vineyard_spec_names[block_index].end()) {
1373-
continue;
1350+
const auto top_host_indices = SelectTopHostIndicesByLocal(host_matches, p2p_host_count);
1351+
reduce_ok =
1352+
top_host_indices.empty() ||
1353+
meta_indexer_->ParallelForQuery(
1354+
top_host_indices.size(),
1355+
[&top_host_indices,
1356+
&candidate_hosts,
1357+
&keys,
1358+
&key_to_host_spec_names,
1359+
&key_to_vineyard_spec_names,
1360+
&full_groups,
1361+
&mamba_state_groups,
1362+
&host_matches,
1363+
use_eagle_pop](std::size_t begin, std::size_t end) {
1364+
for (std::size_t selected_index = begin; selected_index < end; ++selected_index) {
1365+
const size_t host_index = top_host_indices[selected_index];
1366+
const auto &host = candidate_hosts[host_index];
1367+
const V6DPeerSelection no_p2p;
1368+
auto combined_specs =
1369+
MergeHostAndP2PSpecs(host, key_to_host_spec_names, key_to_vineyard_spec_names, no_p2p);
1370+
std::unordered_set<int64_t> fetched_block_keys;
1371+
1372+
for (const auto *full_group : full_groups) {
1373+
auto full_selection = SelectP2PGroupByPrefix(
1374+
host, key_to_host_spec_names, key_to_vineyard_spec_names, *full_group);
1375+
for (size_t block_index : full_selection.covered_indices) {
1376+
const auto peer_it =
1377+
key_to_vineyard_spec_names[block_index].find(full_selection.peer_addr);
1378+
if (peer_it == key_to_vineyard_spec_names[block_index].end()) {
1379+
continue;
1380+
}
1381+
MergeLocationSpecGroup(combined_specs[block_index], peer_it->second, *full_group);
1382+
fetched_block_keys.insert(keys[block_index]);
13741383
}
1375-
MergeLocationSpecGroup(combined_specs[block_index], peer_it->second, *full_group);
1376-
fetched_block_keys.insert(keys[block_index]);
13771384
}
1378-
}
13791385

1380-
size_t full_prefix_len = 0;
1381-
while (full_prefix_len < combined_specs.size() &&
1382-
HasAllLocationSpecGroups(combined_specs[full_prefix_len], full_groups)) {
1383-
++full_prefix_len;
1384-
}
1385-
if (use_eagle_pop && full_prefix_len > 0) {
1386-
--full_prefix_len;
1387-
}
1386+
size_t full_prefix_len = 0;
1387+
while (full_prefix_len < combined_specs.size() &&
1388+
HasAllLocationSpecGroups(combined_specs[full_prefix_len], full_groups)) {
1389+
++full_prefix_len;
1390+
}
1391+
if (use_eagle_pop && full_prefix_len > 0) {
1392+
--full_prefix_len;
1393+
}
13881394

1389-
if (full_prefix_len > 0) {
1390-
auto mamba_selection = SelectP2PGroupsByCoverage(host,
1391-
key_to_host_spec_names,
1392-
key_to_vineyard_spec_names,
1393-
full_prefix_len,
1394-
mamba_state_groups);
1395-
for (size_t query_index : mamba_selection.selection.covered_indices) {
1396-
const size_t block_index = mamba_selection.query_block_indices[query_index];
1397-
const auto *group = mamba_selection.query_groups[query_index];
1398-
const auto peer_it =
1399-
key_to_vineyard_spec_names[block_index].find(mamba_selection.selection.peer_addr);
1400-
if (peer_it == key_to_vineyard_spec_names[block_index].end()) {
1401-
continue;
1395+
if (full_prefix_len > 0) {
1396+
auto mamba_selection = SelectP2PGroupsByCoverage(host,
1397+
key_to_host_spec_names,
1398+
key_to_vineyard_spec_names,
1399+
full_prefix_len,
1400+
mamba_state_groups);
1401+
for (size_t query_index : mamba_selection.selection.covered_indices) {
1402+
const size_t block_index = mamba_selection.query_block_indices[query_index];
1403+
const auto *group = mamba_selection.query_groups[query_index];
1404+
const auto peer_it =
1405+
key_to_vineyard_spec_names[block_index].find(mamba_selection.selection.peer_addr);
1406+
if (peer_it == key_to_vineyard_spec_names[block_index].end()) {
1407+
continue;
1408+
}
1409+
MergeLocationSpecGroup(combined_specs[block_index], peer_it->second, *group);
1410+
fetched_block_keys.insert(keys[block_index]);
14021411
}
1403-
MergeLocationSpecGroup(combined_specs[block_index], peer_it->second, *group);
1404-
fetched_block_keys.insert(keys[block_index]);
14051412
}
1406-
}
14071413

1408-
const int64_t total_match =
1409-
ComputeMambaPrefixMatchBlocks(combined_specs, use_eagle_pop, full_groups, mamba_state_groups);
1410-
host_matches[host_index].p2p_1_fetch = static_cast<int64_t>(fetched_block_keys.size());
1411-
host_matches[host_index].p2p_1_total_match = total_match;
1412-
}
1413-
});
1414+
const int64_t total_match = ComputeMambaPrefixMatchBlocks(
1415+
combined_specs, use_eagle_pop, full_groups, mamba_state_groups);
1416+
host_matches[host_index].p2p_1_fetch = static_cast<int64_t>(fetched_block_keys.size());
1417+
host_matches[host_index].p2p_1_total_match = total_match;
1418+
}
1419+
});
14141420
}
14151421
KVCM_METRICS_COLLECTOR_CHRONO_MARK_END(service_metrics_collector, MetaSearcherHostPrefixReduce);
14161422
if (!reduce_ok) {

0 commit comments

Comments
 (0)