[model-gateway] cache_aware eliminate String allocations in hot path (#16209)
This commit is contained in:
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user