[model-gateway] cache_aware eliminate String allocations in hot path (#16209)

This commit is contained in:
Simo Lin
2025-12-31 00:26:10 -08:00
committed by GitHub
parent 25b48564c3
commit b4ce7a6d71
+29 -27
View File
@@ -202,22 +202,18 @@ impl CacheAwarePolicy {
request_text: &Option<&str>, request_text: &Option<&str>,
healthy_indices: &[usize], healthy_indices: &[usize],
model_id: &str, model_id: &str,
// TODO may skip passing this arg (and compute inside function) if this is not bottleneck
max_load: usize, max_load: usize,
min_load: usize, min_load: usize,
) -> Option<usize> { ) -> Option<usize> {
// Log load balancing trigger // Log load balancing trigger (only compute worker loads if debug enabled)
// TODO may use `&str` if tracing::enabled!(tracing::Level::DEBUG) {
let worker_loads: Vec<(String, usize)> = workers let worker_loads: Vec<(&str, usize)> =
.iter() workers.iter().map(|w| (w.url(), w.load())).collect();
.map(|w| (w.url().to_string(), w.load()))
.collect();
// TODO may change text
debug!( debug!(
"Load balancing triggered | max: {} | min: {} | workers: {:?}", "Load balancing triggered | max: {} | min: {} | workers: {:?}",
max_load, min_load, worker_loads max_load, min_load, worker_loads
); );
}
// Use shortest queue when imbalanced // Use shortest queue when imbalanced
let min_load_idx = healthy_indices let min_load_idx = healthy_indices
@@ -301,31 +297,37 @@ impl LoadBalancingPolicy for CacheAwarePolicy {
result.matched_char_count as f32 / result.input_char_count as f32 result.matched_char_count as f32 / result.input_char_count as f32
}; };
let selected_url = if match_rate > self.config.cache_threshold { // Select worker without String allocation
result.tenant.to_string() let selected_idx = if match_rate > self.config.cache_threshold {
} else { // Cache hit path: find worker by URL (compare &str directly, no allocation)
let min_load_idx = *healthy_indices let tenant_url: &str = &result.tenant;
workers
.iter() .iter()
.min_by_key(|&&idx| workers[idx].load())?; .position(|w| w.url() == tenant_url)
workers[min_load_idx].url().to_string() .filter(|&idx| workers[idx].is_healthy())
} else {
// Low cache match: use worker with minimum load
healthy_indices
.iter()
.min_by_key(|&&idx| workers[idx].load())
.copied()
}; };
// Find the index of the selected worker if let Some(idx) = selected_idx {
if let Some(selected_idx) = workers.iter().position(|w| w.url() == selected_url) { // Update the tree with this request (use worker URL directly, no allocation)
// Only proceed if the worker is healthy tree.insert(text, workers[idx].url());
if workers[selected_idx].is_healthy() {
// Update the tree with this request
tree.insert(text, &selected_url);
// Increment processed counter // Increment processed counter
workers[selected_idx].increment_processed(); workers[idx].increment_processed();
return Some(selected_idx); return Some(idx);
} }
} else {
// Selected worker no longer exists, remove it from tree // Selected worker no longer exists or unhealthy, remove stale tenant from tree
tree.remove_tenant(&selected_url); if match_rate > self.config.cache_threshold {
debug!("Removed stale worker {} from cache tree", selected_url); let tenant_url: &str = &result.tenant;
tree.remove_tenant(tenant_url);
debug!("Removed stale worker {} from cache tree", tenant_url);
} }
// Fallback to first healthy worker // Fallback to first healthy worker